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

核心问题

state-space-models作为 Transformer 的线性复杂度替代方案，在连续信号（音频、视觉）上表现出色，但在离散信息密集数据（如文本）上从未达到 Transformer 质量。Gu & Dao (2024) 识别出根本原因：缺乏内容感知推理——传统 SSM 的动力学是时间不变的（LTI），每个 token 的处理规则完全相同，无法选择性关注或忽略信息。

核心贡献：选择机制（S6）

Mamba 的关键创新是将 SSM 从 LTI（线性时间不变）升级为选择性（Selective）：

S4 (LTI):    B, C, Δ 对所有时间步固定 → 卷积 OR 循环
S6 (Selective): B_t, C_t, Δ_t 是输入 x_t 的函数 → 仅循环 (需 scan)

算法对比（核心直觉）

算法	B	C	Δ	性质
S4	Parameter (D, N)	Parameter (D, N)	τ(Parameter)	LTI, 可用卷积
S6	s_B(x_t): (B, L, N)	s_C(x_t): (B, L, N)	τ(Δ + s_Δ(x_t))	时间变化, 需 scan

效果：S6 模型能根据当前 token 内容决定是传播还是遗忘信息——这就是选择性。

为什么这很重要

在 selective-copy任务中，传统 LTI 模型只能在 token 间距固定时成功（只需时间感知，不需内容感知）。一旦间距随机变化，只有能"看到内容再决定是否记住"的模型才能胜任。Mamba 的选择机制天然支持这种内容感知。

第二个创新：硬件感知算法

选择机制带来了计算挑战：时间变化的 SSM 不能再用卷积（卷积要求 LTI）。Mamba 通过以下方式解决：

并行关联扫描（parallel associative scan / Blelloch scan）：将循环更新展开为前缀和操作，可在 GPU 上并行
IO 感知的 kernel fusion：在 SRAM 中完成 scan 和离散化，避免将扩展状态写入 HBM
重计算（recomputation）：反向传播时不保留中间状态，直接重算

这三个技巧使 Mamba 比所有卷积 SSM 快 3×（A100 GPU）。

架构：极简设计

Mamba block 结构：

x → LayerNorm
  → Linear(d → 2d) → Conv1d → SiLU  [门控分支]
  → Linear(d → 2d_N) → SSM(S6)       [SSM 分支]
  → 逐元素乘法 → Linear(2d → d)
  → + x (残差)

关键设计哲学：

无注意力、无 MLP：单一块类型统管整个模型
扩展比 E=2：计算量可控
同质架构：所有层结构相同，仅参数不同

与 H3 的关系：H3 使用两个 LTI SSM + 门控，Mamba 将门控 SSM 融合为单一选择性 SSM。

实验结果总结

模态	结果
语言	Mamba-3B > Pythia-3B，匹敌 Pythia-7B；5× 推理吞吐
合成任务	Selective Copying + Induction Heads 外推至 >1M tokens
音频	SC09 语音生成 FID 降低 >50%
基因组学	>HyenaDNA，>Transformer

关键概念网络

selective-state-space — S6 选择机制
hardware-aware-algorithm — GPU 优化并行 scan
structured-state-space-models — S4 前身
content-based-reasoning — Mamba 解决的 LTI 弱点
selective-copy — 动机合成任务
induction-heads — LLM 关键机制
hippo — SSM 数学基础
mamba-ssm — 概念主页（已有，需更新）
state-space-models — SSM 家族总览（已有）

代码

https://github.com/state-spaces/mamba

3.8 KiB Raw Blame History Unescape Escape