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

Yi-Lun Liao,Brandon Wood,Abhishek Das,Tess Smidt

2024-03-07

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

Machine Learning,Artificial Intelligence,Computational Physics

What problem does this paper attempt to address?

### What problem does this paper attempt to solve?

This paper primarily addresses the following issues:

1. **Enhancing the representation capability of Equivariant Transformer**:
- Current Equivariant Transformers (such as Equiformer) are limited by computational complexity when handling 3D atomic systems, and can only process low-order equivariant representations (e.g., maximum order \( L_{\text{max}} \leq 3 \)). This paper introduces eSCN convolution to enable the model to efficiently handle higher-order equivariant representations.

2. **Proposing new architectural improvements**:
- Based on the introduction of eSCN convolution, three new architectural improvements are proposed: attention re-normalization, separable \( S^2 \) activation, and separable layer normalization, to better utilize higher-order information.

3. **Improving model performance**:
- With the aforementioned improvements, the proposed EquiformerV2 achieves significant performance enhancements on the large-scale OC20 dataset, particularly in force and energy prediction, with improvements of 9% and 4% respectively compared to previous methods, and offers a better speed-accuracy trade-off.

4. **Increasing data efficiency**:
- EquiformerV2, when trained only on the OC22 dataset, outperforms GemNet-OC trained on both OC20 and OC22 datasets, demonstrating higher data efficiency.

Through these improvements, the paper aims to enhance the representation capability and prediction accuracy of equivariant Transformers in 3D atomic systems, thereby accelerating quantum mechanical computations and achieving breakthroughs in fields such as molecular simulation, materials design, and drug discovery.