Author by: 张嘉瑶
序列建模要处理的是“按顺序来的数据”,比如文本、语音、传感器时间序列等。早期的 RNN 一步一步地读,容易“记不住很久以前的内容”,训练也慢;LSTM 加了门控,记忆好一点,但结构更复杂、更费算力。Transformer 换了思路:不按顺序慢慢看,而是“所有位置两两对比、一起算”,所以训练很快、理解上下文很强。但它的计算量会随着序列长度“按平方增长”,序列一长就又慢又费内存。
为了解决“长序列”这个老大难问题,大家把目光转向结构化状态空间模型(Structured State Space Models, SSMs)。SSMs 借鉴了控制理论和信号处理中的经典状态空间表示,通过结构化的循环和状态空间表征,实现了对长序列的高效处理,其计算复杂度通常为线性或近线性。SSMs 的理论根基可以追溯到- 控制理论,核心思想是:不保存所有历史,而是用一个“压缩过的隐藏状态”不断更新、传递信息。这样做,计算量通常能做到“跟序列长度成正比”,更省时省内存。在这一方向上,S4 是重要的起点;后续的 Mamba、S5、Jamba 等进一步优化了计算效率、内存占用和推理速度,尤其适合特别长的序列。
整体演进是这样:RNN(高效但记忆弱)→ LSTM(记忆强一点但更复杂)→ Transformer(表达力强、训练快,但长序列太贵)→ SSM(在不牺牲表达能力的前提下,把计算成本拉回“线性级别”)。这反映了一个核心追求:既要强的表达能力,又要高效率。
更有意思的是:SSM 借鉴了控制理论(比如状态表示、连续时间到离散时间的转换、稳定性等),是一次“跨学科”带来的新思路。这种跨学科的方法可能为解决机器学习中长期存在的问题(如长程依赖建模和模型可解释性)提供新的途径,通过借鉴其他领域成熟的理论和 Mamba 的成功,也可能激励研究者进一步探索来自工程学、物理学等领域的概念,以设计出更鲁棒和高效的人工智能架构。
总之,S4 到 Mamba 的发展,展示了一个清晰方向:通过更聪明的结构设计,把原本“很贵”的长序列问题,变成“可负担、可扩展”的计算,让模型既快又能“记长事”。这篇文章就是要系统地说明它们的机制、优缺点和影响。
S4(Structured State Space for Sequence Modeling)模型根植于状态空间模型的数学理论。理解其工作原理,必须首先掌握 SSM 从连续时间系统到离散表示,再到其卷积形式的演化路径。
SSM 的基本定义源于现代控制理论,它描述了一个线性时不变(Linear Time-Invariant, LTI)系统。该系统通过一个
其中:
-
$\dot{x}(t)$ 是状态向量$x(t)$ 对时间$t$ 的导数。 -
$A \in \mathbb{R}^{N \times N}$ 是状态矩阵,它决定了系统内部状态的演化动态。 -
$B \in \mathbb{R}^{N \times 1}$ 是输入矩阵,它描述了外部输入如何影响状态。 -
$C \in \mathbb{R}^{1 \times N}$ 是输出矩阵,它描述了如何从状态中读取输出。 -
$D \in \mathbb{R}^{1 \times 1}$ 是直通矩阵,允许输入直接影响输出。在许多序列建模应用中,可以将其视为一个跳跃连接,通常设为$0$ 以便于分析。
这些矩阵
深度学习模型处理的是离散的序列数据(如单词 ID 或音频采样点),而非连续信号。因此,必须将连续时间的 SSM 方程进行离散化,将其转换为一个处理离散序列
其中,我们假定
这个离散形式的方程揭示了 SSM 的循环特性。它与 RNN 的更新规则在形式上高度相似:当前状态
LTI 系统的美妙之处在于其固有的对偶性。除了循环表示,离散化的 SSM 还可以被精确地表示为一个卷积操作。通过将循环方程从一个零初始状态
这正是离散卷积的定义。因此,整个输出序列
这个卷积核
卷积视角的意义极为重大。根据卷积定理,时域上的卷积等价于频域上的逐元素乘法。这意味着,我们可以先通过快速傅里叶变换(FFT)将输入序列
SSM 的这种对偶性,即同时拥有循环和卷积两种等价表示,并非巧合,而是 LTI 系统的一个基本性质。序列建模领域的历史演进在某种程度上体现了循环模型(如 RNN)与卷积模型(如 TCN)之间的范式竞争。RNN 擅长生成但训练缓慢,而 CNN 训练快速但感受野受限。S4 的构建者敏锐地意识到,SSM 的对偶性恰好完美地映射到了深度学习模型的两种核心操作模式:用于推理的序贯生成和用于训练的并行计算。因此,SSM 框架为这两种看似对立的范式提供了一座坚实的数学桥梁,使得模型能够同时继承两者的核心优势:用于推理的
尽管 SSM 在理论上提供了一个优雅的框架,但实践表明,简单地使用随机初始化的 SSM 矩阵
HiPPO(High-order Polynomial Projection Operator,高阶多项式投影算子)是一个通用的数学框架,其核心目标是在线地将一个连续的、无限长的输入历史函数
HiPPO 的核心思想是,在任意时刻
HiPPO 框架最核心的理论成果是证明了:描述这个最优多项式逼近的
这个方程的形式与 SSM 的状态方程
S4 模型具体采用的是 HiPPO 框架下的一个特定变体,称为 HiPPO-LegS(Legendre Scaled)。
该变体有以下两个关键特征:
- 基函数:它使用勒让德多项式(Legendre polynomials)作为逼近历史函数的多项式基。
- 度量:它采用一种随时间缩放的度量,这种度量会给近期的数据点赋予更高的权重,但同时确保整个历史都被纳入考量,而不是像滑动窗口那样丢弃旧信息。
通过 HiPPO-LegS 推导出的
此处的
传统的 RNN 将记忆视为一个离散的、局部的状态转移过程:$$state_{new} = f(state_{old}, input_{current})$$ 这种机制在数学上缺乏深层解释,并且其局部性是导致梯度消失的根源。HiPPO 框架则从一个根本不同的角度重新定义了记忆。它不再关注离散的状态转移,而是提出了一个全局性的问题:“如何用一个固定维度的系数向量,最优地逼近过去全部的连续函数历史?”。
HiPPO 推导出的 ODE,$c'(t) = Ac(t) + Bf(t)$,并非一个启发式的更新规则,而是对“最优系数如何随新数据的到来而演化”这一问题的解析解。因此,HiPPO 矩阵
HiPPO 框架为 SSM 提供了捕捉长距离依赖的理论保障,但直接使用 HiPPO 矩阵在计算上是不可行的,因为其状态更新需要
S4 论文的核心发现是,所有 HiPPO 框架推导出的矩阵都具有一种称为“正规加低秩”(Normal Plus Low-Rank, NPLR)的结构 1。一个矩阵
更具体地,S4 证明了 HiPPO 矩阵
其中:
-
$V \in \mathbb{C}^{N \times N}$ 是一个酉矩阵($V^{*}V = I$) -
$\Lambda \in \mathbb{C}^{N \times N}$ 是一个对角矩阵 -
$P, Q \in \mathbb{R}^{N \times r}$ 是低秩矩阵,其中秩 r 非常小(对于 HiPPO-LegS,r=1)
这一结构性洞察是 S4 所有效率优化的基石。它将一个看似复杂稠密的 HiPPO 矩阵,转化为一个由极其简单的对角部分和微小的低秩扰动构成的结构,为后续的算法优化打开了大门。
在推理阶段,模型采用 SSM 的循环形式,其计算瓶颈在于每一步都需要计算矩阵-向量乘积
首先,由于
关键的挑战在于计算逆矩阵项
伍德伯里恒等式指出:
将
-
$M = I - \frac{\Delta}{2}\Lambda$ 是一个对角矩阵,其逆矩阵$M^{-1}$ 的计算是$O(N)$ 的 - 恒等式右侧的内层逆矩阵
$(I + \dots)^{-1}$ 作用于一个$r \times r$ 的小矩阵上,计算成本极低
最终的表达式表明,$(I - \frac{\Delta}{2}A)^{-1}$ 与一个向量的乘积也可以在
因此,整个
训练阶段的瓶颈在于计算长度为
下表解析了 S4 高效卷积核计算算法的核心数学工具及其在算法中扮演的角色。
表 1:S4 卷积核计算算法的构成要素
| 数学工具 | 在 S4 算法中的作用 | 带来的复杂度效益 |
|---|---|---|
| 生成函数 | 将时域的卷积核计算问题,转化为在频域对生成函数进行求值的问题 | 使得可以利用 FFT/iFFT,将问题核心从矩阵幂转移到函数求值 |
| 伍德伯里矩阵恒等式 | 利用 A 矩阵的 DPLR 结构,将复杂的矩阵求逆运算分解为对角矩阵求逆和低秩项运算 | 避免了 |
| 柯西矩阵 | 识别出在所有单位根上批量求值生成函数的代数结构,将其形式化为柯西矩阵与向量的乘积 | 将问题规约到一个已有快速算法的经典数值问题上 |
| 快速傅里叶变换(FFT/iFFT) | 在计算出频域的核函数值后,高效地($O(L\log L)$)将其转换回时域的卷积核 | 实现了时域与频域之间的高效转换,是整个并行计算流程的最后一环 |
综上所述,通过这一系列精巧的数学变换,S4 成功地将卷积核的计算从
S4 模型的构建过程堪称通过数学等价性进行算法优化的典范。它始于一个计算成本极高的目标——计算一个由矩阵幂定义的卷积核。S4 并未对此进行任何近似,而是通过一连串精确的数学变换,每一步都解锁了一种更高效的计算范式:
- 循环
$\rightarrow$ 卷积:利用 LTI 系统性质,将序贯计算转化为可并行的卷积。 - 卷积
$\rightarrow$ 频域乘积:利用卷积定理,引入 FFT,将问题转移到频域。 - 矩阵逆
$\rightarrow$ DPLR 矩阵逆:利用伍德伯里恒等式,利用$A$ 矩阵的隐藏结构来分解复杂的求逆运算。 - 函数求值
$\rightarrow$ 柯西矩阵乘积:通过观察代数结构,将频域上的批量求值问题与经典的快速算法问题联系起来。
最终得到的算法在计算上与朴素版本的结果完全等价,但速度却快了几个数量级。S4 的卓越之处不仅在于最终的架构,更在于其背后严谨而巧妙的、环环相扣的数学推导与算法设计。
尽管 S4 在理论和算法上取得了巨大成功,但它并非完美无缺。后续的研究和应用揭示了其在数值稳定性、硬件适配性和功能表达能力方面存在的一些根本性局限。
S4 的一个关键问题是,虽然其卷积形式在训练时表现稳定,但在切换到循环形式进行自回归生成时,可能会出现数值不稳定的情况。
论文《It's Raw! Audio Generation with State-Space Models》深入探究了这一现象,并从经典控制理论的角度给出了解释。一个连续时间 LTI 系统稳定的充分必要条件是其状态矩阵
为了解决这个问题,后续研究提出了一些修正方案,例如在 NPLR 分解中施加更强的约束,将 $A = \Lambda - PQ^{}$ 的形式改为 $A = \Lambda - PP^{}$,这有助于增强模型的稳定性 。这一问题凸显了理论模型在有限精度计算机上实现时可能出现的微妙而关键的差异。
S4 训练效率的理论基石是存在近线性的柯西矩阵-向量乘积快速算法。然而,理论上的高效并不总能转化为实践中的高速。这些快速算法通常非常复杂,并且在主流的深度学习框架和硬件加速器(如 GPU)上没有得到原生的、高度优化的支持 。
这意味着,实际的 S4 实现在很大程度上无法达到理论上的
S4 最根本的限制源于其线性时不变(LTI)的本质。LTI 系统意味着其动态特性(由矩阵
论文《Hungry Hungry Hippos: Towards Language Modeling with State Space Models》(H3)通过一系列精心设计的综合性任务,精准地指出了 S4 的功能缺陷。研究发现,S4 难以完成两类对注意力机制而言轻而易举的任务:
-
关联性回忆(Associative Recall):即根据上下文中的提示,回忆起序列早期出现的某个特定信息。
-
跨序列比较(Token Comparison):即比较序列中两个或多个不相邻 token 的关系。
这些任务的共同点是它们都要求模型进行基于内容的推理。例如,回答“法国的首都是什么?”这个问题,需要模型将“法国”和“首都”这两个概念关联起来,并从知识库中检索出“巴黎”。注意力机制通过其基于查询(Query)和键(Key)的动态寻址机制,天然地擅长此类任务。而 S4 的卷积核在序列处理开始前就已经固定,它对所有输入应用的是同一个“滤波器”,无法根据输入内容动态地调整其信息处理方式。这种内容无关性(content-agnostic)的设计,使得纯粹的 S4 模型在许多需要深度语义理解的 NLP 任务上表现不佳 。
S4 模型之所以能实现卓越的计算效率,其根本在于严格遵循了线性时不变(LTI)的系统结构,这使得整个序列到序列的映射可以被表示为一个全局卷积。然而,全局卷积的定义决定了其卷积核必须独立于输入信号,它是一个在处理开始前就已确定的固定滤波器。相比之下,语言理解任务的本质却是高度依赖内容的。例如,识别句子中的主谓宾关系或解决指代消解问题,都要求模型理解词语的语义并动态地建立它们之间的联系,而不是简单地对输入信号应用一个固定的变换。注意力机制通过其内容驱动的查询-键匹配机制解决了这个问题。S4 的设计恰恰牺牲了这种内容感知能力,以换取计算上的易处理性。因此,可以说,赋予 S4 计算效率的 LTI 属性,也正是限制其在语言等复杂符号任务上表达能力的根本原因。这个固有的“表达能力差距”为 S4 的后继者们指明了必须弥补的方向。
-
Gu, Albert, and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces." arXiv preprint arXiv:2312.00752 (2023).
-
Gu, Albert, et al. "Combining recurrent, convolutional, and continuous-time models with linear state space layers." Advances in neural information processing systems 34 (2021): 572-585.
-
Gu, Albert, et al. "Hippo: Recurrent memory with optimal polynomial projections." Advances in neural information processing systems 33 (2020): 1474-1487.
-
Voelker, Aaron, Ivana Kajić, and Chris Eliasmith. "Legendre memory units: Continuous-time representation in recurrent neural networks." Advances in neural information processing systems 32 (2019).
-
Gu, Albert, Karan Goel, and Christopher Ré. "Efficiently modeling long sequences with structured state spaces." arXiv preprint arXiv:2111.00396 (2021).
-
Exploring Language Models. (n.d.). A visual guide to Mamba and state space models. Exploring Language Models. https://exploringlanguagemodels.substack.com/p/a-visual-guide-to-mamba-state-space-models






