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.

  • 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.

  • Results

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

  • Takeaways & Limitations

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

  • Takeaways & Limitations

    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.

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.
  • VMamba adapts Mamba’s linear-complexity selective state-space modeling to visual representation learning.
  • Each image patch gathers contextual information through a compressed hidden state along its scanning path, reducing computational complexity from quadratic to linear.
  • 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.
  • SS2D bridges 1D array scanning and 2D plane traversal through four scanning routes, enabling selective SSMs to process vision data.
  • VMamba reports promising results across image classification, object detection, and semantic segmentation, with computational complexity growing linearly with input length.

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.
  • Vision Transformers (ViTs): ViT research established Transformer-based vision models and later emphasized large-scale pre-training, distillation, and inductive bias.
  • Efficient Attention: Linear-attention variants reformulate attention to reduce computational complexity from quadratic to linear while addressing memory movement and parallelizability.
  • 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.

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.
  • Discretization of SSM: The continuous hidden-state dynamics are discretized over intervals so SSMs can be integrated into deep models.
  • Discretization of SSM: The discretized formulation approximates the zero-order hold method commonly used in SSM-based models.
  • 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.
  • Selective Scan Mechanism: Selective SSMs use input-dependent parameters to capture contextual information, but dynamic weights prevent direct use of convolution for efficient computation.

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.
  • Network Architecture: VSS blocks replace Mamba’s S6 module with SS2D and remove the multiplicative branch because SS2D selectivity already supplies gating effects.
  • 2D-Selective-Scan for Vision Data (SS2D): SS2D adapts selective scanning to non-sequential vision data while preserving input-dependent weights and selective processing.
  • Accelerating VMamba: Vanilla VMamba-T reaches 82.2% classification accuracy, but its 426 images/s throughput and 22.9M parameters expose deployment challenges.
  • Accelerating VMamba: Implementation and architectural refinements target inference speed, including Triton Cross-Scan/Cross-Merge, selective-scan changes, linear transformations, and tensor-layout optimization.

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.
  • 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.
  • Image Classification: 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%.
  • 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.
  • Downstream Tasks: 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.
  • 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.
  • 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.

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.
  • VMamba improves inference speed through a series of architectural and implementation refinements.
  • Linear time complexity makes VMamba advantageous for downstream tasks with large-resolution inputs.
  • Limitations: Existing pre-training methods and VMamba-specific pre-training techniques for SSM-based architectures remain unexplored.
  • Limitations: Limited computational resources prevented exploration of VMamba-Large and fine-grained hyperparameter search.
  • Limitations: The potential of VMamba for more generalized tasks and generalized scanning patterns remains unexplored.

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.

  • The recurrence relation of selective SSM hidden states is derived from the hidden-state expression in Eq. 6.
  • Associative scan algorithms compute the recurrence in parallel, reducing overall computational complexity to linear.
  • The sequence variables V, K, Q, w, H, and Y organize inputs, state-transition terms, hidden states, and outputs across T steps.
  • The parameter A is simplified to R1×Dk, yielding element-wise state evolution and w_i ∈ R^Dk×Dv.
  • The overall response Y is expressed using Q, K, V, and a temporal mask, closely resembling Vision Transformer self-attention.
  • When w has shape (T, Dk), the formulation reduces to Gated Linear Attention, making GLA a special case of Mamba.

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.

  • Cross-Scan uses four scanning paths, with attention maps computed from QK^⊤ and (Q ⊙w)(K/w)^⊤.
  • Each row in a single attention map relates the current patch to all previously scanned foreground tokens impartially.
  • Rearranged diagonal elements produce image-space maps for QK^⊤ and (Q ⊙w)(K/w)^⊤ that distinguish foreground from background pixels.
  • Combining activation maps from all four scanning paths yields a complete visualization for a selected query patch.
  • The weighted attention maps effectively highlight foreground objects and separate them from the background.

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: Vanilla-VMamba uses VSS blocks with a multiplicative branch and no FFN, whereas VMamba VSS blocks remove that branch and add FFN layers.
  • Experiment Setting: VMamba-T/S/B are trained from scratch for 300 epochs with 20-epoch warm-up, batch size 1024, AdamW, and an initial learning rate of 1 × 10^-3.
  • Experiment Setting: Label smoothing, EMA, and scale-specific drop_path ratios are applied, with no additional training techniques employed.
  • Throughput Evaluation: Throughput is evaluated on an A100 GPU with an AMD EPYC 7542 CPU using batch size 128, following the stated Swin protocol.
  • Accelerating VMamba: The acceleration study reports progressive intermediate variants from Vanilla-VMamba-T to VMamba-T.
  • Evolution of ERF: Vanilla-VMamba-T’s effective receptive field evolves from predominantly local to predominantly global across training epochs.

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.

  • VMamba’s downstream evaluation covers object detection and instance segmentation on MSCOCO, plus semantic segmentation on ADE20K.
  • Architectural tables define the Vanilla-VMamba series and identify patch merging, linear layers, DWConv, and SS2D as core components.
  • Multi-scale training and random flipping are used for object detection evaluations with batch size 16.
  • Semantic segmentation uses a UperHead network in MMSegmentation, AdamW, learning rate 6 × 10−5, 160k iterations, and batch size 16.The default input resolution is 512 × 512.

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.

  • VMamba evaluates progressively larger input resolutions from 288 × 288 to 768 × 768 using models trained on 224 × 224 inputs.Performance is measured by ImageNet-1K top-1 accuracy, including a linear-tuning setting for the header network.
  • The VMamba series includes Tiny, Small, and Base architectural scales, with separate tables documenting architecture and acceleration details.
  • At 768 × 768, VMamba reaches 74.7% top-1 accuracy without fine-tuning, 79.2% with linear tuning, and 149 images per second throughput.
  • At the same resolution, Swin reaches 73.1% without fine-tuning and 77.5% with linear tuning, but its throughput falls to 53 images per second.
  • ConvNeXt maintains 103 images per second at 768 × 768, while its directly tested top-1 accuracy drops to 69.5%.
  • VMamba’s FLOPs increase linearly with input size, comparable to CNN-based architectures.

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.
  • Scanning patterns: 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.
  • Scanning patterns: DWConv contributes to establishing 2D priors and the formation of global effective receptive fields.
  • Initialization: The initialization ablation compares the S4D-derived default with random and zero initialization, using simplified skip-connection settings for the alternatives.
  • Downstream evaluation: The study evaluates VMamba and benchmark models on COCO object detection and instance segmentation, and ADE20K semantic segmentation.
  • 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.

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.
  • Influence of the d_state Parameter: Increasing d_state to 8 while reducing ssm-ratio improves accuracy, whereas d_state 16 with ssm-ratio 1 causes performance to decline.
  • Influence of the d_state Parameter: 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.
  • Architecture trade-offs: Reducing ssm-ratio significantly lowers performance but improves inference speed, while increasing layer numbers improves performance and slows the model.
  • Architecture trade-offs: 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.
  • Activation functions: VMamba remains robust across different activation functions, allowing activation choice to reflect computational constraints or other preferences.
Loading 2401.10166v4…