修炼AI技能
AI 教程 / 大模型应用 / 状态空间模型（SSM）：超越注意力机制的高效序列建模

状态空间模型（SSM）：超越注意力机制的高效序列建模

自注意力机制把序列中任意两个位置直接相连，带来了强大的上下文建模能力，却也埋下计算与显存随序列长度平方级膨胀的隐患。当上下文从 2K 拉长到 128K、乃至百万级 token 时，Transformer 的 O(n²) 复杂度逐渐成为训练与推理的硬瓶颈。状态空间模型（State Space Model，SSM）从连续系统控制理论中借来一套成熟的序列建模语言，把序列处理重写为线性时不变（LTI）系统的离散化，从而以 O(n) 的复杂度扫描长序列。本文从 SSM 的数学形式讲起，依次拆解 S4、Mamba、RWKV、RetNet 四条主线，并在复杂度、长序列、推理成本三个维度上与 Transformer 做系统对比，最后给出适用场景与局限。

一、状态空间模型的数学形式

状态空间模型刻画一个连续动力系统：输入信号 u(t) 经过一个隐藏状态 x(t) 映射为输出 y(t)。其连续形式由一对方程给出：

状态方程：  x'(t) = A x(t) + B u(t)
观测方程：  y(t)  = C x(t) + D u(t)

其中 A 是状态转移矩阵（维度 N×N），B 把输入映射到状态，C 把状态映射回输出，D 是直接的跳跃连接（通常可省略）。当序列是离散的 token 流时，需要把上述连续系统离散化：引入时间步长 Δ，用零阶保持（zero-order hold）近似，得到离散参数：

A_bar = exp(Δ · A)
B_bar = (Δ · A)^(-1) · (exp(Δ · A) - I) · Δ · B

离散化之后，系统在时刻 t 的递归形式为：

x_t = A_bar · x_{t-1} + B_bar · u_t
y_t = C · x_t

这与循环神经网络（RNN）的递归结构完全一致：每一步只依赖上一时刻的隐藏状态 x 与当前输入，因此推理时每步计算是常数成本，与序列长度无关。问题在于，这种纯递归形式难以并行训练，而朴素卷积形式又会带来 O(n²) 或 O(n·N) 的开销，因此「如何让 SSM 既高效训练又高效推理」是后续所有工作的主线。

二、S4：结构化状态空间序列模型

S4（Structured State Space Sequence model）由 Albert Gu、Karan Goel 与 Christopher Ré 提出，论文标题为「Efficiently Modeling Long Sequences with Structured State Spaces」，arXiv 编号 2111.00396（ICLR 2022 杰出论文提名）。需要特别更正一处事实：本任务初稿给出的编号 2111.03947 实际对应一篇强化学习论文（「Exponential Bellman Equation and Improved Regret Bounds for Risk-Sensitive Reinforcement Learning」），并非 S4；经核验，S4 的正确编号是 2111.00396。

S4 的核心贡献在于对状态矩阵 A 的结构化参数化。早期连续 SSM（HiPPO 框架）理论上能建模长程依赖，但直接离散化需要 O(N²) 的状态维度计算且数值不稳定。S4 用一个低秩修正（low-rank correction）来参数化 A，使其可以被稳定对角化，从而把 SSM 的计算归结为求解一个 Cauchy 核（Cauchy kernel），配合定制 CUDA 核函数实现高效并行。结果是：S4 在 Long Range Arena 基准上全面领先，并首次解决了长度 16K 的 Path-X 任务（此前所有方法均失败），同时在图像、语言建模上大幅逼近 Transformer，而生成速度快约 60 倍。S4 确立了一个关键范式：用结构化约束换来计算效率，同时保持线性复杂度。需要强调的是，S4 仍是线性时不变系统，即无论输入是什么，A、B、C、Δ 都是固定的，这使它擅长捕捉「位置/频率」层面的长程规律，却难以根据 token 内容做差异化处理——这一点正是 Mamba 后续要突破的核心。

三、Mamba：选择性状态空间

Mamba 由 Albert Gu 与 Tri Dao 提出，论文标题「Mamba: Linear-Time Sequence Modeling with Selective State Spaces」，arXiv 编号 2312.00752（已核验）。Mamba 指出了此前 SSM 类模型的致命弱点：它们大多是线性时不变（LTI）的，即参数 A、B、C、Δ 不随输入变化，因此无法做「基于内容的推理」——模型对每个 token 一视同仁地传播或遗忘信息，难以处理离散、强语义的语言数据。

Mamba 的关键创新是「选择性」（selectivity）：让 SSM 参数成为输入的函数，使模型能根据当前 token 的内容，决定沿序列维度选择性地传播或遗忘信息。这一改动破坏了卷积可并行性，但作者设计了一个硬件感知的并行递归算法（结合核融合与重计算，减少对 HBM 的读写），在递归模式下依然高效。Mamba 还把架构简化为不含注意力、甚至不含 MLP 块的纯 SSM 主干。

根据论文报告，Mamba 推理吞吐量是同规模 Transformer 的约 5 倍，序列长度呈线性扩展，在真实数据上直至百万长度仍有效；在语言建模上，Mamba-3B 在预训练与下游评测中均超过同尺寸 Transformer，并追平两倍大小的 Transformer。值得注意的是，这种「选择性」带来的收益并非免费：参数随输入变化后，模型无法再被等价写成固定卷积核，训练时的并行性从卷积退回到了「并行扫描」（parallel scan），对内核实现与显存布局提出了更高要求，这也是 Mamba 官方把大量精力投入在硬件感知算法上的原因。

四、RWKV 与 RetNet：把 RNN 重新带回大模型

除了 SSM 主线，还有两条「线性循环」路线同样以 O(n) 复杂度挑战 Transformer，值得一并比较。

RWKV（Receptance Weighted Key Value）由 Bo Peng 等人提出，论文标题「RWKV: Reinventing RNNs for the Transformer Era」，arXiv 编号 2305.13048（已核验）。它借鉴线性注意力，用「接受度（Receptance）—权重（Weight）—键（Key）—值（Value）」四个向量重构了时间混合机制，使同一套参数既能以 Transformer 形式做并行训练，又能在推理时退化为 RNN，保持常数级的计算与显存。论文把模型扩展到 140 亿参数，是当时最大的稠密 RNN。

RetNet（Retentive Network）由孙玉涛、董力等人提出，论文标题「Retentive Network: A Successor to Transformer for Large Language Models」，arXiv 编号 2307.08621（已核验）。它提出 retention 机制，并理论推导出递归与注意力的联系，支持三种计算范式：并行（训练）、递归（O(1) 推理）、分块递归（高效长序列）。递归表示让解码吞吐、延迟与显存显著优于 Transformer，且无需 KV Cache。

五、与 Transformer 的复杂度对比

把四类模型放在同一张表里比较，差异一目了然：

模型           训练复杂度     推理每步成本    可否并行训练   长序列友好
Transformer    O(n²)          O(n) 需 KV Cache   是           否（显存爆炸）
S4            O(n·N)          O(1)            是           是
Mamba         O(n)            O(1)            是           是（百万级）
RWKV          O(n)            O(1)            是           是
RetNet        O(n)            O(1)            是           是（分块）

关键区别在两点。其一，注意力对 n 个位置两两计算相似度，训练与预填充阶段是 O(n²)；而 SSM/RWKV/RetNet 在训练时通过卷积或分块并行达到线性或近线性，推理时则是严格递归的 O(1) 每步。其二，Transformer 解码必须缓存全部历史 Key/Value（KV Cache 随 n 线性增长），循环类模型只需维护固定大小的状态向量，显存占用与序列长度解耦。

六、适用场景与局限

SSM 类模型最适合三类场景：一是超长序列建模，如基因组、音频、长文档与持续信号，这类任务中 Transformer 的 O(n²) 直接不可行；二是高吞吐、低延迟的推理服务，O(1) 每步成本与恒定显存对部署极其友好；三是边缘与资源受限环境，固定状态维度让显存可预测。

但局限同样明显。其一，纯递归模型在需要「全局双向注意力」的任务（如掩码语言建模、编码器式理解）上天生吃亏，多数 SSM 是单向的。其二，选择性与内容推理的引入削弱了部分可并行的卷积优势，训练效率依赖高度定制的 CUDA 内核。其三，生态与规模化验证仍落后于 Transformer：Transformer 有成熟的分布式训练、MoE、工具链与海量预训练语料，而 SSM 在超大参数规模下的涌现能力与指令跟随仍有待更多实证。其四，RWKV 等模型在极长上下文下的「状态衰减」可能导致远端信息被渐进遗忘，需要配合位置或重刷机制缓解。

小结

状态空间模型用连续动力系统的语言重构了序列建模：S4 以结构化对角化把 SSM 变成可高效并行计算的线性层，Mamba 进一步引入输入依赖的选择性以弥补内容推理短板，RWKV 与 RetNet 则用线性循环或 retention 机制把 RNN 重新带回大模型舞台。它们共同的卖点是 O(n) 训练复杂度与 O(1) 推理成本，对长序列与高吞吐部署极具吸引力；代价是单向性、生态成熟度与超大规模验证仍不及 Transformer。把它们视为 Transformer 的「高效替代」而非「全面替代」，是在工程选型中更稳妥的判断。

参考与延伸阅读

  • Albert Gu、Karan Goel、Christopher Ré，Efficiently Modeling Long Sequences with Structured State Spaces（S4），arXiv:2111.00396（已核验；本任务初稿所用的 2111.03947 实为强化学习论文，已更正）。
  • Albert Gu、Tri Dao，Mamba: Linear-Time Sequence Modeling with Selective State Spaces，arXiv:2312.00752（已核验）。
  • Bo Peng 等，RWKV: Reinventing RNNs for the Transformer Era，arXiv:2305.13048（已核验）。
  • 孙玉涛、董力等，Retentive Network: A Successor to Transformer for Large Language Models（RetNet），arXiv:2307.08621（已核验）。
  • Albert Gu 等，HiPPO: Recurrent Memory with Optimal Polynomial Projections，arXiv:2008.07669（S4 的理论基础，待核实）。
  • Tri Dao、Albert Gu，Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality，arXiv:2405.21060（SSM 与注意力的对偶性，待核实）。
本文累计阅读 — 次