MaxViT: Multi-Axis Vision Transformer

TL;DR

MaxViT addresses the scalability gap of self-attention in vision by introducing Max-SA, a multi-axis attention that combines blocked local attention with dilated global grid attention in linear time. The authors build a hierarchical backbone by stacking MaxViT blocks that merge MBConv with local and global attention, yielding strong performance across classification, detection, aesthetics assessment, and generation. Across ImageNet, COCO, AVA, and GAN tasks, MaxViT demonstrates state-of-the-art or competitive results with favorable compute/parameter efficiency and scalable behavior to large data regimes. This work offers a practical pathway for integrating convolutional inductive bias with sparse attention, with potential extensions to other modalities and tasks.

Abstract

Transformers have recently gained significant attention in the computer vision community. However, the lack of scalability of self-attention mechanisms with respect to image size has limited their wide adoption in state-of-the-art vision backbones. In this paper we introduce an efficient and scalable attention model we call multi-axis attention, which consists of two aspects: blocked local and dilated global attention. These design choices allow global-local spatial interactions on arbitrary input resolutions with only linear complexity. We also present a new architectural element by effectively blending our proposed attention model with convolutions, and accordingly propose a simple hierarchical vision backbone, dubbed MaxViT, by simply repeating the basic building block over multiple stages. Notably, MaxViT is able to ''see'' globally throughout the entire network, even in earlier, high-resolution stages. We demonstrate the effectiveness of our model on a broad spectrum of vision tasks. On image classification, MaxViT achieves state-of-the-art performance under various settings: without extra data, MaxViT attains 86.5% ImageNet-1K top-1 accuracy; with ImageNet-21K pre-training, our model achieves 88.7% top-1 accuracy. For downstream tasks, MaxViT as a backbone delivers favorable performance on object detection as well as visual aesthetic assessment. We also show that our proposed model expresses strong generative modeling capability on ImageNet, demonstrating the superior potential of MaxViT blocks as a universal vision module. The source code and trained models will be available at https://github.com/google-research/maxvit.

Figures (9)

• Figure 1: Performance comparison of MaxViT with state-of-the-art vision Transformers on ImageNet-1K. Our model shows superior performance in terms of both accuracy vs. computation and accuracy vs. parameters tradeoff.
• Figure 2: MaxViT architecture. We follow a typical hierarchical design of ConvNet practices (e.g., ResNet) but instead build a new type of basic building block that unifies MBConv, block, and grid attention layers. Normalization and activation layers are omitted for simplicity.
• Figure 3: Multi-axis self-attention (Max-SA) (best viewed in color). An illustration of the multi-axis approach for computing self-attention (window/grid size is 4$\times$4). The block-attention module performs self-attention within windows, while the grid-attention module attends globally to pixels in a sparse, uniform grid overlaid on the entire 2D space, with both having linear complexity against input size, as we use fixed attention footage. The same colors are spatially mixed by the self-attention operation.
• Figure 4: Performance comparison on large-scale pre-trained models. MaxViT shows superior scaling performance under both ImageNet-21K and JFT-300M pre-trained settings.
• Figure 5: Vertical layout ablation. Our model scales better than Swin layeout liu2021swin.