Metropolis Augmented Hamiltonian Monte Carlo
Abstract
Hamiltonian Monte Carlo (HMC) is a powerful Markov Chain Monte Carlo (MCMC) method for sampling from complex high-dimensional continuous distributions.
However, in many situations it is necessary or desirable to combine HMC with other Metropolis-Hastings (MH) samplers.
The common HMC-within-Gibbs strategy implies a trade-oﬀ between long
HMC trajectories and more frequent other MH updates.
Addressing this trade-oﬀ has been the focus of several recent works.
In this paper we propose Metropolis Augmented Hamiltonian Monte Carlo (MAHMC), an HMC variant that allows MH updates within HMC and eliminates this trade-oﬀ. Experiments on two representative examples demonstrate
MAHMC’s eﬃciency and ease of use when compared with within-Gibbs alternatives.
1. Introduction
Hamiltonian Monte Carlo (HMC) is a popular Markov Chain Monte Carlo (MCMC) method.
It samples from complex high-dimensional distributions π ( q ) ∝ exp( − U ( q )) on continuous variables q ∈ R n , where U ( q ) : R n → R is commonly referred to as the potential energy.
HMC has enjoyed remarkable empirical success, due to the use of powerful symplectic integrators (Leimkuhler and Reich, 2004)
(e.g. the leapfrog integrator) to maintain high acceptance probabilities for long-range gradients-guided proposals (Neal, 2012b; Betancourt, 2017).
Common HMC implementations involve simulating the Hamiltonian dynamics for multiple leapfrog steps, followed by a Metropolis Hastings (MH) correction.
Chen et al.
(2020) recently rigorously establishes the importance of using multiple leapfrog steps/long trajectories for HMC’s eﬃciency, especially when compared with the Metropolis Adjusted
Langevin Algorithm (MALA) (Dwivedi et al., 2019), a special case of HMC using only one leapfrog step and a widely used algorithm in Bayesian statistics and machine learning.
However, we often face distributions of the form π ( q H , q O )
∝ exp( − U ( q H , q O )), where we can only use HMC for the continuous variables q H ∈ R n , and it is necessary or desirable to use some other MH samplers for the variables q O (Sec. 4.3 of Neal (2012b)).
Some common situations include (1) when q O are discrete, (2) when we have specialized MH samplers for q O that are eﬃcient/easy to use, or (3) when q O are continuous but ∇ q O U ( q H , q O ) is expensive or impossible to compute.
In such cases, we typically adopt an HMC-withinGibbs strategy (Neal, 2012a; Dang et al., 2019; Kelly et al., 2021), where we alternate between HMC updates and other MH updates.
However, this implies a trade-oﬀ between long HMC trajectories and more frequent other MH updates.
Longer HMC trajectories can help suppress random walk behavior, but might hurt overall sampling since other MH updates can only be done infrequently (between HMC updates).
Shorter HMC trajectories allow more frequent other MH updates, but lead to increased random walk behavior.
[Table] [Table] def leapfrog ( q, p, (cid:15) | U ) p ← p − (cid:15) 2 ∇ U ( q ); q ← q + (cid:15)p ; p ← p − (cid:15) 2 ∇ U ( q ); return q, p ; end def MH correction ( q 0 , p 0 , q, p, v | U, K )
E 0 ← U ( q 0 )
+ K ( p 0 );
E ← U ( q ) + K ( q ); if | v | ≤ exp( − E + E 0 ) then v ← v exp( − E 0 + E ) else q, p ← q 0 , p 0 ; return q, p, v ; end def multiple leapfrogs ( q, p, (cid:15), L | U ) for 1 to n do q, p ← leapfrog ( q, p, (cid:15) | U ) ; p ← − p ; return q, p ; end
Several recent works have focused on the above trade-oﬀ. Neal (2020) combines MALA with partial momentum refreshment (Horowitz, 1991) and non-reversible Metropolis accept/reject decisions to allow more frequent other MH updates while suppressing random walk behavior.
Zhou (2020, 2021) proposes mixed HMC (M-HMC) for distributions with mixed discrete and continuous variables to allow making discrete MH updates as part of an HMC trajectory.
In this paper, we propose Metropolis Augmented Hamiltonian Monte
Carlo (MAHMC) to completely eliminate this trade-oﬀ. MAHMC interprets HMC as a deterministic MH proposal, and augments HMC by introducing MH updates as part of the
HMC trajectory.
MAHMC is generally applicable, can be trivially implemented on top of
HMC with no additional overhead, and includes M-HMC with Laplace momentum as a special case.
We demonstrate MAHMC’s advantage over within-Gibbs alternatives (Neal, 2020) and ease of use on two representative examples in Sec. 4.
2. Background 2.1. HMC
HMC introduces auxiliary momentum p ∈ R n associated with kinetic energy K ( p )
= (cid:80) ni =1 p 2 i / 2, and simulates the Hamiltonian dynamics for L steps of size (cid:15) using the leapfrog integrator ( leapfrog in Algo. 1), before making the ﬁnal MH correction ( HMC in Algo. 2).
We observe that we can in fact interpret the multiple leapfrog steps ( multiple leapfrogs in
Algo. 1) as a deterministic MH proposal, and derive the ﬁnal MH correction as the usual
MH acceptance probability for the deterministic proposal (Sec.
5.2 of Betancourt (2017)).
2.2. MALA and MALA variants
Using L = 1 leapfrog step in HMC results in the widely used special case MALA.
MALA uses gradients information, and allows more frequent other MH updates, but suﬀers from random def HMC ( q 0 , (cid:15), L | U, K ) //
Sec. 2.
p 0 ∼ N (0 , I n ); q, p ← q 0 , p 0 ; q, p ← multiple leapfrogs ( q, p, (cid:15), L | U ); // Defined in Algo.q, p, v ← MH correction ( q 0 , p 0 , q, p, Uniform (0 , 1)
| U, K ); // Defined in Algo.return q ; end def MALA-P ( q 0 , p 0 , (cid:15), α | U, K ) //
Sec. 2.
n ∼ N (0 , n ); p 0 ← αp 0 + √ 1 − α 2 n ; q, p ← q 0 , p 0 ; q, p ← multiple leapfrogs ( q, p, (cid:15), 1 | U ); // Defined in Algo.q, p, v ← MH correction ( q 0 , p 0 , q, p, Uniform (0 , 1) | U, K ); // Defined in Algo.return q, − p ; end def MALA-PN ( q 0 , p 0 , v, (cid:15), α, δ | U, K ) //
Sec. 2.
n ∼ N (0 , I n ); p 0 ← αp 0 + √ 1 − α 2 n ; q, p ← q 0 , p 0 ; q, p ← multiple leapfrogs ( q, p, (cid:15), 1 | U ); // Defined in Algo.q, p, v ← MH correction ( q 0 , p 0 , q, p, v | U, K ); // Defined in Algo.v ← ( v + 1 + δ ) mod 2 − 1; return q, − p, v ; end walk behavior due to the use of short trajectories with frequent momentum refreshments.
Partial momentum refreshment (Horowitz, 1991) (MALA-P in Algo. 2) was proposed as a possible remedy.
However, as Neal (2020) explains, a rejection would lead MALA-P to almost double back on itself, making it less eﬃcient than HMC with long trajectories.
Neal (2020) proposes a non-reversible scheme for Metropolis accept/reject decisions (MALA-PN in Algo. 2) to maintain the ability to make frequent MH updates while further suppressing the random walk behavior that comes from MALA-P doubling back on itself due to rejections.
MALA-PN produces long rejection-free runs by encouraging rejections to cluster together, and demonstrates improved performance on multiple problems.
2.3. Within-
Gibbs sampling from π ( q H , q O )
∝
exp
( − U ( q H , q O )) , q H ∈ R n
For distributions π ( q H , q O )
∝ exp( − U ( q H , q O )), q H ∈ R n where we want to use MH updates for q O , we refer to the common strategy of alternating between updates of q H ∈ R n (using a sampler for continuous distributions) and MH updates of the other variables
q
O as within-Gibbs sampling.
In our experiments, we use four types of within-Gibbs samplers as baselines: (1) MALA within Gibbs (MALAwG), (2) HMC within Gibbs (HwG), (3)
MALA-P within Gibbs (MALA-PwG), and (4) MALA-PN within Gibbs (MALA-PNwG).
2.4. M-HMC for distributions with mixed discrete and continuous variables
Zhou (2020, 2021) proposes M-HMC, an HMC variant that evolves the discrete and continuous variables in tandem for distributions with mixed support.
M-HMC naturally supports frequent MH updates within long HMC trajectories, and demonstrates improved perfor- def MAHMC ( q H 0 , p H 0 , q O 0 , (cid:15), L | U, K, Q i , i = 1 , . . .
, N O , P D ) //
Sec. 3.
q H , p H , q O ← q H 0 , p H 0 , q O 0 ; D ∼ P D ( · ); ∆ E ← 0; for j ← 1 to L do if D j = 0 then q H , p H ← leapfrog ( q H , p H , (cid:15)
| U ( · , q O )); // Defined in Algo.else ˜ q O ∼ Q D j ( ·| q H , q O ); if Uniform (0 , 1) ≤ exp( − U ( q H , ˜ q O ))
Q Dj ( q O | q H , ˜ q O ) exp( − U ( q H ,q O ))
Q Dj (˜ q O | q H ,q O ) then q O ← ˜ q O ; ∆ E ← ∆ E + U ( q H , ˜ q O ) − U ( q H , q O ); end
E ← U ( q H , q O ) + K ( p H ); E 0 ← U ( q H 0 , q O 0 )
+ K ( p H 0 ); if Uniform (0 , 1) ≤ exp( − E ) exp( − E 0 ) P D ( D − 1 )exp(∆ E ) P D
( D ) then p H ← − p H ; else q H , p H , q O ← q H 0 , p H 0 , q O 0 ; return q H , p H , q O end mance over strong baselines.
In Sec. 3.2, we show that the practically useful M-HMC implementation (with Laplace momentum) can be seen as a special case of MAHMC.
3. Metropolis Augmented Hamiltonian Monte Carlo (MAHMC) 3.1. Augmenting HMC with MH updates
Motivated by the interpretation of HMC as a deterministic MH proposal (Sec. 2.1), for a given distribution π ( q H , q O )
∝ exp( − U ( q H , q O )), MAHMC allows more frequent MH updates for q O by combining leapfrog steps for q H and MH updates for q O (including the
MH correction) into a single MH proposal, followed by an additional ﬁnal MH correction.
Formally, for some given step size (cid:15) and number of steps L , an MAHMC iteration makes use of the leapfrog integrator leapfrog ( q, p, (cid:15) | U ( · , q O ))
(Algo. 1) and N O MH proposals
Q i (˜ q O | q H , q O ) , i = 1 , . . . , N O to construct an MH proposal making L total updates of q H , q O .
For a sequence D ∈ (cid:8) 0 , 1 , . . .
, N O (cid:9) L of L integers, deﬁne D − 1 =
( D L , D L − 1 , . . .
, D 1 ).
Starting from q H 0 , q O 0 , MAHMC ﬁrst resamples the momentum p H 0 from N (0 , I n ), then samples a sequence of L variable updates represented as a sequence D of L integers from some distribution P D ( D ), applies the variable updates one at a time, before making a ﬁnal
MH correction.
See Algo.
3 for a detailed description of an MAHMC iteration.
Critically, the use of MH correction in each MH update serves as a mechanism to prevent
MAHMC from deviating too much into low-probability regions.
As we empirically verify in
Sec. 4, even with the additional MH updates as part of the trajectory, MAHMC can maintain high acceptance probabilities for long-range proposals, similar to HMC.
This eliminates the need to balance long HMC trajectories and more frequent other MH updates, and contributes to MAHMC’s improved performance over other within-Gibbs alternatives.
3.2. Connections to M-HMC
Zhou (2020, 2021) derives M-HMC using auxiliary Hamiltonian dynamics for the discrete variables.
We note that the practically useful M-HMC implementation using Laplace momentum is in fact a variant of a special case of MAHMC, where P D is implicitly deﬁned using the auxiliary Hamiltonian dynamics and the proposals Q i are for single discrete variables (given all other discrete and continuous variables).
However, M-HMC diﬀers from MAHMC in its use of a persistent kinetic energy k D for the MH corrections in all MH updates.
3.3. MAHMC in Algo. 3 satisﬁes detailed balance with respect to π ( q H , q O )
Proof Sketch To establish detailed balance, we make use of the concept of probabilistic paths , similar to Sec.
2.3 in Zhou (2020, 2021).
Starting from s =
( q H 0 , p H 0 , q O 0 ), a probabilistic path t contains information about all randomness in an MAHMC iteration (Algo. 3):
The sequence D of L integers specifying which updates to use at each of the L steps.
The actual states q H j , p H j , q O j , j = 1 , . . .
, L (before ﬁnal MH correction) at each step.
The sequence of proposed states ˜ q Hj , ˜ p
Hj , ˜ q
Oj , j = 1 , . . .
, L at each step.
• If D j = 0, ˜ q Hj , ˜ p
Hj =
leapfrog ( q Hj , p Hj , (cid:15) | U ( · , q Oj )) , ˜ q
Oj = q
Oj − 1 .
If D j > 0, ˜ q Hj , ˜ p Hj =
q Hj , p Hj , ˜ q
Oj is a sample from Q D j ( ·| q Hj , q Oj ).
The sequence of accept/reject decisions a j ∈ { True , False } , j = 1 , . . .
, L at each step.
Note that we always accept ( a j = True ) for leapfrog updates ( D j = 0).
We can think of an MAHMC iteration as ﬁrst sampling a probabilistic path t which brings s 0 to s L , where s j =
( q Hj , p Hj , q Oj ) , j = 1 , . . .
, L , before making an MH correction to either accept s L or reject and return to s 0 .
We can reverse a probabilistic path t to get probabilistic path t − 1 which brings s L back to s 0 .
t − 1 uses the sequence of L updates speciﬁed by D − 1 , and reverses the sequence of actual and proposed states (with proper momentum negation) as well as the sequence of accept/rejection decisions.
Denote by P ( t | s 0 ) the probability of sampling t starting from s 0 , and P ( s | t ) , s ∈ { s 0 , s L } the
MH correction step in MAHMC, we can derive the transition probability of MAHMC as
P ( s (cid:48) | s )
= (cid:80) t : s → s (cid:48) P ( s (cid:48) | t ) P ( t | s ).
Deﬁne E ( s ) =
U ( q H , q O ) + K ( p H ).
We can establish the desired detailed balance if we can prove exp( − E ( s ))
P ( s (cid:48) | s ) =
exp( − E ( s (cid:48) ))
P ( s | s (cid:48) ).
We show that P ( s L | t ) = min (cid:110) 1 , exp( − E ( s L )) P ( t − 1 | s L ) exp( − E ( s 0 ))
P ( t | s 0 ) (cid:111) is our desired MH acceptance probability.
Note that exp( − E ( s ))
P ( s | s ) =
exp( − E ( s ))
P
( s | s ) is trivially true.
For s (cid:48) (cid:54) =
s , exp( − E ( s ))
P ( s (cid:48) | s ) =
exp( − E ( s )) (cid:88)
[Equation]
) =
[Equation]
min
[Equation] [Equation]
[Equation] [Equation]
) =
[Equation]
min
[Equation] [Equation]
min (cid:8)
[Equation]
s )
[Equation] [Equation]
) (cid:9) =
exp( − E ( s (cid:48) ))
P ( s | s (cid:48) )
For P ( t − 1 |
s L ) P ( t | s 0 ) , sampling D contributes P D
( D − 1 ) P D ( D ) .
Leapfrog updates have no randomness and contribute nothing.
For D j > 0, deﬁne p j = exp( − U (˜ q Hj ,
˜ q Oj ))
Q Dj ( q Oj | q Hj , ˜ q Oj ) exp( − U ( q Hj ,q Oj ))
Q
Dj (˜ q
Oj | q Hj ,q Oj ) .
If a j = True , the MH update contributes Q Dj ( q Oj | q Hj , ˜ q
Oj )min { 1 , 1 /p j } Q
Dj (˜ q
Oj | q Hj ,q Oj )min { 1 ,
p j } = exp( − U ( q Hj ,q Oj )) exp( − U ( q Hj , ˜ q Oj )) , which is captured in the ∆ E updates in Algo.
If a j = False , the MH update again contributes nothing (as in both t and t − 1 we make the same proposal followed by a rejection).
This proves Algo.uses the correct MH acceptance probability, and establishes the desired detailed balance.
(cid:4)
[Table]
4. Experiments
We use the 4 within-Gibbs samplers in Sec. 2.3 as baselines, and follow the setups of Neal (2020) when possible.
For MALA-based samplers, we alternate between N L leapfrog updates for q H and 1 Gibbs update for q O .
For HwG, we alternate between HMC iterations and Gibbs updates.
Although MAHMC can evolve q H , q O in tandem, to make the comparison most informative, we consider an MAHMC within Gibbs (MAHMCwG) sampler, where we make N U − 1 Gibbs updates and N U N L leapfrog updates within each MAHMC iteration, followed by a Gibbs update.
The MAHMC iteration schedules the Gibbs and leapfrog updates similarly to MALA-based samplers, alternating between N L leapfrog updates and
Gibbs update.
This leads to a deterministic P D that always proposes the same (symmetric) sequence of updates.
Note that the use of a deterministic P D makes comparison with Neal (2020) straightforward, but in general we have more ﬂexibilities in picking P D . A simple choice is to randomly pick the update to make at each step with a ﬁxed distribution over the diﬀerent kinds of available updates (i.e. leapfrog and other updates).
The ﬁxed distribution allows us to control the relative frequency of the diﬀerent updates, and the contribution of
P D to the ﬁnal acceptance probability can be easily calculated.
We assume gradients evaluation dominates the computation (Neal, 2020), and evaluate performance using eﬀective sample size (ESS) (calculated with Kumar et al.
per sample per gradients evaluation.
Code to reproduce the results is available at https: //github.com/StannisZhou/mahmc .
4.1. A distribution with mixed discrete and continuous variables (MDC)
We consider the distribution in Sec.
5 of Neal (2020), where q H = ( u, v ) , q O = ( w 1 , . . w 20 ): u ∼ N (0 , 1) , v | u ∼ N ( u, 0 . 04 2 ) , w i | u ∼ Bernoulli (1 / (1 + e u )) , i = 1 , . . .
Correctness
To verify correctness, we compare the histogram of u samples obtained with the 5 samplers, and empirically verify that they match
Eﬃciency
We summarize the results in Tab. 1.
From the results, we can see that by combining longer HMC trajectories with more frequent Gibbs updates, MAHMCwG is 3.85x more eﬃcient than HwG, and 2.4x more eﬃcient than MALA-PNwG.
For HwG and MALA-PNwG, we use the optimal hyperparameters reported in Sec.of Neal (2020)
( L = 40 steps and step size (cid:15) = 0 .
035 for HwG, and N L = 10 , (cid:15) = 0 .
03 , α = 0 .
995 , δ = 0 .
01 for MALA-PNwG).
Our results roughly reproduce the results in Neal (2020) (MALA-PNwG is 1.6x more eﬃcient than HwG, as opposed to 1.83x reporeted in Neal (2020)).
The small discrepancy can be due to the use of a diﬀerent metric (ESS of u instead of ESS of I ( u ∈ ( − 0 .
as in Neal (2020), where I represents the indicator function) and a diﬀerent way to compute ESS (using Kumar et al.
To make the comparison informative, we use N L = 10 and (cid:15) = 0 . 03 for MALAwG and MALA-PwG, and α = 0 .
995 for MALA-PwG. We observe that partial momentum refreshment and non-reversible Metropolis accept/reject decisions are indeed beneﬁcial in improving the performance of MALAwG/MALA-PwG.
However, only MALA-PNwG outperforms HwG.
The optimal performance of MAHMCwG is achieved with N U = 10 and (cid:15)
= 0 .
04, i.e. in each MAHMC we make N U N L = 100 leapfrog updates and N U − 1 = 9 Gibbs updates (uniformly spreaded).
This is far larger than the optimal number of steps L = 40 for HwG, and demonstrates MAHMC’s ability to maintain high acceptance probabilities for longrange proposals that includes MH updates.
We additionally test MAHMCwG with N U =and (cid:15)
= 0 .
035, i.e. making N U − 1 = 3 additional Gibbs updates on top of HwG, and observe that the normalized ESS increases to 6 .
08 × 10 − 3 .
This demonstrates the beneﬁts of more frequent MH updates when we use the same number of leapfrog updates.
4.2. Bayesian Logistic Regression (BLR) with conjugate prior
We apply Bayesian Logistic Regression (BLR) to the breast cancer wisconsin dataset 1 , consisting of 569 pairs of 30 dimensional features and targets y i ∈ { 0 , 1 } .
We standarize the features and append 1 at the end to get x i ∈ R 31 , and specify BLR with conjugate prior as τ ∼ Gamma (1 . 0 , scale = 100) , β ∼ N (0 , 1 τ I 31 ) , y i ∼ Bernoulli
( sigmoid ( x Ti β )) , i = 1 , . . . ,
We consider sampling from the posterior P ( τ, β | X, y ), where q H = β and q O = τ , and we make Gibbs updates for τ using P ( τ | β, X, y ).
Correctness To verify the correctness of the samplers, we ﬁrst apply the samplers to only the prior distribution
[Equation]
and similarly verify the marginal distribution of τ is indeed Gamma (1 . 0 , scale = 100) by looking at the histogram of τ samples.
We additionally verify that all samplers give rough the same posterior means for β , and we can classify the 569 training data points to an accuracy of 98 .
77% using posterior samples from the 5 diﬀerent samplers.
Eﬃciency
We ﬁx N L = 5, and conduct a grid search in
N U ∈ { 2 , 3 , 4 , 5 } , (cid:15) ∈ { 0 . 07 , 0 .
11 } , α ∈ { 0 . 9 , 0 .
98 , 0 . 99 } and δ ∈ { 0 .
03 } , and L ∈ { 10 , 15 , 20 , 25 } for HwG.
Optimal performance is achieved with (cid:15) = 0 .
11 for MALAwG, L = 10 , (cid:15) = 0 .
09 for
HwG, (cid:15) = 0 .
09 , α = 0 .
9 for MALA-PwG, (cid:15)
= 0 .
1 , α = 0 .
9 , δ = 0 .
015 for MALA-PNwG, and
N U
= 2 , (cid:15)
= 0 .
1 for MAHMCwG.
For this example, MAHMCwG only slightly outperforms MALA-PNwG.
However, we comment that MAHMCwG is easier to tune, as we essentially only need to tune the step size and number of steps, same as in HwG, and using the same setups from HwG usually already gives good performance.
However, for MALA-PwG, we need to tune α and δ , which can have signiﬁcant impacts on performance.
For example, for the distribution in 4.1, grid search is done for α ∈ { 0 . 98 , 0 . 99 , 0 .
[-q_hmc0, q_other0, p0 = q_hmc, q_other, p def scan_f(carry, ii): q_hmc, p, q_other, delta_U, key = carry q_hmc, p = take_multiple_leapfrog_steps(q_hmc, p, q_other, epsilon, L) q_other
-]
995 , 0 . 9975 , 0 . 9985 , 0 . 999 } in Neal (2020).
However, in our experiments we empirically observe that all these values give poor performance, and we have to reduce α to 0 .
9 to get performance comparable with MAHMCwG, suggesting the potential challenges in tuning α .
4.3. Additional veriﬁcation of correctness
Since we only used Gibbs updates (which always accept) in our experiments, we additionally verify the correctness of MAHMC by modifying the script used in Zhou (2020, 2021) to make random walk MH updates within HMC for the 1D GMM example using MAHMC.
See https://github.com/StannisZhou/mahmc/blob/main/simplified_mixed_hmc.py for the updated script.