Mamba：少记，为什么反而可能更懂长序列

Thu, 24 Sep 2026 00:00:00 +0000

Transformer 的注意力做了一件昂贵但稳妥的事：不急着决定哪些历史重要，而是把每个旧 token 的 key 和 value 留在 KV cache，等新 token 到来后再逐一比较。Mamba 选择了相反的契约：历史经过当前位置时就被压进一个固定大小的状态，之后不再保留可任意回看的原始清单。

这个差别不是把 $O(L^2)$ 改写成 $O(L)$ 那么简单。压缩历史很容易，难的是压缩后仍保留真正有用的内容。早期线性时不变状态空间模型（LTI SSM）对每个 token 使用同一套动力学，像一只永远以相同方式更新的滤波器；它们适合连续信号，却难以处理语言中“这个词要记住、那个填充词可以略过”的内容判断。Mamba 的突破，是让状态更新的关键参数依赖当前输入，再为这种失去卷积捷径的模型设计硬件感知的并行扫描。

因此，Mamba 最值得理解的不是“它没有 attention”，而是三个相互约束的选择：用有限状态压缩上下文，用输入相关参数决定压缩策略，再用扫描算法让递归结构在 GPU 上仍能高效训练。

两种上下文契约：档案室与工作记忆

设序列长度为 $L$。标准自注意力在训练时显式形成 token 两两交互，计算和注意力矩阵随 $L^2$ 增长；自回归生成虽然每步只计算一个新 query，却要保存并读取随 $L$ 增长的 KV cache。它的优势也来自这里：旧 token 的表示仍在，模型可以按当前问题重新检索它们。

状态空间模型采用另一种形式。离散后，一个最简递归可以写成：

$$h_t=\bar A h_{t-1}+\bar Bx_t,\qquad y_t=Ch_t$$

$h_t$ 是到位置 $t$ 为止的压缩状态。每来一个 token，旧状态与新输入合成下一状态；生成时每层只需保留当前状态，缓存大小不随上下文长度增长，单步计算也不因前缀变长而线性增加。代价同样明确：一旦信息没有进入 $h_t$，未来就无法像注意力那样回头访问原始 token。

图 1：注意力保留可回看的 KV 档案；Mamba 把历史持续压缩进固定状态

这使“长上下文”有了两种不同含义。对注意力模型，窗口更长主要意味着可访问的档案更多，同时付出更大的缓存和读取成本。对循环式 SSM，处理更长序列的计算增长是线性的，但固定状态并不会因为窗口变长自动获得更多容量。能处理一百万 token，不等于能无损记住一百万 token。 原始 Mamba 论文在 DNA 任务上观察到性能一直改善到百万长度，这是特定数据和任务上的实证，不是无限记忆的证明。

选择性不是挑掉 token，而是改变状态如何更新

传统 LTI SSM 的 $Δ,A,B,C$ 对所有位置固定。它可以被看成递归，也可以预先展开成一个全局卷积核，因而容易并行训练；但固定动力学无法根据内容改变记忆方式。原论文用“选择性复制”任务说明这个缺口：需要复制的符号位置会变化，模型必须依据 token 内容决定何时写入，而不能只根据相对位置等待固定时长。

Mamba 保留结构化的 $A$，却让 $Δ_t、B_t、C_t$ 成为输入 $x_t$ 的函数：

$$Δ_t=s_\Delta(x_t),\qquad B_t=s_B(x_t),\qquad C_t=s_C(x_t)$$

离散化后，状态转移也随位置变化：

$$h_t=\bar A_t h_{t-1}+\bar B_t x_t,\qquad y_t=C_t h_t$$

可以把三者理解成不同问题。$B_t$ 决定当前输入以什么方向写入状态，$C_t$ 决定此刻从状态读出什么，$Δ_t$ 调节保留旧状态与接纳当前输入的时间尺度。按论文的机械解释，较大的 $Δ_t$ 更倾向重置旧状态并聚焦当前输入，较小的 $Δ_t$ 更倾向维持状态、忽略这次短暂输入。它不是先给 token 打一个“保留/删除”的离散标签，而是在每一步连续改变信息流。

图 2：输入相关的 Δ、B、C 分别调节遗忘/更新、写入方向与读出内容

这也是“Mamba”与一般 SSM 不能混为一谈的原因。论文消融中，仅把非选择性 SSM 换成选择性 S6 就显著降低语言模型困惑度；$Δ$ 是单项改动里最重要的参数，而同时让 $Δ、B、C$ 具有选择性效果最好。真正补上离散语言能力缺口的，不是“状态空间”四个字，而是内容相关的压缩。

选择性拿走了卷积捷径，扫描算法把并行性补回来

固定参数的 LTI SSM 可以把整段序列写成卷积，训练时一次并行计算。参数随 token 变化后，不再存在一个适用于全序列的固定卷积核；如果朴素地按 $t=1,2,\ldots,L$ 顺序执行递归，GPU 会被大量细小、串行的操作拖慢。

Mamba 的解决办法基于递推组合的结合律。每个位置都可表示成一对局部变换 $(\bar A_t,\bar B_tx_t)$；两个相邻变换能够先组合，再与其他区间组合。于是训练时可以用 parallel scan 在树形结构中并行归并，生成时则继续使用真正的单步递归。同一个方程因此有两种执行计划：批量训练重视序列并行，自回归生成重视常数大小状态。

图 3：同一选择性递推在训练时采用并行扫描，在生成时采用单步状态更新

算法仍不够。选择性 SSM 在概念上会产生形状约为 $B\times L\times D\times N$ 的扩展状态；若把中间量完整写回高带宽显存，线性复杂度也可能被内存流量吃掉。论文的 CUDA 实现把离散化、扫描与后续操作融合在片上 SRAM 中，避免物化庞大中间张量，并在反向传播时重算部分状态。原论文报告：在其测试设置中，选择性扫描在序列超过 2K 后快于 FlashAttention-2，较标准 PyTorch 扫描快 20–40 倍；端到端自回归推理吞吐最高达到同尺度 Transformer 的 5 倍。这些是特定硬件、模型形状和实现下的测量，不应当当作所有部署的固定倍率。

Mamba block 为什么连 MLP 也合并了

原始 Mamba 不只是把 Transformer 的 attention 替换成 SSM。一个 block 先把输入扩展成两路：主路经过短的一维因果卷积、SiLU 和选择性 SSM；旁路形成门控；两路相乘后再投影回模型维度。这个设计把以往“序列混合层 + 独立 MLP”的两段结构合并成一个同质 block，整网不需要 attention，论文所说的“没有 MLP block”指没有单独堆叠的 FFN 子层，而不是内部没有线性投影或非线性。

短卷积负责邻近 token 的局部混合，SSM 负责沿序列传播状态，门控则控制输出通道。它们共同解释了为什么只盯着递推公式会漏掉实际模型的一半：Mamba 的效果来自选择性状态层、局部卷积、门控架构和硬件实现的组合，而非某一个孤立公式。

这种端到端简化也带来工程门槛。官方仓库的高性能路径依赖定制 CUDA kernel，支持的状态维度、数据类型和硬件会影响实际速度；在不匹配的环境里退回参考实现，理论上的线性优势未必自动变成更低延迟。架构复杂度从 KV cache 转移到了 kernel 与扫描实现，并没有凭空消失。

纯 SSM 的边界，后来由混合架构说得更清楚

原始论文的结果很强：Mamba-3B 在其训练与评测设置中优于同规模 Transformer，并可匹配约两倍规模的 Transformer；它还覆盖语言、音频与基因组序列。但后续由 Mamba 作者与 NVIDIA 研究者共同完成的 8B 受控研究给出了更有价值的限制条件。研究让 Mamba、Mamba-2 和 Transformer 使用相同数据、最多训练 3.5T token，发现纯 SSM 在很多任务上可以匹配或超过 Transformer，却在强复制、in-context learning 和长上下文推理任务上落后。

这个结果与上下文契约完全一致。注意力保留原始 KV，可以为一个很晚出现的问题重新指向早期细节；固定状态必须提前判断什么值得留下。选择性能改进这个判断，却不能取消有限状态的信息容量。复制电话号码、逐字复现某段内容或临时学习上下文中的新映射，恰好都是“以后会问什么，在写入时未必知道”的任务。

同一研究中的 Mamba-2-Hybrid 使用 43% Mamba-2、7% attention 和 50% MLP 层。论文报告它在 12 个标准任务上平均超过 8B Transformer 2.65 分，并预测生成时最高可快 8 倍；在额外 23 个长上下文任务上，混合模型平均也能接近或超过 Transformer。这里的“预测最高 8 倍”不是完整线上基准，但方向很有启发：少量 attention 可以作为可回看的精确记忆，其余层用 SSM 承担廉价的顺序压缩。

Mamba-2 又通过 State Space Duality（SSD）说明，结构化 SSM 与某类受约束 attention 并非两个互不相干的世界，而是同一类半可分矩阵的不同计算形式。它的核心层在论文基准中比 Mamba 的选择性 SSM 快 2–8 倍。更准确的产业判断因此不是“SSM 将淘汰 Transformer”，而是：attention 与状态空间提供了不同的记忆接口，模型可以按层、按任务组合它们。

真正改变的是推理系统的记忆预算

昨天讨论 MLA 时，目标仍是在 attention 内部压缩 KV cache；Mamba 更进一步，直接改变缓存对象。它不再为每个历史 token 保存可寻址的 K/V，而是为每层维护一个随序列更新的状态。上下文从“越来越厚的档案”变成“大小固定、不断改写的工作记忆”。这为流式音频、基因组、传感器和超长生成提供了不同的系统设计空间：延迟和状态内存不必随已读长度增长。

但这份预算必须与任务的记忆需求一起看。需要精确引用、任意回看和上下文内临时学习时，attention 的显式档案仍然珍贵；需要持续流式处理、低缓存和稳定单步成本时，选择性状态更有吸引力。混合模型之所以自然，不是折中得保守，而是承认“摘要式记忆”和“可寻址记忆”本来就是两种能力。

Mamba 留下的核心问题也因此比某个榜单更长久：一个序列模型究竟该保存多少过去，又该在什么时候决定遗忘？注意力把决定推迟到读取时，Mamba 把决定提前到写入时。它的贡献，是证明后者在具备内容选择、并行扫描和合适 kernel 后，不再只是效率不错但能力落后的替代品，而成为可以与 attention 共同设计记忆系统的一条主线。

参考资料

Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, 2023/2024.
Dao & Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, 2024.
Waleffe et al., An Empirical Study of Mamba-based Language Models, 2024.
State Spaces Models team, Mamba official implementation.
NVIDIA Megatron-LM, Mamba and Mamba-2 model implementation.

长序列 on 辰远