[論文レビュー] Demystify Mamba in Vision: A Linear Attention Perspective

Dongchen Han, Ziyi Wang+8 more|arXiv (Cornell University)|May 26, 2024

Advanced Vision and Imaging被引用数 17

ひとこと要約

本論文は Mamba が線形アテンション Transformer に密接に関連していることを示し、忘却ゲートとブロック設計を Mamba の視覚タスクでの性能の主な要因として特定し、MLLA（Mamba に触発された線形アテンションモデル）を導出して、分類と密な予測タスクにおいて視覚 Mamba を上回る。

ABSTRACT

Mamba is an effective state space model with linear computation complexity. It has recently shown impressive efficiency in dealing with high-resolution inputs across various vision tasks. In this paper, we reveal that the powerful Mamba model shares surprising similarities with linear attention Transformer, which typically underperform conventional Transformer in practice. By exploring the similarities and disparities between the effective Mamba and subpar linear attention Transformer, we provide comprehensive analyses to demystify the key factors behind Mamba's success. Specifically, we reformulate the selective state space model and linear attention within a unified formulation, rephrasing Mamba as a variant of linear attention Transformer with six major distinctions: input gate, forget gate, shortcut, no attention normalization, single-head, and modified block design. For each design, we meticulously analyze its pros and cons, and empirically evaluate its impact on model performance in vision tasks. Interestingly, the results highlight the forget gate and block design as the core contributors to Mamba's success, while the other four designs are less crucial. Based on these findings, we propose a Mamba-Inspired Linear Attention (MILA) model by incorporating the merits of these two key designs into linear attention. The resulting model outperforms various vision Mamba models in both image classification and high-resolution dense prediction tasks, while enjoying parallelizable computation and fast inference speed. Code is available at https://github.com/LeapLabTHU/MLLA.

研究の動機と目的

• 統一された枠組みの中で、Mamba と線形アテンション Transformer の関係を明らかにする。
• Mamba のどの設計選択が視覚タスクでの優れた性能を生み出しているかを特定する。
• Mamba の核となるアイデアを線形アテンションに統合して競争力のある視覚モデルを形成できることを示す。
• 忘却ゲートが有用となる場面と、視覚タスクでそれを位置エンコーディングで代替できる方法について実践的な指針を提供する。

提案手法

• 選択的状態空間モデル（SSM）と線形アテンションを統一表記で再定式化し、直接比較できるようにする。
• Mamba と線形アテンション Transformer の六つの相違点を特徴づける：入力ゲート、忘却ゲート、ショートカット、アテンションの正規化なし、単一ヘッド vs 複数ヘッド、修正されたブロック設計。
• ImageNet-1K、COCO、ADE20K のタスクで各相違点の影響を評価するアブレーション研究を実施する。
• 忘却ゲートの代替案とブロック設計という二つの核心アイデアを線形アテンションへ移植して、Mamba-Like Linear Attention（MLLA）を提案する。
• MLLA を様々な視覚 Mamba モデルと比較評価し、性能と速度の向上を示す。

実験結果

リサーチクエスチョン

• RQ1 Mamba はコア計算において線形アテンション Transformer とどのように整合し、どのように異なるか。
• RQ2 視覚性能にとって Mamba の設計コンポーネントのうち重要なのはどれで、なぜか。
• RQ3 Mamba の有益な側面を線形アテンションに統合して優れた視覚モデル（MLLA）を得られるか。
• RQ4 入力ゲート、忘却ゲート、ショートカット、正規化、ヘッド数、ブロック設計が視覚ベンチマークの性能と効率にそれぞれどのように影響するか。

主な発見

• Mamba は六つの相違点を持つ線形アテンション Transformer の一種として見ることができる：入力ゲート、忘却ゲート、ショートカット、正規化なし、単一ヘッド、修正されたブロック設計。
• 忘却ゲートとブロック設計が視覚タスクでの Mamba の優れた性能の主な要因である。
• 忘却ゲートを適切な位置エンコーディング（APE, LePE, CPE, RoPE）に置換しても、性能向上を保持しつつ並列計算を可能にする。
• 提案されたMLLAモデルは、二つの核心アイデアを線形アテンションへ組み込んだもので、ImageNet-1K、COCO、ADE20K で様々な視覚 Mamba モデルを上回る成果を達成する。
• MLLA は再帰的な成分に対して推論をより速く維持し、分類と密な予測のベンチマークで競争力のある、または優れた精度を示す。