EquiformerV2: Trasformatore Equivariante Migliorato per il Ridimensionamento a Rappresentazioni di Grado Superiore

EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations

June 21, 2023

Autori: Yi-Lun Liao, Brandon Wood, Abhishek Das, Tess Smidt

cs.AI

Abstract

I Transformer equivarianti come Equiformer hanno dimostrato l'efficacia dell'applicazione dei Transformer al dominio dei sistemi atomistici 3D. Tuttavia, sono ancora limitati a piccoli gradi di rappresentazioni equivarianti a causa della loro complessità computazionale. In questo articolo, indaghiamo se queste architetture possano scalare bene a gradi più elevati. Partendo da Equiformer, sostituiamo prima le convoluzioni SO(3) con convoluzioni eSCN per incorporare in modo efficiente tensori di grado superiore. Poi, per sfruttare meglio la potenza dei gradi più elevati, proponiamo tre miglioramenti architetturali: la rinormalizzazione dell'attenzione, l'attivazione separabile S^2 e la normalizzazione separabile a strati. Mettendo tutto insieme, proponiamo EquiformerV2, che supera i precedenti metodi all'avanguardia sul dataset su larga scala OC20 fino al 12% sulle forze, al 4% sulle energie, offre migliori compromessi velocità-precisione e una riduzione di 2 volte nei calcoli DFT necessari per calcolare le energie di adsorbimento.

English

Equivariant Transformers such as Equiformer have demonstrated the efficacy of applying Transformers to the domain of 3D atomistic systems. However, they are still limited to small degrees of equivariant representations due to their computational complexity. In this paper, we investigate whether these architectures can scale well to higher degrees. Starting from Equiformer, we first replace SO(3) convolutions with eSCN convolutions to efficiently incorporate higher-degree tensors. Then, to better leverage the power of higher degrees, we propose three architectural improvements -- attention re-normalization, separable S^2 activation and separable layer normalization. Putting this all together, we propose EquiformerV2, which outperforms previous state-of-the-art methods on the large-scale OC20 dataset by up to 12% on forces, 4% on energies, offers better speed-accuracy trade-offs, and 2times reduction in DFT calculations needed for computing adsorption energies.