Source-linked AI summary

EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers

Yi-Lun Liao, Alexander J. Hoffman, Sabrina C. Shen, Alexandre Duval, Sam Walton Norwood, Tess Smidt

arXiv:2604.09130v1cs.LGcs.AIphysics.comp-ph

TL;DR

Large-scale 3D atomistic modeling needs equivariant neural networks that improve efficiency, expressivity, and physical consistency. EquiformerV3 addresses this through implementation optimization, architectural refinements, and SwiGLU-S2 activations, achieving state-of-the-art results across three benchmarks. The method also supports smoothly varying potential energy surfaces for energy-conserving simulations and higher-order derivatives.

  • Problem

    Large-scale 3D atomistic modeling requires improvements in the efficiency, expressivity, and physical consistency of SE(3)-equivariant graph neural networks.

  • Method

    EquiformerV3 combines optimized implementation, EquiformerV2 architectural modifications, smooth-cutoff attention, and strictly equivariant SwiGLU-S2 activations.

  • Results

    EquiformerV3 achieves state-of-the-art results on OC20, OMat24, and Matbench Discovery.

  • Takeaways & Limitations

    Combining SwiGLU-S2 activations with smooth-cutoff attention extends EquiformerV3 to smoothly varying potential energy surfaces and tasks requiring energy-conserving simulations and higher-order PES derivatives.

  • Takeaways & Limitations

    Further efficiency improvements may require libraries such as cuEquivariance and OpenEquivariance, while practical deployment requires substantial effort to synthesize and scale materials or molecules.

Abstract

from arXiv · show

As $SE(3)$-equivariant graph neural networks mature as a core tool for 3D atomistic modeling, improving their efficiency, expressivity, and physical consistency has become a central challenge for large-scale applications. In this work, we introduce EquiformerV3, the third generation of the $SE(3)$-equivariant graph attention Transformer, designed to advance all three dimensions: efficiency, expressivity, and generality. Building on EquiformerV2, we have the following three key advances. First, we optimize the software implementation, achieving $1.75\times$ speedup. Second, we introduce simple and effective modifications to EquiformerV2, including equivariant merged layer normalization, improved feedforward network hyper-parameters, and attention with smooth radius cutoff. Third, we propose SwiGLU-$S^2$ activations to incorporate many-body interactions for better theoretical expressivity and to preserve strict equivariance while reducing the complexity of sampling $S^2$ grids. Together, SwiGLU-$S^2$ activations and smooth-cutoff attention enable accurate modeling of smoothly varying potential energy surfaces (PES), generalizing EquiformerV3 to tasks requiring energy-conserving simulations and higher-order derivatives of PES. With these improvements, EquiformerV3 trained with the auxiliary task of denoising non-equilibrium structures (DeNS) achieves state-of-the-art results on OC20, OMat24, and Matbench Discovery.

1. Introduction

EquiformerV3 advances equivariant graph attention Transformers along efficiency, expressivity, and generality through implementation optimization, architectural modifications, and SwiGLU-S2 activations. These changes support strong results across OC20, OMat24, and Matbench Discovery.

  • EquiformerV3 improves EquiformerV2 along efficiency, expressivity, and generality with three complementary advances.The advances are software optimization, architectural modifications, and SwiGLU-S2 activations.
  • 1.75× speedup results from fusing redundant operations and enabling compilation in the EquiformerV2 implementation.
  • SwiGLU-S2 incorporates many-body interactions, preserves strict equivariance, and reduces the complexity of sampling S2 grids.
  • Smooth-cutoff attention combined with SwiGLU-S2 enables modeling smoothly varying potential energy surfaces for energy-conserving simulations and higher-order PES derivatives.
  • EquiformerV3 achieves state-of-the-art results on OC20, OMat24, and Matbench Discovery.
  • On Matbench Discovery, EquiformerV3 improves thermal conductivity results by 18% to 31% compared to eSEN.

2. Related Works

Related work advances equivariant GNNs through more expressive nonlinearities, improved normalization, many-body interactions, and lower-cost tensor products. EquiformerV2 uses eSCN convolutions, separable S2 activation, and separable layer normalization to improve scalability and feature mixing.

  • EquiformerV2 adopts eSCN convolutions to reduce tensor-product complexity and support higher maximum representation degrees on large-scale datasets such as OC20.

3. Background

SE(3)-equivariant networks encode translation and rotation symmetries through relative positions and SO(3) irreducible representations. Their equivariant operations use tensor products and Transformer generalizations, while degree limits control cost and expressivity.

  • SE(3)-equivariant neural networks: SE(3)-equivariance uses relative positions for translation invariance and SO(3) irreducible representations for rotational equivariance.
  • SE(3)-equivariant neural networks: An irrep of degree L has dimension 2L + 1, and higher L encodes finer angular structure useful for direction-sensitive tasks such as force prediction.
  • SE(3)-equivariant neural networks: With uniform channel size C, an irreps feature containing degrees 0 through Lmax has total dimension (Lmax + 1)^2 × C.
  • Equivariant operations: Tensor products combine type-L1 and type-L2 features into type-L3 features through Clebsch–Gordan coefficients.
  • Equivariant operations: Restricting outputs to degrees up to Lmax controls computational cost and prevents feature dimensionality from growing without bound.
  • Equivariant Transformers: Equiformer generalizes Transformer operations to equivariant irreps features, including tensor products, linear layers, and normalization.
  • Equivariant Transformers: EquiformerV2 uses eSCN convolutions to reduce tensor-product complexity, alongside separable S2 activation and separable layer normalization.

4. EquiformerV3

EquiformerV3 improves EquiformerV2 through implementation optimizations, architectural modifications, smooth-cutoff attention, and SwiGLU-S2 activations. These changes target efficiency, training behavior, smooth PES modeling, expressivity, and strict equivariance.

  • Optimizing Software Implementation: 1.75× training speedup comes from fusing redundant operations and enabling compilation while maintaining the same accuracy.The implementation changes include pre-computing constant tensors and specifying scatter output shapes for dynamic-shape compilation.
  • Simple and Effective Modifications to EquiformerV2: Equivariant merged layer normalization shares a merged RMS across degrees, unlike normalization schemes that calculate statistics separately.The modification is intended to preserve relative importance between degrees during normalization.
  • Simple and Effective Modifications to EquiformerV2: The feedforward hidden size increases by 4× because node-wise feedforward networks are substantially cheaper than edge-wise tensor-product components.The larger hidden dimension increases model capacity with minimal additional overhead.
  • Simple and Effective Modifications to EquiformerV2: Smooth radius cutoff applies envelope functions to both messages and softmax attention, preventing denominator changes from causing discontinuities as neighbors cross the cutoff.This modification targets smoothly varying potential energy surfaces with respect to atomic positions.
  • SwiGLU-S2 Activation: SwiGLU-S2 projects irreps features onto S2, applies activation and multiplication to grid features, then returns them to irreps space.ToSphere and FromSphere use sampled longitude-latitude grids, whose density affects equivariance error.
  • SwiGLU-S2 Activation: SwiGLU-S2 introduces many-body interactions while scalar-only nonlinear gating preserves strict equivariance and reduces required S2 grid points.For Lmax = 6, EquiformerV3 reduces attention grid points by 50.6% while maintaining strict equivariance.

5. Experiments

Experiments evaluate EquiformerV3 through OC20 ablations and comparisons on OMat24 and Matbench Discovery. Results show improved efficiency, competitive or stronger accuracy, and broad benchmark performance across energy, force, stress, stability, and thermal-conductivity metrics.

  • 5.1. OC20: The combined improvements decrease energy MAE by 41 meV, force MAE by 1.58 meV/ ˚A, and save 1.58× training time.Reducing attention grid points from 324 to 160 maintains similar training time while allowing more parameters.
  • 5.2. OMat24 Dataset: EquiformerV3 with Lmax = 4 achieves similar OMat24 results to UMA-L while being 23× smaller, and Lmax = 6 further improves performance.The comparison uses direct pre-training and gradient fine-tuning settings with regularization and DeNS.
  • 5.3. Matbench Discovery: EquiformerV3 achieves the best results on all Matbench Discovery metrics among compliant models trained only on MPtrj.Compared with EquiformerV2, κSRME decreases from 1.676 to 0.275.
  • 5.3. Matbench Discovery: EquiformerV3 consistently outperforms UMA-M-1.1 on Matbench Discovery while saving 22.6× training time.It also reduces OMat24 training epochs from 30 to 6 relative to NequIP-OAM-XL and from 15 to 6 relative to PET-OAM-XL.

6. Conclusion

EquiformerV3 enhances SE(3)-equivariant graph attention Transformers across efficiency, expressivity, and generality, achieving state-of-the-art results on three benchmarks.

  • EquiformerV3 combines optimized software, architectural modifications, and SwiGLU-S2 activations to improve efficiency, expressivity, and generality.

Impact Statement

EquiformerV3 supports more accurate approximation of quantum mechanical calculations for chemistry and materials applications, while deployment still requires substantial synthesis and scaling effort.

  • EquiformerV3 enables more accurate approximation of quantum mechanical calculations for accelerating chemistry and materials applications.
  • The methods facilitate identifying molecules or materials with specific properties but require substantial efforts to synthesize and deploy them at scale.

A. Overall Architecture

The architecture combines equivariant attention, smooth radius cutoffs, SwiGLU-S2 activations, and specialized embedding and output modules for graph-level and node-wise predictions.

  • Equivariant Graph Attention: Equivariant graph attention rotates edge features into a local frame, computes scalar attention logits and irreps values, then rotates outputs back.
  • Equivariant Graph Attention: Smooth radius cutoff is incorporated into attention weights, while SwiGLU-S2 transforms value vectors before attention aggregation.
  • Feedforward Network: The feedforward network uses SwiGLU-S2 between ToSphere and FromSphere, with linear layers mixing grid-feature channels.
  • Embedding: Embedding modules encode atom types and edge-degree information using atom embeddings, SO(2) transformations, inverse rotations, and sum aggregation.
  • Output Head: Graph-level heads aggregate transformed node irreps for energy and stress, whereas node-wise force and noise heads use equivariant attention with gate activation.

B. Body-Order Experiment

The body-order experiments test whether activation functions and stacked feedforward networks distinguish geometric graphs requiring higher-order scalarization.

  • Body-Order Experiment: Gate and S2 activations remain limited to 2-body scalarization, yielding 50% accuracy on the counterexamples regardless of the number of FFNs.
  • Body-Order Experiment: Table 5 marks successful graph distinction at 100% accuracy and unsuccessful distinction at 50% accuracy.
  • Body-Order Experiment: A single SwiGLU-S2 activation captures interactions similar to self tensor products and achieves 3-body scalarization.
  • Body-Order Experiment: One SwiGLU-S2 FFN distinguishes the 2-body counterexample with 100% accuracy, while stacking additional FFNs supports higher body orders.

C. Equivariance Errors of Different Activation Functions

This section evaluates SwiGLU-S2 against S2 and gate activations, showing comparable equivariance errors with fewer sampled S2 grid points, including in attention.

  • SwiGLU-S2 achieves equivariance errors similar to gate activation while using more than 2× fewer grid points than S2 activation.For Lmax = 2, it uses 64 grid points, compared with 144 for S2 activation.
  • 160 grid points are required for SwiGLU-S2 in attention under Lmax = 6 and Mmax = 2.This is the default OC20 configuration described in the passage.
  • 50.6% lower sampling complexity is achieved in attention relative to EquiformerV2.EquiformerV2 uses 324 attention grid points, whereas SwiGLU-S2 uses 160.
  • The activation experiments vary maximum degree Lmax and sampled grid resolutions Rϕ and Rθ, using gate activation as the strict-equivariance baseline.Results are marked according to whether their equivariance errors are comparable to the gate baseline.
  • The base model specifies irreps feature dimensions by maximum degree Lmax and channel count C, with SwiGLU-S2 reducing channels after activation.The described configuration changes dimensions from (6, 128) to (6, 64).

E. Details of Experiments on OMat24

The OMat24 experiments use direct pre-training and gradient fine-tuning on approximately 100M AIMD and rattled structures, with regularization and DeNS denoising.

  • Approximately 100M AIMD and rattled OMat24 structures are used for direct pre-training and gradient fine-tuning.Training uses FP32 for these experiments.
  • DeNS denoising is used as an auxiliary task during direct pre-training.The training setup also applies dropout and stochastic depth regularization.

F. Details of Experiments on Matbench Discovery

The Matbench Discovery experiments include MPtrj-only training and fine-tuning of OMat24-pretrained EquiformerV3 models on MPtrj and subsampled Alexandria data.

  • MPtrj-only training uses FP16 for direct pre-training and FP32 for gradient fine-tuning.The setup includes dropout, stochastic depth, and DeNS denoising during direct pre-training.
  • EquiformerV3 models pretrained on OMat24 with Lmax = 4 are fine-tuned on MPtrj and subsampled Alexandria datasets.The fine-tuning setup excludes structures with more than 24000 edges and uses no regularization or DeNS.
  • Fine-tuning initializes all model weights from the OMat24 checkpoint.The attention softmax denominator includes a small ϵ = 10^-8.
  • The experiments define separate hyper-parameter configurations for MPtrj and OMat24-pretrained fine-tuning.These configurations are listed in Tables 10 and 11.
Loading 2604.09130v1…