Axial Attention in Multidimensional Transformers

Jonathan Ho; Nal Kalchbrenner; Dirk Weissenborn; Tim Salimans

Source-linked AI summary

Axial Attention in Multidimensional Transformers

Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, Tim Salimans

arXiv:1912.12180v1Dec 20, 2019cs.CV

TL;DR

Autoregressive models for high-dimensional images and videos must balance computational cost, expressive joint distributions, and implementation practicality. The paper introduces axial attention, which attends along tensor axes and enables parallel context computation with efficient sampling while preserving full context. Axial Transformers match or exceed state-of-the-art results on ImageNet-32, ImageNet-64, and BAIR Robot Pushing.

Problem
High-dimensional autoregressive modeling makes full-context self-attention computationally expensive, while efficiency-oriented alternatives can limit joint-distribution expressiveness or require specialized operations.
Method
Axial Transformers apply masked or unmasked self-attention along individual tensor axes, combining axial layers and shifts to build causal full receptive fields using standard operations.
Results
Axial Transformers match or outperform the state of the art on ImageNet-32 and ImageNet-64 and set a significant new state of the art on BAIR Robot Pushing.
Takeaways & Limitations
The architecture preserves fully expressive joint distributions while reducing multidimensional attention resources and enabling highly parallel context computation and fast sampling.
Takeaways & Limitations
A row-attention-only variant fails to define a full-context model, yielding independent autoregressive models for separate rows.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper introduces Axial Transformers to make self-attention-based autoregressive modeling practical for multidimensional data while preserving full joint-distribution expressiveness. Axial attention reduces resource demands, and the architecture supports highly parallel context computation and fast sampling.

Architecture: The model uses unmasked row and column attention for prior channels and frames, then masked attention and a shift operation to enforce causality during row sampling.Figure 1 describes 8 encoding blocks, 4 integration blocks, and a 4-block inner decoder.
Motivation: Self-attention becomes prohibitively expensive for long multidimensional sequences because memory and computation grow quadratically with sequence length.Images and videos require large receptive fields across multiple dimensions, worsening the resource problem.
Axial attention: Axial attention applies self-attention along one tensor axis without flattening, reducing computation and memory by O(N^(d−1)/d) for d-dimensional tensors with N elements.Each axis is typically much shorter than the total number of tensor elements.
Architecture: The architecture embeds most previous context in parallel and samples one image row autoregressively without re-embedding pixels from earlier rows.This structure preserves dependencies among locations and defines a fully expressive joint distribution.
Results: Axial Transformers achieve state-of-the-art results on ImageNet-32 and ImageNet-64 and significantly outperform previous results on the BAIR Robot Pushing benchmark.The same architecture can be applied to channel-stacked video without nearly any modification.
Implementation: The models avoid specialized accelerator subroutines and use efficient operations widely available in standard deep learning frameworks.The implementation primarily uses dense-dense matrix multiplications and is open sourced.

2 BACKGROUND

The background frames multidimensional self-attention as a trade-off between expressive full-context modeling and feasible computation. Axial Transformers address this trade-off with single-axis attention and standard operations while retaining full conditioning context.

Self-attention: A self-attention layer maps an N×D sequence of embeddings to an N×D output using pairwise query-key similarities and value combinations.The projections into queries, keys, and values use D×D parameter matrices.
Self-attention: Self-attention requires O(N^2) time and space because it computes pairwise similarities and weighted combinations across the sequence.This cost provides expressive power but limits direct application to long image and video sequences.
Computational trade-offs: Flattening even a 32×32×3 image produces 3072 dimensions, making direct self-attention impractical.Existing approaches therefore restrict sequence lengths or conditioning contexts to control cost.
Computational trade-offs: Restricting context reduces attention cost but prevents the model from expressing all joint distributions over the data.Sparse-layer approaches can recover full context, but some require custom sparse GPU kernels.
Axial Transformers: Axial Transformers retain full conditioning context without redundant data copies or custom kernels, using standard deep learning library operations.This design targets both expressive modeling and efficient implementation.

3 AXIAL TRANSFORMERS

Axial Transformers build autoregressive models for multidimensional tensors by applying masked or unmasked attention along individual axes. Layering these operations provides full context while enabling efficient implementation and semi-parallel sampling.

3.1 AXIAL ATTENTION: Axial attention applies attention along one tensor axis without flattening, preserving the original tensor shape while mixing information along that axis.It can be implemented by transposing other axes into the batch dimension, applying standard attention, and transposing back.
3.1 AXIAL ATTENTION: For a d-dimensional tensor with N elements, axial attention saves O(N^(d−1)/d) computation over standard self-attention.On square images, attention operates over multiple sequences of length equal to one image dimension rather than one flattened sequence.
3.1 AXIAL ATTENTION: Masked axial attention restricts each position to earlier positions along its selected axis, providing the causal building block for autoregressive modeling.The model uses standard Transformer components including layer normalization, dense layers, attention blocks, and feedforward blocks.
3.2.1 A MODEL FOR SINGLE-CHANNEL IMAGES: The autoregressive image model follows raster order, transforming embedded pixels into logits that depend only on previous pixels.For a single-channel H × W image, pixel intensities are embedded into H × W × D and mapped to H × W × 256 logits.
3.2.1 A MODEL FOR SINGLE-CHANNEL IMAGES: A row-wise inner decoder uses masked row attention, right-shifting inputs so each position excludes the current pixel from its receptive field.This structure supports parallel context embedding, but row attention alone would produce independent autoregressive models for each row.
3.2.1 A MODEL FOR SINGLE-CHANNEL IMAGES: The outer decoder captures all pixels in preceding rows using unmasked row and masked column layers, shifts that context downward, and adds it to the inner decoder.This converts the row-wise model into a fully expressive autoregressive model with context from the current row and all rows above.
3.2.1 A MODEL FOR SINGLE-CHANNEL IMAGES: The resulting architecture permits semi-parallel sampling because one image row can be sampled by running the autoregressive Transformer only over that row.The reported sampling procedure is N faster than naive sampling that repeatedly evaluates the full network.

4 EXPERIMENTS

The experiments evaluate Axial Transformers on image and video generation benchmarks, including a smaller LSTM inner decoder for semi-parallel sampling. Samples show globally coherent ImageNet scenes and detailed, diverse BAIR video continuations.

4 EXPERIMENTS: The models were benchmarked on downsampled ImageNet and BAIR Robot Pushing using standardized Axial Transformer configurations.The configurations use 8 encoder layers, 8 outer-decoder layers, 4 inner-decoder layers, hidden size 2048, and 16 attention heads.
4 EXPERIMENTS: A single-layer 2048-unit LSTM inner decoder slowed training by about 20% on ImageNet32 and about 80% on ImageNet64.The comparison keeps the number of training steps and other settings fixed.
4 EXPERIMENTS: The Axial Transformer + LSTM inner decoder performed rather well on both ImageNet32 and ImageNet64 benchmarks.This result also demonstrates the feasibility of making the inner decoder smaller for semi-parallel sampling.
4.2 SAMPLES: ImageNet samples were globally coherent and recognizable, while BAIR continuations preserved detailed, diverse motion at temperature 1.0.The BAIR samples use the first dataset frame as input and generate the remaining continuation frames.

5 CONCLUSION

The Axial Transformer applies attention along tensor axes to reduce resource requirements while retaining expressive multidimensional autoregressive modeling. It matches or exceeds state-of-the-art image results and achieves a new state-of-the-art on BAIR Robot Pushing.

5 CONCLUSION: Axial attention achieves O(N (d−1)/d) savings in computation and memory for a d-dimensional tensor with N elements.The operation scales better with input dimensionality and is implemented without custom kernels on modern accelerators.
5 CONCLUSION: Axial Transformer layers and a shift operation build full receptive fields for multidimensional tensors efficiently.The architecture is designed for data organized as high-dimensional tensors.