[論文レビュー] VMamba: Visual State Space Model

Yue Liu, Yunjie Tian+7 more|arXiv (Cornell University)|Jan 18, 2024

Visual Attention and Saliency Detection被引用数 358

ひとこと要約

VMambaは、線形計算量とグローバル受容野を備えた2D視覚状態空間モデルを導入し、Cross-Scan Moduleを用いて多方向の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 Selective Scan (SS2D)を採用する。
• グローバル受容野を線形計算量で確保するため、4隅から2D特徴マップをスキャンするCross-Scan Module (CSM)を導入する。
• 実用的な深層学習利用のために連続状態空間ダイナミクスを離散化する（行列指数と一階近似）。
• ViT/CNNバックボーンに類似した段階的ダウンサンプリングを持つ階層的VMambaバックボーンとしてVSSブロックを積み上げる。
• ImageNet-1KでVMamba variants（Tiny、Small、Base）を訓練し、検出とセマンティック分割の評価のためにCOCOとADE20Kで評価する。

実験結果

リサーチクエスチョン

• RQ1 状態空間ベースの視覚バックボーンは、線形計算量でグローバル受容野を実現できるか？
• RQ2 Cross-Scan Moduleは、効率を犠牲にすることなく効果的な2D情報統合を実現できるか？
• RQ3 ImageNet-1K、COCO、ADE20Kで、スケールを跨いでVMambaはCNNsおよびViTsとどう比較されるか？
• RQ4 VMambaは入力解像度のスケーリングに対するロバスト性を示し、訓練後に適応的なグローバル受容野を示すか？

主な発見

• VMamba系は、一般的なベースラインと同程度のFLOPsで競争力のあるImageNet-1K top-1精度を達成（例: VMamba-Tは4.5G FLOPsで82.2%）。
• COCOで1xおよび3x MSスケジュールの場合、VMamba-T/S/Bは複数の設定でボックスAPとマスクAPでSwinおよびConvNeXtベンチマークを上回る。
• 512x512入力でADE20KのmIoUがVMamba-SおよびVMamba-Bで高く、VMamba-Tは512x512およびMSテストで強力なセグメンテーション性能を発揮。
• Cross-Scan Moduleはグローバルな有効受容野（ERF）を生み出し、クロス形状の長距離依存を可能にし、訓練後にはERFがグローバルになる。
• VMambaは入力解像度の増加に伴いFLOPsが線形に増加しつつグローバルな受容野を維持する。一方ViTは二次的な計算複雑度を要するのに対し、VMambaは入力解像度の増加に伴いFLOPsが線形に増加しつつグローバルな受容野を維持する。