Click the blue text above to follow us
MaxViT: Multi-Axis Vision TransformerZhengzhong Tu, Hossein TalebiComments: ECCV 2022[v1] Mon, 4 Apr 2022 17:59:44 UTC (235 KB)https://github.com/google-research/maxvitby Z Tu · 2022 · Cited by 1045This paper proposes a multi-axis vision Transformer architecture called MaxViT, which innovatively reduces the computational complexity from quadratic to linear through a novel multi-axis self-attention mechanism (Max-SA), efficiently integrating the local feature extraction capabilities of convolution with the global interaction capabilities of Transformers, aiming to address the scalability and computational efficiency bottlenecks of existing vision Transformers when processing high-resolution images.
1. Key Contributions
1. Proposed the Multi-Axis Attention Module: Designed and implemented a novel, plug-and-play multi-axis attention module, which consists of two types of sparse attention mechanisms — Block Attention and Grid Attention.
(1) Block Attention: Performs self-attention within non-overlapping local windows to capture local spatial relationships.
(2) Grid Attention: Performs self-attention on a sparse grid that uniformly covers the entire feature map to capture long-range global spatial relationships.
(3) Linear Complexity: This design ensures that the computational complexity of both global and local spatial interactions is linear with respect to the input image size, significantly enhancing the model’s efficiency and scalability when processing images of any resolution.
2. Designed the MaxViT architecture and its modular building blocks: Proposed a hierarchical visual backbone network named MaxViT, whose core is a standardized building block (MaxViT-Block) that unifies MBConv (convolution), Block Attention (local), and Grid Attention (global). The entire network is constructed by simply stacking this module at different stages. This design allows the network to simultaneously “see” global information from early high-resolution stages to deeper low-resolution stages.
3. Achieved comprehensive SOTA performance: Extensive experiments demonstrate that MaxViT achieves state-of-the-art (SOTA) levels across various visual tasks and data scales.
(1) Image Classification: On the ImageNet-1K/21K and JFT-300M datasets, MaxViT outperforms contemporaneous top models such as CoAtNet, Swin Transformer, and ConvNeXt in both accuracy-FLOPs and accuracy-parameter trade-off curves.
(2) Downstream Tasks: As a backbone network, it has achieved excellent performance in object detection, instance segmentation, and visual aesthetic assessment tasks.
(3) Generative Models: Demonstrated the strong potential of the MaxViT block as a universal visual module, also performing well in image generation tasks.
2. Problem Definition and Research Motivation
This paper precisely focuses on the subproblem of designing a universal visual backbone network, particularly exploring how to efficiently integrate the local inductive bias of convolutional neural networks (ConvNets) with the global information interaction capabilities of vision Transformers (ViT) in a hierarchical architecture. It is not limited to a specific task (such as segmentation) but aims to create a universal foundational model applicable to various computer vision tasks such as image classification, object detection, and generative models.
Challenges and Motivation
1. Quadratic Complexity Bottleneck of Standard Self-Attention: The standard self-attention mechanism used in traditional ViTs has computational and memory requirements that grow quadratically with image size (or the number of tokens), severely limiting its application in high-resolution inputs or early stages of the network, making it impractical and computationally expensive.
2. Loss of Global Information in Local Attention: To address the complexity issue, models like Swin Transformer have adopted a window-based local attention mechanism. While this improves efficiency, it sacrifices long-range global interaction capabilities, leading to poor performance scalability on large-scale datasets (such as ImageNet-21K, JFT) that require a broader receptive field.
3. Trade-off Between Model Capacity and Generalization Ability: How to effectively combine global and local interactions within a limited computational budget to balance the model’s expressive power and generalization ability remains an unresolved challenge.
Introduction
The author argues the necessity of the research through a clear logical chain:
1. Acknowledging the Potential of Transformers: First, acknowledging the tremendous potential of Transformers in the field of computer vision;
2. Pointing Out the Core Contradiction: Next, directly pointing out that the “quadratic complexity” of standard self-attention is the main obstacle to its widespread application;
3. Analyzing the Limitations of Existing Solutions: Subsequently, analyzing the “window attention” scheme represented by Swin Transformer, the author acknowledges its success but immediately points out its defect of “non-locality loss” leading to poor scalability on large-scale data;
4. Introducing the Core Problem: Finally, focusing the research motivation on a clear and critical question: “How to efficiently integrate global and local interactions under controllable computational budgets?”
Argumentation
The argumentation process follows the classic research paper narrative logic of “background-problem-existing solutions-solution defects-paper goals”, making it easy to understand; sufficient citations: key arguments are supported by references (e.g., ViT[22], Swin[56], CoAtNet[19], etc.), enhancing the objectivity and credibility of the argument; hitting the nail on the head: accurately capturing the most core contradiction in the field of vision Transformer research at the time — “computational efficiency” vs. “global receptive field”, making its research motivation very clear and important; foresight: anticipating the trend of models developing towards larger datasets and higher resolutions, and pointing out the limitations of existing methods under this trend, highlighting the necessity and forward-looking value of its work.
Core Technical Hypothesis
The fundamental technical hypothesis underlying the proposed method is: by decomposing the full-size self-attention operation into two computationally sparser and more efficient forms: one performed on the spatially local “Block” axis, and the other on the spatially non-local but sparse “Grid” axis, it can effectively simulate the global and local spatial information interactions that standard self-attention can provide while maintaining linear computational complexity, thus achieving a better balance between model capacity, generalization ability, and computational efficiency.
3. Overall Model Structure and Module Functions
The core architecture of this paper is clearly presented in Figure 2 of the paper, which is divided into two parts: the upper part is the overall hierarchical structure of MaxViT, and the lower part is the detailed structure of its core building unit, MaxViTBlock.
Note: The MaxViT architecture follows the typical hierarchical design of convolutional neural networks (such as ResNet), but constructs a novel foundational building block that unifies MBConv, Block Attention layers, and Grid Attention layers, omitting normalization and activation layers for simplicity.
3.1. Overall Architecture
3.2. MaxViT Block Module Details
MaxViT Block is the soul of the entire architecture, innovatively integrating convolution, local attention, and global attention into a unified module, with its internal data flow as follows:
How Grid Attention achieves a “global perspective” at the operational level
Core Idea: The essence of Grid Attention lies in a clever data reshaping that combines spatially dispersed pixels into a new, shorter sequence, followed by standard self-attention computation on this new sequence.
Steps of Grid Attention Operation
Suppose: There is an input feature map of size H×W=14×14, with a grid size set to G×G=7×7 (this is the default value used in the paper).
Step 1: Conceptual Division – Covering Grid
Step 2: Key Operation – Data Reshaping (The Reshaping Trick)
Note: The multi-axis self-attention mechanism Max-SA, a schematic diagram of the multi-axis method for computing self-attention (window/grid size of 4×4), the Block Attention module performs self-attention operations within the window, while the Grid Attention module focuses globally on pixels in a sparse uniform grid covering the entire two-dimensional space; since we use a fixed attention range, the complexity of both modules is linear with respect to the input size; self-attention operations mix elements of the same color spatially.
Step 3: Applying Self-Attention
Step 4: Reverse Operation
Summary: Why is this method both global and efficient?
** Key Conclusion **
In the operation of Grid Attention: (1) The sequence length N is G×G; the core of Grid Attention is reshaping, which extracts corresponding position pixels from G×G different grid units (for example, all top-left corner pixels of the units) to form a new group (sequence). Therefore, the length of this new sequence, i.e., the number of elements participating in the attention computation N, is the total number of grid units G×G. (2) The channel C is the number of channels in the feature map, this operation only reshapes pixels in the spatial dimension without changing the depth of each pixel’s feature vector; if the input feature map is (112,112,192), then C is 192.
Grid Attention cleverly reduces the quadratic term that is decisive in self-attention from H×W (a huge number that varies with input size) to G (a small, fixed constant), thus achieving a linear computational complexity with a global receptive field.
5. Mathematical Logic and Formula Derivation of the Module
5.1. MBConv Module (MBConv Block)
Formula 1: MBConv (without downsampling)
Formula 2: MBConv (with downsampling)
5.2. Multi-Axis Attention Module
This is the core innovation of MaxViT, which decomposes global self-attention into two computationally cheaper sparse attention operations: Block Attention (local) and Grid Attention (global).
Formula 3: Relative Position Self-Attention
Appendix 1: The Complementary Role of MBConv and Relative Position Self-Attention
Appendix 2: The Size and Functioning of the Relative Position Bias Matrix B
Summarize: The key is to distinguish between the learnable bias parameter table Lookup Table and the dynamically constructed bias matrix B during each attention computation.
Formula 4 & 6: Block Attention
Formula 5 & 7: Grid Attention
6. Loss Function and Optimization Strategy
The paper treats MaxViT as a universal visual backbone network, validated across different visual tasks, and adopts corresponding standard loss functions for each task.
6.1. Loss Functions
The following are the loss functions explicitly mentioned in the paper for different tasks:
1. Softmax Cross-Entropy Loss with Label Smoothing
This loss function is commonly used for single-label, multi-class classification tasks, consisting of two parts: the Softmax function and the cross-entropy loss with label smoothing.
Summarize: The standard cross-entropy loss uses the true label as a “one-hot” vector, for example, if the true category is class c, then the true probability vector y has yc=1, and the rest are all 0. Label smoothing modifies this one-hot vector by distributing a small portion of probability (controlled by the smoothing factor ϵ) to other incorrect categories. This loss function penalizes the model’s “overconfidence” in the correct answer, encouraging better generalization ability.
2. Sigmoid Loss Function
As shown in the screenshot, this loss function is commonly used in multi-label classification scenarios, where an image can belong to multiple categories simultaneously. In this case, the most commonly used “Sigmoid loss function” is binary cross-entropy loss (Binary Cross-Entropy, BCE), calculated independently for each category.
6.2. Training Strategy
7. Experimental Design and Result Analysis
Note: Variants of the MaxViT architecture, B and C represent the number of blocks and channels at each stage, with each attention layer’s attention head set to 32. For MBConv, the expansion rate of 4 and contraction rate of 0.25 are always used in SE, with two Conv layers used in the backbone.
Datasets and Evaluation Metrics
From the standard ImageNet-1K to the large-scale ImageNet-21K, and then to the ultra-large-scale JFT-300M image classification tasks, the scalability of the MaxViT model is strongly demonstrated. The challenges of the datasets are reflected in multiple aspects: the diversity and fine-grained nature of ImageNet categories, the weak labels and enormous scale of JFT, the small objects and complex scenes of COCO, and the subjectivity of AVA. Achieving SOTA or competitive results on these datasets itself proves the model’s strong capabilities.
Evaluation Metrics
For the tasks evaluated, the paper adopts gold standard metrics: classification tasks use Top-1 accuracy; object detection and instance segmentation use AP/AP50/AP75; aesthetic assessment uses PLCC and SRCC; image generation uses FID and IS.
Efficiency and Parameter Count
In multiple graphs, such as Figure 1, Figure 4, and Figure 5, the authors plot accuracy vs. FLOPs and accuracy vs. parameter count curves, visually demonstrating the superiority of MaxViT in the trade-off between performance and computational/parameter costs.
In the result tables, such as Table 2, in addition to accuracy, the parameter count Params, computational load FLOPs, and throughput are clearly listed, providing detailed data support for the model’s efficiency.
Ablation Study
The paper conducts a detailed ablation study in Section 4.5, with the baseline model unified as MaxViT-T:
1. Global Grid Attention [Table 7]
2. Module Order Study [Table 8]
3. MBConv Layer [Table 9]
4. Sequential vs. Parallel Design [Table 10]
8. Critical Evaluation and Future Prospects
Technical Advantages
1. Linear Complexity Global-Local Collaborative Attention: The core technical highlight of this research is the proposed Multi-Axis Attention module, which cleverly decomposes dense full attention into local block attention and sparse grid attention, allowing MaxViT to maintain a global receptive field across the entire feature map while reducing the computational complexity of self-attention from quadratic to linear, making global interactions computationally feasible at all stages of the network (including the early stages with high resolution).
2. Concise, Efficient, and Scalable Hybrid Modular Design: The second advantage of MaxViT is that it constructs a concise yet powerful foundational module (MaxViT Block), which organically unifies the MBConv convolution module and the multi-axis attention module. By simply stacking this module, MaxViT builds a hierarchical and powerful backbone network. This modular design is not only easy to implement but also demonstrates excellent scalability, showing SOTA performance across datasets ranging from small-scale (ImageNet-1K) to ultra-large-scale (JFT-300M), proving its design’s universality for different data scales.
Inspirational: This research significantly inspires the field of computer vision (including image segmentation). It eloquently demonstrates that hybrid architectures can surpass purely convolutional networks or vision Transformers by combining the excellent inductive biases (such as translational invariance) of convolution with the global context modeling capabilities of Transformers. It provides a new paradigm for designing efficient visual backbone networks: constructing networks through a repeatable foundational unit that includes multi-scale/multi-type operators. This inspires researchers to explore how to efficiently integrate various information interaction patterns, such as local and global, sparse and dense, within a foundational module.
Limitations and Potential Defects
1. Limitations of Static Sparse Patterns: Grid Attention adopts a fixed, uniform sparse sampling pattern to simulate global interactions. However, the important information distribution in real-world images is often non-uniform. This static grid may miss two spatially distant but semantically strongly related areas (if they happen to fall outside the grid sampling points), thus limiting its ability to capture arbitrary long-range dependencies.
2. Fixed Information Processing Flow: Ablation studies indicate that the order of `convolution -> local attention -> global attention` is optimal for classification tasks, but not for generative tasks. This suggests that the fixed information processing flow within the MaxViT module may not be optimal for all visual tasks. An ideal design might allow the information flow to adaptively adjust based on task characteristics.
Potential Failure Scenarios
1. Long, slender objects with complex, irregular topologies: For objects like wires, tangled ropes, or the slender tails of animals, local block attention may only capture segments of the object, while sparse grid attention may “skip” certain parts of the object, making it difficult to form a complete and continuous structural understanding of the object.
2. Scenes requiring differentiation of large areas and high-frequency textures: In tasks that require distinguishing objects through subtle but widely distributed texture differences (e.g., distinguishing between two similar fabric materials), sparse global attention may not provide sufficient density of long-range pixel comparisons, leading to recognition failures.
3. Recognition tasks relying on non-uniform key areas: For example, in medical imaging, lesions may occupy only a small part of the image but need to be diagnosed in conjunction with the global anatomical structure. If the critical long-distance contextual information happens to fall within the “blind spots” of grid attention, the model may make incorrect judgments.
Follow-up Research Directions
1. Dynamic or Deformable Multi-Axis Attention: To address the limitations of static grids, a direct improvement is to make it “dynamic.” A mechanism can be designed to allow the sampling point positions of grid attention to be dynamically generated by the input features, similar to deformable convolution. This way, the model can learn to selectively focus global computational resources on the “key” areas of the image, thus more effectively capturing task-relevant long-range dependencies.
2. Task-Adaptive Gated Block: To address the limitations of fixed information flow, the MaxViT block can be improved by introducing gating mechanisms. For example, several learnable “gates” can be set to allow the network to adaptively adjust the contribution weights or connection paths of the convolution, local attention, and global attention components based on input features or task types. This way, the module can tend towards a “local to global” flow for classification tasks while switching to a more optimal mode for other tasks.
3. Decoder Design and Evaluation for Dense Prediction: This is an important experimental supplement that can design a corresponding decoder (e.g., a U-Net-like structure) for the MaxViT backbone network and conduct comprehensive training and evaluation on mainstream semantic segmentation datasets (such as ADE20K, Cityscapes). This not only fills the gaps in the original experiments but also allows for a deeper analysis of the quality of features extracted by MaxViT, and may require targeted adjustments to the attention mechanism to better preserve spatial details.
Appendix: The Real Meaning of FLOPs
