Learning Deep Latent Variable Models via Amortized Langevin Dynamics

1 INTRODUCTION . Latent variable models are widely used for generative modeling ( Bishop , 1998 ; Kingma & Welling , 2013 ) , principal component analysis ( Wold et al. , 1987 ) , and factor analysis ( Harman , 1976 ) . To learn a latent variable model , it is essential to estimate the latent variables , z , from the observations , x. Bayesian inference is a probabilistic approach for estimation , wherein the estimate is represented as a posterior distribution , i.e. , p ( z | x ) = p ( z ) p ( x | z ) /p ( x ) . A major challenge while using the Bayesian approach is that the posterior distribution is typically intractable . Markov chain Monte Carlo ( MCMC ) methods such as Langevin dynamics ( LD ) provide sample approximations for posterior distribution with an asymptotic convergence guarantee . However , MCMC methods converge slowly . Thus , it is inefficient to perform time-consuming MCMC iterations for each latent variable , particularly for large-scale datasets . Furthermore , when we obtain new observations that we would like to perform inference for , we would need to re-run the sampling procedure for them . In the context of variational inference , a method to amortize the cost of datapoint-wise optimization known as amortized variational inference ( AVI ) ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) was recently proposed . In this method , the optimization of datapoint-wise parameters of variational distributions is replaced with the optimization of an inference model that predicts the variational parameters from observations . This amortization enables posterior inference to be performed efficiently on large-scale datasets . In addition , inference for new observations can be efficiently performed using the optimized inference model . AVI is widely used for the training of deep generative models , and such models are known as variational autoencoders ( VAEs ) . However , methods based on variational inference have less approximation power , because distributions with tractable densities are used for approximations . Although there have been attempts to improve their flexibility ( e.g. , normalizing flows ( Rezende & Mohamed , 2015 ; Kingma et al. , 2016 ; Van Den Berg et al. , 2018 ; Huang et al. , 2018 ) ) , such methods typically have constraints in terms of the model architectures ( e.g. , invertibility in normalizing flows ) . 1An implementation is available at : https : //bit.ly/2Shmsq3 Therefore , we propose an amortization method for LD , amortized Langevin dynamics ( ALD ) . In ALD , datapoint-wise MCMC iterations are replaced with updates of an inference model that maps observations into latent variables . This amortization enables simultaneous sampling from posteriors over massive datasets . In particular , when a minibatch training is used for the inference model , the computational cost is constant with data size . Moreover , when inference is performed for new test data , the trained inference model can be used as initialization of MCMC to improve the mixing , because it is expected that the properly trained inference model can map data into the high-density area of the posteriors . We experimentally show that the ALD can accurately perform sampling from posteriors without datapoint-wise iterations . Furthermore , we demonstrate its applicability to the training of deep generative models . Neural networks are used for both generative and inference models to yield Langevin autoencoders ( LAEs ) . LAEs can be easily extended for more flexible generative modeling , in which the latent prior distribution , p ( z ) , is also intractable and defined with unnormalized energy function , by combining them with contrastive divergence learning ( Hinton , 2002 ; Carreira-Perpinan & Hinton , 2005 ) . We refer to this extension of LAEs as contrastive Langevin autoencoders ( CLAEs ) . We experimentally show that our LAEs and CLAEs can generate sharper images than existing explicit generative models , such as VAEs . Moreover , we report their performance of unsupervised anomaly detection . 2 PRELIMINARIES . 2.1 PROBLEM DEFINITION . Consider a probabilistic model with observations x , continuous latent variables z , and model parameters θ , as described by the probabilistic graphical model shown in Figure 1 ( A ) . Although the posterior distribution over the latent variable is proportional to the product of the prior and likelihood : p ( z | x ) = p ( z ) p ( x | z ) /p ( x ) , this is intractable owing to the normalizing constant p ( x ) = ∫ p ( z ) p ( x | z ) dz . This study aims to approximate the posterior p ( z | x ) for all n ob- servations x ( 1 ) , . . .x ( n ) efficiently by obtaining samples from it . 2.2 LANGEVIN DYNAMICS . Langevin dynamics ( LD ) ( Neal , 2011 ) is a sampling algorithm based on the following Langevin equation : dz = −∇zU ( x , z ) dt+ √ 2β−1dB , ( 1 ) where U is a potential function that is Lipschitz continuous and satisfies an appropriate growth condition , β is an inverse temperature parameter , and B is a Brownian motion . This stochastic differential equation has exp ( −βU ( x , z ) ) / ∫ exp ( −βU ( x , z′ ) ) dz′ as its equilibrium distribution . We set β = 1 and define the potential as follows to obtain the target posterior p ( z | x ) as its equilibrium : U ( x , z ) = − log p ( z ) − log p ( x | z ) . ( 2 ) Algorithm 1 Amortized Langevin dynamics ( training time ) φ← Initialize parameters Z ( 1 ) , . . . , Z ( n ) ← ∅ . Initialize sample sets for all n datapoints repeat φ← φ′ ∼ N ( φ′ ; φ− ηφ ∑n i=1∇φU ( x ( i ) , z ( i ) = fz|x ( x ( i ) ; φ ) ) , 2ηφI ) Z ( 1 ) , . . . , Z ( n ) ← Z ( 1 ) ∪ { fφ ( x ( 1 ) ) } , . . . , Z ( N ) ∪ { fφ ( x ( n ) ) } . Add samples until convergence of parameters return Z ( 1 ) , . . . , Z ( n ) Algorithm 2 Amortized Langevin dynamics ( test time ) z ← fz|x ( x ; φ∗ ) . Initialize a sample using a trained inference model Z← ∅ . Initialize a sample set repeat z ← z′ ∼ N ( z′ ; z − η∇zU ( x , z ) , 2ηI ) . Update the sample using traditional LD Z← Z ∪ { z } . Add samples until convergence of parameters return Z We can obtain samples from the posterior by simulating Eq . ( 1 ) using the Euler–Maruyama method ( Kloeden & Platen , 2013 ) as follows : z ← z′ ∼ N ( z′ ; z − η∇zU ( x , z ) , 2ηI ) , ( 3 ) where η is the step size for the discretization . When the step size is sufficiently small , the samples asymptotically move to the target posterior by repeating this sampling iteration . LD can be applied to any posterior inference problems for continuous latent variables provided the potential energy is differentiable on the latent space . However , to obtain samples of the posterior p ( z | x ) for all observations x ( 1 ) , . . .x ( n ) , we should perform an iteration on Eq . ( 3 ) per datapoint as shown in Figure 1 ( B1 ) . It is inefficient particularly if the dataset is large . In the next section , we demonstrate a method that addresses the inefficiency by amortization . 3 AMORTIZED LANGEVIN DYNAMICS . In traditional LD , we perform MCMC iterations for each latent variable per datapoint . This is inefficient particularly if managing massive datasets . As an alternative to performing the simulation of latent dynamics directly , we define an inference model , fz|x , which is a differentiable mapping from observations into latent variables , and consider the dynamics of its parameter φ as follows : dφ = − n∑ i=1 ∇φU ( x ( i ) , z ( i ) = fz|x ( x ( i ) ; φ ) ) + √ 2dB . ( 4 ) Because function fz|x outputs latent variables , the stochastic dynamics on the parameter space induces other dynamics on the latent space and is represented as the total gradient of fz|x : dz ( i ) = dimφ∑ k=1 ∂z ( i ) ∂φk dφk = − dimφ∑ k=1 ∂z ( i ) ∂φk ( ∂ ∂φk U ( x ( i ) , fz|x ( x ( i ) ; φ ) ) ) dt − dimφ∑ k=1 ∂z ( i ) ∂φk n∑ j=1 , j 6=i ∂ ∂φk U ( x ( j ) , fz|x ( x ( j ) ; φ ) ) dt+√2 dimφ∑ k=1 ∂z ( i ) ∂φk dB . ( 5 ) MSE ESS ( 1 std . ) MSCE LD 0.00145 570.24 ( 25.25 ) 0.00136 ALD 0.00233 634.27 ( 40.94 ) 0.00151 Table 1 : Quantitative comparison of the sample quality between traditional LD and our ALD . The mean squared error ( MSE ) between the true mean and the sample average , the effective sample size ( ESS ) , and the Monte Carlo standard error ( MSCE ) are provided as evaluation metrics . The first term of Eq . ( 5 ) approximates −∇z ( i ) U ( x ( i ) , z ( i ) ) dt in Eq . ( 1 ) , and the remaining terms introduce a random walk behavior to the dynamics as in the Brownian term of Eq . ( 1 ) . For the simulation of Eq . ( 4 ) , we use the Euler–Maruyama method , as in traditional LD : φ← φ′ ∼ N ( φ′ ; φ− ηφ n∑ i=1 ∇φU ( x ( i ) , z ( i ) = fz|x ( x ( i ) ; φ ) ) , 2ηφI ) , ( 6 ) where ηφ is the step size . Through the iterations , the posterior sampling is implicitly performed by collecting outputs of the inference model for all datapoints in the training set as described in Algorithm 1 . When we perform inference for new test data , the trained inference model can be used as initialization of a MCMC method ( e.g. , traditional LD ) as shown in Algorithm 2 , because it is expected that the trained inference model can map data into the high-density area of the posteriors . For minibatch training , we can substitute the minibatch statistics of m datapoints for the derivative for all n data in Eq . ( 6 ) : n∑ i=1 ∇φU ( x ( i ) , z ( i ) = fz|x ( x ( i ) ; φ ) ) ≈ n m m∑ i=1 ∇φU ( x ( i ) , z ( i ) = fz|x ( x ( i ) ; φ ) ) . ( 7 ) In this case , we refer to the algorithm as stochastic gradient amortized Langevin dynamics ( SGALD ) . SGALD enables us to sample from posteriors of a massive dataset with a constant computational cost . By contrast , performing traditional LD requires a linearly increasing cost with data size . For minibatch training of LD , adaptive preconditioning is known to be effective to improve convergence , which is referred to as preconditioned stochastic gradient Langevin dynamics ( pSGLD ) ( Li et al. , 2015 ) . This preconditioning technique is also applicable to our SGALD , and we employ it throughout our experiments . Figure 2 shows a simple example of sampling from a posterior distribution , where its prior and likelihood are defined using conjugate bivariate Gaussian distributions ( see Appendix F for more details ) . ALD produces samples that match well the shape of the target distributions . The mean squared error ( MSE ) between the true mean and the sample average , the effective sample size ( ESS ) , and the Monte Carlo standard error ( MSCE ) are provided for quantitative comparison , as shown in Table 1 . It can be observed that the sample quality of ALD is competitive to standard LD , even though ALD does not perform direct update of samples in the latent space . Figure 3 shows the evolution of obtained sample values by traditional LD and our SGALD for posteriors defined by a simple univariate conjugate Gaussian ( see Appendix F.2 for more experimental details ) . SGALD ’ s samples converges much faster than traditional LD . The advantage of our ALD over amortized variational inference ( AVI ) is the flexibility of posterior approximation . Figure 4 is an example where the likelihood p ( x | z ) is defined using a neural network , therefore the posterior p ( z | x ) is highly multimodal . AVI methods typically approximate posteriors using variational distributions , which have tractable density function ( e.g. , Gaussian distributions ) . Hence , their approximation power is limited by the choice of variational distribution family , and they often fail to approximate such complex posteriors . On the other hand , ALD can capture well such posteriors by obtaining samples . The results in other examples are summarized in Appendix F .

Summary: In this paper, the author presented an advanced autoencoder framework LAE. Instead of element-wise MCMC, LAE collected samples from the posterior using the amortized Langevin dynamics of a potential energy distribution. In CLAE, an extended version of LAE, the author used an intractable energy function as the prior, and collected samples using its Langevin function. The author claims that LAE and CLAE are more efficient in large scale data and have better performance compared with traditional autoencoders and variational autoencoders.