状态空间模型与 Mamba：不用注意力读取长序列的方法

为何需要：在长序列面前受阻的现有模型

当今大多数 LLM 都基于 Transformer 运行。Transformer 的核心——注意力——会把上下文中的所有 token 相互比较，细致地交换信息。Mamba 论文（paper）总结说，这一优势伴随着两项代价：一是完全无法处理固定窗口（上下文窗口）之外的内容，二是计算量随窗口长度的平方增长。为降低这一开销，也出现了许多高效注意力变体。但同一篇论文指出，这些变体往往同时失去了让注意力有效的那些性质。

按顺序读取的 RNN（循环神经网络）开销与长度成正比，但存在训练中梯度消失、难以记住遥远过去的问题。据 S4 论文（paper），RNN、CNN 和 Transformer 都有处理长依赖的特殊变体，但要扩展到超过 1 万步的序列仍然困难。在长序列基准 Long Range Arena 的 Path-X 任务（长度 16,384）上，此前所有模型的结果都没有好于随机猜测（50%）。

因此，人们需要的是这样一种模型：训练时能够并行快速运行，生成时每个 token 只需固定开销，同时还能抓住遥远过去的信息。状态空间模型（SSM）正是作为瞄准这些条件的候选者出现的。

核心思想：一张便签代替整份会议记录

先从一个比喻说起。Transformer 就像一个人，把会议内容一行不落地全部记下来，每次有人提问时都把整份记录重新翻一遍。这样很准确，但会议越长，记录越多，翻查的时间也越长。SSM 则更像一个人只拿着一张大小固定的便签，每听到一句新发言就改写一次便签。由于便签大小固定，无论会议多长，每次改写的开销都一样。

这张“便签”就是状态（state）。正如 S4 论文（paper）所说明的，SSM 原本是控制理论、计算神经科学等领域使用的数学模型。输入信号经过隐藏状态连接到输出，矩阵 A·B·C·D 决定状态如何变化以及输出什么。关键在于状态矩阵 A。选好 A，就能在数学上长时间保存遥远过去的信号。

问题在于便签上该写什么。早期的 SSM 无论输入是什么，都用同样的规则改写便签。这种方式很适合音频这类连续信号，但在每个词的含义都很重要的文本上表现较弱。Mamba（paper）正是改进了这一点：让模型看到新 token 时随时判断“这个值得记下来”还是“这个可以忽略”。

工作原理：状态更新、两种计算方式，以及选择

把一个 SSM 层的运作拆成步骤来看会更容易理解。第一，一个 token 进入后，与此前的状态融合成新状态（A 决定保留多少之前的状态，B 决定放入多少新输入）。第二，C 从新状态中读出输出。第三，在下一个 token 上重复同样的过程。正如 Mamba 论文（paper）所述，输入的每个通道都经过一个维度更高的隐藏状态连接到输出。

S4（paper）强调的 SSM 优势在于，同一计算可以从两种方式来看。生成时，可以用逐步更新状态的循环（recurrence）方式运行。训练时，由于参数不随时间变化，可以把整个序列作为一次长卷积（convolution）一次性计算。相当于像 RNN 一样推理，像 CNN 一样并行训练。

Mamba 在此基础上加入了“选择（selection）”。它把决定状态更新幅度的间隔 Δ 以及 B·C 从固定值改为由当前 token 计算得到的值。这样模型就能沿着序列有选择地传递或遗忘信息。这也有代价：由于参数随时间变化，就不能再使用上面的卷积计算方式。Mamba 用考虑硬件的并行扫描算法解决了这个问题。它不在 GPU 的慢速内存中真正展开庞大的扩展状态，从而避免了 GPU 内存层级之间的数据搬运。论文报告，该实现在 A100 GPU 上比以往方法最多快 3 倍。

最后，Mamba 把这种选择性 SSM 与 Transformer 的 MLP 块合并成一个简单的块。这是一种既没有注意力、也没有单独 MLP 块，只堆叠这种块的结构。生成时不需要之前 token 的缓存，因此每个 token 只需固定时间。这是它与背负着不断增大的 KV 缓存的 Transformer 最大的不同。

发展脉络：从 S4 到 Mamba，再到 Mamba-2

S4（paper）解决了先前研究遗留的问题。据论文介绍，其前一阶段的 LSSL（线性状态空间层）表明，使用特殊的状态矩阵 A，深层 SSM 也能处理长程依赖。但对于状态大小 N 和长度 L，它需要 O(N²L) 的计算量和 O(NL) 的内存，实际难以使用。S4 将 A 分解为“低秩项 + 正规项”，使其能够被稳定地对角化。结果，计算被归约为已有充分研究的柯西（Cauchy）核计算，计算量和内存几乎与 N+L 成正比。它比 LSSL 最快可快 30 倍，内存占用少 400 倍。S4 首次解决了 Path-X，取得 88% 的准确率，并在 LRA 总平均分上领先已有模型 20 分以上。它在序列 CIFAR-10 上取得 91% 的准确率，生成速度快 60 倍。

Mamba（paper）针对的是 S4 系列的一个局限：在音频、视觉等连续信号上表现强，但在文本这类离散且信息密集的数据上表现弱。借助前文介绍的选择机制和硬件友好的扫描，它的生成吞吐量比规模相近的 Transformer 高 5 倍。论文报告称，在真实数据上，性能随序列长度增加持续提升，直至 100 万长度的序列。在语言建模方面，Mamba-3B 超过了同等规模的 Transformer，并与两倍规模的 Transformer 相当。例如，在常识推理平均分上，它比 Pythia-3B 高 4 分，也超过了 Pythia-7B。不过，论文展示的缩放定律实验只到 10 亿（1B）参数。

Mamba-2 论文（paper）表明，两类模型比想象中更接近。用多种方式分解一种名为半可分（semiseparable）矩阵的结构化矩阵，就能把 SSM 与各种注意力变体联系起来。这被称为状态空间对偶性（SSD，State Space Duality）。基于这一框架，研究者改进了 Mamba 的选择性 SSM，设计出 Mamba-2，其核心层快了 2~8 倍，同时在语言建模上仍与 Transformer 保持竞争力。

近期走向：为什么不擅长检索，如何撑住长上下文

最大的争议点是上下文内检索（in-context retrieval），即在长文本中准确找到特定内容。Mamba 在这方面弱于 Transformer。通常的解释是固定大小的状态会丢失信息。一项检索瓶颈分析研究（paper）对这一解释提出了反驳。在把问题放在上下文之前的设置中，模型可以预先知道需要记住什么，因此状态大小不应是决定性的限制。然而在这一设置下，Mamba 仍然大幅落后。仅操控选择门控就能让准确率在两个极端之间摆动，说明门控的作用力是足够的。研究者从别处寻找问题所在。token 表示（Representation）中“与问题相关”的信号是在经过较深的层时逐渐形成的，而门控要再经过几层之后才反映这一信号。对检索起决定作用的那些层中，门控尚未与相关性对齐。这意味着选择机制本身足够强大，但它常常在用于判断的信号尚未到达时就已开始工作。

第二条路线是尝试在不进行训练的情况下弥补长上下文中的性能下降。SuPER-Mamba（paper）在每个位置计算已训练好的 Mamba 模型对下一个 token 有多“意外”（Surprisal），并用这一分数决定优先将哪些位置的信息纳入全局状态更新。它既不需要额外训练，也不需要针对问题计算分数。论文报告称，应用于三个 SSM 模型后，LongBench-E 平均分从 8.06 提升到 15.69。

第三条路线是重新审视状态更新方式本身。根据 Mamba-3 振荡动力学研究（paper），Mamba-3 在状态更新中加入了随输入变化的旋转。该研究将这种旋转解释为“欠阻尼振荡”。在受控比较中，振荡动力学在状态追踪任务上优于无振荡的衰减；在语言模型预训练中，它有助于长上下文预测和长度泛化。研究者的解读是，旋转看起来是一种调节已存储信息与新输入的更新如何结合的机制。

局限与争议

第一，压缩的代价。只有一页便笺这一优点，同时也就是弱点。Transformer 把原文原样保存在 KV 缓存中，需要时可以回看；SSM 则无法再取出已经过去的内容。前述检索研究（paper）表明还存在状态大小以外的原因，但也确认了检索是 Mamba 的弱点这一事实本身。SuPER-Mamba 研究（paper）的出发点同样是 SSM 性能常常随上下文变长而下降。

第二，对实现的依赖。一旦加入选择机制，就无法再使用 S4 时期的快速卷积计算。Mamba 通过精心处理 GPU 内存层级的专用扫描内核弥补了这一点（paper）。这意味着，理论上的线性复杂度要转化为实际速度，必须依靠针对硬件的实现。

第三，关于规模的证据。Mamba-2 论文（paper）的表述是 Mamba 在“中小规模”上与 Transformer 相当或领先，Mamba 论文的缩放实验也只到 10 亿参数。在大得多的规模上能否得出相同结论，仅凭本文所依据的资料无法判断。最好不要把基准数据解读为“在所有规模上取代 Transformer”。

延伸学习

建议按时间顺序阅读。先读 S4（paper）的引言和图 1，掌握 SSM 可以用循环和卷积两种方式计算这一概念。如果觉得公式吃力，可以把低秩校正和柯西核部分仅理解为“降低计算成本的技巧”，然后跳过。

接下来是 Mamba（paper）。阅读重点是为什么需要选择机制（选择性复制、induction head 等合成任务），以及硬件友好的扫描避免了什么。随后阅读 Mamba-2（paper），就能看出 SSM 与注意力为何并非完全不同的两个世界。

对于近期动向，建议先读检索瓶颈分析（paper）。它能让你不再只把 Mamba 的弱点归结为“状态太小”。之后依次阅读无需训练即可弥补长上下文的 SuPER-Mamba（paper），以及解读状态更新动力学的 Mamba-3 研究（paper）。