状态空间模型（深挖）

一句话理解

状态空间模型（State Space Model, SSM）是一类将序列建模从 O(n^2) 降到 O(n) 线性复杂度的架构。从经典控制论中的 SSM 出发，经过 S4（结构化状态空间）、H3（Hybrid H3）、Mamba（选择性 SSM）的演进，到 2026 年，SSM 已成为长序列建模（百万级 tokens）中最有竞争力的架构之一，在效率上对 Transformer 形成了实质性的补充。

数学基础：什么是状态空间模型

经典 SSM 定义

状态空间模型源自控制理论，描述一个系统如何随时间演化：

\dot{x}(t) &= A x(t) + B u(t) \quad &\text{(状态方程)} \\ y(t) &= C x(t) + D u(t) \quad &\text{(输出方程)} \end{aligned}$$ 其中： - $u(t)$：输入信号（序列中的 token） - $x(t)$：隐藏状态（模型对历史的压缩表示） - $y(t)$：输出信号 - $A, B, C, D$：可学习的参数矩阵 ### 离散化 在实际应用中，SSM 需要从连续时间离散化到离散时间步。最常用的离散化方法是用零阶保持（ZOH）规则，将连续参数 $(A, B)$ 转换为离散参数 $(\bar{A}, \bar{B})$： $$\bar{A} = \exp(\Delta \cdot A), \quad \bar{B} = (\Delta \cdot A)^{-1}(\exp(\Delta \cdot A) - I) \cdot \Delta B$$ 其中 $\Delta$ 是步长参数。离散化后的 SSM 递归形式为： $$x_k = \bar{A} x_{k-1} + \bar{B} u_k, \quad y_k = C x_k$$ ### 卷积形式 SSM 的优雅之处在于，它可以写成卷积形式，从而在训练时并行计算： $$y = K * u, \quad K = (C\bar{B}, C\bar{A}\bar{B}, \dots, C\bar{A}^{L-1}\bar{B})$$ 这个 kernel $K$ 称为 SSM 卷积核，它可以在训练时一次性计算整个序列的输出（类似卷积的并行性），在推理时又可以切换回递归形式（逐 token 生成，O(1) 每步）。 ## 演进路线：SSM → S4 → H3 → Mamba ### S4（Structured State Space, 2021） S4 的核心贡献是提出了**结构化**的 $A$ 矩阵。传统 SSM 的 $A$ 矩阵是稠密的，训练困难。S4 使用 HiPPO 矩阵（一种特殊的低秩 + 正规 + 低秩分解结构），将 $A$ 初始化为能高效记忆长程依赖的形式： $$A_{nk} = -\begin{cases} \sqrt{(2n+1)(2k+1)} & \text{if } n > k \\ n+1 & \text{if } n = k \\ 0 & \text{if } n < k \end{cases}$$ HiPPO 矩阵的数学性质保证了 SSM 在理论上可以捕获无限长的依赖关系。S4 在 Long Range Arena（LRA）基准上首次用 SSM 击败了所有 Transformer 变体。 ### H3（Hungry Hungry Hippos, 2022） H3 的贡献是填补了 SSM 的一个关键缺陷：**SSM 不能做"回忆"（recall）**——即从历史中精确检索某个 token。H3 在 SSM 层之间插入了一个"记忆"门控机制： $$H3(x) = \text{SSM}(x) \odot \sigma(W_g x) + x$$ 其中 $\odot$ 是逐元素相乘，$\sigma$ 是 sigmoid 门控。这个设计让模型既能利用 SSM 的线性复杂度，又能做精确的 token 级回忆。 H3 在语言建模上首次达到了与 Transformer 可比的困惑度，证明了 SSM 架构在语言任务上的可行性。但 H3 仍然是数据无关的（data-independent）——$A, B, C$ 参数对所有输入序列都一样。 ### Mamba（Selective SSM, 2023） Mamba 是 SSM 家族最重要的突破。核心创新是**选择性**（Selectivity）：让 $B, C, \Delta$ 参数依赖于输入 $u$： $$B_t = s_B(u_t), \quad C_t = s_C(u_t), \quad \Delta_t = s_\Delta(u_t)$$ 其中 $s_B, s_C, s_\Delta$ 是小的线性投影层。 **为什么选择性重要**： - 数据无关的 SSM 对所有输入一视同仁，无法"关注"重要 token、"忽略"无关 token - 选择性 SSM 可以根据输入内容动态调整——类似注意力机制的效果，但保持了线性复杂度 - 实验表明，选择性是 Mamba 在语言建模上超越 Transformer 的关键 **Mamba 的硬件优化**： - 选择性使卷积形式失效（kernel 依赖于输入），Mamba 必须在递归模式下运行 - Mamba 使用 CUDA kernel fusion 和 SRAM 优化，在 GPU 上实现了高效的并行扫描（parallel scan）算法 - 实际吞吐量：Mamba 在推理时比同规模 Transformer 快 5 倍 ### Mamba-2（2024） Mamba-2 进一步揭示了 SSM 和注意力机制之间的深层联系：结构化 SSM 可以理解为一种结构化掩码注意力（Structured Masked Attention, SMA）。这一发现统一了 SSM 和 Transformer 的理论框架，也让 Mamba-2 可以同时利用两种范式的优势。 ## SSM vs Attention 对比 | 维度 | Transformer (Self-Attention) | Mamba (Selective SSM) | |:---|:---|:---| | 计算复杂度 | O(n^2)（训练和推理） | O(n)（训练和推理） | | 内存复杂度 | O(n^2)（注意力矩阵） | O(n)（状态向量） | | 长序列处理 | 需要 KV Cache 压缩 / 稀疏注意力 | 天然支持百万级 tokens | | 全局依赖 | 天然的全局注意力 | 依赖状态向量压缩，可能丢失细节 | | 回忆能力 | 强（精确 token 检索） | 弱（信息被压缩到固定大小的状态中） | | 推理速度 | 随序列长度线性增长（KV Cache） | 恒定每步（O(1) 状态更新） | | 训练并行性 | 优秀（矩阵乘法） | 良好（并行扫描） | | 硬件友好度 | 优秀（矩阵乘法适配 GPU） | 良好（需要专用 kernel） | | 2026 年生态 | 最成熟 | 快速发展中 | | 代表模型 | GPT-5.5, Claude Opus 4.8, Llama 4.1 | 部分混合架构模型 | ## 混合架构：SSM + Attention 2026 年的一个关键趋势是**混合架构**——将 SSM 和 Attention 结合使用： - **SSM 层**处理长程上下文，压缩历史信息 - **Attention 层**处理局部窗口，做精确的信息检索 - 这种架构在保持线性复杂度的同时，弥补了纯 SSM 的回忆能力不足 代表性的混合架构包括 Jamba（AI21 Labs）、Zamba（Zyphra）等。 ## 2026 年定位 1. **长文档处理**：百万级 tokens 的文档分析，纯 Attention 的 O(n^2) 成本不可接受 2. **实时流式推理**：语音助手、实时翻译等需要逐 token 低延迟生成的场景 3. **边缘设备部署**：SSM 的恒定状态大小使其在内存受限设备上更有优势 4. **教学价值**：理解 SSM 是理解"Attention 不是唯一解"的关键——它代表了序列建模的另一条技术路线 ## 相关笔记 - [[Ch06 神经网络基础]] — 序列建模基础 - [[Ch09 Transformer 架构精讲]] — Transformer 架构 - [[概念-注意力机制]] — 注意力机制深入 - 深挖-Mamba — Mamba 深入 ## 来源 - Gu et al., "Efficiently Modeling Long Sequences with Structured State Spaces" (S4, ICLR 2022), https://arxiv.org/abs/2111.00396 - Gu & Dao, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" (2023), https://arxiv.org/abs/2312.00752 - Dao & Gu, "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" (Mamba-2, 2024), https://arxiv.org/abs/2405.21060 - Fu et al., "Hungry Hungry Hippos: Towards Language Modeling with State Space Models" (H3, 2023), https://arxiv.org/abs/2212.14052