【Transformer 与注意力机制】56｜状态空间模型：Mamba、S4 的线性复杂度路径

55｜Transformer 的根本局限 讲过，full attention 的二次关系、KV Cache 的线性增长和自回归串行性共同构成长序列瓶颈。状态空间模型（State Space Model, SSM）试图换一种思路：不让每个 token 看所有历史 token，而是用一个可更新状态携带历史。

这听起来像回到 RNN，但 SSM 不是简单复古。S4、Mamba 等工作把连续系统、结构化矩阵、选择性状态更新和硬件友好 scan 结合起来，试图获得线性复杂度、长序列记忆和可并行训练之间的平衡。

本篇能让你学会三件事：

1. 状态空间模型如何用状态携带历史；
2. S4 和 Mamba 分别解决什么问题；
3. 为什么线性复杂度很有吸引力，但不自动意味着更强。

一、状态如何携带历史

最小状态空间直觉可以写成：

\[ h_t = A h_{t-1} + B x_t \]

\[ y_t = C h_t \]

输入 \(x_t\) 更新状态 \(h_t\)，输出 \(y_t\) 从状态中读出。历史不以 token 列表形式保存，而被压缩进状态向量。

这和 attention 很不同。attention 保留所有历史 K/V，并让当前 token 直接查询；SSM 则把历史不断折叠进状态。优势是序列长度线性增长，代价是状态是否能保留足够信息。

二、从 RNN 到 SSM

RNN 也用状态携带历史，但传统 RNN 难以并行训练，长距离依赖也容易衰减。Transformer 成功的一大原因，就是摆脱了严格递归，训练时可以并行处理序列。

现代 SSM 的目标不是回到普通 RNN，而是在状态更新形式下获得更好的长序列建模和并行计算。S4 使用结构化状态矩阵，让长卷积形式可以高效计算，从而处理长序列。

这条路线的关键是：训练时尽量并行，推理时可以递归更新状态。它试图同时要 Transformer 的训练效率和 RNN 的推理效率。

三、S4：结构化状态空间

S4（Structured State Space Sequence Model）把状态空间模型参数化得足够结构化，使其能高效处理长序列。它在 Long Range Arena 等长序列任务上展示了强能力。

S4 的技术细节涉及 HiPPO、结构化矩阵和高效卷积实现。对本系列来说，重点是它证明了 attention 不是长序列建模的唯一可行路径。只要状态更新设计得足够好，模型也能捕捉长距离信息。

不过 S4 主要在特定长序列任务中出圈，距离通用大语言模型主线还有距离。Mamba 后来把选择性机制和硬件友好实现推到更接近 LLM 的语境。

四、Mamba：selective state space

Mamba 的关键词是 selective。传统 SSM 参数通常不随输入内容动态变化，Mamba 让状态更新的某些参数依赖当前输入。这样模型可以根据 token 内容选择写入、保留或忽略信息。

这很重要。语言不是均匀信号。标点、变量名、实体、指令、无关填充，它们对未来的重要性不同。选择性机制让状态更新更像有条件的记忆管理。

Mamba 还强调硬件友好的 parallel scan。线性时间如果不能高效跑在 GPU 上，实际意义会打折。Mamba 的工程设计正是为了让选择性状态更新可训练、可扩展。

五、SSM 的优势和挑战

SSM 的优势是长序列成本。它不需要保存完整 attention matrix，也不需要让当前 token 查询所有历史 K/V。推理时状态大小可以固定或受控，适合流式场景。

挑战是信息瓶颈。把所有历史压进状态，必然要决定保留什么、丢掉什么。attention 的优势恰恰是历史 token 显式可访问，模型可以按需检索。SSM 必须证明状态压缩不会损害复杂任务所需信息。

另一个挑战是生态。Transformer 有成熟训练配方、硬件 kernel、并行策略和对齐流程。SSM 要成为通用底座，必须在质量、稳定、部署和工具链上都接近或超过 Transformer。

六、混合架构可能更现实

未来不一定是 SSM 完全替代 attention。更现实的路径是混合：部分层用 attention 处理精确检索和全局交互，部分层用 SSM 处理长程状态和流式信息。

这种混合能利用两者优势。attention 保留显式访问历史的能力，SSM 降低长序列成本。MoE、检索、外部记忆也可能加入系统。

后 Transformer 时代也许不是“没有 Transformer”，而是 Transformer 不再独自承担所有序列建模责任。

七、关键概念回顾

• 状态空间模型：用可更新状态表示历史序列。
• S4：结构化状态空间模型，面向长序列高效建模。
• Mamba：带选择性状态更新的线性时间序列模型。
• selective mechanism：根据输入内容动态决定状态如何更新。
• parallel scan：让递归形式在训练中高效并行的计算技术。
• 信息瓶颈：历史被压缩进有限状态时必然面临取舍。

八、常见误解

8.1 “线性复杂度一定更强”

不一定。复杂度低只是成本优势，模型质量、训练稳定和任务适配同样重要。

8.2 “Mamba 就是 RNN 换名字”

不准确。Mamba 使用选择性状态空间和硬件友好 scan，不是传统 RNN 的简单复刻。

8.3 “SSM 已经完全替代 Transformer”

没有。SSM 很有前景，但通用 LLM 生态仍需要大量验证。

九、下一步

SSM 是后 Transformer 路线之一。下一篇横向看 RWKV、RetNet、线性注意力等其他降低复杂度的探索。

十、参考文献

1. Gu, A. et al. “Efficiently Modeling Long Sequences with Structured State Spaces.” ICLR 2022.
2. Gu, A. and Dao, T. “Mamba: Linear-Time Sequence Modeling with Selective State Spaces.” arXiv:2312.00752, 2023.
3. Gu, A. et al. “HiPPO: Recurrent Memory with Optimal Polynomial Projections.” NeurIPS 2020.
4. Dao, T. and Gu, A. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.” ICML 2024.

← 上一篇：55｜Transformer 的根本局限　|　下一篇：57｜RWKV / RetNet / 线性注意力 →