Hamiltonian Monte Carlo (HMC)

Adapted from an appendix of my MS thesis.

Hamiltonian Monte Carlo

Hamiltonian Monte Carlo (HMC) is a type of MCMC method that makes use of gradients to generate new proposed states. HMC attempts to avoid the random walk behavior typical of Metropolis-Hastings by using the gradient to propose new positions far from the current one with high acceptance probability. This allows HMC to better scale to higher dimensions and in principle more complex geometries [1].

In physical systems, the Hamiltonian is a description of the total energy. The total energy can be decomposed into two terms, the kinetic and the potential energy. We can write the Hamiltonian of such a system as the following where $K(\boldsymbol{p},\boldsymbol{q})$ is the kinetic energy and $V(\boldsymbol{q})$ is the potential energy [1].

$$H(\boldsymbol{q},\boldsymbol{p}) = K(\boldsymbol{p},\boldsymbol{q}) + V(\boldsymbol{q}).$$

The probability of being at a particular position with a particular momentum is given as follows [1]. The exponential is a non-negative quantity and can therefore be viewed as an unnormalized probability distribution. The introduction of the minus sign in the exponential is simply a convention, meaning that higher values of energy correspond to lower values of probability [2].

$$p(\boldsymbol{q},\boldsymbol{p}) = \exp\left(-H(\boldsymbol{q},\boldsymbol{p})\right).$$

To simulate such systems the Hamiltonian equations must be solve [1].

$$\begin{aligned} \frac{\mathrm{d}\boldsymbol{q}}{\mathrm{d}t} &= \frac{\partial H}{\partial \boldsymbol{p}} = \frac{\partial K}{\partial \boldsymbol{p}} + \frac{\partial V}{\partial \boldsymbol{p}} \\ \frac{\mathrm{d}\boldsymbol{p}}{\mathrm{d}t} &= -\frac{\partial H}{\partial \boldsymbol{q}} = -\frac{\partial K}{\partial \boldsymbol{q}} - \frac{\partial V}{\partial \boldsymbol{q}}.\end{aligned}$$

The potential energy $V(\boldsymbol{q})$ is given by the probability we wish to sample from $p(\boldsymbol{q})$ . For the momentum $K(\boldsymbol{p},\boldsymbol{q})$ we can invoke an auxiliary variable, and by choosing $p(\boldsymbol{p}|\boldsymbol{q})$ we can write $p(\boldsymbol{q},\boldsymbol{p}) = p(\boldsymbol{p}|\boldsymbol{q})p(\boldsymbol{q})$ . This ensures we can recover the target distribution by marginalizing out the momentum. By introducing the auxiliary variable we can rewrite the equation [1].

$$H(\boldsymbol{q},\boldsymbol{p}) = -\log p(\boldsymbol{p}|\boldsymbol{q}) - \log p(\boldsymbol{q}).$$

We are free to choose the kinetic energy, and if we choose it to be Gaussian and drop the normalization constant, then we have the following where $\boldsymbol{M}$ is the precision matrix [1].

$$K(\boldsymbol{p},\boldsymbol{q}) = \frac{1}{2}\boldsymbol{p}^ \top\boldsymbol{M}^ {-1}\boldsymbol{p} + \log \left|\boldsymbol{M}\right|.$$

If we choose $\boldsymbol{M}$ to be the identity matrix we have $K(\boldsymbol{p},\boldsymbol{q})=\frac{1}{2}\boldsymbol{p}^ \top\boldsymbol{p}$ , making $\frac{\partial K}{\partial \boldsymbol{p}}=\boldsymbol{p}$ and $\frac{\partial K}{\partial \boldsymbol{q}}=0$ , and then we can simplify the Hamiltonian equations [1].

$$\begin{aligned} \frac{\mathrm{d}\boldsymbol{q}}{\mathrm{d}t} &= \boldsymbol{p} \\ \frac{\mathrm{d}\boldsymbol{p}}{\mathrm{d}t} &= -\frac{\partial V}{\partial\boldsymbol{q}}.\end{aligned}$$

Therefore, the HMC algorithm can be summarized as follows [1].

1. Sample $\boldsymbol{p}\sim\mathcal{N}(0,\boldsymbol{I})$

2. Simulate $\boldsymbol{q}_ t$ and $\boldsymbol{p}_ t$ for some amount of time $T$

3. $\boldsymbol{q}_ T$ is the new proposed state

4. Use the Metropolis acceptance criterion to accept or reject $\boldsymbol{q}_ T$ .

In an idealized model, energy is perfectly conserved. However, we use the Metropolis acceptance criterion to correct for errors introduced by the numerical simulation of the Hamiltonian equations. To compute the Hamiltonian equations we have to compute a trajectory. That is, all the intermediate points between one state and the next. In practice this involves computing a series of small integration steps. The most popular method is the leapfrog integrator. Leapfrog integration updates positions $\boldsymbol{q}_ t$ and momentums $\boldsymbol{p}_ t$ at interleaved time points, staggered in such that they leapfrog over each other [1].

An efficient HMC run requires proper tuning of its hyperparameters [1]:

• The time discretization (step size of the leapfrog)

• The integration time (number of leapfrog steps)

• The precision matrix $\boldsymbol{M}$ that parametrizes the kinetic energy.

For example, if the step size is too large then the leapfrog integrator will be inaccurate and too many proposals will be rejected. However, if it is too small we waste computation resources. If the number of steps is too small, the simulated trajectory at each iteration will be too short and sampling will fall back to random walk. But if too large the trajectory might run in circles and we again waste computation resources. If the estimated covariance (inverse of the precision matrix) is too different from the posterior covariance, the proposal momentum will be suboptimal and movement in the position space will be too large or too small in some dimension [1].

References

1. Martin, Osvaldo A., Kumar, Ravin, Lao, Junpeng (2021) Bayesian Modeling and Computation in Python. Chapman and Hall/CRC.
2. Bishop, Christopher M., Bishop, Hugh (2023) Deep Learning: Foundations and Concepts. Springer Cham.