[Paper Review] Equivariant Graph Neural Operator for Modeling 3D Dynamics

Minkai Xu, Jiaqi Han+7 more|arXiv (Cornell University)|Jan 19, 2024

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper proposes Equivariant Graph Neural Operator (EGNO), a novel method that models 3D physical dynamics as continuous trajectories using SE(3)-equivariant neural operators with Fourier-based temporal convolutions. By learning dynamics as a function over time, EGNO achieves superior performance in particle simulations, human motion capture, and molecular dynamics, with up to 52% average improvement over baselines like EGNN.

ABSTRACT

Modeling the complex three-dimensional (3D) dynamics of relational systems is an important problem in the natural sciences, with applications ranging from molecular simulations to particle mechanics. Machine learning methods have achieved good success by learning graph neural networks to model spatial interactions. However, these approaches do not faithfully capture temporal correlations since they only model next-step predictions. In this work, we propose Equivariant Graph Neural Operator (EGNO), a novel and principled method that directly models dynamics as trajectories instead of just next-step prediction. Different from existing methods, EGNO explicitly learns the temporal evolution of 3D dynamics where we formulate the dynamics as a function over time and learn neural operators to approximate it. To capture the temporal correlations while keeping the intrinsic SE(3)-equivariance, we develop equivariant temporal convolutions parameterized in the Fourier space and build EGNO by stacking the Fourier layers over equivariant networks. EGNO is the first operator learning framework that is capable of modeling solution dynamics functions over time while retaining 3D equivariance. Comprehensive experiments in multiple domains, including particle simulations, human motion capture, and molecular dynamics, demonstrate the significantly superior performance of EGNO against existing methods, thanks to the equivariant temporal modeling. Our code is available at https://github.com/MinkaiXu/egno.

Motivation & Objective

• Address the limitation of existing graph neural networks in capturing long-range temporal correlations in 3D physical dynamics.
• Overcome the restriction of next-step prediction by modeling full solution trajectories as continuous functions over time.
• Preserve SE(3)-equivariance—critical for physical generalization—while learning temporal dependencies in 3D spatial dynamics.
• Enable efficient, parallel inference over arbitrary time discretizations without retraining, enhancing deployment flexibility.
• Develop a generalizable framework that integrates with any EGNN architecture for diverse physical dynamics applications.

Proposed method

• Formulate 3D dynamics as a function mapping initial states to full trajectories, enabling continuous-time modeling via neural operators.
• Introduce equivariant temporal convolution layers parameterized in the Fourier domain to preserve SE(3)-equivariance during temporal modeling.
• Leverage the equivariance of Fourier and inverse Fourier transforms by designing kernel integral operators that maintain rotational and translational invariance in frequency space.
• Stack Fourier-based temporal layers with standard EGNN components to build EGNO, enabling end-to-end learning of geometric and temporal dynamics.
• Use time embeddings to enable zero-shot inference at arbitrary temporal resolutions by interpolating between training time steps.
• Integrate geometric features—node features (h), positions (x), and velocities (v)—into the temporal convolution to capture rich dynamic patterns.

Figure 1: Illustration of EGNO . EGNO blocks (green) can be built with any EGNN layers (blue) and the proposed equivariant temporal convolution layers (yellow). Consider discretizing the time window $\Delta T$ into $P$ points $\{\Delta t_{1},\dots,\Delta t_{P}\}$ . Given a current state ${\mathcal{G

Experimental results

Research questions

• RQ1 Can a neural operator framework preserve SE(3)-equivariance while modeling continuous-time 3D dynamics?
• RQ2 How does incorporating temporal convolution in the Fourier domain improve modeling of long-range temporal dependencies compared to autoregressive next-step prediction?
• RQ3 To what extent can EGNO generalize to unseen temporal discretization steps without retraining?
• RQ4 What is the contribution of different geometric features (h, x, v) to the performance of temporal modeling in EGNO?
• RQ5 How does EGNO compare to existing GNN-based models in terms of trajectory accuracy and generalization across diverse physical systems?

Key findings

• EGNO achieves a 36% relative improvement in F-MSE over EGNN on the Aspirin molecular dynamics benchmark.
• On the human motion capture (Mocap-Run) dataset, EGNO achieves a 52% average improvement in prediction accuracy compared to EGNN.
• Abllation studies show that incorporating velocity (v) in temporal convolution significantly improves performance, with the full EGNO variant (using h, x, v) outperforming ablated versions by up to 15% on N-body simulations.
• EGNO generalizes zero-shot to double the temporal resolution (2P) without retraining, producing smooth and accurate trajectories when interpolating between low-resolution time steps.
• The number of Fourier modes (I) has a non-monotonic effect on performance, with optimal performance at I=2 or I=3, beyond which overfitting may occur.
• EGNO enables efficient parallel decoding of future states across any time window in a single forward pass, unlike autoregressive models that require sequential inference.

Figure 2: Ablation studies on the number of modes $I$ on N-body simulation and Mocap- Run datasets.