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: The paper proposes UHA, an original and clever modification of Hamiltonian AIS. The idea is to change the AIS ratio such that it becomes tunable, but can still be computed in closed form (ratios of momentum distributions). This is achieved by dropping the accept-reject step in the HMC update such that a single UHA step only involves (re)sampling momenta and leapfrog integration. This modification allows the authors to tune parameters such as the proposal distribution, the step size and damping coefficients (used in momentum resampling). Many interesting ideas for future extensions are proposed. Experiments on various models demonstrate that UHA obtains a tighter ELBO than other combinations of VI and HMC.