VMamba: Visual State Space Model

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.

한 줄 요약

VMamba는 2D 선택적 스캔과 크로스-스캔 모듈을 결합한 시각 상태 공간 모델로, 높은 성능과 선형 복잡도를 동시에 달성한다.

핵심 기여도

VMamba는 기존 CNN과 ViT를 넘어서는 새로운 시각 기초 모델로, 선형 복잡도를 유지하면서 글로벌 수용 필드를 확보한다.
Cross-Scan Module (CSM)을 도입하여 1D 스캔과 2D 데이터 구조 간의 격차를 해소하고, 2D 선택적 스캔(SS2D)을 가능하게 한다.
COCO에서 VMamba-Tiny/Small/Base는 각각 46.5%/48.2%/48.5%의 mAP를 달성하며, ADE20K에서 47.3%/49.5%/50.0%의 mIoU를 기록한다.
입력 해상도가 증가할수록 VMamba의 성능 우수성이 두드러지며, FLOPs는 선형 성장한다.

핵심 아이디어

VMamba는 기존 ViT의 글로벌 수용 필드와 동적 가중치의 장점을 유지하면서, 계산 복잡도를 줄이기 위해 상태 공간 모델(S6)을 시각 분야로 확장한 모델이다.
기존 1D 선택적 스캔은 순차적 데이터에 적합하지만, 비순차적 2D 이미지 데이터에 적용할 경우 수용 필드가 제한된다. 이를 해결하기 위해 VMamba는 2D 선택적 스캔(SS2D)과 CSM을 결합하여, 4방향 스캔을 통해 모든 위치의 정보를 통합한다.
CSM은 이미지 특징 맵을 단일 방향이 아닌 4개의 코너에서 대각선 방향으로 순회하며, 각 요소가 다양한 방향의 정보를 수집하도록 한다. 이는 글로벌 수용 필드를 확보하면서도 계산 복잡도를 선형으로 유지한다.

기술적 접근법

VMamba는 Visual State-Space (VSS) 블록으로 구성되며, 핵심 연산은 2D 선택적 스캔(SS2D)이다.
SS2D는 4가지 스캔 경로(상단-하단, 좌측-우측, 대각선 등)를 통해 2D 이미지를 순회하며, 각 위치의 정보를 통합한다.
Cross-Scan Module (CSM)은 비인과적 2D 이미지를 순차적 패치 시퀀스로 변환하여, 1D 스캔 기반의 S6를 2D 데이터에 확장한다.
VMamba는 다양한 모델 크기(VMamba-Tiny, Small, Base)로 제공되며, 파라미터 수는 각각 22M, 44M, 75M이다.
학습 및 추론 과정에서 FLOPs는 선형 성장하며, ViT에 비해 입력 해상도 증가에 따른 성능 저하가 적다.

주요 결과

VMamba는 ImageNet-1K에서 기존 CNN과 ViT 대비 경쟁력 있는 성능을 보인다.
COCO 데이터셋에서 VMamba-Tiny/Small/Base는 각각 46.5% (기존 베이스라인 대비 +2.1%), 48.2% (+2.4%), 48.5% (+2.7%)의 mAP를 달성한다.
ADE20K 데이터셋에서 VMamba-Tiny/Small/Base는 47.3% (+1.9%), 49.5% (+2.1%), 50.0% (+2.6%)의 mIoU를 기록한다.
입력 해상도가 224×224에서 384×384로 증가할 때, VMamba-S는 성능이 84%로 상승하며, FLOPs는 선형 성장한다.
VMamba는 ViT보다 FLOPs 증가 속도가 느르고, 높은 해상도에서도 안정적인 성능을 유지한다.

의의 및 한계

VMamba는 기존 CNN과 ViT의 장점을 결합한 새로운 기초 모델로서, 높은 성능과 계산 효율성을 동시에 달성하는 데 기여한다.
특히, 글로벌 수용 필드와 동적 가중치를 유지하면서도 선형 복잡도를 보장하는 점에서 학술적·실용적 가치가 높다.
하지만, VMamba-B 모델의 학습 중 버그가 발생하여 최종 성능이 아직 완전히 검증되지 않았다는 한계가 있다.
또한, CSM의 4방향 스캔 전략은 복잡한 구조를 요구하며, 일부 애플리케이션에서는 구현 난이도가 높을 수 있다.

실용적 활용

VMamba는 고해상도 이미지 처리가 필요한 산업 분야(예: 자율주행, 의료 영상 분석)에서 유용하게 활용될 수 있다.
또한, 계산 자원이 제한된 모바일 및 임베디드 시스템에서도 높은 성능과 낮은 FLOPs를 통해 실용성을 확보할 수 있다.
이 모델은 다양한 시각 인식 작업(분류, 객체 탐지, 세분화)에 적용 가능하며, 기존 CNN과 ViT를 대체하거나 보완하는 역할을 할 수 있다.