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