S4模型详解：结构化状态空间模型剖析

本节目录 Contents

S4模型详解：结构化状态空间模型剖析#

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范式#

SSM的基本定义源于现代控制理论，它描述了一个线性时不变（Linear Time-Invariant, LTI）系统。该系统通过一个$N$维的连续时间隐状态向量$x(t) \in \mathbb{R}^N$来连接一维输入信号$u(t) \in \mathbb{R}$和一维输出信号$y(t) \in \mathbb{R}$。其动态过程由以下一对线性常微分方程（ODE）刻画：

\[ \dot{x}(t) = A x(t) + B u(t) \]

\[ y(t) = C x(t) + D u(t) \]

其中：

$\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$ 以便于分析。

这些矩阵$(A, B, C, D)$是模型的参数，通过数据驱动的方式学习得到。这种连续时间的表述是SSM的一个核心特征，它暗示了模型具有处理不同采样率信号的潜力，因为它在理论上是分辨率无关的。

离散化：连接连续理论与数字数据#

深度学习模型处理的是离散的序列数据（如单词ID或音频采样点），而非连续信号。因此，必须将连续时间的SSM方程进行离散化，将其转换为一个处理离散序列$(u_0, u_1, \dots)$的离散时间系统。S4论文中采用了一种称为双线性变换（Bilinear Transform）的离散化方法。给定一个固定的步长$\Delta$，该变换将连续系统参数$(A, B)$映射到离散系统参数$(\overline{A}, \overline{B})$。离散化后的SSM具有以下循环更新形式

\[ x_k = \bar{A} x_{k-1} + \bar{B} u_k \]

\[ y_k = C x_k \]

其中，我们假定 $\bar{C} = C$ 且 $\bar{D}$ 项被忽略。离散矩阵 $\bar{A}$ 和 $\bar{B}$ 通过以下公式（双线性变换/塔斯廷法）与原始连续矩阵 $A$, $B$ 及采样周期 $\Delta$ 关联：

\[ \bar{A} = (I - \frac{\Delta}{2}A)^{-1}(I + \frac{\Delta}{2}A) \]

\[ \bar{B} = (I - \frac{\Delta}{2}A)^{-1} \Delta B \]

这个离散形式的方程揭示了SSM的循环特性。它与RNN的更新规则在形式上高度相似：当前状态$x_k$由前一时刻的状态$x_{k-1}$和当前输入$u_k$共同决定。这一特性使得SSM能够以自回归的方式进行高效推理，每生成一个新 token，只需进行一次状态更新。

卷积视角：展开循环#

LTI系统的美妙之处在于其固有的对偶性。除了循环表示，离散化的SSM还可以被精确地表示为一个卷积操作。通过将循环方程从一个零初始状态 $x_{-1}=0$ 开始迭代展开，我们可以观察到任意时刻的输出 $y_{k}$ 是所有历史输入 $(u_{0},u_{1},\dots,u_{k})$ 的线性组合：

\[ y_{k}=\overline{C B}u_{k}+\overline{C A B}u_{k-1}+\cdots+\overline{C A}^{k}\overline{B}u_{0} \]

这正是离散卷积的定义。因此，整个输出序列 $y$ 可以被写成输入序列 $u$ 与一个特定卷积核 $\overline{K}$ 的卷积：

\[ y=\overline{K}*u \]

这个卷积核 $\overline{K}$，也称为SSM滤波器或脉冲响应，完全由离散系统矩阵 $(\overline{A},\overline{B},\overline{C})$ 定义。其第 $i$ 个元素（从0开始）为：

\[ \overline{K}_{i}=\overline{C A}^{i}\overline{B} \]

卷积视角的意义极为重大。根据卷积定理，时域上的卷积等价于频域上的逐元素乘法。这意味着，我们可以先通过快速傅里叶变换（FFT）将输入序列$u$和卷积核$\overline{K}$转换到频域，然后进行廉价的逐元素相乘，最后通过逆快速傅里叶变换（iFFT）得到最终的输出序列$y$。由于FFT算法的复杂度为$O(L \log L)$，这使得SSM的训练过程可以像CNN一样高效地并行化，从而避免了RNN训练时的顺序瓶颈。

SSM的这种对偶性，即同时拥有循环和卷积两种等价表示，并非巧合，而是LTI系统的一个基本性质。序列建模领域的历史演进在某种程度上体现了循环模型（如RNN）与卷积模型（如TCN）之间的范式竞争。RNN擅长生成但训练缓慢，而CNN训练快速但感受野受限。S4的构建者敏锐地意识到，SSM的对偶性恰好完美地映射到了深度学习模型的两种核心操作模式：用于推理的序贯生成和用于训练的并行计算。因此，SSM框架为这两种看似对立的范式提供了一座坚实的数学桥梁，使得模型能够同时继承两者的核心优势：用于推理的$O(1)$循环更新和用于训练的基于FFT的并行计算。S4模型的精髓就在于成功地将这一理论桥梁转化为一个实际可行且高效的深度学习架构。

HiPPO 框架：一种有原则的历史压缩机制#

尽管SSM在理论上提供了一个优雅的框架，但实践表明，简单地使用随机初始化的SSM矩阵$(A, B, C)$，其性能往往不尽如人意，甚至会重蹈RNN梯度不稳定的覆辙，无法有效捕捉长距离依赖 [5, 15]。这表明，SSM的性能高度依赖于其参数矩阵的特定结构。S4成功的关键之一，便是引入了HiPPO框架来构建一个具有优良记忆特性的状态矩阵$A$。

HiPPO：连续历史的在线压缩#

HiPPO（High-order Polynomial Projection Operator，高阶多项式投影算子）是一个通用的数学框架，其核心目标是在线地将一个连续的、无限长的输入历史函数$f(t)$压缩成一个固定大小的向量表示。

HiPPO的核心思想是，在任意时刻$t$，它都试图找到一个最优的多项式$g(t)$（例如，阶数为$N-1$），来逼近截至当前时刻的全部历史函数$f_{\le t}$。这种逼近的“最优性”是根据一个随时间变化的度量$\mu(t)$来定义的，该度量决定了历史中不同时间点的重要性。这个过程可以被形式化为一个在线函数逼近问题。

从最优逼近到线性常微分方程#

HiPPO框架最核心的理论成果是证明了：描述这个最优多项式逼近的$N$个系数$c(t)$的演化过程，可以被一个线性常微分方程（ODE）精确刻画：

\[ c^{\prime}(t) = A c(t) + B f(t) \]

这个方程的形式与SSM的状态方程$x'(t) = Ax(t) + Bu(t)$完全一致。这意味着，HiPPO框架为SSM的状态矩阵$A$和输入矩阵$B$提供了一种基于第一性原理的构造方法。通过这种方式得到的矩阵$A$被称为“HiPPO矩阵”。它不再是一个随机矩阵，而是内含了对历史进行最优压缩的动态属性。

HiPPO-LegS矩阵#

S4模型具体采用的是HiPPO框架下的一个特定变体，称为HiPPO-LegS（Legendre Scaled）。

该变体有以下两个关键特征：

基函数：它使用勒让德多项式（Legendre polynomials）作为逼近历史函数的多项式基。
度量：它采用一种随时间缩放的度量，这种度量会给近期的数据点赋予更高的权重，但同时确保整个历史都被纳入考量，而不是像滑动窗口那样丢弃旧信息。

通过HiPPO-LegS推导出的$N \times N$状态矩阵$A$具有一个明确的数学形式：

\[\begin{split} A_{n k} = \begin{cases} (2n+1)^{1/2}(2k+1)^{1/2} & \text{if } n > k \\ n + 1 & \text{if } n = k \\ 0 & \text{if } n < k \end{cases} \end{split}\]

此处的$A$矩阵是一个下三角矩阵。实验证明，仅仅将一个标准SSM中的随机矩阵$A$替换为这个特定的HiPPO-LegS矩阵，就能使其在长距离依赖任务上的性能得到巨大提升。这为S4的构建提供了坚实的经验基础。

传统的RNN将记忆视为一个离散的、局部的状态转移过程：$$state_{new} = f(state_{old}, input_{current})$$这种机制在数学上缺乏深层解释，并且其局部性是导致梯度消失的根源。HiPPO框架则从一个根本不同的角度重新定义了记忆。它不再关注离散的状态转移，而是提出了一个全局性的问题：“如何用一个固定维度的系数向量，最优地逼近过去全部的连续函数历史？”。

HiPPO推导出的ODE，$c'(t) = Ac(t) + Bf(t)$，并非一个启发式的更新规则，而是对“最优系数如何随新数据的到来而演化”这一问题的解析解。因此，HiPPO矩阵$A$不仅仅是一个“好的初始化”，它是一个精确的线性算子，支配着一个最优、连续历史压缩方案的内在动力学。这为SSM的结构提供了比LSTM或GRU的门控机制更为严谨和深刻的理论依据，而S4模型正是建立在这一坚实的理论基石之上。

S4架构：通过结构化矩阵实现计算效率#

HiPPO框架为SSM提供了捕捉长距离依赖的理论保障，但直接使用HiPPO矩阵在计算上是不可行的，因为其状态更新需要$O(N^2)$的计算复杂度。S4模型的真正突破在于，其作者发现并利用了HiPPO矩阵内部隐藏的特殊数学结构，从而设计出了一系列极其高效的算法。

核心洞察：正规加低秩（NPLR）结构#

S4论文的核心发现是，所有HiPPO框架推导出的矩阵都具有一种称为“正规加低秩”（Normal Plus Low-Rank, NPLR）的结构 1。一个矩阵$A$是NPLR的，意味着它可以被分解为一个正规矩阵和一个低秩矩阵的和。

更具体地，S4证明了HiPPO矩阵$A$可以表示为：

\[ A = V\Lambda V^{*} - PQ^{\top} \]

其中：

$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）

$V\Lambda V^{*}$部分代表一个正规矩阵（因为所有正规矩阵都可以被酉对角化），而$-PQ^{\top}$是低秩部分。在SSM中，相似变换（矩阵共轭）是一种等价关系，即参数组$(A, B, C)$和$(V^{-1}AV, V^{-1}B, CV)$定义的系统是等价的。利用这一点，我们可以对系统进行基变换，从而在不失一般性的前提下，将状态矩阵$A$视为一个更简单的形式：“对角加低秩”（Diagonal Plus Low-Rank, DPLR）：

\[ A = \Lambda - P Q^{*} \]

这一结构性洞察是S4所有效率优化的基石。它将一个看似复杂稠密的HiPPO矩阵，转化为一个由极其简单的对角部分和微小的低秩扰动构成的结构，为后续的算法优化打开了大门。

循环机制：实现$O(N)$的快速推理#

在推理阶段，模型采用SSM的循环形式，其计算瓶颈在于每一步都需要计算矩阵-向量乘积$\overline{A}x_{k-1}$。对于一个稠密矩阵$\overline{A}$，这需要$O(N^2)$的计算量。S4利用DPLR结构将这一复杂度降低到了线性级别$O(N)$。

首先，由于$A = \Lambda - PQ^{*}$是DPLR形式，那么$I + \frac{\Delta}{2}A$显然也是DPLR形式。一个DPLR矩阵与向量的乘积可以在$O(N)$时间内完成。

关键的挑战在于计算逆矩阵项$(I - \frac{\Delta}{2}A)^{-1}$与向量的乘积。直接求逆是昂贵的。S4通过应用**伍德伯里矩阵恒等式（Woodbury Matrix Identity）**巧妙地解决了这个问题。

伍德伯里恒等式指出：

\[ (M + UV^{\top})^{-1} = M^{-1} - M^{-1}U(I + V^{\top}M^{-1}U)^{-1}V^{\top}M^{-1} \]

将$I - \frac{\Delta}{2}A = (I - \frac{\Delta}{2}\Lambda) + \frac{\Delta}{2}PQ^{*}$代入恒等式，其中：

$M = I - \frac{\Delta}{2}\Lambda$是一个对角矩阵，其逆矩阵$M^{-1}$的计算是$O(N)$的
恒等式右侧的内层逆矩阵$(I + \dots)^{-1}$作用于一个$r \times r$的小矩阵上，计算成本极低

最终的表达式表明，$(I - \frac{\Delta}{2}A)^{-1}$与一个向量的乘积也可以在$O(N)$时间内完成。

因此，整个$\overline{A}x_{k-1}$的计算复杂度被成功降至$O(N)$，使得S4的自回归推理速度极快。

卷积机制：高效的并行训练#

训练阶段的瓶颈在于计算长度为$L$的SSM卷积核$\overline{K}$。朴素的计算方法（通过$\overline{K}_i = \overline{C}\overline{A}^i\overline{B}$）需要$O(N^2L)$的复杂度，完全不可行。S4设计了一套精妙的算法，将复杂度降低到接近线性的$\tilde{O}(N+L)$。

下表解析了S4高效卷积核计算算法的核心数学工具及其在算法中扮演的角色。

表1：S4卷积核计算算法的构成要素

数学工具	在S4算法中的作用	带来的复杂度效益
生成函数	将时域的卷积核计算问题，转化为在频域对生成函数进行求值的问题	使得可以利用FFT/iFFT，将问题核心从矩阵幂转移到函数求值
伍德伯里矩阵恒等式	利用A矩阵的DPLR结构，将复杂的矩阵求逆运算分解为对角矩阵求逆和低秩项运算	避免了$O(N^{3})$的求逆操作，使得每次函数求值能在$O(N)$内完成
柯西矩阵	识别出在所有单位根上批量求值生成函数的代数结构，将其形式化为柯西矩阵与向量的乘积	将问题规约到一个已有快速算法的经典数值问题上
快速傅里叶变换(FFT/iFFT)	在计算出频域的核函数值后，高效地($O(L\log L)$)将其转换回时域的卷积核	实现了时域与频域之间的高效转换，是整个并行计算流程的最后一环

综上所述，通过这一系列精巧的数学变换，S4成功地将卷积核的计算从$O(N^2L)$的灾难性复杂度降低到了高效的$\tilde{O}(N+L) + O(L \log L)$，从而实现了高效的并行化训练。

S4模型的构建过程堪称通过数学等价性进行算法优化的典范。它始于一个计算成本极高的目标——计算一个由矩阵幂定义的卷积核。S4并未对此进行任何近似，而是通过一连串精确的数学变换，每一步都解锁了一种更高效的计算范式：

循环 $\rightarrow$ 卷积：利用LTI系统性质，将序贯计算转化为可并行的卷积。
卷积 $\rightarrow$ 频域乘积：利用卷积定理，引入FFT，将问题转移到频域。
矩阵逆 $\rightarrow$ DPLR矩阵逆：利用伍德伯里恒等式，利用$A$矩阵的隐藏结构来分解复杂的求逆运算。
函数求值 $\rightarrow$ 柯西矩阵乘积：通过观察代数结构，将频域上的批量求值问题与经典的快速算法问题联系起来。

最终得到的算法在计算上与朴素版本的结果完全等价，但速度却快了几个数量级。S4的卓越之处不仅在于最终的架构，更在于其背后严谨而巧妙的、环环相扣的数学推导与算法设计。

批判性评估：S4的内在局限与实践障碍#

尽管S4在理论和算法上取得了巨大成功，但它并非完美无缺。后续的研究和应用揭示了其在数值稳定性、硬件适配性和功能表达能力方面存在的一些根本性局限。

自回归生成中的数值不稳定性#

S4的一个关键问题是，虽然其卷积形式在训练时表现稳定，但在切换到循环形式进行自回归生成时，可能会出现数值不稳定的情况。

论文《It's Raw! Audio Generation with State-Space Models》深入探究了这一现象，并从经典控制理论的角度给出了解释。一个连续时间LTI系统稳定的充分必要条件是其状态矩阵$A$为赫尔维茨矩阵（Hurwitz matrix），即其所有特征值的实部均为负数。然而，S4所采用的HiPPO-LegS矩阵以及其NPLR参数化方法，并不能保证学习到的矩阵$A$满足这一条件。这导致在循环迭代过程中，状态向量$x_k$的范数可能指数级增长，最终导致数值溢出。

为了解决这个问题，后续研究提出了一些修正方案，例如在NPLR分解中施加更强的约束，将$A = \Lambda - PQ^{*}$的形式改为$A = \Lambda - PP^{*}$，这有助于增强模型的稳定性。这一问题凸显了理论模型在有限精度计算机上实现时可能出现的微妙而关键的差异。

硬件与软件之间的鸿沟#

S4训练效率的理论基石是存在近线性的柯西矩阵-向量乘积快速算法。然而，理论上的高效并不总能转化为实践中的高速。这些快速算法通常非常复杂，并且在主流的深度学习框架和硬件加速器（如GPU）上没有得到原生的、高度优化的支持。

这意味着，实际的S4实现在很大程度上无法达到理论上的$\tilde{O}(N+L)$速度，其性能会受到具体实现和硬件库的限制。这种“硬件鸿沟”使得S4的实际训练速度可能不如理论分析那般理想。这也成为了后续模型（如Mamba）发展的一个重要动机，Mamba的设计从一开始就考虑了GPU的内存层级结构和计算特性，是一种“硬件感知”的架构设计。

内容推理的功能性缺陷#

S4最根本的限制源于其线性时不变（LTI）的本质。LTI系统意味着其动态特性（由矩阵$A, B, C$决定）是固定的，不随输入内容的变化而改变。这使得S4成为一个强大的信号处理器，但却是一个相对较弱的语言理解模型。

论文《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