Because images can be viewed as realizations drawn from a latent variable model, we implemented a variational autoencoder that uses neural networks as the variational family to approximate the Bayesian representation. Unlike standard parametric distributions, neural networks can approximate arbitrary distributions reasonably well. We tested the model on MNIST and on the harder SVHN dataset, comparing architectures to see how the dimension of the latent space affects the learned representation and visualizing the learned manifold for low-dimensional latent spaces. With a convolutional encoder, the variational autoencoder learns SVHN reasonably well, and its generated images improve as the latent dimension grows.
Introduction
This post describes the variational autoencoder of Kingma and Welling (2014) and how it performs on the MNIST and SVHN datasets. We first summarize how the method works and the intuition behind it, starting with the notation used throughout.
- \(X\) is the dataset of interest. Because we work mostly with images, we call \(X\) the image data.
- \(z\) is the latent state variable.
- \(p_{\phi}(z\vert X)\) is the target distribution of the latent state space.
- \(q_{\theta}(z\vert X)\) is the variational family for the latent state space.
The variational autoencoder rests on the premise that hidden features (the latent state variables) generate the image data. We further assume that the latent features for a given set of images (e.g., a set of dog pictures) are drawn from a prior distribution of \(z\), as the figure below illustrates.
We want to model the target distribution of the latent state space given the data \(X\), \(p_{\phi}(z\vert X)\). Because this distribution is typically intractable, we use the variational inference proposed by Kingma and Welling (2014) and measure fit with the Kullback-Leibler (KL) divergence. The rest of the post sets out the methodology and presents results on MNIST and SVHN across several architectures.
Methodology
Variational inference approximates the target distribution \(p_{\phi}(z\vert X)\): we pick the \(q_{\theta}(z \vert X)\) from a variational family that minimizes the KL divergence,
\[\min D_{KL}(q_{\theta}(z\vert X)\|p_{\phi}(z\vert X))\]Following Doersch (2016), the KL divergence satisfies
\[\log p(X) - D_{KL}(q_{\theta}(z\vert X)\|p_{\phi}(z\vert X)) = E_q[\log p_{\phi}(X\vert z)] - D_{KL}(q_{\theta}(z\vert X)\|p(z))\]This identity is the core of the variational autoencoder. Because \(D_{KL}(q_{\theta}(z \vert X)\|p_{\phi}(z\vert X))\) is non-negative, the right-hand side is a lower bound on the log-likelihood of the marginal image distribution. We denote it as
\[\mathcal{L}(X, \phi, \theta) = E_q[\log p_{\phi}(X \vert z)] - D_{KL}(q_{\theta}(z\vert X)\|p(z))\]We therefore maximize the log-likelihood indirectly by maximizing \(\mathcal{L}(X, \phi, \theta)\). \(\mathcal{L}(X, \phi, \theta)\) has two parts: the reconstruction error and the KL divergence of the approximate posterior from the prior.
Monte Carlo EM can estimate \(\theta\) and \(\phi\), but slowly. We instead used stochastic gradient descent with backpropagation through a neural network, which requires a reparameterization trick to supply the inputs \(z\) to \(p_{\phi}(X\vert z)\). Consider two networks: the autoencoding network (encoder) and the reconstruction network (decoder). The encoder takes a set of images and outputs their hidden features; the decoder takes hidden features sampled from the prior distribution of the latent state variables and generates the parameters of \(p_{\phi}(X\vert z)\). Kingma and Welling (2014) show that the sampled input to the decoder, \(\mathbf{z}\), can be expressed as a deterministic variable \(g_{\theta}(x, \varepsilon)\), where \(\varepsilon\) is an auxiliary variable. For example, in the univariate Gaussian case, \(\mathbf{z} = \mu + \sigma \varepsilon\), where \(\varepsilon \sim N(0, 1)\).
Gaussian Encoder
The encoder distribution is typically Gaussian, and we also assume a Gaussian prior for the latent state model. Doersch (2016) derives the KL divergence between two Gaussians as follows:
\[\frac{1}{2}\left(tr(\Sigma\left(X, \theta)^{-1}\Sigma\right)+(\mu(X, \theta) - \mu)'\Sigma(X, \theta)(\mu(X, \theta) - \mu) - k + \log\left(\frac{\vert\Sigma(X, \theta)\vert}{\vert\Sigma\vert}\right)\right)\]where \(k\) is the number of latent features.
With a standard multivariate Gaussian prior, the KL divergence simplifies to
\[\frac{1}{2}\left(tr(\Sigma(X, \theta)) + \mu(X, \theta)'\mu(X, \theta) - k - \log\vert\Sigma(X, \theta)\vert\right)\]For the reparameterization, we sample from \(N(\mu(X, \theta), \Sigma(X, \theta))\) by first sampling \(\varepsilon \sim N(0, I)\) and then computing \(\mathbf{z} = \mu(X, \theta) + \Sigma(X, \theta)^{1/2}\varepsilon\). Because the prior of \(z\) is a standard spherical Gaussian, the posterior covariance is close to diagonal, which simplifies the computation to
\[\mathbf{z} = \mu(X, \theta) + \sqrt{tr(\Sigma(X, \theta))}\varepsilon\]The decoder distribution can be either Bernoulli or Gaussian; for the Gaussian decoder, we used a diagonal covariance matrix.
Empirical Analysis
We ran experiments on MNIST first and then on SVHN. Because the variational autoencoder generally behaves well on MNIST, we start there before testing how well it handles SVHN.
MNIST Data
In the first MNIST experiment, we used a simple fully connected neural network for the Gaussian encoder and the Bernoulli decoder across several latent-space dimensions. Figure 1 shows the decoder’s average output for the sample data in Figure 1a. The decoder improves substantially as the latent dimension increases, a pattern that recurs with SVHN. In the 2D latent space (Figure 1b), only 0, 1, 3, 8, and 9 appear to be recognized. In Figure 1c, 6 and 7 are also recognized, and in Figures 1d and 1e the decoder represents all digits.
Figure 2a shows the manifold our autoencoder learned on MNIST, which is consistent with the 2D case in Figure 1b. Because MNIST images are only \(28\times 28\), a low-dimensional latent space can represent them reasonably well, and a Bernoulli decoder suits their black-and-white pixels. SVHN images are larger, at \(32 \times 32\times 3\) (RGB), and more complex, so low-dimensional latent spaces may represent them less well. The next section describes the architectures we tested on SVHN and how they performed.
SVHN Data
We started from the MNIST model and progressed to more complex architectures, moving from fully connected networks to a convolutional network.
Experiments with Bernoulli Decoder and Fully Connected Neural Network
Although SVHN images are RGB, pixel values are stored between 0 and 1, and matplotlib renders them as color images. We can therefore still model each pixel as a Bernoulli probability, interpreted as relative color. Because SVHN inputs are much larger than MNIST’s, we added layers to the encoder and decoder to increase model capacity. The two-dimensional manifold learned from SVHN (Figure 2b) appears to pick up the digit 8, but the images are otherwise homogeneous apart from shade. The encoder therefore captures two features, shade and the digit 8, which is a decent representation.
We then increased the latent dimension to 5, 20, 50, and 100, the settings used for all remaining SVHN experiments. Figure 3b supports reading the low-dimensional patterns as 8s. As the dimension rises, the model captures more colors and digits, but its output remains far from the sample data; in particular, it never picks up green. This shortfall motivated a different architecture.
Convolutional neural networks are widely used for image processing and classification, so they were a natural next step. In all the experiments below, we used a single convolutional layer with 16 activation maps in the encoder and none in the decoder. Decoders are often built as mirror images of the encoder, using deconvolutional layers to output the generated images, but convolutional networks train slowly without a GPU, so we kept the design simple.
Experiments with Bernoulli Decoder and Convolutional Neural Network
We kept the Bernoulli decoder first. The learned manifold in Figure 2c resembles Figure 2b but also includes blue, which suggests the model captures a color-spectrum feature. Generated images improve markedly as the latent dimension rises (Figure 4); at 100 dimensions, they are almost identical to the sample data.
Experiments with Gaussian Decoder and Convolutional Neural Network
Because RGB values are not probabilities, a Gaussian decoder may be more appropriate than a Bernoulli one. With the same network architecture, however, the Gaussian decoder performs similarly to the Bernoulli version (Figure 5), so it offers no clear advantage. It can also produce negative means, so we applied a sigmoid to force the means between 0 and 1.
The 2D manifold under the Gaussian decoder (Figure 2d) does separate light and dark colors more clearly, and the digit 8 is more apparent. This suggests that a Gaussian decoder learns faster than a Bernoulli one.
Conclusion
With a convolutional encoder, the variational autoencoder of Kingma and Welling (2014) learns the SVHN dataset reasonably well. Performance improves as the number of latent features grows. A Gaussian decoder may suit colored images better than a Bernoulli decoder, although in our tests the two performed similarly.
Our code is available on GitHub.
References
- [stat.ML] Diederik P. Kingma and Max Welling, Auto-Encoding Variational Bayes, 2014.
- [stat.ML] Carl Doersch, Tutorial on Variational Autoencoders, 2016.