この記事の対象読者
- PythonとNumPyが読め、RNNという言葉は聞いたことがあるが、状態空間モデルは初めてという方
対象読者は1レベルに固定しています。制御工学の履修経験は前提にしません。
この記事で得られること
- 状態空間モデルという1つの式が、再帰形式と畳み込み形式という2つの計算方法を持つ理由
- 離散化のやり方を1つ間違えるだけでモデルが発散する様子の実測
- 状態行列の初期化が、記憶できる長さをどれだけ左右するかの実測
- 手元で再現できる100行程度の検証スクリプト
この記事で扱わないこと
- HiPPO理論の導出（直交多項式による射影の数学）
- S4のCauchyカーネルによる高速化アルゴリズムの詳細
- 学習を伴う実験（本記事の計測はすべて順伝播のみです）
導入：貯水池を眺めていた話
自作のミニ言語モデルで、長い系列を扱う層が欲しくなりました。RNNを素直に書いたところ、系列長4096で1バッチに数十秒かかりました。ループがPython側にあるので当然です。
論文を追っていくと、同じ再帰計算が畳み込みとして書き直せて、それがFFTで一発で終わる、と書いてありました。半信半疑で実装して数値を突き合わせたら、誤差が $10^{-16}$ でした。そこで初めて腑に落ちました。
この記事では、状態空間モデルを貯水池にたとえて説明します。
- 入力 $u(t)$ = 流れ込む雨水
- 状態 $x(t)$ = 貯水池の各区画の水位
- 行列 $A$ = 各区画の水の抜けやすさ
- 行列 $B$ = 雨をどの区画にどれだけ配分するか
- 行列 $C$ = 水位計をどう合成して読むか
- ステップ幅 $\Delta$ = 何分おきに水位を更新するか
1. 状態空間モデルの定義
このセクションで分かること：状態空間モデルの連続時間表現と、そこに登場する4つの行列の役割。
状態空間モデル、略してSSMは、もともと制御工学で使われてきた線形システムの表現です。S4の論文では次の形が基本形として示されています。
x'(t) = A x(t) + B u(t), \qquad y(t) = C x(t) + D u(t)
Gu, Goel, Ré の2021年の論文は、状態行列 A を適切に選べば長距離依存を扱える点を出発点に据えています。
貯水池のたとえで読み直すとこうなります。$x'(t) = Ax(t)$ は、雨が降らなくても水位が変化する分、つまり自然な排水です。$Bu(t)$ は降った雨の流入。$Cx(t)$ は水位計から読み取った観測値。$Du(t)$ は、降った雨が貯水池を経由せず直接観測に混ざる分で、深層学習の実装では残差接続に相当するため、以降は $D=0$ として扱います。
重要なのは、$x(t)$ の次元 $N$ が系列長に依存しないことです。系列がどれだけ長くなっても、貯水池の区画数は増えません。固定サイズの容器に過去を圧縮して詰める、これがSSMの基本方針です。
2. 離散化 — 連続時間を計算機に載せる
このセクションで分かること：連続時間の式を離散ステップに変換する2つの方法と、片方だけが発散する実測結果。
上の式は連続時間です。実際のデータは離散的なトークン列なので、ステップ幅 $\Delta$ で離散化します。標準的なのはゼロ次ホールド、略してZOHです。
\bar{A} = \exp(\Delta A), \qquad \bar{B} = A^{-1}(\exp(\Delta A) - I) B
離散化すると、次の再帰式になります。
x_t = \bar{A} x_{t-1} + \bar{B} u_t, \qquad y_t = C x_t
もっと素朴な方法もあります。前進オイラー法です。
\bar{A} = I + \Delta A, \qquad \bar{B} = \Delta B
貯水池のたとえで言えば、ZOHは区間内の雨量を一定とみなして、排水の指数減衰を厳密に積分する方法。オイラーは接線で1歩だけ進める方法です。$\Delta$ が小さければどちらも同じに見えますが、大きくすると差が出ます。
ここからは実際の計測です。先ほどの $\bar{A}$ が、コード中の Abar としてそのまま登場します。状態次元 $N=64$、系列長 $L=4096$、$A$ はS4D-Lin初期化の $-1/2 + i\pi n$ を使いました。
クリックで離散化と再帰計算のコードを展開
import numpy as np
def zoh(A, B, dt):
Abar = np.exp(dt * A)
Bbar = (Abar - 1.0) / A * B # A^{-1}(exp(dt A) - I) B
return Abar, Bbar
def euler(A, B, dt):
return 1.0 + dt * A, dt * B
def ssm_recurrent(Abar, Bbar, C, u):
x = np.zeros_like(Abar, dtype=np.complex128)
y = np.empty(len(u), dtype=np.complex128)
for t in range(len(u)):
x = Abar * x + Bbar * u[t]
y[t] = np.sum(C * x)
return y.real
N, L = 64, 4096
A = -0.5 + 1j * np.pi * np.arange(N) # S4D-Lin 初期化
B = np.ones(N, dtype=np.complex128)
C = (np.random.randn(N) + 1j * np.random.randn(N)) / np.sqrt(N)
u = np.zeros(L); u[0] = 1.0 # インパルス入力
$\Delta$ を段階的に大きくして、インパルス応答の末端値を見ました。
| $\Delta$ の倍率 | 方式 | $\max|\bar{A}|$ | 末端の応答 | 判定 |
|---|---|---|---|---|
| 1倍 | ZOH | 0.99988 | 1.15e-05 | 安定 |
| 1倍 | オイラー | 1.00104 | 1.31e-03 | かろうじて安定 |
| 10倍 | ZOH | 0.99878 | 8.75e-06 | 安定 |
| 10倍 | オイラー | 1.10953 | 5.11e+181 | 発散 |
| 100倍 | ZOH | 0.98787 | 1.33e-23 | 安定 |
| 100倍 | オイラー | 4.93197 | NaN | 発散 |
| 1000倍 | ZOH | 0.88509 | 6.19e-219 | 安定 |
| 1000倍 | オイラー | 48.3284 | NaN | 発散 |
判定の基準は $|\bar{A}|$ が1を超えるかどうかです。1を超えた瞬間、貯水池は水を抜くどころか勝手に増やす装置になります。ZOHは $A$ の実部が負であるかぎり $|\exp(\Delta A)| = \exp(\Delta \mathrm{Re}(A)) < 1$ が数学的に保証されるので、$\Delta$ をどれだけ大きくしても安定側に留まります。
実は最初、この実験を書いたときはオイラーで通ると思っていました。$\Delta$ 倍率10で 5.11e+181 が出て、指数の桁数を二度見しました。$\Delta$ は学習で更新されるパラメータなので、離散化スキームの選択がそのまま学習の安定性に直結します。
3. 2つの顔 — 再帰形式と畳み込み形式
このセクションで分かること：同じSSMが再帰としても畳み込みとしても書けること、そして両者が本当に一致することの数値検証。
再帰式 $x_t = \bar{A}x_{t-1} + \bar{B}u_t$ を、$x_0 = \bar{B}u_0$ から順に展開します。
x_t = \sum_{s=0}^{t} \bar{A}^{\,t-s} \bar{B} u_s
両辺に $C$ を掛けると、出力は畳み込みの形になります。
y_t = \sum_{s=0}^{t} K_{t-s} u_s, \qquad K_j = C \bar{A}^{\,j} \bar{B}
$K$ は長さ $L$ の1次元カーネルです。貯水池のたとえでは、1回だけ雨が降ったときに水位計が示す時系列、つまりインパルス応答そのものです。線形なので、実際の雨はこの応答を時間をずらして足し合わせたものになります。
同じパラメータで両方を計算し、差を測りました。
| 項目 | 値 |
|---|---|
| 最大絶対誤差 | 1.39e-16 |
| 相対誤差 | 4.92e-15 |
| 再帰形式の時間 | 19.7 ms |
| 畳み込み形式の時間 | 0.223 ms |
倍精度の丸め誤差の範囲で完全に一致しました。「等価だと論文に書いてある」と「手元で $10^{-16}$ が出た」の間には、理解の質にけっこうな差があります。
この二面性が実務上どう効くかというと、学習時は畳み込み、推論時は再帰と使い分けられる点です。学習では系列全体が既知なので並列に潰せます。推論では1トークンずつしか来ないので、固定サイズの状態を更新するだけで済みます。
ここまでのまとめ
- SSMは $x' = Ax + Bu$、$y = Cx$ という線形システム。状態次元は系列長に依存しない。
- 離散化はZOHを使う。オイラー法は $\Delta$ が10倍になった時点で発散した。
- 同じSSMが再帰形式と畳み込み形式で書け、実測で相対誤差 4.92e-15 の一致を確認した。
4. 実測：系列長に対するスケーリング
このセクションで分かること：畳み込み形式が長い系列でどれだけ効くか、そして再帰形式が本当に線形時間かの確認。
系列長を256から65536まで変えて、両形式の時間を測りました。
| $L$ | 再帰形式 | 畳み込み形式 | 倍率 |
|---|---|---|---|
| 256 | 1.21 ms | 0.022 ms | 55.0倍 |
| 1024 | 4.88 ms | 0.050 ms | 96.9倍 |
| 4096 | 19.1 ms | 0.177 ms | 108.2倍 |
| 16384 | 77.7 ms | 0.796 ms | 97.6倍 |
| 65536 | 371 ms | 3.88 ms | 95.5倍 |
再帰形式の時間は $L$ に正比例しています。$L$ が4倍になるたびに時間もほぼ4倍です。理論どおり $O(L)$ です。
ただし$O(L)$ であることと、実測で速いことはまったく別の話です。ここが本記事でいちばん伝えたい点かもしれません。再帰形式はPythonのループなので、1ステップあたり数マイクロ秒の固定費がかかります。畳み込み形式は同じ $O(L \log L)$ の仕事をFFTという高度に最適化されたルーチンに丸投げできるため、計算量が理論上は多いのに実測では100倍速い。
貯水池のたとえで言えば、区画ごとに人が桶で水を運ぶのが再帰形式、水路をつないで一斉に流すのが畳み込み形式です。運ぶ水の総量は同じでも、所要時間は段取りで決まります。この構図は、本シリーズのmatmul編で扱った演算強度の話と同じです。
5. 状態行列 $A$ の選び方が記憶を決める
このセクションで分かること：$A$ の初期化を変えると、どれくらい昔まで覚えていられるかが変わることの実測。
$\bar{A} = \exp(\Delta A)$ なので、$A$ の実部が大きく負だと水はすぐ抜けます。実部が0に近ければ長く残ります。つまり$A$ は貯水池の排水口の設計図です。
S4の系譜では、この設計をHiPPOという理論から導きます。実装上よく使われるのがS4D-Lin初期化で、$A_n = -1/2 + i\pi n$ という形です。実部を全区画で $-1/2$ に揃え、虚部で周波数を割り振ります。
3種類の $A$ でインパルス応答を計算し、系列後半にどれだけエネルギーが残るかを測りました。$L = 4096$、$N = 64$ です。
| 初期化 | 系列後半1/4に残るエネルギー比 |
|---|---|
| S4D-Lin、HiPPO由来 | 19.09% |
| ランダムな負の実部 | 8.95% |
| 強い減衰、$\mathrm{Re}(A) = -5$ | 0.082% |
$\mathrm{Re}(A) = -5$ の場合、系列の最後の1024ステップに残るエネルギーは全体の0.08%でした。4096トークンの文章を読ませても、実質的に直近の数百トークンしか見ていないことになります。
一方でS4D-Lin初期化は19%を保持しました。ランダム初期化の2倍以上です。長距離依存を扱えるかどうかは、アーキテクチャよりも初期化で決まる部分が大きいという、地味ですが実務的に重要な事実がここにあります。
6. SSMの系譜と、次に来るもの
このセクションで分かること：SSMがどう発展してきたかの流れと、本記事の次に読むべき話題。
S4は、$A$ への低ランク補正で安定に対角化し、計算をCauchyカーネルに帰着させるという手法でした。その後、対角行列だけでも十分だと分かり、実装は大幅に簡素化されました。本記事で使ったS4D-Lin初期化がその成果です。
ただし、ここまでのSSMには決定的な弱点が残っていました。$\bar{A}$、$\bar{B}$、$C$ が入力によらず一定だという点です。貯水池のたとえで言えば、どんな雨が降っても排水口の開き具合は固定です。この性質があるからこそ畳み込み形式が使えるのですが、同時に「今のトークンは重要だから覚えておく」といった判断ができません。
Mambaが持ち込んだのは、まさにこの点への回答です。Mamba論文の要旨は、SSMのパラメータを入力の関数にすると、トークンに応じた選択的な伝播と忘却が可能になると述べています。
続きは次の記事で扱います。
トラブルシューティング
| 症状 | 原因 | 対処 |
|---|---|---|
| 学習中に損失がNaNになる | 離散化がオイラー法、または $\Delta$ が大きすぎる | ZOHに変える。$\Delta$ をsoftplusで正に制約する |
| 長距離依存がまったく学習されない | $A$ の実部が大きく負で記憶が即座に減衰 | S4D-Lin など HiPPO由来の初期化を使う |
| 再帰形式と畳み込み形式で結果が違う | カーネル長が系列長より短い、またはFFTのパディング不足 | FFT長を $2L$ 以上の2冪にする |
| 推論だけ極端に遅い | 推論でも畳み込み形式を使っている | 推論は再帰形式に切り替える |
A^{-1} の計算でゼロ除算 |
$A$ の要素に0が混入 | 実部を負に制約する。対角要素の初期化を見直す |
| メモリが足りない | カーネル生成で $L \times N$ の中間配列を作っている | チャンクに分けて生成する |
用語集
| 用語 | 説明 |
|---|---|
| SSM | State Space Model。状態空間モデル。線形システムの表現形式 |
| 状態次元 $N$ | 貯水池の区画数にあたる、圧縮された記憶の大きさ |
| 離散化 | 連続時間の微分方程式を、有限のステップ幅の差分式に変換すること |
| ZOH | Zero-Order Hold。区間内の入力を一定とみなす離散化。指数関数を厳密に使う |
| インパルス応答 | 1点だけの入力に対する出力の時系列。畳み込みカーネルそのもの |
| LTI | Linear Time-Invariant。線形かつ時不変。パラメータが時刻によらず一定 |
| HiPPO | 過去の系列を直交多項式で最適に近似するための理論。$A$ の初期化を与える |
| S4 | Structured State Space Sequence model。SSMを実用的な計算量に落とした最初の成功例 |
| S4D | S4を対角行列に簡略化した版。実装が大幅に短い |
学習ロードマップ
入門レベル
- 本記事のコードを写経し、再帰形式と畳み込み形式の誤差を自分で出す
- $\Delta$ を変えて $|\bar{A}|$ を印字し、1を超える瞬間を観察する
- インパルス応答を matplotlib で描き、$A$ の実部との対応を見る
中級レベル
- S4の論文を、まず4章のアルゴリズム部分だけ読む
- S4Dの標準実装を読み、対角化がどれだけコードを短くするかを確認する
- Long Range Arena のタスク定義を読み、なぜ長距離依存の指標になるのかを理解する
発展レベル
- HiPPO論文で、なぜ特定の $A$ が最適な記憶を与えるのかを追う
- 選択機構を入れたSSMを実装し、畳み込み形式が使えなくなることを実際に確認する
- 並列スキャンを実装し、逐次ループとの速度差を測る
参考文献
Gu, Goel, Ré, Efficiently Modeling Long Sequences with Structured State Spaces, arXiv:2111.00396 — S4の原論文。本記事で引用した基本形の式 $x'(t)=Ax(t)+Bu(t)$、低ランク補正による安定な対角化、Long Range Arena での結果が載っています。
Gu, Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752 — 本記事の最終節で引用した、SSMのパラメータを入力の関数にするというアイデアの出典。
state-spaces/mamba 公式リポジトリ — MambaおよびS4系の実装。S4D初期化のコードが読めます。
関連記事
本記事は、Mamba系アーキテクチャを理解するための4本組の2本目です。
- matmulってなんだ？
- 本記事：状態空間モデルってなんだ？
- Mamba-1ってなんだ？
- SSDってなんだ？
親記事はこちらです。
Transformer / Self-Attention / Mamba-2 の各記事もあわせてどうぞ。
貯水池のたとえは、状態次元を増やすことの意味を考えるときにも使えます。区画を増やせば覚えられる情報は増えますが、そのぶん毎ステップの更新コストも増える。この綱引きが、次の記事の主題になります。