Learning Hamiltonian Flow Maps: Mean Flow Consistency for Large-Timestep Molecular Dynamics

TL;DR

This work presents Hamiltonian Flow Maps (HFMs) to accelerate long-time simulations of Hamiltonian systems by learning the mean phase-space evolution over a horizon Δt from single-time, trajectory-free phase-space samples. A trajectory-free Mean Flow consistency objective guides the learning, enabling the model to predict large-timestep dynamics without requiring reference trajectories. The approach is implemented with both translation-invariant and SO(3)-equivariant architectures and augmented with inference-time filters to enforce energy and angular-momentum conservation, as well as drift removal. Empirical results across classical mechanics and MD tasks show HFMs can achieve stable, accurate rollouts at timesteps far beyond classical integrators, including MD with ML force fields, while maintaining training and inference efficiency. The work highlights practical data efficiency and broad applicability, though it notes limitations in preserving exact symplectic structure and in strongly chaotic regimes, motivating future enhancements in physics-informed training and filtering strategies.

Abstract

Simulating the long-time evolution of Hamiltonian systems is limited by the small timesteps required for stable numerical integration. To overcome this constraint, we introduce a framework to learn Hamiltonian Flow Maps by predicting the mean phase-space evolution over a chosen time span $Δt$, enabling stable large-timestep updates far beyond the stability limits of classical integrators. To this end, we impose a Mean Flow consistency condition for time-averaged Hamiltonian dynamics. Unlike prior approaches, this allows training on independent phase-space samples without access to future states, avoiding expensive trajectory generation. Validated across diverse Hamiltonian systems, our method in particular improves upon molecular dynamics simulations using machine-learned force fields (MLFF). Our models maintain comparable training and inference cost, but support significantly larger integration timesteps while trained directly on widely-available trajectory-free MLFF datasets.

Figures (23)

• Figure 1: Hamiltonian Flow Maps (HFMs) for large timesteps in phase space. Top row: Existing approaches rely on trajectory data, typically generated using a teacher through sequential simulation with small timesteps. While this enables training large-timestep models via direct regression, the resulting models are limited to a fixed set of predefined timesteps that cannot be adjusted after training. Bottom row: Our approach learns continuous-time, large-timestep dynamics directly from decorrelated ab-initio samples, without requiring trajectories. The model supports arbitrary timesteps at inference. Our tailored loss combines force matching with a consistency constraint that enforces agreement of the predicted flow across different time horizons (see \ref{['sec:computational-approach']}).
• Figure 2: HFMs vs. classical integration: Symplectic integrators such as Velocity Verlet (VV) advance the system through many local half-steps (left). HFMs instead predict the phase-space displacement over the interval directly (right). By modeling the mean velocity and force $(\bar{{\bv}},\bar{{\boldsymbol{f}}})$ over the interval, HFMs apply a single large update (blue arrows) from the current state.
• Figure 3: Simulation with a trained HFM. The model predicts mean forces ${\bar{\boldsymbol{f}}}$ and mean velocities ${\bar{\bv}}$, conditioned on the current phase-space state $({\bx}_t, {\bp}_t)$, the time interval ${\Delta t}$, and atom information (types and masses). The predictions advance the system by one integration step. To ensure stable rollouts, we refine the updated state using three filters (right): removal of global translation drift, coupled conservation of energy and angular momentum (E&L), and random rotation for approximate rotation equivariance.
• Figure 4: Generated trajectories for two potentials (Barbanis, spring pendulum) from two initial conditions (left/right). We compare the ground truth (dashed), VV (red), and HFM (blue) at two large timesteps. VV deviates even with a small timestep increase, while HFM stays aligned at larger steps.
• Figure 5: Gravitational $N$-body rollout (${t}=3$, ${t^*}=4$). Top: MSE as a function of the number of integration steps $n$, averaged over the test set. The HFM remains accurate at coarse discretization, while VV diverges rapidly. Overall, the HFM requires $4\times$ fewer steps than VV to reach the same accuracy. Bottom: Particle trajectories for 16 and 32 integration steps, colored by per-particle deviation from the ground truth. The HFM yields physically consistent rollouts, whereas VV leads to unphysical scattering.