[Paper Review] 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 Science81 citations

TL;DR

EquiformerV2 scales equivariant Transformers to higher-degree representations using eSCN convolutions and architectural tweaks, achieving state-of-the-art results on OC20 with better speed-accuracy and data efficiency. It also reduces DFT relaxations in AdsorbML and improves OC22 performance over 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.

Motivation & Objective

• Advance 3D atomistic modeling by enabling higher-degree equivariant representations (Lmax) in Transformer-based architectures.
• Improve predictive accuracy for forces and energies on large-scale OC20/OC22 datasets.
• Increase training and inference efficiency to enable practical applications like AdsorbML and DFT-augmented workflows.
• Demonstrate data efficiency and better out-of-distribution performance compared to prior methods.

Proposed method

• Replace SO(3) convolutions with eSCN convolutions to enable higher Lmax (up to 6–8) in EquiformerV2.
• Introduce architectural improvements: attention re-normalization, separable S2 activation, and separable layer normalization.
• Maintain equivariant graph attention and feed-forward structure while using eSCN to mix degree and channel information.
• Employ radial distance embeddings and SO(2) linear layers to realize efficient, equivariant message passing.
• Evaluate on OC20 S2EF (All and All+MD) and OC22 datasets, including AdsorbML scenarios, and compare with prior state-of-the-art.

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

Experimental results

Research questions

• RQ1 How well can higher-degree representations (larger Lmax) be scaled in an equivariant Transformer for 3D atomistic systems using efficient convolutions?
• RQ2 Do architectural improvements (attention normalization, separable S2 activation, separable LN) yield measurable gains over baseline Equiformer when using higher degrees?
• RQ3 Can EquiformerV2 achieve data-efficient and faster learning while surpassing existing methods on OC20/OC22 benchmarks?
• RQ4 What is the impact of EquiformerV2 on adsorption-energy workflows (AdsorbML) and DFT-computation reductions?
• RQ5 How does EquiformerV2 compare to Equiformer and GemNet-OC on QM9 and OC20 S2EF-2M in terms of performance gains from higher degrees?

Key findings

• EquiformerV2 with eSCN convolutions and higher degrees outperforms previous state-of-the-art on OC20 by up to 9% on forces and 4% on energies.
• Attention re-normalization and separable S2 activation improve training stability and force/energy accuracy, with separable LN further boosting force MAE.
• EquiformerV2 trained on OC22 alone outperforms GemNet-OC trained on OC20 and OC22, indicating strong data efficiency.
• In AdsorbML, EquiformerV2 achieves higher success rates and up to 2× reduction in DFT calculations needed for comparable adsorption-energy accuracy.
• EquiformerV2 (λE=4, λF=100) with 153M parameters achieves state-of-the-art results on OC20 S2EF-All+MD, including force MAE improvements and favorable speed-accuracy trade-offs.
• Smaller variant (λE=4, λF=100, 31M) also shows strong performance with favorable training/inference efficiency, indicating scalable deployment potential.

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$ .