[論文レビュー] Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao|arXiv (Cornell University)|Dec 1, 2023

Neural Networks and Applications被引用数 967

ひとこと要約

Mamba は入力依存の選択的状態空間モデルを導入し、線形時間のシーケンスモデリングを実現。Transformersより5倍の推論スループットを達成し、言語・音声・ゲノミクスの領域で高い性能を発揮します。言語タスクでは、Mamba-3Bは同程度の規模のTransformersを上回り、サイズを2倍にすると同等になります。

ABSTRACT

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5$ imes$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

研究の動機と目的

• 長いシーケンスにおける Transformer の非効率性に対処するため、コンテンツ感知型の選択的状態空間（SSMs）を導入する。
• アテンションやMLPなしで線形時間のシーケンスモデリングを実現しつつ、モダリティを横断する精度を維持する。
• 再帰モードでの効率を維持するハードウェア配慮の並列アルゴリズムを選択的SSMsのために開発する。

提案手法

• SSMパラメータを入力アドレスの関数にして、シーケンスに沿った情報フローの離散的で内容ベースの変調を可能にする。
• アテンションやMLPを用いないシンプルなエンドツーエンドのバックボーン（Mambaモデル）に選択的SSMsを統合する。
• 標準的な畳み込み最適化がない状況でも高速推論を維持するよう、ハードウェア配慮のある並列再帰アルゴリズムを設計する。
• 推論時にシーケンス長とともに線形スケーリングを実証し、Transformersより5×のスループット改善を示す。

実験結果

リサーチクエスチョン

• RQ1 従来のサブ二次モデルと比べて、内容ベース（アドレス依存）の選択的SSMsは長大なシーケンス推論を改善できるか。
• RQ2 選択的SSMsは言語・音声・ゲノミクス全領域で競争力のある性能を持つ線形時間のシーケンスモデリングを実現するか。
• RQ3 アテンション/MLPsなしで、再帰型アーキテクチャにハードウェアに優しい並列性を実現し、スループットを維持することは可能か。
• RQ4 同サイズおよびより大きいサイズの Transformers と比較した場合、言語モデリングにおける Mamba の性能はどうなるか。

主な発見

• Mamba は推論を高速化し、Transformersより5×のスループットを達成する。
• モデルはシーケンス長に対して線形にスケールする。
• Mamba は言語・音声・ゲノミクスのモダリティで最先端の成果を示す。
• 言語モデリングにおいて、Mamba-3Bは同サイズのTransformersを上回り、サイズの2倍のTransformersと同等の性能を、事前学習と下流評価の両方で示す。