Mamba / SSM（マンバ / エスエスエム）
Structured State Space Models（構造化状態空間モデル）に基づくシーケンスモデリングアーキテクチャ。Mamba-2では入力依存の選択メカニズムにより、Transformerと同等の品質をO(n)の線形計算量で達成する。
Mamba / SSMとは
Mamba（マンバ）は、Albert GuとTri Daoが2023年12月に発表した、Structured State Space Models（SSM）に基づく革新的なシーケンスモデリングアーキテクチャである。従来のTransformerが系列長に対してO(n²)の計算量を要求するのに対し、MambaはO(n)の線形計算量で同等以上の言語モデリング性能を達成する。
SSMの数学的基盤
SSMは連続時間の状態遷移方程式を基礎とする：
h'(t) = Ah(t) + Bx(t)
y(t) = Ch(t) + Dx(t)
ここでA ∈ R^(N×N)は状態遷移行列、B ∈ R^(N×1)は入力行列、C ∈ R^(1×N)は出力行列。この連続系をZOH（Zero-Order Hold）で離散化すると、効率的な再帰計算が可能になる。
S4（Structured State Spaces for Sequence Modeling、2022年）では、行列Aを対角化し、HiPPO（High-order Polynomial Projection Operator）で初期化することで、長距離依存関係の学習を安定化させた。
Mambaの選択メカニズム
Mamba（S6アーキテクチャ）の最大の革新は、SSMのパラメータB、C、Δを入力に依存させた「選択的SSM」である：
- 従来のSSM（S4等）: A, B, C, Δは全入力に対して固定（Linear Time-Invariant）
- Mamba（S6）: B(x), C(x), Δ(x)が入力xに応じて動的に変化（入力依存選択）
この選択メカニズムにより、モデルは「どの情報を状態に記憶し、どの情報を忘れるか」を動的に制御できる。これはTransformerのAttentionが行う「どのトークンに注目するか」の判断に相当する。
Mamba-2の改良点
2024年5月に発表されたMamba-2は、SSMとAttentionの理論的等価性を示すSSD（Structured State Space Duality）フレームワークに基づく：
- SSDフレームワーク: SSMの再帰計算がsemi-separable行列の乗算と等価であることを証明
- チャンク分割: 入力系列をチャンクに分割し、チャンク内は行列演算（GPUテンソルコア活用）、チャンク間は再帰計算
- マルチヘッド化: Attentionのマルチヘッドに相当する並列状態空間を導入
- 性能: Mamba-1比で学習速度2-8倍向上、2.7Bパラメータでtransformer++と同等品質
| モデル | パラメータ | 学習データ | コンテキスト長 | HellaSwag | PIQA | WinoGrande |
|---|---|---|---|---|---|---|
| Mamba-1 1.4B | 1.4B | 300B tok | 2K | 59.1 | 73.9 | 61.4 |
| Mamba-2 2.7B | 2.7B | 300B tok | 4K | 66.7 | 76.4 | 63.5 |
| Transformer++ 2.7B | 2.7B | 300B tok | 4K | 66.3 | 76.0 | 63.3 |
| Jamba 52B (MoE) | 52B (12B active) | 非公開 | 256K | 87.1 | 83.2 | 82.0 |
実装と利用方法
Mambaの利用方法は複数存在する：
- HuggingFace Transformers:
AutoModelForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf") - mamba-ssm パッケージ:
pip install mamba-ssmでカスタムCUDAカーネル付きインストール - Jamba経由: AI21のJamba 1.5（Mamba+Transformer+MoE）がAPIとオープンウェイトで利用可能
CUDAカーネルの要件：
- NVIDIA GPU: Ampere（A100）以降推奨
- CUDA 11.8以上
- PyTorch 2.0以上
商用採用事例
- AI21 Jamba: 52Bパラメータ（12B active）、256Kコンテキスト、Mamba+Transformer+MoEハイブリッド
- NVIDIA Nemotron-H: Mamba-2ベースのハイブリッドモデル、エンタープライズ推論向け
- Zamba2: Mamba2+共有Attention、2.7Bで7Bクラス性能、エッジデバイス向け
- Codestral Mamba: Mistral AIのコード生成特化Mamba、256Kコンテキスト
Transformerとの使い分け
| 特性 | Transformer | Mamba/SSM |
|---|---|---|
| 短文処理（<4K） | 高品質 | 同等品質 |
| 長文処理（>32K） | メモリ大・低速 | 高効率・高速 |
| In-context learning | 優秀 | やや劣る |
| ストリーミング推論 | KVキャッシュ必要 | 状態ベクトルのみ |
| ハードウェア最適化 | 成熟 | 発展途上 |
| 学習安定性 | 確立済み | 工夫が必要 |
よくある質問（FAQ）
Q1: MambaはGPUなしでも動作するか？ A: カスタムCUDAカーネルなしでも動作するが、速度が大幅に低下する。CPU推論にはmamba.cppプロジェクトが進行中。ONNX変換によるNPU対応も研究されている。実用的にはNVIDIA Ampere以降のGPUが推奨される。
Q2: Mambaの学習にはどれくらいのリソースが必要か？ A: Mamba-2 2.7Bの学習には8×A100 80GBで約3日（300Bトークン）。同規模のTransformerと比較して学習FLOPsは同等だが、長系列での学習はMambaの方が効率的。1.4Bモデルなら4×A100で1日程度。
Q3: MambaとRWKVの違いは？ A: Mambaは連続時間SSMの離散化がベースで、選択メカニズムによる動的状態更新が特徴。RWKVはAttentionの線形近似（WKV演算）がベースで、チャネルミキシングとトークンミキシングの2段構成。理論的背景は異なるが、両者とも線形計算量・定数メモリという性質は共通。
まとめ
- MambaはSSMに入力依存の選択メカニズムを導入した革新的アーキテクチャ
- O(n)の線形計算量でTransformerと同等の言語モデリング品質を達成
- Mamba-2はSSDフレームワークにより学習速度2-8倍向上
- Jamba、Nemotron-H、Zamba2など商用採用が進展
- 長系列処理とストリーミング推論で特に優位性を発揮