重温状态空间模型SSM：HiPPO的高效计算（S4）

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

前面我们用两篇文章《重温SSM：线性系统和HiPPO矩阵》和《重温SSM：HiPPO的一些遗留问题》介绍了 HiPPO 的思想和推导——通过正交函数基对持续更新的函数进行实时逼近，其拟合系数的动力学正好可以表示为一个线性 ODE 系统，并且对于特定的基底以及逼近方式，我们可以将线性系统的关键矩阵精确地算出来。

此外，我们还讨论了 HiPPO 的离散化和相关性质等问题，这些内容奠定了后续的SSM工作的理论基础。

接下来，我们将介绍 HiPPO 的后续应用篇《Efficiently Modeling Long Sequences with Structured State Spaces》[1]（简称 S4），它利用 HiPPO 的推导结果作为序列建模的基本工具，并从新的视角探讨了高效的计算和训练方式，最后在不少长序列建模任务上验证了它的有效性，可谓 SSM 乃至 RNN 复兴的代表作之一。

基本框架

S4 使用的序列建模框架，是如下的线性 ODE 系统：

这里 ， 是转置共轭运算，如果是实矩阵的话，那就是单纯的转置。由于完整的模型通常还会带有残差结构，最后一项 可以整合到残差里边，所以我们可以直接假设 来稍微简化一下形式，但不会降低模型的能力。

该系统具备相似不变性，如果 是 的相似矩阵，即 ，那么代入整理得

将 视为一个整体替换原来的 ，那么新系统的变化是 ，但输出完全不改变。这意味着如果存在 的某个相似矩阵 使得计算更加简单，那么可以完全转到 中分析而不改变结果，这就是后面一系列分析的核心思路。

特别地，S4 将矩阵 A 选取为 HiPPO-LegS 矩阵，即

这个选择的特别之处在于，我们此前推导 LegS 所满足的 ODE 是 的形式，而 LegT 的 ODE 才是 的形式，所以现在就是说 LegT 的 ODE 搭配了 LegS 的 矩阵，因此首先要问的问题是：这样的组合会带来什么影响呢？比如它对历史的记忆是否跟 LegS 一样依然是完整的、平权的？

指数衰减

答案是否定的——S4 所选取的 ODE 系统，关于历史的记忆是指数衰减的，我们可以从两个角度理解这一点。

第一个角度是从《重温SSM：HiPPO的一些遗留问题》讨论过的变换出发，将 LegS 型 ODE 可以等价地写成：

所以设 就可以将 LegS 型 ODE 变成时间变量为 的 LegT 型 ODE，也就是 S4 所用的 ODE。我们知道，LegS 会平等对待每一处历史，但这前提是输入为 ，但 S4 的 ODE 相当于输入直接改为 ，此时对 做均匀离散化的话，结果就是每一处的权重不相等——假设 ，用概率密度的写法就是 ，即 ，即权重是 的指数函数，越新的历史权重越大。

第二个角度则需要多一点线性代数知识。同样在《重温SSM：HiPPO的一些遗留问题》我们说过 HiPPO-LegS 的矩阵 理论上是可以对角化的，并且其特征值为 ，于是存在可逆矩阵 使得 ，其中 ，根据相似不变性，原系统等价于新系统

离散化后（以前向欧拉为例）：

这里的 是每个分量都小于 1 的对角线矩阵，也就意味着每迭代一步，就将历史信息乘以一个小于 1 的数，多步叠加后，就呈现出指数衰减的效应。

离散格式

虽然指数衰减看上去没有 LegS 平等对待每一处历史那么优雅，但实际上没有免费的午餐，对于固定大小的记忆状态 ，在记忆区间越来越大时，LegS 平等对待每一处历史的做法反而会导致每一处历史都比较模糊，对于符合“近大远小”的场景反而得不偿失。此外，S4 型 ODE 右端没有显式地出现时间 t，这也有助于提供训练效率。

对 S4 型 ODE 的记忆性质心中有数之后，我们就可以着手下一步操作了。为了处理实际中的离散序列，我们首先要进行离散化，在上一篇文章中，我们给出了两种精度较高的离散格式，一种是双线性形式

它具有二阶的精度，S4 采用的就是这个离散化格式，也是本文接下来所探讨的格式。另一种是基于精确求解常输入的 ODE，得到

作者后面的作品包括 Mamba 都是用这个格式，此时一般都要假设 为对角矩阵，因为对于 LegS 的矩阵 ，矩阵指数算起来并不友好。

现在我们记：

那么就得到线性 RNN：

其中 是离散化步长，是人为选择的超参数。

卷积运算

在上一篇文章中，我们还提到了 HiPPO-LegS 的矩阵 具备计算高效的特点，具体表现为 或 跟向量 x 相乘，存在计算复杂度为 而不是一般的 的高效算法，但这仅仅意味着式（10）递归计算时比一般的 RNN 高效，而如果想要进行高效训练的话，单纯递归是不够的，需要探究并行计算方法。

线性 RNN 的并行计算有两种思路：一种是在《Google新作试图“复活”RNN：RNN能否再次辉煌？》介绍过的视为 Prefix Sum 问题，直接用 Upper/Lower、Odd/Even、Ladner-Fischer等 Associative Scan 算法进行计算，论文可参考《Prefix Sums and Their Applications》[2]；另一种是转化为矩阵序列和向量序列的卷积运算，利用快速傅里叶变换（FFT）来加速，这是 S4 的思路。但不管哪一种，它们面临共同的瓶颈：幂矩阵 的计算。

具体来说，我们一般会设初始状态 为 0，那么就可以写出：

其中*代表卷积运算，而

注意根据目前的约定， 和 都是标量，所以有 。我们知道，卷积运算可以通过（离散）傅立叶变换转换为频域的加法运算，然后再逆变换回来，它的复杂度为 ， 是序列长度。虽然复杂度看上去比直接递归的 要大，但是傅立叶变换是可以并行的，所以实际上计算速度要更快。

所以，现在问题是如何高效地计算卷积核 ，它需要计算幂矩阵 ，按定义计算的话复杂度还是相当大的。当然，如果只是计算 那倒不是什么问题，因为 是一个常数矩阵，给定 后 也是常数矩阵，不管它的幂多难算，都可以提前算好存起来。然而， 只是中间步骤，我们还要算 ，而 S4 将 视为训练参数，所以 没法提前算好，就是提前算好 效率还是不大够。

生成函数

在进一步分析之前，我们先来插入一个生成函数的概念，这是后面的高效计算的基础步骤之一。同时，对于不大了解卷积运算和离散傅立叶变换的读者，这也可以作为一个科普步骤，从中我们可以大致了解到傅立叶变换加速卷积运算的基本原理。

对于给定序列 ，它的生成函数就是将每个分量当成幂级数的系数来构建幂级数：

如果有两个序列 和 ，那么它们生成函数的乘积：

留意到了没有？ 的第 项系数（即 的系数），正好是 与 的卷积运算。如果我们有快速计算生成函数以及快速提取生成函数某一项系数的方法，那么就可以将卷积运算转换为生成函数，做普通乘法之后然后再提取相应的系数。

离散傅立叶变换（Discrete Fourier Transform，DFT）[3] 正是这样的一种构建生成函数的思路。首先注意到，如果我们只需要对 的不超过前 项做卷积运算，那么生成函数的求和上限不一定非得到正无穷，求和上限改为 也是可以的。针对这种需求，DFT 没有对所有 来计算生成函数，而是选取了特定的 4 进行计算：

提取系数的逆变换（Inverse DFT，IDFT）则是

DFT 和 IDFT 我们都可以通过快速傅里叶变换（Fast Fourier Transform，FFT）[4] 进行高效计算，大部分数值计算框架都已内置了相应函数，所以 DFT 和 IDFT 的计算在效率上没有问题。

但要注意，如果用 DFT 来计算卷积的话，需要稍微微调一下，因为 是周期函数，我们没法区分 和 ，而当我们将两个 L 项求和的 DFT 相乘时，结果会出现 的 项，它会跟 项混合，从而做 IDFT 时实则得到的是两项的系数相加，这样作为卷积结果来说是不正确的。

解决这个问题的方法是将 的 改为 （但求和还是 项求和），也就是增大它的周期，使得乘积结果都是单个周期内，即将 DFT 的定义改为

不过现成的FFT函数基本上都不支持单独调整周期，而是默认周期就是数组长度，所以等价的处理方式是在 后面拼接 个零再做常规的 DFT，得到乘积后做 IDFT，最后只取前 个结果。

从幂到逆

对于卷积核 ，我们有

可以发现，生成函数不仅可以加速卷积的计算，它还将原本的幂矩阵 的计算转化为逆矩阵 的计算。

什么样的矩阵 ，它对应的 比较容易计算呢？首先对角阵肯定没问题，如果 是对角阵，那么 也是对角阵，对角阵的逆直接将对角线元素都取逆即可。其次，如果 可以对角化为 ，即 ，那么 同样容易计算，因为

那 能不能对角化呢？这取决于 能不能对角化。如果 ，根据相似不变性，我们可以完全转到 的新系统去计算，而根据定义新的 为：

显然是一个对角阵。

那么 可以对角化吗？答案是理论上可以，实际上不行。理论上可以，是因为从理论上来说，几乎所有矩阵在复数域内都可以对角化，并且在上一篇文章已经给出了 LegS 的 特征值为 ，也就是连对角化后的对角矩阵我们都知道长什么样了。

实际上不行，是指对数值计算来说很难，因为数值计算要考虑精度、内存、时间等，只要三者之一超出了限度或容忍度，那么理论可行的算法在实际中就不成立。

对于 矩阵，实际上不行的主要原因是对角化 所需要的矩阵 存在数值不稳定问题，说白了也是计算机精度有限导致的。对于这一点，原论文直接不加解释地给出了矩阵 P 的解析解，然后进行验证，这显然不利于读者理解。下面笔者从特征向量计算的角度，给出另一个理解思路。

特征向量

的对角化等价于 - 的对角化，因为 的特征值全是负数，所以简单起见我们转而考虑 - 的对角化，它有 个不同的特征值 ，对角化它所需的矩阵就是其特征向量的堆叠，所以求 本质上是求特征向量。而对于已知特征值的矩阵，求解特征向量的直接方法是求解方程-。

上一篇文章中“计算高效”那一节，我们已经给出了 的第 个分量的计算结果：

所以 - 意味着

记 ，那么 ，稍加整理得

注意- 是一个不定方程，我们有一些灵活调整的自由度（即特征向量不是唯一的），由于 最大是 ，我们可以设 ，然后递归地往回推，直到 得到 ，此后 都有 ，而对于 ，则有

由于我们是想要证明 的数值不稳定性，那么观察一个特征向量即可，我们取 （如果 不是 3 的倍数，简单取个整即可，结论不变），那么

最后的 可以由 String 公式 [5] 得到。由该结果我们可以看到，对于 这个特征值，从 到 存在一个指数级别的衰减过程（反之则爆炸），那么特征向量的分量 到 也存在类似的衰减，在浮点数的有限精度内，是很难精确处理这样的特征向量的。所以，直接对角化 的矩阵 存在数值上的不稳定性。

对角低秩

除了对角阵外，当 可以低秩分解时，同样可以降低 的计算难度。这是因为我们有如下的 Woodbury 恒等式：

这里 ，推导过程利用了 。如果 ，那么理论上 的计算量就比 少得多，因此可以加速计算。特别地，如果 ，那么 就是一个标量的倒数，计算起来最简单。

然而，我们知道 A 是一个下三角阵，且对角线元素没有一个是零，那么它就一定是满秩矩阵。再结合上一节的结论，也就是说 即不低秩，对角化又存在实践上的困难，所以这都不适用，还有什么办法呢？有！利用上面的 Woodbury 恒等式，我们可以推出它更一般的版本：

这个结果告诉我们，如果 M 的逆比较容易算，那么它加/减一个低秩矩阵的逆也容易算。那什么样的矩阵逆比较容易算呢？又回到上一节的答案——对角矩阵。所以，我们可以想办法将 A 或者 往“对角+低秩”的形式上凑。

事实上，仔细观察就会发现， 矩阵本身就有“对角+低秩”的影子。在上一篇文章中，我们将 的定义等价地改写为：

其中 实质就是对角矩阵 ，而 则可以重写为低秩矩阵形式 ，其中 ，换句话说，如果没有 的规定，那么 A 本身就是对角矩阵减去低秩矩阵的形式了。

点睛之笔

虽然有了下三角阵的约束后，这个规律就不再适用了，但我们可以充分利用原本就有的 结构，来辅助构建新的可对角化矩阵。但不得不说，这个技巧相当机智，堪称点睛之笔，让人惊叹，再次为原作者点赞。具体来说，我们考虑 ：

这个新矩阵的对角线元素正好是 -，我们再加上 ，就得到

重点来了，可以看到这是一个反对称矩阵，所以它一定可以（在复数域中）对角化！于是我们就将 分解为了可对角化矩阵与低秩矩阵之和！可能有读者质疑，原本 就 一定是可对角化矩阵，但还是有数值稳定性问题，难道这个反对称矩阵的对角化不用担心数值稳定性问题吗？

重点的重点来了，反对称矩阵不单单一定可以对角化，它一定可以被正交矩阵（复数域叫做酉矩阵）对角化！酉矩阵一般数值稳定性都非常好，所以不用担心这个问题，这也就是为什么我们不直接对角化 ，而绕一圈来构建反对称矩阵的原因。

现在我们得到，存在对角矩阵 和酉矩阵 ，使得 ，从而

抛开脚手架，我们发现最终的结论可以简化为“ 同构于对角阵减去秩1矩阵”：存在酉矩阵 、对角矩阵 、列向量 ，使得：

注意“对角+低秩”的矩阵乘以向量是计算高效的，比如

相当于将 当成向量与x逐位相乘，而 则是 跟 先做内积，然后得到一个标量乘以向量 ，这些都可以在 内完成。

最后冲刺

有了 ，再次根据相似不变性，我们接下来的所有计算都可以转到 中进行，所以下面均设 。首先，对于 ：

留意到 ，其中 是对角阵，于是利用 Woodbury 恒等式得到：

仔细观察，这同样是“对角+低秩”的形式，再乘以 后就能完成 的计算，最终结果是两个“对角+低秩”矩阵的相乘，意味着它同样具有计算高效的特点，这个结果可以在递归推理中用到。

最后是并行训练所需要的卷积核，我们已经将它转化为生成函数（18），现在我们就可以来完成它的计算了。首先通过类似“通分”的操作可以证明：

于是代入 得到

这里 是个对角阵，于是再次利用 Woodbury 恒等式就可以完成计算：

这是关于 z 的标量函数。不过要注意一个细节，傅立叶变换所需要的实际是“截断生成函数”：

也就相当于 的 要换成 ，这里 是提前选定的最大训练长度。接下来，我们只需要代入 进行计算，结果就是 的 DFT，然后 IDFT 就得到 了，这个过程还可以转化为 Cauchy 核问题加速一下，但个人认为不是太核心，就不展开讨论了。

最后的最后，还有一个技巧，就是对于 有 ，此时只是相当于将 要换成 ，而 S4 将 当成训练参数，所以我们可以直接将 当成训练参数，事后再从中解出 用于推理，这样训练时就可以避免计算 了。

这里看上去我们也可以代入 直接计算卷积所用的 的 DFT，而不是迂回地先 IDFT 得到 ，然后拼接零再 DFT，但问题是此时 是一个不定值，我们没法将 看成单个训练参数，这会导致在训练过程中需要计算 ，计算量比较大（当然，如果训练过程中 是完全固定的，那么可以提前算出来，视情况而定）。

草草收尾

经过一通艰难的“长篇大论”，我们总算把 S4 中比较关键的数学细节都捋了一遍，希望能够对有兴趣了解 S4 的读者有所帮助。可以看到，S4 是对 HiPPO 的进一步补充和完成，它的关键一笔是提出了 等价于“对角+低秩”的矩阵形式，为剩余部分的分析奠定了基础。因为一开始 A 是分段定义的形式，而不是矩阵运算形式，这样的定义不利于应用现有的线性代数工具进行一般化分析。

由于 HiPPO 的推导是基于 是一维函数进行的，所以到目前为止，S4 的 也都还是标量。那么 S4 怎么处理向量序列输入呢？

非常暴力，它直接对每个分量独立地应用一遍前述线性 RNN，每个 RNN 使用不同的 参数，然后将结果拼接起来，这个做法直到作者最新的 Mamaba 依然还被应用。

当然，也有简化的做法，直接在单个 RNN 中处理向量输入，只需要相应地将 改为矩阵就行，这就是 S5 [6]（作者不是 Albert Gu 了），这种做法可以理解为单纯借用了 S4 的线性 RNN 形式以及 HiPPO 的矩阵 ，而抛开了 HiPPO 的其他细枝末节，也取得了不错的效果。

让人啼笑皆非的是，S4 提出了诸多精妙的数学技巧来简化和加速 的计算，结果从《Diagonal State Spaces are as Effective as Structured State Spaces》[7] 开始，原作者的后续工作包括 Mamba 基本上都抛弃了这部分内容，而是直接假设 为对角矩阵，这样 RNN 部分就跟《Google新作试图“复活”RNN：RNN能否再次辉煌？》介绍的 LRU 大同小异了。

因此，从当前最新的 SSM 及线性 RNN 的角度看，S4 及 HiPPO 系列工作某种意义上来说已经是“过时”了。很多讲解 Mamba 的文章从 HiPPO、S4 开始说起，从事后来说可谓是“大可不必”了。

当然，对于笔者来说，花那么长的篇幅去学习 HiPPO 和 S4，并不是简单为了理解或使用最新的 SSM 和 RNN 模型，而是通过学习 HiPPO 背后的假设和推导，了解线性系统的记忆方式和瓶颈，为将来构建新模型、新方法积累更多的思路。此外，HiPPO 和 S4 中诸多精妙的数学技巧也让人赏心悦目，并且也不失为提升数学能力的相当不错的练习题。

文章小结

本文介绍了 HiPPO 的后续之作 S4，它的关键之处是提出了“对角矩阵+低秩矩阵”的分解，从而实现了 HiPPO 矩阵的高效并行计算，本文主要对其中比较困难的数学细节做了介绍和推导。

参考文献

[1] https://papers.cool/arxiv/2111.00396

[2] https://www.cs.cmu.edu/~scandal/papers/CMU-CS-90-190.html

[3] https://en.wikipedia.org/wiki/Discrete_Fourier_transform

[4] https://en.wikipedia.org/wiki/Fast_Fourier_transform

[5] https://en.wikipedia.org/wiki/Stirling%27s_approximation

[6] https://papers.cool/arxiv/2208.04933

[7] https://papers.cool/arxiv/2203.14343

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：hr@paperweekly.site 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·

·