Image Transformer

Niki Parmar * 1 Ashish Vaswani * 1 Jakob Uszkoreit 1

Łukasz Kaiser 1 Noam Shazeer 1 Alexander Ku 2 3 Dustin Tran 4

Abstract
Image generation has been successfully cast as
arXiv:1802.05751v3 [cs.CV] 15 Jun 2018

an autoregressive sequence generation or trans-

formation problem. Recent work has shown that
self-attention is an effective way of modeling tex-
tual 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 neighbor-
hoods we significantly increase the size of im- Table 1. Three outputs of a CelebA super-resolution model fol-
ages the model can process in practice, despite lowed by three image completions by a conditional CIFAR-10
maintaining significantly larger receptive fields model, with input, model output and the original from left to right
per layer than typical convolutional neural net-
works. While conceptually simple, our generative
models significantly outperform the current state a tractable likelihood. Beyond licensing the comparatively
of the art in image generation on ImageNet, im- simple and stable training regime of directly maximizing
proving the best published negative log-likelihood log-likelihood, this enables the straightforward application
on ImageNet from 3.83 to 3.77. We also present of these models in problems such as image compression
results on image super-resolution with a large (van den Oord & Schrauwen, 2014) and probabilistic plan-
magnification ratio, applying an encoder-decoder ning and exploration (Bellemare et al., 2016).
configuration of our architecture. In a human eval-
The likelihood is made tractable by modeling the joint dis-
uation study, we find that images generated by
tribution of the pixels in the image as the product of condi-
our super-resolution model fool human observers
tional distributions (Larochelle & Murray, 2011; Theis &
three times more often than the previous state of
Bethge, 2015). Thus turning the problem into a sequence
the art.
modeling problem, the state of the art approaches apply
recurrent or convolutional neural networks to predict each
1. Introduction next pixel given all previously generated pixels (van den
Oord et al., 2016a). Training recurrent neural networks
Recent advances in modeling the distribution of natural to sequentially predict each pixel of even a small image
images with neural networks allow them to generate increas- is computationally very challenging. Thus, parallelizable
ingly natural-looking images. Some models, such as the models that use convolutional neural networks such as the
PixelRNN and PixelCNN (van den Oord et al., 2016a), have PixelCNN have recently received much more attention, and
*
Equal contribution. Ordered by coin flip. 1 Google Brain, have now surpassed the PixelRNN in quality (van den Oord
Mountain View, USA 2 Department of Electrical Engineering and et al., 2016b).
Computer Sciences, University of California, Berkeley 3 Work
done during an internship at Google Brain 4 Google AI, Mountain One disadvantage of CNNs compared to RNNs is their
View, USA. Correspondence to: Ashish Vaswani, Niki Parmar, typically fairly limited receptive field. This can adversely
Jakob Uszkoreit <[email protected], [email protected], affect their ability to model long-range phenomena common
[email protected]>. in images, such as symmetry and occlusion, especially with
Proceedings of the 35 th International Conference on Machine a small number of layers. Growing the receptive field has
Learning, Stockholm, Sweden, PMLR 80, 2018. Copyright 2018 been shown to improve quality significantly (Salimans et al.).
by the author(s). Doing so, however, comes at a significant cost in number
 Image Transformer

of parameters and consequently computational performance 2. Background

and can make training such models more challenging.
There is a broad variety of types of image generation mod-
In this work we show that self-attention (Cheng et al., 2016; els in the literature. This work is strongly inspired by au-
Parikh et al., 2016; Vaswani et al., 2017) can achieve a better toregressive models such as fully visible belief networks
balance in the trade-off between the virtually unlimited and NADE (Bengio & Bengio, 2000; Larochelle & Mur-
receptive field of the necessarily sequential PixelRNN and ray, 2011) in that we also factor the joint probability of
the limited receptive field of the much more parallelizable the image pixels into conditional distributions. Following
PixelCNN and its various extensions. PixelRNN (van den Oord et al., 2016a), we also model the
We adopt similar factorizations of the joint pixel distribu- color channels of the output pixels as discrete values gener-
tion as previous work. Following recent work on model- ated from a multinomial distribution, implemented using a
ing text (Vaswani et al., 2017), however, we propose es- simple softmax layer.
chewing recurrent and convolutional networks in favor of The current state of the art in modeling images on CIFAR-
the Image Transformer, a model based entirely on a self- 10 data set was achieved by PixelCNN++, which models the
attention mechanism. The specific, locally restricted form output pixel distribution with a discretized logistic mixture
of multi-head self-attention we propose can be interpreted likelihood, conditioning on whole pixels instead of color
as a sparsely parameterized form of gated convolution. By channels and changes to the architecture (Salimans et al.).
decoupling the size of the receptive field from the num- These modifications are readily applicable to our model,
ber of parameters, this allows us to use significantly larger which we plan to evaluate in future work.
receptive fields than the PixelCNN.
Another, popular direction of research in image generation
Despite comparatively low resource requirements for train- is training models with an adversarial loss (Goodfellow
ing, the Image Transformer attains a new state of the art et al., 2014). Typically, in this regime a generator net-
in modeling images from the standard ImageNet data set, work is trained in opposition to a discriminator network
as measured by log-likelihood. Our experiments indicate trying to determine if a given image is real or generated. In
that increasing the size of the receptive field plays a sig- contrast to the often blurry images generated by networks
nificant role in this improvement. We observe significant trained with likelihood-based losses, generative adversar-
improvements up to effective receptive field sizes of 256 ial networks (GANs) have been shown to produce sharper
pixels, while the PixelCNN (van den Oord et al., 2016b) images with realistic high-frequency detail in generation
with 5x5 filters used 25. and image super-resolution tasks (Zhang et al., 2016; Ledig
Many applications of image density models require condi- et al., 2016).
tioning on additional information of various kinds: from im- While very promising, GANs have various drawbacks. They
ages in enhancement or reconstruction tasks such as super- are notoriously unstable (Radford et al., 2015), motivating
resolution, in-painting and denoising to text when synthesiz- a large number of methods attempting to make their train-
ing images from natural language descriptions (Mansimov ing more robust (Metz et al., 2016; Berthelot et al., 2017).
et al., 2015). In visual planning tasks, conditional image Another common issue is that of mode collapse, where gen-
generation models could predict future frames of video con- erated images fail to reflect the diversity in the training set
ditioned on previous frames and taken actions. (Metz et al., 2016).
In this work we hence also evaluate two different methods A related problem is that GANs do not have a density in
of performing conditional image generation with the Im- closed-form. This makes it challenging to measure the
age Transformer. In image-class conditional generation we degree to which the models capture diversity. This also
condition on an embedding of one of a small number of complicates model design. Objectively evaluating and com-
image classes. In super-resolution with high magnification paring, say, different hyperparameter choices is typically
ratio (4x), we condition on a very low-resolution image, much more difficult in GANs than in models with a tractable
employing the Image Transformer in an encoder-decoder likelihood.
configuration (Kalchbrenner & Blunsom, 2013). In com-
parison to recent work on autoregressive super-resolution
(Dahl et al., 2017), a human evaluation study found im-
3. Model Architecture
ages generated by our models to look convincingly natural 3.1. Image Representation
significantly more often.
We treat pixel intensities as either discrete categories or ordi-
nal values; this setting depends on the distribution (Section
3.4). For categories, each of the input pixels’ three color
channels is encoded using a channel-specific set of 256
 Image Transformer

Input Gen Truth Input Gen Truth

Table 2. On the left are image completions from our best conditional generation model, where we sample the second half. On the right are
samples from our four-fold super-resolution model trained on CIFAR-10. Our images look realistic and plausible, show good diversity
among the completion samples and observe the outputs carry surprising details for coarse inputs in super-resolution.

d-dimensional embedding vectors of the intensity values

q’
0 − 255. For output intensities, we share a single, separate
set of 256 d-dimensional embeddings across the channels. LayerNorm
For an image of width w and height h, we combine the
Dropout
width and channel dimensions yielding a 3-dimensional
tensor with shape [h, w · 3, d]. LayerNorm FFNN

For ordinal values, we run a 1x3 window size, 1x3 strided Dropout
convolution to combine the 3 channels per pixel to form an
input representation with shape [h, w, d].
To each pixel representation, we add a d-dimensional en-
Local Self-Attention

coding of coordinates of that pixel. We evaluated two dif- MatMul Wv MatMul Wv MatMul Wv
ferent coordinate encodings: sine and cosine functions of
softmax
the coordinates, with different frequencies across different
cmp cmp cmp
dimensions, following (Vaswani et al., 2017), and learned
·
position embeddings. Since we need to represent two co-
+ + + + MatMul Wq MatMul Wk
ordinates, we use d/2 of the dimensions to encode the row
number and the other d/2 of the dimensions to encode the pq p1 p2 p3
q m1 m2 m3
the column and color channel.

3.2. Self-Attention Figure 1. A slice of one layer of the Image Transformer, recom-
puting the representation q 0 of a single channel of one pixel q by
For image-conditioned generation, as in our super-resolution attending to a memory of previously generated pixels m1 , m2 , . . ..
models, we use an encoder-decoder architecture. The en- After performing local self-attention we apply a two-layer position-
coder generates a contextualized, per-pixel-channel repre- wise feed-forward neural network with the same parameters for
sentation of the source image. The decoder autoregressively all positions in a given layer. Self-attention and the feed-forward
networks are followed by dropout and bypassed by a residual
generates an output image of pixel intensities, one channel
connection with subsequent layer normalization. The position
per pixel at each time step. While doing so, it consumes the encodings pq , p1 , . . . are added only in the first layer.
previously generated pixels and the input image represen-
 Image Transformer

tation generated by the encoder. For both the encoder and the transformed memory, weighted by the attention distribu-
decoder, the Image Transformer uses stacks of self-attention tion. In the decoders of our different models we mask the
and position-wise feed-forward layers, similar to (Vaswani outputs of the comparisons appropriately so that the model
et al., 2017). In addition, the decoder uses an attention cannot attend to positions in the memory that have not been
mechanism to consume the encoder representation. For un- generated, yet.
conditional and class-conditional generation, we employ the
To the resulting vector we then apply a single-layer fully-
Image Transformer in a decoder-only configuration.
connected feed-forward neural network with rectified linear
Before we describe how we scale self-attention to images activation followed by another linear transformation. The
comprised of many more positions than typically found in learned parameters of these are shared across all positions
sentences, we give a brief description of self-attention. but different from layer to layer.
Each self-attention layer computes a d-dimensional repre- As illustrated in Figure1, we perform dropout, merge in
sentation for each position, that is, each channel of each residual connections and perform layer normalization after
pixel. To recompute the representation for a given posi- each application of self-attention and the position-wise feed-
tion, it first compares the position’s current representation forward networks (Ba et al., 2016; Srivastava et al., 2014).
to other positions’ representations, obtaining an attention
The entire self-attention operation can be implemented using
distribution over the other positions. This distribution is
highly optimized matrix multiplication code and executed
then used to weight the contribution of the other positions’
in parallel for all pixels’ channels.
representations to the next representation for the position at
hand.
3.3. Local Self-Attention
Equations 1 and 2 outline the computation in our self-
attention and fully-connected feed-forward layers; Figure The number of positions included in the memory lm , or the
1 depicts it. W1 and W2 are the parameters of the feed- number of columns of M , has tremendous impact on the
forward layer, and are shared across all the positions in a scalability of the self-attention mechanism, which has a time
layer. These fully describe all operations performed in every complexity in O(h · w · lm · d).
layer, independently for each position, with the exception of The encoders of our super-resolution models operate on 8×8
multi-head attention. For details of multi-head self-attention, pixel images and it is computationally feasible to attend to
see (Vaswani et al., 2017). 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.
qa = layernorm(q + dropout( Inspired by convolutional neural networks we address this
Wq q(M Wk )T by adopting a notion of locality, restricting the positions
 
softmax √ M Wv )) (1) in the memory matrix M to a local neighborhood around
d
the query position. Changing this neighborhood per query
position, however, would prohibit packing most of the com-
putation necessary for self-attention into two matrix multi-
q 0 = layernorm(qa + dropout(W1 ReLu(W2 qa ))) (2)
plications - one for computing the pairwise comparisons and
another for generating the weighted averages. To avoid this,
In more detail, following previous work, we call the cur- we partition the image into query blocks and associate each
rent representation of the pixel’s channel, or position, to be of these with a larger memory block that also contains the
recomputed the query q. The other positions whose repre- query block. For all queries from a given query block, the
sentations will be used in computing a new representation model attends to the same memory matrix, comprised of all
for q are m1 , m2 , . . . which together comprise the columns positions from the memory block. The self-attention is then
of the memory matrix M . Note that M can also contain q. computed for all query blocks in parallel. The feed-forward
We first transform q and M linearly by learned matrices Wq networks and layer normalizations are computed in parallel
and Wk , respectively. for all positions.
The self-attention mechanism then compares q to each of In our experiments we use two different schemes for choos-
the pixel’s channel representations
√ in the memory with a dot- ing query blocks and their associated memory block neigh-
product, scaled by 1/ d. We apply the softmax function borhoods, resulting in two different factorizations of the
to the resulting compatibility scores, treating the obtained joint pixel distribution into conditional distributions. Both
vector as attention distribution over the pixel channels in are illustrated in Figure 2.
the memory. After applying another linear transformation
Wv to the memory M , we compute a weighted average of
 Image Transformer

Local 1D Attention Local 2D Attention

1D Local Attention For 1D local attention (Section 3.3)
we first flatten the input tensor with positional encodings in Memory Block
Memory Block
raster-scan order, similar to previous work (van den Oord
et al., 2016a). To compute self-attention on the resulting
linearized image, we then partition the length into non-
q q
overlapping query blocks Q of length lq , padding with ze- Query Block

roes if necessary. While contiguous in the linearized image, Query Block

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 lm positions cor-
responding to pixels that have been generated before, which Figure 2. The two different conditional factorizations used in our
can result in overlapping memory blocks. 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
2D Local Attention In 2D local attention models, we
superset of the query block pixels. In every self-attention layer,
partition the input tensor with positional encodings into
each position in a query block attends to all positions in the memory
rectangular query blocks contiguous in the original image block. The pixel marked as q is the last that was generated. All
space. We generate the image one query block after another, channels of pixels in the memory and query blocks shown in white
ordering the blocks in raster-scan order. Within each block, have masked attention weights and do not contribute to the next
we generate individual positions, or pixel channels, again in representations of positions in the query block. While the effective
raster-scan order. receptive field size in this figure is the same for both schemes, in
2D attention the memory block contains a more evenly balanced
As illustrated in the right half of Figure 2, we generate the number of pixels next to and above the query block, respectively.
blocks outlined in grey lines left-to-right and top-to-bottom.
We use 2-dimensional query blocks of a size lq specified by
height and width lq = wq ·hq , and memory blocks extending
the query block to the top, left and right by hm , wm and dence across channels (Salimans et al.). For each pixel, the
again wm pixels, respectively. number of parameters is 10 times the number of mixture
components: 10 for one unnormalized mixture probability,
In both 1D and 2D local attention, we mask attention three means, three standard deviations, and three coefficients
weights in the query and memory blocks such that posi- which capture the linear dependence. For 10 mixtures, this
tions that have not yet been generated are ignored. translates to 100 parameters for each pixel; for 32 × 32 im-
As can be seen in Figure 2, 2D local attention balances hori- ages, the network outputs 102, 400 dimensions, which is a
zontal and vertical conditioning context much more evenly. 7x reduction enabling denser gradients and lower memory.
We believe this might have an increasingly positive effect on
quality with growing image size as the conditioning informa- 4. Inference
tion in 1D local attention becomes increasingly dominated
by pixels next to a given position as opposed to above it. Across all of the presented experiments, we use categorical
sampling during decoding with a tempered softmax (Dahl
3.4. Loss Function et al., 2017). We adjust the concentration of the distribution
we sample from with a temperature τ > 0 by which we
We perform maximum likelihood, in which we maximize divide the logits for the channel intensities.
Ph·w·3
log p(x) = t=1 log p(xt | x<t ) with respect to network
We tuned τ between 0.8 and 1.0, observing the highest
parameters, and where the network outputs all parameters
perceptual quality in unconditioned and class-conditional
of the autoregressive distribution. We experiment with two
image generation with τ = 1.0. For super-resolution we
settings of the distribution: a categorical distribution across
present results for different temperatures in Table 5.
each channel (van den Oord et al., 2016a) and a mixture of
discretized logistics over three channels (Salimans et al.).
5. Experiments
The categorical distribution (cat) captures each intensity
value as a discrete outcome and factorizes across channels. All code we used to develop, train, and evaluate our models
In total, there are 256·3 = 768 parameters for each pixel; for is available in Tensor2Tensor (Vaswani et al., 2018).
32 × 32 images, the network outputs 786, 432 dimensions.
For all experiments we optimize with Adam (Kingma & Ba,
Unlike the categorical distribution, the discretized mixture 2015), and vary the learning rate as specified in (Vaswani
of logistics (DMOL) captures two important properties: et al., 2017). We train our models on both p100 and k40
the ordinal nature of pixel intensities and simpler depen- GPUs, with batch sizes ranging from 1 to 8 per GPU.
 Image Transformer

Table 4. Bits/dim on CIFAR-10 test and ImageNet validation sets.

The Image Transformer outperforms all models and matches Pixel-
CNN++, achieving a new state-of-the-art on ImageNet. Increasing
memory block size (bsize) significantly improves performance.
Model Type bsize NLL
CIFAR-10 ImageNet
(Test) (Validation)
Pixel CNN - 3.14 -
Row Pixel RNN - 3.00 3.86
Gated Pixel CNN - 3.03 3.83
Pixel CNN++ - 2.92 -
PixelSNAIL - 2.85 3.80
Ours 1D local (8l, cat) 8 4.06 -
16 3.47 -
64 3.13 -
256 2.99 -
Ours 1D local (cat) 256 2.90 3.77
Ours 1D local (dmol) 256 2.90 -

Table 3. Conditional image generations for all CIFAR-10 cate- ImageNet is a much larger dataset, with many more cate-
gories. Images on the left are from a model that achieves 3.03 gories than CIFAR-10, requiring more parameters in a gener-
bits/dim on the test set. Images on the right are from our best ative model. Our ImageNet unconditioned generation model
non-averaged model with 2.99 bits/dim. Both models are able
has 12 self-attention and feed-forward layers, d = 512, 8
to generate convincing cars, trucks, and ships. Generated horses,
planes, and birds also look reasonable.
attention heads, 2048 dimensions in the feed-forward lay-
ers, and dropout of 0.1. It significantly outperforms the
Gated PixelCNN and establishes a new state-of-the-art of
3.77 bits/dim with checkpoint averaging. We trained only
5.1. Generative Image Modeling unconditional generative models on ImageNet, since class
labels were not available in the dataset provided by (van den
Our unconditioned and class-conditioned image generation Oord et al., 2016a).
models both use 1D local attention, with lq = 256 and a total
memory size of 512. On CIFAR-10 our best unconditional Table 4 shows that growing the receptive field improves
models achieve a perplexity of 2.90 bits/dim on the test set perplexity significantly. We believe this to highlight a key
using either DMOL or categorical. For categorical, we use advantage of local self-attention over CNNs: namely that
12 layers with d = 512, heads=4, feed-forward dimension the number of parameters used by local self-attention is
2048 with a dropout of 0.3. In DMOL, our best config uses independent of the size of the receptive field. Furthermore,
14 layers, d = 256, heads=8, feed-forward dimension 512 while d > receptivefield, self-attention still requires fewer
and a dropout of 0.2. This is a considerable improvement floating-point operations.
over two baselines: the PixelRNN (van den Oord et al., For experiments with the categorical distribution we evalu-
2016a) and PixelCNN++ (Salimans et al.). Introduced after ated both coordinate encoding schemes described in Section
the Image Transformer, the also self-attention based Pixel- 3.3 and found no difference in quality. For DMOL we only
SNAIL model reaches a significantly lower perplexity of evaluated learned coordinate embeddings.
2.85 bits/dim on CIFAR-10 (Chen et al., 2017). On the
more challenging ImageNet data set, however, the Image
5.2. Conditioning on Image Class
Transformer performs significantly better than PixelSNAIL.
We represent the image classes as learned d-dimensional
We also train smaller 8 layer CIFAR-10 models which have
embeddings per class and simply add the respective em-
d = 512, 1024 dimensions in the feed-forward layers, 8
bedding to the input representation of every input position
attention heads and use dropout of 0.1, and achieve 3.03
together with the positional encodings.
bits/dim, matching the PixelCNN model (van den Oord
et al., 2016a). Our best CIFAR-10 model with DMOL has d We trained the class-conditioned Image Transformer on
and feed-forward layer layer dimension of 256 and perform CIFAR-10, achieving very similar log-likelihoods as in un-
attention in 512 dimensions. conditioned generation. The perceptual quality of generated
 Image Transformer

images, however, is significantly higher than that of our

Table 5. Negative log-likelihood and human eval performance for
unconditioned models. The samples from our 8-layer class-
the Image Transformer on CelebA. The fraction of humans fooled
conditioned models in Table 3, show that we can generate is significantly better than the previous state of the art.
realistic looking images for some categories, such as cars
and trucks. Model Type τ %Fooled
ResNet n/a 4.0
5.3. Image Super-Resolution srez GAN n/a 8.5
Super-resolution is the process of recovering a high resolu- PixelRecursive 1.0 11.0
tion image from a low resolution image while generating (Dahl et al., 2017) 0.9 10.4
realistic and plausible details. Following (Dahl et al., 2017), 0.8 10.2
in our experimental setup we enlarge an 8 × 8 pixel image 1D local 1.0 29.6 ± 4.0
four-fold to 32 × 32, a process that is massively underspec- Image Transformer 0.9 33.5 ± 3.5
ified: the model has to generate aspects such as texture of 0.8 35.94 ± 3.0
hair, makeup, skin and sometimes even gender that cannot 2D local 1.0 30.64 ± 4
possibly be recovered from the source image. Image Transformer 0.9 34 ± 3.5
0.8 36.11 ± 2.5
Here, we use the Image Transformer in an encoder-decoder
configuration, connecting the encoder and decoder through
an attention mechanism (Vaswani et al., 2017). For the
encoder, we use embeddings for RGB intensities for each
pixel in the 8×8 image and add 2 dimensional positional
encodings for each row and width position. Since the input Turk where each worker is required to make a binary choice
is small, we flatten the whole image as a [h × w × 3, d] when shown one generated and one real image. Following
tensor, where d is typically 512. We then feed this sequence the same procedure for the evaluation study as (Dahl et al.,
to our stack of transformer encoder layers that uses repeated 2017), we show 50 pairs of images, selected randomly from
self-attention and feed forward layers. In the encoder we the validation set, to 50 workers each. Each generated and
don’t require masking, but allow any input pixel to attend original image is upscaled to 128 × 128 pixels using the
to any other pixel. In the decoder, we use a stack of local Bilinear interpolation method. Each worker then has 1-2
self-attention, encoder-decoder-attention and feed-forward seconds to make a choice between these two images. In
layers. We found using two to three times fewer encoder our method, workers choose images from our model up
than decoder layers to be ideal for this task. to 36.1% of the time, a significant improvement over pre-
We perform end-to-end training of the encoder-decoder vious models. Sampling temperature of 0.8 and 2D local
model for Super resolution using the log-likelihood ob- attention maximized perceptual quality as measured by this
jective function. Our method generates higher resolution evaluation.
images that look plausible and realistic across two datasets. To measure how well the high resolution samples correspond
For both of the following data sets, we resized the image to to the low resolution input, we calculate Consistency, the
8 × 8 pixels for the input and 32 × 32 pixels for the label L2 distance between the low resolution input and a bicubic
using TensorFlow’s area interpolation method. downsampled version of the high resolution sample. We
observe a Consistency score of 0.01 which is on par with
the models in (Dahl et al., 2017).
CelebA We trained both our 1D Local and 2D Local mod-
els on the standard CelebA data set of celebrity faces with We quantify that our models are more effective than exem-
cropped boundaries. With the 1D Local, we achieve a nega- plar based Super Resolution techniques like Nearest Neigh-
tive log likelihood (NLL) of 2.68 bits/dim on the dev set, bors, which perform a naive look-up of the training data to
using lq = 128, memory size of 256, 12 self-attention and find the high resolution output. We take a bicubic down-
feed-forward layers, d = 512, 8 attention heads, 2048 di- sampled version of our high resolution sample, find the
mensions in the feed-forward layers, and a dropout of 0.1. nearest low resolution input image in the training data for
With the 2D Local model, we only change the query and that sample, and calculate the MS-SSIM score between the
memory to now represent a block of size 8 × 32 pixels and high resolution sample and the corresponding high reso-
16 × 64 pixels respectively. This model achieves a NLL lution image in the training data. On average, we get a
of 2.61 bits/dim. Existing automated metrics like pSNR, MS-SSIM score of 44.3, on 50 samples from the validation
SSIM and MS-SSIM have been shown to not correlate with set, which shows that our models don’t merely learn to copy
perceptual image quality (Dahl et al., 2017). Hence, we con- training images but generate high-quality images by adding
ducted a human evaluation study on Amazon Mechanical synthesized details on the low resolution input image.
 Image Transformer

Input 1D Local Attention 2D Local Attention Original

τ = 0.8 τ = 0.9 τ = 1.0 τ = 0.8 τ = 0.9 τ = 1.0

Table 6. 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.

CIFAR-10 We also trained a super-resolution model on previously proposed (Mansimov et al., 2015), and tasks
the CIFAR-10 data set. Our model reached a negative log- combining modalities such as language-driven editing of
likelihood of 2.76 using 1D local attention and 2.78 using images.
2D local attention on the test set. As seen in Figure 2, our
Fundamentally, we aim to move beyond still images to
model commonly generates plausible looking objects even
video (Kalchbrenner et al., 2016) and towards applications
though the input images seem to barely show any discernible
in model-based reinforcement learning.
structure beyond coarse shapes.

References
6. Conclusion
Ba, Jimmy Lei, Kiros, Jamie Ryan, and Hinton, Geoffrey E.
In this work we demonstrate that models based on self- Layer normalization. arXiv preprint arXiv:1607.06450,
attention can operate effectively on modalities other than 2016.
text, and through local self-attention scale to significantly
larger structures than sentences. With fewer layers, its larger Bellemare, Marc G., Srinivasan, Sriram, Ostrovski, Georg,
receptive fields allow the Image Transformer to significantly Schaul, Tom, Saxton, David, and Munos, Rémi. Unifying
improve over the state of the art in unconditional, probabilis- count-based exploration and intrinsic motivation. CoRR,
tic image modeling of comparatively complex images from abs/1606.01868, 2016. URL http://arxiv.org/
ImageNet as well as super-resolution. abs/1606.01868.
We further hope to have provided additional evidence
that even in the light of generative adversarial networks, Bengio, Yoshua and Bengio, Samy. Modeling high-
likelihood-based models of images is very much a promising dimensional discrete data with multi-layer neural net-
area for further research - as is using network architectures works. In Neural Information Processing Systems, pp.
such as the Image Transformer in GANs. 400–406. MIT Press, 2000.

In future work we would like to explore a broader variety Berthelot, David, Schumm, Tom, and Metz, Luke. BEGAN:
of conditioning information including free-form text, as boundary equilibrium generative adversarial networks.
 Image Transformer

CoRR, abs/1703.10717, 2017. URL http://arxiv. In Empirical Methods in Natural Language Process-
org/abs/1703.10717. ing, 2016. URL https://arxiv.org/pdf/1606.
01933.pdf.
Chen, Xi, Mishra, Nikhil, Rohaninejad, Mostafa, and
Abbeel, Pieter. Pixelsnail: An improved autoregres- Radford, Alec, Metz, Luke, and Chintala, Soumith.
sive generative model. arXiv preprint arXiv:1712.09763, Unsupervised representation learning with deep con-
2017. volutional generative adversarial networks. CoRR,
abs/1511.06434, 2015. URL http://arxiv.org/
Cheng, Jianpeng, Dong, Li, and Lapata, Mirella. Long
abs/1511.06434.
short-term memory-networks for machine reading. arXiv
preprint arXiv:1601.06733, 2016. Salimans, Tim, Karpathy, Andrej, Chen, Xi, Kingma,
Diederik P., and Bulatov, Yaroslav. Pixelcnn++: A pix-
Dahl, Ryan, Norouzi, Mohammad, and Shlens, Jonathan.
elcnn implementation with discretized logistic mixture
Pixel recursive super resolution. 2017. URL https:
likelihood and other modifications. In International Con-
//arxiv.org/abs/1702.00783.
ference on Learning Representations.
Goodfellow, Ian, Pouget-Abadie, Jean, Mirza, Mehdi, Xu,
Bing, Warde-Farley, David, Ozair, Sherjil, Courville, Srivastava, Nitish, Hinton, Geoffrey E, Krizhevsky, Alex,
Aaron, and Bengio, Yoshua. Generative adversarial nets, Sutskever, Ilya, and Salakhutdinov, Ruslan. Dropout: a
2014. simple way to prevent neural networks from overfitting.
Journal of Machine Learning Research, 15(1):1929–1958,
Kalchbrenner, Nal and Blunsom, Phil. Recurrent continuous 2014.
translation models. In Proceedings EMNLP 2013, pp.
1700–1709, 2013. URL http://nal.co/papers/ Theis, Lucas and Bethge, Matthias. Generative image mod-
KalchbrennerBlunsom_EMNLP13. eling using spatial lstms. In Proceedings of the 28th Inter-
national Conference on Neural Information Processing
Kalchbrenner, Nal, van den Oord, Aäron, Simonyan, Systems - Volume 2, NIPS’15, pp. 1927–1935, Cambridge,
Karen, Danihelka, Ivo, Vinyals, Oriol, Graves, Alex, MA, USA, 2015. MIT Press. URL http://dl.acm.
and Kavukcuoglu, Koray. Video pixel networks. CoRR, org/citation.cfm?id=2969442.2969455.
abs/1610.00527, 2016. URL http://arxiv.org/
abs/1610.00527. van den Oord, Aäron and Schrauwen, Benjamin. The
student-t mixture as a natural image patch prior
Kingma, Diederik and Ba, Jimmy. Adam: A method for with application to image compression. Jour-
stochastic optimization. In ICLR, 2015. nal of Machine Learning Research, 15:2061–2086,
2014. URL http://jmlr.org/papers/v15/
Larochelle, Hugo and Murray, Iain. The neural autoregres-
vandenoord14a.html.
sive distribution estimator. In The Proceedings of the 14th
International Conference on Artificial Intelligence and van den Oord, Aäron, Kalchbrenner, Nal, and Kavukcuoglu,
Statistics, volume 15 of JMLR: W&CP, pp. 29–37, 2011. Koray. Pixel recurrent neural networks. ICML, 2016a.
Ledig, Christian, Theis, Lucas, Huszar, Ferenc, Caballero,
van den Oord, Aäron, Kalchbrenner, Nal, Vinyals, Oriol,
Jose, Aitken, Andrew, Tejani, Alykhan, Totz, Johannes,
Espeholt, Lasse, Graves, Alex, and Kavukcuoglu, Koray.
Wang, Zehan, and Shi, Wenzhe. Photo-realistic single
Conditional image generation with pixelcnn decoders.
image super-resolution using a generative adversarial net-
NIPS, 2016b.
work. arXiv:1609.04802, 2016.
Vaswani, Ashish, Shazeer, Noam, Parmar, Niki, Uszkoreit,
Mansimov, Elman, Parisotto, Emilio, Ba, Lei Jimmy, and
Jakob, Jones, Llion, Gomez, Aidan N., Kaiser, Lukasz,
Salakhutdinov, Ruslan. Generating images from cap-
and Polosukhin, Illia. Attention is all you need. 2017.
tions with attention. CoRR, abs/1511.02793, 2015. URL
URL http://arxiv.org/abs/1706.03762.
http://arxiv.org/abs/1511.02793.

Metz, Luke, Poole, Ben, Pfau, David, and Sohl-Dickstein, Vaswani, Ashish, Bengio, Samy, Brevdo, Eugene, Chol-
Jascha. Unrolled generative adversarial networks. CoRR, let, Francois, Gomez, Aidan N., Gouws, Stephan, Jones,
abs/1611.02163, 2016. URL http://arxiv.org/ Llion, Kaiser, Łukasz, Kalchbrenner, Nal, Parmar, Niki,
abs/1611.02163. Sepassi, Ryan, Shazeer, Noam, and Uszkoreit, Jakob.
Tensor2tensor for neural machine translation. CoRR,
Parikh, Ankur, Täckström, Oscar, Das, Dipanjan, and abs/1803.07416, 2018. URL http://arxiv.org/
Uszkoreit, Jakob. A decomposable attention model. abs/1803.07416.
 Image Transformer

Zhang, Han, Xu, Tao, Li, Hongsheng, Zhang, Shaoting,

Huang, Xiaolei, Wang, Xiaogang, and Metaxas, Dim-
itris N. Stackgan: Text to photo-realistic image synthe-
sis with stacked generative adversarial networks. CoRR,
abs/1612.03242, 2016. URL http://arxiv.org/
abs/1612.03242.