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 introduces a gradient-based approach for tuning the step-size and the diagonal mass matrix of HMC together with the parameters of an initial distribution for the Markov chain. They suggest different objectives amenable for SGD: maximize the expected target log-density of the final state of the chain, but also an objective to ensure a somewhat ‘wide’ initial distribution. The approach is illustrated on 2-d toy models, deep latent Gaussian models on (Fashion) MNIST and molecular configurations.