Mamba SSM(マンバエスエスエム)
CMUが2023年に発表した選択的状態空間モデル。入力依存のゲーティング機構により、Transformerの二次複雑度問題を克服しO(L)線形時間で長文脈を処理する。
Mamba SSM とは
Mamba SSMはCMUのAlbert GuとTri Daoが2023年12月に発表した選択的状態空間モデル(Selective State Space Model)。従来のSSM(S4、H3等)が固定パラメータで全系列を一律処理するのに対し、Mambaは入力依存のゲーティングを導入してモデルが動的に何を記憶・忘却するかを選択できる。
Transformerとの根本的な違い
| 項目 | Transformer | Mamba SSM |
|---|---|---|
| 時間複雑度 | O(L²) | O(L) |
| 空間複雑度 | O(L²) KVキャッシュ | O(d) 固定状態 |
| 並列学習 | ✓ アテンション並列 | ✓ スキャン並列 |
| 逐次推論 | KVキャッシュ増大 | 固定サイズ状態 |
Transformerはself-attentionで系列長Lの二乗に比例する計算・メモリを消費するが、MambaはSSMの再帰的状態更新でO(L)に抑える。
Selection Mechanism の仕組み
Mambaの核心は選択メカニズム。SSMのパラメータ B、C、Δ(離散化ステップサイズ)を入力xtの関数として動的に算出する:
- B(入力投影): 入力をどれだけ状態に取り込むか
- C(出力投影): 状態からどれだけ出力に反映するか
- Δ(ステップサイズ): 連続SSMを離散化する際の時間刻み。大Δ→過去を忘れやすい
時系列的に無関係なトークンはΔを小さくして無視し、重要なトークンはΔを大きくして状態に焼き付ける。これにより「選択的コピー」「誘導頭(induction head)」相当の能力をSSMで実現した。
Hardware-Aware Algorithm
選択的SSMはHiPPO理論をベースに離散化が必要で、ナイーブ実装ではGPUのSRAM↔HBM転送がボトルネックになる。MambaはFlashAttentionと同様のカーネルフュージョン(Associative Scan + Kernel Fusion)でHBMアクセスを削減し、実測でフラッシュアテンション並みのスループットを達成した。
性能・実績
- 1.4B〜2.8BパラメータでTransformer同規模モデルと同等のperplexity
- 推論スループット:Transformerの最大5倍(長文脈ほど差が拡大)
- Pile/The Stackでの評価:GPT-Neo相当の精度
- Jamba(12B MoE + Mamba)は256Kトークン文脈でGPT-4同等の長文精度
派生・後継モデル
| モデル | 概要 |
|---|---|
| MambaByte | バイト単位LM、トークナイザー不要 |
| Jamba | Mamba+Transformerハイブリッド、AI21 Labs |
| Zamba | 7Bハイブリッド、Zyphra |
| Falcon Mamba 7B | 純粋SSM、TII UAE |
| Codestral Mamba | コード特化、Mistral AI |
| NVIDIA Mamba-2-Hybrid | Mamba2ブロック+注意ブロック混合 |
2025年の位置づけ
Mambaは「Transformerキラー」の期待を受けたが、実際は長文脈推論でのin-context learningがTransformerより劣るという報告が複数あり、完全置換よりTransformerとのハイブリッドがトレンドになっている。それでも長文書要約・ゲノム配列処理・オーディオモデリングではMamba純粋系が優位な場合がある。