96 lines
3.8 KiB
Markdown
96 lines
3.8 KiB
Markdown
---
|
||
title: "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
|
||
created: 2026-06-18
|
||
updated: 2026-06-18
|
||
type: paper
|
||
tags: ["state-space-models", "linear-complexity", "sequence-modeling", "architecture"]
|
||
sources: ["https://arxiv.org/abs/2312.00752"]
|
||
---
|
||
|
||
# Mamba：选择性状态空间的线性时间序列建模
|
||
|
||
## 核心问题
|
||
|
||
[[state-space-models|状态空间模型（SSM）]]作为 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 通过以下方式解决：
|
||
|
||
1. **并行关联扫描（parallel associative scan / Blelloch scan）**：将循环更新展开为前缀和操作，可在 GPU 上并行
|
||
2. **IO 感知的 kernel fusion**：在 SRAM 中完成 scan 和离散化，避免将扩展状态写入 HBM
|
||
3. **重计算（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
|
||
|
||
## 来源
|
||
|
||
- [arXiv:2312.00752](https://arxiv.org/abs/2312.00752)
|
||
- [原始存档](raw/papers/gu-mamba-2024.md)
|