EquiformerV2: Transformador Equivariante Aprimorado para Escalonamento em Representações de Grau Superior

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

Resumo

Transformers Equivariantes, como o Equiformer, demonstraram a eficácia da aplicação de Transformers ao domínio de sistemas atomísticos 3D. No entanto, eles ainda estão limitados a pequenos graus de representações equivariantes devido à sua complexidade computacional. Neste artigo, investigamos se essas arquiteturas podem escalar bem para graus mais elevados. Partindo do Equiformer, primeiro substituímos as convoluções SO(3) por convoluções eSCN para incorporar eficientemente tensores de grau mais alto. Em seguida, para melhor aproveitar o poder de graus mais elevados, propomos três melhorias arquitetônicas: renormalização de atenção, ativação S^2 separável e normalização de camada separável. Reunindo tudo isso, propomos o EquiformerV2, que supera os métodos state-of-the-art anteriores no conjunto de dados em larga escala OC20 em até 12% nas forças, 4% nas energias, oferece melhores compensações entre velocidade e precisão, e uma redução de 2 vezes nos cálculos de DFT necessários para computar energias de adsorção.

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.