1. 概要
Vision mamba(Vim)は、状態空間モデル(State Space Model:SSM)であるMambaを視覚タスク向けに拡張したモデルです。Vision Transformerと同等以上の性能を達成しながら、高解像度画像の処理において大幅な速度向上とメモリ使用量の削減が可能になりました。ImageNet分類、物体検出、セマンティックセグメンテーションなど様々なタスクで優れた性能を示し、特に高解像度画像の処理において従来モデルを大きく上回る効率を達成しています。
2. 研究背景と課題
近年、Vision Transformer (ViT) は画像認識タスクで大きな成功を収めていますが、自己注意機構 (Self-Attention) の計算量が画像サイズの2乗に比例して増加するという課題がありました。これは特に高解像度画像の処理において深刻な問題となっています。
図1: DeiTとVimの性能と効率の比較。Vimは高解像度画像の処理において大幅な速度向上とメモリ使用量の削減を実現しています。
図1は、従来のTransformerベースのモデル(DeiT)とVimの性能および効率の比較を示しています。
具体的な課題は以下の通り：
-
計算量の増大：Self-Attentionの計算量が画像サイズの2乗に比例して増加するため、高解像度画像の処理が困難。
-
メモリ使用量の増大：大規模なAttentionマップをメモリに保持する必要があり、高解像度画像の処理時にGPUメモリが不足。
-
2次元構造の考慮：画像の2次元構造を効果的に扱うために、Window Attentionなどの追加の工夫が必要。
これらの課題に対し、Vimは状態空間モデル(SSM)を基にした新しいアプローチを提案し、高速かつメモリ効率の良い画像処理を実現します。
3. 提案手法の詳細説明
3.1 Vision Mambaの全体構造
Vimは、全体的な構造はViTとほぼ同じです。
Vimは、ViTの構造を保ちつつ、encoder部分にMambaを取り入れつつ、双方向処理を追加したものとなっています。
図2: Vision mambaの概要。入力画像をパッチに分割し、双方向のSSMで処理します。
Vimの処理フローは以下の通りです：
-
パッチ埋め込み：入力画像 $t \in \mathbb{R}^{H\times W\times C}$ を $J$ 個の非重複パッチ $x_p \in \mathbb{R}^{J\times (P^2 \cdot C)}$ に分割。
-
線形投影：パッチを $D$ 次元のベクトルに投影。
-
位置埋め込みの追加：トークン系列に位置情報を付加。
-
Vimエンコーダ：双方向SSMを用いてトークン系列を処理。
-
分類ヘッド：最終的な予測を行うMLPヘッド。
主要な処理は以下の数式で表されます：
$$
T_0 = [t_{cls}; t^1_p W; t^2_p W; \cdots; t^J_p W] + E_{pos}
$$
$$
T_l = \text{Vim}(T_{l-1}) + T_{l-1}
$$
$$
f = \text{Norm}(T^0_L)
$$
$$
\hat{p} = \text{MLP}(f)
$$
ここで、$T_l$ は $l$ 番目のVimブロックの出力、$L$ はブロックの総数、$f$ は最終的な特徴表現、$\hat{p}$ は予測結果です。
3.2 Vimブロックの詳細
Vimブロックは、従来のMambaブロックを拡張し、画像の2次元構造を効果的に扱えるよう設計されています。
以下の順に処理を行います：
- 入力の正規化
- 線形投影
- 双方向SSM処理
- 特徴の集約
特に重要なのは双方向SSM処理で、以下の式で表されます：
$$
y_o = \text{SSM}(\bar{A}_o, B_o, C_o)(x'_o)
$$
ここで、$o$ は順方向または逆方向を示し、$\bar{A}_o$, $B_o$, $C_o$ はSSMのパラメータです。
この双方向処理により、Vimは以下の能力を獲得します：
- 空間的文脈の把握：画像内の物体の配置や関係性をより良く理解。
- 長距離依存関係の捕捉：画像全体にわたる大域的な特徴を効率的に抽出。
- 異方性特徴の学習：方向に依存する特徴（エッジの方向など）をより正確に捉える。
3.3 効率性の分析
Vimの効率性は主に以下の3つの観点から説明できます：
-
IO効率：高帯域幅メモリ(HBM)とSRAMを効果的に活用し、メモリアクセスを最小限に抑える。
-
メモリ効率：中間状態の再計算を行うことで、必要なGPUメモリを削減。
-
計算効率：SSMの線形時間複雑性により、高解像度画像でも効率的な処理が可能。
具体的に、Self-AttentionとSSMの計算複雑性を比較すると：
$$
\Omega(\text{self-attention}) = 4MD^2 + 2M^2D
$$
$$
\Omega(\text{SSM}) = 3M(2D)N + M(2D)N
$$
ここで、$M$ はシーケンス長、$D$ は特徴次元、$N$ はSSMの次元（固定値）です。SSMは $M$ に対して線形なので、高解像度画像でも効率的に処理できます。
4. 評価実験と結果
Vimの性能評価は、画像分類、セマンティックセグメンテーション、物体検出の3つのタスクで行われました。
4.1 画像分類
ImageNet-1Kデータセットでの実験結果を以下に示します：
表1: ImageNet-1K検証セットでの異なるバックボーンとの比較
主な結果：
- Vim-Tiは76.1%のTop-1精度を達成し、DeiT-Ti (72.2%)を大きく上回る。
- 長シーケンスファインチューニング後、Vim-S†は81.6%に達し、DeiT-B (81.8%)と同等の性能を示す。
- パラメータ効率が高く、同等のパラメータ数で他モデルを上回る性能を達成。
4.2 セマンティックセグメンテーション
ADE20Kデータセットでの結果：
表2: ADE20K検証セットでのセマンティックセグメンテーション結果
主な結果：
- Vim-Sは44.9 mIoUを達成し、DeiT-S (44.0 mIoU)を上回る。
- ResNet-101と同等の性能を、約半分のパラメータ数で実現。
4.3 物体検出とインスタンスセグメンテーション
COCO 2017データセットでの結果：
表3: COCO検証セットでの物体検出とインスタンスセグメンテーション結果
主な結果：
- Vim-Tiは45.7 AP_boxと39.2 AP_maskを達成し、DeiT-Ti (44.4 AP_box, 38.1 AP_mask)を上回る。
- 特に中型・大型物体の検出で顕著な改善を示す。
4.4 効率性の評価
高解像度画像処理時の効率性：
図3: GPUメモリ効率の比較
図4: FPS（フレーム毎秒）の比較
主な結果：
- 1248×1248の高解像度画像処理時、VimはDeiTと比較して：
- 2.8倍高速
- 86.8%のGPUメモリ削減
- 画像サイズが大きくなるほど、VimとDeiTの差が顕著に
これらの結果は、Vimが高解像度画像処理において卓越した効率性を持つことを示しています。
5. 結論と今後の展望
本研究では、Vision Mamba (Vim)という新しい画像認識モデルを提案し、その有効性を実証しました。Vimの主な貢献は以下の通りです：
-
高性能：ImageNet分類、セマンティックセグメンテーション、物体検出など多様なタスクで従来のViTを上回る性能を達成。
-
計算効率：特に高解像度画像処理において、大幅な速度向上とメモリ使用量の削減を実現。
-
スケーラビリティ：モデルサイズの拡大に伴い、一貫した性能向上を示す。
-
汎用性：2D畳み込みなどの画像特有の構造を必要とせず、シーケンスモデリングの枠組みを保持。
今後の研究方向として以下が考えられます：
- より大規模なデータセットやモデルでのVimの評価
- 動画認識など、時空間データへのVimの応用
- マルチモーダル学習におけるVimの活用（例：画像-テキスト学習）
- Vimアーキテクチャのさらなる最適化と効率化
- 医療画像分析や衛星画像解析など、高解像度画像を扱う実世界アプリケーションへのVimの適用
Vimは、高性能と高効率を兼ね備えた画像認識モデルとして、今後のコンピュータビジョン研究や実応用に大きな影響を与える可能性があります。特に、ギガピクセル級の画像処理や、リアルタイム性が求められる応用において、Vimの優位性が発揮されると期待されます。
参考文献
-
Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W., & Wang, X. (2024). Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. arXiv preprint arXiv:2401.09417v2.
-
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., ... & Houlsby, N. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR 2021.
-
Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., & Jégou, H. (2021). Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning (pp. 10347-10357). PMLR.