[논문 리뷰] VMamba: Visual State Space Model

Yue Liu, Yunjie Tian외 7명|arXiv (Cornell University)|2024. 01. 18.

Visual Attention and Saliency Detection인용 수 358

한 줄 요약

VMamba는 Cross-Scan Module을 통해 다방향 2D 정보 통합을 가능하게 하고, 2D 시각 상태 공간 모델을 도입하여 선형 복잡도와 전역 수용 영역을 제공합니다. 이를 통해 ImageNet, COCO, ADE20K에서 경쟁력 있는 성능을 달성합니다.

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.

연구 동기 및 목표

• 전역 수용 필드를 선형 복잡도와 결합한 비전 기반 모델의 필요성을 제고한다.
• 선택적 스캔 접근법으로 상태 공간 모델(SSM)을 활용하여 비전 작업에서 동적 가중치를 가능하게 한다.
• Cross-Scan Module(CSM)을 통해 2D 이미지에 SSM을 적용할 때의 방향 민감도를 해결한다.
• 이미지 분류, 객체 탐지, 의미 분할을 위한 확장 가능한 백본으로 VMamba를 개발한다.

제안 방법

• Visual State Space(VSS) 블록 내에서 2D 선택적 스캔(SS2D)을 채택한다.
• 전역 수용 필드를 선형 복잡도로 보장하기 위해 네 모서리에서 2D 특징 맵을 스캔하는 Cross-Scan Module(CSM)을 도입한다.
• 실용적인 딥러닝 활용을 위해 연속 상태 공간 역학을 이산화한다(행렬 지수 및 1차 근사).
• ViT/CNN 백본과 유사한 단계적 다운샘플링을 갖춘 계층적 VMamba 백본으로 VSS 블록을 쌓는다.
• ImageNet-1K에서 VMamba 변형(Tiny, Small, Base)을 학습하고 COCO 및 ADE20K에서 탐지 및 분할에 대해 평가한다.

실험 결과

연구 질문

• RQ1 상태 공간 기반 시각 백본이 선형 계산 복잡도로 전역 수용 영역을 달성할 수 있는가?
• RQ2 Cross-Scan Module이 효율성을 해치지 않으면서 효과적인 2D 정보 통합을 가능하게 하는가?
• RQ3 VMamba가 ImageNet-1K, COCO, ADE20K에서 크기에 따라 CNNs 및 ViTs와 비교하여 어떤 성능을 보이는가?
• RQ4 VMamba는 입력 해상도 스케일링에 대해 로버스트하며 학습 후 적응적 전역 수용 필드를 나타내는가?

주요 결과

• VMamba 변형은 4.5G FLOPs에서 VMamba-T가 82.2%와 같은 유사한 FLOPs를 가진 인기 백본 대비 경쟁력 있는 ImageNet-1K top-1 정확도를 달성한다.
• COCO에서 1x 및 3x MS 스케줄로 VMamba-T/S/B는 여러 설정에서 상자 및 마스크 AP에서 Swin 및 ConvNeXt 백본 대비 우수하다.
• ADE20K에서 512x512 입력으로 VMamba-S와 VMamba-B가 더 높은 mIoU를 달성하며, VMamba-T는 512x512 및 MS 테스트에서 강력한 분할 성능을 달성한다.
• Cross-Scan Module은 글로벌 효과적 수용 필드(ERF)를 생성하고 교차 모양의 장거리 의존성을 가능하게 하며, 학습 후 ERF가 글로벌이 된다.
• VMamba는 입력 해상도 증가에 따라 선형 FLOPs 증가를 보이며 전역 수용 필드를 유지하는 반면 ViT는 이차 복잡도를 가진다.