EquiformerV2: Transformador Equivariante Mejorado para Escalar a Representaciones de Mayor Grado

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

June 21, 2023

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

cs.AI

Resumen

Los transformadores equivariantes, como Equiformer, han demostrado la eficacia de aplicar transformadores al dominio de los sistemas atomísticos 3D. Sin embargo, todavía están limitados a grados pequeños de representaciones equivariantes debido a su complejidad computacional. En este artículo, investigamos si estas arquitecturas pueden escalar adecuadamente a grados más altos. Partiendo de Equiformer, primero reemplazamos las convoluciones SO(3) con convoluciones eSCN para incorporar eficientemente tensores de mayor grado. Luego, para aprovechar mejor el poder de los grados más altos, proponemos tres mejoras arquitectónicas: re-normalización de atención, activación separable S^2 y normalización de capa separable. Reuniendo todo esto, proponemos EquiformerV2, que supera a los métodos anteriores de última generación en el conjunto de datos a gran escala OC20 en un 12 % en fuerzas, un 4 % en energías, ofrece mejores compensaciones entre velocidad y precisión, y reduce en 2 veces los cálculos de DFT necesarios para calcular energías de adsorció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.