【COLM 2024】Mamba：线性时间序列建模的选择性状态空间｜从高效序列架构演进视角

摘要

本文解读 COLM 2024 杰出论文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》。该论文提出选择性状态空间模型（S6），通过融合选择机制、硬件感知并行扫描与简化同质架构，让没有注意力机制的递归模型在语言建模上首次达到 Transformer 级质量，其特别之处在于把 SSM 参数变成输入的函数、按内容选择性传播与遗忘信息。实验表明Mamba-3B 匹敌 2 倍尺寸 Transformer、推理吞吐 5 倍、训练线性扩展，在语言、音频、基因组三模态均取得 SOTA，并支持百万级上下文，为高效序列架构提供了重要借鉴。

视频讲解：点击观看 B 站视频

论文基本信息

背景与动机

为什么高效序列模型长期打不过 Transformer？

现代基础模型几乎全部建立在 Transformer 的自注意力之上。注意力的核心能力是稠密路由信息：每个 token 与上下文窗口内所有 token 交互，这让它擅长建模复杂数据，但带来两个根本缺陷——无法利用有限窗口之外的信息，以及训练复杂度随窗口长度二次增长；推理时还要维护随上下文线性增长的 KV 缓存。

为克服这些缺陷，学界提出了大量次二次复杂度架构：线性注意力（Linear Attention）、门控卷积、递归模型，以及结构化状态空间模型（SSM）。SSM 家族从 S4 出发，演化出 H3（SSM 层两侧夹门控连接）、Hyena（MLP 参数化全局卷积）、RetNet（并行注意力路径）与 RWKV（LTI 递推，WKV 可看作两个 SSM 之比）。这些模型训练线性、推理恒定时间，但在语言这类信息稠密、离散的模态上始终不如注意力——没有一个被证明能在规模上跨域有效。

论文指出共同病根：这些模型都是时不变（LTI）的，参数不随输入变化，因而无法做内容感知推理——不能根据当前 token 决定记住什么、忘记什么。更根本地，序列建模的本质是把上下文压缩进有限状态：注意力完全不压缩（所以慢），递归模型状态有限（所以快但效果受限），两者之间的桥梁就是选择性：让压缩按内容进行。

研究主线：从问题到结论

图 9：Mamba 论文的研究主线——从注意力的效率缺陷出发，定位 LTI 病根，用选择机制与硬件感知扫描完成从问题到结论的闭环。

基准/方法设计

Mamba 的核心设计围绕三个支柱展开，与既有 SSM 架构（S4/H3/Hyena）形成鲜明对比。

图 1：选择机制总览——先前的 SSM 因时不变而可避免物化大状态，选择性 SSM 把输入相关动态放回模型，靠硬件感知算法控制内存。

• 选择机制（S6）：让 $\Delta$、$B$、$C$ 三个参数成为输入的函数：$B = s_B(x)$、$C = s_C(x)$、$\Delta = \tau_\Delta(\mathrm{Linear}1(x))$（$\tau\Delta$ 为 softplus）。参数沿序列长度维展开，模型由时不变变为时变。
• 硬件感知扫描：参数直接从 HBM 载入 SRAM，离散化与递推在 SRAM 内融合完成；用 work-efficient 并行前缀扫描并行化非线性递推；反传不存中间状态、重计算——内存占用与 FlashAttention 级优化 Transformer 相当。
• 简化架构：把 H3 块与 MLP 块合并为单一 Mamba 块，同质堆叠，扩展因子固定 $E=2$，使用 SiLU 激活。模型既无注意力也无独立 MLP 块，大部分参数（$3ED^2$）在线性投影上。

分类全景

图 10：Mamba 在次二次序列架构谱系中的位置——从 S4 到 H3/Hyena/RetNet/RWKV，再到引入选择机制的 Mamba。

方法细节

从 S4 到 S6：一个参数化的改变

S4 定义连续系统 $h'(t) = Ah(t) + Bx(t)$、$y(t) = Ch(t)$，经零阶保持离散化（$\bar{A} = \exp(\Delta A)$）得到递推 $h_t = \bar{A}h_{t-1} + \bar{B}x_t$。训练用全局卷积（可并行），推理切回递推（恒定时间/步）。

S6 的关键改动：$B$、$C$ 变为输入函数，$\Delta$ 由输入的线性投影经 softplus 得到。$A$ 可以保持静态——因为 $\bar{A} = \exp(\Delta A)$，$\Delta$ 的选择性会自动传导到离散参数。这打破了卷积等价性（时变），但换来了按内容决定记忆/遗忘的能力。

门控定理：统一 SSM 与 RNN

论文证明了选择机制与经典门控的联系：当 $N=1$、$A=-1$、$B=1$ 时，选择性 SSM 递推精确退化为门控 RNN：

$g_t = \sigma(\mathrm{Linear}(x_t))$，$h_t = (1-g_t)h_{t-1} + g_t x_t$

由此得到 $\Delta$ 的机理解释：大 $\Delta$ 重置状态、聚焦当前输入；小 $\Delta$ 保持历史、忽略当前输入。这一视角也解释了为何 $s_\Delta$ 投影到 1 维即可——输入 $x_t$ 该被忽略时，所有通道应一致忽略它。选择机制由此带来三类能力：过滤变间距噪声（如语言中的"um"）、过滤无关上下文（性能随上下文单调提升）、在文档/回合边界重置状态。

并行扫描与硬件实现

递推模式的 FLOPs 为 $O(BLDN)$，低于卷积模式的 $O(BLD\log L)$ 常数因子。但时变递推无法卷积化，必须处理两个挑战：递推的串行性与状态物化。解法是内核融合（离散化+递推在 SRAM 内完成，HBM 只读写 $B\times L\times D$ 的输入输出）+ 并行前缀扫描 + 反向重计算。扫描受内存带宽限制，融合是提速关键：A100 上比此前 SSM 实现快 3 倍，比朴素扫描快 40 倍。

图 2：Mamba 块结构——两堆块对应 Transformer 交错的注意力与 MLP 块的 $12D^2$ 参数，内部 SSM 贡献的参数很少。

实验设计与结果

评测协议

四个设定：合成任务（选择性复制、归纳头）检验内容感知能力；语言建模用 The Pile 300B tokens，覆盖 125M–1.3B 参数缩放律（Chinchilla 协议）与零样本下游评测；基因组用 HG38 预训练 + 大猿物种分类微调（上下文 $2^{10}\to2^{20}$）；音频用 YouTubeMix 波形预训练（BPB）+ SC09 语音生成（NLL/FID/IS）。基线包括 GPT3 配方 Transformer、LLaMa 配方 Transformer++、H3、Hyena、RetNet、RWKV、SaShiMi。

合成任务：选择性的直接证据

图 3：合成任务——选择性复制与归纳头直接检验内容感知能力。

归纳头任务中，模型在长度 $2^8=256$ 上训练，可外推到 $2^{20}=1048576$（4000 倍）保持高准确率，其他方法最多外推 2 倍——选择机制是唯一能外推的关键。

语言建模：匹敌两倍尺寸的 Transformer

Mamba 在每个尺寸档全面胜出，1.4B 平均准确率 59.7 甚至超过同 tokenizer 的 Pythia-2.8B（59.1）。Mamba-3B 在常识推理上比 Pythia-3B 高 4 分，匹敌 2 倍尺寸 Transformer——这是第一个匹配 LLaMa 配方 Transformer++ 的无注意力模型。

图 4：Pile 2K 上下文缩放律——首个匹配 Transformer++ 的无注意力模型。

效率：训练与推理双赢

融合扫描比朴素实现快40 倍；推理时作为递归模型每步恒定时间、无需 KV 缓存，吞吐量达同尺寸 Transformer 的5 倍。

图 5：训练与推理效率基准。

DNA 与音频：百万级长上下文

基因组上，固定模型大小时性能随上下文单调提升至 $2^{20}$（1M），基线持平甚至下降；大猿物种分类（上下文 1M）中 Mamba 准确率领先。音频上，6.1M 参数的 Mamba 在 SC09 上 FID 0.94，对比 SaShiMi 的 1.99（降幅超过一半）；24.3M 版本 FID 0.67、mIS 144.9，超越 WaveGAN/DiffWave 等 GAN 与扩散基线（附录 G 详表）。

图 8：长上下文 DNA 分类——选择机制过滤无关上下文能力的直接验证。

消融：$\Delta$ 最重要，$B,C$ 协同

$\Delta$ 是最重要的选择性参数（门控连接）；状态维数 $N$ 从 1 增到 16 困惑度下降超 1.0、仅增 1% 参数，但只有 $B,C$ 也选择性时才有效（附录 E）。

图 6：归纳头外推曲线——选择性机制带来可无限外推的内容感知能力。

图 7：融合扫描内核的训练效率。

结果对比总结

图 11：结果对比总结——质量匹敌两倍尺寸 Transformer，效率线性扩展，长上下文能力为三模态通用骨干奠定基础。

关键发现

1. 选择机制是性能分水岭：Selective Copying 从 S4 的 18.3% 跃升到 S6 的 99.8%，归纳头可外推 4000 倍（$2^8\to2^{20}$），LTI 模型完全做不到。
2. 首次匹配 Transformer++：125M–1.3B 缩放律上 Mamba 是第一个无注意力模型匹配 LLaMa 配方，且序列越长优势越明显。
3. 匹敌两倍尺寸：Mamba-1.4B 零样本平均 59.7 超 Pythia-2.8B 的 59.1；Mamba-3B 常识推理比 Pythia-3B 高 4 分。
4. 效率数量级提升：推理吞吐 5 倍于 Transformer，融合扫描比朴素实现快 40 倍，训练内存与 FlashAttention 同级。
5. 三模态 SOTA：音频 FID 从 1.99 降至 0.94（减半以上），基因组与音频性能随上下文单调提升至 1M。
6. $\Delta$ 门控理论：$N=1$ 时选择性 SSM 精确退化为门控 RNN，统一了 SSM 离散化与 RNN 门控两套理论。

局限性

• 规模有限：实证仅到约 3B 参数，低于 Llama/RWKV/RetNet 的 7B+；更大规模下是否保持优势未知，规模化需额外工程。
• 连续-离散免费午餐：选择机制牺牲了 LTI 在连续信号（音频/视频）上的强归纳偏置，音频实验需切回复数参数化（附录 G）。
• 生态欠账：微调、prompting、ICL、指令微调、RLHF、量化等 Transformer 生态的成熟适配机制尚未在 Mamba 上建立。
• 硬件依赖：性能依赖定制融合内核（selective scan），新硬件需重新工程化。

常见问题（FAQ）

Mamba 和 Transformer 的核心区别是什么？

Transformer 用自注意力在窗口内稠密路由信息，训练二次、推理需 KV 缓存；Mamba 用选择性状态空间递推，训练线性、推理恒定时间，靠输入相关的参数决定记忆与遗忘，首次在不牺牲质量的前提下实现线性复杂度。

为什么说 Mamba 的训练是线性时间的？

Mamba 的时变递推虽不能卷积化，但硬件感知的并行前缀扫描把 $O(BLDN)$ 的 FLOPs 并行化，且内存带宽受限的操作通过内核融合（SRAM 内完成离散化与递推）保持高效，因此训练随序列长度线性扩展。

选择机制到底"选择"什么？

选择的对象是信息的流入与流出：$\Delta$ 决定当前输入被聚焦还是被忽略（大 $\Delta$ 重置状态、小 $\Delta$ 保持历史），$B,C$ 分别控制输入进状态、状态出输出的细粒度门控——本质上让固定容量的状态按内容做最优压缩。

Mamba 在哪些任务上不如 Transformer？

在连续信号模态（如音频、视频）上，Mamba 的时变选择机制弱于 LTI SSM 的强归纳偏置，需要复数参数化弥补；此外大规模（7B+）验证、生态工具链（微调/量化/RLHF）也落后于 Transformer。

Mamba-2 与 Mamba 是什么关系？

Mamba-2（Dao & Gu, ICML 2024）通过状态空间对偶（SSD）统一了 SSD 与注意力，把选择性扫描进一步映射到类注意力结构，硬件效率再提升约 8 倍，同时保持了 Mamba 的选择性核心。

为什么归纳头任务如此重要？

归纳头（induction heads）被广泛认为是 LLM 上下文学习能力的关键机制。Mamba 在长度 256 训练后外推到 1M 仍保持高准确率，直接证明选择性机制具备可无限外推的内容感知能力，而所有 LTI 对比方法最多外推 2 倍。

参考链接

• 本文：Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752), COLM 2024
• 开源代码：github.com/state-spaces/mamba
• S4：Gu, Goel, Ré, Efficiently Modeling Long Sequences with Structured State Spaces, ICLR 2022
• HiPPO：Gu et al., HiPPO: Recurrent Memory with Optimal Polynomial Projections, NeurIPS 2020
• H3：Dao et al., Hungry Hungry Hippos, ICLR 2023
• FlashAttention：Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022
• RWKV：Peng et al., RWKV: Reinventing RNNs for the Transformer Era, Findings of EMNLP 2023

给大家推荐一款自用写文献综述、无虚构文献的 AI：

🌟复旦大学 FudanNLP 团队自研 切问学术

官网：qiewenpaper.com

覆盖3.6 亿篇可溯源真实中英文文献，能自动整合文献观点生成规范综述

还能挖掘研究创新点、复现实验，配合视频教学，新手快速上手文献综述写作

🍀后记🍀

博客的关键词集中在编程、算法、机器人、人工智能、数学等等，持续高质量输出中。

🌸讨论QQ群：白拾的小屋 (750365700)

⭐B站账号：白拾的物理AI组会（活跃于知识区和动画区）

✨GitHub主页：YhbCode000（工程文件）