State Space Models (SSM)

状态空间模型（SSM）通过隐藏状态向量将输入序列映射为输出序列，并在每一步更新该向量，因此无论序列多长，其内存占用都保持不变。这一思想源自控制理论：在线性系统中，状态更新写作 h' = Ah + Bx，输出写作 y = Ch + Dx；自 1960 年以来，卡尔曼滤波器一直利用带噪测量值估计这类状态。深度学习中的 SSM 会从数据中学习 A、B 和 C，并堆叠多个这样的层；自 2021 年的 S4 和 2023 年的 Mamba 以来，它们已在文本、音频和时间序列等长序列任务上与 Transformer展开竞争，同时计算量随序列长度线性增长。

SSM 的工作原理#

神经 SSM 层会使用学习得到的步长 Δ，通过零阶保持规则 Ā = exp(ΔA)，将连续方程转换为离散递推。每一步都会计算 h_t = Ā h_(t-1) + B̄ x_t 和 y_t = C h_t。当 A、B 和 C 在各个时间步保持固定时，展开递推可以看出，输出是输入与核 (CB̄, CĀB̄, CĀ²B̄, ...) 的卷积。因此，该层在训练时可以作为卷积运行，在 GPU上并行处理整个序列；在生成时则可以作为递推运行，每个 token 只需进行一次固定大小的状态更新。在 PyTorch中，这两种形式会得到相同输出：

import torch

torch.manual_seed(0)
L, N = 64, 16  # 序列长度，状态大小
x = torch.randn(L)

# 对角连续时间 SSM，使用步长 delta（零阶保持）进行离散化
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# 递归形式：每一步更新一次状态，类似生成过程
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# 卷积形式：预先计算核 K_k = C . (A_bar^k * B_bar)，类似训练过程
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

它会输出 Same outputs: True。递推在每一步只保留 16 个状态数值，而预先计算的核可以让训练并行计算全部 64 个输出。

从 S4 到 Mamba#

最初的深度 SSM 训练速度太慢，内存占用也太高，无法大规模训练。S4由 Albert Gu、Karan Goel 和 Christopher Ré 提出；它为状态矩阵引入了一种结构，使卷积核的计算成本更低，并采用专门设计的初始化方式来记住较长历史信息。它在 Long Range Arena 基准的所有任务上都达到了当时的最先进水平，成为第一个解决其 Path-X 任务（16K 步序列）的模型，生成速度比 Transformer 快 60 倍。

固定卷积核会以相同方式处理所有输入，因此无法决定跳过填充内容并保留关键信息。Albert Gu 和 Tri Dao 提出的 Mamba会将 Δ、B 和 C 设为当前 token 的函数，使模型能够选择要写入状态的信息以及要遗忘的信息。这些选择性参数打破了卷积形式，因此 Mamba 会使用感知硬件的并行扫描来计算递推。Mamba 的推理吞吐量达到 Transformer 的五倍；Mamba-3B 的表现优于同等规模的 Transformer，并与其两倍规模的 Transformer 相当。每个 Mamba 块都会用一个短卷积和一个 SiLU门控包裹选择性 SSM：

Mamba-2表明，选择性 SSM 和一种线性注意力是同一计算的两种对偶形式，作者将这一框架称为状态空间对偶性；该方法使核心层的速度提升了 2 到 8 倍。

SSM 与相关架构的对比#

• Transformer：自注意力会比较每个 token 与其他所有 token，因此可以精确复制任意更早的 token，但其成本按平方增长，而且 KV 缓存会随上下文窗口增长。SSM 会将历史信息压缩到固定状态中，成本更低但有损；2024 年的论文 Repeat After Me表明，在复制和检索长输入方面，Transformer 优于固定状态 SSM。
• 循环神经网络（RNN）：RNN 或 LSTM会执行非线性更新，因此训练时必须逐个 token 地处理序列。SSM 的状态更新是线性的，因此可以使用卷积或并行扫描。
• 隐马尔可夫模型（HMM）：HMM 有一组有限的离散状态以及状态转移概率。卡尔曼滤波器背后的线性高斯状态空间模型是其连续对应形式，而深度 SSM 会保留这个连续状态向量。

SSM 的应用场景#

大多数生产应用都采用混合式大型语言模型（LLM），保留少量注意力层以实现精确回忆。AI21 的 Jamba交错使用 Transformer 层、Mamba 层和专家混合（MoE）层，可在单个 80 GB GPU 上运行，并报告了在最长 256K token 的上下文下取得的良好结果；IBM 的混合 Granite 4.0-H 模型则以 9:1 的比例组合 Mamba-2 和 Transformer 层。Vision Mamba会对图像图块执行选择性扫描；SSM 还可以为长到无法进行完整注意力计算的原始音频和 DNA 序列建模。经典形式也在 Ultralytics YOLO26 跟踪中运行：默认的 TrackTrack 跟踪器与 ByteTrack 和 BoT-SORT 一样，会使用卡尔曼滤波器预测每个目标的下一个框；其 8 维状态包含框中心、大小及其速度，并通过恒定速度转移矩阵和读取四个框坐标的线性观测矩阵进行更新。