MambaをState Space Dualityフレームワークで再定式化した後継版。SSMとアテンションの理論的等価性を証明し、GEMMベースの最適化で初代Mamba比最大8倍の学習スループットを実現する。
Mamba-2はTri DaoとAlbert Guが2024年5月にarXivで発表した、初代Mambaの理論的・実装的改良版。核心はState Space Duality(SSD)フレームワーク:SSMとアテンションが同じ数学的構造(セミセパラブル行列積)の異なる表現であることを証明した。
SSDが示したこと:
| 項目 | Mamba-1 | Mamba-2 |
|---|---|---|
| 状態次元 | 16〜64 | 64〜256(大幅拡大) |
| 実装 | カスタムCUDAカーネル | GEMM + cuBLAS |
| 学習速度 | 基準 | 最大8倍高速 |
| マルチヘッド | 非対応 | 対応(アテンション類似) |
| 理論的根拠 | ヒューリスティック | SSD双対性 |
状態次元の拡大により長距離依存性の表現能力が向上し、GEMMベース実装でTensorコアを最大活用できるようになった。
Mamba-2は系列をチャンク分割してスキャンを並列化する。各チャンク内はアテンション形式(高並列性)、チャンク間はSSM再帰形式(線形スケール)で処理し、両者の長所を組み合わせる。これにより単純なシーケンシャルスキャンより大幅に並列効率が向上する。
| モデル | 開発元 | 規模 |
|---|---|---|
| FalconMamba 7B | TII UAE | 7B 純粋SSM |
| Codestral Mamba | Mistral AI | コード特化 |
| NVIDIA Mamba-2-Hybrid | NVIDIA | Transformer+Mamba2ハイブリッド |
| Zamba2-7B | Zyphra | Mamba2+Transformer混合 |
Mamba-2はSSMの主力実装として定着したが、純粋SSMはin-context learning(多ショット学習)でTransformerに劣るという研究結果も増えており、長文脈処理はハイブリッド設計(Attention + SSM交互配置)が主流になりつつある。NVIDIA、MistralなどはMamba2ブロックとアテンションブロックを組み合わせた設計を採用している。