Axial Attention in Multidimensional Transformers

TL;DR

The paper introduces Axial Transformers, a novel self-attention-based model for high-dimensional data that achieves state-of-the-art results while maintaining computational efficiency and ease of implementation.

Contribution

It presents axial attention, a new self-attention mechanism that scales efficiently to high-dimensional data and enables state-of-the-art generative modeling performance.

Findings

Achieved state-of-the-art results on ImageNet-32 and ImageNet-64 benchmarks.

Demonstrated effectiveness on the BAIR Robotic Pushing video benchmark.

Maintained full distribution expressiveness with a semi-parallel decoding structure.

Abstract

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.