[論文レビュー] Multi-Scale VMamba: Hierarchy in Hierarchy Visual State Space Model

Yuheng Shi, Minjing Dong+1 more|arXiv (Cornell University)|May 23, 2024

Image Retrieval and Classification Techniques被引用数 22

ひとこと要約

MSVMamba（Multi-Scale VMamba）を導入し、Multi-Scale 2D scanning（MS2D）、階層的 MS3 ブロック、ConvFFN を組み合わせて、限られたパラメータで長距離依存学習を改善します。ImageNet、COCO、ADE20K におけるSSM ベースのバックボーンの最先端結果を達成します。

ABSTRACT

Despite the significant achievements of Vision Transformers (ViTs) in various vision tasks, they are constrained by the quadratic complexity. Recently, State Space Models (SSMs) have garnered widespread attention due to their global receptive field and linear complexity with respect to the input length, demonstrating substantial potential across fields including natural language processing and computer vision. To improve the performance of SSMs in vision tasks, a multi-scan strategy is widely adopted, which leads to significant redundancy of SSMs. For a better trade-off between efficiency and performance, we analyze the underlying reasons behind the success of the multi-scan strategy, where long-range dependency plays an important role. Based on the analysis, we introduce Multi-Scale Vision Mamba (MSVMamba) to preserve the superiority of SSMs in vision tasks with limited parameters. It employs a multi-scale 2D scanning technique on both original and downsampled feature maps, which not only benefits long-range dependency learning but also reduces computational costs. Additionally, we integrate a Convolutional Feed-Forward Network (ConvFFN) to address the lack of channel mixing. Our experiments demonstrate that MSVMamba is highly competitive, with the MSVMamba-Tiny model achieving 82.8% top-1 accuracy on ImageNet, 46.9% box mAP, and 42.2% instance mAP with the Mask R-CNN framework, 1x training schedule on COCO, and 47.6% mIoU with single-scale testing on ADE20K.Code is available at \url{https://github.com/YuHengsss/MSVMamba}.

研究の動機と目的

• パラメータ制約下の視覚モデルでの長距離忘却問題を SSM で解決する。
• 冗長性を減らし、細粒度情報を維持する階層的でマルチスケールなスキャン戦略を開発する。
• SSM ベースのバックボーン内でチャネル混合と局所特徴抽出を強化するために ConvFFN を統合する。

提案手法

• VMamba の SS2D を、MS2D スキャンと ConvFFN チャンネルミキサーを含む Multi-Scale State Space（MS3）ブロックに置き換える。
• MS2D を、ストライド 1 および s の深さ方向畳み込みによってマルチスケールの特徴マップを作成し、フル解像度とダウンサンプリングマップを S6 ブロックを通して処理し、結果を集約する形で開発する。
• MS2D 後に Squeeze-Excitation（SE）ブロックを組み込み、ConvFFN（深さ方向畳み込み + 2 層全結合）を用いてチャネル間情報交換を強化する。
• Embeddding 次元とブロック数を制御して Comparable FLOPs を確保し、LeViT 的予算と公正な比較を可能にする。
• Nano、Micro、Tiny のモデル variants を提供し、それぞれ 6.9M–33.0M パラメータと 0.9–4.6 GFLOPs のスケーラブルなデプロイを実現する。

Figure 1 : Comparison on ImageNet.

実験結果

リサーチクエスチョン

• RQ1 SSM ベースの視覚バックボーンにおいて、冗長性を減らし長距離依存学習を改善するための多尺度 2D スキャンをどのように設計できるか。
• RQ2 ConvFFN と SE ブロックを統合することが、固定計算予算の下でチャネル間情報交換と全体精度に与える影響はどの程度か。
• RQ3 階層的でマルチスケールな VMamba デザインは、ImageNet、COCO、ADE20K で既存の VMamba バリアントや他のSOTAバックボーンを上回りつつ効率性を維持できるか。

主な発見

• MSVMamba-T は ImageNet-1K で 82.8% Top-1 精度を達成（パラメータ 33M、4.6 GFLOPs）。
• MSVMamba-Nano、-Nano ベースラインは、ほぼ同等の FLOPs で VMamba-Nano を最大 5.5 ポイントの Top-1 精度向上。
• MSVMamba-T は VMamba-T より 0.6% の Top-1 精度を、はるかに低い計算コストで上回る。
• COCO 物体検出では、MSVMamba-T が Swin-T を 1x スケジュール下で +4.2 ボックス AP、+2.9 マスク AP 上回る。
• ADE20K セマンティック分割では、単一スケール検証で 47.6 mIoU、マルチスケール検証で 48.5 を達成。
• ConvFFN と MS2D および SE ブロックを組み合わせると精度向上が顕著（例：SE で Top-1 に +0.5%、 Ablations で ConvFFN で +2.0%）になる。

(a) Horizontal Scan