State Space Models (SSM)

状態空間モデル (SSM)は、各ステップで更新する隠れ状態ベクトルを介して入力系列を出力系列に変換するため、系列がどれだけ長くなってもメモリ使用量は一定です。この考え方は制御理論に由来し、線形システムを状態更新h' = Ah + Bxと出力y = Ch + Dxで表します。また、カルマンフィルターは1960年以来、ノイズを含む測定値からこのような状態を推定してきました。深層学習のSSMはデータからA、B、Cを学習し、これらのレイヤーを多数積み重ねます。2021年のS4や2023年のMamba以降、テキスト、音声、時系列データの長い系列を対象にトランスフォーマーと競合しており、計算量は系列長に対して線形に増加します。

SSMの仕組み#

ニューラルSSMレイヤーは、学習可能なステップサイズΔを使い、ゼロ次ホールド則Ā = exp(ΔA)によって連続方程式を離散的な漸化式に変換します。各ステップではh_t = Ā h_(t-1) + B̄ x_tとy_t = C h_tを計算します。A、B、Cがタイムステップ間で固定されている場合、漸化式を展開すると、出力はカーネル(CB̄, CĀB̄, CĀ²B̄, ...)による入力の畳み込みになります。そのため、同じレイヤーを学習時にはGPU上で系列全体を並列処理する畳み込みとして実行でき、生成時にはトークンごとに一定サイズの状態を更新する漸化式として実行できます。PyTorchでは、どちらの形式でも同じ出力が得られます。

import torch

torch.manual_seed(0)
L, N = 64, 16  # 系列長、状態サイズ
x = torch.randn(L)

# 対角連続時間SSMを、ステップサイズdeltaで離散化（ゼロ次ホールド）
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# 再帰形式: 生成時と同様に、ステップごとに状態を1回更新
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# 畳み込み形式: 学習時と同様に、カーネルK_k = C . (A_bar^k * B_bar)を事前計算
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Same outputs: Trueと出力されます。漸化式では各ステップで16個の状態値だけを保持し、事前計算されたカーネルを使うと、学習時に64個の出力をすべて並列計算できます。

S4からMambaへ#

初期の深層SSMは、大規模な学習には遅すぎ、メモリ消費も過大でした。Albert Gu、Karan Goel、Christopher RéによるS4は、畳み込みカーネルを効率よく計算できる状態行列の構造と、長期の履歴を保持するように設計された初期化を導入しました。Long Range Arenaベンチマークのすべてのタスクで最先端の性能を達成し、16Kステップ系列を扱うPath-Xタスクを解いた最初のモデルとなり、トランスフォーマーの60倍の速度で生成しました。

固定カーネルはすべての入力を同じように扱うため、不要な情報を飛ばして重要な事実を保持する判断ができません。Albert GuとTri DaoによるMambaは、Δ、B、Cを現在のトークンの関数にし、モデルが状態に何を書き込み、何を忘れるかを選べるようにします。これらの選択的パラメーターによって畳み込み形式は使えなくなるため、Mambaはハードウェアを考慮した並列スキャンで漸化式を計算します。Mambaはトランスフォーマーの5倍の推論スループットを達成し、Mamba-3Bは同サイズのトランスフォーマーを上回り、サイズが2倍のトランスフォーマーと同等の性能を示しました。各Mambaブロックでは、選択的SSMを短い畳み込みとSiLUゲートで囲みます。

Mamba-2は、選択的SSMとある種の線形注意が同じ計算を異なる視点から捉えたものだと示しました。著者らはこの枠組みを状態空間双対性と呼び、これを用いてコアレイヤーを2～8倍高速化しました。

SSMと関連アーキテクチャ#

• トランスフォーマー: 自己注意はすべてのトークンを相互に比較するため、過去の任意のトークンを正確にコピーできますが、計算コストは二次的に増加し、KVキャッシュはコンテキストウィンドウに応じて大きくなります。SSMは履歴を固定サイズの状態に圧縮するため効率的ですが、情報は失われます。2024年の論文Repeat After Meは、長い入力のコピーや情報検索で、トランスフォーマーが固定状態SSMを上回ることを示しました。
• リカレントニューラルネットワーク (RNNs): RNNやLSTMは非線形な更新を行うため、学習では系列をトークンごとに順番に処理する必要があります。SSMの状態更新は線形なので、畳み込みや並列スキャンが可能です。
• 隠れマルコフモデル (HMMs): HMMは遷移確率を持つ有限個の離散状態を扱います。カルマンフィルターの基礎となる線形ガウス状態空間モデルは、その連続版です。深層SSMはこの連続状態ベクトルを維持します。

SSMの活用分野#

本番環境での利用の大半は、正確な再現のために少数の注意レイヤーを残すハイブリッド大規模言語モデル (LLMs)です。AI21のJambaは、トランスフォーマーレイヤーとMambaレイヤーをエキスパート混合 (MoE)レイヤーと交互に配置し、単一の80 GB GPUで動作します。また、最大256Kトークンのコンテキストで高い性能を報告しています。IBMのハイブリッドGranite 4.0-Hモデルは、Mamba-2とトランスフォーマーレイヤーを9:1の比率で組み合わせています。Vision Mambaは画像パッチに選択的スキャンを適用し、SSMは全面的な注意機構では長すぎる生音声やDNA系列のモデル化にも使われます。古典的な形式はUltralytics YOLO26の追跡で動作します。デフォルトのTrackTrackトラッカーは、ByteTrackやBoT-SORTと同様に、カルマンフィルターで各対象の次のボックスを予測します。8次元の状態には、ボックスの中心、サイズ、およびそれぞれの速度が含まれます。定速度の遷移行列と、4つのボックス座標を読み取る線形観測行列で状態を更新します。