TransformerはAI大規模モデルの分野を支配してきました。AGI(汎用人工知能)の実現にはTransformerだけで十分だと考える人もいます!
一方、何度かの基盤アーキテクチャの革新が必要だと考える人もいます。例えば、ファーウェイの「インテリジェントワールド 2035」レポートにもそのような記述があります。
モデルの訓練、デプロイ規模、計算需要が爆発的に増加するにつれて、計算能力とエネルギー需要は底なし沼のようになっていることがわかります。AIを賢く、安価に、そして高速に動作させるにはどうすれば良いのでしょうか?
Transformerアーキテクチャの計算複雑性は、シーケンス長の二乗に比例します。処理するテキストの長さが2倍になると、計算量は4倍になります。同時に、メモリ占有量もシーケンス長に応じて線形に増加します。これを用いてAGIを実現するには、エネルギーと計算能力の需要は想像を絶します。
学術界と産業界は解決策を模索しています。
そこでMambaが登場しました。Mamba-3は現在、トップカンファレンスであるICLR 2026でダブルブラインド審査を受けています。
MambaはTransformerの枠組みを修正するのではなく、別の道を選び、より古く、より基礎的な理論である状態空間モデル(SSM)からインスピレーションを得ました。
Mamba-3の物語は、本質的に効率と知性のトレードオフと進化に関するものです。最も実用的で致命的な問題である推論効率から出発し、3つの改善を導入しました。それは、より表現力豊かな再帰的アプローチ、よりスマートな状態更新ルール、そしてハードウェア性能を最大限に引き出すマルチ入力マルチ出力(MIMO)アーキテクチャです。
Mamba-3は具体的にどのような新しい工夫を凝らし、Transformerとどう競合するのでしょうか?
万丈のビルは平地に建つ
Mamba-3の巧妙さを理解するには、まず状態空間モデル(SSM)とは何かを話す必要があります。
この概念は元々、自然言語処理(NLP)のために用意されたものではなく、回路内の信号、気象雲図の変化、運動体の飛行軌跡など、連続的に変化するシステムを予測するのが本来の役割でした。数学的・概念的には、リカレントニューラルネットワーク(RNN)と近縁であり、シーケンスを一歩ずつ処理し、「記憶」を保持して次の決定に影響を与えるモデルです。
名前が示す通り、SSMの核は「状態空間」です。これを、システムが持つすべての主要変数が含まれた、現在の全状況を記述するスナップショットだと想像してください。SSMの役割は、入力シーケンス x(t)を受け取り、それを目に見えない潜在状態 h(t)にマッピングすることです。この h(t)はRNNの隠れ状態のようなものであり、この状態に基づいて出力 y(t)が予測されます。
すべてのSSMは、以下の2つの核となる方程式を中心に機能します:
状態方程式:
h'(t)=A*h(t)+B*x(t)出力方程式:
y(t)=C*h(t)+D*x(t)
ここでのA、B、C、Dの4つのパラメータは通常、システムのダイナミクスを定義する重み行列です。制御理論のような伝統的な分野では、これらの行列は固定されており、既知のシステムを表します。一方、ディープラーニングでは、これらは訓練を通じて最適化される必要のあるパラメータとなり、ニューラルネットワークの学習可能な重みによって表現されます。
SSMの古典的な形態は連続信号向けに設計されていますが、我々が扱うテキスト、画像、音声はコンピュータ内では離散データです。そのため、「離散化」のステップが必要です。
離散化とは、連続的に流れる信号を周期的にサンプリングすることだと想像できます。このプロセスでは、ステップサイズ(Δ)と呼ばれる新しいパラメータが導入され、サンプリング頻度を決定します。離散化の方法はいくつかありますが、Mambaを含むほとんどの最新SSMは、ゼロ次ホールド(ZOH)と呼ばれる単純な方法を使用しています。
離散化を経ることで、SSMはRNNのようにシーケンスデータを処理できるようになります。
しかし、初期の離散SSMは実用的ではありませんでした。なぜなら、RNNのいくつかの欠点、例えば訓練効率の低さや、シーケンス内で遠く離れた2つの要素間の関係、いわゆる「長距離依存性」の問題を捉えにくいという記憶力の悪さを継承していたからです。
転機は2021年に訪れました。Albert Gu氏と彼の同僚たちが、Structured State Space Sequence Model(略称 S4)を提案したのです。この研究が後のMambaへの道を開きました。
S4モデルは2つの大きな成果を達成しました。
一つ目は、畳み込みによる効率的な訓練の実現です。離散SSMは推論時にはRNNのように高速ですが、訓練には非常に時間がかかります。S4の著者らは、SSMが線形な加算と乗算の操作のみを含むため、この一連の再帰操作が一次元の畳み込みカーネルとして展開できることを発見しました。この畳み込みカーネルは、入力シーケンス x を出力 y へと一発でマッピングできます。そして、畳み込み演算は「高速フーリエ変換」と呼ばれるアルゴリズムを用いて効率的に計算できます。
これにより、素晴らしい利点がもたらされました。訓練時、入力シーケンス全体が既知である場合、S4は畳み込みニューラルネットワーク(CNN)のように並列的かつ効率的に計算できます。一方、推論時には、トークンを一つずつ生成する必要がある場合、RNNの形態に戻り、極めて高速な速度と極めて低いメモリ占有量という恩恵を享受できるのです。一石二鳥です。
二つ目は、構造化行列による長期記憶問題の解決です。SSMがより過去の情報を記憶できるようにするため、S4は従来の機械学習モデルのように重み行列 AとBをランダムに初期化するのではなく、HiPPOと呼ばれる技術を採用し、特殊な直交多項式(例:ルジャンドル多項式)から行列の構造を導き出しました。この特殊な初期化方法は、モデルに記憶力強化プラグインを搭載したかのように働き、長大なシーケンスを扱う際の性能を急上昇させました。
その後のS4の変種、例えばDSS、S5、そして今日の主役であるMambaシリーズは、具体的な初期化方法には違いがあるものの、HiPPOの核となる思想を保持しています。それは、行列 AとBに何らかの構造(通常は対角構造)を課すことで、モデルが安定して状態を更新し、長期的な依存関係を記憶できるようにすることです。
Mambaの進化の道のり
2023年、Tri Dao氏とAlbert Gu氏(また彼です)は、論文「Mamba: Linear-Time Sequence Modeling with Selective State Spaces」で初めてMambaアーキテクチャを提案しました。これは、言語モデリングにおいてTransformerと正面から対抗できる最初のアーキテクチャでした。
Mambaの核となる革新は2つあります。
一つ目は、「選択的状態空間モデル」です。これは従来のSSMに「選択」スイッチを取り付けたもので、モデルは現在の入力の重要性に基づいて、どの履歴情報を記憶し、どの情報を無視するかを動的に決定できます。この能力は、これまでTransformerの自己注意メカニズムの専売特許だと考えられていました。
二つ目は、「ハードウェア認識型並列スキャン」です。これは非常にエンジニアリング指向の最適化であり、現代のグラフィックス処理ユニット(GPU)の計算特性に特化し、SSMの再帰計算を処理するための高効率なアルゴリズムを設計し、ハードウェアリソースを最大限に活用します。
1年後、同じ著者らが再び論文を発表し、SSMとTransformerの間の深層的な繋がりをさらに探求し、より高速で強力な改良版、Mamba-2を提案しました。
Mamba-2は、SSMの大部分の計算プロセスが、マスク付き行列乗算と等価であることを発見しました。この発見により、Mamba-2は行列乗算の高効率な実装を利用でき、訓練速度をMamba-1よりも50%向上させました。また、より大きな状態次元をサポートし、特に長大なシーケンスにおいて、モデルがより複雑なタスクを処理できるようにしました。
そして今、物語はMamba-3へと発展しました。
Mamba-3はMamba-2に基づいて、推論効率という観点から、さらなる進化を遂げました。これにより、3つの核となる方法論的な改善がもたらされました。
一つ目は「台形離散化」です。これは、以前のMamba-2の比較的粗い方法(オイラー法)を、より正確な数学的手法(台形法則)に置き換え、連続信号から離散シーケンスへの変換を完了させます。この改善により、モデルの再帰的更新がより表現力豊かになりました。
二つ目は「複素状態空間モデル」です。複素数を導入してSSMを定義することにより、モデルの状態更新能力が大幅に強化され、多くの線形モデルが苦手とする、正確な状態追跡が必要なタスク(例えば、数の偶奇性の判断)における能力不足の問題を解決しました。
三つ目は「多入力多出力SSM」です。これは、デコード速度とハードウェア効率を向上させるためだけに設計されました。元々外積に基づいていた状態更新を、行列乗算に基づく更新に変更することで、計算の「算術強度」を大幅に高め、GPUが「飢餓状態」になるのを防ぎます。
Mamba-3の新しい工夫
より正確な離散化:台形法則
構造化SSMは理論的には連続時間システムとして定義されますが、実際に処理されるデータはすべて離散的です。連続から離散への変換、すなわち離散化が重要なステップです。
Mamba-2が使用したのはオイラー法です。これは、区間の終点値のみを考慮し、矩形の面積で曲線の下の面積を近似すると想像できます。この方法の誤差は O(Δt²) であり、単純ですが精度は十分ではありません。
Mamba-3は、より高度な方法である一般化台形法則を採用しました。これは、単に矩形で近似するのではなく、区間の始点と終点の両方を考慮し、データ依存の凸結合でそれらを接続する台形を使用します。この方法の誤差は O(Δt³) に低下し、精度は桁違いに向上しました。
Mamba-3は状態更新を行う際、現在のタイムステップの入力だけでなく、一つ前のタイムステップの入力も振り返ります。この小さな「振り返り」により、モデルのシーケンスダイナミクスを捉える能力が、より繊細で強力になりました。
この改善はモデルの表現力を高めるだけでなく、Mamba-3が以前の多くの線形モデルが依存していたコンポーネント、すなわち短期因果畳み込みを必要としなくなりました。これにより、モデル全体のアーキテクチャがより簡潔で統一されました。
より賢い状態更新:複素数と回転
現代のSSMは効率を追求するために、核となる状態遷移行列を常に単純化してきました。S4モデルは複雑な「正規行列+低ランク」行列を使用していましたが、Mambaはこれを実数の対角行列に単純化し、Mamba-2ではさらにスカラーにまで単純化されました。これらの単純化は言語モデリングタスクにおいて顕著な性能低下をもたらしませんでしたが、いくつかの単純な状態追跡タスクにおけるモデルの能力を弱めてしまいました。
例えば、二進数シーケンス中の「1」の数が奇数か偶数かを判断するタスク(偶奇性タスク)です。このタスクは単層のLSTM(Long Short-Term Memory)にとっては簡単なことですが、状態遷移行列が実数固有値しか持たないMamba-2にとっては、達成が非常に困難でした。
その理由は、実数固有値は状態の「伸縮」と「反転」のみを表現でき、「回転」を表現できないからです。偶奇性のようなタスクでは、その内在的な状態遷移は周期性があり、「オン」と「オフ」を切り替えるスイッチのように動作します。これは数学的に回転として表現するのが最も自然です。
Mamba-3はこの能力を取り戻すために、複素数を導入しました。
Mamba-3は、複素数値のSSMが離散化された後、状態次元が2倍になった実数値SSMと等価であり、その状態遷移行列が一連の 2x2 回転行列ブロックで構成されることを証明しました。
さらに、この回転操作は、入力および出力の射影行列 BとC に等価的に「吸収」できることも証明されました。これは最終的に驚くべき結論へと導かれます。複素数SSMを使用することは、通常の、スカラー遷移に基づくSSMの入力(B)と出力(C)に、データ依存の回転位置埋め込み(RoPE)を適用することと等価である、というものです。
RoPEは多くの大規模モデル(例:Llama)で利用されており、単語ベクトルに絶対的または相対的な位置情報を注入することで、モデルが語順を理解するのを助けます。Mamba-3がここで実現したのは、RoPEを「データ非依存」で固定された位置エンコーディングから、「データ依存」の動的な状態回転器へと変えることです。
著者らが「RoPEトリック」と呼ぶこの実装により、Mamba-3は極めて少ない計算オーバーヘッドで強力な状態追跡能力を獲得し、Mamba-2では達成できなかった偶奇性やモジュロ演算などのタスクを容易に解決できるようになりました。
究極のハードウェア効率:外積から行列乗算へ
自己回帰生成(つまり、単語を一つずつ出力する)のシナリオでは、性能は通常、1秒あたりの生成トークン数(TPS)で測定されます。この指標において、Mambaのようなモデルは、Transformerのようにシーケンス長に応じて線形に増加するKVキャッシュを維持する必要がなく、固定サイズの隠れ状態を持つため、本質的な優位性があります。
しかし、TPSという指標はハードウェア効率を考慮していません。より基盤となる指標は「算術強度」であり、これは一度の操作における浮動小数点演算回数(FLOPs)と、そのために移動されるデータバイト数の比率として定義されます。
現代のGPUは、超強力な計算工場のようなもので、その計算能力(ops)はデータ転送能力(byte)を遥かに上回っています。もし算術強度が低すぎると、GPUは実際に計算を行うのではなく、メモリからデータを転送してくるのを待つために大量の時間を浪費します。この状況を「メモリ律速」と呼びます。
Mamba-2の状態更新は外積操作です。その算術強度は定数であり、現代のGPUの理想値よりも遥かに低いです。これは、デコード時にMamba-2がGPUの性能を十分に発揮できないことを意味します。
Mamba-3は、一見単純だが非常に効果的な変更を行いました。状態更新を外積から行列乗算へと変更したのです。
これは、信号処理の文脈では、単入力単出力(SISO)システムから多入力多出力(MIMO)システムへの一般化に正確に対応しています。
MIMOの定式化の下では、算術強度は新しく導入されたランク r に比例します。rのサイズを調整することで、算術強度を柔軟に高めることができ、デコードプロセスを「メモリ律速」から「計算律速」へと推移させ、それによってハードウェアをより十分に活用し、より高いTPSを獲得できます。このプロセスは、推論時のメモリ占有量(状態 H のサイズ)を増やしません。
これら3つの主要な改善が組み合わさり、Mamba-3の核となるミキサー(Mixer)プリミティブを構成しています。Mamba-3全体のアーキテクチャもいくつかの調整が行われ、Mamba-3ブロックとSwiGLUブロックを交互に使用し、事前正規化が採用されています。
アーキテクチャ性能の比較
言語モデリングの性能に関して、論文の著者らはFineWeb-Eduデータセットの1000億トークンを使用し、Mamba-3、そしてTransformer、Gated DeltaNet、Mamba-2などのベースラインモデルを、180M、440M、820M、1.5Bの4つの異なるパラメータ規模で事前訓練しました。
結果は、すべてのモデル規模で、Mamba-3があらゆる下流タスクにおいてベースラインモデルを全面的に上回っていることを示しています。
検索能力、つまり長文の中から情報を正確に検索する能力においては、Transformerはすべての履歴情報にロスレスにアクセスできるKVキャッシュメカニズムのおかげで、依然として優位性を持っています。これは、固定状態サイズのモデルすべてに共通する弱点です。
実験では、Mamba-3は関連記憶や質問応答といったタスクでは良好なパフォーマンスを示しましたが、半構造化データや非構造化データから情報を抽出する必要があるタスクではパフォーマンスが芳しくありませんでした。しかし、合成された「干し草の山から針を見つける」(NIAH)タスクでは、Mamba-3のパフォーマンスはベースラインを上回るか同等であり、Mamba-2よりも優れた汎化能力を示しました。
推論効率:
一般的な bf16 精度および状態次元 128 の設定において、Mamba-3のSISOバージョンとMIMOバージョンはいずれも、Mamba-2やGated DeltaNetよりも高速です。
このグラフは、Mamba-3の優位性をより直感的に示しています。横軸は状態サイズ(推論速度の代理として、小さいほど速い)、縦軸は事前訓練のパープレキシティ(モデル性能の代理として、低いほど良い)です。Mamba-3 MIMOバージョンは、状態サイズを増やさず(つまり速度を犠牲にせず)に、性能-効率のパレートフロンティアをさらに一歩前進させました。
最後に、アブレーション実験により、Mamba-3の各改善点の有効性が検証されました。
台形離散化と導入されたバイアス項は相乗効果を発揮し、モデル性能を著しく向上させました。また、状態追跡タスクでは、RoPEを持つMamba-3は偶奇性およびモジュロ演算タスクをほぼ完璧に解決しましたが、RoPEを持たないMamba-3やMamba-2はランダムな推測と大差ない結果となりました。
Mamba-3の物語は、計算効率とモデル能力の間で、より優れたバランスを見つけるための探求に関するものです。
ロスレスな記憶と正確な検索を必要とする長文タスクにおいて、固定サイズの状態記憶メカニズムは、依然としてTransformerと比較した際のMambaの弱点です。著者らも、Mamba-3と外部検索メカニズムを組み合わせ、ハイブリッドアーキテクチャを構築することが、将来の重要な方向性の一つであると認めています。
あなたはMamba-3がTransformerに取って代わると思いますか?それとも有益な補完関係になると思いますか?
参考文献:
https://openreview.net/pdf/a4e02db9a98e8b5cb40d677e00e4c8017a282772.pdf
https://openreview.net/forum?id=HwCvaJOiCj
https://www.ibm.com/think/topics/state-space-model
https://www.ibm.com/think/topics/mamba-model