この記事の対象読者
- 状態空間モデルの再帰形式を理解していて、Mamba-2がなぜ速いのかを式とコードで納得したい方
対象読者は1レベルに固定しています。半分離可能行列という言葉を知らなくても読めるように書きました。
この記事で得られること
- 同じ計算が線形形式と二次形式の2通りで書ける理由と、その一致の数値検証
- 半分離可能行列の階数が状態次元と一致することの実測
- チャンクサイズを1から4096まで振ったときの実行時間曲線と、最適点がどこにあるか
- 行列 $A$ をスカラーに制限したことの効果を測った数字
この記事で扱わないこと
- 構造化マスク注意の一般論（本記事はSSDの特殊ケースに絞ります）
- Mamba-2の言語モデリング性能の評価
- Tensor Coreを使ったGPU実装（本記事の計測はすべてCPUです）
導入：連絡網の話
学級会の連絡事項を全員に伝える方法を考えます。
1つは伝言リレーです。1人が次の1人に伝え、それを繰り返す。1回あたりの手間は小さいのですが、順番待ちが発生します。
もう1つは、誰から誰へ何を伝えるかの対応表を作り、一斉に配る方法です。並列に処理できますが、表の大きさは人数の2乗になります。
この2つが、実は同じ結果を出す。これが構造化状態空間双対性、略してSSDの主張です。そして実務上いちばん効くのは、そこから出てくる第三の方法、班に分けるやり方です。
以降、この連絡網のたとえで通します。
1. SSDは何を主張しているのか
このセクションで分かること：SSDが定義する設定と、そこで現れる半分離可能行列という構造。
SSDは Tri Dao と Albert Gu による2024年5月の論文で提案されました。要旨には次のようにあります。
要旨の骨子は3点です。SSMと注意機構は密接に関連していること。両者が構造化半分離可能行列の様々な分解を通じて結びつくこと。そしてこの枠組みから設計されたMamba-2の中核層がMambaの2倍から8倍高速でありながら、言語モデリングでTransformerと競合すること。
SSDが扱う設定は、状態空間モデルの中でも制限された形です。
x_t = a_t x_{t-1} + B_t u_t, \qquad y_t = C_t^\top x_t
ここで $a_t$ はスカラーです。Mamba-1では $\bar{A}_t$ が $N$ 次元の対角行列でしたが、SSDでは全次元で同じ値に揃えます。連絡網のたとえで言えば、全員の伝達ロスが同じ割合という仮定です。
この制限を入れると、再帰を展開したときに次の形になります。
y_t = \sum_{s \le t} \left( \prod_{r=s+1}^{t} a_r \right) (C_t^\top B_s) \, u_s
つまり $y = M u$ という行列積で書けて、$M$ の成分は次のとおりです。
M_{ts} = \begin{cases} \left( \prod_{r=s+1}^{t} a_r \right) C_t^\top B_s & (s \le t) \\ 0 & (s > t) \end{cases}
これが対応表です。$C_t^\top B_s$ は注意機構のスコア行列と同じ形をしており、$\prod a_r$ は減衰マスクです。マスク付き線形注意そのものになっています。これが双対性の正体です。
2. 実測：線形形式と二次形式は本当に一致するのか
このセクションで分かること：3つの実装が数値的に一致することの確認と、それぞれの素の実行時間。
系列長4096、状態次元64、ヘッド次元32で3つの実装を回しました。ここからは実際のコードです。先ほどの $M_{ts}$ が、コード中の M としてそのまま現れます。
クリックで3実装のコードを展開
import numpy as np
def ssd_linear(a, B, C, X):
L, N = B.shape
S = np.zeros((N, X.shape[1]))
Y = np.empty((L, X.shape[1]))
for t in range(L):
S = a[t] * S + np.outer(B[t], X[t])
Y[t] = C[t] @ S
return Y
def ssd_quadratic(a, B, C, X):
cs = np.cumsum(np.log(a)) # cs[t] = Σ_{s<=t} log a_s
seg = cs[:, None] - cs[None, :] # = Σ_{r=s+1..t} log a_r
M = np.tril(np.exp(seg)) * (C @ B.T)
return M @ X
def ssd_chunked(a, B, C, X, Q):
L, N = B.shape; P = X.shape[1]; nc = L // Q
a_c, B_c = a.reshape(nc, Q), B.reshape(nc, Q, N)
C_c, X_c = C.reshape(nc, Q, N), X.reshape(nc, Q, P)
cs = np.cumsum(np.log(a_c), axis=1)
Lmask = np.tril(np.exp(cs[:, :, None] - cs[:, None, :]))
# 1. 班内: 二次形式
Y = np.einsum('cqs,csp->cqp',
np.einsum('cqn,csn->cqs', C_c, B_c) * Lmask, X_c)
# 2. 班の代表状態
states = np.einsum('cqn,cqp,cq->cnp', B_c, X_c, np.exp(cs[:, -1:] - cs))
# 3. 班間: 線形形式
S = np.zeros((N, P)); prev = np.empty((nc, N, P))
for c in range(nc):
prev[c] = S
S = np.exp(cs[c, -1]) * S + states[c]
# 4. 過去の班からの寄与を足し戻す
Y += np.einsum('cqn,cnp,cq->cqp', C_c, prev, np.exp(cs))
return Y.reshape(L, P)
| 実装 | 線形形式との相対誤差 | 時間 |
|---|---|---|
| 線形形式 | 基準 | 41.2 ms |
| 二次形式 | 5.27e-14 | 197 ms |
| チャンク分解、$Q=64$ | 7.97e-16 | 32.0 ms |
3実装すべてが倍精度の丸め誤差の範囲で一致しました。興味深いのは、チャンク分解が線形形式より誤差が小さいことです。累積積を対数空間でチャンク内に閉じて計算するため、長い区間にわたる誤差の蓄積が起きにくくなっています。
連絡網のたとえで言えば、班に分けたほうが伝言の劣化が少ない、ということです。理屈どおりではありますが、実測で確認できると納得感が違います。
3. 半分離可能行列の階数を測る
このセクションで分かること：対応表 $M$ に隠れている構造の正体を、特異値分解で確かめます。
$M$ は $L \times L$ の下三角行列なので、素朴に見れば $L^2/2$ 個の自由度があります。しかし実際にはそうではありません。
対角線をまたがない任意の部分行列、つまり行の範囲がすべて列の範囲より後にあるようなブロックを取り出すと、その階数は状態次元 $N$ 以下になります。これが半分離可能行列の定義です。
理由は式を見れば明らかです。$s < s_0 \le t_0 \le t$ の範囲では、$\prod_{r=s+1}^{t} a_r$ を $s_0$ と $t_0$ で分割できるので、ブロックは $N$ 本のベクトルの外積の和として書けます。
実際に測りました。$L = 4096$、$N = 64$ です。
| 取り出したブロック | 形状 | 数値的階数 | $\sigma_N / \sigma_0$ |
|---|---|---|---|
| 行2048:2560 × 列1024:1536 | 512×512 | 64 | 7.35e-17 |
| 行3000:3512 × 列100:612 | 512×512 | 64 | 8.01e-17 |
| 行1000:1512 × 列0:512 | 512×512 | 64 | 5.29e-17 |
| 行2048:4096 × 列0:2048 | 2048×2048 | 64 | 7.23e-17 |
どのブロックも階数はぴったり64、つまり $N$ と一致しました。2048×2048のブロックですら階数64です。65番目の特異値は最大特異値の $10^{-17}$ 倍、完全にゼロです。
これが意味するのは、対応表は巨大に見えて、実は $N$ 本の情報しか持っていないということです。連絡網のたとえで言えば、誰から誰への伝言も、64個の共通の中継点を経由して表現できる。状態空間モデルの状態 $x_t$ が、まさにその中継点です。
ここまでのまとめ
- SSDは $a_t$ をスカラーに制限した状態空間モデル。展開すると $y = Mu$ という行列積になる。
- $M$ はマスク付き線形注意と同じ形。線形形式と二次形式の一致を相対誤差 5.27e-14 で確認した。
- $M$ の対角外ブロックの階数は、実測でぴったり $N = 64$。これが半分離可能行列の構造。
4. チャンク分解 — 班に分けるという答え
このセクションで分かること：チャンクサイズを振ったときの実行時間曲線、最適点、そして二次形式が落ちる限界。
線形形式は演算量が少ないが逐次的。二次形式は並列だがメモリが $L^2$。ならば班に分ける。これがSSDアルゴリズムです。
$L = 4096$ で $Q$ を1から4096まで振りました。$Q = 1$ が純粋な線形形式、$Q = L$ が純粋な二次形式に相当します。
| $Q$ | 時間 | 推定演算量 |
|---|---|---|
| 1 | 72.3 ms | 5.11e7 |
| 2 | 40.4 ms | 4.35e7 |
| 4 | 31.3 ms | 4.09e7 |
| 8 | 25.2 ms | 4.19e7 |
| 16 | 24.8 ms | 4.72e7 |
| 32 | 26.6 ms | 5.92e7 |
| 64 | 30.5 ms | 8.41e7 |
| 128 | 40.4 ms | 1.34e8 |
| 512 | 104 ms | 4.36e8 |
| 2048 | 363 ms | 1.64e9 |
| 4096 | 690 ms | 3.25e9 |
$Q = 8$ から $32$ あたりが谷で、両端に向かって時間が増えるU字曲線になりました。試行ごとに最小点は $Q = 8$ と $Q = 16$ の間で揺れるので、この帯域は平坦と読むのが妥当です。
左端が遅いのは、班が小さすぎて逐次ステップ数 $L/Q$ が増えるため。右端が遅いのは、演算量そのものが $Q$ に比例して増えるためです。$Q = 4096$ での推定演算量は $Q = 16$ の約69倍でした。
なお、この最適点はCPU1コアでNumPyを回した場合の値です。GPUでは班内の行列積がTensor Coreに載るため、最適な $Q$ はもっと大きく、64から256程度になります。最適チャンクサイズはハードウェア依存だという点は押さえておく価値があります。
系列長を変えて、二次形式とチャンク分解を比べました。
| $L$ | 二次形式 | チャンク分解 | 倍率 | 対応表のメモリ |
|---|---|---|---|---|
| 1,024 | 9.99 ms | 7.84 ms | 1.27倍 | 8 MB |
| 2,048 | 49.1 ms | 15.9 ms | 3.09倍 | 32 MB |
| 4,096 | 196 ms | 44.9 ms | 4.36倍 | 128 MB |
| 8,192 | 902 ms | 63.7 ms | 14.1倍 | 512 MB |
| 16,384 | 計測不能 | — | — | 2,048 MB |
$L = 16384$ で二次形式はプロセスごと停止しました。OOM Killerに殺されています。対応表そのもので2 GB、それに $C B^\top$ の中間結果でもう2 GB必要で、メモリ3 GBの環境では最初の行列生成すら通りません。
これは実装の粗さではなく、二次形式が構造的に持っている限界です。一方チャンク分解は $Q \times Q$ の小さな表しか作らないので、系列長を伸ばしてもメモリは線形にしか増えません。連絡網のたとえなら、全校生徒分の対応表は物理的に印刷できないが、班ごとの表なら何班あっても紙1枚で足りる、という話です。
5. なぜ $A$ をスカラーに制限したのか
このセクションで分かること：Mamba-1の対角 $A$ とSSDのスカラー $a$ で、二次形式の計算コストがどれだけ変わるか。
ここまで当然のように $a_t$ をスカラーとして扱ってきましたが、Mamba-1では $\bar{A}_t$ は $N$ 次元の対角行列でした。表現力は明らかに下がります。にもかかわらず制限したのはなぜか。
答えは対応表の枚数です。$A$ が $N$ 次元対角だと、減衰の積が次元ごとに違うので、マスクが $N$ 枚必要になります。$C B^\top$ を1回の行列積でまとめられません。
クリックで対角Aの二次形式のコードを展開
def mamba1_quadratic(A_diag, B, C, X):
"""A_t が N 次元対角の場合。次元ごとにマスクが要る"""
L, N = B.shape
cs = np.cumsum(np.log(A_diag), axis=0) # (L, N)
Y = np.zeros((L, X.shape[1]))
for n in range(N):
seg = cs[:, n][:, None] - cs[:, n][None, :]
mask = np.tril(np.exp(seg))
Y += (mask * np.outer(C[:, n], B[:, n])) @ X
return Y
$L = 1024$、$N = 64$ で比較しました。
| 設定 | 必要なマスク枚数 | 時間 |
|---|---|---|
| $A$ が $N$ 次元対角、Mamba-1型 | 64枚 | 508.8 ms |
| $a$ がスカラー、SSD型 | 1枚 | 18.9 ms |
26.9倍の差です。対角 $A$ 版では、階数1の外積を64回作って足すことになり、行列積の形にまとまりません。本シリーズのmatmul編で扱った演算強度の話が、そのまま効いてきます。階数1の演算は読んだ値をほとんど使い回せず、メモリ律速になります。
これが、SSDが表現力を意図的に下げた理由です。行列積の形に収まることの価値が、対角 $A$ の表現力を上回ったという設計判断でした。論文が主張する2倍から8倍という速度向上は、この判断から来ています。
6. SSDが開いたもの
このセクションで分かること：双対性という視点が、アーキテクチャ設計にどう波及したか。
SSDの帰結を3点に整理します。
1つめ。状態空間モデルと線形注意が、同じ構造化行列の別々の分解として理解できるようになりました。線形形式は行ごとに逐次的に、二次形式は行列全体を一気に、チャンク分解はブロックごとに、それぞれ $M$ を計算しているだけです。
2つめ。Transformerで培われた実装ノウハウが移植可能になりました。テンソル並列やシーケンス並列といった学習の分散手法が、行列積の形をしているかぎりそのまま使えます。
3つめ。ハイブリッド設計の指針が得られました。両者が同じ枠組みで書けるなら、層ごとにどちらを選ぶかは表現力とコストのトレードオフで決められます。実際、SSM層と注意層を混ぜたモデルが多数登場しています。
連絡網のたとえで締めるなら、伝言リレーと一斉配布は対立する方式ではなく、同じ連絡網の読み方が2通りあるだけだった、ということです。そう理解できたとき、班に分けるという第三の方法が自然に出てきます。
トラブルシューティング
| 症状 | 原因 | 対処 |
|---|---|---|
| 系列長を伸ばすとOOMで落ちる | 二次形式で $L \times L$ の表を作っている | チャンク分解に切り替える |
| チャンク分解が線形形式より遅い | $Q$ が大きすぎて演算量が膨らんでいる | $Q$ を掃引してU字の谷を探す。CPUなら8から32、GPUなら64から256が目安 |
| 長い系列で結果がずれる | 累積積のアンダーフロー | 対数空間で累積し、exp は差分に対してのみ適用する |
| 二次形式と線形形式の誤差が大きい | マスクの生成で $s > t$ の要素をゼロにし忘れ |
np.tril の適用位置を確認する |
| $Q$ で系列長が割り切れずエラー | チャンク分割の前提が崩れている | 系列を右詰めでパディングし、マスクで無効化する |
| GPUで期待ほど速くならない | 班内の行列積が小さすぎてTensor Coreに載らない | $Q$ とヘッド次元を大きくし、行列積の形を太らせる |
用語集
| 用語 | 説明 |
|---|---|
| SSD | Structured State Space Duality。構造化状態空間双対性 |
| 双対性 | 同じ計算に、性質の異なる2つの等価な表現があること |
| 半分離可能行列 | 対角線をまたがない任意の部分行列の階数が $N$ 以下である行列 |
| 線形形式 | 状態を保持して逐次更新する計算方法。計算量は系列長に比例 |
| 二次形式 | $L \times L$ の行列を作って一括計算する方法。注意機構と同じ形 |
| チャンク分解 | 系列をブロックに割り、ブロック内は二次、ブロック間は線形で計算する方法 |
| 線形注意 | softmaxを取り除き、計算順序の入れ替えで線形時間にした注意機構 |
| 構造化マスク注意 | 注意のマスク行列に構造を課した一般的な枠組み。SSDはその特殊ケース |
| 演算強度 | 転送1バイトあたりのFLOPs。行列積の形に収まると大きくなる |
学習ロードマップ
入門レベル
- 本記事のコードで、3実装の誤差が丸め誤差の範囲に収まることを再現する
- 部分行列の特異値を印字し、$N+1$ 番目から急落することを目で確認する
- $Q$ を振ってU字曲線を描き、自分の環境の最適点を見つける
中級レベル
- SSD論文の3章、半分離可能行列の節を読む
- 線形注意の論文を読み、softmaxを外すと計算順序が入れ替えられる理由を追う
- 公式実装の
ssd_minimal.pyを読み、チャンク分解の4ステップを対応づける
発展レベル
- SSD論文の4章、構造化マスク注意を読み、SSDが一般論のどこに位置するかを把握する
- チャンク分解をTritonで実装し、班内の行列積をTensor Coreに載せる
- SSM層と注意層を混ぜたハイブリッドモデルの設計判断を、本記事の枠組みで読み解く
参考文献
Dao, Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, arXiv:2405.21060、2024年5月31日投稿 — 本記事の主題であるSSDの原論文。半分離可能行列による理論的な接続、チャンク分解アルゴリズム、Mamba-2の設計、そして2倍から8倍という速度向上の根拠が述べられています。
Gu, Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752 — 本記事で比較対象とした対角 $A$ を持つMamba-1の原論文。
state-spaces/mamba 公式リポジトリ — Mamba-2の実装。チャンク分解の最小実装 ssd_minimal.py が含まれています。
関連記事
本記事は、Mamba系アーキテクチャを理解するための4本組の4本目です。
- matmulってなんだ？
- 状態空間モデル(SSM)ってなんだ？
- Mamba-1ってなんだ？
- 本記事：構造化状態空間双対性ってなんだ？
親記事はこちらです。
Transformer / Self-Attention / NumPy / Triton の各記事もあわせてどうぞ。
計測環境はXeon 2.10GHzの1コア、メモリ3GB、NumPy 2.4.4です。GPUを持っている方は、同じスクリプトでU字の谷がどこに移動するか試してみてください。班の適正サイズは、教室の広さで変わります。