Source-linked AI summary
Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, Xinggang Wang
TL;DR
Generic pure-SSM vision backbones had not been explored, while visual modeling requires positional awareness and global context. Vim addresses these challenges with bidirectional SSMs and position embeddings, achieving strong vision performance with lower computation and memory use than ViTs.
Problem
Generic pure-SSM-based backbones had not been explored for visual data, while vision requires positional awareness and global context modeling.
Method
Vim uses bidirectional state space modeling for data-dependent global visual context and position embeddings for location-aware visual understanding.
Results
Vim achieves superior performance to DeiT on ImageNet and improves efficiency for high-resolution image processing, including 2.8× faster inference and 86.8% lower GPU memory use.
Takeaways & Limitations
Vim provides a pure-SSM, sequence-modeling vision backbone with Transformer-level modeling power and lower computation complexity, without image-specific inductive biases.
Takeaways & Limitations
The paper’s impact statement limits its stated consequence to improved practicality for high-resolution visual applications resulting from increased efficiency.
Abstract
from arXiv · showhide
Recently the state space models (SSMs) with efficient hardware-aware designs, i.e., the Mamba deep learning model, have shown great potential for long sequence modeling. Meanwhile building efficient and generic vision backbones purely upon SSMs is an appealing direction. However, representing visual data is challenging for SSMs due to the position-sensitivity of visual data and the requirement of global context for visual understanding. In this paper, we show that the reliance on self-attention for visual representation learning is not necessary and propose a new generic vision backbone with bidirectional Mamba blocks (Vim), which marks the image sequences with position embeddings and compresses the visual representation with bidirectional state space models. On ImageNet classification, COCO object detection, and ADE20k semantic segmentation tasks, Vim achieves higher performance compared to well-established vision transformers like DeiT, while also demonstrating significantly improved computation & memory efficiency. For example, Vim is 2.8$\times$ faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on images with a resolution of 1248$\times$1248. The results demonstrate that Vim is capable of overcoming the computation & memory constraints on performing Transformer-style understanding for high-resolution images and it has great potential to be the next-generation backbone for vision foundation models. Code is available at https://github.com/hustvl/Vim.
1. Introduction
Vision Mamba (Vim) is proposed as a generic pure-SSM vision backbone that addresses Mamba’s unidirectional modeling and lack of positional awareness with bidirectional SSMs and position embeddings. Across vision tasks, Vim is reported to outperform DeiT while improving high-resolution inference efficiency.
- Motivation: Vim is designed to overcome Mamba’s unidirectional modeling and lack of positional awareness for visual data.These challenges arise because visual recognition requires spatial information and global context.
- Implications: Mamba-based visual pretraining can reduce the computational cost of large-scale pretraining.The passage attributes this potential to Mamba’s efficiency.
- Proposed approach: Vim introduces bidirectional SSMs for data-dependent global visual context and position embeddings for location-aware visual understanding.Images are split into patches, projected into vectors, and processed as sequence data in Vim blocks.
- Efficiency: 2.8× faster than DeiT and saves 86.8% GPU memory at 1248×1248 image resolution during batch feature extraction.The claimed efficiency comes with subquadratic-time computation and linear memory complexity without attention.
- Experiments: Vim achieves superior performance to DeiT on ImageNet classification and dense prediction downstream tasks.The paper positions Vim as a pure-SSM sequence model that can serve as a generic vision backbone.
2. Related Work
Related work establishes the progression from convolutional and Transformer vision backbones to state space models for long-sequence and visual applications. It also highlights the trade-off between Transformer modeling capacity and the quadratic complexity imposed by visual token sequences.
- Vision backbones: ViT treats images as patch sequences with a pure Transformer architecture, motivating extensive follow-up work on generic vision backbones.Subsequent approaches include hybrid designs that introduce 2D convolutional priors.
- Vision Transformers: Vanilla Transformer-style models remain attractive because of larger modeling capacity, unified multimodality representation, and compatibility with self-supervised learning.These advantages helped return vanilla Transformer-style models to prominence despite efficiency concerns.
- Vision Transformers: Visual token counts are limited by the quadratic complexity of Transformer processing.The passage connects this constraint to the use of high-resolution visual inputs.
- State space models: S4 and related state space models provide linearly scaling sequence processing for long-range dependency modeling.Follow-up work explores bidirectional, multi-input, gated, and selective SSM designs.
- Visual SSMs: Prior visual SSM research applies S4 to video and multidimensional data, combines S4 with self-attention, or develops selective SSM variants.These studies cover video classification, images, videos, and movie scene detection.
3. Method
Vim adapts Mamba-style state space modeling to vision by combining patch-token representations, positional embeddings, and bidirectional processing. Its hardware-aware implementation targets lower memory traffic and linear sequence-length scaling than self-attention.
- Vision Mamba architecture: Vim transforms images into projected patch tokens, adds positional information, and processes the sequence through stacked Vim blocks before classification.The architecture uses a learnable class token, residual Vim blocks, normalization, and an MLP prediction head.
- Vim block: Bidirectional Vim blocks process normalized token features in forward and backward directions, using convolution, parameter projection, SSM computation, gating, and directional aggregation.The two directional outputs are gated by z and added to produce the output token sequence.
- Efficiency analysis: Vim’s SSM implementation reduces memory I/O from O(BMEN) to O(BME + EN) by loading parameters into faster SRAM before performing the SSM operations.The method writes the output back to HBM after computation in SRAM.
- Efficiency analysis: Vim recomputes intermediate states and activations during backpropagation to reduce GPU memory requirements for long sequences.Recomputed values include states, activation outputs, and convolution outputs.
- Efficiency analysis: SSM computation scales linearly with sequence length M, whereas global self-attention scales quadratically, with SSM dimension N fixed at 16 by default.Both mechanisms provide adaptive global context, but their sequence-length scaling differs.
4. Experiment
Experiments evaluate Vim on ImageNet classification, ADE20K segmentation, COCO detection and instance segmentation, alongside efficiency and bidirectional-design ablations. Vim generally outperforms DeiT while improving high-resolution inference efficiency.
- Efficiency: 2.8× faster and 86.8% less GPU memory at 1248×1248 than DeiT, while performance is similar at 512×512.The comparison uses batch inference with the backbone and FPN; Vim’s efficiency advantage grows with resolution.
- Semantic segmentation: 1.8 mIoU higher for Vim-Ti over DeiT-Ti and 0.9 mIoU higher for Vim-S over DeiT-S on ADE20K segmentation.Vim-S also matches ResNet-101 segmentation performance with nearly 2× fewer parameters.
- Object detection and instance segmentation: 1.3 box AP and 1.1 mask AP higher for Vim-Ti than DeiT-Ti with Cascade Mask R-CNN.For middle-size and large-size objects, Vim-Ti leads DeiT-Ti by 1.6 APbox.
- Ablation study: Bidirectional SSM + Conv1d reaches 73.9 top-1 accuracy and 35.9 mIoU, versus 73.2 top-1 accuracy and 32.3 mIoU for the compared configuration.The paper adopts Bidirectional SSM + Conv1d as Vim’s default block setting.
- Ablation study: 76.1 top-1 accuracy is achieved by the middle class token strategy in the classification-design ablation.The paper associates this strategy with the recurrent nature of SSM and a central object prior in ImageNet.
5. Conclusion and Future Work
Vim is presented as a generic vision backbone built from bidirectional state space modeling and position embeddings, with future applications to unsupervised, multimodal, and high-resolution visual tasks.
- Conclusion: Vim uses bidirectional state space modeling to obtain data-dependent global visual context without image-specific inductive biases.The model learns visual representations through sequence modeling rather than hybrid architectures or equivalent global 2D convolutional kernels.
- Conclusion: Position embeddings make Vim suitable for location-aware visual understanding alongside its bidirectional SSM design.
- Future Work: Vim could support mask image modeling, CLIP-style multimodal pretraining, and downstream analysis of high-resolution medical, remote-sensing, and video data.These applications are identified as future directions rather than evaluated results in the supplied passage.
Impact Statement
The work advances the efficiency of generic vision backbones, increasing their practicality for visual applications involving high-resolution input images.
- Impact: Improved vision-backbone efficiency increases practicality for applications that process high-resolution input images.
A. Visualization
The visualization compares DeiT-Ti and Vim-Ti within Cascade Mask R-CNN, highlighting Vim-Ti's ability to capture a very large object that DeiT-Ti fails to perceive.
- Visualization: Vim-Ti captures the very large object in the image, whereas the DeiT-Ti counterpart fails to perceive it.The comparison is attributed to SSM-based long-range context learning.
B. Additional Setting
The additional settings describe ADE20K semantic-segmentation experiments using UperNet and COCO detection and instance-segmentation experiments using Cascade Mask R-CNN.
- Settings for Semantic Segmentation: ADE20K semantic-segmentation experiments use 150 categories and UperNet as the base framework.The dataset includes 20K training, 2K validation, and 3K testing images.
- Settings for Object Detection and Instance Segmentation: COCO 2017 detection and instance-segmentation experiments use Cascade Mask R-CNN as the canonical base framework.The dataset contains 118K training, 5K validation, and 20K testing images.
C. Extended Comparison on Hierarchical Architecture
Hier-Vim replaces Swin Transformer’s shifted local window attention with a global bidirectional SSM and delivers competitive ImageNet performance across model sizes.
- Method: Hier-Vim replaces shifted local window attention in Swin Transformer with the proposed global bidirectional SSM.The variant is designed for comparison with hierarchical architectures.
- Results: 1.3% higher performance than Swin Transformer at the tiny size demonstrates Hier-Vim’s advantage on ImageNet.
- Results: 0.2% higher performance than Swin Transformer at the small size shows competitive performance among hierarchical architectures.
- Results: 0.4% higher performance than Swin Transformer at the base size further supports Hier-Vim’s competitiveness.