Mamba: Linear-Time Sequence Modeling with Selective State Spaces

背景

Transformer 的两笔账在 2023 年已经到了不能忽视的程度：训练的 O(N²) 注意力矩阵（FlashAttention 缓解了访存但 FLOPs 仍是平方），推理的 KV Cache 随序列长度线性膨胀（百万 token 上下文的 cache 显存让长窗口部署几乎不可行）。状态空间模型（SSM）提供了另一种可能：把历史压缩进固定大小的隐状态，推理成本与序列长度解耦——但当时的最强 SSM（S4，Gu et al. 2021）有个致命限制：线性时不变（LTI）。

LTI 意味着状态转移规则对所有输入一视同仁——同一套「记忆衰减率」用于所有 token。这带来一个漂亮的礼物（LTI 系统等价于卷积，可以并行训练）和一个残酷的代价：模型无法做内容寻址——不能「看到关键词 X 就把当前状态存下来」。注意力的本质恰恰就是内容寻址（query 与 key 匹配决定读写），所以 S4 在「精确检索」类任务上系统性地输给注意力。

动机

Albert Gu（Princeton，S4 作者）与 Tri Dao（Stanford/CMU，FlashAttention 作者）联手的目标：让 SSM 获得选择性（输入相关的记忆读写），同时不放弃线性复杂度与训练并行性。这是一个「既要又要还要」的问题——选择性与卷积形式在数学上直接冲突（参数一旦依赖输入，卷积核就不再固定），两人的解法与 FlashAttention 一脉相承：接受递归形式，然后用硬件感知的算法把递归算出并行度。序列建模的一次「架构级豪赌」：赌选择性 + 线性扩展能在语言建模上追平注意力。

方法

第一步：选择性 SSM（S6）。离散化后的 SSM 递归 h_t = Ā(h_{t-1})·x_t 形式。Mamba 把其中的关键参数（Δ_t 步长、B_t、C_t）全部改为当前输入 x_t 的函数（线性投影 + SiLU）。效果：Δ_t 大 → 状态被重写（记住当前 token）；Δ_t 小 → 状态保持（忽略当前 token）——选择性 = 学到的门控记忆，与 LSTM 的门有血缘关系但形式更简洁。

第二步：硬件感知并行扫描。失去卷积形式后，训练必须沿序列做扫描（scan，前缀积式的递归）。朴素实现是串行的（一个 token 一个 token 过），GPU 利用率为零。Mamba 的解法：(1) 把序列切块（chunk），块内的扫描展开成矩阵乘（可并行）；(2) 块间用一个小规模的 scan 合并（块数远小于 token 数）；(3) 中间状态留在 SRAM、不落 HBM，反向传播用重计算——FlashAttention 的 tiling 哲学原样移植到 scan 上。

第三步：极简架构（Mamba block）。把 SSM 与 MLP 融合成一个统一残差块：输入线性扩展 → 一维卷积（局部特征）→ SiLU 激活 + SSM（选择性记忆）→ 门控输出。对比前代（S4 带复数分支、S5 的分层结构）大幅删减：没有注意力、没有复数、没有多分支——参数预算全部给宽度。整个 block 的设计与现代 MLP 一样「无个性」，个性全在 SSM 核里。

理论侧：作者推导了选择性 SSM 的初始化方案（保持 S4 的稳定动力学），并分析了 Δ、B、C 的角色分工（Δ 是「写入门」，B 是「写什么」，C 是「读什么」）。

实验设置

• 模型矩阵：125M/350M/760M/1.4B/2.8B 五档，与 Pythia/GPT-NeoX（同 tokenizer 同数据）严格对齐——干净对照
• 语言建模：The Pile（验证 loss 曲线）、WikiText-103（困惑度）、PG-19（整本书级长文档）与长序列合成任务（multi-query associative recall——检验精确检索）
• 推理吞吐基准：与等尺寸 Transformer 的 batch 生成吞吐对比（512-1M 序列）
• Scaling 曲线：Chinchilla 口径的「loss vs FLOPs」前沿，检验「同算力下质量不输」
• DNA 基因组建模（长序列杀手锏，前代 S4 的主场）：序列长度百万级

结果

• 质量：2.8B Mamba 在 The Pile 验证 loss 上匹配同尺寸 Transformer（±0.05 内），零星超过；WikiText-103 困惑度持平或略优——「线性架构追平注意力质量」第一次在语言上成立
• 推理吞吐：因无 KV Cache，生成吞吐比同尺寸 Transformer 高 5×（~5200 token/s vs ~1050 token/s 量级，A100），且序列越长差距越大——Transformer 的显存墙在 Mamba 上不存在
• 训练效率：同 token 数训练 wallclock 快 ~（与序列长度相关）约 1.5-3×
• 长序列：PG-19 与 DNA 任务上，Mamba 的「每 token 平均 loss 随位置」曲线比注意力衰减更平——固定大小状态意外地能记很长
• 检索任务的裂缝：合成 multi-query 关联召回任务上，纯 Mamba 在序列很长（16K+）时明显掉分——固定状态终究装不下「任意时刻精确回放」，论文自己也承认这是注意力的领地

有趣的洞见

「记忆的带宽-容量权衡」被 Mamba 摆上台面。注意力的记忆是「全部保留 + 现场检索」（KV Cache = 完整的历史），SSM 的记忆是「固定容量的压缩摘要」——前者精确但昂贵，后者便宜但有损。Transformer 的长上下文困境本质是「带宽无限、容量爆炸」；Mamba 是「容量恒定、带宽有损」。序列建模的核心变量从来不是架构名号，而是「多少历史以什么精度保留」——这个视角把 Mamba、FlashAttention、KV Cache 量化、滑动窗口注意力统一成一个问题家族。

两个研究传统的焊接。S4 来自控制论/信号处理传统（连续时间系统、HiPPO 理论），FlashAttention 来自 GPU 系统工程传统（存储层级、kernel 融合）。Mamba 的成功说明深度学习的很多「新问题」其实有老学科的现成答案——缺的只是把老答案翻译到新硬件上的人。

混合架构是对论文最诚实的继承。Jamba（AI21，2024：每 6-8 层塞一层注意力）、Griffin（DeepMind）、Samba、Zamba……主流路线全是「Mamba 块 + 少量注意力层」——用线性层管日常依赖（便宜），用注意力层管精确检索（少量、够用）。DeepSeek-V2 的 MLA 稀疏注意力从另一头逼近同一目标。终局不是取代，是分工。

批判性评价

最弱的一环：证据规模与任务覆盖。所有对照止步 2.8B，且主要看语言建模 loss——社区放大实验（Falcon-Mamba 7B、Zamba 等）显示结论与任务强相关：记忆密集任务（长上下文检索、代码补全、in-context learning 深拷贝）上注意力仍占优；「Repeat After Me」（Jelassi et al. 2024）直接证明纯 Mamba 在复制类任务上系统性失败。生态成本也被低估：无 KV Cache 意味着整个推理栈（vLLM 等的 paged attention、前缀缓存、投机解码）都要重写。Mamba-2/SSD 框架的「承认式改进」（把 SSM 与注意力统一在一个「结构化掩码注意力」视角下）某种意义上是作者自己给原论文局限画的句号。

依然重要的理由：它是「Transformer 不可撼动」氛围里第一篇让人重新思考基本问题的论文——序列建模到底需要什么？即使在混合架构终局里，Mamba 的选择性扫描也成为标配组件之一；它证明的「硬件感知算法可以救活被放弃的数学形式」（递归）这个方法论，比模型本身更值钱。