Featured image of post Mamba:少记,为什么反而可能更懂长序列

Mamba:少记,为什么反而可能更懂长序列

Mamba 不保存完整 KV cache,而是用选择性状态空间把历史压进固定状态。本文从选择性、扫描算法与实证边界解释它为何成为 Transformer 之外的重要路线。

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 共同设计记忆系统的一条主线。

参考资料

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