Differentiable Annealed Importance Sampling and the Perils of Gradient Noise

1 Introduction . 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 . 2 Background . 2.1 Marginal Likelihood Estimation . 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 . 2.2 Annealed Importance Sampling . 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 →∞ . 3 Differentiable Annealed Importance Sampling . 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 .

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