引言: 高效序列模型的探索之路
对长序列数据 (无论是音频波形、医学信号、文本,还是展平的图像) 进行建模,是机器学习领域的一项基础性挑战。多年来,循环神经网络 (RNNs) 和卷积神经网络 (CNNs) 一直是标准工具。近年来,Transformer 异军突起,取得了惊人的成果。但所有这些模型在应对超长序列时都面临各种权衡。
这时,状态空间模型 (SSMs) 登场了。一种名为 **S4 **(Structured State Space for Sequences,结构化状态空间序列) 的新架构,凭借强大的长程记忆能力,在多项任务中超越了以往方法。S4 建立在经典控制理论坚实的数学基础上,通过一个名为 HiPPO 矩阵的特殊状态矩阵,高效建模连续信号——这是专为长时间保留信息而设计的数学构造。
然而问题在于: HiPPO 矩阵非常复杂。为了在深度学习中使用它,S4 采用了对角加低秩 (DPLR) 结构。这种表示方法虽然强大,但也使得 S4 更难理解、实现和定制——有时甚至像一个上锁的黑盒。
于是我们会问: 能不能简化它?如果直接用一个纯粹的对角状态矩阵,会发生什么?这样数学和代码都会变得极其简单。早期的简化尝试导致性能显著下降。然而,最近的 DSS 模型显示,一个特定的对角矩阵——直接源自 S4 自身的 HiPPO 结构——竟然能有惊人的表现。
这正是论文 《关于对角状态空间模型的参数化与初始化》 的切入点。研究者系统探索了如何构建、参数化并初始化这种更简单的对角 SSM,并推出了 S4D (“S4 on Diagonals”,基于对角矩阵的 S4) ——一个将对角状态矩阵的简单性与 S4 的原则性设计相融合的模型。
最终,他们构建出这样一个模型:
- 简单: 其卷积核计算只需两行代码即可实现。
- 有理论支撑: 首次给出了对角方法为何有效的理论解释。
- 强大: S4D 在图像、音频、医疗时间序列等任务上与原始 S4 性能相当,并在挑战性的 Long Range Arena 基准中取得 85% 平均准确率。
让我们深入看看他们是如何做到的。
图 1: S4D 继承了 S4 的优势,同时更为简洁。 (左) 对角结构使其可视为一组独立的 1D SSM。 (右) 卷积核可以约两行代码实现;不同颜色代表独立的 1D SSM,紫色表示可训练参数。
背景: 状态空间模型速览
从核心来看,状态空间模型通过一个隐藏的“状态”向量 \(x(t)\) 描述系统,该向量随时间按照线性常微分方程 (ODE) 演化:
\[ \begin{array}{l} x'(t) = \mathbf{A} x(t) + \mathbf{B} u(t) \\ y(t) = \mathbf{C} x(t) \end{array} \]其中:
- \(u(t)\) 是输入信号;
- \(y(t)\) 是输出信号;
- \(x(t)\) 是大小为 \(N\) 的隐藏状态;
- \(\mathbf{A} \in \mathbb{C}^{N \times N}\) 为状态矩阵——这是决定内部动态的关键组件。
该连续时间系统也可表示为卷积:
\[ \mathbf{K}(t) = \mathbf{C} e^{t\mathbf{A}} \mathbf{B}, \quad y(t) = (\mathbf{K} * u)(t) \]S4 的精妙之处在于选取了恰到好处的 HiPPO \(\mathbf{A}\) 矩阵,从而得到在长序列中具备优异记忆性能的核 \(\mathbf{K}(t)\)。但该矩阵的结构迫使我们使用复杂的 DPLR 表示。
DSS 的惊人发现是: 从 HiPPO 矩阵出发,先计算其 DPLR 形式,再丢掉“低秩”部分,保留的对角矩阵性能几乎与完整 S4 一样好。这一令人振奋的经验结果引出了两个问题:
- 为什么这个特定的对角矩阵能成功,而随机矩阵却不行?
- 能否比 DSS 的自定义“复数 softmax”方法更简单地获得它?
这篇论文回答了这两个问题,并带来 S4D。
核心方法: 构建 S4D 的步骤
作者将对角 SSM 的设计拆解为三大部分:** 离散化**、卷积核计算、参数化。
1. 离散化: 从连续到离散
SSM 起始于连续时间,但数据是离散的。需将连续参数 \((\mathbf{A}, \mathbf{B})\) 转换为离散参数 \((\overline{\mathbf{A}}, \overline{\mathbf{B}})\),这一过程称为离散化。
常用方法:
- 零阶保持 (ZOH)
- 双线性变换
图 2: 双线性变换和 ZOH 离散化将连续时间 SSM 参数转换为离散时间等价物。
实验表明,对于 S4D,采用哪种方法性能差异不大——两者都适用。这种灵活性进一步简化了设计。
2. 卷积核: 释放简单之力
在离散形式下,SSM 的卷积核为:
\[ \overline{\boldsymbol{K}}_{\ell} = \sum_{n=0}^{N-1} \mathbf{C}_n \,\overline{\mathbf{A}}_n^{\ell} \,\overline{\mathbf{B}}_n \]对于稠密 \(\mathbf{A}\),计算 \(\overline{\mathbf{A}}^\ell\) 成本昂贵——S4 因此采用复杂算法。但对角 \(\mathbf{A}\) 下,只需将每个对角元素提升到 \(\ell\) 次方,计算极其简单。
于是得到一个基于范德蒙矩阵的简洁公式:
图 3: 范德蒙矩阵公式使得对角 SSM 核计算快速高效——在无算法开销的情况下达到与 S4 相当的复杂度。
朴素实现时间复杂度为 \(O(NL)\),利用结构优化可降低到接近 S4 的近线性复杂度 \(\tilde{O}(N+L)\)。
3. 参数化: 细节是关键
简化计算后,关键设计问题包括:
- 参数化 \(\mathbf{A}\): 稳定性要求特征值实部全为负。S4D 通过设 \(\Re(\mathbf{A}) = -\exp(\mathbf{A}_{\mathrm{Re}})\) 强制实部为负;
- \(\mathbf{B}\) 与 \(\mathbf{C}\): 卷积核依赖于它们的逐元素乘积。DSS 直接学习该乘积,而 S4D 将二者分开,初始化 \(B=1\) 并训练 \(C\),若训练 \(B\) 会带来小幅稳定提升;
- 共轭对称性: 为确保实输入产生实输出,复特征值与参数以共轭对形式存储——简化实现并减少存储一半。
S4D: 两全其美
结合 DSS 的对角 \(\mathbf{A}\)、S4 的稳定参数化及高效范德蒙计算,S4D 成为简洁可控的对角 SSM,可直接与 DPLR S4 对比。
图 4: S4D 融合了 DSS 的对角结构、S4 的稳健参数化和高效核计算。
秘诀: 初始化决定成败
对角 SSM 在理论上表达能力很强,但随机初始化表现不佳——瓶颈在于优化。特征值的初始结构至关重要。
S4D-LegS: HiPPO 矩阵的奥秘
DSS 的对角矩阵来自 S4 的 HiPPO-LegS 矩阵: 取 \(\mathbf{A} = \mathbf{A}^{(D)} - \mathbf{P}\mathbf{P}^\top\),只保留对角部分 \(\mathbf{A}^{(D)}\)。
论文中 定理 3 证明: 当状态维度 \(N \to \infty\) 时,对角矩阵 \(\mathbf{A}^{(D)}\) 的动态会收敛到完整 HiPPO \(\mathbf{A}\) 的动态。
图 5: 随着 \(N\) 增大,S4D-LegS 的基函数收敛至 S4 的基函数——这解释了 DSS 的成功。
这种收敛性并非适用于任意低秩扰动——HiPPO 的结构有其独特性。
S4D-Inv 与 S4D-Lin: 更简单的配方
分析 \(\mathbf{A}^{(D)}\) 可发现其虚部服从一种反比缩放律:
\[ \text{S4D-Inv: } \quad A_n = -\frac12 + i\frac{N}{\pi}\left(\frac{N}{2n+1} - 1\right) \]更简单的 S4D-Lin 则使用等距虚部,类似傅里叶频率:
\[ \text{S4D-Lin: } \quad A_n = -\frac12 + i\pi n \]图 6: 结构化且全面分布的虚部是关键;简单改动 (红、紫) 性能较差。
得到两条经验原则:
- 恒定负实部以控制衰减;
- 结构化并分散的虚部以覆盖丰富频率范围。
实验: 验证 S4D 性能
初始化公式极为重要
即便是微小偏差——例如缩放或随机化虚部——也会显著损害性能。
图 7: 偏离推导的特征值公式会降低多项任务性能。
S4D 与 S4: 旗鼓相当
基准任务包括:
- **序列 CIFAR **(图像分类,序列长度 1024)
- **语音命令 **(音频分类,样本长度 16000)
- **BIDMC **(医疗时间序列回归,长度 4000)
图 8: S4D 在性能上高度接近完整 S4;在某些任务 (语音命令) 中,S4D-Inv/Lin 更胜一筹。
终极考验: 长程竞技场 (LRA)
LRA 是专门用于测试长程依赖的基准套件。
图 9: S4D-Inv 平均 85.50%——接近 S4 最佳成绩 (86.09%) ,远超 Transformer 基线 (53.66%) 。
结论: 简单性释放强大潜能
这项研究为用对角状态空间矩阵构建高效、强大的序列模型提供了完整指南。关键结论:
- 简单可行: 用对角 \(\mathbf{A}\) 替代 S4 的 DPLR,并不牺牲性能;核心计算实现极其简便;
- 初始化是核心: 性能秘诀不在于模型容量,而在于结构化、合理的起点 (S4D-LegS、S4D-Inv、S4D-Lin) ;
- 新的默认序列模型: 兼具简洁、理论与性能,S4D 有望在多个领域取代 RNN、CNN、Transformer。
让这些模型更易用,S4D 将推动更广泛的应用和新研究方向。这不仅是一次改进——更是向人人可用的强大实用序列建模迈出的重要一步。