Gradient-based tuning of Hamiltonian Monte Carlo hyperparameters

1 INTRODUCTION . Hamiltonian Monte Carlo ( HMC ) is a popular sampling based method for performing accurate inference on complex distributions that we may only know up to a normalization constant ( Neal , 2011 ) . Unfortunately , HMC can be slow to run in practice as we need to allow time for the simulation to ‘ burn-in ’ and also to sufficiently explore the full extent of the target distribution . Tuning the HMC hyperparameters can help alleviate these issues but this requires domain expertise and must be repeated for every problem HMC is applied to . There have been many attempts to provide an automatic method of tuning the hyperparameters . Some methods use a proxy for the mixing speed of the chain , i.e . the speed at which the Markov chain marginal distribution approaches the target . For example , Levy et al . ( 2018 ) use a variation on the expected squared jumped distance to tune parameters in order to encourage the chain to make large moves within the sample space . Other methods draw upon ideas from Variational Inference ( VI ) . VI ( Jordan et al. , 1999 ) is an optimization based method that is often contrasted to Markov Chain Monte Carlo methods such as HMC . In VI , we approximate the target using a parametric distribution , reducing the approximation bias through optimizing the distribution parameters . The optimization procedure maximises a lower bound on the normalization constant of the target which is equivalent to minimising the KL-divergence between the approximation and the target . To apply this idea to HMC , Salimans et al . ( 2015 ) ; Wolf et al . ( 2016 ) consider the marginal distribution of the final state in a finite length HMC chain as an implicit variational distribution with the intention of tuning the HMC parameters using the VI approach . However , the implicit distribution makes the usual variational lower bound intractable . To restore tractability , they make the bound looser by introducing an auxiliary inference distribution approximating the reverse dynamics of the chain . The looseness of the bound depends on the KL-divergence between the auxiliary inference distribution and the true reverse dynamics . As the chain length increases , the dimensionality of these distributions increases , tending to increase the looseness of the bound . This causes issues during optimization because the increasing magnitude of this extra KL-divergence term encourages the model to fit to the imperfect auxiliary inference distribution as opposed to the target as desired . Indeed , Salimans et al . ( 2015 ) only consider very short HMC chains using their method . In this work , we further investigate the combined VI-HMC approach as this has the potential to provide a direct measure of the chain ’ s convergence without the need to rely on proxies for per- formance . When applied to an implicit HMC marginal distribution , the variational objective can be broken down into the tractable expectation of the log target density and the intractable entropy of the variational approximation . This entropy term prevents a fully flexible variational distribution from collapsing to a point mass maximizing the log target density . Since HMC , by construction , can not collapse to such a point mass , we argue that the entropy term can be dropped provided the initial distribution of the chain has enough coverage of the target . We evaluate our proposed method on a variety of tasks . We first consider a range of synthetic 2D distributions before moving on to higher dimensional problems . In particular , we use our method to train deep latent variable models on the MNIST and FashionMNIST datasets . We also evaluate on a popular statistical mechanics benchmark : sampling molecular configurations from the Boltzmann distribution of a 22 atom molecule , Alanine Dipeptide . Our results show that this method is competitive with or can improve upon alternative tuning methods for HMC on all problems we consider . 2 BACKGROUND . 2.1 HAMILTONIAN MONTE CARLO . HMC is a Markov Chain Monte Carlo method ( Neal , 1993 ) which aims to draw samples from the n-dimensional target distribution p ( x ) = 1Z p ∗ ( x ) where Z is the ( usually unknown ) normalization constant . It introduces an auxiliary variable ν ∈ Rn , referred to as the momentum , which is distributed according to N ( ν ; 0 , diag ( m ) ) , with the resulting method sampling on the extended space ζ = ( x , ν ) . HMC progresses by first sampling an initial state from some initial distribution and then iteratively proposing new states and accepting/rejecting them according to an acceptance probability . To propose a new state , first , a new value for the momentum is drawn from N ( ν ; 0 , diag ( m ) ) , then , we simulate Hamiltonian Dynamics with Hamiltonian , H ( x , ν ) = −log p∗ ( x ) + 12ν T diag ( m ) −1ν arriving at new state ( x′ , ν′ ) . This new state is ac- cepted with probability min [ 1 , exp ( −H ( x′ , ν′ ) +H ( x , ν ) ) ] otherwise we reject the proposed state and remain at the starting state . The Hamiltonian Dynamics are simulated using a numerical integrator , the leapfrog integrator ( Hairer et al. , 2003 ) being a popular choice . To propose the new state , L leapfrog updates are taken , each update consisting of the following equations : νk+ 12 = νk + 1 2 ∇xk log p∗ ( xk ) xk+1 = xk + νk+ 12 1 m νk+1 = νk+ 12 + 1 2 ∇xk+1 log p∗ ( xk+1 ) where 1m = ( 1 m1 , . . . , 1mn ) and denotes element wise multiplication . The step size , , and the mass , m , are hyperparameters that need to be tuned for each problem the method is applied to . We note that in the usual definition of HMC , a single scalar valued is used . Our use of a vector implies a different step size in each dimension which , with proper tuning , can improve performance by taking into account the different scales in each dimension . The use of does mean the procedure can no longer be interpreted as simulating Hamiltonian Dynamics , however , it can still be used as a valid HMC proposal ( Neal , 2011 ) . We do not consider the problem of choosing L in this work . 2.2 VARIATIONAL INFERENCE . VI approximates the target p ( x ) with a tractable distribution qφ ( x ) parameterized by φ . We choose φ as to minimise the Kullback-Leibler divergence with the target , DKL ( qφ ( x ) ||p ( x ) ) . As we only know p ( x ) up to a normalization constant , we can equivalently choose φ as to maximise the tractable Evidence Lower-Bound ( ELBO ) : ELBO = logZ −DKL ( qφ ( x ) ||p ( x ) ) = Eqφ ( x ) [ log p∗ ( x ) − log qφ ( x ) ] . 3 HYPERPARAMETER TUNING THROUGH THE EXPECTED LOG-TARGET . VI tunes the parameters of an approximate distribution to make it closer to the target . We would like to use this idea to tune the hyperparameters of HMC . We can run multiple parallel HMC chains and treat the final states as independent samples from an implicit variational distribution . If each chain starts at initial distribution q ( 0 ) ( x ) and then runs T accept/reject cycles , we can denote this implicit distribution as q ( T ) φ ( x ) , where φ now represents the HMC hyperparameters . Ideally , we would then choose φ as to maximise the ELBO : φ∗ = argmax φ E q ( T ) φ ( x ) [ log p∗ ( x ) − log q ( T ) φ ( x ) ] = argmax φ E q ( T ) φ ( x ) [ log p∗ ( x ) ] +H ( q ( T ) φ ( x ) ) . Whilst the first term in this expression can be estimated directly via Monte Carlo , the entropy term , H ( q ( T ) φ ( x ) ) , is intractable . To get around this , we should consider the purpose of the two terms during optimization . Maximizing the first term encourages q ( T ) φ ( x ) to produce samples that are in the high probability regions of the target , i.e. , ensuring that q ( T ) φ ( x ) is high where log p ∗ ( x ) is high . The entropy term acts as a regularizer preventing q ( T ) φ ( x ) from simply collapsing to a delta function at the mode of p ( x ) . The key observation of our method is that HMC already fulfills this regularization role because the implicit distribution it defines is not fully flexible . If q ( T ) φ ( x ) were to collapse to a delta function , this would require the hyperparameters to be such that the HMC scheme guides samples to the same point in space no matter their starting position , as sampled from q ( 0 ) ( x ) , which is unreasonable for practical problems . Therefore , we propose tuning φ simply by maximizing the expected log target density under the final state of the chain : φ∗ = argmax φ E q ( T ) φ ( x ) [ log p∗ ( x ) ] . ( 1 ) Although HMC does have a regularization effect , removing the entropy term does have some implications that we need to consider . Namely , if the initial distribution , q ( 0 ) ( x ) , is concentrated in a very high probability region of the target , p ( x ) , then optimizing objective ( 1 ) will not encourage HMC to explore the full target . Conversely , it would encourage the chains to remain in this region of high probability , close to their initial sampling point , which is undesirable behaviour . The key to avoiding this problem is to choose an initial distribution that has a sufficiently wide coverage of the target , we will discuss methods for doing this in Section 4 . We perform the optimization in ( 1 ) using stochastic gradient ascent with gradients being obtained using the standard reparameterization trick ( Rezende et al. , 2014 ; Kingma & Welling , 2014 ) . Samples from q ( T ) φ ( x ) can be obtained through a deterministic function of primitive random variables : x0 ∼ q ( 0 ) ( x ) for the initial distribution , γ0 : T−1 , γt ∼ N ( 0 , I ) for the momentum variables and a0 : T−1 , at ∼ U [ 0 , 1 ] for the accept/reject decisions . We denote this function as fφ ( x0 , γ0 : T−1 , a0 : T−1 ) . For minibatches of size N , the gradient is then estimated by 1 N N∑ n=1 ∇φlog p∗ ( fφ ( x ( n ) 0 , γ ( n ) 0 : T−1 , a ( n ) 0 : T−1 ) ) . ( 2 ) Each gradient value can be calculated using automatic differentiation tools . For ( 2 ) to be strictly unbiased , fφ must obey certain smoothness constraints to allow differentiation under the integral sign ( Border , 2016 ) . This does not hold in this case due to the accept/reject step . Other works have found success in making this approximation ( Levy et al. , 2018 ; Thin et al. , 2020 ) and as Thin et al . ( 2020 ) point out , this same issue is encountered when using the very popular ReLU activation function . Our empirical results confirm that ( 2 ) enables effective optimization of ( 1 ) . We also make another approximation when using automatic differentiation . For the leapfrog updates , we stop the backpropagation of the gradient through xk in∇xk log p∗ ( xk ) to prevent the calculation of secondorder gradients . We find this has little impact on the convergence of the algorithm and can lead to 5× speedups in execution time .

Summary: The paper proposes a method to optimize the parameters of the Hybrid Monte Carlo (HMC) algorithm (the step size and the diagonal of momentum's covariance matrix). In order to do that, the authors consider the distribution of samples q_T() obtained after T iterations of the algorithm (T accept/reject steps) starting from some distribution q_0(). Then, a reasonable objective for the optimization would be the KL-divergence between q_T() and the target density p(). However, the evaluation of the KL-divergence includes the entropy of q_T(), whose density is intractable due to numerous accept/reject steps. The proposed solution to this difficulty is to ignore the entropy term and maximize the log density of the target on samples from q_T(). To avoid the degenerate solution (due to ignorance of the entropy), the authors propose to choose q_0() carefully, e.g., to learn q_0() as a normalizing flow approximating the target p() via minimization of mass-covering alpha-divergence. The latter involves the usage of samples from the target distribution.