状态空间模型（SSM）：从 RNN 与 Transformer 的困境出发

在序列建模的发展历程中，RNN 和 Transformer 是两座绕不过去的丰碑。RNN 以递推的方式逐步处理序列，结构简洁但受困于串行计算和长程依赖问题；Transformer 用自注意力机制彻底解决了长程依赖和并行化训练的难题，却引入了序列长度平方级的计算和显存开销。

这两种架构各有不可忽视的缺陷。有没有一种可能，让序列模型同时拥有 RNN 推理时的高效和 Transformer 训练时的并行？

这就是 SSM（State Space Model，状态空间模型）试图回答的问题。本文将从 RNN 和 Transformer 的痛点出发，解释 SSM 是什么、它为什么能做到"鱼与熊掌兼得"，以及它从 2020 年到 2024 年的演进路线。

一、RNN 与 Transformer 的痛点

1. RNN 的核心问题

RNN（包括 LSTM、GRU）的核心是隐藏状态的递推更新：$h_t = f(h_{t-1}, x_t)$。这种递推结构带来了两个根本性问题。

第一，训练无法并行。 由于 $h_t$ 强依赖于 $h_{t-1}$，序列中不同时间步的计算必须严格串行进行。在 GPU 这种为大规模矩阵并行运算设计的硬件上，RNN 的训练效率极低。假设序列长度为 $L$，RNN 的训练需要 $L$ 个串行时间步，无法像矩阵乘法那样在 GPU 上一次算完。

第二，长程依赖衰减。 即使有 LSTM 的门控机制缓解梯度消失问题，当序列长度达到数千甚至数万时，信息从第一个 token 传递到最后一个 token 仍然需要经过漫长的递推链路，每一次传递都会有信息损耗。就像传话游戏，传的人越多，最后听到的内容和最初说的偏差越大。

2. Transformer 的核心问题

Transformer 用自注意力机制解决了 RNN 的两个痛点：通过计算所有 token 之间的两两相关性，实现了全局信息的直接交互和完全并行的训练。但这个解决方案引入了新的代价。

序列长度的平方级复杂度。 自注意力需要计算 $QK^\top$，其中 $Q$ 和 $K$ 的形状都是 $[L, d]$（$L$ 是序列长度，$d$ 是特征维度），矩阵乘法的计算复杂度是 $O(L^2 d)$。内存中也需要存储 $L \times L$ 的注意力矩阵。当序列长度从 1000 增长到 100000 时，计算量增长了 10000 倍。

KV Cache 的显存爆炸。 在自回归生成（逐 token 生成）时，Transformer 需要把之前所有 token 的 Key 和 Value 缓存在显存中，即 KV Cache。对于一个 7B 参数的模型，处理 100K 长度的序列时，KV Cache 本身就可能占用数十 GB 的显存，成为推理的主要瓶颈。

线性 Attention 的妥协。 学术界尝试过多种降低注意力复杂度的方法，例如用核函数近似、低秩分解等手段将复杂度降到 $O(L)$。但这些"线性 Attention"在需要精确记忆、长距离推理的任务上，性能往往大幅下降，难以替代标准的二次方注意力。

二、SSM 的核心承诺

SSM 的出现，正是为了填补 RNN 和 Transformer 之间的空白。它做出了一个极其清晰的承诺：

训练时像 CNN 一样高度并行，推理时像 RNN 一样高效，同时能捕捉长程依赖。

这个承诺如何实现？核心思想可以概括为：同一个 SSM 模型拥有"双重身份"。

训练时，SSM 是一个卷积操作。 离散化后的 SSM 可以被数学等价地展开为一个全局卷积：$Y = X * K$，其中 $K$ 是一个长度为 $L$ 的卷积核。卷积操作可以通过 FFT（快速傅里叶变换）在 GPU 上高度并行地计算，复杂度为 $O(L \log L)$。

推理时，SSM 是一个线性 RNN。 同样的 SSM 可以折叠回递推形式：$h_k = \bar{A} h_{k-1} + \bar{B} x_k$。每生成一个 token，只需要维护一个固定大小的状态向量 $h_k$（维度为 $N$，通常只有 16 到 128），计算量和显存占用都是 $O(1)$（常数级），完全不随序列长度增长。

这两个视角在数学上是严格等价的，只是在不同阶段选择不同的计算方式。这种设计从根源上解决了 RNN 的训练瓶颈和 Transformer 的显存瓶颈。

三、SSM 是什么

1. SSM 的定义

SSM（State Space Model，状态空间模型）并不是深度学习领域原生的发明，它借鉴自经典控制理论和信号处理领域。在连续时间下，一个线性 SSM 由两个简洁的方程定义：

状态方程（记忆如何演化）：

$$h’(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$$

观测方程（如何从记忆中提取输出）：

$$y(t) = \mathbf{C}h(t) + \mathbf{D}x(t)$$

各变量的含义如下：

  • $x(t)$：当前时刻的输入信号，例如一个 token 的 Embedding 向量。
  • $h(t)$：当前时刻的隐藏状态，即模型的"记忆"。
  • $y(t)$：当前时刻的输出信号。
  • $\mathbf{A}$：状态转移矩阵，控制记忆如何随时间衰减或演化。
  • $\mathbf{B}$：输入投影矩阵，控制新输入如何写入记忆。
  • $\mathbf{C}$：输出投影矩阵，控制如何从记忆中读取输出。
  • $\mathbf{D}$：直接穿透矩阵，代表输入不经过记忆直接到达输出的路径。

在深度学习实践中，$\mathbf{D}$ 通常被设为 0 或极小常数，因为它不经过记忆，对序列建模没有实质贡献。我们主要关注 $\mathbf{A}$、$\mathbf{B}$、$\mathbf{C}$ 三个参数。可以用一句话概括三者的分工：$\mathbf{B}$ 负责写入，$\mathbf{A}$ 负责维持和遗忘，$\mathbf{C}$ 负责读出。

2. SSM 与 RNN 的对应关系

SSM 的方程和 RNN 的递推公式在结构上是对应的：

  • $x(t)$ 对应 RNN 的输入 $x_t$。
  • $h(t)$ 对应 RNN 的隐藏状态 $h_t$。
  • $y(t)$ 对应 RNN 的输出。
  • $\mathbf{B}$ 对应 RNN 中输入到隐藏状态的权重矩阵。
  • $\mathbf{A}$ 对应 RNN 中隐藏状态到隐藏状态的权重矩阵。
  • $\mathbf{C}$ 对应 RNN 中隐藏状态到输出的权重矩阵。

两者的本质区别在于：RNN 是离散时间步的、内部带非线性的递推；而 SSM 是连续时间的、内部纯线性的微分方程。SSM 的非线性被移到了层与层之间（类似 Transformer 的做法），而不是像 RNN 那样嵌在状态递推的内部。

3. SSM 不是一种独立的任务，而是一种基础架构

SSM 本身不是一种任务，而是一种序列处理算子，就像 Attention 或卷积一样。当我们将 SSM 作为核心组件堆叠起来，加上 Embedding 层和最后的 Language Modeling Head，它就是一个完整的语言模型。

例如 Mamba 就是一个纯粹基于 SSM 构建的大语言模型，它的任务是 Next-token Generation（下一个词预测），在语言建模基准上直接与 LLaMA 等 Transformer 模型竞争。

SSM 的通用性来自它处理的是"序列"：不仅适用于 NLP 文本，还适用于音频波形序列、视频帧序列、DNA 序列、金融时间序列等。它的核心优势是：无论序列多长，都能高效处理。

4. 两个维度：特征维度 $d$ 与状态维度 $N$

在 SSM 中有两个容易混淆的维度：

  • 特征维度 $d$（Feature Dimension）：输入 $x(t)$ 和输出 $y(t)$ 的维度，例如 768 或 4096。它代表每个 token 的语义表达能力，与 Transformer 中的 $d_{model}$ 对应。
  • 状态维度 $N$（State Dimension）：隐藏状态 $h(t)$ 的维度，通常很小，例如 16 到 128。

$h(t)$ 不是用来存储当前 token 的完整词向量的，它的作用是压缩并摘要历史信息。当新的高维输入 $x(t) \in \mathbb{R}^d$ 到来时，矩阵 $\mathbf{B} \in \mathbb{R}^{N \times d}$ 将它压缩为低维增量，写入状态 $h(t) \in \mathbb{R}^N$；当需要输出时，矩阵 $\mathbf{C} \in \mathbb{R}^{d \times N}$ 将低维摘要重新映射回高维空间，生成 $y(t) \in \mathbb{R}^d$。

$d$ 保证了模型表达单个 token 的能力，而极小的 $N$ 保证了无论序列 $L$ 有多长，推理时只需维护一个极小的状态向量。这就是 SSM 推理时显存占用极低的根本原因。

四、SSM 的发展历史

SSM 并不是凭空出现的，它有一条清晰的演进路线。

1. 2020 年：HiPPO

HiPPO（High-order Polynomial Projection Operators）解决了"如何让模型理论上记住无限长历史"的问题。它将"记忆"重新定义为"用正交多项式逼近历史信号"，并据此推导出了一个最优的记忆矩阵 $\mathbf{A}$。这是 SSM 一切后续发展的理论基石。

2. 2021 年：S4

S4（Structured State Space Sequence Model）将 HiPPO 理论工程化。通过"对角加低秩（DPLR）“的结构化技巧，让 SSM 真正能在 GPU 上高效训练。S4 在 Long Range Arena（LRA）长序列基准上全面超越了所有 Transformer 变体，首次证明了 SSM 在长序列建模上的巨大潜力。

3. 2022 年：S5 与 H3

S5 简化了 S4 的复杂数学结构，使其更易于实现和使用；H3 则尝试将 SSM 与线性注意力结合，首次在语言建模任务上让 SSM 逼近 Transformer 的性能。

4. 2023 年：Mamba

Mamba 是 SSM 的"破圈之作”。它引入了选择性机制（Selectivity），让 SSM 的参数能够根据当前输入的内容动态变化，解决了 S4 作为线性时不变系统无法"选择性关注"的局限。Mamba 在语言建模、基因组学等领域取得了当时的最优表现，并证明了其在真实硬件上的推理速度优势。

5. 2024 年及以后：Mamba-2、Jamba

Mamba-2 从理论上证明了 SSM 与线性 Attention 的深层对偶性，大幅提升了训练速度；Jamba 等混合架构（SSM + 少量 Attention + MoE）开始走向工业级大模型应用，代表了 SSM 从学术研究到工程落地的关键跨越。

五、总结

SSM 试图解决的核心问题可以概括为三个目标的同时达成：

  • 训练并行化：像 CNN 一样利用 FFT 加速，复杂度 $O(L \log L)$。
  • 推理常数级：像 RNN 一样只维护固定大小的状态，每步复杂度 $O(1)$。
  • 长程记忆：通过精心设计的矩阵 $\mathbf{A}$，理论上保证对长序列的稳定记忆。

从 HiPPO 理论到 S4 工程化，再到 Mamba 的选择性机制，SSM 在短短几年内完成了从理论突破到实用模型的跨越。接下来的系列文章将逐步深入 SSM 的数学基础、架构细节和优化技巧，带读者从零开始理解这个正在改变序列建模格局的新范式。