Langevin Autoencoders for Learning Deep Latent Variable Models

Consider a probabilistic model with the observation ${\mathbf{x}}$, the continuous latent variable ${\mathbf{z}}$, and the model parameter ${\bm{\theta}}$ as follows:

To learn this latent variable model (LVM) via maximum likelihood in a gradient-based manner, we need to calculate the expectation over the posterior distribution $p\left({\mathbf{z}}\mid{\mathbf{x}};{\bm{\theta}}\right)$ as follows:

where $\hat{p}_{\mathrm{data}}$ is the empirical distribution define by the training set, and ${\bm{x}}^{(1)},\ldots,{\bm{x}}^{(n)}$ are mini-batch data uniformly drawn from $\hat{p}_{\mathrm{data}}$. However, this expectation cannot be calculated in a closed-form because the posterior is intractable. In this paper, we consider to use Monte Carlo approximation by obtaining samples from the posterior per data point.

Langevin dynamics (LD) (Neal, 2011) is a sampling algorithm based on the following Langevin equation:

where $U$ is a Lipschitz continuous potential function that satisfies an appropriate growth condition, $\beta$ is an inverse temperature parameter, and $B$ is a Brownian motion. This stochastic differential equation (SDE) has $p^{\beta}\left({\bm{z}}\mid{\bm{x}};{\bm{\theta}}\right)\propto\exp\left(-\beta U\left({\bm{x}},{\bm{z}};{\bm{\theta}}\right)\right)$ as its stationary distribution. We set $\beta=1$ and define the potential as follows to obtain the target posterior $p\left({\mathbf{z}}\mid{\bm{x}};{\bm{\theta}}\right)$ as its stationary distribution:

We can obtain samples from the posterior by simulating the SDE of Eq. (3) using the Euler–Maruyama method (Kloeden and Platen, 2013) as follows:

where $\eta$ is a step size for discretization. The initial value ${\bm{z}}_{0}$ is typically sampled from the prior $p\left({\bm{z}}\right)$. When the step size is sufficiently small, the samples asymptotically move to the target posterior by repeating this sampling iteration. To remove the discretization error, an additional Metropolis-Hastings (MH) rejection step is often used. In a MH step, we first calculate the acceptance rate $\alpha_{t}$ as follows:

The sample ${\bm{z}}_{t+1}$ is accepted with probability $\alpha_{t}$, and rejected with probability $1-\alpha_{t}$. If the sample is rejected, we set ${\bm{z}}_{t+1}={\bm{z}}_{t}$. LD can be applied to any posterior inference problems for continuous latent variables, provided the potential energy is differentiable on the latent space. To obtain the posterior samples for all mini-batch data ${\bm{x}}^{(1)},\ldots{\bm{x}}^{(n)}$, we should perform iterations of Eq. (5) per data point, as shown in Figure 1 (b).

After obtaining the samples, the gradient in Eq. (2) is approximated using the samples as follows:

The time averaged gradient in Eq. (8) is sometimes substituted for the one calculated with the final sample as follows:

However, in these traditional approaches, we need to run MCMC iterations from random initialization every time after updating the model parameter ${\bm{\theta}}$. In addition, we also need to run it from stratch to perform inference for new data in test time. This inefficiency hinders the practical use of MCMC for learning DLVMs. A naive approach to mitigate this problem is to initialize MCMC with a persistent sample (Tieleman, 2008; Han et al., 2017) that is the final value of the previous Markov chain. However, this approach is also inefficient especially when the training set is large, because we need to store persistent samples for all training examples.

To alleviate the inefficiency of LD for LVMs, Hoffman (2017) has proposed to use a stochastic encoder $q\left({\mathbf{z}}\mid{\mathbf{x}}\right)$ to initialize the datapoint-wise MCMC as shown in Figure 1 (c). The encoder is typically defined as a Gaussian distribution as in the VAE:

where $\mu$ and $\sigma^{2}$ are mappings from the observation space into the latent space. In the training, LD iterations of Eq. (5) are initialized using a sample from the distribution of Eq. (10), and the model parameter is updated using the stochastic gradient in Eq. (9). Along with it, the encoder is trained via maximization of the evidence lower bound as in VAEs:

Although initializing LD with the encoder can speed up the convergence, this method still relies on datapoint-wise MCMC iterations. Moreover, the encoder has to be trained using the different objective, which also makes its implementation complicated. In Section 3, we demonstrate a method that entirely remove the datapoint-wise iterations by amortization.

As an alternative to the direct simulation of latent dynamics of Eq. (3), we define a deterministic encoder $f_{{\mathbf{z}}\mid{\mathbf{x}}}$, which maps the observation into the latent variable, and consider an SDE over its parameter ${\bm{\phi}}$ as follows:

Because the function $f_{{\mathbf{z}}\mid{\mathbf{x}}}$ outputs the latent variable, the stochastic dynamics on the parameter space induces 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 (d). We can use the Euler–Maruyama method to simulate Eq. (12) with discretization:

As in the traditional LD, the discretization error can be removed by adding a MH rejection step, calculating the acceptance rate as follows:

Through the iterations, the posterior sampling is implicitly performed by collecting outputs of the encoder for each data point as described in Algorithm 1. Note that ${\mathbb{Z}}^{(i)}$ denotes a set of samples of the posterior for the $i$-th data (i.e., $p\left({\mathbf{z}}\mid{\bm{x}}^{(i)}\right)$) obtained using ALD.

If this implicit update iteration has the posterior as its stationary distribution, this sampling procedure is valid as an MCMC algorithm. To obtain the stationary distribution over the latent variables, we first derive the stationary distribution over the parameter ${\bm{\phi}}$ of Eq. (12), then transform it into the latent space by considering the change of random variable by the encoder $f_{{\mathbf{z}}\mid{\mathbf{x}}}$. Based on this strategy, we derive the following theorem:

See the appendix for the solid proof. Theorem 1 suggests that samples obtained by ALD asymptotically converge to the true posterior when we construct the encoder $f_{{\mathbf{z}}\mid{\mathbf{x}}}$ with an appropriate form. Practically, we can implement such a function using a neural network whose parameters are fixed except for the last linear layer. In this implementation, the last linear layer takes a role of the parameter ${\bm{\Phi}}$, and the preceding feature extractor takes a role of the function $g$.

To meet the second condition, the dimensionality of the last linear layer should be larger than the batch size $n$. It is relatively easy to meet this condition because the batch size is about 1,000 at most in practice.

Remark 1: ALD completely removes datapoint-wise iterations.

Because ALD treats the encoder’s outputs themselves as posterior samples, we no longer need to run time-consuming iterations of datapoint-wise MCMC, whereas existing methods, such as Hoffman (2017) only use the encoder for initialization of datapoint-wise MCMC.

Remark 2: ALD is valid as an MCMC algorithm.

Although the sampling procedure of ALD is quite simple, ALD has the true posterior as its stationary distribution under mild assumptions, which guarantee that ALD is valid as an MCMC algorithm. Basically, we can meet the assumptions by using a sufficiently wide neural network for the encoder. In addition, it is worth mentioning that the traditional LD can be seen as a special case of ALD where $g\left({\bm{x}}^{(i)}\right)=\textrm{one-hot}\left(i\right)$. In that case, ${\bm{z}}^{(i)}=f_{{\mathbf{z}}\mid{\mathbf{x}}}\left({\bm{x}}^{(i)};{\bm{\Phi}}\right)$ corresponds to the $i$-th column of ${\bm{\Phi}}$, which is equivalent to running MCMC on the latent space independently for each data point.

Remark 3: The encoder may accelerate the convergence of MCMC.

After running ALD iterations of the encoder for some mini-batch data, it is expected that the encoder can map data into high density area in the latent space. Therefore, the encoder can accelerate the convergence of MCMC for other data by initializing it with the encoder. This characteristics is useful not only in the training of DLVMs to efficiently estimate the gradient of the model parameters but also at test time to infer the latent representation for new test data.

Based on ALD, we propose a novel framework of learning deep latent variable models called the Langevin autoencoder (LAE). Algorithm 2 is a summary of the LAE’s training procedure. First, we prepare an encoder defined by $f_{{\mathbf{z}}\mid{\mathbf{x}}}\left({\bm{x}};{\bm{\Phi}},{\bm{\psi}}\right)\coloneqq{\bm{\Phi}}g\left({\bm{x}};{\bm{\psi}}\right)$. The feature extractor $g$ is typically implemented using a deep neural network. Before each update of the model parameter ${\bm{\theta}}$, the encoder’s final linear layer ${\bm{\Phi}}$ is updated using ALD for $T$ times, and the gradient of ${\bm{\theta}}$ in Eq. (8) is calculated using the time average of $V_{t}$. Along with the update of the model parameter, the encoder’s feature extractor ${\bm{\psi}}$ is also updated using the gradient so that the encoder can map data into high density area of the posteriors. Although we describe the parameter update using a simple stochastic gradient ascent, it can be substituted for more sophisticated optimization methods, such as Adam (DP and Ba, 2015).

It can be seen that the algorithm is very similar to the one for traditional deterministic autoencoders (Hinton and Salakhutdinov, 2006). In particular, when we skip the ALD iterations in lines 4 to 9 and the encoder’s final linear layer ${\bm{\Phi}}$ is also trained via gradient ascent, the algorithm is identical to the training of autoencoders regularized by the latent prior $p\left({\bm{z}}\right)$. In that case, the encoder tends to shrink to the maximum a posteriori (MAP) estimate rather than the posterior; hence the gradient of the model parameter ${\bm{\theta}}$ would be a strongly biased estimate of the marginal log-likelihood gradient. Therefore, the additional ALD iterations can be interpreted as the modification to reduce the bias of the traditional autoencoder by updating the encoder’s parameters along with the noise-injected gradient.

We perform numerical simulation using toy examples to demonstrate that our ALD can properly obtain samples from target distributions. First, we use a LVM where the posterior density can be derived in a closed-form. We initially generate three synthetic data ${\bm{x}}_{1},{\bm{x}}_{2},{\bm{x}}_{3}$, where each ${\bm{x}}_{i}$ is sampled from a bivariate Gaussian distribution as follows:

In this case, we can calculate the exact posterior as follows:

In this experiment, we set ${\bm{\mu}}_{\mathbf{z}}=\left[\begin{array}[]{c}0\\ 0\end{array}\right]$, ${\bm{\Sigma}}_{\mathbf{z}}=\left[\begin{array}[]{cc}1&0\\ 0&1\end{array}\right]$, and ${\bm{\Sigma}}_{\mathbf{x}}=\left[\begin{array}[]{cc}0.7&0.6\\ 0.7&0.8\end{array}\right]$. We simulate our ALD algorithm for this setting to obtain samples from the posterior. We use a neural network of three fully connected layers of 128 units with ReLU activation for the encoder $f_{{\mathbf{z}}\mid{\mathbf{x}}}$; setting the step size to $4\times 10^{-4}$, and update the parameters for 3,000 steps. We omit the first 1,000 samples as burn-in steps and use the remaining 2,000 samples for qualitative evaluation. As we described in Section 3.1, we need to make the last linear layer of the encoder have sufficiently large dimensions to guarantee the convergence to the true posterior. To empirically demonstrate this theoretical finding, we change the dimensionality of the last linear layer of the encoder from $2\ (<n)$ to $128\ (\gg n)$.

The results are summarized in Figure 2. 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\geq 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.

In addition to the simple conjugate Gaussian example, we experiment with a complex posterior, wherein the likelihood is defined with a randomly initialized neural network. For comparison, we also implement the variational inference (VI), in which the posterior is approximated with a Gaussian distribution. Figure 3 shows a typical example, which characterizes the difference between VI and ALD. The mean-filed VI and the full VI use Gaussians with diagonal and full covariance matrices for variational distributions, respectively. The advantage of our ALD over VI is the flexibility of posterior approximation. VI 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. In particular, the mean-field VI, which is widely used for learning DLVMs, cannot capture the correlation between dimensions due to the constraint of the variational distribution. The full VI mitigate the inflexibility, but it still cannot capture the multimodality of the true posterior. Moreover, the full VI requires computational costs proportional to the square of the dimension of the latent variable. On the other hand, ALD can capture such posteriors well. The results in other examples are summarized in the appendix.

To demonstrate the applicability of our LAE to the generative model training, we experiment on image generation tasks using MNIST, SVHN, CIFAR10, and CelebA 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, we compare our LAE with some baseline methods, as shown in Table 1. The VAE (Kingma and Welling, 2013) is one of the most popular deep latent variable models, in which the posterior distribution is approximated using the VI. The VAE-flow is an extension of VAE, in which the flexibility of VI is improved using normalizing flows (Rezende and Mohamed, 2015). In addition to VI-based methods, we use Hoffman (2017) as a baseline method based on Langevin dynamics (LD). As described in Section 2.2, Hoffman (2017) uses a VAE-like encoder to initialize LD, and the encoder is trained by maximizing the evidence lower bound. We use the same fully connected deep neural networks for the model construction of all methods. We set the number of ALD iterations of the LAE to $2$, i.e, $T=2$ in Algorithm 2. Please refer to the appendix for more implementation details.

For evaluation, since the marginal log-likelihood for test data is not tractable, we substitute its evidence lower bound (ELBO) for it using a proposal distribution $q$ as follows:

For the baseline methods, their stochastic encoders are used for the proposal distribution, i.e., $q\left({\bm{z}}\right)\coloneqq q\left({\bm{z}}\mid{\bm{x}};{\bm{\phi}}\right)$. For our LAE, we use a Gaussian distribution whose mean parameter is defined by its encoder, i.e., $q\left({\bm{z}}\right)\coloneqq\mathcal{N}\left({\bm{z}};{\bm{\Phi}}g\left({\bm{x}};{\bm{\psi}}\right),\ \sigma^{2}{\bm{I}}\right)$. We set $\sigma=0.05$ in the experiment.

The results are summarized in Table 1. Note that the negative ELBO is shown in the table, so lower values indicate better results. It can be observed that the LAE consistently outperforms baseline methods, demonstrating that accurate posterior approximation by ALD leads to better results in the training of DLVMs. On training speed, we observe that our LAE is 2.24 and 1.88 times slower than VAE and VAE-flow respectively. This is natural because VAE and VAE-flow do not require MCMC steps. Hoffman (2017) and our LAE are almost identical in terms of training speed.

We also investigate the effect of the MH rejection step and the number of ALD iterations, i.e., $T$ in Algorithm 2, using MNIST dataset. Figure 4 shows a comparison of the learning curves of the negative ELBO for MNIST’s test set. It can be seen that the number of ALD iterations has little effect on performance as long as the number of steps is at least 2. In addition, we observe that the MH rejection step is important to stabilize the training.