PaperlayerView exact revision on arXiv

Source-linked AI summary

Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs

Yi-Lun Liao, Tess Smidt

arXiv:2206.11990v2cs.LGcs.AIphysics.comp-ph

TL;DR

Transformers have struggled to generalize across 3D atomistic graph datasets despite relevant geometric inductive biases. Equiformer combines Transformer architectures with SE(3)/E(3)-equivariant irreps features, tensor products, and equivariant graph attention, achieving effective results across QM9, MD17, and OC20 while inheriting computational limitations of irreps-based methods.

  • Problem

    Transformers have yet to perform well across 3D atomistic graph datasets even when relevant inductive biases are incorporated.

  • Method

    Equiformer combines Transformer operations replaced by equivariant counterparts with irreps features, tensor products, MLP attention, and non-linear message passing.

  • Results

    Equiformer demonstrates effectiveness across QM9, MD17, and OC20, with better results across 12 QM9 regression tasks than each individual comparison model.

  • Takeaways & Limitations

    Equivariant graph attention improves typical Transformer attention through MLP attention and non-linear messages in the evaluated atomistic graph tasks.

  • Takeaways & Limitations

    Higher irreps degrees enlarge features, and tensor products can be computationally expensive, partly because their kernels are not heavily optimized.

Abstract

from arXiv · show

Despite their widespread success in various domains, Transformer networks have yet to perform well across datasets in the domain of 3D atomistic graphs such as molecules even when 3D-related inductive biases like translational invariance and rotational equivariance are considered. In this paper, we demonstrate that Transformers can generalize well to 3D atomistic graphs and present Equiformer, a graph neural network leveraging the strength of Transformer architectures and incorporating SE(3)/E(3)-equivariant features based on irreducible representations (irreps). First, we propose a simple and effective architecture by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empirical results. Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. With these two innovations, Equiformer achieves competitive results to previous models on QM9, MD17 and OC20 datasets.

1 INTRODUCTION

Transformers have struggled on 3D atomistic graphs despite relevant inductive biases. Equiformer addresses this gap by combining Transformer architectures with SE(3)/E(3)-equivariant irreps features and equivariant operations.

  • GNNs model atomistic systems as graphs, updating node features through messages between atoms and exploiting their set-like structure.
  • 3D atomistic models require inductive biases related to E(3), including transformations of coordinates and physical quantities.
  • Transformers have yet to perform well across 3D atomistic graph datasets even when relevant inductive biases are incorporated.
  • Equiformer combines Transformer architectures with SE(3)/E(3)-equivariant irreps features and a novel attention mechanism.
  • Its initial architecture replaces Transformer operations with equivariant counterparts and adds tensor products without complicating graph structures.

3 BACKGROUND

The paper represents 3D geometric information with irreducible representations and propagates it through equivariant operations. Tensor products combine feature types along constrained, learnable paths, while spherical harmonics encode relative positions.

  • E(3) includes translation, rotation, and inversion, while atomistic properties may remain invariant or transform equivariantly under coordinate changes.
  • Relative positions address translation symmetry, and scalar and Euclidean-vector features transform differently under rotations.
  • SO(3) irreps use degree L and dimension 2L+1 to represent feature types with distinct rotational behavior.
  • Irreps features concatenate multiple type-L vectors across channels, with regular scalar features corresponding to type-0 vectors.
  • Spherical harmonics project relative positions into equivariant type-L features, allowing equivariant information to propagate through tensor products.
  • Tensor products combine type-L1 and type-L2 vectors into allowed type-L3 outputs, while discarding degrees above Lmax for efficiency.
  • Each L1 ⊗ L2 → L3 path is independently equivariant and can receive its own learnable weight, including across multiple channels.
  • Equiformer applies equivariant graph attention and feed-forward Transformer blocks to embedded atomistic graphs, using tensor-product operations within the architecture.

4 EQUIFORMER

Equiformer adapts Transformer operations to SE(3)/E(3)-equivariant irreps features and adds tensor products. Its equivariant graph attention combines geometric-content messages, MLP attention, and nonlinear message passing.

  • Architecture: Equiformer replaces Transformer operations with equivariant counterparts and adds tensor products for SE(3)/E(3)-equivariant irreps features.This architecture encodes equivariant information in feature channels rather than complicating graph structures.
  • Equivariant operations: Equivariant linear layers transform each degree-L vector group separately and omit biases for non-scalar features because they can break equivariance.Biases are removed for L > 0 features because they do not depend on inputs.
  • Equivariant operations: Equivariant normalization uses the root mean square of type-L vector norms across channels and removes means and biases for L ≠ 0.This generalizes layer normalization to irreps features while preserving equivariance.
  • Equivariant operations: Gate activation applies nonlinear scalar functions to produce weights that gate non-scalar irreps features while preserving equivariance.SiLU is applied to selected scalars and sigmoid to gating scalars before multiplying corresponding type-L vectors.
  • Equivariant operations: Depth-wise tensor products restrict each output channel to one input channel, reducing weights and memory complexity when weights depend on relative distances.Weights may be input-independent or conditioned on relative distances.
  • Equivariant graph attention: Equivariant graph attention combines target and source features with spherical-harmonic geometric embeddings, MLP attention, and nonlinear message transformations.Attention uses scalar features, while values can use nonlinear gated irreps transformations and tensor products.

5 EXPERIMENT

Equiformer is evaluated on QM9, MD17, and OC20, with results showing strong performance across regression tasks and comparisons against equivariant and conventional models.

  • Experimental setup: Equiformer is benchmarked on QM9, MD17, and OC20, with ablations examining dot product versus MLP attention and linear versus non-linear message passing.The experiments include QM9 quantum-property prediction, MD17 energy and force prediction, and OC20 relaxed-energy prediction.
  • QM9: Across 12 QM9 regression tasks, Equiformer achieves overall better results than each individual compared model.The comparison includes SEGNN and other equivariant Transformers such as SE(3)-Transformer, TorchMD-NET, and EQGAT.
  • MD17: Across 8 MD17 molecules, Equiformer achieves overall better results than each individual compared model.Its equivariant graph attention is more expressive than TorchMD-NET’s alternative and supports vectors of higher degree.
  • MD17: Increasing Lmax from 2 to 3 improves MAE for most MD17 molecules except benzene, where the higher setting results from overfitting.For the last three molecules, adjusting energy and force loss weights can yield lower MAE for both quantities with Lmax = 2.
  • Ablation study: The ablations show that dot product attention with linear message passing is already competitive, while non-linear messages improve results with MLP attention at higher training cost.MLP attention achieves similar results to dot product attention, whereas non-linear messages increase tensor-product operations per block from 1 to 2.

6 CONCLUSION

The paper presents Equiformer as a GNN combining Transformer architectures with irreps-based equivariant features and equivariant graph attention. Experiments on QM9, MD17, and OC20 support its effectiveness and the proposed attention mechanism.

  • Contributions: Equiformer combines Transformer architectures with irreps-based equivariant features through equivariant operations such as tensor products.The model is designed as a graph neural network for 3D atomistic graphs.
  • Contributions: Equivariant graph attention combines multi-layer perceptron attention with non-linear messages.Ablation studies report improvement over typical Transformer attention.
  • Results: Experiments on QM9, MD17, and OC20 demonstrate Equiformer’s effectiveness.The conclusion summarizes results across all three benchmark datasets.

7 ETHICS STATEMENT

The ethics statement frames accurate quantum-property approximation as useful for investigating molecules and materials, while noting deployment barriers.

  • Potential benefits: Equiformer’s more accurate quantum-property approximations are presented as useful for investigating molecules and materials in energy, electronics, and pharmaceuticals.The statement contrasts productive applications with adversarial uses such as creating hazardous chemicals.

8 REPRODUCIBILITY STATEMENT

The paper provides reproducibility resources and mathematical background for its equivariant architecture, including irreps, tensor products, parity, and group symmetries.

  • Reproduction resources: Reproduction details cover architectures, hyper-parameters, and training time for QM9, MD17, and OC20.The paper points readers to dedicated appendix sections for each dataset.
  • Reproduction resources: Code for reproducing Equiformer results on QM9, MD17, and OC20 is publicly available.The repository URL is provided in the reproducibility statement.
  • Irreps features: Equiformer uses type-L or type-(L, p) irreps features, with parity indicating whether features change sign under inversion.Scalar features correspond to type-0 for SE(3) or type-(0, e) for E(3).
  • Tensor products: Tensor products combine irreps features using Clebsch-Gordan coefficients to produce output features with specified degree and parity.For O(3), parity follows e × e = e, o × o = e, and e × o = o.

B.2 DETAILED COMPARISON BETWEEN EQUIVARIANT TRANSFORMERS

Equiformer addresses weaknesses in prior equivariant Transformers with a simpler, more general architecture and a new equivariant graph-attention mechanism. It achieves strong results across QM9, MD17, and OC20 while remaining efficient on large datasets.

  • Prior limitations: Previous equivariant Transformers do not perform well consistently across QM9 and MD17.SE(3)-Transformer underperforms other equivariant networks on QM9, while TorchMD-NET is not comparable to NequIP on MD17.
  • Results: Equiformer achieves the best results simultaneously on MD17, QM9, and OC20.This cross-dataset outcome supports the effectiveness of Transformer architectures for equivariant neural networks on 3D atomistic graphs.
  • Large-scale evaluation: On OC20, Equiformer matches large models such as GNS and Graphormer while saving 2.3× to 15.5× training time.The paper presents this as its first application of equivariant Transformers to large, complicated datasets like OC20.
  • Architecture: The base architecture replaces Transformer operations with equivariant counterparts and adds tensor products without further structural modifications.This design corresponds to Equiformer with dot product attention and linear message passing.
  • Architecture: Equiformer supports vectors of any degree L, whereas some competing equivariant Transformers are limited to L = 0 and 1.Higher degrees such as L up to 2 and 3 have been reported to improve QM9 and MD17 performance.
  • Efficiency: Depth-wise tensor products allow more channels than SE(3)-Transformer without incurring out-of-memory errors.Both architectures parameterize tensor-product weights using relative distances, but depth-wise products are more efficient.
  • Attention: Equivariant graph attention combines MLP attention with non-linear message passing to improve the base architecture.The paper also conducts extensive ablation studies focused on the attention mechanism.

C DETAILS OF ARCHITECTURE

The architecture builds equivariant Transformer modules from irreps-aware operations and depth-wise tensor products. Its graph attention uses scalar-based MLP attention, while non-linear messages improve expressiveness at increased computational cost.

  • Equivariant operations: Equiformer uses equivariant versions of Transformer operations together with depth-wise tensor products as its core building blocks.These operations are analyzed through the equivariance condition f(DX(g)x) = DY(g)f(x).
  • Equivariant operations: Linear layers preserve equivariance by combining type-L vectors that transform with the same representation matrix.The output type-L vector therefore transforms according to the same matrix as its inputs.
  • Equivariant operations: Layer normalization remains equivariant because scalar features are invariant and non-scalar vector norms are invariant under E(3) transformations.Scaling by RMS norms and learnable parameters therefore does not change the transformation behavior.
  • Tensor products: Depth-wise tensor products restrict each output channel to one input channel while retaining interactions between different vector degrees.This gives an efficiency pattern analogous to depth-wise convolution.
  • SE(3) and E(3): The main experiments use SE(3)-equivariant irreps, while E(3)-equivariance can be incorporated with straightforward representation and operation changes.E(3) handling distinguishes even and odd parity types throughout linear, normalization, tensor-product, and gate operations.
  • Attention: Dot-product attention derives weights from scaled products of queries and keys, whereas MLP attention uses scalar features for attention weights.The key and value are split from irreps features, while the query is obtained with a linear layer.
  • Computational complexity: Non-linear messages increase tensor-product operations from 1 to 2, inevitably increasing training and inference time.Tensor products are compute-intensive, so the added expressiveness carries a direct computational cost.

D.3 COMPARISON OF TRAINING TIME AND NUMBERS OF PARAMETERS

Equiformer’s training cost reflects its use of efficient depth-wise products and expressive non-linear messages. It can reach competitive results with substantially fewer epochs than TorchMD-NET, though higher expressiveness increases per-epoch cost.

  • Training cost: Equiformer training for 300 and 600 epochs takes 61 and 122 GPU-hours, respectively.These measurements are reported for QM9.
  • Comparison: Equiformer with MLP attention and non-linear messages is faster than SEGNN despite having more channels and parameters.The comparison attributes this to depth-wise tensor products replacing SEGNN’s more compute-intensive fully connected tensor products.
  • Comparison: Equiformer reaches competitive results after 300 or 600 epochs, compared with TorchMD-NET’s 3000 epochs.Its epochs take longer because non-linear messages double the number of tensor products and higher degrees such as Lmax = 2 slow training.
  • Experimental setup: Training configurations normalize energy targets by mean and standard deviation and force targets by the standard deviation of ground-truth energy.The reported QM9 experiments use 6 Transformer blocks, Lmax = 2 or 3, and about 3.50M parameters.
  • Experimental setup: For MD17, Equiformer with Lmax = 2 is compared with TorchMD-NET using adjusted force-loss weights to reduce both energy and force MAE.The table reports energy and force in meV and meV/Å.

E.3 COMPARISON OF TRAINING TIME AND NUMBERS OF PARAMETERS

On MD17, Equiformer generally outperforms NequIP while using a smaller maximum degree in the faster comparison. Its two tensor products per attention block nevertheless make training longer.

  • Training cost: Equiformer with Lmax = 2 is faster than NequIP with Lmax = 3 because it uses a smaller maximum degree.The paper also reports overall better results for this Equiformer configuration.
  • Training cost: Equiformer training takes longer than NequIP because each equivariant graph-attention layer uses two tensor products instead of one linear-convolution tensor product.The additional operation reflects the use of non-linear messages.

F.2 TRAINING DETAILS

The OC20 experiments compare Equiformer variants, equivariance choices, training costs, and error distributions across validation sub-splits. SE(3)-equivariant Equiformer generally outperforms the E(3) variant, while stronger models improve different sub-splits unevenly.

  • Training setup: Equiformer with IS2RS auxiliary training uses deeper networks and an additional equivariant graph attention module to predict relaxed structures.The reported configuration uses 18 Transformer blocks.
  • Training cost and model comparison: Training Equiformer without IS2RS takes about 43.6 hours, compared with 58.3 hours for E(3)-Equiformer.Linear-message variants take about 30.4 and 33.1 hours, respectively, under the reported setup.
  • Reported comparisons: The OC20 validation results and ablation comparisons are summarized in Tables 15–17.These tables cover validation performance, SE(3)/E(3) ablations, and training time with parameter counts.
  • SE(3) versus E(3) equivariance: Including inversion improves MAE on the ID and OOD Cat sub-splits but worsens other sub-splits, making E(3)-Equiformer slightly inferior overall.The authors suggest inversion may not be the key bottleneck and that E(3) features may require different hyper-parameters.
  • Training cost and model comparison: Equiformer achieves better results than SEGNN with comparable training time when IS2RS auxiliary training is omitted.The comparison also reports training time and parameter counts.
  • Error distributions: Error improvements are non-uniform across OC20 sub-splits: MLP attention and non-linear messages help ID more than OOD Ads.Without IS2RS auxiliary training, stronger models mainly reduce errors above 0.02 eV, which may not improve EwT.

G LIMITATIONS

Equiformer inherits computational and scalability constraints from irreps-based equivariant networks, and its stronger attention mechanisms do not improve every task equally. The proposed attention is local, but its cost depends on graph edges and may require adaptation outside 3D atomistic graphs.

  • Irreps and tensor-product costs: Higher irreps degrees can enlarge features, while tensor products can be computationally expensive because their kernels are not heavily optimized.The authors leave path pruning and other efficiency gains for future work.
  • Task and dataset dependence: The benefit of equivariant graph attention depends on the task and dataset: MLP attention helps OC20 clearly but not QM9 significantly.The paper attributes this difference to QM9 having fewer atoms and less diverse atom types.
  • Computational overhead: Equivariant graph attention requires more computation than typical graph convolution.Softmax adds a sum aggregation, and non-linear messages increase tensor products from one to two.
  • Computational overhead: Under a constrained training budget, stronger attention may be suboptimal when its improvement is small and training becomes slower.The relevant stronger components are MLP attention and non-linear messages.
  • Scope beyond atomistic graphs: The attention complexity scales with the product of channel and edge counts because attention is restricted to local neighborhoods.Using the mechanism in domains such as computer vision would require further modifications.
Loading 2206.11990v2…