A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
Andrew Campbell * 1 Wenlong Chen * 2 Vincent Stimper * 3 4 Jose´ Miguel Hernandez-Lobato´ 3 Yichuan Zhang 5
Abstract behaviour and highly correlated samples. Furthermore, the step sizes may need to be tuned on a dimension by dimen- Hamiltonian Monte Carlo (HMC) is one of the sion basis depending on the scale and shape of the target most successful sampling methods in machine distribution (Neal, 2011). learning. However, its performance is signif- cantly affected by the choice of hyperparameter Traditionally, samples are produced by running a single values. Existing approaches for optimizing the HMC chain for a long time compensating for imperfect hy- HMC hyperparameters either optimize a proxy perparameter choices through ergodicity. However, it is be- for mixing speed or consider the HMC chain as comingly increasingly attractive to run many short MCMC an implicit variational distribution and optimize chains in parallel to make better use of parallel compute a tractable lower bound that can be very loose hardware (Hoffman & Ma, 2020). In this case, it is even in practice. Instead, we propose to optimize an more important to choose good hyperparameters that en- objective that quantifes directly the speed of con- courage fast mixing. This approach also offers the novel vergence to the target distribution. Our objective opportunity to choose different hyperparameter values for can be easily optimized using stochastic gradient each step in the chain, providing more tuning fexibility. descent. We evaluate our proposed method and Fully exploiting this opportunity in practice is a challenge. compare to baselines on a variety of problems MCMC hyperparameters are commonly tuned according including sampling from synthetic 2D distribu- to the average acceptance probability but it is unclear how tions, reconstructing sparse signals, learning deep this can scale up to tuning every hyperparameter individ- latent variable models and sampling molecular ually. Looking to backpropagation’s success in training confgurations from the Boltzmann distribution neural networks, a gradient based method would seem most of a 22 atom molecule. We fnd that our method appropriate. Unfortunately, there is no universal tractable is competitive with or improves upon alternative metric to quantify the performance of HMC that we can baselines in all these experiments. optimize for. We must therefore make a choice about which approximate metric is best suited for this application. 1. Introduction One approach is to use a proxy for the mixing speed of the chain. Levy et al. (2018) make use of a variation of the ex- Hamiltonian Monte Carlo (HMC) is a very popular Markov pected squared jumped distance (Pasarica & Gelman, 2010) Chain Monte Carlo (MCMC) method for generating approx- to encourage proposals to make large moves in space. Alter- imate samples from complex probability distributions. It natively, one can draw upon ideas from Variational Inference fnds wide use in machine learning (Neal, 2011) and across (VI) (Jordan et al., 1999), which matches an approximate the broader statistical community (Carpenter et al., 2017). distribution q to the target distribution p by maximizing Unfortunately, HMC’s performance depends heavily on the the Evidence Lower Bound or ELBO. This is equivalent to choice of hyperparameters such as the proposal step size. minimizing the KL-divergence between q and p. A step size that is too large can lead to unstable dynamics, while a step size that is too small may result in random walk Salimans et al. (2015); Wolf et al. (2016) use VI to obtain a training objective in which q is the joint distribution of all *Equal contribution 1Department of Statistics, University HMC samples along the chain. For tractability, they intro- of Oxford 2Baidu, Inc. 3Department of Engineering, Univer- duce an auxiliary inference distribution approximating the 4 sity of Cambridge Max Planck Institute for Intelligent Systems reverse dynamics of the chain. The looseness of their ELBO 5Boltzbit Ltd. Correspondence to: Andrew Campbell
−1.0 maximizing the expected log target density under the fnal state of the chain:
−1.5 ∗ ∗ φ = (T ) p (x) . argmax Eq (x) log (1) φ φ −2.0 HMC with ID1 pre-training Although HMC does have a regularization effect, remov- E[log p(x)] HMC with ID1 post-training ing the entropy term does have some implications that we −2.5 HMC with ID2 pre-training must consider. Namely, if the initial distribution, q(0)(x), is HMC with ID2 post-training True E[log p(x)] concentrated in a very high probability region of the target, −3.0 p(x), then optimizing the objective in (1) will not encourage 0 2 4 6 8 10 HMC to explore the full target. Conversely, it would encour- Step in HMC chain age 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 0.8 Initial distribution 1 Initial distribution 2 initial distribution that has a suffciently wide coverage of Target p(x) 0.6 the target. We discuss methods for doing this in Section 3.3.
3.1. Demonstration on a Toy Problem 0.4 p(x) We now demonstrate the ideas of the previous section on a 0.2 very simple toy problem. Here the target is a 1-dimensional normal distribution N (0, 1) which we attempt to sample 0.0 from using a 10 step HMC chain with each step consisting −4 −2 0 2 4 of 5 leapfrog updates. We initialize the chain either with x a narrow initial distribution N (0, 0.25) or a wide initial Figure 1. (Upper) Expected log target as a function of step t in distribution N (0, 4), as shown in the bottom plot in Figure 1. the HMC chain for initial distribution 1, N (0, 0.25), and initial We keep m constant for all steps but train one step size t for distribution 2, N (0, 4), before and after training. The ‘true’ value, each HMC step. The top of Figure 1 plots the progression of Ep log p(x) is also plotted. (Lower) The target pdf along with E (t) log p(x) during sampling for these two cases, before qφ the two initial distributions. and after hyperparameter training. Before training, when using the narrow initial distribution, E (t) log p(x) initially maximise the ELBO: qφ starts above the true value but converges from above as ∗ ∗ (T ) (t) φ = (T ) p (x) − q (x) the marginal HMC distribution, q , spreads to cover the argmax Eq (x) log log φ φ φ φ target. However, after training according to (1), all the step
∗ h(T ) i = (T ) p (x) + H q (x) . sizes have become very small causing the HMC chains to argmax Eq (x) log φ φ φ remain at their initial sampled positions which is obviously detrimental for convergence. To avoid this, we can use a Whilst the frst term in this expression can be estimated (T ) wide initial distribution. The top plot in Figure 1 shows that, directly via Monte Carlo, the entropy term, H qφ (x) , in this case, tuning hyperparameters according to (1) greatly is intractable. To get around this, we should consider the speeds up convergence. purpose of the two terms during optimization. Maximizing the frst term encourages q(T )(x) to produce samples that φ 3.2. Estimating the Gradients are in the high probability regions of the target, i.e., ensuring (T ) ∗ that qφ (x) is high where log p (x) is high. The entropy We solve (1) using stochastic gradient descent (SGD), with (T ) term acts as a regularizer preventing qφ (x) from simply gradients computed using a Monte Carlo approximation and collapsing to a point mass at the mode of p(x). The key the reparameterization trick (Rezende et al., 2014; Kingma observation of our method is that HMC already fulflls this & Welling, 2014). Here, we are forced to make another regularization role because the implicit distribution it defnes approximation since the discontinuous accept/reject step in- (T ) is not fully fexible. If qφ (x) were to collapse to a point troduces a non-differentiability in the Monte Carlo estimate mass, this would require the hyperparameters to be such of the objective. We can either ignore this and use biased that the HMC scheme always guides the chain state to the gradients or remove the accept/reject step altogether as other same point in space, irrespective of what initial position is works have suggested (Salimans et al., 2015; Caterini et al., sampled from q (0)(x). This is unreasonable for practical 2018). We opt for the former since we wish to work with problems. Therefore, we propose optimizing φ simply by HMC in its original form with the accept/reject step. A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
1 (T ) Our empirical results confrm that we are able to success- SKSD between the fnal state distribution qφ (x) and p(x) (T ) fully optimize the hyperparameters whilst using biased gra- The SKSD requires only samples from qφ and gradients of ∗ dients. The intuition behind this is that the bias originates the target, rxlog p (x). This objective encourages suitable from the effect of the hyperparameters on the acceptance values of s because if s is too small then the tuning of φ probability through the leapfrog discretization error. In prac- by solving (1) (which occurs jointly with the tuning of s) (T ) tice, this discretization error and its gradient with respect to will result in a degenerate qφ (x) far from the target. The the hyperparameters is small due to the leapfrog integrator SKSD measures this discrepancy and provides a learning having second order accuracy (Leimkuhler & Reich, 2004). signal for increasing s. Conversely, if s is too large then (T ) This results in a small overall bias. More concrete arguments qφ (x) will also be far from the target since the HMC chain justifying this can be found in the supplement. Finally, to im- will not be able to compensate for the poor initialization. prove computational effciency, we avoid the calculation of The SKSD will then favor to decrease s. Finally, as in second-order gradients by stopping backpropagation of the the leapfrog updates, we stop the gradient computations ∗ ∗ gradient through xk in rxk log p (xk) during the leapfrog through x in rx log p (x) when evaluating SKSD to avoid updates. We fnd this has little impact on convergence and the calculation of second-order gradients. can lead to 5× speedups in execution time. Given that the SKSD is a tractable objective that measures the discrepancy between q(T )(x) and p(x), it is theoretically 3.3. Tuning the Initial Distribution φ feasible to use the SKSD to optimize φ too. However, we For our method to be useful, the initial distribution q (0)(x) found that the SKSD does not scale well when optimizing must provide a suffcient coverage of p(x). However, a many parameters, which is why (1) is used instead. Note, q(0)(x) that is overly spread out is undesirable because the however, that the SKSD works very well in practice when burn-in time is then unnecessarily long. We describe below a we only tune the single scalar parameter s. (0) method for choosing q (x) that aims to achieve an optimal We have empirically evaluated the method described in this trade-off between these two requirements. section on a variety of applications and have found that it The main idea is to use a variational approximation to the gives consistently good results. target as the initial distribution, as done by Hoffman (2017). This should be a distribution that can be easily sampled from 3.4. Final optimization procedure and easily tuned to ft the target, e.g. a Gaussian or a normal- Algorithm 1 summarizes our optimization strategy, where izing fow (Tabak & Vanden-Eijnden, 2010; Rezende & Mo- Adam update(η, r L, i) returns the new value for η given hamed, 2015). Rather than use the standard ELBO for tun- η (0) by the i-th iteration of the Adam optimizer using gra- ing q (x), we use α-divergence minimization (Hernandez-´ (0) dient rηL, Dα(q (x) || p(x)) is an estimate of the Lobato et al., 2016). The α value dictates the mass covering ψi−1 α-divergence whose gradient is computed using doubly behaviour of the resulting approximation, with α = 0 corre- sponding to the standard mode seeking D �q(0)(x)||p(x) reparameterized gradient estimators (Tucker et al., 2019), KL (T ) minimization and α = 1 corresponding to the mass cover- SKSD(x1:N , score(x)) estimates the sliced kernelized Stein 0 � (0) (0) ing DKL p(x)||q (x) minimization. We compare both α discrepancy and HMCφi−1 (xn , score(x)) runs an HMC values in our experiments. The α-divergence is very useful (0)0 chain with initial state xn , target score function score(x) in this context as it can provide a mass covering approxima- and hyperparameters φi−1. Details on the computation tion without the use of samples from the target. However, if (0) (T ) of Dα(q (x) || p(x)), i) and SKSD(x , score(x)) are samples from p(x) are available, then q(0)(x) can alterna- ψi−1 1:N given in the Supplementary Material. tively be tuned on those samples via maximum likelihood (ML), which is also mass-covering. 4. Experiments The previous approaches will produce a q(0)(x) that fts the target, but do not guarantee that it will be broad enough. To Different experiments are performed to confrm that we can address this, we allow our method to automatically adjust discover good hyperparameter settings by using the previ- the width of the initial distribution as necessary to keep ously described method. We provide code for reproducing (T ) 2 qφ (x) as closely matched to p(x) as possible. This is all our experiments on github . achieved by applying a scalar scale factor s centered around 1 the mean µ of q(0)(x) to each sample x from this distribu- We use the metric called maxSKSD in Gong et al. (2021) i 2https://github.com/VincentStimper/ tion, i.e. xˆi = s(xi − µ) + µ. We tune s by minimizing the hmc-hyperparameter-tuning Sliced Kernelized Stein Discrepancy (Gong et al., 2021) or A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
Gaussian Laplace Dual Moon Mixture Wave 1 Wave 2
maxELT α = 0 0.0677 0.0005 0.2370 0.0004 0.0525 0.0462 maxELT α = 1 0.0009 0.0004 0.8637 0.0010 0.0158 0.0801 maxELT α = 0 SKSD 0.0008 0.0016 0.1684 0.0004 0.0020 0.0217 maxELT α = 1 SKSD 0.0009 0.0014 0.2528 0.0004 0.0019 0.0317 Hoffman (2017) 0.0364 0.0005 1.1553 0.0846 0.9447 0.0465 Ruiz & Titsias (2019) 0.0003 0.0003 1.6290 0.0003 0.0024 0.2375 NUTS 0.0044 0.0016 0.2326 0.0023 0.0260 0.0965
Figure 2. Histograms of 2D targets gener- Table 1. KSD between the HMC samples and the target distribution for the baselines and the ated by rejection sampling. 4 variations of our method on each of the synthetic target distributions (equations are given in the Supplementary Material).
Algorithm 1 Optimization Procedure 100 Input: Initial ψ0, φ0 and s0, number of iterations I1 and I2 and number of samples N ∗ HMC before training (x) = (r p (x)) −1 Defne score stop gradient xlog 10 HMC after training # Train initial distribution KSD Initial Distribution for i = 1 to I1 do (0) 10−2 ψi ← Adam update(ψi−1, rψ Dα(q (x) || p(x)), i) i−1 ψ i−1
end for −4 −3 −2 −1 0 1 µ ← mean of q(0) (x) 10 10 10 10 10 10 ψi −1 Initialization stepsize # Train HMC hyperparameters for i = 1, to I2 do Figure 3. KSD versus step size used as initialization point for opti- for n = 1, to N do # This loop is vectorized in practice mization. Wave 1 is the target distribution. (0) (0) Draw xn ∼ qψ (x) (0)0 (0) xn ← si−1(xn − µ) + µ # Rescale samples tic. The second baseline is the method from Ruiz & Titsias (T ) (0)0 (2019) which initializes chains with a variational distribu- xn ← HMCφi−1 (xn , score(x)) end for tion, trained using a novel divergence metric. The HMC 1 PN ∗ (T ) φi ← Adam update(φi−1, rφi−1 N n=1 log p (xn ), i) parameters are not tuned however. The fnal baseline is the (T ) si ← Adam update(si−1, rsi−1 SKSD(x1:N , score(x)), i) popular No-U-Turn Sampler (Hoffman & Gelman, 2014). end for Full details regarding the baselines are given in the Supple- Return ψI1 , φI2 , sI2 mentary Material. Results are shown in Table 1. Our method with automatic s 4.1. 2D distributions scaling using the SKSD performs consistently well across We frst focus on drawing approximate samples from a range distributions. For some simple targets, e.g. Gaussian and of synthetic 2D target densities (Figure 2). We use 30- Laplace, it performs slightly worse than Ruiz & Titsias step HMC chains and a factorized Gaussian q(0)(x) that (2019), however, this method breaks down on the more com- is trained by minimizing the α-divergence. We optimize plex targets, Dual Moon and Wave 2 whereas ours remains step sizes and masses using (1), a procedure we refer to as consistent. Furthermore, we ft the targets better when we ‘maxELT’ for maximizing the expected log target. Addi- tune s by minimizing the SKSD than when we do not. For tionally, as an ablation study, we consider different initial some distributions, tuning s effectively helps prevent mode distribution training strategies: α = 0 or 1 and whether or seeking behaviour. We confrm this quantitatively in the not to tune the scaling s by minimizing the SKSD (s = 1 Supplementary Material by comparing expected log target when not tuned). To quantify convergence to the target, we values. We fnd that narrow initial distributions (α = 0) of- used the Kernelized Stein Discrepancy (KSD) (Liu et al., ten lead to excessively high log target values, but the tuning 2016; Chwialkowski et al., 2016) between the generated of s can prevent this pathology. samples and the targets. We also investigate the method’s robustness to the initial- We include three baselines for reference. The frst one is ization point. Figure 3 plots the KSD between the HMC taken from Hoffman (2017) and initializes the HMC chains samples and the target distribution (Wave 1) before and after with an α = 0 trained Gaussian distribution and tunes step training the hyperparameters for a large range of initializa- sizes according to a minimum acceptance probability heuris- tion step sizes. We fnd our method is largely invariant to A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization the initialization point provided it is not excessively large, same initial distribution. The frst two are parallel samplers in which case every step in the chain is rejected and there is where for each sample we frst sample the initial distribution no gradient signal for learning. However, we do no worse and then run 20 HMC accept/reject steps. The step sizes than the initial distribution as each sample remains at its and masses are tuned either by a grid search to maximize initial sampling point. the log marginal likelihood on a validation set or the step sizes are tuned such that the average acceptance probability 4.2. Sparse Signal Recovery is 0.65. The fnal baseline is sequential—the No-U-Turn sampler. To generate samples, we frst sample the initial dis- We now consider higher dimensional target distributions. tribution once then run 15000 accept/reject steps, discarding Specifcally, posterior distributions in Bayesian compressed the frst 5000 samples. The hyperparameters are tuned using sensing (Donoho, 2006; Ji et al., 2008). The aim is to dual averaging during the burn-in period. Full experimental d recover the sparse signal w ∈ R from measurements y ∈ details are in the Supplementary Material. Rn where n < d. The observation model is y = Xw + n 2 The results are shown in Figure 4 and we see that our method e where e ∈ R , e ∼ N (0, σ0I), is additive Gaussian noise and X ∈ Rn×d is a measurement matrix obtained by generally outperforms the baselines. It is also noticeable sampling its entries from a standard Gaussian distribution that the sequential sampler (NUTS) has a very high variance and then normalizing the rows to have unit Euclidean norm. between repeats. This is because it is very sensitive to the We place a sparsity enforcing horseshoe prior on w and initialization point. NUTS performs well in the very few 2 cases where the initialization is close to a posterior mode consider the posterior for w given y, X and σ0. When n << d this posterior is multi-modal representing multiple with good properties. However, if it is initialized close a potential explanations for the available observations. In sub-optimal mode, it will likely stay there for the whole run these experiments we fx n = 6 and d = 64. and fnd hyperparameters only suited to this local region giving a very poor log likelihood. The parallel samplers do We draw posterior samples using a 20 step HMC chain with not have this issue as they have a new initialization point for 5 leapfrog iterations per step, initialized with a factorized each sample. We also note that, in this case, using SKSD to Gaussian distribution that is trained with α = 1. We do not tune s results in a slight degradation in performance. On this use α = 0 here since this method performs signifcantly problem, α = 1 is already well suited to fnd a good initial worse in this problem due to its mode seeking tendency distribution and, indeed, the value for s tuned by SKSD was (Hernandez-Lobato´ et al., 2015). We evaluate the quality very close to 1. of our samples by calculating their log marginal likelihood on test data generated using the same sparse signal as in 4.3. Deep Latent Gaussian Models the training set. The highest test log marginal likelihood is obtained when the generated samples come from the true We now focus on training Variational Autoencoders (VAEs) posterior distribution. (Kingma & Welling, 2014; Rezende et al., 2014). These are generative models of the form p(x, z) = pθ(x|z)N (z; 0, I),
−4000 explaining observed data x with a latent variable z and a like- lihood pθ(x|z) parameterized by a neural network NNθ(z) −6000 with parameters θ and input z. The standard training in −8000 VAEs is to use a factorized Gaussian distribution qψ(z|x)
−10000 parameterized by NNψ(x) to approximate pθ(z|x) and train (θ, ψ) jointly by maximizing the ELBO. −12000 Our approximation to pθ(z|x) is the HMC chain output −14000 (T )
Test log likelihood distribution q φ (z|x). The initial state of the HMC chain (0) −16000 is sampled from qψ (z|x), which is given by a factorized −18000 Gaussian distribution parameterized by NNψ(x) and appro- priately scaled by s as before. The parameters ψ are trained −20000 Initial maxELT maxELT Grid pā = 0.65 NUTS by minimizing the α-divergence and s is trained using the distribution (ours) & SKSD search (ours) SKSD. We train φ and θ by jointly maximizing the objective in (1) with respect to these two parameters. We consider Figure 4. Violin-Box plot of test log marginal likelihoods on the HMC chains of length 30 with 5 leapfrog iterations per step sparse signal recovery problems. We draw 10,000 posterior sam- and tune a different step size parameter per dimension and ples to estimate the log marginal likelihood of 1000 test observa- per step in the HMC chain while the mass parameters are all tions and repeat 20 times to obtain confdence bands. kept constant and equal to 1 all throughout the HMC chain.
We compare our method with three baselines that use the Note that we use only one set of HMC hyperparameters for A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
the Supplementary Material. Table 2. Average test marginal log-likelihood and its standard error (SE) estimated using HAIS for different methods for MNIST and On both datasets, the HMC based methods achieve better Fashion MNIST. For methods with a scale factor s, we report the performance than VAE or IWAE, showing that reducing the fnal scale after training. We also report the best log-likelihood approximation bias of the variational distribution with HMC values from previous works. greatly helps. Furthermore, we see that adding the scale MNIST Fashion MNIST factor to our method signifcantly improves performance as Model Scale Mean SE Scale Mean SE this avoids degenerate behaviour when training the HMC hyperparameters. We note that, without any scaling s, α = 1 VAE - -85.08 0.22 - -108.54 0.60 outperforms α = 0 due to α = 0 resulting in a too narrow DReG-IWAE - -83.73 0.21 - -104.48 0.58 maxELT α = 0 1.0 -83.48 0.21 1.0 -104.08 0.58 initial distribution. With scaling, the SKSD automatically maxELT α = 1 1.0 -82.46 0.21 1.0 -103.57 0.58 widens the initial distribution making the performance of maxELT α = 0 SKSD 6.79 -81.91 0.20 5.58 -103.18 0.58 both methods similar. Finally, we observe that HMC based maxELT α = 1 SKSD 3.90 -81.94 0.20 3.59 -102.29 0.57 methods top out at similar log-likelihood values (within Hoffman - -81.74 0.20 - -103.04 0.58 around one standard error). We believe this is due to the Ruiz & Titsias - -82.45 0.21 - -105.13 0.59 methods reaching the limits of the chosen neural network Salimans et al. - -81.94 - - -104.44 0.59 architecture on these datasets, with no more gains to be Caterini et al. - -82.62 - - -104.26 0.58 made from more accurate posterior approximations. all pθ(z|x) targets, independently of x. One could make 4.4. Molecular Confgurations the hyperparameters depend on x through an amortization Finally, we evaluate our method on the complex real-world network, but we found that this did not improve performance. problem of sampling equilibrium molecular confgurations We consider two benchmark datasets: MNIST and Fashion from the Boltzmann distribution of the molecule Alanine MNIST. As is common, we use binarized images and a Dipeptide. The unnormalized target distribution for the Bernoulli likelihood p (x|z) parameterized with the same θ atom coordinates x is e−u(x), where u denotes the potential convolutional architecture as Salimans et al. (2015). Full energy of the system, which can be obtained using the laws experimental details are in the Supplementary Material. of physics. This problem is usually tackled via Molecular We evaluate the quality of our trained models using the Dynamics (MD) simulations. Here, we aim to produce −u(x) marginal log-likelihood log pθ(x) on the test set estimated independent samples from e using our trained short using Hamiltonian Annealed Importance Sampling or HAIS HMC chains. We do not operate directly on the Cartesian (Sohl-Dickstein & Culpepper, 2012). Results are shown in coordinates but apply the coordinate transform presented by Table 2. We also report log-likelihood values for multiple Noe´ et al. (2019), to map some of the Cartesian coordinates baselines. Using the same neural architecture, we imple- to bond lengths, bond angles, and dihedral angles. The fnal mented the standard VAE and IWAE3 models. We also dimensionality for x is 60. For more details, we refer to the implemented another method for tuning φ (Hoffman, 2017) Supplementary Material. where the step sizes are adjusted to make the minimum For the initial distribution, we use a normalizing fow based average acceptance probability in each minibatch equal to on real-valued non-volume preserving (RNVP) transforma- 0.25. Furthermore, we implemented the method for tuning tions (Dinh et al., 2017), followed by 50 HMC steps with 10 ψ from Ruiz & Titsias (2019) while the HMC step sizes leapfrog iterations per step. We used different methods to were tuned to ensure the mean average acceptance proba- train the fow and the HMC hyperparameters. The fow was bility in each minibatch is equal to 0.65. We update θ as trained with α = {0, 1}-divergence and by ML. The latter in our method and use the same number of leapfrog steps was done using 105 training data samples obtained via a for a fair comparison. Finally, we report results for MNIST MD simulation5. As in the case with 2D densities and with from Salimans et al. (2015) and Caterini et al. (2018) which sparse signals but unlike in the VAE case, we frst trained include HMC hyperparameter tuning during training and the fows and then kept them fxed when tuning the HMC use the same network architecture as us. These authors only hyperparameters and the scale factor. evaluated on MNIST, so we reimplemented their methods for Fashion MNIST4. We confrm signifcant differences For each initial distribution type, the HMC hyperparameters between the models using paired t-tests. Full results are in were tuned according to maxELT alone or by maxELT with SKSD scale training. As a baseline, we optimized the HMC 3 We used the DReG estimator from (Tucker et al., 2019) for parameters via grid search, keeping step sizes and masses the IWAE. 4 We reimplemented Salimans et al. (2015) ourselves 5The same dataset was used to obtain the mean and variances and used the implementation from https://github.com/ required in the normalization step of the coordinate transform. anthonycaterini/hvae-nips for Caterini et al. (2018) A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
α = 0 0.08 α = 1 ML
0.06
0.04 KL divergence
0.02 (a) Ground truth (MD) (b) maxELT & SKSD
0.00 Figure 5. Ramachandran plot of (a) the ground truth determined Initial maxELT maxELT Grid search p̄a = 0.65 via a MD simulation and (b) samples from the model with the distribution (ours) & SKSD (ours) proposal trained with ML and HMC hyperparameters being tuned by maxELT & SKSD. Figure 6. Visualization of the KL-divergences of the proposal and the models where different HMC hyperparameter tuning schemes constant across dimensions and HMC step, varying these were used. two parameters in a grid and picking the combination that gave the lowest median marginal KL-divergence to the MD Table 3. P-values of the Wilcoxon test with the alternative hypoth- training data. We also considered another baseline by ad- esis that the model with HMC parameters tuned by maxELT & justing the step size value such that the average acceptance SKSD leads to lower KL-divergences of the marginals with respect probability was 0.65 (referred to as pa = 0.65). Further to the ground truth than the respective baseline. implementation details are given in the Supplementary Ma- terial. Grid search pa = 0.65 A new MD simulation was run to obtain 106 ground truth α = 0 0.0030 0.39 samples for evaluating the performance of the different α = 1 0.0030 0.0040 methods. For tractability, this is done by comparing the ML 0.010 5.1e-5 marginal distributions of the generated and ground truth samples. First, we use Ramachandran plots (Ramachandran one case maxELT & SKSD leads a signifcant improvement et al., 1963), which are 2D histograms for the two dihedral over the baselines and in the remaining case the two methods angles in the bonds connecting an amino acid to the pro- are on par with each other. All the other remaining p-values tein backbone. These plots are frequently used to analyse can be found in the Supplementary Material. how proteins fold locally. In the Alanine Dipeptide case we can obtain a Ramachandran plot for the bond connecting the two amino acids forming this molecule. Two sample 5. Discussion and Related Work plots are shown in Figure 5. We compute KL divergences Our experiments show a general trend that, when we solely between the 2D histograms (Ramachandran plots) for the optimize the objective in (1), the value of α used signif- ground truth samples and for the samples generated by the cantly affects performance. However, when applying the different methods. The results are given in Figure 6. All SKSD scaling, this difference becomes smaller, showing the corresponding Ramachandran plots are shown in the that this technique is useful for automatically fnding a suit- Supplementary Material. Our method outperforms the base- able initial distribution. lines for all proposals except for the one trained with the α = 0-divergence, where it improves upon grid search but A limitation of the method is that the hyperparameter op- timization procedure must be run before samples can be performs worse than the pa = 0.65 baseline. produced. Therefore, our method is most applicable in Finally, for each dimension of x and using kernel density cases where the optimization time is outweighed by the time estimation, we compute the KL-divergences between the savings gained through better hyperparameter choices and 1D densities produced by the ground truth samples and a faster mixing speed. This is the case for diffcult prob- by the samples generated by the different methods. We lems where choosing reasonable hyperparameters is highly perform a Wilcoxon test to check whether the resulting non-trivial and also for problems where a large number of 60 divergence values (one per dimension) are consistently samples need to be produced for a downstream task. lower for one method or another. Table 3 shows p-values for the case in which maxELT & SKSD is being tested for Previous works have also used VI to obtain an objective for having lower KL-divergences than the baselines. In all but the gradient based tuning of HMC hyperparameters. As dis- cussed previously, Salimans et al. (2015); Wolf et al. (2016) A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization consider the entire joint distribution of the HMC chain as HMC. Though, they must still use gradient-free optimiza- the variational distribution. The higher dimensionality in- tion on a heuristic to tune the HMC parameters, one could creases the looseness of the variational bound, making the investigate performance if hyperparameters were instead method highly sensitive to the length of the chain and the tuned using our differentiable objective. accuracy of a reverse dynamics approximation. Our exper- Our work also builds upon methods from statistical mechan- iments in section 4.3 show we can consider much longer ics. The Boltzmann Generator (Noe´ et al., 2019) opened up chains and perform just as well without needing a reverse this line of research by using a normalizing fow to sample approximation. Caterini et al. (2018) also construct an al- molecular confgurations. We found we can improve upon ternative ELBO for HMC. However, they only sample the this by using the fow as the initial distribution for HMC, auxiliary momentum variables once at the start of the chain fne tuning the fow samples with our short chains. Our which reduces the empirical performance of HMC. Alanine Dipeptide experiment comes from the recent work In contrast to these methods, some gradient based tuning of Wu et al. (2020) on stochastic normalizing fows, which techniques do not use ideas from VI but instead optimize a consist of stochastic steps interspersed between determin- proxy for mixing speed. Levy et al. (2018) generalize the istic steps in a normalizing fow. It is possible to combine standard leapfrog integrator used in HMC with multi-layer such models with our approach, using our objective to tune perceptrons which are then trained by maximizing a modi- hyperparameters within the stochastic layers. We leave this fed version of the expected squared jumped distance. We extension to future work. improve upon this objective by directly optimizing conver- gence speed by using gradient information from the target 6. Conclusion distribution itself. It would be an interesting direction to use our objective to train this generalised leapfrog oper- In this work, we presented a new objective motivated by ator. Finally, Titsias & Dellaportas (2019) consider the VI that can be easily used for the gradient-based optimiza- gradient based tuning of the Markov transition operator, tion of HMC hyperparameters. We provided a fully auto- pφ(xt|xt−1), in the case of a single long MCMC chain. matic method for choosing an initial distribution for the They optimize with respect to the expected acceptance prob- HMC chain that reduces burn-in time and aids optimization. ability for the next step in the chain, regularized with the Evaluating on multiple real-world problems, we found our entropy of pφ(xt|xt−1). Unfortunately, as this entropy is in- method is competitive with or improves upon existing meth- tractable when using the leapfrog algorithm as the transition ods for tuning hyperparameters. We hope this encourages operator, it cannot be applied to HMC in its current form. further work applying this idea to other methods that use HMC and that would beneft from increased convergence There are also many non-gradient based heuristics for tuning speed. HMC hyperparameters. The popular No-U-Turn Sampler (Hoffman & Gelman, 2014) can adaptively set the number of leapfrog steps L to avoid U-turns and fnd a global constant Acknowledgements for the step sizes by adjusting the average acceptance rate. In Andrew Campbell acknowledges support from the EPSRC section 4.1, we found we can outperform NUTS even though CDT in Modern Statistics and Statistical Machine Learning we do not adaptively set L. With our objective, we can tune (EP/S023151/1). Jose´ Miguel Hernandez´ Lobato acknowl- individual step sizes (and masses) for each dimension and edges support from a Turing AI Fellowship under grant step in the chain, allowing for a much higher degree of EP/V023756/1. Yichuan Zhang acknowledges support from granular control over the algorithm. Furthermore, we do not Samsung Research, Samsung Electronics Co., Seoul, Re- need to rely on ‘rules of thumb’ such as standard acceptance public of Korea and from Boltzbit Ltd. rate targets used in many algorithms (Hoffman & Gelman, 2014; Hoffman, 2017) but we can automatically tune all continuous hyperparameters using information from the References target distribution directly. Carpenter, B., Gelman, A., Hoffman, M. D., Lee, D., Other works use MCMC as part of a hybrid inference Goodrich, B., Betancourt, M., Brubaker, M., Guo, J., scheme. Ruiz & Titsias (2019) propose a novel objective to Li, P., and Riddell, A. Stan: A probabilistic programming improve the training of the initial distribution using MCMC language. Journal of statistical software, 76(1), 2017. samples. However, their objective cannot be used for tuning MCMC as it encourages fnal samples to be close to the Caterini, A. L., Doucet, A., and Sejdinovic, D. Hamiltonian initial distribution. Hoffman et al. (2019) use normalizing variational auto-encoder. Advances in Neural Information fows to warp the target distribution such that it is close to Processing Systems (NeurIPS), 2018. an isotropic Gaussian and thus easy to sample from using Chwialkowski, K., Strathmann, H., and Gretton, A. A kernel A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
test of goodness of ft. Proceedings of the International Leimkuhler, B. and Reich, S. Simulating Hamiltonian dy- Conference on Machine Learning (ICML), 2016. namics. Number 14. Cambridge university press, 2004.
Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density esti- Levy, D., Hoffman, M. D., and Sohl-Dickstein, J. General- mation using Real NVP. International Conference on izing Hamiltonian Monte Carlo with neural networks. In Learning Representations (ICLR), 2017. International Conference on Learning Representations (ICLR), 2018. Donoho, D. L. Compressed sensing. IEEE Transactions on Information Theory, 52(4):1289–1306, 2006. doi: Liu, Q., Lee, J., and Jordan, M. A kernelized stein dis- 10.1109/TIT.2006.871582. crepancy for goodness-of-ft tests. Proceedings of the International Conference on Machine Learning (ICML), Gong, W., Li, Y., and Hernandez-Lobato,´ J. M. Sliced 2016. kernelized stein discrepancy. International Conference on Learning Representations (ICLR), 2021. Neal, R. M. Probabilistic Inference using Markov Chain Monte Carlo Methods. Technical report, University of Hairer, E., Lubich, C., and Wanner, G. Geometric numerical Toronto, Department of Computer Science, 1993. integration illustrated by the Stormer–Verlet¨ method. Acta Numerica, 12:399–450, 2003. Neal, R. M. MCMC using Hamiltonian dynamics. Hand- book of markov chain monte carlo, 2(11):2, 2011. Hernandez-Lobato,´ J., Hernandez-Lobato,´ D., and Suarez,´ A. Expectation propagation in linear regression models Noe,´ F., Olsson, S., Kohler,¨ J., and Wu, H. Boltz- with spike-and-slab priors. Machine Learning, 99:437– mann generators: Sampling equilibrium states of many- 487, 2015. body systems with deep learning. Science, 365(6457), Hernandez-Lobato,´ J. M., Li, Y., Rowland, M., Bui, T., September 2019. ISSN 0036-8075, 1095-9203. doi: Hernandez-Lobato, D., and Turner, R. Black-box alpha 10.1126/science.aaw1147. divergence minimization. Proceedings of the Interna- Pasarica, C. and Gelman, A. Adaptively scaling the tional Conference on Machine Learning (ICML), 2016. Metropolis algorithm using expected squared jumped Hoffman, M. and Gelman, A. The no-u-turn sampler: Adap- distance. Statistica Sinica, 20(1):343–364, 2010. tively setting path lengths in Hamiltonian Monte Carlo. Ramachandran, G. N., Ramakrishnan, C., and Sasisekharan, Journal of Machine Learning Research (JMLR), 2014. V. Stereochemistry of polypeptide chain confgurations. Hoffman, M. and Ma, Y. Black-box variational inference Journal of Molecular Biology, 7(1):95–99, July 1963. as a parametric approximation to Langevin dynamics. ISSN 0022-2836. doi: 10.1016/S0022-2836(63)80023-6. Proceedings of the International Conference on Machine Rezende, D. and Mohamed, S. Variational Inference with Learning (ICML), 2020. Normalizing Flows. Proceedings of the International Hoffman, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, Conference on Machine Learning (ICML), 2015. D., and Vasudevan, S. Neutra-lizing bad geometry in Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic hamiltonian monte carlo using neural transport. arXiv Backpropagation and Approximate Inference in Deep preprint arXiv:1903.03704, 2019. Generative Models. Proceedings of the International Hoffman, M. D. Learning deep latent Gaussian models Conference on Machine Learning (ICML), 2014. with Markov Chain Monte Carlo. Proceedings of the Ruiz, F. and Titsias, M. A contrastive divergence for com- International Conference on Machine Learning (ICML), bining variational inference and MCMC. Proceedings 2017. of the International Conference on Machine Learning Ji, S., Xue, Y., and Carin, L. Bayesian compressive sens- (ICML), 2019. ing. IEEE Transactions on signal processing, 56(6):2346– 2356, 2008. Salimans, T., Kingma, D., and Welling, M. Markov Chain Monte Carlo and variational inference: Bridging the gap. Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., and Saul, Proceedings of the International Conference on Machine L. K. An introduction to variational methods for graphical Learning (ICML), 2015. models. Machine learning, 37(2):183–233, 1999. Sohl-Dickstein, J. and Culpepper, B. J. Hamiltonian an- Kingma, D. P. and Welling, M. Auto-encoding Variational nealed importance sampling for partition function estima- Bayes. In International Conference on Learning Repre- tion. Technical report, Redwood Center for Theoretical sentations (ICLR), 2014. Neuroscience University of California, Berkeley, 2012. A Gradient Based Strategy for Hamiltonian Monte Carlo Hyperparameter Optimization
Tabak, E. G. and Vanden-Eijnden, E. Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1):217–233, 2010. ISSN 15396746, 19450796. doi: 10.4310/CMS.2010.v8.n1. a11. Titsias, M. and Dellaportas, P. Gradient-based adaptive Markov Chain Monte Carlo. Advances in Neural Infor- mation Processing Systems (NeurIPS), 2019. Tucker, G., Lawson, D., Gu, S., and Maddison, C. Dou- bly reparameterized gradient estimators for Monte Carlo objectives. International Conference on Learning Repre- sentations (ICLR), 2019. Wolf, C., Karl, M., and van der Smagt, P. Variational in- ference with Hamiltonian Monte Carlo. arXiv preprint arXiv:1609.08203, 2016. Wu, H., Kohler,¨ J., and Noe,´ F. Stochastic Normalizing Flows. Advances in Neural Information Processing Sys- tems (NeurIPS), 2020.