EquiformerV2 : Transformeur équivariant amélioré pour la mise à l'échelle vers des représentations de degré supérieur

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

Résumé

Les Transformers équivariants tels que l'Equiformer ont démontré l'efficacité de l'application des Transformers au domaine des systèmes atomistiques 3D. Cependant, ils restent limités à de faibles degrés de représentations équivariantes en raison de leur complexité computationnelle. Dans cet article, nous étudions si ces architectures peuvent bien s'adapter à des degrés plus élevés. En partant de l'Equiformer, nous remplaçons d'abord les convolutions SO(3) par des convolutions eSCN pour incorporer efficacement des tenseurs de degré supérieur. Ensuite, pour mieux exploiter la puissance des degrés plus élevés, nous proposons trois améliorations architecturales : la re-normalisation de l'attention, l'activation séparable S^2 et la normalisation de couche séparable. En combinant tout cela, nous proposons l'EquiformerV2, qui surpasse les méthodes de pointe précédentes sur le jeu de données OC20 à grande échelle avec une amélioration allant jusqu'à 12 % sur les forces, 4 % sur les énergies, offre de meilleurs compromis vitesse-précision, et une réduction par deux des calculs DFT nécessaires pour calculer les énergies d'adsorption.

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.