EquiformerV2: Усовершенствованный эквивариантный трансформер для масштабирования до представлений более высокой степени

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

June 21, 2023

Авторы: Yi-Lun Liao, Brandon Wood, Abhishek Das, Tess Smidt

cs.AI

Аннотация

Эквивариантные трансформаторы, такие как Equiformer, продемонстрировали эффективность применения архитектур Transformer в области 3D атомных систем. Однако они по-прежнему ограничены использованием представлений низкой степени эквивариантности из-за их вычислительной сложности. В данной работе мы исследуем, могут ли такие архитектуры масштабироваться до более высоких степеней. Начиная с Equiformer, мы сначала заменяем свертки SO(3) на свертки eSCN, чтобы эффективно включать тензоры более высокой степени. Затем, чтобы лучше использовать потенциал более высоких степеней, мы предлагаем три архитектурных улучшения — повторную нормализацию внимания, разделяемую активацию S^2 и разделяемую послойную нормализацию. Объединив все это, мы предлагаем EquiformerV2, который превосходит предыдущие передовые методы на крупном наборе данных OC20, улучшая точность предсказания сил на 12%, энергий на 4%, предлагая лучший компромисс между скоростью и точностью, а также сокращая количество необходимых расчетов DFT для вычисления энергий адсорбции в 2 раза.

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.