Mamba と状態空間モデル(SSM) — 注意機構なしで系列を扱う

注意機構は「全部を覚えて毎回見返す」方式で、その代償が文長の2乗の計算量です。状態空間モデルは逆に「固定サイズのメモを更新し続ける」方式で、系列長に対して線形に伸びます。連続時間の線形システムから出発し、S4がなぜ畳み込みに化けるのか、Mambaが何を「選択的」にして畳み込みを捨てたのか、そしてどこで注意機構に負けるのかまでを前提知識ゼロから追います。

Efficiently Modeling Long Sequences with Structured State Spaces

「全部を見返す」以外のやり方

Transformerの自己注意は、原理としては驚くほど素朴です。新しい単語を1つ処理するたびに、それまでの全単語をもう一度見に行く。だから離れた語との関係も1ステップで届きますし、精度もよく出ます。代わりに、文が2倍になれば見比べる組み合わせは4倍になります。この2乗の壁の正体はAttention機構を1から理解するで扱いました。

では、人間はどうしているでしょうか。長い会議の議事録を取るとき、私たちは発言のたびに冒頭から読み返したりしません。手元のメモを1枚だけ持ち、新しい発言が来るたびにそれを更新する。メモの大きさは会議が3時間でも5時間でも変わりません。

この「1枚のメモを更新し続ける」方式は、古くはRNNがやっていたことです。ところがRNNは、同じ行列を何百回も掛けるうちに昔の情報が消えるか爆発するかしてしまい、しかも1ステップずつしか進めないのでGPUで並列化できませんでした（詳しくはRNNとLSTMを1から理解する）。

状態空間モデル（State Space Model, SSM）は、この「1枚のメモ」方式を制御工学の言葉で立て直したものです。そしてMambaは、そこに1つの仕掛けを足すことで、言語という気まぐれなデータでもメモが機能するようにしました。以下、そのメモの正体を式まで降りて見ていきます。

直感: 状態とは「圧縮された過去」

まず登場人物は3つだけです。

• 入力 $x_t$: 今このステップで入ってきた信号（例: 1トークンぶんの数値）
• 状態 $h_t$: 手元のメモ。過去すべてを固定サイズに押し込んだもの
• 出力 $y_t$: メモから読み出した答え

やることも1行です。新しいメモ＝古いメモを少し薄めたもの＋今回の入力。そして答えはメモを覗いて作る。

ここで決定的に重要なのは、メモのサイズが系列の長さに依存しないことです。1万トークン読んでもメモは同じ大きさのまま。だから1トークンを処理するコストは常に一定で、系列全体では長さに比例します。注意機構が過去のKとVを全部持ち続ける（＝KVキャッシュが長さに比例して膨らむ）のとは、記憶の作り方が根本的に違います。

もちろんタダではありません。固定サイズに押し込むということは、押し込みきれない情報は捨てているということです。どこで得をしてどこで損をするのかは、記事の後半で正面から扱います。

仕組み: 連続時間の線形システム

SSMの出発点は、制御工学で100年近く使われてきた線形システムです。時間 $t$ を連続だと思って書きます。

式(1)を日本語に直すと、「メモの変化の速さ $h'(t)$ は、今のメモ $h(t)$ を行列 $A$ で作り替えたものと、今の入力 $x(t)$ を行列 $B$ で作り替えたものの足し算で決まる。答え $y(t)$ は、メモを行列 $C$ で読み出したもの」となります。$A$ は「メモがどう自然に薄れていくか」、$B$ は「入力をメモのどこに書き込むか」、$C$ は「メモのどこを読むか」を担当します。

この式が言っているのは要するに、メモが今この瞬間にどう変わるかは、「今メモに入っているもの」と「今入ってきたもの」の2つだけで決まる、ということです。過去の入力を別に取っておく必要はありません。昔の入力の影響はすべて、すでに $h(t)$ の中に畳み込まれているからです。冒頭の議事録の比喩がそのまま式になっただけ、と思って構いません。

ばねと重りの運動でも、部屋の温度変化でも、この形で書けます。SSMはそれを「文章」に持ち込んだわけです。

ただしテキストは連続ではなく、トークンが1個ずつ並んだ離散の列です。そこで離散化します。刻み幅 $\Delta$ で区切り、その区間では入力が一定だとみなす（ゼロ次ホールド）と、式(1)は次の漸化式になります。

式(2)は「1ステップ進めるときの行列 $\bar{A}$ は、$A$ に刻み幅 $\Delta$ を掛けて指数を取ったもの」と言っています。$\bar{B}$ も同様に $\Delta$ から決まります。

言い換えると、微分方程式を解く話はここで終わりです。1トークン進めるという操作は、「今のメモに決まった行列を1回掛けて、新しい入力を足す」だけになりました。式(1)が連続時間の設計図で、式(2)がそれをコードにしたもの、という関係です。

ここで $\Delta$ は「1トークンぶんがどれだけの時間に相当するか」を表すつまみです。$\Delta$ が小さければメモはほとんど変わらず（＝入力を無視して昔を保つ）、大きければメモが大きく書き換わります。このつまみが後でMambaの主役になります。

なぜ $A$ の固有値が寿命を決めるのか

式(2)を $k$ ステップ回すと、最初に入った情報には $\bar{A}^k$ が掛かります。同じ行列の累乗です。行列の累乗の振る舞いを決めるのは固有値で、その絶対値が1より小さければ情報は指数的に薄れて消え、1より大きければ発散します。RNNの勾配消失・爆発と同じ現象が、ここでは「メモの寿命」として現れます。

だからSSMの成否は $A$ の設計に懸かっています。ランダムに初期化した $A$ では長距離の記憶がまるで残らないことが知られており、S4の系譜はここにHiPPOという理論を持ち込みました。HiPPOは「過去の信号を多項式で近似して保持するには $A$ をどう置くべきか」を導いたもので、その解として得られる特定の行列を初期値に使います。「昔をよく覚える初期値」が理論から出てくる、というのがこの系統の面白いところです。

S4: 再帰は畳み込みでもある

ここでS4（Structured State Space Sequence model）の核心的なトリックが出てきます。式(2)は逐次計算に見えますが、$\bar{A}, \bar{B}, C$ が全ステップで同じ（時不変）なら、展開して1本の式にまとめられます。

$$y_t = \sum_{k=0}^{t} C\bar{A}^{k}\bar{B}\, x_{t-k}$$

総和の記号に怯む必要はありません。この式が言っているのは要するに、時刻 $t$ の出力は、それまでに入ってきた入力すべての重み付き和であり、$k$ ステップ前の入力に掛かる重みは $C\bar{A}^{k}\bar{B}$ という決まった値だ、ということです。重要なのは、その重みが「何ステップ前か」だけで決まり、「それが文中のどこだったか」には一切依らない点です。だから重みの列を先に1本用意しておけば、あとはそれを入力列に沿って滑らせるだけで全時刻の出力が出ます。

これは「カーネル $\bar{K} = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \dots)$ を入力列に畳み込んだもの」です。つまり同じモデルが、再帰としても畳み込みとしても書ける。

この二面性が効きます。学習時は入力列が全部そろっているので、畳み込みとして扱えばFFTを使って $O(L \log L)$ で一気に計算でき、ステップ間の依存がないのでGPUを埋め尽くせます。推論時は再帰として扱えば、1トークンあたり定数時間・定数メモリで進めます。RNNの「並列化できない」という弱点を、学習時だけ別の顔に切り替えることで回避したわけです。

Mamba: 状態を「選択的」にする

S4には残った弱点がありました。$\bar{A}, \bar{B}, C$ が全ステップ共通ということは、どのトークンも同じ扱いで処理されるということです。音声や信号のように性質が一様なデータでは問題ありませんが、言語は違います。「the」も「,」も、固有名詞も、同じ強さでメモを書き換えてしまう。逆に「この単語だけは覚えておけ」「ここからは前の話題を忘れろ」という判断ができません。

Mamba（Gu & Dao, 2023）がやったのは一点です。$\Delta$、$B$、$C$ を入力 $x_t$ の関数にする。

# 疑似コード: 各時刻で B, C, Δ をトークンから作る（Aは学習パラメータのまま）
B_t     = x_t @ W_B                      # このトークンを状態のどこへ書くか
C_t     = x_t @ W_C                      # 状態のどこを読むか
delta_t = softplus(x_t @ W_d + bias)     # このトークンをどれだけ強く反映するか
A_t     = exp(delta_t * A)               # 時刻ごとに変わる減衰
h       = A_t * h + delta_t * B_t * x_t  # 状態を更新
y_t     = C_t @ h

これが選択的状態（Selective State）です。$\Delta_t$ が大きければそのトークンは強く書き込まれ、直前の状態は大きく上書きされる。$\Delta_t$ が小さければそのトークンはほぼ無視され、状態はそのまま保たれる。「今のは重要だから覚える／今のは埋め草だから流す」という取捨選択が、モデル自身の判断でできるようになりました。LSTMの入力ゲート・忘却ゲートが果たしていた役割を、線形システムの刻み幅として再発明した形です。

論文はこの効果を、意図的に作った2つの合成課題で示しています。1つは、ランダムな間隔で挟まれたトークンの中から特定のものだけを拾う課題。もう1つは、文中で一度見たパターンを後で再現する誘導ヘッド的な課題です。時不変なS4は前者で苦戦し、選択を入れたMambaは解けるようになる、という対比です。

代償と、それを払う方法（並列スキャン）

いいことばかりではありません。$\bar{A}$ が時刻ごとに違うということは、もう畳み込みで書けないということです。S4が持っていた「学習時はFFT」という切り札を失います。

Mambaはここを並列スキャン（parallel scan）で解きました。「前から順に累積する」処理は逐次に見えますが、累積和を思い出してください。$[1,2,3,4]$ の累積和は、まず隣同士を足して $[1+2,\ 3+4]$ を作り、その結果を組み合わせれば求まります。トーナメント表のように木構造で計算すれば、深さは $\log L$、総仕事量は $O(L)$ です。状態更新 $h_t = \bar{A}_t h_{t-1} + u_t$ も同じ形（結合則が成り立つ演算）に書けるので、同じ手が使えます。

さらにMambaの実装は、GPUのメモリ階層を意識した作りになっています。状態は各チャンネルごとに拡張された配列で、素直に書くと巨大な中間結果がHBM（GPUの主記憶）に載ってしまう。そこで離散化・スキャン・出力の掛け算を1つのカーネルに融合し、状態は高速なSRAM上だけで回して、逆伝播では保存する代わりに再計算します。「演算を増やしてでもメモリ往復を減らす」という発想はFlashAttentionとまったく同じで、実際どちらもTri Daoが関わっています。

計算量: 系列長との付き合い方

注意機構とSSMの違いを、系列長 $L$ の関数として並べます。

表の右2列が実務上いちばん効きます。Transformerは1トークン出すたびに過去全部を見に行くので、文脈が伸びるほど1トークンが遅くなり、KVキャッシュがメモリを食って同時実行数が落ちます。SSMは文脈が10倍になっても1トークンの計算も状態の大きさも変わりません。

ただし $O(L)$ と $O(L^2)$ の差が効くのは $L$ が十分大きいときだけ、という点は誤解されがちです。短い系列では定数項（実装の効率、GPUの使い切り方）が支配的で、高度に最適化されたFlashAttentionのほうが速いことは普通にあります。オーダーは傾きの話であって、切片の話ではありません。

どこで勝ち、どこで負けるか

勝つところは、素性がはっきりしています。系列が非常に長く、しかも情報が全体に広く分布しているデータ。ゲノム配列、音声波形、センサーの時系列などです。推論の1トークンあたりコストが一定なので、エッジ端末や長時間のストリーミング処理とも相性がよい。言語モデルとしても、Mamba論文は同規模のTransformerを上回るスケーリングを報告しています。

負けるところも、原理から予想がつきます。固定サイズの状態に押し込むという前提そのものが弱点です。文中のある文字列を正確にコピーする、長い文書から特定の1文を一字一句取り出す——こうした逐語的な想起の課題では、注意機構が明確に有利であることが理論と実験の両面から示されています（Jelassi et al., 2024）。注意機構は過去を丸ごと保持しているので「見に行けば必ずある」のに対し、SSMは状態に残っていなければ復元できません。長文脈での「干し草の中の針」型の評価で差が出やすいのはこのためです。

だから実務での主流はハイブリッドになりました。層の大半をMambaにして、要所に注意機構の層を数枚だけ挟む。Jamba（AI21, 2024）はこの構成を公開モデルとして示したもので、線形コストの利点を保ちながら逐語的な想起の弱さを補います。

もう一つ、理論側の進展としてMamba-2（Dao & Gu, 2024）が「構造化状態空間双対性（SSD）」を示しました。特定の構造を持つ注意機構と選択的SSMは同じ計算の別の見方であり、両者は行き来できる——という結果です。この見方に立つと、Mambaは注意機構の対抗馬というより、同じ空間の別の座標に見えてきます。

現場ではこう使う

誰がいつ触るか。 推論コストの担当者が「文脈長を伸ばしたいがKVキャッシュでGPUが埋まる」と言い出したとき、長時間の音声・センサー・ログを扱うチームがTransformerの2乗コストに詰まったとき、そしてオンデバイス推論でメモリ上限が固定のとき。逆に、社内文書からの正確な引用が要件なら、まずRAG＋注意機構ベースを疑うべきで、SSM単体を選ぶ場面ではありません。

実際に触る名前。 実装は mamba-ssm パッケージ（MambaLHS 相当のブロックと selective_scan カーネル）を使うのが一般的です。設計時に動かすつまみは4つに絞られます。

• d_state（式の $N$ に相当、状態の幅）: 大きくすると記憶容量は増えるがメモリと速度を食う
• d_conv: SSMの前段に入る短い1次元畳み込みの幅。局所的なn-gramを拾う担当で、これを外すと性能が目に見えて落ちる
• expand: ブロック内部でチャンネルを何倍に広げるか
• dt_min / dt_max（$\Delta$ の初期化範囲）: 最重要。ここが実効的な記憶の時定数を決めるので、扱う系列長に対して短すぎると長距離依存が最初から学習できない

知らないと事故になる落とし穴。

1. CUDAカーネルへの依存。選択的スキャンの高速実装は専用カーネルで、環境やdtypeが噛み合わないと純PyTorchの参照実装に静かにフォールバックし、桁で遅くなります。ベンチマーク前に、融合カーネルが実際に使われているかを必ず確認してください。
2. 「状態」の運用がKVキャッシュと違う。Transformerでは共通プレフィックスのKVを使い回す（prefix caching）のが定番ですが、SSMで同じことをするにはその時点の状態をスナップショットして保存・復元する必要があります。サーバの実装がそれを持っていないと、分岐のたびにプロンプトを最初から流し直すことになり、線形コストの利点を自分で捨てます。
3. 評価指標の選び方。パープレキシティだけを見ると差が出ないのに、逐語的な想起を要求するタスクで急に崩れる、という壊れ方をします。長文脈の検索・引用が要件なら、その形の評価を必ず入れてください。
4. 「Transformerを置き換える」と考えない。実運用の多くはハイブリッドです。何層に1枚注意機構を挟むかは設計判断であり、そこがまさに議論の対象になります。

面接や設計レビューで問われること。 「なぜMambaはS4のようにFFTで畳み込めないのか」——選択によって $\bar{A}$ が時刻ごとに変わり、時不変でなくなるから。「ではどうやって並列化するのか」——結合則を使った並列スキャンで、深さ $\log L$ に落とすから。この2問に即答できれば、この系統の設計思想を掴んでいると見なされます。

まとめ

• SSMは「固定サイズの状態を更新し続ける」線形システムで、系列長に対して線形時間・定数メモリ
• S4は時不変性を利用して、学習時は畳み込み（FFT）・推論時は再帰、と顔を使い分けた
• Mambaは $\Delta, B, C$ を入力依存にして選択を獲得し、失った畳み込みを並列スキャンとカーネル融合で埋めた
• 勝つのは長く一様な系列、負けるのは逐語的な想起。だから実務はハイブリッドに落ち着いている

次は、この選択的状態が「注意機構の一種」として書き直せるという構造化状態空間双対性の中身を、行列の形から追う予定です。