Ergodic Inference: Accelerate Convergence by Optimisation

1 INTRODUCTION . Statistical inference methods in machine learning are dominated by two approaches : simulation and optimisation . Markov chain Monte Carlo ( MCMC ) is a well-known simulation-based method , which promises asymptotically unbiased samples from arbitrary distributions at the cost of expensive Markov simulations . Variational inference ( VI ) is a well-known method using optimisation , which fits a parametric approximation to the target distribution . VI is biased but offers a computationally efficient generation of approximate samples . There is a recent trend of hybrid methods of MCMC and VI to achieve a better balance between computational efficiency and bias . Hybrid methods often use MCMC or VI as an algorithmic component of the other . In particular , Salimans et al . ( 2015 ) proposed a promising modified VI method that reduces approximation bias by using MCMC transition kernels . Another technique reduces the computational complexity of MCMC by initialising the Markov simulation from a pretrained variational approximation ( Hoffman , 2017 ; Han et al. , 2017 ) . Levy et al . ( 2018 ) proposed to improve MCMC using flexible non-linear transformations given by neural networks and gradientbased auto-tuning strategies . In this work , we propose a novel hybrid method , called ergodic inference ( EI ) . EI improves over both MCMC and VI by tuning the hyper-parameters of a flexible finite-step MCMC chain so that its last state sampling distribution converges fast to a target distribution . EI optimises a tractable objective function which only requires to evaluate the logarithm of the unnormalized target density . Furthermore , unlike in traditional MCMC methods , the samples generated by EI from the last state of the MCMC chain are independent and have no correlations . EI offers an appealing option to balance computational complexity vs. bias on popular benchmarks in machine learning . Compared with previous hybrid methods , EI has following advantages : • EI ’ s hyperparameter tuning produces sampling distributions with lower approximation bias . • The bias is guaranteed to decrease as the length of the MCMC chain increases . • By stopping gradient computations , EI has less computational cost than related baselines . We also state some disadvantages of our method : • The initial state distribution in EI ’ s MCMC chain has to have higher entropy than the target . • The computational complexity per simulated sample of EI is in general higher than in VI . 2 BACKGROUND . 2.1 MONTE CARLO STATISTICAL INFERENCE . Monte Carlo ( MC ) statistical inference approximates expectations under a given distribution using simulated samples . Given a target distribution π , MC estimations of an expectation Eπ [ f ( x ) ] are defined as empirical average of the evaluation of f on samples from π . To generate samples from π , we assume that the unnormalized density function π∗ ( x ) can be easily computed . In a Bayesian setting we typically work with π∗ ( x|y ) given by the product of the prior p ( x ) and the likelihood p ( y|x ) , where y denotes observed variables and x denotes the model parameters specifying p ( y|x ) . 2.2 MARKOV CHAIN MONTE CARLO . Markov chain Monte Carlo ( MCMC ) casts inference as simulation of ergodic Markov chains that converge to the target π . The MCMC kernel M ( x′|x ) is characterised by the detailed balance ( DB ) property : π ( x ) M ( x′|x ) = π ( x′ ) M ( x|x′ ) . Given an unnormalised target density π∗ , an MCMC kernel can be constructed in three steps : first , sample an auxiliary random variable r from an auxiliary distribution qφ1 with parameters φ1 ; second , create a new candidate sample as ( x ′ , r′ ) = fφ2 ( xt−1 , r ) , where fφ2 is a deterministic function with parameters φ2 ; finally , accept the proposal as xt = x ′ with probability pMH = min { 0 , π∗ ( x′ ) qφ1 ( r′ ) / [ π∗ ( xt−1 ) qφ1 ( r ) ] } , otherwise duplicate the previous sample as xt = xt−1 . The last step is well known in the literature as the Metropolis-Hastings ( M-H ) correction step ( Robert & Casella , 2005 ) and it results in MCMC kernels that satisfy the DB condition . In the following , we denote the joint MCMC parameters ( φ1 , φ2 ) by φ . If fφ2 does not preserve volume , then it requires a Jacobian correction factor in the ratio in pMH . Hamiltonian Monte Carlo ( HMC ) is a successful MCMC method which has drawn great attention . A few recent works based on this method are Salimans et al . ( 2015 ) ; Hoffman ( 2017 ) ; Levy et al . ( 2018 ) . In HMC , the auxiliary distribution qφ1 is often chosen to be Gaussian with zero-mean and a constant diagonal covariance matrix specified by φ1 . The most common fφ2 in HMC is a numeric integrator called the leapfrog algorithm , which simulates Hamiltonian dynamics defined by log π∗ ( Neal , 2010 ) . The leapfrog integrator requires the gradient of log π∗ and a step size parameter given by φ2 . Given any initial state x0 , MCMC can generate asymptotically unbiased samples x1 : n. For this , MCMC iteratively simulates the next sample xt through the application of the MCMC transition kernel to the previous sample xt−1 . It is well known in the literature that MCMC is computationally demanding ( MacKay , 2002 ; Bishop , 2006 ) . In particular , it is often necessary to run sufficiently long burn-in MCMC simulations to reduce simulation bias . Another drawback of MCMC is sample correlation , which increases the variance of the MC estimator ( Neal , 2010 ) . To avoid strong sample correlation , the common practice in MCMC to tune hyper-parameters manually using sample quality metrics like effective sample size , ( Hoffman , 2017 ; Robert & Casella , 2005 ) , which has been developed into automated gradient-based tuning strategies in recent work ( Levy et al. , 2018 ) . 2.3 MC ESTIMATION USING VARIATIONAL INFERENCE . Variational inference ( VI ) is a popular alternative to MCMC for generating approximate samples from π . Unlike MCMC reducing sample bias by long burn-in simulation , VI casts the sample bias reduction as an optimisation problem , where a parametric approximate sampling distribution P is fit to the target π . In particular , VI optimises the evidence lower bound ( ELBO ) given by LELBO ( P ‖ π∗ ) = Ep [ log π∗ ( x ) ] +H ( P ) , ( 1 ) whereH ( P ) = −Ep [ log p ( x ) ] , also known as the entropy , must be tractable to compute . LELBO ( PT ‖ π∗ ) is a lower bound on the log normalising constant logZ = log ∫ π∗ ( x ) dx . This bound is tight when P = π. therefore , the approximation bias in VI can be defined as the gap between LELBO ( P ‖ π∗ ) and logZ , that is , ∆bias ( P ) = logZ − LELBO ( P ‖ π ) = DKL ( P ‖ π ) ≥ 0 , ( 2 ) where DKL ( P ‖ π ) denotes the Kullback-Leibler ( KL ) divergence . Variational approximations often belong to simple parametric families like the multivariate Gaussian distribution with diagonal covariance matrix . This results in computationally efficient algorithms for bias reduction and sample generation , but may also produce highly biased samples in cases of over-simplified approximation that ignores correlation . Designing variational approximation to achieve low bias under the constraint of tractable entropy and efficient sampling procedures is possible using flexible distributions parameterised by neural networks ( NNs ) ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ) . However , how to design such NNs for VI is still a research challenge . 2.4 HYBRID METHODS AND VARIANTS OF MCMC AND VI . The balance between computational efficiency and bias is a challenge at the heart of all inference methods . MCMC represents a family of simulation-based methods that guarantee low-bias samples at cost of expensive simulations ; VI represents a family of optimisation-based methods that generate high-bias samples at a low computational cost . Many recent works seek a better balance between efficiency and bias by combining MCMC and VI . Salimans et al . ( 2015 ) proposed to reduce variational bias by optimising an ELBO specified in terms of the tractable joint density of short MCMC chains . The idea seems initially promising , but the proposed ELBO becomes looser and looser as the chain grows longer . Caterini et al . ( 2018 ) construct an alternative ELBO for HMC that still has problems since the auxiliary momentum variables are sampled only once at the beginning of the chain , which reduces the empirical performance of HMC . Inspired by contrastive divergence , Ruiz & Titsias ( 2019 ) proposed a novel variational objective function to optimise variational parameters by adding additional term that minimise the KL between a MCMC distribution and variational approximation to reduce variational bias . Hoffman ( 2017 ) and Han et al . ( 2017 ) proposed to replace expensive burn-in simulations in MCMC with samples from pre-trained variational approximations . This approach is effective at finding good initial proposal distributions . However , it does not offer a solution for tuning HMC parameters ( Hoffman , 2017 ) , which are critical for good empirical performance . Another line of research has focused on improving inference using flexible distributions , which are transformed from simple parametric distributions by non-linear non-volume preserving ( NVP ) functions . Levy et al . ( 2018 ) proposed to tune NVP parameterised MCMC w.r.t . a variant of the expected squared jumped distance ( ESJD ) loss proposed by Pasarica & Gelman ( 2010 ) . Song et al . ( 2017 ) proposed a similar auto-tuning for NVP parameterised MCMC using an adversarial loss . 3 ERGODIC INFERENCE . Ergodic inference ( EI ) is motivated by the well-known convergence of MCMC chains ( Robert & Casella , 2005 ) : MCMC chains converge in terms of the total variation ( TV ) distance between the marginal distribution of the MCMC chain and the target π . Inspired by the convergence property of MCMC chains , we define an ergodic approximation PT to π with T MCMC steps as following . Given a parametric distribution P0 with tractable density p0 ( x0 ; φ0 ) parameterlized by φ0 and an MCMC kernel M ( x′|x ; φ ) constructed using the unnormalised target density π∗ and with MCMC hyperparameter φ , an ergodic approximation of π is the marginal distribution of the final state of an T -step MCMC chain initialized from P0 : pT ( xT ; φ0 , φ ) = ∫ T∏ t=1 M ( xt|xt−1 ; φ ) p0 ( x0 ; φ0 ) dx0 : T−1 . ( 3 ) We call φ0 and φ the ergodic parameters of PT . Well known in MCMC literature like ( Robert & Casella , 2005 ; Murray & Salakhutdinov , 2008 ) , the ergodic approximation pT converges to π after every MCMC transition and with sufficiently long chain pT is guaranteed to be arbitrarily close to π with arbitrary φ and φ0 . It is important to clarify that ergodic approximation is different from the modified variational methods like ( Ruiz & Titsias , 2019 ) which only optimise the variational parameters φ0 , but the optimisation objective functions involve MCMC similation . In the following section , we show how EI can tune the ergodic parameters to minimise the bias of PT as an approximation to the target π with finite T .

Summary: This paper proposes a new combination of Markov chain Monte Carlo (MCMC) and variational inference (VI) for improving approximate inference. The main contribution is the optimization objective that allows improving the quality of samples obtained from the combination of VI and MCMC. Specifically, the authors minimize the "approximate" version of the Kullback-Leibler (KL) divergence between the distribution of MCMC + VI and the true distribution. The authors validate the effectiveness of their formulation through experiments on 6 synthetic benchmarks and generative modeling of MNIST (experiments on Bayesian neural networks are also provided in the appendix).