画像分野のMamba, "Vision Mamba" #DeepLearning

はじめに

画像処理分野では、アテンションマップを用いたViT（Vision Transformer）の登場によりさらなる発展を遂げてきました。しかし、ViTによる高解像度画像の処理における計算コストとメモリ使用量の大きさは、依然として大きな課題です。

この問題に対処するため、「Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model」という論文では、新しいアーキテクチャ「Vision Mamba（Vim）」を提案しています。このモデルは、状態空間モデル（State Space Model、SSM）を基にしており、特に高解像度画像の効率的な処理能力に焦点を当てています。
Vimは、従来の自己注意機構に依存せず、双方向の状態空間モデルを活用して表現学習を行います。このアプローチにより、Vimは画像の大局的な情報を効率的に捉え、高い性能を実現しつつ計算とメモリの使用を大幅に削減します。本記事では、Vision Mambaの概要、主な貢献、そして視覚タスクにおけるその有効性について解説していきます。

状態空間モデルの基本

状態空間モデル（State Space Models, SSM）は、動的システムをモデル化するために使用される強力な数学的フレームワークです。このモデルは、システムの内部状態と観測データの関係を定義する二つの主要な方程式から構成されます。

状態方程式

状態方程式は、システムの内部状態が時間によってどのように変化するかを記述します。

x_{t+1} = Ax_t + Bu_t + w_t

ここで、

(x_t): 時刻(t)におけるシステムの状態ベクトル。
(A): 状態遷移行列、現在の状態から次の状態への変化を示します。
(u_t): 時刻(t)における制御入力ベクトル、システムへの外部からの入力を表します。
(B): 制御入力行列、入力が状態にどのように影響するかを示します。
(w_t): プロセスノイズ、モデルの不確実性や外部からのランダムな影響を表します。

観測方程式

観測方程式は、システムの内部状態に基づいて、どのような観測が得られるかをモデル化します。

y_t = Cx_t + v_t

ここで、

(y_t): 時刻(t)における観測ベクトル。
(C): 観測行列、内部状態から得られる観測へのマッピングを示します。
(v_t): 観測ノイズ、測定の不確実性を表します。

状態空間モデルは、その汎用性により、気象予測、株価の分析、脳波データの解析、ロボティクスなど、様々な分野で応用されています。特に、時系列データの複雑な依存関係を捉える能力により、深層学習と組み合わせて、より精度の高い予測モデルを構築するために用いられます。

Vision Mambaにおける双方向SSMの数式的説明

Vision Mambaでは、双方向の状態空間モデル（SSM）を利用して、画像から得られるパッチトークンシーケンスを前方向と後方向の両方から効率的に分析します。以下にその数式的な処理を示します。

SSMの基本形

状態方程式:
$ x_{t+1} = Fx_t + Bu_t + w_t $
観測方程式:
$ y_t = Hx_t + v_t $

ここで、$(x_t)$は時刻$(t)$における状態、$(y_t)$は時刻$(t)$における観測、$(F)$は状態遷移行列、$(B)$は入力行列、$(H)$は観測行列、$(u_t)$は制御入力、$(w_t)$はプロセスノイズ、$(v_t)$は観測ノイズを表します。

Vision Mambaにおける双方向SSMの数式的説明

SSMの基本形

状態方程式:
$x_{t+1} = Fx_t + Bu_t + w_t$
観測方程式:
$y_t = Hx_t + v_t$

Vision Mambaにおける双方向SSM処理

前方向処理:
$\overrightarrow{x_{t+1}} = F\overrightarrow{x_t} + \overrightarrow{w_t}$
$\overrightarrow{y_t} = H\overrightarrow{x_t} + \overrightarrow{v_t}$
後方向処理:
$\overleftarrow{x_{t-1}} = G\overleftarrow{x_t} + \overleftarrow{w_t}$
$\overleftarrow{y_t} = I\overleftarrow{x_t} + \overleftarrow{v_t}$

ここで、$(\overrightarrow{x_t})$と$(\overleftarrow{x_t})$はそれぞれ前方向と後方向の状態、$(F)$と$(G)$は前方向と後方向の状態遷移を表す行列、$(H)$と$(I)$は前方向と後方向の観測行列、$(\overrightarrow{w_t})$と$(\overleftarrow{w_t})$、$(\overrightarrow{v_t})$と$(\overleftarrow{v_t})$はそれぞれ前方向と後方向のプロセスノイズと観測ノイズを表します。

メリット

この双方向のSSM処理により、各トークンは前後の全ての情報を総合的に考慮することが可能になり、より深いレベルでの画像理解が実現します。また、SSMの効率的な計算構造は、大規模なデータセットや高解像度の画像に対する高速な処理を可能にします。

モデルの概要図

この図では、入力画像からパッチトークンを生成し、それらをVimエンコーダに供給するプロセスを示しています。ImageNet分類タスクを実行するために、パッチトークンシーケンスに追加の学習可能な分類トークンを連結します。Mambaモデルがテキストシーケンスモデリングに使用されるのとは異なり、Vimエンコーダはトークンシーケンスを前方向および後方向の両方で処理し、それにより、視覚タスクに必要な空間的認識を取り入れています。

アルゴリズムの詳細

Vimブロック内の処理詳細

Vision Mamba (Vim) のアルゴリズムは、双方向の状態空間モデル（SSM）を使用して、画像から抽出されたトークンシーケンスを効率的に処理するためのものです。ここでは、そのプロセスを詳細に説明します。

アルゴリズムのプロセス

入力シーケンスの正規化: まず、入力トークンシーケンス $(Tl-1)$ を正規化層で正規化します。
線形投影: 正規化されたシーケンスを次元 $(E)$ の $(x)$ と $(z) $に線形投影します。
前方向と後方向の処理:
- $(x)$ を前方向と後方向の両方で処理します。
- 各方向で、最初に1次元畳み込みを適用し、$(x')$ を得ます。
- $(x')$ を $(B_o)$, $(C_o)$, そして $(∆_o)$ に線形投影します。
SSM操作:
- $(∆_o)$ を使用して $(A_o)$ と $(B_o)$ を変換します。
- $(A_o)$, $(B_o)$, $(C_o)$ を使用して SSM 操作を実行し、$(y_o)$ を計算します。
ゲーティングと結合:
- $(y_{forward})$ と $(y_{backward})$ を $(z)$ でゲーティングし、加算して出力トークンシーケンス $(Tl)$ を得ます。
残差接続:
- $(Tl)$ に残差接続を加え、最終的なトークンシーケンスを出力します。

評価実験

ImageNet分類タスク、ADE20Kデータセットにおけるセマンティックセグメンテーション、およびCOCO 2017データセットにおけるオブジェクト検出とインスタンスセグメンテーションの結果を通じて、Vimの有効性が検証されました。

ImageNet分類タスク

Vimは、画像分類においても印象的な性能を示しています。特に、ImageNetデータセットにおける評価では、Vimが既存のモデルと比較して高い精度を達成しています。

Vimモデルは、同じパラメータ規模の他のモデルと比較して、Top-1精度で顕著な改善を示しました。

これは、Vimが高解像度の画像処理においても、計算効率とメモリ使用のバランスを保ちつつ、精度を向上させることができることを示しています。

ADE20Kデータセットにおけるセマンティックセグメンテーション

Vimは、セマンティックセグメンテーションタスクにおいて、異なるモデルサイズでDeiTに対して一貫して優れた性能を示しました。

Vim-TinyはDeiT-Tinyを1.8 mIoUで上回り、
Vim-SmallはDeiT-Smallに比べて0.9 mIoU高い結果を達成しています。

COCO 2017データセットにおけるオブジェクト検出とインスタンスセグメンテーション

オブジェクト検出およびインスタンスセグメンテーションのタスクでは、以下の成果が得られました。

Vim-TinyはDeiT-Tinyをオブジェクト検出で1.3 box AP、インスタンスセグメンテーションで1.1 mask AP上回りました。
中規模および大規模のオブジェクトに対して、Vim-Tinyはそれぞれ1.6 APbox_m/1.3 APmask_mと1.4 APbox_l/1.8 APmask_lでDeiT-Tinyを上回っています。

結論

Vision Mambaは、ImageNetの画像分類からセマンティックセグメンテーション、オブジェクト検出に至るまで、幅広い視覚タスクにおいて、既存のモデルと比較して計算効率と性能の両方で優れていることが実証され、特にVimが高解像度画像の処理を改善することが示されています。

参考文献

Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model, Lianghui Zhu and Bencheng Liao and Qian Zhang and Xinlong Wang and Wenyu Liu and Xinggang Wang