Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

Lianghui Zhu,Bencheng Liao,Qian Zhang,Xinlong Wang,Wenyu Liu,Xinggang Wang

2024-02-10

Abstract: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

Computer Vision and Pattern Recognition,Machine Learning

What problem does this paper attempt to address?

The paper mainly aims to address the following issues:

### Research Background and Objectives

- **Challenges of Visual Representation Learning**: When dealing with high-resolution images, existing Transformer models based on self-attention mechanisms (such as ViT) face challenges in terms of speed and memory usage.
- **Advantages of State Space Models (SSM)**: SSMs have shown great potential in long-sequence modeling, particularly with the successful application of the Mamba model in language modeling.

### Problems Addressed

1. **Visual Representation Learning Without Self-Attention**: A new visual backbone network, Vision Mamba (Vim), is proposed. It does not rely on self-attention mechanisms to learn visual representations and can effectively handle high-resolution images.

2. **Efficient Handling of High-Resolution Images**: By combining positional embeddings and bidirectional state space models (SSM), Vim can overcome the computational and memory limitations when processing high-resolution images, achieving faster speed and lower memory consumption.

3. **Application as a General Visual Backbone**: Vim is shown to outperform existing models (such as DeiT) not only in image classification tasks but also in downstream dense prediction tasks (such as semantic segmentation, object detection, and instance segmentation).

### Main Contributions

- **Proposing Vision Mamba (Vim)**: It combines bidirectional SSM for data-dependent global visual context modeling and positional embeddings to enhance position-aware visual understanding.
- **Efficiency**: Vim has sub-quadratic time complexity and linear memory complexity. Specifically, compared to DeiT, Vim is 2.8 times faster and saves 86.8% of GPU memory when performing batch inference on 1248×1248 resolution images.
- **Performance Advantage**: Vim demonstrates superior performance over highly optimized pure visual Transformer models (such as DeiT) in both ImageNet classification tasks and dense prediction downstream tasks.

In summary, the paper aims to develop a new visual backbone network, Vim, which not only addresses efficiency issues in high-resolution image processing but also performs excellently in various visual tasks, promising to become a core component of the next-generation visual foundational models.