この記事の対象読者
- 状態空間モデルの再帰形式と畳み込み形式を知っていて、Mambaが何を足したのかを知りたい方
対象読者は1レベルに固定しています。状態空間モデルそのものが初めての方は、先に本シリーズの状態空間モデル編を読むと前提がそろいます。
この記事で得られること
- 選択機構という1つの変更が、なぜ畳み込み形式を捨てさせたのかの実測
- 並列スキャンが逐次ループの何倍出るか、実測とその頭打ちの理由
- 推論時のメモリと1トークンあたりの時間が、Transformerとどれだけ違うかの実数
- 手元で再現できる検証スクリプト
この記事で扱わないこと
- 学習を伴う言語モデリング実験（本記事の計測はすべて順伝播とメモリ試算です）
- CUDAカーネルの実装詳細
- Mamba-2以降で導入されたSSDの枠組み（本シリーズの別記事で扱います）
導入：全部メモする書記の限界
議事録を取る書記を想像してください。方法は2つあります。
1つは、全発言を録音しておいて、質問が来るたびに録音を最初から全部聞き直す方法。正確ですが、会議が長くなるほど聞き直しの時間が伸びます。これがTransformerのKVキャッシュです。
もう1つは、A4一枚の要約メモだけを持って、発言があるたびに書き足していく方法。何時間の会議でもメモは一枚のままです。これが状態空間モデルです。
ただし、二番目の方法には落とし穴があります。何を聞いても機械的に同じ調子で書き足す書記だと、雑談も重要な決定事項も同じ重みでメモに混ざり、肝心なことが薄まります。
Mambaがやったのは、この書記に今の発言はメモに残すか、聞き流すかという判断力を与えることでした。以降、この書記のたとえで通します。
1. Mambaが変えたのは何か
このセクションで分かること：選択機構の定義。従来のSSMと何が違うのか、パラメータの形の変化として理解します。
Mambaは Albert Gu と Tri Dao による2023年12月の論文で提案されました。論文の要旨では、変更点が次のように述べられています。
要旨の主張は2段構えです。SSMのパラメータを入力の関数にすると、トークンに応じた選択的な伝播と忘却が可能になる。そのうえで、この変更で効率的な畳み込みが使えなくなるため、再帰モードで動くハードウェア対応の並列アルゴリズムを設計したとしています。
引用の中に、この記事で扱うことがほぼ全部入っています。整理すると次のようになります。
| 従来のSSM | Mambaの選択的SSM | |
|---|---|---|
| $\Delta$ | 学習で決まる定数 | $\Delta_t = \mathrm{softplus}(W_\Delta u_t)$ |
| $B$ | 学習で決まる定数 | $B_t = W_B u_t$ |
| $C$ | 学習で決まる定数 | $C_t = W_C u_t$ |
| $A$ | 定数 | 定数のまま |
| 時不変性 | あり | なし |
| 畳み込み形式 | 使える | 使えない |
書記のたとえで読むと、$\Delta_t$ が今の発言をメモに残す度合い、$B_t$ がメモのどの欄に書くか、$C_t$ が答えるときにメモのどこを読むかです。従来のSSMではこの3つが会議の最初から最後まで固定でした。
$\Delta_t = 0$ にすると $\bar{A} = \exp(0 \cdot A) = 1$、$\bar{B} = 0$ となり、状態はまったく変化しません。つまり完全に聞き流すことができます。この「何もしない」を選べることが、選択機構の本質です。
2. 実測：時不変な書記は何を取りこぼすか
このセクションで分かること：ノイズに埋もれた少数の重要トークンを固定サイズ状態に残せるか、時不変と選択的で比べた数字。
論文が例に挙げるのは選択的コピーというタスクです。系列の大半はノイズで、まれに覚えるべき値が現れます。それを固定サイズの状態にどれだけ正確に残せるかを見ます。
学習を伴う実験ではなく、機構そのものの挙動を測る実験として組みました。系列長1024、うち信号8個、状態次元32、乱数シード20通りで計測しています。
クリックで計測コードを展開
import numpy as np
def trial(L=1024, n_signal=8, N=32, seed=0):
rng = np.random.default_rng(seed)
u = rng.normal(0, 1.0, L)
pos = np.sort(rng.choice(np.arange(L), n_signal, replace=False))
u[pos] = rng.normal(0, 3.0, n_signal)
is_signal = np.zeros(L, bool); is_signal[pos] = True
A = -0.5 - 1j * np.pi * np.arange(N) / N
B = np.ones(N, dtype=np.complex128)
def run(dt_seq):
x = np.zeros(N, dtype=np.complex128)
for t in range(L):
Ab = np.exp(dt_seq[t] * A)
Bb = (Ab - 1.0) / A * B
x = Ab * x + Bb * u[t]
return x
base = 20.0 / L
x_lti = run(np.full(L, base)) # 時不変
x_ideal = run(np.where(is_signal, base, 0.0)) # 選択的 = 参照
cos = np.abs(np.vdot(x_lti, x_ideal)) / (
np.linalg.norm(x_lti) * np.linalg.norm(x_ideal))
return float(cos)
print(np.mean([trial(seed=s) for s in range(20)]))
参照とするのは、信号だけを通した状態です。選択的SSMは $\Delta_t$ をノイズ位置で0にすればこの参照と一致するので、比較の焦点は時不変SSMが参照からどれだけずれるかになります。
| 方式 | 参照状態とのコサイン類似度 |
|---|---|
| 時不変 SSM | 0.649 ± 0.212 |
| 選択的 SSM | 1.000（構成上の参照） |
時不変の場合、コサイン類似度の平均は0.649でした。標準偏差0.212というばらつきも見逃せません。信号がどこに落ちたかによって、状態の汚染度が試行ごとに大きく変わります。書記のたとえで言えば、同じ会議でも雑談の混ざり方次第でメモの質が乱高下する状態です。
なお、選択的側が1.000なのは構成上そう定義したためであり、性能を測った数字ではありません。ここで示したのは「時不変では原理的に取りこぼしが起きる」という一点だけです。学習によって時不変SSMがこの差をどこまで埋められるかは、この実験からは分かりません。
3. 実測：入力依存にすると畳み込みが壊れる
このセクションで分かること：時不変性が崩れると畳み込み形式がどれだけ誤るか、崩れ具合を連続的に変えて測った結果。
畳み込み形式が成立するのは、カーネル $K_j = C\bar{A}^j\bar{B}$ が時刻 $t$ によらないからです。$\Delta_t$ が時刻ごとに変わると、$t$ ごとに違うカーネルが必要になり、畳み込み1回では表現できません。
これは理屈としては当たり前ですが、どれくらいのばらつきで実用にならなくなるのかは測ってみないと分かりません。$\Delta_t$ に相対的な揺らぎを加え、平均 $\Delta$ を使った畳み込み形式と、正しい逐次計算の誤差を比べました。系列長2048、状態次元32です。
| $\Delta_t$ の相対的な揺らぎ | 畳み込み形式の相対誤差 |
|---|---|
| 0%（時不変） | 1.77e-15 |
| 5% | 5.29% |
| 20% | 22.2% |
| 50% | 53.8% |
| 100% | 78.5% |
揺らぎ5%で誤差5.3%。ほぼ線形に効いています。わずか5%の入力依存性を入れただけで、畳み込み近似は使い物にならなくなるということです。近似で逃げる余地がないことが、この表で確認できます。
ここまでのまとめ
- Mambaは $\Delta$、$B$、$C$ を入力の関数にした。$A$ は定数のまま。
- 時不変SSMは、ノイズ混じりの系列で参照状態とのコサイン類似度が0.649まで落ちた。
- 入力依存性が5%入るだけで、畳み込み形式の相対誤差は5.3%に達する。近似での代替は成立しない。
4. 並列スキャン — 失った並列性を取り戻す
このセクションで分かること：逐次に見える再帰を並列化する原理と、その実測速度、そして頭打ちが起きる理由。
$x_t = a_t x_{t-1} + b_t$ という形の再帰は、逐次にしか見えません。しかしこの演算には結合則が成り立ちます。2つの区間をまとめる演算を次のように定義します。
(a_1, b_1) \circ (a_2, b_2) = (a_1 a_2,\; a_2 b_1 + b_2)
結合則が成り立つので、区間をどうまとめてもよい。したがって二分木状にまとめれば、$O(\log L)$ 段で全区間の結果が得られます。書記のたとえでは、複数の書記が会議を区間ごとに分担してメモを取り、あとで区間同士を合成するやり方です。
NumPyで実装して測りました。逐次ループと並列スキャンの結果は最大絶対誤差 $10^{-15}$ 台で一致しています。
| $L$ | 逐次ループ | 並列スキャン | 倍率 | 段数 |
|---|---|---|---|---|
| 1024 | 0.262 ms | 0.068 ms | 3.83倍 | 10 |
| 4096 | 1.005 ms | 0.164 ms | 6.12倍 | 12 |
| 16384 | 4.072 ms | 0.566 ms | 7.20倍 | 14 |
| 65536 | 15.61 ms | 3.698 ms | 4.22倍 | 16 |
| 262144 | 62.81 ms | 21.87 ms | 2.87倍 | 18 |
$L = 16384$ で7.2倍がピークで、そこから先は倍率が落ちています。並列スキャンは総演算量を増やすためです。逐次は $L$ 回の演算で済みますが、並列スキャンは $O(L \log L)$ の演算を行い、各段で配列全体を読み書きします。$L$ が大きくなるとメモリ帯域が律速し、並列化の利得を食い潰します。
ここが、Mambaの論文が並列スキャンではなくハードウェアを意識した並列アルゴリズムと表現している理由です。実装上の勝負どころは段数を減らすことではなく、中間状態を低速なHBMに書き戻さず、高速なSRAM上で処理し切ることにあります。書記のたとえなら、メモ帳を毎回鞄にしまうか、机の上に置いたまま書き続けるかの差です。
5. Mambaブロックの構造
このセクションで分かること：選択的SSMが実際のニューラルネットのブロックにどう組み込まれているか。
Mamba論文の要旨は、選択的SSMを注意機構もMLPブロックも持たない簡素化されたアーキテクチャに統合したと説明しています。Transformerブロックが注意層とFFN層の2段構成なのに対し、Mambaブロックは1種類のブロックを積むだけです。
短い畳み込みが入っている点は見落とされがちですが、これは局所的な文脈を拾う役割を担っています。状態空間モデルが長距離を、幅4程度の畳み込みが近傍を、それぞれ担当する分業です。
6. 実測：推論時のコスト
このセクションで分かること：長文を読むときのメモリと1トークンあたりの時間が、録音方式と要約メモ方式でどれだけ違うか。
24層、$d_{model}=2048$、状態次元16、拡張率2、fp16という構成で、推論時に保持すべきメモリを比較しました。Transformer側はKVキャッシュ、Mamba側はSSM状態と畳み込み状態の合計です。
| 系列長 | KVキャッシュ | Mamba状態 | 倍率 |
|---|---|---|---|
| 1,024 | 192 MB | 3.56 MB | 53.9倍 |
| 8,192 | 1,536 MB | 3.56 MB | 431倍 |
| 32,768 | 6,144 MB | 3.56 MB | 1,725倍 |
| 131,072 | 24,576 MB | 3.56 MB | 6,899倍 |
| 1,048,576 | 192 GB | 3.56 MB | 55,188倍 |
Mamba側は系列長によらず3.56 MBで一定です。100万トークンでも変わりません。これが固定サイズの容器に過去を圧縮するということの実際の意味です。
1トークン生成にかかる時間も測りました。注意側は全キャッシュに対するスコア計算とsoftmax、SSM側は状態の1回更新です。
| 系列長 | 注意の1ステップ | SSMの1ステップ | 倍率 |
|---|---|---|---|
| 512 | 0.870 ms | 0.136 ms | 6.4倍 |
| 2,048 | 3.308 ms | 0.143 ms | 23.1倍 |
| 8,192 | 13.41 ms | 0.132 ms | 101倍 |
| 32,768 | 90.90 ms | 0.143 ms | 638倍 |
SSM側は系列長が64倍になっても時間が変わりません。0.132から0.143 msの間で一定です。論文が挙げるTransformerの5倍のスループットという数字は、この性質に由来します。
ただし、公平を期すために書いておくと、この計測は最適化されたTransformer実装との比較ではありません。NumPyで書いた素朴な注意計算です。FlashAttentionのような実装を使えば注意側の絶対値は大きく改善します。それでも系列長に対する傾きの違いは残り、長文になるほど差が開くという構図は変わりません。
そして、圧縮には代償があります。3.56 MBに詰め込む以上、録音を全部残す方式より情報は落ちます。長文の中の特定の一文を正確に引用させるようなタスクでは、この差が効いてきます。要約メモの書記は、会議の流れを把握するのは得意でも、逐語引用は苦手なのです。
トラブルシューティング
| 症状 | 原因 | 対処 |
|---|---|---|
| 学習中に損失がNaNになる | $\Delta_t$ が負または過大 | softplusで正に制約し、初期バイアスを小さめに設定する |
| 長距離依存が学習されない | $\Delta$ の初期値が大きく状態が毎ステップ更新され過ぎ | $\Delta$ の初期範囲を 0.001 から 0.1 程度に狭める |
| 逐次実装が極端に遅い | Pythonループで再帰を回している | 並列スキャン、または公式の融合カーネルを使う |
| 並列スキャンが逐次より遅い | $L$ が小さく段数のオーバーヘッドが上回る | $L$ が数千未満なら逐次のまま使う |
| 学習は速いが推論が遅い | 推論でも系列全体を再計算している | 状態をキャッシュし、1ステップ更新に切り替える |
| 特定の語の逐語再現ができない | 固定サイズ状態による情報圧縮の原理的な限界 | 注意層を一部混ぜたハイブリッド構成を検討する |
用語集
| 用語 | 説明 |
|---|---|
| 選択機構 | $\Delta$、$B$、$C$ を入力の関数にする仕組み。Mambaの中核 |
| S6 | Selective Scan SSM の通称。Mambaの中核レイヤーの呼び名 |
| 時不変 | パラメータが時刻によらず一定であること。畳み込み形式の前提 |
| 並列スキャン | 結合則を使い、逐次の再帰を対数段数で計算する手法 |
| KVキャッシュ | 注意機構で過去のキーと値を保持する領域。系列長に比例して増える |
| SRAM / HBM | GPU上の高速な小容量メモリと、低速な大容量メモリ |
| ゲート経路 | ブロック内で本流と並走し、要素ごとの積で出力を調整する経路 |
| 拡張率 | ブロック内部で次元を何倍に広げるか。Mambaでは2が標準 |
学習ロードマップ
入門レベル
- 本記事のコードで、時不変SSMのコサイン類似度を自分の環境で再現する
- $\Delta_t$ の揺らぎを変え、畳み込み形式の誤差がどう伸びるかを描画する
- 並列スキャンを実装し、逐次ループとの一致を確認する
中級レベル
- Mamba論文の3章を読み、選択機構の動機づけを追う
- 公式リポジトリの
mamba_simple.pyを読み、ブロック構造をコードで確認する - HuggingFace上の小さなMamba系モデルを動かし、長文でのメモリ使用量を実測する
発展レベル
- 選択的スキャンのCUDAカーネルを読み、SRAM上での処理がどう実装されているかを見る
- 注意層とMamba層を混ぜたハイブリッド構成の論文を読み、どこに注意を残すかの設計判断を追う
- Mamba-2で導入されたSSDの枠組みへ進む
参考文献
Gu, Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752、2023年12月1日投稿、2024年5月31日改訂 — 本記事で引用した選択機構の説明、ハードウェアを意識した並列アルゴリズム、注意もMLPも持たないアーキテクチャという記述の出典です。
state-spaces/mamba 公式リポジトリ — 著者らによる実装。ブロック定義と選択的スキャンのカーネルが読めます。
Gu, Goel, Ré, Efficiently Modeling Long Sequences with Structured State Spaces, arXiv:2111.00396 — Mambaの前提となるS4の原論文。状態空間モデルの基本形と離散化の扱いはこちらに。
関連記事
本記事は、Mamba系アーキテクチャを理解するための4本組の3本目です。
- matmulってなんだ？
- 状態空間モデル(SSM)ってなんだ？
- 本記事：Mamba-1ってなんだ？
- SSDってなんだ？
親記事はこちらです。
LLM / Transformer / Self-Attention / GPU の各記事もあわせてどうぞ。
選択機構は畳み込み形式を捨てる決断でした。次の記事では、その捨てたものを別の形で取り戻す試み、つまり再帰計算を行列積として書き直す枠組みを扱います。書記のメモ取りが、そのまま表計算になる話です。