Axial Attention in Multidimensional Transformers
Authors: Jonathan Ho
Nal Kalchbrenner
Dirk Weissenborn
Tim Salimans
Venue:  (2020)
Keywords: self-attention
transformer
images
videos

Abstract: Self-attention effectively captures large receptive fields with high information bandwidth, but its computational resource requirements grow quadratically with the number of points over which attention is performed. For data arranged as large multidimensional tensors, such as images and videos, the quadratic growth makes self-attention prohibitively expensive. These tensors often have thousands of positions that one wishes to capture and proposed attentional alternatives either limit the resulting receptive field or require custom subroutines. We propose Axial Attention, a simple generalization of self-attention that naturally aligns with the multiple dimensions of the tensors in both the encoding and the decoding settings. The Axial Transformer uses axial self-attention layers and a shift operation to efficiently build large and full receptive fields.  Notably the proposed structure of the layers allows for the vast majority of the context to be computed in parallel during decoding without introducing any independence assumptions. This semi-parallel structure goes a long way to making decoding from even a very large Axial Transformer broadly applicable. We demonstrate state-of-the-art results for the Axial Transformer on the ImageNet-32 and ImageNet-64 image benchmarks as well as on the BAIR Robotic Pushing video benchmark. We open source the implementation of Axial Transformers.

Under review as a conference paper at ICLR 2020
AXIAL ATTENTION IN MULTIDIMENSIONAL TRANSFORMERS
Anonymous authors Paper under double-blind review
ABSTRACT
We propose Axial Transformers, self-attention-basee autoregressive model for images and other data organized as high dimensional tensors. Existing autoregres- sive models either suffer from excessively large computational resource require- ments for high dimensional data, or make compromises in terms of distribution expressiveness case of implementation in order to decrease resource require- ments. Our architecture, by contrast, maintains both full expressiveness over joint distributions over data and ease of implementation with standard deep learning frameworks, while requiring reasonable memory and computation and achieving state-of-the-art results on standard generative modeling benchmarks. Our models are based on axial attention, a simple generalization of self-attention that natu- rally aligns with the multiple dimensions of the tensors in both the encoding and the decoding settings. Notably the proposed structure of the layers allows for the vast majority of the context be computed in parallel during decoding without introducing any independence assumptions. This semi-parallel structure goes a long way to making decoding from even very large Axial Transformer broadly applicable. We demonstrate state-of-the-art results for the Axial Transformer on the ImageNet-32 and ImageNet-64 image benchmarks as well as on the BAIR Robotic Pushing video benchmark. We open source the implementation of Axial Transformers.
INTRODUCTION
Autoregressive models are a family of exact likelihood-based generative models that represent the joint distribution of data r = (Z1.... ZNN) as product of conditionals (z) = -I, (0 x<)))
Neural network models in this family have achieved Ktte-offthheaat log likelihoods on high-
dimensional image and video datasets (van den Oord et al.] 201 CCen et al. 2018 Menick &
Kalchbrenner 2018: Parmar et al.) [0018] Child et al. 2019; Weissenborn ct al.) 2019 Salimans
ct al. 2017 Kalchbrenner et al., 2017 Uria et al. 2016 Parikh et al. 2016 Theis & Bethge 2015
van den Oord et al. 2016b) due to architectural innovations that enable the following capabilities:
Large, high information bandwidth receptive fields for each pixel x;, capable of expressing long-range dependencies over previous pixels 2<c, and
2. Computationally efficient, vectorizable computation of the log likelihood and its gradient.
Autoregressive model architectures that can read long-range dependencies over large receptive fields are able express all joint distributions over the data. Meanwhile, architectures that admit fast log likelihood gradient computation are suitable for training using stochastic gradient method on maximum likelihood objective-a straightforward, stable training procedure for generative models.
These desiderata make self-attention compelling building block for autoregressive model architec- tures. Self-attention a neural network operation that is able to transform a sequence y1...., Y/ into a sequence v/..... ... V'x. where each 1/ depends on all yı by way of single vectorizable com- putation (Vaswani et al. (017). Self-attention is remarkably effective at learning long-rang depen- dencies between data dimensions and neural networks that incorporate self-attention in their designs are state-of-the-arr on many tasks from language modelling and machine translation to image and video modelling (Parmar al. 2018 Child al. 2019).

---

Under review as a conference paper at ICLR 2020
DURING SAMPLING
per
pet
down
The Axial Transformer model for 2-dimensional tensors. Before sampling a channel we encode all previous channels and frames with 8 blocks of unmasked and unmasked column at- tention (left). Then, for each row, we apply 4 blocks of unmasked row and masked column attention integrate the previously sampled rows for the active channels into our encoded representation (middle). Finally, we shift the encoded representation up to make sure the conditioning information satisfies causality, and we run the inner decoder consisting of 4 blocks of masked row attention to sample a new row in the image (right).
Figure I:
But the power of self-attention comes at the price of computational complexity. The memory and computation it consumes grow quadratically with the sequence length N making it prohibitively expensive to directly apply self-attention to long sequences. In case of autoregressive models of multidimennional tensors such as images or videos, the aim to capture large receptive fields in multiple dimensions further exacerbates the problem as even a modest number of receptive field steps in each dimension encompass large total number of locations. Various approaches have been proposed to alleviate this difficulty at the cost of either limiting the receptive field or requiring operations that may not be broadly available on GPUs or TPUs.
We propose the Axial Transformer, a simple yet effective seff-atennionn based autoregressive model for data organized multidimensional tensors. Rather than applying attention flattened string of tensor elements, our model instead applies attention along single axis of the tensor without flattening- refer to this "axial attention." Since the length of any single axis (that the height or width of an image) is typically much smaller than the total number of elements, an axial attention operation enjoys sigmificant saving in computation and memory over standard self-attention: for d-dimensional tensor with shape N 1//d factor of resources over standard self-attention.
N1/d axial attention saves (dN(d-1)/d)
Our Axial Transformer architecture allows for the majority of the context oi to be embedded with high degree of parallelism without introducing conditional independence assumptions among any of the locations, but has an interesting property that is amenable pling procedure. To sample one row of an image, the Axial Transformer only runs an autoregressive Transformer over that one row only, without re-embedding pixels from previous rows. We structure the Axial Transformer, however, so that it always defines a fully expressive joint distribution. No dependencies on previous pixels are ever lost.
comple-to-implemen
fast sam-
We evaluate Axial Transformers on image and video modelling benchmarks. We show that Axial Transformer achieves state-of-the-ar results on ImageNet-32 and on ImageNet-64. We also show that, simply by stacking video along the channel dimension, the Axial Transformer video without nearly any modification. On the BAIR Robot Pushing
be directly
applicd to the channel-stacke benchmark, the Axial Transformer significantly outperforms previous results without using archi- tecture specially designed for videos. The generated samples on these datasets are of the expected high quality.
Axial Transformers do not require subroutines for GPUs or TPUs that may exhibit unfavorable memory bandwidth and computation trade-offs. Axial Transformers simple implement using
2

---

Under review as a conference paper at ICLR 2020
Pall
Columan
Figure 2: Types of axial attention layers that are the building blocks of the Axial Transformer. The blue locations correspond to the receptive field of the output red location.
efficient operations that are widely available in deep learning frameworks (primarily dense-dense MatMuls). An open source implementation of our models is available at anonymi zed URL.
BACKGROUND
To set the stage for our discussion, we first review self-attention and its computational resource requirements in the context of autoregressive modeling. A self-attention layer takes as input length N sequence of D-dimensional embeddings X (a N x D matrix) and produces an output sequence Y (also N x D matrix) via:
Q XWQ: K= XWR, == XW,
softmax (QKT/VD): =
Wo. WK. and Wv are x D parameter matrices responsible for projecting the entries of the sequence X into keys, queries, and values, respectively. Each entry of the output sequence is linear combination of values in V weighted by the attention matrix A, which itself is computed from similarities between all pairs of query and key vectors. Both the expressive power and the resource cost of self-attention come from computing A and Y: it takes O(N²) time and space to compute the pairwise similarities between Q and K and to compute the linear combination of V vectors.
This quadratic complexity makes it impractical to apply self-attention to images and videos directly as flattened vectors: small 32 32 image has 3072 dimensions. Sequences such as these are too long for self-attention, attempts to scale self-attention to these modalities generally involve re- stricting these sequence lengths in a modality-aae performance.
manner while attempting to preserve modeling
One strategy is to restrict the conditioning context Zzi to a carefully designed small subset of the data dimensions. While this reduces the cost of attention, which is only performed over these small subsets instead of the full data, the model can no longer express all joint distributions over the data. Parmar al. (2018) propose image models with conditioning context Ies restricted to small window of the full image, but the implementation requires redundant data copies to extract and process these windows. Weissenborn al. (2019) similarly scale video autoregressive models by restricting the context, again preventing their model from expressing all joint distributions over pixels. Our models do not restrict context and hence we obtain better log likelihoods, as we will see in section[]
A different strategy is stack multiple sparse attention layers, each with restricted context for computational efficiency, but in manner that overlapping these layers yields a full-context model. Child al. (2019) propose two sparse attention patterns with this property. However, the architecture they propose that works best for images (the Strided Sparse Transformer) requires custom sparse attention GPU kernels to implement a specific block-sparse variant of patrix-matrix-mulltii The model cannot be easily implemented on other hardware such TPUs.
See table for a summary of these architecture design tradeoffs. Our goal in this paper to de- sign attention-baas -base autoregressive models that attain the best of all worlds. Our Axial Transformer, described in subsequent sections, has full conditioning context, so its ability to express joint distri- butions is never limited. The Axial Transformer also does not require any redundant data copies or custom kernels to implement in an efficient way. Indeed, we designed, and will make open source, an efficient implementation that uses only standard operations in deep learning libraries.
3

---

Under review as a conference paper at ICLR 2020
Table 1: Trade-off's of recently proposed multicimensional Transformer architectures.
Full receprive
Attention faster than ((N2)
Needs no custom kernels
Semi-porallel context aggregation
Model
field
Transformer Vaswani et al
2017
yes
n0
yos yes
Image Transformer Parmar e
2018
no
yes
Block Translomner Wersenborr
2019
yes
yes
Strided Sparse Transtamer Chill el Axial Transformer (ours)
2019
yes
yes
no
no
yes
y35
yes
yes
AXIAL TRANSFORMERS
We now describe Axial Transformers, self-attention-basel autoregressive models for high- dimensional data tensors. We describe its basic building block in section 3.1 and then we complete the description into a full autoregressive model in section 3.2
3.1 AXIAL ATTENTION
We first introduce our basic building block for developing -e--attentionnbseeee autoregressive models for high-dimensional data tensors. The proposed approach does not change the original shape of the multidimensional data tensor and performs masked or unmasked attention over a single axis of the tensor at a time. We call this operation axial attention, denoted by Attention, (x). It performs attention over axis k of the tensor mixing information along axis & while keeping information along other axes independent. be implemented by transposing all axes except lc to the batch axis, calling standard attention subroutine, then undoing the transpose (an alternative is to use the einsum operation available in most deep learning libraries).
is straightforward to implement: axial attention over axis & can
When the data is an image, we call Attention, column attention, as it mixes information within columns while keeping separate columns independent. We call Attentiong row attention for anal- ogous reasons. Axial attention on a square image of size N = S x S performs attention on S sequences of length S-this is a total of O(S S²) O /N) computation-an O(vN) savings in computation over standard self-attention. In general. for d-dimensional tensor with N Sa axial attention saves O(N(d-1)/d) computation over standard attention. Of course, single layer of axial attention along some axis k does not have the full receptive field since it covers a single axis, but will see in section 3.2 that stacking two axial attention layers allows the model to obtain global receptive field.
will be important for us also define MaskodAttenttion, be the causally masked variant of Attention,: component i of the result of MaskedAAttentionn (r) along axis & depends on only com- ponents ,,..., of æ along axis k. The receptive fields of these attention patterns, both unmasked and masked, are illustrated in fig. 2 We will use these masked blocks to build our autoregressive model section 3.2
Axial attention can be used within standard Transformer layers in a straightforward manner to pro- duce Axial Transformer layers. The basic building blocks are the same as those found in the standard Transformer architecture:
Lay)rNorm(z): layer normalization (Ba et al. 2016). and
Densen (a)): dense layer operating over the last axis of the input æ. The letter denotes the dimension of the output activations. If the input has shape H x W x C, then this operation is identical to a 1 x convolution, and the output has shape H x W x D.
We use these to define ResNet axial attention blocks operating on tensors of D-dimensional embed- dings (Vaswani et al. 2017 Child et al. 2019):
FeedforwardBlock(()) = r Densep (owneeooononnnnnnooiiieo
(aanliaartttDDDneeee_aaaaveemm)) (LayerNorm(z)))
AttentionBlock,(()) = -x + (AteeeeAAttetionnaaayya
TransformerBlock; =
4

---

Under review as a conference paper at ICLR 2020
D' is chosen to be some constant factor larger than D, from 1 to 4 Vaswani et al. 2017). We also define a MaskedTrranffommerBocck using MaskedAttention, in place of Attentiong.
Operations similar to unmasked axial attention have been proposed in other contexts in computer vision (Huang et al. 2019). Our focus in forthcoming sections is the use of masked axial attention and its utility in autoregressive image modeling, which is not explored in these works.
3.2 AXTAL TRANSFORMERS
We now describe Axial Transformers, our axial attention autoregressive models for images and videos. We will use the axial attention operations described in section [37 building blocks in multi-layer autoregressive model of the form po(x) (a:) = I po( xe.s) following the raster sean ordering of pixels. We will accomplish this by building an autoregressive model over rows (sec- tion B.2.I). then conditioning each row on previous rows (section [.2.I), then further conditioning on previous channels and frames (section 3.2.2). Decomposing the model in this manner also leads to simple fast and partly parallel sampling procedure (section 3.2.1.
3.2.1 A MODEL FOR CCLLNEL IMAGES
We begin with autoregressive model for single-channel image with shape H x W, with each pixel taking an integer value in [0, 255] representing its intensity. As is standard practice with Transformers, pixel intensities embeddings, which we call h. The architecture's 's responsibility is to transform into a x W x 256 tensor of logits suitable for classification or sampling. These logits must depend only on previous pixels in the input x along the raster scan ordering to ensure that the architecture defines a valid autoregressive model.
first embedded into H x W x D tensor of D-dimensional
Inner Decoder: a row-wise model a "row-wise" model:
Our idea is to begin with masked row attention layers to create
h + Embed
h < ShiftRight(h) PositionEmmbeddings
h <- Maskee/TranssormmeBBooo////
X Lrow
Here, Trow is the number of masked row attention blocks applied to h. PositionEmbeddings is a H x W x D tensor of position embeddings that inform the attention layers of the position. For parameter efficiency we use "additively factorized" position embeddings, meaning that we parameterize them as broadcasted sum of I 1 x D embeddings for rows and 1 W D embeddings for columns.
The operation ShiftRight shifts the input right by one pixel, which has the effect of shifting the receptive field left by one pixel. This ensures that the masked row attention layers exclude the current pixel from their receptive field, which is crucial for architecture to define a correct autoregressive model.
As this model employs row attention only, it enjoys the computational efficiency benefits described in section 3.1 However, it clearly does not define full-context model because each location in the output does not depend on input pixels in previous rows. If we were to use the resulting h as logits for pixel intensity prediction, we would obtain a set of H independent autoregressive models p(rii ( ,...,,,... 11) for each row i E [1, H), not single autoregressive model with full context. We address this issue next.
capturing the rows above Each pixel Zy in the aforementioned model already depends on previous pixels in its own row rıKjj- We just need to make it depend on all previous rows i<i. too. So, we insert unmasked row and masked column layers in the beginning of the model as follows (newly inserted operations are underlined):
Outer Decoder:
h +- Embed(z)
u < PosiiiinEEmbeddings
u <- NakkeffTaasoomeeBBBccttaaaaaaaa (Transforner/llckk((()
x Trppger/ /2
h < Shift(Don(()) + ShiftRight(h) + PositionEmbeddings
h MaskedTraassommeBBlcc..
X Lrow
5

---

Under review as a conference paper at ICLR 2020
Channel iidex plara
RG3 chamata for hrmne T-3
RG3 chamats frane T-
chanele hrme
Provocus Hel sheroos
Figure 3: Arrangement of inputs to the encoding network of the Axial Transformer. Previously available or generated channels an image or video are sequentially stacked in the input. A variable number of padding planes are used as placeholders for future generated channels. A final integer plane signals the Axial Transformer the channel that is being generated that step.
The tensor u represents context captured above the current pixel. It is computed by unmasked row and masked column attention layers, repeated to total of Tupper layers increase model capacity, which make cover the receptive field at all rows above and including the current pixel. The ShiftDown operation shifts u down one pixel, which shifts its receptive field up one pixel. Thus we have context which captures all pixels above while excluding the current row, which we add to h as input to the masked row layers. We have thus converted the row-wise model into fully expressive autoregressive model that captures not only pixels in the current row but also those above.
Following standard practice, we pass the final h through layer normalization and a final dense layer to produce logits with shape H x W x 256. The logits each location depend on all previous pixel locations in the raster scan ordering.
Semi-Paralle Sampling Naive implementations of sampling from sequential models are notori- ously slow because they require re-evaluating the entire network to sample each location. In the case of our model for VN VN square image, each network evaluation takes O(N (Iupper + Lyww) time, so sampling the whole image would take WVLVVN((LLL + Lrow)), which is far too large.
Fortunately, our architecture is amenable to a particularly simple implementation of a faster sampling that is able compute large sections of the model in parallel (see Figure I Pseudocode is as follows:
For each row E H]:
(a) Compute the upper context u including information about all Iei, using the upper layers
(b) For each column j W]]:
Sample v5 conditioned on u and prior elements of row (xi,<)).
Because the Loow row-wise layers are independent over rows (they depend on other rows only through the upper context, explained in section B.2.1p, sampling one row can be accomplished by evaluating the row-wise layers for that one row only, completely ignoring other rows. Thus, in one row of VN pixels, each pixel can be sampled in O(NL) so pixels can be sampled in L. Before each of the VN rows can be sampled, the upper context must be com- puted DN for a total O 2 over the course all rows. Thus we arrive at (T Tupper [row)) )) in total, which is vN faster than the naive implementation. To our knowledge, sampling speedups of this type are not possible with contemporary work on scaling Transformers to
images and videos (Child et al. 2019 Weissenborn al. 2019
3.2.2 CHANNEL ENCODER FOR MULTI-CHANNER IMAGES AND VIDEOS
We have just described architecture for single-channel image of shape H x W. Here, we show how to extend the architecture to multi-channel images or videos of shape H x W x C (here C is either the number of channels in multi-channel image, or the product of the number of channels and timesteps in a video). One way to model such data of shape H W x C is to simply stack the channels on top of each other into single-channel image of shape (H (H-C) x W or H (W.C). This

---

Under review as a conference paper at ICLR 2020
Table 2: Unconditional and class-conditional image modeling results (bits/dim)
Model
ImageNet 32x32
ImageNet 64x64
Multiscale PixelCNN (Reed et al.
2017)
3.95
3.70
PixelCNN/RNN (van den Cord cf
2016a
3.86
3.63
Gated PixcICNN (van Cord ct al.
2016b
3.83
3,57
PixeISNAIL (Chen et al. (2018)
3.80
3,52
SPN (Menick & Katchbrenner 2018)
3.79
3.52
Image Transformer (Parmar et al. 2018
3.77
Strided Sparse Translormer (Child al.
2019
3.44
Axial Translormer LSTM inner decoder
3.77
3,46
Axial Transformer
3.76 (3.758)
3.44 (3.439)
Table 3:
Video modeling results (bits/dim) on the BAIR Robotic Pushing dataset (Ebert et al.
2017
We condition on single video frame and model the next 15 frames, similar to Weissenborn Kumar et al. (2019) instead condition the prior frames of the video.
al. 2019).
Model
bits/dim next 15 frames
VideoFlow Kumar et al. 2019)
1.87
Video Transtormer (Weissenborn et al. 2019) Axial Transformer (ours)
1.35
1.29
simple to implement, but does increase the sequence length for column attention or row attention, which can be undesirable for large C. We instead opt model one channel a time as a single- channel image, but now conditioned on previous channels using an extra set of unmasked row and unmasked column attention layers. This means that have a model of the form p( [ [.,,,c)), where previous channels <e are processed into H x W x D tensor of context information, which is then added into the first encoding blocks of the model in section 3.2. (Figure 3).
We do not share any parameters among any of these layers. At training time, we train on a random channel slice of each image: we process the previous slices using these unmasked attention layers to produce context tensor, and maximize the likelihood of the randomly chosen slice conditioned on this context. This amounts to training on an unbiased estimate of log likelihood for the whole data tensor. See fig. for an illustration of this complete model.
EXPERIMENTS
We benchmarked our models on standard datasets for generative image and video models: down- sampled ImageNet (van den Oord al., 2016a) and BAIR Robot Pushing (Ebert al.] 2017). All Axial Transformers have 8 total layers in the encoder, 8 layers in the outer decoder and 4 layers in the inner decoder. We use hidden size of 2048 neurons throughout and for all setups and 16 heads with 128 neurons each for the attention component. We train for approximately 200k steps on ImageNet32 and ImageNet64 and for 200k steps on BAIR Robot Pushing. Our models can overfit on ImageNet32, but on the other datasets the models keep gradually improving with more steps. See table 2 and table for our results.
4.1 ABLATION STUDY
To push the limits of the semi-parae marilelrllll sampling by making the inner decoder as small as possible, we train an Axial Transformer with the inner decoder replaced by a single LSTM layer of 2048 units. This slows down training time by about 20% on ImageNet32 and about 80% on ImageNet64 when maintaining the number of steps and all else fixed. We find that the Axial Transformer + LSTM inner decoder performs rather well on the ImageNet32 and ImageNet64 benchmarks (table 2}. thereby also showing the effectiveness of the remaining parts of the Axial Transformer that capture the context of the rows above. We also find however that the full four layers of the inner decoder of the Axial
7

---

Under review as a conference paper at ICLR 2020
Transformer provide additional boost in performance well significantly faster training. The Axial Transformer + LSTM inner decoder has the advantage of requiring only a couple of matrix- vector products to compute the layers at each autogressive step, comparing favourably with about the 12 matrix-vector products required by the Axial Transformer, but the slower training time would make the LSTT inner decoder quickly impractical for larger tensors.
4.2 SAMPLES
In fig. 4 and fig- 5 we show samples from our 64 x 64 and 32 x 32 ImageNet models. The samples are globally coherent and show visibly recognizable scenes, meaning that our Axial Transformer architecture successfully captures long-range dependencies across thousands of data dimensions in these image datasets. The samples also don't show any architecture-correlaeed artefacts. In addition, in fig. [ we show samples from the BAIR Robotic Pushing dataset. The first frame is each row is given by the dataset and the rest are continuation. We note the high quality exactness of details and the very large diversity (at temperature 1.0).
CONCLUSION
We proposed the Axial Transformer, an self-attention based autoregressive model for data organized as high dimensional tensors. is based on axial attention, a simple generalization of self-attention that scales better with the dimension of input data, achieving a O(N(d-1)/4) savings in compu- tation and memory for d-dimensional input tensor with N elements. Axial attention is easy to implement and does not require custom kernels to run efficiently on modern accelerators. Axial Transformers use axial self-attention layers and shift operation naturally and efficiently build full receptive fields of multidimenssonal tensors. Our model matches or outperforms the state-of-the- art NetgeN and ImageNet-64 image benchmarks and sets significant new state-of-the-art on the BAIR Robot Pushing video benchmark.
Figure 4: 64 x 64 ImageNet samples at temperature 1.0
8

---

Under review as a conference paper at ICLR 2020
Figure 5: 32 x 32 ImageNet samples at temperature 0.99
Figure 6: 15 x 64 x 64 BAIR Robot Pushing samples at temperature 1.0

---

Under review as a conference paper at ICLR 2020
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: arXiv:1607.0644 2016.
Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. PixeISNAIL: An improved au- toregressive generative model. In International Conference on Machine Learning, pp. 863-871, 2018.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint
2019.
Frederik Ebert, Chelsea Finn, Alex X Lee, and Sergey Levine. Self-supervised visual planning with temporal skip connections. In Conference on Robot Learning, pp. 344-356, 2017.
Zilong Huang, Xinggang Wang, Lichao Huang, Chang Huang, Yunchao Wei, and Wenyu Liu. Cc- net: Criss-cross attention for semantic segmentation. In Proceedings of the IEEE International Conference on Computer Vision, pp. 603-612, 2019.
Nal Kalchbrenner, Aäron Oord, Karen Simonyan, Ivo Danihelka, Oriol Vinyals, Alex Graves, and Koray Kavukcuoglu. Video pixel networks. In International Conference on Machine Learning, pp. 1771-1779, 2017.
Manoj Kumar, Mohammad Babaeizadeh, Dumitru Erhan, Chelsea Finn, Sergey Levine, Laurent Dinh, and Durk Kingma. Videoflow: A flow-based generative model for video. arXiv preprint arXiv:1903.01434 2019.
Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. arXiv preprint arXiv: vrXiiv/8/2.000 18/2.01688, 2018.
Ankur P Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model for natural language inference. arXiv preprint
2016.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pp. 4052- 4061, 2018.
Scott Reed, Aäron van den Oord, Nal Kalchbrenner, Sergio Gómez Colmenarejo, Ziyu Wang. Yutian Chen, Dan Belov, and Nando Freitas. Parallel multiscale autoregressive density estimation. In Proceedings of the 34th International Conference Machine Learning Scarning-Volume 70, pp. 2912- 2921. JMLR. org, 2017.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. PixeICN+++: Improving the PixeICNN with discretized logistic mixture likelihood and other modifications. In International Conference on Learning Representations (ICLR), 2017.
Lucas Theis and Matthias Bethge. Generative image modeling using spatial Istms. In Advances in Neural Information Processing Systems, pp. 1927-1935, 2015.
Benigno Uria, Mare-Alexandre Côté, Karol Gregor, lain Murray, and Hugo Larochelle. Neural autoregressive distribution estimation. The Journal of Machine Learning Research, 17(1):7184- 7220, 2016.
Aaron den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. International Conference on Machine Learning (ICML), 2016a.
Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Ko- ray Kavukcuoglu. Conditional image generation with PixelCNN decoders. arXiv preprint arXiv:1606.05322 05328, 2016b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems, pp. 5998-6008. 2017.
Dirk Weissenborn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. arXiv preprint
2019.
10


Reviews:
[{"review_id": "ryeuXv839B", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "It is known that the standard self-attention method is computationally expensive and cost a significantly large amount of storage when the number of points to be attended is large.\n\nThis paper attempts to solve this problem and proposed the Axial Attention method. It is claimed to be able to save an O(N^(d-1)/d) factor of resources over standard self-attention.\n\nThe proposed method looks novel to me, but some of the related works are missing and the experiment session is insufficient. \n\n1)  The author should at least include the following works which also aim to reduce the cost of self-attention. Since the author did not mention these works which also focus on solving the same problem, It is hard for me to judge if the proposed method is better than existing works.\n[a] CCNet: Criss-Cross Attention for Semantic Segmentation\n[b] A^2-Nets: Double Attention Networks\n\n2) self-attention has shown its effectiveness on a broad range of computer vision tasks, including image generation, detection, segmentation, and classification. I do not get why the proposed method is only benchmarked for generative models. Is it because the proposed method cannot be adopted on other popular CV tasks, such as detection, segmentation, and classification? Extra experiments should be included if the proposed method is not only designed for generative models.\n\n3) The ablation study is missing. The author directly compared its own method with other existing methods that are implemented and trained with different hyperparameters. It is hard to know which indeed benefits the accuracy gain and how significant is the proposed method. \n\n4) In table 2 and 3, I do not see a clear advantage of the proposed method over the SOTA methods.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "S1lPeozCYH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes axial attention as an alternative of self-attention for data arranged as large multidimensional tensors, which costs too much computational resource since the complexity of traditional self-attention is quadratic in order to capture long-range dependencies for full receptive fields. The axial attention is applied within each axis of the data separately while keeping information along other axes independent. Therefore, for a d-dimensional tensor with N = S^d, axial attention saves O(N^{(d−1)/d}) computation over standard attention. The proposed axial attention can be used within standard Transformer layers in a straightforward manner to produce Axial Transformer layers, without changing the basic building blocks of traditional Transformer architecture.  The authors did experiments on two standard datasets for generative image and video models: down-sampled ImageNet and BAIR Robot Pushing, and they claim that their proposed method matches or outperforms the state-of-the-art on ImageNet-32 and ImageNet-64 image benchmarks and sets a significant new state-of-the-art on the BAIR Robot Pushing video benchmark. \n\nReasons to accept:\n\n1.\tSimple, easy-to-implement yet effective approach to adapt self-attention to large multidimensional data, which can save considerable computation for efficiency, while still have competitive performance.\n2.\tClear writing, with sufficient but not redundant introduction of background knowledge and explanation of both the advantages and drawbacks of existing models (too large computational complexity on high-dimensional data).\n\nSuggestions for improvement:\n\n1.\tIt would be better if the authors can provide more analysis or case study to show the reason why Axial attention (Axial Transformer) can reach good performance even if it omits considerable operations compared to traditional Transformers, or to show why the attention operations within axis are important instead of attention operations between axis. \n2.\tDefinition of “axis” should be more clear in section 3 (there could be some ambiguities of “axis”).\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "BylznUt3qr", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper claims to propose a new approach to solve the computational problems of self-attention. However, the paper mainly focuses on adapting Transformer for image generation, which has far less applications. The whole paper needs to be rewritten to make their target and contribution clearer.\n\n1. The authors overclaim that they provide a new approach for accelerating self-attention. However, they only adapted Transformer for image generation. In fact, Transformer does not equal to self-attention. Currently, two directional self-attention like Bert has much wider applications compared with Transformer like sequential self-attention. \n\n2. For a paper claim to improve self-attention, they should show its effectiveness on a broad range of tasks, with comprehensive experimental evaluation. However, authors mainly reported the image generation on several datasets. \n\nOverall, the authors need to rewrite the paper. They should either show more applications with the proposed self-attention approach or treat it as a new approach for image generation.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "Bkp7J9Mor", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a novel approach to deal with the computational problems of self-attention without introducing independence assumptions. The proposed approach is simple, easy to understand, and easy to implement.\n\nHowever, evaluation for this paper is severely lacking. As it is, there is not enough information provided to adequately assess the proposed method's strengths in practice. The following should be added:\n\nEvaluation on a variety of different tasks, such as image segmentation, temporally consistent object detection, object tracking, etc. Why are the evaluations limited to generative modeling? To prove the generality of the method (as claimed), it needs to be applied to various tasks.\nRuntime (in inference) comparisons for each of the datasets and for each of the baselines. Additionally, a theoretical analysis for runtime in terms of the size of the input should be given (the column in Table 1 should have runtimes for each method clearly specified, and this should be done for each dataset and baseline)\nAblation study. What is the baseline architecture used without axial attention? There is only comparison to previous work which may have used a different architecture.\n\nIf these concerns are thoroughly addressed, I would be happy to increase my score.", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a self-attention-based autoregressive model called Axial Transformers for images and other data organized as high dimensional tensors. The Axial Attention is applied within each axis of the data to accelerate the processing.

Most of the authors claim that main idea behind Axial Attention is widely applicable, which can be used in many core vision tasks, such as detection and classification. However, the revision fails to provide more application for Axial attention.

Overall, the idea behind this paper is interesting but more convincing experimental results are needed.

Decision: Reject