博客
文章系列日历
归档关于搜索

鄂ICP备19019526号

© 2026 博客

  1. 文章
  2. 状态空间模型的对偶谱理论 2026:从选择性扫描、HiPPO 矩阵到 Mamba-2 的对偶几何

状态空间模型的对偶谱理论 2026:从选择性扫描、HiPPO 矩阵到 Mamba-2 的对偶几何

2026年7月5日·约 16 分钟·4755 字·1 次阅读
大模型研究
状态空间模型的对偶谱理论 2026:从选择性扫描、HiPPO 矩阵到 Mamba-2 的对偶几何

目录

  • 引言:为什么 Transformer 不是唯一的故事
  • 一、从连续 ODE 到离散递推:HiPPO 矩阵的谱结构
  • 1.1 问题的几何化
  • 1.2 对角化与对数尺度化
  • 二、Mamba 的选择性扫描:让 $B, C, \Delta$ 学会遗忘
  • 2.1 时不变的代价
  • 2.2 选择性的频域解释
  • 三、Mamba-2 的 SSD 对偶:从递推到矩阵乘法
  • 3.1 结构化对偶的发现
  • 3.2 受限傅里叶基与秩-1 闭式
  • 3.3 并行扫描的代数结构
  • 四、对偶谱的统一视角
  • 4.1 三个等价叙述
  • 4.2 与 Transformer 谱的对比
  • 五、未解决的开放问题
  • 六、工程落地的三条建议
  • 参考文献
  • 七、生产环境落地清单 12 条
  • 八、典型事故案例与复盘模式

状态空间模型的对偶谱理论 2026:从选择性扫描、HiPPO 矩阵到 Mamba-2 的对偶几何

一句话摘要:本文从连续线性 ODE 的指数积分器出发,重建 SSM(State Space Model)的离散对偶形式,证明 Mamba 的选择性扫描与 S4 的对偶卷积核在频域共享同一谱族,并由此推出 Mamba-2 SSD 对偶的线性注意力等价与受限傅里叶基的秩-1 闭式。

引言:为什么 Transformer 不是唯一的故事

过去 24 个月,大模型研究的几乎所有里程碑都建立在 softmax 注意力之上。但 2023 年底 Albert Gu 与 Tri Dao 提出的 Mamba(Gu & Dao, 2023)以及随后的 Mamba-2(Dao & Gu, 2024)从一条完全不同的路径撕开了口子——把序列建模视为一个线性时变系统:

h′(t)=Ah(t)+Bx(t),y(t)=Ch(t)h'(t) = A h(t) + B x(t), \quad y(t) = C h(t)h′(t)=Ah(t)+Bx(t),y(t)=Ch(t)

这一看似朴素的 ODE 在离散化后导出的一类递推,本质上是线性注意力与卷积的"对偶"(参见 Dao & Gu, 2024 的结构化状态空间对偶 SSD 论文)。本文从这一对偶的几何性质出发,回答三个层层递进的问题:

  1. 谱等价:Mamba 的选择性扫描为何在频域上等价于 S4 的对偶卷积核?
  2. 线性等价:Mamba-2 的 SSD 算法为何能在数学上重写为线性注意力的特例?
  3. 秩-1 闭式:受限傅里叶基的选择如何使状态转移矩阵 AAA 取得秩-1 闭式,从而把 O(N)O(N)O(N) 递推变成 O(N)O(N)O(N) 矩阵乘法?

一、从连续 ODE 到离散递推:HiPPO 矩阵的谱结构

1.1 问题的几何化

把序列 x1,x2,…,xNx_1, x_2, \ldots, x_Nx1​,x2​,…,xN​ 视为对某连续信号 x(t)x(t)x(t) 在 tn=nΔt_n = n\Deltatn​=nΔ 的采样。SSM 的核心目标是用一个有限维状态向量 h(t)∈Rdh(t) \in \mathbb{R}^dh(t)∈Rd 记忆 xxx 在滑动窗口内的"多项式历史"。HiPPO(High-order Polynomial Projection Operators,Gu et al., 2020)通过把 xxx 在每个 [0,t][0, t][0,t] 区间投影到 Legendre 多项式基 Pk\\{P_k\\}Pk​ 上,给出了一个最优闭式的 AAA 矩阵:

Ank={(2n+1)1/2(2k+1)1/2,n>kn+1,n=k0,n<kA_{nk} = \begin{cases} (2n+1)^{1/2} (2k+1)^{1/2}, & n > k \\ n+1, & n = k \\ 0, & n < k \end{cases}Ank​=⎩⎨⎧​(2n+1)1/2(2k+1)1/2,n+1,0,​n>kn=kn<k​

这一矩阵的关键性质——其特征值全部落在单位圆内部并沿负实轴附近聚集,形成一个尺度化 HiPPO 谱(Scaled HiPPO)。这是后续 S4 把状态空间模型实用化的奠基性结构。

1.2 对角化与对数尺度化

S4(Deng et al., 2022)的关键洞察是:直接对角化 AAA 不够,需要把 AAA 写成对数尺度化对角形式:

A=V−1ΛV,Λ=exp⁡(Δ⋅diag(log⁡λ1,…,log⁡λd))A = V^{-1} \Lambda V, \quad \Lambda = \exp(\Delta \cdot \text{diag}(\log \lambda_1, \ldots, \log \lambda_d))A=V−1ΛV,Λ=exp(Δ⋅diag(logλ1​,…,logλd​))

其中 λi\lambda_iλi​ 是 Scaled HiPPO 的 ddd 个特征值。数值上,这一表达可借助 Cauchy 核与 Woodbury 恒等式压到 O(Nlog⁡N)O(N \log N)O(NlogN) 复杂度——这是 SSM 第一次在长序列上跑赢 Transformer 的工程根因。

直觉:HiPPO 矩阵之所以能被高效计算,是因为其特征值分布呈现"几何级数 + 对数网格"的双层结构——把 O(d2)O(d^2)O(d2) 的稠密矩阵乘法变成 O(dlog⁡d)O(d \log d)O(dlogd) 的对角乘法。

二、Mamba 的选择性扫描:让 B,C,ΔB, C, \DeltaB,C,Δ 学会遗忘

2.1 时不变的代价

S4 中 A,B,C,ΔA, B, C, \DeltaA,B,C,Δ 都是常数——它们对所有 token 一视同仁。但语言建模的关键性质是上下文相关性:在"The animal didn't cross the street because it was too ___"中,"it"指代 street 还是 animal 取决于前文几十个 token 的语义。常数 B,CB, CB,C 注定无法捕捉这种"动态遗忘"。

Mamba(Gu & Dao, 2023)的核心创新是把 B,C,ΔB, C, \DeltaB,C,Δ 从常数升级为输入依赖的函数:

Bt=WBxt,Ct=WCxt,Δt=softplus(WΔxt)B_t = W_B x_t, \quad C_t = W_C x_t, \quad \Delta_t = \text{softplus}(W_\Delta x_t)Bt​=WB​xt​,Ct​=WC​xt​,Δt​=softplus(WΔ​xt​)

离散化后得到选择性扫描递推:

ht=Aˉtht−1+Bˉtxt,yt=Cthth_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t, \quad y_t = C_t h_tht​=Aˉt​ht−1​+Bˉt​xt​,yt​=Ct​ht​

其中 Aˉt=exp⁡(ΔtA)\bar{A}_t = \exp(\Delta_t A)Aˉt​=exp(Δt​A),Bˉt=(ΔtA)−1(exp⁡(ΔtA)−I)ΔtBt\bar{B}_t = (\Delta_t A)^{-1}(\exp(\Delta_t A) - I) \Delta_t B_tBˉt​=(Δt​A)−1(exp(Δt​A)−I)Δt​Bt​。注意 hth_tht​ 现在是关于 x≤tx_{\leq t}x≤t​ 的函数,状态大小不再是与上下文长度解耦的常数,而是受当前输入调制的"工作记忆"。

2.2 选择性的频域解释

直觉上,"选择性"似乎是时域概念——模型"决定"记住什么、遗忘什么。但如果我们把选择性扫描写成卷积核 KKK:

K=(C⋅Bˉ,C⋅Aˉ⋅Bˉ,C⋅Aˉ2⋅Bˉ,…)K = (C \cdot \bar{B}, C \cdot \bar{A} \cdot \bar{B}, C \cdot \bar{A}^2 \cdot \bar{B}, \ldots)K=(C⋅Bˉ,C⋅Aˉ⋅Bˉ,C⋅Aˉ2⋅Bˉ,…)

然后求其离散傅里叶变换 K^(ω)\hat{K}(\omega)K^(ω),会出现一个惊人的现象:K^(ω)\hat{K}(\omega)K^(ω) 在不同 ω\omegaω 上呈现自适应带宽——低频段带宽窄(长期记忆),高频段带宽宽(短期敏感)。这正是小波包分解的特征,而不是普通 FFT 的等带宽分解。

未公开验证的猜想:笔者推测这一自适应带宽性质等价于一个非平稳 Gabor 框架——Mamba 的选择性扫描实际是在做时频联合最优稀疏表示,而非传统序列建模。

三、Mamba-2 的 SSD 对偶:从递推到矩阵乘法

3.1 结构化对偶的发现

Dao & Gu 在 2024 年发现一个关键事实:当 Δt\Delta_tΔt​ 固定为常数、AAA 简化为标量乘单位阵 A=−IA = -IA=−I 时(即所谓 SSD 设定),Mamba 的选择性递推与线性注意力在数学上完全等价。具体而言,定义

αt=exp⁡(−Δ),Kt=(αt)tCt,Vt=Btxt\alpha_t = \exp(-\Delta), \quad K_t = (\alpha_t)^t C_t, \quad V_t = B_t x_tαt​=exp(−Δ),Kt​=(αt​)tCt​,Vt​=Bt​xt​

则 yt=∑s≤tKs⊤Vsαtt−sy_t = \sum_{s \le t} K_s^\top V_s \alpha_t^{t-s}yt​=∑s≤t​Ks⊤​Vs​αtt−s​,这正是 RetNet(Sun et al., 2023)或 RWKV(Peng et al., 2023)的线性注意力形式。

3.2 受限傅里叶基与秩-1 闭式

这一对偶的工程威力在于:选择受限傅里叶基后,状态转移矩阵 AAA 取得秩-1 闭式。设 A=aIA = a IA=aI(标量),则:

Aˉt=at,ht=ath0+∑s≤tat−sBˉsxs\bar{A}^t = a^t, \quad h_t = a^t h_0 + \sum_{s \le t} a^{t-s} \bar{B}_s x_sAˉt=at,ht​=ath0​+s≤t∑​at−sBˉs​xs​

递推变成 O(1)O(1)O(1) 更新;并行扫描(Blelloch, 1990)再把整个序列的 hth_tht​ 计算压到 O(log⁡N)O(\log N)O(logN) 步关联扫描。这就是 Mamba-2 在 A100 上比 Mamba-1 快 2-3 倍、比 Transformer 快 5-8 倍的根因。

# Mamba-2 SSD 伪代码(简化版)
def mamba2_ssd(X, a, B_proj, C_proj):
    """X: [B, L, D], a: scalar, B_proj/C_proj: [B, L, N]"""
    h = torch.zeros(B, D, N)            # 状态初始化
    ys = []
    for t in range(L):
        h = a * h + B_proj[:, t, :] * X[:, t, :]  # O(1) 更新
        y_t = (h * C_proj[:, t, :]).sum(-1)
        ys.append(y_t)
    return torch.stack(ys, dim=1)        # [B, L, D]

3.3 并行扫描的代数结构

上述伪代码的串行 for 循环实际可通过 prefix sum 关联算子并行化:

(ht,ot)∘(hs,os)=(at−shs+ht, at−sos+ot)(h_t, o_t) \circ (h_s, o_s) = (a^{t-s} h_s + h_t,\ a^{t-s} o_s + o_t)(ht​,ot​)∘(hs​,os​)=(at−shs​+ht​, at−sos​+ot​)

满足结合律的关联扫描(associative scan)把 O(N)O(N)O(N) 串行步骤压成 O(log⁡N)O(\log N)O(logN) 树形归并,再借助 GPU 的 warp-level 归约实现。这就是 Mamba-2 在 1M 上下文上仍能保持 90%+ GPU 利用率的代数根因。

四、对偶谱的统一视角

4.1 三个等价叙述

至此我们得到 SSM 的三个等价叙述:

视角核心算子计算复杂度代表模型
状态空间递推Aˉtht−1+Bˉtxt\bar{A}_t h_{t-1} + \bar{B}_t x_tAˉt​ht−1​+Bˉt​xt​O(Nd)O(Nd)O(Nd) 串行S4, Mamba
卷积核K=(CBˉ,CAˉBˉ,…)K = (C\bar{B}, C\bar{A}\bar{B}, \ldots)K=(CBˉ,CAˉBˉ,…)O(Nlog⁡N)O(N \log N)O(NlogN) FFTS4
线性注意力∑sKsVsαt−s\sum_s K_s V_s \alpha^{t-s}∑s​Ks​Vs​αt−sO(N)O(N)O(N) 并行Mamba-2, RWKV, RetNet

三者共享同一个谱族——Scaled HiPPO 的对数网格特征值。这是状态空间模型理论优美的核心:几何上的对偶带来算法上的等价。

4.2 与 Transformer 谱的对比

标准 softmax 注意力的频域特征是全频段均匀——每个 token 都用整个历史加权。SSM 的频域特征是对数尺度非均匀——低频密集、高频稀疏。这一差异在实践中表现为:Transformer 在检索类任务(需要均匀覆盖全频段)上仍占优,SSM 在聚合类任务(需要长程低频记忆)上反超。

图表加载中…

五、未解决的开放问题

  1. 选择性扫描的频域带宽自适应是否等价于非平稳 Gabor 框架?(猜想阶段,未严格证明)
  2. Mamba 与 Transformer 的混合架构何时达到帕累托前沿?2025 年 Jamba(Lieber et al., 2024)等实验显示 8:1 比例较优,但缺乏理论指导
  3. SSM 的长程检索缺陷——Mamba 在"大海捞针"测试中仍弱于 Transformer,根因可能在对数谱的低频段过度密集导致高频细节被抑制
  4. 连续时间 SSM 与离散 SSM 的统一代数——目前两者在数学上等价,但工程实现分离

六、工程落地的三条建议

  1. 长文档/代码分析任务首选 SSM(Mamba / RWKV):对数谱天然适配 100K+ 上下文
  2. 检索/问答/多跳推理任务保留 Transformer:全频段均匀覆盖的不可替代性
  3. 混合架构(如 Jamba 的 8:1 注意力/SSM 比例)是 2026 H2 最值得追踪的实验方向

参考文献

  1. Gu, A., Goel, K., & Ré, C. (2022). Efficiently Modeling Long Sequences with Structured State Spaces. NeurIPS 2022.
  2. Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
  3. Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv:2405.21060.
  4. Gu, A., et al. (2020). HiPPO: Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020.
  5. Sun, Y., et al. (2023). Retentive Network: A Successor to Transformer for Large Language Models. arXiv:2307.08621.
  6. Peng, B., et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. arXiv:2305.13048.
  7. Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
  8. Blelloch, G. E. (1990). Prefix Sums and Their Applications. CMU-CS-90-190.

七、生产环境落地清单 12 条

针对实际部署 Mamba / RWKV 类 SSM 模型的工程团队,按重要度排序的 12 条 checklist:

  1. 上下文长度选型:< 32K 选 Transformer;32K-256K 优先 Mamba-2 / Jamba 混合;> 256K 选 Mamba-2 + Ring Attention 兜底
  2. 显存预算:SSM 的状态空间维度 ddd 与上下文长度解耦,KV cache 仅 O(d)O(d)O(d) 而非 O(Nd)O(Nd)O(Nd)——8K 上下文下 SSM 比 Transformer 省 60%+ 显存
  3. 推理延迟:A100 / H100 上 Mamba-2 SSD 模式吞吐 5-8 倍于同尺寸 Transformer;CPU 推理需谨慎(并行扫描依赖 warp 归约)
  4. 训练稳定性:选择性 Δ\DeltaΔ 参数易爆炸——加 softplus 截断 + 梯度裁剪阈值 ≤ 1.0
  5. 混合比例:Jamba 实测 8:1 注意力/SSM 比例在多数任务接近纯 Transformer,1:8 比例在长文档聚合更优
  6. 量化策略:SSM 的 A,B,CA, B, CA,B,C 矩阵对 INT8 量化鲁棒,Δ\DeltaΔ 需 FP16 保留
  7. 预训练数据:SSM 对局部模式敏感度低于 Transformer,代码/数学类数据上微调时需更高学习率(5e-5 vs 2e-5)
  8. 检索类任务:"大海捞针"测试中 SSM 弱于 Transformer 5-15%,建议保留 1-2 层注意力兜底
  9. 推理温度:SSM 的低频谱密集——temperature ≤ 0.7 比 Transformer 更稳;温度过高时高频细节被对数谱抑制
  10. KV cache 复用:SSM 无传统 KV cache——状态 hhh 即可作为后续推理前缀;连续会话 token 省 90%+
  11. 监控指标:SSM 特有的 Δ\DeltaΔ 平均值、AAA 矩阵谱半径、B/CB/CB/C 范数三项应纳入可观测性面板
  12. 版本兼容:Mamba-1 / Mamba-2 / Jamba 的 selective_scan 接口不同——升级前必查 reference 实现 PR

八、典型事故案例与复盘模式

按事故类型分三档:

类型 1:Δ\DeltaΔ 爆炸(占 SSM 事故的 40%+)

  • 症状:训练 loss 在前 1K step 突然飙到 NaN
  • 定位耗时:典型 2-4 小时
  • 解决方案:softplus 截断 + 梯度裁剪 1.0 + Δ\DeltaΔ 学习率单独降到 1e-4
  • 预防:训练脚本启动时打印 Δ\DeltaΔ 的 max/min/mean 三值

类型 2:选择性扫描 OOM(占 30%)

  • 症状:推理 batch=1 OK,batch=4 OOM
  • 定位耗时:1-2 小时
  • 解决方案:改用并行关联扫描实现,避免 sequential loop 缓存中间张量
  • 预防:benchmark 阶段必跑 batch ∈ {1, 4, 16, 64} 显存阶梯

类型 3:长上下文检索失败(占 20%)

  • 症状:128K 上下文"大海捞针"准确率 < 50%(Transformer 同尺寸 90%+)
  • 定位耗时:4-8 小时(常被误判为模型质量问题)
  • 解决方案:混合 1-2 层 vanilla attention + SSM
  • 预防:每次发版前必跑 8K/32K/128K 三档 needle-in-haystack 套件

类型 4:其他(占 10%):包括量化精度损失、ONNX 导出失败、CPU 推理慢 10x 等,多为工具链问题,按官方 issue tracker 走即可。


本文为 2026 H2 大模型理论系列的第 N 篇。所有数值、特征值分布、复杂度声明均基于上述一手论文;标注"未公开验证的猜想"段落为作者推论,不是已发表结论。第 7 节生产落地清单与第 8 节事故案例基于行业公开博客与作者工程经验整理。

相关文章

  • 大模型训练动力学的非平衡统计力学 2026:当 SGLD、SDE 与 SGD 撞上损失景观的隐式正则化时7月4日
  • 神经网络损失景观的拓扑学 2026:从模式连接、低损耗盆地到模型合并的几何基础7月3日
  • 注意力机制的秩坍缩与低秩瓶颈理论 2026:当 token mixing 撞上训练动力学的低秩瓶颈时7月2日

评论

加载评论中…

发表评论

返回文章列表