Ergodic Measure Preserving Flows

Training probabilistic models with neural network components is intractable in most cases and requires to use approximations such as Markov chain Monte Carlo ( MCMC ) , which is not scalable and requires significant hyper-parameter tuning , or mean-field variational inference ( VI ) , which is biased . While there has been attempts at combining both approaches , the resulting methods have some important limitations in theory and in practice . As an alternative , we propose a novel method which is scalable , like mean-field VI , and , due to its theoretical foundation in ergodic theory , is also asymptotically accurate , like MCMC . We test our method on popular benchmark problems with deep generative models and Bayesian neural networks . Our results show that we can outperform existing approximate inference methods . 1 INTRODUCTION . Approximate statistical inference with unnormalised density functions is fundamentally important problem both Bayesian and frequentist inference . In particular , the successes of many sophisticated generative models in machine learning rely on power inference algorithms . Markov chain Monte Carlo ( MCMC ) methods and variational inference ( VI ) , originally developed in statistics and physics , are two most important approximate inference methods in machine learning , which has been widely used in all kinds of probabilistic models , like latent topic models ( Blei et al. , 2003 ) , Boltzmann machines ( Hinton , 2002 ; Salakhutdinov & Larochelle , 2010 ) , Bayesian non-parametric models ( Neal , 2000 ; Kurihara et al. , 2007 ) . However , they are facing great challenges in the recent research on probabilistic modelling with deep neural networks ( NNs ) . In particular , Bayesian deep neural networks become popular in recent works , because it exploits the Bayesian framework to overcome the overfitting and data demanding problems in deep learning ( Neal , 2012 ) . Another interesting research direction is to use deep neural networks in latent variable models to transform the simple latent random variables into complex distribution , which is also known as to Deep Generative Models ( DGM ) . DGMs have been proved to be very powerful generative models . Inspired by DGMs , many recent works on variational inference adopted NN to construct flexible approximate distributions . In particular , variational autoencoders ( Kingma & Welling , 2014 ) and normalising flows ( Rezende & Mohamed , 2015 ) are two most influential works in this direction . However , due to lack of understanding of the convergence of specific NNs , the research of NNbased inference is focused on engineering the architecture of inference NNs based on experiments and heuristics . In this work , we propose a novel approximate inference method based on the classic inference theory of MCMC . Our method is inspired by the idea of parallel simulations of MCMC and the recent advances in variational inference with NNs . Like these variational methods , it is straightforward to accelerate the computation of our method using parallelised simulations on Graphical Processing Units . More importantly , with solid theoretical foundations in the theory of MCMC , the proposed method guarantees asymptotic convergence to arbitrary distributions of interest . It is a great advantage over variational inference , because of the approximation bias in variational methods . Our method is also attractive to a wide range of probabilistic models without NNs and Bayesian NNs . 2 BACKGROUND . 2.1 BAYESIAN NEURAL NETWORKS . Given data D = { xn , yn } Nn=1 formed by feature vectors xn and corresponding scalar targets yn , we can assume that each yn is obtained as yn = f ( xn ; θ ) + n , where f ( · ; θ ) is the output of a deep neural network with weights θ and the n are independent noise variables with n ∼ N ( 0 , σ2 ) . This model specifies a likelihood function p ( y1 , . . . , yn|x1 , . . . , xn , θ ) which can be combined with a Gaussian prior on θ to obtain a posterior distribution p ( θ|D ) . Predictions for the y ? corresponding to a new feature vector x ? are then obtained by using the predictive distribution p ( y ? |x ? , D ) =∫ p ( y ? |x ? , θ ) p ( θ|D ) dθ . However , integrating with respect to the posterior distribution p ( θ|D ) is intractable and approximations have to be performed in practice , with the most popular methods for this being VI and MCMC . 2.2 DEEP GENERATIVE MODELS . Generative models extract intrinsic structure from data by making use of latent variables . LetD be a dataset with n data points { xn } Nn=1 . Given a latent representation z , the data point x is assumed to be sampled from the conditional distribution pθ ( x|z ) , which is specified in terms of some parameters θ . This conditional distribution is often refered to as the decoder . Given a prior distribution p ( z ) over the latent variables , the joint distribution of data and latent variables is pθ ( x , z ) = p ( z ) pθ ( x|z ) . The marginal probability of data x under the model is then pθ ( x ) = ∫ p ( z ) pθ ( x|z ) dz . Until recenlty , pθ ( x|z ) was typically specified using a simple distributional family , e.g . generalized linear models ( Murphy , 2012 ; Bishop , 2006 ) . However , more recently , deep generative models ( DGMs ) use deep neural networks with weights θ to specify the decoder ( Kingma & Welling , 2014 ; Goodfellow et al. , 2014 ) . Maximum likelihood is a straightforward way to train probabilistic models with latent variables . However , the marginal likelihood pθ ( x ) is intractable to compute in DGMs and approximations are needed in practice . As before , the most popular methods for this are based on VI and MCMC . We briefly describe these methods in the following sections . 2.3 VARIATIONAL INFERENCE AND NORMALIZING FLOWS . VI approximates a complex posterior distribution with another simpler parametric distribution which is found by optimizing a lower bound on the marginal likelihood . Let the complex posterior be pθ ( z|x ) , with associated marginal likelihood p ( x ) , and let qφ ( z ) , parameterized by φ , be a simpler tractable distribution . The lower bound of the marginal likelihood is then defined as log pθ ( x ) ≥ Eqφ [ log pθ ( x , z ) − log qφ ( z ) ] , ( 1 ) which is often known as the evidence lower bound ( ELBO ) . The more flexible the parametric family qφ , the better the approximation quality to the true posterior and the tighter the value of the ELBO . Mean-field VI uses a form for qφ ( z ) which assumes independence between random variables . This reduces the computational cost of the optimization problem but often leads to poor performance with complex posterior distributions , such as the ones arising in DGMs or in Bayesian neural networks . Amortization can be used to accelerate convergence and reduce computational cost when multiple inference problems have to be solved simultaneously . The optimization of ( 1 ) can be amortized by making qφ depend explicitly on the data x . In this case , qφ ( z ) is replaced with qφ ( z|x ) , where φ are now the weights of a neural network which computes the parameteres of a tractable parametric distribuion on z from x . In this manner , for any new value of x , we can readily obtain a correspoding variational approximation given by qφ ( z|x ) . Variational auto-encoders ( VAEs ) ( Kingma & Welling , 2014 ) are DGMs trained by using mean-field VI with a Gaussian parametric distribuion and amortization . Rezende & Mohamed ( 2015 ) improve over this method by using a more flexible variational family called normalizing flows ( NFs ) . The NF family is obtained by applying L invertible non-linear transformations f1 , . . . , fL to a random variable z0 with tractable density q0 ( z0 ) and exact simulation . The resulting output is a random variable zL = fL ◦ · · · ◦ f1 ( z0 ) with density qφ ( zL|x ) = q0 ( z0|x ) ∏L l=1 |∂fl ( zl−1 ) /∂zl−1| −1 and with φ being the parameters of f1 , . . . , fL . Stochastic gradient descent ( SGD ) , in combination with the reparameterisation trick ( Kingma & Welling , 2014 ) , can be used for the scalable optimization of the ELBO in VAEs and NFs . However , the main limitation of VAEs and NFs is the bias present in their variational approximations . This bias can be quite high , even in the case of NFs , since the transformations f1 , . . . , fL have to be rather simple to ensure invertibility and to reduce computational costs . 2.4 MARKOV CHAIN MONTE CARLO . Markov chain Monte Carlo ( MCMC ) is an approximate inference method which does not have the aforementioned bias problem . MCMC works by simulating a stationary Markov chain that generates asymptotically unbiased samples from the distributions of interest . Formally , a Markov chain is a sequence of random variables z0 , z1 , . . . in which the transition from state zl to zl+1 is defined by the conditional probability distribution of zl+1 given zl , denoted by K ( zl , zl+1 ) . Markov chains in MCMC methods have the following strong stationary property : if one state zl of the chain follows the stationary distribution π , so does the next state zl+1 . In particular , π ( zl+1 ) = ∫ π ( zl ) K ( zl , zl+1 ) dzl . ( 2 ) If zl follows a distribution πl that is different from the stationary one , then the distribution of zl+1 πl+1 ( zl+1 ) = ∫ πl ( zl ) K ( zl , zl+1 ) dzl , ( 3 ) is guaranteed to be closer to π than πl . This property implies that , with sufficiently many transitions , the distribution πl of zl converges to π irrespectively of the distribution the initial state z0 . Despite beign asymptotically unbiased , MCMC methods are less popular than VI for two reasons . First , they are computationally more expensive and second , they typically include hyper-parameters in the Markov kernel K which are highly problem dependent and are hard to tune in practice . 3 ERGODIC MEASURE PRESERVING FLOWS . In this section , we describe an inference method that combines the strengths of MCMC and VI and avoids their drawbacks . The idea is to use the output distribution of a MCMC chain , given by ( 3 ) , as the variational distribution and optimize a simple to evaluate objective function for tuning MCMC parameters . Since MCMC converges to the target asymptotically , our variational approximation can be arbitrarily accurate . The result is a computationally efficient method which avoids the bias of parametric approximations and which can do automatic tuning of hyper-parameters . 3.1 DEFINITIONS . Given the target distribution π with unnormalised density function π∗ , we define an approximate distribution q by a mixture of sequential deterministic transformations that preserves the measure π∗ . We call such approximate distributions measure preserving flows ( MPFs ) . The transformations preserving a given measure π∗ are formally defined as follows ( Billingsley , 1986 ) . Definition 3.1 . Measure Preserving Transformation ( MPT ) . Let ( Ω , F , P ) be a probability space and µ be a consistent measure with P . A mapping T : Ω→ Ω is a measure preserving transformation if T is measurable in both the input filed F and the output field F and µ ( A ) = µ ( T−1 ( A ) ) for allA ∈ F . If T is a one-to-one mapping onto Ω , then T preserves µ : µ ( A ) = µ ( T−1TA ) = µ ( TA ) . An example of MPT is any transformation with Jacobian determinant equal to 1 , which preserves the Lebesgue measure . In practice , it is straightforward to verify if a transformation T preserves the measure with density function π∗ with the following conditions : ( i ) Bijection : T is invertible , ( ii ) Preservation of density function : π∗ ( z ) = π∗ ( T ( z ) ) for all z . ( iii ) Preservation of base measure : the Jacobian determinant is one if the Lebesgue measure is the base measure . In probability theory , MPTs are often used within the area of ergodic stochastic processes since many of these processes can be reformulated as a composition of MPTs . In particular , MCMC kernels are MPTs and stationary Markov chains in MCMC are ergodic processes ( Robert & Casella , 2005 ) . The joint probability of states in a MCMC chain is q ( z0 , z1 , . . . , zL ) = q0 ( z0 ) ∏L l=1K ( zl−1 , zl ) , where q0 is the distribution of the initial state z0 . The density of the last state zL is then obtained by integrating out all the previous states : qL ( zL ) = ∫ q0 ( z0 ) L∏ l=1 K ( zl−1 , zl ) dz0 dz1 . . . dzL−1 . ( 4 ) If the Markov chain is ergodic , qL converges to the stationary distribution π in total variation distance as the length L of the chain increases ( Robert & Casella , 2005 ) . We define a measure preserving flow ( MPF ) as a representation of ( 4 ) in which the kernelK becomes a deterministic transformation Tr : Z → Z with stochastic auxiliary input r following distribution µ . We can then define zL as the result of applying these deterministic transformations to z0 , that is , zL = TrL ◦ · · · ◦ Tr1 ( z0 ) . By following the rule of changing variables , it is then straightforward to derive the density of zL as qL ( zL ) = ∫ q ( zL , r1 : L ) dr1 : L = ∫ q ( z0 ) µ ( r1 : L ) δ ( zL − TrL ◦ · · · ◦ Tr1 ( z0 ) ) dz0 dr1 : L , ( 5 ) where δ denotes the Dirac delta function . Note that there is no Jacobian term in ( 5 ) because of the preservation of the Lebesgue measure . Because MPFs are equivalent to ergodic Markov chains , the density obtained at the output of an MPF , that is , qL , will converge to the stationary distribution π as L increases . Hamiltonian Monte Carlo ( HMC ) is one of the most successful MCMC methods , which also can be interpreted as an MPF . Given the target random variable z ∈ Rd with unnormalised density function π∗ , the HMC kernel is essentially applying a deterministic transformationH to the previous state zi with an auxiliary random variable r ∈ Rn with the density function µ ( · ) . The transformation H is given by the solution of Hamiltonian dynamics ż ( t ) = ∂rK ( r ) , ṙ ( t ) = −∂zU ( z ) , ( 6 ) where ż denotes the derivative of z w.r.t . the time of the dynamics t , U ( z ) = − log π∗ ( z ) and K ( r ) = − logµ ( r ) . The preservation of total Hamiltonian energy , H ( z , r ) = U ( z ) + K ( r ) , is the characteristic property of Hamiltonian dynamics , which can be easily verified by Ḣ ( z , r ) = 0 using ( 6 ) . It is straightforward to see that H preserves the joint measure π ( z ) µ ( r ) . In particular , H is a bijective transformation because the dynamics are deterministic and time reversible ( Neal , 2010 ) and the preservation of the Hamiltonian energy implies the preservation of density . Finally , the volume preservation of H in the space of ( z , r ) is a well known property of Hamiltonian dynamics , which can be proved by Liouville ’ s theorem ( Leimkuhler & Reich , 2004 ; Neal , 2010 ) . We can write the marginal distribution of the last sample generated by HMC as an MPF : q ( zL ) = ∫ q ( zL , r1 : L ) dr1 : L = ∫ q ( z0 ) µ ( r ) δ ( zL −HrL ◦ · · · ◦ Hr1 ( z0 ) ) dz0 dr1 : L . ( 7 ) We call the MPF generated by Hamiltonian dynamics a Hamiltonian MPF ( HMPF ) .

Summary: This paper proposes training latent variable models (as in VAE decoders) by running HMC to approximate the posterior of the latents, and then estimating model parameters by maximizing the complete data log-likelihood. This is not a new idea by itself and is used e.g. as a baseline in Kingma and Welling's original VAE paper. The novelty in this paper is that it proposes tuning the parameters of the HMC inference algo by maximizing the likelihood achieved by the final sample in the MCMC chain. This seems to work well in practice and might be a useful method, but it is not clear under what conditions it should work.