[Paper Review] VMamba: Visual State Space Model

Yue Liu, Yunjie Tian+7 more|arXiv (Cornell University)|Jan 18, 2024

Visual Attention and Saliency Detection358 citations

TL;DR

VMamba introduces a 2D visual state space model with linear complexity and global receptive fields, using a Cross-Scan Module to enable multi-directional 2D information integration, achieving competitive ImageNet, COCO, and ADE20K performance.

ABSTRACT

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.

Motivation & Objective

• Motivate a vision foundation model that combines global receptive fields with linear complexity.
• Leverage state space models (SSMs) with a selective scan approach to enable dynamic weights in vision tasks.
• Address direction-sensitivity in applying SSMs to 2D images via a Cross-Scan Module (CSM).
• Develop VMamba as a scalable backbone for image classification, object detection, and semantic segmentation.

Proposed method

• Adopt a 2D Selective Scan (SS2D) within a Visual State Space (VSS) block.
• Introduce Cross-Scan Module (CSM) that scans 2D feature maps from four corners to ensure global receptive fields with linear complexity.
• Discretize continuous-state space dynamics for practical deep learning use (matrix exponentials and first-order approximations).
• Stack VSS blocks into a hierarchical VMamba backbone with staged down-sampling similar to ViT/CNN backbones.
• Train VMamba variants (Tiny, Small, Base) on ImageNet-1K and evaluate on COCO and ADE20K for detection and segmentation.

Experimental results

Research questions

• RQ1 Can a state-space–based visual backbone achieve global receptive fields with linear computational complexity?
• RQ2 Does the Cross-Scan Module enable effective 2D information integration without sacrificing efficiency?
• RQ3 How does VMamba compare to CNNs and ViTs on ImageNet-1K, COCO, and ADE20K across scales?
• RQ4 Does VMamba exhibit robustness to input resolution scaling and demonstrate adaptive global receptive fields after training?

Key findings

• VMamba variants achieve competitive ImageNet-1K top-1 accuracy with similar FLOPs to popular baselines (e.g., VMamba-T 82.2% at 4.5G FLOPs).
• On COCO with 1x and 3x MS schedules, VMamba-T/S/B surpass Swin and ConvNeXt baselines in box and mask AP across several settings.
• VMamba-S and VMamba-B achieve higher mIoU on ADE20K with 512x512 inputs, and VMamba-T attains strong segmentation performance with 512x512 and MS testing.
• The Cross-Scan Module yields a global effective receptive field (ERF) and enables cross-shaped, long-range dependencies, with ERF becoming global after training.
• VMamba demonstrates linear FLOPs growth with increasing input resolution, while maintaining global receptive fields, unlike ViTs which incur quadratic complexity.