Title: IMAGE TRANSFORMER

Abstract: Image generation has been successfully cast as an autoregressive sequence generation or transformation problem. Recent work has shown that self-attention is an effective way of modeling textual sequences. In this work, we generalize a recently proposed model architecture based on self-attention, the Transformer, to a sequence modeling formulation of image generation with a tractable likelihood. By restricting the self-attention mechanism to attend to local neighborhoods we significantly increase the size of images the model can process in practice, despite maintaining significantly larger receptive fields per layer than typical convolutional neural networks. We propose another extension of self-attention allowing it to efficiently take advantage of the two-dimensional nature of images. While conceptually simple, our generative models trained on two image data sets are competitive with or significantly outperform the current state of the art in autoregressive image generation on two different data sets, CIFAR-10 and ImageNet. We also present results on image super-resolution with a large magnification ratio, applying an encoder-decoder configuration of our architecture. In a human evaluation study, we show that our super-resolution models improve significantly over previously published autoregressive super-resolution models. Images they generate fool human observers three times more often than the previous state of the art.

:

Table 1 : Three outputs of a CelebA super-resolution model followed by three image completions by a conditional CIFAR-10 model, with input, model output and the original from left to right

1. INTRODUCTION:

Recent advances in modeling the distribution of natural images with neural networks allow them to generate increasingly natural-looking images.

Some models, such as the PixelRNN and PixelCNN (van den Oord et al., 2016) , have a tractable likelihood. Beyond licensing the comparatively simple and stable training regime of directly maximizing log-likelihood, this enables the straightforward application of these models in problems such as image compression (van den Oord & Schrauwen, 2014) and probabilistic planning and exploration (Bellemare et al., 2016) .

The likelihood is made tractable by modeling the joint distribution of the pixels in the image as the product of conditional distributions (Larochelle & Murray, 2011; Theis & Bethge, 2015) . Having thus turned the problem into a sequence modeling problem, the state of the art approaches apply recurrent or convolutional neural networks, predicting each next pixel given all previously generated pixels (van den Oord et al., 2016) . Training recurrent neural networks to sequentially predict each pixel of even a small image is computationally very challenging. Thus, models based on much more parallelizable convolutional neural networks such as the PixelCNN have recently received much more attention, and have now surpassed the PixelRNN in quality PixelCNN.

One disadvantage of CNNs compared to RNNs is their typically fairly limited receptive field. This can adversely affect their ability to model long-range phenomena common in images, such as symmetry and occlusion, especially with a small number of layers. Growing the receptive field has been shown to improve quality significantly (Salimans et al.) . Doing so, however, like deepening the network, comes at a significant cost in number of parameters and consequently computational performance and can make training such models more challenging.

In this work we aim to find a better balance in the trade-off between the virtually unlimited receptive field of the necessarily sequential PixelRNN and the limited receptive field of the much more parallelizable PixelCNN and its various extensions.

We adopt similar factorizations of the joint pixel distribution as previous work. Following recent work on modeling text (Vaswani et al., 2017) , however, we propose eschewing recurrent and convolutional networks in favor of the Image Transformer, a model based entirely on a self-attention mechanism (Cheng et al., 2016; Parikh et al., 2016) . The specific, locally restricted form of multihead self-attention we propose could also be interpreted as a sparsely parameterized form of gated convolution, allowing for significantly larger receptive fields than CNNs at the same number of parameters.

Despite comparatively low resource requirements for training, the Image Transformer attains a new state of the art in modeling images from the standard ImageNet data set, as measured by loglikelihood. Our experiments indicate that increasing the size of the receptive field plays a significant role in this improvement.

Many applications of image density models require conditioning on additional information of various kinds: from images in enhancement or reconstruction tasks such as super-resolution, in-painting and denoising to text when synthesizing images from natural language descriptions (Mansimov et al., 2015) . In visual planning tasks, conditional image generation models could predict future frames of video conditioned on previous frames and taken actions.

In this work we hence also evaluate two different methods of performing conditional image generation with the Image Transformer. In image-class conditional generation we condition on an embedding of one of a small number of image classes. In super-resolution with high magnification ratio, we condition on a very low-resolution image, employing the Image Transformer in an encoder-decoder configuration (Kalchbrenner & Blunsom, 2013) . In comparison to recent work on autoregressive super-resolution (Dahl et al., 2017) , a human evaluation study found images generated by our models look convincingly natural significantly more often.

2. BACKGROUND:

There is a broad variety of types of image generation models in the literature. This work is most strongly inspired by autoregressive models such as fully visible belief networks and NADE (Bengio & Bengio, 2000; Larochelle & Murray, 2011) in that we also factor the joint probability of the image pixels into conditional distributions. Following PixelRNN (van den Oord et al., 2016) , we also model the color channels of the output pixels as discrete values generated from a multinomial distribution, implemented using a simple softmax layer.

The current state of the art in modeling images in the CIFAR-10 data set was achieved by the PixelCNN++, modeling the output pixel distribution with a discretized logistic mixture likelihood, conditioning on whole pixels instead of color channels and changes to the architecture (Salimans et al.) . Most of these modifications can also be applied to our model which we plan to evaluate in future work.

Another, currently wildly popular direction of research in image generation is training models with an adversarial loss (Goodfellow et al., 2014) . Typically, in this regime a generator network is trained

softmax FFNN + + + + cmp cmp cmp q' q m 1 m 2 m 3 p q p 1 p 2 p 3 MatMul W v MatMul W v • MatMul W q MatMul W k MatMul W v LayerNorm Dropout Local Self-Attention

Figure 1 : A slice of one layer of the Image Transformer, recomputing the representation q of a single channel of one pixel q by attending to a memory of previously generated pixels m 1 , m 2 , . . .. We apply a two-layer feed-forward neural network to the weighted average produced by the selfattention mechanism, perform layer normalization and sum the result with a residual connection. The position encodings p q , p 1 , . . . are added only in the first layer.

in opposition to a discriminator network trying to determine if a given image is real or generated.

In contrast to the often blurry images generated by networks trained with likelihood-based losses, such generative adversarial networks (GANs) have been shown to generate sharper images with realistic high-frequency detail in generation and image super-resolution tasks (Zhang et al., 2016; Ledig et al., 2016) .

While very promising, GANs have various drawbacks. They are notoriously unstable (Radford et al., 2015) , motivating a large number of methods attempting to make their training more robust (Metz et al., 2016; Berthelot et al., 2017) . Another common issue is that of mode collapse, where generated images fail to reflect the diversity in the training set (Metz et al., 2016) .

A related problem is that GANs do not readily offer a probabilistic interpretation of their outputs, making it very challenging to measure the degree to which the models capture diversity. In contrast to models with a tractable likelihood, it also complicates optimizing model design, as objectively comparing different parameterizations or hyperparameter choices in this setting is considerably more difficult than comparing log-probabilities assigned to a validation set.

3. MODEL ARCHITECTURE:



3.1. IMAGE REPRESENTATION AND 2D POSITIONAL INFORMATION:

We treat both the input and predicted pixel RGB intensities as categorical variables rather than real numbers. Each input pixel's channel is encoded using a channel-specific set of 256 d-dimensional embedding vectors of the channel intensity values 0 -255. For output intensities, we share a single, separate set of 256 d-dimensional embeddings across channels.

We then combine the width and channels dimensions, yielding, for an image of width w and height h, a 3-dimensional tensor with shape [h, w • 3, d].

To each pixel representation, we add a d-dimensional encoding of the coordinates of that pixel. Following Vaswani et al. (2017) , the encoding consists of sine and cosine functions of the coordinates, with different frequencies across different dimensions. Since we need to represent two coordinates, we use d/2 of the dimensions to encode the row number and the other d/2 of the dimensions to encode the the column and color channel.

The resulting tensor forms the input to our 2D local attention models (Section 3.3). For 1D local attention (Section 3.3) and the input to our super-resolution models we flatten this tensor in rasterscan order, similar to previous work (van den Oord et al., 2016) . This yields a [h • w • 3, d] tensor.

3.2. SELF-ATTENTION:

Like the Transformer (Vaswani et al., 2017) , the Image Transformer uses stacks of self-attention and position-wise feed-forward layers. Before we describe how we scale self-attention from sentences to images, which contain many more positions, we give a brief description of the self-attention layer.

Each self-attention layer computes a new d-dimensional representation for each position, that is each channel of each pixel. To recompute the representation for a given position, it first compares the position's current representation to other positions' representations, obtaining an attention distribution over the other positions. This distribution is then used to weight the contribution of the other postions' representations to the next representation for the position at hand.

Equation 1 and Figure1 fully describe all operations performed in every layer, independently for each position, with the exception of multi-head attention. For a detailed description of multi-head self-attention the reader is referred to (Vaswani et al., 2017) .

EQUATION

In more detail, following previous work, we call the current representation of the pixel's channel, or position, to be recomputed the query q. The other positions whose representations will be used in computing a new representation for q are m 1 , m 2 , . . . which together comprise the columns of the memory matrix M . Note that M can also contain q. We first transform q and M linearly by learned matrices W q and W k , respectively.

The self-attention mechanism then compares q to each of the pixel's channel representations in the memory with a dot-product, scaled by 1/ √ d. We apply the softmax function to the resulting compatibility scores, treating the obtained vector as attention distribution over the pixel channels in the memory. After applying another linear transformation W v to the memory M , we compute a weighted average of the transformed memory, weighted by the attention distribution. In the decoders of our different models we mask the outputs of the comparisons appropriately so that the model cannot attend to positions in the memory that have not been generated, yet.

To the resulting vector we then apply a single-layer fully-connected feed-forward neural network with rectified linear activation followed by another linear transformation. The learned parameters of these are shared across all positions but different from layer to layer. Lastly, we perform layer normalization followed by dropout (Ba et al., 2016; Srivastava et al., 2014) .

The entire self-attention operation can be implemented using highly optimized matrix multiplication code and executed in parallel for all pixels' channels.

3.3. LOCAL SELF-ATTENTION:

The number of positions included in the memory l m , or the number of columns of M , has tremendous impact on the scalability of the self-attention mechanism, which has a time complexity in

O(h • w • l m • d).

The encoders of our super-resolution models operate on 8 × 8 pixel images and it is computationally feasible to attend to all of their 192 positions. The decoders in our experiments, however, produce 32 × 32 pixel images with 3072 positions, rendering attending to all positions impractical.

Inspired by convolutional neural networks we address this by adopting a notion of locality, restricting the positions in the memory matrix M to a local neighborhood around the query position. Changing this neighborhood per query position, however, would prohibit packing most of the computation

Memory Block:



Local 1D Attention:

Memory Block q q Query Block Query Block

Local 2D Attention:



Figure 2 : The two different conditional factorizations used in our experiments, with 1D and 2D local attention on the left and right, respectively. In both, the image is partitioned into non-overlapping query blocks, each associated with a memory block covering a superset of the query block pixels. In every self-attention layer, each position in a query block attends to all positions in the memory block. The pixel marked as q is the last that was generated. All channels of pixels in the memory and query blocks shown in white have masked attention weights and do not contribute to the next representations of positions in the query block. While the effective receptive field size in this figure is the same for both schemes, in 2D attention the memory block contains a more evenly balanced number of pixels next to and above the query block, respectively. necessary for self-attention into two matrix multiplications -one for computing the pairwise comparisons and another for generating the weighted averages. To avoid this, we partition the image into query blocks and associate each of these with a larger memory block that also contains the query block. For all queries from a given query block, the model attends to the same memory matrix, comprised of all positions from the memory block.

The self-attention is then computed for all query blocks in parallel, while the feed-forward networks and layer normalizations are computed in parallel for all positions.

In our experiments we use two different schemes for choosing query blocks and their associated memory block neighborhoods, resulting in two different factorizations of the joint pixel distribution into conditional distributions. Both are illustrated in Figure 2 .

1D Local Attention:

To compute self-attention on raster-scanned linearized images, we partition the length into non-overlapping query blocks Q of length l q , padding with zeroes if necessary. While contiguous in the linearized image, these blocks can be discontiguous in image coordinate space. For each query block we build the memory block M from the same positions as Q and an additional l m positions from pixels that have been generated before, which can result in overlapping memory blocks.

2D Local Attention:

In 2D local attention models, we partition the image into query blocks rectangular and contiguous in the original image space. We generate the image one query block after another, ordering the blocks in raster-scan order. Within each block, we generate individual positions, or pixel channels, again in raster-scan order.

As illustrated in the right half of Figure 2 , we generate the blocks outlined in grey lines left-to-right and top-to-bottom. We use 2-dimensional query blocks of a size l q specified by height and width l q = w q • h q , and memory blocks extending the query block to the top, left and right by h m , w m and again w m pixels, respectively.

In both 1D and 2D local attention, we mask attention weights in the query and memory blocks such that positions that have not yet been generated are ignored.

As can be seen in Figure 2 , 2D local attention balances horizontal and vertical conditioning context much more evenly. We believe this might have an increasingly positive effect on quality with growing image size as the conditioning information in 1D local attention becomes increasingly dominated by pixels next to a given position as opposed to above it.

Table 2 : Conditional image generations for all CIFAR-10 categories. Images on the left are from a model that achieves 3.03 bits/dim on the test set. Images on the right are from our best nonaveraged model with 2.99 bits/dim. Both models are able to generate convincing cars, trucks, and ships. Generated horses, planes, and birds also look reasonable.

4. INFERENCE:

Across all of the presented experiments, we sample from the various models with a tempered softmax (Dahl et al., 2017) . We adjust the concentration of the distribution we sample from with a temperature τ > 0 by which we divide the logits for the channel intensities.

We tuned τ between 0.8 and 1.0, observing the highest perceptual quality in unconditioned and classconditional image generation with τ = 1.0. For super-resolution we present results for different temperatures in Table 4 .

5. EXPERIMENTS:

For all of our experiments we optimize with Adam (Kingma & Ba, 2015) , and vary the learning rate as specified in Vaswani et al. (2017) . We train our models on both p100 and k40 GPUs, with batch sizes ranging from 1 to 4 per GPU. As Table 3 shows, our models improve over various previously proposed models including the Pix-elRNN and the gated PixelCNN. On ImageNet we establish a new state of the art of 3.78, which we can improve to 3.77 by averaging the last ten checkpoints.

While the PixelCNN++ achieved significantly better log-likelihoods on CIFAR-10 (Salimans et al.), we expect that many of the modifications in the PixelCNN++ carry over to the Image Transformer. We further believe our curated images for various classes to be of reasonable perceptual quality.

5.2. CONDITIONING ON IMAGE CLASS:

We represent the image classes as learned d-dimensional embeddings per class and simply add the respective embedding to the input representation of every input position together with the positional encodings.

We trained the class-conditioned Image Transformer on CIFAR-10 and ImageNet data sets, achieving very similar log-likelihoods as in unconditioned generation. The perceptual quality of generated images, however, is significantly higher than that of our unconditioned models. We present some samples in Table 2 .

5.3. IMAGE SUPER-RESOLUTION:

Super-resolution is the process of recovering a high resolution image from a low resolution image while generating realistic and plausible details. Following (Dahl et al., 2017) , in our experimental setup we enlarge an 8×8 pixel image four-fold to 32×32, a process that is massively underspecified: the model has to generate aspects such as texture of hair, makeup, skin and sometimes even gender that cannot possibly be recovered from the source image.

Here, we use the Image Transformer in an encoder-decoder configuration, connecting the encoder and decoder through an attention mechanism (Vaswani et al., 2017) . Since the input is an 8 × 8 image, it is practical to use 1D attention with only one query and one memory block, each covering the entire image. We further use model dimension d = 512, 1024 hidden units in the position-wise Input 1D Local Attention 2D Local Attention Original τ = 0.8 τ = 0.9 τ = 1.0 τ = 0.8 τ = 0.9 τ = 1.0 Table 5: Images from our 1D and 2D local attention super-resolution models trained on CelebA, sampled with different temperatures. 2D local attention with τ = 0.9 scored highest in our human evaluation study.



Figures and Tables:

Negative log-likelihoods on the CIFAR-10 test and ImageNet validation sets. The Image Transformer outperforms all models but PixelCNN++, achieving a new state of the art on ImageNet. Larger memory blocks significantly improve its performance.

Table content: <table><tr><td>Model Type</td><td>Memory Block Size</td><td/><td>NLL</td></tr><tr><td/><td/><td colspan="2">CIFAR-10 (Test) ImageNet (Validation)</td></tr><tr><td>Pixel CNN</td><td>-</td><td>3.14</td><td>-</td></tr><tr><td>Row Pixel RNN</td><td>-</td><td>3.00</td><td>3.86</td></tr><tr><td>Gated Pixel CNN</td><td>-</td><td>3.03</td><td>3.83</td></tr><tr><td>Pixel CNN++</td><td>-</td><td>2.92</td><td>-</td></tr><tr><td colspan="2">Image Transformer 1D local 8</td><td>4.06</td><td>-</td></tr><tr><td/><td>16</td><td>3.47</td><td>-</td></tr><tr><td/><td>64</td><td>3.13</td><td>-</td></tr><tr><td/><td>256</td><td>2.99</td><td>3.78</td></tr><tr><td>with checkpoint averaging</td><td>256</td><td>2.98</td><td>3.77</td></tr><tr><td colspan="2">5.1 GENERATIVE IMAGE MODELING</td><td/><td/></tr><tr><td colspan="4">Our unconditioned and class-conditioned image generation models both use 1D local attention, with</td></tr><tr><td>l</td><td/><td/><td/></tr></table>

Peer review (iclr2018):
 Summary This paper extends self-attention layers (Vaswani et al., 2017) from sequences to images and proposes to use the layers as part of PixelCNNs (van den Oord et al., 2016). The proposed model is evaluated in terms of visual appearance of samples and log-likelihoods. The authors find a small improvement in terms of log-likelihood over PixelCNNs and that super-resolved CelebA images are able to fool human observers significantly more often than PixelRNN based super-resolution (Dahl et al., 2017). Review Autoregressive models are of large interest to the ICLR community and exploring new architectures is a valuable contribution. Using self-attention in autoregressive models is an intriguing idea. It is a little bit disappointing that the added model complexity only yields a small improvement compared to the more straight-forward modifications of the PixelCNN++. I think the paper would benefit from a little bit more work, but I am open to adjusting my score based on feedback. I find it somewhat surprising that the proposed model is only slightly better in terms of log-likelihood than a PixelRNN, but much better in terms of human evaluation – given that both models were optimized for log-likelihood. Was the setup used with Mechanical Turk exactly the same as the one used by Dahl et al.? These types of human evaluations can be extremely sensitive to changes in the setup, even the phrasing of the task can influence results. E.g., presenting images scaled differently can mask certain artifacts. In addition, the variance between subjects can be very high. Ideally, each method included in the comparison would be re-evaluated using the same set of observers. Please include error bars. The CelebA super-resolution task furthermore seems fairly limited. Given the extreme downsampling of the input, the task becomes similar to simply generating any realistic image. A useful baseline would be the following method: Store the entire training set. For a given query image, look for the nearest neighbor in the downsampled space, then return the corresponding high-resolution image. This trivial method might not only perform well, it also highlights a flaw in the evaluation: Any method which returns stored high-resolution images – even if they don’t match the input – would perform at 50%. To fix this, the human observers should also receive the low-resolution image and be asked to identify the correct corresponding high-resolution image. Using multiplicative operations to model images seems important. How does the self-attention mechanism relate to “gated” convolutions used in PixelCNNs? Could gated convolutions not also be considered a form of self-attention? The presentation/text could use some work. Much of the text assumes that the reader is familiar with Vaswani et al. (2017) but could easily be made more self-contained by directly including the definitions used. E.g., the encoding of positions using sines and cosines or the multi-head attention model. I also felt too much of the architecture is described in prose and could be more efficiently and precisely conveyed in equations. On page 7 the authors write “we believe our cherry-picked images for various classes to be of higher perceptual quality”. This is a meaningless result, not only because the images were cherry-picked. Generating realistic images is trivial - you just need to store the training images. Analyzing samples generated by a generative model (outside the context of an application) should therefore only be used for diagnostic purposes or to build intuitions but not to judge the quality of a model. Please consider rephrasing the last sentence of the abstract. Generating images which “look pretty cool” should not be the goal of a serious machine learning paper or a respected machine learning conference.