PaperlayerView exact revision on arXiv

Source-linked AI summary

MambaVision: A Hybrid Mamba-Transformer Vision Backbone

Ali Hatamizadeh, Jan Kautz

arXiv:2407.08083v2cs.CV

TL;DR

Vision backbones must balance Transformers’ global-context capability against their quadratic complexity and Mamba’s limitations for spatially integrated vision processing. MambaVision redesigns Mamba and combines it with later-stage self-attention in a hierarchical backbone, achieving strong accuracy-throughput results across classification and downstream vision tasks.

  • Problem

    Transformers are computationally expensive because attention has quadratic sequence-length complexity, while autoregressive Mamba formulations can limit global-context modeling for vision.

  • Method

    MambaVision redesigns the Mamba block and combines MambaVision Mixer and MLP components with Transformer blocks in a hierarchical, multi-resolution architecture.

  • Results

    MambaVision achieves a new SOTA Pareto front for ImageNet-1K Top-1 accuracy and image throughput and outperforms comparably sized backbones across detection and segmentation tasks.

  • Takeaways & Limitations

    Self-attention in final layers significantly improves global-context and long-range-dependency modeling while the hybrid backbone provides favorable throughput and broad vision-task performance.

Abstract

from arXiv · show

We propose a novel hybrid Mamba-Transformer backbone, MambaVision, specifically tailored for vision applications. Our core contribution includes redesigning the Mamba formulation to enhance its capability for efficient modeling of visual features. Through a comprehensive ablation study, we demonstrate the feasibility of integrating Vision Transformers (ViT) with Mamba. Our results show that equipping the Mamba architecture with self-attention blocks in the final layers greatly improves its capacity to capture long-range spatial dependencies. Based on these findings, we introduce a family of MambaVision models with a hierarchical architecture to meet various design criteria. For classification on the ImageNet-1K dataset, MambaVision variants achieve state-of-the-art (SOTA) performance in terms of both Top-1 accuracy and throughput. In downstream tasks such as object detection, instance segmentation, and semantic segmentation on MS COCO and ADE20K datasets, MambaVision outperforms comparably sized backbones while demonstrating favorable performance. Code: https://github.com/NVlabs/MambaVision

NVIDIA

The passage contains an NVIDIA contact email.

  • The listed contact email uses the NVIDIA domain.

1. Introduction

MambaVision addresses limitations of Mamba and Transformer backbones for vision by combining a redesigned Mamba formulation with Transformer blocks in a hierarchical architecture. The resulting models target efficient visual processing while improving global-context modeling and ImageNet accuracy-throughput tradeoffs.

  • Transformers provide flexible attention-based modeling but incur quadratic complexity with sequence length.
  • Mamba offers linear-time sequence processing, but its autoregressive formulation can limit global-context modeling for vision tasks.
  • MambaVision redesigns the Mamba block and places self-attention blocks in later layers to capture global context and long-range spatial dependencies.
  • Its hierarchical design uses CNN-based residual blocks for fast higher-resolution feature extraction and combines MambaVision with Transformers in later stages.
  • MambaVision achieves a new SOTA Pareto front on ImageNet-1K for Top-1 accuracy and image throughput, outperforming Mamba-, CNN-, and ViT-based models.

2. Related work

Prior vision backbones span CNNs, Vision Transformers, hybrid Conv-Transformer models, and Mamba-based architectures. MambaVision is positioned within this landscape as a hierarchical design combining convolutional, Mamba, and Transformer components.

  • Conv-Based: Modern CNNs improve efficiency and accuracy through wider networks, larger kernels, normalization, systematic design, and progressive learning.
  • Transformer-Based: Vision Transformers enlarge receptive fields through self-attention but initially require extensive training data and lack CNNs’ inherent advantages.
  • Conv-Transformer: Hybrid Conv-Transformer models combine convolutional processing with self-attention to balance feature learning, efficiency, and throughput.
  • Mamba-Based: Vision Mamba variants use bidirectional or multi-directional scanning to capture global context, but these approaches increase computational overhead and latency.
  • Mamba-Based: MambaVision uses CNN layers in its first two stages and MambaVision and Transformer blocks in later stages within a four-stage hierarchy.

3. Methodology

MambaVision combines a hierarchical CNN front end with redesigned MambaVision mixers and Transformer blocks in later stages. Its design replaces restrictive causal processing with spatially suitable operations and places self-attention in final layers to capture long-range dependencies.

  • 3.1. Macro Architecture: MambaVision uses four hierarchical stages: CNN-based layers in stages 1–2, followed by MambaVision and Transformer blocks in stages 3–4.The CNN stages support fast feature extraction at higher input resolutions.
  • 3.1. Macro Architecture: The stem forms overlapping patches using two consecutive 3 × 3 stride-2 CNN layers, while inter-stage downsamplers halve spatial resolution.The stem projects features into a C-dimensional embedding space.
  • 3.2.2. Layer Architecture: A MambaVision block uses residual normalization, token mixing, and MLP updates, with the first N/2 layers using MambaVision mixers and the remaining N/2 using self-attention.The framework also supports windowed attention, with query, key, value, and head dimensions defined for the attention mechanism.
  • 3.2.2. Layer Architecture: The MambaVision mixer replaces causal convolution with regular convolution and adds a symmetric non-SSM branch to combine sequential and spatial information.The two branch outputs are projected to C/2 dimensions, concatenated, and passed through a final linear projection.
  • 3.2.2. Layer Architecture: The redesigned mixer uses selective scan operations and SiLU activations, with the authors reporting richer representations, better generalization, and improved computer-vision performance.The paper provides PyTorch-like pseudocode for the MambaVision mixer and experimentally validates its design choices.

4. Experiments

Image classification experiments use ImageNet-1K and a standardized training setup to support comparable performance analysis across models.

  • Image classification: Models are trained for 300 epochs using 32 A100 GPUs under a standard recipe from previous work.Self-attention uses window sizes of 14 in stage 3 and 7 in stage 4 across all MambaVision variants.

5. Results

MambaVision achieves strong classification, detection, and segmentation results while scaling across model sizes and resolutions. Ablations show that its redesigned token mixer and late-stage self-attention are central to these results.

  • Image classification: 84.2% Top-1 accuracy lets MambaVision-B exceed ConvNeXt-B and Swin-B on ImageNet-1K while delivering significantly better image throughput.It also exceeds VMamba-B at the same accuracy comparison and uses 56% less GFLOPs than MaxViT-B.
  • Semantic segmentation: +1.5, +0.6, and +1.0 mIoU are MambaVision-T, -S, and -B improvements over Swin counterparts on ADE20K.These gains were obtained without extensive downstream-task hyperparameter optimization.
  • Large-scale training: 87.3% and 88.1% Top-1 accuracy are achieved by MambaVision-L3 at 256 and 512 resolutions after ImageNet-21K scaling.MambaVision-B rises from 84.2% to 84.9% at 224 resolution, while MambaVision-L rises from 85% to 86.1%.
  • Token mixer design: +1.0% ImageNet Top-1, +1.1 box AP, +0.8 mask AP, and +0.9 mIoU result from concatenating SSM and non-SSM branch outputs.The ablation attributes these gains to richer feature representations and enhanced global context understanding.
  • Hybrid pattern: 82.3% Top-1 accuracy is obtained when self-attention occupies the last N/2 layers of each stage, outperforming other tested hybrid patterns.Self-attention in only the last N/4 layers reaches 81.9%, whereas random integration reaches 81.3%.
  • Interpretability: Final-stage attention maps focus on semantically meaningful regions, including aircraft bodies, bird head and tail features, and interacting people or objects.These visualizations indicate attention to object boundaries and fine-grained details without explicit supervision.

6. Conclusion

MambaVision is a vision-tailored hybrid Mamba-Transformer backbone that combines redesigned Mamba modeling with final-layer self-attention to improve global-context learning. It achieves strong accuracy-throughput performance across classification and downstream vision tasks.

  • MambaVision redesigns the Mamba formulation to enhance global context representation learning for vision applications.
  • Self-attention blocks in the final layers significantly improve long-range dependency capture while maintaining efficiency.
  • MambaVision achieves a new SOTA Pareto front for ImageNet-1K Top-1 accuracy and image throughput.
  • Experiments across classification, detection, and segmentation demonstrate the approach’s versatility and effectiveness.
  • ImageNet-21K pretraining produces strong performance matching SOTA models for large-scale vision applications.

G. Ablation Study

The ablation study evaluates attention window sizes across classification, detection, and instance segmentation. A window size of 14 improves benchmark performance with nearly unchanged throughput compared with 7.

  • 82.3% ImageNet top-1 accuracy and 41.8% COCO mask AP are achieved with the larger window size.These improvements come with nearly identical throughput: 6298 img/s versus 6318 img/s.
  • The analysis selects 14 and 7 as default window sizes because the combination improves vision performance while preserving efficiency.The reported throughput decrease is 0.3%.

H. Architecture Details

MambaVision uses a four-stage hierarchical backbone with progressive convolutional down-sampling. Stages 3 and 4 combine Mamba sequence modeling with self-attention, while variants scale model capacity for different trade-offs.

  • The hierarchical backbone has 4 stages, each using convolutional down-sampling to reduce spatial resolution by a factor of two.
  • Stages 3 and 4 combine Mamba-based sequence modeling with self-attention mechanisms.This design uses Mamba for efficient sequence processing and self-attention for global context modeling.
  • Variants T, S, B, and L scale channel dimensions and layer counts to provide different complexity-performance trade-offs.

I. Training Details

Image classification is trained on ImageNet-1K using a 300-epoch, 32-A100-GPU setup. Pretrained models are also used as backbones for detection and segmentation tasks.

  • Image classification experiments use ImageNet-1K training for 300 epochs on 32 A100 GPUs.
  • Training uses the LAMB optimizer, batch size 4096, and learning rate 4e-3.
  • Self-attention uses window sizes 14 and 7 in stages 3 and 4, respectively, across MambaVision variants.
  • Pre-trained models serve as backbones for object detection, instance segmentation, and semantic segmentation.

J. Interpretability

MambaVision’s interpretability analysis visualizes learned attention patterns across diverse object categories using paired input, heat-map, and overlay examples. The examples indicate consistent pattern recognition and adaptability across contrasting scenarios.

  • J. Interpretability: Figure S.1 presents paired examples visualizing MambaVision’s learned attention patterns across diverse object categories.Each case uses an original image, attention heat map, and attention overlay.
  • J. Interpretability: The visualizations indicate consistent attention patterns across different instances of similar object categories.
  • J. Interpretability: Contrasting paired examples show the model’s adaptability across manmade objects, natural subjects, and varied activities.
Loading 2407.08083v2…