Source-linked AI summary
SE(3) Equivariant Graph Neural Networks with Complete Local Frames
Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Bin Shao, Tie-Yan Liu
TL;DR
SE(3)-equivariant GNNs seek expressive geometric modeling without the computational burden of some high-dimensional equivariant methods. ClofNet uses equivariant local complete frames to encode geometric tensors for graph message passing. It achieves best or competitive reported performance on Newtonian dynamics and molecular conformation tasks.
Problem
Equivariant GNN design must balance expressive geometric modeling against the computational cost of high-dimensional equivariant operations, while radial-only approaches can miss complex geometric quantities.
Source: §1 · p. 2Method
Results
ClofNet achieves best or competitive performance across Newtonian many-body modeling and molecular conformation tasks, including 26.5% higher COV-mean and 26.6% higher COV-median than ConfGF on Drugs.
Source: §1 · p. 2, §7.1 · p. 7, §7.2 · p. 9Takeaways & Limitations
The results support using complete local frames to model geometric quantities efficiently in the evaluated many-body dynamics and molecular conformation settings.
Source: §1 · p. 2, §8 · p. 9The method assumes non-degenerate local frames, with all neighboring frames degenerating when particles are restricted to a straight line.
Source: §4 · p. 3
Abstract
from arXiv · showhide
Group equivariance (e.g. SE(3) equivariance) is a critical physical symmetry in science, from classical and quantum physics to computational biology. It enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which has been shown to improve the generalization performance and data efficiency for downstream tasks. Constructing an equivariant neural network generally brings high computational costs to ensure expressiveness. Therefore, how to better trade-off the expressiveness and computational efficiency plays a core role in the design of the equivariant deep learning models. In this paper, we propose a framework to construct SE(3) equivariant graph neural networks that can approximate the geometric quantities efficiently. Inspired by differential geometry and physics, we introduce equivariant local complete frames to graph neural networks, such that tensor information at given orders can be projected onto the frames. The local frame is constructed to form an orthonormal basis that avoids direction degeneration and ensure completeness. Since the frames are built only by cross product operations, our method is computationally efficient. We evaluate our method on two tasks: Newton mechanics modeling and equilibrium molecule conformation generation. Extensive experimental results demonstrate that our model achieves the best or competitive performance in two types of datasets.
1. Introduction
SE(3)-equivariant GNNs must balance expressive geometric modeling with computational efficiency. ClofNet addresses this trade-off with complete local frames and is evaluated on many-body dynamics and molecular conformation generation.
1. Introduction
- Equivariant GNNs face a trade-off between expressive modeling and computational cost: high-dimensional equivariant methods are costly, while radial-only approaches can miss higher-order geometric quantities.Radial-only methods also face direction degeneration and discard higher-order tensor inputs.Source: §1 · pp. 1–2
- ClofNet achieves best or competitive results across the evaluated scenarios, including competitive many-body performance with 0.4% of training samples.The paper evaluates synthetic many-body dynamics and real-world molecular conformation generation.Source: §1 · p. 2
- ClofNet uses complete orthonormal local frames to scalarize geometric tensors before graph message passing, with cross-product frame construction intended to keep computation efficient.The paper also extends its frame strategy to high-order tensor inputs.Source: §1 · p. 2
2. Preliminaries
The preliminaries establish notation for many-body systems, geometric tensors, symmetry transformations, and graph-based message passing. They frame the paper’s inputs as particle positions and features subject to geometric and permutation symmetries.
2. Preliminaries
- A many-body system is represented by particle positions in R3, with neighborhoods defined for each particle and cross and tensor products used as operations.The notation identifies N as the number of particles and xi(t) as particle i's position at time t.Source: §2 · p. 2
- SE(3) comprises translations and 3D rotations that preserve distances, while an SO(3)-invariant scalar is unchanged under rotation.The preliminaries define equivariance through compatible group actions on a function’s input and output.Source: §2 · p. 2
- Tensor fields represent geometric or physical quantities whose transformations satisfy group equivariance; an equivariant vector field is given as a dynamical system’s velocity.The paper calls these objects tensor or geometric tensors under its stated definition.Source: §2 · p. 2
- Graph message passing is used because many-body modeling is independent of particle labels, and EGNN provides an efficient way to learn equivariant representations.Node features, positions, and edge attributes are the graph inputs described here.Source: §2 · p. 2
3. Motivation from many body interactions.
A two-particle example shows why radial-only force representations are insufficient when external fields produce forces in other directions. The paper motivates a complete orthogonal frame to represent arbitrary force directions.
3. Motivation from many body interactions.
- External-field forces can point in arbitrary 3D directions, so a radial-only prediction is insufficient when the potential gradient is not aligned with the particle displacement.The potential is decomposed into a distance-dependent term and an external contribution that may depend on distances and angles.Source: §3 · p. 3
- The proposed two-body frame takes the relative displacement as its first direction, the cross product of particle positions as its second, and their cross product as its third.These three orthogonal directions allow force vectors to be expressed without direction degeneration, as described in the example.Source: §3 · p. 3
- The example contrasts this frame-based representation with mainstream approaches that express predicted forces only along the radial direction.The passage identifies Schütt et al. (2017) and Satorras et al. (2021b) as examples.Source: §3 · p. 3
4. Methodology
ClofNet centralizes particle positions, scalarizes geometric tensors using equivariant local frames, processes invariant scalars with graph message passing, and vectorizes them when equivariant tensor outputs are needed. The frames and the three-block workflow are designed to preserve the stated symmetries.
ClofNet architecture
- ClofNet centralizes positions, scalarizes tensors into invariant coefficients, applies graph message passing, and vectorizes embeddings when vector or tensor outputs are required.For scalar outputs, the paper says a pooling layer after message passing is sufficient.Source: §4 · pp. 3–4
Local equivariant frame construction
- Each edge frame is built from a normalized relative displacement and cross-product directions; a small ϵ is added to normalization if particle positions collapse.The frame is assumed non-degenerate; all neighboring frames degenerate only when the particles lie on a straight line, which the passage calls a measure-zero set.Source: §4 · p. 3
Scalarization block
- Scalarization projects geometric inputs onto the local frame, yielding SO(3)-invariant coefficients; the described scalarization is invertible and thus retains tensor information.The paper also describes extending frames by tensor products for higher-order tensors.Source: §4 · p. 4
Expressive scalar processing
- The paper presents ClofNet as compatible with flexible graph neural networks because message passing operates on invariant scalars rather than equivariant tensor representations.The approach also describes extending scalarization to high-order geometric tensors.Source: §4 · p. 4
Graph message-passing block
- The graph message-passing block learns invariant edge embeddings from scalarized geometry and node or edge attributes, allowing nonlinear activations and attention on invariant scalars.The block propagates and aggregates information on the graph.Source: §4 · p. 4
Vectorization block
- The vectorization block pairs edge embeddings with their complete frames and aggregates the resulting vectors at nodes to estimate an equivariant vector field.The paper extends the pairing approach to other tensor outputs using the corresponding tensor product of the frame.Source: §4 · pp. 4–5
Implementation choice
- The scalarization operation can be inserted into each message-passing block, but the authors report no significant gains in their two experiments and exclude it to reduce computation.This is an empirical design choice reported for the two experiments in the paper.Source: §4 · p. 5
5. Theoretical properties of ClofNet
ClofNet’s construction guarantees SE(3) equivariance by combining translation handling with equivariant frames and invariant scalarization. Its universality depends on the expressiveness of the graph message-passing block, with a practical efficiency condition for common choices.
Equivariance guarantee
- ClofNet’s frame, scalarization, and vectorization composition is equivariant under SO(3), while translation symmetry is handled by centering the particle system.The frames transform equivariantly and scalarization is invariant under SO(3).Source: §5 · p. 5
Expressiveness
- ClofNet is universal for continuous SE(3)- and permutation-equivariant functions when its graph block approximates permutation-equivariant graph polynomials and connects linearly to vectorization.The scalarization and vectorization are bijective, so the graph block’s expressiveness determines the remaining approximation capacity.Source: §5 · p. 5
- The universality argument requires multiplication of invariant scalars across nodes, whereas common GCN or transformer blocks are more efficient when higher-hop contributions are negligible.The theoretical requirement means scalarization must include features from all nodes, not only direct neighbors.Source: §5 · p. 5
6. Related Work
Prior equivariant methods either lift data into higher-dimensional spaces or operate in the original space. ClofNet follows the latter direction, aiming to preserve complete geometric information without complex vector-level transformations.
Related approaches
- Lifting methods use spaces such as Lie groups or spherical harmonics and can be expressive, but spherical-harmonic calculations or Lie-group integration add computational cost.For a restricted linear equivariant function class, group convolution in message-passing form is a legal implementation.Source: §6 · p. 6
- Original-space approaches are computationally efficient, but methods such as EGNN may lose geometric tensor information by using only invariant distances.EGNN’s messages are nonlinear, yet the cited description feeds only the invariant distance feature into the model.Source: §6 · p. 6
- ClofNet operates in the original space and uses equivariant frames to preserve complete geometric information without complex vector-level transformations.The paper contrasts this design with spherical harmonics and Clebsch–Gordan tensor products.Source: §6 · p. 6
7. Experiments
The experiments test ClofNet on synthetic Newtonian trajectory prediction and molecular conformation and distance-distribution tasks. Across these settings, results assess benchmark performance, data efficiency, and the contribution of model components.
7.1. Newtonian many-body system
- ClofNet significantly outperforms all baselines on every Newtonian dataset and is faster than TFN and SE(3)-Transformer in forward time.The comparison covers electrostatic, gravity-plus-electrostatic, and Lorentz-like-plus-electrostatic force fields, including non-radial force directions.Source: Table 1 · p. 7, §7.1 · p. 7
- On G+ES, ClofNet outperforms EGNN across training-data sizes, beats GNN in the small-data regime, and is comparable to GNN in the large-data regime.The sweep ranges from 200 to 50,000 training samples; ClofNet remains competitive with 200 samples, or 0.4% of 50,000.Source: Figure 3 · p. 7, §7.1 · p. 7
7.2. Molecular Systems
- The molecular experiments estimate equivariant gradient fields from perturbed atomic positions and evaluate both conformation generation and distance-distribution discrepancies.The study uses GEOM-QM9 and GEOM-Drugs for conformation generation and ISO17 for distance modeling.Source: §7.2 · pp. 7–8
- ClofNet achieves the best or second-best performance on almost all conformation-generation metrics and datasets; on Drugs, its COV-mean and COV-median exceed ConfGF by 26.5% and 26.6%.COV thresholds are 0.5 Å for QM9 and 1.25 Å for Drugs.Source: Table 2 · p. 8, §7.2 · p. 9
- ClofNet dramatically outperforms ConfGF on distance-distribution accuracy, reducing MMD by a magnitude on both Single-median and Pair-Median metrics.The task evaluates discrepancy between generated and reference distance distributions.Source: Table 3 · p. 8, §7.2 · pp. 8–9
- Ablations show that scalarization improves Drugs COV-mean and COV-median by 23.0% and 21.9%, respectively; the authors also report a contribution from attention in gradient-field modeling.The vectorization block was not ablated because it guarantees an equivariant vector-field output.Source: Table 4 · p. 8, §7.2 · p. 9
8. Conclusion and Future Work
ClofNet uses complete local frames to use tensor information without high-dimensional equivariant function embeddings, and can work with any GNN without breaking equivariance. The authors report theoretical and empirical support for its effectiveness and plan to extend the strategy to other local symmetry groups.
8. Conclusion and Future Work
- ClofNet aims to use tensors losslessly without high-dimensional equivariant function embeddings.The framework is designed to express geometric information efficiently and flexibly.Source: §8 · p. 9
- ClofNet can cooperate with any graph neural network without breaking equivariance.Source: §8 · p. 9
- The authors report that theoretical analyses and extensive empirical results verify ClofNet’s effectiveness.Source: §8 · p. 9
- The authors plan to extend their strategy to other local symmetry groups.Source: §8 · p. 9
A.1. The special Eulidean group SE(3)
SE(3) is defined as the semidirect product of 3D rotations and translations. Its elements act on 3D points by rotating and then translating them.
A.1. The special Eulidean group SE(3)
- SE(3) is the semidirect product of the 3D rotation group SO(3) and 3D translations.Source: §A.1 · p. 13
- An SE(3) element is a pair (R, a), where R is in SO(3) and a is in R^3.Source: §A.1 · p. 13
- SE(3) acts on a point by applying rotation R followed by translation a: (R, a) · x = Rx + a.Source: §A.1 · p. 13
- The passage parameterizes a rotation in SO(3) by an angle ϕ in [0, 2π].Source: §A.1 · p. 13
A.2. Torsion force field
The section considers a four-body system consisting of one node and three other positions.
A.2. Torsion force field
- The system is written X = (x_l, x_i, x_j, x_k), with four positions.Source: §A.2 · p. 13
A.2.1. DIHEDRAL ANGLE
For a node and three neighbors, the dihedral angle is defined using the normals of two planes. The angle depends on all four position vectors and cannot be determined from the listed distances alone.
A.2.1. DIHEDRAL ANGLE
- The dihedral angle θ is given by the inner product of the normals to the planes spanned by (x_l − x_i, x_l − x_j) and (x_l − x_i, x_l − x_k), normalized by their magnitudes.The normals are formed by cross products, as expressed in the cosine formula.Source: §A.2.1 · p. 13
- The dihedral angle depends on x_l, x_i, x_j, and x_k, not solely on the three listed pairwise distances.Source: §A.2.1 · p. 13
A.2.2. TORSION ANGLE AND TORSION FORCE
This appendix section details torsion-force construction and ClofNet’s message-passing and score-based conformation-generation workflows, alongside analyses of computational cost, equivariance, information preservation, and universality. It also describes experimental setups and evidence on partially observed dynamics and molecular conformation data.
A.2.2. TORSION ANGLE AND TORSION FORCE
- The torsion angle is the dihedral angle between two planes, and the torsion force is its energy gradient after high-frequency terms are cut off at n ≤2.The planes share the vector xl − xi, with the other defining vectors given by the neighboring particle positions.Source: §A.2.2 · p. 13
A.3.1. GRAPH MESSAGE-PASSING BLOCK
- The graph message-passing block embeds geometric scalars and node and edge features, while its transformer variant refines node embeddings with attention and updates edgewise messages residually.The naive block is described as an EGNN equipped with ClofNet’s equivariant local frames.Source: §A.3.1 · p. 14
A.3.2. LEARNING GRADIENT FIELDS AS SCORE-BASED MODELS
- ClofNet models the Hamiltonian gradient, or score function, as an SO(3)-equivariant vector field and connects it to an evolving block that numerically integrates Langevin diffusion.The reverse process uses Langevin dynamics, and the molecular experiment uses the same annealed Langevin sampling algorithm as Shi et al. (2021).Source: §A.3.2 · p. 15
A.3.3. SCALABILITY
- ClofNet’s scalarization adds 9.6% real-time computational cost and 17.4% memory cost for molecular conformation generation, with O(3 × 3 × E) extra scalar-and-frame computation.The passage characterizes this computation as the same order as one-hop message passing and much less than neural-network backpropagation cost.Source: §A.3.3 · p. 15
A.4.1. PERMUTATION AND TRANSLATION EQUIVARIANCE
- Recentering positions at the initial centroid handles translation while adding that centroid back restores translation-equivariant predicted positions; shared message functions and vector aggregation provide permutation equivariance.The appendix notes that graph-based algorithms automatically guarantee permutation equivariance.Source: §A.4.1 · p. 17
A.4.2. EQUIVARIANCE OF CLOFNET
- The appendix establishes SO(3) equivariance of the complete frame and ClofNet, combines this with translation equivariance for SE(3) equivariance, and states that ClofNet is reflection anti-symmetric.The frame’s equivariance is Proposition A.1; the model’s reflection behavior follows from how its frame vectors transform under central reflection.Source: §A.4 · p. 16, §A.4.2 · p. 18
A.5. Information loss
- ClofNet’s input transformation is invertible and preserves all information about the prediction target, whereas EGNN’s distance-only input loses information needed to determine a dihedral-angle-dependent force.The appendix states I(Y, X) = I(Y, X̂_ClofNet), while the EGNN distance representation does not determine the dihedral angle.Source: §A.5 · p. 19
A.6. Universal approximation
- The stated universal-approximation conditions require feature functions to represent translation-invariant, permutation-equivariant polynomials and pooling functions to contain all linear SO(3)-equivariant mappings to the output tensor space.The appendix frames these as two conditions for a function class to have the universal approximation property.Source: §A.6 · p. 19
- The tensorized graph neural network equipped with equivariant bases is universal for continuous equivariant functions.Equivariant-frame scalarization turns tensor products into multiplication of invariant scalars; the appendix also constructs equivariant linear functionals using frame coefficients and an unrestricted linear map. It cautions that using only the radial direction causes direction degeneration.Source: §A.6 · pp. 20, 21, 23
A.7.1. NEWTONIAN MANY-BODY SYSTEM
- Newtonian-system trajectories span four systems, with 5,000 timesteps per trajectory sliced from 3,000 to 4,000; an additional G+ES training set has 50,000 trajectories.The additional G+ES experiment retains the original validation and test datasets.Source: §A.7.1 · p. 23
A.7.2. PARTIALLY OBSERVED N-BODY EXPERIMENT
- POS hides two of six gravitating particles, making the observed four-particle subsystem non-Markov; evaluation tests interpolation, extrapolation, and equivariance error.The trajectory prediction uses initial positions and velocities, with mean squared error for trajectory prediction and a transformation-based equivariance error.Source: §A.7.2 · p. 24
- ClofNet outperforms all default equivariant baselines on POS interpolation and extrapolation, while all three models show small, nonzero equivariance error.The comparison concerns the partially observed N-body (POS) dataset.Source: §A.7.2 · p. 24, Table 5 · p. 25
A.7.3. MOLECULAR CONFORMATION GENERATION
- The molecular datasets contain 200,000 training conformations, with 22,408 GEOM-QM9 and 14,324 GEOM-Drugs test conformations.For each dataset, 40,000 molecules contribute five likely conformations each; the distances-over-distributions task uses ISO17.Source: §A.7.3 · p. 25
- Following Shi et al. (2021), ClofNet estimates atomic-position gradients with a score-based model and uses annealed Langevin dynamics to generate conformations (Song & Ermon, 2019; Song et al., 2020).Source: §A.7.3 · p. 25
- The implementation uses noise scales from 0.01 to 10 split into 50 levels, adds node/edge types and relative distances as attributes, and follows Shi et al. (2021) for other forward and reverse-process hyperparameters.The reported Table 2 results are copied from Shi et al. (2021) under the same benchmark and data-split setting.Source: §A.7.3 · p. 25
- Figure 4 visualizes best-aligned Drugs conformations and illustrates ClofNet’s capacity to generate high-quality molecular conformations.The figure shows best-aligned samples for randomly selected molecules from GEOM-Drugs.Source: §A.7.3 · p. 25, Figure 4 · p. 26