[Long Review] Axial Attention in Multidimensional Transformers

Tutorials

[Long Review] Axial Attention in Multidimensional Transformers

Full self-attention over an image or a video is prohibitively expensive, so how do you keep the expressiveness of a Transformer without the quadratic blowup?

Full self-attention over an image or a video is prohibitively expensive, so how do you keep the expressiveness of a Transformer without the quadratic blowup? Ho, Kalchbrenner, Weissenborn, and Salimans propose axial attention: only attend along one axis of the tensor at a time, stack the axes, and you recover full context at a fraction of the compute.

This long-form review of the Axial Transformers paper walks through the architecture end to end, how axial attention layers compose, why the semi-parallel decoding structure actually works without independence assumptions, and how the model lands state-of-the-art results on ImageNet-32, ImageNet-64, and BAIR Robotic Pushing. Speech engineers should take note: the same trick applies to any high-dimensional tensor, including spectrogram-time-frequency grids, and the paper is open-sourced. The deep-dive walks through every design choice if you're thinking about efficient attention for audio or multimodal work.