Hamiltonian normalizing flows

Nov 2026 normalizing flowsHamiltonian dynamicsMCMC

How to build an invertible neural map out of Hamiltonian dynamics, why the Jacobian determinant disappears, how the idea extends to curved geometries, and how the result can make Hamiltonian Monte Carlo cheaper.

1. What a normalizing flow does, and what it costs

We often need a probability model that we can both sample and evaluate: draw new points, and compute their density. Two situations are typical. In density estimation we have samples of an unknown distribution and want a model of it. In Bayesian inference we know a target density \(\pi\) only up to a constant, \(\tilde\pi \propto \pi\), and want samples from it.

A normalizing flow handles both. Take a simple base density \(p_Z\), such as a standard Gaussian, and an invertible map \(f_\theta\). If \(z \sim p_Z\), then \(x = f_\theta(z)\) has the density given by the change-of-variables formula [1]:

The determinant is the bill. For a generic map in dimension \(d\) it costs \(\mathcal{O}(d^3)\), so much of the work on flows goes into architectures that make it cheap: coupling layers, autoregressive maps, triangular Jacobians [2][3][4]. Physics offers a different way out: use a map whose determinant is known to be exactly one.

2. A particle in a landscape

Think of \(q \in \mathbb{R}^d\) as the position of a particle moving on a surface whose height is a potential \(V(q)\). Give the particle a momentum \(p \in \mathbb{R}^d\) and a mass matrix \(M\). Its total energy is

and its motion obeys Hamilton’s equations:

The particle rolls down valleys, gains speed, climbs the next slope and slows down; energy is conserved along the way. Call \(\Phi_t\) the map that sends a starting point \((q_0,p_0)\) to its position in phase space at time \(t\). Two of its properties are exactly what a normalizing flow needs.

It is invertible. Running time backwards undoes the motion.

It preserves volume. This is Liouville’s theorem: \(\det \nabla \Phi_t = 1\). A cloud of particles can stretch and fold, but the phase-space volume it occupies never changes. A density on phase space is therefore simply carried along the trajectories, with no correction term:

Compared with the change-of-variables formula above, the log-determinant has vanished, and physics, not architecture design, removed it.

3. From equations to a network: the Euclidean case

“Euclidean” here means that the mass matrix \(M\) is constant: the particle has the same inertia everywhere. To compute, we need a discrete version of \(\Phi_t\) that keeps both properties. The leapfrog integrator does it. With step size \(\varepsilon\):

Each line changes one variable by an amount that depends only on the other. Such a map is a shear: its Jacobian is block-triangular with identity blocks on the diagonal, so its determinant is exactly one. The inverse is the same scheme run with \(-\varepsilon\). This is the integrator of Hamiltonian Monte Carlo [5].

A neural Hamiltonian flow (NHF) [6] makes the potential, and possibly the kinetic energy, learnable. The recipe has three steps:

1. Draw \((q_0,p_0)\) from a simple base density \(\rho_0\), for instance \(\mathcal{N}(0,I) \otimes \mathcal{N}(0,M)\).
2. Apply \(L\) leapfrog steps with a neural potential \(V_\theta\).
3. Keep \(q_L\) as the sample. The density of the pair is known exactly: \(\log \rho_L(q_L,p_L) = \log \rho_0(q_0,p_0)\).

def nhf_sample(V_theta, base, n, L, eps, M_inv):
    q, p = base.sample(n)              # (q0, p0) ~ rho_0
    log_rho = base.log_prob(q, p)      # exact, and it never changes
    for _ in range(L):                 # L leapfrog layers
        p = p - 0.5 * eps * grad(V_theta, q)
        q = q + eps * p @ M_inv
        p = p - 0.5 * eps * grad(V_theta, q)
    return q, p, log_rho               # no log-det term anywhere

How the model is trained depends on the situation. The momentum is an auxiliary variable: the model lives on \((q,p)\) while the quantity of interest is \(q\) alone. Both cases use a distribution \(r(p \mid q)\) over momenta, often simply \(\mathcal{N}(0,M)\).

Sampling a target known up to a constant. We fit the flow to the extended target \(\pi(q)\, r(p \mid q)\) by minimising the reverse Kullback–Leibler divergence. Since the density is transported unchanged, the loss reads

which only requires \(\log \tilde\pi\) and back-propagation through the integrator.

Learning from data. For a data point \(x\), draw \(p \sim r(\cdot \mid x)\), integrate backwards to the base, and maximise the lower bound

Two choices matter in practice. Toth et al. [6] learn both \(V\) and \(K\); symmetries can also be built into the dynamics [7]. In [8] we fix \(K\) to a quadratic form, which removes a network, reduces the number of parameters, and leaves \(V_\theta\) as the only learned object, which makes it easier to interpret. The same structure also fits kinetic PDEs, whose characteristics are Hamiltonian: in [9] we use it to solve the 1D Vlasov–Poisson system.

4. Curving the space: the Riemannian case

A constant mass matrix gives the particle the same inertia everywhere. Many targets do not have one scale: a posterior can be narrow in one region and wide in another, or bend along a curved ridge. With a fixed \(M\), the step size must suit the hardest region, and many leapfrog layers are needed.

The remedy, borrowed from Riemannian manifold HMC [10], is to let the mass depend on position through a metric \(G(q)\), a symmetric positive-definite matrix field:

Where \(G\) is large the particle is heavy and moves slowly; where \(G\) is small it moves fast. The geometry adapts locally to the target. Without the potential terms, the trajectories of \(H = \tfrac12 p^\top G(q)^{-1} p\) are the geodesics of \(G\): the particle follows the straight lines of a curved space.

The price is that \(H\) no longer splits into a function of \(q\) plus a function of \(p\), so the explicit leapfrog no longer applies. The generalised leapfrog replaces it [10]:

The first two equations are implicit and are solved by a few fixed-point iterations. The scheme is still symplectic and reversible, so the determinant is still one and everything in section 3 carries over. What changes is what the network can learn: with a learnable metric \(G_\theta\), the flow learns a geometry as well as a landscape, and the metric itself shows where the transport stretches or compresses space. This is the direction I study in [11].

5. From flow to sampler: a surrogate for HMC

Hamiltonian Monte Carlo samples \(\pi\) by setting \(U = -\log \pi\), drawing a momentum, running \(L\) leapfrog steps on \(H = U + K\), and accepting the end point \((q',p')\) with probability

Every leapfrog step needs \(\nabla U\). When \(\pi\) involves an expensive forward model (a cosmological simulation, a PDE solve), these gradients dominate the cost.

The surrogate idea [12][13][14] is to run the trajectory with a cheap approximation \(\tilde U\), and to keep the exact \(U\) only in the acceptance step. This remains exact: the leapfrog map built on \(\tilde U\) is still reversible and volume-preserving, so the Metropolis correction with the true \(H\) leaves \(\pi\) invariant. The surrogate only affects how often proposals are accepted. Each proposal then costs \(L\) cheap gradients and a single evaluation of \(U\). Li et al. [14] use a neural network as the surrogate, which brings us directly to the flows above.

A trained Hamiltonian flow provides such a surrogate. With a fixed kinetic energy, its potential \(V_\theta\) is a cheap, differentiable function, learned so that Hamiltonian dynamics carries mass towards \(\pi\). Plugged into the proposal, it keeps the chain exact; the closer it is to \(U\) up to a constant, the higher the acceptance rate. The Riemannian version could play the same role for Riemannian manifold HMC, where computing the metric is itself a large part of the cost.

6. What I take from this

Hamiltonian flows turn an engineering constraint (a cheap Jacobian) into a physical property (Liouville’s theorem). What remains to learn has a meaning: a potential, and possibly a geometry. And because these flows use the same integrator as HMC, a model trained as a generator can be reused inside an exact sampler.

Choose the dynamics so that the hard part of the computation vanishes, and spend the network’s capacity on the physics.

References

1. D. J. Rezende, S. Mohamed. Variational inference with normalizing flows. ICML, 2015. arXiv:1505.05770
2. L. Dinh, J. Sohl-Dickstein, S. Bengio. Density estimation using Real NVP. ICLR, 2017. arXiv:1605.08803
3. I. Kobyzev, S. J. D. Prince, M. A. Brubaker. Normalizing flows: an introduction and review of current methods. IEEE TPAMI, 2021. arXiv:1908.09257
4. G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. JMLR 22(57), 2021. arXiv:1912.02762
5. R. M. Neal. MCMC using Hamiltonian dynamics. Handbook of Markov Chain Monte Carlo, 2011. arXiv:1206.1901
6. P. Toth, D. J. Rezende, A. Jaegle, S. Racanière, A. Botev, I. Higgins. Hamiltonian generative networks. ICLR, 2020. arXiv:1909.13789
7. D. J. Rezende, S. Racanière, I. Higgins, P. Toth. Equivariant Hamiltonian flows. 2019. arXiv:1909.13739
8. V. Souveton, A. Guillin, J. Jasche, G. Lavaux, M. Michel. Fixed-kinetic neural Hamiltonian flows for enhanced interpretability and reduced complexity. AISTATS, PMLR 238, 2024. PMLR
9. V. Souveton, S. Terrana. Hamiltonian normalizing flows as kinetic PDE solvers: application to the 1D Vlasov–Poisson equations. ECAI workshop, PMLR 277, 2025. PMLR code
10. M. Girolami, B. Calderhead. Riemann manifold Langevin and Hamiltonian Monte Carlo methods. J. R. Stat. Soc. B 73(2):123–214, 2011. RePEc
11. V. Souveton. Riemannian neural Hamiltonian flows: geodesic symplectic transport and interpretability. Preprint, 2026. arXiv:2609.21647 code
12. C. E. Rasmussen. Gaussian processes to speed up hybrid Monte Carlo for expensive Bayesian integrals. Bayesian Statistics 7, 2003. MPG
13. C. Zhang, B. Shahbaba, H. Zhao. Hamiltonian Monte Carlo acceleration using surrogate functions with random bases. Statistics and Computing 27(6), 2017. arXiv:1506.05555
14. L. Li, A. Holbrook, B. Shahbaba, P. Baldi. Neural network gradient Hamiltonian Monte Carlo. Computational Statistics 34(1), 2019. arXiv:1711.05307