---
title: "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"
authors: ["Albert Gu", "Tri Dao"]
date: 2023-12-01
arxiv_id: "2312.00752v2"
categories: ["cs.LG", "cs.AI"]
affiliations: ["Carnegie Mellon University", "Princeton University"]
paper_type: "conference"
code: "https://github.com/state-spaces/mamba"
---

# Mamba: Linear-Time Sequence Modeling with Selective State Spaces

## 摘要

Foundation model 几乎全部基于 Transformer 架构，但其注意力机制的二次复杂度在处理长序列时效率极低。各种次二次复杂度架构（线性注意力、门控卷积、结构状态空间模型）试图取代注意力，但在语言等核心模态上始终达不到 Transformer 质量。本文识别出这些模型的根本弱点——**缺乏内容感知推理能力（content-based reasoning）**——并通过两个关键创新解决：(1) 让 SSM 参数成为输入的函数（选择机制，S6），使模型能根据当前 token 选择性传播或遗忘信息；(2) 设计硬件感知的并行算法，在循环模式下高效计算。最终形成极简架构 Mamba——无注意力甚至无 MLP 块。Mamba 推理吞吐量是 Transformer 的 5 倍，序列长度线性扩展，在语言、音频、基因组学等多个模态达到 SOTA。Mamba-3B 性能超过同规模 Transformer 并匹敌两倍规模的 Transformer。

## 核心贡献

1. **选择机制（Selection Mechanism / S6）**：将 SSM 参数（Δ, B, C）变为输入依赖，从时间不变（LTI）升级为时间变化
2. **硬件感知算法**：通过并行关联扫描（parallel associative scan）在 SRAM 中计算，避免 GPU HBM 之间的 IO 瓶颈
3. **极简架构 Mamba**：将 H3 架构中的 SSM 层与 MLP 门控融合为单一同质块
4. **选择复制（Selective Copying）和归纳头（Induction Heads）合成任务**：Mamba 不仅轻松解决，且能无限外推（>1M tokens）

## 方法框架

### 从 S4 到 S6

传统 S4 的关键局限是 **线性时间不变性（LTI）**：参数 (Δ, A, B, C) 对所有时间步固定。这意味着状态更新规则不随输入内容改变——模型无法"选择性"关注或忽略特定 token。

Mamba 的选择机制（S6）将 B, C, Δ 变为输入 x 的函数：
```
B_t = s_B(x_t)      # 输入 → 输入投影
C_t = s_C(x_t)      # 输入 → 输出投影
Δ_t = τ_Δ(Δ + s_Δ(x_t))  # 输入依赖的步长
```

核心差异：
| 特性 | S4 (LTI) | S6 (Selective) |
|------|---------|---------------|
| 参数 | 时间不变 | 时间变化（输入依赖） |
| 计算模式 | 卷积 OR 循环 | 仅循环（需 scan） |
| 选择性 | 无 | 有（过滤/保留） |
| 内容感知 | 否 | 是 |

### 硬件感知并行 Scan

选择机制消除了卷积等价性——模型必须是时间变化的，无法用卷积并行计算。Mamba 通过**并行关联扫描（parallel associative scan / Blelloch scan）**解决：

1. 将状态更新展开为前缀和操作
2. 在 GPU SRAM 中做 kernel fusion，避免将扩展状态写入 HBM
3. 输入在 HBM → 加载到 SRAM → scan + 离散化 → 写回 HBM

结果：比所有基于卷积的 SSM 快 3×（A100 GPU）。

### Mamba 架构

```
Input → Mamba Block → ... (×L) → Output

Mamba Block:
  x → LayerNorm → [Linear(expand) → Conv1d → SiLU → SSM(S6)] → LayerNorm → Linear → + (residual)
```

关键设计：
- **无注意力、无 MLP**：用选择性 SSM 取代二者
- **扩展因子 E=2**：Linear 将 d_model 扩展到 2× 再投影回
- **残差连接 + SiLU 激活**
- **H3 简化**：将 H3 的两个门控 SSM 融合为单一选择性 SSM

## 实验结果

- **合成任务**：Selective Copying 和 Induction Heads → Mamba 可以泛化到 >1M token 序列
- **语言建模**：Mamba-3B 在 pretraining perplexity 和 0-shot 评估上超过 Pythia-3B，匹敌 Pythia-7B；5× 推理吞吐
- **音频**：在 SC09 语音生成上将 FID 降低一半以上
- **基因组学**：在 DNA 建模上超过 HyenaDNA 和 Transformer

## 关键概念

- [[selective-state-space]] — S6 选择机制，输入依赖的 SSM 参数化
- [[hardware-aware-algorithm]] — GPU 层次优化的并行 scan
- [[structured-state-space-models]] — S4 前身，HiPPO 矩阵 + 对角结构
- [[selective-copy]] — 需要内容感知的选择性复制任务
- [[induction-heads]] — 解释 LLM in-context learning 能力的机制
- [[hippo]] — SSM 的数学基础（High-order Polynomial Projection Operators）
- [[content-based-reasoning]] — Mamba 识别并解决的核心弱点

## 参考

- 代码：https://github.com/state-spaces/mamba
- S4 (Gu et al., 2022)
- H3 (Dao et al., 2023)
- 选择复制任务 (Arjovsky et al., 2016)
- 归纳头 (Olsson et al., 2022)
