L ANGEVIN A UTOENCODERS FOR L EARNING
D EEP
L ATENT V ARIABLE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Markov chain Monte Carlo (MCMC), such as Langevin dynamics, is valid for approximating intractable distributions.
However, its usage is limited in the context of deep latent variable models since it is not scalable to data size owing to its datapoint-wise iterations and slow convergence.
This paper proposes the amortized Langevin dynamics (ALD), wherein datapoint-wise MCMC iterations are entirely replaced with updates of an inference model that maps observations into latent variables.
Since it no longer depends on datapoint-wise iterations, ALD enables scalable inference from large-scale datasets.
Despite its efﬁciency, it retains the excellent property of MCMC; we prove that ALD has the target posterior as a stationary distribution under some assumptions.
Furthermore, ALD can be extended to sampling from an unconditional distribution such as an energy-based model, enabling more ﬂexible generative modeling by applying it to the prior distribution of the latent variable.
Based on ALD, we construct a new deep latent variable model named the Langevin autoencoder (LAE).
LAE uses ALD for autoencoder-like posterior inference and sampling from the latent space EBM.
Using toy datasets, we empirically validate that ALD can properly obtain samples from target distributions in both conditional and unconditional cases, and ALD converges signiﬁcantly faster than traditional LD.
We also evaluate LAE on the image generation task using three datasets (SVHN, CIFAR-10, and CelebA-HQ).
Not only can LAE be trained faster than non-amortized MCMC methods, but LAE can also generate better samples in terms of the Fr´echet Inception Distance (FID) compared to AVI-based methods, such as the variational autoencoder 1 .
I NTRODUCTION
Variational inference (VI) and Markov chain Monte Carlo (MCMC) are two practical tools to approximate intractable distributions.
Recently, VI has been dominantly used in deep latent variable models (DLVMs) to approximate the posterior distribution over the latent variable
z given the observation x , i.e., p ( z | x ) .
At the core of the success of VI is the invention of amortized variational inference (AVI) (Kingma et al., 2014;
An & Cho, 2015; Su et al., 2018; Eslami et al., 2018; Kumar et al., 2018), which replaces optimization of datapoint-wise variational parameters with an inference model that predicts latent variables from observations.
The advantage of AVI over traditional VI is that minibatch training can be used for its optimization, which enables efﬁcient posterior inference on large-scale datasets. In addition, we can also leverage the optimized inference model to perform inference for new data in test time.
However, the approximation power of AVI (or VI itself) is limited because it relies on distributions with tractable densities for approximations.
Although there have been attempts to improve their ﬂexibility (e.g., normalizing ﬂows (Rezende & Mohamed, 2015; Kingma et al., 2016; Van Den Berg et al., 2018; Huang et al., 2018)), such methods typically have architectural constraints (e.g., invertibility in normalizing ﬂows).
Compared to VI, MCMC can approximate complex distributions because it does not rely on any tractable distributions.
Instead, MCMC repeats sampling from the target distribution and uses obtained samples to approximate the posterior distribution.
Langevin dynamics is a typical example of
MCMC for sampling from a continuous distribution.
However, despite its high approximation ability, MCMC has received relatively little attention in learning DLVMs.
It is because MCMC methods [Figure] take a long time to converge, making it difﬁcult to be used in the training of DLVMs.
When learning
DLVMs with MCMC, we need to run MCMC iterations for sampling from each posterior per data point, i.e., p (cid:0) z | x ( i ) (cid:1) ( i = 1 , . .
. , n ), where n is the number of training data, as shown in Figure 1 (B-1).
It is problematic, mainly when training with a large-scale ( n > 10K ) dataset, because it is time-consuming to run massive MCMC iterations for all data points.
Furthermore, we need to re-run the sampling procedure when we obtain new observations in test time.
As in VI, there have been some attempts to introduce the concept of amortized inference to MCMC.
For example, Hoffman (2017) initializes MCMC sampling using an inference model that predicts latent variables from observations.
However, as they use inference models only for the initialization of MCMC, these methods still rely on datapoint-wise sampling iterations.
Not only is it timeconsuming, but implementations of such partially amortized methods also tend to be complicated compared to the simplicity of AVI.
To make MCMC more suitable for the inference of DLVMs, a more straightforward and sophisticated framework of amortization is needed.
This paper proposes the amortized Langevin dynamics (ALD), which replace datapoint-wise MCMC iterations with updates of an inference model that maps observations into latent variables (Figure(B-2)).
Since latent variables depend on the inference model, the updates of the inference model can be regarded as implicit updates of latent variables, which enables us to perform posterior inference without datapoint-wise MCMC iterations.
Notably, our ALD treats outputs of the inference model themselves as samples from the target distribution, whereas existing amortization methods use the outputs only as initialization of MCMC.
Therefore, ALD can be implemented straightforwardly like
AVI.
Moreover, despite the simplicity, we can theoretically guarantee that ALD has the true posterior as a stationary distribution under some assumptions, which is a critical requirement for valid MCMC algorithms.
Although we have introduced ALD as a posterior sampling algorithm, the application of ALD is not limited to sampling from posterior distributions.
Recent studies have demonstrated that applying an energy-based model (EBM) into the prior distribution over the latent variable enables more ﬂexible generative modeling for DLVMs.
When we train an EBM, we have to obtain samples from the
EBM by running costly MCMC iterations.
By extending our ALD to sampling from unconditional distributions, we can apply ALD into sampling from such EBMs.
When sampling from an EBM with ALD, we prepare a function that maps ﬁxed inputs into latent variables, and updates of EBM samples are replaced with updates of the function’s parameters.
In the same way with the posterior case, the updates of the sampler function can be regarded as implicit updates of samples from EBMs.
Using our ALD for sampling from both the posterior and the EBM over the latent variable, we derive a novel framework of learning DLVMs, which we refer to as the Langevin autoencoder (LAE).
Interestingly, the learning algorithm of LAE naturally takes the combined form of an autoencoder-like architecture and adversarial training.
Our experiments show that ALD can properly obtain samples from target distributions using toy datasets.
Subsequently, we perform numerical experiments of the image generation task using the SVHN, CIFAR-10, and CelebA-HQ datasets.
Not only can LAE be trained faster than non-amortized MCMC methods, but LAE can also generate better samples in terms of the Fr´echet Inception Distance compared to AVI-based methods, such as VAE.
P RELIMINARIES 2. P ROBLEM D EFINITION
Consider a probabilistic model with the d x -dimensional observation x , the d z -dimensional continuous latent variable z , and the model parameter Θ , as described by the probabilistic graphical model shown in Figure 1 (A).
Although the posterior distribution over the latent variable is proportional to the prior and the likelihood: p ( z | x ) = p ( z ) p ( x | z ) /p
( x ) , this is intractable due to
the normalizing constant p ( x ) =
(cid:82) p ( z ) p ( x | z ) d z .
This study aims to approximate the posterior p ( z | x ) for all n observations x (1) , . . .
x ( n ) efﬁciently by obtaining samples from it.
2. L ANGEVIN D YNAMICS
Langevin dynamics (LD) (Neal, 2011) is a sampling algorithm based on the following Langevin equation:
[Equation]
where U is a Lipschitz continuous potential function that satisﬁes an appropriate growth condition, β is an inverse temperature parameter, and B is a Brownian motion.
This stochastic differential equation has p β ( z | x ) ∝ exp ( − βU ( x , z )) as its equilibrium distribution.
We set β = 1 and deﬁne the potential as follows to obtain the target posterior p ( z | x ) as its equilibrium:
[Equation]
We can obtain samples from the posterior by simulating Eq.
(1) using the Euler–Maruyama method (Kloeden & Platen, 2013) as follows:
[Equation]
) where η is a step size for discretization.
When the step size is sufﬁciently 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 the posterior samples for all observations x (1) , . . .
x ( n ) , we should perform iterations of Eq.
(3) per data point, as shown in Figure 1 (B-1).
It is inefﬁcient, mainly if the dataset is large.
In addition, we need to re-run the time-consuming iterations for new observations in test time.
In the next section, we demonstrate a method that addresses the inefﬁciency by amortization.
A MORTIZED L ANGEVIN D YNAMICS 3. G ENERAL I DEA
As an alternative to the direct simulation of latent dynamics, we deﬁne an inference model f z | x , which maps the observation into the latent variable.
Formally, the dynamics of its parameter Φ is
[Equation]
Because the function f z | x outputs latent variables, the stochastic dynamics on the parameter space induce dynamics on the latent space.
The main idea of our amortized Langevin dynamics (ALD) is to regard the transition on this induced dynamics as a sampling procedure for the posterior distributions, as shown in Figure 1 (B-2).
We can use the Euler–Maruyama method to simulate Eq.
(4) like traditional LD:
[Equation]
(cid:17)(cid:17)
[Equation] [Equation]
[Table]
Through the iterations, the posterior sampling is implicitly performed by collecting outputs of the inference model for each data point as described in Algorithm 1.
Note that Z ( i ) denotes a set of samples of the posterior for the i -th data (i.e., p (cid:0) z | x ( i ) (cid:1) ) obtained using ALD.
When we perform inference for new test data, the trained inference model can be used to initialize an MCMC method (e.g., traditional LD) because it is expected that the trained inference model can map data into the high-density area of the posteriors.
By this amortization, we replace the direct update of latent variables ( z (1) , . . .
, z ( n ) ) with the update of the global parameter Φ .
A signiﬁcant advantage of amortization is that the cost of MCMC can be reduced by using minibatch training.
For minibatch training, we substitute the minibatch statistics of m data points for the derivative for all n data.
[Equation]
We refer to the minibatch version of ALD as stochastic gradient amortized Langevin dynamics (SGALD).
SGALD enables us to sample from posteriors of a massive dataset efﬁciently.
Moreover, in the context of stochastic gradient LD (SGLD), it is known that adaptive preconditioning effectively improves convergence compared to the naive SGLD (Li et al., 2016) 2 .
This preconditioning technique is also applicable to our SGALD, and we employ it throughout our experiments.
3. T HEORETICAL A NALYSIS
To justify our ALD as a posterior sampling algorithm, we provide a theoretical analysis of the stationary distribution of our ALD algorithm.
Here, X and Z denote matrices with x ( i ) and z ( i ) in rows X i, : and Z i, : , respectively.
Our main result is as follows:
Theorem 1.
Let q ( Z | X ) be a stationary distribution of the latent variables induced by Eq.
When the mapping f z | x meets the following conditions, q ( Z | X ) satisﬁes q ( Z | X ) ∝ exp ( − U ( X , Z )) :
= exp (cid:0) − (cid:80) n i =1 U (cid:0) x ( i ) , z ( i ) (cid:1)(cid:1) .
The mapping has the form of f z | x ( x ; Φ ) =
Φ g ( x ) , where Φ is a d z × d matrix, g is a mapping from R d x to R d , and d is the dimensionality of g ( x ) .
The rank of G is n , where G is a matrix with g (cid:0)
x
( i ) (cid:1) in row
G i, : .
See Appendix A for the proof.
Theorem 1 suggests that samples obtained by ALD asymptotically converge to the true posterior when we construct the inference model f z | x with an appropriate form.
Practically, we can implement such a function using a neural network whose parameters are ﬁxed except for the last linear layer.
In this implementation, the last linear layer takes a role of the parameter Φ , and the preceding feature extractor takes a role of the function g .
In our experiments, we randomly initialize weights of the feature extractor and freeze them throughout the training.
In addition, the dimensionality of the last linear layer should be sufﬁciently large to meet the second condition.
Therefore, the second condition derives a trade-off between approximation quality and computational costs.
It is worth noting that a similar trade-off is also known in the context of AVI.
In AVI, the approximation quality is inﬂuenced by the capacity of the inference model, and the gap between the optimal variational distribution and the amortized distribution is often denoted as the amortization gap (Cremer et al., 2018).
In experiments, we conﬁrm that preserving the condition does not become a signiﬁcant computational overhead in practice.
In addition, we should note that decaying the step size η is needed to ensure convergence when we perform the simulation with discretization as described by Welling & Teh (2011).
3. E XTENSION TO U NCONDITIONAL C ASES
Currently, we have introduced ALD as a sampling algorithm for conditional posterior distributions.
We can also apply ALD to sampling from unconditional unnormalized distributions, namely energybased models (EBMs).
Consider an unconditional distribution over a random variable z deﬁned by an energy function f
z as follows:
[Equation]
where f z maps the variable z into a scalar value.
To obtain samples from this EBM using ALD, we prepare a sampler function f z | u ( u ; Ψ ) that maps its input u into the variable z .
Here, the input vector u is ﬁxed, whereas observations are used in the posterior case.
To run multiple MCMC chains in parallel, we prepare k ﬁxed inputs u (1) , . .
, u ( k ) and update the parameter of the sampler function as follows:
[Equation]
Typically, the ﬁxed input vectors u (1) , . .
, u ( k ) are chosen from a standard Gaussian distribution.
As in the posterior case, we can guarantee the stationary distribution matches the EBM by choosing an appropriate form for the function f z | u , and such a function can be implemented using a neural network whose parameters are ﬁxed except for the last linear layer.
For minibatch training, we can substitute the gradient for all k chains with the stochastic gradient of m minibatch chains:
[Equation]
The advantage of using amortization in the unconditional case is that we can run massive chains parallel using minibatch training.
L ANGEVIN A UTOENCODERS
Using ALD for sampling from both the posterior and the energy-based prior, we derive a novel framework for learning DLVMs.
We here consider a latent variable model deﬁned as follows:
[Equation]
where σ is a variance parameter of a Gaussian distribution 3 .
To learn the latent variable model, we need to estimate both the model parameter Θ and the latent variable z .
In Bayesian learning, the estimation is represented as the joint posterior distribution p (cid:0) Θ , z (1) , . . .
, z ( n ) | x (1) , . . .
, x ( n ) (cid:1) 4 .
To obtain samples from the posterior, we can combine traditional LD and ALD as follows:
[Equation]
where f z | x ( X ; Φ ) is a matrix with f z | x (cid:0)
x ( i ) ; Φ (cid:1) in its i -th row.
If we omit the Gaussian noise injection in Eq.
(10), it corresponds to gradient ascent for maximum a posteriori (MAP) estimation
Θ , Φ , Ψ ← Initialize parameters repeat
[Equation]
Update the generative model
[Equation]
(cid:46) Update the inference model
[Equation]
(cid:46) Update the sampler model until convergence of parameters return Θ , Φ , Ψ of Θ ; if we additionally use a ﬂat prior for p ( Θ ) , it yields the maximum likelihood estimation (MLE).
In this study, we assume a ﬂat prior for p ( Θ ) and omit the notation for simplicity.
In Eq.
(10), we cannot calculate the derivative of the potential function ∇ Θ U in a closed-form because the latent prior p ( z | Θ ) is deﬁned using an unnormalized energy function.
However, we can obtain the unbiased estimator of the derivative by obtaining samples from the prior as follows:
[Equation]
where ˜ z (1) , . . .
˜
z ( n ) are sampled from the latent prior p ( z | Θ ) (see Appendix C for the derivation).
To get samples from the latent prior, we can also use ALD by preparing a sampler function f z | u that takes ﬁxed inputs u (1) , . . . , u ( k ) , as described in Section 3.3.
Here, we set the number of chains equal to the number of data points for simplicity, i.e., k = n .
In summary, the encoder f
z | x , the decoder f x | z , and the latent energy function f
z are trained by minimizing the following loss function L , whereas the latent sampler f z | u is trained by maximizing it, while stochastic noise of the Brownian motion is injected in their update in order to avoid shrinking to MAP estimates:
[Equation]
We refer to this framework of learning DLVMs as the Langevin autoencoder (LAE).
We summarize the algorithm of the LAE in Algorithm 2.
LAE is closely related to the traditional autoencoder (AE) and other deep generative models, such as the variational autoencoder (VAE) and the generative adversarial network (GAN).
We discuss the relationship in the next section in detail.
R ELATED W ORKS
Amortized inference is well-investigated in the context of variational inference.
It is often referred to as amortized variational inference (AVI) (Rezende & Mohamed, 2015; Shu et al., 2018).
The basic idea of AVI is to replace the optimization of the datapoint-wise variational parameters with the optimization of shared parameters across all data points by introducing an inference model that predicts latent variables from observations.
The AVI is commonly used in generative models (Kingma & Welling, 2013), semi-supervised learning (Kingma et al., 2014), anomaly detection (An & Cho, 2015), machine translation (Su et al., 2018), and neural rendering (Eslami et al., 2018; Kumar et al., 2018).
However, in the MCMC literature, there are few works on such amortization.
Han et al.
(2017) use traditional LD to obtain samples from posteriors to train deep latent variable models.
Such Langevin-based algorithms for deep latent variable models are called alternating back-propagation (ABP) and are applied in several ﬁelds (Xie et al., 2019; Zhang et al., 2020; Xing et al., 2018; Zhu et al., 2019).
However, ABP requires datapoint-wise Langevin iterations, causing slow convergence.
Moreover, when we perform inference for new data in test time, ABP requires
MCMC iterations from randomly initialized samples again.
Although Li et al.
(2017) and Hoffman (2017)
propose amortization methods for MCMC, they only amortize the initialization cost in
MCMC by using an inference model.
Therefore, they do not entirely remove datapoint-wise MCMC iterations.
Autoencoders (AEs) (Hinton & Salakhutdinov, 2006) can be seen as a particular case of LAEs, wherein the Gaussian noise injection to the update of the inference model (encoder) and the generative model (decoder) is omitted in Eqs.
(10) and (11), and a ﬂat prior is used for p ( z | Θ ) .
When a different distribution is used as a latent prior, it is known as sparse autoencoders (SAEs) (Ng et al., 2011).
In these cases, the dynamics in Eqs.
(10) and (11) are dominated by gradient ∇ U ; hence, both the latent variables and the model parameter converge to MLE or MAP estimates (or other stationary points).
Therefore, AEs (and SAEs) can be considered MLE (and MAP) algorithms for the parameter Θ and the latent variables Z .
Variational Autoencoders (VAEs) are based on AVI, wherein an inference model (encoder) is deﬁned as a variational distribution
q ( z | x ; Φ ) using a neural network.
Its parameter Φ is optimized by maximizing the evidence lower bound.
Interestingly, there is a contrast between VAE and LAE when stochastic noise is used in posterior inference.
In VAE, noise is used to sample from the stochastic inference model in calculating the potential U , i.e., in the forward calculation.
However, in LAE, the inference model itself is deterministic, and stochastic noise is used for its parameter update along with the gradient calculation ∇ φ U , i.e., in the backward calculation.
The advantage of LAE over VAE is that LAE can ﬂexibly approximate complex posteriors by obtaining samples, whereas VAE’s approximation ability is limited by choice of variational distribution
q ( z | x ; Φ ) because it requires a tractable density function.
Although there are several considerations in the improvement of the approximation ﬂexibility, these methods typically have architectural constraints (e.g., invertibility and ease of Jacobian calculation in normalizing ﬂows (Rezende & Mohamed, 2015; Kingma et al., 2016; Van Den Berg et al., 2018; Huang et al., 2018; Titsias & Ruiz, 2019)), or they incur more computational costs (e.g., MCMC sampling for the reverse conditional distribution in unbiased implicit variational inference (Titsias & Ruiz, 2019)).
Energy-based Models ’ training is challenging, and many researchers have been studying methodology for its stable and practical training.
A signiﬁcant challenge is that it requires MCMC sampling from EBMs, which is challenging to perform in high dimensional space.
Our LAE avoids this difﬁculty by deﬁning the energy function in latent space rather than data space.
A similar approach is taken in several works (Pang et al., 2020a;b), but they use traditional LD to obtain latent samples without amortization.
Generative adversarial networks (GANs) are closely related to our LAE because both are trained using adversarial loss functions.
For a detailed discussion, see Appendix D.
E XPERIMENT
In our experiment, we ﬁrst test our ALD algorithm on toy examples to investigate its behavior, then we show the results of its application to the training of deep generative models.
6. T OY E XAMPLES
We perform numerical simulation using toy examples to demonstrate that our ALD can properly obtain samples from target distributions in conditional and unconditional cases.
First, we use examples where the posterior density can be derived in a closed-form.
We initially generate three synthetic data x 1 , x 2 , x 3 , where each x i is sampled from a bivariate Gaussian distribution as follows:
[Equation]
In this case, we can calculate the exact posterior as follows:
[Equation]
= N (cid:16) z ; (cid:0) Σ − 1 z + Σ − 1 x (cid:1) − 1 (cid:0) Σ − 1 z µ z + Σ − 1 x x (cid:1) , (cid:0) Σ − 1 z + Σ − 1 x (cid:1) − 1 (cid:17) ,
In this experiment, we set µ z = (cid:20)
(cid:21) , Σ z = (cid:20) 1 0 0 1 (cid:21) , and Σ x = (cid:20) 0 .
8 (cid:21) .
We simulate our ALD algorithm for this setting to obtain samples from the posterior.
We use a neural network [Figure] of three fully connected layers of 128 units with ReLU activation for the inference model f z | x ; setting the step size to 4 × 10 − 4 , and update the parameters for 3,000 steps.
We omit the ﬁrst 1, samples as burn-in steps and use the remaining 2,000 samples for qualitative evaluation.
The result is shown in Figure 2 (A).
ALD produces samples that match the shape of the target distributions well, even though ALD does not perform direct updates of samples in the latent space.
We also performed a similar experiment in a univariate setting to see the convergence speed of our SGALD, the minibatch version of ALD (see Appendix F.1 for the detailed experimental setting).
Figureshows the evolution of obtained sample values by traditional LD and our SGALD.
It can be observed that SGALD’s samples converge much faster than traditional LD.
GT
AVI
ALD
In addition to the simple conjugate Gaussian example, we experiment with a complex posterior, wherein the likelihood is deﬁned with a randomly initialized neural network.
For comparison, we also implement the amortized variational inference (AVI) method, in which the posterior is approximated with a Gaussian distribution parameterized by a neural network (see Appendix F.2 for more experimental details).
Figure 4 shows a typical example, which characterizes the difference between AVI and ALD.
The advantage of our ALD over AVI is the ﬂexibility of posterior approximation.
AVI methods typically approximate posteriors using variational distributions, which have tractable density functions.
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 such posteriors well.
The results in other examples are summarized in Figure 6 in the appendix.
Furthermore, we also test our ALD for sampling from an unconditional distribution.
In this experiment, we use a mixture distribution of eight Gaussians and obtain samples using ALD, as shown in
Figure 2 (B).
We can observe that ALD adequately captures the actual density’s multimodality and works well in the unconditional case. [Table] [Table]
I MAGE G ENERATION
To demonstrate the applicability of our LAE to the generative model training
[Table]
we experiment on image generation tasks using SVHN, CIFAR
[Table]
and CelebA-HQ datasets.
Note that our goal here is not to provide the state-of-the-art results on image generation benchmarks but to verify the effectiveness of our ALD as a method of approximate inference in deep latent variable models.
For this aim [Table] we compare our LAE with ﬁve baseline methods, as shown in Table 1.
VAE (Kingma &
Welling, 2013) is one of the most popular deep latent variable models in which the posterior distribution is approximated using the AVI.
VAE-ﬂow is an extension of VAE in which the ﬂexibility of
AVI is improved using normalizing ﬂows.
In addition to AVI-based methods, we use three methods based on Langevin dynamics (LD).
The alternating back-propagation (ABP) uses traditional LD to approximate the posterior, and the deep latent Gaussian model (DLGM) uses a VAE-like inference model to initialize LD.
The latent energy-based model (LEBM) uses an EBM for the latent prior, and the EBM and posterior sampling is performed via traditional LD.
LEBM can be regarded as a non-amortization version of our LAE.
We apply a commonly used convolutional neural network-based architecture for all models and a multi-layer perceptron for an energy-based model in the latent space of LAE and LEBM.
Please refer to Appendix F.3 for more detailed experimental settings.
For quantitative evaluation of the sample quality, we report the Fr ´ echet Inception Distance (FID) (Heusel et al., 2017).
The results are summarized in Table 1.
It can be observed that LAE outperforms VI-based methods in terms of FID, although it does not reach LEBM’s performance in SVHN and CelebA-HQ.
In training speed, LAE takes 34.12 seconds per epoch on average to train with CIFAR-10, while LEBM and DLGM take 84.29 and 60.66 seconds per epoch, respectively.
This result shows that LAE is approximately 2 . 47 times faster than the non-amortized LEBM and 1 .
78 times faster than the partially amortized DLGM.
C ONCLUSION
This paper proposed amortized Langevin dynamics (ALD), an efﬁcient MCMC method for deep latent variable models.
The ALD amortizes the cost of datapoint-wise iterations by using inference models.
We showed that our ALD algorithm could accurately approximate posteriors with both theoretical and empirical studies.
Using ALD, we derived a novel scheme of deep generative models called the Langevin autoencoder (LAE).
We demonstrated that our LAE performs better than VIbased methods in sample quality and can be trained faster than non-amortized LD methods.
This study will be the ﬁrst step to further work on efﬁcient MCMC for latent variable models with large-scale datasets.
For instance, deriving a Metropolis-Hastings rejection step for ALD and algorithms based on Hamiltonian Monte Carlo methods is an exciting direction of future work. Moreover, developing more sophisticated way of choosing the feature extractor of the inference model is also important.
In our experiments, we used a randomly initialized neural network that is ﬁxed throughout the training, but there could be a better way to improve the performance of LAE.
A P ROOF OF T HEOREM
First, we prepare some lemmas.
Lemma 2.
Let h : R d z × d → R d z × n be a linear map deﬁned by h ( Φ ) :
= Φ G .
When the rank of G is n , there exists an orthogonal linear map τ : R d z × d → R d z × d such that τ ( Φ ) =
[ ˜ Φ (1) , ˜ Φ (2) ] satisﬁes ker ˜ h
= span (cid:104)
d z × n , ˜ Φ (2) (cid:105) , where ˜ h :
= h ◦ τ − 1 , ˜
Φ (1) :
= (cid:104) ˜
φ 1 , . .
˜ φ n (cid:105) , ˜ Φ (2) :
= (cid:104) ˜
φ n +1 , . . .
˜
φ d (cid:105) and ˜ φ i ∈ R d z for i = 1 , . . .
, d .
Proof.
The singular value decomposition of G is represented by T G ˜ T (cid:62) =
(cid:20) Λ 0 ( d − n × n ) (cid:21) , where Λ is a n × n diagonal matrix Λ = diag ( λ 1 , . . .
, λ n ) , and T and ˜ T are orthogonal matrices.
Since the rank of G is n , λ 1 , . .
λ n are non-zero.
When we set τ ( Φ ) :
= Φ T (cid:62) , we obtain
[Equation]
From the above equation, ker ˜ h
= span (cid:104)
d (cid:48) × n , ˜ Φ (2) (cid:105) holds.
Lemma 3.
For ˜ Φ ∈ R d z × d , let Φ satisfy:
[Equation]
˜ Φ (1) , ˜ Φ (2) (cid:105) =
τ [Equation] Φ ) . A map ˜ h (1) : R d z × n → R d z × n deﬁned by ˜ h (1) (cid:16) ˜
Φ (1) (cid:17) :
=
˜ h
(cid:16)(cid:104) ˜
Φ (1) , 0 (cid:105)(cid:17) satisﬁes Φ G =
˜ h
(1) (cid:16) ˜
Φ (1) (cid:17) and is linear isomorphic.
Proof.
From the deﬁnition, we have
[Equation]
By the deﬁnition, ˜ h (1) is linear.
Here, ˜ h (1) is injective, since ker ˜
h
= span (cid:104)
d (cid:48) × n , ˜ Φ (2) (cid:105) , and hence, dim (cid:16) Im ˜ h
(1) (cid:17) ≥ d z × n .
Since Im ˜ h
⊂ R d z × n , ˜ h
(1) is surjective.
Lemma 4.
For V : R D (cid:51) φ (cid:55)→ V ( Φ ) :
=
U (cid:0) X , f z | x ( X ; Φ ) (cid:1) ∈ R , ˜ Φ = (cid:104) ˜
Φ (1) , ˜ Φ (2) (cid:105) :
= τ ( Φ ) , ˜ V :=
V ◦ τ − 1 and ˜ V (1) (cid:16) ˜
Φ (1) (cid:17) :=
˜ V (cid:16)(cid:104) ˜
Φ (1) , 0 d z × ( d − n ) (cid:105)(cid:17) , Eq.
(4) is equivalent to
[Equation]
Proof.
By direct calculation
[Equation]
we obtain
[Equation]
Then, the following equivalence holds:
[Equation]
where we used τ ◦ τ (cid:62)
=
id because τ is orthogonal.
From Eq.
(18), the dynamics in Eq.
(19) is equivalent to Eq.
(16) and Eq.
In the following, we prove Theorem 1 using the above lemmas.
The latent variables Z := ˜ h
(1) (cid:16) ˜
Φ (1) (cid:17) is independent of ˜
Φ (2) , and the probability distribution q ( Z | X ) of Z is given by the pushforward measure (cid:16) ˜ h
(1) (cid:17) # (cid:16) p (1) ∗ (cid:17) ( Z ) of the probability distribution p (1) of ˜ Φ by ˜ h
The amortized Langevin dynamics has q ( Z | G ) as its stationary distribution of Z .
Then, we have
[Equation]
where we used that d ˜ h (1) d ˜ Φ (1)
=
˜ h (1) because of the linearity of ˜ h
(1) and is constant with respect to
Z .
The last equation is derived as follows.
From Lemma 3, Φ G =
˜
h
(1) (cid:16) ˜
Φ (1) (cid:17) holds when z ← f z | x ( x ; Φ ∗ ) (cid:46) Initialize a sample using a trained inference model Z ← ∅ (cid:46) Initialize a sample set repeat z ← [Equation] I
[Equation] (cid:46) Update the sample using traditional LD Z ← Z ∪ { z } (cid:46) Add samples until convergence of parameters return Z ˜ Φ = [Equation]
˜ Φ (1) , ˜ Φ (2) (cid:105) =
τ ( Φ ) .
Thus, when Φ = τ − 1 (cid:16)(cid:104) ˜
Φ (1) , 0 (cid:105)(cid:17) , we obtain ˜ h
(1) (cid:16) ˜
Φ (1) (cid:17) =
Φ G = τ − 1 (cid:16)(cid:104) ˜
Φ (1) , 0 (cid:105)(cid:17) G .
In particular, for ˜ Φ (1) = (cid:16) ˜ h
(1) (cid:17) − 1 ( Z ) , we have [Equation]
B M AXIMUM L IKELIHOOD T RAINING FOR LAE
In Section 4, we derive Bayesian learning algorithm of LAE, where the whole model is formulated as a Bayesian neural network.
We can also think of the frequentist approach, where the training is deﬁned as maximum likelihood.
In this case, the objective is to maximize the evidence log p ( x | Θ ) , and its derivative is as follows.
[Equation]
We can approximate the derivative by obtaining samples from the posterior p ( z | x , Θ ) and the latent prior p ( z | Θ ) using LAE as in Eq.
(5) and (7).
Hence, the maximum likelihood version of
LAE is summarized as shown in Algorithm 4.
C D ERIVATION OF E Q . (13)
[Equation]
Θ , Φ , Ψ ← Initialize parameters repeatrepeat
[Equation]
Update the inference model [Equation] (cid:46) Update the sampler model until convergence of Φ and Ψ Θ ← Θ − η ∇ Θ L ( Θ , Φ , Ψ ) (cid:46) Update the generative model until convergence of Θ return Θ , Φ , Ψ where C ( Θ ) = (cid:82) exp ( − f z ( z ; Θ ))
d z .
By direct calculation, we obtain
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
where ˜ z (1) , . . .
˜
z ( n ) are samples drawn from p ( z | Θ ) .
D A DDITIONAL R ELATED W ORK
Generative Adversarial Network (GAN) (Goodfellow et al., 2014) is similar to LAE in that both are trained by minimax game between two functions (i.e., the energy function and the sampler function in LAE; the discriminator and the generator in GAN).
However, there are some differences between them.
First, the minimax game is performed in the latent space in LAE, while it is performed in the observation space in GAN.
In other words, the latent variable is identical to the observation (i.e., p ( x | z ) = 1 x = z ) in GAN.
Note that the latent variable z is different from the input of GAN’s generators.
Here, the input of the GAN’s generators is denoted as u for the analogy with LAE.
Second, the loss function is slightly different.
In GAN, the loss function is as follows:
[Equation]
where f x | u denotes the generator that maps its inputs u into the observation space, and f x denotes the discriminator that maps from the observation space into (0 , 1) , and u ( i ) ∼ N ( u ; 0 , I ) .
The discriminator is trained to minimize this loss function, whereas the generator is trained to maximize it.
The main difference to the loss function of LAE is the second term.
When we substitute it with − log f x (cid:0) f x | u (cid:0) u ( i ) ; Ψ (cid:1) ; Θ (cid:1) , it becomes more similar.
This modiﬁcation is often used to alleviate gradient vanishing and stabilize the training of GAN’s generator (Goodfellow et al., 2014; Johnson & Zhang, 2018).
In this formulation, the counter parts of the energy function and the sampler function are − log f x ( · ; Θ ) and f x | u ( · ; Ψ ) , respectively.
Another difference between LAE and GAN is that the input vector of the sampler function is ﬁxed through the training in LAE, whereas the input of the generator changes per iteration by sampling from N ( u ; 0 , I ) in GAN.
Furthermore, LAE is trained using noise injected gradient, whereas GAN [Equation] is trained with a standard stochastic optimization method like SGD.
In GAN, the discriminator is also trained with standard stochastic optimization, which corresponds to the MLE case where noise injection is omitted.
Although there are some investigations to apply Bayesian approach to GAN (Saatci & Wilson, 2017; He et al., 2018), their discriminators are not deﬁned as energy functions.
Wasserstein GANs (WGANs)
(Arjovsky et al., 2017) also have a loss function similar to LAE’s:
[Equation]
where D denotes the discriminator of WGANs that maps from the observation space into the real space R .
In this case, the counter part of the energy function is − D ( x ; Θ ) , although D has a constraint of 1-Lipschitz continuity, which the energy function of LAE does not has. E A DDITIONAL E XPERIMENT
We perform an additional experiment to investigate the amortization gap when the capacity of the inference model is not enough to meet the second condition of Theorem 1.
We use the same experimental setting with the bivariate Gaussian example in Section 6.1, and change the dimensionality of the last linear layer of the inference model from 2 ( < d ) to 128 ( > d ) .
The results are summarized in Figure 5.
It can be observed that the sample quality is good when the dimensionality of the last linear layer is equal to or greater than the number of data points (i.e., d ≥ n ). When the dimensionality is smaller than the number of data points, the samples for some data points shrink to a small area, while good samples are obtained for the remaining data points. [Equation] F. C ONJUGATE U NIVARIATE G AUSSIAN E XAMPLE
In the experiment of conjugate univariate Gaussian example, we initially generate 100 synthetic data x (1) , . . .
, x (100) , where each x ( i ) is sampled from a univariate Gaussian distribution as follows:
[Equation]
In this experiment, we set µ z = 0 , σ 2 z = 1 , σ 2 x = 0 . 01 .
In this case, we can calculate the exact posterior as follows:
[Equation]
In this experiment, we obtain 20,000 samples using SGALD.
We use four fully-connected layers of 128 units with tanh activation for the inference model and set the step size η φ to 0 .
We set the batch size to 10.
[Figure]
AVI (
[Figure]
ALD (hist.) ALD (sample) [Table] F. N EURAL LIKELIHOOD EXAMPLE
We perform an experiment with a complex posterior, wherein the likelihood is deﬁned with a randomly initialized neural network f θ .
Particularly, we parameterize f θ by four fully-connected layers of 128 units with ReLU activation and two dimensional outputs like p ( x | z ) =
N (cid:0) f
θ ( z ) , σ 2 x I (cid:1) .
We initialize the weight and bias parameters with N (0 , 0 .
2 I ) and N (0 , 0 .
1 I ) , respectively.
In addition, we set the observation variance σ x to 0.25.
We used the same neural network architecture for the inference model f φ .
Other settings are same as the previous conjugate Gaussian experiment.
The results are shown in Figure 6.
The left three columns show the density visualizations of the ground truth or approximation posteriors of AVI methods; the right two columns show the visualizations of 2D histograms and samples obtained using ALD.
For AVI method, we use two different models.
One uses diagonal Gaussians, i.e., N (cid:0) µ ( x ; φ ) , diag (cid:0) σ 2 ( x ; φ ) (cid:1)(cid:1) , for the variational distribution, and the oher uses Gaussians with full covariance N ( µ ( x ; φ ) , Σ ( x ; φ )) .
From the density visualization of GT, the true posterior is multimodal and skewed; this leads to the failure of the
Gaussian AVI methods notwithstanding considering covariance.
In contrast, the samples of ALD accurately capture such a complex distribution, because ALD does not need to assume any tractable distributions for approximating the true posteriors.
The samples of ALD capture well the multimodal and skewed posterior, while Gaussian AVI methods fail it even when considering covariance.
[Figure]
F. I MAGE G ENERATION
Our implementation of baseline models is based on the ofﬁcial public code of Pang et al.
(2020a) which is available at https://github.com/bpucla/latent-space-EBM-prior .
The architecture of neural networks for LAE is summarized in Table 2.
Conv k x s x p x c denotes a convolutional layer with k × k kernel, s × s stride, p × p padding, and c output channels.
Deconv k x s x p x c denotes a transposed convolutional layer with k × k kernel, s × s stride, p × p padding, and c output channels.
Linear d is a fully connected layer of output dimension d .
We apply the swish function (Ramachandran et al., 2017) as activation after each convolution, transposed convolution or linear layer except the last one.
d x , d z and d u are the dimensionality of x , z and u respectively.
We ﬁx the parameters in the encoder and the sampler except for the last linear layer to meet the condition of Theorem 1.
For all datasets, we set the minibatch size m to 100 .
The latent dimensionality d z is set to 64 for SVHN and 128 for CIFAR10 and CelebA-HQ.
We set d u = 2048 throughout the experiments.
The training of LAE is performed via the Langevin sampling iterations as explained in Section 4. We use preconditioned stochastic gradient Langevin dynamics (pSGLD) (Li et al., 2016) as the sampling algorithm.
The implementation is available at https://bit.ly/2Swow0F