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

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.

한 줄 요약

Vision Mamba(Vim)는 양방향 SSM과 위치 임베딩을 결합한 시각적 표현 학습 모델로, DeiT 대비 2.8× 빠르고 86.8% GPU 메모리 절약.

핵심 기여도

Vim은 양방향 Mamba 블록과 위치 임베딩을 결합하여, 시각적 표현 학습에서 Transformer와 동등한 성능을 유지하면서도 계산 복잡도를 줄임.
ImageNet 분류에서 Vim은 DeiT 대비 2.8× 빠르고, 1248×1248 해상도 이미지 추론 시 86.8%의 GPU 메모리 절약.
Vim은 ADE20k 세분화 분할에서 35.9 mIoU 달성, 이는 양방향 SSM + Conv1d 설계 덕분.
Vim은 순수 SSM 기반으로 2D 인덕티브 바이어스 없이 시퀀스 방식으로 이미지 처리.

핵심 아이디어

기존 시각 모델은 Self-attention을 통해 글로벌 컨텍스트를 학습하지만, 이는 고해상도 이미지 처리 시 메모리와 계산 효율성에 한계가 있다. Vim은 Mamba 모델의 하드웨어 인식 설계와 SSM의 효율성을 활용하여, Self-attention 없이도 Transformer와 동등한 표현력을 유지하는 새로운 접근법을 제시한다. Vim은 양방향 SSM을 도입하여 데이터 종속적 글로벌 컨텍스트를 학습하고, 위치 임베딩을 통해 시공간 정보를 인식한다. 특히, 양방향 SSM + Conv1d 설계는 분류 및 세분화 분할 성능을 동시에 향상시킨다.

기술적 접근법

**입력 처리**: 16×16 커널로 이미지를 패치로 분할 후 1D 시퀀스로 변환.
**Vim 블록**: 양방향 Mamba 블록을 사용하여 시퀀스를 처리. 기본적으로 24개 블록을 쌓고, SSM 차원은 16, 히든 차원은 192, 확장 차원은 384로 설정.
**양방향 SSM**: 각 블록에 역방향 SSM과 Conv1d를 추가하여, 시퀀스의 양방향 정보를 활용.
**분류 설계**: 중간 클래스 토큰(Middle class token) 전략을 사용해 76.1%의 최고 정확도 달성.
**하드웨어 효율성**: Mamba의 하드웨어 인식 알고리즘을 활용해, 고해상도 이미지 처리 시 메모리와 계산 효율성 향상.

주요 결과

**ImageNet 분류**: Vim은 DeiT 대비 2.8× 빠르고, 1248×1248 해상도 이미지 추론 시 86.8%의 GPU 메모리 절약.
**ADE20k 세분화 분할**: Vim은 35.9 mIoU 달성, 이는 DeiT 기반 모델 대비 3.6 mIoU 개선.
**COCO 객체 탐지**: Vim은 기존 Vision Transformer 대비 높은 성능을 유지하면서도 메모리 사용량을 줄임.
**양방향 SSM + Conv1d**: 분류 정확도 73.9%, 세분화 분할 35.9 mIoU로, 기존 Mamba 블록 대비 1.6%의 정확도 향상.

의의 및 한계

Vim은 순수 SSM 기반으로, 2D 인덕티브 바이어스 없이 시퀀스 방식으로 이미지를 처리함으로써, 다양한 시각적 표현 학습에 유연하게 적용 가능하다. 특히, 고해상도 이미지 처리 시 Transformer 대비 뛰어난 메모리 및 계산 효율성을 보여주며, Vision Foundation Model의 차세대 백본으로서의 가능성을 제시한다. 그러나 현재는 ImageNet 기반의 감독 학습만 실험되었으며, 비감독 학습 및 멀티모달 태스크에 대한 평가가 부족한 점이 한계로 지적된다.

실용적 활용

Vim은 고해상도 이미지 처리가 필요한 의료 영상 분석, 원격 감시, 장기 동영상 분석 등에 적용 가능하다. 또한, 하드웨어 효율성 덕분에 모바일 및 임베디드 시스템에서도 활용 가능하며, 대규모 비감독 학습을 통해 멀티모달 CLIP 스타일 프리트레이닝에도 유용할 것으로 기대된다.