S4

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$。 这意味着什么？ ...

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 Model 简洁：一个 SSM Block 内部到底在做什么 上一篇文章介绍了 SSM 的数学基础和典型 Block 结构。本文将聚焦于 S4（Structured State Space Sequence Model）这个 SSM 家族中第一个真正取得工程成功的模型，逐个拆解 S4 Block 中每个组件的任务、输入、输出和设计目的。 本文的讲解方式类似讲解 Transformer Block：先讲整体架构，再逐个组件说明"它是什么、它做什么、它为什么存在"，不涉及数学优化和工程优化的细节。 一、S4 的整体定位 S4 于 2021 年提出，是第一个在长序列建模任务上全面超越 Transformer 的 SSM 模型。它在 Long Range Arena（LRA）基准上取得了全面领先的成绩——按原始 LRA 六项任务的算术平均来看，S4 约为 81.84%，而当时的 Transformer 变体约为 45%。 数据注记：有的总结材料里会把 S4 在 LRA 上的平均准确率写成约 86.1%。这个数字和上文按六项明细算出的 81.84% 并不完全吻合，可能来源于不同的报告版本、是否包含特定任务变体、加权平均还是算术平均等差异。读者如果在写作或复习时引用这一数据，建议优先以逐任务明细为准，并注意区分“某一版本汇总的平均值”和“从明细算出的算术平均”。 S4 的核心贡献不在于改变了 SSM 的基本方程，而在于解决了两个实际问题： 如何初始化矩阵 $\mathbf{A}$： 通过 HiPPO 理论推导出理论上最优的 $\mathbf{A}$ 矩阵，使 SSM 从一开始就具备长程记忆能力。 如何高效计算： 通过 DPLR（对角加低秩）结构和 Z 变换，将卷积核计算的复杂度从 $O(LN^2)$ 压缩到 $O((L+N)\log(L+N))$。 但这些优化细节将在后续的优化篇中讲解。本文先专注于 S4 Block 的架构本身。 ...