引言: 高效序列模型的探索之路

对长序列数据 (无论是音频波形、医学信号、文本,还是展平的图像) 进行建模,是机器学习领域的一项基础性挑战。多年来,循环神经网络 (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 的原则性设计相融合的模型。

最终,他们构建出这样一个模型:

  1. 简单: 其卷积核计算只需两行代码即可实现。
  2. 有理论支撑: 首次给出了对角方法为何有效的理论解释。
  3. 强大: S4D 在图像、音频、医疗时间序列等任务上与原始 S4 性能相当,并在挑战性的 Long Range Arena 基准中取得 85% 平均准确率。

让我们深入看看他们是如何做到的。

S4D 是一种简单却强大的对角状态空间模型。左图展示了其循环视图,可视为一组独立的一维系统。右图展示了其可解释的卷积核,只需两行代码即可实现。

图 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 一样好。这一令人振奋的经验结果引出了两个问题:

  1. 为什么这个特定的对角矩阵能成功,而随机矩阵却不行?
  2. 能否比 DSS 的自定义“复数 softmax”方法更简单地获得它?

这篇论文回答了这两个问题,并带来 S4D。


核心方法: 构建 S4D 的步骤

作者将对角 SSM 的设计拆解为三大部分:** 离散化**、卷积核计算参数化

1. 离散化: 从连续到离散

SSM 起始于连续时间,但数据是离散的。需将连续参数 \((\mathbf{A}, \mathbf{B})\) 转换为离散参数 \((\overline{\mathbf{A}}, \overline{\mathbf{B}})\),这一过程称为离散化。

常用方法:

  • 零阶保持 (ZOH)
  • 双线性变换

双线性变换和零阶保持 (ZOH) 离散化方法的公式。它们将连续 SSM 参数转换为离散版本。

图 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\) 次方,计算极其简单。

于是得到一个基于范德蒙矩阵的简洁公式: 对角 SSM 卷积核可高效地用范德蒙矩阵求解。

图 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 对比。

该表总结了 S4、DSS 和 S4D 的设计选择。S4D 融合了其他两者的优势。

图 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}\) 的动态。

定理 3 的可视化。随着 N 增大,对角近似 (S4D-LegS,图 b 和 c) 的基函数收敛到原始 S4 模型 (S4-LegS,图 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 \]

不同 S4D 初始化的特征值虚部。S4D-LegS (蓝) 为原版,S4D-Inv (橙) 使用简单公式近似,S4D-Lin (绿) 采用傅里叶频率模式。其他简单缩放律 (红、紫) 效果不佳。

图 6: 结构化且全面分布的虚部是关键;简单改动 (红、紫) 性能较差。

得到两条经验原则:

  1. 恒定负实部以控制衰减;
  2. 结构化并分散的虚部以覆盖丰富频率范围。

实验: 验证 S4D 性能

初始化公式极为重要

即便是微小偏差——例如缩放或随机化虚部——也会显著损害性能。 S4D 初始化的消融实验。微小改动 (如缩放或随机化虚部) 亦导致性能下降明显。

图 7: 偏离推导的特征值公式会降低多项任务性能。

S4D 与 S4: 旗鼓相当

基准任务包括:

  • **序列 CIFAR **(图像分类,序列长度 1024)
  • **语音命令 **(音频分类,样本长度 16000)
  • **BIDMC **(医疗时间序列回归,长度 4000)

在更大模型上的消融数据集结果。对角 S4D 在各任务中与完整 DPLR S4 高度竞争。

图 8: S4D 在性能上高度接近完整 S4;在某些任务 (语音命令) 中,S4D-Inv/Lin 更胜一筹。

终极考验: 长程竞技场 (LRA)

LRA 是专门用于测试长程依赖的基准套件。

长程竞技场 (LRA) 基准结果。S4D 在性能上与 S4 接近,S4D-Inv 平均分达 85.50%,显著优于 Transformer。

图 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 将推动更广泛的应用和新研究方向。这不仅是一次改进——更是向人人可用的强大实用序列建模迈出的重要一步。