Gradient-based tuning of Hamiltonian Monte Carlo hyperparameters
Authors: Andrew Campbell
Wenlong Chen
Vincent Stimper
José Miguel Hernández-Lobato
Yichuan Zhang
Venue:  (2021)
Keywords: Hamiltonian Monte Carlo
HMC
MCMC
Variational Inference

Abstract: Hamiltonian Monte Carlo (HMC) is one of the most successful sampling methods in machine learning. However, its performance is significantly affected by the choice of hyperparameter values, which require careful tuning. Existing approaches for automating this task either optimise a proxy for mixing speed or consider the HMC chain as an implicit variational distribution and optimize a tractable lower bound that is too loose to be useful in practice. Instead, we propose to optimize an objective that quantifies directly the speed of convergence to the target distribution. Our objective can be easily optimized using stochastic gradient descent.  We evaluate our proposed method and compare to baselines on a variety of problems including synthetic 2D distributions, the posteriors of variational autoencoders and the Boltzmann distribution for molecular configurations of a 22 atom molecule. We find our method is competitive with or improves upon alternative baselines on all problems we consider.

Under review as a conference paper at ICLR 2021
GRADIENT-BASED TUNING OF HAMILTONIAN MONTE CARLO HYPERPARAMETERS
Anonymous authors Paper under double-blind review
ABSTRACT
Hamiltonian Monte Carlo (HMC) is one of the most successful sampling meth- ods in machine learning. However, its performance is significantly affected by the choice of hyperparameter values, which require careful tuning. Existing ap- proaches for automating this task either optimize proxy for mixing speed or consider the HMC chain tractable lower bound that is too loose to be useful in practice, Instead, we pro- pose to optimize an objective that quantifies directly the speed of convergence the target distribution. Our objective can be easily optimized using stochastic gra- dient descent. We evaluate variety of problems including synthetic 2D distributions, the posteriors of varia- tional autoencoders and the Boltzmann distribution for molecular configurations of a 22 atom molecule. We find our method is competitive with or improves upon alternative baselines on all problems we consider.
implicit variational distribution and optimize
proposed method and compare to baselines a
INTRODUCTION
Hamiltonian Monte Carlo (HMC) is popular sampling based method for performing accurate infer- ence on complex distributions that we may only know up to 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. Tuming 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 proxy for the mixing speed of the chain, i.e. the speed at which the Markov chain marginal distribution approaches the target. For example, (Leyy et al.) (2018) use 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 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 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 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 KI.-divergence between the auxiliary inference distri- bution and the true reverse dynamics. As the chain length increases. the dimensionality of these distributions increases, tending to increase the looseness the bound. This causes issues during op- timization 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-

---

Under review as a conference paper at ICLR 2021
formance. When applied to 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 fully flexible variational distribution from collapsing to a point mass maximizing the log target density. Since HMC, by construction, cannot collapse to such point mass, we argue that the entropy term be dropped provided the initial distribution of the chain has enough coverage of the target. We evaluate our proposed method on variety of tasks. We first consider 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 bench- mark: sampling molecular configurations from the Boltzmann distribution of 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 Markov Chain Monte Carlo method Neal 1993) which aims draw samples from the n-dimensional target distribution p(æ) = zp" p where 2 is the (usually unknown) nor- malization constant. It introduces auxiliary variable E IR", referred
the momentum,
which is distributed according to [v; O, diag(m)). with the resulting method sampling on the extended space
(x. v). HMC progresses by first sampling an initial state from some ini- tial distribution and then iteratively proposing new states and accepting/rejecting them accord- ing to an acceptance probability. To propose a new state, first, a new value for the momentum drawn from N (1; 0, diag(m)). then, we simulate Hamiltonian Dynamics with Hamiltonian, (2, = loggp" (a:) + _ - diag(m)
arriving at new state (x', This new state is ac- cepted with probability min [1, cxp( exp(-H (a,,, v/) + (z,,]) otherwise we reject the proposed state and remain at the starting state. The Hamiltonian Dynamics are simulated using a numerical inte- grator, the leapfrog integrator (Hairer et al. 2003) being a popular choice. To propose the new state, T. leapfrog updates taken, each update consisting of the following equations:
U = Vk +
④ Va,logy logp" (Tx)
1
Fk+1 Izk + k
m
Uk+1 = / + +
2 c V Vanzs lagpplo""""xx:
and denotes element wise multiplication. The step size, e, and the mass, m, 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 c is used. Our use of a vector e implies 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 e does mean the procedure can longer be interpreted simulating Hamiltonian Dynamics, however, it can still be used as valid HMC proposal (Neal 2011 We do not consider the problem of choosing L in this work.
where m m. my.
2.2 VARIATIONAL INFERENCE
VI approximates the target p(x) with a tractable distribution (a(x) parameterized by d. We choose to minimise the Kullback-Leible divergence with the target, Dx (1s (x:) (=))((x))) As we only know p(z) up to normalization constant, we equivalently choose @ to maximise the tractable Evidence Lower-Bound (ELBO):
ELBO log Z DD (=s())lp(()) Bgg(2) [log (=x) log 9, (x)]
3 HYPERPARAMETER TUNING THROUGH THE EXPECTED LOG-TA
VI tunes the parameters of an approximate distribution to make it closer to the target. We would like use this idea tune the hyperparameters of HMC. We multiple parallel HMC chains and
2

---

Under review as a conference paper at ICLR 2021
treat the final states as independent samples from an implicit variational distribution. If each chain starts at initial distribution of (xr) and then runs T accept/reject cycles, we can denote this implicit (T) (z)) where & now represents the HMC hyperparameters. Ideally, we would then distribution as Lo
choose ① as to maximise the BLBO:
argmax E (w) llog (z) log 4 (x)] = argmax E. (w) [logp" [x)] + H1 ( (o)) 9
Whilst the first term in this expression can be estimated directly via Monte Carlo, the entropy term, H (gs (000((())))))))() (I)), is intractable. To get around this, we should consider the purpose of the two terms during optimization. Maximizing the first term encourages 4 (2)
(T)) (x) to produce samples that are
in the high probability regions of the target, i.e., ensuring that Is. (x) is high where log p" px() is high. The entropy term acts as a regularizer preventing 9 (z) from simply collapsing to 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 1s 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 matter their starting position, sampled from (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:
(f) (x) were
argmax E (T)) (c) [log (lgp"(zz). (x)]
(1)
Although HMC does have a regularization effect, removing the entropy term does have some im- plications that we need to consider. Namely, if the initial distribution, (z), is concentrated in very high probability region of the target, p(æ). then optimizing objective 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
We perform the optimization in I1 using stochastic gradient ascent with gradients being ob- tained using the standard reparameterization trick (Rezende al. 2014; Kingma & Welling 2014). Samples from 9 (') (x) can be obtained through a deterministic function of primitive ran- dom variables: Zo (xx) for the initial distribution, 70 7s ~ NO I) for the momen- tum variables and 00:T-1; 1, at (0, 1] for the accept/reject decisions. We denote this function as fz 70:T-1. -1). For minibatches of size N, the gradient is then estimated by
(n) Σ Vologr" (Jo (z6") - -
(())
Each gradient value can be calculated using automatic differentiation tools. For (2) to be strictly unbiased, fa 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.) (020) and as Thin 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 (T. We also make another approximation when using automatic differentiation. For the leapfrog updates, we stop the backpropagation of the gradient through in Vz, log p" (ax) to prevent the calculation of second- order gradients. We find this has little impact on the convergence of the algorithm and can lead to 5x speedups in execution time.
3.1 DEMONSTRATION ON TOY PROBLEM
We now demonstrate the ideas of the previous section on a very simple toy problem. Here the target is a I-dimensional normal distribution AO 1) which we attempt sample from using a 10 step HMC chain with each step consisting of 5 Ieapfrog updates. We initialize the chain either with a narrow initial distribution (0, 0.25) or a wide initial distribution A(0 4). We keep m constant for all steps but train one step size for each HMC step. Figure plots the progression of E 0o [log p()) during sampling for these two cases, before and after training. Before training,
3

---

Under review as a conference paper at ICLR 2021
-1.0
Initaa distribution Initial distribution Targetato
-1.5
0.6
E(logp(e)
-2.0
0.4
HMC with [D] pre-training
-25
HMC with ID1 post-training
0.2
HMC with ID2 pre-training
HMC with ID2 post-training
0.0
-3.0
8
10
-2
Step HMC chain
200 [logp(z)] as function of step t in the HMC chain for initial distribution 1. (00 0.25), and initial distribution (O 4), before and after training. The 'true" value, Ep [log(())) is also plotted. (Right) The target pdf along with the two initial distributions.
Figure 1: (Left) E
when using the narrow initial distribution, Bogo [log p(x) initially starts above the true value but converges from above as the marginal HMC distribution, 9 spprads to cover the target. However, after training by optimizing eq ( all the step sizes have become very small causing the HMC chains to remain at their initial sampled positions which is obviously detrimental for convergence. To avoid this, we can use a wide initial distribution. Figure Figure[[shows shows that, in this case, optimizing (I) greatly speeds up convergence to the true distribution. We give the code for all our experiments in supplementary material.
METHODS FOR FINDING A SUITABLE INITIAL DISTRIBUTION
In order for the method be useful, we need use an initial distribution that sufficiently cov- the mass of the target. On the other hand, it would be unhelpful to distribution that was overly spread out relative to the target because we are focusing on short parallel HMC chains so we would like keep the burn-in time neering problem, we provide an automatic method in this section that we have empirically evaluated variety of applications and have found give consistently good results. The main idea is to use a variational approximation to the target as the initial distribution, as done by Hoffman (2017). This should be distribution that be easily sampled from and easily tuned fit the target, e.g. Gaussian distribution or normalizing flow TTaakk & Vanden Vanden-Eijnden 2010 Rezende & Mohamed, 2015). Rather than use the standard ELBO, we use a-divergence minimization (Hemánnezz-Lobat et al. 2016) for training. The a value dictates the mass covering behaviour of the resulting approx- imation, with
minimum. Although ultimately this trade-off is engi-
corresponding to the standard mode seeking DKL L. (g (z))))))(())) minimization and a = 1 corresponding to the mass covering Dx L (p(z)) (1)))))) minimization. We compare both values in experiments. The a-divergence is very useful in this context as it pro- vide us with a mass covering approximation without the use of samples from the target, it requires only the (unnormalized) target density, regardless of the value of a. However, if samples from the target are available, then alternatively (x) can be tuned via maximum likelihood which is also mass-covering.
Note that the previous approaches will produce an initial distribution that fits the target, but do not guarantee that this initial distribution will be broad enough. To address this, we also allow our method to automatically adjust the width of the initial distribution as necessary to keep g (z) as closely matched to p(x) as possible. This achieved by applying a sealar scale factor centered around the mean y to each sample z; from the initial distribution, i.e. using f = s(x; - 14) + as our sample, In particular, we train s by minimizing the Sliced Kernelized Stein Discrepancy (Gong et al. 2020) or ! between the final state distribution and the target, SKSD(g (x),
(7)
We the metric referred the maxSKSD Gong al. 2020

---

Under review as a conference paper at ICLR 2021
Table 1:
KSD between the HMC samples and the target distribution for the basclines and the variations of our method on each of the synthetic target distributions.
Gaussian
Laplace
Dual Moon
Mixture
Wave 1
Wave 2
Wave 3
0
0.0677
0.0005
0.2370
0.0004
0.0224
0.0525
0.0462
Q =1
0.0009
0.0004
0.8637
0.0010
0.0067
0.0158
0.0801
SKSD & a 0
0.0008
0.0016
0.1684
0.0004
0.0017
0.0020
0.0217
SKSD & a = 1
0.0009
0.0014
0.2528
0.0004
0,0017
0.0019
0,0317
min 0.25
0.0364
0.0005
1.1553
0.0846
0.0645
0.9447
0.0465
NUTS
0.0044
0.0016
0.2326
0.0023
0.0124
0.0260
0.0965
The SKSD is a differentiable scalable discrepancy measure requiring only samples from q and gradients of the target, V.,logy" (v)). This objective encourages suitable values of because, say 8 is too small, then the & training (which occurs jointly with s training) will result in a degenerate (2) (z) far from the target. The SKSD measures this discrepancy and provides learning signal do
("'(x)) for increasing s. Conversely, if is too large then Is (x) will also be far from the target since the chain will not be able to account for the extremely poor initialization. The SKSD will then result in learning signal to decrease to more reasonable value. Given that the SKSD is tractable objective that measures the discrepancy between d/2 (x) and p(x), is theoretically feasible to use the SKSD to optimize & too, but we found that the SKSD does not perform well when optimizing too many parameters, which is why (T used instead. Note, however, that the SKSD works very well in practice when only tune the single scalar parameter
EXPERIMENTS
5.1 2D DISTRIBUTIONS
We evaluate our tuning method on range of synthetic 2D target densities shown in Figure 2a The equations for these distributions are listed in Appendix We use 30-step HMC chains initialized with a factorized Gaussian approximation, which is trained by minimizing the a-divergence. We consider variations of our method in which we optimize individual step sizes and masses for each dimension and HMC step using (1) but vary the value of a E {0, and whether not to include the tuning of the scaling s by minimizing the SKSD (s = when not tuned). We include two baselines for reference. The first one taken from Hoffman (2017). where the step size in each dimension l: is given by 0и with d being the standard deviation in dimension k; as estimated by a Gaussian fitted by minimizing the a O-divergence and cp being adjusted as to control the minimum acceptance probability over batch of parallel chains. In line with Hoffman (017 we set this minimum acceptance probability to 0.25. The second baseline is the popular No- O-Turn Sampler (Hoffman & Gelman 2 We use the dual averaging variant of the No-U-Ture Sampler, which includes a method for tuning the step size as to encourage an equivalent of the HMC average acceptance probability towards target 8 E (0, 1). We set ó 0.2 in experiments. To quantify convergence to the target, we used the Kernelized Stein Discrepancy (KSD) (Liu et al.) 2016 Chwialkowski al. 2016) between the generated samples and the targets. The results shown in Table [ Our method consistently outperforms the baselines. Furthermore, the methods that tune s by minimizing the SKSD tend fit the targets better than those methods that do not tune this scale factor. For some distributions, the scale factor effectively helps prevent mode seeking. In Appendix A.2 we confirm this quantitatively by comparing expected log target values and find that narrow initial distributions (a = 0) often lead to artificially high values with the SKSD preventing this pathology.
5.2 DEEP LATENT GAUSSIAN MODELS
We now use our method to train Deep Latent Gaussian Models or DLGMs (Kingma & Welling 2014 Rezende et al. 2014). DI.MMs are popular generative models that describe observed data by the following generative process: first sample a latent variable z ~ NO, I) and then sample x
We use the implementation from http:///a//gggggt....ooo//o//anr
our experiments.
5

---

Under review as a conference paper at ICLR 2021
5.0
a=0
4.5
a=1
Average rank
Gaussian
Laplace
Dual Moon
3.5
www
ww
3.0
Mixture
Wave
Wave 2
2.5
2.0
Initial distribution
maxELT (ours)
maxELT P, 0.65 Grid search
Wave
SKSD
(ours)
(a)
(b)
Figure 2: (a) Histograms of targets generated by rejection sampling. (b) Average ranking of marginal KI.-divergences for each method in the Alanine Dipeptide experiment (lower is better).
from some distribution parameterized by the output of deep neural network that receives z as input, i.c., x ~ p=(x))) with 0 being the neural network parameters. The Variational Auto-Encoder (VAE) popular method for training these models. This method works by fitting a factorized Gaussian approximation q (z(ar) to the posterior Po(=(r), (z). The means and variance parameters of (=(2))) are given by a deep neural network that receives x as input. The parameters and e are then optimized jointly by maximising
ELBO objective. In our experiments, we consider the HMC sampling distribution, 1o (z [z]xr), as a variational approximation with the initial distribution being a factorized Gaussian distribution (0)) (=|zr) whose parameters are given by a deep neural network receiving x as input. The parameters y are trained by minimizing an a-divergence. We train (ó, 8) by maximizing the objective given by ① Note that train only one of HMC hyperparameters for all x values, iee, we assume the same hyperparameter values will work on all po (=(2)r) targets, independently of x. One could make the hyperparameters depend æ through amortization network, but we found
(T)
that this did not improve performance. We also include a scaling of g (2|z) with scale factor s trained by minimizing the SKSD. We consider two benchmark datasets: MNIST and Fashion MNIST. As it is common in the literature, use binarized versions of these datasets, with p0(x)2) giving the parameters of a Bernoulli distribution over the pixels. The likelihood is parameterized with the same convolutional architecture used by Salimans al. (2015). Full experiment details can be found in Appendix [..] In these experiments, we found that an appropriate value of s can be successfully found by minimizing the SKSD. The values of obtained are large enough prevent our HMC tuning method to overfit to regions of high posterior density, but also small enough to ensure the generation of accurate samples. More details can be found in Appendix B.2
We evaluate the quality of the models trained with our method on both the MNIST and Fashion MNIST datasets. For each model, we estimate the data marginal log-likelihood, log po(p), on each of the test images using Hamiltonian Annealed Importance Sampling or HAIS (Sohl-Dickstein & Culpepperi 2012) and report the average value and standard error Table For comparison, we also report logllikelihood values for multiple other methods. Using the same neural architecture, implemented the standard VAE and WWAFlmodels We also implemented another method of tuning (Hoffman) 2017) where the step sizes are adjusted make the minimum average acceptance probability across a batch of images equal to 0.25. We update 0 as in our method and use the same number of leapfrog steps for fair comparison. Finally, we report the best lig-likeldhood values from Salimans et al, 2015): Caterini et al. (2018) which include HMC hyperparameter tuning network architecture
during training and use the on the MNIST dataset). We confirm the difference in average log-likelihood between models is significant by performing a paired t-test for each pairing models, we report the results in Appendix We find that on both MNIST and Fashion MNIST, the HMC based methods generally achieve
(note that these authors only evaluated
B.3
We used the DReG estimator from Tucker al. 2019 for the IWAE.

---

Under review as a conference paper at ICLR 2021
Table 2: Average test set marginal log-likelihood and its standard error for different models for MNIST and Fashion MNIST estimated using HAIS. For models with a scale factor, we also report the final scale after training. For comparison, we report the best log- likelihood values from previous works that tune HMC hyperparameters and use the same architecture.
MNIST
Fashion MNIST
Model VAE DReG-IWAE a= 0 a =
Scale
Mean
Standard Error
Scale
Mean
Standard Error
-85.08
0.2172
108.54
0.6010
-83.73
0.2109
-04,448
0,5841
1.0
-83.48
0.2101
1.0
-104.08
0.5834
1.0
-82.46
0.2073
1.0
-103.57
0.5826
a 0 & SKSD
6.79
-81.91
0.2042
5.58
-103.18
0,5802
a = 1 & SKSD
3.90
-81.94
0,2045
3.59
-102.29
0,5748
Hoffman (2017)
-81.74
0.2046
-103.04
0.5804
Salıaan et al. (2015)
-81.94
Catermni et al. (0118)
-82.62
significantly better performance than VAE IWAE, showing that reducing the approximation bias of the variational distribution with HMC greatly helps model learning. Furthermore, we see that adding the scale factor to our model significantly improves performance as this avoids degenerate behaviour when training the HMC hyperparameters. We note that, without any scaling, a = outperforms a = due to a resulting in a too narrow initial distribution. With scaling, the SKSD automatically widens the initial distribution making the performance between the a values similar. Finally, we observe that HMC based methods top out at similar oog-likeldood values (within around one standard error), we believe this is due to reaching the limits of the architecture on these datasets with no more gains to be made from more accurate posterior approximations.
5.3 MOLECULAR CONFIGURATIONS
For our final experiment, we evaluate our method on the complex real-world problem of sam- pling cquilibrium molecular configurations from the Boltzmann distribution of the molecule Ala- nine Dipeptide. The unnormalized target distribution given the atom coordinates is e e=(()), where u is the potential energy of the system, which can be obtained using the laws of physics. While this problem is usually tackled through Molecular Dynamics (MD) simulations, giving a sequence of highly correlated samples, we aim to produce samples using our trained short HMC chains. We do not operate directly on the Cartesian coordinates but apply the coordinate transform presented by Noć et al. (2019), see also Appendix C.2 to map some of the Cartesian coordinates to bond lengths, bond angles, and dihedral angles giving dimensionality of 60. Our methods use normalizing flow based on real-valued non-volume preserving (RNVP) transformations (Dinh et al.) 2017) as initial distribution, followed by 50 HMC steps. We used various methods to train the initial distribution and HMC hyperparameters (using individual masses and step sizes for each dimension and HMC step in this case). The RNVP models were trained with a = 10, diverrgenc and maximum likeli- hood. The latter was done using 10 training data samples obtained via a MD simulation In these experiments, we first trained the initial sampling distributions (via ML or a-divergence) and then kept them fixed when tuning the HMC hyperparameters and the scale factor. To evaluate the perfor- mance of the different methods, a new MD simulation was run to obtain 10 test samples. Since we the model likelihood is intractable, we approximated the 60 marginal distributions by kernel density estimates from the MD test samples and compared them using the KL-divergence with the corre- sponding density estimates for the samples of the HMC tuning methods. This performance measure also has practical relevance as the marginals of proteins, especially those of the dihedral angles, determine important material properties such how the protein folds (Ramachandran al. 1963). The HMC hyperparameters were tuned by maximising the expected log-target without adjusting the initial distribution (referred to as maxELT) or by maximising the expected log-target whilst also tuning the scale factor of the initial distribution by minimizing the SKSD (referred to as maxELT & SKSD). As baseline, we also optimized the HMC parameters via a grid search, keeping step sizes
'The same dataset was used to obtain the mean and variances required in the normalization step of the coordinate transform.
7

---

Under review as a conference paper at ICLR 2021
Table 3: P-values of the Wilcoxon test comparing the model with HMC parameters tuned by maxELT & SKSD with the baseline models based on the KL. divergences of the marginals. Bold means maxELT & SKSD improves on the baseline at the given p-value.
Grid search
P. = 0.65
a = 0
0.158
0.627
1
0.0102
0.0247
ML
0.883
2.63 10
and masses constant across dimension and HMC step and varying these two constants in grid, picking the combination that gave the lowest median marginal KL-divergence to the MD training data. We also implemented another baseline, adjusting the step size constant such that the average acceptance probability was 0.65 (referred to as P. = 0.65). Further details about the implementation are given in Appendix[C.I
The marginal KI -ileeivegeences vary greatly in magnitude so we use rank based metric to summarize the results. For each initial distribution type and marginal, we assign each HMC method a rank (1-5) according to the ordering of KL values on that marginal. We then average over the 60 marginals for each of the 15 methods, resulting in Figure 2b For the a = trained initial distribution, maxELT & SKSD clearly results in the lowest average ranked marginal KL-divergences, meaning its sam- ples are closest to the target. For the other two initial distributions, maxELT & SKSD is generally on-par with the best baseline. To confirm this observation, we performed two-sided Wilcoxon test [Wilcoxon] 1945) to determine whether the differences between average rank across KI KI--divergences significant not. The resulting p-values for comparisons between maxELT & SKSD and the baselines are shown in Table 3) further details are given in Appendix C.2 The pair of models show- ing significant difference, i.c., the respective p-value is smaller than 0.05, are those where our method outperforms the baseline. For the other cases, it is on par with them.
6 DISCUSSION AND RELATED WORK
Our experiments show general trend that, when we solely optimize (I), the value of used sig- nificantly affects performance. However, when applying the SKSD scaling, this difference becomes smaller, showing that this technique is useful for automatically finding a suitable initial distribution. In practice, it may be simple and effective to just use the standard ELBO method (i.e. = 0) to find an initial distribution for HMC with the SKSD scaling to ensure that it is wide enough.
Comparing to other methods for tuning the HMC hyperparameters, there are some methods also inspired by variational inference. [Salimans et al. 2015): Wolf et al. (2016) use an approximation to the reverse dynamics of the chain to put a lower bound on the standard ELBO and make the optimization tractable. The method is then dependent on the accuracy of this reverse approximation with the concern that this new lower bound gets looser and looser as the the HMC chain length gets larger and larger. As our experiments in (..2 show, we can do just as well as these methods by directly optimizing a measure of convergence without the need for extra approximations. Caterini ct al. (2018) also construct an alternative ELBO for HMC, however, they only sample the auxiliary momentum variables once at the start of the chain which reduces the empirical performance of HMC. In contrast to these methods, some gradient based tuning techniques do not use ideas from variational inference but instead optimize a proxy for mixing speed. Levy al. (2018) generalize the standard leapfrog integrator used in HMC with multi-layer perceptrons which are then trained by maximising modified version of the expected squared jumped distance. We improve upon this objective by directly optimizing convergence speed and using gradient information from the target distribution itself. It would be an interesting direction to use our objective to train this generalised leapfrog operator: Finally, there is the method of Tiitsias & Dellaportas) (2019) who consider the gradient based tuning of the Markov transition operator, P0 11) in the case of a single long MCMC chain. For the objective, they use the expected acceptance probability for the next step in the chain, regularized with the entropy of pa ps ( 11). Unfortunately, as this entropy is intractable when using the leapfrog algorithm the Markov transition operator, it cannot be applied to HMC in its current form.
There are also many non- gradient based heuristics for tuning the HMC hyperparameters. The popu- lar No-U-Turn Sampler (Hoffman & Gelman 201T) can adaptively set the number of leapfrog steps L to avoid U-turns and find global constant for the step sizes by adjusting the average acceptance
8

---

Under review as a conference paper at ICLR 2021
outperform NUTS even though do not adaptively set L. With our objective, we can tune individual step sizes (and masses) for each dimension and step in the chain, allowing for much higher degree of granular control over the algorithm. Furthermore, we do not need to rely on 'rules of thumb' such as standard acceptance rate targets used in many al- gorithms (Hoffman & Gelman 2011 [Hofmman| (017) but
rate. In section 5.1
found
automatically tune all continuous hyperparameters using information from the target distribution directly.
Our work also builds upon methods from statistical mechanics. The Boltzmann Generator (ooé et al. 2019) opened up this line of research by applying a normalizing flow to the problem of sam- pling molecular configurations. We found we can improve upon this by using the flow as the initial distribution for HMC, fine tuning the flow samples with our short chains. Our Alanine Dipeptide ex- periment comes from the recent work of et al. (2220) on stochastic normalizing flows consisting of stochastic steps interspersed between deterministic steps in normalizing flow. An interesting fu- ture extension would be to combine the approaches and tune parameters within the stochastic layers with our objective.
CONCLUSION
In this work, presented new objective motivated by the variational inference ELBO that can be easily used for the gradient- gradient-based optimization of HMC hyperparameters. We provided fully au- tomatic method for choosing an initial distribution for the HMC chain that reduces burn-in time and aids optimization. Evaluating with or improves upon existing tuning methods. We hope this encourages further work applying this idea to other methods that use HMC and that would benefit from increased convergence speed.
multiple real-world problems, we find our method is competitive
REFERENCES
Kim C. Border. Differentiating an integral: Leibniz' rule, 2016.
Anthony Caterini, Arnaud Doucet, and Dino Sejdinovic. Hamiltonian variational auto- encoder. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 8167-8177. Curran Associates, Inc., 2018. URL httt:://paperr nips ctppsper/ www.oonnaaaaatttnccceeeeeeeeeeee .pf
Kacper Chwialkowski, Heiko Strathmann, and Arthur Gretton. A kernel test of goodness of fit. volume 48 of Proceedings of Machine Learning Research, pp. 2606-2615, New York, New York, USA, 20-22 Jun 2016. PMLR.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In- ternational Conference on Learning Representations, 2017.
Wenbo Gong, Yingzhen Li, and José Miguel Herannde-ooeaaneeaaoboooooa Sliced kernelized stein discrepancy. 1653
arXiv:2006.1653 16531, 2020.
Ernst Hairer, Christian Lubich, and Gerhard Wanner. Geometric numerical integration illustrated by the stömmer-verle method. Acta Numerica, 09:399-450 2003.
José Miguel Hemándezz-Lobtt. Alpha-divergence minimization for bayesian deep learning. Bayesian Deep Learning Workshop NIPS 2016. Centre Convencions Internacional Barcelona, Barcelona, Spain, 2016.
José Miguel Hernández-Lobate Yingzhen Li, Mark Rowland, Thang Bui, Daniel Hernandez- Lobato, and Richard Turner. Black-box alpha divergence minimization. volume 48 Proceed- ings of Machine Learning Research, pp. 1511-1522 New York. New York, USA, 20-22 Jun 2016. PMLR. URL httt:://proceeet ngs. mm press htt////////h/////eeeee// obatob 6..
htm
Matthew Hoffman and Andrew Gelman. The no-u-turn sampler: Adaptively setting path lengths in hamiltonian monte carlo. Journal of Machine Learning Research, 15, 11 2011.

---

Under review as a conference paper at ICLR 2021
Matthew D Hoffman. Learning deep latent gaussian models with markov chain monte carlo. In International conference on machine learning, pp. 1510-1519, 2017.
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183 233, 1999.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1:Xiv:44/2.6900 2014.
Diederik P. Kingma and Prafulla Dhariwal. Glow: Generative Flow with Invertible Ixl Convolu-
tions. Advances in Neural Information Processing Systems, 2018.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. In International Conference on Learning Representations, 2014.
Daniel Levy, Matthew D. Hoffman, and Jascha Sohl-Dickstein. Generalizing hamiltonian monte carlo with neural networks. In International Conference on Learning Representations, 2018.
Qiang Liu, Jason Lee, and Michael Jordan. A kerellized stein discrepancy for goodness fit tests. volume 48 of Proceedings of Machine Learning Research, pp. 276-284, New York, New York, USA, 20-22 Jun 2016. PML.R
Radford M. Neal. Probabilistic inference using markov chain monte carlo methods. Technical report, University of Toronto, Department of Computer Science, 1993.
Radford M Neal. Meme using hamiltonian dynamics. Handlboo of markov chain monte curlo, 2 (11):2, 2011.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457), September 2019. ISSN 0036-8075, 1095-9203, doi: 10.1 10.112/sciencce.aw..1...
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. arXiv:1912.02762 [cs, stat), December 2019.
G. N. Ramachandran, C. Ramakrishnan, and V. Sasisekharan. Stereochemistry of polypeptide chain configurations. Journal of Molecular Biology, 7(1):95-99, July 1963. ISSN 0022-2836. doi: PMLR. 10.1016/500222-886663))0000
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. volume 37 of Proceedings of Machine Learning Research, pp. 1530-1538. Lille, France, 07-09 Jul 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. volume 32 of Proceedings of Machine Learning Research, pp.
2014.
Herbert Robbins and Sutton Monro. A stochastic approximation method. Annals of Mathematical Statistics, 3304):000-077, 1951.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, pp. 1218-1226, 2015.
Jascha Sohl-Dickstein and Benjamin J. Culpepper: Hamiltonian annealed importance sampling for partition function estimation. Technical report, Redwood Center for Theoretical Neuroscience University of Californa,, Berkeley, 2012.
Esteban G. Tabak and Eric Landen-Eijnder Density estimation by dual ascent of the log- lodelio.d..od. Communications in Mathematical Sciences. 8(1):217-233, 2010. ISSN 15396746, 19450796. doi: 10.4310/CMS. 2010........000..11111.. al...
10

---

Under review as a conference paper at ICLR 2021
Achille Thin, Nikita Kotelevskii, lean-Stanislas Denain, Leo Grinsztajn, Alain Durmus, Maxim Panov, and Eric Moulines. Metflow: A new efficient method for bridging the gap between markov chain monte carlo and variational inference. arXhv:2002.12255, 2020. URL https://arxiv
org/ 2002. 12253
Michalis Titsias and Petros Dellaportas. Gradient-based adaptive markov chain monte carlo. In Advances in Neural Information Processing Systems, pp. 15730-15739, 2019.
George Tucker, Dieterich Lawson, Shixiang Gu, and Chris Maddison. Doubly reparameterized gradient estimators for monte carlo objectives. In International Conference Learning Repre- sentations, 2019.
Frank Wilcoxon. Individual comparisons by ranking methods. Biometrics Bulletin, 1(6):80-33, 1945.
Christopher Wolf, Maximilian Karl, and Patrick van der Smagt. Variational inference with hamilto- nian monte carlo. arXiv preprint
2016.
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic Normalizing Flows. arXiv:2002.06707 (physics, star), February 2020.

---

Under review as a conference paper at ICLR 2021
A 2D DISTRIBUTIONS
A.1 EQUATIONS FOR 2D TARGETS
Table Unnormalized log densities for the 2D distributions used in the first experiment.
Name
Unnormalized log density, log x x=(x)
Gaussian
2
60
19
21
40 19 19-2)
Laplace
2 5
1/2 - 5
Dual Moon -3.125
Vv o + 23
-2 +
log
exp (-0.5 2--
exp
-0.5 (2-2)
2)
Mixture
log
E 1= exp
-0..5 (x1
5cos
(2km)
(x2
5sin
26x
27
7
2
Wave
-0.5
0.4
log
exp
-0.5
IP2 02n+sinn 550a11)
Wave 2
exp
00.5
sin (0.5 +3xxpp
0:30
0.5
(2)
1)9)
0.35
2
log
exp
-0.5
(0..5mn
0.4
+
Wave 3
2g sin (05mm1)))
14as =--1
exp
-0.5
A.2 COMPARISON BETWEEN -E, log (x)]
AND -E, Em
Here, we compare EE, log (x)]
(x))
with -F [logp" (x)] quantify the mode seeking be- %
haviour of the different methods. We display our results in Table 5 We that the SKSD helps the method better match the ground truth and on some distributions c.g. Gaussian, helps prevent -F, log p" (x:] from underestimating -pp log (x))- This may oceur with a narrow initial distribution (a 0) because the optimization will encourage the chains to remain in the region of high log target, artificially inflating the E (r) log p" (x)] value.
Table 5: -E (r) [xx)] values for the baselines and the variations of our method on each of the synthetic test distributions. The ground truth value -E, logg" (xr) is found using rejection sampling.
Gaussian
Laplace
Dual Moon
Mixture
Wave
Wave 2
Wave 3
Ground-truth
2.8083
2.0075
0.8511
0.9282
0.4994
-0.1703
0.1483
0
2.4911
1,9937
1.0315
0.9216
0.4845
0.0463
0.1314
0 =
2.7861
2.0431
2.9218
0.9197
0.4982
-0.1050
0.6938
SKSD & a 2.8390
2.0074
0.8439
0.9174
0.5142
-0.1293
0.2500
SKSD & a =
2.8183
2.0281
0.7987
0.9198
0.4993
-0.1814
1.2002
min 0.25
3,0148
1.9140
4.1515
2.4894
1.4081
1.4220
1.0397
NUTS
2.7862
1.9955
0.7828
0.9379
0.4959
-0.1979
0.0911
A.3 TRAINING TIME
We provide the training times for the 2D experiments in Table We note that the SKSD times are almost exactly double the SKSD times. This is because the majority of training time is taken up by sampling from the HMC chains. Each iteration, we use one batch of 100 HMC samples to estimate the expected log target objective and another separate batch of 100 is used to estimate the
12

---

Under review as a conference paper at ICLR 2021
SKSD resulting in the approximate doubling of total training time. In practice, this can be made more efficient by using the same batch of samples to estimate both objectives resulting in similar training times whether not the SKSD is included.
Table 6: Training time for synthetic 2D problems CPU (sec / 100 iters).
Gaussian
Laplace
Dual Moon
Mixture
Wave
Wave 2
Wave 3
22.62
25.33
33.44
61.17
21.75
65.91
65.08
1
23.60
25.92
34.15
62.35
21.98
67.62
67.02
SKSD a = 0
45.44
50.01
63.8
128.71
41.67
131.15
129.79
SKSD & a
46.44
50.28
65.35
132.53
42.2
133.09
133.48
min p = 0.25
4.96
4.16
3.66
16.83
2.8
6.78
7.14
B DEEP LATENT GAUSSIAN MODEL EXPERIMENTS
B. DETAILS ABOUT THE SETUP
In these experiments, we set the dimension of the latent variable 2 to 32. For our HMC variational distribution, (T) (2(ar) we set T 30 and use 5 leapfrog updates. We only consider training the step sizes here and leave all masses 1. As there are multiple parameters being optimized jointly, we summarise the entire method in this section for clarity. The parameters being optimized are: the decoder neural network parameters, é the HMC step sizes (a total of 30 x 32 960 scalar values), & the encoder (q,) (2lz)) neural network parameters and the scale factor used to scale (0) samples from q (z|xr). There are then optimization objectives:
L1 E (
% (a|=) (oggo(x;
Our tuning heuristic
(3)
C2 SKSD (T) (2|x)) po(2)=))
(4)
C3 E (0) (=|z)
log (2, 2) log (0) (z|=))
Standard ELBO
(5)
2
w' po(x, 2²)
C1 21 2 (=(w)
log w²
DReG IWAE objective
1= wi
q (0) (2*1z)
(6)
Using these 4 objectives, we then follow Algorithm during training. We use Tj-11 +- Adam, (C.)
to denote one gradient step using the Adam optimizer (Kingma & Ba 2014) maximising objective L. with respect to parameter 7). We note that we introduce 10" pre-training steps before starting HMC optimization as these updates are very quick to do and make sure the HMC optimization has reasonable starting point. We see that when a = O, we train v by maximising (5) which is equivalent minimising the = 0 divergence with the target. When = the objective for zi which is the doubly reparameterized version the IWAE objective (Tucker al. 2019). This is equivalent to minimizing the Hernandzz-Lood Lobato (2016).Fo For the case where we do not use the SKSD then we simply omit from model and omit the updates from Algorithm I For the VAE and IWAE baselines we simply
use (6) as
divergence with the target, shown by
13

---

Under review as a conference paper at ICLR 2021
the pre-training steps for the full 1.5 x 102 steps using the 0 or 1 updates respectively.
Algorithm 1: DLGM Training Algorithm
Pre-training steps; for f = l...., 10a do
a then
D.+1 < Adamo, (L3); 1++ + Adamg, (Ca);
end
a then Oc+1 & Adamo, (La);
1++1 f A (L4);
end end
IMC-Iraining steps; for 10a 1.5 10 do
if' = then
d Adam,, (C3);
end a then
1 + Adam, (L4);
end
0++1 Adams, (L1);
8.41 f Adame, (C1);
8141 Adam.. L2);
end
B.2 EPFECTIVENESS OF THE SCALING
a=0
Epizia [logp(x,)) 2)] Ea/ica [log [lo 2)1
12
a=
SKSD scale
10
SKSD scale a=1
128
256
Scale factor
Figure Ep(ste) po =)] (x(a) [log po(, =)] averaged over 200 randomly chosen MNIST test images for range of fixed scalings used during training. The ground truth poste- rior p(z[z) is estimated using 100 HAIS samples. The final scales found when running the training with the SKSD are also plotted.
To demonstrate the scaling's effectiveness, we run our optimization scheme on the MNIST dataset
for range of fixed scale factors and then compute [p(=]=] log po =)] E g (=|=) [og pa (æ, =)]
with samples from p(=|zr) found through HAIS (Sohl-Dickstein & Culpepper 2012). The results for
14

---

Under review as a conference paper at ICLR 2021
and are shown in Figure For small scales, the metric is negative implying 4a oversampling high probability regions of the target with a higher scale factor alleviating this issue. Figure BasureBlalso also shows the scale found by SKSD training when this is included the optimization run, it find an appropriate scale factor that is large enough to prevent this pathology whilst also ensuring stable performance.
(=(xx)
(T)
B.3 DEEP LATENT GAUSSIAN MODEL PAIRED T-TESTS
We carry out paired t-test for each model pairing with the null hypothesis that the two population means of the log- ligllikel log p(z) are equal. Log-likelihood values are paired between models by observed data point I. The p-values for the tests on the MNIST dataset are given in Table 7 and the p-values for the Fashion MNIST dataset are given in Table g A 0 value represents that the p-value is numerically indistinguishable from O.
Table 7: p-values for paired I-tests on test log- likelikel values on the MNIST dataset.
DReG-IWAE
Hoffman
2017)
& SKSD & SKSD
VAE DReG-IWAE
0
0
4.93e-18
0
0
a 0
0
0
-
1.07e-115
1.50e-104 0.2405
2.04e-173
a & SKSD & SKSD
1,45e-14 14
2.95e-19
Table 8: p-values for paired t-tests on test log likelihood values on the Fashion MNIST dataset.
DReG-IWAE
0
Hoffman (011)
& SKSD & SKSD
VAE DReG-IWAE
2.51e-22 1.25c-94
2.12c-205
4.54c-217
a=0 0 a=1
6.46c-42
8.48c-124
5,90e-129
8.45e-28
1.05e-262 1.35e-135
2.04e-40
a = 0 && SKSD
4.81e-4
a && SKSD
1.27e-79
B.4 TRAINING TIME
We give the training times for the DLGM training algorithms examined in Table 9 that the VAE and IWAE baselines are much faster as they do not involve sampling from any HMC chains. We also see that the times when the SKSD loss included are much longer than when is omitted. This is because, in our experiments, we use single HMC sample to estimate the expected log target objective, L, each gradient step, however, we estimate the SKSD objective, C2, using a further 30 samples resulting in much longer times per training step. We decided to not use the same batch of samples estimate both L, and L2 our experiment in order to clearly identify the effect of the SKSD. Keeping total training iterations the same, if we had used these extra samples to estimate &, then it would be unclear if any improvement was down to the SKSD or just to reduced variance in the estimate for L1. In practice, it would be more efficient to estimate both Lı and L2 using the same batch of HMC samples and reduce the total number of training steps as the reduced variance of L, will help speed up training.
We first note
C MOLECULAR CONFIGURATIONS SAMPLING EXPERIMENTS
C.1 DETAILS ABOUT THE SETUP
All the RNVP models we used as initial distributions have the same architecture. They consist of five coupling blocks composed of two alternating coupling layers. The scale and shift iith the
15

---

Under review as a conference paper at ICLR 2021
Table 9: Training Time for DLGM on GPU (sec 1000 iters).
MNIST
Fashion MNIST
VAE
2.95
3.34
DReG-IWAE
4.34
4.45
0=0 0
102.73
103.13
103.66
104.02
& SKSD
1514.93
1520.27
a = & SKSD
1527.11
1539.52
Hoffman (017)
126.84
133.13
coupling layers is given by fully connected networks having three layers with 128 hidden units each. Before each coupling block, we applied actívation normalization (Kingma & Dhariwal 2018).
For use in scaling the initial distribution samples, we use an empirical estimate of the nor- malizing flow sample mean using 10 samples. To calculate the SKSD metric, to save computation, we use fixed vectors for what Gong et al. (2020) refer to as gr instead of optimizing them. This was because found g. ended up very close to one-hot vectors during optimization anyway so this approximation does not make a large difference to the method.
Training the RNVP model with the a-divergence when a was achieved by using the re- verse KL-divergence as training objective (Papamakarios et al. 2019 Hernández-Lobato et al.) 2016).
As baselines for tuning HMC parameters we used grid search as well as training the ac- ceptance probability. For the grid search, 25 different parameter settings were tested for each initial distribution. In each setting the step size and log mass is held constant over all HMC layers but different step size constants are tested in combination with different log mass constants giving 25 total combinations. The metric used to find the best combination was the median KL-divergence over all 60 marginals computed using 10a samples against the 10a molecular dynamics training data. When training the acceptance probability, the log mass was held at 0 for all layers and the step size was a constant for all layers with this step size constant being adjusted so that the average acceptance probability was 0.65. The constant was updated each training iteration with update of the form 6++1 f ar (0.65 Pa). where is the average acceptance probability and ae being a parameter decreasing according the Kobbins-Momet conditions (Robbins & Monro 1951)
When testing the models, the KL-divergence was computed by first computing a kernel den- sity estimate for each marginal using Gaussian kernel with bandwidth chosen using Scott's rule and then finding the KL-divergence between the kernel density estimates using numerical integration.
C.2 COORDINATE GROUPS AND MODEL COMPARISON
The coordinate transformation. introduced in Noć et al. 2019) which we used splits the feature dimensions in four different groups: 17 bond angles, 17 bond lengths, 17 dihedral angles, and Cartesian coordinates, adding up to 60 dimensions in total. Due to their differing physical mean- ing and relevance, they follow different distributions. For each group, three sample marginals are shown in Figure 4. Figure 5) and Figure 6] We used the MacCallum lab's implementation for our coordinate transform which can be found httts:://itth// http://ttttbb..m..aaa.coc..o ccal 1 umll lab/ smitnnnnneeerrttrr
The bond lengths and angles follow mostly unimodal, almost Gaussian, distributions which is due to the regular vibrations of the atoms within the molecule. The dihedral angles can have multiple modes while the Cartesian coordinates are quite irregular. As evaluating the KL-divergences on one- dimensional distributions is much easier than approximating it for higher dimensional distributions, we compared our models using the KL-divergences of these marginals. Figure visualizes the distribution of all Kl--divergences.
16

---

Under review as a conference paper at ICLR 2021
Ground truth
maXELT
Bond angles
maxELT SKSD
Grid search
0.65
Bond lengths
Dihedral angles
Cartesian coordinates
-2
-2
2
-2
Figure 4: Sample distributions of marginals from the four coordinate groups. The graphs compare the ground truth with The RNVP was trained with the a = O-divergence and the HMC parameters were tuned with the indicated methods.
models having RNVP as initial distribution followed by 50 HMC steps.
17

---

Under review as a conference paper at ICLR 2021
Ground truth maXELT
Bond angles
maxELT SKSD
Grid search 0.65
Bond lengths
Dihedral angles
Cartesian coordinates
-2
-2
2
-2
2
Figure 5: Sample distributions of marginals from the four coordinate groups. The graphs compare the ground truth with The RNVP was trained with the a = 1-divergence and the HMC parameters were tuned with the indicated methods.
models having a RNVP as initial distribution followed by 50 HMC steps.
18

---

Under review as a conference paper at ICLR 2021
Ground truth
maXELT
Bond angles
maxELT SKSD
Grid search
0.65
Bond lengths
Dihedral angles
Cartesian coordinates
-2
2
-2
2
Figure 6: Sample distributions marginals from the four coordinate groups. The graphs compare the ground truth RNVP was trained via maximum likelihood and the HMC parameters were tuned with the indicated methods.
with models having RNVP as initial distribution followed by 50 HMC steps. The
19

---

Under review as a conference paper at ICLR 2021
Cartesian maxELT & SKSD
Bond
0.004
0.6
0.003
0.002
0.001
0.0
0= 0
a=1
ML
0
a=1
ML
Angle
Dihedral
0.004
0.2
0.002
0.1
0.000
Figure 7: Violin plot of the Kl--divergence of the marginals between the model tuned using maxELT & SKSD and the test set. The results are split up into the four different variable groups, i.e. Cartesian coordinates, bond lengths, bond angles, and dihedral angles. The median marked by an x and the inner bars are the upper and lower quartiles.
Table 10: P-values of the Wilcoxon test comparing models with initial distribution trained with O-divergence. Values in bold mean that the showed the top row model has lower KL- divergences than the left column model, i.c. the top row model is the better model at the significance level given by the p-value.
maxELT
maxELT & SKSD
Grid search
P, = 0.65
maxELT
0.00056
0.00064
0.00065
maxELT & SKSD
0.00056
0.158
0.627
Grid search
0.00065
0.158
0.489
Pa 0.65
0.00065
0.627
0.489
Clearly, there a few outliers to both the high and the low end. To not let them distort results, we decided to use the rank based metric to summarize our results. Furthermore, we performed a paired two-sided Wilcoxon rank test for pairs of models assess the statistical significance of our results. The null hypothesis for this test is that the difference of the KL-divergences has a symmetric distribution around 0, i.e. there is no model with marginals consistently closer to the ground truth than the other. The resulting p-values are listed in the tables Tabll 1O Table 11] Table 12 and Table 13 Here, we also included an initial distribution that was trained with the 2-divergence. Due to the poor performance, we did not include it in the main text. Overall, we see again that our method of tuning HMC parameters, i.c maxELT SKSD, is competitive or outperforms the baselines, i.e. the grid search and the training to get an average acceptance probability of 0.65.
Table 11: P-values of the Wilcoxon test comparing models with initial distribution trained with 1-divergence. Values bold mean that the showed the top row model has lower KL- divergences than the left column model, i.e. the top row model is the better model at the significance level given by the p-value.
maxELT
maxELT & SKSD
Grid search
P, 0.65
maxELT
0.0017
0.024
0.015
maxELT & SKSD
0.0017
0.010
0.025
Grid search
0.024
0.010
0.023
P, 0.65
0.015
0.025
0.023
20

---

Under review as a conference paper at ICLR 2021
Table 12: P-values of the Wilcoxon test comparing models with initial distribution trained with 2-divergence. Values in bold mean that the test showed the top row model has lower KL- divergences than the left column model, i.e. the top row model is the better model the significance level given by the p-value.
maxELT
maxELT & SKSD
Grid search
0.65
maxELT
2.4e-8
0.14
0.0066
maxELT & SKSD
2.4e-8
0.00069
6.6e-6
Grid search
0.14
0,00069
0.83
Pa 0.65
0.0066
6.6e-6
0.83
Table 13: P-values of the Wilcoxon test comparing models with initial distribution trained with maximum likelihood. Values bold mean that the test showed the top row model has lower KL- divergences than the left column model, i.e. the top row model is the better model the significance level given by the p-value.
maxELT
maxELT SKSD
Grid search
P, = 0.65
maxELT
0.60
0.16
0.0094
mT & SKSD
0.60
0.883
2.6e-5
Grid search
0.16
0.88
6.0c-6
Pa 0.65
0.0094
2.6e-5
6.0e-6
C.3 SCALE PROGRESSION DURING TRAINING
To improve the overlap of the initial distribution with the target, we scale the former and learn the scaling parameter through the SKSD. Figure 8 shows the progression of the scale parameter during training for different initial distributions. The distributions trained using the a-divergence with a tend be mode seeking. Hence, we expect the scale to larger than so the proposal covers the whole target distribution, and this is indeed the case. With needs be shrunken. The model trained with maximum likelihood already
the distribution tends to
be mode covering, has a low KL-divergence (see e.g. FFiuure 2b) so there is not much modification needed, i.e. the scale close to
21

---

Under review as a conference paper at ICLR 2021
1.15
1.10
1.05
Scale
1.00
0.95
0.90
==2
ML
200
400
600
800
1000
Training iteration
Figure 8: Progression of scale factor for maxELT & SKSD models during training. RNVP models 0, 2 was used as initial distribu-
trained by maximum likelihood and the a-divergence with tion.
22


Reviews:
[{"review_id": "UgKA9qNWfJw", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe 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.\n\nPositives:\nThe submission suggests a practical approach for tuning HMC that remains a challenging problem. The combination of the different objectives is new as far as I am aware.  Empirical experiments are provided to justify the approach on standard benchmark problems, where it is seems to be competitive with state of the art methods, and a more extensive study on sampling molecular configurations. \n\nNegatives:\nI feel that further arguments are needed to justify why the entropy of the proposed state can be ignored when adapting the hyperparameters of the sampler. The paper argues that “Since HMC, by construction, cannot 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”. I am not convinced by this: take a standard normal target, then a leapfrog-integrator with 2 steps, unit mass matrix and step size of sqrt(2) proposes deterministically from a point mass distribution and this happens everywhere on the state space. While this might be an unrealistic example, it is not clear to me how such situations can be avoided in general.\nIt is also not clear to me why the Sliced Kernelized Stein Discrepancy objective automatically adjusts the width of the initial distribution. In equation (4) the discrepancy is between the final state and the target and I fail to see how this relates to the width of the initial density.\n\nRecommendations:\nI vote for a weak reject at the moment. The ideas proposed in the paper are indeed interesting. However, I am not yet convinced that the objectives yield HMC kernels that explore the state space well (so the HMC proposal does not become close to deterministic/completes a U-turn so that entropy comes largely from the initial distribution which is however trained with a different objective). Also the use of the Sliced Kernelized Stein Discrepancy specifically should be better motivated. I am happy to increase my score if the authors better clarify these points.\n\nFurther comments/issues:\nThe authors claim in the abstract that existing approaches “optimize a tractable lower bound that is too loose to be useful in practice”. Can this be backed up more concretely? I understand that such methods (such as Thin et al., 2020) use a looser bound, but not that these types of bounds are useless in practice.\nIn section 3.1, how do the acceptance rates compare for the narrow vs the wide initial distribution? My intuition would be that the acceptance rates for the narrow one are smaller than for the wide one. Would it then be possible to get a better exploration even in this case by including an objective to target an acceptance rate (say increase the stepsize if the acceptance rate is above 0.65)?\n\n\nMinor comments:\nIs it obvious that equation (6) minimizes the 1-divergence? For k=1, is this not the standard VAE/0-divergence, while for k>1 the IWAE objective can be seen as a 0-divergence on an extended space?\nWhat are the \\gamma variables simulated from N(0,I) exactly? Are they really the momentum variables? Are the initial momentum variables not from N(0,diag(m))?\nIn the experiments from Section 5.1, why do you target a minimum acceptance rate of 0.25 and not an average rate of 0.65, which seems a more common choice in the adaptive MCMC literature?\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "OdoXuwwPFBN", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\n### Summary:\n\nThis paper proposes a variational inference based framework to tune some of the hyper-parameters of HMC algorithms automatically. The authors drop the entropy term from regular ELBO formulation, which facilitates a gradient-based approach. However, dropping this term requires extra care for which authors offer an automated-method. Finally, the authors demonstrate empirical validation on several problems.\n\n### Strength:\n\nThe authors do an excellent job of articulating their intuition behind the idea both (see section 3.) While dropping the entropy term from ELBO decomposition is heuristic-based, the explanations are well-formulated, and Figure 1 does an excellent job of getting the point across. \n\nMore so, since dropping the entropy term can cause pathological behaviors, the authors propose a method to ensure wider initial distributions. I commend the authors for the non-trivial engineering that was required to make their ideas work. I also, commend the author's effort of conducting statistical tests and extensive empirical evaluations. \n\n### Concerns:\n\nMy main concern with the work is that it is often on-par with the competing methods--I understand that a new method doesn't need to be SOTA on every benchmark--and the SKSD enabled variants that achieve this performance are prohibitively slow (see Tables 6 and 9.) I could not help but feel concerned when no discussion was offered for an almost tenfold increase in the computational time for training DLGMs. To convince me, I will suggest offering an honest discussion on the run-times of the approaches.\n\nI find the discussion in section B.1important, and believe it should be more formal.  Specifically, I will suggest algorithimizing what objective is used at which stage. Alternatively, authors can choose to restructure this some other way; however, it is too important to be left in its current form. \n\n### Updates after the rebuttal\n\nI like the paper and found the revised version more transparent. I support the engineering approach of the paper; however, as we all know, these papers often require authors to go to greater lengths to convince. After reading the other discussion and reviews, I think the authors can consider a few additional experiments. I would suggest investing in a more involved toy-experiment to better motivate the engineering solutions. If possible, authors can also consider a more careful ablation study to establish the relevance of each component on this toy-model. Further, the authors offered explanations for the training time aberrations; if possible, authors can consider including the equally-fast-variants in the revision to be more convincing. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "OXuCXTBvEqC", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n========\n\nthe article proposes to tune an HMC sampler by maximising E_\\param[\\log target(X_T)] over the parameters of the HMC sampler. Furthermore, the article studies the influence of the initial distribution. While the approach is certainly interesting, I have not found the empirical studies satisfying enough. \n\nComments:\n=========\n1. The article considers a vector \\epsilon as well as a mass matrix. Usually, the parameter epsilon is chosen as a scalar number: choosing epsilon as a vector can indeed also be seen as a particular type of preconditioning (or choice of mass matrix). I have found this part of the paper not extremely well explained.\n\n2. It is indeed also difficult to choose L, and that is mainly what the no-U-turn method tries to automate. In practice, dynamically adapting L can make a lot of difference in high-dimensional settings and/or different parts of the state space exhibit different scales. It would have been very interesting to investigate how the proposed method can be used **in conjunction with** no-U-turn type strategies. Furthermore, it was not entirely clear to me how the \\epsilon was tuned when the no-U-turn was used.\n\n3. In the 2D example, since the authors have used rejection sampling to produce the plots, it is also easy to accurately estimate the mean/covariance of the target distribution. It would have been interesting to use these statistics [although, it is not possible to do so in more complex scenarios] and see if this leads to improved performances.\n\n4. in the \"\\min \\bar{p}\" method, why choose a target acceptance rate of 0.25? My experience says that the number is usually chosen much higher.\n\n5. While reporting the KSD, I think it would have been very interesting to report the ESS [or variations of it], since it is the standard measure of efficiency in the MCMC literature.\n\n6. Finally, while the 2D examples are certainly very interesting, I am not convinced that directly going from 2D to super-difficult-target is the right approach to understand the properties of the proposed methods. There are many settings that are more difficult than these 2D distributions, but much more tractable than the DLG/molecular targets.\n\nIn summary, I think that the authors are proposing an interesting line of research, but more careful numerical investigations are necessary to really understand the worth of the methodology.", "extra_scores": {}, "extra_text": {}}, {"review_id": "5fGzsvKcotc", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nMajor concerns:\n1. The method is an engineering trick rather than a grounded approach to the optimization of sampling algorithms. Indeed, in many cases, people use MCMC methods to obtain guarantees for the sampling procedure. The proposed method removes all these guarantees by relying on the choice of the initial distribution q_0(). Moreover, the optimization of q_0() via mass-covering objectives is a notoriously hard problem since samples from the target distribution are not given in a usual setting.\n\n2. I think the paper lacks an essential comparison with the method proposed by Titsias (Gradient-based Adaptive Markov Chain Monte Carlo, 2019). This paper proposes a more general objective for parameter optimization explicitly fostering high entropy of the proposal. Moreover, in contrast with the learning step of q_0(), it operates in an adaptive manner, not requiring any pretraining steps.\n\n3. Given the limited theoretical novelty, I would expect the ICLR paper to demonstrate highly successful empirical results. However, it is not the case for the current submission. I'm quite confident that the results on CV tasks are out of practical interest. Also, for the molecular dynamics, the metrics' choice hinders the assessment of the practical significance.\n\nMinor comments:\n1. I don't find the comparison of marginal distributions on the 60d problem to be a convincing way to compare samplers' performance. I would suggest considering either another metric or another problem.\n2. I also would suggest to include the description (at least the formula for the density) of the problem \"molecular configurations.\" It would provide the reader with an additional intuition on its difficulty.\n3. I think section 4 would benefit from the clear description of the choice of s, for instance, from the description of the variable mu, which appears there for the first time.\n\nAdditional comments:\nAfter rereading the review, I feel that it may sound a bit harsh for the authors. Therefore, I want to say aloud that I find the paper's subject to be of great interest, consider any work in this direction valuable, and encourage the authors to continue their studies. My criticism is only an attempt to approach the review process objectively.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a tuning strategy for Hamiltonian Monte Carlo (HMC). The proposed algorithm optimizes a modified variational objective over the T step distribution of an HMC chain. The proposed scheme is evaluated experimentally.

All of the reviewers agreed that this is an important problem and that the proposed methods is promising. Unfortunately, reviewers had reservations about the empirical evaluation and the theoretical properties of the scheme. Because the evaluation of the scheme is primarily empirical, I cannot recommend acceptance of the paper in its current form.

I agree with the following specific reviewer concerns. The proposed method does not come with any particular guarantees, and particularly no guarantees regarding the effect of dropping the entropy term and using an SKSD training scheme to compensate. While guarantees are not necessary for publication, the paper should make up for this with comprehensive and convincing experiments. I agree with R1 that more careful ablation studies on toy models are needed, if nothing else to reveal the strengths and weaknesses of the proposed approach. I would also recommend a more careful discussion about the computational cost of this method and how it can be fairly compared to baselines. I don't agree that "deliberately wasteful" experiments reveal much, especially if running more realistic experiments reduces the relative impact of the proposed method.
Decision: Reject