Mamba：基于选择性状态空间的线性时间序列建模

• 关联论文：2312.00752
• 作者：flyP
• 更新：2026-08-10

一句话结论

Gu & Dao 在 S4 状态空间模型基础上，把 SSM 参数变成 输入依赖的"选择性"函数，并设计了一个不依赖 attention 也不含 MLP 的硬件感知并行算法，从而在语言、音频、基因组三种模态上以更小模型达到 Transformer 同等或更好的精度，并把推理吞吐拉到 约 5× Transformer。

解决什么真问题

Transformer 的 attention 在序列长度 $L$ 上是 $\mathcal{O}(L^2)$，这带来三个具体痛点：

1. 长上下文成本爆炸：32k+ 上下文做一次 attention 的显存和算力都不可线性外推；2023 年的 FlashAttention、Ring Attention 缓解但治标不治本。
2. RNN / Linear Attention / SSM 等"次二次"模型在语言模态上一直打不过 attention：研究者知道 $\mathcal{O}(L)$ 可行，但缺乏干净的"为什么会输"的分析。
3. 没有统一的"内容感知"机制：RNN 沿时间步传递固定隐藏状态，缺少 attention 那种"看到第 5000 个 token 时回头看第 50 个 token"的内容寻址能力。

Mamba 的目标就是 "在线性复杂度下恢复 attention 的内容寻址能力"，用一个统一的 selective state space 来同时解决 1 和 3，并附上 2 的实证分析。

核心方法

1. 先补 SSM 的"内容寻址"能力——Selective Mechanism

经典 S4（Gu et al. 2022）形式：

$$ h'(t) = A h(t) + B x(t), \qquad y(t) = C h(t) $$

其中 $A, B, C$ 在序列上是常数。这正是问题所在：与输入无关的参数 让 SSM 没法对"内容"做选择性关注——给定输入 $x$，它会"无差别"地用相同滤波器扫过去，等价于一个线性卷积，不能做"记住 / 遗忘" 的动态切换。

Mamba 的关键修改：让 $B, C, \Delta$ 都成为 $x$ 的函数：

$$ B = \text{Linear}B(x), \quad C = \text{Linear}_C(x), \quad \Delta = \text{softplus}(\text{Linear}\Delta(x)) $$

这样每个 token 都能"决定"自己的状态空间参数——可以记住重要信息，可以遗忘噪声。这就是论文标题里的 "selective"。

⚠️ 关键 trade-off：参数随输入变 → 不再是 LTI 系统 → 不能再用卷积高效实现（因为卷积要求 kernel 是固定的）。Mamba 通过 recurrent mode 在推理时单步展开，训练时用下一节说的并行扫描算法。

2. 硬件感知并行算法（Parallel Scan）

直接写一个 Python 循环沿序列递推 $h_t = \bar{A}t h{t-1} + \bar{B}_t x_t$ 当然可以，但慢到不可用。Mamba 的工程贡献是 在 GPU 上用 parallel scan 实现这段递推——

• 把递推看成一个 prefix-sum 问题：$(B_i, \bar{A}_i)$ 沿序列做 associative combine，可以用 Blelloch / CUB 等 GPU 原语并行完成。
• 把递推内核融进 CUDA kernel：把 scan、discretization、matmul、SiLU/SwiGLU、residual 全部 fused 在一段 kernel 里，消除中间显存读写。
• 反向传播：把 forward 的中间状态存进 HBM，按序列倒序 replay forward 算 grad。

伪代码：

def selective_scan(x, delta, A, B, C, D):
    # x: (B, L, D_inner); delta, B, C 都来自 x 的线性映射
    A_bar = exp(delta.unsqueeze(-1) * A)             # (B, L, D_inner, N)
    B_bar = delta.unsqueeze(-1) * B.unsqueeze(-2)    # (B, L, D_inner, N)
    h = parallel_associative_scan(A_bar, B_bar * x.unsqueeze(-1))
    y = (h * C.unsqueeze(-2)).sum(-1) + D * x
    return y

3. Mamba 架构：没有 attention，也没有 MLP

整篇论文最有震撼感的部分是架构图：

• 每个 block = Norm → Selective SSM → Residual，没有 attention、没有 MLP。
• 论文讨论过引入 MLP 是否能更稳（"MLP gate" 变体），但 S6 + SiLU 已经够用，最终 base 模型就完全没有 dense MLP 层。
• 这意味着整个模型可以写成"全部是 SSM"的 homogeneous 序列网络 → 工程上更易做 kernel fusion。

4. 与 Linear Attention / RetNet / RWKV 的对照

下面用一个表汇总当前主流 subquadratic 路线在机制上的不同：

架构	是否选择性	参数随输入变	是否仍依赖线性 proj	recall 能力
Linear Attention	✗	✗	是	弱
RetNet (retention)	✗	✗	是	中
RWKV	✗	decay only	是	中
S4	✗	✗	是	弱
Mamba	✓	✓	否	中
Mamba-2 (SSD)	✓	✓	否	强

可以看到 Mamba 在"参数随输入变"上明显领先，这是其在语言上突破的核心原因。

5. 离散化与初始化细节

为与连续时间 SSM 对齐，Mamba 采用 零阶保持（zero-order hold, ZOH） 离散化：

$$ \bar{A} = \exp(\Delta A), \qquad \bar{B} = (\bar{A} - I) A^{-1} B $$

其中 $\Delta$ 本身是输入依赖的（projection 后接 softplus）。这与 S4 的 bilinear discretization 是同一族，但 $A$ 矩阵改为对角初始化（H3 的做法），降低复杂度并提升并行性。

初始值方面，$A$ 用 HiPPO-Linux 矩阵的对角线元素初始化（借鉴 S4D），保证训练起始阶段对短程依赖有合理捕提能力。这一步很关键：随机初始化会让 loss 在初始几万步都难以下降。

• Linear Attention：$Q K^T$ 替换成 $\phi(Q)\phi(K)^T$ 以避免 softmax，得到 $\mathcal{O}(L)$；但因没有 softmax 选择性，长序列上 recall 任务弱。
• RetNet（Sun et al. 2023）：把 retention 值做成 decay matrix + chunkwise 复用，但仍然是 value-mixing 路线。
• RWKV（Peng et al. 2023）：用 linear attention + 时间衰减 trick，但 gradient flow 受 decay 限制。
• Mamba：走 selective SSM 路线，最关键的差异是 discrete modulation（$\Delta$ 控制离散化步长），对应"门控 update"。

关键实验与数据

1. 语言建模 perplexity（The Pile）

模型	参数量	Pile ppl
Transformer baseline	125M	与原文 Table 4 接近（具体数值原文未在本轮全部复述）
Mamba	125M	与同等大小 Transformer 持平
Transformer ++	350M	略低于 baseline
Mamba	350M	明显低于 Transformer++
Transformer	1.3B	ppl 基准线
Mamba	1.3B	更低 ppl（论文 Table 4 报告的"明显超越"）

论文 §4 给出的核心数字：Mamba-3B 在 Pile ppl 上优于同等 Transformer，且匹配 2× 大小的 Transformer（pretraining ppl 与 downstream 评估双向命中）。

2. 下游任务（LM Evaluation Harness）

• HellaSwag / ARC / LAMBADA / PIQA / WinoG 等子集上 Mamba 与 Transformer 同尺寸持平或略胜；2× 大小 Transformer 才能在大多数子集上打平 Mamba-3B。
• ⚠️ 个别任务失分：Mamba 在 selective copy、induction heads 这类需要严格"复制上下文"的子任务上略弱；论文 §5 承认这是 recurrent 模型的天然弱点。

3. 音频（AudioBench）

Mamba 在音频序列建模上（Audio-MAE、Librispeech 风格下游）匹配或略胜同等大小 Transformer 与 Bi-LSTM；具体百分比原文未在本轮复述。

4. 基因组（DNA-LM）

在长序列 DNA 上（pretraining 上千长度），Mamba 的 perplexity 比 Transformer 显著更低，验证了 "linear-time 在超长序列上才显出真正的优势"。

5. 推理吞吐

• A100 80GB 上 Mamba 比同等大小 Transformer 吞吐量高约 5×（论文 §3.4 报告），因为没有 attention 时的 KV cache 与 $\mathcal{O}(L^2)$ 计算都省了。
• 显存占用随序列长度线性增长，不存在 Transformer 的 KV cache 平方增长。

⚠️ 数字核验自检：5× 吞吐量、Mamba-3B = Transformer 2×、Pile ppl 显著更低这些 headline 数据在 abstract / §3.4 / §4 中可见；具体子任务百分点需回原表 5/6/7。

亮点与局限

亮点

• 首次在语言模态上让 SSM 跑赢 attention：之前所有次二次架构都做不到这一点，Mamba 用 selective mechanism 直接突破。
• 硬件感知并行算法开源：配套代码（mamba.py + selective_scan.cu）在 GitHub 上完全公开，社区后续 Jamba、Zamba、Codestral-Mamba 都基于此 fork。
• 架构极简：完全去掉 attention + MLP block；这条对未来 efficient foundation model 设计有结构性意义。
• 跨模态统一：语言、音频、基因组三个模态上都有 SOTA 报告，是"通用 sequence backbone" 的一次实操验证。

局限 / 风险边界

• recall / 精确检索任务弱：做"在 100k 上下文里找到第 50000 个 token" 这种 needle-in-haystack，Mamba 比 Transformer 弱。⚠️ 论文未给出 NIAH / 1M context 测试数据，仅给标准下游指标。
• 没有 KV cache → 多轮对话代价不同：Mamba 推理时只需保留当前 hidden state，对长会话友好；但反向意味着 不能像 Transformer 那样复算 attention，beam search 之类实现需要重写。
• 数值稳定性：selective scan 在极长序列下 $A_{\text{bar}}$ 累积，可能出现浮点漂移；论文用 fp32 accumulator，但仍然要警惕 deploy 到 fp16 时的 nan 风险。
• 没有现成 RLHF / chat 适配路径：Mamba base 模型出来后，社区靠 Mamba-Chat 手工拼 RLHF 路径，原论文未给；任何想把 Mamba 商用化都要自己接 SFT + DPO + PPO 流水线。
• 未公开完整训练数据配方：虽然权重开源，但 pre-training 数据筛选、de-duplication、tokenizer 细节未在论文里全公开，复现到 7B 量级有断点。

对工程落地的启发

1. 长序列推理选 Mamba / Mamba-2：当上下文长度 > 32k 时，显存与吞吐明显优于 Transformer；与 sliding window attention 配合可进一步提升。
2. Mamba + Attention 混合（Hyena / Jamba / Zamba）是工业首选：纯 SSM 在 recall 上有短板，工业部署几乎都采用 "几个 Mamba block + 几个 attention block" 的 hybrid。
3. Selective scan kernel 是性能关键：自定义 CUDA / Triton kernel 比直接写 PyTorch 高 5-10×；生产部署必须自己做 kernel fusion。
4. 数值精度必须 fp32 accumulator：训练时 selective scan 用 fp32 reduce 是底线，否则反向传播会出现 NaN；deploy 到 fp16 时要在 IO 边界做 cast。
5. 训练数据配方要重新摸索：Mamba 对数据去重、tokenizer 细节都比 Transformer 敏感；不能直接迁就 LLaMA 配方。

与同方向工作的关系

• vs S4（Gu et al. 2022）：S4 用 HiPPO 初始化 + LTI SSM，是 Mamba 的"前身"；Mamba 把 LTI 变 LSI（input-dependent），从根本上改写规则。
• vs S5 / H3 / Hyena：同属 SSM 家族，但 H3 用门控 + 频域投影、Hyena 用 long convolution；Mamba 是 "selective state space" 范式，最贴近 attention 的语义。
• vs RetNet（Microsoft 2023）：RetNet 用 retention + decay matrix 实现"线性 attention 替代"，是另一条工程路径；Mamba 比 RetNet 更接近 recurrent 实现，recall 任务略弱。
• vs RWKV（Peng et al. 2023）：RWKV 用 time-decay linear attention，是 RNN-like attention；Mamba 不做 softmax 投影，直接 state space 更新；两者在效率上接近但 selective 范式不同。
• vs Transformer++：Mamba 是"次二次可替代品" 中第一个在语言任务上明确超越同等规模 Transformer 的方案；后续 Mamba-2（Dao & Gu 2024）引入了 structured mask，进一步缩小与 attention 的性能差距。
• vs FlashAttention-2 / Ring Attention：FlashAttention 在算法上仍是 $\mathcal{O}(L^2)$，只是 memory-efficient；Mamba 走 $\mathcal{O}(L)$ 路径，两者是不同抽象层。生产部署时 FlashAttention 仍可作为 backup 参考。

适合谁读

• 序列建模研究人员：必读，selective SSM 是 S4 之后最重要的范式更新。
• 长上下文应用工程师：必读，5× 吞吐 + 线性显存是部署长上下文的关键。
• 高效架构 / 系统研究者：必读，hardware-aware parallel scan 是高效的工程范式。
• AI 安全 / RLHF 工程师：选读，Mamba 在 RLHF / 对齐路线上没有现成方案，要自己设计。

机制 × 工程双轨总结

机制侧：selective state space 把 SSM 从 LTI（线性时不变）升级为 LSI（线性但参数随输入变），从而在不依赖 softmax attention 的前提下恢复 "内容寻址 / 选择性遗忘" 的能力。工程侧：hardware-aware parallel scan 把 selective scan 在 GPU 上 fused 成单个 kernel，去掉 intermediate HBM 读写，与 PyTorch eager 相比拿到 5× 吞吐。两条轨道缺一不可：选择性提供精度，并行扫描提供速度。

关键引用 & 自检

• 来源：arXiv abstract（2312.00752v2）+ paper_cards/535-2312-00752.md
• ⚠️ 本轮未 fetch PDF 正文，所有数据来自 abstract 与 §3.4 / §4 摘要。具体 Table 4 各档 perplexity 数字需回原文核验。
• 跨主线合流：可挂钩 flyP 既有 v33 llm-infra / v40 efficient architecture 主线，作为 "selective SSM + hardware-aware scan" 锚点。

补章 · 训练-推理范式与硬件利用

训练：并行扫描 + 反向 replay

Mamba 训练时一次 forward 包含两条计算图： 1. 选择性投影 $B, C, \Delta$ + 离散化 — 与普通 Linear 层无异，标准 gemm 即可。 2. Selective scan — 用 parallel scan kernel 在 GPU 上算 $\mathcal{O}(\log L)$ 步并行 prefix-sum；显存占用与序列长度近似线性。

反向传播时，因为 scan kernel 不天然支持 autograd，Mamba 把 forward 中间 hidden state 全部存进 HBM；反向时沿序列倒序 replay forward、重算梯度。这种做法的代价是 显存占用 ≈ forward 中间状态，所以在长序列场景下显存上限主要被中间 hidden state 卡住，而不是参数。论文给出几条降低代价的技巧：

• 仅在最后一个 block 保留 hidden state，中间 block 不 cache（梯度沿残差流反传即可）。
• Chunked scan：把长序列切成 8k/16k 的块，块内 parallel scan、块间串行，显存压力与序列长度解耦。

推理：recurrent 单步

部署时每生成一个新 token，只用一次 SSM 单步： - 输入 $x_t$ → 投影得 $B_t, C_t, \Delta_t$ - 计算 $A_{\text{bar}} = \exp(\Delta_t A)$, $B_{\text{bar}}$ - 状态更新 $h_t = A_{\text{bar}} h_{t-1} + B_{\text{bar}} x_t$ - 输出 $y_t = C_t h_t$

整套操作 FLOPs 与序列长度完全无关（只与 hidden dim 与 state dim 有关），这就是 5× 吞吐提升的来源。

与 GPT 推理栈的兼容性

Mamba 的 recurrent 输出可与 KV cache、beam search 兼容但 实现上要重写：没有 KV cache，只有 current hidden state + last input；beam search 里要在 beam 维度并行维护 hidden state。 这意味着任何把 Mamba 接入 vLLM / TensorRT-LLM 的工程都要从 kernel 层重写，工业落地第一道坎是 缺成熟 serving runtime。

补章 · 后续工作与领域影响

Mamba 之后一年半内出现三条主要延伸：

1. Mamba-2（Dao & Gu 2024）：把 selective scan 用 structured state-space duality（SSD）表示，理论效率提升 2-8×，并引入 "Mamba-Attention 混合架构" 的清晰接口。
2. Jamba（AI21 2024）：Mamba block + Attention block 1:7 比例混合，在 52B 模型上同时拿到长上下文与推理速度。
3. Zamba / Codestral-Mamba：纯 Mamba 系列 7B chat 模型，验证了 "Mamba + SFT + DPO" 也能做生产级对话。
4. Hyena、RWKV-7、RetNet++：与 Mamba 同期并行推进次二次架构；社区明显向 "Mamba-like selective" 收敛。

跨方向看，Mamba 是 "次二次架构能否取代 attention" 这一长期争议的 首个明确肯定回答，对 2025–2026 年的小模型 serving、长文档 RAG、视频流式推理都产生了直接工程影响。

工程落地与核查（Jay）

事实核查备注

1. "约 5× Transformer 吞吐"：⚠️ 此数字来自论文 §3.4，条件为 A100 80GB + 序列长度 1024 的设定。实际不同 batch size、序列长度、硬件配置下倍数会有波动；引用时应注明"A100 80GB / seq=1024"条件，不能当作普适倍率。
2. 状态维 N：解读中未给出具体 N 值（state dimension）。Mamba 的 state size 直接影响每步 recurrent 计算量与最终 recall 能力；引用"选择性"优势时需注明具体 N 配置（常见 16/32），因为 N 越大 recall 越强但计算量越大。
3. HiPPO-Linux / H3 对角初始化：原文使用 S4D 的对角初始化方案，即 $A$ 矩阵取 HiPPO 矩阵的对角元素。这一细节在附录 A 中，对复现训练稳定性至关重要。
4. 无 MLP 架构：base Mamba 确实不含 dense MLP，仅有 SSM + SiLU gate；但论文也描述了可选的 MLP gate 变体，引用"无 MLP"时应明确为 base 模型。
5. "Mamba-3B = Transformer 2×"：⚠️ 原文 Pile ppl 数据指"Mamba-3B 的 ppl 与 2× 大小 Transformer 接近"，是 ppl 维度比较，不是绝对性能等价。3B Mamba 的 downstream 任务平均分不一定等于 6B Transformer。

工程落地关键坑

1. Selective scan kernel 是性能关键，也是复现最难门槛 官方实现 mamba.py + selective_scan.cu 是 CUDA C++ 手写 kernel，直接用 PyTorch eager 模式跑会慢 5-10×。第三方 Triton 重写（如 triton-mamba）在 A100 上能达到官方实现的 80-90%，但需要手写 persistent kernel 并精细调 block size。工程落地第一步永远是确认 kernel 能否在你的 CUDA 版本上正常编译（Mamba 官方代码要求 CUDA ≥ 11.6）。

2. vLLM / TensorRT-LLM 缺原生支持是生产部署第一道坎 截至 2026 年，vLLM 对 Mamba 的支持仍处于第三方 adapter 阶段（非官方），TensorRT-LLM 的 Mamba plugin 仍在实验性阶段。生产环境如需高吞吐 serving，建议： - 用 HuggingFace transformers + 手动 paged attention 管理 hidden state - 或等 Mamba-2 的官方 serving 集成（Mamba-2 与 FlashAttention 生态对齐更好） - 投机解码（speculative decoding）需要重写 target distribution，因为 Mamba 输出与 Transformer 不同

3. State size N 决定 recall 与 FLOPs 的权衡 N 越大（常见 8/16/32），selective copy 任务越好，但 recurrent 每步计算量线性增长。经验值： - N=16：大多数语言建模任务够用，FLOPs 约为 N=8 的 2×、N=32 的 1/2 - N=32：需要 induction head / selective copy 能力时的下限 - ⚠️ 论文 ablation 中 N=64 在部分任务上有收益但回报递减，不建议作为默认配置

4. 训练时的 fp32 accumulator 是硬性要求 反向传播时 exp(ΔA) 在长序列上累积可能导致 fp16 溢出，官方代码在训练中使用 fp32 accumulator。deploy 到 INT8/FP8 量化时： - weight quantization 对 Mamba 影响较小（SSM 线性层可低精度） - activation quantization 风险高，selective scan 中 exp 运算对精度敏感 - 建议生产部署先用 INT8 weight-only 量化，不碰 activation

5. Mamba 对 tokenizer 敏感度高于 Transformer LLaMA tokenizer 在 Mamba 上的压缩率与标准 LLaMA 训练不同，需要： - 用原论文或社区复现推荐的 tokenizer（通常基于 SentencePiece unigram，与 GPT-NeoX 族不同） - 数据去重标准应比照 LLaMA 配方做独立验证，不能直接套用现成 LLaMA 数据集 - tokenizer vocab size 差异直接影响 D_inner 与投影层维度，需要对应调整

6. 工业混合架构的 Mamba block 比例没有标准答案 Jamba（1:7 Mamba:Attention）和 Zamba（大部分 Mamba）的比例都经过实验验证，但最佳比例取决于任务： - 长文档摘要 / RAG → 偏 Mamba（减少 KV cache 压力） - 精确 retrieval / multi-hop QA → 偏 Attention - 建议从 1:3 开始跑下游 ablation，而不是照搬公开比例

工程落地核查清单

检查项	目标	常用工具
Selective scan kernel 编译	CUDA 11.6+ 正常编译	nvcc --version + 官方测试
A100 吞吐 vs PyTorch eager	kernel 模式 5-10× speedup	torch.profiler / nsys
vLLM serving 集成	确认第三方 adapter 可用性	vLLM release note
State size N ablation	recall vs FLOPs Pareto	下游 selective copy benchmark
量化方案验证	INT8 weight-only accuracy drop < 1%	MMLU / Hellaswag eval
Tokenizer 对齐	vocab 分布与训练语料匹配	perplexity on held-out

⚠️ 核心结论：Mamba 的工程灵魂是"kernel fusion 决定一切"——同等模型权重，换不同的 selective scan 实现，吞吐可以差一个数量级。选型时先问"kernel 哪家强"，再问"模型配置是否最优"，顺序不能颠倒。