Betancourt 2015 - Riemannian HMC and the SoftAbs Metric

Summary

Riemannian HMC lets the momentum covariance depend on position, $\pi (p\mid q)=\mathcal{N}(0,\Sigma (q))$, adding a $\frac{1}{2}\log |\Sigma (q)|$ term that acts as an energy reservoir, so potential-energy variation is no longer capped at $d/2$. A good metric resembles the Hessian; the SoftAbs transformation makes the Hessian positive definite. With it, trajectories span the entire funnel and give nearly independent transitions.

Overview

Euclidean HMC fails on hierarchies because of limited density variation and step-size sensitivity (Betancourt 2015 - Euclidean HMC Weaknesses - Step Size and Density Variation). Riemannian HMC (Girolami and Calderhead 2011, ref [21]) fixes both.

Main Content

Riemannian HMC

Allow position-dependent covariance: $\pi (p\mid q)=\mathcal{N}(p\mid 0,\Sigma (q))$ with kinetic energy

$$T(p,q)=\frac{1}{2}{p}^{\top }{\Sigma }^{-1}(q)p+\frac{1}{2}\log |\Sigma (q)|.$$

The Hamiltonian now generates dynamics on a Riemannian manifold with metric $\Sigma$. The dynamic metric acts as a local correction to the target; if well chosen it compensates position-dependent correlations, reducing computational burden and sensitivity to step size.

The log-determinant energy reservoir

The term $\frac{1}{2}\log |\Sigma (q)|$ is nominally the normalisation of the momentum density, but it functions as a reservoir that absorbs and releases energy along the trajectory, allowing much larger variations in potential energy (Fig. 10: $V$ is essentially unconstrained in practice even though $V$ variation is still bounded by $T$ variation, because $T$ now varies greatly).

Choosing the metric. To correct position-dependent geometry, the metric should leave the target locally isotropic, which suggests the Hessian of the target. The Hessian itself is not guaranteed positive definite.

)

Apply the SoftAbs transformation to the eigenvalues of the Hessian to obtain a smooth, positive-definite metric; it yields a generic and efficient RHMC implementation. (Footnote 4: the Fisher-Rao metric from information geometry is an alternative regulariser, but it regularises by integrating out exactly the correlations needed for effective corrections in hierarchical models.) The paper does not restate the formula; the standard SoftAbs map (Betancourt 2013, arXiv:1212.4693) replaces each Hessian eigenvalue $\lambda$ by $\lambda \coth (\alpha \lambda )$ with a regularisation constant $\alpha$.

Figure 11: funnel exploration

RHMC with SoftAbs on the 101-dimensional funnel: trajectories sweep the entire funnel and the chain of $v$ (only 100 iterations plotted, compared with 2500 for Figure 7) has mean $0.42$ and standard deviation $2.76$ against the true $(0,3)$: nearly independent transitions and drastically smaller autocorrelations.

Why $\lambda \coth (\alpha \lambda )$ is positive definite

For $\lambda \to \pm \infty$, $\lambda \coth (\alpha \lambda )\to |\lambda |$; at $\lambda \to 0$, it tends to $1/\alpha$. Hence every eigenvalue maps to a positive number at least $1/\alpha$, behaving like $|\lambda |$ for large curvature. A numpy check:

import numpy as np
alpha = 1e6
soft = lambda lam: lam/np.tanh(alpha*lam)
print(soft(np.array([-5., -1e-3, 1e-3, 5.])))  # ~[5, 1e-3(approx |l|), 1e-3, 5]; no negatives

Cost. RHMC needs Hessians (second derivatives) and an implicit integrator in general, so it is considerably more expensive per step; the paper notes that RHMC was not run in the Stan benchmarks because its Stan implementation was still under development. In practice (2020s) Stan/PyMC use Euclidean NUTS with reparameterisation.

Connections

Contrast with ensemble tuning in Hoffman 2022 - MEADS - Overview, which preconditions with marginal standard deviations only. Position-dependent preconditioning of the same kind appears in normalizing-flow-adapted NUTS (E12, Stochastic Volatility and the Inference Frontier).

See Also

Betancourt 2015 - HMC for Hierarchical Models - Overview, Efficient MCMC