Axial Attention in Multidimensional Transformers

1 INTRODUCTION . Autoregressive models are a family of exact likelihood-based generative models that represent the joint distribution of data x = ( x1 , . . . , xN ) as a product of conditionals pθ ( x ) = ∏N i=1 pθ ( xi |x < i ) . Neural network models in this family have achieved state-of-the-art log likelihoods on highdimensional image and video datasets ( van den Oord et al. , 2016a ; Chen et al. , 2018 ; Menick & Kalchbrenner , 2018 ; Parmar et al. , 2018 ; Child et al. , 2019 ; Weissenborn et al. , 2019 ; Salimans et 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 : 1 . Large , high information bandwidth receptive fields for each pixel xi , capable of expressing long-range dependencies over previous pixels x < i , 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 to express all joint distributions over the data . Meanwhile , architectures that admit fast log likelihood gradient computation are suitable for training using a stochastic gradient method on a maximum likelihood objective—a straightforward , stable training procedure for generative models . These desiderata make self-attention a compelling building block for autoregressive model architectures . Self-attention is a neural network operation that is able to transform a sequence y1 , . . . , yN into a sequence y′1 , . . . , y ′ N , where each y ′ i depends on all yi by way of a single vectorizable computation ( Vaswani et al. , 2017 ) . Self-attention is remarkably effective at learning long-range dependencies between data dimensions and neural networks that incorporate self-attention in their designs are state-of-the-art on many tasks from language modelling and machine translation to image and video modelling ( Parmar et al. , 2018 ; Child et al. , 2019 ) . 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 the case of autoregressive models of multidimensional 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 can encompass a 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 self-attention-based autoregressive model for data organized as multidimensional tensors . Rather than applying attention to a flattened string of tensor elements , our model instead applies attention along a single axis of the tensor without flattening—we refer to this as “ axial attention. ” Since the length of any single axis ( that is , the height or width of an image ) is typically much smaller than the total number of elements , an axial attention operation enjoys a significant saving in computation and memory over standard self-attention : for a d-dimensional tensor with shape N = N1/d × · · · × N1/d , axial attention saves a O ( N ( d−1 ) /d ) factor of resources over standard self-attention . Our Axial Transformer architecture allows for the majority of the context x < i to be embedded with a high degree of parallelism without introducing conditional independence assumptions among any of the locations , but has an interesting property that it is amenable to a simple-to-implement fast sampling 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 . We evaluate Axial Transformers on image and video modelling benchmarks . We show that Axial Transformer achieves state-of-the-art results on ImageNet-32 and on ImageNet-64 . We also show that , simply by stacking a video along the channel dimension , the Axial Transformer can be directly applied to the channel-stacked video without nearly any modification . On the BAIR Robot Pushing benchmark , the Axial Transformer significantly outperforms previous results without using an architecture 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 are simple to implement using efficient operations that are widely available in deep learning frameworks ( primarily dense-dense MatMuls ) . An open source implementation of our models is available at anonymized URL . 2 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 a length N sequence of D-dimensional embeddings X ( a N ×D matrix ) and produces an output sequence Y ( also a N ×D matrix ) via : Q = XWQ , K = XWK , V = XWV A = softmax ( QK > / √ D ) , Y = AV WQ , WK , and WV are D × D parameter matrices responsible for projecting the entries of the sequence X into keys , queries , and values , respectively . Each entry of the output sequence Y is a 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 ( N2 ) 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 : a small 32×32×3 image has 3072 dimensions . Sequences such as these are too long for self-attention , so attempts to scale self-attention to these modalities generally involve restricting these sequence lengths in a modality-aware manner while attempting to preserve modeling performance . One strategy is to restrict the conditioning context x < i 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 et al . ( 2018 ) propose image models with conditioning context x < i restricted to a small window of the full image , but the implementation requires redundant data copies to extract and process these windows . Weissenborn et 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 4 . A different strategy is to stack multiple sparse attention layers , each with restricted context for computational efficiency , but in a manner that overlapping these layers yields a full-context model . Child et 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 matrix-matrix-multiply . The model can not be easily implemented on other hardware such as TPUs . See table 1 for a summary of these architecture design tradeoffs . Our goal in this paper is to design attention-based autoregressive models that attain the best of all worlds . Our Axial Transformer , described in subsequent sections , has a full conditioning context , so its ability to express joint distributions 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 AXIAL TRANSFORMERS . We now describe Axial Transformers , our self-attention-based autoregressive models for highdimensional 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 self-attention-based autoregressive models for high-dimensional data tensors . The proposed approach does not change the original shape of the multidimensional data tensor and performs a masked or unmasked attention over a single axis of the tensor at a time . We call this operation axial attention , denoted by Attentionk ( x ) . It performs attention over axis k of the tensor x , mixing information along axis k while keeping information along other axes independent . It is straightforward to implement : axial attention over axis k can be implemented by transposing all axes except k to the batch axis , calling standard attention as a subroutine , then undoing the transpose ( an alternative is to use the einsum operation available in most deep learning libraries ) . When the data is an image , we call Attention1 column attention , as it mixes information within columns while keeping separate columns independent . We call Attention2 row attention for analogous reasons . Axial attention on a square image of size N = S × S performs attention on S sequences of length S—this is a total of O ( S · S2 ) = O ( N √ N ) computation—an O ( √ N ) savings in computation over standard self-attention . In general , for a d-dimensional tensor with N = Sd , axial attention saves O ( N ( d−1 ) /d ) computation over standard attention . Of course , a single layer of axial attention along some axis k does not have the full receptive field since it covers a single axis , but we will see in section 3.2 that stacking two axial attention layers allows the model to obtain a global receptive field . It will be important for us to also define MaskedAttentionk to be the causally masked variant of Attentionk : component i of the result of MaskedAttentionk ( x ) along axis k depends on only components 1 , . . . , i of x 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 in section 3.2 . Axial attention can be used within standard Transformer layers in a straightforward manner to produce Axial Transformer layers . The basic building blocks are the same as those found in the standard Transformer architecture : • LayerNorm ( x ) : layer normalization ( Ba et al. , 2016 ) , and • DenseD ( x ) : a dense layer operating over the last axis of the input x . The letter D denotes the dimension of the output activations . If the input has shape H × W × C , then this operation is identical to a 1× 1 convolution , and the output has shape H ×W ×D . We use these to define ResNet axial attention blocks operating on tensors of D-dimensional embeddings ( Vaswani et al. , 2017 ; Child et al. , 2019 ) : • FeedforwardBlock ( x ) = x+ DenseD ( Nonlinearity ( DenseD′ ( LayerNorm ( x ) ) ) ) • AttentionBlockk ( x ) = x+ DenseD ( Attentionk ( LayerNorm ( x ) ) ) • TransformerBlockk ( x ) = FeedforwardBlock ( AttentionBlockk ( x ) ) D′ is chosen to be some constant factor larger than D , from 1 to 4 ( Vaswani et al. , 2017 ) . We also define a MaskedTransformerBlockk using MaskedAttentionk in place of Attentionk . 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 .

Summary: 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.