EquiformerV2：用于扩展到更高阶表示的改进等变Transformer

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

June 21, 2023

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

cs.AI

摘要

等变换器（Equivariant Transformers）如Equiformer已经证明了将Transformer应用于三维原子系统领域的有效性。然而，由于计算复杂性，它们仍然局限于小范围的等变表示。在本文中，我们调查了这些架构是否能够很好地扩展到更高的程度。从Equiformer开始，我们首先用eSCN卷积替换SO(3)卷积，以有效地整合更高阶的张量。然后，为了更好地利用更高阶的能力，我们提出了三种架构改进--注意力重归一化、可分离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.