[Paper Review] SE(3) Equivariant Graph Neural Networks with Complete Local Frames

Weitao Du, He Zhang+5 more|arXiv (Cornell University)|Oct 26, 2021

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology22 citations

TL;DR

This paper proposes ClofNet, a computationally efficient SE(3)-equivariant graph neural network that uses complete local frames—constructed via cross products—to enable faithful, reference-free representation of geometric tensors. By projecting higher-order geometric quantities onto orthonormal frames, the model achieves state-of-the-art performance in both many-body dynamics modeling and molecular conformation generation, with strong sample efficiency and low equivariance error.

ABSTRACT

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.

Motivation & Objective

• Address the challenge of balancing expressiveness and computational efficiency in SE(3)-equivariant deep learning for 3D many-body systems.
• Overcome direction degeneration in existing equivariant GNNs that rely on incomplete radial-only frames.
• Enable lossless representation of geometric quantities (e.g., vectors, torsion forces) in a reference-free, equivariant manner.
• Develop a framework that supports arbitrary-order tensors and maintains group equivariance without high-dimensional lifting or complex tensor operations.
• Demonstrate strong generalization and data efficiency in physical modeling and molecular conformation generation tasks.

Proposed method

• Introduce local complete frames constructed from cross products of particle vectors to form orthonormal bases in 3D space, ensuring completeness and avoiding direction degeneration.
• Project geometric tensors (scalars, vectors, higher-order) onto these local frames via a scalarization block, converting them into scalar coefficients for message passing.
• Use a standard graph neural network for message passing on the scalarized coefficients, preserving SE(3) equivariance due to the equivariant nature of the frames.
• Construct frames using only cross product operations, which are computationally efficient compared to methods relying on Clebsch-Gordan decomposition or spherical harmonics.
• Extend the frame construction to handle higher-order tensor inputs by recursively applying cross products and maintaining orthonormality.
• Integrate the framework into a score-based generative modeling setup for molecular conformation generation, using annealed Langevin dynamics for sampling.

Experimental results

Research questions

• RQ1 Can a computationally efficient SE(3)-equivariant GNN be designed that avoids the high cost of high-dimensional lifting or tensor decomposition?
• RQ2 Does the use of complete local frames improve the expressiveness and accuracy of geometric representation compared to radial-only or incomplete frame methods?
• RQ3 To what extent does the proposed method generalize to out-of-distribution and low-data regimes in physical systems?
• RQ4 Can the model maintain strong equivariance while achieving state-of-the-art performance in both dynamics modeling and molecular conformation generation?
• RQ5 How does the frame construction via cross products compare in efficiency and stability to alternative equivariant function space representations?

Key findings

• ClofNet achieves the best performance on the POS dataset, with an interpolation MSE of 0.138 and extrapolation MSE of 2.502, significantly outperforming Radial Field (7.665) and EGNN (6.449).
• The model maintains a low equivariance error ($\Delta_{EQ} = 4.13 \cdot 10^{-6}$), demonstrating strong numerical preservation of SE(3) equivariance.
• On the GEOM-QM9 and GEOM-Drugs datasets, ClofNet generates high-quality molecular conformations with low RMSD, as shown in visualizations of best-aligned structures.
• ClofNet achieves competitive results even with only 0.4% of the training data, highlighting its superior sample efficiency.
• The ablation study confirms that the use of complete local frames is critical for modeling complex geometric quantities like torsion forces.
• The framework’s computational efficiency is validated by the use of only cross product operations for frame construction, avoiding costly tensor product decompositions.