Source-linked AI summary
Geometric and Physical Quantities Improve E(3) Equivariant Message Passing
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, Max Welling
TL;DR
Equivariant graph networks have largely restricted node and edge information to invariant scalars, despite tasks requiring geometric and physical quantities. SEGNNs address this gap with steerable MLP-based message passing that incorporates covariant information into messages and updates. Across computational physics and chemistry tasks, the paper reports benefits from steerable and nonlinear message passing, with SEGNNs outperforming all competitors on OC20 ISRE.
Problem
Equivariant graph networks often restrict node and edge attributes to invariant scalars, although computational physics and chemistry tasks use quantities such as force, velocity, and spin.
Method
SEGNNs use steerable node features and steerable MLPs to process scalar, vector, and tensor attributes in equivariant message and update functions.
Results
SEGNNs outperform all competitors on the OC20 ISRE task, while ablations show benefits of steerable over invariant messages and nonlinear over linear convolutions.
Takeaways & Limitations
Steerable geometric and physical information can be integrated into equivariant graph networks across a wide range of computational physics and chemistry tasks.
Takeaways & Limitations
Applications relying on learned simulations or simulated quantities require rigorous cross-checks and monitoring.
Abstract
from arXiv · showhide
Including covariant information, such as position, force, velocity or spin is important in many tasks in computational physics and chemistry. We introduce Steerable E(3) Equivariant Graph Neural Networks (SEGNNs) that generalise equivariant graph networks, such that node and edge attributes are not restricted to invariant scalars, but can contain covariant information, such as vectors or tensors. This model, composed of steerable MLPs, is able to incorporate geometric and physical information in both the message and update functions. Through the definition of steerable node attributes, the MLPs provide a new class of activation functions for general use with steerable feature fields. We discuss ours and related work through the lens of equivariant non-linear convolutions, which further allows us to pin-point the successful components of SEGNNs: non-linear message aggregation improves upon classic linear (steerable) point convolutions; steerable messages improve upon recent equivariant graph networks that send invariant messages. We demonstrate the effectiveness of our method on several tasks in computational physics and chemistry and provide extensive ablation studies.
1 INTRODUCTION
The paper extends equivariant graph message passing beyond scalar node and edge attributes, enabling geometric and physical information to steer computation. SEGNNs use steerable representations and nonlinear message passing to exploit these cues while preserving transformation-aware structure.
- The authors frame SEGNNs and related models as equivariant non-linear convolutions, connecting graph aggregation with convolutional message passing.
- SEGNNs generalise equivariant graph networks by allowing node and edge attributes to be scalar-, vector-, or tensor-valued.
- Steerable node attributes let message and update functions incorporate geometric cues such as relative orientation, force, velocity, and spin.
- The framework uses steerable MLPs to provide a new class of equivariant activation functions for steerable feature fields.
- Extensive ablations compare steerable with invariant message passing and nonlinear with linear convolutions.
2 GENERALISED E(3) EQUIVARIANT STEERABLE MESSAGE PASSING
The section develops E(3)-equivariant message passing using steerable features, whose transformations preserve geometric structure while supporting scalar, vector, and tensor attributes. Steerable MLPs use geometric or physical attributes to condition equivariant message and node-update functions.
- Equivariant message passing: Equivariance makes graph networks robust to rotations, reflections, translations, and permutations while supporting invariant energy and equivariant force prediction.Equivariant message and update functions guarantee equivariance of the entire network.
- Steerable features: Steerable vectors transform by representation matrices D(g), with type-l spaces having dimension 2l + 1 and combined spaces formed by direct sums.Wigner-D matrices act separately on the corresponding type-l vector spaces.
- Steerable MLPs: Steerable MLPs interleave equivariant linear maps and nonlinearities, with Clebsch-Gordan products ensuring mappings between steerable vector spaces.The conditioning steerable vector can encode geometric information such as relative atom positions.
- Geometric embeddings: Spherical harmonic embeddings convert vectors into type-l steerable vectors, allowing geometric and physical information to enter steerable MLPs.The paper uses these embeddings for quantities such as relative orientation, velocity, or force.
- Steerable activation functions: Steerable MLPs provide a new activation-function class that directly leverages local geometric cues through physical or predicted steerable node attributes.The MLPs can be applied node-wise to steerable feature fields.
- Steerable E(3) graph networks: SEGNN layers steer messages with edge attributes and node updates with node attributes, so geometric and physical quantities influence both stages.Additional attributes may include distances, forces, spins, or velocities.
3 MESSAGE PASSING AS CONVOLUTION, RELATED WORK
This section frames message passing and point convolutions within a unified non-linear convolution perspective. It distinguishes linear or pseudo-linear equivariant convolutions from methods whose messages are nonlinear and whose features carry covariant geometric information.
- Non-linear convolution: Equivariant linear operators are group convolutions, and equivariant message-passing operators written in simple message-passing form are non-linear convolutions.This framing unifies related architectures and highlights equivariant and nonlinear layers as important design aspects.
- Point convolutions: Point convolutions implement linear message passing whose transformations are conditioned on relative positions and may use spherical-harmonic basis coefficients.Continuous kernels are needed for non-uniform point-cloud grids.
- Steerable convolutions: Steerable convolutions represent rotated kernels through steerable transformations, producing feature responses associated with the full E(3) group.Explicit kernel rotations and steerable transformations provide equivalent convolutional views.
- Related equivariant methods: Attention-based equivariant methods make convolution weights input-dependent, creating pseudo-linear transformations because feature transformations remain linear.The attention weights multiply the steerable convolution operator.
- Equivariant message passing: EGNNs use nonlinear but invariant messages conditioned on node features and squared pairwise distances, making them isotropic non-linear convolutions.Their messages are nonlinear in neighboring features but depend on distance rather than directional geometry.
- Steerable message passing: SEGNN-style messages condition steerable MLPs on attributes such as spherical-harmonic embeddings, enabling messages that carry covariant geometric information.This makes the messages more general than those based only on invariant distances.
4 EXPERIMENTS
Experiments evaluate SEGNNs on N-body, QM9, and OC20 tasks, comparing steerable features, attributes, graph connectivity, and competing architectures. Across these evaluations, the reported results support steerable, nonlinear message passing and geometric or physical attribute injection.
- Experiments: SEGNN experiments cover charged N-body position prediction, QM9 molecular-property regression, and OC20 IS2RE energy prediction.The experiments use MSE for N-body and MAE-based evaluation for QM9 and OC20.
- Comparisons: SEGNNs are compared with EGNN, steerable linear models, steerable nonlinear point convolutions, and OC20 leaderboard competitors.EGNN is recovered with lf = la = 0, while SElinear represents a linear steerable implementation.
- QM9: SEGNNs outperform an architecturally comparable EGNN on QM9 while removing attention modules and reducing connectivity to atoms within 2 Å.The paper notes that runtime remains a relevant consideration despite the performance comparison.
5 CONCLUSION
The conclusion presents SEGNNs as a generalisation of equivariant graph networks that supports vector- and tensor-valued node and edge information. Experiments and ablations report benefits from steerable over invariant message passing and from nonlinear over linear convolutions.
- 5 CONCLUSION: SEGNNs generalise equivariant graph networks by allowing node and edge information to be scalar-, vector-, or tensor-valued.The architecture is presented as supporting geometric and physical cues in node updates through steerable features.
- 5 CONCLUSION: Ablation studies report benefits for steerable over non-steerable message passing and for nonlinear over linear convolutions.The conclusion summarizes these findings across the paper’s experiments.
- 5 CONCLUSION: On the OC20 ISRE task, SEGNNs outperform all competitors.This is the paper’s explicit conclusion for the OC20 benchmark.
6 REPRODUCIBILITY STATEMENT
The reproducibility statement describes error bars, repeated-seed checks, ablations, verified comparisons, and mathematical documentation. It also states that the paper introduces no new mathematical results.
- 5 REPRODUCIBILITY STATEMENT: N-body experiments report means and standard deviations from eight runs with different initial seeds, alongside reproduced prior results and feature-order ablations.The authors also describe including error bars and reproducibility checks where considered necessary.
- 5 REPRODUCIBILITY STATEMENT: For OC20, the paper reports official numbers from the Open Catalyst Project challenge for comparisons with other methods.The statement identifies these as the basis for the OC20 comparisons.
- 5 REPRODUCIBILITY STATEMENT: The architecture is described in Section 2.1 with further implementation details in Appendix Section C, including proofs of equivariance and properties of Wigner D matrices and spherical harmonics.The appendix also provides a self-contained mathematical exposition.
- 5 REPRODUCIBILITY STATEMENT: The paper does not introduce new mathematical results but gives a unifying non-linear-convolution view of equivariant graph neural networks.The statement frames the mathematical material as exposition and analysis of existing concepts.
7 ETHICAL STATEMENT
The ethical statement says SEGNNs may support physical and chemical modeling and potentially reduce simulation costs, while emphasizing the need for rigorous checks when learning simulated quantities.
- 7 ETHICAL STATEMENT: SEGNNs are suited to physical and chemical modeling and may serve as shortcuts for computationally expensive simulations.The statement presents potential carbon-footprint reduction only as a possible indirect consequence.
- 7 ETHICAL STATEMENT: Learned simulations or simulated quantities require rigorous cross-checks and monitoring.The statement identifies this as a downside of relying on simulations.
A MATHEMATICAL BACKGROUND
The appendix develops the group-theoretic foundations for modeling translations, rotations, and reflections, then formalizes representations, equivariance, and invariance. It also introduces steerable vectors transformed by O(3) representations.
- Steerable vectors and representations: Steerable feature vectors transform through Wigner-D matrices, which represent orthogonal-group actions on feature components.Steerable MLPs use feature vectors composed of such steerable vectors rather than only scalar elements.
- The Euclidean group E(3): E(3) models three-dimensional translations, rotations, and reflections as transformations parametrized by translation vectors and orthogonal matrices.E(3) decomposes into translations and O(3) transformations.
- The orthogonal group O(3) and special orthogonal group SO(3): O(3) contains rotations and reflections, whereas SO(3) contains only determinant-1 rotations.
- Homogeneous spaces: The sphere S2 is treated as a homogeneous space because group elements can map any point to any other point, although S2 itself is not a group.
- Group representations: A representation maps group elements to invertible linear transformations on a vector space while preserving the group structure.Representations can act on finite-dimensional vectors or functions.
- Invariance and equivariance: An equivariant operator transforms its output predictably under input transformations, while an invariant operator leaves its output unchanged.
A.5 CLEBSCH-GORDAN PRODUCT AND STEERABLE MLPS
This section explains how Clebsch-Gordan products combine steerable inputs into steerable outputs and form the linear component of steerable MLPs. Sparsity, band-limiting, and conditioning order constrain the resulting interactions and computation.
- Steerable MLPs: Steerable MLPs use a learnable Clebsch-Gordan tensor product in place of a regular MLP’s weight matrix.One input can be fixed, making the product act as a linear map between steerable spaces.
- Clebsch-Gordan tensor product: The Clebsch-Gordan tensor product maps steerable inputs to outputs whose components transform according to O(3) representations.The resulting tensor-product representation can be reorganized into subspaces of different steerable types.
- Clebsch-Gordan tensor product: The Clebsch-Gordan product is sparse because many coefficients vanish, reducing the set of permitted type interactions.
- Clebsch-Gordan tensor product: Clebsch-Gordan coefficients enforce allowable output types, with l restricted by |l1−l2| ≤ l ≤ l1+l2.Examples include scalar, dot, and cross products.
- Computational considerations: Clebsch-Gordan tensor products become slow for higher-order irreducible representations because multiple operations are required for each connection.
- Computational considerations: Hidden representations are band-limited to lmax = lf, and conditioning order la determines which output types can interact.Using an attribute order much larger than the hidden-feature order is described as not sensible.
B STEERABLE GROUP CONVOLUTIONS
The appendix connects steerable vectors and spherical functions to group convolutions, framing steerable message passing as non-linear group convolution. Steerable representations avoid discretization limitations and provide exact O(3) equivariance.
- Steerable functions: Steerable functions transform under a group through linear combinations of a fixed finite basis, linking them to steerable vectors.
- Group convolutions: Group convolutions operate on feature maps over positions and group transformations, with kernels represented as functions on the group.
- Steerable convolution kernels: Spherical harmonics provide a continuous basis for expanding steerable convolution kernels and representing their transformation behavior.The basis coefficients can depend on radial distance while remaining O(3)-invariant.
- Steerable convolution kernels: Steerability lets convolution responses at arbitrary rotations be obtained by transforming basis-coefficient vectors before reduction.
- Steerable convolution kernels: When only m = 0 kernel components are used, axial symmetry yields point-wise spherical signals and improves computational efficiency.The paper uses this restriction in weighted Clebsch-Gordan products, although full steerable group convolutions need not impose it.
- Exact equivariance: Regular group convolutions cannot be perfectly equivariant to all O(3) transformations because of discretization, whereas steerable group convolutions are exactly O(3)-equivariant.
- Nonlinearities: The paper works in the steerable domain with gated nonlinearities rather than applying standard point-wise activations directly to steerable coefficients.
C.1 PSEUDOCODE OF SEGNN AND ABLATED ARCHITECTURES
The pseudocode shows how SEGNN layers embed geometric and physical attributes, build non-linear steerable messages, aggregate them, and update steerable node features with gated nonlinearities.
- Ablated architectures: The SEnon-linear layer forms messages with two steerable MLP-style tensor-product stages, while SElinear uses a single linear message transformation.
- SEGNN: SEGNN embeds relative positions and additional geometric or physical attributes through spherical harmonics before message computation.Examples include velocity, acceleration, spin, and force.
- Attribute handling: Multiple steerable attributes may be added or concatenated, but concatenation introduces significantly more weights without improving conducted-experiment performance.
- Gated nonlinearities: Gated nonlinearities activate scalar irreps directly and multiply higher-order irreps by additional activated scalar irreps.
- SEGNN: SEGNN concatenates node features, neighbor features, and squared relative distance before applying gated tensor-product message layers.
- SEGNN: After message aggregation, node features are updated through gated non-linear tensor-product layers conditioned on steerable attributes.
C.2 EXPERIMENTAL DETAILS
The experiments evaluate SEGNNs and steerable ablations across N-body and QM9 settings, varying steerable orders, architectures, training data, and interaction cutoffs. Results include performance, runtime, and message-count comparisons.
- N-body experiments: The N-body experiments compare SEGNNs with steerable linear and non-linear point convolutions using ablations over feature and attribute orders.Table C.1 reports MSE and runtime for different maximum orders of steerable features and attributes.
- N-body experiments: SEGNNs are significantly more data efficient than EGNNs on the N-body system experiment.Table C.2 compares MSE across different numbers of available training samples.
- N-body experiments: The gravitational 100-body task predicts position or force at t = 4 from the state at t = 3 using simulated trajectories.The dataset contains 10,000 training trajectories and 2,000 trajectories for evaluation; an example trajectory appears in Fig. C.1.
- N-body experiments: SEGNNs generalise much better than default MPNNs or EGNNs on the gravitational 100-body problem.Table C.3 measures positional and force MSE together with forward time for batches of 20 samples on an NVIDIA GeForce RTX 3090 GPU.
- QM9 experiments: A 2 Å cutoff sends approximately 6× fewer messages per layer than a 5 Å cutoff on the QM9 training partition, with smaller variance.Figure C.2 plots mean message count against cutoff radius and marks the transition from disconnected to connected graphs.
- QM9 experiments: On QM9, models use atom-type inputs and compare message-passing architectures with steerable convolutional ablations under matched parameter budgets.The QM9 setup uses 256-dimensional atom-type representations and four SEGNN message-passing layers, while the convolutional ablations use seven layers.
D LICENSES
The project uses open-source scientific software and datasets under the licenses listed in the appendix, alongside CUDA’s proprietary license.
- Software licenses: The codebase uses Python, NumPy, PyTorch, PyTorch Geometric, e3nn, and Weights & Biases under the licenses specified for each dependency.The listed licenses include PSF, BSD 3-Clause, BSD, MIT, and CUDA’s proprietary license.
- Dataset licenses: QM9 is used under CC0 1.0, OC20 under Creative Commons Attribution 4.0, and adapted QM9 dataloaders under an Educational and Not-for-Profit Research License.These dataset and loader licenses are stated explicitly in the appendix.