状態空間モデルとMamba：アテンションを超えるか

なぜ「アテンションを超える」話になるのか

Self-Attention の計算 は系列内の全ペアを比較するため、系列長 n に対して計算量とメモリが O(n²) で増えます。Transformer のアーキテクチャ はこの全結合的な情報のやり取りで強力な性能を得た一方、長文ではこの二乗が重くのしかかります。一方、LSTM/GRU のような再帰モデルは O(n) で済みますが、逐次依存のため学習を並列化しにくく、長距離の勾配が消えやすい弱点がありました。

状態空間モデル（State Space Model, SSM）は、この対立を別の角度から解こうとする系列モデルです。狙いは明快で、学習時は畳み込みのように並列化でき、推論時は再帰のように定数メモリで進められ、しかも系列長に線形 という三つを同時に満たすことです。

状態空間モデルの基本式

SSM はもともと制御理論の連続系から来ています。隠れ状態 h(t) が入力 x(t) を受けて時間発展し、そこから出力 y(t) を読み出します。

連続系:
  h'(t) = A h(t) + B x(t)     ← 状態の時間発展
  y(t)  = C h(t)              ← 状態からの読み出し

A は状態遷移行列（状態をどう保持・減衰させるか）、B は入力をどう状態に書き込むか、C は状態をどう出力に射影するかを決めます。ディープラーニングで扱うのは離散系列なので、ステップ幅 Δ で離散化し、連続の A・B を離散版 Ā・B̄ に変換します（ゼロ次ホールドなどの規則を使う）。

離散系（時刻 t = 1, 2, ...）:
  h_t = Ā h_{t-1} + B̄ x_t
  y_t = C h_t

ここで重要なのは、この再帰が 線形 だという点です。LSTM の各ステップに入る非線形ゲートがなく、状態更新は行列演算だけで閉じています。この線形性こそが、次に見る二重表現を成立させる鍵になります。

二重表現：再帰と畳み込みは同じもの

線形時不変（Ā・B̄・C が時刻に依らず固定）の SSM では、再帰を展開すると出力が入力の 畳み込み で書けます。h_0 = 0 として再帰を順に展開すると、出力 y_t は過去の入力 x_t, x_{t-1}, ... に固定の係数を掛けて足したものになります。

再帰を展開すると:
  y_t = C B̄ x_t + C Ā B̄ x_{t-1} + C Ā² B̄ x_{t-2} + ...

つまり畳み込みカーネル K と入力 x の畳み込み:
  K = (C B̄, C Ā B̄, C Ā² B̄, ..., C Ā^{n-1} B̄)
  y = K * x

この K を SSM 畳み込みカーネル と呼びます。同じモデルが二つの顔を持つわけです。

この二重性が SSM の核心です。学習では畳み込みとして GPU で並列に流し、推論では再帰に切り替えて1トークンずつ生成する。同じパラメータで表現だけを使い分ける ため、両者の利点を同時に取れます。

HiPPO：長距離記憶を安定に持つための初期化

素朴に A をランダム初期化すると、Ā^k が指数的に減衰または発散し、長距離の情報がすぐ失われます。S4（Structured State Space Sequence model）が解いたのがここで、HiPPO（High-order Polynomial Projection Operator） という理論に基づいて A を構成します。

HiPPO の発想は、状態 h_t を「これまでの入力履歴を直交多項式（ルジャンドル多項式など）で最適に近似した係数」とみなすことです。新しい入力が来るたびに、過去全体を最良近似し続けるように状態を更新する、その更新則を満たす特定の行列が HiPPO 行列です。これで A を初期化すると、状態が過去を圧縮した要約として機能し、長い系列でも情報を失いにくくなります。

S4 の限界：線形時不変ゆえの「選べなさ」

S4 系は長系列で強い一方、本質的な弱点がありました。Ā・B̄・C が 入力に依らず固定（線形時不変, LTI） であることです。畳み込みカーネル K が全時刻で同一ということは、モデルは「どの位置の入力も同じ重みで状態へ書き込み、同じ重みで読み出す」しかできません。

これが問題になるのは、文脈に応じて情報を取捨選択 したい場合です。たとえば「重要なトークンだけ状態に強く書き込み、無関係なフィラーは無視する」「特定のキーが来たら過去をリセットする」といった内容依存の操作は、係数が固定では原理的に表現できません。アテンションが内容（クエリとキーの一致）で重みを動的に決められるのと対照的です。これが「SSM は連想的な想起タスク（特定トークンの記憶・コピー）が苦手」と言われた理由でした。

Mamba の選択機構：パラメータを入力の関数にする

Mamba（S6 とも）はこの限界を、B・C・離散化幅 Δ を入力 x_t の関数にする ことで突破します。各時刻でこれらを動的に生成するため、もはや時不変ではありません（選択的 SSM, selective SSM）。

Mamba の選択機構（各時刻 t で動的に決まる）:
  Δ_t = softplus(Linear_Δ(x_t))    ← ステップ幅も入力依存
  B_t = Linear_B(x_t)
  C_t = Linear_C(x_t)
  Ā_t, B̄_t = 離散化(A, B_t, Δ_t)   ← Δ_t に応じて遷移・入力ゲートが変わる
  h_t = Ā_t h_{t-1} + B̄_t x_t
  y_t = C_t h_t

直観的には、Δ_t がゲートとして働きます。Δ_t が大きいと現在の入力を強く状態に取り込み（注目）、小さいと現在の入力を無視して過去の状態を保つ（スキップ）。B_t・C_t が入力依存になることで、内容に応じた書き込み・読み出しが可能になり、S4 が苦手だった選択的な想起ができるようになります。

選択的スキャンとハードウェア対応

逐次に見える再帰でも、h_t = Ā_t h_{t-1} + B̄_t x_t は 結合的（associative） な演算の連鎖なので、並列スキャン（prefix scan, Blelloch スキャン）で O(log n) の深さに並列化できます。Mamba はこれを 選択的スキャン（selective scan） として実装します。

さらに Mamba は、FlashAttention と同じ IO-aware の発想を持ち込みます。状態 h_t（系列長 × 状態次元 N で大きくなりうる）を GPU の主記憶（HBM）に実体化せず、高速な SRAM 上でスキャンを完結させ、必要な出力だけを書き戻します。離散化された Ā_t・B̄_t もその場で計算し（カーネル融合）、巨大な中間テンソルを HBM に置かないことで、選択機構を入れても線形時間と高い実効速度を両立させます。

推論側の差は実用上とても大きい点です。Transformer は生成のたびに過去全トークンの K/V を参照するため KV キャッシュ が系列長に比例して膨らみます。対して SSM/Mamba は固定サイズの状態 h さえ持てばよく、どれだけ長く生成してもメモリと1ステップの計算が一定 です。

「超えるか」への答え：適材適所

Mamba がアテンションを「超える」かは、タスク次第というのが正確な答えです。長文要約・音声・ゲノムなど 超長系列で内容依存がそこまで尖らない 領域では、線形コストと定数メモリ推論が効き、Transformer を上回る効率と同等以上の性能を示すことがあります。一方、複数の事実を厳密に照合・コピーするような連想的タスクでは、全ペアを直接見られるアテンションがなお強い場面も残ります。

実際の潮流は二者択一ではなく、ハイブリッド です。多くの層を Mamba（線形・定数メモリ）にしつつ、要所に少数のアテンション層を挟むことで、長距離効率と精密な想起の両取りを狙う構成が有力になっています。SSM が示したのは、アテンションの代替というより「系列モデリングには O(n²) の全比較とは別の、線形で長距離を扱う設計軸がある」という選択肢の拡張です。アテンションの数理は Self-Attention の計算 と Transformer のアーキテクチャ で、その推論コストの実態は KV キャッシュ で押さえると、両系統の長短が線でつながります。