Transformerは自然言語処理の覇者ですが、致命的な弱点があります。Self-Attentionの計算量が系列長 $N$ に対して $O(N^2)$ であるため、長い系列の処理が困難なのです。FlashAttentionでメモリ効率は改善されましたが、計算量自体は $O(N^2)$ のままです。
もし系列長に対して線形 $O(N)$ の計算量で、Transformerに匹敵する性能を達成できるモデルがあれば、長い系列の処理が根本的に改善されます。この目標に最も近づいたのが、State Space Models（SSM）、特に2023年にGu & Daoが提案したMambaです。
Mambaは、制御理論の状態空間モデルを深層学習に応用した手法です。連続時間の線形システムを離散化し、入力に依存した選択的なメカニズムを導入することで、Transformerに匹敵する性能を線形計算量で達成しました。
SSMとMambaを理解することは、以下のような場面で直接役立ちます。
- 次世代アーキテクチャの理解: Mamba、Mamba-2、Jamba（MoEとMambaのハイブリッド）、Zamba など、SSMベースのモデルが急速に増えています
- 長系列処理のアプリケーション: DNA配列解析、音声処理、時系列予測など、系列長が数万〜数百万に達するタスクで威力を発揮します
- Transformerの限界の理解: Attentionが本質的に必要な場面と、SSMで代替可能な場面を判断できるようになります
本記事の内容
- 連続時間の状態空間モデルの基礎
- 離散化（Zero-Order Hold）の導出
- S4: HiPPO初期化と対角化による効率的な計算
- Mamba: 選択的SSMと入力依存パラメータ
- ハードウェア効率の良いアルゴリズム設計
- Pythonによるスクラッチ実装
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
連続時間の状態空間モデル
制御理論からの出発
状態空間モデル（SSM）は、制御理論と信号処理で数十年にわたって使われてきた枠組みです。連続時間の線形時不変（LTI）システムは、以下の微分方程式で記述されます。
$$ \frac{d\bm{h}(t)}{dt} = \bm{A}\bm{h}(t) + \bm{B}x(t) $$
$$ y(t) = \bm{C}\bm{h}(t) + Dx(t) $$
ここで： – $x(t) \in \mathbb{R}$: 入力信号 – $\bm{h}(t) \in \mathbb{R}^n$: 状態ベクトル（$n$ は状態次元、隠れ次元） – $y(t) \in \mathbb{R}$: 出力信号 – $\bm{A} \in \mathbb{R}^{n \times n}$: 状態遷移行列 – $\bm{B} \in \mathbb{R}^{n \times 1}$: 入力行列 – $\bm{C} \in \mathbb{R}^{1 \times n}$: 出力行列 – $D \in \mathbb{R}$: 直接伝達項（通常省略またはスキップ接続で実現）
この式を直感的に理解しましょう。状態ベクトル $\bm{h}(t)$ は「系の記憶」です。現在の状態は、過去の状態の時間発展（$\bm{A}\bm{h}(t)$）と新しい入力の影響（$\bm{B}x(t)$）の和で更新されます。出力は状態の線形変換（$\bm{C}\bm{h}(t)$）です。
RNNとの類似性は明らかです。RNNも隠れ状態 $\bm{h}_t$ を持ち、入力に基づいて更新し、出力を生成します。違いは、SSMが連続時間を前提としており、行列 $\bm{A}$ の構造に理論的な制約を課せることです。
なぜ連続時間から始めるか
離散時間のRNNを直接使えばよいように思えますが、連続時間から出発することには重要な利点があります。
- 初期化の理論: 連続時間の行列 $\bm{A}$ に対して、長期記憶を保持する理論的に最適な初期化（HiPPO）が存在します
- 解像度の柔軟性: 連続時間モデルは任意のサンプリングレートで離散化でき、学習時と推論時で異なる解像度が使えます
- 畳み込みとの等価性: 離散化後、LTIシステムは畳み込みと等価になり、学習時の並列計算が可能です
では次に、連続時間のSSMをどのように離散化するかを見ましょう。
離散化
Zero-Order Hold（ZOH）
連続時間の微分方程式をデジタルコンピュータで扱うために、離散化が必要です。最も一般的な離散化手法はZero-Order Hold（ZOH）です。
ZOHは「入力信号がサンプル間で一定」と仮定します。ステップサイズ $\Delta$ を導入し、微分方程式を解くと：
$$ \bm{h}(t + \Delta) = e^{\bm{A}\Delta}\bm{h}(t) + (\bm{A})^{-1}(e^{\bm{A}\Delta} – \bm{I})\bm{B}x(t) $$
離散時間の記法に書き直すと：
$$ \bm{h}_k = \bar{\bm{A}}\bm{h}_{k-1} + \bar{\bm{B}}x_k $$
$$ y_k = \bm{C}\bm{h}_k $$
ここで離散化されたパラメータは：
$$ \bar{\bm{A}} = e^{\bm{A}\Delta} $$
$$ \bar{\bm{B}} = (\bm{A})^{-1}(e^{\bm{A}\Delta} – \bm{I})\bm{B} $$
$\bar{\bm{A}}$ は行列指数関数であり、一般には計算が重いですが、$\bm{A}$ が対角行列であれば $\bar{\bm{A}} = \text{diag}(e^{a_1 \Delta}, \ldots, e^{a_n \Delta})$ と要素ごとの計算で済みます。これが後述するS4の対角化の動機です。
畳み込みとの等価性
離散化されたSSMの漸化式を展開すると：
$$ \bm{h}_k = \bar{\bm{A}}^k \bm{h}_0 + \sum_{j=0}^{k-1} \bar{\bm{A}}^{k-1-j} \bar{\bm{B}} x_j $$
$\bm{h}_0 = \bm{0}$ とし、出力を計算すると：
$$ y_k = \bm{C}\bm{h}_k = \sum_{j=0}^{k-1} \bm{C}\bar{\bm{A}}^{k-1-j}\bar{\bm{B}} \, x_j = \sum_{j=0}^{k-1} \bar{K}_{k-1-j} \, x_j $$
ここで $\bar{K}_i = \bm{C}\bar{\bm{A}}^i\bar{\bm{B}}$ はSSMカーネルです。この式は $y = \bar{K} * x$（畳み込み）に他なりません。
つまり、離散化されたLTI-SSMは畳み込みフィルタと数学的に等価です。これにより：
- 学習時: 畳み込み（FFTで $O(N \log N)$）で並列に計算
- 推論時: 漸化式で $O(1)$ ステップで逐次計算（RNNモード）
この2つのモードを切り替えられることがSSMの大きな利点です。Transformerは学習時は並列ですが、推論時はKVキャッシュが必要でメモリコストが高い。RNNは推論時は効率的ですが、学習時の並列化が困難です。SSMは両方の利点を持ちます。
S4: 長距離依存の学習
長期記憶の問題
素朴なSSM（ランダム初期化の $\bm{A}$）は、RNNと同様に長距離依存を学習できない問題を抱えています。離散化後の $\bar{\bm{A}}^k$ は、$k$ が大きくなると指数的に減衰（$|\bar{A}| < 1$）または発散（$|\bar{A}| > 1$）するため、数百ステップ以上前の情報を保持することが困難です。
HiPPO初期化
S4（Structured State Spaces for Sequence Modeling, Gu et al., 2022）は、HiPPO（High-order Polynomial Projection Operators）という理論的に最適な初期化を提案しました。
HiPPOの基本アイデアは、状態ベクトル $\bm{h}(t)$ が「過去の入力信号の最適な多項式近似の係数」を保持するように $\bm{A}$ を設計することです。
HiPPO-LegS（Legendre polynomials, Scaled）の場合、$\bm{A}$ の要素は：
$$ A_{nk} = \begin{cases} -(2n+1)^{1/2}(2k+1)^{1/2} & \text{if } n > k \\ -(n+1) & \text{if } n = k \\ 0 & \text{if } n < k \end{cases} $$
この特殊な構造により、状態ベクトルの各要素は過去の入力信号のルジャンドル多項式展開の係数に対応します。これにより、理論的に最適な長期記憶が実現されます。
対角化による効率化
HiPPO行列 $\bm{A}$ は密行列であり、$\bm{h}_k = \bar{\bm{A}}\bm{h}_{k-1} + \bar{\bm{B}}x_k$ の計算に $O(n^2)$ が必要です（$n$ は状態次元）。これでは系列長 $N$ に対して $O(N n^2)$ となり、$n$ が大きいと非効率です。
S4は $\bm{A}$ を対角化することでこれを解決します。$\bm{A} = \bm{V}\bm{\Lambda}\bm{V}^{-1}$ と対角化すると：
$$ \bar{\bm{A}} = e^{\bm{A}\Delta} = \bm{V}e^{\bm{\Lambda}\Delta}\bm{V}^{-1} $$
変換 $\tilde{\bm{h}} = \bm{V}^{-1}\bm{h}$ を導入すると、対角化された座標系では：
$$ \tilde{\bm{h}}_k = e^{\bm{\Lambda}\Delta} \tilde{\bm{h}}_{k-1} + \tilde{\bm{B}}x_k $$
$e^{\bm{\Lambda}\Delta}$ は対角行列なので、各状態変数が独立に更新され、計算量は $O(n)$ です。全系列の計算量は $O(Nn)$ で、$n = 64$ 程度の場合は非常に高速です。
S4は長距離依存ベンチマーク（Long Range Arena）で当時の最高性能を達成しましたが、自然言語処理でTransformerに追いつくことはできませんでした。その理由と解決策を次に見ましょう。
Mamba: 選択的SSM
S4（LTI-SSM）の限界
S4のSSMは線形時不変（LTI）です。行列 $\bm{A}, \bm{B}, \bm{C}$ は入力に依存しない固定パラメータであり、離散化後のカーネル $\bar{K}$ も入力に依存しません。
これは畳み込みとの等価性（学習時の並列化）を可能にする一方、入力依存の処理ができないことを意味します。例えば、「このトークンは重要なので記憶し、あのトークンは無視する」という選択的な処理は、LTIシステムでは原理的に不可能です。
TransformerのSelf-Attentionは本質的に入力依存です。各トークンが他の全トークンとの関連度をデータに基づいて動的に計算します。この入力依存性が、Transformerがin-context learning（few-shot推論）や情報の選択的処理に長けている理由です。
Mambaの選択的メカニズム
Mamba（Gu & Dao, 2023）は、SSMのパラメータ $\bm{B}, \bm{C}, \Delta$ を入力に依存させることで、この問題を解決します。
$$ \bm{B}_k = \text{Linear}_B(\bm{x}_k) \in \mathbb{R}^{n} $$
$$ \bm{C}_k = \text{Linear}_C(\bm{x}_k) \in \mathbb{R}^{n} $$
$$ \Delta_k = \text{softplus}(\text{Linear}_\Delta(\bm{x}_k)) \in \mathbb{R}^{+} $$
$\bm{A}$ は入力に依存しない学習可能パラメータのままですが、$\Delta_k$ を通じて離散化後の $\bar{\bm{A}}_k = e^{\bm{A}\Delta_k}$ は入力に依存します。
$\Delta$ の直感的理解
ステップサイズ $\Delta_k$ は「入力 $x_k$ をどれだけ記憶に取り込むか」を制御するゲートとして機能します。
$\Delta_k$ が大きい場合、$\bar{\bm{A}}_k = e^{\bm{A}\Delta_k} \to \bm{0}$（$\bm{A}$ が安定な負の固有値を持つ場合）となり、$\bar{\bm{B}}_k$ の項が支配的になります。つまり、過去の記憶をリセットし、現在の入力を強く取り込む操作です。
$\Delta_k$ が小さい場合、$\bar{\bm{A}}_k \to \bm{I}$ となり、過去の記憶をほぼそのまま保持し、現在の入力はほとんど取り込まない操作です。
これはLSTMのゲーティング機構と本質的に同じ役割を果たしています。LSTMの忘却ゲート $f_t$ が「過去をどれだけ覚えているか」を制御するのと同様に、$\Delta_k$ が「過去と現在のバランス」を入力に応じて動的に調整します。
畳み込みモードの喪失と代替
入力依存のパラメータを導入すると、SSMはもはやLTIではなくなり、畳み込みとの等価性が失われます。つまり、学習時にFFTベースの高速畳み込みが使えません。
Mambaは、代わりにハードウェア効率の良いアルゴリズムを設計しました。具体的には：
- 並列スキャン（Parallel Scan）: 漸化式 $\bm{h}_k = \bar{\bm{A}}_k \bm{h}_{k-1} + \bar{\bm{B}}_k x_k$ をGPU上で並列に計算する。work-efficient parallel prefix sumのアルゴリズムにより、$O(N)$ の計算量で $O(\log N)$ の深さ（レイテンシ）を実現
- カーネル融合: 離散化、SSM更新、出力計算をSRAM上で1つのカーネルに融合し、HBMへの中間結果の書き戻しを回避（FlashAttentionと同様の戦略）
Pythonによる実装
離散化SSMの実装
import numpy as np
import matplotlib.pyplot as plt
np.random.seed(42)
def discretize_zoh(A, B, delta):
"""Zero-Order Holdによる離散化。
対角行列Aを仮定して効率的に計算する。
"""
# A が対角の場合、行列指数は要素ごと
A_bar = np.exp(A * delta) # 対角要素ごと
B_bar = (1.0 / A) * (A_bar - 1.0) * B # 対角の場合の簡略化
return A_bar, B_bar
def ssm_forward_recurrent(A, B, C, x, delta):
"""SSMの再帰的な前方計算（推論モード）。"""
seq_len = len(x)
n = len(A)
h = np.zeros(n)
outputs = np.zeros(seq_len)
for k in range(seq_len):
A_bar, B_bar = discretize_zoh(A, B, delta)
h = A_bar * h + B_bar * x[k] # 要素ごとの積（対角行列）
outputs[k] = np.dot(C, h)
return outputs
def ssm_forward_conv(A, B, C, x, delta):
"""SSMの畳み込みモード（学習時の並列計算）。"""
seq_len = len(x)
n = len(A)
A_bar, B_bar = discretize_zoh(A, B, delta)
# SSMカーネルの計算: K_i = C @ A_bar^i @ B_bar
kernel = np.zeros(seq_len)
A_power = np.ones(n)
for i in range(seq_len):
kernel[i] = np.dot(C, A_power * B_bar)
A_power = A_power * A_bar
# 畳み込み（因果的）
outputs = np.convolve(kernel, x)[:seq_len]
return outputs, kernel
# パラメータ設定
n = 16 # 状態次元
seq_len = 256
# HiPPO風の初期化（対角近似）
A = -np.arange(1, n + 1, dtype=np.float64) # 負の実数（安定）
B = np.ones(n) * 0.5
C = np.random.randn(n) * 0.1
delta = 0.01
# テスト信号: 低周波 + 高周波 + ノイズ
t = np.linspace(0, 1, seq_len)
x = np.sin(2 * np.pi * 3 * t) + 0.5 * np.sin(2 * np.pi * 15 * t) + 0.2 * np.random.randn(seq_len)
# 再帰モードと畳み込みモードで計算
y_recurrent = ssm_forward_recurrent(A, B, C, x, delta)
y_conv, kernel = ssm_forward_conv(A, B, C, x, delta)
# 可視化
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
# 左上: 入力と出力
axes[0, 0].plot(t, x, alpha=0.5, label='Input $x(t)$')
axes[0, 0].plot(t, y_recurrent, label='Output (recurrent)', linewidth=2)
axes[0, 0].set_xlabel('Time')
axes[0, 0].set_ylabel('Value')
axes[0, 0].set_title('SSM Input/Output')
axes[0, 0].legend()
axes[0, 0].grid(True, alpha=0.3)
# 右上: 再帰モードと畳み込みモードの一致確認
axes[0, 1].plot(t, np.abs(y_recurrent - y_conv))
axes[0, 1].set_xlabel('Time')
axes[0, 1].set_ylabel('|Recurrent - Convolution|')
axes[0, 1].set_title(f'Mode Equivalence (max diff: {np.max(np.abs(y_recurrent - y_conv)):.2e})')
axes[0, 1].set_yscale('log')
axes[0, 1].grid(True, alpha=0.3)
# 左下: SSMカーネル
axes[1, 0].plot(kernel[:100], color='#ffa726')
axes[1, 0].set_xlabel('Lag (k)')
axes[1, 0].set_ylabel('Kernel Weight')
axes[1, 0].set_title('SSM Kernel $\\bar{K}_k = C\\bar{A}^k\\bar{B}$')
axes[1, 0].grid(True, alpha=0.3)
# 右下: 状態次元と記憶能力
dims = [4, 8, 16, 32, 64]
memory_lengths = []
for d in dims:
A_test = -np.arange(1, d + 1, dtype=np.float64)
B_test = np.ones(d) * 0.5
C_test = np.random.randn(d) * 0.1
_, kernel_test = ssm_forward_conv(A_test, B_test, C_test, np.ones(500), delta)
# 90%エネルギーに達するまでのラグ
energy = np.cumsum(kernel_test**2)
energy /= energy[-1] + 1e-10
mem_len = np.searchsorted(energy, 0.9)
memory_lengths.append(mem_len)
axes[1, 1].plot(dims, memory_lengths, 'o-', color='#66bb6a', markersize=8)
axes[1, 1].set_xlabel('State Dimension $n$')
axes[1, 1].set_ylabel('Effective Memory Length')
axes[1, 1].set_title('Memory Capacity vs State Dimension')
axes[1, 1].grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('ssm_basics.png', dpi=150, bbox_inches='tight')
plt.show()
左上のグラフでは、SSMが入力信号（青）をフィルタリングして出力（オレンジ）を生成していることが確認できます。SSMは線形フィルタとして機能し、特定の周波数成分を選択的に通過させています。
右上のグラフでは、再帰モード（推論用）と畳み込みモード（学習用）の出力が浮動小数点精度の限界（$10^{-14}$以下）まで一致していることが確認できます。これにより、学習と推論で異なる計算モードを安全に切り替えられることが保証されます。
左下のSSMカーネルは、時間とともに減衰するインパルス応答を示しています。カーネルの形状は $\bm{A}$ の固有値で決まり、HiPPO初期化された $\bm{A}$ は長い減衰特性を持ちます。
右下のグラフから、状態次元 $n$ を増やすほど記憶容量（有効記憶長）が増加することがわかります。これは多項式の次数が増えるほど複雑な関数を近似できるのと同じ原理です。
Mamba風の選択的SSMの実装
class SelectiveSSM:
"""Mamba風の選択的SSM（入力依存パラメータ）。"""
def __init__(self, d_model, d_state=16):
self.d_model = d_model
self.d_state = d_state
# 固定パラメータ
self.A_log = -np.arange(1, d_state + 1, dtype=np.float64)
# 入力依存パラメータの射影
self.W_B = np.random.randn(d_model, d_state) * 0.01
self.W_C = np.random.randn(d_model, d_state) * 0.01
self.W_delta = np.random.randn(d_model, 1) * 0.01
self.delta_bias = np.ones(1) * -3.0 # softplus後に小さな値
def forward(self, x):
"""選択的SSMの前方計算。
Parameters
----------
x : np.ndarray, shape (seq_len, d_model)
Returns
-------
y : np.ndarray, shape (seq_len, d_model)
"""
seq_len, d_model = x.shape
n = self.d_state
# 入力依存パラメータ
B = x @ self.W_B # (seq_len, d_state)
C = x @ self.W_C # (seq_len, d_state)
delta_logit = x @ self.W_delta + self.delta_bias # (seq_len, 1)
delta = np.log1p(np.exp(delta_logit)) # softplus, (seq_len, 1)
A = self.A_log # (d_state,)
# 各次元独立にSSMを適用（簡略化）
y = np.zeros_like(x)
h = np.zeros((d_model, n))
for k in range(seq_len):
# 離散化（入力依存の delta）
A_bar = np.exp(A[None, :] * delta[k]) # (1, d_state)
B_bar = delta[k] * B[k:k+1, :] # 簡略化された離散化
# 状態更新（各d_model次元で独立）
h = h * A_bar + x[k:k+1, :].T @ B_bar # (d_model, d_state)
# 出力
y[k] = np.sum(h * C[k:k+1, :], axis=1) # (d_model,)
return y, delta
# テスト
d_model = 32
seq_len = 128
d_state = 16
ssm = SelectiveSSM(d_model, d_state)
x = np.random.randn(seq_len, d_model) * 0.5
# 選択的SSMの前方計算
y, delta = ssm.forward(x)
# δの可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# 左: 出力
axes[0].imshow(y.T, aspect='auto', cmap='RdBu', vmin=-0.5, vmax=0.5)
axes[0].set_xlabel('Sequence Position')
axes[0].set_ylabel('Feature Dimension')
axes[0].set_title('Selective SSM Output')
# 右: ステップサイズ Δ の分布
axes[1].hist(delta.flatten(), bins=50, color='#00d4ff', alpha=0.7, edgecolor='white')
axes[1].set_xlabel('Step Size Δ')
axes[1].set_ylabel('Count')
axes[1].set_title('Distribution of Input-Dependent Δ')
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('mamba_selective.png', dpi=150, bbox_inches='tight')
plt.show()
print(f"入力形状: {x.shape}")
print(f"出力形状: {y.shape}")
print(f"Δの範囲: [{delta.min():.4f}, {delta.max():.4f}]")
print(f"Δの平均: {delta.mean():.4f}")
左のヒートマップでは、選択的SSMの出力が系列位置ごとに異なるパターンを示しており、入力に応じた動的な処理が行われていることが確認できます。LTI-SSMであれば固定のフィルタリングしかできませんが、選択的SSMは各位置で異なるフィルタ特性を持ちます。
右のヒストグラムでは、ステップサイズ $\Delta$ の分布を示しています。$\Delta$ が入力に依存して変化しており、小さい値（記憶を保持）から大きい値（記憶をリセット）まで幅広い範囲を取っていることがわかります。これが「どの入力を記憶し、どの入力を忘れるか」を選択的に制御するMambaの核心メカニズムです。
Mamba vs Transformer
計算量の比較
| 特性 | Transformer | Mamba |
|---|---|---|
| 学習計算量 | $O(N^2 d)$ | $O(N d n)$ |
| 推論計算量（1トークン） | $O(Nd)$（KVキャッシュ） | $O(dn)$（状態更新のみ） |
| メモリ（推論） | $O(Nd)$（KVキャッシュ増大） | $O(dn)$（固定） |
| 長距離依存 | ★★★★★ | ★★★★☆ |
| In-context learning | ★★★★★ | ★★★☆☆ |
Mambaの最大の利点は推論時のメモリが系列長に依存しないことです。Transformerは系列長に比例するKVキャッシュが必要ですが、Mambaは固定サイズの状態ベクトルのみです。
ハイブリッドアーキテクチャ
Transformerの強み（情報の選択的取得、in-context learning）とMambaの強み（線形計算量、固定メモリ推論）を組み合わせたハイブリッドアーキテクチャも登場しています。
- Jamba（AI21, 2024）: MambaとTransformerの層を交互に配置
- Mamba-2（Gu & Dao, 2024）: SSMとAttentionの統一的な理論を構築し、計算効率を改善
- Zamba（Zyphra, 2024）: Mamba + Shared Attention
まとめ
本記事では、State Space ModelsとMambaの理論を解説しました。
- SSMは制御理論の状態空間モデルを深層学習に応用した枠組みで、連続時間の微分方程式をZOHで離散化して使う
- 離散化されたLTI-SSMは畳み込みと等価であり、学習時はFFTで並列計算、推論時は漸化式で逐次計算が可能
- S4はHiPPO初期化により長期記憶を理論的に保証し、対角化で計算を効率化した
- Mambaは $\bm{B}, \bm{C}, \Delta$ を入力依存にすることで、LTIの制約を外し、Transformerに匹敵する性能を$O(N)$の計算量で達成した
- $\Delta$ は「記憶の保持と更新のバランス」を入力に応じて動的に制御するゲートとして機能する
- 推論時のメモリが系列長に依存しない点が、Transformerに対する構造的な利点
次のステップとして、以下の記事も参考にしてください。