Modified Hamiltonian Monte Carlo for Bayesian inference

Abstract

The Hamiltonian Monte Carlo (HMC) method has been recognized as a powerful sampling tool in computational statistics. We show that performance of HMC can be significantly improved by incorporating importance sampling and an irreversible part of the dynamics into a chain. This is achieved by replacing Hamiltonians in the Metropolis test with modified Hamiltonians and a complete momentum update with a partial momentum refreshment. We call the resulting generalized HMC importance sampler Mix & Match Hamiltonian Monte Carlo (MMHMC). The method is irreversible by construction and further benefits from (i) the efficient algorithms for computation of modified Hamiltonians; (ii) the implicit momentum update procedure and (iii) the multistage splitting integrators specially derived for the methods sampling with modified Hamiltonians. MMHMC has been implemented, tested on the popular statistical models and compared in sampling efficiency with HMC, Riemann Manifold Hamiltonian Monte Carlo, Generalized Hybrid Monte Carlo, Generalized Shadow Hybrid Monte Carlo, Metropolis Adjusted Langevin Algorithm and Random Walk Metropolis–Hastings. To make a fair comparison, we propose a metric that accounts for correlations among samples and weights and can be readily used for all methods which generate such samples. The experiments reveal the superiority of MMHMC over popular sampling techniques, especially in solving high-dimensional problems.

Appendices

Appendix A: Invariance of the PMMC step

The Partial Momentum Monte Carlo step of the MMHMC method leaves the importance target distribution $\tilde{\pi }$ (Eq. 12) invariant if for a transition kernel $T(\cdot |\cdot )$ the following condition is satisfied

$$\begin{aligned} \tilde{\pi }(\varvec{\theta }',\mathbf {p}') =\int \tilde{\pi }(\varvec{\theta },\mathbf {p})T \left( (\varvec{\theta }',\mathbf {p}')|(\varvec{\theta }, \mathbf {p})\right) \mathrm{d}\varvec{\theta } \mathrm{d}\mathbf {p} \end{aligned}$$

for all $n=1,\dots ,N$.

The PMMC step is sampling on a space augmented with a noise vector $\mathbf {u} \sim \mathcal {N}(0,M)$ with the extended density $\hat{\pi }$ (defined in Eq. 17), for which

$$\begin{aligned} \tilde{\pi }(\varvec{\theta },\mathbf {p})=\int \hat{\pi } (\varvec{\theta },\mathbf {p},\mathbf {u})\mathrm{d}\mathbf {u} . \end{aligned}$$

Therefore, we want to show that

$$\begin{aligned}&\hat{\pi }(\varvec{\theta }',\mathbf {p}',\mathbf {u}')\nonumber \\&\quad =\int \hat{\pi }(\varvec{\theta },\mathbf {p},\mathbf {u}) T\left( (\varvec{\theta }',\mathbf {p}',\mathbf {u}')| (\varvec{\theta },\mathbf {p},\mathbf {u})\right) \mathrm{d} \varvec{\theta } \mathrm{d}\mathbf {p} \mathrm{d}\mathbf {u}, \end{aligned}$$

(28)

for the transition kernel defined as

$$\begin{aligned}&T\left( (\varvec{\theta }',\mathbf {p}',\mathbf {u}')| (\varvec{\theta },\mathbf {p},\mathbf {u})\right) \\&\quad =\mathcal {P}\cdot \delta \left( (\varvec{\theta }', \mathbf {p}',\mathbf {u}')-\mathcal {R}(\varvec{\theta }, \mathbf {p},\mathbf {u}) \right) \\&\qquad +\,(1-\mathcal {P})\cdot \delta \left( (\varvec{\theta }', \mathbf {p}',\mathbf {u}')-\hat{\mathcal {F}}(\varvec{\theta }, \mathbf {p},\mathbf {u})\right) , \end{aligned}$$

where $\mathcal {P}=\min \left\{ 1,\hat{\pi }(\mathcal {R} (\varvec{\theta },\mathbf {p},\mathbf {u}))/\hat{\pi } (\varvec{\theta },\mathbf {p},\mathbf {u})\right\}$ is the Metropolis probability, $\delta$ is the Delta function, $\mathcal {R}$ is the proposal function (Eq. 16) and $\hat{\mathcal {F}}(\varvec{\theta },\mathbf {p},\mathbf {u}) =(\varvec{\theta },\mathbf {p},-\mathbf {u})$ is the flipping function. Note that the map $\mathcal {R}$ is volume preserving; hence, the Metropolis probability $\mathcal {P}$ does not include the Jacobian factor. For the sake of clarity, we denote $\mathbf {x}=(\varvec{\theta },\mathbf {p},\mathbf {u})$ and write the right-hand side of the expression (28) as

$$\begin{aligned}&\int T(\mathbf {x}'|\mathbf {x})\hat{\pi }(\mathbf {x}) \mathrm{d}\mathbf {x}\\&\quad =\underbrace{\int \min \left\{ \hat{\pi }(\mathbf {x}),\hat{\pi } (\mathcal {R}(\mathbf {x})) \right\} \cdot \delta \left( \mathbf {x}' - \mathcal {R}(\mathbf {x}) \right) \mathrm{d}\mathbf {x}}_\text {1st term}\\&\qquad \underbrace{+\int \hat{\pi }(\mathbf {x}) \cdot \delta \left( \mathbf {x}' - \hat{\mathcal {F}}(\mathbf {x}) \right) \mathrm{d}\mathbf {x}}_\text {2nd term}\\&\qquad \underbrace{-\int \min \left\{ \hat{\pi }(\mathbf {x}),\hat{\pi }(\mathcal {R} (\mathbf {x})) \right\} \cdot \delta \left( \mathbf {x}'- \hat{\mathcal {F}} (\mathbf {x}) \right) \mathrm{d}\mathbf {x}}_\text {3rd term}. \end{aligned}$$

Applying change of variables $\mathbf {x}=\hat{\mathcal {F}} \circ \mathcal {R}(\bar{\mathbf {x}})$, which is volume preserving, to the first term in the sum, omitting the bars, and using the fact that $\hat{\mathcal {F}} =\mathcal {R}\circ \hat{\mathcal {F}}\circ \mathcal {R}$, one obtains

$$\begin{aligned} \text {1st term}= & {} \int \min \left\{ \hat{\pi }(\hat{\mathcal {F}} \circ \mathcal {R}(\mathbf {x})),\hat{\pi }(\hat{\mathcal {F}} (\mathbf {x})) \right\} \\&\cdot \delta \left( \mathbf {x}' -\hat{\mathcal {F}}(\mathbf {x}) \right) \mathrm{d}\mathbf {x}. \end{aligned}$$

Since $\hat{\pi }\circ \hat{\mathcal {F}}=\hat{\pi }$, the first and third terms cancel out. Employing change of variables $\mathbf {x}=\hat{\mathcal {F}}(\bar{\mathbf {x}})$ to the second term and again omitting the bars lead to

$$\begin{aligned} \text {2nd term}=\int \hat{\pi }(\mathbf {x}) \cdot \delta \left( \mathbf {x}'- \mathbf {x} \right) \mathrm{d}\mathbf {x} =\hat{\pi }(\mathbf {x}'), \end{aligned}$$

which proves the equality (28).

Appendix B: Modified Hamiltonians for splitting integrators

The coefficients for the two-stage integrator family (14) and modified Hamiltonians (22)–(24) are the following:

$$\begin{aligned} c_{21}=&\frac{1}{24}\Big (6b-1\Big ) \nonumber \\ c_{22}=&\frac{1}{12}\Big (6b^2-6b+1\Big )\nonumber \\ c_{41} =&\frac{1}{5760}\Big (7-30b\Big ) \nonumber \\ c_{42} =&\frac{1}{240}\Big (-10b^2+15b-3\Big ) \nonumber \\ c_{43} =&\frac{1}{120}\Big (-30b^3+35b^2-15b+2\Big ) \nonumber \\ c_{44} =&\frac{1}{240}(20b^2-1). \end{aligned}$$

(29)

Using (29), one can also obtain the modified Hamiltonian for the Verlet integrator, since two steps of Verlet integration are equivalent to one step of the two-stage integrator with $b=1/4$. The coefficients are therefore

$$\begin{aligned} c_{21}= & {} \frac{1}{12}, \quad c_{22}=- \frac{1}{24} \\ c_{41}= & {} -\frac{1}{720}, \quad c_{42} = \frac{1}{120}, \quad c_{43} = -\frac{1}{240}, \quad c_{44} =\frac{1}{60}. \end{aligned}$$

For three-stage integrators (15) (a two-parameter family), the coefficients are

$$\begin{aligned} c_{21}= & {} \frac{1}{12} \Big (1-6a(1-a)(1-2b)\Big )\\ c_{22}= & {} \frac{1}{24} \Big (6a(1-2b)^2-1\Big )\\ c_{41}= & {} \frac{1}{720}\Big ( 1 + 2 (a-1) a (8 + 31 (a-1) a) (1 - 2 b) - 4 b\Big )\\ c_{42}= & {} \frac{1}{240} \Big (6 a^3 (1 - 2 b)^2 -a^2 (19 - 116 b + 36 b^2 + 240 b^3)\\&+\,a (27 - 208 b + 308 b^2) - 48 b^2 + 48 b -7\Big ) \\ c_{43}= & {} \frac{1}{180}\Big (1 + 15 a (1- 2 b) (-1 + 2 a (2 - 3 b + a (4 b-2)))\Big ) \\ c_{44}= & {} \frac{1}{240} \Big (-1 + 20 a (1 - 2 b) (b + a (1 + 6 (b-1) b))\Big ). \end{aligned}$$

The coefficients for the modified Hamiltonians (25)–(26) are calculated as

$$\begin{aligned} k_{21}= & {} c_{21}, \quad k_{22}=c_{22}, \\ k_{41}= & {} c_{41}, \quad k_{42}=3c_{41} +c_{42}, \\ k_{43}= & {} c_{41} +c_{44}, \quad k_{44}=3c_{41}+c_{42}+c_{43}. \end{aligned}$$

For the fourth-order modified Hamiltonian (25), we use the second-order centered finite difference approximations of time derivatives of the gradient of the potential function

$$\begin{aligned} \mathbf {U}^{(1)}=\frac{\mathbf {U}(t_{n+1}) -\mathbf {U}(t_{n-1})}{2\varepsilon }, \end{aligned}$$

(30)

with $\varepsilon =h$ for the Verlet, $\varepsilon =h/2$ for two-stage and $\varepsilon =ah$ for three-stage integrators with a being the integrator’s coefficient advancing position variables. The sixth-order modified Hamiltonian (26), here considered only for the Verlet and two-stage integrators, is calculated using fourth-order approximation for the first derivative and second-order approximations for the second and third derivatives

$$\begin{aligned} \mathbf {U}^{(1)}= & {} \frac{\mathbf {U}(t_{n-2})-8\mathbf {U}(t_{n-1}) +8\mathbf {U}(t_{n+1})-\mathbf {U}(t_{n+2})}{12\varepsilon } \\ \mathbf {U}^{(2)}= & {} \frac{\mathbf {U}(t_{n-1})-2\mathbf {U}(t_n) +\mathbf {U}(t_{n+1})}{\varepsilon ^2}\\ \mathbf {U}^{(3)}= & {} \frac{-\mathbf {U}(t_{n-2})+2\mathbf {U}(t_{n-1}) -2\mathbf {U}(t_{n+1})+\mathbf {U}(t_{n+2})}{2\varepsilon ^3}, \end{aligned}$$

where $\varepsilon$ depends on the integrator as before. The interpolating polynomial in terms of the gradient of the potential function $\mathbf {U}(t_i)=U_{\varvec{\theta }} (\varvec{\theta }^i), \; i=n-k,\dots ,n,\dots ,n +k, n \in \{0, L\}$ is constructed from a numerical trajectory $\{U_{\varvec{\theta }}(\varvec{\theta }^i\}^{L+k}_{i=-k}$ where $k=1$ and $k=2$ for the fourth- and sixth-order modified Hamiltonians, respectively.

Appendix C: Modified PMMC step

In the modified PMMC step proposed for MMHMC, a partial momentum update is integrated into the modified Metropolis test; i.e., it is implicitly present in the algorithm. This reduces the frequency of derivative calculations in the Metropolis function. To implement this idea, we recall that the momentum update probability

$$\begin{aligned} \mathcal {P}=\min \left\{ 1,\frac{\exp \big (-\hat{H}(\varvec{\theta }, \mathbf {p}^{*},\mathbf {u}^{*})}{\exp \big (-\hat{H}(\varvec{\theta }, \mathbf {p},\mathbf {u})\big )}\right\} \end{aligned}$$

(31)

depends on the error in the extended Hamiltonian (18). Let us first consider the fourth-order modified Hamiltonian (22) with analytical derivatives of the potential function. It is easy to show that the difference in the extended Hamiltonian (18) between a current state and a state with partially updated momentum is

$$\begin{aligned} \varDelta {\hat{H}}= & {} {U(\varvec{\theta })} + \frac{1}{2} (\mathbf {p}^{*})^\mathrm{T} M^{-1} \mathbf {p}^{*}\nonumber \\&+\, h^2c_{21}(\mathbf {p}^{*})^\mathrm{T} M^{-1}U_{\varvec{\theta } \varvec{\theta } }(\varvec{\theta } ) M^{-1} \mathbf {p}^{*}\nonumber \\&+\, {h^2c_{22}U_{\varvec{\theta } }(\varvec{\theta } ) M^{-1} U_{\varvec{\theta }}(\varvec{\theta })} + \frac{1}{2} (\mathbf {u}^{*})^\mathrm{T} M^{-1}\mathbf {u}^{*} \nonumber \\&-\, {U(\varvec{\theta })} - \frac{1}{2} \mathbf {p}^\mathrm{T} M^{-1} \mathbf {p} - h^2c_{21}\mathbf {p}^\mathrm{T} M^{-1}U_{\varvec{\theta } \varvec{\theta } }(\varvec{\theta } ) M^{-1} \mathbf {p} \nonumber \\&-\, {h^2c_{22}U_{\varvec{\theta }}(\varvec{\theta }) M^{-1} U_{\varvec{\theta }}(\varvec{\theta })} - \frac{1}{2} \mathbf {u}^\mathrm{T} M^{-1} \mathbf {u} \nonumber \\= & {} {h^2}c_{21}\Big (\varphi A + 2\sqrt{\varphi (1-\varphi )}B \Big ) \end{aligned}$$

(32)

with

$$\begin{aligned} A= & {} (\mathbf {u}-\mathbf {p})^\mathrm{T} U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) (\mathbf {u}+\mathbf {p}) \nonumber \\ B= & {} \mathbf {u}^\mathrm{T}U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) \mathbf {p}. \end{aligned}$$

(33)

For the sixth-order modified Hamiltonian (24) for Gaussian problems, the error in the extended Hamiltonian (18) can be calculated in a similar manner

$$\begin{aligned} \varDelta {\hat{H}}= & {} {h^2}c_{21}\Big (\varphi (A-B) + 2\sqrt{\varphi (1-\varphi )}C \Big )\nonumber \\&+\,{h^4}c_{44}\Big (\varphi (D-E) + 2\sqrt{\varphi (1-\varphi )}F \Big ), \end{aligned}$$

(34)

with

$$\begin{aligned} A&=\mathbf {u}^\mathrm{T} U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } )\mathbf {u} \\ B&=\mathbf {p}^\mathrm{T}U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) \mathbf {p}\\ C&=\mathbf {u}^\mathrm{T}U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) \mathbf {p}\\ D&=(U_{\varvec{\theta } \varvec{\theta } }(\varvec{\theta } ) \mathbf {u})^\mathrm{T}U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) \mathbf {u}\\ E&=(U_{\varvec{\theta } \varvec{\theta } }(\varvec{\theta } ) \mathbf {p})^\mathrm{T}U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) \mathbf {p}\\ F&=(U_{\varvec{\theta } \varvec{\theta } }(\varvec{\theta } ) \mathbf {u})^\mathrm{T}U_{\varvec{\theta } \varvec{\theta } } (\varvec{\theta } ) \mathbf {p}. \end{aligned}$$

Therefore, if the modified Hamiltonians (22)–(24) with analytical derivatives are used, a new momentum can be determined as

$$\begin{aligned} \bar{\mathbf {p}}=\left\{ \begin{array}{l l} \sqrt{1-{\varphi }}\mathbf {p}+\sqrt{{\varphi }}\mathbf {u} &{} \text{ with } \text{ probability } \\ &{} \qquad \mathcal {P}= \min \{1,\exp (-\varDelta \hat{H})\}\\ \mathbf {p} &{} \text{ otherwise, } \end{array}\right. \end{aligned}$$

(35)

where $\mathbf {u} \sim \mathcal {N}(0,M)$ is the noise vector, $\varphi \in \left( 0,1\right]$ and $\varDelta \hat{H}$ is defined as in (32) or (34).

Consequently, for models with no hierarchical structure, there is no need to calculate gradients within the PMMC step, second derivatives can be taken from the previous Metropolis test within the HDMC step and there is no need to generate $\mathbf {u}^{*}$.

If the modified Hamiltonians are calculated using numerical time derivatives of the gradient of the potential function, for the Verlet, two- and three-stage integrators as in (25)–(26), the difference in the fourth-order extended Hamiltonian becomes

$$\begin{aligned} \varDelta {\hat{H}}= {h}k_{21}\Big ((\mathbf {p}^{*})^\mathrm{T} P^{*}_1 -\mathbf {p}^\mathrm{T}P_1 \Big ), \end{aligned}$$

(36)

whereas for the sixth-order extended Hamiltonian it is

$$\begin{aligned} \varDelta {\hat{H}}= & {} hk_{21}\Big ((\mathbf {p}^{*})^\mathrm{T} P^{*}_1 -\mathbf {p}^\mathrm{T}P_1 \Big )\\&+\,hk_{41}\Big ((\mathbf {p}^{*})^\mathrm{T} P^{*}_3 -\mathbf {p}^\mathrm{T}P_3 \Big )\\&+\,h^2k_{42}\Big (U_{\mathbf {x}}^\mathrm{T} P^{*}_2 - U_{\mathbf {x}}^\mathrm{T}P_2 \Big )\\&+\,h^2k_{43}\Big ((P^{*}_1)^\mathrm{T} P^{*}_1 -P_1^\mathrm{T}P_1 \Big ). \end{aligned}$$

Here, $P^{*}_1, P^{*}_2,P^{*}_3$, are the first-, second- and third-order scaled time derivatives of the gradient, respectively (see Sect. 2.3.3), calculated from the trajectory with updated momentum $\mathbf {p}^{*}$. The computational gain of the new PMMC step, in this case, results from skipping a calculation of the terms multiplying $k_{22}$ in (25) and $k_{44}$ in (26). It has to be admitted that the term multiplying $k_{22}$ in (25) is of negligible cost, and thus the gain from using the new momentum update is not as significant as in the case of modified Hamiltonians with analytical derivatives. On the contrary, the saving in computation arising from the absence of the term multiplying $k_{44}$ in the sixth-order modified Hamiltonian (26) is essential.

In summary, in the case of the sixth-order modified Hamiltonian, with derivatives calculated either analytically or numerically, the proposed momentum refreshment enhances computational performance of MMHMC. This also applies to the cases when the fourth-order modified Hamiltonian with analytical derivatives is used. In this situation, however, if the Hessian matrix of the potential function is dense, instead of using the modified Hamiltonian with analytical derivatives, we recommend using numerical derivatives, for which the saving is negligible. On the other hand, if the computation of the Hessian matrix is not very costly (e.g., being block diagonal, sparse, close to constant), it might be more efficient to use analytical derivatives, for which the new formulation of the Metropolis test leads to computational saving.

Appendix D: Algorithmic summary

We provide the algorithmic summary for HMC (Algorithm 1) and two alternative algorithms for the MMHMC method. One (Algorithm 2) uses the modified Hamiltonians defined through analytical derivatives of the potential function and is recommended for the problems with sparse Hessian matrices. The other algorithm (Algorithm 3) relies on the modified Hamiltonians expressed through numerical time derivatives of the gradient of the potential function. This algorithm, although including additional integration step, is beneficial for cases where higher order derivatives are computationally demanding.

Table 5 Parameter values used for the multivariate Gaussian model experiments

Appendix E: Experimental setup

See Tables 5, 6 and 7.

Table 6 Parameter values used for the Bayesian logistic regression model experiments

Table 7 Parameter values used for the stochastic volatility model experiments