Mamba-3: 推論効率を最優先に設計された新しい状態空間モデル

tl;dr

Mamba-3 は、推論効率を最優先の目標として設計された新しい状態空間モデル（SSM）です。これは、トレーニング速度を最適化した Mamba-2 からの大きな飛躍です。主な改良点として、より表現力豊かな再帰数式、複素数値による状態追跡、そしてデコード速度を落とさずに精度を向上させる MIMO（多入力・多出力）バリアントの採用が挙げられます。

その結果、Mamba-3 SISO は、1.5B スケールにおいて、すべてのシーケンス長で Mamba-2、Gated DeltaNet、さらには Llama-3.2-1B（Transformer）を上回るプリフィル＋デコードのレイテンシを達成しました。

また、チームはハードウェアのパフォーマンスを最大限に引き出すため、Triton、TileLang、CuTe DSL を組み合わせて構築されたカーネルをオープンソース化しました。

本ブログは Goomba Lab ブログにも同時投稿されており、カーネギーメロン大学、プリンストン大学、Cartesia AI、Together AI の研究者たちとの共同研究による成果です。

2024 年半ばの Mamba-2 リリース以来、ほとんどのアーキテクチャが Mamba-1 から移行しました。その理由は、Mamba-2 が「トレーニング効率こそが状態空間モデル（SSM）の最大のボトルネックである」という賭けに出たことにあります。これにより基盤となる SSM メカニズムが簡素化され、前身に比べて 2〜8 倍のトレーニング速度を実現し、より広い採用につながりました。

それ以来、LLM の状況は変化し始めています。事前学習も依然として極めて重要ですが、現在は学習後の調整やデコードにより多くの注目が集まっており、これらは極めて推論負荷の高い工程です。特に、コード生成や数学における検証可能な報酬付き強化学習（RLVR）などの学習後手法のスケールアップには、膨大な量のアンドロールアウトが必要となります。さらに最近では、Codex、Claude Code、あるいは OpenClaw などのエージェント型ワークフローが推論需要を天井知らずに押し上げています。

推論の重要性が明確に高まっているにもかかわらず、多くの線形アーキテクチャ（Mamba-2 を含む）は、依然として「トレーニングファースト」の視点で開発されてきました。事前学習を高速化するため、基盤となる SSM は段階的に単純化されてきました（例：対角遷移行列がスカラー×単位行列に削減されるなど）。これによりトレーニング速度は向上しましたが、その代償として推論ステップが「単純すぎ」てしまい、メモリ律速に陥っています。つまり、GPU の演算能力は使われず、メモリの転送に時間を費やしている状態です。

この新しい「推論の時代」において、私たちが重視しているのは、品質と効率のフロンティアを押し広げること、つまりより良いモデルを、より速く動かすことです。

ここで自然な疑問が湧きます。

推論を意識して設計された SSM とは、どのようなものだろうか？

Mamba-3 モデルについて

何が欠けているのか？ 線形モデルの主な魅力はその名が示す通り、固定サイズのステート（状態）を持つため、計算量がシーケンス長に対して線形にスケールすることです。しかし、タダより高いものはありません。効率的な計算を可能にする固定ステートサイズは、過去のすべての情報を 1 つの表現に圧縮することをモデルに強います。これは、過去のすべての情報を（KV キャッシュという）継続的に成長するステートに保存する Transformer とは根本的に異なる点です。では、ステートを大きくできない場合、いかにしてその固定されたステートにより多くの仕事をさせることができるでしょうか？

私たちは、初期の設計がトレーニングを高速化するために再帰や遷移行列を単純化したことに着目しました。しかし、この変更はダイナミクスの豊かさを減少させ、デコードをメモリ律速にしました。各トークンの更新において、メモリの移動に対して行われる計算が極めて少なかったのです。ここから、私たちが引くことのできる 3 つのレバーが見えてきます。(1) 再帰そのものをより表現力豊かにする、(2) より豊かな遷移行列を使用する、(3) 各更新内に、より多くの並列的（かつほぼ無料の）作業を追加する、の 3 点です。

これらの洞察に基づき、Mamba-2 を以下の 3 つの中核的な方法で改良しました。

1. 指数台形離散化スキームに由来するより一般的な再帰を用いて、SSM メカニズムの表現力を向上させる。
2. 複素数値の SSM システムをモデル化することで、状態追跡能力を拡張する。
3. 単一入力・単一出力（SISO）の SSM に代わり、複数の SSM を並列にモデル化する多入力・多出力（MIMO）SSMを採用することで、デコードレイテンシにほとんど影響を与えることなく、モデルの全体的なパフォーマンスを向上させる。

これら 3 つの変更により、Mamba-3 は同様の推論レイテンシを維持しながら、パフォーマンスのフロンティアを押し広げました。

特筆すべきは、これら 3 つのすべての変更が、より「古典的」な制御理論および状態空間モデルの文献にインスパイアされている点です。

私たちの研究は、線形アテンションやテストタイムトレーニングなど、再帰の代替的な解釈を用いる多くの現代的な線形アーキテクチャの潮流に逆らうものです。なぜなら、それらはこれらの概念を容易には捉えられないからです。

アーキテクチャ

Mamba-2 レイヤーで何が変わったのでしょうか？上記のコア SSM に対する 3 つの方法論的アップグレードに加え、現代的な言語モデルにより合致させるため、アーキテクチャ全体を刷新しました。

図を見てわかるとおり、いくつかの変更点があります。大まかには以下の通りです。

Norm（正規化）: SSM 用語で「BCNorm」と呼ばれる QKNorm を追加しました。これは経験的に Mamba-3 モデルのトレーニングを安定させます。この Norm の追加により、Mamba-3 は現代の Transformer や Gated DeltaNet (GDN) モデルと並びました。QKNorm により、Mamba-2 の RMSNorm はオプションになりますが、長さの一般化能力を高めるのに役立つため、ハイブリッドモデルでは維持する価値があることが経験的にわかっています。

ショートコンボリューションとの決別: Mamba-1/2 の厄介なショート因果コンボリューションを廃止できました。これは、(1) BCNorm 後の B および C への単純なバイアスと、(2) 新しい離散化に基づく再帰を組み合わせることで実現しました。新しい再帰は入力から隠れステートへコンボリューションを暗黙的に適用します。

実証結果

最終的な Mamba-3 モデルを、他の人気のある線形代替案や Transformer ベースラインに対して評価しました。

言語モデリング

新しい Mamba-3 モデルは、さまざまな事前学習済みモデルスケールにおいて、言語モデリングのタスクで、従来の Mamba-2 モデルや GDN などの強力な線形アテンション代替案を上回る性能を示しました。Mamba-3-SISOは、従来の線形モデルと直接比較可能です。例えば、アーキテクチャ形状（モデル次元、ステートサイズなど）は Mamba-2 と正確に一致し、トレーニング時間も同程度です。また、Mamba-3 のMIMOバリアントは、1B スケールで通常の Mamba-3 に比べて下流タスクの精度を 1 ポイント以上向上させました。ただし、MIMO はより長いトレーニング時間を必要とする一方、デコードレイテンシは長くならないという点に注意が必要です。

なぜトレーニングコストは上がるのに、推論はそうならないのか？

これについてはブログの第 2 部で詳しく説明しますが、読者の皆様に少しだけ先取りしてお見せしましょう。

この二律背反は、トレーニングと推論がそれぞれ計算律速かメモリ律速かという性質の違いに起因します。現在の線形モデルは、高速なトレーニングのために大量のGPU テンソルコアを使用するように設計されています（Mamba-2 の主な貢献の 1 つです）。しかし、デコード中は、各タイムステップに必要な計算量が非常に少ないため、ハードウェアはほとんどの時間アイドル状態のままです。

したがって、各タイムステップに必要な FLOPs の量を増やすようにアーキテクチャを設計すれば、アイドル中のコアを一部使用できるため、推論レイテンシはほぼ一定に保たれます。しかし、トレーニングはそうはいきません。

検索タスク

固定サイズのステートを持つ線形モデルは、検索ベースのタスクにおいて、Transformer 対応モデルに自然と劣ります。予想通り、純粋なモデル内では Transformer が検索タスクで優位ですが、Mamba-3 は 2 次未満の代替モデルの中では健闘しています。興味深いことに、MIMO の追加はステートサイズを増やすことなく、さらに検索パフォーマンスを向上させました。

この本質的な不足がありながらも全体的に強力なモデリング性能を考慮すると、

将来、線形レイヤーはグローバルな自己アテンションレイヤーと併用して主に使用されると予測されます。

言語モデリングに関しては少なくともそうです。

線形レイヤーの一般的なメモリのような性質と、自己アテンションの KV キャッシュによる正確なデータベースのような保存機能を組み合わせたハイブリッドモデルは、純粋なモデルを上回り、大幅なメモリおよび計算コストの削減を可能にすることが経験的に示されています。そして実際、線形レイヤーと自己アテンションの組み合わせは、バニラの Transformer と比較して、より良い検索を可能にすることを確認しました。

ただし、これらの線形モデルが自己アテンションとどのように相互作用するかという正確な方法は完全には理解されていません。例えば、Mamba-3 のオプションの出力前プロジェクションを使用すると、合成的な NIAH タスクでの長さの一般化パフォーマンスが向上しますが、文脈内での現実世界の検索タスクにはわずかなコストがかかります。さらに、返されるノームの詳細、例えばその配置（プレゲートかポストゲートか）や種類（グループ化か通常か）でさえも、FDA や SWDE などの半構造化および非構造化データで構成されるタスクの精度に無視できない影響を与えます。

至る所にあるカーネル

Mamba-3 を使って人々が何を構築してくれるのか、楽しみにしています。これを促進するため、オリジナルの Mamba-2 Triton カーネルと速度が同等のカーネルをオープンソース化します。

次へ

パート 1 の最後までお読みいただきありがとうございます！カーネルや実験結果、アブレーションに関する詳細の多くは、この投稿では時間が足りずカバーしきれませんでしたが、ご安心ください！それらのすべては私たちの論文に記載されており、カーネルはmamba-ssmでオープンソース化されています。

次回のシリーズ第 2 弾（最終回）では、Mamba-3 への 3 つの中核的な改良とその SSM の基盤、そして私たちが特に興味を持っている方向性について掘り下げます。

参考文献

1. Mamba: Linear-Time Sequence Modeling with Selective State Spaces [PDF]
   Gu, A. and Dao, T., 2024.
2. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality [PDF]
   Dao, T. and Gu, A., 2024.
3. Gated Delta Networks: Improving Mamba2 with Delta Rule [PDF]
   Yang, S., Kautz, J. and Hatamizadeh, A., 2025.
4. Learning to (Learn at Test Time): RNNs with Expressive Hidden States [PDF]
   Sun, Y., Li, X., Dalal, K., Xu, J., Vikram, A., Zhang, G., Dubois, Y., Chen, X., Wang, X., Koyejo, S., Hashimoto, T. and Guestrin, C., 2025.
5. Hungry Hungry Hippos: Towards Language Modeling with State Space Models [PDF]
   Fu, D.Y., Dao, T., Saab, K.K., Thomas, A.W., Rudra, A. and Ré, C., 2023.
6. In-context Learning and Induction Heads
   Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds, Z., Hernandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S. and Olah, C., 2022. Transformer Circuits Thread.
7. RWKV: Reinventing RNNs for the Transformer Era [PDF]
   Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Grella, M., GV, K.K., He, X., Hou, H., Lin, J., Kazienko, P., Kocon, J., Kong, J., Koptyra, B., Lau, H., Mantri, K.S.I., Mom, F., Saito, A., Song, G., Tang, X., Wang, B., Wind, J.S., Wozniak, S., Zhang, R., Zhang, Z., Zhao, Q., Zhou, P., Zhou, Q., Zhu, J. and Zhu, R., 2023.
8. Test-time regression: a unifying framework for designing sequence models with associative memory [PDF]
   Wang, K.A., Shi, J. and Fox, E.B., 2025.
9. An Empirical Study of Mamba-based Language Models [PDF]
   Waleffe, R., Byeon, W., Riach, D., Norick, B., Korthikanti, V., Dao, T., Gu, A., Hatamizadeh, A., Singh, S., Narayanan, D., Kulshreshtha, G., Singh, V., Casper, J., Kautz, J., Shoeybi, M. and Catanzaro, B., 2024.