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。

核心贡献

选择机制（Selection Mechanism / S6）：将 SSM 参数（Δ, B, C）变为输入依赖，从时间不变（LTI）升级为时间变化
硬件感知算法：通过并行关联扫描（parallel associative scan）在 SRAM 中计算，避免 GPU HBM 之间的 IO 瓶颈
极简架构 Mamba：将 H3 架构中的 SSM 层与 MLP 门控融合为单一同质块
选择复制（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）**解决：

将状态更新展开为前缀和操作
在 GPU SRAM 中做 kernel fusion，避免将扩展状态写入 HBM
输入在 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)

4.6 KiB Raw Blame History Unescape Escape