今回の論文
今回取り上げるのは、Tri Dao と Albert Gu による論文「Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality」です。2024年5月31日に arXiv で公開され、ICML 2024 に採択された研究です。公開元は arXiv / Proceedings of Machine Learning Research、URL は https://arxiv.org/abs/2405.21060、DOI は https://doi.org/10.48550/arXiv.2405.21060 です。研究分野は、状態空間モデル、attention、長文LLM、効率的な系列モデリングです。
この論文を選んだ理由は、Mamba 系の状態空間モデルを「Transformerの代替」として別物扱いするのではなく、attention と同じ数学的な土台で捉え直しているからです。理論的な整理にとどまらず、その見方から Mamba-2 という高速なアーキテクチャと、GPUで動きやすいSSDアルゴリズムまで導いている点が、開発や事業のヒントにつながります。
どんな技術か
Mamba-2 は、状態空間モデル（SSM）を使ったLLMアーキテクチャです。Transformer の self-attention が全トークン間の関係を直接計算するのに対し、SSMは系列を状態として順に更新しながら処理します。そのため、長い系列に対して理論上は線形にスケールし、生成時には巨大なKVキャッシュを持たずに済む可能性があります。
ただし、従来のSSMには弱点がありました。attention のようにGPUの行列積ユニットを使いやすい形ではなく、Mamba の selective scan も専用の実装を必要とします。また、attention で発展してきた multi-head、tensor parallel、sequence parallel のような設計語彙をそのまま移しにくい問題もありました。
この論文の中心にあるのが Structured State Space Duality、略して SSD です。SSDは、SSMの再帰的な計算を、semiseparable matrix という構造化行列の掛け算として見直します。すると、SSMは「再帰で線形に計算する形」と「attentionに近いブロック行列計算の形」の両方を持つことが分かります。この二重性を使うことで、長文に強いSSMの性質を保ちながら、GPUが得意な行列積で高速に学習・推論できるようにします。
課題
この技術が解こうとしている課題は、長い系列を扱うLLMで、性能・速度・メモリ効率を同時に満たすことです。
Transformer は非常に強力ですが、attention は系列長に対して二乗の計算が必要になります。推論時にも、過去トークンの key/value をKVキャッシュとして保持するため、長文になるほどメモリ使用量が増えます。RAGで大量文書を読む、コードベースを横断して解析する、長い対話履歴を持つエージェントを動かす、といった用途では、このコストがそのまま制約になります。
一方、SSMは系列長に対して線形にスケールし、生成時には状態だけを持てばよいという魅力があります。しかし、従来のSSMは言語のような情報密度の高いデータでattentionに劣ることが多く、Mamba はその弱点を selective SSM で改善したものの、計算の中心である scan がGPUの行列積ほど扱いやすいわけではありませんでした。
つまり課題は、「Transformerの表現力と実装エコシステム」と「SSMの長文効率」をどう接続するかです。単に新しいモデルを作るだけではなく、学習時の並列化、推論時の速度、既存のGPUクラスタでの動かしやすさまで含めて設計する必要があります。
用語解説
- 状態空間モデル（SSM）
- 入力を受け取りながら内部状態を更新し、その状態から出力を作る系列モデルです。この論文では、SSMを再帰計算だけでなく構造化行列として見直すため、SSMの「状態を通じて過去を圧縮する」性質を理解しておくことが重要です。
- Selective SSM
- Mambaで使われた、入力に応じてSSMのパラメータを変える仕組みです。固定されたフィルタで系列を処理するだけでなく、トークンごとに何を残すかを変えられるため、言語のような情報密度の高いデータで重要になります。
- Semiseparable Matrix
- 行列全体をそのまま持たず、特定の低ランク構造で効率的に表せる構造化行列です。この論文では、SSMの系列変換が semiseparable matrix と対応することを示し、SSDアルゴリズムの土台にしています。
- Structured State Space Duality（SSD）
- SSMを再帰的な線形時間計算としても、attentionに近いブロック行列計算としても扱えるという二重性です。Mamba-2では、このSSDによりSSMの長文効率とGPUの行列積効率を両立させます。
- Tensor Parallelism
- 巨大モデルの行列演算を複数GPUに分割する並列化手法です。Mamba-2はブロック内の投影を並列化しやすい形に整理しており、Transformerで発展した分散学習の考え方をSSMへ持ち込みやすくしています。
技術の仕組み
Mamba-2 の重要な点は、単にMambaの層を少し変えたことではありません。SSMを構造化行列として表し、計算方法とモデル設計を一緒に作り直したことです。
基本アイデア
通常のSSMは、時刻 t の状態 h_t を、前の状態 h_{t-1} と現在の入力 x_t から更新します。単純化すると、h_t = A_t h_{t-1} + B_t x_t、y_t = C_t h_t のような形です。ここで A_t、B_t、C_t が入力に応じて変わると、Mambaの selective SSM になります。
この再帰は、1トークンずつ処理すれば線形時間です。しかし学習では長い系列をまとめて処理したいので、逐次処理だけではGPUを十分に使い切れません。Mambaでは associative scan により並列化しますが、それでもTransformerのattentionやMLPのような大きな行列積ほどGPUに自然ではありません。
SSDは、このSSM計算を「系列方向の行列を入力に掛ける処理」として見ます。その行列は一般の密行列ではなく、semiseparable matrix という特殊な構造を持ちます。構造があるため、全要素を明示的に持たなくても効率的に計算できます。そして、ブロックに分解すると、対角付近は通常のブロック計算、ブロック間は低ランクな状態の受け渡しとして扱えます。
線形形式と二次形式を使い分ける
SSDの面白いところは、同じ計算に2つの見方があることです。
1つ目は、SSMらしい再帰形式です。系列を左から右へ流し、状態を更新します。この形は長い系列に対して線形にスケールし、生成時には状態だけを持てばよいため、KVキャッシュを大きくしない利点があります。
2つ目は、attentionに近い二次形式です。query、key、value に似たテンソルの組み合わせとして、ブロック内の関係をまとめて計算します。系列全体で見ると二次的な計算に見えますが、ブロック分解と低ランク構造を使うことで、GPUの行列積を活かせます。
論文の主張は、「SSMかattentionか」という二択ではなく、両者の共通部分を構造化行列として扱えば、計算方法を選べるというものです。長文では再帰的な効率を使い、ブロック内では行列積の効率を使う。この組み合わせがSSDアルゴリズムの核です。
Mamba-2のブロック設計
Mamba-2では、元のMambaブロックも分散学習しやすい形に整理されています。Mambaでは、入力から一部のパラメータを作り、その結果をさらに別の投影に使うような逐次的な構造がありました。Mamba-2では、A、B、C、X などのデータ依存パラメータをブロックの冒頭で並列に投影します。
この変更により、投影をまとめて実行しやすくなり、テンソル並列にも乗せやすくなります。また、attentionのheadに相当する考え方をSSMにも導入し、head dimension をTransformerに近い 64 や 128 の単位で扱います。論文では、Mambaを multi-input SSM と見なし、attentionでいう multi-value attention に近い構造として整理しています。
さらに、Mamba-2ではSSD層だけでモデルを組む構成に加えて、attention層やMLP層を混ぜるハイブリッド構成も検証しています。これは実務的に重要です。すべてをSSMに置き換えるのではなく、長文処理が得意なSSD層と、検索的な参照が得意なattention層、計算効率の良いMLP層を組み合わせる設計が可能になります。
推論方法
生成時には、Mamba-2はSSMの再帰形式を使えます。Transformerでは各層で過去トークンのKVを保存し、次トークン生成のたびに参照します。系列が長くなるほどKVキャッシュが増えるため、メモリが重くなります。
SSMでは、過去の情報は状態に圧縮されます。そのため、理想的には系列長に比例するKVキャッシュを持たず、固定サイズに近い状態で生成できます。もちろん、状態に圧縮するということは、過去トークンをそのまま取り出すattentionとは性質が違います。そこで論文では、純粋なMamba-2だけでなく、少数のattention層を混ぜると性能が上がることも示しています。
データの扱い方
実験では、主に The Pile を使って言語モデルを学習しています。小規模なスケーリング則の評価から、2.7Bパラメータ規模のモデルを300Bトークンで学習する比較まで行っています。評価では Pile の validation perplexity に加え、LAMBADA、HellaSwag、PIQA、ARC、WinoGrande、OpenBookQA などのゼロショットタスクを使っています。
この構成は、「新しい層が速い」だけでなく、「同じような学習条件で言語モデルとして本当に競争力があるか」を見るためのものです。
実験と結果
論文では、SSDアルゴリズムとMamba-2アーキテクチャを、合成タスク、言語モデリング、ゼロショット評価、速度ベンチマーク、アブレーションで検証しています。
Associative Recall
まず、多数のキーと値の対応を系列内に埋め込み、後で指定されたキーに対応する値を答える multi-query associative recall タスクで評価しています。このタスクは、長い系列の中から特定情報を保持・参照する能力を見るためのものです。
結果として、Mamba-2は状態サイズを大きくしたときにMambaやattention系の比較対象より強い性能を示しました。これは、SSDにより大きな状態サイズを現実的な速度で扱えるようになったことと関係しています。SSMは過去を状態に圧縮するため、状態容量が足りないと情報を落とします。Mamba-2はその状態を大きくしやすくした点が効いています。
The Pileでの言語モデリング
言語モデリングでは、Mamba-2、Mamba、Transformer++ を同じような条件で比較しています。論文では、Mamba-2が perplexity と wall-clock time の両方で Mamba と Transformer++ を Pareto dominate すると説明しています。つまり、同じ計算時間ならより低い perplexity、同じ perplexity ならより短い時間を狙える傾向が示されています。
2.7Bパラメータ、300Bトークン学習の比較では、Mamba-2は Pile validation perplexity で 6.09 を記録し、Transformer++ の 6.13 と同等以上でした。ゼロショット評価の平均では、Transformer++ と Mamba-2 はどちらも 60.2 で同水準です。純粋なSSMモデルが、同条件のTransformerに近い実力を示している点が重要です。
Attentionを少し混ぜる効果
興味深い結果として、Mamba-2に少数のattention層を混ぜると性能が改善しています。350Mモデル、48層の実験では、attention層なしのMamba-2が perplexity 8.60 だったのに対し、6層前後のattentionを入れると 8.26 付近まで改善しました。論文では、約10%程度のattention層がよいバランスだと示しています。
2.7B規模でも同様です。Mamba-2-Attention は Pile perplexity 5.95、ゼロショット平均 61.0 で、純粋なMamba-2の 6.09 / 60.2 と Transformer++ の 6.13 / 60.2 を上回りました。
この結果は、SSMが一般的な系列変換に強く、attentionが過去トークンへの直接参照や検索的な処理に効く、という役割分担を示唆しています。実務でモデルを設計する場合も、すべてを一種類の層で統一するより、タスクに応じて層を混ぜる発想が有効そうです。
速度ベンチマーク
速度面では、SSDアルゴリズムがMambaの最適化された selective scan より 2〜8倍高速だと報告されています。理由は、SSDがブロック分解によりGPUのtensor core、つまり行列積ユニットを活かせるためです。
また、SSDは系列長に対して線形にスケールするため、FlashAttention-2との比較でも、系列長が2Kを超えるあたりから高速になり、16Kでは大きな差が出るとされています。これは長文LLMでは特に重要です。短いチャットではTransformerの最適化が強い一方、長い文書やログを扱う場面では、系列長に対するスケールの違いが効いてきます。
ただし、論文は短い系列長ではMamba-2全体がTransformerより常に速いとは限らないとも述べています。Transformerはattention層とMLP層が混在し、MLPはGPUで非常に効率的です。一方、純粋なMamba-2は全層がSSDです。そのため、短文ではMLPを混ぜたハイブリッド設計のほうが実用的になる可能性があります。
何に使える？
Mamba-2は、長い系列を効率よく扱いたいAIシステムで使い道があります。特に、単に推論を速くするだけでなく、モデル構造そのものを長文向けに設計したい場合に参考になります。
長文LLMとRAG
RAGでは検索で候補文書を絞ったあと、複数の長い文書を読ませる段階でコンテキスト長とコストが問題になります。Mamba-2のようなSSM系モデルは、長い系列を線形に処理できるため、大量の検索結果や長い社内文書を扱うモデル基盤に向いています。
ただし、Mamba-2は過去情報を状態に圧縮するため、特定トークンを厳密に探す用途ではattentionの直接参照が有利な場合があります。論文のハイブリッド結果を踏まえると、RAG向けにはSSD層を中心にしつつ、少数のattention層で検索的な参照を補う構成が有望です。
長時間エージェント
AIエージェントでは、ユーザー指示、ツール実行結果、環境状態、途中判断が長く積み上がります。Transformerだけで長い履歴を抱えるとKVキャッシュが重くなります。Mamba-2のように状態で履歴を圧縮できるモデルは、長時間タスクの継続や、低メモリなエージェント実行に役立つ可能性があります。
これは論文からの応用推測ですが、Mamba-2にattention層を少し混ぜる結果を考えると、エージェントでも「普段の履歴圧縮はSSM、重要な過去イベントの直接参照はattentionや外部メモリ」という分担が考えられます。
コード・ログ・時系列データ
コードベース、ログ、センサーデータ、金融時系列、ユーザー行動履歴のようなデータは、長い系列として扱うほど文脈が効きます。Mamba-2は系列長に対するコストを抑えやすいため、こうしたデータをまとめて学習・推論するモデルの候補になります。
特にログ解析や監視では、直近だけでなく長い履歴のパターンが重要です。Transformerで全履歴attentionを取るのが重すぎる場合、SSM系の状態更新で長期傾向を持ち続ける設計が使えそうです。
軽量なオンデバイス・低コスト推論
生成時にKVキャッシュが膨らみにくい点は、GPUメモリが限られる環境でも魅力です。モバイル、エッジ、社内の小さなGPUサーバーなどでは、KVキャッシュが大きいモデルは運用しにくくなります。Mamba-2の考え方は、低メモリで長めの文脈を扱うモデル設計に応用できます。
開発や事業へのヒント
この論文から得られる大きなヒントは、AIモデルの改善を「新しい層を足す」だけでなく、数学的な表現、GPUでの計算、分散学習のしやすさまで一体で考えることです。
モデル設計は実装効率込みで考える
Mamba-2は、SSMを理論的に整理しただけでなく、SSDアルゴリズムとしてGPUの行列積に乗せています。プロダクト開発でも、精度だけでなく「この処理はGPU/CPU/DB/検索基盤のどこで効率よく動くか」を早い段階で考えると、スケール時の差が出ます。
ハイブリッド構成を恐れない
論文では、純粋なMamba-2よりも少数のattention層を混ぜたモデルが良い結果を出しています。これは、AIアプリ開発にも通じます。すべてをLLMの長文コンテキストに押し込むのではなく、検索、圧縮、状態管理、再ランキング、ツール実行を役割分担させるほうが強い場合があります。
長文処理は「圧縮」と「直接参照」の設計です
SSMは過去を状態に圧縮します。attentionは過去を直接参照します。RAGやエージェントを作るときも、この2つの発想を分けて考えると設計しやすくなります。頻繁に使う背景情報は状態や要約に圧縮し、正確に取り出す必要がある情報は検索やattentionに任せる、という分担です。
小規模プロダクトでも使える発想がある
Mamba-2を自前で事前学習するのは簡単ではありません。しかし、SSDの考え方から得られる「同じ処理を別の計算形式に変換すると速くなる」という視点は、小規模なAI機能にも使えます。たとえば、逐次的なLLM呼び出しをバッチ化する、長い履歴を状態オブジェクトに圧縮する、検索結果の一部だけを直接参照する、といった設計です。
限界
Mamba-2には注意点もあります。
まず、SSMは過去情報を状態に圧縮するため、すべての過去トークンを直接見られるattentionとは性質が違います。長い文書のどこかにある文字列をそのままコピーする、特定の過去発話を厳密に参照する、といったタスクでは、純粋なSSMだけでは不利になる可能性があります。論文でも、少数のattention層を混ぜると性能が改善することが示されています。
次に、実装の難しさがあります。SSDアルゴリズムは理論上きれいですが、高速に動かすには専用カーネル、ブロック分解、tensor parallel、sequence parallel との統合が必要です。既存のTransformer推論基盤にそのまま差し替えられる技術ではありません。
計算コストについても、短い系列では常に有利とは限りません。TransformerのMLP層やFlashAttentionは非常に最適化されています。Mamba-2は長文で強みが出やすい一方、短いチャットや小さなバッチでは、既存Transformerのほうが運用しやすい場合があります。
また、評価は主に言語モデリングと代表的なゼロショットタスクに基づいています。RAG、コード生成、長時間エージェント、マルチモーダル処理など、実アプリに近いタスクでどこまで強いかは、追加検証が必要です。
最後に、エコシステムの問題があります。Transformer向けには学習ライブラリ、量子化、推論サーバー、LoRA、評価ツールが豊富にあります。Mamba-2系モデルを本番投入するには、これらに相当する周辺技術の成熟も必要になります。
よくある質問
Q. Mamba-2はTransformerを置き換える技術ですか？
A. 完全な置き換えというより、長文効率を重視した有力な代替・補完技術です。論文では純粋なMamba-2がTransformer++と同等水準の性能を示し、少数のattention層を混ぜた構成ではさらに改善しています。用途によってSSM、attention、MLPを組み合わせる方向が現実的です。
Q. Mamba-2と元のMambaの違いは何ですか？
A. 元のMambaは selective SSM によって言語モデリング性能を高めました。Mamba-2は、SSMとattentionの関係をSSDとして整理し、SSM計算をGPUの行列積に乗せやすいブロックアルゴリズムへ作り直しています。ブロック設計もテンソル並列しやすい形に変更されています。
Q. なぜSSDで速くなるのですか？
A. SSMの再帰計算を semiseparable matrix の構造化行列計算として捉え、ブロック分解します。これにより、ブロック内の重い計算をGPUが得意な行列積として実行できます。論文では、Mambaの最適化されたscan実装より2〜8倍高速だと報告されています。
Q. RAGアプリに直接使えますか？
A. 既存のLLM API利用だけでは、Mamba-2の内部構造を直接使うことはできません。ただし、長文をすべてattentionで抱えるのではなく、状態として圧縮する部分と、検索で正確に取り出す部分を分ける設計思想はRAGに応用できます。将来的にMamba-2系の長文モデルを選択肢に入れる価値もあります。
Q. Mamba-2の弱点は何ですか？
A. 過去情報を状態に圧縮するため、正確なコピーや特定位置への直接参照ではattentionに劣る可能性があります。また、高速実装には専用カーネルや分散学習の知識が必要です。短い入力ではTransformerの最適化実装のほうが扱いやすい場面もあります。
今日の学び
この論文は、長文LLMでTransformerのattentionが抱える二乗計算とKVキャッシュの重さ、そしてSSMが抱える実装効率と設計語彙の不足を扱いました。そこで、SSMを semiseparable matrix として捉える Structured State Space Duality を導入し、再帰的な線形計算とattentionに近いブロック行列計算をつなげました。
そこから生まれたMamba-2は、長文に強いSSMの性質を保ちながら、GPUの行列積を活かして高速に動くアーキテクチャです。開発や事業へのヒントは、長文AIを作るときに「圧縮する処理」と「直接参照する処理」を分け、モデル構造・計算方式・インフラを一体で設計することです。