PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2306.12059v3cs.LGcs.AIphysics.comp-ph

TL;DR

Equivariant Transformers such as Equiformer are limited by the cost of higher-degree representations, motivating investigation of scalable high-degree architectures. The paper combines eSCN convolutions with three architectural improvements to create EquiformerV2, which improves OC20/OC22 performance and efficiency while showing dataset- and task-dependent gains.

  • Problem

    Tensor-product complexity limits equivariant Transformers to small representation degrees, despite higher degrees capturing angular information important for atomic-energy and force prediction.

  • Method

    EquiformerV2 replaces Equiformer’s SO(3) convolutions with eSCN convolutions and adds attention re-normalization, separable S2 activation, and separable layer normalization.

  • Results

    EquiformerV2 outperforms state-of-the-art methods across OC20 and OC22 tasks, with OC20 gains of up to 9% on forces and 4% on energies.

  • Takeaways & Limitations

    The combined efficient convolutions and architectural changes provide better speed-accuracy trade-offs and the best AdsorbML success rate, including a 2× reduction in DFT calculations for comparable adsorption-energy accuracy.

  • Takeaways & Limitations

    Performance gains from higher degrees and the proposed improvements depend on datasets and tasks, with gains not significant on small datasets such as QM9.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Equivariant models use symmetry-aware representations for 3D atomistic prediction, but tensor-product costs restrict Equiformer and related models to low degrees. EquiformerV2 addresses this by combining efficient eSCN convolutions with three architectural changes for higher-degree representations.

  • Motivation: Equivariant ML models reduce quantum-mechanical computation costs while accounting for atomic discreteness and Euclidean and permutation symmetries.These methods support applications including molecular simulation, material design, and drug discovery.
  • Problem: Tensor-product complexity grows with maximum irrep degree Lmax, restricting Equiformer and related equivariant GNNs to small values such as Lmax ≤ 3.The low-degree restriction consequently limits performance.
  • Motivation: Higher degrees capture angular resolution and directional information important for accurate atomic-energy and force prediction.eSCN reduces SOp3q tensor products to SOp2q linear operations and enables scaling to values such as Lmax up to 8.
  • Approach: EquiformerV2 adapts eSCN convolutions to Equiformer and adds attention re-normalization, separable S2 activation, and separable layer normalization.The paper investigates whether equivariant Transformer designs can scale effectively to higher-degree representations.
  • Results: EquiformerV2 outperforms previous state-of-the-art methods on OC20, improving forces by up to 9% and energies by up to 4%.It also offers better speed-accuracy trade-offs and achieves a 2× reduction in DFT calculations for comparable adsorption-energy accuracy.

2 BACKGROUND

The background defines equivariant irreducible-representation features and tensor-product message passing for 3D atomistic graphs. It then situates Equiformer and eSCN as Transformer-based and computationally efficient approaches, respectively.

  • SE(3)/E(3)-equivariant neural networks: Equivariant neural networks encode 3D rotation behavior using irreducible-representation features, whose degree L corresponds intuitively to angular frequency.Higher degrees are important for tasks sensitive to angular information, including force prediction.
  • SE(3)/E(3)-equivariant neural networks: Equivariant GNNs pass transformed irreducible-representation features between graph nodes using tensor products to interact different type-L vectors.Tensor products combine representations through Clebsch-Gordan coefficients and restrict output degrees according to representation-selection rules.
  • SE(3)/E(3)-equivariant neural networks: Equivariant convolutions perform tensor products between input irreducible features and spherical harmonics of relative-position vectors.Vectors above the chosen maximum degree Lmax are typically discarded to control feature dimensionality.
  • Equiformer: Equiformer combines equivariant irreducible features and operations with equivariant graph attention, nonlinear attention and message-passing functions, and Transformer regularization.Its operations include tensor products, equivariant linear operations, layer normalization, and gate activation.
  • eSCN convolution: eSCN convolutions replace costly SO(3) tensor-product computation with SO(2) linear operations for more efficient equivariant message passing.The method rotates relative positions so spherical-harmonic projections become sparse, simplifying the remaining tensor-product paths.

3 EQUIFORMERV2

EquiformerV2 combines eSCN convolutions with three architectural changes to make higher-degree equivariant representations practical and more effective. The design replaces costly tensor-product operations while addressing attention normalization, activation stability, and normalization across degrees.

  • eSCN convolutions replace Equiformer’s costly tensor products, enabling maximum degrees of 6 or 8 on OC20.Traditional SO(3) convolution complexity scales unfavorably with Lmax, while eSCN reduces the cost enough to support larger degrees.
  • Attention re-normalization adds layer normalization before nonlinear functions and softmax attention-weight computation.The extra normalization keeps attention inputs within a comparable range when Lmax increases.
  • Separable S2 activation splits degree-0 activation from higher-degree activation to prevent large gradients and training instability.It applies SiLU to one scalar-feature partition and S2 activation to another partition together with higher-degree vectors, then concatenates the outputs.
  • Separable layer normalization uses mean and standard deviation for degree-0 vectors and RMS statistics for other degrees.Unlike independently normalizing each degree, SLN preserves relative magnitudes between nonzero degrees, which improves performance.
  • Equivariant graph attention rotates concatenated node embeddings using relative-position rotations, incorporates radial distance embeddings, and uses SO(2) linear layers.The resulting attention weights and values are combined, rotated back to the original coordinate frames, and aggregated through parallel attention functions.
  • The feed-forward network and embedding modules also replace gate activations or depth-wise tensor products with separable S2 activation and SO(2) linear operations.Embedding additionally multiplies transformed outputs by edge distance embeddings.

4 EXPERIMENTS

Experiments show that EquiformerV2’s architectural changes are necessary to exploit higher-degree representations, improving accuracy and speed-accuracy trade-offs across OC20, OC22, AdsorbML, QM9, and OC20 S2EF-2M evaluations.

  • 4.1.1 ABLATION STUDIES: Attention re-normalization improves energy MAE by 2.4%, while separable S2 activation and separable layer normalization improve force MAE.Directly adding eSCN convolutions does not outperform the original eSCN baseline; the proposed architectural changes are needed.
  • 4.1.1 ABLATION STUDIES: Higher degrees consistently help EquiformerV2, whereas increasing Lmax from 6 to 8 or Mmax from 3 to 4 degrades eSCN energy predictions.Adding more Transformer blocks and training for more epochs also improve performance.
  • 4.1.2 MAIN RESULTS: 9% lower S2EF force MAE and 4% lower S2EF energy MAE are achieved by the 153M-parameter EquiformerV2 on OC20 S2EF-All+MD versus previous state-of-the-art results.The model also improves IS2RS AFbT by 4 percentage points and IS2RE energy MAE by 4%.
  • 4.1.3 ADSORBML RESULTS: 2× fewer DFT calculations are required by EquiformerV2 at k=2 than by competing models at k=5 for comparable AdsorbML adsorption-energy accuracy.Its AdsorbML accuracy improves over SCN by 8 and 5 percentage points at k=1 and k=2, respectively.
  • 4.2 OC22 DATASET: OC22-only EquiformerV2 outperforms GemNet-OC trained on both OC20 and OC22, while improving over eSCN by 18.9% in energy MAE and 8.9% in force MAE.OC22 contains about 8.4M structures versus about 130M in OC20, indicating better data efficiency.
  • 4.3 COMPARISON WITH EQUIFORMER: EquiformerV2’s gains are dataset- and task-dependent: they are not significant on small QM9, while force predictions benefit more than energy predictions.On OC20 S2EF-2M with Lmax=2, EquiformerV2 is 2.3× faster than Equiformer and achieves better force MAE with similar energy MAE.

5 CONCLUSION

EquiformerV2 scales equivariant Transformers to higher-degree representations by combining eSCN convolutions with three architectural improvements. It outperforms state-of-the-art methods across OC20 and OC22 tasks, improves speed-accuracy trade-offs, and performs strongly in AdsorbML.

  • EquiformerV2 replaces Equiformer’s SO(3) convolutions with eSCN convolutions and adds attention re-normalization, separable S2 activation, and separable layer normalization.
  • EquiformerV2 outperforms state-of-the-art methods on all reported OC20 and OC22 tasks.
  • EquiformerV2 improves speed-accuracy trade-offs and achieves the best success rate when used in AdsorbML.

6 ETHICS STATEMENT

EquiformerV2 more accurately approximates quantum mechanical calculations and may support practical applications in chemistry and materials science. The paper emphasizes that identifying promising molecules or materials does not remove the substantial challenges of synthesis and large-scale deployment.

  • EquiformerV2 achieves more accurate approximations of quantum mechanical calculations and moves toward machine-learned force fields for practical chemistry and materials applications.
  • The method facilitates identifying molecules or materials with specific properties, but substantial hurdles remain for synthesizing and deploying them at scale.

7 REPRODUCIBILITY STATEMENT

The paper provides mathematical, architectural, experimental, and training details across its appendices. It also reports that the implementation and reproduction code are publicly available for key datasets.

  • Architectural, hyper-parameter, training, dataset, and experiment details are organized across the OC20, OC22, and QM9 appendices.
  • Code for reproducing EquiformerV2 results on OC20 S2EF-2M and QM9 is available in the authors’ public GitHub repository.
  • The appendices provide background on group theory, equivariance, eSCN convolutions, S2 activation, and related equivariant neural-network methods.
  • The paper details tensor-product message passing, eSCN’s reduction of SO(3) convolutions to SO(2) linear operations, and S2 activation’s sphere-sampling procedure.

C DETAILS OF ARCHITECTURE

The architecture specifies dimensions for equivariant features, attention, feed-forward networks, radial functions, and sphere-sampled activations. It mainly targets SE(3)-equivariance and applies regularization to internal modules but not the output head.

  • The implementation mainly considers SE(3)-equivariance, combining three-dimensional rotation and translation equivariance without inversion equivariance.
  • Equivariant graph attention is parameterized by hidden, scalar-attention, value, and parallel-head dimensions, with sphere-sampled activation resolution R.
  • Feed-forward networks use intermediate dimension dffn, while radial functions generate edge-distance embeddings from radial bases and atomic-number embeddings.
  • EquiformerV2 uses maximum degree Lmax with C channels per degree, and embedding dimension dembed for most irreps features.
  • The model applies dropout to attention weights and stochastic depth to graph-attention and feed-forward outputs, excluding the output head.
  • For OC20, the core S2EF task predicts structure energy and per-atom forces from adsorbate-catalyst structures.

D.2 TRAINING DETAILS

The experiments evaluate EquiformerV2 on OC20, including ablations, large training splits, AdsorbML, and comparisons of accuracy, speed, and computational cost. Additional OC20 IS2RE results show that auxiliary-task training and Noisy Nodes affect whether higher-degree representations improve performance.

  • OC20 experiments: OC20 experiments cover S2EF-2M ablations, S2EF-All and S2EF-All+MD training, and AdsorbML evaluations.The study also reports model training time, inference speed, and parameter counts.
  • Relaxation procedure: Structural relaxations use LBFGS and stop after 200 steps or when the maximum predicted force per atom reaches 0.02 eV/Å.The cited relaxation procedure covers OC20 IS2RE and IS2RS, with separate AdsorbML settings.
  • AdsorbML: AdsorbML ranks ML-relaxed structures, evaluates the top k candidates with DFT single-point calculations, and computes success and speedup metrics.The procedure considers roughly 90k initial structures and varies k from 1 to 5.
  • AdsorbML: AdsorbML speed-accuracy comparisons average success rates over k=1 to 5 and measure average GPU-seconds per structure relaxation.The comparison is shown in Figure 5, alongside example relaxed structures in Figure 6.
  • IS2RE: Without IS2RS auxiliary training, EquiformerV2 overfits and performs worse than Equiformer, whereas IS2RS plus Noisy Nodes yields better energy MAE.The reported results are on the OC20 IS2RE validation split.

E DETAILS OF EXPERIMENTS ON OC22

The OC22 experiments train two EquiformerV2 models with different energy and force loss coefficients and report their computational requirements.

  • Training setup: Two EquiformerV2 models are trained on OC22 with different energy coefficients λE and force coefficients λF.Each model has 121.53M parameters and costs 4552 GPU-hours to train.
  • Computational cost: OC22 IS2RE relaxations require 38 GPU-hours.

F.1 ADDITIONAL RESULTS OF TRAINING WITH NOISY NODES

Additional QM9 experiments examine whether Noisy Nodes strengthens the performance gains associated with higher-degree representations. EquiformerV2 performs better than Equiformer on more tasks when trained with this augmentation, and the gain increases relative to standard training.

  • Noisy Nodes: Noisy Nodes adds coordinate noise and a node-level denoising task to training.The augmentation is used to test higher-degree models under an auxiliary node-level objective.
  • QM9 results: With Noisy Nodes, EquiformerV2 outperforms Equiformer on more QM9 tasks than without the augmentation.The comparison includes prior Noisy Nodes and denoising-pretraining methods.
  • Data efficiency: The more expressive model matches performance with significantly less data, despite QM9 being more than 30× smaller than the referenced dataset.

F.2 TRAINING DETAILS

QM9 training follows Equiformer’s data partition while varying task-specific optimization settings and adding Noisy Nodes hyperparameters for coordinate denoising.

  • Base training settings: Most QM9 tasks use batch size 64, 300 epochs, learning rate 5 × 10^-4, 128 radial bases, and six Transformer blocks.R2 uses a separate batch size and learning rate.
  • Hyperparameters: The QM9 hyperparameter table specifies optimizer, scheduling, regularization, model dimensions, maximum degree, and point-sample resolution.
  • Noisy Nodes settings: Noisy Nodes training introduces noise standard deviation, denoising coefficient, denoising probability, and corrupt ratio as additional hyperparameters.The denoising coefficient controls the auxiliary task’s relative importance, while the noise standard deviation applies Gaussian coordinate noise.

F.3 ABLATION STUDY ON ARCHITECTURAL IMPROVEMENTS

The ablation study evaluates proposed architectural improvements on QM9’s ∆ε task against Equiformer. Directly increasing Lmax from 2 to 4 and using eSCN convolutions degrades performance, attributed to overfitting on the smaller dataset.

  • Ablation setup: The study compares proposed architectural improvements with the Equiformer baseline on QM9’s ∆ε task.The results are summarized in Table 14.
  • Ablation findings: Increasing Lmax from 2 to 4 and using eSCN convolutions directly degrades performance relative to the baseline.
  • Ablation findings: The reported degradation is attributed to overfitting because QM9 is smaller, has fewer atoms, less diverse atom types, and less angular variation.
Loading 2306.12059v3…