SSM 与 Mamba

状态空间模型(SSM, 如 S4/Mamba)把序列建模写成线性时不变系统: 连续形式 $\dot{h}(t) = Ah(t) + Bx(t)$, $y(t) = Ch(t)$; 用 ZOH(零阶保持)离散化得 $\bar{A} = e^{A\Delta}$, $\bar{B} = (e^{A\Delta} - I)A^{-1}B$, 于是推理是线性递推 $h_t = \bar{A}h_{t-1} + \bar{B}x_t$, 输出 $y_t = Ch_t$——每步 $O(1)$ 状态更新, 无 KV 缓存。时不变时($\bar{A}, \bar{B}, C$ 与输入无关)可展开为全局卷积 $y = \bar{K} * x$(其中 $\bar{K} = (C\bar{B}, C\bar{A}\bar{B}, \ldots)$), 训练可整段并行。Mamba 的关键是选择性机制: 让 $\Delta_t, B_t, C_t$ 依赖输入 $x_t$(如 $\Delta_t = \mathrm{softplus}(W_\Delta x_t)$), 类似门控, 能按内容决定“记住”还是“忽略”——代价是失去卷积形式, 改用并行扫描(associative scan)训练。复杂度: 自注意力 $O(L^2)$(时间与显存), Mamba $O(L)$(序列长度线性), 长序列下显存与速度优势显著。

长上下文建模(50k+ token)、边缘设备上的线性时间推理; 面试对比题高频出现: Mamba vs Transformer(复杂度、并行性、长程依赖)、Mamba vs LSTM(选择性 vs 门控)、SSM 离散化推导。

Transformer 的 $O(L^2)$ 注意力随序列长度平方增长, 长上下文训练/推理昂贵; RNN(LSTM)虽 $O(L)$ 但串行无法并行训练。Mamba 用线性递推保持 $O(L)$ 复杂度, 又借并行扫描实现训练并行化, 以选择性机制弥补“固定参数递推”表达力不足——在与同规模 Transformer 相当的质量下, 推理吞吐与长序列显存占用都大幅改善。

写出 SSM 离散化递推 $h_t = \bar{A}h_{t-1} + \bar{B}x_t$ 并推导连续→离散($\bar{A} = e^{A\Delta}$ 的 ZOH 推导), $\Delta$ 起什么作用?

训练并行化: 时不变 SSM 为什么等价于全局卷积 $y = \bar{K} * x$? 引入选择性后如何用并行扫描(associative scan)恢复并行训练?

选择性机制是什么: $\Delta_t, B_t, C_t$ 为什么必须依赖输入 $x_t$? 它与 LSTM 门控的异同(内容感知 vs 固定参数)?

复杂度对比: 自注意力 $O(L^2)$ 的时间与显存 vs Mamba 的 $O(L)$; 在 $L = 100k$ 长上下文下分别意味着什么?

为什么流行的是 Transformer + Mamba 混合架构(每几层插入 Mamba 块)? 纯 Mamba 的潜在弱点是什么?