Auxiliary Variational MCMC

1 INTRODUCTION . Markov Chain Monte Carlo ( MCMC ) and Variational Inference ( VI ) are well-established approaches to approximating expectations under the complex distributions p that frequently arise in machine learning and statistics ( Wainwright & Jordan , 2008 ; Brooks et al. , 2011 ) . VI is usually fast and cheap but often places strong restrictions on the class of approximating distributions leading to irreducible bias . MCMC on the other hand is asymptotically exact but may require an infeasible amount of computation to converge . Given the complimentary strengths and weaknesses of the two methods it is natural to wish to combine them ( De Freitas et al. , 2001 ; Salimans et al. , 2015 ) . The most naive pairing of the two methods , simply using a variational approximation q as the proposal distribution in a Metropolis-Hastings sampler ( see for example Gamerman & Lopes ( 2006 ) or De Freitas et al . ( 2001 ) ) , is known to scale poorly with the dimension of the target distribution ( De Freitas et al. , 2001 ) and mix inefficiently when p is multi-modal . In this paper we suggest an alternative approach inspired both by the successes of black box variational inference ( Ranganath et al. , 2014 ) and auxiliary variable MCMC methods , such as Hamiltonian Monte Carlo ( HMC ) ( Duane et al. , 1987 ; Girolami & Calderhead , 2011 ) . The key contributions of our work are : • A general framework for marrying variational inference with Markov Chain Monte Carlo in a way likely to produce efficient samplers . • The use of the auxiliary variational method to capture latent low-dimensional structure in our target distributions and exploit this structure to suppress random walk behaviour . • The extension of the Metropolis-Hastings algorithm to continuous mixture proposals . • The introduction and demonstration of a specific instance of our framework , the Aux- iliary Variational Sampler ( AVS ) . Our sampler takes advantage of flexible distributions parameterized by neural networks that can be trained in a fully black-box manner . 2 AUXILIARY VARIATIONAL MCMC . The key idea behind Auxiliary Variational MCMC is to exploit structure present in the target distribution p ( x ) by first fitting a parameterized variational approximation in an augmented space . This allows the sampler to leverage learned low-dimensional structure . In the subsequent sections we introduce the auxiliary variational method and describe how it can be combined with a carefully chosen class of proposal distributions to construct an efficient sampler . 2.1 MIXTURE PROPOSAL MCMC . To develop a valid MCMC algorithm we need to construct an ergodic Markov chain whose stationary distribution is our target distribution p ( x ) . In order to do this we introduce a Metropolis-Hastings ( Gamerman & Lopes , 2006 ) like algorithm with a specially chosen form of proposal distribution . This proposal can be naturally combined with the auxiliary variational method which we will introduce in section 2.2 . We first consider a mixture proposal distribution of the following form1 : q̃ ( x′|x ) = ∫ q̃ ( x′|a ) q̃ ( a|x ) da ( 1 ) We prove in supplement ( 7 ) that the following forms a valid MCMC sampling step : 1 . Sample a from q̃ ( a|x ) 2 . Sample x′ from q̃ ( x′|a ) 3 . Accept the candidate sample x′ with probability min { 1 , q̃ ( x|a ) q̃ ( a|x′ ) p ( x′ ) q̃ ( x′|a ) q̃ ( a|x ) p ( x ) } ( 2 ) otherwise reject x′ and define the new sample as a copy of the current x , namely x′ = x . It is worth noting that the above procedure is not equivalent to simply performing Metropolis-Hastings in the joint ( x , a ) space , but is a sampler in x alone , specified by the marginalized proposal distribution given by ( 1 ) . The two become equivalent only in the case that q̃ ( a|x ) is constant , see supplement ( 7 ) . The mixture proposal can be extended to an arbitrary number of auxiliary variables as long as the acceptance ratio is adjusted accordingly . For example with two auxiliary variables , the proposal becomes : q̃ ( x′|x ) = ∫ q̃ ( x′|a′ ) q̃ ( a′|a ) q̃ ( a|x ) dada′ ( 3 ) where q̃ ( a′|a ) is another arbitrary proposal distribution . The acceptance probability is now given by : min { 1 , p ( x′ ) q̃ ( x|a′ ) q̃ ( a′|a ) q̃ ( a|x′ ) p ( x ) q̃ ( x′|a ) q̃ ( a|a′ ) q̃ ( a′|x ) } ( 4 ) For a proposal q̃ ( a′|a ) that is symmetric , q̃ ( a′|a ) = q̃ ( a|a′ ) , this simplifies to : min { 1 , p ( x′ ) q̃ ( x|a′ ) q̃ ( a|x′ ) p ( x ) q̃ ( x′|a ) q̃ ( a′|x ) } ( 5 ) We now show how the mixture proposal can be naturally combined with the auxiliary variational method . 2.2 THE AUXILIARY VARIATIONAL METHOD . The auxiliary variational method ( Agakov & Barber , 2004 ; Ranganath et al. , 2016 ) is a strategy for creating more expressive families of approximating distributions for use in variational inference . Instead of minimizing the Kullback-Leibler divergence between our approximating distribution q ( x ) and our target p ( x ) , we instead minimize the divergence in an augmented space ( x , a ) with additional auxiliary variables a . In doing this , we first define a joint qφ ( x , a ) in the augmented space and a joint p ( x , a ) = pθ ( a|x ) p ( x ) where θ and φ are parameters . Marginalizing p ( x , a ) over a recovers p ( x ) by construction . We are free to decide on the dimension and form ( e.g . continuous , discrete or mixed ) for a and our objective is the Kullback-Leibler divergence between the joint approximation q and joint target p : ( φ∗ , θ∗ ) = arg min φ , θ KL ( qφ ( x , a ) ||p ( x ) pθ ( a|x ) ) ( 6 ) 1In general we use q̃ to denote a proposal distribution and q to denote a variational distribution . Algorithm 1 : Auxiliary variational sampler Initialize φ and θ . Fit the variational distribution while Not Converged do { a1 , . . . , aN } ∼ qφ ( a ) { x1 , . . . , xN } ∼ qφ ( x|an ) . re-parameterized L← 1 N ∑ n log qφ ( xn|an ) qφ ( an ) p ( xn ) pθ ( an|xn ) φ← φ− η∇φL θ ← θ − η∇θL end a ∼ qφ ( a ) . Mixture-Model MCMC sampling x0 ∼ qφ ( x|a ) for t = 0 , . . . , T do a ∼ pθ ( a|xt ) a′ ∼ N ( a′ a , σ2aI ) x′ ∼ qφ ( x|a′ ) A← p ( x ′ ) pθ ( a ′|x′ ) qφ ( xt|a ) p ( xt ) pθ ( a|xt ) qφ ( x′|a′ ) . acceptance ratio u ∼ U [ 0 , 1 ] . sample u from uniform distribution if u ≤ A then xt+1 = x ′ else xt+1 = xt end end The Auxiliary Variational Sampler . The algorithm begins by fitting a mixture variational distribution to the target p ( x ) by stochastic gradient descent , based on the auxiliary variational method . After the auxiliary variational model is fitted to p ( x ) the MCMC sampling phase begins and is a modified form of Metropolis-Hastings sampling . Moving to the joint space allows us to tractably learn complex approximating distributions qφ ( x , a ) whose marginal , qφ ( x ) = ∫ qφ ( x|a ) qφ ( a ) da , may be intractable . Whilst sampling in high-dimensional spaces is difficult , in many cases of interest the high probability or typical regions lie close to much lower dimensional manifolds . A key point of fitting an auxiliary variational distribution is that , by constraining a to have dimension much lower than x , we can then exploit this low dimensional manifold to form a more efficient proposal distribution , staying close to the manifold of significant probability . 2.3 COMBINING AUXILIARY VARIATIONAL INFERENCE AND MCMC . After fitting our variational approximation to p ( x ) , we will have three variational distributions : q∗ ( x|a ) , q∗ ( a ) and p∗ ( a|x ) that are the solution of the optimization problem given in ( 6 ) . These distributions approximately satisfy the following relationships : ∫ p ( x ) p∗ ( a|x ) dx ≈ q∗ ( a ) , ∫ q∗ ( x|a ) q∗ ( a ) da ≈ p ( x ) ( 7 ) which become exact iff the divergence in ( 6 ) becomes zero . Here p∗ ( a|x ) is a learned stochastic mapping from the high-dimensional target space to the low-dimensional auxiliary space and q∗ ( x|a ) is a mapping in the opposite direction . These learned mappings can be composed in a variety of ways to form marginal proposal distributions of the kind discussed in section 2.1 and also joint proposals . We now discuss some natural combinations and argue that these may form good proposal distributions in practice . 2.3.1 NAIVE PROPOSAL DISTRIBUTIONS . A simple proposal constructed from our variational distribution is to perform Metropolis-Hastings in the joint ( x , a ) space with an independent proposal given by : q̃ ( x′ , a′|x , a ) = q∗ ( x′|a′ ) q∗ ( a′ ) ( 8 ) If the variational approximation is accurate such that q ( x , a ) ≈ p ( x ) p ( a|x ) then one might expect this to have high acceptance probability . A potential downside of this scheme is that proposals are independent between time-steps and in high dimensions this may cause the acceptance ratio to become impractically low . Another natural MCMC proposal distribution to consider is : q̃ ( x′|x ) = ∫ q∗ ( x ′|a ) p∗ ( a|x ) da ( 9 ) where we have replaced the arbitrary proposal distribution of ( 1 ) with the optimal variational distributions learned by minimizing the joint divergence given in ( 6 ) . We might expect this to be a promising proposal distribution , with high acceptance probability , as it already approximately satisfies the stationarity criterion of our MCMC chain . That is∫ q̃ ( x′|x ) p ( x ) dx = ∫ q∗ ( x ′|a ) p∗ ( a|x ) p ( x ) dxda ≈ ∫ q∗ ( x ′|a ) q∗ ( x|a ) q∗ ( a ) dxda = ∫ q∗ ( x ′|a ) q∗ ( a ) da ≈ ∫ p∗ ( a|x′ ) p∗ ( a ) da = p ( x′ ) However , if the variational distribution truly captures underlying structure in p ( x ) , it has been our experience that p∗ ( a|x ) and q∗ ( x|a ) become approximate inverses causing q̃ ( x′|x ) to resemble the identity mapping and thereby slowing mixing in the chain .

Summary: This paper proposes an auxiliary variable MCMC scheme involving variational inference for efficient MCMC. Given a target distribution p(x), the authors introduce an auxiliary variable a, and learn conditional distributions p(a|x) and q(a|x) by minimizing the KL divergence between p(x)p(a|x) and q(a)q(x|a), with q(a) something simple (the authors use Gaussian). A MH proposal step involves simulating x givea the current MCMC sample x (from p(a|x), taking a step in A-space, and then returning back to the X space (using q(x|a)).  The authors show how to calculate the acceptance probability.