PIXELVAE: A LATENT VARIABLE MODEL FOR NATURAL IMAGES
Natural image modeling is a landmark challenge of unsupervised learning. Variational Autoencoders (VAEs) learn a useful latent representation and model global structure well but have difficulty capturing small details. PixelCNN models details very well, but lacks a latent code and is difficult to scale for capturing large structures. We present PixelVAE, a VAE model with an autoregressive decoder based on PixelCNN. Our model requires very few expensive autoregressive layers compared to PixelCNN and learns latent codes that are more compressed than a standard VAE while still capturing most non-trivial structure. Finally, we extend our model to a hierarchy of latent variables at different scales. Our model achieves state-of-the-art performance on binarized MNIST, competitive performance on 64 × 64 ImageNet, and high-quality samples on the LSUN bedrooms dataset.
1 INTRODUCTION
Building high-quality generative models of natural images has been a long standing challenge. Although recent work has made significant progress (Kingma & Welling, 2014; van den Oord et al., 2016a;b), we are still far from generating convincing, high-resolution natural images.
Many recent approaches to this problem are based on an efficient method for performing amortized, approximate inference in continuous stochastic latent variables: the variational autoencoder (VAE) (Kingma & Welling, 2014) jointly trains a top-down decoder generative neural network with a bottom-up encoder inference network. VAEs for images typically use rigid decoders that model the output pixels as conditionally independent given the latent variables. The resulting model learns a useful latent representation of the data and effectively models global structure in images, but has difficulty capturing small-scale features such as textures and sharp edges due to the conditional independence of the output pixels, which significantly hurts both log-likelihood and quality of generated samples compared to other models.
PixelCNNs (van den Oord et al., 2016a;b) are another state-of-the-art image model. Unlike VAEs, PixelCNNs model image densities autoregressively, pixel-by-pixel. This allows it to capture fine details in images, as features such as edges can be precisely aligned. By leveraging carefully constructed masked convolutions (van den Oord et al., 2016b), PixelCNNs can be trained efficiently in parallel on GPUs. Nonetheless, PixelCNN models are still very computationally expensive. Unlike typical convolutional architectures they do not apply downsampling between layers, which means that each layer is computationally expensive and that the depth of a PixelCNN must grow linearly with the size of the images in order for it to capture dependencies between far-away pixels. PixelCNNs also do not explicitly learn a latent representation of the data, which can be useful for downstream tasks such as semi-supervised learning.
∗Corresponding author; igul222@gmail.com
Our contributions are as follows:
• We present PixelVAE, a latent variable model which combines the largely complementary advantages of VAEs and PixelCNNs by using PixelCNN-based masked convolutions in the conditional output distribution of a VAE.
• We extend PixelVAE to a hierarchical model with multiple stochastic layers and autoregressive decoders at each layer. This lets us autoregressively model not only the output pixels but also higher-level latent feature maps.
• On MNIST, we show that PixelVAE: (1) establishes a new state-of-the-art likelihood, (2) performs comparably to PixelCNN using far fewer computationally expensive autoregressive layers, (3) learns more compressed latent codes than a standard VAE while still accounting for most non-trivial structure, and (4) learns a latent code which separates digits better than a standard VAE.
• We evaluate hierarchical PixelVAE on two challenging natural image datasets (64 × 64 ImageNet and LSUN bedrooms). On 64× 64 ImageNet, we report likelihood competitive with the state of the art at significantly less computational cost. On LSUN bedrooms, we generate high-quality samples and show that hierarchical PixelVAE learns to model different properties of the scene with each of its multiple layers.
2 RELATED WORK
There have been many recent advancements in generative modeling of images. We briefly discuss some of these below, especially those that are related to our approach.
The Variational Autoencoder (VAE) (Kingma & Welling, 2014) is a framework to train neural networks for generation and approximate inference jointly by optimizing a variational bound on the data log-likelihood. The use of normalizing flows (Rezende & Mohamed, 2015) improves the flexibility of the VAE approximate posterior. Based on this, Kingma et al. (2016) develop an efficient formulation of an autoregressive approximate posterior model using MADE (Germain et al., 2015). In our work, we avoid the need for such flexible inference models by using autoregressive priors.
The idea of using autoregressive conditional likelihoods in VAEs has been explored in the context of language modeling in (Bowman et al., 2016), however in that work the use of latent variables fails to improve likelihood over a purely autoregressive model.
Simultaneously to our work, Chen et al. (2016) present a VAE model for images with an an autoregressive output distribution. In constrast to Chen et al. (2016), who focus on models with a single layer of latent variables, we also investigate models with a hierarchy of latent variables (and corresponding autoregressive priors) and show that they enable us to scale our model to challenging natural image datasets.
Another promising recent approach is Generative Adversarial Networks (GANs) (Goodfellow et al., 2014), which pit a generator network and a discriminator network against each other. Recent work has improved training stability (Radford et al., 2015; Salimans et al., 2016) and incorporated inference networks into the GAN framework (Dumoulin et al., 2016; Donahue et al., 2016). GANs generate compelling samples compared to our work, but still exhibit unstable training dynamics and are known to underfit by ignoring modes of the data distribution (Dumoulin et al., 2016). Further, it is difficult to accurately estimate the data likelihood in GANs.
3 PIXELVAE MODEL
Like a VAE, our model jointly trains an “encoder” inference network, which maps an image x to a posterior distribution over latent variables z, and a “decoder” generative network, which models a distribution over x conditioned on z. The encoder and decoder networks are composed of a series of convolutional layers, respectively with strided convolutions for downsampling in the encoder and transposed convolutions for upsampling in the decoder.
As opposed to most VAE decoders that model each dimension of the output independently (for example, by modeling the output as a Gaussian with diagonal covariance), we use a conditional PixelCNN in the decoder. Our decoder models x as the product of each dimension xi conditioned on all previous dimensions and the latent variable z:
p(x|z) = ∏ i p(xi|x1, . . . , xi−1, z)
We first transform z through a series of convolutional layers into feature maps with the same spatial resolution as the output image and then concatenate the resulting feature maps with the image. The resulting concatenated feature maps are then further processed by several PixelCNN masked convolutional layers and a final PixelCNN 256-way softmax output.
Unlike typical PixelCNN implementations, we use very few PixelCNN layers in our decoder, relying on the latent variables to model the structure of the input at scales larger than the combined receptive
field of our PixelCNN layers. As a result of this, our architecture captures global structure at a much lower computational cost than a standard PixelCNN implementation.
3.1 HIERARCHICAL ARCHITECTURE
The performance of VAEs can be improved by stacking them to form a hierarchy of stochastic latent variables: in the simplest configuration, the VAE at each level models a distribution over the latent variables at the level below, with generation proceeding downward and inference upward through each level (i.e. as in Fig. 3). In convolutional architectures, the intermediate latent variables are typically organized into feature maps whose spatial resolution decreases toward higher levels.
Our model can be extended in the same way. At each level, the generator is a conditional PixelCNN over the latent features in the level below. This lets us autoregressively model not only the output distribution over pixels but also the prior over each set of latent feature maps. The higher-level PixelCNN decoders use diagonal Gaussian output layers instead of 256-way softmax, and model the dimensions within each spatial location (i.e. across feature maps) independently. This is done for simplicity, but is not a limitation of our model.
The output distributions over the latent variables for the generative and inference networks decompose as follows (see Fig. 3).
p(z1, · · · , zL) = p(zL)p(zL−1|zL) · · · p(z1|z2) q(z1, · · · , zL|x) = q(z1|x) · · · q(zL|x)
We optimize the negative of the evidence lower bound (sum of data negative log-likelihood and KL-divergence of the posterior over latents with the prior).
−L(x, q, p) = −Ez1∼q(z1|x) log p(x|z1) +DKL(q(z1, · · · zL|x)||p(z1, · · · , zL)) = −Ez1∼q(z1|x) log p(x|z1) + ∫
z1,··· ,zL
L∏ j=1 q(zj |x) L∑ i=1 log q(zi|x) p(zi|zi+1) dz1...dzL
= −Ez1∼q(z1|x) log p(x|z1) + L∑
i=1 ∫ z1,··· ,zL L∏ j=1 q(zj |x) log q(zi|x) p(zi|zi+1) dz1...dzL
= −Ez1∼q(z1|x) log p(x|z1) + L∑
i=1 ∫ zi,zi+1 q(zi+1|x)q(zi|x) log q(zi|x) p(zi|zi+1) dzidzi+1
= −Ez1∼q(z1|x) log p(x|z1) + L∑
i=1
Ezi+1∼q(zi+1|x) [ DKL(q(zi|x)||p(zi|zi+1)) ]
Note that when specifying an autoregressive prior over each latent level zi, we can leverage masked convolutions (van den Oord et al., 2016b) and samples drawn independently from the approximate posterior q(zi | x) (i.e. from the inference network) to train efficiently in parallel on GPUs.
4 EXPERIMENTS
4.1 MNIST
We evaluate our model on the binarized MNIST dataset (Salakhutdinov & Murray, 2008; Lecun et al., 1998) and report results in Table 1. We also experiment with a variant of our model in which each PixelCNN layer is directly conditioned on a linear transformation of latent variable, z (rather than transforming z first through several upsampling convolutional layers) (as in (van den Oord et al., 2016b) and find that this further improves performance, achieving an NLL upper bound comparable with the current state of the art. We estimate the marginal likelihood of our MNIST model using the importance sampling technique in Burda et al. (2015), which computes a lower bound on the likelihood whose tightness increases with the number of importance samples per datapoint. We use N = 5000 samples per datapoint (higher values don’t appear to significantly affect the likelihood estimate) and achieve state-of-the-art likelihood.
4.1.1 USING FEW PIXELCNN LAYERS
The masked convolutional layers in PixelCNN are computationally expensive because they operate at the full resolution of the image and in order to cover the full receptive field of the image, PixelCNN typically needs a large number of them. One advantage of our architecture is that we can achieve strong performance with very few PixelCNN layers, which makes training and sampling from our model significantly faster than PixelCNN. To demonstrate this, we compare the performance of our model to PixelCNN as a function of the number of PixelCNN layers (Fig. 4a). We find that with fewer than 10 autoregressive layers, our PixelVAE model performs much better than PixelCNN. This is expected since with few layers, the effective receptive field of the PixelCNN output units is too small to capture long-range dependencies in the data.
We also observe that adding even a single PixelCNN layer has a dramatic impact on the NLL bound of PixelVAE. This is not surprising since the PixelCNN layer helps model local characteristics which
are complementary to the global characteristics which a VAE with a factorized output distribution models.
4.1.2 LATENT VARIABLE INFORMATION CONTENT
Because the autoregressive conditional likelihood function of PixelVAE is expressive enough to model some properties of the image distribution, it isn’t forced to account for those properties through its latent variables as a standard VAE is. As a result, we can expect PixelVAE to learn latent representations which are invariant to textures, precise positions, and other attributes which are more efficiently modeled by the autoregressive decoder. To empirically validate this, we train PixelVAE models with different numbers of autoregressive layers (and hence, different PixelCNN receptive field sizes) and plot the breakdown of the NLL bound for each of these models into the reconstruction term log p(x|z) and the KL divergence term DKL(q(z|x)||p(z)) (Fig. 4b). The KL divergence term can be interpreted as a measure of the information content in the posterior distribution q(z|x) (in the sense that in expectation, samples from q(z|x) require KL(q||p) fewer bits to code under a code optimized for q than under one optimized for p (Burnham & Anderson, 2003)) and hence, models with smaller KL terms encode less information in their latent variables.
We observe a sharp drop in the KL divergence term when we use a single autoregressive layer compared to no autoregressive layers, indicating that the latent variables have been freed from having to encode small-scale details in the images. Since the addition of a single PixelCNN layer allows the decoder to model interactions between pixels which are at most 2 pixels away from each other (since our masked convolution filter size is 5×5), we can also say that most of the non-trivial (long-range) structure in the images is still encoded in the latent variables.
4.1.3 LATENT REPRESENTATIONS
On MNIST, given a sufficiently high-dimensional latent space, VAEs have already been shown to learn representations in which digits are well-separated (Sønderby et al., 2016). However, this task becomes more challenging as the capacity of the latent space is decreased. PixelVAE’s flexible output distribution should allow it to learn a latent representation which is invariant to small details and thus better models global factors of variation given limited capacity.
To test this, we train a PixelVAE with a two-dimensional latent space, and an equivalent VAE. We visualize the distribution of test set images in latent space and observe that PixelVAE’s latent representation separates digits significantly better than VAE (Figure 5). To quantify this difference, we train a K-nearest neighbors classifier in the latent space of each model and find that PixelVAE
significantly outperforms VAE, achieving a test error of 7.2% compared to VAE’s 22.9%. We also note that unlike VAE, PixelVAE learns a representation in which digit identity is largely disentangled from other generative factors.
4.2 LSUN BEDROOMS
To evaluate our model’s performance with more data and complicated image distributions, we perform experiments on the LSUN bedrooms dataset (Yu et al., 2015). We use the same preprocessing as in Radford et al. (2015) to remove duplicate images in the dataset. For quantitative experiments we use a 32× 32 downsampled version of the dataset, and we present samples from a model trained on the 64× 64 version. We train a two-level PixelVAE with latent variables at 1×1 and 8×8 spatial resolutions. We find that this outperforms both a two-level convolutional VAE with diagonal Gaussian output and a singlelevel PixelVAE in terms of log-likelihood and sample quality. We also try replacing the PixelCNN layers at the higher level with a diagonal Gaussian decoder and find that this hurts log-likelihood, which suggests that multi-scale PixelVAE uses those layers effectively to autoregressively model latent features.
4.2.1 FEATURES MODELED AT EACH LAYER
To see which features are modeled by each of the multiple layers, we draw multiple samples while varying the sampling noise at only a specific layer (either at the pixel-wise output or one of the latent layers) and visually inspect the resulting images (Fig. 6). When we vary only the pixellevel sampling (holding z1 and z2 fixed), samples are almost indistinguishable and differ only in precise positioning and shading details, suggesting that the model uses the pixel-level autoregressive distribution to model only these features. Samples where only the noise in the middle-level (8 × 8) latent variables is varied have different objects and colors, but appear to have similar basic room geometry and composition. Finally, samples with varied top-level latent variables have diverse room geometry.
4.3 64× 64 IMAGENET
The 64×64 ImageNet generative modeling task was introduced in (van den Oord et al., 2016a) and involves density estimation of a difficult, highly varied image distribution. We trained a heirarchical PixelVAE model (with a similar architecture to the model in section 4.2) on 64× 64 ImageNet and report validation set likelihood in Table 2. Our model achieves a likelihood competitive with van den Oord et al. (2016a;b), despite being substantially less computationally complex. A visual inspection of ImageNet samples from our model (Fig. 7) also reveals them to be significantly more globally coherent than samples from PixelRNN.
5 CONCLUSIONS
In this paper, we introduced a VAE model for natural images with an autoregressive decoder that achieves strong performance across a number of datasets. We explored properties of our model, showing that it can generate more compressed latent representations than a standard VAE and that it can use fewer autoregressive layers than PixelCNN. We established a new state-of-the-art on binarized MNIST dataset in terms of likelihood on 64× 64 ImageNet and demonstrated that our model generates high-quality samples on LSUN bedrooms.
The ability of PixelVAE to learn compressed representations in its latent variables by ignoring the small-scale structure in images is potentially very useful for downstream tasks. It would be interesting to further explore our model’s capabilities for semi-supervised classification and representation learning in future work.
ACKNOWLEDGMENTS
The authors would like to thank the developers of Theano (Theano Development Team, 2016) and Blocks and Fuel (van Merriënboer et al., 2015). We acknowledge the support of the following agencies for research funding and computing support: Ubisoft, Nuance Foundation, NSERC, Calcul Quebec, Compute Canada, CIFAR, MEC Project TRA2014-57088-C2-1-R, SGR project 2014- SGR-1506 and TECNIOspring-FP7-ACCI grant.
C MNIST RECONSTRUCTIONS
B MNIST SAMPLES
D MORE SAMPLES FOR HIERARCHICAL LATENT SPACE VISUALIZATIONS
E MODEL ARCHITECTURE
E.1 MNIST
For our quantitative MNIST experiments, the architectures of our encoder and decoder are as follows. Unless otherwise specified, all convolutional layers use ReLU nonlinearity. We also make an open-source implementation of this model available at https://github.com/igul222/ PixelVAE.
Encoder x→ (µ, σ) Kernel size Stride Output channels
Convolution 3x3 1 32 Convolution 3x3 2 32 Convolution 3x3 1 32 Convolution 3x3 2 64
Pad 7×7 feature maps to 8×8 Convolution 3x3 1 64 Convolution 3x3 2 64 Convolution 3x3 1 64 Convolution 3x3 1 64 Convolution 3x3 1 64
Flatten
Linear - - 2×latent dimensionality
Decoder z → x Kernel size Stride Output channels
Linear - - 4×4×64 Reshape to (64, 4, 4)
Convolution 3x3 1 64 Convolution 3x3 1 64 Transposed convolution 3x3 2 64 Convolution 3x3 1 64
Crop 8×8 feature maps to 7×7 Transposed convolution 3x3 2 32 Convolution 3x3 1 32 Transposed convolution 3x3 2 32 Convolution 3x3 1 32 PixelCNN gated residual block 7x7 1 32
PixelCNN gated residual block(s) [ 5x5 ] ×N 1 32 PixelCNN gated convolution 1x1 1 32 PixelCNN gated convolution 1x1 1 32 Convolution 1x1 1 1
E.2 LSUN BEDROOMS AND 64×64 IMAGENET
The LSUN and ImageNet models use the same architecture: all encoders and decoders are residual networks; we use pre-activation residual blocks with two 3 × 3 convolutional layers each and ELU nonlinearity. Some residual blocks perform downsampling, using a 2 × 2 stride in the second convolutional layer, or upsampling, using subpixel convolution in the first convolutional layer. Weight normalization is used in masked convolutional layers; in all other layers, batch normalization is used. We optimize using Adam with learning rate 1e-3. Training proceeds for 400K iterations using batch size 48.
For further architectural details, please refer to our open-source implementation at https:// github.com/igul222/PixelVAE.
Bottom-level Encoder x→ h1 Kernel size Resample Output channels
Embedding - - 48 Convolution 1x1 - 192 Residual block [ 3x3 ] × 2 - 192 Residual block [ 3x3 ] × 2 Down ×2 256 Residual block [ 3x3 ] × 2 - 256 Residual block [ 3x3 ] × 2 Down ×2 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512
Bottom-level Decoder z1 → x Kernel size Resample Output channels
Convolution 1x1 1 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 Up ×2 256 Residual block [ 3x3 ] × 2 - 256 Residual block [ 3x3 ] × 2 Up ×2 192 Residual block [ 3x3 ] × 2 - 192 Embedding - - 48 PixelCNN gated residual block [ 3x3 ] × 2 - 384 PixelCNN gated residual block [ 3x3 ] × 2 - 384 PixelCNN gated residual block [ 3x3 ] × 2 - 384
Top-level Encoder h1 → h2 Kernel size Resample Output channels
Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 Down ×2 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 Down ×2 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512
Top-level Decoder z2 → z1 Kernel size Resample Output channels
Linear - - 4×4×512 Reshape to (512, 4, 4)
Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 Up ×2 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 Up ×2 512 Residual block [ 3x3 ] × 2 - 512 Residual block [ 3x3 ] × 2 - 512 PixelCNN convolution 5x5 - 512 PixelCNN gated residual block [ 3x3 ] × 2 - 512 PixelCNN gated residual block [ 3x3 ] × 2 - 512 PixelCNN gated residual block [ 3x3 ] × 2 - 512 Convolution 1x1 - 256
UPDATE: The authors addressed all my concerns in the new version of the paper, so I raised my score and now recommend acceptance.
--------------
This paper combines the recent progress in variational autoencoder and autoregressive density modeling in the proposed PixelVAE model. The paper shows that it can match the NLL performance of a PixelCNN with a PixelVAE that has a much shallower PixelCNN decoder.
I think the idea of capturing the global structure with a VAE and modeling the local structure with a PixelCNN decoder makes a lot of sense and can prevent the blurry reconstruction/samples of VAE. I specially like the hierarchical image generation experiments.
I have the following suggestions/concerns about the paper:
1) Is there any experiment showing that using the PixelCNN as the decoder of VAE will result in better disentangling of high-level factors of variations in the hidden code? For example, the authors can train a PixelVAE and VAE on MNIST with 2D hidden code and visualize the 2D hidden code for test images and color code each hidden code based on the digit and show that the digits have a better separation in the PixelVAE representation. A semi-supervised classification comparison between VAE and the PixelVAE will also significantly improve the quality of the paper.
2) A similar idea is also presented in a concurrent ICLR submission "Variational Lossy Autoencoder". It would be interesting to have a discussion included in the paper and compare these works.
3) The answer to the pre-review questions made the architecture details of the paper much more clear, but I still ask the authors to include the exact architecture details of all the experiments in the paper and/or open source the code. The clarity of the presentation is not satisfying and the experiments are difficult to reproduce.
4) As pointed out in my pre-review question, it would be great to include two sets of MNIST samples maybe in an appendix section. One with PixelCNN and the other with PixelVAE with the same pixelcnn depth to illustrate the hidden code in PixelVAE actually captures the global structure.
I will gladly raise the score if the authors address my concerns.