MCMC Variational Inference via Uncorrected Hamiltonian Annealing

Given an unnormalized target distribution we want to obtain approximate samples from it and a tight lower bound on its ( log ) normalization constant logZ . Annealed Importance Sampling ( AIS ) with Hamiltonian MCMC is a powerful method that can be used to do this . Its main drawback is that it uses non-differentiable transition kernels , which makes tuning its many parameters hard . We propose a framework to use an AIS-like procedure with Uncorrected Hamiltonian MCMC , called Uncorrected Hamiltonian Annealing . Our method leads to tight and differentiable lower bounds on logZ . We show empirically that our method yields better performances than other competing approaches , and that the ability to tune its parameters using reparameterization gradients may lead to large performance improvements . 1 Introduction . Variational Inference ( VI ) [ 4 , 41 , 45 ] is a method to do approximate inference on a target distribution p ( z ) = p̄ ( z ) /Z that is only known up to the normalization constant Z . The basic insights are , first , that the evidence lower bound ( ELBO ) Eq ( z ) [ log p̄ ( z ) − log q ( z ) ] lower-bounds logZ and , second , that maximizing the ELBO is equivalent to minimizing the KL-divergence from q to p. The simplest VI method chooses a parameterized family for q and optimizes its parameters to maximize the ELBO . A recent direction involves combining VI with Markov chain Monte Carlo ( MCMC ) [ 34 , 43 ] . These methods can be seen as an instance of the auxiliary VI framework [ 2 ] – they create an augmented variational distribution that represents all intermediate random variables generated during the MCMC procedure . An augmented target distribution that attempts to capture the inverse MCMC dynamics is optimized jointly with this variational distribution . However , it has been observed that capturing inverse dynamics is challenging [ 43 , §5.4 ] ( further discussion in Section 4 ) . Annealed Importance Sampling ( AIS ) [ 22 , 27 ] is a powerful technique used to build augmented distributions without the need of learning inverse dynamics . While it was originally proposed to estimate expectations using importance sampling , it can be easily used to build lower bounds on normalization constants of intractable densities [ 18 , 44 ] . AIS creates a sequence of densities that bridge from a tractable initial approximation q to the target p̄ . Then , the augmented variational distribution is given by a sequence of MCMC kernels targeting each bridging density , while the augmented target uses the reversals of those kernels . It turns out that the ratio of these augmented distributions can be computed using only evaluations of the bridging densities . Combining Hamiltonian MCMC kernels with AIS has been observed to produce strong lower bounds [ 35 , 44 ] . However , these bounds are sensitive to numerous parameters , such as the initial distribution , bridging schedule , and parameters of the MCMC kernels . It would be desirable to optimize these parameters 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . to tighten the bound . Unfortunately , the presence of Metropolis-Hastings acceptance steps means that the the final estimator is non-differentiable , and thus reparameterization gradients can not be used . In this work , we propose Uncorrected Hamiltonian Annealing ( UHA ) , a differentiable alternative to Hamiltonian AIS . We define an augmented variational distribution using Hamiltonian MCMC kernels , but dropping the accept-reject steps . This is motivated by the fact that Hamiltonian dynamics sometimes have high acceptance rates . Since these uncorrected MCMC kernels do not exactly hold the bridging densities invariant , an augmented target distribution can not be defined in terms of reversals . Instead , we define our augmented target by deriving an algorithm for the exact reversal of the original ( corrected ) MCMC kernel and dropping the accept-reject step . Surprisingly , this yields a very simple expression for the resulting lower bound . We use reparameterization gradients to tune various parameters involved in the lower bound produced by UHA , including the initial approximation q , parameters of the uncorrected MCMC kernel , and the bridging densities . Experimentally , tuning all these leads to large gains . For example , in several inference tasks we observe that tuning UHA with K = 64 bridging densities gives better results than traditional Hamiltonian AIS with K = 512 . Finally , we use UHA to train VAEs [ 24 , 31 ] . In this case we observe that using UHA leads to higher ELBOs . In addition , we observe that increasing the number of bridging densities with UHA consistently leads to better results , and that for a large enough number of bridging densities the variational gap ( difference between ELBO and true log-likelihood ) becomes small , and models with higher log-likelihood are obtained . 2 Preliminaries . Variational inference and augmentation . Suppose that p ( z ) = 1Z p̄ ( z ) is some target density , where p̄ is unnormalized and Z = ∫ p̄ ( z ) dz is the corresponding normalizer , and let ELBO ( q ( z ) , p̄ ( z ) ) = E q ( z ) log p̄ ( z ) q ( z ) ( 1 ) be the `` ELBO operator '' . Variational inference ( VI ) is based on the fact that for any q ( z ) we have [ 4 ] logZ = ELBO ( q ( z ) , p̄ ( z ) ) + KL ( q ( z ) ‖p ( z ) ) . ( 2 ) In VI , the parameters of q are tuned to maximize the `` evidence lower bound '' ( ELBO ) . Since the KL-divergence is non-negative , this is always a lower bound on logZ . Also , maximizing the ELBO is equivalent to minimizing the KL-divergence from q to p. To get tighter bounds and better approximations recent work has made use of augmented distributions [ 2 , 21 ] . Let z1 : M = ( z1 , · · · , zM ) and suppose that p̄ ( z1 : M ) = p̄ ( zM ) p ( z1 : M−1|zM ) augments the original target density while preserving its normalization constant . Then , for any q ( z1 : M ) we have logZ = ELBO ( q ( z1 : M ) , p̄ ( z1 : M ) ) + KL ( q ( z1 : M ) ‖p ( z1 : M ) ) . ( 3 ) The first term is called the `` augmented '' ELBO and again lower bounds logZ . By the chain rule of KL-divergence [ 12 ] , the KL-divergence from q to p over z1 : M upper-bounds the KL-divergence over zM . This justifies using the marginal of q over zM to approximate the original target distribution . Annealed Importance Sampling . A successful approach for creating augmented distributions is Annealed Importance Sampling ( AIS ) [ 27 ] . It creates an augmented proposal distribution q by applying a sequence of transition densities Tm ( zm+1|zm ) , and an augmented target by defining transition densities Um ( zm|zm+1 ) . This gives the augmented densities q ( z1 : M ) = q ( z1 ) M−1∏ m=1 Tm ( zm+1|zm ) and p̄ ( z1 : M ) = p̄ ( zM ) M−1∏ m=1 Um ( zm|zm+1 ) . ( 4 ) Naively , the ratio of these densities is p̄ ( z1 : M ) q ( z1 : M ) = p̄ ( zM ) q ( z1 ) M−1∏ m=1 Um ( zm|zm+1 ) Tm ( zm+1|zm ) . ( 5 ) To define the transitions Tm and Um , AIS creates a sequence of unnormalized densities π̄1 , . . . , π̄M−1 that “ bridge ” from a starting distribution q to the target p̄ , meaning that π̄1 is close to q and π̄M−1 is close to p̄ . Then , for each intermediate distribution , Tm ( zm+1|zm ) is chosen to be a Markov kernel that holds πm invariant , and Um to be the reversal of Tm with respect to πm , defined as Um ( zm|zm+1 ) = T ( zm+1|zm ) πm ( zm ) πm ( zm+1 ) . ( 6 ) This choice produces a simplification so that eq . 5 becomes p̄ ( z1 : M ) q ( z1 : M ) = p̄ ( zM ) q ( z1 ) M−1∏ m=1 π̄m ( zm ) π̄m ( zm+1 ) . ( 7 ) This can be easily evaluated without needing to evaluate the transition densities . The ratio from eq . 7 can be used to get an expression for the lower bound ELBO ( q ( z1 : M ) , p̄ ( z1 : M ) ) . Research has shown that the AIS augmentation may lead to extremely tight lower bounds [ 18 , 17 , 35 , 44 ] . Hamiltonian Dynamics . Many MCMC methods used to sample from p ( z ) are based on Hamiltonian dynamics [ 3 , 8 , 29 , 42 ] . The idea is to create an augmented distribution p ( z , ρ ) = p ( z ) S ( ρ ) , where S ( ρ ) is a distribution over a momentum variable ρ ( e.g . a Multivariate Gaussian ) . Then , one can define numerical integration schemes where z and ρ evolve while nearly holding p ( z , ρ ) constant . When corrected by a Metropolis-Hastings acceptance step , this can be made to exactly hold p ( z , ρ ) invariant . This is alternated with a scheme that resamples the momentum ρ while holding S ( ρ ) invariant . When Hamiltonian dynamics work well , z can quickly move around , suppressing random-walk behavior . There are a variety of different Hamiltonian MCMC methods , corresponding to different integration schemes , momentum distributions , and ways of resampling the momentum . For instance , HMC and Langevin dynamics use the leapfrog integrator , a Gaussian for the momentum variables and a full resampling of the momentum variables at each step [ 29 , 42 ] . On the other hand , if the momentum variables are only partially resampled , the under-damped variants of HMC and Langevin dynamics are recovered [ 29 ] . It was observed that partial resampling may lead to improved perfomance [ 9 ] . It is easy to integrate Hamiltonian dynamics into AIS . First , define an augmented target p̄ ( z , ρ ) = p̄ ( z ) S ( ρ ) and an augmented starting distribution q ( z , ρ ) = q ( z ) S ( ρ ) . Then , create a series of augmented densities π̄1 ( z , ρ ) , . . . , π̄M−1 ( z , ρ ) bridging the two as π̄m ( z , ρ ) = π̄m ( z ) S ( ρ ) . Finally , define the forward transition Tm ( zm+1 , ρm+1|zm , ρm ) to be an iteration of a Hamiltonian MCMC method that leaves πm ( z , ρ ) invariant . We will describe a single transition Tm as a sequence of three steps : ( 1 ) resample the momentum ; ( 2 ) simulate Hamiltonian dynamics and apply an accept-reject step ; and ( 3 ) negate the momentum . The precise process that defines the transition is shown in Alg . 1 . Note that this algorithm is quite general , and compatible with HMC , Langevin dynamics and their underdamped variants ( by selecting an appropriate integrator and resampling method ) . Algorithm 1 Corrected Tm ( zm+1 , ρm+1|zm , ρm ) 1 . Sample ρ′m from some s ( ρ ′ m|ρm ) that leaves S ( ρ ) invariant . Set z′m ← zm . 2 . Simulate Hamiltonian dynamics as ( z′′m , ρ′′m ) ← Tm ( z′m , ρ′m ) . Calculate an acceptance probability α = min ( 1 , π̄m ( z′′m , ρ′′m ) /π̄m ( z′m , ρ′m ) ) . With probability α , set ( z′′′m , ρ′′′m ) ← ( z′′m , ρ′′m ) . Otherwise , set ( z′′′m , ρ′′′m ) ← ( z′m , ρ′m ) . 3 . Reverse the momentum as ( zm+1 , ρm+1 ) ← ( z′′′m , −ρ′′′m ) . return ( zm+1 , ρm+1 ) Representing Tm this way makes it easy to show it holds the density πm ( z , ρ ) invariant . The overall strategy is to show that each of the steps 1-3 holds πm invariant , and so does the composition of them [ 29 , §3.2 ] . For steps 1 and 3 this is trivial , provided that S ( ρ ) = S ( −ρ ) . For step 2 , we require that the simulation Tm has unit Jacobian and satisfies T −1m = Tm . Then , Tm can be interpreted as a symmetric Metropolis-Hastings proposal , meaning the Metroplis-Hastings acceptance probability α is as given . A typical choice for Tm that satisfies these requirements is the leapfrog integrator with a momentum reversal at the end . ( This reversal then gets `` un-reversed '' in step 3 for accepted moves . ) Since Tm holds πm invariant , we can define Um as the reversal of Tm wrt πm . Then , eq . 7 becomes p̄ ( z1 : M , ρ1 : M ) q ( z1 : M , ρ1 : M ) = p̄ ( zM , ρM ) q ( z1 , ρ1 ) M−1∏ m=1 π̄m ( zm , ρm ) π̄m ( zm+1 , ρm+1 ) . ( 8 ) Using this ratio we get an expression for the lower bound ELBO ( q ( z1 : M , ρ1 : M ) , p̄ ( z1 : M , ρ1 : M ) ) obtained with Hamiltonian AIS . While this method has been observed to yield strong lower bounds on logZ [ 35 , 44 ] ( see also Section 5.2 ) , its performance depends on many parameters : initial distribution q ( z ) , momentum distribution S , momentum resampling scheme , simulator Tm , and bridging densities . We would like to tune these parameters by maximizing the ELBO using reparameterization-based estimators . However , due to the accept-reject step required by the Hamiltonian MCMC transition , the resulting bound is not differentiable , and thus reparameterization gradients are not available .

Summary: This paper studies the annealed important sampling (AIS) using the Hamiltonian dynamics kernels in each interpolant. To circumvent the problem that the Hamiltonian dynamics involves the non-differentiable operations, this paper proposes to discard the accept-reject operations in the kernel, and derives the resulting density ratios in a simple form related to the momentum variables. The resulting approach, named Uncorrected Hamiltonian Annealing (UHA), is a fully differentiable AIS method. This paper validates the effectiveness of UHA through extensive experiments and highlights the importance of differentiability.