《Mamba：带选择性状态空间的线性时间序列建模》

要解决的问题

序列模型一直要权衡两件事：模型能用上多少先前读到的内容，以及处理一条长序列要花多少计算。注意力机制让每个位置与所有其他位置两两比较，序列一变长，代价就迅速上升。为了把这份代价压下来，出现了一批次二次时间（subquadratic-time）的架构：线性注意力（linear attention）、门控卷积（gated convolutions）、结构化状态空间模型（structured state space model）。由 Albert Gu 和 Tri Dao 于 2023 年 12 月 1 日提交到 arXiv 的 Mamba 论文指出，这些更便宜的架构有一个共同弱点：无法做基于内容的推理（content-based reasoning）^[1]。

所谓基于内容的推理，是依据当前输入本身是什么来决定保留还是丢弃信息，而不只是依据它在序列中的位置。如果模型永远用同一套固定规则压缩历史，那么一个填充词和一个关键人名在它眼里没有区别。这确实是个要害问题，因为语言模型要做的许多事——回忆前文提到的事实、忽略干扰项、原样复制某个词——都取决于这些词的内容。

选择性的状态空间

起点是状态空间模型（state space model，SSM），这是一种经典的序列描述方式：随着新词不断到来，用一个隐状态一步一步地更新。这样读序列的代价随长度线性增长，每读一个新词的工作量大致相同。问题在于，早期 SSM 的更新规则是固定的：无论输入是什么，状态如何演化、有多少新输入进入状态，都由同一组矩阵决定。

Mamba 的核心做法是让 SSM 的参数成为输入的函数。既然更新依赖于眼前的这个词，模型就可以依据当前词，在序列长度这一维度上选择性地传播或遗忘信息^[1]。换句话说，模型给自己的记忆加了一道由数据决定的闸门：无关紧要的词可以几乎不留痕迹地通过，重要的词则可以被写进状态并一直带下去。

架构

这些选择性 SSM 被整合进一个简化的端到端神经网络架构，其中没有注意力，甚至连 MLP 模块也没有，作者把它叫作 Mamba^[1]。结构保持简单是有意为之：这项工作的主张是，承担建模重任的应当是序列混合机制本身，而不是额外的深度或额外的组件。

递推式的写法对现代硬件通常并不友好，因为它只能一步一步地算。Mamba 用一套面向硬件的并行算法在递推模式下解决了这个问题^[1]，使得递推在训练时可以被高效地并行求值，而不必串行推进。同样的结构也让生成变得便宜：每产出一个词只需要更新一次状态，不必对整个上下文重新做一遍比较。

结果

论文报告说，Mamba 的推理很快，吞吐量（throughput）是 Transformer 的 5 倍，并且随序列长度线性扩展，在真实数据上，序列长到百万级别时性能仍在提升^[1]。后一点在实践中尤其值得注意：许多长序列模型在上下文继续加长后，效果要么停滞要么变差，而这里在真实数据上，随着序列向百万词元增长，结果持续变好。

在语言建模上，作者报告 Mamba-3B 在同规模比较中优于 Transformer，并追平了规模两倍于它的 Transformer，预训练（pretraining）和下游评测两方面都是如此^[1]。

意义

这项工作的贡献既是一组数字，也是一个设计上的论证。论文把「基于内容的推理」指认为先前次二次架构真正欠缺的能力，说明让状态空间参数依赖输入就能补上这一能力，随后又报告说，一个完全由这类模块搭建的网络在同规模下可与基于注意力的模型竞争^[1]。对于注意力代价最难承受的长序列场景来说，这使得 Mamba 成为「注意力之外还有什么选择」这一问题上的重要答案之一。