Spatial PixelCNN: Generating Images from Patches

Variational autoencoders (Kingma & Welling, 2013) are a form of latent variable model. They are trained to encode their input $\mathbf{x}$ into a latent code $\mathbf{z}$. The true posterior $p(\mathbf{z}|\mathbf{x})$ is often intractable, so an approximate posterior $q(\mathbf{z}|\mathbf{x})$ is computed by optimizing a variational lower bound on the log-likelihood of the data. Often this latent code is interpretable and encodes a global representation of $\mathbf{x}$.

(Gulrajani et al., 2016; Chen et al., 2016) have evaluated integrating a PixelCNN decoder into a VAE framework. Gulrajani et al., 2016 note that it is critical to ensure the receptive field of the PixelCNN is smaller than the input image $\mathbf{x}$. Otherwise, a powerful PixelCNN decoder alone can model the entire image distribution, rendering the latent code $\mathbf{z}$ from the VAE unused. When evaluated on binarized MNIST (Larochelle & Murray, 2011), both report the successful combination yields state-of-the-art results.

Spatial PixelCNN indeed fits the criteria needed to ensure utilization of the latent code $\mathbf{z}$, as it only trains on a patch of an image at a time. In fact the inclusion of a latent variable model, in this case a VAE, is necessary to ensure the model has a global representation of the images.

Compositional Pattern Producing Networks (CPPNs) (Stanley, 2007) are feedforward networks that utilize a wide array of transfer functions and are often trained via evolutionary algorithms. CPPNs can encode 2-D images (Nguyen et al., 2015), 3-D objects (Clune & Lipson, 2011) and even weights of another target network (Stanley et al., 2009).

To encode a 2-D color image, a CPPN performs a transformation $f:\mathbb{R}^{2}\to\mathbb{R}^{3}$ parameterized by a network which takes as input a pair of coordinates $(i,j)$ and outputs 3 color values (e.g. RGB) (Secretan et al., 2008). A pair of coordinates $(i,j)$ is often computed by evenly sampling a pre-defined range e.g. $[-1,1]$. The coordinates are assembled into a grid, each corresponding to a pixel in the training image. At image generation time, an image can be rendered at an arbitrarily large resolution by simply sampling more coordinates within the specified range, and querying the CPPN for the associated color value.

CPPNs have been shown to impose a strong spatial prior yielding images of highly regular patterns (Secretan et al., 2008; Nguyen et al., 2015). It is this spatial regularity from CPPNs that we exploit in this work. We find that conditioning Spatial PixelCNN on a grid of coordinates is crucial for ensuring a coherent image is generated, and also confers an ability to upscale images.

Since directly modeling high resolution images is challenging, an effective approach is to generate an image in hierarchical stages of increasing resolutions (Zhang et al., 2016; Denton et al., 2015; Karras et al., 2017). While producing impressive results, this approach outputs fixed-sized images, and requires storing the entire image in GPU memory at once—this is problematic when the training images exceed GPU memory capacity. Spatial PixelCNN instead allows for choosing the target image size at generation time, and only processes a small patch of the image at a time.

To cut down on GPU memory requirements, Tyka, 2017 devises upscaling image tiles, and stitching them together to produce the final result. While requiring less memory, the approach still retains the reliance on generating a fixed-sized output.

Our framework most closely resembles Ha, 2016, which combines spatial coordinates, and the latent code from a VAE, trained as a GAN end-to-end on full-sized images. Our model differs in two important ways: (1) Spatial PixelCNN conditions each pixel on the latent code and a local neighborhood of preceding pixels rather than assuming conditional independence given the latent code; (2) Spatial PixelCNN is trained on patches rather than full-sized images.

Our proposed model is based on a modified version of PixelCNN++ (Salimans et al., 2017). In order to reduce cumbersome mathematical notation, we restate Eq. 1 as:

where $i$ denotes the virtual index of the pixel, if the image were flattened into a 1-D array. Rather than training the model on full-sized $n\times n$ images $\mathbf{x}$, we instead choose random patches $\mathbf{y}$ of size $m\times m$ taken from the images:

Given that our goal is to model the set of images $\mathbf{X}$, rather than merely a collection of all patches ${\mathbf{Y}_{m\times m}}$, we condition on a normalized coordinate grid ${\mathbf{G}_{n\times n}}$ $\in\mathbb{R}^{2\times n\times n}$ that has the same number of coordinate pairs as pixels in the full-sized image. Specifically, ${\mathbf{G}_{n\times n}}$ is composed of $n\times n$ evenly spaced coordinate pairs $(i,j)$ within the range $[-1,1]$. We choose patches $\mathbf{g}$ from ${\mathbf{G}_{n\times n}}$ corresponding to the given patch $\mathbf{y}$. To condition each pixel $y_{i}$ on its corresponding coordinate pair $g_{i}$, we make use of the gated spatial conditioning as outlined in van den Oord et al., 2016 arriving at the conditional probability:

We found such spatial conditioning crucial to imparting coherence to the patches $\mathbf{y}$. Without conditioning on $\mathbf{g}$, the model may assign a high probability to distinct patches of the generated image, allowing juxtapositions (Fig. 3). However, as the model is trained on patches $\mathbf{y}$ of fewer dimensions, the extra spatial information provided by $\mathbf{g}$ is not enough to ensure global coherence (Fig. 3(c)).

In order to provide this global coherence, we also condition on a latent code, provided by a VAE. Previous treatments (Gulrajani et al., 2016; Chen et al., 2016) combinining PixelCNN and VAE have made use of PixelCNN as a decoder. We found that by only training on patches, the VAE was unable to capture global features (data not shown). Instead, we jointly train the two models: VAE on images $\mathbf{X}$ and Spatial PixelCNN on patches ${\mathbf{Y}_{m\times m}}$.

Specifically, we minimize the sum of an image loss and a patch loss:

The image loss $\mathcal{L}^{x}$ is the typical VAE loss, that of the negative log-likelihood of the images and the KL-divergence of the approximate posterior from the prior:

The patch loss $\mathcal{L}^{y}$ is the negative log-likelihood of the patches, given the coordinate grid $\mathbf{g}$ and latent $\mathbf{z}$:

Note that instead of co-training both VAE and Spatial PixelCNN, it is also possible to pre-train the VAE first, and then train the Spatial PixelCNN. This pre-trained VAE typically has a lower KL-divergence than the co-trained approach. Global features extracted from the pre-trained VAE also allow for successful reconstructions, though we do not evaluate its efficacy in this paper.

As our model is based on patches of images, a sliding window is used during image generation. In order to ensure that each pixel being generated has the maximal number of preceding pixels to condition on, the sliding window moves one pixel at a time from left to right, top to bottom (Fig. 4). This is the same ordering defined by the model for pixels to condition on (Eq. 1). Only the maximally conditioned pixels are generated from each patch of the sliding window.

One of the unique aspects of our model is its ability to upscale images to a higher resolution, while only being trained on lower resolution images. To accomplish this, we condition on coordinate patches from a larger coordinate grid during the generation process. That is, when we generate a $56\times 56$ image, we subdivide the grid into $56\times 56$ evenly spaced steps. Then as we slide the window over the image we wish to generate, we condition on the associated patch from ${\mathbf{G}_{56\times 56}}$.

We use a modified version of PixelCNN++ (Salimans et al., 2017). Our model similarly makes use of six blocks of ResNet (He et al., 2016) layers. Spatial PixelCNN is conditioned on both a latent vector and a regularized coordinate grid. For each change in dimensions between blocks, the coordinate grid is resampled to the new layer’s dimensions. This resampling is performed as a simple linear interpolation. Given a grid patch defined by $m\times m$ coordinates in the range $[(i_{1},j_{1}),(i_{m},j_{m})]$, we linearly interpolate the range into $\frac{m}{2}\times\frac{m}{2}$ equal sized steps. This regularization is key to preserving spatial coherence between ResNet blocks.

We additionally make use of ResNet blocks for the encoder and decoder of our VAE. For each ResNet layer of the decoder, we condition on the latent vector $\mathbf{z}$.

For most of our MNIST (LeCun et al., 1998) experiments we use two ResNet layers for each block, with a convolution filter size of $25$. We term this the small network. In order to determine the trade-off between the model’s ability to compress data (as determined by having a smaller expected bits per dimension), versus its ability to effectively upscale images, we also conduct an experiment with a large network. This large network utilizes five ResNet layers for each block, with a convolution filter size of $140$. The size of the latent code is fixed at $60$ for both network sizes.

We make use of the Tensorflow (Abadi et al., 2015) framework for training and evaluating our models. We train using the Adam (Kingma & Ba, 2014) optimizer, with a batch size of $128$, and an initial learning rate of of $0.001$ which is annealed using exponential decay at a rate of $0.999995$ every batch. Additionally a dropout rate of $0.5$, along with L2 regularization of $0.0001$ is utilized. We train each model until convergence, as determined by a lack of improvement over one hundred epochs in the model’s bits per dimension.

At testing and image generation time we use the exponential moving average over the model parameters (Polyak & Juditsky, 1992), calculated with a decay of $0.9995$ at each parameter update. We report our negative log-likelihood values in bits per dimension. The reported bits per dimension is calculated as the average bits per dimension over all possible patches for the images in the test set.

As previously noted (Theis et al., 2015), the negative log-likelihood may not be a great qualitative measure for evaluating generative models. This can be seen clearly in our experiments, where bits per dimension does not accurately correlate with the ability of the model to generate coherent images (Fig. 3).

MS-SSIM: We first assess the diversity of the images generated by utilizing Multiscale Structural Similarity (Wang et al., 2003; Odena et al., 2016) (MS-SSIM). To calculate the MS-SSIM for a given model, we begin by randomly sampling $500$ images from the model. We produce an MS-SSIM score for all unique pairs of images. We then report the mean and standard deviation of these scores. The lower the MS-SSIM, the more diverse the images in the set. The ideal MS-SSIM of a given model should closely resemble that of the underlying data, so we also report an MS-SSIM for $500$ images from the MNIST test set for comparison.

While a measure of diversity shows the model does not exhibit mode collapse, it is unable to address the subjective assessment of image quality. Such an assessment should include the ability of the model to accurately reconstruct a given input, as well as ensuring generated samples reflect the underlying data distribution. Thus we devise two additional measures.

LeNet Accuracy: As the combination of PixelCNN and VAE allows for conditioning on an interpretable latent representation, we can assess reconstruction accuracy for these models. We take inspiration from the use of the Inception model (Odena et al., 2016) for assessing accuracy. As our model is trained on MNIST, we instead train a version of LeNet (LeCun et al., 1998), which has demonstrated strong classification ability for this dataset. We reconstruct $1000$ images from the test set using our models, then measure the classification accuracy for these reconstructions. For comparison, we also report the classification accuracy for the original $1000$ images from the test set.

LeNet Confidence: In order to assess random samples generated from the model, we devise a simple confidence score inspired by the Inception Score (Salimans et al., 2016). As part of the Inception Score measures diversity, we reformulate our metric to only measure the conditional probability of the label, given an image. For a given image, we calculate the softmax of the LeNet logits. The index of the largest value indicates the predicted class. We take the largest value multiplied by one hundred as a confidence measure, indicating how confident the model is in the prediction. We report the mean and standard deviation of this value for $1000$ random samples from each model.

Given that our model additionally shows strong ability to upscale images to higher resolutions, we also report results for generating $56\times 56$ images. For both classification accuracy and confidence, we first downsample images to $28\times 28$ before running the LeNet classifier. We note that individually these metrics are imperfect, especially when assessing downsampled images, but when taken in aggregate they yield a more holistic assessment.

We first assess the importance of spatial conditioning on the generative ability of PixelCNN++ (Table 2). We train PixelCNN++ on various patch sizes, both with and without spatial conditioning. We note that even when trained on full-sized images, conditioning on a grid of coordinates boosts both MS-SSIM and LeNet Confidence (Table 2 Row 2). This implies the addition of coordinates may help capture structure versus a baseline PixelCNN++.

As we sweep across patch sizes, we see the addition of coordinates keeps the MS-SSIM within range of the underlying dataset. Though, we note as patch size decreases, we observe a decrease in confidence scores for the $28\times 28$ samples, and associated coherence of the resultant images (Fig. 3(c)).

An important observation from the experiments is that the trained bits per dimension value is not a great representation of the quality of random samples generated from the models (Figs. 3(b) & 3(c)). In the case of training on $4\times 4$ patches (Table 2 Row 2), the model achieves a similar ability to compress the data (as denoted by a low bits per dimension), as training on $20\times 20$ patches (Table 2 Row 2). Though, comparing the LeNet Confidence of $28\times 28$ generated images, the $20\times 20$ patches clearly model the underlying dataset more accurately (training on $20\times 20$ patches achieves a Confidence of $93.7\pm 1.7$, while training on $4\times 4$ patches only results in a score of $79.1\pm 5.5$).

We next consider the effects provided by the addition of a latent code. Even without spatial conditioning, the model trained on $16\times 16$ patches produces higher confidence scores with the addition of a VAE (Table 2 Row 2 & Table 2 Row 2). Though, the model has trouble with scale (Fig. 3(d)). This follows our intuition that spatial conditioning is key to capturing regularity across images.

We now examine Spatial PixelCNN conditioned on both coordinates and a latent code. While a direct comparison is difficult, we note that Spatial PixelCNN produces samples qualitatively similar to PixelCNN++ (Fig. 1). Additionally Spatial PixelCNN trained on $8\times 8$ patches receives a similar Confidence score to the PixelCNN++ baseline when generating $28\times 28$ images (Table 2 Row 2 & Table 2 Row 2).

An interesting result can be observed when comparing the small network trained on $8\times 8$ patches versus the small network trained on $16\times 16$ patches. Given an equivalent network size, training on smaller patches produces superior results (Table 2 Rows 2 & 2). Only with the addition of more parameters, through the use of the large network, does training on $16\times 16$ patches produce a higher confidence and accuracy (Table 2 Rows 2 & 2) when generating $28\times 28$ images.

Finally we assess the ability for the network to generate high resolution images while only training on a low resolution dataset. Looking at the generated $56\times 56$ images, the small network trained on $8\times 8$ patches works exceedingly well for upscaling (Fig. 1(b)). It even tends to retain structural detail better than the large network trained on $16\times 16$ patches (Fig. 1(c)). Not only does it surpass the accuracy and confidence scores of the large network (Table 2 Rows 2 & 2 and Table 3 Rows 3 & 3), it only loses a small amount of accuracy and confidence, despite upscaling $2\times$. It even shows remarkable ability for large upscaling factors (Figs. 2, 5(e), & 5(f)). See Appendix A for examples of $20\times$ and $50\times$ upscaling.