PaperlayerView exact revision on arXiv

Source-linked AI summary

VMamba: Visual State Space Model

Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, Jianbin Jiao, Yunfan Liu

arXiv:2401.10166v4cs.CV

TL;DR

Vision backbones need global modeling with better efficiency than quadratic self-attention, especially for long or high-resolution inputs. VMamba adapts selective state-space modeling through VSS blocks and four-way SS2D scanning, and reports strong performance across visual tasks with linear input-scaling complexity. The study also identifies unexplored pre-training compatibility and limited large-scale exploration as boundaries.

  • Problem

    Efficient vision architectures must retain global receptive fields and dynamic weighting without the quadratic complexity or task-level degradation associated with existing attention approaches.

    Source: §1 · p. 1

  • Method

    VMamba uses VSS blocks with SS2D, which traverses 2D image data along four routes and applies selective SSM processing to bridge 1D scans with spatial vision structure.

    Source: §1 · p. 2, §4.1 · p. 4, §4.2 · p. 4

  • Results

    VMamba achieves promising performance across image classification, object detection, and semantic segmentation while exhibiting linear growth in computational complexity with input length.

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

  • Takeaways & Limitations

    VMamba’s linear time complexity makes it advantageous for downstream tasks with large-resolution inputs.

    Source: §6 · p. 10

    Compatibility between existing large-scale unsupervised pre-training methods and SSM-based architectures remains unexplored, and limited resources prevented Large-scale models and fine-grained hyperparameter searches.

    Source: §6 · p. 10

Abstract

from arXiv · show

Designing computationally efficient network architectures remains an ongoing necessity in computer vision. In this paper, we adapt Mamba, a state-space language model, into VMamba, a vision backbone with linear time complexity. At the core of VMamba is a stack of Visual State-Space (VSS) blocks with the 2D Selective Scan (SS2D) module. By traversing along four scanning routes, SS2D bridges the gap between the ordered nature of 1D selective scan and the non-sequential structure of 2D vision data, which facilitates the collection of contextual information from various sources and perspectives. Based on the VSS blocks, we develop a family of VMamba architectures and accelerate them through a succession of architectural and implementation enhancements. Extensive experiments demonstrate VMamba's promising performance across diverse visual perception tasks, highlighting its superior input scaling efficiency compared to existing benchmark models. Source code is available at https://github.com/MzeroMiko/VMamba.

1 Introduction

VMamba adapts selective state-space modeling to vision by introducing SS2D, addressing the mismatch between 1D sequential scans and 2D image structure while retaining linear complexity. Across classification, detection, and segmentation, the VMamba family reports strong performance and input-scaling efficiency.

  • Existing efficient-attention approaches either restrict the effective receptive field or incur performance degradation across tasks.Source: §1 · p. 1
  • VMamba adapts Mamba’s linear-complexity selective state-space modeling to visual representation learning.Source: §1 · pp. 1–2
  • Each image patch gathers contextual information through a compressed hidden state along its scanning path, reducing computational complexity from quadratic to linear.Source: §1 · p. 2
  • VMamba consistently achieves higher ImageNet-1K classification accuracy than benchmark CNN, ViT, and SSM models across model scales.VMamba-Base reaches 83.9% top-1 accuracy and exceeds Swin by +0.4%, with throughput of 646 versus 458.Source: §1 · p. 2
  • SS2D bridges 1D array scanning and 2D plane traversal through four scanning routes, enabling selective SSMs to process vision data.Source: §1 · pp. 1–2
  • VMamba reports promising results across image classification, object detection, and semantic segmentation, with computational complexity growing linearly with input length.Source: §1 · p. 2

2 Related Work

Related work spans CNNs, ViTs, efficient-attention methods, and state-space models. These approaches seek stronger visual modeling and better computational scaling, especially for long or high-resolution inputs.

Convolutional Neural Networks (CNNs)

  • CNN research improves visual modeling and efficiency using operators such as depth-wise and deformable convolution.Source: §2 · p. 2

Vision Transformers (ViTs)

  • ViT research established Transformer-based vision models and later emphasized large-scale pre-training, distillation, and inductive bias.Source: §2 · pp. 2–3

Efficient Attention

  • Linear-attention variants reformulate attention to reduce computational complexity from quadratic to linear while addressing memory movement and parallelizability.Source: §2 · p. 3

State Space Models (SSMs)

  • State-space models emerged as alternatives to Transformers because self-attention has quadratic complexity for long sequences such as high-resolution images.Source: §2 · p. 3

3 Preliminaries

The preliminaries define continuous and discretized state-space models, then describe selective scanning as an input-dependent mechanism for contextual modeling. Selective scans retain efficient recurrence computation despite dynamic parameters.

Formulation of SSMs

  • Continuous-time SSMs map an input signal to an output response through a hidden state using a linear time-invariant system.Source: §3 · p. 3

Discretization of SSM

  • The continuous hidden-state dynamics are discretized over intervals so SSMs can be integrated into deep models.Source: §3 · p. 3
  • The discretized formulation approximates the zero-order hold method commonly used in SSM-based models.Source: §3 · p. 3

Selective Scan Mechanism

  • SS2D traverses image patches along four paths, processes the resulting sequences with separate S6 blocks, and merges them into a 2D feature map.Source: Figure 2 · p. 4, §3 · p. 4
  • Selective SSMs use input-dependent parameters to capture contextual information, but dynamic weights prevent direct use of convolution for efficient computation.Source: §3 · p. 4

4 VMamba: Visual State Space Model

VMamba is a hierarchical vision backbone built from VSS blocks that replace Mamba’s one-dimensional selective scan with SS2D for spatial processing. Architectural and implementation changes improve practical inference efficiency while preserving strong classification performance.

Network Architecture

  • VMamba uses Tiny, Small, and Base scales with hierarchical stages that progressively reduce spatial resolution and stack VSS blocks.Source: §4.1 · p. 4
  • VSS blocks replace Mamba’s S6 module with SS2D and remove the multiplicative branch because SS2D selectivity already supplies gating effects.Source: §4.1 · p. 4

2D-Selective-Scan for Vision Data (SS2D)

  • SS2D adapts selective scanning to non-sequential vision data while preserving input-dependent weights and selective processing.Source: §4.2 · pp. 4–5

Accelerating VMamba

  • Vanilla VMamba-T reaches 82.2% classification accuracy, but its 426 images/s throughput and 22.9M parameters expose deployment challenges.Source: §4.3 · p. 5
  • Implementation and architectural refinements target inference speed, including Triton Cross-Scan/Cross-Merge, selective-scan changes, linear transformations, and tensor-layout optimization.Source: §4.3 · p. 5

5 Experiments

VMamba is evaluated against benchmark vision backbones across classification, dense prediction, scalability, and selective-scan analyses. It combines strong task performance with efficient throughput and linear resource growth as input resolution increases.

  • VMamba is evaluated across image classification, object detection, instance segmentation, semantic segmentation, input-resolution scaling, receptive fields, activation maps, and scanning patterns.The experiments compare VMamba with benchmark models and analyze SS2D behavior and scalability.Source: §5 · p. 6

Image Classification

  • 82.6% top-1 accuracy is achieved by VMamba-T on ImageNet-1K, exceeding DeiT-S by 2.8% and Swin-T by 1.3% at similar FLOPs.Source: §5.1 · p. 6
  • 1,686 images/s is the VMamba-T throughput, while VMamba-S and VMamba-B reach 877 and 646 images/s, respectively.VMamba-T is 1.47× faster than S4ND-Conv-T and 1.08× faster than Vim-S, while retaining performance leads of 0.4% and 2.1%.Source: §5.1 · p. 6

Downstream Tasks

  • 47.3%/48.7%/49.2% object-detection mAPs are achieved by VMamba-T/S/B under 12-epoch fine-tuning, exceeding Swin-T/S/B by 4.6%/3.9%/2.3%.Instance-segmentation mAPs also exceed Swin and ConvNeXt, and the advantages persist with 36-epoch multi-scale training.Source: §5.2 · p. 7
  • 3.4% higher mIoU than Swin-T and 1.9% higher than ConvNeXt-T are achieved by VMamba-T on ADE20K in the single-scale setting.VMamba-S/B also outperform NAT-S/B, with advantages persisting under multi-scale input.Source: §5.2 · p. 7

Scalability

  • Linear FLOPs and memory growth are maintained as input resolution increases, alongside high throughput and stable classification performance.At 768 × 768 resolution, VMamba reaches 74.7% top-1 accuracy without fine-tuning and 79.2% with linear tuning.Source: §5.2 · pp. 7–8

Analysis

  • Cross-Scan outperforms unidirectional, bidirectional, and cascade scanning in computational efficiency and classification accuracy at comparable parameter and FLOP budgets.The four-way scan captures 2D contextual information, and removing DWConv further enhances the advantage.Source: §5.3 · pp. 9–10

6 Conclusion

VMamba is an efficient SSM-based vision backbone that bridges 1D selective scanning and 2D image traversal through SS2D. Its linear time complexity and extensive experiments support downstream processing of large-resolution inputs, while several research directions remain open.

  • VMamba integrates selective SSMs into visual processing through SS2D, bridging ordered 1D scanning and non-sequential 2D traversal.Source: §6 · p. 10
  • VMamba improves inference speed through a series of architectural and implementation refinements.Source: §6 · p. 10
  • Linear time complexity makes VMamba advantageous for downstream tasks with large-resolution inputs.Source: §6 · p. 10

Limitations

  • Existing pre-training methods and VMamba-specific pre-training techniques for SSM-based architectures remain unexplored.Source: §6 · p. 10
  • Limited computational resources prevented exploration of VMamba-Large and fine-grained hyperparameter search.Source: §6 · p. 10
  • The potential of VMamba for more generalized tasks and generalized scanning patterns remains unexplored.Source: §6 · p. 10

B Derivation of the Recurrence Relation of Selective SSMs

This section derives a recurrence for selective SSM hidden states and connects its sequence formulation to attention-like matrix operations. Associative scan algorithms enable the recurrence to be computed in linear time.

D Visualization of Attention and Activation Maps

The visualizations examine how SS2D’s four scanning paths construct attention and activation maps. They show traversal-aware foreground context and focused separation between foreground objects and background pixels.

E Detailed Experiment Settings

The experiments specify VMamba’s architectural variants, ImageNet training protocol, throughput evaluation, acceleration stages, and receptive-field analysis. The setup compares progressive design and implementation changes under controlled training and hardware conditions.

Network Architecture

Experiment Setting

Throughput Evaluation

Accelerating VMamba

Evolution of ERF

F Performance of the VMamba Family on Downstream Tasks

VMamba is evaluated on object detection, instance segmentation, and semantic segmentation using protocols built around Swin and standard vision libraries. The supplied passages specify the optimization schedules, training settings, and architectural table references for these downstream evaluations.

G Details of VMamba’s Scale-Up Experiments

VMamba’s scale-up experiments examine architecture details, implementation acceleration, classification under larger input resolutions, and computational scaling. At 768 × 768, VMamba retains stronger accuracy-throughput behavior than the compared models, while FLOPs increase linearly.

H Ablation Study

The ablation study examines scanning patterns, initialization choices, and downstream evaluation setups. Cross-Scan provides global effective receptive fields and 2D priors, while initialization changes do not significantly affect image-classification performance.

Scanning patterns

  • Cross-Scan is compared with Unidi-Scan, Bidi-Scan, and Cascade-Scan to assess alternative 2D feature-map traversal strategies.Source: §H.1 · p. 21
  • Among the analyzed scanning methods, Bidi-Scan, Cascade-Scan, and Cross-Scan show global effective receptive fields, while Cross-Scan and Cascade-Scan show 2D priors.Source: §H.1 · p. 21
  • DWConv contributes to establishing 2D priors and the formation of global effective receptive fields.Source: §H.1 · p. 21

Initialization

  • The initialization ablation compares the S4D-derived default with random and zero initialization, using simplified skip-connection settings for the alternatives.Source: §H.2 · pp. 21–22

Downstream evaluation

  • The study evaluates VMamba and benchmark models on COCO object detection and instance segmentation, and ADE20K semantic segmentation.Source: Table 7 · p. 22, Table 8 · p. 23

Initialization

  • For image classification with SS2D blocks, model performance is not significantly affected by the initialization method.The paper also visualizes effective receptive fields for the different initialization methods.Source: §H.2 · p. 23

H.3 Influence of the d_state Parameter

The d_state and architecture ablations characterize trade-offs among state dimension, throughput, layer depth, channel mixing, token mixing, and activation functions. Results favor jointly selecting these hyperparameters rather than increasing d_state alone.

Influence of the d_state Parameter

  • Increasing d_state from 1 to 4 slightly improves performance but substantially decreases throughput.Source: §H.3 · p. 23
  • Increasing d_state to 8 while reducing ssm-ratio improves accuracy, whereas d_state 16 with ssm-ratio 1 causes performance to decline.Source: §H.3 · p. 24
  • The d_state results indicate that modest increases do not necessarily improve performance and that d_state and ssm-ratio must be selected jointly for speed-performance trade-offs.Source: §H.3 · p. 24

Architecture trade-offs

  • Reducing ssm-ratio significantly lowers performance but improves inference speed, while increasing layer numbers improves performance and slows the model.Source: §H.4 · p. 24
  • The ssm-ratio and layer-number trade-off balances channel-mixing and token-mixing, while mlp-ratio and ssm-ratio combinations are also important under fixed FLOPs.Source: §H.4 · p. 24

Activation functions

  • VMamba remains robust across different activation functions, allowing activation choice to reflect computational constraints or other preferences.Source: §H.5 · p. 24

Related papers

Loading 2401.10166v4…