[论文解读] EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations

Yi-Lun Liao, Brandon C. Wood+2 more|arXiv (Cornell University)|Jun 21, 2023

Machine Learning in Materials Science被引用 81

一句话总结

EquiformerV2 将等变 Transformers 扩展到更高阶表示，使用 eSCN 卷积和架构改进，在 OC20 上实现最前沿的结果，具备更快的速度-准确性和数据效率。它还减少 AdsorbML 的 DFT 松弛，并在 OC22 上的性能超越 GemNet-OC。

ABSTRACT

Equivariant Transformers such as Equiformer have demonstrated the efficacy of applying Transformers to the domain of 3D atomistic systems. However, they are 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 large-scale OC20 dataset by up to $9\%$ on forces, $4\%$ on energies, offers better speed-accuracy trade-offs, and $2 imes$ reduction in DFT calculations needed for computing adsorption energies. Additionally, EquiformerV2 trained on only OC22 dataset outperforms GemNet-OC trained on both OC20 and OC22 datasets, achieving much better data efficiency. Finally, we compare EquiformerV2 with Equiformer on QM9 and OC20 S2EF-2M datasets to better understand the performance gain brought by higher degrees.

研究动机与目标

• 通过在基于 Transformer 的架构中实现更高阶的等变表示（Lmax）来推进3D原子级建模。
• 提高对大规模 OC20/OC22 数据集上的力和能量的预测准确性。
• 提升训练和推理效率，以实现如 AdsorbML 和基于DFT的增强工作流等实际应用。
• 展示数据效率以及相对于先前方法在分布外表现上的更好。

提出的方法

• 将 SO(3) 卷积替换为 eSCN 卷积，以在 EquiformerV2 中实现更高的 Lmax（最高可达 6–8）。
• 引入架构改进：注意力再归一化、可分离的 S2 激活，以及可分离的层归一化。
• 在保持等变图注意力和前馈结构的同时，使用 eSCN 来混合阶次信息和通道信息。
• 使用径向距离嵌入和 SO(2) 线性层来实现高效、等变的消息传递。
• 在 OC20 S2EF（All 和 All+MD）和 OC22 数据集上进行评估，包括 AdsorbML 场景，并与此前的最先进方法进行比较。

Figure 1: Overview of EquiformerV2. We highlight the differences from Equiformer (Liao & Smidt, 2023 ) in red . For (b), (c), and (d), the left figure is the original module in Equiformer, and the right figure is the revised module in EquiformerV2. Input 3D graphs are embedded with atom and edge-deg

实验结果

研究问题

• RQ1 如何在使用高效卷积的情况下，将更高阶的表示（更大的 Lmax）在用于3D 原子系统的等变 Transformer 中扩展？
• RQ2 在使用更高阶时，架构改进（注意力归一化、可分离的 S2 激活、可分离的 LN）是否能带来相对于基线 Equiformer 的可衡量提升？
• RQ3 EquiformerV2 是否能够实现数据效率更高、学习更快，同时在 OC20/OC22 基准上超越现有方法？
• RQ4 EquiformerV2 对 AdsorbML 的吸附能量工作流和 DFT 计算减少有何影响？
• RQ5 与 Equiformer 和 GemNet-OC 相比，EquiformerV2 在 QM9 和 OC20 S2EF-2M 上对更高阶带来的性能提升有何差异？

主要发现

• 采用 eSCN 卷积和更高阶表示的 EquiformerV2 在 OC20 上的表现领先于此前的最先进方法，力的提升最多达 9%，能量提升达 4%。
• 注意力再归一化和可分离的 S2 激活提升了训练稳定性与力/能量精度，可分离的 LN 进一步提升力的 MAE。
• 仅在 OC22 上训练的 EquiformerV2 就超越在 OC20 和 OC22 上训练的 GemNet-OC，表明强数据效率。
• 在 AdsorbML 中，EquiformerV2 实现了更高的成功率，并且为达到可比的吸附能量精度，所需的 DFT 计算量降低最多 2 倍。
• EquiformerV2（λE=4, λF=100）/ 153M 参数，在 OC20 S2EF-All+MD 上实现了最先进的结果，包括力 MAE 的改进和有利的速度-精度权衡。
• 较小的变体（λE=4, λF=100, 31M）也显示出强劲的表现，具有有利的训练/推理效率，表明具备可扩展的部署潜力。

Figure 2: Illustration of different activation functions. $G$ denotes conversion from vectors to point samples on a sphere, $F$ can typically be a SiLU activation or MLPs, and $G^{-1}$ is the inverse of $G$ .