[論文レビュー] LocalMamba: Visual State Space Model with Windowed Selective Scan

Tao Huang, Xiaohuan Pei+4 more|arXiv (Cornell University)|Mar 14, 2024

Data Visualization and Analytics被引用数 31

ひとこと要約

LocalMamba は Vision Mamba に窓付き局所スキャンと層ごとのスキャン方向探索を導入し、効率的な計算で ImageNet、COCO、ADE20K で優れた結果を達成します。

ABSTRACT

Recent advancements in state space models, notably Mamba, have demonstrated significant progress in modeling long sequences for tasks like language understanding. Yet, their application in vision tasks has not markedly surpassed the performance of traditional Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). This paper posits that the key to enhancing Vision Mamba (ViM) lies in optimizing scan directions for sequence modeling. Traditional ViM approaches, which flatten spatial tokens, overlook the preservation of local 2D dependencies, thereby elongating the distance between adjacent tokens. We introduce a novel local scanning strategy that divides images into distinct windows, effectively capturing local dependencies while maintaining a global perspective. Additionally, acknowledging the varying preferences for scan patterns across different network layers, we propose a dynamic method to independently search for the optimal scan choices for each layer, substantially improving performance. Extensive experiments across both plain and hierarchical models underscore our approach's superiority in effectively capturing image representations. For example, our model significantly outperforms Vim-Ti by 3.1% on ImageNet with the same 1.5G FLOPs. Code is available at: https://github.com/hunto/LocalMamba.

研究の動機と目的

• 視覚の状態空間モデルにおける局所依存関係を保持するための局所的な 2D 対応スキャンの利用を動機づける。
• 窓付き局所スキャン機構を提案し、グローバル文脈を維持しつつ画像内の局所構造をより良く捉える。
• ネットワーク深さにわたるスキャンパターンを最適化するための学習可能な層別スキャン方向探索を導入する。
• Plain および階層型モデル variants（LocalVim および LocalVMamba）を開発し、スケーラビリティと有効性を検証する。
• 分類、検出、セマンティックセグメンテーションの各タスクで Vim、VMamba、CNN、ViT を上回る改善を実証する。

提案手法

• 入力特徴を並列の局所窓内で処理し Spatial-Channel Attention (SCAttn) によって集約する四枝局所スキャンブロックを備えた LocalMamba を導入する。
• 水平・垂直および 2x2/7x7 の局所窓を各々標準形と反転形で用い、層ごとに四方向を選択する differentiable search（à la DARTS）を拡張する。
• トレーニング時に複数の SSM を層ごとに連続緩和して組み合わせ、推論時にトップ4の方向を選択する。
• Vim/VMamba ブロックを LocalMamba ブロックに置換した LocalVim（plain）と LocalVMamba（hierarchical）の2つのアーキテクチャ変種を提供する。
• ImageNet 分類、COCO オブジェクト検出/セグメンテーション、ADE20K セマンティックセグメンテーションの結果を報告し、利得を示す。

Figure 1 : Illustration of scan methods. (a) and (b): Previous methods Vim [ 60 ] and VMamba [ 32 ] traverse the entire row or column axis, resulting in significant distances for capturing dependencies between neighboring pixels within the same semantic region ( e.g. , the left eye in the image). (c

実験結果

リサーチクエスチョン

• RQ1 窓付きの局所スキャン戦略は、グローバル文脈を損なうことなく Vision Mamba モデルにおける局所 2D 依存関係の保持を改善できるか。
• RQ2 最適なスキャン方向の層別探索は、固定または単一方向スキャンに比べて有意な利得をもたらすか。
• RQ3 Plain および階層型 LocalMamba の変種は、Vim、VMamba、CNN、ViT と比較して分類、検出、セグメンテーションの各タスクでどの程度の性能を示すか。

主な発見

手法	画像サイズ	Params (M)	FLOPs (G)	Top-1 ACC (%)
LocalVim-T	224^2	8	1.5	76.2
LocalVim-S	224^2	28	4.8	81.2
VMamba-T	224^2	22	5.6	82.2
VMamba-S	224^2	44	11.2	83.5
LocalVMamba-T	224^2	26	5.7	82.7
LocalVMamba-S	224^2	50	11.4	83.7

• ImageNet-1K では LocalVim-T が Top-1 精度 76.2%、FLOPs 1.5G で DeiT-Ti (72.2%) を上回る。
• Hierarchical LocalVMamba-T は 82.7% の精度に達し、Swin-T を 1.4% 上回る。
• LocalVim-S と LocalVMamba-S は分類において Vim および VMamba のベースラインを強く上回る。
• COCO オブジェクト検出では LocalVMamba-T が 46.7 APb と 42.2 APm を達成し、Swin-T を上回る。
• ADE20K セグメンテーションでは LocalVim-S が 46.4 mIoU (SS)、LocalVMamba-S が 50.0 mIoU (SS) / 51.0 mIoU (MS) を達成する。
• アブレーションにより局所スキャンが Vim-T を 1.0% 向上させ、SCAttn が ImageNet で約 0.6% の利得を追加する。

Figure 2 : By extending the original scan with our local scan mechanism, our method significantly improves the ImageNet accuracies of Vim [ 60 ] while keeping similar FLOPs.