[論文レビュー] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Tri Dao, Albert G. Gu|arXiv (Cornell University)|May 31, 2024

Neural Networks and Applications被引用数 66

ひとこと要約

本論文は構造化状態空間モデル（SSM）とアテンションの間に形式的な結びつきを確立し、Structured State Space Duality（SSD）を導入して、Language modelingにおいてTransformerに匹敵する効率的なSSMベースのアーキテクチャ（Mamba-2）を導出し、スケーラビリティの利点を提供する。

ABSTRACT

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling.

研究の動機と目的

• SSMとTransformerを橋渡しして長いシーケンスのスケーラビリティと効率を改善する動機付け。
• SSMと構造化行列およびアテンション変種を結ぶ理論的枠組み（SSD）の開発。
• 計算効率のために線形再帰と二次的双対形を統合する新しいアルゴリズム（SSD）の提示。
• SSDと構造化アテンションを活用したより良いスループットを実現するアーキテクチャ設計（Mamba-2）の提案。

提案手法

• SSMを行列変換として形式化し、半独立（semiseparable）行列（SSS表現）へ結びつける。
• 構造化状態空間デュアリティ（SSD）を導入して線形（再帰）と二次的（アテンション様）形を統一する。
• 構造化行列を用いて線形アテンションを構造化マスクドアテンション（SMA）へ一般化する。
• 半独立行列のブロック分解に基づくSSDアルゴリズムをハードウェア効率のために開発する。
• grouped-value attentionと並列データ射影を備えたMamba-2アーキテクチャを提案する。
• 計算上のトレードオフを分析し、ハードウェア指向の実装指針を提供する。

実験結果

リサーチクエスチョン

• RQ1 SSMを半独立行列として表現し、アテンションとのアルゴリズム的デュアル性を明らかにするにはどうすればよいか。
• RQ2 統一されたSSDフレームワークは、Transformerに匹敵する効率的なSSMの学習と推論を実現できるか。
• RQ3 現代のハードウェアに適合しつつ言語モデリングの性能を保つようにSSMを最適化するための修正（例：Mamba-2）は何か。
• RQ4 構造化マスクドアテンションの変種は線形アテンションのフレームワークとどのように関連し、一般化するか。

主な発見

• SSMsは半独立行列に同等であり、適切な構造を用いるとO(TN)計算が可能である。
• SSDは線形再帰とアテンション様の二次形の橋渡しを提供し、効率的なハイブリッドを可能にする。
• 新しいSSDアルゴリズムはMambaにおける選択的スキャンよりも性能を発揮しつつ、より大きな再帰状態を可能にする。
• パラレルデータ射影とSMA変種を備えたMamba-2アーキテクチャは、標準の下流タスクでMambaおよびオープンソースのTransformersと同等かそれを上回る。
• 言語モデリング、効率指標、および難易度の高いリコールタスクでの実証的検証は競争力のある性能を支持する。