EquiformerV2: Verbeterde Equivariante Transformer voor Schaalbaarheid naar Representaties van Hogere Graad

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

June 21, 2023

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

cs.AI

Samenvatting

Equivariante Transformers zoals Equiformer hebben de effectiviteit aangetoond van het toepassen van Transformers op het domein van 3D-atomistische systemen. Ze zijn echter nog steeds beperkt tot kleine graden van equivariante representaties vanwege hun rekencomplexiteit. In dit artikel onderzoeken we of deze architecturen goed kunnen schalen naar hogere graden. Uitgaande van Equiformer vervangen we eerst SO(3)-convoluties door eSCN-convoluties om hogere-graadstensoren efficiënt te integreren. Vervolgens stellen we, om het potentieel van hogere graden beter te benutten, drie architectuurverbeteringen voor: aandacht-renormalisatie, separabele S^2-activering en separabele laagnormalisatie. Door dit alles samen te voegen, stellen we EquiformerV2 voor, dat de vorige state-of-the-art methoden op de grootschalige OC20-dataset overtreft met tot wel 12% op krachten, 4% op energieën, betere snelheid-nauwkeurigheid-compromissen biedt, en een 2-voudige reductie in DFT-berekeningen die nodig zijn voor het berekenen van adsorptie-energieën.

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.