論文 · arXiv 2312.00752
Mamba: 選択的状態空間による線形時間シーケンスモデリング
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
3行まとめ
- 🎯 Transformerは系列長に対して二乗で計算量が増加し、自己回帰推論時に巨大なKVキャッシュを必要とする一方、従来のLTI状態空間モデルなどの劣二次時間モデルは言語モデリングのような離散的でコンテンツ推論を要するタスクで苦戦していました。
- 🔧 SSMパラメータを入力依存（選択的）にして情報を動的にフィルタリング・伝播できるようにし、ハードウェアを意識した融合型並列結合スキャンと再計算アルゴリズムを組み合わせることで、拡張された隠れ状態がGPUのHBMに展開されるのを防ぎます。
- 📈 Mamba-3Bは下流のゼロショットタスクにおいて2倍のサイズのTransformerに匹敵またはそれを上回り（例: 平均精度でPythia-2.8Bの59.1%、Pythia-6.9Bの61.7%に対し63.3%）、最大5倍高い推論スループットを提供します。
概要
現在ディープラーニングにおける多くの革新的なアプリケーションを支えている基盤モデルは、ほぼ例外なくTransformerアーキテクチャとその中核であるアテンションモジュールに基づいています。長系列に対するTransformerの計算非効率性を解消するため、線形アテンション、ゲート付き畳み込み、再帰モデル、構造化状態空間モデル（SSM）など、多くの劣二次時間（subquadratic-time）アーキテクチャが開発されてきましたが、言語などの重要なモダリティにおいてアテンションと同等の性能を達成できていませんでした。本研究では、これらモデルの決定的な弱点が「コンテンツに基づく推論」を行えない点にあることを特定し、いくつかの改善を提案します。第一に、SSMのパラメータを入力の関数に設定するだけで、離散モダリティに対する弱点が克服され、現在のトークンに応じて系列長次元に沿って情報を選択的に伝播または忘却できるようになります。第二に、この変更により効率的な畳み込みは使用できなくなりますが、再帰モードにおけるハードウェア認識型の並列アルゴリズムを設計しました。これらの選択的SSMを、アテンションやMLPブロックすら持たないシンプルなエンドツーエンドのニューラルネットワークアーキテクチャ（Mamba）に統合しました。Mambaは高速な推論（Transformerの5倍高いスループット）と系列長に対する線形スケーリングを実現し、実データにおいて100万長トークンまでの系列で性能向上が確認されました。汎用的なシーケンスモデルのバックボーンとして、Mambaは言語、音声、ゲノミクスなどの複数のモダリティにおいてSOTAの性能を達成しています。言語モデリングにおいて、Mamba-3Bモデルは事前学習および下流タスクの評価の両方で同規模のTransformerを上回り、その2倍のサイズのTransformerに匹敵する性能を示しています。
英語の原文概要
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 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アーキテクチャと自己アテンションに依存しています。密なアテンションは系列コンテキストウィンドウ全体で情報を効果的に伝達しますが、系列長 に対して二乗の計算コスト が発生すること、および過去のキーと値（KVキャッシュ）を保存する必要があるためステップあたりの推論が線形時間となりメモリのボトルネックが生じるという、2つの根本的な課題を抱えています。
線形アテンション、再帰アーキテクチャ、S4などの構造化状態空間モデル（SSM）を含むこれまでの劣二次時間アプローチは、時間に関して線形または疑似線形にスケールします。しかし、構造化SSMは歴史的に、テキストのような離散的で情報密度の高いモダリティにおいてTransformerの品質に匹敵することができませんでした。著者らはその根本原因を突き止めました。従来のSSMは線形時不変（LTI）であるという点です。遷移ダイナミクスが入力コンテンツに応じて変化しないため、コンテキストを選択的に圧縮したり、無関係なノイズをフィルタリングしたり、可変の間隔にわたって保持を調整したりすることができませんでした。
主要なアイデア
- コンテキスト圧縮としての選択機構: シーケンスモデリングには、履歴を状態へと圧縮することが求められます。従来のLTI SSMは固定の遷移パラメータを使用していたため、すべてのタイムステップを均一に処理せざるを得ませんでした。現在の入力トークンに基づいてSSMの行列をパラメータ化することで、モデルは情報を選択的にフィルタリングまたは記憶するために必要なコンテンツ認識能力を獲得します。
- 非畳み込みボトルの克服: パラメータを入力依存にすると、線形再帰と大域的畳み込みの数学的等価性が崩れます。畳み込みが利用できない場合、単純な再帰計算ではGPUの高帯域幅メモリ（HBM）上に サイズの拡張された潜在状態を展開（実体化）する必要があります。Mambaは、離散化、並列結合スキャン、出力射影を高速なGPUオンチップSRAM内で融合（ヒューズ）して計算し、HBMとのトラフィックを最小限に抑えることでこれを解決します。
- 統合アーキテクチャ: Mambaは、交互に配置されるマルチヘッドアテンション（MHA）とMLPブロックを、H3のゲート付き設計とSwiGLUスタイルのゲート付きMLPを組み合わせた単一の均一なブロックに置き換えています。
手法
連続時間状態空間モデルは、暗黙的な潜在状態 を介して入力 を にマッピングします。
ステップサイズ を持つゼロ次ホールド（ZOH）離散化を用いることで、連続システムは離散的な再帰パラメータに変換されます。
選択機構（S6）
を時間によらず一定に保つ代わりに、Mambaではこれらを入力依存にします。
- 、ただし 、
これは古典的なRNNのゲーティングを一般化したものです。 の場合、再帰はゲート付き更新 （ここで ）へと直接単純化されます。
ハードウェア認識型の状態拡張
HBM内に の状態を展開することなく高いスループットを維持するため、以下の最適化を行います。
- パラメータ をHBMからオンチップSRAMに直接ロードします。
- への離散化と作業効率の良い並列結合スキャンを、すべてSRAM内部で計算します。
- HBMへ書き戻す前に、潜在状態を によって縮約し の形状にします。
- 逆伝播時の中間状態は、HBMに保存するのではなくSRAM内で再計算します。
実験結果
- 人工データでの外挿性: Induction Headsタスクにおいて、系列長256で学習されたMambaは系列長 （4000倍の系列長）まで完璧に外挿（精度100%）できたのに対し、アテンションベースラインは を超えると失敗するかメモリ不足（OOM）になりました。Selective Copyingタスクでは、H3が57.0%、S4が18.3%にとどまる中、Mambaは99.8%の精度を達成しました。
- 言語モデリング: The Pile上で最大1.3BパラメータまでのChinchillaスケーリング則のもとで評価した結果、MambaはTransformer++（Rotary Embeddings、SwiGLU、RMSNorm）と同等以上の性能を示しました。下流のゼロショットタスク（LAMBADA、PIQA、HellaSwag、WinoGrande、ARCなど）において、Mamba-2.8Bは平均精度63.3%を達成し、Pythia-2.8B（59.1%）やPythia-6.9B（61.7%）を上回り、6B〜7Bパラメータのモデルに匹敵または凌駕しました。
- ゲノミクスおよび音声: 系列長100万（）での大型類人猿DNA分類において、MambaはHyenaDNAの54.87%に対し、71.67%（1.4Mパラメータ）および81.31%（7Mパラメータ）の精度を達成しました。SC09音声生成では、MambaはFIDを1.99（SaShiMi）から0.94（6.1Mパラメータ）および0.67（24.3Mパラメータ）へと大幅に改善しました。
- 効率性: 融合スキャンカーネルは標準的なPyTorchのスキャン実装よりも20〜40倍高速に動作し、32KコンテキストにおいてFlashAttention-2より最大7倍高速です。推論スループットは同等のTransformerに比べて最大5倍向上しています。
限界と未解決の課題
- 連続信号での性能: 入力依存の選択機構は離散的な推論（テキスト、DNA）を向上させますが、生音声波形の事前学習に関するアブレーションでは、均一にサンプリングされた連続データに対しては純粋な線形時不変（LTI）モデル（標準的なS4など）のほうが優れていることが示されています。
- モデル規模: 評価されたモデルは約3Bパラメータおよび300B事前学習トークンにとどまっており、7B〜70B+規模以上での実証的性能についてはさらなる検証が必要です。
- 下流のエコシステム: Transformerには、プロンプトエンジニアリング、インコンテキスト操作、KVキャッシュの量子化、アダプタ調整などの広範なツール群が存在します。SSMがこれらの対話的・応用手法にどのように適合するかは、現在も活発に研究が進められている分野です。
重要用語
- Structured State Space Models (SSMs)（構造化状態空間モデル）
- 暗黙的な潜在状態を介して1次元入力をマッピングする、連続線形動的システムに由来するディープシーケンスモデルの一種。
- Linear Time Invariance (LTI)（線形時不変性）
- システムの遷移パラメータがすべてのタイムステップで静的に保たれる数学的特性。高速な大域的畳み込み表現を可能にするが、コンテンツ依存の推論を制限する。
- Selection Mechanism（選択機構）
- SSMの遷移ダイナミクス（Δ、B、C）を入力トークンの関数としてパラメータ化し、系列長に沿ってモデルが動的にコンテキストを伝播、リセット、または無視できるようにする手法。
- S6
- 畳み込みの代わりに結合的並列スキャンを介して計算される選択的構造化状態空間モデルの略称。
- Parallel Associative Scan（並列結合スキャン）
- アクセラレータハードウェア上で系列長全体にわたる線形再帰を並列に評価する、作業効率の高いプレフィックス計算アルゴリズム。
- Hardware-Aware State Expansion（ハードウェア認識型の状態拡張）
- 巨大な中間状態が高帯域幅メモリ（HBM）を往復するのを避けるため、高速なオンチップSRAM内で離散化とスキャンを実行するGPUカーネル融合および再計算戦略。
- Induction Heads（インダクションヘッド）
- モデルが先行するコンテキストトークンに続いて以前に観察されたトークンを再現できるかどうかを評価する、合成連想記憶メカニズム。
- Mamba Block（Mambaブロック）
- 個別のアテンションやMLP層を持たず、入力射影、1D畳み込み、選択的SSM、および乗算ゲーティングを統合したシンプルな均一シーケンスモデルブロック。
よくある質問
選択的SSMを高速なFFT畳み込みで学習できないのはなぜですか？
高速な大域的FFT畳み込みは、システムダイナミクス（行列A、B、C）が時間を通じて一定である（線形時不変性）ことを前提としています。Mambaでは各タイムステップの入力トークンに応じてB、C、Δが変化するため、畳み込みカーネルが位置ごとに異なってしまい、畳み込みとしての等価性が失われるため再帰的な並列スキャンが必要となります。
Mambaは学習時の再帰的な状態拡張によるメモリ爆発をどのように防いでいますか？
状態拡張次元Nを持つ単純な再帰では、GPUのHBM上に(B, L, D, N)の形状の中間状態を展開する必要があります。Mambaはカーネル融合を採用し、低次元の入力のみをSRAMにロードして高速なSRAM内で離散化と並列スキャンを計算し、縮約された(B, L, D)の出力のみをHBMに書き戻します。さらに逆伝播時には中間活性化を再計算します。
MambaのΔパラメータと再帰ゲーティングの関係は何ですか？
ステップサイズパラメータΔは、一般化されたゲーティング機構として機能します。N=1かつA=-1のゼロ次ホールド（ZOH）のもとで離散化すると、選択的再帰は数学的に標準的なRNNゲートへと帰着します。大きなΔはモデルに現在の入力へ注目させて状態をリセットさせ、小さなΔは過去の状態を維持して入力を無視するように作用します。
Mambaは生成時にKVキャッシュを使用しますか？
いいえ、使用しません。Mambaは生成中に自己回帰的な状態空間の再帰として動作するため、過去のすべてのトークンのキーと値をキャッシュするのではなく、固定サイズ（層あたり次元D × N）の隠れ状態のみを維持します。これにより、ステップあたりO(1)の定数時間推論が可能になります。
このページの解説は AI が論文から作成したもので、誤りを含む場合があります。詳細は原文をご確認ください。