SEGNO: Generalizing Equivariant Graph Neural Networks with Physical Inductive Biases

TL;DR

SEGNO addresses the generalization gap in equivariant GNNs for physical dynamics by enforcing continuity of trajectories and incorporating second-order motion via a Neural ODE. It leverages an $O(3)$-equivariant GNN to model acceleration and uses a neural ODE solver to obtain a continuous latent trajectory between observed states, with theoretical guarantees of trajectory uniqueness and bounded error under Lipschitz conditions. Empirically, SEGNO outperforms state-of-the-art baselines across synthetic N-body dynamics, MD22 molecular dynamics, and CMU motion capture, demonstrating better long-horizon generalization and rollout stability. By preserving $E(3)$-equivariance and enabling plug-and-play backbones, SEGNO provides a robust, physically informed framework for simulating complex multi-object dynamics at scale.

Abstract

Graph Neural Networks (GNNs) with equivariant properties have emerged as powerful tools for modeling complex dynamics of multi-object physical systems. However, their generalization ability is limited by the inadequate consideration of physical inductive biases: (1) Existing studies overlook the continuity of transitions among system states, opting to employ several discrete transformation layers to learn the direct mapping between two adjacent states; (2) Most models only account for first-order velocity information, despite the fact that many physical systems are governed by second-order motion laws. To incorporate these inductive biases, we propose the Second-order Equivariant Graph Neural Ordinary Differential Equation (SEGNO). Specifically, we show how the second-order continuity can be incorporated into GNNs while maintaining the equivariant property. Furthermore, we offer theoretical insights into SEGNO, highlighting that it can learn a unique trajectory between adjacent states, which is crucial for model generalization. Additionally, we prove that the discrepancy between this learned trajectory of SEGNO and the true trajectory is bounded. Extensive experiments on complex dynamical systems including molecular dynamics and motion capture demonstrate that our model yields a significant improvement over the state-of-the-art baselines.

Figures (9)

• Figure 1: Learned trajectories of models with different inductive bias. All models can map input to output. However, discrete and first-order continuous models fail to learn the true intermediate states due to the lack of considering continuity and second-order laws.
• Figure 2: Illustration of learned trajectories from EGNN (left) and SEGNO (right). They are trained to predict the positions of N-body charged systems after 1000ts (See Section \ref{['sec:nbody']}). The green, red, and dotted grey lines are the true, average, and predicted trajectories. $t_0, t_1$ is the observed states. $t_{0.5}$ is the predicted latent state The blue area denotes the variance.
• Figure 3: Effects of iteration number $\tau$.
• Figure 4: Visualization of Motion Capture with 50 ts. Left to Right: initial position, GMN, SEGNO (all in blue). Ground truths are in red.
• Figure 5: Example trajectories of 5-body charged system. From left to right, the number of positive and negative charges are 1, 3, 0.

Theorems & Definitions (7)

• Proposition 3.1
• Lemma 4.1
• Proposition 4.2
• Theorem 4.3
• Corollary 4.4
• Lemma A.1
• proof