Differentiable Annealed Importance Sampling and the Perils of Gradient Noise
Authors: Guodong Zhang
Kyle Hsu
Jianing Li
Chelsea Finn
Roger Baker Grosse
Venue: NeurIPS 2021 Poster (2021)
Keywords: Annealed Importance Sampling
Marginal Likelihood Estimation
Variational Inference

Abstract: Annealed importance sampling (AIS) and related algorithms are highly effective tools for marginal likelihood estimation, but are not fully differentiable due to the use of Metropolis-Hastings correction steps. Differentiability is a desirable property as it would admit the possibility of optimizing marginal likelihood as an objective using gradient-based methods. To this end, we propose Differentiable AIS (DAIS), a variant of AIS which ensures differentiability by abandoning the Metropolis-Hastings corrections. As a further advantage, DAIS allows for mini-batch gradients. We provide a detailed convergence analysis for Bayesian linear regression which goes beyond previous analyses by explicitly accounting for the sampler not having reached equilibrium. 
Using this analysis, we prove that DAIS is consistent in the full-batch setting and provide a sublinear convergence rate. Furthermore, motivated by the problem of learning from large-scale datasets, we study a stochastic variant of DAIS that uses mini-batch gradients. Surprisingly, stochastic DAIS can be arbitrarily bad due to a fundamental incompatibility between the goals of last-iterate convergence to the posterior and elimination of the accumulated stochastic error. This is in stark contrast with other settings such as gradient-based optimization and Langevin dynamics, where the effect of gradient noise can be washed out by taking smaller steps. This indicates that annealing-based marginal likelihood estimation with stochastic gradients may require new ideas.


Differentiable Annealed Importance Sampling and the Perils of Gradient Noise
Guodong Zhang Kyle Hsu Jianing Li', Chelsea Finn², Roger Grosse 1:2 UUniversity of Toronto, Vector Institute, SStanford University {gdzhang, rgrosse toronto. edu (kylehsu, cbfinn)0cs. wttssssaaffrrr. edu, jrobert jrobert..i@maal..............l.t uttoonoo..oooooocc
Abstract
Annealed importance sampling (AIS) and related algorithms are highly effective tools for marginal likelihood estimation, but not fully differentiable due to the use of Metropolis- -asttings correction steps. Differentiability is a desirable property would admit the possibility of optimizing marginal likelihood an objective using gradient-based methods. To this end, we propose Differentiable AIS (DAIS), variant of AIS which ensures differentiabilityy by abandoning the Metropolis- Hastings corrections. As further advantage, DAIS allows for mini-batch gradients. We provide detailed convergence analysis for Bayesian linear regression which goes beyond previous analyses by explicitly accounting for the sampler not having reached equilibrium. Using this analysis, we prove that DAIS is consistent in the full-batch setting and provide sublinear convergence rate. Furthermore, motivated by the problem of learning from large-scale datasets, we study a stochastic variant of DAIS that uses mini-batch gradients. Surprisingly, stochastic DAIS can be arbitrarily bad due to fundamental incompatibility between the goals of last-iterate convergence to the posterior and elimination of the accumulated stochastic error. This stark contrast with other settings such as gradient- based optimization and Langevin dynamics, where the effect of gradient noise can be washed out by taking smaller steps. This indicates that annealing- based marginal likelihood estimation with stochastic gradients may require new ideas.
Introduction
Marginal likelihood (ML.,, sometimes called evidence, is a central quantity in Bayesian learning it measures how well model can describe a particular dataset. It is commonly used to select hyperparameters for Gaussian processes [Rasmussen, 2003], where either closed-form solutions or accurate, tractable approximations are available. However, it is more often the case that computing ML is computationally intractable, as it involves summation or integration over high-dimensional model parameters or latent variables. In this case, one must resort to numerical methods or other approximations [Kass and Raftery, 1995]. In the context of model comparison (e.g., evaluating gener- ative models [Wu al., 2016, Huang et al., 2020]), annealed importance sampling (AIS) [Neal, 2001] is one of the most popular and effective algorithms. Notably, AIS is closely related to other generic MI. estimators that yield accurate estimation [Grosse et al., 2015), including Sequential Monte Carlo (SMC) [Doucet et al., 2001] and nested sampling [Skilling et al., 2006]. Under some assumptions, AIS able to produce accurate estimates of marginal likelihood given enough computation time (it converges to the true ML value quickly by adding more intermediate distributions).
AIS alternates between Markov chain Monte Carlo (MCMC) transitions and importance sampling updates. where the MCMC step typically involves non-differentiable Metropolis -aaiisss (MH) correction. Unfortunately, the non-differentiabiliyy precludes gradient- based optimization of the sampler and complicates theoretical analysis. To deal with this, we marry AIS with Hamiltonian
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Monte Carlo (HMC) [Neal, 2011] and derive an unbiased yet differentiable ML estimator named differentiable AIS (DAIS) by removing the MH correction step, which further unlocks the possibility of mini-batch computation. Moreover, DAIS can be made memory efficient by caching noise and simulating Hamiltonian dynamics in reverse [Maclaurin et al., 2015]. We analyze the convergence of DAIS in the setting of Bayesian linear regression. Our analysis goes beyond prior analyses of AIS in that we account for the sampler not having reached equilibrium. In the full-batch setting, we show that DAIS retains the consistency guarantee of AIS (despite the lack of MH steps) and has sublinear convergence rate.
Furthermore, motivated by the problem of learning from large-scale datasets, we study a stochastic variant of our algorithm that uses gradients estimated from subset of the dataset. Given the success of stochastic optimization [Robbins and Monro, 1951, Bottou and Bousquet, 2011 1] and stochastic gradient MCMC algorithms [Welling and Teh, 2011, Chen et al., 2014], one may presume that stochastic gradient DAIS would perform well. Surprisingly, the stochastic version of DAIS can be arbitrarily bad. In partieular, we show that the log ML estimates of DAIS with stochastic gradients are inconsistent due to a fundamental incompatibility between the goals of last-iterate convergence to the posterior and elimination of the accumulated stochastic error. This is in stark contrast with other settings such as gradient-based optimization and Langevin dynamics, where the gradient noise can be washed out by taking smaller steps. This indicates that annealing aanselnng-based ML estimation with minibatch gradients may require new ideas.
We validate our theoretical analysis with simulations. We also demonstrate empirically that DAIS can be applied to variational autoencoders (VAEs) [Kingma and Welling, 2013, Rezende et al., 2014] for a tighter evidence lower bound, which in turn leads to improved performance compared to the vanilla VAE. We also compare to importance weighted autoencoders (IWAE) [Burda et al., 2016]. While IWAE is more effective with a low compute budget, we show that DAIS eventually outperforms IWAE compute increases. Finally, like AIS, DAIS can be used to evaluate generative models. We show that it performs on par with AIS despite the removal of the MH correction step and outperforms the IWAE bound by large margin.
2 Background
2.1 Marginal Likelihood Estimation
For model M and observed data D, one can define the marginal likelihood (ML) as
M(DMM = P(D, M)M)do= = f h M(D0)M))(0)M.....
(1)
where O denotes the parameters of the model. ML estimation can be regarded as an instance of estimating the partition function of an unnormalized distribution. Given a distribution defined as p(0)= = / Z where the unnormalized density s(0) can be efficiently computed, we are interested in estimating the partition function Z f())do. Here, f(0) corresponds to p(D, 0 .M) in (1). In this paper, we focus on the setting of ML estimation because the stochastic version of DAIS is most naturally viewed as arising from data subsampling, but otherwise our analysis applies to the more general setting.
is often the case that computing ML is computationally intractable. One approach is to approximate (1) with Monte Carlo methods, e.g., one can approximate the integration using importance sampling:
P(DIM) BBy (D) (())
1 S
O( MI ((,))
with 02 ~ (())
(2)
However, this estimation exhibit high variance for small or medium S when the target distribution p(D, M) and proposal distribution (0) are dissimilar.
2.2 Annealed Importance Sampling
Annealed importance sampling (AIS) is an algorithm which estimates the ML by gradually changing, or "annealinn"". a distribution. Formally, the algorithm takes a sequence of distributions ,K. with ps (0) f0) Z& and Z fo)()d In the context of ML estimation, the starting

---

distribution fo is the tractable prior distribution M()) with Zu = 1, while the target distribution fK p(D. M) with 2x = p(D]M). For each Pk: one must also specify an MCMC transition operator Tk which leaves pı: invariant.
The output of AIS is an unbiased estimate 2x of the exact ML Zx. Importantly, unbiasedness holds for any finite K, as shown in Neal [2001]. Moreover, AIS can be viewed importance sampling over an extended space [Neal. 2001]. In particular, we have 2x Equa /Twaa/ /qwal] with the target and proposal distributions defined as
(iwd Oo: K = po (o((0) ... (x(0x (xx1)
(3)
Qowd (0o.k) = (0= (00)0))))))))
where (0)0") is forward MCMC kernel and T. (0'10) = (o Ip: (f')p( (0) is the correspond- ing reverse kernel. Here, Qlwd represents the chain of states generated by AIS, and lowd is fictitious (unnormalized) reverse chain which begins with a sample from pK and applies the transitions in reverse order. In practice, the intermediate distributions have be chosen carefully for a low variance estimate 2x. One typically uses geometric averages of the initial and target distributions:
(5) where 0- = 90 < B1 < ... < 3x 1 is the annealing schedule. Indeed, AIS gives an unbiased estimate 2 of Z. However, as Z can vary over many orders of magnitude, it is often more meaningful to talk about estimating log Z. Unfortunately, unbiased estimators 2 can result in biased estimators
Pk (s.(0) (0) f3. (0) fo(0)'~f+ (0) /Zg. p(O M))( MY
of log Z because E log 2 < log E& by Jensen's 's inequality, resulting in only a lower bound. In particular, we have the AIS bound
E Bygaa log 2x
Σ
Egun [log f&z (0k-1) log fa.-, (0k-1)]
(6)
k=1
K
(Bk - 1k- [log PD(00s_11 Σ
where (5) facilitated the simplification from (6) (7). Of course, it is not enough have lower bound; we would also like the estimates to be close to the true value. Fortunately, AIS is consistent in that the estimate log 2 converges to the correct value in the limit of infinitely many intermediate distributions [Neal, 2001], under the very idcalized assumption of perfect transitions (i.e. that each transition To generates an exact sample from pk. independent of the previous state). To give some intuition, the bound (6) can be simplified as
E Egn log 3x log 3K- DKL lppk), pk),
(8)
and the sum of the KL. divergence terms diminishes as K xo.
3 Differentiable Annealed Importance Sampling
In this section, we motivate and derive a dif- ferentiable AIS (DAIS) algorithm for marginal likelihood (ML) estimation. We also discuss its application to variational inference for a tighter ELBO and memory efficient implementation.
Algorithm Differentiable AIS (DAIS)
Go vo sample from Pa (o()), () LDAIS = - log pn( (p)) for k 1..... K do
2 <- 0 + MM-Vv
Ideally, assuming continuously parameterized model class (e.g.. variational autoencoder), we would like differentiate through the lower bound (7) to find an optimal model M. How- ever, AIS must be instantiated with an MCMC transition kernel Tj that satisfies detailed bal- ance to ensure it leaves pk invariant. practice, this is typically achieved by using MH step,
Vi V4-1 + + log Sa. (0 00=
0: + 0z- + + + -- lik
V& f- 7ルn + / - 728, E ~ N(O.
L += log (v+)) log π (V=n1)
end for return LDAIS += p( 0x (M)
3

---

which is generally not differentiable. We thus remove the MH correction and, in particular, specify each transition to consist of deterministic leapfrog integration step followed by stochastic par- tial momentum refreshment [Horowitz, 1991]. Algorithm 1 details the simulation of Hamiltonian dynamics using such transitions. With ? = O, the algorithm resemables unadjusted Langevin dynam- ics [Roberts et al., 1996] but computes the MI. bound on the fly. In practice, choosing 0 < < < (7 = 0.9 a common default) helps avoid random walk behavior and accelerates mixing [Neal, 201 I, Chen et al., 2014]. Importantly we retain the formalism of performing importance sampling an extended space despite the loss of detailed balance." To show this, we can define the forward and (unnormalized) backward distributions
qwd (ow((o..: Vo.K) = Po(Oo) v1 00. vo) ... Tx(0k V 1K
(9)
Powd lowa(On:K: Vo.k) = (x )T
(x -1 0x vx) ... (10. vol0
(10)
where the transition operator (Ox V& 1 = (x 0 vk-1) T T (vxlV*) is the composition of a leapfrog step and momentum refreshment step. We define the reverse chain by starting with an exact sample and executing each of the above steps of Algorithm in the reverse order, which leads to a surprisingly simple expression for our estimator. In particular, the backward transition operator is defined by Te(O Vk-1 [0,, v&) = (
Vk-1 V& ). Note that we need to flip the sign of VA the reverse chain to account for time reversal. As consequence of the above definitions, we have
Th (VVv/VVs)) = (VVV
(1I)
This is because (V"(va [Vva/Va VR.) ((yvn 72)M) and (V' vV) = -N(v, (1 - -2 M)) Fur- thermore, since T% a deterministie leapfrog update, is reversible and volume preserving, so we have Ty (0 Vk- 02, vva) = (O,, vx10x-1, Va-1). With this, we can derive the DAIS bound:
DDAI Equa [log gowal (0:KK Vo.k) log qtwa (oo:K, Vo:k]]
K
x xvx))
(12)
Eggus logp(D,,x| - (0 +
Σ log n (V
=
We remark that this bound supports the computation of pathwise derivatives.
It's useful to consider an intuition for the final term of (12), since this will help to clarify our convergence analysis. Observe that log (Vvx) simply the negative kinetic energy plus constant, so this term will be negative if the kinetic energy increases over the course of a leapfrog step and positive if it decreases. For small enough step sizes, leapfrog steps approximately conserve the total energy. If the posterior distribution is becoming more concentrated over the course of annealing (as typically the case), the start of the leapfrog step is likely to have atypically high potential energy, so the leapfrog step will convert the potential energy to kinetic energy, and this term will be negative in expectation. Conversely, if the distribution becoming more spread out, kinetic energy will be converted to potential energy, and the term will be positive in expectation. Hence, when summed over the whole trajectory, this term helps estimate the volume of the support of the posterior distribution.
3.1 Differentiable Annealed Variational Inference
DAIS can be applied to variational inference for a tighter bound; we name this differentiable annealed variational inference (DAVT). We note that the general idea of incorporating auxiliary MCMC states into a variational approximation was discussed in Salimans et al. [2015], but their formulation requires the specification and learning of reverse transition model, whereas ours does not.
Recall that we can lower bound the log ML by choosing a tractable variational distribution and optimizing the bound. This has been widely adopted in variational autoencoders [Kingma and Welling, 2013, Rezende et al., 2014] and Bayesian neural networks [Blundell et al., 2015, Zhang et al., 2018]. The lower bound has the following form:
L Egg [log pDD, 0 0lM) - log q. (0)
(13)
'The discontinuity introduced by the MH step makes it hard to the reparameterization trick. though this can done using a delicate gradient estimator [Naesseth et al., 2017]. This orthogonal our work and our fix is simpler and easier for us to analyze.
This true even with mini-bacch gradients.

---

However, the lower bound can be quite loose if the variational posterior family 9s (0) is restrictive, e.g. Gaussian. To improve the bound. we can define a new variational distribution on an extended space as in (9), but starting from de rather than po:
giwd O Vo.k) (vo))))(V)TT(((O V1 Oo vo) TK((K vK OKI, VK-1).
(14)
We also define associated intermediate distributions p& 2o (0) MMM This gives new lower bound:
log (vx) π (V
DDAVI
Egun logp(D,0x|M) log go
(15)
We can maximize this lower bound over model parameters of M. all parameters of AIS (c.g.. annealing schedule 3ạ) as well as variational parameters Φ.
3.2 Memory-Efficient Implementation
Naively optimizing instantiations of (12) or (15) w.r.t. parameters using reverse-mode differentiation in- volves storing the entire sequence of sampled states oo Vo..... .o....,0.. VK- This can be problematic in cases when K is large due to the large memory overhead. However, DAIS compatible with the idea of re- versible learning [Maclaurin et al., 2015], which ame- liorates this problem. Instead of storing the states in memory, we can compute the previous state given the current state by reversing the dynamics. Recall that each DAIS transition is deterministic and reversible other than the use of noise EG for momentum refreshment. The exact noise samples can also be computed in reverse if one generator) for managing pseudorandom number generator seeds. Assuming exact arithmetic (in practice, this is impossible), this means that the memory footprint of DAIS can be made constant with respect to the number of intermediate distributions K. Similar memory-efficiency tricks have also been used in other applications [Li et al., 2020, Ruan al., 2021].
Table I: Memory and time usage of DAIS imple- mentations. / is 32 for point format.
floating-
Scheme
Memory
Time
Naive
O(BK)
O(K) O(K)
Rev. Learning
(11g2(1/7)
deterministic and reversible scheme (e.g. the linear congruential
However, as discussed by Maclaurin et al. [2015], reversible learning with finite arithmetic precision requires some storage counteract compounding round-off error. For y # 0 (ry 0.9 is common default), we need on average log, (1/7) bits per parameter per step. which is still small compared to naive storage. We defer further exposition neeoory-effcciee DAIS to Appendix D. We remark that reversible learning a potentially crucial property DAIS as affords some degree of scaling to longer chain lengths and, indirectly, bigger models.
4 Convergence Analysis for Bayesian Linear Regression
Neal [2001] has pointed that AIS is consistent, i.e. that it converges to the true log ML value in the limit of infinitely many intermediate distributions. However, these consistency results depend on the idealized assumption of perfect transitions (where each transition returns an independent exact sample), and therefore don't 't account for the time required for the samples reach equilibrium. Here, we analyze DAIS for a Bayesian linear regression model with realistic (imperfect) transitions. Accounting for the convergence of the sampler is essential for separating the behaviors in the bull-bacchaae and mini-batch regimes.
In particular, focus the Bayesian linear regression setting and adopt the following model:
prior: 0 ~ ~N 1
likelihood: y ~ N(00, I'I) = 0 N Nu,,, A where And = X and p.x = (X -XX 0²
posterior: 0 (H A where plos = A (ApHtp Anale.) and Apos Ap + Aud
with E R denoting the input features and y RRax1 the targets. We choose Bayesian linear
regression because it enables us to analyze the dynamics analytically in similar manner as done by the noisy quadratic model (NQM) [Zhang et al., 2019] in the context of optimization. We adopt the leapfrog step (we assume identity mass matrix without loss of generality because can absorb
5

---

M into the input matrix X in Algorithm 1) and obtain the following update rule (see Appendix B. for derivation):
0%: -
1
Th
A
84 pos
0
The
A
8x pos
Vk
+
Th
A
9 pos poss
2
4
2
(16)
θk
A
Bk pas
Vk-1 + 7) A
Bx pos
2
where A B pas A + 3a And and Maos we can compute the expectation and covariance of Og and Vy at any time which suffices to compute the lower bound in closed-form.
Pk
(A
84 pos
(A pplp + Br-Anates)) With these iterative updates,
4.1 Sublinear Convergence in the Full-Batch Setting
With the model defined, we now show that our algorithm is asymptotically consistent, i.e., the bound (12) converges to exact log ML as K goes to infinity. For Bayesian linear regression, the update rules in (16) are affine transformations of Gaussian random variables, so the distribution of 0g is also Gaussian in the form of (us, E₁). We can compute the gap between the log ML and our lower bound closed-form (see Appendix B.2 for derivation):
logp(D) P(D) LAAI
1 +
Epos
K
(Vv)) π (VV=-1)
pk
Mpos Agen
Tr
+ log
Egg
log
(17)
2
2
Σ
where d is the feature dimension. Here, 1 and ② measure the error of last-iterate Markov chain convergence and will both vanish as long as μk + Mpps and ΣK - Epss- We will show later that they converge with rate of O (G The key is to show that (resp. A&) lags behind plpos (resp. A poss) with roughly 1 steps. Formally, we have the following.
8
Lemma 1. Given equally spaced Bg. running DAIS with 7 0 and 1 where 2 4 yields Re
Pk-1
Ba 14 pas O(K20-1), ||A A Ah l12 = O(K2o-1).
(18)
We remark that the assumption of 3%. being equally spaced is not essential and can be relaxed as long as they are chosen by scheme that leads to A5. - -8k-1 going down approximately in inverse proportion to K. In addition, note that the assumption of full momentum refreshment is for convenience and we believe a similar result holds for 7 > O.
Importantly, this lemma implies that both L and 2 vanish sublinearly if we choose c <<< 1 The 2' analysis of error term 3 is more nuanced. In particular, this error could either come from using transitions for each these intermediate distributions that do not bring the distribution close to equilibrium, from using finite number of distributions to anneal from Po Pk. Surprisingly, the error 3 decays as fast as the other two terms if the step size scales as 1/K with 2 1. In summary, we have the following theorem.
Theorem 1. Given equally spaced BK. running DAIS with 7 = and 17 Ke where c 2 yields
logp(D) - LDAIS = O(K2e-))
With 1. we have the convergence rate (LL/VK).
We remark that with perfect transitions, the requirement of e 2 1/4 is not necessary and we can achieve (1//K) convergence, as also shown Grosse et al. [2013]. The gap between O(1/VR and O11/K) highlights the importance of considering convergence the stationary distribution.
4.2 Inconsistency in the Stochastic Setting
Standard AIS involves MH steps which require a costly computation using all of the data, thus defeating the potential computational efficiency benefit of stochastic gradients [Chen et al., 2014].

---

This is likely why the convergence properties of stochastic gradient AIS were previously unknown. In contrast, the only part of DAIS that requires full-batch computation is the gradient term in the leapfrog step of Algorithm 1, which may be amenable to estimation via mini-batching.
We have shown that our algorithm is asymptotically consistent in the full-batch setting. Often, consistent/convergent algorithm in the deterministic setting readily implies similar convergence result in the stochastic setting. For example, SGD [Robbins and Monro, 1951] and SGMCMC [Chen et al., 2014, Ma et al., 2015] are both convergent in the presence of noise. This begs the question of whether DAIS is consistent when we only have access to stochastic gradients. Here, we adopt an additive noise model V log fx (0)
log (0) + E. This model commonly used in the stochastic approximation literature, and such a model has also been adopted in Chen et al. [2014]. With such a noise model, we have the following dynamics:
O2: f- (I-
A
Fk pos
OK-
OF
I
Tk
3
A
Vk
+
ms
8y pos
M
+
e
2
4
2
2
(19)
2 7E
1/ 4- 7) A
Bk pos
0
I
A
8% pos
Vk-1 1 + 176 A 8k pis
where stochastic noise E has variance lowered bounded by E.. Further, we let 1e E and Σ be the covariance of v/. Surprisingly, we find that DAIS is incompatible with stochastic gradients. even though it admits the same importance sampling interpretation the reverse transition To conditions on the same matc- of data as Ti. We summarize the result in the following theorem.
Theorem 2. For stochastic DAIS with full momentum refreshment (? in Algorithum 1) and any stepsize scheme, we have
lim inf [log P(D) - -Daisl > O.
(20)
K-500
Here, we give some intuition why DAIS fails in the stochastic setting. To ensure convergence of 0K to (H po a major requirement is for the step sizes to satisfy lim Σx= no = oo [Robbins and Monro, 1951]. However, the randomness of mini-batching sampling would contribute to the variance Vz. In particular, we have the following recursion:
2
RE Σ = Th- Ao pos
1
BK +
I
Th
8 pos
m Ee.
(21)
2
For notational convenience, we let Σx A ni Σy and Σg avoid confusion with ΣΣ and Eg in the full-batch setting. In this case, if we follow Stephan et al. [2017] and Chen et al. [2014] in assuming is Gaussian," we have
A par
8r Ex
8 pos
+ (I-
)2 Here, we used Σy,
2
Ey
K
log
π ((vs)
Σ k=1
K
1
| 2 2
1 IT
d
Σ
K
2
The T
(22)
(V
2
The second term of (22) goes to infinity as lim Ek v = xo. Intuitively, the gradient noise adds the kinetic energy, and the size of this contribution proportional to n Since this effect is cumulative over all K steps, The has to be reduced at least as 1/vK for the kinetic energy term to go down. However, this contradicts the requirement that lim Σ ni = oc, needed for last-iterate convergence. In summary, the convergence of 0Ι requires us to make sure the sum of step sizes goes to infinity, which in turn results in con-convergent kinetic energy term (22).
One may wonder why gradient noise does not hurt the convergence of SGLD [Welling and Teh, 2011] or SGMCMC [Ma et al., 2015). Generally speaking, these algorithms are only concerned with the last iteration convergence to the true posterior, hence one can eliminate the stochastic error by taking more steps of a smaller size. In contrast, our bound (and potentially other AIS-style algorithms) relies on all intermediate distributions, and so the error induced by stochastic gradient noise accumulates over the whole trajectory. Simply taking smaller steps fails to reduce the error. We conjecture that AIS-style algorithms are inherently fragile to gradient noise.
2 II reality, the noise consists of two parts: multiplicative input subsampling noise and additive label noise. We can assume that the noise is lower bounded by the additive part. See Appendix B.4 for justifications. "This assumption is non-essential, though Gaussian noise is reasonable by the central limit theorem.
7

---

5
Related Works
For ML estimation (or partition function estimation), Sequential Monte Carlo (SMC) [Doucet et al., 2001, Del Moral et al., 2006] is another popular method which is derived from particle filtering. While SMC is based on different intuition from AIS, the underlying mathematics is equivalent. In SMC, the intermediate distributions are defined by conditioning on a sequence of increasing subsets of data. Both AIS and SMC are closely related to broader family of techniques for partition function estimation, all based on the following identity from statistical physics: log Zx log 2o-
So Ee-po approximation. In comparison, thermodynamic integration (TT) [Frenkel and Smit, 2001] estimates this integration using numerical quadrature, and path sampling [Gelman and Meng. 1998] does so with Monte Carlo integration. Recently, Masrani et al. [2019] connected TI and variational inference for a tighter bound on the log ML, but they computed each intermediate term using importance sampling rather than annealing- based sampling algorithms. More recently?, Thin et al. [2021] proposed a similar algorithm based on sequential importance sampling with unadjusted Langevin kernels in the context of variational auto-encoders. Concurrently, Geffner and Domke [2021] proposed the same algorithm as our DAVI with a focus on the empirical side.
d d8 log f0( dB. In particular, the weight update in AIS can be seen a finite difference
In the context of variational inference, many papers have also investigated tighter lower bounds for the log MI. Burda al. [2016] proposed strietly tighter log-likelihood lower bound derived from importance weighting and Luo et al. [2019] recently extended this idea of importance sampling to derive an unbiased (but potentially high variance) estimator of log ML using the Russian roulette estimator [Kahn, 1955]. Salimans et al. [2015] proposed to incorporate MCMC iterations into the variational approximation. The central idea that we can interpret the stochastic Markov chain as part of a variational approximation in an extended space, as we did in the paper. However, the proposed methods require learning reverse kernels, which has a large impact on performance. The same authors also briefly discussed annealed variational inference, which combines variational inference and AIS. However, their derivation relies on the detailed balance assumption and is therefore not amenable to gradient-based optimization. Later, Caterini et al. [2019] proposed the Hamiltonian VAF, which improves Hamiltonian variational inference [Salimans al., 2015] with an optimally chosen reverse MCMC kernel. In particular, they removed the momentum sampling step and used deterministic transitions. The resulting algorithm can be thought as a normalizing flow scheme in which the flow depends explicitly on the target distribution. Along this line, Le et al. [2018] 8J. Naesseth et al. [2018], Maddison et al. [2017] proposed to meld variational inference and SMC for time-series models.
Finally, stochastic gradient variants of several MCMC algorithms [Welling and Teh, 201 1, Chen al.. 2014, Ma al., 2015] have been proposed over the last decade. In particular, these works showed that adding the "right amount" of noise to the parameter updates leads to samples from the target posterior as long as the step size is annealed. Importantly, the convergence rates of these algorithms are established in both the full-batch setting [Dalalyan, 2017, Cheng et al., 2018] and the stochastic setting [Chen et al., 2015, Teh et al., 2016, Raginsky et al., 2017, Zou et al., 2020]. By contrast, the convergence properties for AIS and related algorithms were largely unknown even for the deterministic case, and it remains largely unexplored whether AIS can be made compatible with stochastic gradients.
6 Simulations
In this section, we discuss the experiments used to validate our algorithm and theory. Importantly, we do not aim achieve state-of-the-ar on these tasks.
6.1 Bayesian Linear Regression
In Section 4, we proved for the Bayesian linear regression setting that while DAIS is asymptotically consistent with full-batch gradient, the noise injected into the system via stochastic gradients precludes convergence. Here, we verify our theory with numerical simulations. The n input vectors X E Rnx and targets y E R² respectively consist of entries sampled from (0 0.01) and A((0,)) In particular, we choose n 10, 000 and d = 10 for our simulations (the results are qualitatively same
'The paper appeared arXiv after we submitted the paper.
8

---

0.250
Suboptimelty Gap
Subeplimeiity Gep
Suboptimalty Gap
0.335
1.0
0.500
1.1
20'
106
0.250
sampll
10
0.500
Uheory
10
12"
Num of Intermediate Distributions
Num of Intermediate Distributions
Num of Intermediate Distributions
(a) Full refreshment
(b) Partial refreshment
(c) Mini-batch gradient
Figure Cap between log ML and DAIS bound function of number of intermediate distributions. Solid lines are exact computation our DAIS bound; dotted lines are sample-base simulation (Monte Carlo method with 100 samples); dashed lines are theoretical predictions based on Theorem with slope 2c the rightmost figure, we use a batch size of 100.
For
Table 2: Test negative log likelihood of the trained model, estimated using AIS with 10,000 intermediate distribution and 10 particles. For VАE/WWAE, we used $ x K samples. The numbers reported are averaged over three runs. The standard deviations fairly small over three runs (< 0.06).
S K S K
K
50
500
Objective
=
K
= K = 10 K = K 10
= 50
VAE
86.93 86.93
86.95
86.94
86.89
IWAE
85.43
84.46
83.87
DAVI
86.51
84.49
84.45
85.23
83.62
83.65
DAVI (adapt)
86.49
84.42
84.39
85.00
83.56
83.69
with different and d). In addition, we set the observation variance 6² 1. For convenience, we set the linear annealing scheme 3 K
In Figure 1, we report the gap between exact log ML and our bound as a function of number of intermediate distributions. With full-batch gradients, the simulations (solid and dotted lines) align well with our theoretical predictions (dashed lines) for different step-size scaling schemes, suggesting our bound in Theorem tight. In addition, we observe in Figure Ic that the gap fails vanish with batch gradients for all step-size scaling schemes. Interestingly, with = 1/2, the gap stays constant. This matches our predictions that the deterministic error decays as O(K2 (()) while the stochastic error is proportional to Σ n/ O(K2==1) O(1).
6.2 Variational Autoencoder
compare the performance DAVI to vanilla VAE [Kingma and Welling, 2013] and IWAE [Burda et al., 2016] on density modeling tasks. We use the dynamically binarized MNIST [LeCun et al., 1998] dataset. We use the same architecture as in IWAE paper. The prior p(z) is 50-dimensional standard Gaussian distribution. The conditional distributions p(x; (z) independent Bernoulli, with the decoder parameterized by two hidden layers, each with 200 tanh units. The variational posterior q(z|x) is also 50-dimensional Gaussian with diagonal covariance, whose mean and variance both parameterized by two hidden layers with 200 tanh units (see other details in Appendix C.1).
In the first set of experiments, we investigate the effect of number of intermediate distributions K and combine with importance sampling (as done IWAE) with samples DAVI. To be specific, we define the bound as follows:
po (x, z
(V T
log
(23)
Lo (2
where we sample (z/- vg. vi vi....) independently from qiwa- By default, we use partial momentum refreshment with y 0.9 and equally spaced annealing parameters 8g K. We compare it to vanilla VAE and IWAE bounds with K samples. As shown in Table 2, increasing K gives strictly better models with lower test negative log-likelihood. However, IWAE achieves slightly better performance with roughly the same computation if K is small. the other hand, DAVI is more effective with more compute budget (i.c., S x K is large) and eventually outperforms IWAE.

---

20.3
MMAE
55
HAIS
075
HAIS
60.0
HAIS
DAIS
07.0
10.3
60.2
10
Nam of gradiant updates nartclea
Nam of cradient undatis particiss
Num of gradient ucdates particuss
(a) standard VAE
(b) IWAE with 10
(c) DAIS with K 10
Figure 2: Results of different algorithms in evaluating VAE-, IWAE- DAIS-trained models. DAIS performs
on par with AIS/HAIS but without requiring the MH correction steps.
In the second set of experiments, we learn the annealing scheme of DAVI together with the parameters of encoder and decoder. Comparing the third and fourth rows of Table 2, see that learning the annealing scheme improves the performance slightly.
I Lastly, we also compare our algorithm with IWAE, AIS, and Hamiltonian AIS (HAIS) [Sohl- Dickstein and Culpepper, 2012] in evaluating the log-likelihood of trained models. Notably, IWAE and AIS have been widely used in VAE evaluation, sec e.g. Wu et al. [2016], Huang et al. [2020]. For HAIS and DAIS, we employ the optimal step-size scaling scheme derived in Theorem with = 1/4 and only tune the step size for the case of K 10. For all implementation details, please Appendix C.2. In particular, we compare the evaluation algorithms on models trained using the VAE, IWAE, and DAIS objectives. In Figure 2, report the estimated negative log- likelihood as function of the number of particles (for IWAE) or gradient updates (for AIS, HAIS and DAIS). Interestingly, observe that IWAE performs better when have limited computation and AIS/HAIS/DAIS win out by a big margin we increase K. Indeed, it is known that naive importance sampling can have exponential sample complexity in some problem parameters (e.g. the dimension), and that AIS can overcome this to efficiently give accurate results (see, for instance, the analysis Neal [2001]) J). This may explain the superior performance AIS-based algorithms when more compute is used. Importantly, we observe that DAIS performs on par with AIS/HAIS but without requiring the MH correction steps.
7 Conclusion
In this paper, proposed differentiable AIS (DAIS) algorithm for marginal likelihood estimation. We provided a detailed convergence analysis for Bayesian linear regression which goes beyond existing analyses. Using this, we proved sublinear convergence rate of DAIS in the full-batch setting. However, we showed that DAIS inconsistent when mini-batch gradients are used due to fundamental incompatibility between the goals of last-iterate convergence to the posterior and elimination of the pathwise stochastic error. This comprises an interesting counterexample to the general trend of algorithms consistent in the deterministic setting remaining consistent in the stochastic setting. Our negative result helps explain the difficulty of developing practically effective AIS-like algorithms that exploit mini-batch gradients. Our numerical experiments validate our claims.
Acknowledgements
We thank Ricky Tian Qi Chen, Murat A. Erdogdu, Sergey Levine, Xuechen Li, Mufan Li, Chris J. Maddison, and Matthew D. Hoffman many helpful discussions. particular, we thank Xuechen Li for his suggestion regarding memory-efficient implementation. We also thank Shengyang Sun, Xuechen and Yangjun Ruan for detailed comments on early drafts, and Radford M. Neal for insightful feedback on arXiv preprint. Finally, we thank the anonymous NeurlPS reviewers for their useful feedback on earlier versions of this manuscript.
GZ was supported by Ontario Graduate Fellowship. KH was supported by Sequoia Capital Stanford Graduate Fellowship. RG was supported by the CIFAR AI Chairs program. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
HHAIS is an algorithm similar to ours which combines Hamiltonian dynamics with AIS, but with accept-reject steps included.
10

---

References
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference Machine Learning, pages 1613-1622. PMLR. 2015.
Léon Bottou and Olivier Bousquet. 13 the tradeoffs large-scale learning. Optimization for machine learning, page 351, 2011.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. ICLR (Poster), 2016.
AL Caterini, A Doucet, and D Sejdinovic. Hamiltonian variational auto-encoder. Advances in Neural Information Processing Systems, 31, 2019.
Changyou Chen, Nan Ding, and Lawrence Carin. On the convergence of stochastic gradient meme algorithms with high-order integrators. In Proceedings of the 28th International Conference Neural Information Processing Systems-Volee Volume 2, pages 2278 2286, 2015.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient Hamiltonian monte carlo. In International conference on machine learning, pages 1683-1691, PMLR, 2014.
Xiang Cheng, Niladri Chatterji. Peter L Bartlett, and Michael Jordan. Underdamped langevin meme: A non-asymptotic analysis. In Conference on Learning Theory, pages 300-323, PMLR, 2018.
Arnak $ Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. Journal of the Royal Statistical Society: Series B (Statistical Methodology). 79(3): 651-676, 2017.
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential monte carlo samplers. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 68(3):411 1-66,, 2006.
Arnaud Doucet, Nando De Freitas, and Neil Gordon. An introduction to sequential monte carlo
methods. In Sequential Monte Carlo methods in practice, pages 3-14. Springer, 2001.
Daan Frenkel and Berend Smit. Understanding molecular simulation: from algorithms to applications, volume Elsevier, 2001.
Tomas Geffner and Justin Domke. MCMC variational inference via uncorrected Hamiltonian annealing. Advances in Neural Information Processing Systems, 2021.
Andrew Gelman and Xiao-Li Meng. Simulating normalizing constants: From importance sampling to bridge sampling to path sampling. Statistical science, pages 163-185, 1998.
Roger B Grosse, Chris J Maddison, and Ruslan Salakhutdinov. Annealing between distributions by averaging moments. In NIPS, pages 2769-2777. Citeseer, 2013.
Roger Grosse, Zoubin Ghahramani, and Ryan P Adams. Sandwiching the marginal likelihood using bidirectional monte carlo. arXiv preprint 1rXiv:1511.022 2015.
Alan M Horowitz. A generalized guided monte carlo algorithm. Physics Letters B, 268(2):47--252, 1991.
Sicong Huang, Alireza Makhzani, Yanshuai Cao, and Roger Grosse. Evaluating lossy compression rates of deep generative models. In International Conference on Machine Learning, pages 4444- 4454. PMLR, 2020.
Herman Kahn. Use of different Monte Carlo sampling techniques. Rand Corporation, 1955.
Robert E Kass and Adrian E Raftery. Bayes factors. Journal of the american statistical association, 90(430):773-798 1995.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference Learning Representations, 2015.

---

Diederik P Kingma and Max Welling. arXiv:1312.6114,200 2013.
Auto-encoding variational
bayes.
arXiv preprint
Tuan Anh Le, Maximilian Igl. Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding sequential monte carlo. In International Conference Learning Representations, 2018.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 (6611))2278-2222 1998.
Xuechen Li, Ting-Kam Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pages 3870-3882. PMLR, 2020.
Yucen Luo, Alex Beatson, Mohammad Norouzi, Jun Zhu, David Duvenaud, Ryan P Adams, and Ricky TQ Chen. Sumo: Unbiased estimation of log marginal probability for latent variable models. In International Conference on Learning Representations, 2019.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme, Advances in Neural Information Processing Systems, 28:2917-2925, 2015.
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient- based hyperparameter optimization through reversible learning. In International conference on machine learning. pages 2122-2122. PMLR, 2015.
Chris J Maddison, Dieterich Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Whye Teh. Filtering variational objectives. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pages 6576-656 6586. 2017.
Vaden Masrani, Tuan Anh Le, and Frank Wood. The thermodynamic variational objective. arXiv preprint arXiv:
2019.
Christian Naesseth, Francisco Ruiz, Scott Linderman, and David Blei. Reparameterization gradients through acceptance-rrjection sampling algorithms. In Artificial Intelligence and Statistics, pages 489-498. PMLR, 2017.
Christian Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei. Variational sequential monte carlo. In International Conference on Artificial Intelligence and Statistics, pages 968-977. PMLR, 2018.
Radford M Neal. Annealed importance sampling. Statistics and computing, 11(2):125-139, 2001.
Radford M Neal. MCMC using Hamiltonian dynamics. Handbbook of Markov Chain Monte Carlo, page 113, 2011.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient langevin dynamics: nonasymptotic analysis. In Conference Learning Theory, pages 1674-1703. PMLR, 2017.
Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning, pages 63-71. Springer, 2003.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning. pages 1 1286. PMLR, 2014.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pages 400-407, 1951.
Gareth O Roberts, Richard Tweedie, et al. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli, 2(4):341-36) 1996.
Yangjun Ruan, Karen Ullrich, Daniel Severo, James Townsend, Ashish Khisti, Amaud Doucet, Alireza Makhzani, and Chris Maddison. Improving lossless compression rates via monte carlo bits-back coding, arXiv preprint ar arXiv:2102.11086 2021.
12

---

Tim Salimans, Diederik Kingma, and Max Welling, Markov chain monte carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, pages 1218-1226, 2015.
John Skilling et al. Nested sampling for general bayesian computation. Bayesian analysis, 1(4): $33-859, 2006.
Jascha Sohl-Dickstein and Benjamin Culpepper. Hamiltonian annealed importance sampling for partition function estimation. arXiv preprint arXiv: 2012.
M Stephan, Matthew D Hoffman, David M Blei, et al. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18(134):1-35 2017.
Yee Whye Teh, Alexandre H Thiery. and Sebastian J Vollmer. Consistency and fluctuations for stochastic gradient langevin dynamics. Journal of Machine Learning Research, 17, 2016.
Achille Thin, Nikita Kotelevskii, Arnaud Doucet, Alain Durmus, Eric Moulines, and Maxim Panov. Monte carlo variational auto-encoue encoderss In International Conference on Machine Learning, pages 10247-10257. PMLR, 2021.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pages 681-688. Citeseer, 2011.
Yuhuai Wu, Yuri Burda, Ruslan Salakhutdinov, and Roger Grosse. On the quantitative analysis of decoder-based generative models. arXiv preprint 1..12:1611.04273, 2016.
Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International Conference on Machine Learning, pages 5852-5861. PMLR, 2018.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E Dahl, Christo- pher Shallue, and Roger Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. arXiv preprint arXiv: 1907,04 044,, 2019.
Difan Zou, Pan Xu, and Quanquan Gu. Faster convergence of stochastic gradient langevin dynamics for non-log-concave sampling. arXiv preprint arXiv:2010.09597, 2020.
13


Reviews:
[{"review_id": "nfjFjY6H4m_", "reviewer": "Reviewer_8rku", "summary": "This paper develops and studies a variant of annealed importance sampling with Hamiltonian Monte Carlo transitions that does not require running a Metropolis-Hastings accept/reject step after each HMC proposal. \n\nThe proposed algorithm, like AIS, produces unbiased estimates of a model's marginal likelihood. As such, it can be used to define a Monte Carlo variational objective $L_{DAVI}$, as the expected log marginal likelihood estimate. Without the discrete MH accept/reject step, the objective's gradient can be estimated unbiasedly using the reparameterization trick. In an experiment, the authors exploit this to train VAEs using the $L_{DAVI}$ objective, achieving results competitive with other Monte Carlo variational objectives, namely IWAE.\n\nA downside of the proposed algorithm is that without the MH corrections, the consistency of the AIS algorithm (in the long-chain, fine-annealing-schedule limit) is no longer guaranteed. The paper nonetheless proves a consistency result for the modified algorithm when applied to a particular Bayesian linear regression model, and demonstrates empirically that on the VAE task, performance does improve with longer chains. The paper also shows that the algorithm is *not* consistent for the Bayesian linear regression model if the HMC transitions use noisy gradients. These theoretical results are tested and illustrated in an experiment on the Bayesian linear regression model.", "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": "This is a nice piece of work: it introduces a clever modification of AIS that makes it suitable for gradient-based variational optimization, presents some evidence that the method works, and makes a couple interesting theoretical observations.\n\nFirst, what I like about the paper: DAIS strikes me as interesting in part because, at first glance, removing the Metropolis-Hastings correction in AIS should be dangerous. For one, in many MCMC transition kernels, the MH correction plays an essential role, and without it, the AIS extended proposal distribution could be quite far from the target, yielding high-variance marginal likelihood estimates. Furthermore, the stationarity of the MCMC kernels in AIS leads to a natural choice of extended target distribution, and a simple, easy-to-work-with expression for the marginal likelihood estimate—ditching the MH correction risks losing that. But the paper's clever choice to use Hamiltonian Monte Carlo transitions sidesteps both concerns: well-tuned HMC proposals are almost always accepted anyway, so removing the MH step should not really affect the inference distribution; and the authors show that even without the MH step it is possible to choose a simple extended target $q_{bwd}$, for which the resulting importance weight is a \"surprisingly simple expression\" (L132). (I was surprised!)\n\nIt is also nice to see an analysis of the algorithm, even if only for a specific model (the Bayesian linear regression model)—consistency in the long-chain limit is another thing you might expect to lose by dropping the MH steps, and it's interesting to have an example where it still holds.\n\nA potential weakness of the paper is that its aims are somewhat unclear, as partly evidenced by the title: is the goal to present a new method for variational inference, illuminate \"perils of gradient noise,\" or something else altogether? Because of this, it's unclear by what standards the submission should be judged. For instance:\n\n* I'm not sure what exactly the takeaway message is regarding gradient noise. Even outside an AIS setting, it is known that naively using stochastic gradients in Hamiltonian Monte Carlo can be arbitrarily bad (https://arxiv.org/abs/1402.4102). Furthermore, using stochastic gradients changes the distribution $q_{fwd}$; is there a corresponding change to $q_{bwd}$ that preserves their density ratio $L_{DAIS}$? If not, I'd expect $L_{DAIS}$ to be biased, for any chain length. (If there *is* an importance sampling justification for $L_{DAIS}$ in the presence of gradient noise, it could be useful to present it.)\n\n* If the goal of the paper were primarily to introduce a new method, it would be nice to see more empirical validation or at least discussion of the method's value. What does DAVI add over IWAE? Intuitively, AIS encodes a much more sophisticated inference process than IWAE does (MCMC can solve many inference problems that would be difficult to solve with vanilla importance sampling); are there interesting model families that require this flexibility in the variational posterior? Are there other interesting applications of a differentiable AIS algorithm?\n\nOverall, however, I am inclined to accept the paper because it makes an interesting and technically sound contribution that I can imagine others in the NeurIPS community appreciating and building on.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes; the paper does not overstate its contributions and I think its focus is sufficiently theoretical that if negative societal impact results, it will be very indirectly.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "XxDQtOM0Cz", "reviewer": "Reviewer_SRxf", "summary": "In this paper, the authors introduce a new *differentiable* annealed importance sampling scheme (DAIS), which provides a lower bound to the log marginal likelihood. Crucially, DAIS *does not* require a Metropolis-Hasting (MH) correction step which is what allows it to be easily differntiated through. This is also combined with variational inference which they title DAVI. ", "questions": "", "limitations": "", "rating": 9, "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": "# Strengths\n\nThis is a *fantastic* and very well-written paper. Firstly, the significance of this work is huge: AIS is a popular way to estimate the log marginal likelihood of posteriors but they rely on using a sequence of transition kernels, $\\mathcal{T}_k$ that has $p_k$ as an invariant distribution. The invariant distribution condition is usually enforced by a MH correction step; it is the MH step that prevents differentiation through AIS. The authors completely side-step this by using Hamiltonian Monte Carlo without a correction step, thus violating the invariant distribution condition. As AIS is a form of IS (which provides an unbiased estimate of the marginal likelihood), it is no surprise that DAIS produces a lower bound to the log marginal likelihood; crucially though, the use of HMC without the MH correction allows for a closed-form and simple expression for the actual estimate of the log marginal likelihood. This allows DAIS to be simply combined with variational inference to produce DAVI.\n\nWhat I think makes this paper really great is that the authors don't just stop here. The authors investigate the theoretical properties of DAIS on Bayesian linear regression. Amazingly, the authors demonstrate that the log marg likelihood estimate produced by DAIS is consistent as the $K \\rightarrow \\infty$ for Bayesian linear regression *without the restrictive assumptions used in Neal's AIS paper which relies on the invariance condition of the transition kernel and that each of samples are independent from one another*. Surprisingly, the authors demonstrate that this same result does not apply if stochastic gradients are used (where the stochasticity is from using mini-batches) and their empirical experiments demonstrate that it can provide an arbitrarily bad estimate. This is a huge result and contrasts with the convergence results for stochastic Langevin MCMC and stochastic HMC.\n\nLastly, the authors applied DAVI on learning a VAE on MNIST where they demonstrated that it is competitive with a vanilla VAE and IWAE. The authors also demonstrated that DAIS can produce competitive log marginal likelihood estimates.\n\n# Weaknesses\n\nIn my opinion, there are no glaring weaknesses in this work. There are a couple of things that can be changed to make the paper better but I will list them in the section below.\n\n# Questions/Comments\n1) As the variational approximation used in DAVI is more expressive than a simple Gaussian, it would be cool to see how well DAVI does at estimating complex distributions.\n2) For the VAE experiments, no mini-batches were used, correct? If not, it would be interesting to see how much worse the learned VAE is when using mini-batches.\n3) Errors bars are missing for figures 1 and 21\n4) I think some tweaks can be done to the derivations in the appendix to make it more readable. For example, there are many inequalities that would make the derivation easier to follow if they were either said outright or if a reference was provided for instance $\\Vert AB\\Vert_2 \\leq \\Vert A \\Vert_2 \\Vert B \\Vert_2$. I also thinking color coding would help a lot as well.\n\n# Conclusion\n As I stated earlier, I think this is a great paper and a valuable contribution to the research community.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7 hours"}, "extra_text": {}}, {"review_id": "7IA6rIhTKXd", "reviewer": "Reviewer_m6Nw", "summary": "The authors develop a variant of annealed importance sampling (AIS) that allows for pathwise derivatives to be computed. This method can be used for (approximate) maximum likelihood learning of model hyperparameters. The authors prove that the proposed algorithm is consistent when applied to Bayesian linear regression (even when transitions are not perfect) and leads to a rate of convergence of $O(1/\\sqrt{K})$, where $K$ is the number of transitions performed. The authors show that, although it at first appears the proposed algorithm can be minibatched, the algorithm is inconsistent in the presence of additive noise when applied to Bayesian linear regression. ", "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": "## Summary of review\nThe proposed algorithm appears to address a signficant drawback of AIS in a sensible way. The exposition in the paper is generally very good. The analysis of the algorithm in the cases of Bayesian linear regressoin provides some theoretical evidence that the algorithm is practial. The experiments are a useful illustration of the proposed approach, and the authors are very explicit that the experiments are intended as illustrative. Most of my comments, given below, are relatively minor. I am not very familiar with some of the related work, but my impression is that the work contains an interesting contribution, is well-presented, and considers interesting questions relating to the proposed method. \n\n## Notation\n- At points the authors refer to $p(\\theta, D| \\mathcal{M})$ as the target distribution. It seems to me that $p(\\theta|D, \\mathcal{M})$ is the target distribution (of course these are the same up to a normalizing constant). \n- Punctuation is missing from some equations (11, 12).\n\n## Remarks\n- I found the justification for the additive noise model slightly unclear. While it seems to have precedent in the literature, it would be nice if the authors could be more explicit about why neglecting the multiplicative term is reasonable. Also is the noise known to satisfy sufficient criteria for the CLT to apply? It seems there is at least some weak dependence between sampled points due unless mini-batches are sampled with replacement. I ask the latter question more out of curiousity, since the precedent for this model in the literature seems quite strong.\n- I found the citation to Robbins and Munro in line 244 confusing. If I understand correctly, in your proof it is needed that the step size sequence is not in $\\ell_2$, where Robbins and Munro require that it is not in $\\ell_1$ but is in $\\ell_2$. Please clarify if this is not correct. The placement of this citation led me to think you were claiming that Robbins and Munro require that the sequence of step sizes is in not in $\\ell_2$.\n- The description of data generating for BLR is confusing. It sounded to me as though $X$ and $Y$ are generated independently, but glancing at the code provided, it seems that $Y$ was generated from the model (i.e. by multiplying $X$ by a random weight vector). \n\n## Questions\n- Does minimizing the variational variant of the bound (eqn 14) with respect to variational parameters $\\phi$ have an interpretation as appproximate inference? For example, is the difference between the log marginal likelihood and this bound a divergence between $q_\\phi$ and the posterior?\n- Is there a technical obstacle to extending your proof to $\\gamma > 0$, or is the choice $\\gamma=0$ made only for simplicity?\n- What is meant by \"optimal\" in theorem 1? Is this matching a known lower bound? Or is it simply indicating that $c=1/4$ leads to the tightest bound for any $c$?\n- Could the authors clarify what is meant by equation 42 in the appendix? Is the integral over $R^{d\\times d}$ as suggested by $d\\Sigma$, or is it over some parameter that parameterizes the matrix $\\Sigma$? (This may be standard notation in some areas, if so I would be happy for the authors to point to a reference on the subject). ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors address limitations of their work well including:\n- Limitations in the stochastic setting\n- Potential memory limitations in finite precision\n- Experiments are not claimed to be state of the art, but illustrate the method and theory well\n\nThe authors claim that since the contribution is methodological and theoretical, determining any direct societal impact would be difficult. This seems reasonable for this paper. \n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "5 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper suggests a method for using annealed importance sampling (AIS) together with Hamiltonian Monte Carlo (HMC). The central idea is that if the accept/reject step is dropped, once can still define a ratio of augmented distributions that is tractable to estimate. The advantage of dropping the accept/reject step is that the estimator is now differentiable. A variational objective is defined in terms of this estimator, and reparameterization is used to optimize the objective. One might worry that dropping the accept/reject step could have serious consequences by changing the properties of the estimator. This is addressed, at least for the case of a simple Bayesian linear regression model, by an analysis that shows a convergence rate of 1/\sqrt{K} can be obtained, where K is the number of annealing distributions. Another analysis is given where gradients are corrupted with additive noise. Here, a more negative result is given that even in the limit that K goes to infinity, the variational bound never becomes tight. This suggests that a significant price is paid for computing gradients on subsets of data. Strictly speaking, this negative result also applies only to the Bayesian linear regression setting, however this shows that any stronger guarantee would have to exclude that-typically considered "easiest"-setting.

Generally speaking, all reviewers felt the algorithm was novel, correct, and potentially significant. There were some fairly minor technical concerns about the experiments, and some questions about the intended "scope" of the contribution (e.g. regarding the title). The authors appear to have good answers for these questions that could easily be integrated into the paper for a final submission.

Decision: Accept (Poster)