SEGNO: Generalizing Equivariant Graph Neural Networks with Physical
  Inductive Biases
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.
Yang Liu, Jiashun Cheng, Haihong Zhao, Tingyang Xu, Peilin Zhao, Fugee Tsung, Jia Li, Yu Rong
# Physics-Inspired Neural Graph ODE for

###### Abstract

Simulating and modeling the long-term dynamics of multi-object physical systems is an essential and challenging task. Current studies model the physical systems utilizing Graph Neural Networks (GNNs) with equivariant properties. Specifically, they model the dynamics as a sequence of discrete states with a fixed time interval and learn a direct mapping for all the two adjacent states. However, this direct mapping overlooks the continuous nature between the two states. Namely, we have verified that there are countless possible trajectories between two discrete dynamic states in current GNN-based direct mapping models. This issue greatly hinders the model generalization ability, leading to poor performance of the long-term simulation. In this paper, to better model the latent trajectory through discrete supervision signals, we propose a **P**hysics-**I**nspired **N**eural **G**raph **O**DE (PINGO) algorithm. In PINGO, to ensure the uniqueness of the trajectory, we construct a Physics-Inspired Neural ODE framework to update the latent trajectory. Meanwhile, to effectively capture intricate interactions among objects, we use a GNN-based model to parameterize Neural ODE in a plug-and-play manner. Furthermore, we prove that the discrepancy between the learned trajectory of PIGNO and the true trajectory can be theoretically bounded. Extensive experiments verify our theoretical findings and demonstrate that our model yields an order-of-magnitude improvement over the state-of-the-art baselines, especially on long-term predictions and roll-out errors.

## 1 Introduction

It is a vital problem to simulate and model the complex dynamics of multi-object physical systems, i.e. N-body systems. This problem is relevant to numerous fundamental scientific domains, including molecular dynamics [19], protein folding [11], drug and catalyst virtual screening [32], robot motion planning/control [30], and cosmological simulation [34].

Because of the complex interaction of multiple objects in the N-body system, recent studies propose to use Graph Neural Networks (GNNs) [10; 28] to model the N-body systems. Specifically, they model the objects in the physical system as nodes, the physical relations as edges, and use a message-passing network to learn interactions among nodes. Recently, some works try to encode the physical symmetry into GNNs to ensure the translating/rotating/reflecting equivariant between the geometric input and output of GNN. These models are known as equivariant GNNs [3; 14; 15; 28] and have emerged as a type of leading approaches for N-body system modeling. Most of the existing studies
on equivariant GNNs treat the dynamic process as a sequence of discrete states, encompassing the positions, velocities, and forces of each object. Then given an input state, a direct mapping is learned by constraining the model output to approximate the next adjacent state after the input one, over a fixed time interval. We refer to this type of method as direct-mapping models.

However, we question the validity of the direct-mapping models on two fronts. Firstly, models derived from this paradigm exhibit limited generalizability over time. Specifically, it is possible to apply models trained on short-term data for predictions on longer time intervals through rollout, but not vice versa. Secondly, a direct-mapping model can not capture the continuous physical knowledge between discrete states. To verify this, we train multiple EGNN models with different random seeds on a 3-body system with two adjacent stats \(t_{0},t_{1}\). We extract the object positions from their hidden layers, as the estimation of intermediate motion trajectories between two status \(t_{0.5}\), and perform one-step rollouts to make the prediction of next state \(t_{2}\). We visualize all predicted trajectories (in dotted grey), the mean trajectory (in blue), and the mean and variance of predicted stats in the left part of Figure 1 (a). As shown in Figure 1 (a), the predictions of EGNN have a high variance at both intermediate and rollout states, indicating that the direct-mapping models are unable to learn a uniquely determined motion trajectory. This is a crucial reason why existing models struggle with long-term physical simulations.

In this paper, to better model the latent motion trajectory under discrete supervision signals, we propose **P**hysics-**I**nspired **N**eural **G**raph **O**rinary Differential Equation, dubbed PINGO. Unlike the direct-mapping models which use GNNs to fit all kinematic states, PINGO introduces a Neural ODE framework built upon motion equations to update the position and velocity of the physical system. Theoretically, we prove the uniqueness of the learned latent trajectory of this framework and further provide an upper bound on the discrepancy between the learned and the actual latent trajectory. Figure 1 (a) also depicts an example of prediction results of the latent trajectory of PINGO, following the same training protocol of EGNN. PINGO achieves a significant small variance on both the intermediate and rollout state, and it coincides with the true latent trajectory. Meanwhile, we employ a GNN model to parameterize the force relationships in the Neural ODE framework. Owing to the GNNs' expressive power, PINGO can adeptly address the force prediction issues for each object in the system. Figure 1 (b) shows the key difference between PINGO and direct-mapping methods. Moreover, we provide theoretical proof that PINGO exhibits equivariance properties identical to those of the input GNN model. This property offers the flexibility to adapt various GNN backbones in PINGO to suit different downstream tasks. We conduct extensive experiments on both synthetic and real world physical systems. The results shows that PINGO have an order-of-magnitude improvement over the state-of-the-art baselines on both direction prediction and rollout setting, especially for the extreme long-term simulation.

Figure 1: **(a)** Comparison of the trajectory prediction of EGNN and PINGO2. The green line is the true trajectory extracted from a 3-body system. The dotted grey line is the predicted trajectory from the same model with different random seeds. The red line is the average trajectory aggregated by the predicted trajectories. The green circle (\(t_{0},t_{1}\)) is the ground truth state used for training. The red circle (\(t_{0.5}\)) is the predicted latent state from the average trajectory. The dark blue circle (\(t_{2}\)) is the predicted rollout state. The blue area is the variance of the predicted latent state and rollout state. **(b)** The overview of PINGO. \(\mathbf{q},\dot{\mathbf{q}}\) and \(\ddot{\mathbf{q}}\) represent the position, velocity, and acceleration (force), respectively.
Related Works

GNNs for Modeling the Dynamics of N-body SystemsInteraction Network (IN)[1] is a pioneer work that models objects and relations as graphs and applies multi-step rollouts to make predictions. After that, several studies [21; 24; 27] leverage GNNs and their variants to explicitly reason about the dynamics behind object interactions. Recently, researchers have introduced physical symmetry into models of interacting objects of physical system. For example, TFN [31] and SE(3) Transformer [9] employ spherical harmonics to construct models with 3D rotation equivariance in the Euclidean group for higher-order geometric representations. LieConv [7] and LieTransformer [18] leverage the Lie convolution to extend equivariance on Lie groups. In addition to these methods, which rely on irreducible representations of certain symmetry groups, recent studies [15; 28; 29] apply scalarization techniques to introduce equivariance into the message-passing process in GNNs. Furthermore, SEGNN [3] generalizes EGNN [28] and extends invariant scalars on nodes/edges to covariant vectors and tensors. Finally, EGHN [14] extends the framework of GMN [15] to design equivalent pooling and up-pooling modules for hierarchical modeling of large-scale dynamic systems, such as proteins. Nevertheless, these methods model dynamics in physical systems solely by learning direct mappings between discrete states and conducting long-term simulations via the rollout method.

Physics-Inspired Neural NetworksInfusing physical knowledge has been shown to improve learning neural networks for dynamical system modeling. Broadly speaking, energy conservation, symplectic nature, and ODE are three common physical biases. Our work focuses on ODE bias which models derivatives of the state rather than the states directly. Representation models include Lagrangian Neural Networks (LNN) [8; 23], Hamiltonian neural networks (HNN) [12], and Neural ODE [4; 13]. Recent studies [2; 25; 26] have Integrated Neural ODE and GNNs to learn the motion of interactive particles. Graph Neural Ordinary Differential Equation [25] uses graph convolutional networks to parameterize the first- and second-order derivatives of system states. A concurrent work [2] directly uses actual acceleration to train GNNs. However, how they approximate the system trajectory is still an open problem. Motivated by the physical laws, we theoretically justify how our approach can improve the generalization ability in the time dimension and quantify the error brought by the discrete ODE solver. Another research line [16; 17; 33; 35; 36] employs GNN and the first-order Neural ODE to produce smooth trajectories of multi-agent systems. They validate these methods on complex irregularly sampled and partially observed systems in COVID-19 and social scenarios. In contrast, our work lies in the strong approximation ability of the second-order Neural ODE and equivariant GNNs to learn the dynamics of physical systems.

## 3 Preliminary

N-body SystemWe study N-body systems [15; 21] with a set of objects \(\mathcal{P}=\{P_{i}\}_{i=1}^{N}\). At time \(t\), the state of the system is represented by its geometric feature \((\mathbf{q}^{(t)},\mathbf{\dot{q}}^{(t)})\), where \(\mathbf{q}^{(t)}\in\mathbb{R}^{N\times 3}\) and \(\mathbf{\dot{q}}^{(t)}\in\mathbb{R}^{N\times 3}\) are the position and the velocity vector, respectively, for the object \(P_{i}\). Additionally, objects \(P\) are associated with non-geometric attributes such as mass or charge, which is denoted by \(\mathbf{h}\in\mathbb{R}^{N\times d}\). We use a graph \(\mathcal{G}=\{\mathcal{P},\mathcal{E}\}\) to represent the spatial connections in the system where \(\mathcal{E}\) is an edge set that is constructed via geometric distance cutoff or physical connectivity. The attributes of edge \(e_{ij}\in\mathcal{E}\) (e.g., object distances) are denoted by \(a_{ij}\). The system state at time \(t\) is abbreviated as \(\mathcal{S}^{(t)}=(\mathcal{G},\mathbf{q}^{(t)},\mathbf{\dot{q}}^{(t)},\mathbf{h})\). Since the predictions at different times share the same model, without ambiguity, we will omit the temporal superscript \(t\) for all variables for brevity when necessary.

Physical BackgroundAccording to Newton's second law, _"When a body is acted upon by force, the time rate of change of its momentum equals the force"_. For physical systems that follow Newtonian mechanics, their motion can be described by the following general ordinary differential equation

\[\mathbf{\ddot{q}}^{(t)}=\frac{d^{2}\mathbf{q}^{(t)}}{dt^{2}}=f(t,\mathbf{q}^{(t)},\mathbf{\dot {q}}^{(t)},\mathbf{h}). \tag{1}\]

In this work, we focus on dynamical systems that can be formulated as:

\[\mathbf{\ddot{q}}^{(t)}=f(\mathbf{q}^{(t)},\mathbf{h}), \tag{2}\]

where \(\mathbf{\ddot{q}}^{(t)}\) is the acceleration at time \(t\). If we have a closed-form solution of \(f\), given the initial position \(\mathbf{q}^{(t_{0})}\) and velocity \(\mathbf{\dot{q}}^{(t_{0})}\) at time \(t_{0}\), the position at time \(t_{1}>t_{0}\) can be obtained by integrating
the above differential equation:

\[\mathbf{q}^{(t_{1})}=\mathbf{q}^{(t_{0})}+\int_{t_{0}}^{t_{1}}(\dot{\mathbf{q}}^{(t_{0})}+ \int_{t_{0}}^{t}\ddot{\mathbf{q}}^{(m)}\,dm)\,dt=\mathbf{q}^{(t_{0})}+\int_{t_{0}}^{t_{1} }g(\mathbf{q}^{(t)},\mathbf{h})\,dt. \tag{3}\]

Under the aforementioned assumptions, we have

\[\dot{\mathbf{q}}^{(t)}=g(\mathbf{q}^{(t)},\mathbf{h}). \tag{4}\]

But such a solution is not generally available, especially when the system is complex. Thus, current works are interested in learning a graph neural network to directly approximate \(\mathbf{q}^{(t_{1})}\) on observed system trajectories.

GNN-based Direct-mapping ModelExisting studies seek to use a GNN to directly approximate the above unavailable integration results with training pairs \(\big{(}\mathbf{q}^{(t_{0})},\mathbf{q}^{(t_{1})}\big{)}\). Namely, given the system state at time \(t_{0}\), modern GNN simulators \(\psi_{\theta}\) with parameters \(\theta\) predicts \(\mathbf{q}^{(t_{1})}\) and update node features via message passing. Specifically, each layer of \(\psi_{\theta}\) computes

\[\mathbf{q}^{\prime}_{i},\dot{\mathbf{q}}^{\prime}_{i},\mathbf{h}^{\prime}_{i}=\varphi(\bm {q}_{i},\dot{\mathbf{q}}_{i},\mathbf{h}_{i},\sum_{j\in\mathcal{N}_{i}}\mathbf{m_{ij}}), \qquad\mathbf{m_{ij}}=\phi(\mathbf{q}_{i},\mathbf{q}_{j},\dot{\mathbf{q}}_{i},\dot{\mathbf{q}}_{j},\mathbf{h}_{i},\mathbf{h}_{j},a_{ij}), \tag{5}\]

where \(\phi\) and \(\varphi\) are the edge message function and node update function, respectively, which typically are MLPs. \(\mathbf{q}_{i},\dot{\mathbf{q}}_{i}\in\mathbb{R}^{3}\) and \(\mathbf{h}_{i}\in\mathbb{R}^{d}\) are the features of particle \(P_{i}\). \(\mathbf{m_{ij}}\) defines the message between node \(i\) and \(j\). \(\mathcal{N}_{i}\) collects the neighbors of node \(i\). The prediction is obtained by applying several iterations of message passing. Although the direct method precisely predicts the state of object dynamics with the interval \(T=t_{1}-t_{0}\), its prediction with other intervals is notably inconsistent and suboptimal.

Problem DefinitionThis study concentrates on the temporal generalization capacity of the model. Adopting established approaches, the model is trained to predict the subsequent position \(\mathbf{q}^{(t_{1})}\) as accurately as feasible, considering the system state at time \(t_{0}\). Additionally, the model is capable of (1) accurately predicting unobserved intermediate time points by further uniformly partitioning the time interval \(T=t_{1}-t_{0}\) into increments of timestep \(\Delta t\) (i.e., \(t_{0}<t_{0}+\Delta t<t_{0}+2\Delta t<\cdots<t_{1}\)); (2) simulating rollout trajectories \(\mathbf{q}^{(t_{0}:t_{k})}=(\mathbf{q}^{(t)},\ldots,\mathbf{q}^{(t_{k})})\) where \(t_{k}=t_{0}+k(t_{1}-t_{0}),k\in\mathbb{N}^{+}\).

## 4 Physics-Inspired Neural Graph ODE

In this section, we introduce how the proposed Physics-Inspired Neural Graph ODE (PINGO) works. For notation, we employ \(\mathbf{q}^{(t)}_{\theta},\dot{\mathbf{q}}^{(t)}_{\theta},\ddot{\mathbf{q}}^{(t)}_{\theta} \in\mathbb{R}^{N\times 3}\) to denote GNN approximations of the entire system, whereas \(\mathbf{q}^{(t)},\dot{\mathbf{q}}^{(t)},\ddot{\mathbf{q}}^{(t)}\in\mathbb{R}^{N\times 3}\) represent the actual trajectories in classical mechanics. Note that, trajectories based on classical mechanics are continuous and typically governed by second-order derivatives, describing how objects change their position and velocity over time.

### Physics-Inspired ODE Framework

Leveraging the insights from physics, we propose the parameterization of \(f\) in Eq. 2 using GNNs:

\[\ddot{\mathbf{q}}^{(t)}_{\theta}=f_{\theta}(\mathbf{q}^{(t)}_{\theta},\mathbf{h}), \tag{6}\]

where \(f_{\theta}\) represents a GNN model. Given the initial position \(\mathbf{q}^{(t_{0})}_{\theta}=\mathbf{q}^{(t_{0})}\) and velocity \(\dot{\mathbf{q}}^{(t_{0})}_{\theta}=\dot{\mathbf{q}}^{(t_{0})}\), we should be able to compute the position \(\mathbf{q}^{(t_{1})}_{\theta}\) at time \(t_{1}\) using Eq. 3 with a perfect learned \(f_{\theta}\). If we take \(t_{0}\) and \(t_{1}\) as the input and target timesteps respectively, and in line with prior studies [3, 28], \(f_{\theta}\) is trained to minimize the discrepancy between the exact and approximated positions:

\[\mathcal{L}_{\text{train}}=\sum_{s\in\mathcal{D}_{\text{train}}}||\mathbf{q}^{(t_{ 1})}_{\theta,s}-\mathbf{q}^{(t_{1})}_{s}||^{2}, \tag{7}\]

where \(\mathcal{D}_{\text{train}}\) denotes the training set and \(\mathbf{q}^{(t_{1})}_{\theta,s},\mathbf{q}^{(t_{1})}_{s}\) denote the GNN prediction and actual trajectory \(s\) respectively. We can now examine the model's generalizability to other time steps (e.g., \(t_{0}+3\Delta t\)) given sufficient training. In an ideal scenario, where \(\mathbf{q}^{(t_{1})}_{\theta}\) is computationally feasible and the loss \(L\) is minimized to zero, the subsequent proposition holds:
**Proposition 4.1**.: _Given that \(f_{\theta}\) is continuous on \(t\in[t_{0},t_{1}]\) and in the absence of an external field, if \(\mathbf{q}_{\theta}^{(t_{0})}=\mathbf{q}^{(t_{0})}\) and \(\dot{\mathbf{q}}_{\theta}^{(t_{0})}=\dot{\mathbf{q}}^{(t_{0})}\), and if \(\mathbf{q}_{\theta}^{(t_{1})}=\mathbf{q}^{(t_{1})}\), then it follows that \(f_{\theta}(\mathbf{q}^{(t)},\mathbf{h})=f(\mathbf{q}^{(t)},\mathbf{h}),\forall t\)._

The detailed proof is given in Appendix. While the proof draws upon ODE theory, the underlying intuition is straightforward. Per the Picard-Lindelof theorem, a unique trajectory exists that passes \(\mathbf{q}^{(t_{1})}\) with initial \(\mathbf{q}^{(t_{0})},\dot{\mathbf{q}}^{(t_{0})}\). For \(\mathbf{q}_{\theta}^{(t_{1})}\) to match \(\mathbf{q}^{(t_{1})}\), \(f_{\theta}(\mathbf{q}^{(t)},\mathbf{h})\) must equate to \(f(\mathbf{q}^{(t)},\mathbf{h})\). If our model \(f_{\theta}\) accurately approximates \(\ddot{\mathbf{q}}^{(t)}\), the system trajectory is recovered. This proposition shows that the proposed framework is able to train across diverse timesteps and generalize to others, showcasing its empirical value and utility.

Discretization of Integration in the Framework.Despite Proposition 4.1 is attractive, the integral of a graph neural network is intractable. A common solution is adopting an ODE solver to produce a discrete trajectory to approximate the continuous one. Thus it is necessary to study whether the model is still consistent with the design after applying ODE solvers and how error varies according to the chosen timestep. In this work, we focus on a symplectic Euler integrator which is computationally efficient and effective in modeling dynamical systems. Specifically, it computes

\[\dot{\mathbf{q}}_{\theta}^{(t_{1})}=\dot{\mathbf{q}}_{\theta}^{(t_{0})}+\sum_{k=0}^{ \tau-1}\ddot{\mathbf{q}}_{\theta}^{(t_{0}+k\Delta t)}\Delta t,\qquad\mathbf{q}_{\theta }^{(t_{1})}=\mathbf{q}_{\theta}^{(t_{0})}+\sum_{k=1}^{\tau}\dot{\mathbf{q}}_{\theta}^{ (t_{0}+k\Delta t)}\Delta t, \tag{8}\]

where \(t_{0}+\tau\Delta t=t_{1}\). Figure 2 illustrates the entire computation process. During each iteration, PINGO first computes the acceleration based on the GNN and then updates the velocity and position in order. Based on our update strategy, an estimation of \(g\) is defined as

\[\dot{\mathbf{q}}_{\theta}^{(t)}=g_{\theta}(\mathbf{q}_{\theta}^{(t)},\mathbf{h}). \tag{9}\]

The error analysis is provided in the next section.

### A Bounded Approach to Real Trajectories

Since \(\mathbf{q}_{\theta}^{(t)}\) is computed through discrete numerical integration, optimizing Eq. 7 will not yield a neural network that can perfectly reconstruct the true \(\mathbf{q}^{(t)}\), even if the loss is zero. Existing theoretical findings [37] illustrate that training using discrete integration trajectories results in a bounded approximation of their first-order derivative, where the bound is linked to the chosen timestep. Given \(\mathcal{B}(\mathbf{q},r)\subset\mathbb{C}^{3}\) as the complex ball of radius \(r>0\) centered at \(\mathbf{q}\in\mathbb{C}^{3}\), and time interval \(T=t_{1}-t_{0}\), the following holds:

**Theorem 4.2**.: _Let \(\mathbf{q}_{1},\mathbf{q}_{2}\in\mathbf{R}^{3}\) and \(r_{1},r_{2}>0\), the ODE solver is Euler solver and it iterates \(\tau\) times with timestep \(\Delta t\). Suppose that the target \(f\) and the learned \(f_{\theta}\) are analytic and bounded by \(m_{2}\) on \(\mathcal{B}(\mathbf{q}_{2},r_{2})\), and the target \(g\) and the learned \(g_{\theta}\) are analytic and bounded by \(m_{1}\) on \(\mathcal{B}(\mathbf{q}_{1},r_{1})\). Then, there exist constants \(T_{0}\) such that, if \(0<T<T_{0}\),_

\[||f_{\theta}(\mathbf{q},\mathbf{h})-f(\mathbf{q},\mathbf{h})||_{\infty}\leq O(\Delta t+\frac{ \mathcal{L}}{\Delta t}),\quad\forall\mathbf{q}\in\mathbf{R}^{3},\]

_where \(\mathcal{L}=\frac{1}{T}||\mathbf{q}_{\theta}^{(t_{1})}-\mathbf{q}^{(t_{1})}||_{\infty}\) represents prediction error in \(l_{\infty}\) - norm._

Figure 2: PINGO schematic. The forward step (left box) first uses GNN \(f_{\theta}\) (We omit \(\theta\) subscript in the figure for brevity) to decide second-order derivatives of particles and then updates their velocities and positions in order. The right part is the trajectory obtained by iteratively repeating the forward process in a 3-body system.
The proof is reported in the Appendix. Based on Theorem 4.2, we analyze the error introduced by the Euler solver. We use two metrics common in classical numerical analysis, namely, local and global truncation error. The local truncation error \(\epsilon_{t+\Delta t}\) of PINGO is defined as follows:

\[\epsilon_{t+\Delta t}=||\mathbf{q}^{(t+\Delta t)}-\mathbf{q}^{(t)}-\hat{\mathbf{q}}^{(t)} \Delta t-f_{\theta}(\mathbf{q}^{(t)},\mathbf{h})\Delta t^{2}||_{2}, \tag{10}\]

which represents the error accumulated in a single step. And the global truncation error \(\mathcal{E}_{t+k\Delta t}\) is defined as follows:

\[\mathcal{E}_{t+k\Delta t}=||\mathbf{q}^{(t+k\Delta t)}-\mathbf{q}_{\theta}^{(t+k \Delta t)}||_{2}, \tag{11}\]

which denotes the error accumulated in the first k steps. Then we have

**Corollary 4.3**.: _Given the same conditions as in Theorem 4.2, if the training loss \(L\) is adequately minimized and \(g_{\theta}\) satisfies the Lipschitz condition, the local truncation error \(\epsilon_{t+\Delta t}\) and the global truncation error \(\mathcal{E}_{t+k\Delta t}\) are \(O(\Delta t^{2})\) and \(O(\Delta t)\), respectively._

The proof is reported in Appendix. Corollary 4.3 shows that for the prediction of \(\tau\)-th iteration, its error depends on the chosen \(\Delta t\). These statements imply that PINGO can be trained by minimizing Eq. 7 and generalize to other timesteps.

### Equivariance of PINGO

Given particle positions at time \(t\), computing the resultant force of each body is the key to modeling its dynamic changes. In general, PINGO is flexible and can incorporate various GNNs as the backbone, among which equivariant graph neural networks have shown effectiveness in preserving the symmetry of physical systems. For example, if we rotate the input system, the model output will be rotated to the same degree as well. There exist multiple symmetry groups for 3-dimensional systems such as SO(3) (i.e., rotational equivariance) and SE(3) (i.e., rotational and translational equivariance). We can prove that PINGO would not break the equivariance property of backbone GNNs.

**Proposition 4.4**.: _Suppose the backbone GNN \(f_{\theta}\) of PINGO is equivariant to group \(\mathcal{G}\), then the trajectory \(\mathbf{q}_{\theta}\) is equivariant to group \(\mathcal{G}\)._

The proof and detailed discussion of equivariance is provided in Appendix. Intuitively, the integration step Eq. 8 only involves the linear computation of equivariant terms, thus PINGO would preserve the same equivariant property as the used GNNs. Specifically, we choose the widely used EGNN [28] as the example of backbone, assuming the acceleration depends on positions, the message passing is defined by

\[\mathbf{m_{ij}}=\phi_{e}\left(||\mathbf{q}_{\theta,i}-\mathbf{q}_{\theta,j}||^{2},\mathbf{h}_ {i},\mathbf{h}_{j},a_{ij}\right),\quad\mathbf{\tilde{q}}_{\theta,i}=\frac{1}{N-1} \sum_{j\in\mathcal{N}_{i}}(\mathbf{q}_{\theta,i}-\mathbf{q}_{\theta,j})\phi_{q}(\mathbf{ m}_{ij}). \tag{12}\]

Here \(\phi_{e},\phi_{x}\) denotes Multi-Layer Perceptrons (MLP) whose output is a scalar and the output of \(\phi_{h}\) is a vector. The non-geometric features are updated via skip connections. Analogous to neural ODE methods, the model parameters are shared among all iterations.

## 5 Experiments

### N-body system

We build upon the experimental setting introduced in [28] where the task is to estimate all particle positions after a fixed timestep. We consider two types of N-body systems, charged and gravity particles, which are driven by electromagnetic [21] and gravitational forces [3] between every pair of particles respectively. Each system consists of 5 particles that have initial position, velocity, and attributes like positive/negative charge or mass. We sample 3000 trajectories for training, 2000 for validation, and 2000 for testing.

Implementation detailsWe compare our method with GNN and equivariant methods: Radial Field [22], TFN [31], SE(3) Transformer [9], EGNN[28], GMN [15], and SEGNN [3]. In addition, we also compare to Graph Neural Ordinary Differential Equation (GDE) [25]. Without specification,
we use EGNN as the backbone of PINGO. The layer number is searched within 1000 time step (1000 ts), we add two additional settings (1500 ts and 2000 ts) to evaluate the performance of long-term prediction. Meanwhile, we report the average forward time in seconds for 100 samples of each method. From Table 4 we can observe that:

* It is evident that PINGO, equipped solely with EGNN, outperforms all baselines across all datasets and settings. Notably, compared to the best baseline SEGNN, the average error improvement on Charged and Gravity datasets is \(0.254\) and \(0.504\) respectively, indicating significant improvement.
* As the time step increases, PINGO's performance improvement becomes more pronounced. Compared to the best baseline SEGNN, the average error improvement increases from \(0.074\) at 1000 time steps to \(0.660\) at 2000 time steps, demonstrating the efficacy of PINGO in handling long-term prediction scenarios.
* As expected, PINGO's forward time (\(0.0277\)s) is slower than that of EGNN (\(0.0126\)s) due to its additional integration operation. Nevertheless, PINGO's forward time remains competitive compared to the best baselines (\(0.0315\)), indicating its efficiency.

Generalization capability from long-term to short-term.It is interesting to see how PINGO can generalize from long-term training to short-term testing. Accordingly, we train models on 1000ts on two datasets and make the test on shorter time steps by performing PINGO on the smaller \(\tau\) step with the same ratio. For the baselines, we extract the object position information from their hidden layers as the prediction of its intermediate steps. Table 2 reports the mean and standard deviation of each setting. From Table 2 we can observe that:

* Clearly, PINGO outperforms all other baselines across all settings by a large margin. Notably, when there is a lack of supervised signals at 250/500/750ts, the performance of all other baselines

\begin{table}
\begin{tabular}{c c c c c c c|c} \hline \hline \multirow{2}{*}{**Method**} & \multicolumn{3}{c}{**Charged**} & \multicolumn{3}{c|}{**Gravity**} & \multirow{2}{*}{**Time [s]**} \\  & 1000 ts & 1500 ts & 2000 ts & 1000 ts & 1500 ts & 2000 ts \\ \hline
**Linear** & 6.830\(\pm\)0.016 & 20.012\(\pm\)0.029 & 39.513\(\pm\)0.061 & 7.928\(\pm\)0.001 & 29.270\(\pm\)0.003 & 58.521\(\pm\)0.003 & 0.0002 \\
**GNN** & 1.077\(\pm\)0.004 & 5.059\(\pm\)0.250 & 10.591\(\pm\)0.352 & 1.400\(\pm\)0.071 & 4.691\(\pm\)0.238 & 10.508\(\pm\)0.432 & 0.0064 \\
**GDE** & 1.285\(\pm\)0.004 & 4.026\(\pm\)0.164 & 8.708\(\pm\)0.145 & 1.412\(\pm\)0.005 & 2.793\(\pm\)0.083 & 6.291\(\pm\)0.153 & 0.0088 \\
**TFN** & 1.544\(\pm\)0.231 & 11.116\(\pm\)1.285 & 23.823\(\pm\)0.383 & 3.536\(\pm\)0.067 & 37.705\(\pm\)0.298 & 73.472\(\pm\)0.661 & 0.0440 \\
**SE(3)-Tr.** & 24.83\(\pm\)0.009 & 18.891\(\pm\)0.237 & 36.730\(\pm\)0.381 & 4.401\(\pm\)0.005 & 52.134\(\pm\)0.008 & 98.243\(\pm\)0.647 & 0.2661 \\
**Radial Field** & 1.060\(\pm\)0.007 & 12.514\(\pm\)0.009 & 26.388\(\pm\)0.331 & 1.860\(\pm\)0.075 & 7.021\(\pm\)0.150 & 16.474\(\pm\)0.003 & 0.0052 \\
**EGNN** & 0.711\(\pm\)0.029 & 2.998\(\pm\)0.089 & 6.836\(\pm\)0.003 & 0.766\(\pm\)0.011 & 3.661\(\pm\)0.055 & 9.039\(\pm\)0.216 & 0.0126 \\
**GMN** & 0.824\(\pm\)0.002 & 3.436\(\pm\)0.156 & 7.409\(\pm\)0.214 & 0.620\(\pm\)0.043 & 2.801\(\pm\)0.194 & 6.756\(\pm\)0.427 & 0.0137 \\
**SEGNN** & 0.448\(\pm\)0.003 & 2.573\(\pm\)0.053 & 5.972\(\pm\)0.168 & 0.471\(\pm\)0.026 & 2.110\(\pm\)0.004 & 5.819\(\pm\)0.335 & 0.0315 \\ \hline
**PINGO** & **0.433\(\pm\)**0.013 & **2.183\(\pm\)**0.048 & **5.614\(\pm\)**0.128 & **0.338\(\pm\)**0.027 & **1.693\(\pm\)**0.217 & **4.857\(\pm\)**0.580 & 0.0277 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Mean squared error (\(\times 10^{-2}\)) of the N-body system, and forward time in seconds for a batch size of 100 samples on Tesla T4 GPU. The header of each column is the time step for training and prediction. Bold font indicates the best result and Underline is the strongest baseline. Results are averaged across 5 runs. We report both mean and standard deviation in the table.

\begin{table}
\begin{tabular}{c c c c c|c c c c} \hline \hline \multirow{2}{*}{**Method**} & \multicolumn{3}{c}{**Charged**} & \multicolumn{3}{c}{**Gravity**} \\  & 250 ts & 500 ts & 750 ts & 1000 ts & Avg & \(250\) ts & \(500\) ts & \(750\) ts & \(1000\) ts & Avg \\ \hline
**GNN** & 73.40\(\pm\)0.400 & 31.79\(\pm\)2.28 & 12.862\(\pm\)0.008 & 29.72 & 181.92\(\pm\)0.21 & 90.33\(\pm\)1.93 & 30.66\(\pm\)1.13 & 0.746\(\pm\)0.008 & 75.93 \\
**GDE** & 92.65\(\pm\)2.50 & 43.94\(\pm\)1.09 & 12.20\(\pm\)0.212 & 0.652\(\pm\)0.008 & 37.36 & 136.05\(\pm\)0.158 & 56.80\(\pm\)0.162 & 12.12\(\pm\)1.048 & 5.588\(\pm\)0.15 & 51.39 \\
**EGNN** & 6.756\(\pm\)2.35 & 3.816\(\pm\)4.48 & 3.668\(\pm\)0.34 & 0.568\(\pm\)0.009 & 3.702 & 7.146\(\pm\)0.786 & 29.70\(\pm\)0.24 & 9.712\(\pm\)0.560 & 0.382\(\pm\)0.11 & 11.89 \\
**GNN** & 10.44\(\pm\)2.74 & 10.92\(\pm\)4.47 & 4.518\(\pm\)1.36 & 0.512\(\pm\)0.166 & 6.598 & 7.430\(\pm\)0.197 & 5.402\(\pm\)0.157 & 5.730\(\pm\)0.004 & 0.349\(\pm\)0.46 & 5.762 \\
**SEGNN** & 21.78\(\pm\)0.000 & 25.74\(\pm\)1.36 & 3.14\(\pm\)0.34 & 0.342\(\pm\)0.24 & 27.25 & 10.58\(\pm\)0.48 & 49.63\(\pm\)0.27 & 2.852\(\pm\)0.70 & 4.484\(\pm\)0.002 & 21.62 \\
**PINGO** & 0.188\(\pm\)0.03 & 0.312\(\pm\)0.008 & 3.60\(\pm\)0.046 & 0.309\(\pm\)0.11 & 0.292 & 0.064\(\pm\)0.002 & 0.128\(\pm\)0.005 & 0.176\(\pm\)0.004 & 0.210\(\pm\)0.007 & 0.145 \\ \hline \hline \end{tabular}
\end{table}
Table 2: The generalization from long-term to short-term. All models are trained on 1000ts and test on 250/500/750/1000 ts. Mean squared error (\(\times 10^{-2}\)) and the standard deviation are reported. Results averaged across 5 runs.
decreases significantly. By contrast, PINGO achieves similar results as in 1000ts, demonstrating its robust generalization to short-term prediction.
* Another interesting point is that PINGO's error exhibits a distinct trend compared to other baselines. While the errors of other baselines significantly increase with decreasing time steps, PINGO achieves even smaller errors with shorter time steps. This observation justifies our theoretical results that the error is bound by the chosen timestep.
* Additionally, the standard deviation of PINGO is much smaller than that of other baselines, indicating the numerical stability of PINGO. This result further confirms our theoretical finding that PINGO can obtain the unique latent trajectory between two discrete states.

Generalization capability for the long-term simulation.In this part, we evaluate the generalizability of models for extreme long-term simulation. Specifically, we train all models on 1000ts and use rollout to make the prediction for the longer time step (over 40 rollout steps, indicating over 40000ts.). Figure 3 depicts the mean squared error of all methods on two datasets. As shown in Figure 3, all baselines experience numerical explosion due to error accumulation during the rollout process, leading to a quick drop in prediction performance. In contrast, PINGO demonstrates an order-of-magnitude error improvement over other baselines for the extreme long-term simulation. This numerical stability can be attributed to the Neural ODE framework for modeling position and velocity. Nevertheless, we also noticed an increase in errors during the rollout process in PINGO. We conjecture this error may be introduced by the incorrect force estimation by GNN. Therefore, a more powerful GNN with better force estimation ability could further enhance the performance of PINGO.

Compared with numerical methodsTo verify the effectiveness of PINGO, we further compare it with symplectic Euler solvers that have a closed-form solution. The results of the gravity system are illustrated in Figure 4, and we provide results of charged systems in Appendix. Euler-\(\Delta t\) denotes results using \(\Delta t\) as the forward timestep. From the figure, we can find that under the same timestep (i.e., Euler-125), PINGO achieves better long-term performance. According to Corollary 4.3, PINGO is at least a first-order method if the prediction error is adequately minimized. Since the timestep of PINGO is set to 125, under appropriate conditions, PINGO will at least achieve comparable performance with Euler-125.

Ablation studyWe conduct two ablation studies on PINGO: (1) To validate the effectiveness of second-order modeling, we compare PINGO with its first-order variant PINGO\({}^{\text{1st}}\), where we use the same backbone and experimental settings to learn the first-order derivatives. The results are shown

Figure 4: Mean squared error for PINGO and symplectic Euler solver on the gravity system, which is calculated between them and groundtruth (i.e., Euler-1).

Figure 3: Mean squared error for the long-term simulation by rollout. All models are trained on 1000ts.
in Figure 6, we can observe that PINGO consistently enhances performance across all scenarios, with particularly notable improvements in long-term simulations. This validates the efficacy of incorporating second-order derivatives in modeling physical systems, emphasizing the significant advantages of integrating physical knowledge into the learning of dynamical systems; (2) Then we study the effect of the chosen timestep via increasing the iteration times \(\tau\). Since the target timestep is fixed, a larger iteration time indicates a smaller timestep. The results are displayed in Figure 5. It is obvious that better results can be achieved by choosing a small timestep. Additionally, the performance would not increase after a sufficient iteration, which is around 10 in both datasets. According to Theorem 4.2, these errors are mainly attributed to learning loss which is related to the representative ability of GNNs.

### CMU Motion Capture

For real-world applications, we evaluate our model on CMU Motion Capture Database [5], which contains various trajectories of human motions. Our main focus lies in the walking motion of a single object (subject #35) [21]. We adopt a random split strategy introduced by [15] where train/validation/test data contains 200/600/600 frame pairs. For comprehensive evaluations of long-term performance, we broaden our assessment scope to include scenarios with intervals of 40 ts and 50 ts, in addition to the default settings with 30 ts.

Implementation detailsIn this task, we use GMN as the backbone of PINGO. The norm of velocity and the coordinates of the gravity axis (z-axis) are set as node features to represent the motion dynamics. Note that the human body operates through joint interactions, we augment the edges with 2-hop neighbors. These operations are implemented across all applicable baselines. Following [15], 6 key bones are selected as sticks and the rest are isolated objects for GMN configurations.

ResultsTable 3 reports the performance of PINGO and various compared models. It is evident that PINGO outperforms all baseline models by a significant margin across all scenarios. Notably, the improvements are more pronounced in long-term simulations, with PINGO achieving 18.619 \(\times\) 10\({}^{-2}\) lower MSE than the runner-up model GMN. To gain further insights into the superior performance of PINGO, we illustrate the predicted motion of GMN and PINGO in Figure 7. Interestingly, it can be observed that the predictions of GMN appear to lag behind the ground truths, while PINGO demonstrates a closer match. This discrepancy may be attributed to the lack of constraints imposed by modeling the rollout trajectories. We provide more visualizations in Appendix.

\begin{table}
\begin{tabular}{c c c c c c c|c} \hline \hline
**Model** & **TFN** & **SE(3)-Tr.** & **RF** & **EGNN** & **GMN** & **PINGO** & **Abs. Imp.** \\ \hline \(30\) ts & 24.932\(\pm\) 1.023 & 24.655\(\pm\) 0.870 & 149.459\(\pm\) 0.750 & 24.013\(\pm\) 0.462 & 16.005\(\pm\) 0.386 & **14.462\(\pm\) 1.06** & 1.543 \\ \(40\) ts & 49.976\(\pm\) 1.664 & 44.279\(\pm\) 0.355 & 306.311\(\pm\) 1.100 & 39.792\(\pm\) 1.129 & 38.193\(\pm\) 0.067 & **22.229\(\pm\)** 1.49** & 15.964 \\ \(50\) ts & 73.716\(\pm\) 4.343 & 68.796\(\pm\) 1.139 & 549.476\(\pm\) 3.461 & 50.930\(\pm\) 2.675 & 47.883\(\pm\) 0.599 & **29.264\(\pm\)** 0.946 & 18.619 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Mean squared error (\(\times 10^{-2}\)) on CMU motion capture dataset. Results are averaged across 5 runs.

Figure 7: Visualization of Motion Capture with 50 time step. Left to Right: initial position, GMN, PINGO (all in blue). Ground truths are in red.
Conclusions

In this work, we highlight the problem of long-term simulation of N-body systems and introduce PINGO, a flexible neural graph ODE framework that is capable of learning physical dynamics from observed trajectories. Its core idea is to learn latent physical dynamics and employ numerical methods to infer system states. Theoretical findings show that our model can generalize to other timesteps via the same training criteria as existing studies. We demonstrate the potential of PINGO by applying it to a wide range of physical systems. PINGO outperforms all competitors in all cases. Extensive ablation studies have further verified the generalization ability of PINGO and the effectiveness of its physical design.
