MCMC Variational Inference via Uncorrected Hamiltonian Annealing
Authors: Tomas Geffner
Justin Domke
Venue: NeurIPS 2021 Poster (2021)
Keywords: Variational Inference
Annealed Importance Sampling
MCMC
Hamiltonian Monte Carlo
differentiable

Abstract: Given an unnormalized target distribution we want to obtain approximate samples from it and a tight lower bound on its (log) normalization constant log Z. Annealed Importance Sampling (AIS) with Hamiltonian MCMC is a powerful method that can be used to do this. Its main drawback is that it uses non-differentiable transition kernels, which makes tuning its many parameters hard. We propose a framework to use an AIS-like procedure with Uncorrected Hamiltonian MCMC, called Uncorrected Hamiltonian Annealing. Our method leads to tight and differentiable lower bounds on log Z. We show empirically that our method yields better performances than other competing approaches, and that the ability to tune its parameters using reparameterization gradients may lead to large performance improvements.

MCMC Variational Inference via Uncorrected Hamiltonian Annealing
Tomas Geffner
Justin Domke
College of Information and Computer Science University of Massachusetts, Amherst Amherst, MA tgeffnerccs umass. edu
College of Information and Computer Science University Massachusetts, Amherst Amherst, MA domke@cs. umass edu edu
Abstract
Given an unnormalized target distribution want to obtain approximate samples from it and tight lower bound on its (log) normalization constant log Z. Annealed Importance Sampling (AIS) with Hamiltonian MCMC is powerful method that can be used to do this. Its main drawback is that it uses non-differentiable transition kernels, which makes tuning its many parameters hard. We propose framework to use an AIS-like procedure with Uncorrected Hamiltonian MCMC, called Uncor- rected Hamiltonian Annealing. Our method leads tight and differentiable lower bounds on log Z. We show empirically that our method yields better performances than other competing approaches, and that the ability tune its parameters using reparameterization gradients may lead to large performance improvements.
Introduction
Variational Inference (VI) [4, 41, 45] is method to do approximate inference target distribution p(z) pZz)/Z that is only known up to the normalization constant Z. The basic insights are, first, that the evidence lower bound (ELBO) Egt:) [log p(2) log q[=)] lower-bounds log Z and, second, that maximizing the ELBO is equivalent to minimizing the KL-divergence from & to p. The simplest VI method chooses a parameterized family for and optimizes its parameters to maximize the ELBO.
recent direction involves combining VI with Markov chain Monte Carlo (MCMC) [34, 43]. These methods can be seen as an instance of the auxiliary VI framework [2] they create an augmented variational distribution that represents all intermediate random variables generated during the MCMC procedure. An augmented target distribution that attempts to capture the inverse MCMC dynamics optimized jointly with this variational distribution. However, it has been observed that capturing inverse dynamics is challenging [43, $5.4] (further discussion in Section 4).
Annealed Importance Sampling (AIS) [22, 27] is powerful technique used build augmented distri- butions without the need of learning inverse dynamics. While it was originally proposed to estimate expectations using importance sampling, can be easily used to build lower bounds on normalization constants of intractable densities [18, 44J. AIS creates a sequence of densities that bridge from tractable initial approximation q to the target p. Then, the augmented variational distribution given by a sequence of MCMC kernels targeting each bridging density, while the augmented target uses the reversals of those kernels. turns out that the ratio of these augmented distributions can be computed using only evaluations of the bridging densities. Combining Hamiltonian MCMC kernels with AIS has been observed to produce strong lower bounds [35, 44].
However, these bounds are sensitive to numerous parameters, such as the initial distribution, bridging schedule, and parameters of the MCMC kernels. It would be desirable to optimize these parameters
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

to tighten the bound. Unfortunately, the presence of Metropolis- Hastings acceptance steps means that the the final estimator is con-differentiable, and thus reparameterization gradients cannot be used.
In this work, we propose Uncorrected Hamiltonian Annealing (UHA), a differentiable alternative to Hamiltonian AIS. We define an augmented variational distribution using Hamiltonian MCMC kernels, but dropping the accept-reject steps. This is motivated by the fact that Hamiltonian dynamics sometimes have high acceptance rates. Since these uncorrected MCMC kernels do not exactly hold the bridging densities invariant, an augmented target distribution cannot be defined in terms of reversals, Instead, we define our augmented target by deriving an algorithm for the exact reversal of the original (corrected) MCMC kernel and dropping the accept-reject step. Surprisingly, this yields a very simple expression for the resulting lower bound.
We use reparameterization gradients to tune various parameters involved the lower bound produced by UHA, including the initial approximation 4 parameters of the uncorrected MCMC kernel, and the bridging densities. Experimentally. tuning all these leads to large gains. For example, in several inference tasks we observe that tuning UHA with K 64 bridging densities gives better results than traditional Hamiltonian AIS with K = 512.
Finally, we use UHA train VAES [24, 31J. In this case we observe that using UHA leads to higher ELBOs. addition, we observe that increasing the number of bridging densities with UHA consistently leads to better results, and that for large enough number of bridging densities the variational gap (difference between ELBO and true log-likelihood) becomes small, and models with higher ligllikel obtained.
2 Preliminaries
Variational inference and augmentation. Suppose that p(z) (2((2 is some target density, where is unnormalized and Z f f(z)dz is the corresponding normalizer, and let
log p(2)
(LBO(g())
E
((=)
g(z)
be the "ELBO operator". Variational inference (VI) is based on the fact that for any g(2) we have [4]
log Z = ELBO(g(2), p()) + KL(g(>)ll((z)),
(2)
In VI, the parameters of q are tuned to maximize the "evidence lower bound" (ELBO). Since the KL-divergence non-negative, this is always lower bound on log Z. Also, maximizing the ELBO is equivalent to minimizing the KL--diveggecce from g to p-
To get tighter bounds and better approximations recent work has made use of augmented distributions [2. 21]. Let 21.M = (21,... 2 and suppose that p(21.M) p(2Md (((2)) -1|m)) augments the original target density while preserving its normalization constant. Then, for any ((21:>)) we have
log Z (LLBO(((E.M)). ( ((()))) + (U(g(Zaa))lb(<x..))
The first term is called the "augmentedd" ELBO and again lower bounds log Z. By the chain rule of KL-divergence [12], the KL-divergence from q to p over ž1:Mr upper-bounds the KL-divergence over ZM. This justifies using the marginal of q over 2M to approximate the original target distribution.
Annealed Importance Sampling. A successful approach for creating augmented distributions is Annealed Importance Sampling (AIS) [27]. It creates an augmented proposal distribution by applying a sequence of transition densities Tr (2m+1l2m) and an augmented target by defining transition densities Uml(mm km-1). This gives the augmented densities
M-1
M-1
(( - -(==)
Ti (2m+11lzzn) and p(ZuM) = p(ZNs) II m=l
II
Um (mm)))+1).
(4)
Naively, the ratio of these densities is
p(21:1)) p(22r) (21:,))
M-1
U (Zm 2m+z)
(5)
((11)
Tm 2m Zm
m=1

---

To define the transitions Tm and Um. AIS creates sequence of unnormalized densities T ....... AM that "bridge" from starting distribution q to the target p. meaning that म is close to q and - is close to p. Then, for each intermediate distribution, (= [2m) is chosen to be a Markov kernel that holds Tye invariant, and Un to the reversal of Tm with respect to Ty, defined as
Tm (2m)
Um (2ml #m+ T(2m+1/2m)
(6)
Tm Tm (2m+1)
This choice produces a simplification so that eq. 5 becomes
p(ZI.M) ((LMM) = g(x1)
p(2M)
M-1
T' (2m)
II m=1
(7)
Tm (2m+1))
This can be easily evaluated without needing to evaluate the transition densities. The ratio from eq. 7 can be used to get expression for the lower bound ELBOO (()1:)))). Research has shown that the AIS augmentation may lead to extremely tight lower bounds [18, 17, 44].
Hamiltonian Dynamics. Many MCMC methods used to sample from p(z) are based on Hamiltonian dynamics [3. 8, 29, 42]. The idea is to create an augmented distribution p(z, p) = p(=)S(p). where S()) is a distribution over a momentum variable p (e.g. a Multivariate Gaussian). Then, one can define numerical integration schemes where 2 and p evolve while nearly holding p(z. p) constant. When corrected by Metropolis- Hastings acceptance step, this can be made exactly hold hozdp(2, p) invariant. This is alternated with a scheme that resamples the momentum p while holding Sp) invariant. When Hamiltonian dynamics work well, z can quickly move around, suppressing random-wall behavior.
There are a variety of different Hamiltonian MCMC methods, corresponding to different integration schemes, momentum distributions, and ways of resampling the momentum. For instance, HMC and Langevin dynamics use the leapfrog integrator, a Gaussian for the momentum variables and full resampling of the momentum variables at each step [29, 42]. On the other hand, if the momentum variables are only partially resampled, the under-damped variants of HMC and Langevin dynamics are recovered [29]. It was observed that partial resampling may lead to improved perfomance [9].
is easy to integrate Hamiltonian dynamics into AIS. First, define an augmented target p(z,p)
(p))(pp) and an augmented starting distribution q(z, p) (()) Then, create a series of augmented densities म1 (1 p). ..... I2, p) bridging the two as 7m (z. p) Am (p)))p)). Finally, define the forward transition Tm. (2m+1, Pm-11 f 2mm Pm.) to be an iteration of Hamiltonian MCMC method that leaves TT. (2, p) invariant. We will describe a single transition Tm as a sequence of three steps: (I) resample the momentum; (2) simulate Hamiltonian dynamics and apply an accept-reject step; and (3) negate the momentum. The precise process that defines the transition is shown in Alg. 1. Note that this algorithm is quite general, and compatible with HMC, Langevin dynamics and their underdamped variants (by selecting an appropriate integrator and resampling method).
Algorithm 1 Corrected Tm (2m+1 1 Pmo+ (.mm Pm)
1. Sample p'm, from some s sp'm IPm) that leaves S(p) invariant. Set z'm f- Am.
2. Simulate Hamiltonian dynamics as (2m. Pin) + Tm Fm((ie (z'm.P'))
Calculate an acceptance probability a min (1, Rmm (tm
(2"h Pmm)) //////m/m p'm)).
With probability a, set (= PT))) + (2ils p%)). Otherwise, set (2 pmn) t (ohms Plm)).
3. Reverse the momentum (2m41, Pmn+1) (=m pm)) return (2m+1.pm+1)
Representing Tm this way makes it easy to show it holds the density n. (2, p) invariant. The overall strategy is to show that each of the steps 1-3 holds 7m invariant, and so does the composition of them [29, $3.2]. For steps I and 3 this is trivial, provided that S(p) = 3(-p). For step 2, we require that the simulation Tm has unit Jacobian and satisfies T--1 = Ti. Then, T. can be interpreted as a symmetric Metropolis- Heatopolis-Hastingss proposal, meaning the Metroplis- Hastings acceptance probability is as given. A typical choice for Tm that satisfies these requirements is the leapfrog integrator with a momentum reversal at the end. (This reversal then gets "un-reversed" in step 3 for accepted moves.)
Since T., holds n. invariant, we can define Um as the reversal of Tin wrt Ae, Then, eq. 7 becomes
M
p (21..M P1:a)) q( (1:M; PI:M)
p(eMs PM))
Fm (Zmn Pm)
(8)
((1..P))
Fm (2m=+1 Pm++1
3

---

Using this ratio we get an expression for the lower bound ELBO(q( 21:MM P1:M) ) (I.M)) obtained with Hamiltonian AIS. While this method has been observed to yield strong lower bounds on log Z [35, 44] (see also Section 5.2), its performance depends on many parameters: initial distribution q(2), momentum distribution S, momentum resampling scheme, simulator Tm and bridging densities. We would like to tune these parameters by maximizing the ELBO using reparameterization based estimators. However, due to the accept-reject step required by the Hamiltonian MCMC transition, the resulting bound is not differentiable, and thus reparameterization gradients are not available.
3 Uncorrected Hamiltonian Annealing
The contribution of this paper is the development of uncorrected Hamiltonian Annealing (UHA). This method is similar to Hamiltonian AIS (eq. 8), but yields differentiable lower bound. The main idea simple. For any transitions T. and Un. by the same logic as in eq. 5, we can define the ratio
P (EL.M.PIMM. p(ZM, PM) (1),
M-1
Um (Zm: Pm #m++1 pm++1)
II
g PLEM
Tm (Zm+1s Pm++ Imm Pm))
1==1
Hamiltonian AIS defines Tm as a Hamiltonian MCMC kernel that holds 7. invariant, and Un as the reversal of T. with respect to Tm- While this leads to nice simplification, there is no requirement that these choices made. We can use any transitions as long as the ratio U is tractable.
We propose to use the "uncorrected" versions of the transitions Tm and Un used by Hamiltonian AIS, obtained by dropping the accept-reject steps. To get an expression for the uncorrected Un we first derive the reversal Un used by Hamiltonian AIS (Alg. 2). These uncorrected transitions no longer reversible with respect the bridging densities (2. and thus we cannot use the simplification used by AIS to get eq. 8. Despite this. we show that the ratio /T for the uncorrected transitions can still be easily computed (Thm. 2). This produces differentiable estimator, meaning the parameters can be tuned by stochastic gradient methods designed to maximize the ELBO.
We start by deriving the process that defines the transition U,. used by Hamiltonian AIS. This is shown in Alg. 2. It can be observed that Um follows the same three steps of Tm (resample momentum, Hamiltonian simulation with accept-reject, momentum negation), but in reverse order.
Algorithm 2 Corrected Um (Zin Pm [mm 2mm 1 Pm=+1)
1. Set (am, P!") + (2m+1, -)
2. Simulate Hamiltonian dynamics as Pmm) + Tm (2m:Pm))
Calculate acceptance probability a min (1, Rm Pm)) / am (2m pm)).
With probability a, set (z' pmm) + (2 plm). Otherwise, set (oims Pmm) + (2m Pm)).
3. Sample Pm from Srev (Pm (m)). the reversal of sp'm (mm) with respect to S(mm) Set Zm t z/m- return (2m+Pm)
Lemma 1. The corrected Um (Alg. 2) is the reversal of the corrected Tm (Alg. 1) with respect to Tim-
Sketch). First, we claim the general result that Tı, T2 and T3 have reversals U1, U2 and U3. respectively, then the composition T Ti o T2 T3 has reversal U = U3 o U2 U1 (all reversals with respect to same density т). Then, we apply this to the corrected Im and Um: Tm. is the composition of three steps that hold T. invariant. Thus, its reversal Um is given by the composition of the reversals of those steps, applied in reversed order. A full proof is in Appendix F.
We now define the "uncorrected" transitions used by UHA. shown in Algs. 3 and 4. These just the transitions used by Hamiltonian AIS but without the accept-reject steps. (If Hamiltonian dynamics simulated exactly, the acceptance rate is one and the uncorrected and corrected transitions are equivalent.) We emphasize that, for the "uncorrected" transitions, Tm does not exactly hold m. invariant and Um is not the reversal of Tm Thus, their ratio does not give simple expression in terms of 7. as in eq. 8. Nevertheless, the following result shows that their ratio has a simple form.
Theorem 2. Let Tm and Um be the uncorrected transitions defined in Algs. 3 and 4, and let the dynamics simulator Tm p) be volume preserving and self inverting. Then,
Uml Peol 2m pm++1)
S(pm)
(10)
T'm (2m+11 Pm+t1 Rmrs Pmn)
S((m)))

---

Algorithm 3 Uncorrected Tm (Zm++1 +1 F.mm pm)
1. Sample P', from some s(m pm)) that leaves S() invariant. Set s/m t- Zm.
2. Simulate Hamiltonian dynamics as (2% pmm) Tm Bm))
3. Reverse the momentum as (2m+1, Por+1) ) + (#mm - pmm). return 1m++1/mmm)
Algorithm 4 Uncorrected Um (Zm Pm 2m++ pm+1)
1. Set (zm Pm)) (2m+11, -pm+1))
2. Simulate Hamiltonian dynamics as (z' Pmm) + Tm (2 P'a).
3. Sample Pm from Srev (pm Pm), the reversal of s( mm Pm)) with respect to S(pm). Set 2m 2ym return (2ms Pm)
where P'm is the second component of (1m+1:1111 Pmm+1). (That is, from Algs. 3 and 4.)
(ketch)) We consider variants of Algs. 3 and in which each time z is assigned we add Gaussian noise with some variance a/. We then derive the densities for Tm and Um using the rule for transformation of densities under invertible mappings, using that Tm is self-inverting and volume preserving. Taking the ratio gives eq. 10. Since this is true for arbitrary the limit as + A full proof is in Appendix G.
take the stated result as
As an immediately corollary of eq. 9 and Theorem we get that for UHA
((:1:M)P.MM) p(ZM) g (21.M, P1:M)
M-1
5(pm+1)
II
(11)
g(=1)
S(p'))
This ratio can be used get an expression for the lower bound ELLOO PI:M), p(Z Pla))) obtained with UHA. As mentioned in Section 2, the parameters of the augmented distributions are tuned to maximize the ELBO, equivalent to minimizing the KL-divergence from q to p. While computing this ELBO exactly is typically intractable, an unbiased estimate can be obtained using sample from ( P1:M) as shown in Alg. If sampling is done using reparameterization, then unbiased reparameterization gradients may be used together with stochastic optimization algorithms to optimize the lower bound. In contrast, the variational lower bound obtained with Hamiltonian AIS (see Alg. 6 in Appendix A) does not allow the computation of unbiased reparameterization gradients.
Algorithm 5 Generating the (differentiable) uncorrected Hamiltonian annealing variational bound.
Sample ~ g and P1 S. Initialize estimator as +- log ((21). for m 1,2, .M- 1 do
Run uncorrected Tm (Alg. 3) on input (Zm: Pm), storing P'm and the output (zm+1: Pm+1) 1).
Update estimator as L f- L + log ((pm+1))S((pp))
Update estimator as C t log p(zM).
return R
3.1 Algorithm Details
Simulation of dynamics. We use the leapfrog operator with step-size simulate Hamiltonian dynamics. This has unit Jacobian and satisfies T. = T (if the momentum is negated after the simulation), which are the properties required for eq. 11 be correct (see Theorem 2).
Momentum distribution and resampling. We set the momentum distribution S(p) = A (N(p0, )) to be a Gaussian with mean zero and covariance Σ. The resampling distribution s(p' ()) must hold this distribution invariant. As is common we use s(p'lp) = ( (1 - ))²) where E [0. 1) is the damping coefficient. If 17 = 0, the momentum is completely replaced with a new sample from S in each iteration (used by HMC and Langevin dynamics [29, 42]). For larger 7,, the momentum becomes correlated between iterations, which may help suppress random walk behavior and encourage faster mixing [9] (used by the underdamped variants of HMC and Langevin dynamics [29]).
Bridging densities. We set कm (2,
(2,p) p) = q(2, p)² 1- p(z, pm where 8m E [0, 1] and Bm 8mt1
5

---

Computing gradients. We set the initial distribution g(<1) to be a Gaussian, and perform all sampling operations in Alg. using reparameterization [24, 31, 39]. Thus, the whole procedure is differentiable and reparameterization-oosss gradients may be used to tune parameters by maximizing the EL.BO. These parameters include the initial distribution g(21), the covariance Σ of the momentum distribution, the step-size e of the integrator, the damping coefficient 7 of the momentum resampling distribution, and the parameters the bridging densities (including 3), among others. As observed in Section 5.2.1 tuning all of these parameters may lead to considerable performance improvements.
4 Related Work
UHA and slight variations have been proposed in concurrent work by Thin et al. [38], who use uncorrected Langevin dynamics together with the uncorrected reversal to build variational lower bounds, and by Zhang et al. [46], who proposed UHA with under-damped Langevin dynamics together with a convergence analysis for linear regression models.
There are three other lines of work that produce differentiable variational bounds integrating Monte Carlo methods. One is Hamiltonian VI (HVT) [34, 43J. It uses eq. 9 to build lower bound on log Z, with Tm set to an uncorrected Hamiltonian transition (like UHA but without bridging densities) and Um set conditional Gaussians parameterized by learnable functions. Typically, single transition is used, and the parameters of the transitions are learned by maximizing the resulting ELBO.'
A second method is given by Hamiltonian VAE (HVAE) [7], based on Hamiltonian Importance sampling [28]. They augment the variational distribution with momentum variables, and use the leapfrog integrator to simulate Hamiltonian dynamics (a deterministic invertible transformation with unit Jacobian) with a tempering scheme as a target-informed flow [30, 37J.
The third method is Importance Weighting (IW) [6, 13, 15]. Here, the idea is that ELBO(q(2), p(2)) < Eog R Ex ((2)) and that the latter bound can be optimized, rather than the traditional ELBO. More generally, other Monte-Carlo estimators can be used [16].
Some work defines novel contrastive- weretiiiveddvveeeeee-eeee like objectives in terms of the final iteration of an MCMC chain [32, 26]. These do not provide an ELBO-like variational bound. While in some cases the initial distribution can be optimized to minimize the objective [32], gradients do not flow through the MCMC chains, meaning MCMC parameters cannot be optimized by gradient methods.
For latent variable models, Hoffman [19] suggested to run few MCMC steps after sampling from the variational distribution before computing gradients with respect to the model parameters, which is expected to "debias" the gradient estimator to be closer the true likelihood gradient. The variational distribution is simultaneously learned to optimize a standard ELBO. (AIS can also be used [14].)
Experiments and Results
This section presents results using UHA for Bayesian inference problems on several models of varying dimensionality and for VAE training. We compare against Hamiltonian AIS, IW, HVI and HVAE, We report the performance of each method for different values of K, the number of likelihood evaluations required to build the lower bound (e.g. number of samples used for IW, number of bridging densities plus one for UHA). Note that, for a fixed K, all methods have the same oracle complexity (i.e. number of target/target's ' gradient evaluation), and that for K = 1 they all reduce to plain VI.
For UHA and Hamiltonian AIS we use under-damped Langevin dynamics, that is, we perform just one leapfrog step per transition and partially resample momentum. We implement all algorithms using Jax [5].
'The formulation of HVI allows the use of more than one transition. However, this leads an increased number of reverse models that must be leamned, and thus not typically used in practice. Indeed, experiments by Salimans et al. [34] use only one HMC step while varying the number of leapfrog integration steps, and results from Wolf et al. [43] show that increasing the number of transitions may actually yield worse bounds (they conjecture that this is due to the difficulty of learning inverse dynamics.).

---

Table I: Our method (UHA) yields better bounds than importance weighting (IW) for mod- erate or high dimensions. ELBO achieved by different methods when using a Student-t target distribution of varying dimensionality, higher is better. Since the target is normalized, perfect inference algorithm would achieve the true value of log Z
Target
Plain VI
UHA
IW
dimension
K =
K
4
K
=6
K
128
K 128
K 1024
20
-0.82
-0.55
-0.36
-0.19
-0.14
-0.14
-0.088
200
-8.1
-5.5
-3.5
-1.9
-1.4
-3.7
-2.9
500
-20.5
-339
-9.0
-5.2
-3.8
-12.0
-10.4
5.1 Toy example
This section compares results obtained with UHA and IW when the target is set to factorized Student- t with mean zero, scale one, and three degrees of freedom. We tested three different dimensionalities: 20, 200 and 500. all cases we have log Z = O, so we can exactly analyze the tightness of the bounds obtained by the methods. We set the initial approximation be a mean-field Gaussian, and optimize the objective using Adam [23] with a step-size of 0.001 for 5000 steps. For UHA we tune the initial approximation g(z), the integrator's step-size and the damping coefficient 77.
We ran UHA for E [4, 16, 64, 128) and IW for K E {128, 1024). Table 1 shows the results for the three dimensionalities considered. It can be observed that UHA performs significantly better than IW as the dimensionality increases; for the target with dimension 500, UHA with K = 16 yields better bounds than IW with K 1024. On the other hand, the methods perform similarly for the low dimensional target. Finally, in this case both methods have similar time costs. For instance, for K 128 UHA takes 14.2 seconds optimize and IW takes 13.9.
5.2 Inference tasks
This section shows results using UHA for Bayesian inference tasks. For this set of experiments, for UHA we tune the initial distribution (=),, ), the integrator's step-size and the damping coefficient ry. We include detailed results tuning more parameters in Section 5.2.1.
Models. We consider four models: Brownian motion (d 32), which models Brownian Motion process with a Gaussian observation model; Convection Lorenz bridge (d = 90), which models nonlinear dynamical system for atmospheric convection; and Logistic regression with the ala (d = 120) and madelon (d = 500) datasets. The first two obtained from the "Inference gym" [36].
Baselines. We compare UHA against IW, HVAE, simple variant of HVI, and Hamiltonian AIS (HAIS). For all methods which rely on HMC (i.c. all except IW) we use a singe integration step-size common to all dimensions and fix the momentum distribution we learn the initial distribution (=2,, integration step-size c and the reverse dynamics Um, (set to a factorized Gaussian with mean and variance given by affine functions), and for HVAE we learn g(2), e and the tempering scheme (we use the quadratic scheme parameterized by single parameter).
standard Gaussian. For HVI
Training details. We set g(=) to be mean-field Gaussian initialized to maximizer of the ELBO, and tune the parameters of each method by running Adam for 5000 steps. We repeat all simulations for different step-sizes in (10-3 10-1 10-5), and select the best one for each method. Since Hamiltonian AIS' parameters cannot be tuned by gradient descent, we find a good pair (E) by grid search. We consider 7) {0.5, 0.9, 0.99) and three values of e that correspond to three different rejection rates: 0.05, 0.25 and 0.5. We tested all 9 possible combinations and selected the best one.
Results are shown in Fig. 1. Our method yields better lower bounds than all other competing approaches for all models considered, and that increasing the number of bridging densities consistently leads to better results. The next best performing method is Hamiltonian AIS. IW also shows good performance for the lower dimensional model Brownian motion. However, for models of higher dimensionality IW leads to bounds that are several nats worse than the ones achieved by UHA. Finally, HVI and HVÁE yield bounds that are much worse than those achieved by the other three methods, and do not appear to improve consistently for larger K. For HVAE, these results consistent with
7

---

the ones in the original paper [7, $4], in that higher K may sometimes hurt performance. For HVI, we believe this is related to the use of just one HMC step and suboptimal inverse dynamics.
Optimization times for Plain VI, IW and UHA (the latter two with K = =32) are 2.4,3.. 2.4 and 4.4 seconds for the Brownian motion dataset, 2.5, 6.8 and 6.9 seconds for Lorenz convection, 2.8, 8.3 and 19.9 seconds for Logistic regression (AIA), and 4.6, 16.6 and 121.2 seconds for Lo- gistic regression (Madelon). While IW and UHA have the same oracle complexity for the same K, we see that the difference between their time cost depends on the specific model under consideration. All other methods that use HMC have essentially the same time cost as UHA.
Brownian motion
Lorenz convection
UHA (ours)
HAIS
-180
BW
ELBO
HVAF
HMVI
-1185
-1175
1180
Logistic regression (a1a)
Logistic regression (madelon)
-680
-2410
-625
2422
-700
-2430
-706
-710
-2440
-716
-245
-720
-2400
10'
10
10
16'
10'
10'
K
K
Figure I: Our method achieves much better bounds than other competing methods. K repre- sents the number of likelihood evaluations to build the lower bound required by each method, The leftmost point of all lines coincide because, for K = 1, all methods reduce to plain VI. Vertical bars indicate one standard deviation obtained by running simulations with four different random seeds.
5.2.1 Tuning More Parameters with UHA
A basic version of UHA involves fitting variational distribution using plain VI, and then tuning the integration step-size e and the damping coefficient ry. However, more parameters could be tuned:
Moment distribution cov Σ: We propose to learn a diagonal matrix instead of using the identity. M.. We propose to learn the sequence B, with the restrictions Bn- = 0, BM = 1, 3m < Bm++1 and Bi, E [o, 1).
Bridging densities' coefficients 3m Typically 8m m/M.
Initial distribution (z): Instead of fixing g(2) to be maximizer of the typical ELBO, we propose to learn to maximize the augmented ELBO obtained using UHA.
Integrator's step-size e: Instead of learning a unique step-size e, we propose learn step-size that is a function of B, i.e. e(3). In our experiments we use an affine function.
Bridging densities parameters uh: Instead of setting the m-th bridging density to be ye --13m we 1-8m påm,, where (p()...) is mean-field Gaussian with a mean and diagonal propose to set it to (p(Am)
covariance specified as affine functions of B.
We consider the four models described previously and compare three methods: UHA tuning all parameters described above, UHA tuning only the pair (e, 17), and Hamiltonian AIS with parameters (e. 1)) obtained by grid-search. We perform the comparison for K ranging from 2 to 512. (For K 2 64 we tune the UHA's parameters using K- 64 and extrapolate them as explained in Appendix D.)
Results are shown in Fig. 2. It can be observed that tuning all parameters with UHA leads to significantly better lower bounds than those obtained by Hamiltonian AIS (or UHA tuning only and n)). Indeed, for the Logistic regression models, UHA tuning all parameters for K = 64 leads to results comparable to the ones obtained by Hamiltonian AIS with K = 512.
To verify what parameters lead to larger performance improvements, we tested UHA with K = 64 tuning different subsets of (,, E, 8, q(=), e(3), (B)). Fig. 3 shows the results. It can be observed
8

---

Browrian motion
Lorenz convection
Log reg (a18)
Log (madelon)
6.3
-1150
-1154
-1820
-105
-8900
2992
0.2
64
6-2
0.2
UHA
UHA
-1154
-1157
-1503
-152
-1150
-2323
-552
-590
402
14115
-2150
-2100
2420
0.4
-392
-2172
500
410
-24800
K
K
Figure 2: UHA tuning all parameters leads to better performance than other methods.
that tuning the bridging parameters /3 and the initial approximation g(=) leads to the largest gains in performance, and that tuning all parameters always outperforms tuning smaller subsets of parameters. We show a more thorough analysis, including more subsets and values of K in Appendix B.
Brownian motion
Lorenz convection
Log reg (a1a)
Log reg (madelon)
02
1185.2 29
2039
430
2410
1155..0 79
1180-00 00
-158 30
-14
2614
11511:30
1 4
5
3 u.
Al
u
8
8
8
u
Figure 3: Tuning all parameters leads to better results than tuning subsets of them. Largest gains obtained by tuning bridging coefficients 3 and initial distribution q. ELBO achieved as function of parameters tuned (x-axis), for K == 64. The subsets are ordered in terms of increasing performance (same ordering is used for all four models). Parameters are step-size c, damping coefficient f). moment covariance Σ, bridging densities parameters 3 and vt initial distribution q.
Finally, Appendix E shows results comparing UHA (tuning several parameters) against HMC, mean field VI and IW in terms of the approximation accuracy achieved on logistic regression model with a fixed computational budget.
5.3 VAE training
Our method can be used to train latent variable models, such as Variational [24, 31J. In this case the initial approximation g(=[xx) and the model p(x, z) are parameterized by two neural networks (encoder and decoder), whose parameters are trained by maximizing the ELBO. UHA can be used to train VAES by augmenting these two distributions as described in Section 3.
(VAE)
Datasets.
We use three datasets: mnist [25] (numbers 1-9), emnist-letters [11] (letters A-Z), and kmnist [10] (cursive Kuzushiji). All consist on greyscale images of 28 28 pixels. In all cases we use stochastic binarization [33] and a training set of 50000 samples, a validation set of 10000 samples, [1].
and a test set of 10000 samples. All datasets are available in tensorflow-
Baselines. We compare against Importance Weighted Auto-encoders [6] and plain VAE training [24].
Architecture details. We set ((=|æz) to a diagonal Gaussian, p(z) to a standard Normal, and p(x|=)
to a Bernoulli. We consider two architectures for the encoder and decoder: (1) Feed forward networks with one hidden layer of size 450 and Relu non-linearities, with a latent space dimensionality of 64;
(2) Architecture used by Burda et al. [6], feed forward networks with two hidden layers of size 200 with tanh non-lincarssies, with latent space dimensionality of 50.
Training details. In all cases the encoder and decoder are initialized parameters that maximize the ELBO. For IW we tune the encoder and decoder parameters (using the rearyyreparameetrrizee estimator [40]), and for UHA we tune the integration step-size c, damping coefficient 1), bridging parameters 8, momentum covariance Σ (diagonal), and the decoder parameters. Following Caterini et al. [7] we constrain e E (0, 0.05) to avoid unstable behavior of the leapfrog discretization. We use Adam with step-size of 10-1 to train for 100 epochs and use the validation set for early stopping. We repeated all simulations for three different random seeds. In all cases the standard deviation of the results was less than 0.1 nats (not shown in tables).

---

All methods achieved better results using the architecture with one hidden layer. These results are shown in Tables and 3. The first one shows the ELBO on the test set achieved for different values of K, and the second one the log-likelihood on the test set estimated with AIS [44]. It can be observed that UHA leads to higher ELBOOs, higher log-likelihoods, and smaller variational gaps (difference between ELBO and log-likelihoode than IW for all datasets, with the difference between both methods' performance increasing for increasing K. Notably, for = 64, the variational gap for UHA becomes quite small, ranging from 0.8 to 1.4 nats depending on the dataset.
Results for the architecture from Burda et al. [6] (two hidden layers) are shown in Tables 4 and 5 (Appendix C). Again, we observe that UHA consistently leads to higher ELBOs and the best test ligeliied was consistently achieved by UHA with K 64. However, for smaller K, IW sometimes had better log-likelinge likelihoods than UHA (despite worse ELBOs).
Table 2: ELBO on the test set (higher is better). For K both methods reduce plain VI.
==1 ==8
K 16
K 32
K 64
UHA
993.4
89.8
-88.8
88.1
-77.6
mnist
IW
99.44
90.55
-9.99
89.4
-9.0
UHA
1337.9
133.5
-13.33
131.5
130.9
letters
IW
-137.9
133.66
-133.9
133.2
-132.7
UHA
-18422
176.6
-174.6
173.2
-171.6
kmnist
IW
-184.2
179.77
-178.7
117.8
-177.0
Table 3: Log-likelihood on the test set (higher is better). This is estimated using AIS with under- damped HMC using 2000 bridging densities, HMC iteration with 16 leapfrog steps per bridging density, integration step-size 0.06, and damping coefficient 0.8.
K =1
K= 8
K = 16
K = 32
K = 64
UHA
-88.5
-87.5
-87.2
-87.0
-86.9
mnist
IW
-88.5
-7..6
-8755
-7.33
-87.2
UHA
11319
130.7
-130.3
130.1
1129.9
letters
IW
-131.9
130.99
-130.7
133.66
-130.4
UHA
-14.33
172.2
-1716
171.2
-170.2
kmnist
IW
-174.3
173.00
-172.6
172.44
-172.2
Discussion
Since UHA yields differentiable lower bound, one could tune other parameters not considered in this work. For instance, different momentum distribution per bridging density could be used, that is, (2, p) .. S=) (p). We believe additions such as this may yield further gains. Also, method can be used to get tight and differentiable upper bounds on log Z using the reversed AIS procedure described by Grosse et al. [18].
Finally, removing accept-reject steps might sometimes lead to instabilities during optimization the step-size e becomes large. We observed this effect when training VAEs on some datasets for the larger values of K. We solved this by constraining the range of (previously done by Caterini et al. [7]). While this simple solution works well, we believe that other approaches (e.g. regularization, automatic adaptation) could work even better. We leave the study of such alternatives for future work.
Acknowledgments and Disclosure of Funding
This material is based upon work supported part by the National Science Foundation under Grant No. 1908577.
10

---

References
[1] TensorFlow Datasets, a collection ready-to-use datasets. htttps:/www.teesoortt org/
datasets.
[2] Felix V Agakov and David Barber. An auxiliary variational method. In International Conference on Neural Information Processing, pages 561-566. Springer, 2004.
[3] Michael Betancourt, Simon Byrne, Sam Livingstone, Mark Girolami, et al. The geometric foundations of hamiltonian monte carlo. Bernoulli, 23(4A):2257-2288 2017.
[4] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: review for statisticians. Journal of the American Statistical Association, 1 112(518)) (2(5118)859-87, 2017.
[5] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- -Minne, and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018.
[6] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In Proceedings of the International Conference on Learning Representations, 2016.
[7] Anthony L Caterini, Arnad Doucet, and Dino Sejdinovic. Hamiltonian variational auto encoder.. In Advances in Neural Information Processing Systems, 2018
[8] Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference machine learning, pages 1683-1691, PMLR, 2014.
[9] Xiang Cheng, Niladri S Chatterji, Peter L Bartlett, and Michael I Jordan. Underdamped langevin meme: A non-asymptotic analysis. In Conference on Learning Theory, pages 300-323. PMLR, 2018.
[10] Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature. arXiv preprint arXiv:1812.01718, 2018.
[11] Gregory Cohen, Saced Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. In 2017 International Joint Conference on Neural Networks (IJCNN, pages 2921-2926. IEEE, 2017.
[12] Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
[13] Chris Cremer, Quaid Morris, and David Duvenaud. Reinterpreting importance-weighted autoencoders. arXiv preprint arXiv:1704.02900 2017.
[14] Xinqiang Ding and David J Freedman. Learning deep generative models with annealed impor- tance sampling, arXiv preprint arXiv:1906.04990 2019.
[15] Justin Domke and Daniel Sheldon. Importance weighting and variational inference. In Advances Neural Information Processing Systems, 2018.
[16] Justin Domke and Daniel Sheldon. Divide and couple: Using monte carlo variational objectives for posterior approximation. In Advances in Neural Information Processing Systems, 2019.
[17] Roger B Grosse, Siddharth Ancha, and Daniel M Roy. Measuring the reliability of meme inference with bidirectional monte carlo. Advances in Neural Information Processing Systems, 2016.
[18] Roger B Grosse, Zoubin Ghahramani, and Ryan Adams. Sandwiching the marginal likelihood using bidirectional monte carlo. arXiv preprint arxiv:1551.0220 2015.
[19] Matthew D Hoffman. Learning deep latent gaussian models with markov chain monte carlo. In International conference on machine learning, pages 1510-1519. PMLR, 2017.
[20] Matthew D Hoffman, Andrew Gelman, et The no-u-turn sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15(1):1593-1623 2014.
11

---

[21] Chin-War Wei Huang, Shawn Tan, Alexandre Lacoste, and Aaron Courville. Improving explorability in variational inference with annealed variational objectives. In Advances in Neural Information Processing Systems, 2018.
[22] Christopher Jarzynski. Equilibrium free-energy differences from nonequilibrium measurements: master-equation approach. Physical Review E, 56(5): 1997.
[23] Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. In Proceedings of the International Conference Learning Representations, 2015.
[24] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In Proceedings of the International Conference on Learning Representations, 2013.
[25] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-basee learning applied to document recognition. Proceedings of the IEEE, (6(11))2278-233 1998.
[26] Yingzhen Li, Richard E Tumnrr and Qiang Liu. Approximate inference with amortised meme. arXiv preprint arXiv: (eo::030223333 2017.
[27] Radford M Neal. Annealed importance sampling. Statistics and computing, 11(2):125-139, 2001.
[28] Radford M Neal. Hamiltonian importance sampling. In talk presented at the Banff International Research Station (BIRS) workshop on Mathematical Issues in Molecular Dynamics, 2005.
[29] Radford M Neal et al. Memc using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
[30] Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning (ICML-15), 2015.
[31] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference deep generative models. In Proceedings of the 3 International Conference on Machine Learning (ICML-14), pages 1278 1286, 2014.
[32] Francisco Ruiz and Michalis Titsias. A contrastive divergence for combining variational inference and meme. In International Conference on Machine Learning, pages 5537-5545. PMLR, 2019.
[33] Ruslan Salakhutdinov and lain Murray. On the quantitative analysis of deep belief networks. In Proceedings of the 25th international conference on Machine learning. pages 1772-79, 2008.
[34] Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International Conference Machine Learning, pages 1218- 1226, 2015.
[35] Jascha Mohl-Dicksteein and Benjamin J Culpepper. Hamiltonian annealed importance sampling for partition function estimation. arXiv preprint arXiv: /20:. 2012.
[36] Pavel Sountsov, Alexey Radul, and contributors. Inference gym, 2020.
[37] Esteban G Tabak and Cristina V Turner. family of nonparametric density estimation algo- rithms. Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
[38] Achille Thin, Nikita Kotelevskii, Arnaud Doucet, Alain Durmus, Eric Moulines, and Maxim Panov. Monte carlo variational auto-encoders. In International Conference on Machine Learning, pages 10247-10257. PMLR, 2021.
[39] Michalis Titsias and Miguel GireillGr Doubly stochastic variational bayes for non- conjugate inference. Proceedings of the 31st International Conference on Machine Learning (ICML-14), pages 1971-1979, 2014.
[40] George Tucker, Dieterich Lawson, Shixiang Gu, and Chris J Maddison. Doubly reparameterized gradient estimators for monte carlo objectives. In Proceedings of the International Conference on Learning Representations, 2019.
12

---

[41] Martin J Wainwright, Michael Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning. 1(1-2):1-305, 2008.
[42] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-II), pages 681-688. Citeseer, 2011.
[43] Christopher Wolf, Maximilian Karl, and Patrick van der Smagt. Variational inference with hamiltonian monte carlo. arXiv preprint (eXiv:/609.006 2016.
[44] Yuhuai Wu, Yuri Burda, Ruslan Salakhutdinov, and Roger Grosse. On the quantitative analysis of decoder-based generative models. In Proceedings of the International Conference on Learning Representations, 2017.
[45] Cheng Zhang. Judith Butepage, Hedvig Kjellstrom, and Stephan Mandt. Advances in variational inference. arXiv preprint arXiv:1711.05533 2017.
[46] Guodong Zhang, Kyle Hsu, Jianing Li, Chelsea Finn, and Roger Grosse. Differentiable annealed importance sampling and the perils of gradient noise. arXiv preprint arXiv:2107.10211 2021.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [No]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No] Mentioning them here. Jax [5]: Apache
License 2.0. Inference gym [36]: Apache License 2.0. Tensorflow-datasers [1]: Apache
License 2.0.
(c) Did you include any new assets either in the supplemental material URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you' 're using/curating? [No]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "cUVceKz45Ml", "reviewer": "Reviewer_uWL9", "summary": "The paper propose to use uncorrected HMC kernels (i.e. HMC kernels without accept-reject steps) in AIS to make the estimate of log normalisation constant differentiable w.r.t. to various parameters (e.g. step-size for numerical integration).\nThe author(s) name(s) the method UHA.\nThe promise is that one can then using reparameterisation gradient to tune these parameters.\nThrough experiments on inference tasks,\nthe benefit of tuning parameters in UHA is demonstrated against other methods that use HMC in AIS.\nFor experiments on VAE training,\nUHA with tuned parameters has shown better test ELBO than IW.", "questions": "", "limitations": "", "rating": 8, "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": "Originality: \nUHA is a novel method derived by a neat modification to the common way of using HMC in AIS.\nIt solves an important limitation of previous methods: they are sensitive to hyper-parameters but there is no easy way to tune them.\nDifferences against previous methods are made clear through the motivation as well as the dedicated related work section.\n\nQuality: \nThe method is technically sound and the benefit of being able to tuning all parameters is clearly demonstrated by experiments in section 5.1 and 5.2.\n\nClarity: \nThe paper is very well-written and easy to follow.\nTechnical details are also enough for reproducing the results.\n\nSignificance: \nThe fact that the proposed method leads to a differentiable estimate is a solid contribution to the community.\nOne of the biggest challenges to use HMC in AIS is about hyper-parameter tuning and UHA and I can see how UHA would popularise the use of AIS as now the hyper-parameter tuning is made easy.\n\nMisc:\n- An interesting baseline that is not included is tuning parameters using the REINFORCE estimator (either in UHA or HAIS), although I could expect it would be worse than the reparameterisation estimator.\n- The grid search for HAIS is insufficient. One interesting thing to check from a more fine-grained search of HAIS is that we could see how much we lose by getting rid of the accept-reject step. Another interesting thing to check is how different the optimal parameters are for HAIS and UHA.\n- For the last sentence from L291, how do we know the true $\\log Z$ to compute the variational gap?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are discussed in section 6.\nThere is no societal impact discussed but I don't think it's needed for this technical paper.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "ae1pS2qBKJe", "reviewer": "Reviewer_BV2F", "summary": "This paper studies the annealed important sampling (AIS) using the Hamiltonian dynamics kernels in each interpolant. To circumvent the problem that the Hamiltonian dynamics involves the non-differentiable operations, this paper proposes to discard the accept-reject operations in the kernel, and derives the resulting density ratios in a simple form related to the momentum variables. The resulting approach, named Uncorrected Hamiltonian Annealing (UHA), is a fully differentiable AIS method. This paper validates the effectiveness of UHA through extensive experiments and highlights the importance of differentiability. ", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "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": "## Novelty and Significance\nVariational inference and MCMC are mainstream approaches for posterior inference, with both merits and shortcomings. Using the MCMC kernels in the variational posterior is an appealing direction to combine both merits together. And this paper makes an important step. By discarding the accept-reject step in the Hamiltonian kernel, UHA enables full differentiability. Compared to the standard AIS where the density ratios can be computed by the interpolated distributions, this paper presents an expression that can be simply obtained from the momentum variables in the Hamiltonian dynamics. Therefore, I think the proposed UHA is novel and  important in the field.\n\n## The tightness of the bound\nBy using more samples, IWAE [1] provably tighten the lower bound compared to the standard ELBO, which enables better model learning. In contrast, the UHA introduces an augmented space for the inference. Is there any guarantees on the tightness of the bound ?\n\n## Test-Time Encoder\nDoes the encoder at test-time also involve the UHA procedure ? If so, how large is the computational overhead compared to a standard NN encoder ? \n\n\n## Clarity\nThis paper is well written. This paper clearly introduces the background knowledge and presents the improvement of the proposed approach. Detailed algorithm details are provided.\n\n## Experiment\nI do not have the expertise to fully evaluate the experiments. But I particularly enjoy seeing the improvements by turning more parameters in Figure 2 and Figure 3. I think they are strong evidences on the benefits of the proposed approach.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I do not see a big limitation of the approach.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "YEoQ3rRXYsz", "reviewer": "Reviewer_x5iz", "summary": "The paper proposes UHA, an original and clever modification of Hamiltonian AIS. The idea is to change the AIS ratio such that it becomes tunable, but can still be computed in closed form (ratios of momentum distributions). This is achieved by dropping the accept-reject step in the HMC update such that a single UHA step only involves (re)sampling momenta and leapfrog integration. This modification allows the authors to tune parameters such as the proposal distribution, the step size and damping coefficients (used in momentum resampling). Many interesting ideas for future extensions are proposed. Experiments on various models demonstrate that UHA obtains a tighter ELBO than other combinations of VI and HMC. \n", "questions": "", "limitations": "", "rating": 7, "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": "# Contributions\n\n* Variational inference with Uncorrected Hamiltonian Annealing (UHA)\n\n* Achieves tighter ELBO and allows optimization of parameters \n\n# Originality\n\nThe paper proposes UHA, an original and clever modification of Hamiltonian AIS. The idea is to change the AIS ratio such that it becomes tunable, but can still be computed in closed form (ratios of momentum distributions). This is achieved by dropping the accept-reject step in the HMC update such that a single UHA step only involves (re)sampling momenta and leapfrog integration. This modification allows the authors to tune parameters such as the proposal distribution, the step size and damping coefficients (used in momentum resampling). Many interesting ideas for future extensions are proposed. \n\n# Quality\n\nThe paper is technically sound and the experiments are instructive. \n\n# Clarity\n\nIn general, the paper is well written. \n\nAlthough the notation is generally consistent, there are a few ambiguities:\n\n- Page 2, Lines 47/48: Here you denote the number of bridging distributions by $K$ although in the remainder $M$ is the number of bridging distributions (and $K$ the number of likelihood evaluations)\n\n- Page 5, Line 172: How are the $\\beta_m$ chosen? When tuning also the $\\beta_m$ (Figure 2), it would be interesting to see the tuned annealing schedule. \n\n- Page 6, Lines 210-213: What is $d$? I guess the dimension of $z$ space? \n\n# Significance\n\nThe experiments show that UHA outperforms other variants of Hamiltonian sampling and variational inference. \n\nHowever, to better judge the improvements achieved with UHA it would be useful to also show the correct log evidence (or rather an accurate estimate obtained, for example, with standard AIS using many bridging distributions and multiple parallel trajectories). It would be helpful to show these estimates in Figure 1-3. \n\n# Comments\n\n* Line 204: You perform only a single leapfrog step. Have you run tested with more leapfrog steps?\n\n* In standard AIS, multiple annealing sequences (\"trajectories\") are simulated independently. Have you considered doing something similar with UHA (resulting in yet another augmentation of the sample space)?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "-", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "OWW_4aqlDte", "reviewer": "Reviewer_5mVb", "summary": "The authors introduce a variational inference method that leverages Annealed Importance Sampling (AIS) and Langevin/HMC MCMC. In particular HMC dynamics are used to target a sequence of intermediate bridging densities in order to generate approximate samples from the target posterior distribution. Importantly, since the method is uncorrected (i.e. does without a accept/reject step) the resulting lower bound is fully differentiable. This makes various hard-to-set parameters (e.g. temperatures for bridging densities) learnable. In experiments the authors demonstrate that their method, Uncorrected Hamiltonian Annealing (UHA), leads to tight lower bounds on the log evidence. ", "questions": "", "limitations": "", "rating": 8, "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": "This work presents an elegant approach for combining AIS and HMC within the confines of variational inference. While various authors have proposed methods that marry MCMC with variational inference over the years, some of the resulting constructions are somewhat clunky, requiring e.g. learning reverse models. In contrast the proposed method, UHA, is comparatively simple and does not require learning complex neural networks or the like. What is particularly attractive about this approach is that it effectively defines an implicit/semi-parametric variational distribution that makes extensive use of the target's gradient information. For this reason one would expect UHA to be an effective inference algorithm in a wide variety of scenarios without the need for model specific modifications (at least if the method is reasonably numerical stable). So while this work is in some sense incremental in that it combines existing algorithms from the literature, it does so in a principled way that cleverly addresses an important class of problems. The experiments are reasonably comprehensive and convincing.\n\nDetailed questions and comments:\n\n- The related work section is rather cursory. I think it would be very valuable if the authors extended the discussion there to give a more comprehensive comparison to existing methods that combine MCMC with variational inference. Even if there is not sufficient room in the main text, this discussion could be included in the supplemental materials. Apart from this shortcoming, however, the paper is generally clearly written.\n\n- \"In all cases, for UHA and Hamiltonian AIS we use under-damped Langevin dynamics, that is, we perform just one leapfrog step per transition and partially resample momentum.\" Can you comment on the choice of one leapfrog step? For a fixed amount of computation cost is the regime with more leapfrog steps and fewer bridging densities less effective? Or is this choice driven by numerical stabilities or some other concern?\n\n- Can you comment on using different momentum distributions for each bridging density? It seems like this could be important. Is there a simple parametric scaling that suggests itself?\n\n- Can you comment on numerical stability of the method? It seems you initialize with a vanilla mean field distribution obtained from ELBO maximization. Does everything break down without this initialization scheme, e.g. if you just initialize q(z) to be the unit normal distribution?\n\n- How well can UHA deal with multi-modality? It would be informative to e.g. consider a 2-dimensional multi-modal example and depict the learned variational distribution for increasing values of K.\n\n- \"However, for smaller K, IW sometimes had better log-likelihoods than UHA (despite worse ELBOs)\" Any hypothesis as to why this is?\n\n- There is no discussion of runtime or computational complexity. It would be useful to get an idea of how expensive UHA becomes for large K.\n\n- \"For K ≥ 64 we tune the UHA’s parameters using K = 64 and extrapolate them as explained in Appendix D.\" Is this entirely because of computational cost or are there other factors in play here as well (e.g. stability)?\n\n## After author response\n\nI find that the authors have answered questions raised by the reviewers adequately and continue to believe that this paper should be accepted for publication. The elegance and (relative) simplicity of UHA make it a valuable contribution to the literature, and I imagine it is quite likely that follow-up work will exploit the ideas introduced here in other fruitful contexts.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors should more explicitly address some of the (potential) limitations of their work. For example: computational complexity, performance on multi-modal problems, non-applicability to discrete latent variables, etc\n\nThe authors do not address potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.25"}, "extra_text": {}}, {"review_id": "-ghLj22EARw", "reviewer": "Reviewer_kH8D", "summary": "The paper proposed to use the uncorrected HMC kernel in an AIS-type variational scheme. This enables optimizing the ELBO via reparameterization gradients for many algorithm tuning parameters, including step size, momentum covariance, and annealing schedule. The method is shown to outperform competitors in most cases. ", "questions": "", "limitations": "", "rating": 7, "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": "The proposed method appears to be sound and potentially quite useful. A few suggestions:\n\n1. The proof of Thm 2 lacks some rigor – it would be better to argue measure theoretically.\n\n2. Many of the papers cited as being on arXiv are published either in workshops or conferences. The author(s) should check all arXiv citations carefully. One tricky case is [17], where there's a related published paper https://papers.nips.cc/paper/2016/file/0e9fa1f3e9e66792401a6972d477dcc3-Paper.pdf\n\n3. There are a few missing pieces and areas for improvement in the experiments:\n\n- Please include training runtimes for the different methods\n- I would have liked to see a \"toy\" example with known log Z to demonstrate the tightness of the lower bounds produced by the proposed method\n- Given that the proposed method produces a posterior approximation, some comparison of the approximation accuracy compared to, say, Stan's dynamic HMC implementation would have been very interesting. I realize posterior approximation per se is not the focus of the paper, but I wonder how competitive the proposed method would be with dynamic HMC when measured in terms of number of gradient evaluations. Even if it were not quite as good, the fact that is also produces an estimate of the marginal likelihood could make the method a valuable alternative when the user wants to do model selection ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There was some discussion of limitations", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers all came to a clear consensus that this paper should be accepted. The authors addressed all points raised by reviewers thoroughly. In the camera ready, please make sure to follow through on the edits discussed during the review process. In particular, multiple reviewers pointed out that the experiments should include runtimes, and that there should also be an experiment that provides a comparison to the true log marginal -- please ensure these are included in the camera-ready. Please also improve the rigour of the proof of Theorem 2 using measure-theoretic arguments.

Decision: Accept (Poster)