EquiformerV2: Verbesserter äquivarianter Transformer für die Skalierung auf höhergradige Darstellungen

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

June 21, 2023

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

cs.AI

Zusammenfassung

Äquivariante Transformer wie Equiformer haben die Wirksamkeit der Anwendung von Transformern auf den Bereich von 3D-atomaren Systemen demonstriert. Sie sind jedoch aufgrund ihrer rechnerischen Komplexität immer noch auf geringe Grade von äquivarianten Darstellungen beschränkt. In dieser Arbeit untersuchen wir, ob diese Architekturen gut auf höhere Grade skaliert werden können. Ausgehend von Equiformer ersetzen wir zunächst SO(3)-Faltungen durch eSCN-Faltungen, um effizient höhergradige Tensoren einzubeziehen. Um die Leistung höherer Grade besser zu nutzen, schlagen wir dann drei architektonische Verbesserungen vor – Attention-Renormalisierung, separable S^2-Aktivierung und separable Layer-Normalisierung. All dies zusammenfassend, schlagen wir EquiformerV2 vor, das bisherige state-of-the-art Methoden auf dem groß angelegten OC20-Datensatz um bis zu 12 % bei Kräften, 4 % bei Energien übertrifft, bessere Geschwindigkeits-Genauigkeits-Kompromisse bietet und eine 2-fache Reduzierung der für die Berechnung von Adsorptionsenergien benötigten DFT-Berechnungen ermöglicht.

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.