Hamiltonian Graph Networks with ODE Integrators
Alvaro Sanchez-Gonzalez Victor Bapst Kyle Cranmer Peter Battaglia


[
    [
    December 30, 2023
=====================




We introduce an approach for imposing physically informed inductive biases in learned simulation models. We combine graph networks with a differentiable ordinary differential equation integrator as a mechanism for predicting future states, and a Hamiltonian as an internal representation. We find that our approach outperforms baselines without these biases in terms of predictive accuracy, energy accuracy, and zero-shot generalization to time-step sizes and integrator orders not experienced during training. This advances the state-of-the-art of learned simulation, and in principle is applicable beyond physical domains.





§ INTRODUCTION

Learning to simulate complex physical processes has been shown to benefit from rich inductive biases about the structure of the system's state, such as modeling particles and their interactions via nodes and edges in a graph <cit.>, and optimizing internal representations to predict known physical quantities <cit.>. Here we explore another class of physically informed inductive biases: that a system's dynamics can be modeled as an ordinary differential equation (ODE) and formulated as Hamiltonian mechanics. These have also been individually studied recently by <cit.> (ODE bias), and <cit.> (Hamiltonian bias), but here we combine them and incorporate both into graph-network-based neural architectures <cit.> for learning simulation.






Our experiments show that our trained models have greater predictive accuracy than baselines, as well as better energy conservation (via the Hamiltonian inductive bias) and stronger generalization to novel time-steps and integrator orders (via the ODE integrator inductive bias). These inductive biases also carry trade-offs: with lower order integrators and coarse time-steps, both our learned Hamiltonian model as well as a model which uses the true Hamiltonian derived from physics struggle, because the Hamiltonian puts hard constraints on how the higher order structure of the dynamics must be modeled, whereas the less constrained models can learn ways of approximating such structure.




§ BACKGROUND




  
Graph networks

We represent a particle system as a graph whose nodes correspond to particles, and with edges connecting all nodes to each other. All of our models use a graph network (GN) <cit.>, which operates on graphs G = (𝐮, V, E) with global features, 𝐮, and variable numbers of nodes, V, and edges, E. Here we focus on GNs which can compute per-node outputs, V' = GN_V(G), and global outputs 𝐮' = GN_𝐮(G). Preliminary experiments show our GN model outperformed an MLP-based approach by several orders of magnitude, consistently with previous work <cit.>.




  
Numerical integrators for solving ODEs

Given a first-order ODE and initial conditions, numerical integration can be used to approximate solutions to the initial value problem,

    𝐲≡𝐲(t)
      ,  ẏ≡d𝐲/dt = f_ẏ(t, 𝐲)
      ,  𝐲(t_0) = 𝐲_0   .

The family of Runge-Kutta (RK) integrators are fully differentiable and can generate trajectories via iterations of the form, 𝐲_n+1 = RK(t_n, Δ t, 𝐲_n, f_ẏ), where Δ t = t_n+1-t_n. 

The lowest order RK integrator (Euler integrator, RK1) iterates as, 𝐲_n+1 = 𝐲_n + Δ t· f_ẏ(t_n,𝐲_n). Higher order RK integrators produce more accurate trajectories by composing multiple queries to the function f_ẏ. Here we explore first- through fourth-order RK integrators: RK1, RK2, RK3, and RK4 (See Supp. Sec. <ref> for results with symplectic integrators).




  
Hamiltonian mechanics

In Hamiltonian mechanics, the Hamiltonian, ℋ(𝐪, 𝐩), is a function of the canonical position, 𝐪, and momentum, 𝐩, coordinates, and usually corresponds to the energy of the system
[Our methods are also compatible with time-dependent Hamiltonians, ℋ(t, 𝐪, 𝐩). However, since our dataset does not require time dependency, for simplicity we will omit t everywhere from here on.]. 

The dynamics of the system follow Hamilton's equations, two first-order ODEs (analogous to Eq. <ref> with 𝐲 = (𝐪,𝐩)),

    𝐪̇≡d𝐪/dt = ∂ℋ/∂𝐩  ,  ṗ≡d𝐩/dt = -∂ℋ/∂𝐪  ⇒  
    (𝐪̇,ṗ)≡d(𝐪,𝐩)/dt = f_𝐪̇, ṗ(𝐪, 𝐩)   .





§ MODELS






  
Delta graph network () 
Our main baseline, a “delta graph network” (), replicates previous approaches for learning to simulate <cit.>, and directly predicts changes to 𝐪 and 𝐩,

    (𝐪, 𝐩)_n+1   = 
    (𝐪, 𝐩)_n + 
    (Δ𝐪, Δ𝐩)_n  , where   (Δ𝐪, Δ𝐩)_n←GN_V(Δ t, 𝐪_n, 𝐩_n,  𝐜; ϕ)    .

The 𝐜 are static parameters (masses, spring constants) of the system, and ϕ are the neural network parameters. The GN's signature matches the integrator's, and so is analogous to learning an integrator.




  
ODE graph network ()

Our “ODE graph network” () imposes an ODE integrator as an inductive bias in the GN, by assuming that the dynamics of (𝐪, 𝐩) follow a first-order ODE (Eq. <ref>). We train a neural network that learns the ODE, that is, learns to produce the time derivatives (𝐪̇, ṗ) (which are independent from Δ t). We again use the node output of a GN to model the per-particle time derivatives, and provide the GN, together with the initial conditions and Δ t, to an RK integrator, 

    (𝐪, 𝐩)_n+1   = RK(Δ t, (𝐪, 𝐩)_n, f^OGN_𝐪̇, 𝐩̇) 
    
    f^OGN_𝐪̇, 𝐩̇(𝐪, 𝐩)    ≡GN_V(𝐪, 𝐩, 𝐜; ϕ) = (𝐪̇,ṗ)    .

The f_𝐪̇, 𝐩̇ is a function that the integrator can use to operate on any (𝐪, 𝐩) and query more than once. For fixed Δ t and when using the RK1/Euler integrator (see Sec. <ref>) this is equivalent (up to a scale factor) to the . 




  
Hamiltonian ODE graph network ()

Our “Hamiltonian ODE graph network” () imposes further constraints by using a GN_𝐮 to compute a single scalar for the system—the Hamiltonian—via the global output, analytically differentiating it with respect to its inputs, 𝐪 and 𝐩, and, in accordance with Hamilton's equations (Eq. <ref>), treating these gradients as -ṗ and 𝐪̇, respectively,

    ℋ_GN(𝐪, 𝐩)    = GN_𝐮(𝐪, 𝐩, 𝐜; ϕ) 
     
    f^HOGN_𝐪̇, 𝐩̇(𝐪, 𝐩)
    
       ≡(∂ℋ_GN/∂𝐩
    ,-∂ℋ_GN/∂𝐪) = (𝐪̇, ṗ)    .

The resulting ODE defined by f^HOGN_𝐪̇, 𝐩̇ can be integrated similarly to the  by passing those functions to the integrator (Eq. <ref>). Crucially, the ℋ_GN is not supervised directly, but instead learned end-to-end through the integrator.




§ RESULTS





We train and test our approach on datasets consisting of particle systems where particle j exerts a spring force on particle i, as defined by Hooke's law, 𝐅^ij = -k^ij·(𝐪^i - 𝐪^j),
where k^ij is the spring constant. All systems contained between 4 and 9 particles, and simulations were computed with time-steps of 0.005. To form the data from the raw simulations, we sub-sampled the trajectories at a fixed time-step of 0.1, except in the time-step generalization conditions described below. We trained all models to make next-step predictions of all particles' positions and momenta, with a loss function which penalized mean squared error between the true and predicted values. 












Fig. <ref>a shows that for models trained and tested on RK4, the  and  have lower rollout error than the , and even lower than integrating the true Hamiltonian at the same time step. Fig. <ref>b shows that the average energy of the system also stays closer to the initial energy for the  and . The  energy seems to preserve energy better than the , however this is not surprising: as shown below, the  matches the  well, and both have imperfect energy conservation when integrated with RK4, which is not symplectic[Symplectic integrators are designed to have more accurate energy conservation.]. Similar experiments with a third-order symplectic integrator<cit.> yielded similar rollout errors for the , with 3 times better energy conservation than the  (See Supp. Sec. <ref>). Figs. <ref>c-g (and https://tinyurl.com/hogn-example-trajectoriesvideo link) illustrate the rollout accuracy of the different models for a long 500-step trajectory.



  
Generalizing to untrained time-steps




















Both the  and  exhibit better zero-shot generalization than the  on time-steps which were never experienced during training. Fig. <ref>a shows each model's performance when trained on a time-step of 0.1 and tested on time-steps between 0.005 and 0.5. We also trained models on multiple time-steps, from 0.02 to 0.2, to study whether that could improve generalization. We varied the time-steps of the ODE models via the integrator, while for the  we provided the time-step as input to the model. Fig. <ref>b shows that the  is on par with the other two models within the training range, but becomes worse outside this range. Note, predictions for long time-steps (> 0.3) become very inaccurate for all learned models, as well as the , which suggests that the models are determining the dynamics on the basis of features only appropriate for shorter time-steps. Results for other integrators in Supp. Sec. <ref>.











  
Generalization across integrators



























We varied the integrators used by the model to examine how the quality of the integrator interacts with our different models. Figs. <ref>c-d show that for lower order integrators (RK1-3), the  and  have higher rollout and energy errors than the non-Hamiltonian models ( and ).

We speculate the non-Hamiltonian models are more accurate because they are not obligated to respect the constraint the Hamiltonian imposes between the 𝐪's and 𝐩's dynamics, and can instead, for instance, learn approximations to the time derivatives which are not consistent with the 𝐪 and 𝐩 gradient vector fields of any Hamiltonian.

However, Fig. <ref>e shows that when the  is tested on integrators different from those on which it was trained, its performance is always significantly worse. This indicates that the  is not learning accurate (𝐪̇, ṗ), but rather some other function, which interfaces with the training-time integrator to produce accurate trajectories, but is inappropriate for other integrators. By contrast, Fig. <ref>f shows that the  model usually generalizes to greater accuracy when used with higher order integrators at test time, and specifically when trained with RK4 (dark purple line in Fig. <ref>f), perfectly matches the behavior of the true Hamiltonian (grey line in Fig. <ref>c). 

<cit.>, in their Hamiltonian Neural Networks, use finite differences to approximate the derivatives of 𝐪 and 𝐩 when not available analytically, use those approximated values during training to supervise the Hamiltonian gradients, and then generate test trajectories with an RK4 integrator. In our experiments this is analogous to training the  using RK1 and testing on RK4, and indeed, our findings are generally consistent with theirs, but also indicate that directly learning through higher order integrators is generally better.








§ DISCUSSION AND FUTURE WORK


We incorporated two physically informed inductive biases—ODE integrators and Hamiltonian mechanics—into graph networks for learning simulation, and found they could improve performance, energy, and zero-shot time-step generalization. We also analyzed the impact of varying the ODE integrator independently between training and test, and found that the Hamiltonian approaches benefit most from training with an RK4 integrator, while non-Hamiltonian ones can learn more accurate predictions when trained on lower order integrators. However the  generalizes better across integrators, and approximates the  best in terms of performance and energy accuracy. 

Similar to the graph network's inductive bias for representing complex systems as nodes interacting via edges, the ODE integrator and Hamiltonian inductive biases are informed by physics, but are not specific to physics. Our approach may be applicable to many complex multi-entity systems which can be modeled with ODEs or have some notion of canonical position and momentum coordinates. Moreover, there are other physically informed inductive biases to explore in future work, such as imposing separable potential and kinetic energy terms, Lagrangian mechanics, time-reversibility, or entropic constraints.





§ ACKNOWLEDGMENTS


We acknowledge Irina Higgins, Danilo Rezende, Razvan Pascanu, Shirley Ho and Miles Cranmer for helpful discussions and comments on the manuscript.













unsrtnat












§ DATA GENERATION
 



 §.§ Initial conditions


Each particle i had initial conditions (all in figures in international system of units) drawn from independent uniform distributions:



  * Mass m^i ∈ [0.1, 1]

  * Spring constant k^i∈ [0.5, 1]

  * Initial position q^i_0 ∈ [-1, 1]^2

  * Initial velocity v^i_0 ∈ [-3, 3]^2, p^i_0 = m^iv^i_0




 §.§ Physics simulator


Our dataset consists of simulations for particle-spring systems where particle j exerts a force on particle i (Hooke's force) 𝐅^ij = -k^ij·(𝐪^i - 𝐪^j)
where k^ij is the spring constant, calculated as k^ij = k^i· k^j. 

Trajectories were generated using RK4 with a generation time step of 0.005s, with trajectories of up to 4s in length.



 §.§ Dataset


The datasets were obtained by subsampling the trajectories generated with the physics simulator at different time intervals.

We generated data for systems with between 2 and 15 particles, although we only used data with 4, 5, 6, 8 and 9 particles for training.

The training dataset consisted of 10000 one-step pairs of states (for each number of particles) separated by either a fixed time step of 0.1 s, or by a random time step drawn uniformly from the [0.02, 0.2] s interval and rounded to the nearest exact multiple of the generation time-step. In the latter case, to avoid having discrete values of the time step, instead of using a fixed generation time-step of 0.005 s, we used random values form the interval 0.005 s ± 10%. Each pair was sampled from a different randomly generated trajectory. Additionally we generated 1000 validation and 1000 test pairs for each number of particles.

We also stored validation and test trajectories of length 20 sampled at 0.005, 0.01, 0.03, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5 seconds (1000 trajectories for each number of particles and time step). All results presented in the paper correspond to the full test dataset of 20-step test trajectories for systems with 4, 5, 6, 8 and 9 particles.



§ GRAPH NETWORK
 


Each of the MLPs used on the edge model, node model and global model of the graph network, have output sizes of [64, 64], and softplus activations after each layer including the last one. The graph network always uses sum as the aggregation function for edges and nodes.

The output of the network is obtained by applying an additional non-activated linear layer with the desired output size to the output globals (, output size of 1 to represent a single scalar value per graph) or to the output nodes ( and , output size of 4 to represent values associated to 4 canonical coordinates for each node/particle).

We chose softplus because relu activation did not work well with the Hamiltonian model. We hypothesize this is a consequence of Eq. <ref>: a function approximator with relu activation is a piecewise linear function (linear almost everywhere), so after taking the derivatives of the output of the network with respect to its inputs, it becomes a function that is constant almost everywhere, and the gradient-based optimization process struggles to optimize it. This is consistent with <cit.> where they find better results for tanh activation than for relu activation. In our case softplus seemed to produce better results than tanh.

We always remove the mean position of the objects from the inputs to the GraphNetwork. This does not prevent the models from predicting correct absolute positions since all dynamics are translation invariant and models (Including RK integrators) only predict relative differences with respect to the previous state. This allows the model to produce correct predictions even when particles have travelled far away from their initial positions.

Also, because the ground truth Hamiltonian in spring systems is time-independent (it is a conservative system), we omit feeding absolute time to the GraphNetworks to avoid unnecessary dependencies.




§ TRAINING DETAILS
 


We used TensorFlow and the DeepMind Graph Nets library to build our  neural networks, calculate the analytic gradients of the Hamiltonian, and train our models. We trained with a batch size of 100, for a million training steps. We used AdamOptimizer to minimize the loss.

Due to the different nature of the models (which may require different learning rates) we ran each of the models with 13 initial learning rates uniformly spanning an interval between 10^-1 and 10^-4 in log scale. Learning rate was decayed exponentially at a rate of 0.1 every 2· 10^5, with a lower limit of 10^-7.

For each model, we report median values and min-max range for the 4 learning rates with the smallest rollout error. In practice the variance across seeds given the same learning rate was negligible (except on generalization results). 

Rollout error is defined as the RMS position error averaged across all examples, dimensions, particles, and sequence axis. Energy error is calculated as the RMS of the deviation between the mean energy of a trajectory and the initial energy (normalized by the initial energy, to yield relative errors), averaged across all examples.




§ ADDITIONAL RESULTS
 



 §.§ Symplectic integrators



We attempted using symplectic integrators of first order (Symplectic Euler, S1), second order (Verlet Integrator, S2) and third order (S3) with coefficients as described in <cit.>[We also tried a fourth order symplectic integrator, as described in <cit.>, however it did not improve results, possibly due to the linear nature of the Hooke's force<cit.>.]. It is worth noting that these integrators are only symplectic for separable Hamiltonians ℋ(𝐪, 𝐩) = T(𝐩) + V(𝐪), or equivalently, for systems for which ṗ = f_ṗ(𝐪) and 𝐪̇ = f_𝐪̇(𝐩). However, we do not impose this constraint, neither to the  (by adding the scalar output of two independent networks, with 𝐩 and  𝐪 inputs respectively) nor the  (by only feeding 𝐪 to the network that outputs ṗ and vice versa), so there are not any guarantees that the learned models will correctly preserve energy for any Hamiltonian.

Fig. <ref>a shows the predictive accuracy including results for models trained with symplectic integrators. Contrary to the RK integrators, in this case the  seems to be able to produce lower errors than the  even for low order integrators (S1 and S2), in a similar way to how the  was able to produce very low errors for low order RK integrators. While this seems like an advantage of using symplectic integrators it also indicates that the  model is not really learning something close to the  in these cases, and, as we will see in the next section, it will cause worse generalization to unseen test time-steps. Similarly, Fig. <ref>d shows that the models trained with the S1, S2 integrators do not generalize at all to higher order symplectic integrators, and only the model trained with S3 shows good generalization across all integrators that is comparable to the generalization across integrators of the model trained with RK4 (Fig. <ref>b) and the behavior of the  (Fig. <ref>a).







We speculate these differences in accuracy between training the  with the RK integrators and the symplectic integrators are related to fundamental differences between the algorithms behind the two types of integrators. Firstly while the output of RK integrators is usually a weighted average of the results obtained across several internal iterations for intermediate time-steps, the symplectic integrators feed each internal iteration with just the output of the previous internal iteration, and output just the result of the very last internal iteration. This means that the trained models can learn to encode additional information in the outputs of the intermediate internal iterations, which may allow the neural networks to distinguish whether the gradients are being evaluated for the first, the second, etc. or the final internal iteration. Secondly, in RK integrators, the derivatives of the position and momentum are always evaluated at the same point in each internal iteration, but in the symplectic integrators the derivatives of the position and the momentum are evaluated at different points, in alternating fashion. This has two consequences, the first consequence is that an n-th order symplectic integrator performs 2n network evaluations, while a n-th order RK integrator performs only n network evaluations, as it shares the evaluations of the position and momentum gradients at the same point. The additional network evaluations may explain why apparently lower order symplectic integrators, when used with learned models, can produce lower errors (e.g. S2 vs RK2). The second consequence is that, if as speculated before, the model can identify which calls corresponds to which internal iterations, it may also identify when it is being evaluated to produce gradients for the momentum or for the position. This would allow the model to get around the Hamiltonian constraint (that used to force it to produce a position and momentum vector field of gradients consistent with a common Hamiltonian), and allow it to start behaving more similarly to the , where the derivatives of the position and momentum can be fully independent.

With respect to energy conservation, Fig. <ref>b, shows that models trained with higher order symplectic integrators (S1 and S2) preserve energy much better than the RK integrators, and, as mentioned in the main text, allow the  to preserve energy much better than the . It is also worth noting, that the  trained with RK2, RK3 and RK4 integrators also improves in energy conservation when integrated with symplectic integrators (Fig. <ref>b), confirming that the models trained with high order RK integrators learn something very similar to to the ground truth Hamiltonian that even generalizes to a different family of integrators.



 §.§ Time generalization for other integrators



Figs. <ref>, <ref>, and <ref> show that for RK integrators of lower order (RK1-RK3) the  overfits heavily to the time-step used during training, however, the  presents consistent generalization to time-steps not seen during training, yielding approximately similar errors. This again is evidence that the  learns something more consistent with a Hamiltonian, and that in fact follows more closely the  than the .






In the case of symplectic integrators, we observe that even the  overfits to the time-step used during training when using first order (S1) and second order (S2) integrators (Figs. <ref> and <ref>), and starts generalizing well only with the third order (S3) integrator (Fig. <ref>), although even for S3 the generalization to longer (> 0.2) time-steps is comparatively worse than with RK4 (Fig. <ref>a-b). This is consistent with the previous discussion (Section <ref>) in that the  is not really learning an accurate Hamiltonian when trained with a symplectic integrator.







