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

Rather than regressing future states from reference trajectories, we learn the cumulative dynamics directly from a single phase-space configuration and its instantaneous time derivative.

Background

Hamiltonian

#classical-mechanics

v = \dot{x} = \frac{\partial H (x, p)}{\partial p}, f = \dot{p} = - \frac{\partial H (x, p)}{\partial x}

$x$ : position, $p$ : momenta, $v$ : instantaneous velocity, $f$ : instantaneous force
The state at a later time $t^{*}$ can be obtained by integrating over time interval $[t, t^{*}]$

[\begin{matrix} x_{t^{*}} \\ p_{t^{*}} \end{matrix}] = [\begin{matrix} x_{t} \\ p_{t} \end{matrix}] + \int_{t}^{t^{*}} [\begin{matrix} v_{τ} \\ f_{τ} \end{matrix}] d τ

Molecular Dynamics

Concrete instantiation of Hamiltonian mechanics in which Hamiltonian is separable ie

H (x, p) = T (p) + V (x), T (p) = \sum_{i = 1}^{N} \frac{| | p^{(i)} | |_{2}^{2}}{2 m^{(i)}}

Equation of motion reduces to

v = \frac{\partial T (p)}{\partial p} = \frac{p}{m}, f = - \frac{\partial V (x)}{\partial x}

Method

Model the Hamiltonian evolution directly in phase space over a finite time interval
Hamiltonian flow map $u_{t \to t^{*}}$ learns

[\begin{matrix} x_{t^{*}} \\ p_{t^{*}} \end{matrix}] = [\begin{matrix} x_{t} \\ p_{t} \end{matrix}] + \int_{t}^{t^{*}} [\begin{matrix} v_{τ} \\ f_{τ} \end{matrix}] d τ \to \bar{u} (x_{t}, p_{t}, t^{*} - t) = \frac{1}{t^{*} - t} \int_{t}^{t^{*}} [\begin{matrix} v_{τ} \\ f_{τ} \end{matrix}] d τ

Uses Mean Flow and self-distillation
This time averaged velocity and force captures the non-trivial, integration-dependent component of the flow map.
By multiplying $(t^{*} - t)$ and take time derivative $\frac{d}{d t}$ the equation on the right becomes

\bar{u} (x_{t}, p_{t}, t^{*} - t) - (t^{*} - t) \frac{d}{d t} \int_{t}^{t *} [\begin{matrix} v_{τ} \\ f_{τ} \end{matrix}] d τ

\frac{d}{d t} (\bar{u} (x_{t}, p_{t}, t^{*} - t)) = (v_{t} \cdot \partial_{x_{t}}) \bar{u} + (f_{t} \cdot \partial_{p_{t}}) \bar{u} - \partial_{t^{*} - t} \bar{u}

Above two equations are the consistency condition, which can be used for regression objective.

L (θ, x, p, v, f) = E_{Δ t} [| | {\bar{u}}^{θ} (x, p, Δ t) - {\bar{u}}_{t g t} | |_{2}^{2}]

{\bar{u}}_{t g t} = [\begin{matrix} v \\ f \end{matrix}] + Δ t [(v \cdot \partial_{x}) {\bar{u}}^{θ} + (f \cdot \partial_{p}) {\bar{u}}^{θ} - \partial_{Δ t} {\bar{u}}^{θ}]

Given N particles and d dimensions ${\bar{u}}^{θ} = R^{2 d N + 1} \to R^{2 d N}$

Force matching: in case $Δ t = 0$ ,

L (θ, x, p, v, f)_{d N + 1} = E_{Δ t} [| | {\bar{u}}^{θ} (x, p, Δ t = 0) - f | |_{2}^{2}]

Network learns to predict interatomic force.
-> self distillation across time. instantaneous force prediction is the base signal, pred at larger $Δ t$ are constrained to match the accumulation of local dynamics.

Since MLFF dataset only provides molecular geometry and force labels (x, f), momenta is sampled from Boltzmann distribution

p_{i} N (0, m_{i} k_{B} T I)

Remove center of mass drift by ensuring $\sum_{i = 1}^{N} p_{i} = 0$
Ensure the kinetic energy stays the same after drift removal by rescaling momenta $p_{i} \leftarrow \sqrt{\frac{T_{t g t}}{T_{c u r r}}} p_{i}$

Appendix B3: Coupled Conservation Filter set such that the corrected momenta is as close to the model's predicted momenta as possible, while

Angular momentum conservation: the total angular momentum $L = \sum r_{i} \times p_{i}$ is conserved.
Energy conservation: Total kinetic energy $K = \sum \frac{| p_{i} |^{2}}{2 m_{i}}$ is conserved.