I MAGE T RANSFORMER
A BSTRACT
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 based on self-attention, the Transformer, to a sequence modeling formulation of image generation with a tractable likelihood.
By applying a self-attention mechanism with a limited receptive ﬁeld multiple times in parallel to different parts of the sequence, we signiﬁcantly increase the length of sequences the model can process efﬁciently.
We propose another simple extension of self-attention to allow it to 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 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 magniﬁcation ratio, applying an encoder-decoder conﬁguration of our architecture.
In a human evaluation study, we show that our super-resolution models improve signiﬁcantly over previously published autoregressive super-resolution models.
Images they generate fool human observers three times more often than the previous state of the art.
Lastly, we provide examples of images generated or completed by our various models which, following previous work, we also believe to look pretty cool.
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
I NTRODUCTION
Recent advances in modeling the distribution of natural images with neural networks allow them to generate increasingly natural-looking images.
Code
Some models, such as the PixelRNN and PixelCNN (van den Oord et al., 2016a), 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).
Typically, 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 then apply recurrent or convolutional neural networks, predicting each next pixel given all previously generated pixels (van den Oord et al., 2016a;b; Salimans et al.).
In this paper we adopt similar factorizations, but following recent work on modeling text (Vaswani et al., 2017), 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).
Despite comparatively low resource requirements for training, the generative Image Transformer reaches a new state of the art in modeling images from the standard ImageNet data set, as measured by log-likelihood.
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 magniﬁcation ratio, we condition on a very low-resolution image, employing the Image Transformer in an encoder-decoder conﬁguration (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 signiﬁcantly more often.
B ACKGROUND
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., 2016a), 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.).
Some of these modiﬁcations can also apply to our model and we plan to evaluate them in future work.
Another, currently wildly popular direction of research in image generation is to train models with an adversarial loss (Goodfellow et al., 2014).
Typically, in this regime a generator network is trained 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 reﬂect the diversity in the training set (Metz et al., 2016).
softmax
FFNN + + + + cmp cmp cmp q’ q mmmp q ppp
MatMul Wv MatMul Wv ·
MatMul Wq MatMul Wk
MatMul Wv
LayerNorm
Dropout
L o ca l S e l f - A tt e n t i o n
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 ﬁrst layer.
A related problem is that GANs do not readily offer a probabilistic interpretation of their outputs.
Not only does this complicate evaluation as it makes it difﬁcult 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 comparing different parameterizations or hyperparameter choices is typically much more involved than comparing log-probabilities assigned to a validation set.
M ODEL A RCHITECTURE 3. I MAGE R EPRESENTATION AND 2D P OSITIONAL I NFORMATION
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-speciﬁc 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 ﬂatten this tensor in rasterscan order, similar to previous work (van den Oord et al., 2016a).
This yields a [ h · w · 3 , d ] tensor.
3. S ELF -A TTENTION
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 ﬁrst 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.
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 ﬁrst 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 rectiﬁed 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).
Equation 1 and Figure1 summarize all operations performed for a single position in each layer.
[Equation]
The entire self-attention operation can be implemented using highly optimized matrix multiplication code and executed in parallel for all pixels’ channels.
In all of our experiments we used multihead attention with a varying number of heads.
For a detailed description of multihead self-attention the reader is referred to (Vaswani et al., 2017).
3. L OCAL S ELF -A TTENTION
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 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.
[Figure]
In our experiments we use two different ways of 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 local 2D attention models, we generate the image query block by query block in raster-scan order, but for query blocks rectangular in the original image space.
Within each block, we generate individual positions again in raster-scan order.
As illustrated on the right side of
Figure 2, we generate the blocks outlined in grey lines left-to-right and top-to-bottom.
This allows us to use 2-dimensional query blocks of a size l q speciﬁed 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 for pixels in the query and memory blocks that have not yet been generated.
As shown in Figure 2, 2D local attention balances horizontal and vertical context much more evenly.
We believe this will have an increasingly positive effect on quality with growing image size.
I NFERENCE
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.
E XPERIMENTS
For all of our experiments we optimize with Adam (Kingma & Ba, 2015), and vary the learning rate as speciﬁed in Vaswani et al.
We train our models on both p100 and k40 GPUs, with batch sizes ranging from 1 to 4 per GPU.
5. G ENERATIVE I MAGE M ODELING
Our unconditioned and class-conditioned image generation models both use 1D local attention, with l q = 256 and a total memory size of 512 .
On CIFAR-10 our best class-conditioned model ( 2 .bits/dim) uses 8 self-attention and feed-forward layers, d = 1024 , 16 attention heads, 2048 dimensions in the feed-forward layers, and a dropout of 0 . 3 .
Our smaller CIFAR-10 models ( 3 .bits/dim) have d = 512 , 1024 dimensions in the feed-forward layers, 8 attention heads and use
[Table]
checkpoint averaging
[Table]
dropout = 0 .
Our state of the art ImageNet unconditioned generation model is signiﬁcantly larger, with 12 self-attention and feed-forward layers, d = 1024 , 4096 -dimensional feed-forward layers, 16 attention heads, and dropout = 0 .
As Table 3 shows, our models improve over various previously proposed models including the PixelRNN and the gated PixelCNN.
On ImageNet we establish a new state of the art of 3.82, which we can improve to 3.81 by averaging the last ten checkpoints.
While the PixelCNN++ achieved signiﬁcantly better log-likelihoods on CIFAR-10 (Salimans et al.), we believe our cherry-picked images for various classes to be of higher perceptual quality.
This further corroborates their ﬁnding that log-likelihood and perceptual quality often do not correlate.
5. C ONDITIONING ON I MAGE C LASS
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 signiﬁcantly higher than that of our unconditioned models.
We show examples in Table 2.
5. I MAGE S UPER -R ESOLUTION
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), our experimental setup is massively under-speciﬁed, enlarging an 8 × 8 pixel image four-fold to 32 × 32 .
This means 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.
We use the Image Transformer in an encoder-decoder conﬁguration, connecting them through attention as described in (Vaswani et al., 2017).
Since the input is an 8 × 8 image, we use 1D attention with one query and one memory block covering the entire image, model dimension d = 512 ,hidden units in the position-wise feed-forward network, four encoder layers and 12 decoder layers.
We train end-to-end, maximizing likelihood.
For both of the following data sets, we resized the image to 8 × 8 pixels for the input and 32 ×pixels for the label using TensorFlow’s area interpolation method.
CelebA We trained on the standard CelebA data set of celebrity faces with cropped boundaries.
Existing automated metrics like pSNR, SSIM and MS-SSIM have been shown to not correlate with perceptual image quality (Dahl et al., 2017).
Instead, we conducted a human evaluation study on
Amazon Mechanical Turk.
Each worker is required to make a binary choice when shown one generated and one real image.
Following (Dahl et al., 2017), we show 50 pairs of images, selected randomly, to 50 workers each.
In our method, workers choose images from our model up to 36 .
The fraction of humans fooled is signiﬁcantly better than the previous state of the art.
[Table] [Table]
of the time
[Table]
a signiﬁcant improvement over previous models.
Sampling temperature of 0 .
9 and 2D local attention maximized perceptual quality as measured by this evaluation.
CIFAR-10 We also trained a super-resolution model on the CIFAR-10 data set.
Our model reached a negative log-likelihood of of 2 .
74 on the test set using 1D local attention.
As seen in Table 6, our
[Table]
model commonly generates plausible looking objects even though the input images seem to barely show any discernible structure beyond coarse shapes.
C ONCLUSION
In this work we demonstrate that models based on self-attention can operate effectively on modalities other than text and through local self-attention scale to signiﬁcantly larger structures than sentences.
We further hope to have provided additional evidence that even in the light of generative adversarial networks, autoregressive generation of images is very much a promising area for further researchas is using network architectures such as the Image Transformer in GANs.
In future work we would like to explore a broader variety of conditioning information including free-form text, as previously proposed (Mansimov et al., 2015) and tasks combining modalities such as language-driven editing of images.
Fundamentally we aim to move beyond still images to video (Kalchbrenner et al., 2016) and towards applications of such models in more model-based reinforcement learning approaches.
C ELEB A S UPERRESOLUTION
Local 2D > Original Original > Local 2D
Local 2D
Original
Local 2D
Original