Structured State Space Models(構造化状態空間モデル)に基づくシーケンスモデリングアーキテクチャ。Mamba-2では入力依存の選択メカニズムにより、Transformerと同等の品質をO(n)の線形計算量で達成する。
Mamba(マンバ)は、Albert GuとTri Daoが2023年12月に発表した、Structured State Space Models(SSM)に基づく革新的なシーケンスモデリングアーキテクチャである。従来のTransformerが系列長に対してO(n²)の計算量を要求するのに対し、MambaはO(n)の線形計算量で同等以上の言語モデリング性能を達成する。
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(S6アーキテクチャ)の最大の革新は、SSMのパラメータB、C、Δを入力に依存させた「選択的SSM」である:
この選択メカニズムにより、モデルは「どの情報を状態に記憶し、どの情報を忘れるか」を動的に制御できる。これはTransformerのAttentionが行う「どのトークンに注目するか」の判断に相当する。
2024年5月に発表されたMamba-2は、SSMとAttentionの理論的等価性を示すSSD(Structured State Space Duality)フレームワークに基づく:
| モデル | パラメータ | 学習データ | コンテキスト長 | 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の利用方法は複数存在する:
AutoModelForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf")pip install mamba-ssmでカスタムCUDAカーネル付きインストールCUDAカーネルの要件:
| 特性 | Transformer | Mamba/SSM |
|---|---|---|
| 短文処理(<4K) | 高品質 | 同等品質 |
| 長文処理(>32K) | メモリ大・低速 | 高効率・高速 |
| In-context learning | 優秀 | やや劣る |
| ストリーミング推論 | KVキャッシュ必要 | 状態ベクトルのみ |
| ハードウェア最適化 | 成熟 | 発展途上 |
| 学習安定性 | 確立済み | 工夫が必要 |
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段構成。理論的背景は異なるが、両者とも線形計算量・定数メモリという性質は共通。