Title: PIXELVAE: A LATENT VARIABLE MODEL FOR NATURAL IMAGES

Abstract: 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. Pix-elCNNs also do not explicitly learn a latent representation of the data, which can be useful for downstream tasks such as semi-supervised learning. 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. 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. VAEs, which assume (conditional) independence among pixels, are known to suffer from blurry samples, while PixelCNN, modeling the joint distribution, produces sharp samples, but lack a latent representation that might be more useful for downstream tasks. PixelVAE combines the best of both worlds, providing a meaningful latent representation, while producing sharp samples.

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 x i conditioned on all previous dimensions and the latent variable z:

p(x|z) = i p(x i |x 1 , . . . , x i-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(z 1 , • • • , z L ) = p(z L )p(z L-1 |z L ) • • • p(z 1 |z 2 ) q(z 1 , • • • , z L |x) = q(z 1 |x) • • • q(z L |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) = -E z1∼q(z1|x) log p(x|z 1 ) + D KL (q(z 1 , • • • z L |x)||p(z 1 , • • • , z L )) = -E z1∼q(z1|x) log p(x|z 1 ) + z1,••• ,z L L j=1 q(z j |x) L i=1 log q(z i |x) p(z i |z i+1 ) dz 1 ...dz L = -E z1∼q(z1|x) log p(x|z 1 ) + L i=1 z1,••• ,z L L j=1 q(z j |x) log q(z i |x) p(z i |z i+1 ) dz 1 ...dz L = -E z1∼q(z1|x) log p(x|z 1 ) + L i=1 zi,zi+1 q(z i+1 |x)q(z i |x) log q(z i |x) p(z i |z i+1 ) dz i dz i+1 = -E z1∼q(z1|x) log p(x|z 1 ) + L i=1 E zi+1∼q(zi+1|x) D KL (q(z i |x)||p(z i |z i+1 ))

Note that when specifying an autoregressive prior over each latent level z i , we can leverage masked convolutions (van den Oord et al., 2016b) and samples drawn independently from the approximate posterior q(z i | x) (i.e. from the inference network) to train efficiently in parallel on GPUs. et al. (2016b) . In "Gated PixelVAE without upsampling", a linear transformation of latent variable conditions the (gated) activation in every PixelCNN layer instead of using upsampling layers.

4. EXPERIMENTS:

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 D KL (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 z 1 and z 2 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. 

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.

:



Figures and Tables:

Figure 1: Samples from hierarchical PixelVAE on the LSUN bedrooms dataset.

Figure2: Our proposed model, PixelVAE, makes use of PixelCNN to model an autoregressive decoder for a VAE. VAEs, which assume (conditional) independence among pixels, are known to suffer from blurry samples, while PixelCNN, modeling the joint distribution, produces sharp samples, but lack a latent representation that might be more useful for downstream tasks. PixelVAE combines the best of both worlds, providing a meaningful latent representation, while producing sharp samples.

Figure 3: We generate top-down through a hierarchical latent space decomposition. The inference network generates latent variables by composing successive deterministic functions to compute parameters of the stochastic random variables. Dotted lines denote contributions to the cost.

Figure 4: (a) Comparison of Negative log-likelihood upper bound of PixelVAE and NLL for Pixel-CNN as a function of the number of PixelCNN layers used. (b) Cost break down into KL divergence and reconstruction cost.

Figure 5: Visualization of the MNIST test set in the latent space of (a) convolutional VAE and (b) PixelVAE with two latent dimensions. PixelVAE separates classes more completely than VAE.

Figure 6: We visually inspect the variation in image features captured by the different levels of stochasticity in our model. For the two-level latent variable model trained on 64 × 64 LSUN bedrooms, we vary only the top-level sampling noise (top) while holding the other levels constant, vary only the middle-level noise (middle), and vary only the bottom (pixel-level) noise (bottom). It appears that the top-level latent variables learn to model room structure and overall geometry, the middle-level latents model color and texture features, and the pixel-level distribution models low-level image characteristics such as texture, alignment, shading.

Figure 7: Samples from hierarchical PixelVAE on the 64x64 ImageNet dataset.

Figure 8: Reconstructions for (a) LSUN Bedrooms and (b) 64×64 ImageNet.Left-most columns are images from the test set, and the following 5 columns are top-down generations from the highest level of latent variables. We see that the reconstructions capture high-level semantic properties of the original images while varying in most of the details. We also visualized similar reconstructions by generations from the lower level of latent variables, and in this case the reconstructions were visually indistinguishable from the original images.



Table content: <table><tr><td>4.1 MNIST</td><td/></tr><tr><td>Model</td><td>NLL Test</td></tr><tr><td>DRAW (Gregor et al., 2016)</td><td>≤ 80.97</td></tr><tr><td>Discrete VAE (Rolfe, 2016)</td><td>= 81.01</td></tr><tr><td>IAF VAE (Kingma et al., 2016)</td><td>≈ 79.88</td></tr><tr><td colspan="2">PixelCNN (van den Oord et al., 2016a) = 81.30</td></tr><tr><td colspan="2">PixelRNN (van den Oord et al., 2016a) = 79.20</td></tr><tr><td>VLAE (Chen et al., 2016)</td><td>= 79.03</td></tr><tr><td>Convolutional VAE</td><td>≤ 87.41</td></tr><tr><td>PixelVAE</td><td>≤ 80.64</td></tr><tr><td colspan="2">Gated PixelCNN (our implementation) = 80.10</td></tr><tr><td>Gated PixelVAE</td><td>≈ 79.48 (≤ 80.02)</td></tr><tr><td>Gated PixelVAE without upsampling</td><td>≈ 78.96 (≤ 79.58)</td></tr></table>

ImageNet generative modeling task was introduced in (van denOord 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 Table2. Our model achieves a likelihood competitive with van denOord 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. Model performance on 64 × 64 ImageNet. We achieve competitive NLL at a fraction of the computational complexity of other leading models.5 CONCLUSIONSIn 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.

Table content: <table><tr><td>Model</td><td>NLL Validation (Train)</td><td>FLOPs</td></tr><tr><td>Convolutional DRAW (Gregor et al., 2016)</td><td>≤ 4.10 (4.04)</td><td>-</td></tr><tr><td>Real NVP (Dinh et al., 2016)</td><td>= 4.01 (3.93)</td><td>-</td></tr><tr><td>PixelRNN (van den Oord et al., 2016a)</td><td>= 3.63 (3.57)</td><td>154 × 10 9</td></tr><tr><td>Gated PixelCNN (van den Oord et al., 2016b)</td><td>= 3.57 (3.48)</td><td>134 × 10 9</td></tr><tr><td>Hierarchical PixelVAE</td><td>≤ 3.62 (3.55)</td><td>63 × 10 9</td></tr></table>

Peer review (iclr2017):
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.