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

Mamba’s autoregressive processing can limit global-context use in vision tasks that need it for accurate local predictions. MambaVision redesigns Mamba and adds Transformer blocks in final layers; it achieves a new ImageNet-1K Pareto front for Top-1 accuracy and image throughput, with experiments across multiple vision tasks.

  • Problem

    Mamba’s autoregressive processing is inefficient for spatial data and limits global-context use in one forward pass, although vision tasks often require global context for accurate local predictions.

    Source: §1 · pp. 1–2

  • Method

    MambaVision redesigns the Mamba block and combines it with Transformer blocks, finding that self-attention at final stages enhances global-context and long-range-dependency capture.

    Source: §1 · p. 2

  • Results

    Experiments across classification, detection, and segmentation demonstrate versatility, while MambaVision achieves a new ImageNet-1K Pareto front for Top-1 accuracy and image throughput.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The study supports Mamba-Transformer hybrids as a viable vision-backbone design and identifies final-stage self-attention as effective for capturing global context.

    Source: §1 · p. 2

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

1. Introduction

MambaVision is a redesigned, hierarchical hybrid Mamba-Transformer backbone for vision. The paper motivates the design by limitations in attention and autoregressive Mamba, and reports strong ImageNet-1K and downstream-task performance.

1. Introduction

  • Transformers [31] have quadratic attention complexity, while Mamba [7] offers linear-time sequence processing; autoregressive Mamba limits global-context capture in vision.Bidirectional SSMs such as Vim [40] seek to address context limitations but can add latency and complexity; the introduction says CNN and ViT backbones still outperform the best Mamba-based vision models.Source: §1 · pp. 1–2
  • MambaVision establishes a new ImageNet-1K Pareto front for Top-1 accuracy and image throughput, outperforming Mamba-, CNN-, and ViT-based models.The figure reports throughput measured on an A100 NVIDIA GPU with batch size 128.Source: Figure 1 · p. 1, §1 · p. 2
  • MambaVision backbones outperform comparably sized counterparts on MS COCO object detection and instance segmentation, and ADE20K semantic segmentation.Source: §1 · p. 2
  • The authors redesign the Mamba block for vision and propose a hybrid of MambaVision Mixer and MLP components with Transformer blocks.The hierarchical architecture uses multiple resolutions and CNN-based residual blocks for fast extraction of larger-resolution features.Source: §1 · p. 2
  • Their integration study finds that self-attention blocks in the final stages improve capture of global context and long-range spatial dependencies.They examine Transformer blocks in earlier, middle, and final layers, and at every l layers.Source: §1 · p. 2

2. Related work

Related work moves from convolutional and Transformer backbones to hybrids and vision SSMs. MambaVision is positioned as a hierarchical hybrid designed to combine efficient feature extraction with stronger short- and long-range spatial modeling.

Conv-Based

  • CNN work modernizes convolutional backbones through ConvNeXt, RegNetY [26], and EfficientNetV2 [27], but CNNs still lack the global receptive field needed for long-range dependencies.ConvNeXt redesigns ResNet [12] with increased width, larger kernels, and layer normalization; RegNetY uses design-space analysis, while EfficientNetV2 uses neural architecture search and progressive learning.Source: §2 · p. 2

Transformer-Based

  • ViTs [5] expanded receptive fields with self-attention, but quadratic complexity and early data demands motivated training and efficiency adaptations.DeiT [28] uses distillation for smaller datasets; Swin Transformer [21] uses shifted windows, while Twins [3] and PVT [32] use separable attention or hierarchical structures.Source: §2 · p. 2

Conv-Transformer

  • Hybrid models combine convolution and attention: CoAT, CrossViT [1], NextViT [17], Efficient-Former [18], and FasterViT [10] target feature learning or efficiency–accuracy trade-offs.Efficient-Former and FasterViT emphasize competitive performance with high throughput through hybrid architecture design.Source: §2 · p. 2

Mamba-Based

  • MambaVision’s hierarchy uses residual convolutional blocks in stages 1–2, then MambaVision and Transformer blocks in stages 3–4, with final-layer attention recovering global context.The final Transformer layers are described as capturing long-range spatial dependencies.Source: Figure 2 · p. 3
  • Vim [40] uses bidirectional SSMs for global context but adds overhead and slows training and inference; VMamba [20] uses four-way cross-scans, whereas MambaVision’s simpler mixer captures short- and long-range dependencies.The paper also reports that MambaVision achieves higher ImageNet Top-1 accuracy and throughput than Vim, and higher performance with significantly higher throughput than VMamba.Source: §2 · pp. 2–3
  • EfficientVMamba [25] uses SSMs at larger resolutions and CNNs at lower ones, unlike MambaVision’s opposite allocation; SiMBA [24] addresses stability but leaves spatial limits and trails it.MambaVision places CNN layers at higher resolutions and SSM/self-attention at lower resolutions.Source: §2 · p. 3

3. Methodology

MambaVision combines a four-stage hierarchical backbone with a redesigned vision mixer and self-attention in stages 3 and 4. Its method pairs CNN-based early feature extraction with sequential and spatial feature modeling in later stages.

3.1. Macro Architecture

  • A four-stage hierarchy uses CNN layers in stages 1–2 for fast high-resolution feature extraction, followed by MambaVision and Transformer blocks in stages 3–4.The stem uses two stride-2 3 × 3 CNN layers to embed overlapping patches, and stage downsamplers halve spatial resolution.Source: §3.1 · p. 3

3.2. Micro Architecture

  • The micro-architecture section revisits Mamba and state-space-model preliminaries before detailing the design of stages 3 and 4.Source: §3.2 · p. 3

3.2.1. Mamba Preliminaries

  • Mamba’s selection mechanism dynamically adjusts B, C, and Δ to process inputs selectively, while discretization converts continuous SSM parameters into discrete ones for computational efficiency.The MambaVision mixer uses selective scan and adds a symmetric branch without SSM to complement sequential modeling with spatial information.Source: §3.2.1 · pp. 3–4, §3.2.2 · p. 5

3.2.2. Layer Architecture

  • For N layers, the first N/2 use MambaVision mixer blocks and the remainder use self-attention; each layer has residual mixer and MLP layers, and its mixer combines regular-convolution and symmetric convolution-SiLU branches.Source: §3.2.2 · p. 4

4. Experiments

Experiments evaluate ImageNet-1K classification and downstream detection, instance segmentation, and semantic segmentation using pretrained MambaVision backbones. The section specifies training and task-specific evaluation configurations.

4. Experiments

  • ImageNet-1K classification follows prior training recipes [11] [21], with 300 epochs on 32 A100 GPUs and attention windows of 14 and 7 in stages 3 and 4.ImageNet-1K is cited as.Source: §4 · p. 5
  • Pretrained models are evaluated for detection and instance segmentation on MS COCO [19] using Cascade Mask-RCNN with a ×3 LR schedule.Source: §4 · p. 5
  • Semantic segmentation on ADE20K uses a UperNet head [34], with all experiments run on 8 A100 GPUs.Source: §4 · p. 5

5. Results

MambaVision reports strong ImageNet-1K classification and downstream segmentation results, scales to ImageNet-21K, and uses ablations to motivate its token-mixer design and final-stage self-attention.

5.1. Image classification

  • MambaVision-B reaches 84.2% ImageNet Top-1 accuracy, versus 83.8% for ConvNeXt-B, 83.5% for Swin-B, and 83.9% for VMamba-B, with higher image throughput than these comparators.MambaVision-B also has 56% less GFLOPs than MaxViT-B.Source: §5.1 · p. 5, Table 1 · p. 6

5.3. Ablation

  • ImageNet-21K pretraining raises MambaVision-B Top-1 accuracy from 84.2% to 84.9% and MambaVision-L from 85% to 86.1% at 224 resolution; MambaVision-L3 reaches 87.3% at 256 and 88.1% at 512.MambaVision-L3 has 739.6M parameters; the authors describe this as the first scaling of a Mamba-based vision architecture to ImageNet-21K.Source: §5.3 · p. 6, Figure 4 · p. 7

5.2. Object Detection and Segmentation

  • For MS COCO, MambaVision exceeds ConvNeXt and Swin variants in box and mask AP; on ADE20K, it beats Swin-T/S/B by +1.5, +0.6, and +1.0 mIoU, respectively.The reported downstream results were achieved without extensive hyperparameter optimization; MambaVision also attains higher mIoU than Focal Transformers across scales at comparable model sizes.Source: §5.2 · p. 6, Table 2 · p. 7, Table 3 · p. 7

5.3. Ablation

  • The token-mixer ablation found that concatenating the SSM and non-SSM branch outputs improved ImageNet Top-1 by +1.0%, COCO box AP by +1.1, mask AP by +0.8, and ADE20K mIoU by +0.9.The study used the MambaVision-T configuration; the authors attribute the gains to richer feature representations and enhanced global-context understanding.Source: §5.3 · p. 7, Table 4 · p. 8
  • Attention visualizations show final-stage self-attention focusing on semantically meaningful regions, including object boundaries, fine-grained features, and interacting subjects and objects.The examples include an aircraft, a bird, and a person holding an object; the authors interpret the visualizations as support for using final-stage self-attention to capture global context and long-range dependencies.Source: Figure 5 · p. 8, §5.3 · p. 8
  • Among hybrid patterns, self-attention in the last N/2 layers of each stage achieved the best Top-1 accuracy at 82.3%, compared with 81.9% for the last N/4 layers and 81.3% for random placement.These isoparameter MambaVision-T experiments support placing self-attention in later layers.Source: §5.3 · pp. 7–8, Table 5 · p. 8

6. Conclusion

MambaVision is a vision-tailored Mamba-Transformer hybrid that redesigns Mamba for global context representation. Its results support the hybrid’s effectiveness across vision tasks and highlight final-layer self-attention as a way to capture long-range dependencies efficiently.

6. Conclusion

  • MambaVision achieves a new SOTA Pareto front for Top-1 accuracy and image throughput, outperforming Transformer- and Mamba-based models by a significant margin.The conclusion reports this result without giving specific metric values.Source: §6 · p. 8
  • MambaVision is the first Mamba-Transformer hybrid backbone tailored for vision, with a redesigned Mamba formulation to enhance global context representation learning.Source: §6 · p. 8
  • Systematic analysis found that placing self-attention blocks in the final layers significantly improves long-range dependency capture while maintaining efficiency.Source: §6 · p. 8
  • Experiments across classification, detection, and segmentation demonstrate MambaVision’s versatility and effectiveness across multiple vision tasks.Source: §6 · p. 8
  • With ImageNet-21K pretraining, MambaVision achieves strong performance matching SOTA models, demonstrating potential for large-scale vision applications.Source: §6 · p. 8

Appendix · G. Ablation Study

The ablation finds that a stage-3 window size of 14 improves vision-task performance over size 7 with nearly unchanged throughput; the authors select 14 and 7 as their default window sizes. The study evaluates MambaVision-T on classification, detection, and instance segmentation.

G. Ablation Study

  • The study compares stage-3 window sizes of 7 and 14 in MambaVision-T across image classification, object detection, and instance segmentation.Source: G. Ablation Study · p. 11
  • At window size 14, ImageNet top-1 accuracy reaches 82.3% and COCO mask AP reaches 41.8%, with throughput of 6298 vs. 6318 img/s for the two window sizes.The authors describe these as consistent improvements with nearly identical throughput.Source: G. Ablation Study · p. 11, Table S.1 · p. 11
  • When the input size exceeds the window size, self-attention is computed within local windows.Source: G. Ablation Study · p. 11
  • The authors select 14 and 7 as their default window sizes based on the performance and computational-cost trade-off.Source: G. Ablation Study · p. 11
  • The authors characterize the throughput decrease as negligible at 0.3% and justified by gains across vision tasks.They note the trade-off on modern hardware such as the NVIDIA A100 GPU.Source: G. Ablation Study · p. 11
  • Throughput is measured for image classification with batch size 128 on a single NVIDIA A100 GPU; COCO experiments use a Mask-RCNN head and ×1 LR schedule.Source: Table S.1 · p. 11

H. Architecture Details

MambaVision uses a four-stage hierarchical backbone with convolutional down-sampling, and introduces a Mamba–self-attention hybrid in Stages 3 and 4. Its T, S, B, and L variants scale channel dimensions and layer counts to offer different complexity–performance trade-offs.

H. Architecture Details

  • The backbone has four stages, each using convolutional down-sampling to reduce spatial resolution by a factor of two.Source: H. Architecture Details · p. 11
  • Stages 3 and 4 combine Mamba-based sequence modeling with self-attention for efficient sequence processing and global context modeling.Source: H. Architecture Details · p. 11
  • Variants T, S, B, and L scale channel dimensions and layer counts to provide different complexity–performance trade-offs.Source: H. Architecture Details · p. 11

I. Training Details

MambaVision was trained on ImageNet-1K with a 300-epoch, large-batch setup, then evaluated as a backbone on detection and segmentation tasks. The training details also specify stage-wise attention windows and task-specific downstream heads.

I. Training Details

  • ImageNet-1K training used 300 epochs on 32 A100 GPUs, LAMB, batch size 4096, and learning rate 4e-3.Source: §I · p. 11
  • Downstream evaluation used pre-trained backbones on MS COCO [19] for object detection and instance segmentation, and ADE20K for semantic segmentation.All downstream tasks used AdamW and batch size 16; Cascade Mask-RCNN with a ×3 LR schedule served detection and instance segmentation, while UperNet [34] served semantic segmentation.Source: §I · p. 11
  • Self-attention in stages 3 and 4 used window sizes of 14 and 7, respectively, across all MambaVision variants.Source: §I · p. 11

J. Interpretability

MambaVision’s attention visualizations illustrate its interpretability across diverse object categories, highlighting semantic regions and object boundaries. Paired examples show consistent patterns, contextual adaptation, and fine-grained feature attention.

J. Interpretability

  • Paired examples show consistent attention patterns across similar objects and adaptability across contrasting domains.The passage describes paired scenarios including manmade objects and natural subjects.Source: J. Interpretability · p. 11
  • Attention heat maps highlight fine-grained discriminative features and accurately trace object boundaries across varied positions and backgrounds.Examples include animal-fur textures and structural elements of vehicles and containers.Source: J. Interpretability · p. 11, Figure S.1 · p. 12
  • Across diverse categories, MambaVision’s attention maps localize key semantic regions and object boundaries.Figure S.1 shows paired input, heat-map, and overlay examples across six rows of varied subjects.Source: J. Interpretability · p. 11, Figure S.1 · p. 12

Related papers

Loading 2407.08083v2…