Source-linked AI summary
MaxViT: Multi-Axis Vision Transformer
Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, Yinxiao Li
TL;DR
Vision self-attention has limited scalability with image size, motivating an efficient alternative that retains global interactions. The paper introduces Max-SA, which combines blocked local and dilated global attention at linear complexity, and builds MaxViT by stacking these operations with convolutions. MaxViT achieves state-of-the-art performance across a broad range of vision tasks and data regimes, while experiments at billion-parameter scale remain future work.
Problem
Vision Transformers’ self-attention does not scale efficiently with image size, while local attention can sacrifice non-locality and model capacity.
Method
Max-SA decomposes dense attention into blocked local and dilated global attention, and MaxViT hierarchically stacks Max-SA with convolutions.
Results
MaxViT achieves state-of-the-art results across image classification, object detection, segmentation, aesthetic assessment, and image generation under various data regimes.
Takeaways & Limitations
MaxViT provides a unified vision backbone that combines efficient convolution with sparse attention and maintains local and global spatial interactions throughout the network.
Takeaways & Limitations
Experiments on billion-parameter models trained on planet-scale datasets such as JFT-3B are left for future work.
Abstract
from arXiv · showhide
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.
1 Introduction
MaxViT addresses the scalability limits of vision self-attention by combining local and global interactions with linear complexity. Its hierarchical hybrid backbone achieves strong performance across diverse vision tasks and efficiency trade-offs.
- MaxViT outperforms recent Transformer-based models on both accuracy-versus-computation and accuracy-versus-parameter trade-offs.
- Vision Transformers can underperform without extensive pre-training because their capacity and weaker inductive bias can lead to overfitting.
- Max-SA combines local and global spatial interactions, adapts to different input lengths, and maintains linear complexity even in high-resolution stages.
- MaxViT hierarchically stacks repeated blocks that unify Max-SA with convolutions, capturing local and global interactions throughout the network.
- MaxViT achieves state-of-the-art results across classification, detection, segmentation, aesthetic assessment, and image generation under diverse data regimes.
2 Related work
Vision research evolved from dominant convolutional networks toward Transformer-based architectures. Related work adapts Transformers to visual locality, two-dimensional structure, efficiency, and hierarchical processing.
- ConvNets established the dominant vision architecture and improved through residual and dense connections, encoder-decoder schemes, feature pyramids, and other efficiency-oriented designs.
- ViT showed that pure Transformer architectures can solve vision problems by treating image patches as visual words, stimulating extensive research.
- Swin Transformer introduced shifted-window attention in a hierarchical architecture to address locality and the two-dimensional structure of images.
- Subsequent vision Transformers explored sparse attention, improved locality, pyramidal designs, and more efficient training strategies.
3 Method
MaxViT decomposes dense attention into blocked local and grid-based global operations, then combines Max-SA with MBConv in a hierarchical backbone. The design preserves global receptive fields while keeping attention complexity linear.
- Max-SA decomposes dense attention into window and grid attention, reducing vanilla attention’s quadratic complexity to linear without losing non-locality.
- MaxViT Architecture: MaxViT repeats identical blocks across four hierarchical stages, combining Max-SA with MBConv while progressively reducing resolution and increasing channels.
- Relative self-attention is used as the key operator in MaxViT to provide content-dependent spatial mixing with learned relative positional bias.
- Multi-axis Attention: Block attention partitions feature maps into non-overlapping P × P windows and performs local interactions within each window.
- Multi-axis Attention: Grid attention applies self-attention over a fixed uniform grid, producing dilated global spatial mixing with linear complexity.
- Multi-axis Attention: Using fixed window and grid sizes balances local and global computation, while Max-SA can replace Swin attention with the same parameters and FLOPs.
4 Experiments
MaxViT is evaluated across classification, detection, aesthetic assessment, generation, and ablation settings, with strong results across tasks and data regimes. Experiments also examine scaling, attention placement, block ordering, and architectural components.
- Large-scale pre-training: 88.38% accuracy is achieved by MaxViT-B with ImageNet-21K pre-training, outperforming CoAtNet-4 by 0.28% using 43% of its parameters and 38% of its FLOPs.MaxViT-XL reaches 88.70% at 512 × 512, while JFT-300M training yields 89.53% with 475 million parameters.
- Downstream tasks: MaxViT-T outperforms existing image aesthetic assessment methods on AVA at similar resolutions, with performance improving as input resolution increases.The model also shows better linear correlation than the cited state-of-the-art method using multi-resolution inputs.
- Image generation: MaxViT achieves better FID and IS than HiT for 128 × 128 unconditional ImageNet generation with significantly fewer parameters.The result supports the effectiveness of MaxViT blocks for generation tasks.
- Ablation studies: Ablations show that earlier global grid-attention improves performance over local attention or convolutions, while sequential multi-axis stacking outperforms parallel counterparts with fewer parameters and computation.The study also compares MBConv usage and block orders; global-to-local ordering is reported as best for generation.
5 Discussion and Conclusion
The paper presents MaxViT as a unified architecture combining efficient convolution and sparse attention, with state-of-the-art performance across vision tasks and strong scalability to massive datasets.
- MaxViT unifies efficient convolution and sparse attention in a simple design that achieves state-of-the-art performance across varied vision tasks.
- The appendix documents architectures for image classification, detection and segmentation, aesthetic assessment, and image generation.
- The appendix reports comprehensive experiments on ImageNet-1K, ImageNet-21K, JFT, and ImageNet-1K image generation visualizations.
A.1 Backbone Details
The MaxViT backbone combines MBConv with block and grid attention, using relative attention and complementary local-global spatial layouts to build a hierarchical vision model.
- MBConv is MaxViT’s main convolution operator, using a pre-activation structure to promote homogeneity between convolutional and Transformer blocks.
- Relative attention combines learned location-aware bias with input-adaptive attention and supports translation equivariance and global interactions.
- Block attention partitions features into non-overlapping P × P blocks for local spatial interactions, while grid attention performs dilated global spatial mixing.
- Max-SA applies local attention followed by sparse global attention, unlike axial attention’s columnwise-then-row-wise sequence, while retaining a global receptive field with O(N) complexity.
- The MaxViT block contains MBConv, block attention, and grid attention and is repeated across a four-stage hierarchical backbone.
- For detection and segmentation, MaxViT supplies multiscale S2, S3, and S4 features to an FPN and detection head.
A.3 Image Aesthetics Model
The image aesthetics model is designed to capture both local visual quality and global semantic information, using normalized Earth Mover’s Distance to predict score distributions.
- Image aesthetic assessment requires modeling pixel-level qualities such as sharpness and contrast alongside semantic qualities such as composition and depth-of-field.
- The model uses normalized Earth Mover’s Distance between ground-truth and predicted score distributions as its training loss.
- The loss uses cumulative distribution functions over N = 10 score bins with r = 2.
- For aesthetics assessment, the MaxViT classification head is replaced by a fully connected layer with 10 neurons followed by softmax.
- The GAN experiment uses a latent code, progressive hierarchical upsampling, cross-attention, and MaxViT blocks across five stages.
B.1 ImageNet Classification
The ImageNet-1K experiments evaluate MaxViT variants under pretraining and fine-tuning settings, with model-specific stochastic-depth regularization and documented training hyperparameters.
- The ImageNet-1K experiments provide separate pretraining and fine-tuning settings for MaxViT variants.
- Stochastic-depth rates are customized for each MaxViT model size to regularize the models separately.
- The reported hyperparameters include AdamW optimization, model-specific learning rates, training epochs, batch sizes, augmentation, and loss settings.
B.2 Coco Detection and Segmentation
MaxViT is evaluated on COCO2017 object detection and instance segmentation using ImageNet-1K-pretrained backbones, with task-specific fine-tuning at 896 × 896 resolution.
- COCO2017 provides 118K training and 5K validation samples for object detection and instance segmentation evaluation.
- ImageNet-1K-pretrained MaxViT backbones initialize the downstream detection and segmentation models.
- 896 × 896 resized inputs, batch size 256, AdamW, and model-specific stochastic depth are used for both tasks.Learning rates are 1e-3, 3e-3, and 3e-3 for MaxViT-T/S/B, with stochastic depths 0.8, 0.3, and 0.3.
B.3 Image Aesthetics Assessment
The paper evaluates MaxViT on visual aesthetic assessment and describes a GAN-based image-generation setup using a ResNet discriminator and R1-regularized logistic losses.
- The AVA benchmark contains 255K images rated by amateur photographers, with each image receiving an average score from 1 to 10 from about 200 raters.
- AVA ground truth is represented by each image’s histogram of human ratings, where higher scores indicate better visual aesthetic quality.
- The image-generation model uses a ResNet-based discriminator and standard non-saturating logistic GAN loss with R1 gradient penalty.
- GAN training uses Adam at learning rate 1e-4 for both networks, batch size 256, one million TPU steps, and gradient-penalty weight 10.
- The discriminator objective combines real-sample and generated-sample logistic terms with an R1 gradient penalty weighted by γ.
C Complete Experimental Results
The paper provides complete performance comparisons for ImageNet-1K, ImageNet-21K, and JFT-pretrained models, alongside additional unconditional ImageNet-1K generation results.
- Additional visual results for unconditional ImageNet-1K image generation are provided in Figure 8.
- Figure 8 presents unconditional generation results on ImageNet-1K at 128 × 128 resolution.
- Table 13 reports complete performance comparisons under the ImageNet-1K-only setting.
- Table 14 reports complete performance comparisons for ImageNet-21K- and JFT-pretrained models.