摘要本文解读 COLM 2024 杰出论文《Mamba: Linear-Time Sequence Modeling with Selective State Spaces》。该论文提出选择性状态空间模型S6通过融合选择机制、硬件感知并行扫描与简化同质架构让没有注意力机制的递归模型在语言建模上首次达到 Transformer 级质量其特别之处在于把 SSM 参数变成输入的函数、按内容选择性传播与遗忘信息。实验表明Mamba-3B 匹敌 2 倍尺寸 Transformer、推理吞吐 5 倍、训练线性扩展在语言、音频、基因组三模态均取得 SOTA并支持百万级上下文为高效序列架构提供了重要借鉴。视频讲解点击观看 B 站视频摘要论文基本信息背景与动机为什么高效序列模型长期打不过 Transformer研究主线从问题到结论基准/方法设计分类全景方法细节从 S4 到 S6一个参数化的改变门控定理统一 SSM 与 RNN并行扫描与硬件实现实验设计与结果评测协议合成任务选择性的直接证据语言建模匹敌两倍尺寸的 Transformer效率训练与推理双赢DNA 与音频百万级长上下文消融$\Delta$ 最重要$B,C$ 协同结果对比总结关键发现局限性常见问题FAQMamba 和 Transformer 的核心区别是什么为什么说 Mamba 的训练是线性时间的选择机制到底选择什么Mamba 在哪些任务上不如 TransformerMamba-2 与 Mamba 是什么关系为什么归纳头任务如此重要参考链接论文基本信息项目内容标题英文Mamba: Linear-Time Sequence Modeling with Selective State Spaces标题中文Mamba线性时间序列建模的选择性状态空间作者Albert Gu, Tri Dao机构卡内基梅隆大学机器学习系 · 普林斯顿大学计算机科学系会议COLM 2024Outstanding Paper AwardarXivhttps://arxiv.org/abs/2312.00752项目网站https://github.com/state-spaces/mamba背景与动机为什么高效序列模型长期打不过 Transformer现代基础模型几乎全部建立在 Transformer 的自注意力之上。注意力的核心能力是稠密路由信息每个 token 与上下文窗口内所有 token 交互这让它擅长建模复杂数据但带来两个根本缺陷——无法利用有限窗口之外的信息以及训练复杂度随窗口长度二次增长推理时还要维护随上下文线性增长的 KV 缓存。为克服这些缺陷学界提出了大量次二次复杂度架构线性注意力Linear Attention、门控卷积、递归模型以及结构化状态空间模型SSM。SSM 家族从 S4 出发演化出 H3SSM 层两侧夹门控连接、HyenaMLP 参数化全局卷积、RetNet并行注意力路径与 RWKVLTI 递推WKV 可看作两个 SSM 之比。这些模型训练线性、推理恒定时间但在语言这类信息稠密、离散的模态上始终不如注意力——没有一个被证明能在规模上跨域有效。论文指出共同病根这些模型都是时不变LTI的参数不随输入变化因而无法做内容感知推理——不能根据当前 token 决定记住什么、忘记什么。更根本地序列建模的本质是把上下文压缩进有限状态注意力完全不压缩所以慢递归模型状态有限所以快但效果受限两者之间的桥梁就是选择性让压缩按内容进行。研究主线从问题到结论图 9Mamba 论文的研究主线——从注意力的效率缺陷出发定位 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 块同质堆叠扩展因子固定 $E2$使用 SiLU 激活。模型既无注意力也无独立 MLP 块大部分参数$3ED^2$在线性投影上。分类全景图 10Mamba 在次二次序列架构谱系中的位置——从 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论文证明了选择机制与经典门控的联系当 $N1$、$A-1$、$B1$ 时选择性 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 倍。图 2Mamba 块结构——两堆块对应 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。合成任务选择性的直接证据架构内部层Selective Copying 准确率S4S4LTI18.3%H3S457.0%H3Hyena30.1%MambaS456.4%MambaS6选择性99.8%图 3合成任务——选择性复制与归纳头直接检验内容感知能力。归纳头任务中模型在长度 $2^8256$ 上训练可外推到 $2^{20}1048576$4000 倍保持高准确率其他方法最多外推 2 倍——选择机制是唯一能外推的关键。语言建模匹敌两倍尺寸的 Transformer模型Pile ppl ↓LAMBADA acc ↑HellaSwag ↑平均 acc ↑Mamba-130M10.5644.335.344.7Pythia-160M29.6433.030.240.6Mamba-370M8.2855.646.550.0Pythia-410M9.9551.440.648.2Mamba-1.4B6.8064.959.159.7Pythia-1.4B7.5161.752.155.2Mamba 在每个尺寸档全面胜出1.4B 平均准确率 59.7 甚至超过同 tokenizer 的 Pythia-2.8B59.1。Mamba-3B 在常识推理上比 Pythia-3B 高 4 分匹敌 2 倍尺寸 Transformer——这是第一个匹配 LLaMa 配方 Transformer 的无注意力模型。图 4Pile 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$ 协同架构内部层PPL ↓H3S4real10.34H3S68.95MambaS4real10.56MambaS68.69选择性 $\Delta$选择性 $B$选择性 $C$PPL ↓✗✗✗10.93✓✗✗9.81✓✓✓8.71$\Delta$ 是最重要的选择性参数门控连接状态维数 $N$ 从 1 增到 16 困惑度下降超 1.0、仅增 1% 参数但只有 $B,C$ 也选择性时才有效附录 E。图 6归纳头外推曲线——选择性机制带来可无限外推的内容感知能力。图 7融合扫描内核的训练效率。结果对比总结图 11结果对比总结——质量匹敌两倍尺寸 Transformer效率线性扩展长上下文能力为三模态通用骨干奠定基础。关键发现选择机制是性能分水岭Selective Copying 从 S4 的 18.3% 跃升到 S6 的 99.8%归纳头可外推 4000 倍$2^8\to2^{20}$LTI 模型完全做不到。首次匹配 Transformer125M–1.3B 缩放律上 Mamba 是第一个无注意力模型匹配 LLaMa 配方且序列越长优势越明显。匹敌两倍尺寸Mamba-1.4B 零样本平均 59.7 超 Pythia-2.8B 的 59.1Mamba-3B 常识推理比 Pythia-3B 高 4 分。效率数量级提升推理吞吐 5 倍于 Transformer融合扫描比朴素实现快 40 倍训练内存与 FlashAttention 同级。三模态 SOTA音频 FID 从 1.99 降至 0.94减半以上基因组与音频性能随上下文单调提升至 1M。$\Delta$ 门控理论$N1$ 时选择性 SSM 精确退化为门控 RNN统一了 SSM 离散化与 RNN 门控两套理论。局限性规模有限实证仅到约 3B 参数低于 Llama/RWKV/RetNet 的 7B更大规模下是否保持优势未知规模化需额外工程。连续-离散免费午餐选择机制牺牲了 LTI 在连续信号音频/视频上的强归纳偏置音频实验需切回复数参数化附录 G。生态欠账微调、prompting、ICL、指令微调、RLHF、量化等 Transformer 生态的成熟适配机制尚未在 Mamba 上建立。硬件依赖性能依赖定制融合内核selective scan新硬件需重新工程化。常见问题FAQMamba 和 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-2Dao 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/mambaS4Gu, Goel, Ré, Efficiently Modeling Long Sequences with Structured State Spaces, ICLR 2022HiPPOGu et al., HiPPO: Recurrent Memory with Optimal Polynomial Projections, NeurIPS 2020H3Dao et al., Hungry Hungry Hippos, ICLR 2023FlashAttentionDao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022RWKVPeng 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工程文件