Axial Attention in Multidimensional Transformers

By Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, Tim Salimans. 2019

666 citations recorded by the source corpus.

We propose Axial Transformers, a self-attention-based autoregressive model for images and other data organized as high dimensional tensors. Existing autoregressive models either suffer from excessively large computational resource requirements for high dimensional data, or make compromises in terms of distribution expressiveness or ease of implementation in order to decrease resource requirements. 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 naturally 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 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.

Subfields: Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (2020, neutral): “An alternative way to scale attention is to apply it in blocks of varying sizes (Weissenborn et al., 2019), in the extreme case only along individual axes (Ho et al., 2019; Wang et al., 2020a).”
Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation (2020, methodology): “Besides discriminative tasks, self-attention is also applied to generative modeling of images [89,8,31].”
VideoGPT: Video Generation using VQ-VAE and Transformers (2021, positive): “…et al., 2016; Mittal et al., 2017; Marwah et al., 2017; Vahdat & Kautz, 2020; Child, 2020), Flows (Dinh et al., 2014; 2016; Kingma & Dhariwal, 2018; Ho et al., 2019a), Oord et al., 2017; Razavi et al., 2019; Ramesh et al., 2021), and lately diffusion and score matching models (Sohl-Dickstein et…”
Perceiver: General Perception with Iterative Attention (2021, methodology): “2020a; others self-attend independently on rows and columns (Ho et al., 2019; Wang et al., 2020a); others replace convolution with patchwise self-attention (Parmar et al.”
Taming Transformers for High-Resolution Image Synthesis (2020, methodology): “[12] and [24] retain the full receptive field but can reduce costs for a sequence of length n only from n(2) to n √ n, which makes resolutions beyond 64 pixels still prohibitively expensive.”
Deformable DETR: Deformable Transformers for End-to-End Object Detection (2020, background): “To compensate, Child et al. (2019); Huang et al. (2019); Ho et al. (2019) attend key elements at ﬁxed intervals to signiﬁcantly increase the receptive ﬁeld on keys.”
Long Range Arena: A Benchmark for Efficient Transformers (2020, methodology): “Early work in this area has typically considered combinations of fixed patterns (Child et al., 2019; Ho et al., 2019; Beltagy et al., 2020; Zaheer et al., 2020) that learn sparse attention by considering combinations of fixed strides or local windows.”
Efficient Transformers: A Survey (2020, methodology): “…et al., 2019) O ( N √ N ) FP Reformer (Kitaev et al., 2020) O ( N log N ) LP Routing Transformer (Roy et al., 2020) O ( N N ) LP Axial Transformer (Ho et al., 2019) O ( N √ N ) FP Compressive Transformer (Rae et al., 2020) O ( N 2 ) RC Sinkhorn Transformer (Tay et al., 2020b) O ( B 2 ) LP…”
MaX-DeepLab: End-to-End Panoptic Segmentation with Mask Transformers (2020, background): “Both classes of methods have boosted various vision tasks such as image classification [18, 5, 72, 37, 57, 89, 28], object detection [91, 80, 72, 36, 10, 108], semantic segmentation [15, 106, 39, 29, 109, 107], video recognition [91, 18], image generation [73, 35], and panoptic segmentation [89].”
MSA Transformer (2021, methodology): “Feedforward Layers We deviate from Ho et al. (2019) in the interleaving of the feedforward layers.”
Transformers in Vision: A Survey (2021, methodology): “Colorization Transformer [109] is a probabilistic model based on conditional attention mechanism [83].”
Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers (2020, background): “Instead, [46] adopts a specially designed axial-attention [20] which is less scalable to standard computing facilities.”
Is Space-Time Attention All You Need for Video Understanding? (2021, background): “A decomposed attention over the two spatial axes of the image was proposed in (Ho et al., 2019; Huang et al., 2019; Wang et al., 2020b) and our (T+W+H) adds a third dimension (time) for the case of video.”
ViViT: A Video Vision Transformer (2021, neutral): “We note that factorising attention over input dimensions has also been explored in [29, 78], and concurrently in the context of video by [4] in their “Di-vided Space-Time” model.”
LambdaNetworks: Modeling Long-Range Interactions Without Attention (2021, background): “…memory footprint of global attention was quickly solved by multiple works which proposed to limit the size of the attention contexts such as local attention (Ramachandran et al., 2019; Hu et al., 2019) and axial attention (Ho et al., 2019; Wang et al., 2020a; Shen et al., 2020) (See Section D.2).”
Bottleneck Transformers for Visual Recognition (2021, methodology): “Being a hybrid model, BoTNet differs from pure attention models such as SASA [49], LRNet [33], SANet [68], Axial-SASA [31, 62] and ViT [15].”
Scaling Local Self-Attention for Parameter Efficient Visual Backbones (2021, background): “Other forms of non-global pixel interaction can also be implemented efficiently [6, 18, 54, 1].”
Multi-Scale Vision Longformer: A New Vision Transformer for High-Resolution Image Encoding (2021, background): “The first is the sparse attention mechanism, including contentindependent sparsity [28, 8, 29, 14] and content-dependent sparsity [17, 32, 37, 51].”