SSM 优化（中）：DPLR 结构与高效离散化

上一篇介绍了 HiPPO 理论如何推导出理论上最优的 $\mathbf{A}$ 矩阵，同时也指出了两个工程问题：HiPPO-LegS 的时变性和密集矩阵 $O(N^3)$ 的计算代价。

本文将详细介绍 S4 如何通过三个关键技术解决这些问题：用 DPLR 结构将密集矩阵转化为稀疏表示，用双线性变换替代 ZOH 离散化以保持 DPLR 结构，以及用 Z 变换将卷积核序列的计算转化为矩阵求逆问题。

一、DPLR 结构：将密集矩阵变为稀疏表示

1. 问题的起点

HiPPO 推导出的 $\mathbf{A}$ 是一个 $N \times N$ 的密集下三角矩阵。如果直接使用它，离散化后的 $\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$ 和卷积核 $\bar{K}_i = \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}}$ 的计算都需要密集矩阵运算，复杂度为 $O(N^3)$。

我们需要一种特殊的矩阵结构，使得幂运算和矩阵求逆都能在 $O(N)$ 内完成。

2. 什么矩阵的幂运算最快

在高等代数中，对角矩阵的幂运算是最简单的。如果 $\mathbf{\Lambda} = \text{diag}(\lambda_0, \lambda_1, \ldots, \lambda_{N-1})$，那么 $\mathbf{\Lambda}^i = \text{diag}(\lambda_0^i, \lambda_1^i, \ldots, \lambda_{N-1}^i)$，复杂度仅为 $O(N)$。

但纯对角矩阵意味着各个状态维度之间完全独立，没有信息交互，无法表达复杂的序列模式。

2.1 S4 内部的状态组织方式：每个特征维度独立维护一个 $N$ 维状态

S4 并非只使用“一个 $N$ 维状态”，这一点在与 Transformer 的对比中非常关键。

具体来说，S4 层通常作用于 $d$ 个特征维度（例如 $d = 768$ 或更大），而对每个特征维度 $j$，都会独立初始化并维护一个 $N$ 维的状态空间（即一个独立的 $\mathbf{A}_j \in \mathbb{C}^{N\times N}$、$\mathbf{B}_j$、$\mathbf{C}_j$ 体系，实际参数化后通常共享结构但有独立可学习变量）。

这意味着：

• 总的状态规模约为 $d \times N$。
• 但对于每一个特征维度而言，状态维度 $N$ 仍然很小（例如 16～128）。
• 各特征维度的状态演化在 SSM 内部是相互独立的——它们之间不像 Transformer 中的 Attention 那样，通过 $QK^\top$ 直接基于内容交叉交互。
• 特征维度之间的信息混合主要发生在 SSM 层之外的线性投影、门控等组件中，而不是在线状态演化本身。

这一设计带来两个后果：

1. 并行友好：$d$ 个独立通道天然可在特征维度上并行计算。
2. 与 Transformer 的对比鲜明：Transformer 中所有特征维度共同参与 Attention 的交叉计算；S4 则将“跨特征混合”从状态演化中解耦出来，交给线性层 / 门控处理。这是 SSM 实现高效序列混合的重要结构基础，也是后续 Mamba 等模型进一步发展的出发点之一。

读者若将“SSM 的状态”误解为单一 $N$ 维向量，容易低估其表达能力；若误解为与 Transformer 类似的特征间即时交互，又会高估其结构相似性。准确的理解是：SSM 在每个特征维度独立运行一个低维状态空间，总体状态量约为 $dN$，但内部耦合主要通过低秩项（在 DPLR 中）和外部混合组件完成。

3. DPLR 的定义

S4 提出了 DPLR（Diagonal Plus Low-Rank，对角加低秩）结构：

$$\mathbf{A} = \mathbf{\Lambda} - \mathbf{p}\mathbf{q}^\top$$

其中：

• $\mathbf{\Lambda} \in \mathbb{C}^{N \times N}$ 是复数对角矩阵，对角线元素 $\lambda_n$ 控制各独立状态的衰减率和振荡频率。
• $\mathbf{p}, \mathbf{q} \in \mathbb{C}^N$ 是两个 $N$ 维列向量，$\mathbf{p}\mathbf{q}^\top$ 是秩为 1 的矩阵，作为对角矩阵的"低秩修正"。

4. DPLR 的优势

参数量从 $O(N^2)$ 降到 $O(N)$。 不再需要存储和更新 $N \times N$ 个参数，只需存储 $\mathbf{\Lambda}$ 的 $N$ 个对角元以及 $\mathbf{p}$ 和 $\mathbf{q}$ 的 $2N$ 个元素。加上共轭对称的约束（详见下文），实际参数量进一步减半。

保留了全局耦合能力。 虽然主体是对角的，但低秩项 $\mathbf{p}\mathbf{q}^\top$ 引入了所有状态维度之间的交互。秩 1 矩阵的任意一行/列都与其他行/列线性相关，因此所有状态维度通过这个低秩项间接"通信"。

复数对角元的表达力。 在实数域中，对角元素只能表示纯指数衰减（实部）或纯振荡（虚部）。复数对角元允许同时表达衰减和振荡：$\lambda_n = a_n + i\omega_n$，实部 $a_n$ 控制衰减速度，虚部 $\omega_n$ 控制振荡频率。这使得 SSM 能同时捕捉序列中的趋势（低频）和周期性模式（高频）。

5. 共轭对称条件

由于输入 $x(t)$ 和输出 $y(t)$ 都是实数（或实数向量），但中间状态在复数域计算，为了保证输出为实数，DPLR 的参数必须满足共轭对称条件。

具体做法是将 $N$ 个状态分为 $N/2$ 对。对于第 $k$ 对：

$$\lambda_{2k} = a_k + i b_k, \quad \lambda_{2k+1} = a_k - i b_k = \overline{\lambda_{2k}}$$

$$p_{2k} = u_k + i v_k, \quad p_{2k+1} = u_k - i v_k = \overline{p_{2k}}$$

$$q_{2k} = s_k + i t_k, \quad q_{2k+1} = s_k - i t_k = \overline{q_{2k}}$$

可以验证：在这种共轭对称约束下，矩阵 $\mathbf{A} = \mathbf{\Lambda} - \mathbf{p}\mathbf{q}^\top$ 的实数域等价表示是一个实矩阵，输出 $y(t) = \mathbf{C}h(t)$ 也是实数。

在工程实现中，我们只存储 $N/2$ 个复数对角元素和 $N/2$ 个复数 $p$、$q$ 元素，另一半通过对称性自动生成，进一步减半了存储开销。

6. DPLR 如何保留 HiPPO 的记忆特性

S4 论文通过严格的矩阵分析证明（Theorem 3.2）：通过正交多项式的根和权重直接构造出的 DPLR 矩阵，其特征值与原始密集的 HiPPO 矩阵极其接近，并且其转移函数（Transfer Function）在低频段完美逼近 HiPPO 的理论最优解。

具体构造方法是：取 $N$ 阶勒让德多项式 $P_N(x)$ 的 $N$ 个实数根 $r_0, r_1, \ldots, r_{N-1}$ 作为对角矩阵 $\mathbf{\Lambda}$ 的元素（或在复数域推广中取复平面上的根），再根据高斯-勒让德求积公式（Gauss-Legendre Quadrature）的权重构造低秩向量 $\mathbf{p}$ 和 $\mathbf{q}$。

备注：本节涉及的 DPLR 构造以及后续的复数域运算，是 S4 系列博客中第一次较为完整地展开“可对角化结构 + 低秩修正”这一思想。对此不熟悉的读者可以暂时跳过构造细节，仅保留结论：“S4 用 DPLR 把 $O(N^3)$ 的密集矩阵运算转化为 $O(N)$ 的对角+低秩运算”。后文会在第五篇和第六篇中继续引用这一结论，但不会反复重复构造过程。

核心结论是：S4 不是把密集矩阵"分解"成了 DPLR，而是利用正交多项式的根和权重，直接从头构造了一个在数学性质上等价、但计算上极其高效的 DPLR 矩阵。

6.1 NUDT 与 DPLR 的关系：一个术语问题（避免混淆）

你在一些材料里会看到“NUDT”（Normal Plus Low-Rank，有时也写作 NPLR / Normal Plus Diagonal Plus Transpose 等）这个词。它和 DPLR 到底是什么关系？

简短结论：本质上讲的是同一类做法——只不过视角和侧重点略有不同。在 S4 实践中，最后真正落地的形式通常是 DPLR：$\mathbf{A} = \mathbf{\Lambda} - \mathbf{p}\mathbf{q}^\ast$。

• NUDT / NPLR 这类说法的核心意思：把 $\mathbf{A}$ 分成“结构良好的主体”（例如对角、正规、或者可对角化的部分）加上一个低秩修正，目的是在保留表达力的同时，让矩阵的幂运算 / 求逆变得可加速。
• S4 中最常用的具体形式是 DPLR：$\mathbf{A} = \mathbf{\Lambda} - \mathbf{p}\mathbf{q}^\ast$，其中 $\mathbf{\Lambda}$ 是复数对角矩阵，$\mathbf{p}\mathbf{q}^\ast$ 是秩为 $r$ 的低秩项，S4 中通常取 $r=1$，有时会推广到 $r=2$。
• 因此，“NUDT 是 DPLR 的统称”这种说法大致可以理解为：它们属于同一类“可对角化/正规结构 + 低秩修正”的矩阵家族，只是在不同材料里强调的命名不一样；而在深度学习 SSM 实践中，最自然、最常用的那个具体形式，往往就是 DPLR。

6.2 构造 DPLR 的真实动机：为什么选择“从正交多项式的根与权重出发”

这一点在速成资料里经常被跳过，导致读者产生误解：“S4 是把 HiPPO 密集矩阵分解成 DPLR 的”。

更准确的说法其实是：S4 并没有直接去分解原始的密集 LegS 矩阵，而是利用正交多项式的根和权重，从头构造了一个在关键性质上与 HiPPO 等价、但结构上对我们友好的 DPLR 矩阵。

为什么会这么做？

• 原始 HiPPO 矩阵（比如 HiPPO-LegS）是个 $N \times N$ 的实数、密集、下三角矩阵，数学性质很好但计算代价很高：涉及它的矩阵指数、矩阵求逆、矩阵幂运算都可能是 $O(N^3)$ 量级。
• 一种自然的想法是：“与其死磕那个稠密矩阵，不如换一个参数化方式，让它既有相似的频谱 / 记忆特性，又天生长在我们想要的代数结构里。”

S4 选择的方向大致是：

1. 用正交多项式的根作为对角元（正交多项式的根在数值上稳定，且与求积公式、谱近似密切相关，意味着对角部分从一开始就是“天然高效”的结构）；
2. 用求积权重构造低秩部分（高斯求积中每个根对应一个权重，S4 把权重信息吸收进低秩向量 $\mathbf{p}, \mathbf{q}$，从而在“对角 + 低秩”的形式里编码原本属于稠密矩阵的关键信息）；
3. 得到一个结构上容易加速的矩阵（对角部分 $O(N)$，低秩部分秩很小，矩阵求逆、求值、幂运算都能用 Woodbury / Sherman–Morrison 转化为对角运算加少量修正，大幅降复杂度）。

换句话说：

• S4 的关键不在于“先把旧矩阵拆解开”，而在于“用一套新的构造方式，直接得到一个在输入—输出行为上逼近旧矩阵、但结构上对加速极度友好的矩阵”。
• 这也就是为什么读者常常会觉得“S4 = 把密集矩阵分解为 DPLR”这个说法有点奇怪：实际上更准确的描述是“构造一个 DPLR，使它在关键性质上与 HiPPO 矩阵一致”。

二、双线性变换：替代 ZOH 的离散化方法

1. 离散化的回顾

SSM 的核心计算依赖离散化：将连续方程 $h’(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$ 转化为离散递推 $h_k = \bar{\mathbf{A}}h_{k-1} + \bar{\mathbf{B}}x_k$。

第二篇文章介绍了 ZOH 离散化，给出 $\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$。但 ZOH 在 S4 中有一个严重的问题：矩阵指数 $e^{\mathbf{A}\Delta}$ 不保持 DPLR 结构。即使 $\mathbf{A}$ 是 DPLR，$e^{\mathbf{A}\Delta}$ 通常会变成密集矩阵，所有优化都白费。

2. 双线性变换的推导

双线性变换（Bilinear Transform / Tustin’s Method）是另一种离散化方法，它的推导来自梯形积分法则。

将微分方程在区间 $[t_k, t_{k+1}]$ 上积分：

$$h_{k+1} - h_k = \int_{t_k}^{t_{k+1}} \left(\mathbf{A}h(t) + \mathbf{B}x(t)\right) dt$$

用梯形法则近似右边的积分（假设 $x(t) \approx x_k$ 在区间内不变）：

$$\int_{t_k}^{t_{k+1}} \mathbf{A}h(t) dt \approx \frac{\Delta}{2}\left(\mathbf{A}h_k + \mathbf{A}h_{k+1}\right)$$

$$\int_{t_k}^{t_{k+1}} \mathbf{B}x(t) dt \approx \Delta \mathbf{B}x_k$$

代入并整理，把含 $h_{k+1}$ 的项移到左边：

$$\left(\mathbf{I} - \frac{\Delta}{2}\mathbf{A}\right)h_{k+1} = \left(\mathbf{I} + \frac{\Delta}{2}\mathbf{A}\right)h_k + \Delta\mathbf{B}x_k$$

两边同乘 $\left(\mathbf{I} - \frac{\Delta}{2}\mathbf{A}\right)^{-1}$，得到：

$$h_{k+1} = \underbrace{\left(\mathbf{I} - \frac{\Delta}{2}\mathbf{A}\right)^{-1}\left(\mathbf{I} + \frac{\Delta}{2}\mathbf{A}\right)}{\bar{\mathbf{A}}} h_k + \underbrace{\left(\mathbf{I} - \frac{\Delta}{2}\mathbf{A}\right)^{-1}\Delta\mathbf{B}}{\bar{\mathbf{B}}} x_k$$

这就是双线性变换公式的来源。 它不是凭空出现的，而是梯形积分法则的直接结果。

3. 前向 Euler 法的完整推导（作为对比基准）

前向 Euler 法是最简单的离散化方法，它直接用前向差分近似导数。

推导过程： 回忆导数的定义： $$h’(t_k) = \lim_{\Delta \to 0} \frac{h(t_k+\Delta) - h(t_k)}{\Delta}$$ 取步长 $\Delta = t_{k+1} - t_k$，用前向差分代替导数： $$h’(t_k) \approx \frac{h_{k+1} - h_k}{\Delta}$$ 将其代入连续状态方程 $h’(t_k) = \mathbf{A}h(t_k) + \mathbf{B}x(t_k)$，并假设 $x(t_k) = x_k$： $$\frac{h_{k+1} - h_k}{\Delta} \approx \mathbf{A}h_k + \mathbf{B}x_k$$ 整理得到递推公式： $$h_{k+1} = (\mathbf{I} + \Delta\mathbf{A})h_k + \Delta\mathbf{B}x_k$$ 因此：$\bar{\mathbf{A}}{Euler} = \mathbf{I} + \Delta\mathbf{A}$，$\bar{\mathbf{B}}{Euler} = \Delta\mathbf{B}$。

精度与稳定性：

• Euler 法的精度是一阶的（截断误差 $O(\Delta)$）。
• 其稳定性是条件性的：只有当对所有特征值 $\lambda_n$ 都满足 $|1 + \lambda_n\Delta| < 1$ 时，离散化后的系统才是稳定的。
• 对于 HiPPO 矩阵，其特征值范围很广（从较接近 0 的衰减模式到非常负的值），很难保证全部特征值都落在该稳定条件中，因此 Euler 法在此场景下很难保证全部稳定。

备注：之所以在此完整列出 Euler 法，是为了后文与双线性变换和 ZOH 的对比更清晰——三种方法在“公式来源、精度、稳定性、是否需要矩阵指数、是否保持 DPLR 结构”上各有取舍。

双线性变换的精度是二阶的（误差 $O(\Delta^2)$），且具有无条件稳定性：只要连续系统 $\mathbf{A}$ 是稳定的（特征值实部为负），无论步长 $\Delta$ 取多大，离散化后的 $\bar{\mathbf{A}}$ 的特征值都在单位圆内。

4. 双线性变换为何与 DPLR 天然配套

双线性变换只涉及矩阵的加法和求逆，不涉及矩阵指数。而矩阵求逆恰好可以被 Woodbury 恒等式加速。

具体来说，$\bar{\mathbf{A}}$ 的计算核心是 $(\mathbf{I} - \frac{\Delta}{2}\mathbf{A})^{-1}$。由于 $\mathbf{A} = \mathbf{\Lambda} - \mathbf{p}\mathbf{q}^\top$：

$$\mathbf{I} - \frac{\Delta}{2}\mathbf{A} = \left(\mathbf{I} - \frac{\Delta}{2}\mathbf{\Lambda}\right) + \frac{\Delta}{2}\mathbf{p}\mathbf{q}^\top$$

这是一个"对角矩阵加秩 1 修正"的形式，完美匹配 Sherman-Morrison 公式的使用条件。求逆复杂度从 $O(N^3)$ 降到 $O(N)$。

三、三种离散化方法的对比

前文依次介绍了前向 Euler 法（§2.3）、ZOH（第二篇博客，作为背景知识）和双线性变换（§2.2）。为了便于横向比较，下表汇总了这三种方法在关键维度上的差异。

离散化方法	公式核心	精度	稳定性	需要矩阵指数	保持 DPLR 结构	实际使用
前向 Euler	$\bar{\mathbf{A}} = \mathbf{I} + \Delta\mathbf{A}$	一阶 $O(\Delta)$	条件稳定（$\max_n|1+\lambda_n\Delta| < 1$）	否	是（但结构过于简单，单独使用意义有限）	几乎不用
ZOH（零阶保持）	$\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$	理论上精确（假设输入在 $\Delta$ 内恒定）	无条件稳定（连续系统稳定则离散系统稳定）	是（矩阵指数）	否（矩阵指数通常破坏 DPLR 结构）	经典控制论常用；Mamba 中亦会用到（见下文说明）
双线性变换（Tustin）	$\bar{\mathbf{A}} = (\mathbf{I}-\frac{\Delta}{2}\mathbf{A})^{-1}(\mathbf{I}+\frac{\Delta}{2}\mathbf{A})$	二阶 $O(\Delta^2)$	无条件稳定（只要 $\mathrm{Re}(\lambda_n) < 0$）	否（只需矩阵求逆）	是（可通过 Woodbury/Sherman-Morrison 高效计算）	S4、S5、H3 等以 DPLR 为核心的模型

关于 Euler 精度的常见误解澄清：

坊间偶有资料将 Euler 法的精度表述为“一阶精度（误差 $O(\Delta^2)$）”，这在术语使用上是不准确的。

• Euler 法确实是一阶方法（方法的阶数为 1），但其每步截断误差是 $O(\Delta^2)$，积分区间固定时总误差为 $O(\Delta)$。
• 在本文中，我们按照惯例表述为“一阶精度（误差 $O(\Delta)$）”；若读者看到其他材料中称“误差 $O(\Delta^2)$”，那通常指的是单步局部误差，而非整体方法阶数。

为什么 ZOH 列为“经典控制论”，但后面又会和 Mamba 挂钩？

• ZOH 的物理直觉非常清晰：“假设输入在 $\Delta$ 内保持不变”，所以在经典控制论和很多连续系统建模场景里都被广泛使用。
• 矩阵指数 $e^{\mathbf{A}\Delta}$ 的计算通常会破坏 DPLR 结构（变成稠密矩阵），因此在 S4 这一类依赖 DPLR 加速的模型里，它不作为首选。
• 到了 Mamba（见后续博客），情况会变：Mamba 放弃了卷积视角，改用选择性机制 + 并行扫描，这时候离散化方法的选择逻辑就不完全是“要不要保持 DPLR 结构”了。
• 简单提前说一句：在 Mamba 的实际实现里，ZOH 或其变体会在某些环节出现，尤其是配合选择性机制的时候；但这不在本篇用 DPLR 加速为主线的讨论范围内。如果读者是冲着 Mamba 来的，可以参考后续章节里关于 Mamba 离散化的讲解。

S4 选择双线性变换不是因为它“数学上更优雅”，而是因为它与 DPLR 结构是天然配套的：把“矩阵指数问题”转化为了“矩阵求逆问题”，而后者可以被 Woodbury 恒等式（结合 DPLR 结构）高效加速。这一点将在后文 Z 变换与 Woodbury 加速部分进一步展开。

四、Z 变换：把序列问题转化为求值问题

1. 问题的转化

训练 SSM 需要计算长度为 $L$ 的卷积核序列 $\bar{K}_i = \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}}$（$i = 0, 1, \ldots, L-1$）。朴素方法是逐步计算 $\bar{\mathbf{A}}^i$，复杂度 $O(LN^2)$（因为 $\bar{\mathbf{A}}$ 是 $N \times N$ 矩阵）。

补充说明：推理时状态 $h$ 的本质与 KV Cache 的类比

推理阶段使用的循环视角 $h_k = \bar{\mathbf{A}}h_{k-1} + \bar{\mathbf{B}}x_k$ 中，每一步都只需要上一步的状态向量 $h_{k-1}$ 与当前输入 $x_k$。换句话说，推理时缓存的 $h$，在功能上与 Transformer 推理时缓存的 KV Cache 是同一类东西——都是对“过去信息的压缩快照”。 两者的区别主要在于规模与增长方式：

• Transformer 的 KV Cache 通常随生成长度线性增长，缓存内容是每一步的 Key 与 Value，总规模随序列长度 $L$ 扩张，容易在长序列生成中导致显存瓶颈。
• S4/S4 类模型推理时维护的状态 $h_k$ 是一个固定大小的向量（维度约为 $N$，且在每一特征通道独立），无论已生成多少 token，其规模大致保持不变。

因此，常有资料强调 SSM 模型“告别了 KV Cache”，这主要是针对 “不随序列长度线性增长的、极小状态缓存” 这一特性而言的；但若从功能抽象层面看，$h$ 本质上仍是一种极小的、用于摘要历史的缓存结构。

这一对比有助于澄清一个常见误解：SSM 并非完全没有“缓存”，而是将其替换为一个维度远小于 KV Cache 且不随生成长度扩张的状态向量。这也是 SSM 在推理阶段显存占用呈常数级（相对 $L$ 而言）的核心原因之一。

Z 变换的核心思想是：不逐项计算 $\bar{\mathbf{A}}^i$，而是利用生成函数将整个序列转化为一个封闭形式的表达式。

2. 生成函数的闭式求和

对卷积核序列做 Z 变换：

$$\hat{K}(z) = \sum_{i=0}^{\infty} \bar{K}i z^{-i} = \sum{i=0}^{\infty} \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}} z^{-i}$$

利用矩阵等比级数求和公式：

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

这个公式把"逐项计算 $\bar{\mathbf{A}}^i$“的问题转化为了"计算矩阵 $(z\mathbf{I} - \bar{\mathbf{A}})^{-1}$“的问题。

3. 从生成函数恢复时域序列

要从 $\hat{K}(z)$ 恢复时域序列 ${\bar{K}_0, \bar{K}1, \ldots, \bar{K}{L-1}}$，在 $L$ 个均匀分布的单位根 $z_k = e^{2\pi ik/L}$ 上求值 $\hat{K}(z)$，然后做逆 FFT：

$$\bar{K}i = \frac{1}{L}\sum{k=0}^{L-1} \hat{K}(z_k) e^{2\pi iki/L}$$

4. 计算瓶颈的转移

现在问题变成了：如何在 $L$ 个不同的 $z$ 值上高效求值 $\hat{K}(z) = \mathbf{C}(z\mathbf{I} - \bar{\mathbf{A}})^{-1}\bar{\mathbf{B}}$？

对于每个 $z_k$，核心运算是矩阵 $(z_k\mathbf{I} - \bar{\mathbf{A}})^{-1}$。由于 $\bar{\mathbf{A}}$ 继承了 $\mathbf{A}$ 的 DPLR 结构：

$$z_k\mathbf{I} - \bar{\mathbf{A}} = (z_k\mathbf{I} - \bar{\mathbf{\Lambda}}) + \bar{\mathbf{p}}\bar{\mathbf{q}}^\top$$

这完美匹配 Sherman-Morrison 公式。每个 $z_k$ 上的求值可以在 $O(N)$ 内完成，$L$ 个点的总复杂度是 $O(LN)$。

五、完整计算流水线

综合以上技术，S4 训练时计算卷积核的完整流水线如下：

第一步：双线性变换离散化。 利用 Woodbury 加速，将参数 $\mathbf{\Lambda}$、$\mathbf{p}$、$\mathbf{q}$、$\mathbf{B}$、$\mathbf{C}$、$\Delta$ 转化为离散化后的 $\bar{\mathbf{\Lambda}}$、$\bar{\mathbf{p}}$、$\bar{\mathbf{q}}$、$\bar{\mathbf{B}}$。复杂度 $O(N)$。

第二步：在 $L$ 个单位根上求值 $\hat{K}(z)$。 利用 DPLR + Sherman-Morrison 公式，对每个单位根 $z_k$ 在 $O(N)$ 内完成求值。$L$ 个点的总复杂度 $O(LN)$。

第三步：逆 FFT 恢复时域卷积核。 复杂度 $O(L \log L)$。

总复杂度为 $O(LN + L \log L)$。当 $L = 16384$、$N = 256$ 时，$LN \approx 4 \times 10^6$，在现代 GPU 上是完全可以接受的。对比朴素方法的 $O(LN^2) \approx 10^9$，加速了约 250 倍。

六、关于时变性问题的解决：测度选择、变量替换与“换测度”的完整逻辑

上一篇提到 HiPPO-LegS 的方程中含有 $\frac{1}{t}$ 因子，系统是时变的。读者自然会问：S4 到底是怎么解决这个问题的？是做变量替换把 $\frac{1}{t}$ 吸收掉，还是直奔一个时不变的测度？

事实上，两种思路都曾出现在相关讨论中，下面把它们的位置与逻辑关系补全。### 6.1 不同测度对应的记忆策略（四测度对比）

在 HiPPO 框架中，测度 $\mu_t$ 决定了对历史不同时刻 $\tau$ 的关注权重，因而也决定了系统是时变还是时不变。下表汇总了四个常被提及的测度变体及其特征。

测度名称	窗口	度量 $d\mu$	归一化因子	记忆衰减策略	对应基函数	时变/时不变
LegS（Scaled Legendre）	$[0, t]$（随时间增长）	$\frac{1}{t}d\tau$	$\frac{1}{t}$（随 $t$ 变化）	均匀关注整个历史（不衰减）	缩放后的勒让德多项式	时变
LegT（Translated Legendre）	$[t-T, t]$（固定长度 $T$）	$d\tau$	$\frac{1}{T}$（常数）	只关注最近 $T$ 个时间步（有限窗口）	平移后的勒让德多项式	时不变
LagT（Translated Laguerre）	$[0, \infty)$（半无限）	$e^{-(t-\tau)}d\tau$	不依赖 $t$ 的绝对值，只依赖时间差 $(t-\tau)$	指数衰减，越近越重要	平移后的拉盖尔多项式	时不变
FouT（Translated Fourier）	$[t-T, t]$（固定长度 $T$）	$d\tau$	$\frac{1}{T}$（常数）	均匀关注，衰减可选，更强调周期结构	复指数函数（傅里叶基）	时不变

选择测度 = 选择“记忆衰减曲线”。

• LegS：不衰减，所有历史同等重要（适合全局趋势，但引入时变性）。
• LegT：固定窗口，近期记忆长度受限，适合定长历史建模。
• LagT：指数衰减，越近越重要。
• FouT：不（或弱）衰减，改用三角基，更适合具有周期性的信号。

提示：LegT 的完整推导（包括窗口映射 $s = \frac{2(\tau-t)}{T}+1$、最优系数 $c_n(t)$ 的形式、对 $t$ 求导时上/下限边界项 $f(t-T)P_n(-1)$ 的出现与在 S4 实现中常被忽略的处理、LegT $\mathbf{A}$ 矩阵的严格下三角结构与全零特征值、以及 DPLR 重构版为何能稳定）已在本篇较前面的「6.2 LegT 的完整推导」一节展开。感兴趣的读者可跳回该节阅读；此处仅保留对照表与设计含义。

6.3 另一个方向：把 $\frac{1}{t}$ 吸收进 $\mathbf{A}$ 的变量替换（LegS 的备选处理方式）

另一种思路不是“换测度”，而是对 LegS 的时变方程做变量替换，把 $\frac{1}{t}$ 吸收进一个新的 $\mathbf{A}$ 对角线中，从而得到一个时不变的系统。

例如，可令 $$h(t) = t^{-\alpha} c(t)$$ （具体幂次 $\alpha$ 依赖于所研究的 HiPPO 变体及期望的归一化方式），通过对 $c(t)$ 的时变方程进行代入与求导，重新整理后有时可以得到关于 $h(t)$ 的时不变形式。

这一类做法的思想是：

• 不改变测度本身（仍使用 LegS 类型的窗口与度量）；
• 而是通过状态变量的缩放，把原本显式依赖时间 $t$ 的因子转移到系统矩阵的结构中。

这种处理在某些理论分析与特定变体中是有意义的，但在 S4 当时的工程实现主线中，更常见的选择是直接使用天然时不变的测度（LegT / LagT），从而避免额外引入变量替换带来的複雑性与潜在副作用。

6.4 “换测度”与“变量替换”之间的结论

归纳起来，面对 HiPPO-LegS 的时变性，主要有两条路：

1. 换测度：改用 LegT / LagT / FouT 等天然时不变的测度，从源头消除 $\frac{1}{t}$。这是 S4 工程实现中主要走的路径。
2. 变量替换：保留 LegS 测度，通过类似 $h(t)=t^{-\alpha}c(t)$ 的变换把时变因子吸收进新的 $\mathbf{A}$ 对角线。

两者并不互相排斥，而是代表了不同的设计偏好：

• 换测度通常更直接、更易与后续的 DPLR 构造与卷积视角对接；
• 变量替换则保留了对原有测度的分析，可能在特定理论场景或变体中使用。

再补充一点，很多读者会有一个隐含假设：“S4 肯定是用 LegS，然后想办法把时变性修好了”。其实不是这样的。

更准确的说法是：S4 实际上更倾向于不用 LegS，而是选 LegT、LagT 这些从一开始就天然时不变的测度。 因为一旦选了 LegT/LagT，后面就可以直接走卷积视角，不需要额外再去处理 $\frac{1}{t}$。

所以教程里经常出现“变量替换消掉 $\frac{1}{t}$”这种说法时，可以把它理解为：这是一种可能的思路，尤其在理论分析里有意义；但在 S4 当时的工程主线里，更常见的做法是直接选一个天然时不变的测度，省得再多绕一圈。

这里还有一个细节值得单独说清楚，就是为什么有些人会觉得 LegS 更“正统”。

• LegS 的窗口是 $[0,t]$，保证了对整个历史的均匀关注，非常适合理论分析。
• 但代价就是引入了时变性。
• LegT 的窗口是 $[t-T, t]$，固定长度，工程实现里反而更方便，但它是有限窗口的记忆。
• LagT 则是指数衰减的记忆，越近的越重要。
• FouT 则更强调周期结构，适合有明显频率成分的信号。

所以**“选哪个测度”本身就是在选记忆衰减策略**：你是想记全部历史（LegS）、最近 $T$ 步（LegT）、按指数衰减（LagT），还是更偏向周期结构（FouT）。S4 在工程上选 LegT/LagT，并不代表 LegS 在所有场合都应该被丢弃，而是说：如果你的目标是先把卷积视角稳稳地拿下来，天然时不变的测度就是一条更短的路。

七、S4 的可学习参数总结

在 S4 的代码实现中，被优化器更新的可学习参数（Leaf Tensors）是：

• $\mathbf{\Lambda}$ 的实部和虚部（$N$ 个复数对角元，共 $2N$ 个实参数，共轭对称后减半为 $N$ 个）。
• $\mathbf{p}$ 和 $\mathbf{q}$ 向量（共 $2N$ 个实参数，共轭对称后减半为 $N$ 个）。
• $\mathbf{B}$ 和 $\mathbf{C}$ 向量。
• $\Delta$（离散化步长）。
• 外部的线性层权重（输入投影 $\mathbf{W}{in}$、输出投影 $\mathbf{W}{out}$）。
• 小核卷积的权重。

密集的 $\mathbf{A}$ 和 $\bar{\mathbf{A}}$ 不是直接存储的参数，而是在前向传播时由上述参数"现场组装"出来的中间变量。所有数学技巧（DPLR、Sherman-Morrison、Z 变换）都是为了在组装和计算这些中间变量时，把复杂度压在可控范围内。

八、总结

本文介绍了 S4 的三个核心优化技术：

DPLR 结构： 将 HiPPO 的密集 $N \times N$ 矩阵转化为对角矩阵加秩 1 修正，参数量从 $O(N^2)$ 降到 $O(N)$，同时保留了全局耦合能力和多尺度记忆特性。复数对角元可以同时表达衰减和振荡。

双线性变换： 替代 ZOH 的离散化方法。它来自梯形积分法则，具有二阶精度和无条件稳定性，且不涉及矩阵指数，与 DPLR 结构天然配套。

Z 变换： 将卷积核序列的计算转化为有理函数求值问题。结合 DPLR + Sherman-Morrison 公式，每个求值点的复杂度为 $O(N)$，总复杂度从 $O(LN^2)$ 降到 $O(LN)$。

这三个技术的组合，使得 HiPPO 理论的优雅数学性质真正能够在 GPU 上高效实现，为 S4 在长序列建模上的突破奠定了基础。

下一篇将继续深入 S4 的高效计算细节，介绍 Cauchy 矩阵加速、FFT 卷积的工程实现，以及 S4 的反向传播策略和显存优化。