MambaVision: A Hybrid Mamba-Transformer Vision Backbone

In plain words

Image models have mostly been built from two kinds of layers. Convolutions look at small neighbourhoods and are cheap. Attention lets every patch of an image look at every other patch, which sees the whole picture but gets expensive as the picture grows. A newer layer type, Mamba, handles long sequences cheaply, but it reads its input one step at a time, so it has a harder time taking in the whole image at once.

MambaVision is a hybrid. The early layers are plain convolutions, which are cheap at high resolution. The later layers mix Mamba-style blocks with attention blocks. The authors found that putting attention in the final layers gave the model much better understanding of long-range structure. Each Mamba-style block also gets a second, convolutional branch next to the sequence branch, and the two are joined before the output.

On ImageNet classification, the authors report a better accuracy-versus-speed trade-off than comparable Mamba, convolutional and transformer models. They also use the same network as a backbone for object detection and segmentation.

TL;DR

• Hierarchical four-stage backbone. Stages 1-2 are convolutional residual blocks; stages 3-4 mix a redesigned Mamba mixer with multi-head self-attention §3.1 §3.2.2.
• The Mamba mixer runs a selective scan on one half of its channels and a convolutional gate on the other half, then concatenates them eq. 7.
• Author-reported ImageNet-1K top-1: 82.3% (Tiny, 31.8M params, 6298 img/s) up to 85.3% (L2, 241.5M params, 1021 img/s), on an A100 at batch 128 tab. 1.
• Official code is non-commercial (NVIDIA Source Code License-NC). Implement from the paper and this card; do not copy code.

Problem

Transformers are costly on long inputs because self-attention scales quadratically with sequence length §1. Mamba, a state space model, scales linearly and matches or beats Transformers on language tasks §1. For vision, the authors argue that Mamba's sequential, autoregressive processing fits images poorly: pixels have no single natural order, and a step-by-step scan has limited access to global context in one pass §1. Bidirectional variants such as Vim add latency and do not always improve accuracy, according to the authors §1 §2. The gap the paper targets is the accuracy and throughput of pure Mamba vision backbones relative to ViT and CNN backbones §1.

Method

Macro architecture

MambaVision has four stages §3.1 app. B.

• Stem. Two consecutive 3×3 convolutions with stride 2, each followed by BatchNorm and ReLU, map the image to H/4 × W/4 × C §3.1 mamba_vision.py.
• Stages 1 and 2. Residual convolutional blocks eq. 1. Each block is conv 3×3, BN, GELU, conv 3×3, BN, plus the input eq. 1 mamba_vision.py.
• Stages 3 and 4. Layers of pre-normalised residual blocks eq. 6. Tokens are cut into non-overlapping windows before the layers and reassembled after them mamba_vision.py.
• Between stages. A 3×3 convolution with stride 2 halves the resolution §3.1 mamba_vision.py.

In stages 3 and 4, with N layers in a stage, the first N/2 layers use the MambaVision mixer and the last N/2 use self-attention §3.2.2. The ablation in tab. 5 finds that placing self-attention in the last N/2 layers of each stage works best (82.3%) §5.3.

The paper's supplementary architecture table tab. S.2 is useful but does not agree with the code in several places. See Paper vs code.

Mamba preliminaries

A discrete state space model updates a hidden state with a zero-order-hold discretisation of learned continuous parameters §3.2.1 eq. 3. Mamba makes the input and step-size parameters input-dependent (selectivity) §3.2.1. The paper gives the discretisation formulas but does not give implementation details of the scan; it refers to the original Mamba paper §3.2.1.

\bar{A} = \exp(\Delta A), \quad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I)\cdot(\Delta B), \quad \bar{C} = C eq. 3

MambaVision mixer

The mixer replaces Mamba's causal convolution with a regular convolution, since a causal convolution restricts influence to one direction §3.2.2. It adds a symmetric branch with no SSM, made of one convolution and a SiLU activation, to recover content the sequential scan may lose §3.2.2. Each branch projects to C/2 channels, the branches are concatenated, and a final linear layer maps back to C §3.2.2 eq. 7.

\begin{aligned} X_1 &= \mathrm{Scan}\big(\sigma(\mathrm{Conv}(\mathrm{Linear}(C,\tfrac{C}{2})(X_{in})))\big) \\ X_2 &= \sigma\big(\mathrm{Conv}(\mathrm{Linear}(C,\tfrac{C}{2})(X_{in}))\big) \\ X_{out} &= \mathrm{Linear}(\tfrac{C}{2}, C)\big(\mathrm{Concat}(X_1, X_2)\big) \end{aligned} eq. 7

Here σ is SiLU and Scan is the selective scan of §3.2.1 §3.2.2. The paper gives a PyTorch-like listing of the mixer alg. 1. It is not runnable as written. Use the steps below and the code citations.

# MambaVision mixer  [eq. 7]
# Input: tokens (B, L, C), where L tokens form one window in row-major order
# Output: tokens (B, L, C)

def mambavision_mixer(tokens):
    # One linear map from C to C, split into two halves of C/2 each  [eq. 7]
    # [code: mambavision/models/mamba_vision.py#L404-L406]
    projected = linear(tokens)                    # (B, L, C)
    seq_half, gate_half = split(projected, 2)     # (B, L, C/2) each
    # Depthwise 1D convolution over the L tokens, then SiLU, per branch  [eq. 7]
    seq_half = silu(depthwise_conv1d(seq_half))   # (B, L, C/2)
    gate_half = silu(depthwise_conv1d(gate_half)) # (B, L, C/2)
    # Selective scan on the sequence branch only  [eq. 7] [§3.2.1]
    seq_half = selective_scan(seq_half)           # (B, L, C/2)
    # Concatenate both branches, then project back to C  [eq. 7]
    mixed = concat(seq_half, gate_half)           # (B, L, C)
    return linear(mixed)                          # (B, L, C)

Layer

Each layer is pre-normalised. The mixer and the MLP each sit inside a residual connection eq. 6. Layer normalisation is used for the norm §3.2.2. The MLP expansion ratio is 4 in the code mamba_vision.py mamba_vision.py.

\begin{aligned} \hat{X}^{n} &= \mathrm{Mixer}(\mathrm{Norm}(X^{n-1})) + X^{n-1} \\ X^{n} &= \mathrm{MLP}(\mathrm{Norm}(\hat{X}^{n})) + \hat{X}^{n} \end{aligned} eq. 6

# One layer of stage 3 or 4  [eq. 6]
# The mixer is the MambaVision mixer or self-attention (see Hyperparameters)

def layer(x):
    # Token mixing with a residual connection  [eq. 6]
    x = x + mixer(layer_norm(x))
    # Channel MLP with a residual connection  [eq. 6]
    x = x + mlp(layer_norm(x))
    return x

Self-attention

Attention is standard multi-head self-attention, computed within windows §3.2.2 eq. 8. The paper's text says d_h is the number of heads; the scaling in the code uses the per-head dimension instead (see Pitfalls).

\mathrm{Attention}(Q, K, V) = \mathrm{Softmax}\!\left(\frac{Q K^{\mathsf{T}}}{\sqrt{d_h}}\right) V eq. 8

Stage-level sketch

# MambaVision forward pass  [§3.1] [app. B]
# Input: image (B, 3, H, W)

def mambavision(image):
    # Stem: two stride-2 3x3 convolutions  [§3.1]
    x = stem(image)                               # (B, C, H/4, W/4)
    for stage in stages:
        if stage.uses_convolution:
            # Stages 1-2: residual convolutional blocks  [eq. 1]
            x = conv_blocks(x)
        else:
            # Stages 3-4: windowed layers  [eq. 6]
            x = windowed_layers(x)
        # Downsample between stages, not after the last stage  [§3.1]
        if not stage.is_last:
            x = downsample(x)                     # 3x3 conv, stride 2
    # Head: batch norm, global average pool, linear classifier  [code: ...#L691]
    x = batch_norm(x)
    x = global_average_pool(x)
    return linear(x)

Hyperparameters

Model variants

Values are taken from the constructors in the code. The paper gives the same family in tab. S.2 but does not list depths or head counts in its text. mamba_vision.py

Stage widths are C, 2C, 4C, 8C for each variant mamba_vision.py. Stages 1 and 2 are convolutional (conv = i in (0, 1)) mamba_vision.py. Window sizes for classification are 8, 8, 14, 7 for the four stages mamba_vision.py app. A. At 224 input, stage 3 is 14×14 and stage 4 is 7×7, so each fits in one window (derived from the code).

Attention blocks are placed in the last ⌊N/2⌋ layers of stages 3 and 4, where the code takes range(depth//2, depth) for even depth and range(depth//2+1, depth) for odd depth mamba_vision.py.

Mamba mixer and attention (code-defined)

The paper does not state these. They come from the code.

ImageNet-1K training

Downstream training

Results

All numbers are author-reported. Throughput is measured on one A100 GPU at batch 128 for classification tab. 1.

ImageNet-1K classification tab. 1

Setup: 224×224 input unless noted; throughput on an A100 at batch 128 tab. 1.

ImageNet-21K pre-training (values from the text only; the plot in fig. 4 was not read): MambaVision-B rises from 84.2% to 84.9% at 224; L rises from 85.0% to 86.1% at 224; L3 (739.6M params) reaches 87.3% at 256 and 88.1% at 512 §5.3.

Object detection and instance segmentation, MS COCO tab. 2

Cascade Mask R-CNN, 3× schedule, crop 1280×800 tab. 2.

Semantic segmentation, ADE20K tab. 3

UperNet, crop 512×512 tab. 3.

Ablations

Token mixer design tab. 4. Setup: MambaVision-T layout; COCO with Mask R-CNN and 1× schedule tab. 4 §5.3.

Hybrid integration pattern tab. 5. S = self-attention, M = MambaVision mixer. All variants have 31.8M params tab. 5.

Window size tab. S.1 app. A. MambaVision-T, classification, COCO and throughput.

Paper vs code

Official code: https://github.com/NVlabs/MambaVision, linked from the abstract. Read at commit 7860a506b2eb844eaaae676f08461ce8c3c26f43. The repository is under the NVIDIA Source Code License-NC, which excludes commercial use LICENSE. The pretrained weights are CC BY-NC-SA 4.0 README.md. This card cites the code by path and line only. It quotes no code and uses no weights.

Differences found by comparing the paper's text, tables and appendix with the code:

1. Stem. Text: two stride-2 3×3 convolutions §3.1. Code: the same mamba_vision.py. Table tab. S.2 gives the second convolution no stride and lists the stem at 112×112. Follow the text and code.
2. Stage widths and block counts in tab. S.2. The table's widths are twice the code's widths for every variant. For T, the table lists 160, 320, 640, 1280 channels, and the code uses 80, 160, 320, 640 mamba_vision.py. Each table row pairs a stage's blocks with the downsampling convolution that feeds it. The code's stage at 56×56 is at width 80, stage at 28×28 at 160, and so on.
3. Last-stage block counts. For T, tab. S.2 lists 4 MambaVision + 4 self-attention blocks in the last stage. The code's depth there is 4, and its rule gives 2 + 2 mamba_vision.py. For S (depth 5) and B (depth 5), the table lists 6 blocks tab. S.2; the code has 5. The parameter count in tab. 1 (31.8M for T) can settle which is intended, by instantiating the constructor and counting. Not done here.
4. Activation. Paper eq. 1 writes GELU. The code uses the tanh approximation mamba_vision.py.
5. Mixer convolution. The paper says "regular convolution" §3.2.2. The code uses depthwise 1D convolution (groups equal the channel count) with kernel 3 mamba_vision.py mamba_vision.py.
6. Δ activation and bias. The paper's equations do not show a softplus or a learned Δ bias. The code applies softplus to Δ and adds the dt projection bias mamba_vision.py.
7. Windowing of the Mamba blocks. The paper describes windowed attention §3.2.2. The code windows every block of stages 3 and 4, including the Mamba mixers, so the SSM scan runs within each window mamba_vision.py.
8. Algorithm 1. The listing uses self.d_model, dt_min, dt_max, D and n -> d n without defining them, and omits the softplus and dt bias. The code mamba_vision.py is the reference.
9. Attention scale. The appendix says d_h is the number of heads app. A. The code scales by the per-head dimension mamba_vision.py.
10. ImageNet recipe. Paper: 300 epochs, LR 4e-3 app. C. Tiny config: 310 epochs, LR 0.005 mambavision_tiny_1k.yaml mambavision_tiny_1k.yaml. The example script uses LR 5e-4 and batch 256 with 8 GPUs train.sh; it is an example, not the recipe. The Base config also has 310 epochs and LR 0.005 mambavision_base_1k.yaml mambavision_base_1k.yaml.
11. Training additions not in the paper. Distillation against an EMA teacher, weight 0.5, from 25% of training mambavision_tiny_1k.yaml train.py. EMA decay 0.9998 mambavision_tiny_1k.yaml. Mixup, CutMix, RandAugment and random erasing also appear there and are not in the paper.
12. Crop. Training script default crop 0.875 train.py; Tiny config 1.0 mambavision_tiny_1k.yaml. Not stated in the paper.
13. Drop path. The Tiny config has drop_path: null mambavision_tiny_1k.yaml, and the script's default is None train.py. The model function defaults to 0.2 only if the key is absent mamba_vision.py. Set the value explicitly.
14. Segmentation. The paper says 8 A100 GPUs §4. The config is trained on 16 GPUs with one image per GPU mamba_vision_160k_ade20k-512x512_tiny.py.
15. Detection batch. The paper says batch 16 app. C. The base dataset config uses 2 per GPU [code: object_detection/configs/base/datasets/coco_instance.py#L38]. The GPU count is not stated.
16. Window sizes per task. Classification uses 14/7 for stages 3/4 mamba_vision.py. Detection uses 112/56 cascade_mask_rcnn_mamba_vision_tiny_3x_coco.py. Segmentation uses 64/32 mamba_vision_160k_ade20k-512x512_tiny.py. The paper gives window sizes only for classification.
17. Dependency. The code calls selective_scan_fn from the external mamba_ssm package mamba_vision.py. The scan is not implemented in this repository. Pure-PyTorch scans are not provided here.

Text claims that disagree with the paper's own tables:

• Ablation text: causal-conv baseline ADE20K mIoU "44.2% (−1.4)" app. A. Table 4 shows 44.2 against 46.0, a gap of −1.8 tab. 4.
• Ablation text: baseline mask AP "40.2 (−1.6)" §5.3. Table 4 says 40.4 tab. 4.
• Ablation text: concatenation gives "+0.9 in mIoU" §5.3. Table 4 gives 46.0 against 45.7, a gain of +0.3 tab. 4.
• Text: MambaVision-B has "56% less GFLOPs than MaxViT-B" §5.1. Table 1 gives 15.0 against 23.4 G, a reduction of about 36%.
• Text: Swin-B 83.5% top-1 §5.1 does not appear in tab. 1.

No discrepancy was found in the detection and segmentation mIoU/AP values versus the text §5.2.

Pitfalls

• Global attention and scan at 224 input. Stages 3 and 4 are 14×14 and 7×7 at 224 input, and their windows are 14 and 7 mamba_vision.py. Each stage is one window, so attention and the SSM scan see the whole map. This follows from the code and window partitioning mamba_vision.py. Window size matters only at other resolutions.
• Padding is not masked. When a feature map is not a multiple of the window size, zeros are added mamba_vision.py. Attention is called without a mask mamba_vision.py, and the SSM scan runs over the padded tokens. Padded positions influence outputs. Detection with 1280×800 input, after downsampling, triggers this padding cascade_mask_rcnn_mamba_vision_tiny_3x_coco.py.
• Token order. Tokens inside a window are scanned in row-major order mamba_vision.py. The paper does not state the order.
• Convolution bias. The mixer's convolutions are built with bias=conv_bias//2 mamba_vision.py. With conv_bias=True the expression is 0, so no bias is created. Adding a bias would deviate from the code.
• Algorithm 1 is not a drop-in implementation. See Paper vs code, item 8.
• Mamba placement. The first N/2 layers of each of stages 3 and 4 use the Mamba mixer and the last N/2 use attention §3.2.2 tab. 5. For odd depth the rule is range(depth//2+1, depth), so the extra layer goes to the Mamba side mamba_vision.py.
• Throughput depends on the scan kernel. Throughput figures are measured on an A100 at batch 128 tab. 1. The code uses the external mamba_ssm CUDA kernel mamba_vision.py. Expect different numbers with another scan implementation.
• Equation numbers. They were counted from the LaTeX source order. The compiled PDF was not read. Check them against the PDF.
• Figures. Figures are described in words only: fig. 1 is the Pareto plot of top-1 against throughput, fig. 2 the macro architecture, fig. 3 the MambaVision block, fig. 4 ImageNet-21K scaling, fig. 5 self-attention maps, and fig. S.1 attention visualisations §5.3 app. A.
• Licensing of the code. The code is NC. Reimplement from the paper and the parameter values on this card, citing paths for facts only.

Lineage

• Builds on Mamba: Gu and Dao, 2023, arXiv:2312.00752 §3.2.1 [bib gu2023mamba]. No card exists for it, so only its arXiv link is given.
• Related Mamba vision backbones compared in the tables: VMamba (Liu et al., arXiv:2401.10166) §2 tab. 1; Vim (Zhu et al.) §2 tab. 1; EfficientVMamba §2 tab. 1; SiMBA §2 tab. 1.
• Transformer and convolutional baselines: Swin Transformer §2, ConvNeXt §2, NextViT §2, FasterViT §2, MaxViT tab. 1. None has a card here.
• Superseded by: not known from this wiki's content.