Differentiable Annealed Importance Sampling and the Perils of Gradient Noise
Year: 2021

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 MetropolisHastings 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.

heading: 1 Introduction
text: Marginal likelihood (ML), sometimes called evidence, is a central quantity in Bayesian learning as it measures how well a 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 generative models [Wu et 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 ML 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 is 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 a non-differentiable Metropolis-Hastings (MH) correction. Unfortunately, the non-differentiability 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 (NeurIPS 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 a 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 a subset of the dataset. Given the success of stochastic optimization [Robbins and Monro, 1951, Bottou and Bousquet, 2011] 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 particular, 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-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 as 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 a large margin.

heading: 2 Background

heading: 2.1 Marginal Likelihood Estimation
text: For a modelM and observed data D, one can define the marginal likelihood (ML) as p(D|M) = ∫ p(D,θ|M)dθ = ∫ p(D|θ,M)p(θ|M)dθ, (1)
where θ 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(θ) = f(θ)/Z where the unnormalized density f(θ) can be efficiently computed, we are interested in estimating the partition function Z = ∫ f(θ)dθ. Here, f(θ) corresponds to p(D,θ|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.
It 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(D|M) = Eq(θ) [ p(D|θ,M)p(θ|M)
q(θ)
] ≈ 1 S S∑ i=1 p(D|θi,M)p(θi|M) q(θi) with θi ∼ q(θ) (2)
However, this estimation can exhibit high variance for small or medium S when the target distribution p(D,θ|M) and proposal distribution q(θ) are dissimilar.

heading: 2.2 Annealed Importance Sampling
text: Annealed importance sampling (AIS) is an algorithm which estimates the ML by gradually changing, or “annealing”, a distribution. Formally, the algorithm takes in a sequence of distributions p0, . . . , pK , with pk(θ) = fk(θ)/Zk and Zk = ∫ fk(θ)dθ. In the context of ML estimation, the starting
distribution f0 is the tractable prior distribution p(θ|M) with Z0 = 1, while the target distribution fK is p(D,θ|M) with ZK = p(D|M). For each pk, one must also specify an MCMC transition operator Tk which leaves pk invariant.
The output of AIS is an unbiased estimate ẐK of the exact ML ZK . Importantly, unbiasedness holds for any finite K, as shown in Neal [2001]. Moreover, AIS can be viewed as importance sampling over an extended space [Neal, 2001]. In particular, we have ZK = Eqfwd [qbwd/qfwd] with the target and proposal distributions defined as
qfwd(θ0:K) = p0(θ0)T1(θ1|θ0) · · · TK(θK |θK−1) (3) qbwd(θ0:K) = fK(θK)T̃K(θK−1|θK) · · · T̃1(θ0|θ1), (4)
where Tk(θ|θ′) is a forward MCMC kernel and T̃k(θ′|θ) = Tk(θ|θ′)pk(θ′)/pk(θ) is the corresponding reverse kernel. Here, qfwd represents the chain of states generated by AIS, and qbwd is a fictitious (unnormalized) reverse chain which begins with a sample from pK and applies the transitions in reverse order. In practice, the intermediate distributions have to be chosen carefully for a low variance estimate ẐK . One typically uses geometric averages of the initial and target distributions:
pk(θ) = pβk(θ) = fβk(θ)/Zβk = f0(θ)1−βkfK(θ)βk/Zβk = p(θ|M)p(D|θ,M)βk/Zβk (5)
where 0 = β0 < β1 < · · · < βK = 1 is the annealing schedule. Indeed, AIS gives an unbiased estimate Ẑ of Z . However, as Z can vary over many orders of magnitude, it is often more meaningful to talk about estimating logZ . Unfortunately, unbiased estimators ofZ can result in biased estimators of logZ because E log Ẑ ≤ logEẐ by Jensen’s inequality, resulting in only a lower bound. In particular, we have the AIS bound
Eqfwd log ẐK = K∑ k=1 Eqfwd [ log fβk(θk−1)− log fβk−1(θk−1) ] (6)
= K∑ k=1 (βk − βk−1)Eqfwd [log p(D|θk−1,M)] (7)
where (5) facilitated the simplification from (6) to (7). Of course, it is not enough to have a lower bound; we would also like the estimates to be close to the true value. Fortunately, AIS is consistent in that the estimate log Ẑ converges to the correct value in the limit of infinitely many intermediate distributions [Neal, 2001], under the very idealized assumption of perfect transitions (i.e. that each transition Tk generates an exact sample from pk, independent of the previous state). To give some intuition, the bound (6) can be simplified as
Eqfwd log ẐK = logZK − K∑ k=1 DKL(pk−1 ‖ pk), (8)
and the sum of the KL divergence terms diminishes as K →∞.

heading: 3 Differentiable Annealed Importance Sampling
text: Algorithm 1 Differentiable AIS (DAIS)
θ0,v0 sample from p0(θ), π , N (0,M) LDAIS = − log p0(θ0) for k = 1, . . . ,K do θk− 12 ← θk−1 + η 2M
−1vk−1 v̂k ← vk−1 + η∇ log fβk(θk− 12 ) θk ← θk− 12 + η 2M −1v̂k
vk ← γv̂k + √ 1− γ2ε, ε ∼ N (0,M)
LDAIS += log π(v̂k))− log π(vk−1) end for return LDAIS += log p(D,θK |M)
In this section, we motivate and derive a differentiable AIS (DAIS) algorithm for marginal likelihood (ML) estimation. We also discuss its application to variational inference for a tighter ELBO and a memory-efficient implementation.
Ideally, assuming a continuously parameterized model class (e.g., variational autoencoder), we would like to differentiate through the lower bound (7) to find an optimal modelM. However, AIS must be instantiated with an MCMC transition kernel Tk that satisfies detailed balance to ensure it leaves pk invariant. In practice, this is typically achieved by using a MH step,
which is generally not differentiable.1 We thus remove the MH correction and, in particular, specify each transition to consist of a deterministic leapfrog integration step followed by a stochastic partial momentum refreshment [Horowitz, 1991]. Algorithm 1 details the simulation of Hamiltonian dynamics using such transitions. With γ = 0, the algorithm resemables unadjusted Langevin dynamics [Roberts et al., 1996] but computes the ML bound on the fly. In practice, choosing 0 < γ < 1 (γ = 0.9 is a common default) helps avoid random walk behavior and accelerates mixing [Neal, 2011, Chen et al., 2014]. Importantly, we retain the formalism of performing importance sampling on an extended space despite the loss of detailed balance.2 To show this, we can define the forward and (unnormalized) backward distributions as
qfwd(θ0:K ,v0:K) = p0(θ0)π(v0)T1(θ1,v1|θ0,v0) · · · TK(θK ,vK |θK−1,vK−1) (9) qbwd(θ0:K ,v0:K) = fK(θK)π(vK)T̃K(θK−1,vK−1|θK ,vK) · · · T̃1(θ0,v0|θ1,v1) (10)
where the transition operator Tk(θk,vk|θk−1,vk−1) = T ′k(θk, v̂k|θk−1,vk−1)T ′′k (vk|v̂k) 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 1 in the reverse order, which leads to a surprisingly simple expression for our estimator. In particular, the backward transition operator is defined by T̃k(θk−1,vk−1|θk,vk) = T ′′k (v̂k|vk)T ′k(θk−1,vk−1|θk,−v̂k). Note that we need to flip the sign of v̂k in the reverse chain to account for time reversal. As a consequence of the above definitions, we have
T ′′k (vk|v̂k) = T ′′k (v̂k|vk)π(vk)/π(v̂k). (11) This is because T ′′k (vk|v̂k) = N (γv̂k, (1 − γ2)M) and T ′′k (v̂k|vk) = N (γvk, (1 − γ2)M). Furthermore, since T ′k is a deterministic leapfrog update, it is reversible and volume preserving, so we have T ′k(θk−1,vk−1|θk,−v̂k) = T ′k(θk, v̂k|θk−1,vk−1). With this, we can derive the DAIS bound:
LDAIS = Eqfwd [log qbwd(θ0:K ,v0:K)− log qfwd(θ0:K ,v0:K)]
= Eqfwd [ log p(D,θK |M)− log p0(θ0) + K∑ k=1 log π(v̂k) π(vk−1) ] .
(12)
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 π(vk) is simply the negative kinetic energy plus a 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 is 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 is 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 to estimate the volume of the support of the posterior distribution.

heading: 3.1 Differentiable Annealed Variational Inference
text: DAIS can be applied to variational inference for a tighter bound; we name this differentiable annealed variational inference (DAVI). 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 a 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 ≡ Eqφ [log p(D,θ|M)− log qφ(θ)] (13) 1The discontinuity introduced by the MH step makes it hard to use the reparameterization trick, though this can be done using a delicate gradient estimator [Naesseth et al., 2017]. This is orthogonal to our work and our fix is simpler and easier for us to analyze.
2This is true even with mini-batch gradients.
However, the lower bound can be quite loose if the variational posterior family qφ(θ) 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 qφ rather than p0:
qfwd(θ0:K ,v0:K) = qφ(θ0)π(v0)T1(θ1,v1|θ0,v0) · · · TK(θK ,vK |θK−1,vK−1). (14)
We also define associated intermediate distributions pk(θ) = qφ(θ)1−βkp(D,θ|M)βk . This gives a new lower bound:
LDAVI ≡ Eqfwd [ log p(D,θK |M)− log qφ(θ0) + K∑ k=1 log π(v̂k) π(vk−1) ] . (15)
We can maximize this lower bound over model parameters of M, all parameters of AIS (e.g., annealing schedule βk) 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 involves storing the entire sequence of sampled states θ0,v0, . . . ,θK ,vK . This can be problematic in cases when K is large due to the large memory overhead. However, DAIS is compatible with the idea of reversible 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 εk for momentum refreshment. The exact noise samples can also be computed in reverse if one uses a deterministic and reversible scheme (e.g. the linear congruential 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 et al., 2021].
However, as discussed by Maclaurin et al. [2015], reversible learning with finite arithmetic precision requires some storage to counteract compounding round-off error. For γ 6= 0 (γ = 0.9 is a common default), we need on average log2(1/γ) bits per parameter per step, which is still small compared to naive storage. We defer further exposition on memory-efficient DAIS to Appendix D. We remark that reversible learning is a potentially crucial property of DAIS as it affords some degree of scaling to longer chain lengths and, indirectly, bigger models.

heading: 4 Convergence Analysis for Bayesian Linear Regression
text: 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 account for the time required for the samples to 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 full-batch and mini-batch regimes.
In particular, we focus on the Bayesian linear regression setting and adopt the following model:
prior: θ ∼ N (µp,Λ−1p )
likelihood: y ∼ N (Xθ, σ2I)⇒ θ ∼ N (µ∗,Λ−1lld ) where Λlld = X>X σ2 and µ∗ = (X >X)−1X>y
posterior: θ ∼ N (µpos,Λ−1pos) where µpos = Λ−1pos(Λpµp + Λlldµ∗) and Λpos = Λp + Λlld
with X ∈ Rn×d denoting the input features and y ∈ Rn×1 the targets. We choose Bayesian linear regression because it enables us to analyze the dynamics analytically in a 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 an identity mass matrix without loss of generality because we can absorb
M into the input matrix X in Algorithm 1) and obtain the following update rule (see Appendix B.1 for derivation):
θk ← ( I− η 2 k
2 Λβkpos
) θk−1 + ( ηkI−
η3k 4 Λβkpos
) vk−1 +
η2k 2 Λβkposµ βk pos
v̂k ← −ηkΛβkposθk−1 + ( I− η 2 k
2 Λβkpos
) vk−1 + ηkΛ βk posµ βk pos
(16)
where Λβkpos = Λp + βkΛlld and µ βk pos = (Λ βk pos)−1(Λpµp + βkΛlldµ∗). With these iterative updates, we can compute the expectation and covariance of θk and vk at any time k, which suffices to compute the lower bound in closed-form.

heading: 4.1 Sublinear Convergence in the Full-Batch Setting
text: 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 θk is also Gaussian in the form of N (µk,Σk). We can compute the gap between the log ML and our lower bound in closed-form (see Appendix B.2 for derivation):
log p(D)− LDAIS =
1 2 ‖µK − µpos‖2Λpos︸ ︷︷ ︸
1
+ 1
2 Tr(ΛposΣK)−
d
2︸ ︷︷ ︸ 2
+ 1 2 log |Σpos| |Σp| − Eq [ K∑ k=1 log π(v̂k) π(vk−1) ] ︸ ︷︷ ︸
3
(17)
where d is the feature dimension. Here, 1 and 2 measure the error of last-iterate Markov chain convergence and will both vanish as long as µK → µpos and ΣK → Σpos. We will show later that they converge with a rate of O( 1η2K ). The key is to show that µk (resp. Λk) lags behind µ βk pos (resp. Λβkpos) with roughly 1 η2 steps. Formally, we have the following.
Lemma 1. Given equally spaced βk, running DAIS with γ = 0 and η ∼ 1Kc where c ≥ 1 4 yields
‖µk−1 − µβkpos‖2 = O(K2c−1), ‖Λk−1 −Λβkpos‖2 = O(K2c−1). (18)
We remark that the assumption of βk being equally spaced is not essential and can be relaxed as long as they are chosen by a scheme that leads to βk − βk−1 going down approximately in inverse proportion to K. In addition, we note that the assumption of full momentum refreshment is for convenience and we believe a similar result holds for γ > 0.
Importantly, this lemma implies that both 1 and 2 vanish sublinearly if we choose c < 12 . The analysis of error term 3 is more nuanced. In particular, this error could either come from using transitions for each of these intermediate distributions that do not bring the distribution close to equilibrium, or from using a finite number of distributions to anneal from p0 to pK . Surprisingly, the error 3 decays as fast as the other two terms if the step size scales as 1/Kc with c ≥ 14 . In summary, we have the following theorem. Theorem 1. Given equally spaced βk, running DAIS with γ = 0 and η ∼ 1Kc where c ≥ 1 4 yields
log p(D)− LDAIS = O(K2c−1).
With c = 14 , we have the convergence rate O(1/ √ K).
We remark that with perfect transitions, the requirement of c ≥ 1/4 is not necessary and we can achieve O(1/K) convergence, as also shown in Grosse et al. [2013]. The gap between O(1/ √ K) and O(1/K) highlights the importance of considering convergence to the stationary distribution.

heading: 4.2 Inconsistency in the Stochastic Setting
text: 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, a consistent/convergent algorithm in the deterministic setting readily implies a 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 model3 ∇̃ log fk(θ) = ∇ log fk(θ) + ε. This model is 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:
θk ← ( I− η 2 k
2 Λβkpos
) θk−1 + ( ηkI−
η3k 4 Λβkpos
) vk−1 +
η2k 2 Λβkposµ βk pos + η2k 2 ε
v̂k ← −ηkΛβkposθk−1 + ( I− η 2 k
2 Λβkpos
) vk−1 + ηkΛ βk posµ βt pos + ηkε
(19)
where stochastic noise ε has variance lowered bounded by Σε. Further, we letµvk = E[v̂k] and Σvk be the covariance of v̂k. Surprisingly, we find that DAIS is incompatible with stochastic gradients, even though it admits the same importance sampling interpretation if the reverse transition T̃k conditions on the same mini-batch of data as Tk. We summarize the result in the following theorem. Theorem 2. For stochastic DAIS with full momentum refreshment (γ = 0 in Algorithm 1) and any stepsize scheme, we have
lim inf K→∞
|log p(D)− LDAIS| > 0. (20)
Here, we give some intuition why DAIS fails in the stochastic setting. To ensure convergence of θK to N (µpos,Σpos), a major requirement is for the step sizes to satisfy limK→∞ ∑K k=1 η 2 k = ∞ [Robbins and Monro, 1951]. However, the randomness of mini-batching sampling would contribute to the variance of v̂k. In particular, we have the following recursion:
Σ̃vk = η 2 kΛ βk posΣ̂k−1Λ βk pos +
( I− η 2 k
2 Λβkpos
)2 + η2kΣε. (21)
For notational convenience, we let Σ̂vk , η 2 kΛ βk posΣ̂k−1Λ βk pos + (I − η
2 k
2 Λ βk pos)2. Here, we used Σ̃vk ,
Σ̂vk and Σ̂k to avoid confusion with Σ v k and Σk in the full-batch setting. In this case, if we follow Stephan et al. [2017] and Chen et al. [2014] in assuming ε is Gaussian,4 we have
Eq [ K∑ k=1 log π(v̂k) π(vk−1) ] = K∑ k=1 [ −1 2 ‖µvk‖22 − 1 2 Tr(Σ̂vk) + d 2 ] − K∑ k=1 [ 1 2 η2kTr(Σε) ] . (22)
The second term of (22) goes to infinity as limK→∞ ∑K k=1 η 2 k =∞. Intuitively, the gradient noise adds to the kinetic energy, and the size of this contribution is proportional to η2k. Since this effect is cumulative over all K steps, ηk has to be reduced at least as 1/ √ K for the kinetic energy term
to go down. However, this contradicts the requirement that limK→∞ ∑K k=1 η 2 k = ∞, needed for last-iterate convergence. In summary, the convergence of θK requires us to make sure the sum of step sizes goes to infinity, which in turn results in a non-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.
3In 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.
4This assumption is non-essential, though Gaussian noise is reasonable by the central limit theorem.

heading: 5 Related Works
text: 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 a 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 a broader family of techniques for partition function estimation, all based on the following identity from statistical physics: logZK − logZ0 =∫ 1 0 Eθ∼pβ [ d dβ log fβ(θ) ] dβ. In particular, the weight update in AIS can be seen as a finite difference approximation. In comparison, thermodynamic integration (TI) [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 recently5, 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.
In the context of variational inference, many papers have also investigated tighter lower bounds for the log ML. Burda et al. [2016] proposed a strictly 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 is 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 VAE, which improves Hamiltonian variational inference [Salimans et 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 of as a normalizing flow scheme in which the flow depends explicitly on the target distribution. Along this line, Le et al. [2018], 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, 2011, Chen et al., 2014, Ma et 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.

heading: 6 Simulations
text: In this section, we discuss the experiments used to validate our algorithm and theory. Importantly, we do not aim to achieve state-of-the-art on these tasks.

heading: 6.1 Bayesian Linear Regression
text: 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 ∈ Rn×d and targets y ∈ Rn respectively consist of entries sampled from N (0, 0.01) and N (0, 1). In particular, we choose n = 10, 000 and d = 10 for our simulations (the results are qualitatively same
5The paper appeared on arXiv after we submitted the paper.
with different n and d). In addition, we set the observation variance σ2 = 1. For convenience, we set the linear annealing scheme βk = kK .
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 1 is tight. In addition, we observe in Figure 1c that the gap fails to vanish with mini-batch gradients for all step-size scaling schemes. Interestingly, with c = 1/2, the gap stays constant. This matches our predictions that the deterministic error decays as O(K2c−1) = O(1) while the stochastic error is proportional to ∑K k=1 η 2 k = O(K2c−1) = O(1).

heading: 6.2 Variational Autoencoder
text: We compare the performance of 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 a 50-dimensional standard Gaussian distribution. The conditional distributions p(xi|z) are independent Bernoulli, with the decoder parameterized by two hidden layers, each with 200 tanh units. The variational posterior q(z|x) is also a 50-dimensional Gaussian with diagonal covariance, whose mean and variance are 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 it with importance sampling (as done in IWAE) with S samples in DAVI. To be specific, we define the bound as follows:
log 1
S S∑ i=1
( pθ(x, z i K)
qφ(zi0|x) K∏ k=1 π(v̂ik) π(vik−1)
) , (23)
where we sample (zi0,v i 0, v̂ i 1, . . . ) independently from qfwd. By default, we use partial momentum refreshment with γ = 0.9 and equally spaced annealing parameters βk = k/K. We compare it to vanilla VAE and IWAE bounds with S ×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 S ×K is small. On the other hand, DAVI is more effective with more compute budget (i.e., S ×K is large) and eventually outperforms IWAE.
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, one can see that learning the annealing scheme improves the performance slightly.
Lastly, we also compare our algorithm with IWAE, AIS, and Hamiltonian AIS (HAIS)6 [SohlDickstein and Culpepper, 2012] in evaluating the log-likelihood of trained models. Notably, IWAE and AIS have been widely used in VAE evaluation, see 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 1 with c = 1/4 and only tune the step size for the case of K = 10. For all implementation details, please see Appendix C.2. In particular, we compare the evaluation algorithms on models trained using the VAE, IWAE, and DAIS objectives. In Figure 2, we report the estimated negative log-likelihood as a function of the number of particles (for IWAE) or gradient updates (for AIS, HAIS and DAIS). Interestingly, we observe that IWAE performs better when we have limited computation and AIS/HAIS/DAIS win out by a big margin if 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 in Neal [2001]). This may explain the superior performance of 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.

heading: 7 Conclusion
text: In this paper, we proposed a 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 a sublinear convergence rate of DAIS in the full-batch setting. However, we showed that DAIS is inconsistent when mini-batch gradients are used due to a 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.

heading: Acknowledgements
text: We thank Ricky Tian Qi Chen, Murat A. Erdogdu, Sergey Levine, Xuechen Li, Mufan Li, Chris J. Maddison, and Matthew D. Hoffman for many helpful discussions. In particular, we thank Xuechen Li for his suggestion regarding a memory-efficient implementation. We also thank Shengyang Sun, Xuechen Li and Yangjun Ruan for detailed comments on early drafts, and Radford M. Neal for insightful feedback on our arXiv preprint. Finally, we thank the anonymous NeurIPS reviewers for their useful feedback on earlier versions of this manuscript.
GZ was supported by Ontario Graduate Fellowship. KH was supported by a 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.
6HAIS is an algorithm similar to ours which combines Hamiltonian dynamics with AIS, but with accept-reject steps included.

Review summaries:
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.   The 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.  A 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.

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. 

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.