""" ============================================================================= 第 4 讲 · 代码 3：状态空间模型（SSM）入门 —— 从 S4 到 Mamba 的核心思想 ============================================================================= 对应讲义：lectures/lecture04_attention_alternatives_moe.md 第 3.2 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 状态空间模型是另一条通向"线性复杂度序列模型"的路。它来自控制理论，最简单的离散形式是： h_t = A h_{t-1} + B x_t （状态更新） y_t = C h_t （输出） 这个文件讲清楚三件事： 1. **两种算法**：同一个线性时不变（LTI）SSM，既可以逐步递推（像 RNN，适合推理）， 也可以写成一次卷积（适合并行训练）—— 这是 S4 等早期模型的关键。 2. **选择性（selectivity）**：Mamba 让 A、B、C 随输入变化，模型可以"选择"记住什么、忽略什么。 代价是卷积形式不再成立，需要"并行扫描"算法。 3. **殊途同归**：Mamba-2 的 SSM（标量衰减）与第 01 号文件的"门控线性注意力"是同一个东西。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_state_space_models.py """ import importlib import math import torch _la = importlib.import_module("01_linear_attention") gated_linear_attention_recurrent = _la.gated_linear_attention_recurrent def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：线性时不变 SSM 的两种算法 # ============================================================================= def ssm_recurrent(a, b, c, x): """ 对角 SSM 的递推算法（单个输入通道）： h_t = a ⊙ h_{t-1} + b · x_t （a、b、c 都是 N 维向量，a 是对角矩阵 A 的对角线） y_t = c · h_t 【参数】 a, b, c : (N,) 张量，N 是状态维度；|a_i| < 1 保证稳定（不会爆炸） x : (T,) 输入序列 【返回】(T,) 输出序列 【复杂度】O(T N)，每一步只依赖上一步的状态 —— 推理时每个 token 的代价是常数。 """ h = torch.zeros_like(a) ys = [] for t in range(len(x)): h = a * h + b * x[t] ys.append((c * h).sum()) return torch.stack(ys) def ssm_convolution(a, b, c, x): """ 同一个 SSM 的卷积算法。 【推导】展开递推：h_t = Σ_{s≤t} a^{t-s} b x_s，所以 y_t = Σ_{s≤t} (c · a^{t-s} b) x_s = Σ_{j≥0} K_j x_{t-j}， 其中 K_j = Σ_i c_i a_i^j b_i 也就是说，输出是输入与一个固定"卷积核" K 的（因果）卷积。 核 K 只依赖参数、不依赖输入，可以先算好，再对整个序列做一次卷积 —— 完全并行 （S4 等模型用 FFT 在 O(T log T) 内完成；这里为了简单直接用矩阵形式）。 【边界】只有当 a、b、c **不随时间/输入变化**（线性时不变，LTI）时，这个卷积形式才成立。 """ T = len(x) j = torch.arange(T, dtype=x.dtype) K = (c[None, :] * a[None, :] ** j[:, None] * b[None, :]).sum(-1) # (T,) 卷积核 # 构造 Toeplitz 矩阵 M[t, s] = K[t - s]（s ≤ t），y = M x idx = torch.arange(T)[:, None] - torch.arange(T)[None, :] M = torch.where(idx >= 0, K[idx.clamp(min=0)], torch.zeros(())) return M @ x def demo_two_algorithms(): """演示 1：递推 = 卷积。""" print_header("演示 1：线性时不变 SSM —— 递推与卷积两种算法") torch.manual_seed(0) N, T = 16, 200 a = torch.rand(N) * 0.9 + 0.05 # 0.05 ~ 0.95：不同的"记忆长度" b, c = torch.randn(N), torch.randn(N) x = torch.randn(T) y_rec = ssm_recurrent(a, b, c, x) y_conv = ssm_convolution(a, b, c, x) print(f" 状态维度 N = {N}，序列长度 T = {T}") print(f" 递推 vs 卷积：最大相对差异 {((y_rec - y_conv).abs().max() / y_rec.abs().max()).item():.1e}") print(""" 和线性注意力一样，这里也有"对偶性"：训练时整段序列已知，用卷积（并行）； 推理时逐个生成，用递推（每步常数代价，状态大小固定）。 a 的每个分量控制一种"记忆长度"：a_i = 0.95 时，一个输入在 20 步后还剩 0.95^20 ≈ 0.36； a_i = 0.05 时，一步之后就几乎忘光了。""") # ============================================================================= # 第 2 部分：选择性 —— Mamba 的关键改动 # ============================================================================= def selective_ssm(x, delta, a_base, b, c): """ 选择性 SSM（Mamba 思路的极简版本，单通道、对角 A）： a_t = exp(-Δ_t · a_base) （Δ_t 由输入决定，Δ_t 越大，衰减越强、写入越多） h_t = a_t ⊙ h_{t-1} + Δ_t · b · x_t y_t = c · h_t 【直觉】Δ_t 像一个"聚焦程度"： Δ_t ≈ 0：a_t ≈ 1，几乎不写入 → 忽略当前输入，原样保留记忆； Δ_t 很大：a_t ≈ 0，大量写入 → 忘掉过去，记住当前输入。 因为 a_t、写入强度都随输入变化，模型就能**选择**记住什么、忽略什么。 【参数】 x : (T,) 输入 delta : (T,) 每个位置的 Δ_t（真实模型中由输入经过一个线性层 + softplus 算出） a_base : (N,) 正数 b, c : (N,) 【边界】a_t 随时间变化 → 不再是 LTI → 没有固定的卷积核，演示 2 会看到卷积算法给出错误答案。 """ h = torch.zeros_like(a_base) ys = [] for t in range(len(x)): a_t = torch.exp(-delta[t] * a_base) h = a_t * h + delta[t] * b * x[t] ys.append((c * h).sum()) return torch.stack(ys) def demo_selectivity(): """演示 2：一个"记住标记位置的值"的任务 —— LTI 模型做不到，选择性模型可以。""" print_header("演示 2：选择性 —— 记住重要的、忽略无关的") torch.manual_seed(0) T = 100 x = torch.randn(T) marker = 10 # 第 10 个位置的值是"重要信息" target = x[marker].item() print(f" 任务：序列中第 {marker} 个输入是重要的（值 {target:+.3f}），其余都是噪声；") print(f" 要求在最后一步（第 {T - 1} 步）输出这个值。") # 选择性模型：在标记位置 Δ 很大（清空并写入），其他位置 Δ ≈ 0（不写入、不遗忘） delta = torch.full((T,), 1e-4) delta[marker] = 20.0 a_base, b, c = torch.ones(1), torch.ones(1), torch.ones(1) y_sel = selective_ssm(x, delta, a_base, b, c) / delta[marker] # 除以写入强度，恢复原值 print(f" 选择性 SSM：最后一步输出 {y_sel[-1].item():+.3f}（目标 {target:+.3f}）") # LTI 模型：不管 a 取什么值，都无法同时做到"长期保留"和"忽略噪声" print(" 线性时不变 SSM（固定 a，每一步都同等地写入输入）：") for a in [0.5, 0.9, 0.99, 1.0]: y = ssm_recurrent(torch.tensor([a]), torch.ones(1), torch.ones(1), x) print(f" a = {a:<5}：最后一步输出 {y[-1].item():+.3f}") print(""" LTI 模型对每个输入一视同仁：a 小则很快忘掉第 10 个值，a 大则把后面 89 个噪声也累加了进来。 选择性模型根据输入决定写不写、忘不忘，轻松完成任务。这就是 Mamba 相对 S4 的关键改进， 也是为什么"数据依赖的门控"在所有现代线性序列模型（Mamba-2、Gated DeltaNet ……）中都是核心。 【边界】一旦参数随输入变化，就没有固定的卷积核了：""") y_wrong = ssm_convolution(torch.exp(-delta.mean() * a_base), delta.mean() * b, c, x) print(f" 用'平均参数'强行做卷积，最后一步输出 {y_wrong[-1].item():+.3f}（错误）") print(" 所以选择性模型的训练需要另一种并行算法 —— 并行扫描（下一个演示）。") # ============================================================================= # 第 3 部分：并行扫描 # ============================================================================= def sequential_linear_recurrence(a, b): """顺序计算 h_t = a_t h_{t-1} + b_t（h_{-1} = 0），返回所有 h_t。""" h = torch.zeros(()) out = [] for t in range(len(a)): h = a[t] * h + b[t] out.append(h) return torch.stack(out) def parallel_scan_linear_recurrence(a, b): """ 用并行扫描（Hillis-Steele 算法）计算同一个递推，只需 ⌈log2 T⌉ 轮。 【关键观察】把每一步看成一个"变换" h ↦ a h + b，记作 (a, b)。两个变换先后作用的结果仍是同类变换： 先 (a1, b1) 再 (a2, b2)：h ↦ a2 (a1 h + b1) + b2 = (a1 a2) h + (a2 b1 + b2) 也就是 (a1, b1) ∘ (a2, b2) = (a1·a2, a2·b1 + b2)。这个组合运算满足**结合律**， 而满足结合律的"前缀累积"都可以并行计算： 第 1 轮：每个位置与前面第 1 个位置组合； 第 2 轮：与前面第 2 个位置组合； 第 3 轮：与前面第 4 个位置组合 …… 每轮距离翻倍。 每一轮内部的所有位置可以同时计算（在 GPU 上并行），总共 log2(T) 轮。 【返回】与 sequential_linear_recurrence 相同 """ A, B = a.clone(), b.clone() T = len(a) step = 1 while step < T: # 位置 t 与位置 t - step 组合：(A_{t-step}, B_{t-step}) ∘ (A_t, B_t) A_prev = torch.cat([torch.ones(step), A[:-step]]) B_prev = torch.cat([torch.zeros(step), B[:-step]]) A, B = A_prev * A, A * B_prev + B step *= 2 return B def demo_parallel_scan(): """演示 3：并行扫描与顺序计算结果相同。""" print_header("演示 3：并行扫描 —— 用 log2(T) 轮算完一个长度为 T 的递推") torch.manual_seed(0) T = 1000 a = torch.rand(T) * 0.5 + 0.5 b = torch.randn(T) seq = sequential_linear_recurrence(a, b) par = parallel_scan_linear_recurrence(a, b) print(f" T = {T}：顺序计算需要 {T} 步；并行扫描只需要 {math.ceil(math.log2(T))} 轮") print(f" 两者的最大相对差异：{((seq - par).abs().max() / seq.abs().max()).item():.1e}") print(" （Mamba 在 GPU 上用高度优化的扫描核函数；Mamba-2 则把它改写成分块的矩阵乘法，更好地利用 Tensor Core。）") # ============================================================================= # 第 4 部分：殊途同归 —— SSM 与门控线性注意力 # ============================================================================= def demo_ssm_equals_linear_attention(): """演示 4：Mamba-2 式的 SSM = 门控线性注意力。""" print_header("演示 4：Mamba-2 的 SSM 就是门控线性注意力") torch.manual_seed(0) T, N, P = 50, 8, 4 # 序列长度、状态维度、输入通道数 B = torch.randn(T, N) # 随时间变化的 B_t（相当于 key） C = torch.randn(T, N) # 随时间变化的 C_t（相当于 query） X = torch.randn(T, P) # 输入（相当于 value） gamma = torch.rand(T) * 0.5 + 0.5 # 标量衰减 a_t（Mamba-2 把 A 限制为"标量 × 单位矩阵"） # SSM 写法：状态 H_t (N × P) = a_t H_{t-1} + B_t X_t^T；输出 y_t = C_t^T H_t H = torch.zeros(N, P) ys = [] for t in range(T): H = gamma[t] * H + torch.outer(B[t], X[t]) ys.append(C[t] @ H) y_ssm = torch.stack(ys) y_gla = gated_linear_attention_recurrent(C, B, X, gamma) print(f" SSM（C→query，B→key，x→value）vs 门控线性注意力：最大差异 {(y_ssm - y_gla).abs().max().item():.1e}") print(""" 把 C_t 看成 query、B_t 看成 key、x_t 看成 value，两者的公式一模一样。 Mamba-2 的论文把这种对应关系称为"状态空间对偶"（SSD）。 来自 RNN、控制理论、注意力三个方向的研究，最终汇聚到了同一个形式： S_t = （衰减/擦除）× S_{t-1} + （写入）； y_t = 读取(S_t) 区别只在"衰减/擦除/写入"的具体设计 —— 这正是讲义第 3 节那张演进表的主线。""") def main(): demo_two_algorithms() demo_selectivity() demo_parallel_scan() demo_ssm_equals_linear_attention() print("\n全部演示运行完毕。") if __name__ == "__main__": main()