14.3 状态空间模型与混合架构：注意力的挑战者

自注意力机制六年来统治了序列建模领域，但其平方复杂度和 KV 缓存的线性增长使其在超长序列场景下面面临根本性的效率限制。状态空间模型（State Space Model，SSM）提供了一种截然不同的建模范式。

14.3.1 SSM 的基本思想

SSM 源于控制理论中的线性时不变系统。其核心是用线性递推替代注意力机制的全连接计算：

$h_t = Ah_{t-1} + Bx_t$ $y_t = Ch_t + Dx_t$

其中 $h_t$ 是隐藏状态，$A$、$B$、$C$、$D$ 是系统矩阵。每步计算只涉及矩阵-向量乘法，复杂度为 $O(1)$（相对于注意力的 $O(n)$）。整个序列的处理复杂度为 $O(n)$，无需存储 $O(n^2)$ 的注意力矩阵。

S4（Structured State Space for Sequence Modeling，Gu 等人，2022 年）通过对 $A$ 矩阵的结构化参数化（HiPPO 初始化）解决了 SSM 在长距离建模上的困难，首次在长序列基准上达到了 Transformer 的水平。

14.3.2 Mamba：选择性状态空间

Mamba（Gu & Dao，2024 年）是 SSM 方向的里程碑。它的核心创新是选择性机制——让系统矩阵 $B$、$C$ 和 $\Delta$（离散化步长）依赖于输入：

$B_t = f_B(x_t), \quad C_t = f_C(x_t), \quad \Delta_t = f_\Delta(x_t)$

这意味着模型可以根据当前输入动态决定“记住什么、忘记什么”——类似于注意力机制的选择性关注能力。

Mamba 在以下方面展现了优势：

• 推理效率：不需要 KV 缓存，状态大小固定，推理时间与序列长度无关

• 训练效率：可以通过并行扫描算法在 GPU 上高效训练

• 长序列能力：在处理超长序列（数十万词元）时比标准 Transformer 快数倍

14.3.3 Mamba 2 与后续改进

Mamba 2（2024 年）将选择性 SSM 与结构化矩阵联系起来，展示了 SSM 和注意力之间深层的数学联系——在某种意义上，选择性 SSM 可以被视为一种“结构化的线性注意力”。

Mamba 2 还优化了 GPU 上的计算效率，使其在标准硬件上的吞吐量进一步提升。

14.3.4 混合架构：取两者之长

越来越多的研究发现，纯 SSM 模型和纯 Transformer 各有优劣——SSM 在效率上占优但在精细的上下文检索（“大海捞针”测试）上不如注意力；Transformer 在精确关注上占优但长序列效率低。

混合架构将两者结合：

Jamba（AI21，2024 年）：交替使用 Transformer 层和 Mamba 层，在 256K 上下文窗口内保持了优秀的检索精度和高效的推理速度。

Griffin（Google，2024 年）：结合局部注意力和门控线性循环层，在长序列任务上超越了纯 Transformer 基线。

Zamba（Zyphra，2024 年）：使用共享注意力层和 Mamba 块的混合架构，在参数效率上表现出色。

混合架构代表了一种务实的演进方向——不是完全替代 Transformer，而是在其基础上引入更高效的组件来处理特定场景。