E NHANCING THE I NDUCTIVE B IASES OF G RAPH N EURAL ODE FOR M ODELING P HYSICAL S YSTEMS
Anonymous authors
Paper under double-blind review
A BSTRACT
Neural networks with physics-based inductive biases such as Lagrangian neural networks (L NN ), and Hamiltonian neural networks (H NN ) learn the dynamics of physical systems by encoding strong inductive biases.
Alternatively, Neural
ODEs with appropriate inductive biases have also been shown to give similar performances.
However, these models, when applied to particle based systems, are transductive in nature and hence, do not generalize to large system sizes.
In this paper, we present a graph-based neural ODE, G NODE , to learn the time evolution of dynamical systems.
Further, we carefully analyse the role of different inductive biases on the performance of G NODE .
We show that, similar to L NN and H NN , encoding the constraints explicitly can significantly improve the training efficiency and performance of G NODE significantly.
Our experiments also assess the value of additional inductive biases, such as Newton’s third law, on the final performance of the model.
We demonstrate that inducing these biases can enhance the performance of model by orders of magnitude in terms of both energy violation and rollout error.
Interestingly, we observe that the G NODE trained with the most effective inductive biases, namely MCG NODE , outperforms the graph versions of
L NN and H NN , namely, Lagrangian graph networks (L GN ) and Hamiltonian graph networks (H GN ) in terms of energy violation error by ∼ 4 orders of magnitude for a pendulum system, and ∼ 2 orders of magnitude for spring systems.
These results suggest that N ODE -based systems can give competitive performances with energy conserving neural networks by employing appropriate inductive biases.
I NTRODUCTION AND RELATED WORKS
Learning the dynamics of physical systems is a challenging problem that has relevance in several areas of science and engineering such as astronomy (motion of planetary systems), biology (movement of cells), physics (molecular dynamics) and engineering (mechanics, robotics)(LaValle, 2006; Goldstein, 2011).
The dynamics of a system is typically expressed as a differential equation, solutions of which may require the knowledge of abstract quantities such as force, energy, and drag(LaValle, 2006; Zhong et al., 2020; Sanchez-Gonzalez et al., 2020).
From an experimental perspective, the real observable for a physical system is its trajectory represented by the position and velocities of the constituent particles.
Thus, learning the abstract quantities required to solve the equation, directly from the trajectory, can extremely simplify the problem of learning the dynamics (Finzi et al., 2020).
Infusing physical laws as prior has been shown to improve the learning in terms of additional properties such as energy conservation, and symplectic nature (Karniadakis et al., 2021;
Lutter et al., 2019;
Liu et al., 2021).
To this extent, three broad approaches have been proposed, namely, Lagrangian neural networks (L NN ) (Cranmer et al., 2020a; Finzi et al., 2020; Lutter et al., 2019), Hamiltonian neural networks (H NN )
(Sanchez-Gonzalez et al., 2019; Greydanus et al., 2019; Zhong et al., 2020; 2021), and neural ODE (N ODE ) (Chen et al., 2018;
Gruver et al., 2021).
The learning efficiency of
L NN s and H NN s is shown to enhance significantly by employing explicit constraints (Finzi et al., 2020) and their inherent structure Zhong et al.
In addition, it has been shown that the superior performance of H NN s and L NN s is mainly due to their second order bias, and not due to their symplectic or energy conserving bias (Gruver et al., 2021).
More specifically, an H NN with separable potential ( V ( q ) ) and kinetic ( T ( q, ˙ q ) ) energies is equivalent to a second order N ODE of the form ¨ q = ˆ F ( q, ˙ q ) .
Thus, a N ODE with a second order bias can give similar performances to that of an H NN (Gruver et al., 2021).
Recently, a Bayesian-symbolic approach, which assumes the knowledge of Newtonian mechanics along with symbolic features, have been used to learn the dynamics of physical systems Xu et al.
An important limitation pervading these models is their transductive nature; they need to be trained on each system separately before simulating their dynamics.
For instance, a N ODE trained on 5pendulum can work only for 5 -pendulum.
This approach is referred to as transductive.
In contrast, an inductive model, such as G NODE , when trained on 5 -pendulum “learns” the underlying function governing the dynamics at a node and edge level.
Hence, a G NODE , once trained, can perform inference on any system size that are significantly smaller or larger than the ground truth as we demonstrate later in § 4.
This inductive ability is enabled through Graph neural networks (Scarselli et al., 2008) due to their inherent topology-aware setting, which allows the learned system to naturally generalize.
Although the idea of graph-based modeling have been suggested for physical systems (Cranmer et al., 2020b; Greydanus et al., 2019), the inductive biases induced due to different graph structures and their consequences on the dynamics remain poorly explored. In this context, there are a few recent works that aim to develop better representations of physical systems through equivariant graph neural networks Satorras et al. ; Han et al. ; Huang et al.
However, the present work focuses on the inductive biases for physical systems rather than focusing on better representation learning. Some of the proposed inductive biases may also be a natural consequence of forcing equivariance. An an example, enforcing invariance to the Lagrangian w.r.t. space translation and rotation also enforces Newton’s third law, which we achieve through an inductive bias (§ 3).
In the context of generalizability to unseen systems, there exists work on few-shot based on meta-learning (Lee et al., 2021).
Meta-learning assumes the availability of limited amount of training data to adapt to an unseen system (task).
This is different from our objective where we perform zero-shot generalizability, i.e., inference on an unseen system without any further training.
To summarize, our key contributions are as follows:
Topology-aware modeling , where the physical system is modeled using a graph-based N ODE (G NODE ), enabling zero-shot generalizability to unseen system sizes.
Decoupling the dynamics and constraints , where the equation of motion of the G NODE is decoupled into the individual terms such as forces, Coriolis-like term and explicit constraints.
Decoupling the internal and body forces , where the internal forces due to the interaction among particles are decoupled from the forces due to external fields.
Newton’s third law , where the forces of interacting particles are enforced to be equal and opposite.
We theoretically prove that this inductive bias enforces the conservation of linear momentum, in the absence of an external field, for the predicted trajectory exactly.
We analyze the role of these biases on n -pendulum and spring systems.
Interestingly, depending on the nature of the system, we show that some these biases can either significantly improve the performance or have marginal effects.
We also show that the final model with the appropriate inductive biases significantly outperform the simple version of G NODE with no additional inductive biases and even the graph versions of L NN and H NN .
B ACKGROUND ON DYNAMICAL SYSTEMS
A dynamical system can be defined by ordinary differential equations (ODEs) of the form, ¨ q =
F ( q, ˙ q, t ) , where F represents the dynamics of the system.
Thus, the time evolution of the system can be obtained by integrating this equation of motion.
In Lagrangian mechanics, the dynamics of a system is derived based on the Lagrangian of the system, which is defined as (Goldstein, 2011):
[Equation]
Once the Lagrangian of the system is computed, the dynamics can be obtained using the EulerLagrange (EL) equation as ddt ( ∇ ˙ q L ) =
∇ q L .
Note that the EL equation is equivalent to the
D’Alembert’s principle written in terms of the force as (cid:80) ni
=1 ( N i − m i ¨ q i ) · δq i = 0 , where δq i represents a virtual displacement consistent with the constraints of the system.
Thus, an alternative to the EL equation can be written in terms of the forces and other components as (LaValle, 2006;
Murray et al., 2017)
[Equation]
where M ( q ) represents the mass matrix, C ( q, ˙ q ) the Coriolis-like forces, N ( q ) the conservative forces, Υ( ˙ q ) represents the non-conservative forces such as drag or friction, and Π represents any external forces.
We note that Lagrangian dynamics is traditionally applied in systems with conservative forces, i.e., where a Lagrangian is well-defined.
However, Eq. 2 is generic and can be applied to model any dynamical system irrespective of whether the system has a well-defined Lagrangian or not.
Examples of system where a well-defined Lagrangian may not exist include that of a colloidal system undergoing brownian motion.
The constraints of the system are given by A ( q ) ˙ q = 0 , where A ( q ) ∈ R kD represents k constraints.
Thus, A represents the non-normalized basis of constraint forces given by A = ∇ q ( ϕ ) , where ϕ is the matrix corresponding to the constraints of the system and λ is the Lagrange multiplier corresponding to the relative magnitude of the constraint forces.
Note that M, C, and N are directly derivable from the Lagrangian of the system L as M = ∇ ˙ q ˙ q L, C = ∇ q ˙ q L, and N = ∇ q L .
Thus, decoupling the dynamic equation simplifies the representation of the dynamics F with enhanced interpretability.
To learn the dynamics, Eq. 2 can be represented as
[Equation]
To solve for λ , the constraint equation can be differentiated with respect to time to obtain A ¨ q + ˙ A ˙ q = 0 .
Substituting for ¨ q from Eq.
3 and solving for λ , we get
[Equation]
Accordingly, ¨ q can be obtained as
[Equation]
I NDUCTIVE BIASES FOR NEURAL ODE
To learn the dynamical systems, N ODE s parameterize the dynamics F ( q t , ˙ q t ) using a neural network and learn the approximate function ˆ F ( q t , ˙ q t ) by minimizing the loss between the predicted and actual trajectories, that is, L =
||
q t +1 − ˆ q t
+1 || .
Thus, a N ODE essentially learns a cumulative ˆ
F , that consists of the contribution of all the terms contributing the dynamics.
In contrast, in an
L NN , the learned L is used in the EL equation to obtain the dynamics.
Similarly, in L NN with explicit constraints, the learned Lagrangian is used to obtain the individual terms in the Eq.
5 by differentiating the L NN , which is then substituted to obtain the acceleration.
However, this approach is transductive in nature.
To overcome this challenge, we present a graph-based version of the
N ODE namely, G NODE , which models the physical system as a graph.
We show, that G NODE can generalize to arbitrary system size after learning the dynamics ˆ F .
3. G RAPH NEURAL ODE (G NODE ) FOR PHYSICAL SYSTEMS
In this section, we describe the architecture of G NODE .
The graph topology is used to learn the approximate dynamics ˆ F by minimizing the loss on the predicted positions.
Graph structure.
An n -particle system is represented as a undirected graph G =
{V , E} , where nodes represent particles and edges represent the connections or interactions between them.
For instance, in pendulum or springs system, the nodes correspond to bobs or balls, respectively, and the edges correspond to the bars or springs, respectively.
Input features.
Each node is characterized by the features of the particles themselves, namely, the particle type ( t ), position ( q i ), and velocity ( ˙ q i ).
The type distinguishes particles of differing characteristics, for instance, balls or bobs with different masses.
Further, each edge is represented by the edge features w ij =
( q i − q j ) , which represents the relative displacement of the nodes connected by the given edge.
Neural architecture:
A detailed description of the various neural layers within G NODE is provided in App. A. Here, we provide a summary of the various neural layers within G NODE .
Fig.
8 in the appendix provides a pictorial description.
In the pre-processing layer, we construct a dense vector representation for each node v i and edge e ij using MLPs.
In many cases, internal forces in a system that govern the dynamics are closely dependent on the topology of the structure.
We model this through a deep, message-passing G NN .
Let z i
= h
Li denote the final representation constructed by the G NN over each node.
These representations are passed through another MLP to predict the acceleration of each node.
Trajectory prediction and training.
Based on the predicted ¨ q , the positions and velocities are predicted using the velocity Verlet integration.
The loss function of G NODE is computed by using the predicted and actual accelerations at timesteps 2 , 3 , . .
, T in a trajectory T , which is then backpropagated to train the MLPs.
Specifically, the loss function is as follows.
[Equation]
Here, (ˆ¨ q T ,t i ) is the predicted acceleration for the i th particle in trajectory T at time t and ¨ q T ,t i is the true acceleration.
T denotes a trajectory from T , the set of training trajectories.
Note that the accelerations are computed directly from the ground truth trajectory using the Verlet algorithm as: [Equation]
[Equation]
Since the integration of the equations of motion for the predicted trajectory is also performed using the same algorithm as: q ( t + ∆ t )
= 2 q ( t ) − q ( t − ∆ t ) + ¨ q (∆ t ) 2 , this method is equivalent to training from trajectory/positions.
It should be noted that the training approach presented here may lead to learning the dynamics as dictated by the Verlet integrator.
Thus, the learned dynamics may not represent the “true” dynamics of the system, but one that is optimal for the Verlet integrator 1 .
G NODE learns the dynamics directly without decoupling the constraints and other components governing the dynamics from the equation ¨ q = ˆ F ( q, ˙ q, t ) .
Next, we analyze the effect of providing the constraints explicitly in the form of inductive biases and enhance the performance of G NODE . In the subsequent sections, we iteratively stack inductive biases on top of G NODE . Thus, the final proposed architecture, MCG NODE , is the eventual model that considers all biases. 3. D ECOUPLING THE DYNAMICS AND CONSTRAINTS (CG NODE ).
To this extent, while keeping the graph architecture exactly the same, we use Eq.
5 to compute the acceleration of the system.
In this case, the output of G NODE is N i =
squareplus ( MLP V ( z i )) instead of ¨ q i , where N i is the conservative force on particle i (refer Eq. 2).
Then, the acceleration is predicted using Eq.
5 with explicit constraints.
This approach, termed as constrained G NODE (CG NODE ), enables us to isolate the effect of imposing constraints explicitly on the learning, and the performance of G NODE .
It is worth noting that for particle systems in Cartesian coordinates, the mass matrix is a diagonal one, with the diagonal entries, m ii s, representing the mass of the i th particle .
Thus, the m ii s are learned as trainable parameters.
As a consequence of this design, the graph directly predicts N i of each particle.
Note that in the case of systems with drag, the graph would contain cumulative of N i and Υ i . Examples of such systems are demonstrated later (see App. E).
3. D ECOUPLING THE INTERNAL AND BODY FORCES (CDG NODE )
In a physical system, internal forces arise due to interactions between constituent particles.
On the contrary, the body forces are due to the interaction of each particle with an external field such as gravity or electromagnetic field.
While internal forces are closely related to the topology of the system, the effect of external fields on each particle is independent of the topology of the structure.
In addition, while the internal forces are a function of the relative displacement between the particles, the external forces due to fields depend on the actual position of the particle.
To decouple the internal and body forces, we modify the graph architecture of the G NODE .
Specifically, the node inputs are divided into local and global features.
Only the local features are involved in the message passing.
The global features are passed through an MLP to obtain an embedding, which is concatenated with the node embedding obtained after message passing.
This updated embedding is passed through an MLP to obtain the particle level force, N i .
Here, the local node features are the particle types and the global features are position and velocity.
Note that in the case of systems with drag, the Υ i can be learned using a separate node MLP using particle type and velocity as global input features.
This architecture is termed as constrained and decoupled G NODE (CDG NODE ).
The architecture of CDG NODE is provided in Fig.
3. N EWTON ’ S THIRD LAW (MCG NODE )
According to Newton’s third law, “every action has an equal and opposite reaction” .
In particular, the internal forces on each particle exerted by the neighboring particles should be equal and opposite so that the total system is in equilibrium.
That is, f ij = − f ji , where f ij represents the force on particle i due to its neighbor j .
Note that this is a stronger requirement than the EL equation itself, i.e., ddt ( ∇ ˙ q i L ) = ∇ q i L .
EL equation enforces that the rate of change of momentum of each particle is equal to the total force on the particle, N i =
(cid:80) n i j =1 f ij , where n i is the number of neighbors of i — it does not enforce any restrictions on the individual component forces, f ij from the neighbors that add to this force.
Thus, although the total force on each particle may be learned by the G NODE , the individual components may not be learned correctly.
We show that G NODE with an additional inductive bias enforcing Newton’s third law, namely momentum conserving G NODE (MCG NODE ), conserves the linear momentum of the learned dynamics exactly.
Theorem
In the absence of an external field, MCG NODE exactly conserves the momentum.
Proof of the theorem is provided in the Appendix B. To enforce this bias, we modify the graph architecture as follows.
In particular, the output of the graph after message passing is taken from the edge embedding instead of taking from the node embedding.
The final edge embedding obtained after message passing is passed through an MLP to obtain the force f ij .
Then the nodal force is computed as the sum of all the forces from the edges incident on the node.
Note that while doing the sum, the symmetry condition f ij = − f ji is enforced.
Further, the global forces coming from any external field is added separately to obtain the total nodal forces.
Finally, the parameterized masses are learned as the masses in a particle system in Cartesian coordinates are diagonal in nature.
E MPIRICAL E VALUATION
In this section we benchmark G NODE and the impact of the inductive biases. Our implementation is available at https://anonymous .
4open .
science/r/graph neural ODE-8B3D. Details of the simulation environment is provided in App.
H
4. E XPERIMENTAL SETUP
• Baselines: To compare the performance of G NODE and its associated versions MCG NODE , CDG NODE , and CG NODE , we compare with graph-based physics-informed neural networks, namely, (1) Lagrangian graph network (L GN )
(Cranmer et al., 2020a), and (2) Hamiltonian graph network (Sanchez-Gonzalez et al., 2019).
Since the exact architecture and details of the graph architecture are not provided in the original works, a full graph architecture (Cranmer et al., 2020b; Sanchez-Gonzalez et al., 2020) is used for modeling L GN and H GN 2 .
Since Cartesian coordinates are used, the parameterized masses are learned as a diagonal matrix as performed in (Finzi et al., 2020).
We note that while in G NODE , the predicted acceleration is the direct output of the graph neural network, in L GN (or H GN )
the output is the Lagrangian (or Hamiltonian) and further differentiation of the Lagrangian (or Hamiltonian) provides acceleration. These additional computational steps makes L GN or H GN computationally expensive. We empirically establish this in Fig. ?? . This result is consistent with earlier results that notes neural ODEs (N ODE ) as being more efficient and generalizable than Hamiltonian neural networks if appropriate inductive biases are injected Gruver et al.
Thus, we also compare with N ODE (Gruver et al., 2021). • Datasets and systems: To compare the performance G NODE s with different inductive biases and with the baseline models, we selected systems on which N ODE s are shown to give good performance (Zhong et al., 2021; Finzi et al., 2020), viz , n -pendulums and springs, where n = (3 , 4 , 5) .
Further, to evaluate zero-shot generalizability of G NODE to large-scale unseen systems, we simulate 5, 20, and 50-link spring systems, and 5-, 10-, and 20-link pendulum systems.
The details of the experimental systems are given in Appendix.
The training data is generated for systems without drag for each particle.
The timestep used for the forward simulation of the pendulum system is 10 − 5 s with the data collected every 1000 timesteps and for the spring system is 10 − 3 s with the data collected every 100 timesteps.
The detailed data-generation procedure is given in App.
I. • Evaluation Metric:
Following the work of (Finzi et al., 2020), we evaluate performance by computing the following three error metrics, namely, (1) relative error in the trajectory, defined as the rollout error, (RE(t)) , (2) Energy violation error (EE(t)) , (3) relative error in momentum conservation, (ME(t)) , given by:
[Equation]
Note that all the variables with a hat, for example ˆ x , represent the predicted values based on the trained model and the variables without hat, that is x , represent the ground truth.
• Model architecture and training setup: For all versions of G NODE , L GN , and H GN , the MLPs are two-layers deep.
We use 10000 data points generated from 100 trajectories to train all the models.
This dataset is divided randomly in 75:25 ratio as training and validation set.
Detailed training procedure and hyper-parameters are provided in App.
I. All models were trained till the decrease in loss saturates to less than 0.001 over 100 epochs.
The model performance is evaluated on a forward trajectory, a task it was not explicitly trained for, of 10 s in the case of pendulum and 20 s in the case of spring.
Note that this trajectory is 2-3 orders of magnitude larger than the training trajectories from which the training data has been sampled.
The dynamics of n -body system is known to be chaotic for n ≥ 2 .
Hence, all the results are averaged over trajectories generated from 100 different initial conditions.
4. E VALUATION OF INDUCTIVE BIASES IN G NODE
• Performance of (G NODE ).
Fig. 2 shows the performance of G NODE on pendulum and spring systems of sizes 3 to 5.
The G NODE for pendulum is trained on a 3-pendulum system and that of spring is trained on a 5-spring system.
Hence, the 4- and 5-pendulum systems and 3- and 4-spring systems are unseen , and consequently, allows us to evaluate zero-shot generalizability.
We observe that energy violation and rollout errors of G NODE for the forward simulations are comparable, both in terms of the value and trend with respect to time, across all of the systems.
Furthermore, when we compare with N ODE , which has been trained on each system separately, while the errors are comparable in 5-pendulum, G NODE is better in 5-spring.
It is worth highlighting that N ODE has an unfair advantage on 5-pendulum since it is trained and tested on the same system, while 5-pendulum is unseen for G NODE (trained on 3-pendulum).
This confirms the generalizability of G NODE to simulate the dynamics of unseen systems of arbitrary sizes.
More experiments comparing N ODE with G NODE to further highlight the benefit of topology-aware modeling are provided in App.
D. • G NODE with explicit constraints (CG NODE ).
Next, we analyze the effect of incorporating the constraints explicitly in G NODE .
To this extent, we train the 3-pendulum system using CG NODE model with acceleration computed using Equation 5.
Further, we evaluate the performance of the model on 4- and 5-pendulum systems.
Note that the spring system is an unconstrained system and hence CG NODE is equivalent to G NODE for this system.
We observe that (see Fig. 2) CG NODE exhibits ∼ 4 orders of magnitude superior performance both in terms of energy violation and rollout error.
This confirms that providing explicit constraints significantly improves both the learning efficiency and the predictive dynamics of G NODE models where explicit constraints are present , for example, pendulum systems, chains or rigid bodies.
• Decoupling internal and body forces (CDG NODE ).
Now, we analyze the effect of decoupling local and global features in the G NODE by employing the CDG NODE architecture.
It is worth noting that the decoupling the local and global features makes the G NODE model parsimonious with only the relative displacement and node type given as input features for message passing to the graph architecture.
In other words, the number of features give to the graph architecture is lower in CDG NODE in G NODE .
Interestingly, we observe that the CDG NODE performs better than CG NODE and G NODE by ∼ 1 order of magnitude for pendulum and spring systems, respectively.
• Newton’s third law (MCG NODE ).
Finally, we empirically evaluate the effect of enforcing Newton’s third law, that is, f ij = − f ji on the graph using the graph architecture namely, MCG NODE .
Figure 2 shows the performance of MCG NODE for both pendulum and spring systems.
Note that it was theoretically demonstrated that this inductive bias enforces the exact conservation of momentum, in the absence of an external field.
Figure 4 shows the ME (see Equation 8) of MCG NODE in comparison to N ODE , G NODE , CG NODE , CDG NODE , and MCG NODE for pendulum and spring systems.
We observe that in the spring system the ME of MCG NODE is zero.
In the case of a pendulum, although lower than N ODE , G NODE , and CG NODE , the ME of MCG NODE is comparable to that of CDG NODE .
Correspondingly, MCG NODE exhibits improved performance for the spring system in comparison to other models.
However, in pendulum system, we observe that the performance of MCG NODE is similar to that of CDG NODE .
This is because in spring systems, the internal forces between the balls govern the dynamics of the system, whereas in pendulum the dynamics is governed mainly by the force exerted on the bobs by the external gravitational field.
Interestingly, we observe that the energy violation error of MCG NODE remains fairly constant, while this error diverges for other models.
This confirms that a strong enforcement of momentum conservation in turn ensures that the energy violation error stays constant over time, ensuring a stable long-term equilibrium dynamics of the system.
In addition, MCG NODE exhibits significantly improved training efficiency than N ODE (see App.F).
Finally, we also evaluate the ability of MCG NODE to learn the dynamics of systems with drag and external forces (see App. E).
We observe that MCG NODE clearly outperforms
N ODE in this case as well. MCG NODE will, thus, be particularly useful for the statistical analysis of chaotic systems where multiple divergent trajectories correspond to the same energy state.
4. C OMPARISON WITH BASELINES
Now, we compare the G NODE that exhibited the best performance, namely, MCG NODE with the baselines namely, L GN and H GN .
Note that both L GN and H GN are energy conserving graph neural networks, which are expected to give superior performance in predicting the trajectories of dynamical systems (Sanchez-Gonzalez et al., 2019).
Especially, due to their energy conserving nature, these models are expected to have lower energy violation error than other models, where energy conservation bias is not explicitly enforced.
Figure 3 shows the energy violation and rollout error of G NODE , L GN , and H GN trained on 5-link system and tested on 3, 4, 5-link pendulum and spring systems.
Interestingly, we observe that MCG NODE outperforms both L GN and H GN in terms of rollout error and energy error for both spring and pendulum systems.
Specifically, in case of pendulum systems, we observe that the energy violation and rollout errors are ∼ 5 orders of magnitude lower in the case of MCG NODE in comparison to L GN and H GN .
Further, we compare the ME of
MCG NODE with L GN and H GN (see Fig.
4 right palette).
We observe in both pendulum and spring systems that the ME of MCG NODE is significantly lower than that of L GN and H GN .
Thus, in the case of spring systems, although the performance of H GN and MCG NODE are comparable for both RE and EE , we observe that MCG NODE exhibits zero ME , while in H GN and
L GN ME diverges.
The superior performance of MCG NODE could be attributed to the carefully constructed inductive biases, namely, (i) explicit constraints—improves data efficiency and reduces the error in pendulum systems, (ii) decoupling internal and body forces—leads to superior performance in pendulum systems which act under gravity, (iii) momentum conservation—leads to superior performance in spring systems by enforcing an additional conservation law strongly.
We note that these inductive biases may also be adapted suitably for LGN and HGN and thereby improve their performance.
Thus, the contribution of our work is not the neural architecture in isolation, but the core design principles that are generic enough to be abstracted beyond neural ODEs.
4. G ENERALIZABILITY
Now, we focus on the zero-shot generalizability of G NODE , L GN , and H GN to large systems that are unseen by the model.
To this extent, we model
5-, 10-, and 20-link pendulum systems and 5-, 20-, 50-link spring systems.
Note that all the models are trained on the 5-link systems.
In congruence with the comparison earlier, we observe that MCG NODE performs better than both L GN and H GN , when tested on systems with large number of links.
However, we observe that the variation in the error represented by the 95% confidence interval increases with increase in the number of particles.
This suggests that the increase in the number of particles essentially results in an increased noise in the prediction of the forces.
However, error in the MCG NODE is significantly lower than the error in L GN and H GN for both pendulum and spring systems, as confirmed by the 95% confidence interval, which could again be attributed to the three inductive biases infused in MCG NODE , that are absent in L GN and H GN .
This suggests that MCG NODE can generalize to larger system sizes with increased accuracy.
To demonstrate the applicability of MCG NODE to learn the dynamics of complex systems, we consider a challenging problem, that is, inferring the dynamics of an elastically deformable body.
Specifically, a 5 × 5 × 5 solid cube, discretized intoparticles, is used for simulating the ground truth.
3D solid system is simulated using the peridynamics framework.
The system is compressed isotropically and then released to simulate the dynamics of the solid, that is, the contraction and expansion in 3D.
We observe that MCG NODE learns the dynamics of the system and perform inference much superior to G NODE (see Fig. 6).
It should also be noted that although the momentum error in MCG NODE seems high, this is due to the fact that the actual momentum in these systems are significantly low.
This is demonstrated by plotting the absolute momentum error in these systems which in MCG NODE is close to zero, while in G NODE it explodes (see Fig. 20).
These results suggest that the MCG NODE could be applied to much more complex systems such as solids or other particle-based systems where momentum conservation is equally important as energy conservation. 4. R OBUSTNESS TO N OISE Finally, we show the robustness to noise of the MCG NODE in comparison to G NODE , and additional baselines, namely, H GN and L GN . Fig.
7 shows the performance of these models on 5-spring and 5-pendulum systems for varying noise values. Note that the noise is given in terms of % standard deviation of the distribution of the original output values. We observe that MCG NODE is robust against noise in comparison to other models.  C ONCLUSIONS
the original output values. In this work, we carefully analyze various inductive biases that can affect the performance of a G NODE .
Specifically, we analyze the effect of a graph architecture, explicit constraints, decoupling global and local features that makes the model parsimonious, and Newton’s third law that enforces a stronger constraint on the conservation of momentum.
We observe that inducing appropriate inductive biases results in an enhancement in the performance of G NODE by ∼ 2-3 orders magnitude for spring systems and ∼ 5-6 orders magnitude for pendulum systems in terms of energy violation error.
Further, MCG NODE outperforms the graph versions of both Lagrangian and Hamiltonian neural networks, namely, L GN and H GN .
Altogether, the present work suggests that the performance of G NODE -based models can be enhanced and even made superior to energy conserving neural networks by inducing the appropriate inductive biases.
Limitations and Future Works: The present work focuses on particle-based systems such as pendulum and spring systems.
While these simple systems can be used to evaluate model performances, real-life are more complex in nature with rigid body motions, contacts, and collisions.
Extending G NODE to address these problems can enable modeling of realistic systems.
A more complex set of systems involve plastically deformable systems such as steel rod, and viscoelastic materials where the deformation depends on the strain rate of materials.
Further, instead of explicitly providing the constraints, learning the constraints directly from the trajectory can enhance the application of the models to unseen systems.
Note that the present work is limited to physical systems.
Similar approach could be employed in other dynamical systems as well.
From the neural modeling perspective, the G NODE is essentially a graph convolutional network.
Currently, messages from all neighbors are treated as equally important.
Using attention over each message passing head may allow us to better differentiate neighbors based on the strength of their interactions.
We hope to pursue these questions in the future.
A D ETAILS OF G NODE
The architecture of G NODE is depicted in Fig.
G NODE constitutes of the following neural layers.
Pre-Processing.
In the pre-processing layer, we construct a dense vector representation for each node v i and edge e ij using MLP em as:
[Equation]
squareplus is an activation function.
Note that the MLP em corresponding to the node and edge embedding functions are parameterized with different weights.
Here, for the sake of brevity, we simply mention them as MLP em .
Acceleration prediction.
In many cases, internal forces in a system that govern the dynamics are closely dependent on the topology of the structure.
To capture this information, we employ multiple layers of message-passing between the nodes and edges.
In the l th layer of message passing, the node embedding is updated as:
[Equation]
where, N i = { v j ∈ V | e ij ∈ E} are the neighbors of v i .
W l V is a layer-specific learnable weight matrix.
h lij represents the embedding of incoming edge e ij on v i in the l th layer, which is computed as follows.
[Equation]
Similar to W l V , W l E is a layer-specific learnable weight matrix specific to the edge set.
The message passing is performed over L layers, where L is a hyper-parameter.
The final node and edge representations in the L th layer are denoted as z i
=
h Li and z ij = h
Lij respectively.
Finally, the acceleration of the particle ¨ q i is predicted as:
[Equation]
Why Verlet integrator: The canonical coordinates employed in Newtonian mechanics are the position ( q ) and velocity ( ˙ q ) (or momentum) as these quantities are the direct observables.
Verlet algorithm is obtained by substracting the Taylor series expansion of r ( t + ∆ t ) with r ( t − ∆ t ) .
The expression thus obtained is time-reversible and has error to the order of ( O (∆ t 4 ) ).
Due to the timereversible nature, Verlet algorithm is energy conserving (symplectic), and hence performs better than other traditional integrators such as RK4 in predicting the long term trajectory of dynamical systems without drift.
B P ROOF OF T HEOREM
P ROOF .
Consider a system without any external field where the dynamics is governed by internal forces, for instance, a n -spring system. The conservation of linear momentum for this system implies that
[Equation] ˆ˙ q i [Equation] = 0 , where ˆ˙ q i represents the predicted velocity, ˆ m ii represents the learned mass of the particle i , and ˙
q i represents the ground truth velocity.
Assume the error in the rate of change of momentum of the predicted trajectory to be ˆ ε t .
To prove that MCG NODE conserves the momentum exactly, it is sufficient to prove that ˆ ε t = 0 .
Consider,
[Equation]
Note that in equation 3, all terms except N i are either known or are zero.
Thus, without loss of generality, (cid:80) ni =1 N i = 0 , in the absence of an external field.
Hence,
[Equation]
(cid:88)
[Equation]
f
[Equation]
f [Equation] The inductive bias f ij = − f ji implies that ( f ij + f ji )= 0 .
Thus, n
[Equation]
□
C P ERFORMANCE OF L GN AND H GN WHEN IMPLEMENT ON THE G NN OF G NODE .
In Fig.
9, we compare the performance of L GN and H GN implemented on the full graph network (Cranmer et al., 2020b; Sanchez-Gonzalez et al., 2020) against the versions implemented on the G NN used in G NODE .
As visible, the full graph network version provides superior performance.
D Z ERO - SHOT GENERALIZABILITY : N ODE VS G NODE VS MCG NODE
To showcase zero-shot generalizability, we compare the performance of N ODE , G NODE , MCG NODE in Figure 10.
In the case of N ODE , the system is trained and evaluated on each of the pendulum and spring systems separately.
In the case of G NODE and MCG NODE , the training is performed on 5-spring and 3-pendulum system and evaluated on all other systems.
First, we observe that the performance of N ODE and G NODE are comparable for both rollout and energy error.
This suggests that the performance of G NODE , with no additional inductive bias, is comparable to that of N ODE with the additional advantage that G NODE can generalize to any system size.
Now, we compare the performance of N ODE with MCG NODE , which has all the additional inductive biases discussed earlier.
We observe that MCG NODE performs significantly better than N ODE , as expected.
It is worth noting that while the inductive bias of explicit constraints can be infused in N ODE in a similar fashion as in G NODE , other inductive biases such as decoupling the global and local features and Newton’s third law are unique to the graph structure.
These inductive biases cannot be infused in a N ODE structure.
This, in turn, demonstrates the superiority of G NODE which allows additional inductive biases due to the explicit representation of the topology through the graph structure. E D RAG AND EXTERNAL FORCE
Here, we evaluate the ability of MCG NODE to learn the dynamics of systems subjected to drag linear in velocity (drag force given by − 0 .
1 ˙ q ) and an external force.
Fig. 11 shows the performance of MCG NODE and N ODE on a 5-spring system with linear drag and external force. We observe that the MCG NODE exhibits superior performance in both the cases. F T RAINING EFFICIENCY
Data Efficiency: Here, we compare the training efficiency of MCG NODE with the N ODE and H GN .
Despite our best efforts, L GN was unable to learn the dynamics from the small amount of data provided and was exploding up on forward simulation leading to NaN values.
Similar behavior was observed on H GN with low values data points.
We focus on the best graph-based architecture and compare the performance with the baseline N ODE 5-spring and 3-pendulum systems.
Figureshows the performance of MCG NODE in comparison to N ODE for the dataset sizes of 100, 500, 1000, 5000, and 10000 in terms of the geometric mean of the energy error and rollout error.
We observe that the learning of MCG NODE is significantly simplified in comparison to that of N ODE — MCG NODE exhibits orders of magnitude better performance in terms of both energy and rollout error for both pendulum and spring systems. Further, we observe that MCG NODE performs orders of magnitude better than H GN . Computational Efficiency: We further compare the computation times of the training and inference phases on G NODE , MCG NODE , H GN and L GN .
Tables 1-2 present the results.
As visible,
MCG NODE , is the fastest on average. The difference with L GN and H GN is the highest in spring systems. MCG NODE benefits from the inductive biases as well as the fact that the neural architecture directly predicts the force (acceleration is computed from this) on each particle. In contrast,
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
H GN and L GN predict the Hamiltonian and Langrangian, which needs to be further differentiated to obtain acceleration.
L GN is the slowest since it requires double differentiation.
G E RRORS ON I NTERNAL AND E XTERNAL F ORCES
In this section, we report the errors in the predicted internal and external forces on each particle of a system.
Figs.
13-16 presents the results on G NODE , MCG NODE , L GN and H GN .
As visible, the predicted force (represented using dashed lines) overlaps closely with the actual forces (solid line) in MCG NODE (Fig. 14).
While the error is also low in G NODE (Fig. 13), it is inferior to MCG NODE showcasing the usefulness of the inductive biases.
We also note that the error is the highest in L GN (Fig. 15).
H E XPERIMENTAL SYSTEMS H. S IMULATION E NVIRONMENT
All the simulations and training were carried out in the JAX environment (Schoenholz & Cubuk, 2020; Bradbury et al., 2020).
The graph architecture was developed using the jraph package (Godwin* et al., 2020).
The experiments were conducted on a machine with Apple M1 chip having 8GB
RAM and running MacOS Monterey.
In an n -pendulum system, n -point masses, representing the bobs, are connected by rigid bars which are not deformable.
These bars, thus, impose a distance constraint between two point masses as
[Equation]
(16) where, l i represents the length of the bar connecting the ( i − 1)
th and i th mass.
This constraint can be differentiated to write in the form of a Pfaffian constraint as
[Equation]
Note that such constraint can be obtained for each of the n masses considered to obtain the A ( q ) .
The mass of each of the particles is m i and the force on each bob due to the acceleration due to gravity in m i g .
This force will be compensated by the constraint forces, which in turn governs the motion of the pendulum.
For L GN , the Lagrangian of this system has contributions from potential energy due to gravity and kinetic energy.
Thus, the Lagrangian can be written as
[Equation]
And, for H GN , the Hamiltonian is given by
[Equation]
19) where g represents the acceleration due to gravity in the y direction.
H. n - SPRING SYSTEM
In this system, n -point masses are connected by elastic springs that deform linearly with extension or compression.
Note that similar to a pendulum setup, each mass m i is connected only to two masses m i − 1 and m i +1 through springs so that all the masses form a closed connection.
Thus, the force experienced by each bob i due to a spring connecting i to its neighbor j is f ij = 1 / 2 k ( || q
i − q j ||− r i ) .
For L GN , the Lagrangian of this system is given by
[Equation]
And,for H GN , the Hamiltonian by:
[Equation]
where r 0 and k represent the undeformed length and the stiffness, respectively, of the spring.
I I MPLEMENTATION DETAILS I. D ATASET GENERATION
Software packages: numpy-1.22.1, jax-0.3.0, jax-md-0.1.20, jaxlib-0.3.0, jraph-0.0.2.dev
Hardware: Chip: Apple M1, Total Number of Cores: 8 (4 performance and 4 efficiency), Memory: 8 GB, System Firmware Version: 7459.101.3, OS Loader Version: 7459.101.
For NODE, all versions of G NODE , LGN: All the datasets are generated using the known Lagrangian of the pendulum and spring systems, along with the constraints, as described in Section H. For each system, we create the training data by performing forward simulations with 100 random initial conditions.
For the pendulum system, a timestep of 10 − 5 s is used to integrate the equations of motion, while for the spring system, a timestep of 10 − 3 s is used.
The velocity-Verlet algorithm is used to integrate the equations of motion due to its ability to conserve the energy in long trajectory integration.
While for HGN, datasets are generated using Hamiltonian mechanics.
Runge-Kutta integrator is used to integrate the equations of motion due to the first order nature of the Hamiltonian equations in contrast to the second order nature of L GN , N ODE , and all versions of G NODE .
From the 100 simulations for pendulum and spring system, obtained from the rollout starting from 100 random initial conditions, 100 data points are extracted per simulation, resulting in a total of 10000 data points.
Data points were collected every 1000 and 100 timesteps for the pendulum and spring systems, respectively.
Thus, each training trajectory of the spring and pendulum systems are 10 s and 1 s long, respectively.
Here, we do not train from the trajectory.
Rather, we randomly sample different states from the training set to predict the acceleration.
I. A RCHITECTURE
For N ODE , we use a fully connected feed-forward neural network with 2 hidden layers, each having 16 hidden units with a square-plus activation function.
For all versions of G NODE , graph architecture is used with each MLP consisting of two layers with 5 hidden units.
One round of message passing is performed for both pendulum and spring.
For L GN and H GN , full graph architecture is used with each MLPs consist of two layers with 16 hidden units and one round of message passing were performed.
The hyper-parameters were optimized based on the good practices of the model development.
[Equation]
VS VELOCITY
Figure 17 shows the loss curve with relative position and relative velocity as edge features.
We observe that when relative velocity is provided as the edge feature, model struggles to learn and saturates quickly.
I. T RAINING DETAILS
The training dataset is divided in 75:25 ratio randomly, where the 75% is used for training and 25% is used as the validation set.
Further, the trained models are tested on its ability to predict the correct trajectory, a task it was not trained on.
Specifically, the pendulum systems are tested for 1 s , that is 10 5 timesteps, and spring systems for 20 s , that is 2 × 10 4 timesteps on 100 different trajectories created from random initial conditions.
All models are trained for 10000 epochs.
A learning rate of 10 − 3 was used with the Adam optimizer for the training.
• NODE Parameter Value
Hidden layer neurons (MLP) Number of hidden layers (MLP) Activation function squareplus Optimizer ADAM Learning rate 1 .
e −Batch size • G NODE , CGNODE, CDGNODE, MCG NODE
{+Parameter Value Node embedding dimension Edge embedding dimension Hidden layer neurons (MLP) Number of hidden layers (MLP) Activation function squareplus Number of layers of message passing Optimizer ADAM Learning rate 1 . 
+}
e −Batch size • LGN, HGN
[Table]