SEGNO: Generalizing Equivariant Graph Neural Networks with Physical Inductive Biases
Authors: Yang Liu
Jiashun Cheng
Haihong Zhao
Tingyang Xu
Peilin Zhao
Fugee Tsung
Jia Li
Yu Rong
Venue: ICLR 2024 spotlight (2024)
Keywords: Equivariant Graph Neural Network
Graph Neural Network
Primary area: learning on graphs and other geometries & topologies

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.

Published as a conference paper at ICLR 2024
SEGNO: GENERALIZING EQUIVARIANT GRAPH NEU- RAL NETWORKS WITH PHYSICAL INDUCTIVE BIASES
Yang Liu Li; Jiashun Cheng!,2, Haihong Zhao', Tingyang Xu², Peilin Zhao?, Fugee Tsung".
Jia Li Yu Rongăt
'The Hong Kong University of Science and Technology (Guangzhou) "hhe Hong Kong University of Science and Technology *Tencent Al Lab
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 transi- tions among system states, opting to employ several discrete transformation layers to learn the direct mapping between two adjacent states; (2) Most models only ac- count for first-order velocity information, despite the fact that many physical sys- tems are governed by second-ordes motion laws. To incorporate these inductive biases, propose the Second-order Equivariant Graph Neural Ordinary Differ- ential Equation (SEGNO). Specifically, we show how the second-order continuity be incorporated into GNNs while maintaining the equivariant property. Fur- thermore, we offer theoretical insights into SEGNO, highlighting that it can learn unique trajectory between adjacent states, which is crucial for model generaliza- tion. Additionally, we prove that the discrepancy between this learned trajectory of SEGNO and the true trajectory is bounded. Extensive experiments on complex dy- namical systems including molecular dynamics and motion capture demonstrate that our model yields a significant improvement over the state-of-the- baselines.
INTRODUCTION
Equivariant Graph Neural Networks (Equiv- GNNs) (Satorras et al., 2021; Han et al., 2022b;
Brandstetter et al., 2021; Huang et al., 2022; Wu et al., 2024) have emerged as essential tools for simulating the multi-object physical system, i.e., N-body systems, which is relevant to nu- merous fundamental scientific domains, includ- ing molecular dynamics (Karplus & McCam- mon, 2002), protein folding (Gligorijević al., 2021), robot motion planning/control (Siciliano al., 2009). Specifically, given the input state, they learn to predict the output state after specific timestep. To achieve these, Equiv- GNNs model the whole system geometric graph, which treats physical objects as nodes, and physical relations edges, and encode the symmetry into a message- message-passing network to ensure their outputs equivariant with respect to any Paunslatooooriienaatiiiiieleee inputs. Such property makes them well-suited for learning the unknown dynamics of complex physical systems that cannot be described analyti- cally (Han et al., 2022a;b).
Wrong
Noa Discette
Uniqur Flini-under Continans
Unigue Sccond-oddrr Continaoas
Input Stute
Output State
Correct
Learned Trajectory
Figure 1: Learned trajectories of models with dif- ferent inductive bias. All models map input to output. However, discrete and first-order con- tinuous models fail to learn the true intermediate states due to the lack of considering continuity and second-order laws.
of the
"qqual contribution. Work is done when Yang Liu and Jiashun Cheng worked as interns in Tencent Al Lab. rong@hommtaa com)
1 "ooresponniig authors (ji ee@ust.hk, hk, yu.

---

Published as a conference paper at ICLR 2024
Learning to model such interacting systems is challenging. Given the vast parameter space of GNNs and finite observations of transitions of system states, there would exist multiple solutions that satisfy observed data (Curry & Morgan, 2006). Therefore, learning the real dynamic function from these solutions is crucial to model generalization ability Gruver et al. (2022). Though Equiv-GNNs (Han al., 2024) have partially addressed this challenge by eliminating models that lack symmetry, ex- isting Equiv-GNNs have yet to incorporate sufficient physical inductive bias to model the physical dynamics.
In particular, two essential inductive biases have not been well investigated in this field. First, exist- ing models (Satorras al., 2021; Brandstetter et al., 2021) are composed of sequence of discrete state transformation layers, which learn a direct-mapping between adjacent states with discrete tra- jectories. We refer to them as discrete models. They are inconsistent with the continuous nature of system trajectories and fail to learn correct intermediate states. Second, most models only ac- count for first-order information. Many physical dynamical systems, such as Newton's 's equations of motion, are governed by second-order laws (Norcliffe et al., 2020). Therefore, these methods learn the system's s state and fail to capture the underlying dynamics of the
incomplete representations physical systems. In Figure 1, we illustrate the comparison of learned trajectories of models with different types of inductive bias.
In this work, we take deep insight into the continuity and second-order inductive bias in Equiv- GNNs and propose a framework dubbed Second-order Equivariant Graph Neural Ordinary Dif- ferential Equation (SEGNO'). Different from previous models that use Equiv-GNNs fit discrete kinematic states, SEGNO introduces Neural Ordinary Differential Equations (Neural ODE) to ap- proximate a continuous trajectory between two observed states. Furthermore, better estimate the underlying dynamics, SEGNO is built upon second-order motion equations to update the position and velocity of the physical systems. Theoretically, we prove the uniqueness of the learned latent tra- jectory of SEGNO and further provide an upper bound on the discrepancy between the learned and the actual latent trajectory. Meanwhile, we prove that SEGNO can maintain equivariance properties identical to the backbone Equiv-GNNs. This property offers the flexibility to adapt various back- bones in SEGNO to suit different downstream tasks in plug-and-play manner. We conduct extensive experiments on both synthetic and real-world physical systems. Our results reveal that SEGNO has a better generalization ability over the state-of- -fhth--rt baselines and second-order inductive bias is beneficial to learn complex multi-object dynamics.
2 BACKGROUND
N-body System We study N-body systems (Kipf al., 2018: Huang et al., 2022) with set of N particles P = [ At time t, the state of each particle in the system is represented by:
geometric features including the position vector g² E R and the velocity vector q(") R3; non-geometric features such as mass or charge. denoted by h4 Rd; 3. spatial connection with others, where an edge Eij is constructed via geometric distance cutoff or physical interaction (e.g., chemical bonds) and the edge attributes (e.g.. object distances, bond type) are denoted by aij- For simplicity, we denote (1))), ())) and (h e = {eıg], a = (miz)) as dynamic and static state information at system level correspondingly. In this work, we focus dynamical systems that can be formulated as:
d²q())
g(t)
dia l(( h))
where q(()) is the acceleration time Given trajectory q² with the initial system states (a(t)), (tt) at time to and static states (h, a), our goal is to predict the subsequent position ((ti) within fixed time interval T t to-
E(3) Equivariance In 3-dimensional Euclidean space, the laws of physics remain invariant re- gardless of E(3) transformations. including translation. rotation, and reflection. Formally, a func- tion F X P + Y, where X, / C TR is F(3)-iquivariant,, if for any transformation o E F(3),
Flo x .) = o F(x,.....
X.
SEGNO is also musical term in Italian meaning "sign", marking the beginning end of musical repeat. 2To avoid ambiguity, (t) denotes the position of the i-th particle at time t, whereas 7 denotes the trajectory over the entire time interval.
2

---

Published as a conference paper at ICLR 2024
In general, given the system state at time tu, modern Equiv-GNN d with parameters & directly predicts gltx) by leveraging several message passing layers on an interaction graph. To maintain brevity and avoid ambiguity, we omit the temporal superscript t here since the predictions at different times share the same model 178. Specifically, each layer of Te computes
Equivarant GNNs
m plqe qz qi qz. h4, he aig), d: q/ h = v(q. q: h..
JEN;
where and V/ are the edge message function and node update function, respectively. maj defines the message between node and j. AG collects the neighbors of node The prediction is obtained by applying several iterations of message passing. To construct equivariant layers, and could be both equivariant (e.g., SEGNN (Brandstetter et al., 2021)) or alternatively, μ is equivariant and is invariant (e.g., EGNN (Satorras et al., 2021) and GMN (Huang al., 2022)).
3 SEGNO FRAMEWORK
In this section, we introduce how the proposed SEGNO works and its equivariant properties, In our dynamic systems, we incorporate the ODE formulation to model the latent continuous trajectory
with initial system states (t)) ((to)). The position q (to++t) for any l' E [0, 7 is calculated as
pott'
(glto) + I f(a(m)) h) dm) dt
ΦU lo (gllto) = a (lo+t) = =q(0))
to
(t)) poott
g h) dt,
where y represents mapping from the trajectory to its first-order derivative q()) In this vein, we can denote (( as (g Note that most physical dynamical systems follow the second-order motion law. The velocity ((t++t") for any t' [0, T] can further be computed as
(t)) h)) + /
pto-t'
Ve s (bo) == (lotte) = g(q
f(q h)) dt.
(5)
To incorporate the second-order inductive bias, SEGNO parameterizes the acceleration function fo(g h);
(6)
where fo, approximation of ʃ. represents Equiv-GNN with parameters which computes:
m p(q: (1) (1) h, h2, an), q h4 = v(g:" (t)
hi, Σ m
JEN:
where and // are determined by this Equiv-GNN backbone. Let qe denote the approximated trajectory generated by SEGNO satisfying the initial conditions do following Eq. 4, the predicted position of SEGNO at time t can be represented by
(to) = ((ot)) gbaa) -t(t) Then,
rot (0( dt)dt
@r (q (ta)) = 4b (t1) = (t0)
(g + S fo( h)) dm) dt to
(to)
49
+
(8)
to
(to)
to (0) dt,
qa +
where parameterized version of is determined by the integral of the GNN fa in Eq. 4. Since it is infeasible to directly calculate the integration in &r, parameterized go. fo. SEGNO, we utilize an ODE solver to generate a discrete trajectory that serves as an approximation of the latent continuous trajectory. Specifically, we divide the entire time interval T' into 7 equal sub-intervals with timestep At = T/T. We then denote V.r.o. and At.g as the numerical integrators that approach v'st.g.f and dac. using the following equations
(g (tt) and d' (f)))) with
(t)
h) V Dtt.go.fo (g = q + Gi(fe (h), At), (++At) & 2s (q) (t) = q + Gz At (1) At),
go (t+At)
3

---

Published as a conference paper at ICLR 2024
where 91 and 22 are the increment functions that approximate the increment of continuous integral given the initial value of the integrand and the integration width Åt. For instance, with the increment y. the numerical integrators become the Euler integrators
functions (z, y) G2(2,)) v)
(1~
de (1) + fo At,
48 (++A))
9e + 2
((+50) At.
(10)
For approximation, SEGNO composites integrator a & Neural ODE solver following
(11) ④T (q
(to)) = 2 St.gg o Pr.. (q (f)) = ( (t..) (q("")))
(11)
49
As result, SEGNO offers way reuse existing Equiv-GNNs build second-order Neural ODEs.
However, a natural question arises: does a Neural ODE solver compromise the equivariance of backbone GNNs? To address this question, we show that the equivariance property of backbone GNNs could be maintained in SEGNO.
Proposition 3.1. Suppose the backbone GNN fo of SEGNO is (e3--veaivariant and translation- invariant, and the integrators' increment functtionn Gr, 92 are 0(3) -qquivvau,a then the output tra- jectory qo is -qqiivaria.
The proof and example illustrations are provided in Appendix A. where we show that the general numerical integrators including symplectic Fuler, Velocity Verlet, and Leapfrog satisfy Proposi- tion 3. 1. Meanwhile, since Equiv-GNNs such as EGNN (Satorras et al., 2021) and GMN (Huang et al., 2022) are built upon preserve the same equivariant property as the backbone GNNs.
and translation-invarial functions, SEGNO would
SEGNO ANALYSIS
Besides equivariance, another essential problem is how SEGNO learns from observed system states. In this section, we examine the approximation quality of SEGNO.
4.1 SOLUTION UNIQUENESS
is known that continuous dynamics have unique solution under specific continuous conditions, according to Picard's 's existence theorem (Coddington et al., 1956).
Lemma 4.1. For the system q(0) = (D h), with given initial position qat)) and velocity q(to), if f is Lipschitz continuous, then this system has unique solution q0)) over the interval E [to, f1).
The proof is provided in Appendix A.2. In addition, under the SEGNO framework, we can obtain the following results:
Proposition 4.2. Given the same conditions as in Lemma 4.1, if the realistic measurement gff)) given, there exists a fo- obtained by minimizing the discrepancy between g((a)) and qtt,), such that for (g h) = (") h)) holds over the interval E [to,
The detailed proof is given in Appendix A.3. This proposition remarks that SEGNO be trained in line with prior studies (Satorras et al., 2021; Brandstetter et al., 2021). That is, given to and t as the input and target timesteps. SEGNO is trained to minimize the discrepancy between the exact and approximated positions:
Lunin
Σ BEDiiin
g.
(t1)
(12)
(fi) (4r) where Durain denotes the training set. With slight abuse of notation, here we denote 98. q as the model prediction and actual position of trajectory s. If our learned model adequately approxi- mates the system, Proposition 4.2 shows that becomes possible for SEGNO to recover the latent trajectories of [to, &1] between input and output system states via Neural ODE. In contrast, in the ab- sence of continuous constraints, multiple discrete trajectories can exist between the input and output, making it challenging for discrete models to accurately learn the underlying dynamic functions.
4

---

Published as a conference paper at ICLR 2024
4.2 APPROXIMATION ABILITY
In this section, theoretically show that SEGNO is capable of learning trajectory that remains bounded to the real solution. To achieve this, we first derive the boundness of the learned second- order derivative fg. Existing theoretical findings (Zhu et al.. 2022) have illustrated that training using Neural ODE solvers results in a bounded approximation of their first-order derivatives. In this work, we generalize this theorem the second-order Neural ODE employed with Euler integrator.
4.2.1 BOUND OF ACCELERATION
To begin with, we first introduce relevant notations. Let B(p,r) R denotes a real ball of radius r > 0 centered at E R for a subset IC C I2 and an analytic function 17, we define
B(C,r) = U B(p,r), ||il = sup |17())]
(13)
PEK
pEK
We denote set of points on the exact trajectory associated with the time increment interval [a, as V2
f (g S S b).
(14)
Given the time interval T = 4 to uniformly partitioned into 7 sub-intervals of timestep At (i.e.. lo + Al <.... < (1), the following theorem holds:
Theorem 4.3. For , > O, the ODE solver is compositions of an Euler Integrator & St.g with At T/T, and we denote the supremum norm of the approximation error for the ODE solver within B(V,", .11
CJT = + -1( ΦT B(V3" .r)
Suppose that the target f and the learned fe analytic and bounded by m2 B(V .T2), and the target and the learned go are analytic and bounded by my on B(V?T, rr). Then, there exist constants To such that, i T < To. Vt
[to, t1).
C2T At
llfo(q")) h) f(q"), h)lle S S(At+
The proof is reported in Appendix A.4. The assumption that the ʃ.g g and the learned fo; analytic and bounded within specific balls is valid and grounded in real-world physical observations. Coneretely, the trajectories of most dynamic physical systems exhibit smoothness, implying that there are no sudden changes in velocity or acceleration. This smooth property ensures that the system behaves in physically realistic and predictable manner, which widely employed practice in control theory and dynamic systems analysis (Zhu et al., 2022).
4.2.2 BOUND OF TRAJECTORY
Based on Theorem 4.3, we can now analyze the error introduced by the Euler integrator. We in- troduce two metrics common in classical numerical analysis, namely, local and global truncation error (Poli et al., 2020). The local truncation error C++A. of SEGNO is defined as:
C+ ||9 (1+A0 t0) n0) At (o h)Af*12.
(15)
which represents the error for single timestep. The global truncation error EL+AAe is defined
ELTAA ||q (++kk)) (++tAt))
(16)
which denotes the enror accumulated in the first k steps. Then we have
Corollary 4.4. Given the same conditions as in Theorem 4.3, if our learned model adequately approximates the system and go and fo satisfy the Lipschitz continuity, then, Vl. E [lo.1] and k = -,1 ... the local truncation enror E141t and the global truncation error Ertkst are the order OAt2) and O(t)), respectively.
The proof is reported in Appendix A.5. Corollary 4.4 shows that for the prediction of 7-th itera- tion, its error depends on the chosen At. These statements imply that SEGNO can be trained by minimizing Eq. 12 and generalize to other timesteps.
Though there exist studies (Sanchez-Gonzall Gonzalez et al., 2020; Bishnoi et al., 2022) that employ accel- erations to train models, remark that SEGNO is different from them and better. As Theorem 4.3 states, we aim to learn the instantaneous acceleration of each timestep by minimizing the position discrepancy, while previous works the average acceleration to train the model and then adopt
5

---

Published as a conference paper at ICLR 2024
semi-implicit Euler method to update the next state well. The average acceleration is computed from the observed trajectory (i.e., q(t.)) = g(aa) 2g(h)) + (fa)). Thus, we can find that their local truncation error is ((T) even if their loss zero, while that of SEGNO is ((Ai)(T = -Li).
SEGNO theoretically achieves lower error without the need to obtain average acceleration by differ-
entiating future positions.
Empirical verification To verify our theoretical findings, we train multiple models of EGNN (Sator- ras et al., 2021) and SEGNO with EGNN backbone, utilizing different random seeds on a 3-body sys- tem with two adjacent states ( o (t))) We derive the intermediate state (ffo..) from their intermediate layers/iterations, which serve as estimations of la- tent motion trajectories between the two states. In Figure 2, we visualize all predicted trajectories (in dotted grey line), the mean trajectory (in red line), and the mean and variance of predicted states. We can observe that EGNN predictions exhibit high er- (red circle) and large variance (blue area) the intermediate state q (to iinicatinn that the discrete models are incapable of learning the underlying real dynamics from observed system states. Conversely, the learned trajectory of SEGNO demonstrates a sig- nificantly lower error and small variance. The ad- ditional numerical comparisons are shown in Ap- pendix C.I.
trajectory prodicted traeetoory Predicted tnjeedory
Sampled ground truch state
Lateni intermediate state Variance prodicted states
fas
f1
EGNN
SEGNO
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 Sec- tion 5.1). The green, red, and dotted grey lines are the true, average, and predicted tra- jectories. fo, & is the observed states. fo.s is the predicted latent state The blue area de- notes the variance.
5 EXPERIMENTS
Datasets To validate the effectiveness of SEGNO, we first evaluate it on two simulated N-body systems, namely Charged particles and Gravity particles, which are driven by electromagnetic (Kipf et al., 2018) and gravitational forces (Brandstetter et al., 2021) between each pair of particles, re- spectively. Subsequently, we compare our model with state-of-the-art models in two challenging datasets: (1) MD22 (Chmiela et al., 2023) which includes the trajectories of seven large and dif- ferent types of molecules generated via molecular dynamics simulation: (2) CMU motion cap- ture (CMU, 2003), which contains various trajectories of human motions in real-world. Note that all these dataset are symmetric, N-body systems and MD22 are E33 ((3)---uivariant, and CMU motion capture is -qqiivaairrr.
Baselines We compare SEGNO with various models: (1) Equivariant models including Radial
Field (Köhler et al., 2019), TFN (Thomas et al., 2018), SE(3) Transformer (Fuchs et al., 2020), EGNN (Satorras et al., 2021), GMN (Huang et al., 2022), SEGNN (Brandstetter et al., 2021); (2) Non-equivariant models including GNN, Linear model, and Graph Neural Ordinary Differential Equation (GDE) (Poli et al., 2019).
5.1 SIMULATED N- B-BODY SYSTEM
Implementation details We build upon the experimental setting presented in (Satorras et al., 2021) where the task is to estimate all particle positions after a fixed timestep. Each system consists of 5 particles, each with initial positions, velocities, and attributes like positive/negative charge or mass. The graph is fully connected and the initial velocity norm is provided as additional input node features. We employ EGNN (Satorras et al., 2021) as the GNN backbone (fa) of SEGNO in these settings. For each system, besides the common setting (Satorras et al., 2021; Brandstet- ter et al., 2021) with 1000 timesteps (1000 ts), we introduce two extra targets-1500 ts and 2000 is-which render the learning of latent trajectories more challenging. Other settings including the are introduced in Appendix B.
*Symmetry is partially hroken by gravity.

---

Published as a conference paper at ICLR 2024
Table 1: Mean squared error (x10-2) of the N-body system. Bold font indicates the best result and underline is the strongest baseline. Results averaged across runs. We report both mean and standard deviation in the table.
Method
Charged 1500
Gravity 1500
1000
2000 Es
1000
2000 is.
Linear
6.830 wns
200012
39.513: can
7.928
29.270
58.521
GNN
L.777
6.059
10.591 039
1.400)
1.691gamm
10.508
GDE
1.285
4,026
8.708
1.412-48ss
2.793
6.291
11.
21,823 10s
3.536_4a0
37.705
7
SE(3)-Th
2.4832
18.89 2
36,730= 0:31
4,401=:00
52.1342.us4
98. 2
Rudial Field
1 0600
12.514g
26.388-2011
1.8
2.021
16,474gmmm
EGNN
2.998:2.0m com
6.836= une
0.766
3,66 quss
9.039
GMN
0.824 uuzz
3.436 0.19
7.409 c214
0.620 0.94
2,801 0.15
6.756 0427
SEGNN
2
5.972
=
2.110
5.819
SEGNO
0.433 w0.3
1.858 code
4,285 660
0.338-000 980
1.362 2007
4.017.0007
Table 2: Ablation studies (x10 2) simulated N- N-body systems. Results averaged across 5 runs.
Charged
Gravity
Order
Continuity
1000 is
1500 is
2000 is
1000 Ls
1500 Ls
2000 Ls
First Second
Discrete
0.798
2.215 wass
4.996 0:48
0.466: oer
2.342: eu24
5.501 0254
Discrete
2.125= 0.738=
4.948 use
0.362 oon
2.249 non
5,343
First
Continuous
0.537+ a944
1.972 wase
4.405 0129
0.352 oues
1.468 nowe
4.895 c.02
Second
Continuous
0.433
1,858
4,285
0.338
1.362 n017
4.017
0007
Results Table depicts the overall results of all models on two datasets. is evident that SEGNO, equipped solely with EGNN, outperforms all baselines across all datasets and settings. Specifically, compared to the best baseline SEGNN, the average eтror improvement (x10-3) on Charged and Gravity datasets is 0.805 and 0.894 respectively, demonstrating significant improvement. Addition- ally, as timesteps increase, the performances baselines largely drop while SEGNO still can model the latent trajectory. Thus, SEGNO's performance improvement becomes more pronounced. For example, compared to the best baseline SEGNN, the relative improvement Gravity datasets in- creases from 28. 28% 1000 time steps to 35.45% 1500 time steps, demonstrating the strong generalization ability of SEGNO.
5.1.1 ABLATION STUDY
In this section, we conduct several ablation experiments on simulated N-body systems to scrutinize model design, and provide empirical validation for our theoretical findings.
To validate the effect of inductive biases incorporated in SEGNO, we construct three variants of SEGNO, each featuring unique combination of physical inductive biases. Table 2 reports the results where the term 'First' indicates that the model employ fo parameterize velocity rather than accelerations. 'Discrete' implies that SEGNO does not share the parameters across iterations, akin to discrete models. The original version of SEGNO denoted as 'Second' and 'Continuous'. From Table 2, models consistently surpass discrete models and the second-orden bias consistently enhances the performance of first-order models. These findings serve to corroborate the efficacy of SEGNO's construction and further emphasize the significant advantages of integrating physical inductive bi- ases into the learning process of dynamics.
Physical inductive biases in SEGNO
observe that across all scenarios, continuous
Effeets of iteration times т We further investigate the impact of the iteration number in SEGNO empirically verify Corol- lary 4.4. Given that the target time interval T remains constant, a larger iteration number indicates smaller interval At. The re- sults are displayed in Figure 3. is obvious that superior outcomes be attained by opting for smaller interval Al. Additionally, the performance would not increase after reaching a sufficient num- ber of iterations, which is approximately 10 for both datasets. As per Theorem 4.3, the plateau in performance improvement can be attributed to learning errors which is related to the representative ability of GNN backbone in SEGNO.
N-body system
Gravity
Charged
0.2
MSE
0.0
T. times
Figure 3: Effects of iteration number T.

---

Published as a conference paper at ICLR 2024
Table 3: Mean squared error (x10-3) on MD22 dataset. Bold font indicates the best result and
Underline is the strongest baseline. Results averaged across 3 runs.
Molecule
Type
Atom
TEN
RF
EGNN
GMN
SEGNO
Ac-Ala3NNHMe
Protcin
42
1,434
1.340g o012
0.539g m
0.960=
0.418: noes
DHA
Lipid
0.577
1.487 woe
0.782
0
0.370
AT-AT
Nuclese acid
60
a.407
1.270, one
0.583
0.568
nore
0.440 cues
Stachyuse
Carbohydrate
2,069_
0.629
0.583
pres
AT-AT-CG-CG
Nuclesc acid
118
2.596 use
0.609
0,58
2:00
0.394 000
Buckyball Catcher
Supramolecule
0,440
0.554
0.309
0.199 0103
Double-walled Nanotube
Supramolecule
370
0.382
0.349
0,32 2
0.225
Table 4: Mean squared error (x (x10-3) on CMU motion capture dataset. Bold font indicates the best result and Underline is the strongest baseline. Results averaged across 3 runs.
Model
TFN
SE(3)-Tr.
RF
EGNN
GMN
SEGNO
Abs, Imp.
24.932+ 102
24.655
149.459: usse
24.01
16,005 uss
14.462
1.543
40
49.976. 1884
44.279
306.31 1.30
39.792,
38.193 user
22.229:
15,964
50 ts 73.7161430
68.796. 1.35
549.4761 1461
50.930 2075
47.883 n
29.2641.9t
18.619
5.2 MOLECULAR DYNAMIC
Implementation details We use the atomic number and initial velocity norm input node fea- tures. Two atoms are neighbors if their distance is less than a threshold. Our goal is to predict atom positions after 10 data frames. We run experiments on NVIDIA RTX A6000 GPU. TFN and Srans-Tran run out of memory, thus we omit results of SE(3)-Transformer and part of TFN results. We use EGNN (Satorras et al., 2021) as the GNN backbone (fo) of SEGNO. Other settings including the pprrmeerrsmeterss are introduced in Appendix B.
Results Table 3 summarizes the results of all models. It is evident that SEGNO outperforms the baselines across molecules, even on the Double-walled Nanotube comprising 370 atoms, support- ing the general effectiveness of encoding physical inductive biases. It is worth noting that SEGNO utilizes GMN as its backbone. In comparison to the original GMN, the errors are considerably di- minished in all instances. The average relative improvement of SEGNO over GMN 7 molecules is 15.6%. Such results demonstrate the effectiveness of enhancing physical inductive biases on equivariant GNNs in empirical applications.
5.3 CMU MOTION CAPTURE
Implementation details CMU Motion Capture (CMU, 2003) contains the trajectories of human motion under several scenarios. Following previous studies (Kipf et al., 2018), we focus on the walking motion of single object (subject #35). The goal of this task is predict the data frame after 30 timesteps. Similar to N-body systems, we broaden our assessment scope to include scenarios with intervals of 40 ts and 50 ts, in addition the default settings with 30 ts. We use GMN as the backbone of SEGNO. 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, two joints are I-hop neighbors if they are connected naturally and we augment the edges with 2-hop neighbors. Other settings are introduced in Appendix B.
Results Table reports the performance of SEGNO and various compared models. We can ob- serve that SEGNO outperforms all baseline models by significant margin across all scenarios. Notably, the improvements are more pronounced in long-term simulations, with SEGNO achieving 18.619 10-2 lower MSE than the runner-up model GMN. To gain further insights into the supe- rior performance of SEGNO, we illustrate the predicted motion of GMN and SEGNO in Figure 4.
12
25
Figure 4: Visualization of Motion Capture with 50 ts. Left to Right: initial position, GMN, SEGNO (all in blue). Ground truths are in red.
8

---

Published as a conference paper at ICLR 2024
Interestingly, it can be observed that the predictions of GMN appear lag behind the ground truths, while SEGNO demonstrates a closer match. This discrepancy may be attributed to the lack of con- straints imposed by modeling latent trajectories. We provide more visualizations in Appendix C.8.
6 RELATED WORK
Graph neural networks (Li et al., 2019; Tang et al., 2022: Gao et al., 2023; Tang et al., 2023) have shown promising performance in learning complex physical dynamics. IN (Battaglia et al., 2016), NRI (Kipf et al., 2018), and HRN (Mrowca et al., 2018) are pioneer works that model physical objects and relations as graphs and learn their interaction and evolution. Recently, re- searchers have taken into account the underlying physical symmetry of systems. TEN (Thomas al., 2018) and SE(3) Transformer (Fuchs al., 2020) employ spherical harmonics construct models with 3D rotation equivariance in the Euclidean group for higher-order geometric represen- tations. LicConv (Finzi al., 2020a) and LieTransformer (Hutchinson et al., 2021) leverage the Lie convolution to extend equivariance on Lie groups. SEGNN (Brandstetter et al., 2021) proposes to use steerable vectors and their equivariant transformations to represent and process node features. In addition to these methods, series of studies (Satorras et al., 2021; Huang et al., 2022; Schütt et al., 2021; Wang & Chodera, 2023) apply scalarization techniques to introduce equivariance into the message-passaae passing process in GNNs. Nevertheless, these methods model dynamics in physical systems solely by learning direct mappings between discrete states. They ignore the second-order and continuous nature of observed trajectories, leading to suboptimal generalization performance.
Another research line (Thangamuthu et al., 2022) leverages ODE and Hamiltonian mechanics to capture the interactions in the systems such Lagrangian Neural Networks (LNN) (Finzi et al., 2020b; Lutter et al., 2019: Bhattoo et al., 2023), Hamiltonian neural networks (HNN) (Greydanus al., 2019), and Neural ODE (Chen al., 2018: Gruver et al., 2022; Norcliffe al., 2020). Recent studies have also enhanced these methods with GNNs. GDE (Poli et al., 2019) and HOGN (Sanchez- Gonzalez al., 2019) combine GNNs with a differentiable ODE integrator. GNODE (Bishnoi al., 2022) incorporates graph-based Neural ODE with additional inductive biases, such as Newton's third law. Compared with these works, we provide theoretical insights that show a second-order Graph Neural ODE can obtain bounded error of instantaneous acceleration and position through minimizing position discrepancy. These findings are further validated in complex applications in- cluding molecular and human motion dynamics. In particular, GNS "Sanchez-Gonzale et al., 2020; Pfaff et al., 2021) also optimizes models via accelerations. However, they only learn the average accelerations that are calculated from the observed trajectories. In contrast, our study focuses learning instantancous accelerations and we show it theoretically achieves lower errors than GNS. Furthermore, it is worth mentioning that GNS does not consider equivariance, which is critical in- ductive bias that captures symmetries in physical systems. We provide an experimental comparison between them in Appendix C.2.
Besides the above studies, a series of works have combined GNNs and the first-order Neural ODE to learn multi-agent systems (e.g., social networks), including CG-ODE (Huang et al., 2021), LG- ODE (Huang et al., 2020), and HOPE (Luo et al., 2023). However, due to the application difference, they do not consider equivariance and focus on historical state encoding. Thus, they are hard to extend to our task where equivariance and second-order information are vital.
7 CONCLUSION
In this work, we address the generalization problem of learning N-body systems and introduce SEGNO, which incorporates the equivariant property from GNN backbones and second-order phys- ical inductive bias. We theoretically prove the uniqueness and boundedness of the trajectories in- ferred by SEGNO and empirically demonstrate the potential of SEGNO by applying it to wide range of physical systems. Extensive ablation studies have further substantiated the generalization ability of SEGNO. For future works, we are interested in (1) extending our framework to solve stochastic (Salvi et al., 2022) and partial differential equations (Fortunato et al., 2022: Strönisch et al., 2023; Cao et al., 2023; Xue et al., 2023); (2) jointly considering trajectory forecasting and static tasks such as molecular property (Satorras et al., 2021) or force field predictions (Batzner et al., 2022); (3) integrating more sophisticated physical principles through advanced techniques such as pre-training (Rong et al., 2020; Cheng et al., 2023; Jiao et al., 2023), prompt tuning (Sun et al., 2023) or the utilization of large language models (Li et al., 2023).

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
The research of Li was supported by NSFC Grant No. 62206067, HKUST--KUST(GZ) 20 for 20 Cross-campus Collaborative Research Scheme C019 and Guangzzhu-HHKUST(GZZ) Joint Funding Scheme 2023A03J0673. The research of Tsung was supported by NSFC Grant No. 72371271 and Guangzhou Nansha District Key Project 20232D003. Particularly. Yu Rong expresses heartfelt gratitude to his wife, Yunman Huang, for the birth of their daughter Xing Rong, which brought immense joy during the pursuit of this research.
REFERENCES
Peter W Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. arXiv preprint arXiv: 2016.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korn- bluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. EQ)evvaurr graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Ravinder Bhattoo, Sayan Ranu, and N. M. Anoop Krishnan. Learning the dynamics of particle- based systems with lagrangian graph neural networks. Mach. Learn. Sei. Technol., 4(1):15003, 2023.
Suresh Bishnoi, Ravinder Bhattoo, Sayan Ranu, and NM Krishnan. Enhancing the inductive biases of graph neural ode for modeling dynamical systems. arXiv preprint arXiv: 10740, 2022.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling. Ge- ometric and physical quantities improve (3) equivariant message passing. arXiv preprint arXiv:2110.02905, 2021.
Yadi Cao, Menglei Chai, Minchen Li, and Chenfanfu Jiang. Efficient learning of mesh-based phys- ical simulation with bi-stride multi-scale graph neural network. In ICML, volume 202, pp. 3541- 3558, 2023.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In NeurIPS, pp. 6572 6583, 2018.
Jiashun Cheng, Man Li, Jia Li, and Fugee Tsung. Wiener graph deconvolutional network improves graph self-supervised learning. In AAAI, pp. 7131-7139, 2023.
Stefan Chmiela, Valentin Massilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9())eadf0873, 2023.
CMU. Carnegie-mellon motion capture database. 2003. URL
as cmu edu.
Earl A Coddington, Norman Levinson, and Teichmann. Theory of ordinary differential equations, 1956.
Bruce Curry and Peter H Morgan. Model selection in neural networks: some difficulties. European Journal of Operational Research, 170(2):567-577, 2006.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolu- tional neural networks for cquivariance to lie groups arbitrary continuous data. In ICML, pp- 3165-3176 3176. PMLR, 2020a.
Mare Finzi, Ke Alexander Wang, and Andrew Gordon Wilson. Simplifying hamiltonian and la- grangian neural networks via explicit constraints. In NeurIPS, pp. 13880 13889, 2020b.
Meire Fortunato, Tobias Pfaff, Peter Wirnsherger, Alexander Pritzel, and Peter W. Battaglia. Multi- scale meshgraphnets. CoRR, abs/2210.00612, 2022.
Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling, Se transforrrss:::: 3d roto- translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
10

---

Published as a conference paper at ICLR 2024
Ziqi Gao, Chenran Jiang, Jiawen Zhang, Xiaosen Jiang, Lanqing Li, Peilin Zhao, Huanming Yang, Yong Huang. and Jia Li. Hierarchical graph learning for protein-protein interaction. Nature Communications, 14(1):1093, 2023.
Vladimir Gligorijević, P Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Beren- berg. Tommi Vatanen, Chris Chandler, Bryn C Taylor, Ian M Fisk, Hera Vlamakis, al. Structure- based protein function prediction using graph convolutional networks. Nature communications, 12(1):3168, 2021.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. NeurIPS, pp. 15353-15363, 2019.
Nate Gruver, Mare Anton Finzi, Samuel Don Stanton, and Andrew Gordon Wilson. Deconstructing the inductive biases of hamiltonian neural networks. In ICLR, 2022.
Jiaqi Han, Wenbing Huang, Hengbo Ma, Jiachen Li, Josh Tenenbaum, and Chuang Gan. Learning physical dynamics with subequivariant graph neural networks. In NeurlPS, 2022a.
Jiaqi Han, Wenbing Huang. Tingyang Xu, and Yu Rong. Equivariant graph hierarchy-basee neural networks. In NeurlPS, 2022b.
Jiaqi Han, Jiacheng Cen, Liming Wu, Zongzhao Li, Xiangzhe Kong, Rui Jiao, Ziyang Yu, Tingyang Xu, Fandi Wu, Zihe Wang, Hongteng Xu, Zhewei Wei, Yang Liu, Yu Rong, and Wenbing Huang. A survey of geometric graph neural networks: Data structures, models and applications, 2024.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural networks, 2(5): -366, 1989.
Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant graph mechanics networks with constraints. In International Conference on Learning Represen- tations, 2022. URL https: /tttpp:://peereeeee wwwp//weew..e..eeeeeepee. net/ forum htt.:///ww.......oooooooaaaaaaaa 2d-
Zijie Huang, Yizhou Sun, and Wei Wang. Learning continuous system dynamics from irregularly- sampled partial observations. In NeurlPS, volume 33, pp. 16177-16187, 2020.
Zijie Huang, Yizhou Sun, and Wei Wang. Coupled graph ode for learning interacting system dy- namics. In KDD, pp. 705-715, 2021.
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: equivariant self-attention for lie groups In ICML, pp. 4533 4543. PMLR, 2021.
Rui Jiao, Jiaqi Han, Wenbing Huang, Yu Rong, and Yang Liu. Energy-motivated equivariant pre- training for 3d molecular graphs. AAAI, 2023. ISBN 978-1-57735-880-0
Martin Karplus and J Andrew McCammon. Molecular dynamics simulations of biomolecules. Na- ture structural biology, 99):646-652, 2002.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. arXiv preprint arXiv: /802./8800.4440 04687, 2018.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi- body systems with symmetric energies. arXiv preprint 111100191. 2019.
Jia Li, Yu Rong. Hong Cheng, Helen Meng, Wenbing Huang, and Junzhou Huang. Semi-supervised graph classification: A hierarchical graph perspective. In WWW, pp. 972-982, 2019.
Yuhan Li, Zhixun Li, Peisong Wang, Jia Li, Xiangguo Sun, Hong Cheng, and Jeffrey Xu Yu. A survey of graph meets large language model: Progress and future directions. arXiv preprint
mrXiv:2311.12399 2023.
Xiao Luo, Jingyang Yuan, Zijie Huang, Huiyu Jiang, Yifang Qin, Wei Ju, Ming Zhang, and Yizhou Sun. Hope: High-order graph ode for modeling interacting dynamics. In ICML. 2023.
11

---

Published as a conference paper at ICLR 2024
Michael Lutter, Christian Ritter, and Jan Peters. Deep lagrangian networks: Using physics model prior for deep learning. In ICLR, 2019.
Damian Mrowea, Chengxu Zhuang, Elias Wang. Nick Haber, Li Fei-Fei, Joshua B Tenenbaum, and Daniel LK Yamins. Flexible neural representation for physics prediction. arXiv preprini (eXiv:18066.. 08047, 2018.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Ko- mbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Alexander Norcliffe, Cristian Bodnar, Ben Day, Nikola Simidjievski, and Pietro Liò. On second order behaviour in augmented neural odes. In NeurlPS, volume 33, pp. 5911-59922 2020.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W. Battaglia. Learning mesh- based simulation with graph networks. In ICIR, 2021.
Michael Poli, Stefano Massaroli, Junyoung Park, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Graph neural ordinary differential equations. arXiv preprint arXiv:1911.07533 2019.
Michael Poli, Stefano Massaroli, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Hyper- solvers: Toward fast continuous-depth models. In NeurlPS, volume 33, pp. 21105-21117, 2020.
Yu Rong. Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenhing Huang, and Junzhou Huang. Self-supervised graph transformer large-scale molecular data. Advances in neural information processing systems, 1355991112571, 2020.
Cristopher Salvi, Maud Lemercier, and Andris Gerasimovics. Neural stochastic pdes: Resolution- invariant learming of continuous spatiotemporal dynamics. In NeurIPS, 2022.
Alvaro
Victor Bapst, Kyle Cranmer, and Peter W. Battaglia. Hamiltonian graph networks with ODE integrators. arXiv preprint arXiv: 127900, 2019.
Alvaro
Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. In ICML, pp. 8459-8468. PMLR, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural net- works. In ICML, pp. 9323-9332. PMLR, 2021.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In JCML, pp. 9377-9388. PMLR, 2021.
Bruno Siciliano, Lorenzo Sciavicco, Luigi Villani, and Giuseppe Oriolo. Roboties. Springer London, 2009, doi: 10.10077978---846288 978-1- 84 6288-642-1
URL /ttt:s://cci.org//////
Sebastian Strönisch, Maximilian Sander, Andreas Knüpfer, and Marcus Meyer. Multi gpu approach for training of graph ML models on large CFD meshes. CoRR, 0ttp22337.1359... 2023.
Xiangguo Sun, Hong Cheng, Jia Li, Bo Liu, and Jihong Guan. All in one: Multi-task prompting for graph neural networks. KDD, pp. 2120-2131, 2023.
Jianheng Tang, Jiajin Li, Ziqi Gao, and Jia Li. Rethinking graph neural networks for anomaly detection. ICML, pp. 21076-21089. PMLR, 2022.
Jianheng Tang, Fengrui Hua, Ziqi Gao, Peilin Zhao, and Jia Li. Gadbench: Revisiting and bench- marking supervised graph anomaly detection. In NeurIPS Datasets and Benchmarks Track, 2023.
Abishek Thangamuthu, Gunjan Kumar, Suresh Bishnoi, Ravinder Bhattoo, N M Anoop Krishnan, and Sayan Ranu. Unravelling the performance of physics-informet graph neural networks for dynamical systems. In NeuriPS Datasets and Benchmarks Track, 2022.
12

---

Published as a conference paper at ICLR 2024
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and cunssation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:/802.0822 2018.
Yuanqing Wang and John Chodera. Spatial attention kinetic networks with (n)-equivariance. arXiv preprint arXiv:2301.080: 00883, 2023.
Liming Wu, Zhichao Hou, Jirui Yuan, Yu Rong, and Wenbing Huang. Equivariant spatio-temporal attentive graph networks simulate physical dynamics. NeurIPS, 36, 2024.
Rui Xue, Haoyu Han, MohamadAli Torkamani, Jian Pei, and Xiaorui Liu. Lazygnn: Large-scale graph neural networks via lazy propagation. In ICML, volume 202, pp.
2023.
Aiqing Zhu, Pengzhan Jin, Beibei Zhu, and Yifa Tang. On numerical integration in neural ordinary differential equations. In ICML, pp. 27527-27547, 2022.
13

---

Published as a conference paper at ICLR 2024
A
PROOFS
Note that h is identical for both real and learned systems. Unless otherwise specified, we omit h in the following discussions.
A.1 PROOF OF PROPOSITION 3.1
According to the constraints of the initial conditions, the initial position (fo) is equuiiauivai and the initial velocity ((t)) is (3))eqqivariant and translation-invarinn Given that backbone GNN So are -euuiqruirr and translation-invariint, and 91 is -quivvvarrr,, for any translation vector b E R3, orthogonal matrix E R3x3 and E [0, At), we have
Le 90. gAgg (ta)
(ta)
+ b) = go (Aq
(ta)
+ b) + 91 (fo (ta) + b)),t')
Ag + (i ( ),t'
(17)
Ad + AGi (Jo(((( (10)), ),()
=Aq (a(lo+(') A At A go. (9) (00)
Thus, Ve go (g) (10)) is proved to be -e3)-equivariant and translation-invarian. Note that Ø2 is (3) eeguivarraa then, for qo (fo+t') we have
Le go (Ag (ta) b) Aqe (fo) + b (V 10 (Aq (to) + b),
Age (to) + G2(A 9o So 90 (10)) )[() +
(18)
Ag Aggl) + AG2 ge. (4o (10)), )
),(') + b
Aqo (to+t)) +
(to) is -q3iivuuirrriant. Note that the composition of E(3)- cquivariant functions is still (33)icquivariant, For any time t E [fo, tıl. q which is generated iteratively calling integrator 4 r.. (Eq. 11) with suitable t' E [0, At), -qquivvrrar Over- all, the approximated trajectory qo is -3)_qiuviraiiant.
Therefore, 99
(+)
0r (9o
In real-world applications, the assumption of this proposition is generally satisfied. In terms of the equivariance of backbone GNN fe. for example, if fo is EGNN whose message passing is defined by
m = de (110b7)
12, dog 10 12, h, by ass:
T0
1
Σ
( (m)))
N
JEN.
(19)
h (+++)) (1) + on (h
Σ
)
JENG
Here oy denotes Multi-Layer Perceptron (MIP) whose output is a scalar and the output of des Φn are vectors. The non-geometric features are updated via skip connections. Analogous neural ODE methods, the model parameters are shared among all iterations. By (Satorras et al., 2021), it can be easily shown that fo (eqi-varrr and causlation-invariaa
In addition, the increment functions G1 G2 of several widely used numerical integrators for motion dynamics are -qqiivaaa,,,,,,,,, ensuring that the approximated trajectory is -qquivarrrri For example, symplectic Fuler integrator computes
(++ At) (t) g(t+ At, ((++t)) = g(()) + q(() At,
(20)
where (1, = G2( v) is (3) equivariant. It is straightforward to show that
A (1++)0)) At A (4+0) Aq (t) b Aq (c+a-))) At Aq (++ + b.
(21)
This property also holds for Velocity Verlet
q (++bt) (*) + (1) At + (01)) At2 ((++ = q(1)) (g(1) + (((+a)))
(22)
2
14

---

Published as a conference paper at ICLR 2024
and Leapfrog
q ((+a()) = n0hh + + 2 A0 g((+a)) = q()) + +g(++()) A..
(23)
The proof is the same as symplectic Euler.
A.2 PROOF OF LEMMA 4.1
We first prove the uniqueness of the solution in our dynamic system. The key this proof is to convert the high-order non-linear ODE to first-order non-linear ODE. Let's define new variable Q(o) (q(() ()) and (()) as its first-order derivative with respect t. Then in terms of this new variable, the second-order non-linear ODE becomes
(9) = (F(Q )), ), ()))))),
(24)
where F(Q(1)) f((()) and G(Q())
q(0 Furthermore, let's define HTQ(!))
((Q((),((Q)))) and Eq. 24 is reformatted
Q ( = (Q)))). Q (q(()) . (1)))
(25)
which exactly first-order non- non-linear ODE with a known initial condition. Note that and are both Lipschitz continuous, thus H is Lipschitz continuous. Then, based on Picard's existence theorem (Coddington et al., 1956), the aforementioned non-linear ODE has unique solution (t)),
Vt E [fo, f1]. Subsequently, our system has a unique solution (()), ,t E [to,t]]
A.3 PROOF OF PROPOSITION 4.2
Due the uniqueness of the solution, if the realistic measurement (1:) is given, then the trajectory (+)) is fully determined the time interval [to, t1]. Under SEGNO framework, define the discrepancy between 40 (ti) and q (f1) as
d(g (11) (()))) 11 (t:) 9 - (qq(t))) - ((
(26)
where ||. represents the p-norm. According to Eq. 9, & Se is fully determined by Se. Without loss of generality, we take SEGNO with Euler integrators as an example. Given the known initial position and velocity, along with neural ODE update scheme in SEGNO, we can convert Eq. 26 to the following form
d(qb) (t1) . (()) (19b)) (tti)lp
to (k+(+))
(q (1(o+ka1t) fd
(m)(m))dm)) (fotkkt)) At fo (totkAt) (
At?)
toth
lotkar
((ottt))))) (( (lott))))) a thotha))) fo (otkttt))))
(27)
where D(-) denotes the discrepancy function. Considering that there are no sudden changes in velocity or acceleration for smooth trajectory, with sufficiently small At, we can approximate D(-) as follows
(o (1q d (l++AA1) [J((9) ((((00-))))))
ko-)))) do
(10 (10 ) (f(a (o+kki))) - fo(g) (ookklt) At
(28)
(f(q (ottitt)) fo (0 (lotiae)) At?
Note that Lemma 4.1 guarantees the uniqueness of target trajectory as well as its acceleration. There- fore, according to the aforementioned analysis, SEGNO tends to let fo (9) (too+kA))) approximate the unique accelaration ((q(to+)))) via minimizing position loss. Then, under the assumption of widely used universal approximation theorem (Hornik et al., 1989), there exists a fe- such that for (1 h) (t) h). 17 [..
15

---

Published as a conference paper at ICLR 2024
A.4 PROOF OF THEOREM 4.3
The proof sketch is as follows: We first revisit the core definitions pertaining to neural ODE in SEGNO and introduce its variant with Euler integrator, then derive the bound for first-order approx- imation eπrоr [19e~ gl and finally extend the results to the second-order case I|fe fl to finish the proof.
A.4.1 EULER INTEGRATOR IN SEGNO
As mentioned in Eq. 10, the Euler integrator @ Ptt,go as mentioned in Eq. 20, approaches @ar.go via
go 4g (1))) go (t) fo (2) )
(29)
D go (9) (0) = q + go g ((+2() Att.
Considering Eq. 11 and 12, is straightforward to reframe our learning objective the context of a neural ODE:
Lwain
Σ
11
1
Σ
(r. 112
(30)
DEDmm
BEDii
A.4.2 APPROXIMATION ERROR OF fo
In our dynamical system, go and are entirely determined by fo and respectively. Thus, we first establish the boundedness of I|ge gl then demonstrate the approximation etror of Ilfo flloo
Lemma A.1. For (ta) E R as the initial position of trajectory. Ta, Tb, T, > and k E Z+ given ODE solver that is T compositions of an Euler Integrator Dst.go with At = T/T, set of points on exact trajectory associated with time increment interval [kAt, T + &&t]
T+kAt = (ov kAt
ST
(31)
and denote
-
( ΦT. Bl ".))"
(32)
KAt
complex balls centered at E +T++Ai with radius ra + Then, there exist constants To and that and suppose that g and go are analytic and bounded by m within B(VT+kát ,a +r), the union of
kAt
depends on ra/m, Tb/ m, T, and + Patgg such that, if0 < T To, VI [lo + kAI, T + k&I).
(g Ill S CmAt
T- KA:
(33)
E
where e the base of the natural logarithm.
Proof. This result can be directly derived from Theorem 3.1, 3.2 and Corollary 3.3 in (Zhu et al., 2022) via replacing the integrator from Runge- Kutta integrator with Euler integrator since Euler integrator has been proven satisfy the prerequisites of theorems in Appendix B.2 in (Zhu al., 2022).
Since B(VT TT C. B(V r1). per our assumption, then and are both analytic and bounded by my in ((V V T1). To utlize Lemma A.I, we k and T t to. With suitable rı and mi, we have T < To and, vt [tu, t.l.
e (t)) ()) S Atm
e CT.
(34)
where 4J = + p "...)"" - Φr or .n)) and C1 a control constant.
To establish the connection between g and f. we only focus on the first time step At instead of the entire time interval T. Recall that g is obtained by integrating f. we have
Vst. = (t++ = (t) +
f(
(35)
16

---

Published as a conference paper at ICLR 2024
the flow map for the first-order derivative of the exact trajectory in a single step. As defined in Eq. 29, its corresponding Euler integrator has the form
(o+ V Dr. (o (t go qo
go (g) " At.
(36)
Note that B(VO r2) B( r2). and fo are both analytic and bounded by m2 in B(V r2). In lemma A.I, we substitute gbt.go by V Var.go.fo and set T At with = 1 and k = 0, then the corresponding loss becomes
1 At
sup
V Vanigoooo ( mn
At
QEB(V (2)
(37)
V War.go.fo fq")) )
Al
1
gg f -
(a)
Al
where q" denote the point that maximizes the Eq. 37 and q()) is another trajectory with @(te)= q¹ because q" Vo may not holds. For both q(t) and q², the actual position and estimation from SEGNO at to + At have the form
(lo+A1)= q((n)
9(q
to
gg ( (t)) q0(() go( (10)
) At (o
(38)
(1o (b+ =0(² b00l 16)
At + lo (fo) At))
q
Then, with suitable At and r1, g ( (1(()) go (0 +() and go (g66 (o) are bounded by my since g(to+ 1t)
Att) (tt+ E B(V,, 71). Given such, the aforementioned loss is transformed into
4g
CT,
At [llga (4)
1 (1)) go (g(lo+ + 1lgg (Ht+t()))) (to+t))))))
+ 1190 ((lott))) و(q (t+t1) + [l (1++)))) g (((
(39)
1
(6m1 +im,Att -
e-160)
At
Subsequently, given that At < T < To. we have, Vt E [o. to + At).
()))()())) fl) CtmzAt
0,11
CtmeAt
1
(6m1 + CimiAt+
e (b),,
e
E
At
(40)
S At
CT
e
6m1 Af + Cimı),
At
CT At
$ O(A+ +
To extend the boundness up to f, we can easily utilize Lemma A.I with different -,..., (Eq. 34) to repeat the above derivation, and obtain, Vt. E [to + kAt, to + (Is + 1)At].
,
THAA
(L - ((q S O(AL+
FRA A
(41)
where
C++At kAt
=110 &r B(VT-RAr, .r)"" Therefore, Vt E [1,
T+kA
(1) ()) < sup 000,,.,711
O KAt
Al
(42)
C2T
s0la+ +
At
where 2JT =11 or -rrllls(33r and it concludes the proof.
17

---

Published as a conference paper at ICLR 2024
A.5 PROOF OF COROLLARY 4.4
A.5.1 LOCAL TRUNCATION ERROR
We use Taylor series expansion to approximate the trajectory at time & + At as
q (((+A)) ②0)) ((g + f((")) + )
(43)
2
Then the local truncation error C++ Vt [t], equals to
A+ (1t0) q (0) (( AL lo )
= 11a + 9(9 ) + 1
2 A( +(
80) -99 "0000)))4 -
1 =11 - - + +0(At))ll
= 11 1
=12(((000) - -o(( 1
2 (1) + O(At))
(44)
1 |1 2 (f((( (o( + 11 fo(g A ((At))
V3A/2
1
(15(g()) So(g + 1
-2 12 0(419)
A
2
V3At
2CT
O
+
At ) 0(
2
Here the last inequality is due to Theorem 4.3. If the learned model achieves near-perfect approx- imation of the true trajectory, resulting in a significantly diminished loss LTT. Then, Vt (1,,
CAA S ((AP))
(45)
A.5.2 GLOBAL TRUNCATION ERROR
Considering & E (4o, ta] and k = 1..., T 1, the boundness of the global truncation error can be derived in a recursive way via
11+(+1)44 ((((+((())))) 9o
((+(++))41)1 112
=Ig(++((41)0)) do (1+A))) go (++RA() ) fola) fo(g )
((Q())
(++RDI) (++(())) -q (4RA()) 9(g (+tRE))))
49
+ )1 + - -ool At
To(g )
(46)
At (4h
+ g ((+AA)) At go ((+A())) Addl2
< ELARD + 1+((-1)A+ + [l (++At))) (ttha())))))) - folg (tt(k))kk)))2222 A12 )|2 + [Il9o(q go ( ((tk2e)) Ill + 119(g (+H(A))))))) go (( (+tkn)))l2)
Note that go and fo satisfy the Lipschitz continuity, we have
(++kee)) - go g ((+ba))) 112 I Ly Lyllg
4
(++
||2 = Lg LgllEttadel2.
(47)
[l (1+h00) fe
th (.* l|la S Ly ((+hk01))
(e+ ((+) ||| L)llE++krella
where Lg and Ly denote Lipschitz constant for go and fo respectively. Given that the learned model achieves a near-perfect approximation of the true trajectory, by Lemma A. and Eq. 34, obtain
[ - go ( (++A00) I|2 S v3(Cmm + e C/ $ O(At).
(48)
Thus, the global truncation error in Eq. 46 is transformed into
& (1 T.g Al + ) + O(A())
(49)
18

---

Published as a conference paper at ICLR 2024
Figure 5: Example trajectories of 5-body charged system. From left to right, the number of positive and negative charges are 1, 3, 0.
X
Figure 6: Example trajectories of 5-body gravity system. The system formed by gravitational force can either aggregate (i.e., left) or disperse (i.e., right).
Note that the global truncation error for the first step EAAD where h: 1, is exactly the local truncation error ét+ mentioned in Eq. 44. Given that k S we can derive, VI. E [4,,4] and
T,,
EHAA S (1 + Lg A/ 1 0(A()) +.... (1+ 1,
T
(50)
50(TA() 0 At
A/ 0(1)).
B IMPLEMENTATION DETAILS
B.1 MORE DETAILS ON SIMULATED N-BODY SYSTEMS
N-body charged system We use the same N-body charged system code* with previous work (Satorras al., 2021; Brandstetter et al., 2021). They inherit the 2D implementation of (Kipf' et al., 2018) and extend it to dimensions. System trajectories are generated in 0.001 timestep and unbounded with virtual boxes. The initial location is sampled from a Gaussian distribution (mean = O, standard deviation a = 0.5), and the initial velocity is a random vector of norm 0.5. Ac- cording the charge types, three types of systems exist where the difference between the number of positive and negative charges are 1, 3, and O. Example trajectories of these three types of systems provided in Figure 5.
N-body gravity system The code of gravitational N-body systems is provided by (Brandstetter et al., 2021). They implement it under the same framework as the above charged N-body systems. System trajectories are generated in 0.001 timestep. using gravitational interaction and no boundary conditions. Particle positions are initialized from a unit Gaussian, particle velocities are initialized with a norm equal to one, random direction, and partiele mass is set The system formed
19

---

Published as a conference paper at ICLR 2024
Carbohydrates
Supramolecules
Lipids
Nucleic acids
Number
pairwise correlations
Proteins
Availability of ab initio reference data
200
400
600
Atoms
Figure Molecular structures of MD22 dataset, which is borrowed from Fig.1 original pa- per (Chmiela et al., 2023).
by gravitational force can either aggregate or disperse. Example trajectories of these two types of systems provided in Figure 6.
Hyperparameter We empirically find that the following hyperparameters generally work well, and use them across most experimental evaluations: Adam optimizer with learning rate 0.001, the number of epochs 500, hidden dim 64, weight decay 1 10-12 12, and layer number 4. We set the iteration time of SEGNO to 8. The representation degrees of Sran)-rra and TFN are set to 3 and 2. The number of training, validation, and testing sets are 3000, 2000, and 2000, respectively.
B.2 MORE DETAILS ON MD22
Molecules and Hy Hyperparameter The molecular structures of MD22 are displayed in Figure 7, which is borrowed from their paper (Chmiela et al., 2023). We use the following hyperparameters across all experimental evaluations: Adam optimizer with learning rate 0.001, the number of epochs 5000 with early stopping 100, hidden dim 64, weight decay 1 10-12 and layer number 4. The iteration time of SEGNO is searched from 4 to 6. The representation degrees of SE(3) trans-rrrn and TFN are set to 2. Due to the limited memory, the batch size of TFN is set to 10. The number of training. validation, and testing sets are 500, 2000, and 2000, respectively. The threshold for graph construction is set to 2.5 for all molecules.
B.3 MORE DETAILS ON MOTION CAPTURE
Hyperparameter We the following hyperparameters across all experimental evaluations: Adam optimizer with learning rate 0.001, the number of epochs 3000, hidden dim 64, weight decay 1 10 12, 12 and layer number The iteration time SEGNO set to 4. We adopt random split strategy introduced by Huang et al. (2022) where train/validation/test data contains 200/600/600 frame pairs.
C ADDITIONAL EXPERIMENTS
C.1 ACCURACY OF LEARNED LATENT TRAJECTORY
interesting to see how models learn the latent trajectory between the input and output states. Accordingly, train models 1000ts datasets and make the test shorter time steps
20

---

Published as a conference paper at ICLR 2024
Table 5: The generalization from long-term short-term. All models are trained 1000ts and test on 250/500/750/1000 ts. Mcan squared error (x10 and the standard deviation are reported. Results averaged across runs.
Charged
Gravity
Method
250
500 is
750
1000
Avg
250
500
750 ts
1000
Avg
GNN
73,40
21.79
12. 26.861a1
0.826
29.72
181 9=06
90.33a155
30.66
0,746
75.93
GDE
92.65
43.94
12.20
0.6522m05
37.36
136.0g x
56.801as
12/21
0.588
51.39
EGNN
6.756
3.816144
3.668
0.568
3.702
7.146
29:70
9.712:3w
0.382
11.89
GMN
10.
10.92
4.51811.
0.512.mus
6.598
1.430
9.540.i21
5.7301sm
0.349
5.762
SEGNN
21.78+
52.74+116
34.13-149
0.342-ux4
27.25
10.58
49.63+m
25.82+22u
0.448
21.62
SEGNO
18
0.312
0.360=00w
0.309
0.292
0.064
0.128
0.176ee0
0.210
0.145
by performing SEGNO the smaller step with the same ratio. For the baselines, we treat the forward timestep of each hidden layer as the same and extract their object position information as the prediction. Table reports the mean and standard deviation of each setting, From Table 5 we can observe that:
Clearly, SEGNO outperforms all other baselines across all settings by large margin. Notably, when there is a lack of supervised signals at 250/500/750ts, the performance of all other baselines decreases significantly. By contrast, SEGNO achieves similar results as in 1000ts, demonstrating its robust generalization to short-term prediction.
Another interesting point that SEGNO's emro exhibits distinct trend compared to other baselines. While the errors of other baselines significantly increase with decreasing time steps, SEGNO achieves even smaller errors with shorter time steps. This observation justifies our theo- retical results that the error is bound by the chosen timestep.
Additionally, the standard deviation of SEGNO is much smaller than that other baselines, indi- cating the numerical stability of SEGNO. This result further confirms our theoretical finding that SEGNO can obtain better latent trajectory between two discrete states.
C.2 COMPARISON WITH GNS
We conducted additional evaluations of GNS and SEGNO-avg.. which are learned by minimizing average acceleration, two N-body systems. The results of these evaluations are presented in Table We can observe that SEGNO outperforms both GNS and SEGNO-avg. In all cases, show- ing that training second-order neural odes on position loss outperforms training models on average acceleration.
Table Comparsion (x10-2) between SEGNO and GNS on simulated N-body systems. Results averaged across 3 runs.
Charged 1500 is
Gravity 1500 is
Method
1000 ts
2000 Es
1000 ts
2000
GNS SEGNO-avg. SEGNO
3.245=
11.689 uasi
31.632+ a2re
4.204 uost
17.095
50.275 02211
2.04
10.145 10.145
24.244
1.431
19. 18.4888...m
54.370+
0.433_am
1,858 200e
4.285 now
0.338 L 6037
1.362_
4.017.1 9rw
C.3 ROLLOUT COMPARISON ON N-BODY SYSTEMS
We evaluate the generalizability of models for rollout 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 8 depicts the mean squared error of all methods on two datasets. We can observe that all baselines experience numerical explosion due to error accumulation during the rollout process, leading quick drop in prediction performance. In contrast, SEGNO demon- strates an order-of- oidder-f-mmanitudee error improvement over other baselines. This numerical stability can be attributed the Neural ODE framework for modeling position and velocity.
21

---

Published as a conference paper at ICLR 2024
Charged
Gravity
MSE
MSE
5
10
20
25
30
35
0
10
15
20
25
30
35
The number of rollout step
The number of rollout step
GNN
GMN
SECNN SECNO
CDE
EGNY
Figure 8: Mean squared error of rollout. Each rollout step is equal to 1000 ts. All models are trained on 1000ts.
C.4 COMPARISON WITH PHYSICS-INFORMER GNNS
Following the experimental settings (Thangamuthu et al., 2022), we thoroughly compare the rollout position. energy, and momentum errors of SEGNO with those of GNODE (Bishnoi et al., 2022) and HGNN (Thangamuthu et al., 2022) on N-body systems. GNODE and HGNN are trained by minimizing acceleration and Hamilton respectively. The target timestep is set to 100ts. The results are demonstrated Table 7. We can observe that SEGNO consistently outperforms GNODE and HGNN in rollout errors. Although SEGNO is only trained on the system positions, it achieves comparable performance on energy and momentum errors with GNODE, which further validates its effectiveness.
Table 7: Rollout position, energy, and momentum errors. All models are trained on 100us.Results averaged across 3 runs.
Charged
Gravity 1000 is
Method
100
500
1000
1500
2000
100
500
1500
2000
HGNN
2.16
11.47
21.95
30.40
36.89
3.26
16.78
30.83
40.07
45.94
Rollout
GNODE
0.36
5.41
14.05
21.61
27.39
0.32
6.42
17.53
26.17
32.16
SEGNO
0.06
4.26
12.26
19.42
25.08
0.04
5.62
16.84
25.71
31.82
HGNN
1.94
7.49
11.36
14,10
15.53
2.78
10.73
16.10
19.50
22.34
Energy
GNODE
0.82
5.82
10.87
14.07
15.78
0.68
7.46
14.40
18.39
20.85
SEGNO
0.79
5.74
9.93
12.66
14.05
0.11
7.38
13.58
16.94
18,89
HGNN
10.12
33.37
45.98
$1.67
54,38
6.99
25.65
39.12
46.08
49,
Momemtum
GNODE
4.17
28.65
43.97
50.84
54.06
2.35
21.43
36.93
44.92
49.19
SEGNO
2.17
28.39
43.81
50.46
53.54
0.75
21.40
36.88
44.65
48.42
C.5 COMPARISON WITH ADVANCED FORCE FIELD PREDICTION MODELS
For more comprehensive evaluation, add empirical comparison on the MD22 dataset with NeuqlP (Batzner et al., 2022) and Allegro (Musaelian et al., 2023)?. Both models take atom po- sitions and numbers as input and optimize the errors between geometric outputs (generally treated as predicted forces) and true positions. We try our best to tune the models and the results are demonstrated in Table 8. It can be observed both methods do not perform well in predicting the time-dependent evolution of molecular dynamics, which is mainly attributed to the lack of dynamic information.
hhtps:/github.co//iinggrooo
22

---

Published as a conference paper at ICLR 2024
Table 8:
Comparison with equivariant GNNs in the domain of advanced force field prediction.
Mean squared error (x x)10 on MD22 dataset averaged across runs are reported.
Molecule
NequiP
Allgero
SEGNO
Ac-Ala3-NHMe
12.060 0.220
11.785+ 0.193
0.418 +: 0.088
DHA
13.275 0.164
13.126
0.370 0.020
Stachyose
11.375 0.104
11.164+ 0.135
0.440 0.0%
AT-AT
9.178=
9.032
0.548 0.066
AT-AT-CG-CG
8.959 0,066
8.866 0.125
0.394 0.033
Buckyball Catcher
5,418
5.331 0007
0.199 1: 0.020
Double-walled Nanotube
3.852 0,077
3.794 0.022
0.225 :: 0.008
C.6 RUNTIME COMPARISON ON N-BODY SYSTEMS
We evaluate the running time of each model on N-body systems with Telsa T4 GPU and report the average forward time in seconds for 100 samples. The results are listed Table We can observe that SEGNO's forward time (0.0227s) remains competitive compared to the best baseline SEGNN (0.0315s), indicating its efficiency.
Table 9: Forward time in seconds for a batch size of 100 samples on a Tesla T4 GPU.
Linear
GNN
GDE
TFN
-E(3)-Tr:
RF
EGNN
GMN
SEGNN
SEGNO
0.0002
0.0064
0.0088
0.0440
0.2661
0.0052
0.0126
0.0137
0.0315
0.0277
C.7 GENERALIZATION CAPABILITY TO LARGE SYSTEMS
In this part, we evaluate the generalizability of models for larger system sizes. Specifically, we train all models on 5-body gravity systems and then test them on 10-body and 20-body gravity systems. Table 10 shows their results. We compare three strong baselines, i.e., EGNN, GMN, and SEGNN. The results show that the performance of all baselines significantly drops when testing on larger systems. In contrast, SEGNO still demonstrates a marked improvement over other baselines, especially 20-body systems.
C.8 MORE RESULTS ON CMU MOTION CAPTURE
This section illustrates more visualizations of GMN and SEGNO on modeling object motions. From Figure we can observe that SEGNO is able track the ground-true truth trajectories accurately, which is consistent with the performance in Table 4.
23

---

Published as a conference paper at ICLR 2024
30
30
25
10
10
115.15.25.155
15
40.35
15
-3
15
10
35. 25
15 -5
15
-
-5
-3
Figure 9: Example snapshots of Motion Capture with 50 time steps. Top: Prediction of GMN. Bottom: Prediction of SEGNO. Ground truth in red, and prediction in blue.
The results on the larger systems. Mean squared error and the standard deviation are reported. Results are averaged across 3 runs.
Table 10:
Method 10-Body
20-Body
EGNN
0.566 u.316 0.16
1,985
GMN
0.716 0.314
1.323 + 0.400
SEGNN
0.333 0.036
3,937 2.121
SEGNO
0.152 0.021
0.850 :: 0.055
24


Reviews:
[{"review_id": "kqJsJpeHqA", "reviewer": "Reviewer_GrTs", "summary": "In this work, authors combine neural ODE with equivariant GNN to propose a second order equivariant GNODE (SEGNO). Authors prove that the SEGNO maintains equivariance. The study prove that the discrepancy between the trajectory learned by SEGNO and the ground truth is bounded. Results show that the SEGNO provides relatively superior performance on various dynamical systems.", "questions": "1. Authors may rephrase the text ``accurately approximate'' in the proof of theorem 4.2?\n2. Theorem 4.1 is a standard proof for any dynamical system and is not specifically related to SEGNO. It is not clear why this is a property of SEGNO and discussed under section 4. Authors may comment.\n3. Proposition 4.2 is trivially dependent on the universal approximation theorem and is not needed to be stated as a proposition. Essentially, it proves that a NN can approximate the function $f(q,h)$ by minimizing the error between $q_{\\theta}$ predicted by a NN and the ground truth $q$. This how every NN is trained. Authors may explain why this warrants a separate proof. \n4. The present work focusses on combining equivariant GNNs with GNODEs. There are several ways of considering baselines. One set of baselines would be different equivariant architectures directly employing a data-driven approach to predict the dynamics. This is how some of the baselines are included in the present work. However, extensive literature shows that physics-informed inductive bias such GNODE, Hamiltonian GNN, and Lagrangian GNN can provide superior performance. Comparison with other such physics-informed GNNs such as Hamiltonian and Lagrangian GNN would have been insightful. \n5. Another set of baselines would be GNODE with different equivariant backbone architectures. This has not been extensively evaluated in the present work. In the present work, authors employ GMN as the backbone. While it makes sense for the motion capture video, it is not clear why GMN would be a good backbone for the MD simulation dataset. Backbones such as NequIP would be better for such cases. Authors are encouraged to evaluate this.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Learning dynamical systems with physics-based inductive bias is an important problem to be studied.\n2. Combining equivariance with NODE is an interesting approach. \n3. Several interesting and varied datasets are considered for empirical examples including n-body systems, MD22, and motion capture datasets. \n4. The results show that the SEGNO outperforms baselines for several tasks including molecular dynamics, and motion capture.", "weaknesses": "1. The idea of combining GNODE with an equivariant GNN is fairly straightforward. In fact, there are several works on GNODE that employ a second-order bias (Gruver et al., Bishnoi et al., both cited in the paper). While the backbone GNN can have some effect on the GNODE (as shown in the work: Thangamuthu, A., Kumar, G., Bishnoi, S., Bhattoo, R., Krishnan, N.M. and Ranu, S., 2022. Unravelling the performance of physics-informed graph neural networks for dynamical systems. Advances in Neural Information Processing Systems, 35, pp.3691-3702.), there are no restrictions on employing an equivariant GNN backbone with GNODE. Thus, the novelty of the work may be limited.\n\n2. In the present work, GMN is used as the backbone GNN along with the GNODE based inductive bias. The choice of the specific backbone may be appropriate for some datasets such as motion capture. However, the architecture does not perform well for molecular datasets or even n-body systems. Evaluation on additional backbone architectures is missing.\n\n3. The baselines considered are simple and not necessarily SOTA. For the MD22, there are several equivariant GNNs such as DimeNet, NequIP, Allegro, MACE, BOTNet, Equiformer, etc. (and many more newer architectures) that have shown SOTA performance. It is not clear why such advanced baselines or graph architectures are not used.\n\n4. The metrics used for comparison does not necessarily capture the complete picture regarding the dynamics. There are several metrics such as energy violation error, momentum error, rollout error etc. in the literature. Moreover, for atomic structures additional metrics such as JS divergence of the radial distribution functions capture how well the predicted the structure is similar to the ground truth. Such metrics have not been included in the present work leading to an incomplete evaluation. \n\n5. There are several works on articulated rigid body in the literature including physics-informed GNNs such as Hamiltonian NNs, and Lagrangian graph neural networks for articulated body. Comparison with such baselines are not included in the work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qCvgybi8Lx", "reviewer": "Reviewer_rsLd", "summary": "The authors present an innovative graph neural network framework that introduces second-order continuity into the modeling process. Unlike traditional discrete updates in other graph neural networks, they employ neural ordinary differential equations to capture continuous trajectories. This approach preserves the equivariance property of the underlying graph neural networks. Their model is rigorously evaluated across a range of standard test systems, including time evolution of N-body systems, molecular trajectories, and human motion trajectories. In all experiments, their model consistently outperform existing methods.", "questions": "- Can SEGNO be included in a Markov Chain Monte Carlo (MCMC) framework aimed at accelerating Molecular Dynamics simulations, a common application for predicting future trajectory states? Typically, such applications require Metropolis-Hastings corrections to asymptotically sample from the equilibrium distribution.\n\n- Graph neural networks are known for their transferability, as demonstrated by other methods on the QM9 dataset. While acknowledging the difference in tasks, do the authors envision the potential for SEGNO to be employed in a similarly transferable manner?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper is well written and structured\n- The authors introduce a novel framework for graph neural networks that incorporates second-order continuity by utilizing neural ODEs. This innovation enables more accurate modeling of physical systems, as many are governed by second-order laws.\n- Their approach stands out by allowing the modeling of continuous trajectories, in contrast to existing graph neural networks that rely on discrete updates.\n- They rigorously prove the claims they make, such as the preservation of equivariance and the derivation error bounds. \n- Their model undergoes extensive testing on various time series prediction tasks, consistently outperforming previous models and showcasing its superior performance.\n-  The insightful ablation study underscores the significance of incorporating second-order dynamics and utilizing a continuous model instead of a discrete one.\n\nIn conclusion, I consider this paper to be a significant and valuable addition to the community, and I recommend its acceptance.", "weaknesses": "- Run times for the various methods are not reported. Given the neural ODE nature of SEGNO, it's reasonable to assume that this model might have longer execution times compared to some of its competitors. It would be advantageous to include these run times in the appendix, with brief mentions in the main text for clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DAGw1W9bZr", "reviewer": "Reviewer_mXJL", "summary": "## Summary\n\nThe paper investigates the integration of physical inductive biases, equivariant Graph Neural Networks (GNNs), and Neural Ordinary Differential Equations (ODEs) to improve model prediction accuracy. The authors validate their approach on 3 datasets: N-body, molecular dynamics, and captured human motion data. The technical contributions are sound though not groundbreaking. Given the paper's strengths and incremental contribution, I recommend a borderline acceptance.\n\n## Detailed Comments\n\n### Contributions\n\n1. The paper shows that the error bound between the predictions and the ground truth outperforms previous methods focused on average accelerations.\n2. It establishes that when the incremental function is either E(3) or O(3), the combined Neural ODE and equivariant GNN also possess these properties.\n3. It shows enough empirical validation and ablation studies for later scholars to save their time and make proper choices.\n\n### Comments and Suggestions\n\na) Lack of Visual Illustration: Apart from human motion data, the paper does not provide visual illustrations for any of the datasets. This makes it challenging to grasp the scope of the method's applicability, especially for readers not specialized in fields like molecular dynamics or N-body systems.\n\nb) Figure 4 Inconsistency: The first column in Figure 4 is hard to interpret due to inconsistent color coding of initial conditions and predictions across different columns. It would be helpful to maintain color consistency.\n\nc) Unclear experiment in Figure 2: It is unclear which experiment or dataset Figure 2 refers to, leading to possible confusion for the readers.\n\nd) Inconsistent Sectioning for Human Motion Experiment: Unlike the other experiments, the human motion experiment is not provided a standalone section, making it confusing for readers.\n\ne) Future Direction on Scalability: The paper could benefit from discussing scalability issues for large graphs. Since all examples in the paper involve graphs with fewer than 1,000 nodes, it would be interesting to explore the method's applicability to larger networks. Some relevant references include:\n\n- [MultiScale MeshGraphNets](https://arxiv.org/abs/2210.00612)\n- [Multi-GPU Approach for Training of Graph ML Models on large CFD Meshes](https://arxiv.org/abs/2307.13592)\n- [Efficient Learning of Mesh-Based Physical Simulation with BSMS-GNN](https://arxiv.org/abs/2210.02573)\n- [LazyGNN: Large-Scale Graph Neural Networks via Lazy Propagation](https://arxiv.org/abs/2302.01503)\n\n## Conclusion\n\nThe paper presents a sound technical approach integrating physical inductive biases, equivariant GNNs, and higher-order Neural ODEs to improve model prediction accuracy. It's a sound but not groundbreaking paper. Therefore, I recommend a borderline acceptance for this paper.", "questions": "See above", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "See above", "weaknesses": "See above", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)