Geometric Integration & Hamiltonian Monte Carlo

Euler and Runge-Kutta treat the state space as flat $\mathbb{R}^{2n}$ and chase down local truncation error and they wipe out the global invariants along the way. Run them on a Hamiltonian system and energy drifts until the simulation blows up.

In HMC that energy drift breaks Detailed Balance and wrecks the sampler.

We want integrators that hold onto structure like symplecticity and volume and momentum by design and not by accident, and that is what geometric integration is. These notes build up the geometry of cotangent bundles and shadow Hamiltonians and then poke at the practical edge cases you hit when you actually implement the thing.

1. The Symplectic Setup

A Hamiltonian system lives on a symplectic manifold $(\mathcal{M}, \omega)$ and usually that manifold is the cotangent bundle $T^*Q$ with coordinates $(q, p)$ and the standard 2-form $\omega = \sum dq_i \wedge dp_i$.

Given $H: \mathcal{M} \to \mathbb{R}$ the vector field $X_H$ satisfies $\omega(X_H, \cdot) = dH(\cdot)$ and in coordinates this gives Hamilton’s equations.

Or you write it as $\dot{z} = J \nabla H(z)$ with $J = \begin{pmatrix} 0 & I \\ -I & 0 \end{pmatrix}$.

The symplectic structure hands $C^\infty(\mathcal{M})$ a Lie algebra through the Poisson bracket.

A flow $\Psi_t$ is symplectic when $\Psi_t^* \omega = \omega$ and this is the same as saying $(\nabla \Psi_t)^T J (\nabla \Psi_t) = J$. Every Hamiltonian flow is symplectic and by Liouville’s theorem it holds the volume form $\Omega = \omega^n / n!$ fixed too.

2. Splitting and Why Leapfrog Works

For $H(q, p) = U(q) + K(p)$ we cannot solve the flow exactly. But we can split it. If $A$ and $B$ are each integrable then their exact flows $\Phi_A$ and $\Phi_B$ are symplectic and stacking symplectic maps together gives you another symplectic map.

Leapfrog is the symmetric splitting.

Each sub-map is a shear with a triangular matrix and unit diagonal so the Jacobian determinant is just 1. Symplecticity holds for any step size $h$ inside the stability region. The geometry is exact and only the trajectory gets approximated.

3. Lie Group Integrators (Munthe-Kaas)

The standard update $z_{n+1} = z_n + h\dot{z}_n$ assumes the state space is closed under addition and that falls apart on a Lie group like $SO(3)$ because adding a tangent vector to a rotation matrix spits out something that is no longer a rotation matrix.

The fix is to use the exponential map $\exp: \mathfrak{g} \to G$ from the Lie algebra.

where $\sigma_n \in \mathfrak{g}$. Munthe-Kaas or RKMK methods push Runge-Kutta onto manifolds by solving the ODE inside the Lie algebra and mapping back. Joint angles stay unit length and rigid bodies keep their geometric constraints.

4. Variational Integrators

Lagrangian mechanics minimizes action.

Discretize the action directly and not the ODE.

Apply Hamilton’s Principle and you get the Discrete Euler-Lagrange equations.

The key point is that every variational integrator is symplectic because it chops up the variational principle instead of the differential equation. The discrete symplectic form $\omega_L$ falls out for free. Variational integration is the standard in orbital mechanics and celestial simulation and I think it is underappreciated in the ML community.

5. Shadow Hamiltonians

What does Leapfrog actually solve? Not $H$ exactly and instead there is a modified Hamiltonian $\tilde{H}$ whose exact flow runs through the numerical points $(q_k, p_k)$.

Write the Leapfrog operator using Lie operators as $\Psi_h = e^{h/2 L_A} e^{h L_B} e^{h/2 L_A}$.

By Baker-Campbell-Hausdorff this equals $e^{h L_{\tilde{H}}}$ and for symmetric splitting the error terms show up only in even powers of $h$.

For separable $H(q, p) = \frac{1}{2}p^2 + U(q)$ this reads

$\tilde{H}$ is formally an infinite series and it diverges. But for small $h$ truncation works fine. For a harmonic oscillator with frequency $\omega$ stability asks for

Past that threshold the shadow Hamiltonian stops existing and the dynamics turn chaotic.

6. The Ge-Marsden No-Go Theorem

Can you build a symplectic integrator that also holds energy exactly?

Theorem (Ge-Marsden 1988): If a fixed-step integrator is symplectic and energy-preserving with $H(z_{k+1}) = H(z_k)$ and the system is non-integrable then the integrator is just the identity map.

You have to pick between exact geometry and exact energy. Symplectic integrators pick geometry and this is the better trade because the trajectory stays near energy contours anyway and preserving phase-space volume is more useful statistically than hitting exact energy conservation.

7. HMC and the Kepler Problem

HMC in practice runs three steps.

1. Sample $p \sim \mathcal{N}(0, M)$.
2. Run Leapfrog for $L$ steps.
3. Metropolis accept/reject: $\alpha = \min(1, \exp(H_{old} - H_{new}))$.

If Leapfrog held $H$ perfectly then $\alpha$ would always be 1. It holds $\tilde{H}$ instead and so the acceptance error scales as $\epsilon^2$.

A good test case is the Kepler problem with $H = \frac{1}{2}\|p\|^2 - \frac{1}{\|q\|}$. Leapfrog gives a stable orbit that precesses a little. RK4 bleeds energy until the planet spirals into the sun or pumps energy in and flings it out.

import numpy as np
import matplotlib.pyplot as plt
 
def kepler_potential(q):
    return -1.0 / np.linalg.norm(q)
 
def kepler_grad(q):
    return q / np.linalg.norm(q)**3
 
def leapfrog_kepler(steps=2000, h=0.01):
    q = np.array([1.0, 0.0])
    p = np.array([0.0, 0.5]) # Low velocity for elliptical orbit
    path = []
    energy = []
    
    for _ in range(steps):
        p -= (h/2) * kepler_grad(q)
        q += h * p
        p -= (h/2) * kepler_grad(q)
        
        path.append(q.copy())
        energy.append(0.5*np.dot(p, p) + kepler_potential(q))
        
    return np.array(path), np.array(energy)
 
def rk4_kepler(steps=2000, h=0.01):
    z = np.array([1.0, 0.0, 0.0, 0.5]) # q1, q2, p1, p2
    path = []
    energy = []
    
    def f(state):
        q = state[:2]
        p = state[2:]
        return np.concatenate([p, -kepler_grad(q)])
    
    for _ in range(steps):
        k1 = f(z)
        k2 = f(z + h/2 * k1)
        k3 = f(z + h/2 * k2)
        k4 = f(z + h * k3)
        z += (h/6) * (k1 + 2*k2 + 2*k3 + k4)
        
        path.append(z[:2].copy())
        p = z[2:]
        energy.append(0.5*np.dot(p, p) + kepler_potential(z[:2]))
        
    return np.array(path), np.array(energy)
 
def plot_kepler_results():
    path_lf, e_lf = leapfrog_kepler()
    path_rk, e_rk = rk4_kepler()
    
    plt.figure(figsize=(12, 5))
    
    # Orbits
    plt.subplot(1, 2, 1)
    plt.plot(path_lf[:,0], path_lf[:,1], label='Leapfrog (Symplectic)')
    plt.plot(path_rk[:,0], path_rk[:,1], 'r--', label='RK4 (Energy Drift)')
    plt.scatter([0], [0], color='orange', s=100, label='Sun')
    plt.title('Kepler Orbits (2000 steps)')
    plt.legend()
    
    # Energy error
    plt.subplot(1, 2, 2)
    plt.plot(e_lf - e_lf[0], label='Leapfrog Error')
    plt.plot(e_rk - e_rk[0], label='RK4 Error')
    plt.title('Energy Fluctuation vs Drift')
    plt.ylabel('H - H0')
    plt.legend()
    
    # plt.show()

8. RATTLE (Constrained Systems)

Systems with holonomic constraints $g(q) = 0$ like fixed bond lengths in a molecule break plain Leapfrog within a few steps. The fix is to throw in Lagrange multipliers.

RATTLE solves for $\lambda$ to make both $g(q_{n+1}) = 0$ and $\nabla g \cdot v_{n+1} = 0$ hold. Its particular staggering of multiplier updates keeps the symplectic form on the constrained manifold. The proof is not obvious but the result holds.

9. Adaptive Step Sizes and Symplecticity

Naively varying $h$ based on the state wipes out symplecticity. If $h$ depends on $z$ then the map $\Psi_{h(z)}$ is no longer symplectic and energy drift comes right back.

The fix is to reparametrize time by setting $dt = g(q)\,d\tau$ and solving a modified Hamiltonian $\mathcal{H} = g(q)(H(q, p) - E_0) = 0$. Fixed steps in $\tau$ give you variable resolution in $t$ and the geometric invariants stay intact.

10. Neural Hamiltonian Flows

If you parametrize a Hamiltonian $H_\theta(q, p)$ as a neural network and run the flow with Leapfrog layers then you get a normalizing flow that is exactly reversible and volume-preserving and physics-aware all at once. Exact reversibility means you can reconstruct the latent noise from the output with no approximation. Volume-preserving means the Jacobian determinant is just 1 and the change-of-variables formula simplifies. And physics-aware means the network learns the conserved quantities of whatever dynamical system it is modeling.

This is the idea behind Hamiltonian Neural Networks and Symplectic Adjoint Networks. The symplecticity is what makes the normalizing flow tractable because without it computing the Jacobian determinant would be too expensive to bother with.

11. The Edge of Stability

The shadow Hamiltonian is an asymptotic series. For small $h$ it nails the numerical flow almost perfectly. But in practice you want the largest stable $h$ you can get and as $h$ grows the error terms grow too. Past a threshold the series diverges and orbits leave the invariant KAM tori and fill phase space and HMC acceptance rates collapse.

The practical problem is always the same and it is to push $h$ as close to the convergence boundary as you can without crossing it.

Liouville’s Theorem (Proof)

Theorem: If $\dot{z} = f(z)$ and $\text{div } f = 0$ then volume is conserved.

Let $V(t)$ be the volume of a set $\Omega_t$ carried along by the flow.

Set $J(t) = \det \nabla \Psi_t$. Jacobi’s formula gives $\frac{d}{dt} \det A(t) = \det A(t) \text{Tr}(A^{-1} \dot{A})$.

With $A = \nabla \Psi_t$ the chain rule gives $\dot{A} = (\nabla f) A$ and so $A^{-1} \dot{A} = \nabla f$ and $\text{Tr}(A^{-1} \dot{A}) = \text{div } f$.

For a Hamiltonian $f = (H_p, -H_q)$ the divergence is $\text{div } f = \frac{\partial^2 H}{\partial q \partial p} - \frac{\partial^2 H}{\partial p \partial q} = 0$.

So $\dot{J} = 0$ and since $J(0) = 1$ we get $J(t) = 1$ for all time. $\square$

Symplectic Euler

Order 1 symplectic method and simple but worth knowing.

1. $p_{n+1} = p_n - \epsilon \nabla U(q_n)$
2. $q_{n+1} = q_n + \epsilon p_{n+1}$ (use the new momentum)

The Jacobian is

The determinant is $(1 - \epsilon^2 H_{qq}) - (-\epsilon^2 H_{qq}) = 1$ and that confirms the map is symplectic.

It is not time-reversible though. Symmetrizing it gives you Leapfrog with $\Psi_{Leap} = \Phi_{Euler}^* \circ \Phi_{Euler}$.

KAM and Long-Term Stability

How long does shadow Hamiltonian stability actually last?

KAM theory says that for small perturbations of integrable systems most invariant tori survive. A symplectic integrator is exactly that kind of perturbation. Backward error analysis shows the discrete trajectory stays $O(h^p)$ close to $\tilde{H}$ for times $T \propto \exp(1/h)$.

That is exponentially long. In practice symplectic integrators track a nearby stable system for astronomical timescales literally. This is why people use them to study Solar System stability over billions of years where RK4 would have diverged within a few millennia.

Exterior Calculus Perspective

Symplecticity is area preservation said carefully in the language of exterior calculus.

$\omega = \sum dq_i \wedge dp_i$ measures oriented area in phase space and $\Psi$ is symplectic iff $\Psi^* \omega = \omega$. By Stokes

For HMC this means the integrator holds phase-space structure with no compression and no expansion. And without this the Markov chain would pick up a bias toward high-density regions.

Generating Functions

Every symplectic map $\Psi$ can locally be written through a generating function $S(q_0, q_1)$ with $p_0 = -\frac{\partial S}{\partial q_0}$ and $p_1 = \frac{\partial S}{\partial q_1}$.

When $S$ satisfies Hamilton-Jacobi

it pins down the exact flow. Symplectic integrators match up with approximate generating functions $S_d \approx S$. The approximation still satisfies the symplectic group’s consistency conditions and that is why the method works even though $S_d$ is inexact.

Stochastic Hamiltonian Dynamics

For Bayesian inference we often want Hamiltonian Langevin dynamics.

The friction term $-\gamma p$ breaks exact volume preservation but the Gibbs measure $e^{-H/T}$ is still preserved. Discretization uses a stochastic splitting.

The Hamiltonian part stays exactly symplectic and the Langevin part is an OU process and gets solved exactly. This hybrid is the workhorse behind large-scale posterior sampling in deep learning.

Geodesic Integration

When the metric $g_{ij}(q)$ depends on the state as it does in general relativity and robotics the kinetic energy becomes $K(q, p) = \frac{1}{2} p^T g^{-1}(q) p$ and the Hamiltonian is no longer separable. Standard Leapfrog breaks.

The options are implicit methods or geodesic splitting and the latter means decomposing the metric when it has symmetry or iterating a fixed-point equation for $p_{n+1/2}$. The trajectory then follows the curvature instead of flying off the tangent plane.

Reversibility and Debugging

Symplectic integrators are time-reversible. Negate $p$ and run $N$ steps and you should land back at $q_0$ up to floating-point noise.

In large physics simulations people use integer arithmetic to make this exactly reversible at the bit level. This lets you do reversible debugging and rewind from a crash to find the exact moment things went singular without storing checkpoints. The bit-level symmetry is just $\Psi_{-h} \circ \Psi_h = I$ taken seriously.

The Symplectic Group $Sp(2n, \mathbb{R})$

The symplectic matrices form a Lie group and $M \in Sp(2n, \mathbb{R})$ iff $M^T J M = J$.

Three properties worth noting. First $\det M = 1$ so volume and orientation are preserved. Second the eigenvalues come in quartets and if $\lambda$ is an eigenvalue then so are $1/\lambda$ and $\bar{\lambda}$ and $1/\bar{\lambda}$. This spectral symmetry stops eigenvalues from drifting into the unstable region without first colliding pairwise at the unit circle. And third the dimension is $n(2n+1)$.

Geometric integration at its core is about keeping the numerical update on this group or near it despite the discretization error.

Multi-symplectic Methods for PDEs

Hamiltonian PDEs like the wave equation and Schrodinger often have multi-symplectic structure with symplectic forms in both time and space.

Standard symplectic integrators only handle the temporal part. Multi-symplectic integrators like the Preissman box scheme hold a local energy-momentum conservation law per spacetime cell and kill the unphysical spatial oscillations that otherwise mess up the wave-packet shape over long propagation.

Optimal Transport

Hamiltonian flows connect to optimal transport in a way I think deserves more attention. Brenier’s theorem says the optimal map pushing $\mu$ to $\nu$ is the gradient of a convex function. In the Benamou-Brenier dynamic picture the optimal paths are geodesics in the Wasserstein space of probability measures.

Integrating these Wasserstein flows geometrically means holding onto displacement convexity of the energy. Symplectic discretizations of the OT problem keep the transport maps from leaking probability mass into regions where it should not go.

Why Automatic Differentiation Matters Here

Leapfrog needs $\nabla U$. Finite differences like $(f(x+h) - f(x))/h$ bring in $O(h^2)$ errors that break the symplectic structure. Even a small gradient error fires off as a non-conservative force and wipes out the shadow Hamiltonian’s stability guarantees.

Automatic differentiation in JAX and PyTorch computes gradients to machine precision and keeps the numerical vector field gradient-like by construction. This is arguably the single biggest practical advance that makes geometric integration work for high-dimensional models.

Historical notes

Hamiltonian mechanics goes back to William Rowan Hamilton’s formulation in 1833. The first explicit symplectic integrator which is Leapfrog or Stormer-Verlet showed up in De Vogelaere’s 1956 work and Loup Verlet independently rediscovered it for molecular dynamics in 1967 which gave it broad practical impact. Hybrid Monte Carlo came along in 1987 from Duane et al. and tied symplectic integration to Bayesian sampling. The next year Ge and Marsden proved their No-Go theorem which says no fixed-step integrator can be both symplectic and energy-conserving for non-integrable systems. The definitive textbook treatment came in 2006 with Hairer and Lubich and Wanner’s “Geometric Numerical Integration.” More recently Hoffman and Gelman introduced the No-U-Turn Sampler in 2014 and Greydanus et al. proposed Hamiltonian Neural Networks in 2019 and pulled geometric integration into deep learning.

Definitions

1. Cotangent bundle ($T^*Q$): the space of all possible positions and momenta for a mechanical system. Each point $(q, p) \in T^*Q$ specifies a complete state.
2. Holonomic constraint: a restriction expressed purely in terms of coordinates, e.g. $q_1^2 + q_2^2 = 1$. RATTLE enforces these within the symplectic framework.
3. Poisson bracket ($\{f, g\}$): the Lie algebra operation on phase-space functions, defined as $\sum_i \frac{\partial f}{\partial q_i}\frac{\partial g}{\partial p_i} - \frac{\partial f}{\partial p_i}\frac{\partial g}{\partial q_i}$. Conservation laws correspond to functions whose Poisson bracket with $H$ vanishes.
4. Shadow Hamiltonian ($\tilde{H}$): the exact conserved quantity of the discrete map. It differs from $H$ by $O(h^2)$ for leapfrog.
5. Symplectic 2-form ($\omega$): the geometric object $\omega = \sum dq_i \wedge dp_i$ measuring oriented area in phase space. A map preserving $\omega$ is symplectic.
6. Symplectic manifold: a smooth manifold equipped with a closed, non-degenerate 2-form. Phase space $(T^*Q, \omega)$ is the prototypical example.

References

1. Hairer, E., Lubich, C., & Wanner, G. (2006). “Geometric Numerical Integration”. The definitive reference for structure-preserving algorithms.

2. Leimkuhler, B., & Reich, S. (2004). “Simulating Hamiltonian Dynamics”. Bridge between molecular dynamics and numerical analysis.

3. Neal, R. M. (2011). “MCMC using Hamiltonian dynamics”. The foundational HMC paper for the machine learning community.

4. Marsden, J. E., & West, M. (2001). “Discrete mechanics and variational integrators”. The seminal paper on deriving integrators from the discrete action principle.

5. Munthe-Kaas, H. (1998). “High order Runge-Kutta methods on manifolds”. Introduced the RKMK methods for Lie group integration.

6. Bridges, T. J., & Reich, S. (2001). “Multi-symplectic integrators: numerical methods for Hamiltonian PDEs”. The primary source for multi-symplectic geometry in wave systems.

7. Ambrosio, L., Gigli, N., & Savare, G. (2008). “Gradient Flows in Metric Spaces”. Connection between Wasserstein geometry and optimal transport.