今回の論文
今回取り上げるのは、Albert Gu、Tri Dao による 2023 年の論文「Mamba: Linear-Time Sequence Modeling with Selective State Spaces」です。公開元は arXiv で、その後 COLM 2024 に採択されています。研究分野は、系列モデリング、基盤モデルアーキテクチャ、状態空間モデルです。URL は https://arxiv.org/abs/2312.00752 です。
この論文を選んだ理由は、Transformer の代替候補として話題になっただけでなく、「長い入力をもっと安く扱いたい」「でも単純な軽量化では attention の強みが消える」という難題に、かなり筋の良い答えを出しているからです。RAG の長文コンテキスト、ログ解析、音声や時系列処理など、開発現場で応用の想像がしやすい技術です。
どんな技術か
Mamba は、長い系列を線形時間で処理できる状態空間モデル、つまり State Space Model をベースにしながら、入力内容に応じて「何を覚え、何を忘れるか」を動的に変えられるようにした技術です。
わかりやすく言うと、Transformer の self-attention は「全トークン同士を毎回見比べる」ことで強い表現力を得ていますが、そのぶん入力が長くなると計算量とメモリ使用量が急増します。Mamba はそこを、固定長の状態を順番に更新する方式へ寄せつつ、単なる RNN のように機械的に流すのではなく、現在の入力に応じて状態更新の仕方そのものを変えることで、内容依存の処理能力を持たせています。
つまり Mamba は、「attention を使わずに、できるだけ attention 的な振る舞いを取り戻す」ためのアーキテクチャです。特に長文や長時間系列を扱う基盤モデルの設計として重要な論文です。
課題
この技術が解決しようとしている課題は、長い系列を高精度に扱いたいのに、Transformer では計算コストが重すぎることです。
何が難しいのかというと、言語やコードのような離散的で情報密度の高いデータでは、「どの過去情報が重要か」を入力ごとに選び分ける必要があることです。単に古い情報を圧縮して覚えるだけでは足りず、たとえば前半に出た識別子、参照関係、固有名詞、条件分岐などを、必要なタイミングで引き出せなければなりません。
既存手法ではここに限界がありました。線形 attention、畳み込み系モデル、従来の状態空間モデルは、計算量の面では有利でも、内容に応じたルーティングが弱く、特にテキストのような離散系列で Transformer ほど安定して強くありませんでした。論文では、この弱点を「content-based reasoning の不足」と捉えています。
なぜこの課題を解く必要があるかというと、実際の AI システムはどんどん長い入力を扱う方向へ進んでいるからです。RAG では多数文書をつなげますし、エージェントでは長い履歴やツール実行ログを持ちます。音声認識やセンサ時系列、ゲノム解析でも、系列長の増加はそのまま性能やユースケースの拡張に直結します。長文に強いだけでなく、長くなっても計算が破綻しにくいバックボーンが必要です。
実務では、長いプロンプトによるレイテンシ増加、GPU メモリ圧迫、バッチ処理効率の低下、長履歴での性能劣化として現れます。Mamba はこの問題に対して、モデル構造そのものから別解を出した技術です。
用語解説
- 状態空間モデル(SSM)
- 系列を順番に読みながら、内部状態を更新していくモデルです。Mamba はこの枠組みを使うことで、系列長に対して線形にスケールする計算を実現しています。この記事では、attention を使わず長文を扱う土台として重要です。
- Selective SSM
- Mamba の核になる考え方で、状態更新に使うパラメータを入力依存に変える仕組みです。従来の SSM は更新則が固定的でしたが、Selective SSM によって「今のトークンは強く覚える」「この情報は流す」といった選別が可能になります。
- 線形時間計算量
- 入力長を 2 倍にしたとき、計算量もおおむね 2 倍に増える性質です。Transformer の self-attention は一般に系列長に対して二乗で重くなりますが、Mamba はこの点を改善し、長文処理で有利になります。
- 並列スキャン(parallel scan)
- 再帰的な状態更新を GPU 向けに並列化する実装上の工夫です。Selective SSM は単純な畳み込みに落とせないため、そのままでは遅くなりがちです。Mamba は並列スキャンとカーネル融合を使って、線形時間の理論を実効性能につなげています。
- induction heads
- 過去に見た対応関係を再利用する能力を測る代表的な synthetic task です。Transformer の in-context learning を説明する文脈でよく出てきます。Mamba 論文では、この種の選択的コピー課題を解けることが、内容依存の推論能力を示す重要な裏付けになっています。
技術の仕組み
Mamba の提案は、「SSM をもっと賢くし、そのうえで GPU で速く動かす」という二段構えです。単なる理論モデルではなく、アーキテクチャと実装の両方をセットで設計しています。
基本アイデア
従来の SSM は、系列に沿って状態を更新しながら情報を運びます。この形式は長い系列に強い一方で、更新ルールが固定だと「入力内容に応じてどこを見るか」を変えにくい欠点がありました。Mamba はここに対して、状態更新に関わるパラメータ Δ, B, C を入力 x の関数にします。
これにより、同じ位置の更新でも、現在のトークンが何かによって挙動が変わります。重要なトークンなら強く状態へ書き込み、不要なトークンなら影響を小さくする、といった選別が可能になります。論文ではこれを selection mechanism と呼んでいます。
なぜ「選択」が必要なのか
論文では selective copy や induction heads のような課題を使って、固定的な更新では難しい例を示しています。たとえば、系列中の特定のマーカーの後ろにある値だけを覚えるには、「何が来たか」に応じて記憶動作を切り替えなければなりません。
attention はこの切り替えを比較的自然に行えますが、従来の線形時間モデルはそこが弱かったです。Mamba は、選択的な状態更新を入れることで、そのギャップを埋めようとしています。
モデル構造
Mamba ブロックは、従来の Transformer ブロックとはかなり見た目が違います。self-attention 層や通常の MLP 層をそのまま積むのではなく、入力を投影し、短い畳み込みと selective SSM を通し、ゲーティング付きで出力へ戻す構成です。
ポイントは、attention を完全に消していることです。しかも論文では、Mamba ブロックを均一に積むシンプルな構成でよい結果を出しています。これは、長文性能のために特殊な補助層を足したというより、バックボーン自体を置き換えうることを意味します。
Selective SSM の中身
状態空間モデルは、ざっくり言えば内部状態 h_t を前時刻 h_{t-1} と現在入力 x_t から更新し、そこから出力 y_t を作る仕組みです。通常はこの更新則が時間不変に近いのですが、Mamba では更新の粒度を決める Δ と、入力から状態・状態から出力への関わりを決める B, C を入力依存にします。
この変更の意味はかなり大きいです。固定係数の SSM は「常に同じルールで流すパイプ」に近いですが、Mamba は「入力ごとに弁の開き方が変わるパイプ」になります。そのため、テキストのように重要箇所が疎に現れるデータでも、必要な情報を選んで長く保持しやすくなります。
学習方法
学習自体は通常の次トークン予測や系列モデリングの枠組みで行えます。特別な教師信号を追加するというより、アーキテクチャ側で長距離依存を扱いやすくして、通常の pretraining から性能を引き出す設計です。
重要なのは、selection mechanism が end-to-end で学習されることです。どのトークンで状態を強く更新するか、どこで過去情報を維持するかは、人手ルールではなく学習で獲得されます。このため、言語だけでなく音声やゲノムのような別モダリティにも展開しやすい構造になっています。
推論方法
推論時は再帰的に状態を更新していくため、自己回帰生成との相性が良いです。attention の KV キャッシュのように系列長に応じて履歴全体を保持し続ける必要が薄く、理論上は一定サイズの状態を持ちながら順次生成できます。
これが長いコンテキストで効いてきます。論文では Mamba が Transformer より最大 5 倍高い推論スループットを示したと報告しています。長文を読む・長文を生成するタスクでは、この差がそのままコストや待ち時間に効きます。
実装上の重要な工夫
Selective SSM は入力依存になったぶん、従来 SSM のように畳み込みへ単純変換できません。ここが Mamba の難所です。論文ではこの問題に対して、kernel fusion、parallel scan、recomputation を組み合わせた hardware-aware algorithm を設計しています。
要するに、「理論上は線形時間でも GPU で遅ければ意味がない」ので、再帰計算をハードウェアに合わせて再構成したわけです。Mamba の実務的な価値は、この実装込みで成立しています。
実験と結果
論文では、Mamba が本当に Transformer の代替になりうるかを、合成課題、言語、音声、ゲノムなど複数の観点から検証しています。単に 1 つのベンチマークで勝ったというより、「長い系列をうまく扱える一般的なバックボーンか」を見ています。
何を検証したのか
主な検証は 3 つあります。1 つ目は、selection mechanism によって内容依存の推論が可能になるかです。2 つ目は、言語モデリングで Transformer 系や既存の線形時間モデルと比べてどこまで戦えるかです。3 つ目は、音声やゲノムのような長い系列でも性能が伸びるかです。
どんなデータセットや評価指標を使ったのか
言語では The Pile を使った事前学習で perplexity を比較し、下流評価として LAMBADA、HellaSwag、PIQA、ARC-E、ARC-C、Winogrande などの zero-shot ベンチマークを使っています。加えて、selective copy や induction heads の synthetic task で長距離依存の能力も確認しています。音声では SC09、ゲノムでは HG38 系のデータで評価しています。
言語モデリングでは同規模 Transformer を上回った
論文の要点のひとつは、Mamba が attention-free でありながら、言語モデリングでかなり強いことです。著者らは、Mamba が同規模の Transformer を上回り、3B 級モデルではおおむねその 2 倍規模の Transformer に近い性能だと報告しています。
具体的には、zero-shot 評価の平均スコアで Mamba-2.8B は 63.3 を記録し、Pythia-2.8B の 59.1 を明確に上回りました。さらに Pythia-6.9B の 61.7 も超えており、同論文の主張どおり「同規模比較で強く、より大きい Transformer 級にも届く」傾向が見えます。
長い系列ほど有利になった
The Pile の scaling law 実験では、Mamba は他の attention-free モデルより良くスケールし、系列長 2048 と 8192 の条件で強い結果を示しています。論文では、長くなるほど Transformer++ に迫る、あるいは競合しやすい傾向が示されています。
これは実務上かなり大事です。短文のベンチマークだけで強いのではなく、系列長が伸びたときに優位性が出るなら、RAG や長履歴エージェントの基盤として意味があります。
induction heads では 100 万トークン超まで外挿できた
合成課題の induction heads では、学習長 256 のモデルを、テスト時に 64 から 1,048,576 トークンまで伸ばして評価しています。論文では、Mamba がこの課題を解けるだけでなく、極端に長い系列へも外挿できることを示しました。
もちろん synthetic task だけで実務性能は決まりません。ただし、「選択的に記憶して再利用する」という性質が本当にモデル内で実現されていることの強い傍証になります。
音声やゲノムでも有効だった
Mamba は言語専用モデルではありません。論文では音声波形生成や DNA 系列でも先行法を上回る結果が示されており、長い連続系列や生体配列のような領域でも有望です。基盤モデルのバックボーンとして、モダリティ横断で使える可能性がある点は大きな意義です。
結果から何が言えるのか
この結果から言えるのは、Mamba は「Transformer を少し軽くした亜種」ではなく、長い系列を安く扱いながら内容依存の推論もある程度確保する、新しい設計軸だということです。特に、長文でのスケーリングと推論効率を重視するなら、今でも学ぶ価値の高い論文です。
何に使える?
Mamba の価値は、単なる研究トレンドではなく、「長い系列を扱う AI システムをどう作るか」という現実の設計判断にあります。
長文 RAG や長履歴エージェントの基盤
大量の文書断片、会話履歴、ツールログを抱えるエージェントでは、コンテキスト長がすぐ肥大化します。Mamba 系の考え方は、こうした長い履歴を二乗コストで処理し続ける限界に対する別解になります。現時点でそのまま商用 LLM を置き換えるのは慎重に見るべきですが、長履歴メモリ部品や補助エンコーダとしては十分検討余地があります。
ストリーミング音声やセンサ時系列
音声、IoT、監視ログ、株価、機械稼働データのように、入力が途切れず流れる系では、Mamba の再帰的な状態更新が扱いやすいです。過去全体への attention を毎回計算しなくてよいため、低遅延処理や長時間監視と相性が良いです。
コード補完や大規模ログ解析
コードベース全体の文脈や長い実行ログを読む用途では、遠く離れた依存関係を安く保持できることが効きます。特に、ログ異常検知や開発支援では「すべてを完全比較する」より、「重要な状態だけ選択的に維持する」方が現実的な場面があります。
マルチモーダルの長系列エンコーダ
動画フレーム列、長時間音声、時系列センサ、医療波形などでは、系列長がすぐ巨大になります。Mamba はモダリティ横断で使える可能性が示されているので、長系列側のエンコーダとして採用する余地があります。特に Transformer だけではコストが合わない部分の代替候補です。
開発や事業へのヒント
この論文から得られるヒントは、「性能を上げるには attention を増やすしかない」という発想を崩せることです。プロダクト設計にもそのまま応用できます。
長文処理はモデル選定の基準を変えるべき
AI アプリを作るとき、短いベンチマークの精度だけでモデルを選ぶと、長文入力時のレイテンシとコストで苦しくなります。Mamba が示しているのは、アーキテクチャ次第で長文時のコスト構造を変えられるということです。長文ワークロードが中心なら、精度だけでなく系列長に対する計算量を最初から評価軸に入れるべきです。
「全部参照」ではなく「重要情報だけ残す」発想
selection mechanism の本質は、全履歴を均等に扱わないことです。この考え方は、RAG の chunk 圧縮、会話メモリの要約更新、ログ監視の状態保持などにも応用できます。プロダクト全体でも、「毎回全部見る」より「状態を持ち、更新を選別する」設計が効く場面は多いです。
小規模開発でも差別化ポイントになる
Mamba をゼロから大規模事前学習するのは簡単ではありませんが、既存 OSS 実装を使って長系列分類やエンコーダ用途に試すことは現実的です。長文処理コストが重い業務システムでは、バックボーン変更だけで運用コストの差が出る可能性があります。
今後注目すべき方向性
今後は、Transformer か Mamba かの二択ではなく、役割ごとに構造を分ける流れが強まると考えられます。たとえば、短距離の精密参照は attention、長距離の保持は state space、というハイブリッド構成です。論文でも Mamba 単体の強さが示されていますが、実務ではこの設計思想をどう組み合わせるかが重要になります。
限界
Mamba にも明確な限界があります。まず、attention を完全に置き換えられるかはタスク次第です。論文では強い結果が出ていますが、すべての知識集約タスクや複雑な参照推論で常に Transformer より優れるとまでは言えません。
また、実装は見た目ほど単純ではありません。Selective SSM を高速に動かすには並列スキャンや専用カーネルが重要で、理論上の線形時間がそのまま誰の環境でも出るわけではありません。実装依存性やハードウェア依存性は注意点です。
データ依存性もあります。論文の付録では、selection mechanism が常に有利とは限らず、長い音声波形では非選択的な SSM の方が良い場面も示されています。つまり、「選択的にすれば何でも強くなる」わけではなく、モダリティの性質に合うかを見る必要があります。
さらに、実運用では既存エコシステムも無視できません。Transformer はツール、最適化、学習済みモデル、推論基盤が非常に充実しています。Mamba 系は今後広がる余地がありますが、導入コストや検証コストはまだ高めです。
よくある質問
Q. Mamba は Transformer を置き換える技術ですか?
A. 一部の用途では有力な代替候補ですが、全面置換と見るのは早いです。特に長い系列やストリーミング処理では強みがありますが、既存の学習済み資産や周辺ツールの成熟度はまだ Transformer が優位です。実務では、長系列部分に限定して試すのが現実的です。
Q. Mamba の何が従来の状態空間モデルと違うのですか?
A. 最大の違いは、状態更新のパラメータを入力依存にしている点です。従来の SSM は更新則が固定寄りで、内容に応じた選別が弱い傾向がありました。Mamba はこの制約を崩し、必要な情報を残しやすくしています。
Q. RAG に直接使えますか?
A. 直接 LLM 本体として置く方法も考えられますが、現実的にはまず長文エンコーダや履歴圧縮モジュールとして試す方が導入しやすいです。長い入力を安く処理したい部分に Mamba 的な構造を入れる、という使い方が有望です。
Q. 生成時の高速化にも効きますか?
A. 効く可能性があります。論文では Transformer より高い推論スループットを報告しています。特に自己回帰生成で、履歴を大きな KV キャッシュとして持ち続ける負担を減らせる点は魅力です。ただし、実際の速度は実装品質とハードウェア依存です。
Q. 小規模チームが学ぶ価値はありますか?
A. あります。Mamba を本番導入しなくても、「長文処理を全部 attention で解かない」「状態を持って選別する」という設計思想は、RAG、会話メモリ、ログ処理、時系列解析に広く応用できます。論文を読む価値は、その発想の転換にあります。
今日の学び
この論文は、Transformer が長い系列で重くなりすぎるという課題を扱いました。そこに対して、入力に応じて記憶更新を切り替える Selective SSM と、GPU 上で速く動かす並列スキャン実装を組み合わせ、線形時間で強い系列モデリングを目指しました。
開発の観点では、「長文を全部比較する」以外の設計余地があることが大きな学びです。長履歴エージェント、長文 RAG、時系列処理を作るなら、モデル本体だけでなく、情報をどう保持し、どこで選別するかを設計対象として見るべきだとわかります。