论文精读：Mamba——线性时间序列建模的选择性状态空间

为什么精读这一篇

2023 年 12 月 1 日，Albert Gu（CMU）与 Tri Dao（FlashAttention 作者）把 arXiv 2312.00752 挂上 arXiv，题目直白：Mamba——线性时间序列建模的选择性状态空间。三年过去，这篇论文被引过万（Semantic Scholar 口径），是「Transformer 替代品」里被验证最充分的一支。本站已有架构综述与 FlashAttention、KV Cache 的专题文章，本文不铺历史，专注把这一篇讲透：动机、机制、算法、数字，再补截至 2026-10-07 的后续演进。

动机：Attention 的账单与 SSM 的路线

自注意力之所以有效，论文引言归结为「上下文窗口内信息的密集路由」，代价写在明处：训练开销随序列长度 N 二次增长，推理时要背着与上下文等长的 KV cache，长序列场景两头难受。

状态空间模型（SSM）走了另一条路。它源自控制论的连续系统 h′(t) = A·h(t) + B·x(t)、y(t) = C·h(t)，离散化后变成线性递推 h_t = Ā·h_{t-1} + B̄·x_t、y_t = C·h_t——每一步只维护一个固定维度的状态 h，序列长度线性扩展，推理状态恒定。

S4（Gu 等，ICLR 2022）是这个路线的代表作：用 HiPPO 理论（对输入历史做最优多项式投影的记忆框架）初始化 A，让小状态尽量留住历史；且离散化后的系统既能按递归算、也能按全局卷积算（两种数学等价的形式），训练走卷积分支即可并行。但 SSM 有个多年没解决的老毛病：对文本这类离散、信息稠密的数据「不够有效」。根因在于参数是时不变的——Δ、A、B、C 对每个 token 一视同仁，不会按内容取舍。

核心创新：让参数跟着输入走

论文 §3.1 给出一个干净的压缩视角：序列建模的本质是把上下文压缩进有限状态。Attention 完全不压缩（KV cache 全量保存历史），RNN 压缩太狠（固定小状态，没有取舍的裁量权）。好的模型应当「该记的记、该忘的忘」——这就是选择性（selectivity）。

做法落在算法 2（作者称为 S6，即 S4 + selection + scan）：把 Δ、B、C 三个参数变成输入 x 的函数（各自过一个小线性投影，Δ 再经 softplus 保证为正），A 保持不动。两个直觉值得记住：

Δ 是门控的推广。论文定理 1 证明：当 N=1、A=−1 时，递推退化成 h_t = (1−g_t)·h_{t-1} + g_t·x_t 的经典门控 RNN。Δ 大则看重当前输入，Δ 小则忽略它、保住历史——GRU 式的门控思想被装进了 S4 的骨架。
状态可随时重置。遇到噪声 token 就调小 Δ 视而不见；遇到文档边界就清空状态，拼接序列互不串味。

代价同样清楚：参数随 token 变化后系统不再时不变，卷积形式消失，只剩递归——而递归看起来只能串行。这是下一节要拆的矛盾。

硬件感知算法：把矛盾装进 SRAM

选择性带来表达力，也让并行化看似无解。论文 §3.3 用三件套化解：

并行扫描（parallel scan）。递推算子满足结合律，O(N) 的串行递推可以重排成 O(log N) 深度的树形扫描，时间维度不再串行。
kernel 融合。扫描是内存带宽瓶颈：把离散化、逐步递推、与 C 相乘整个流程融合进一个 CUDA kernel，状态只在 GPU 的快速 SRAM 里存活，从不物化到慢速 HBM，最后只把输出写回——这正是作者做 FlashAttention 的同款思路。
重计算。反向传播不保存中间状态，需要时从 HBM 重载输入重算，显存开销与 FlashAttention 级别的优化 Transformer 相当。

「选择性与并行化看似矛盾却兼顾」的直觉在：正因为有选择性，状态才可以小而有效（消融显示状态维 N=16 就够，只增约 1% 参数），小状态才塞得进 SRAM、撑得起融合 kernel；时间维度的串行交给树形扫描。选择性管质量，硬件算法管把它跑起来。效果是训练中这套扫描比标准 PyTorch 实现快 20–40×，A100 上比此前 SSM 内核最快 3×，序列超过 2K 后训练速度快于 FlashAttention-2。

最终架构反而更简：Mamba 块把 S6 与门控分支融合（H3 风格），整个模型没有 attention 块、也不再有独立的 MLP 块。

flowchart LR
    subgraph ATTN["Attention：全量上下文，二次开销"]
        XA["输入 x1..xN"] --> QA["投影出 Q、K、V"]
        QA --> MA["N×N 两两打分"]
        MA --> OA["输出 y1..yN，KV cache 随长度增长"]
    end
    subgraph S6["Mamba：状态压缩，线性开销"]
        XB["输入 x_t"] --> PB["一次投影同时产出 Δ、B、C"]
        PB --> DB["Δ 离散化 A：Ā = exp(ΔA)"]
        DB --> SB["并行扫描：h_t 由 h_t-1 与 B̄·x_t 递推"]
        SB --> OB["输出 y_t = C_t·h_t"]
    end

论文数字：以小胜大，线性到百万

Scaling laws（Pile，125M–1.3B，Chinchilla 协议）：Mamba 是第一个在困惑度与下游平均分上匹配「Transformer++」（PaLM/LLaMA 一代配方）的 attention-free 架构，序列越长优势越大。
以小胜大（Table 1）：Mamba-2.8B 零样本平均 63.3，高于同规模 Pythia-2.8B（59.1）、RWKV-3B（59.6），也超过参数量约 2.5 倍的 Pythia-6.9B（61.7）与 RWKV-7.4B（62.5）——对应摘要那句「3B 匹敌两倍大小的 Transformer」。
吞吐：摘要明确给出推理吞吐比 Transformer 高 5×——没有 KV cache，batch 可以开得更大。
长度外推（合成任务）：induction heads 任务训练于长度 256（2⁸），测试直接推到 2²⁰ ≈ 105 万 token，Mamba 完美外推（训练长度的 4096 倍），对照方法不超过 2 倍；selective copying 上含 S6 层的架构做到 99.8%，S4 层只有 56.4%、Hyena 层 28.4%——没有选择性的层在这个任务上根本学不动。
真实长序列：DNA 建模（人类基因组 HG38，约 45 亿 token）上下文从 1024 一路推到 100 万，困惑度持续下降，而同条件的 HyenaDNA 随长度变差；音频生成（YouTubeMix 钢琴，16kHz）序列长度推到约 10⁶ 仍随长度改善。摘要里那句「性能随序列长度提升至百万级」由这三组实验共同背书。

Mamba 之后：SSD、混搭与 2026 年的现实

Mamba-2（2024 年 5 月，两位作者署名对调）补上了理论拼图：**状态空间对偶（SSD）**框架证明，状态维 N 的 SSM 变换等价于乘一个「半可分离矩阵」；它的二次形式等价于去掉 softmax 的注意力 (L∘QKᵀ)·V，其中 L 是数据依赖的衰减掩码——SSM 与（线性）attention 从此是同一族计算的两种视角。架构上 Mamba-2 把 A 简化为标量乘单位阵、引入类似 multi-head 的头结构，专用 kernel 比一代 selective scan 快 2–8×，并允许 8× 更大的递归状态；Pile 300B token 上 Mamba-2-2.7B 超过 Mamba-2.8B 与 Pythia-6.9B。大状态正是冲着召回短板去的。

混搭（hybrid）则是长上下文的主力配方。AI21 的 Jamba 把 Transformer、Mamba 与 MoE 揉在一起：每 8 层只放 1 层 attention、7 层 Mamba（1:7），总参 52B、活跃 12B，发布时支持 256K 上下文（训练时验证到 1M）；256K 下 KV cache 只占 4GB（同规模 Mixtral 是 32GB），128K 长上下文吞吐约 3× 于 Mixtral。直觉很清楚：Mamba 层用线性成本「流过」绝大部分 token，少数 attention 层专职精确召回。

截至 2026-10-07 的现状（以本次检索为准）：Attention 在主流大模型里仍是底座，Mamba 没有取代它，但已成为长序列与多模态分支的常规选项——混合架构被大量长上下文模型采用，纯 SSM 在基因组、音频、时序等超长序列领域站稳。系列本身仍在演进：Mamba-3（ICLR 2026，arXiv 2603.15569，代码已并入 state-spaces/mamba 官方仓库）引入 MIMO 式状态与复数状态扩展，继续朝「更大状态、更高硬件利用率」推进。

边界：什么时候别用、什么时候该用

召回与精确复制是固有短板。固定大小的状态注定是有损压缩，attention 的 KV cache 却能无损直达历史——「电话簿式检索」「精确复制随机串」这类任务上纯 SSM 系统性落后，后续研究把这叫作 recall gap。机制研究（Arora 等）发现 Mamba 能学出类似 induction heads 的机制，但需要不同的组件与训练条件；缓解路线是加大状态（Mamba-2 的 8× 状态、Mamba-3 的 MIMO），或干脆混入 attention 层。论文自己也不回避：§5 明说选择性帮了离散模态（文本、DNA），反而可能拖累 LTI 擅长的音频、视频等连续信号。
生态与硬件差距。论文 §5 自述：Transformer 生态（微调、prompting、in-context learning、RLHF、量化）对 SSM 是否成立当时未知——这些能力后来被混合架构与社区工作补了大半，但纯 SSM 的工具链仍薄。硬件侧 attention 有 Tensor Core 与 FlashAttention 全家桶，SSM 依赖定制 CUDA kernel：官方 mamba-ssm 包要求 Linux、CUDA 11.6+ 并现场编译，ROCm 还要打补丁，工程摩擦肉眼可见。
选型直觉（本文作者判断）：超长序列（10 万 token 以上）、高吞吐推理、流式或近乎无限的上下文、基因组与音频信号，是 SSM 路线的主场；需要强检索、强召回与复杂工具调用的通用对话模型，attention（或至少混合）仍是稳妥选择。

为什么精读这一篇

动机：Attention 的账单与 SSM 的路线

核心创新：让参数跟着输入走

硬件感知算法：把矛盾装进 SRAM

论文数字：以小胜大，线性到百万

Mamba 之后：SSD、混搭与 2026 年的现实

边界：什么时候别用、什么时候该用

参考资料

读者留言

论文精读：Mamba——线性时间序列建模的选择性状态空间

为什么精读这一篇

动机：Attention 的账单与 SSM 的路线

核心创新：让参数跟着输入走

硬件感知算法：把矛盾装进 SRAM

论文数字：以小胜大，线性到百万

Mamba 之后：SSD、混搭与 2026 年的现实

边界：什么时候别用、什么时候该用

参考资料

读者留言

相关阅读