SSM 优化(下):Cauchy 核、FFT 加速与工程优化

上一篇介绍了 S4 通过 DPLR 结构、双线性变换和 Z 变换将卷积核计算的复杂度从 $O(LN^2)$ 降到 $O(LN)$。但 $O(LN)$ 在超长序列($L = 16384$)和大状态维度($N = 256$)下仍然有约 $4 \times 10^6$ 的运算量。本文将介绍 S4 如何利用 Cauchy 矩阵的特殊结构和 FFT 进一步加速,以及训练时的反向传播策略和显存优化。最后讨论 S4 的根本局限性,以及它如何引向 Mamba 的诞生。

读者在阅读本篇时,可能会对一个结构细节有疑问:S4 中的状态到底是一个总的 $N$ 维向量,还是别的什么东西?这一点在跟 Transformer 对比时特别容易混淆,所以本文先把它澄清掉,避免后续关于“选择性机制”和 Mamba 的讨论建立在错误的前提上。

一、Cauchy 矩阵与快速求值

1. 一个先决问题:S4 真的只有一个 $N$ 维状态吗?

在很多地方,S4 被简略地说成是“一个 $N$ 维状态空间”。这其实是对整个模型结构的过度简化,容易让人误以为 S4 和 Transformer 的内部机制很像。

更准确的说法是:

  • S4 层通常作用于 $d$ 个特征维度(例如 $d = 768$ 或更大)。
  • 对每一个特征维度 $j$,都会独立地维护一个 $N$ 维状态空间。换句话说,每个特征维度有自己的 $\mathbf{A}_j$、$\mathbf{B}_j$、$\mathbf{C}_j$ 体系(实际参数化后可能会共享结构,但变量是独立可学的)。
  • 所以总的状态规模大约是 $d \times N$,但在每一个特征维度内部,状态维度 $N$ 仍然很小,典型值大概是 $16 \sim 128$。

这意味着什么?

  1. $d$ 个通道天然可并行。 每个特征维度的状态演化彼此独立,非常适合在特征维度上并行计算。
  2. SSM 内部不做“跨特征维度的内容交叉”。 这一点跟 Transformer 很不一样。在注意力机制里,所有特征维度通过 $QK^\top$ 直接参与交互;而在 S4 的状态演化阶段,各个特征维度是彼此独立演化的,特征之间的混合主要发生在 SSM 层之外的线性投影、门控之类的组件里,而不是在线状态演化本身。

读者如果把“SSM 的状态”误解为单个总的 $N$ 维向量,容易低估它的表达能力;如果误解为和 Transformer 一样做特征间的即时交互,又会高估结构上的相似性。准确的理解是:S4 在每个特征维度独立运行一个低维状态空间,总的状态量约为 $dN$,但内部耦合主要靠低秩项(在 DPLR 中)和外部混合组件来完成。

顺便把文章里早先提过的“$d$ 与 $N$ 的典型配置”再加一句:

  • 在这类模型里,特征维度 $d$ 通常会比较大,常见的范围大概是从 $2560$ 到 $4096$ 甚至更高,尤其是在 Mamba 那种偏向语言建模的配置里。
  • 而状态维度 $N$,即每个特征通道内部那个小状态的大小,通常只在 $16 \sim 128$ 之间。

这个反差本身就是 SSM 设计里很关键的一块:传统 RNN 为了存得下足够多的历史信息,往往会让隐藏维 $h$ 的大小跟输入特征维度 $d$ 差不多大,甚至更大;而 SSM 能把 $N$ 压到远小于 $d$,是因为 $\mathbf{A}$ 那边有数学上的设计,让它能用很小的状态去做历史信息的压缩和维持。

所以如果你在某处看到“SSM 的状态很小”这种说法,要理解的重点不仅是“有一个小状态”,而是:每个特征维度内部有一个小状态,但这样的小状态又有 $d$ 个,所以整体上状态规模大约是 $dN$,只是每个通道里的 $N$ 很小。 这也是为什么 SSM 在推理时能做到每步显存开销看起来很低——因为每个 token 真正需要额外维护的状态,主要是这个固定大小的 $h$,而不是随着序列长度膨胀的东西。

2. 问题的具体形式

上一篇结尾,我们将卷积核计算归结为:在 $L$ 个单位根 $z_k = e^{2\pi ik/L}$ 上求值

$$\hat{K}(z) = \mathbf{C}(z\mathbf{I} - \bar{\mathbf{A}})^{-1}\bar{\mathbf{B}}$$

当 $\bar{\mathbf{A}}$ 是对角矩阵 $\bar{\mathbf{\Lambda}} = \text{diag}(\bar{\lambda}_0, \bar{\lambda}1, \ldots, \bar{\lambda}{N-1})$ 时(先忽略低秩修正项),求值的核心运算为:

$$\hat{K}(z_k) = \sum_{n=0}^{N-1} \frac{C_n \bar{B}_n}{z_k - \bar{\lambda}_n}$$

令 $w_n = C_n \bar{B}_n$,对 $L$ 个不同的 $z_k$ 值求值,写成矩阵形式:

$$\begin{pmatrix} \hat{K}(z_0) \ \hat{K}(z_1) \ \vdots \ \hat{K}(z_{L-1}) \end{pmatrix} = \underbrace{\begin{pmatrix} \frac{1}{z_0 - \bar{\lambda}0} & \frac{1}{z_0 - \bar{\lambda}1} & \cdots & \frac{1}{z_0 - \bar{\lambda}{N-1}} \ \frac{1}{z_1 - \bar{\lambda}0} & \frac{1}{z_1 - \bar{\lambda}1} & \cdots & \frac{1}{z_1 - \bar{\lambda}{N-1}} \ \vdots & & \ddots & \vdots \ \frac{1}{z{L-1} - \bar{\lambda}0} & \frac{1}{z{L-1} - \bar{\lambda}1} & \cdots & \frac{1}{z{L-1} - \bar{\lambda}{N-1}} \end{pmatrix}}{\text{Cauchy 矩阵 } \mathbf{M} \in \mathbb{C}^{L \times N}} \begin{pmatrix} w_0 \ w_1 \ \vdots \ w{N-1} \end{pmatrix}$$

矩阵 $\mathbf{M}$ 的元素为 $M_{kn} = \frac{1}{z_k - \bar{\lambda}_n}$,这就是一个 Cauchy 矩阵。

3. 朴素计算的复杂度

直接做矩阵-向量乘法,对每个 $k$($L$ 个)和每个 $n$($N$ 个)做一次除法和加法,总复杂度为 $O(LN)$。

当 $L = 16384$、$N = 256$ 时,$LN \approx 4 \times 10^6$。对于单层单通道,这还可以接受。但 S4 有 $d$ 个并行通道(如 $d = 768$),且有多层,总计算量就不可忽视了。

4. 快速 Cauchy Matvec 的核心思想

Cauchy 矩阵有特殊的代数结构,可以利用快速算法加速。这里有两种主要方法:

方法一:快速多极子方法(FMM)。 FMM 将 $z_k$ 视为“观测点”,$\bar{\lambda}_n$ 视为“电荷”,$\frac{w_n}{z_k - \bar{\lambda}_n}$ 视为电荷在观测点产生的“电势”。对于距离很远的观测点和电荷群,不需要逐对计算,而是把一群远处的电荷近似为一个“等效电荷”(多极展开),一次计算它们对远处所有观测点的总贡献。总复杂度从 $O(LN)$ 降到 $O((L+N)\log(L+N))$。FMM 适用于 $z_k$ 和 $\bar{\lambda}_n$ 任意分布的情况。

方法二:FFT 加速。 在 S4 中,$z_k$ 恰好是均匀分布的单位根 $e^{2\pi ik/L}$。利用单位根的代数性质 $z_k^L = 1$,在单位根上求值有理函数等价于计算一个循环卷积,而循环卷积可以直接用 FFT 在 $O(L \log L)$ 内完成。

5. S4 实际使用的方法——以及 FMM 到底在哪里

S4 的官方实现选择的是 FFT 加速,而非 FMM。原因很简单:$z_k$ 是均匀分布的单位根,这是 FFT 加速的天然条件。FFT 算法实现简单、速度快、与 PyTorch 的 torch.fft 库无缝集成。

这里有个常见的困惑:既然 FMM 能把 Cauchy Matvec 做到 $O((L+N)\log(L+N))$,那它到底有没有被用在真实的 SSM 模型里?

更准确的说法是:

  • FMM 是 Cauchy Matvec 的通用快速算法,适合 $z_k$ 和 $\bar{\lambda}_n$ 任意分布的场景,在理论分析或非均匀采样的场合有意义。
  • 但在 S4 的实际代码中,由于 $z_k$ 恰好是均匀分布的单位根,反而不一定需要通用 FMM。单位根有一个特殊的代数性质 $z_k^L = 1$,这使得在单位根上求值有理函数,等价于计算一个循环卷积;而循环卷积可以直接用 FFT 在 $O(L\log L)$ 内完成。
  • 所以 S4 的代码里大量用的是 torch.fft 之类的 FFT 工具,而不是通用 FMM。

后续的一些模型甚至进一步绕开了这个问题:

  • S4D(Simple Diagonal) 直接让 $\mathbf{A}$ 为纯对角矩阵,去掉低秩修正项,此时卷积核 $\bar{K}_i = \sum_n C_n \bar{\lambda}_n^i B_n$,根本不需要 Cauchy Matvec,直接算就可以了。
  • Mamba 则完全放弃卷积视角,改用并行扫描,离散化后的递推直接计算。

结论:FMM 是存在的、也是有意义的,但在 S4 的主流实现里,真正用的是基于单位根的 FFT 加速;而在更后续的模型里,这个问题往往被简化掉了或绕过去了。

5. 完整计算流程

将所有步骤串联起来:

  1. 双线性变换离散化: 利用 Sherman-Morrison 公式在 $O(N)$ 内完成。
  2. 构造求值所需的"生成元": 单位根 $z_k$、对角元 $\bar{\lambda}_n$、权重 $w_n = C_n \bar{B}_n$。复杂度 $O(N)$。
  3. Cauchy Matvec(FFT 加速): 计算 $\hat{K}(z_k)$ 在所有单位根上的值。复杂度 $O(L \log L)$。
  4. 低秩修正项的求值: 类似地计算 Woodbury 公式中的第二项(rank-1 correction)。复杂度 $O(L \log L)$。
  5. 合并: $\hat{K} = \text{Term 1} - \text{Term 2}$。复杂度 $O(L)$。
  6. 逆 FFT: 恢复时域卷积核 $\bar{K}$。复杂度 $O(L \log L)$。

总复杂度为 $O(L \log L + N)$,对比上一篇的 $O(LN)$,在 $L \gg N$ 时有显著加速。

当 $L = 16384$、$N = 256$ 时,$L \log L \approx 16384 \times 14 \approx 2.3 \times 10^5$,比 $LN \approx 4 \times 10^6$ 小了一个数量级以上。

二、训练时的反向传播

1. 梯度计算的整体流程

训练时需要计算损失对所有可学习参数($\mathbf{\Lambda}$、$\mathbf{p}$、$\mathbf{q}$、$\mathbf{B}$、$\mathbf{C}$、$\Delta$)的梯度。

由于前向传播被实现为卷积操作 $Y = X * \bar{K}$,反向传播自然分为两部分:

对输入 $X$ 的梯度: 这是标准的卷积反向传播,深度学习框架自动处理。

对卷积核 $\bar{K}$ 的梯度: 框架先计算 $\frac{\partial \mathcal{L}}{\partial \bar{K}}$,然后通过链式法则传回基础参数:

$$\frac{\partial \mathcal{L}}{\partial \mathbf{A}} = \sum_i \frac{\partial \mathcal{L}}{\partial \bar{K}_i} \cdot \frac{\partial \bar{K}_i}{\partial \mathbf{A}}$$

对 $\mathbf{B}$、$\mathbf{C}$、$\Delta$ 同理。

3. 显存优化:重计算策略

问题: 如果在前向传播中存储完整的卷积核 $\bar{K} \in \mathbb{R}^L$(对于每个通道),显存开销为 $O(d \cdot L)$。当 $L = 16384$、$d = 768$ 时,单层就需要约 50 MB 的额外显存。多层累加后,显存消耗不可接受。

解决方案: 前向传播时不存储 $\bar{K}$,只存储参数 $\mathbf{\Lambda}$、$\mathbf{p}$、$\mathbf{q}$、$\mathbf{B}$、$\mathbf{C}$、$\Delta$。反向传播时重新计算 $\bar{K}$ 及其梯度。

这种策略的代价是计算量翻倍(前向和反向各算一次卷积核),但显存从 $O(dL)$ 大幅降到 $O(dN)$。在 GPU 上,由于算力远大于显存带宽,用时间换空间通常是很划算的交易。

4. 反向传播与重计算——再补充一个容易混的点

有读者可能会问:既然 SSM 这么像线性 RNN,那反向传播是不是要像传统 RNN 一样,沿时间步做 BPTT(Backpropagation Through Time)?

在 S4 的主流实现里,并不是这样。原因是前向传播被实现成了卷积 $Y = X * \bar{K}$,所以反向传播时,深度学习框架直接把它当作一个普通的一维因果卷积层来求梯度,而不是去显式地沿时间步展开反传。

这一点其实也呼应了第 3 篇提过的“双视角”思想:

  • 如果只用循环视角(逐步递推),反向传播就会退化为沿时间步的 BPTT,慢,而且容易梯度消失/爆炸。
  • 而 S4 把状态演化转换成了卷积视角,训练时的前向和反向都可以借助卷积相关的高效算子来做。

也就是说,“换成卷积视角”不只是为了训练速度快,背后还有一个工程上的反传策略差别:卷积形式下,框架能用成熟的卷积梯度路径,而不需要重新实现一套沿时间步的 BPTT。

顺便提一下,在更后续的一些实现(例如 mamba-ssm 这类工程库)里,还会用到**重计算(Recomputation)**之外的另一类思想:伴随状态法(Adjoint Method)。

  • 伴随状态法的核心思想是:不显式存储前向过程中间变量,而是引入“伴随变量”,通过求解某种伴随微分方程(或其离散对应)来得到梯度。
  • 在 SSM 这类问题里,它提供了一种不完全依赖于存储大量中间状态、也不完全依赖于直接反向展开的方式,来获得对参数的梯度。
  • 不过,本文目前讲的 S4 主线里,更常见的是“重计算 + 卷积反传”这条路线;伴随状态法是另一类思路,适合某些连续时间 / 数值求解场景,在 SSM 的不同变体或不同实现里可能会出现。

所以读者如果看到“重计算”这个词,不要把它和伴随状态法混为一谈:它们都是“不原样存前向中间量”的思路,但出发点和做法不完全一样。

三、S4 在各任务上的表现

1. Long Range Arena(LRA)

S4 在 LRA 全部六个任务上全面领先:

模型 ListOps Text Retrieval Image Pathfinder Path-X
Transformer 36.37 64.27 57.46 42.44 71.40 0.00
Performer 18.01 65.40 53.82 42.77 77.05 0.00
BigBird 36.07 64.02 59.24 40.83 74.87 0.00
S4 58.35 76.02 87.09 88.65 94.20 86.73
  • ListOps(序列长度 2000):S4 58.35%,Transformer 36.37%。
  • Text(序列长度 1000-4000):S4 76.02%,Transformer 64.27%。
  • Retrieval(序列长度 4000-8000):S4 87.09%,Transformer 57.46%。
  • Image(序列长度 1024):S4 88.65%,Transformer 42.44%。
  • Pathfinder(序列长度 1024):S4 94.20%,Transformer 71.40%。
  • Path-X(序列长度 16000):S4 86.73%,Transformer 0.00%。

在 Path-X 这个超长序列任务上,不仅标准 Transformer 完全崩溃,Performer、BigBird 这一类“高效注意力”方法也一样达到 0%;而 S4 仍然达到 86.73%,充分证明了 SSM 在超长序列建模上的优势。

2. 其他任务

  • 语音(SC09): S4 在原始音频波形生成上取得当时的最优结果。
  • 时间序列预测: 在 ETTh、ETTm 等基准上超越 Informer。
  • 像素级图像分类(CIFAR-10 sequential): S4 达到 90%+ 准确率,接近 ViT。

四、S4 的根本局限:线性时不变系统

1. 什么是 LTI

S4 是一个线性时不变(Linear Time-Invariant, LTI)系统。这意味着卷积核 $\bar{K}_i = \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}}$ 与输入内容完全无关。

无论输入是"请记住第一个词"还是"请忽略前面的所有信息",S4 都使用完全相同的记忆策略——相同的衰减速率、相同的关注模式。

2. LTI 的具体限制

无法选择性关注。 设想一个任务:从一段 10000 字的文本中,找到所有提到"爱因斯坦"的句子并复述。S4(LTI)无法根据输入内容动态调整关注策略,只能均匀地衰减所有历史信息。Transformer 的注意力机制可以通过 $QK^\top$ 直接"定位"到特定位置。

无法控制信息流。 在自回归生成中,有些信息(如系统指令、前文的约束条件)需要被长期保留,而有些信息(如过渡性填充词)可以快速遗忘。LTI 系统对所有信息使用相同的衰减曲线,无法区分这两类信息。

3. LTI 的数学根源

LTI 限制的根源在于卷积核 $\bar{K}_i = \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}}$ 的定义。$\mathbf{A}$、$\mathbf{B}$、$\mathbf{C}$ 是固定的参数,不随输入变化,所以卷积核也是固定的。

对比 Transformer:注意力权重 $\alpha_{ij} = \text{softmax}(q_i k_j^\top)$ 是输入依赖的(input-dependent),模型可以根据内容动态调整"关注哪里"。

4. LTI 的根本矛盾

SSM 能够使用卷积加速训练,恰恰是因为它是 LTI 的(参数不随时间步变化,卷积核可以预计算)。如果让参数依赖输入,就会破坏时不变性,卷积视角失效。

这是 S4 面临的根本矛盾:训练效率(需要 LTI)与表达能力(需要 input-dependent)之间的取舍。

五、从 S4 到 Mamba:打破 LTI 的天花板

Mamba(2023 年)通过引入选择性机制(Selectivity),巧妙地解决了这个矛盾。

1. Mamba 的核心思想

Mamba 让 $\mathbf{B}$、$\mathbf{C}$、$\Delta$ 成为输入 $x$ 的函数:

$$\mathbf{B}_k = \text{Linear}_B(x_k), \quad \mathbf{C}_k = \text{Linear}C(x_k), \quad \Delta_k = \text{softplus}(\text{Linear}\Delta(x_k))$$

这意味着在每个时间步 $k$,模型可以根据当前输入的内容,动态调整:

  • $\mathbf{B}_k$(写入方式):决定当前输入以什么方式和强度写入状态。
  • $\mathbf{C}_k$(读取方式):决定当前从状态中读取什么信息。
  • $\Delta_k$(步长):决定模型"关注"当前输入的程度——$\Delta$ 大意味着多关注当前输入,$\Delta$ 小意味着跳过当前输入。

2. 选择性的直觉

$\Delta_k$ 的作用类似于一个"信息阀门":

  • 当输入内容重要时(如关键词、实体名),$\Delta_k$ 变大,模型"打开阀门",将更多信息写入状态。
  • 当输入内容不重要时(如停用词、过渡句),$\Delta_k$ 变小,模型"关闭阀门",状态保持不变,不被无关信息污染。

这使得 SSM 具备了类似 Attention 的"选择性关注"能力,但计算复杂度仍然是线性的。

3. Mamba 如何保持高效训练

参数依赖输入意味着系统不再是 LTI,卷积视角失效。Mamba 的解决方案是放弃卷积视角,改用**并行扫描(Parallel Scan)**算法。

并行扫描是一种 $O(\log L)$ 深度、$O(L)$ 工作量的并行前缀和算法。它将 $L$ 个递推步骤 $h_k = \bar{\mathbf{A}}k h{k-1} + \bar{\mathbf{B}}_k x_k$ 组织成一棵二叉树,通过两轮遍历(上行聚合 + 下行分发)在 $O(L)$ 的总工作量内完成所有计算,且可以完全在 GPU 上并行执行。

4. Mamba 的推理优势

Mamba 在推理时的优势更加显著。由于每步只需要维护一个固定大小的状态 $h_k \in \mathbb{R}^N$($N$ 通常为 16 到 128),显存占用为 $O(1)$。对比 Transformer 的 KV Cache 显存占用为 $O(L)$,当序列长度达到 100K 时,Mamba 的显存效率优势是压倒性的。

这里有个细节,很多说法会把 SSM 描绘成“彻底没有 KV Cache”。更准确一点的说法是:SSM 推理时缓存的状态 $h$,本质上和 KV Cache 是同一类东西——都是对过去信息的压缩快照。 真正的区别不在于是不是缓存,而在于规模:

  • Transformer 的 KV Cache 通常是随生成长度线性增长的,每个时间步都把 Key 和 Value 存下来,长序列下容易变成显存瓶颈。
  • 而 SSM 推理时维护的 $h$,是一个固定大小的向量(维度大约是 $N$,并且按特征通道独立维护),不管已经生成了多少 token,它的规模都大致不变。

所以 SSM 的优势并不是“完全不用缓存”,而是把缓存压缩到了一个极小的、不会随序列长度膨胀的状态向量。这也是为什么它在长序列生成里显存占用能做到常数级(相对于 $L$ 而言)的核心原因之一。

实验表明,Mamba 在推理速度上比同等规模的 Transformer 快 5 倍以上,且序列越长优势越大。

六、S4 的后续发展

1. S4D(Simple Diagonal)

S4D 是 S4 的简化版本,直接令 $\mathbf{A}$ 为纯对角矩阵(去掉低秩修正项)。这进一步简化了计算,不需要 Cauchy Matvec,直接计算 $\bar{K}_i = \sum_n C_n \bar{\lambda}_n^i B_n$ 即可。实验表明,S4D 的性能与 S4 相当,但实现简单得多。

2. Mamba-2

Mamba-2(2024 年)从理论上证明了 SSM 与线性 Attention 的深层对偶性,大幅提升了训练速度。它还引入了更高效的硬件感知算法,使得 SSM 在现代 GPU 上的训练效率进一步提高。

3. Jamba 等混合架构

Jamba 等模型将 SSM 与少量 Transformer Attention 层和 MoE(Mixture of Experts)结合,走向工业级大模型应用。这种混合架构在长序列处理效率和模型表达能力之间取得了更好的平衡。

七、SSM 优化的完整知识图谱

回顾整个优化系列,SSM 从理论到工程的优化路径可以概括为:

理论层面:

  • HiPPO 理论:推导出理论上最优的 $\mathbf{A}$ 矩阵初始化,使 SSM 具备多尺度长程记忆能力。
  • LegT 测度选择:消除 HiPPO 方程的时变性,恢复卷积视角。

数学优化:

  • DPLR 结构:将密集 $N \times N$ 矩阵转化为对角加秩 1 修正,参数量从 $O(N^2)$ 降到 $O(N)$。
  • 双线性变换:替代 ZOH 的离散化方法,与 DPLR 结构天然配套。
  • ZOH(零阶保持):经典控制论中常用的方法,假设输入在 $\Delta$ 内恒定,离散化结果为 $\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$。它在 Mamba 中也会出现——特别是在选择性机制对应环节,Mamba 的实现里会涉及 ZOH 或其变体;但本系列的 S4 主线(依赖 DPLR 加速)并不以 ZOH 为主,其中因为矩阵指数通常会破坏 DPLR 结构,无法直接带来加速。关于这一点,第五篇博客的“三种离散化方法对比”表里已有注记,可供交叉参考;Mamba 篇后续若展开,会再次说明 ZOH 在选择性 SSM 语境下的位置。
  • Z 变换 + Sherman-Morrison:将卷积核计算转化为矩阵求逆,利用 Woodbury 恒等式在 $O(N)$ 内完成。

计算优化:

  • Cauchy Matvec + FFT:将求值复杂度从 $O(LN)$ 降到 $O(L \log L)$。
  • 重计算策略:反向传播时节省显存,用时间换空间。

架构演进(打破 LTI):

  • Mamba 的选择性机制:让 $\mathbf{B}$、$\mathbf{C}$、$\Delta$ 依赖输入,获得内容感知能力。
  • 并行扫描算法:替代卷积视角,保持高效并行训练。

七(续)、补一块容易被跳过的细节:推理时到底开不开 KV Cache

前面第六节提到 Mamba 的推理优势时,说它是 $O(1)$、比 Transformer 省显存。这句话本身没错,但容易让人得出一个不准确的结论:SSM 是“完全没有缓存”。

实际上,更准确的理解是:

  • Transformer 的 KV Cache 是一个随序列长度增长的缓存,里面存的是每一步的 Key 和 Value,长序列下显存会被这个缓存吃掉大部分。
  • SSM 的缓存是一个固定大小的状态向量 $h$,每一步只要存这个状态,长度为 $N$(并且按特征通道独立维护),不随已生成长度线性膨胀。

所以 SSM 并不是“没有缓存”,而是把缓存做成了一个极小的、基本不增长的状态向量。这也是为什么有些文章喜欢说“告别 KV Cache”——它的真正意思是“不用那种随长度线性膨胀的 KV Cache”,而不是字面上的完全没有缓存。

这一点在和 Mamba 的工程实现对接的时候也很重要:mamba-ssm 那类库里有很多重计算(Recomputation)相关的设计,目的之一恰恰是控制这类中间状态的显存开销,只不过 SSM 这边存的是状态向量,规模天然小得多。

八、总结

本文完成了 SSM 优化篇的最后拼图:

  • Cauchy 矩阵: SSM 高效计算中出现的特殊矩阵结构。利用 S4 中单位根的特殊性,可以用 FFT 在 $O(L \log L)$ 内完成矩阵-向量乘法,比朴素的 $O(LN)$ 快一个数量级。
  • 反向传播: SSM 的反向传播完全兼容现代深度学习框架,前向传播时的卷积操作使得反向传播也是标准的卷积梯度计算。
  • 重计算策略: 前向传播不存储卷积核,反向传播时重新计算,用两倍的计算量换取大幅的显存节省。
  • S4 的局限: LTI 系统无法根据输入内容动态调整记忆策略,这是 SSM 在语言建模等需要"选择性关注"的任务上不如 Transformer 的根本原因。
  • Mamba 的突破: 通过选择性机制打破 LTI 限制,同时用并行扫描算法保持高效训练,最终在语言建模等任务上达到了与 Transformer 相当的性能,同时拥有显著的推理效率优势。

从 HiPPO 的理论奠基,到 S4 的工程化实现,再到 Mamba 的选择性突破,SSM 在短短几年内完成了从学术创新到实用模型的跨越。这条优化之路仍在继续,Mamba-2、Jamba 等后续工作正在将 SSM 推向更广阔的应用场景。