Source-linked AI summary
Hamiltonian Graph Networks with ODE Integrators
Alvaro Sanchez-Gonzalez, Victor Bapst, Kyle Cranmer, Peter Battaglia
TL;DR
Learned simulation models benefit from physically informed structure, but the effects of combining ODE integration and Hamiltonian mechanics in graph networks require evaluation. The paper combines these biases and finds better accuracy, energy behavior, and zero-shot generalization, while identifying limitations for lower-order integration and coarse time-steps.
Problem
The paper addresses the limited use of combined ODE and Hamiltonian inductive biases in graph-network architectures for learned simulation.
Method
The approach combines graph networks with a differentiable ODE integrator and uses a learned Hamiltonian to generate dynamics through Hamilton’s equations.
Results
The models improve predictive accuracy, energy conservation, and zero-shot generalization to unseen time-steps and integrator orders relative to baselines.
Takeaways & Limitations
ODE-integrator and Hamiltonian biases can improve learned simulation, while the Hamiltonian approach generalizes better across integrators and approximates the true Hamiltonian more closely.
Takeaways & Limitations
With lower-order integrators and coarse time-steps, Hamiltonian models can have higher rollout and energy errors because Hamiltonian constraints restrict the learned dynamics.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper combines graph networks with ODE integrators and Hamiltonian mechanics to impose physically informed inductive biases on learned simulation. Experiments show improved predictive accuracy, energy conservation, and generalization to unseen time-steps and integrator orders.
- 1 Introduction: The approach combines ODE and Hamiltonian inductive biases with graph-network architectures for learning simulation.ODE and Hamiltonian biases had previously been studied individually; this work incorporates both into graph-network-based neural architectures.
- 1 Introduction: The trained models achieve greater predictive accuracy than baselines without these biases.
- 1 Introduction: The Hamiltonian bias improves energy conservation, while the ODE-integrator bias strengthens generalization to novel time-steps and integrator orders.
- 1 Introduction: Lower-order integrators and coarse time-steps create trade-offs because Hamiltonian constraints restrict how higher-order dynamics can be modeled.Both the learned Hamiltonian model and a model using the true Hamiltonian struggle in these settings, whereas less constrained models can approximate such structure.
2 Background
The paper represents particle systems as graphs and uses differentiable Runge-Kutta integration to advance states governed by first-order ODEs. Hamiltonian mechanics provides a formulation in which derivatives of the system energy determine position and momentum dynamics.
- Graph networks: Particle systems are represented as graphs whose nodes are particles and whose edges connect every pair of nodes.The graph network operates on global features, variable numbers of nodes, and edges, producing per-node or global outputs.
- Model comparison: The baseline DeltaGN predicts state changes directly from a graph state and time-step using a graph-network node output.Figure 1 contrasts this direct state-change prediction with the ODE-based models.
- Numerical integrators for solving ODEs: Runge-Kutta integrators generate trajectories by repeatedly querying a differentiable time-derivative function.The paper examines RK1, RK2, RK3, and RK4; higher-order methods compose multiple function queries.
- Hamiltonian mechanics: In Hamiltonian mechanics, H(q, p) is a function of canonical position q and momentum p that usually corresponds to system energy.The dynamics follow Hamilton’s equations, expressed as two first-order ODEs.
3 Models
The models progress from directly predicting state changes to learning ODE derivatives and then deriving those derivatives from a learned Hamiltonian. The learned derivative functions are integrated end-to-end, with the Hamiltonian model imposing the strongest structural constraint.
- Delta graph network (DeltaGN): DeltaGN directly predicts position and momentum changes with a graph network.Its update adds predicted changes to the current state, using the time-step, state, system parameters, and neural-network parameters.
- ODE graph network (OGN): OGN learns time derivatives independently of the time-step and supplies them, with initial conditions and a time-step, to an RK integrator.For fixed time-step and RK1/Euler, this is equivalent to DeltaGN up to a scale factor.
- Hamiltonian ODE graph network (HOGN): The methods are compatible with time-dependent Hamiltonians, but time dependence is omitted because the dataset does not require it.
- ODE graph network (OGN): The OGN derivative function can be evaluated at arbitrary intermediate states and queried multiple times by the integrator.
- Hamiltonian ODE graph network (HOGN): HOGN uses a graph-network global output to compute one scalar Hamiltonian, then differentiates it with respect to q and p to obtain the dynamics.The Hamiltonian network is learned end-to-end through the integrator rather than supervised directly.
4 Results
The physically informed models improve rollout accuracy, energy conservation, and zero-shot generalization, but their benefits depend on the integrator and test regime.
- The experiments use particle-spring systems and train models for next-step prediction of all particles’ positions and momenta with mean squared error.The systems contain 4–9 particles and are generated with 0.005-second simulation steps before subsampling.
- OGN and HOGN achieve lower RK4 rollout error than DeltaGN and integrating the true Hamiltonian, while their energy remains closer to the initial value.The OGN preserves energy better than HOGN in this RK4 setting; a third-order symplectic integrator gives HOGN similar rollout error with 3 times better energy conservation than OGN.
- OGN and HOGN generalize better than DeltaGN to unseen time-steps, although all learned models and the true Hamiltonian become very inaccurate beyond 0.3 seconds.DeltaGN is competitive within its training range but degrades outside it; ODE models vary time-steps through the integrator.
- For lower-order RK1–RK3 integrators, HOGN and the true Hamiltonian have higher rollout and energy errors than DeltaGN and OGN.The authors speculate that non-Hamiltonian models are more accurate because they need not enforce the Hamiltonian constraint linking q and p dynamics.
- OGN performance is always significantly worse with test integrators different from the training integrator, whereas HOGN generally becomes more accurate with higher-order test integrators.When trained with RK4, HOGN matches the true Hamiltonian’s behavior.
5 Discussion and future work
The study finds that ODE-integrator and Hamiltonian biases improve learned simulation across several evaluation axes, with integrator choice shaping the trade-offs. The authors also identify broader applicability beyond physics and several future bias families.
- The approach combines ODE-integrator and Hamiltonian mechanics biases in graph networks to improve performance, energy accuracy, and zero-shot time-step generalization.
- Hamiltonian approaches benefit most from RK4 training, while non-Hamiltonian models can predict more accurately when trained with lower-order integrators.
- HOGN generalizes better across integrators and best approximates the true Hamiltonian in performance and energy accuracy.
- The proposed biases are physics-informed but not physics-specific, potentially extending to multi-entity systems modeled by ODEs or canonical position and momentum coordinates.Future directions include separable potential and kinetic energies, Lagrangian mechanics, time-reversibility, and entropic constraints.
A Data generation
The data-generation setup simulates particle-spring systems with Hooke forces, RK4 integration, and independently sampled particle initial conditions.
- Each particle’s initial velocity vi 0 is independently sampled from [−3, 3]2, and its initial momentum is pi 0 = mivi 0.
- Particle j exerts Hooke’s force Fij = −kij · (qi −qj) on particle i, with spring constant kij = ki · kj.
- Trajectories are generated with RK4 using a 0.005-second time-step and extend up to 4 seconds.
A.3 Dataset
The dataset is formed by subsampling simulated trajectories across particle counts and time intervals, with graph-network architecture and input choices designed for the modeled dynamics.
- Simulated trajectories are subsampled at different time intervals to create the datasets.
- Systems contain 2–15 particles, but training uses only systems with 4, 5, 6, 8, and 9 particles.
- Training uses 10,000 one-step state pairs per particle count, separated by either 0.1 seconds or uniformly sampled time-steps from 0.02–0.2 seconds.
- Validation and test trajectories contain 20 steps sampled at time-steps from 0.005 to 0.5 seconds, with 1,000 trajectories per particle count and time-step.
- The graph-network MLPs use [64, 64] outputs, softplus activations, and sum aggregation for edges and nodes.
- HOGN outputs one scalar Hamiltonian per graph, while OGN and DeltaGN output four canonical-coordinate values per particle.
- Softplus is preferred because ReLU performed poorly for HOGN, whose input derivatives are needed for gradient-based optimization.The authors report that softplus also performed better than tanh in their experiments.
- Removing mean particle position preserves correct absolute-position predictions because the dynamics and integrators are translation invariant.
C Training details
The models were trained with fixed optimization settings, while learning-rate selection accounted for differences among model types. Results summarize the four learning rates with the smallest rollout error.
- Training used TensorFlow, Graph Nets, batch size 100, one million steps, and AdamOptimizer.
- Thirteen initial learning rates spanned 10^-1 to 10^-4 uniformly in log scale, with exponential decay and a 10^-7 floor.
- Reported values are medians and min-max ranges over the four learning rates with the smallest rollout error.
- Rollout error is RMS position error averaged across examples, dimensions, particles, and sequence positions.
- Energy error is normalized RMS deviation of trajectory mean energy from initial energy, averaged across examples.
D.1 Symplectic integrators
Symplectic integrators improve energy conservation but create distinct accuracy and generalization behavior. HOGN performs well with higher-order symplectic integration, whereas lower-order variants can undermine Hamiltonian fidelity.
- S1, S2, and S3 are first-, second-, and third-order symplectic integrators, respectively, applied without imposing separable Hamiltonian structure.
- HOGN achieves lower predictive error than the true Hamiltonian with S1 and S2, but this may indicate it is not learning the true Hamiltonian.
- S1- and S2-trained models fail to generalize to higher-order symplectic integrators, while S3 generalizes comparably to RK4-trained models and the true Hamiltonian.
- A fourth-order symplectic integrator did not improve results, possibly because the Hooke’s-force system is linear.
- Higher-order symplectic integrators preserve energy better than RK integrators, and HOGN preserves energy better than OGN under this setting.
- HOGN models trained with RK2, RK3, or RK4 also improve energy conservation when tested with symplectic integrators.
D.2 Time generalization for other integrators
Time-step generalization depends strongly on integrator family and order. HOGN generalizes across unseen time-steps with RK1-RK3, but symplectic S1 and S2 overfit and S3 remains weaker than RK4 for longer steps.
- With RK1-RK3, OGN heavily overfits the training time-step, whereas HOGN yields approximately similar errors on unseen time-steps.
- For symplectic integration, HOGN overfits with S1 and S2 and begins generalizing well only with S3.
- S3 generalization to time-steps longer than 0.2 is comparatively worse than RK4 generalization.
- Figures D.7-D.9 present time generalization for S1, S2, and S3 integrators, respectively.