LLM2023· CLASSIC

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Gu, Dao

选择性状态空间模型——O(N) 复杂度，媲美 Transformer 的语言建模能力。可能是 post-Transformer 时代的候选之一。

arXiv:2312.00752

#architecture#ssm#long-context#efficiency

核心贡献

01Selective State Space Model (SSM)：递归形式天然 O(N) 复杂度
02Input-dependent 参数：每个 token 动态决定状态转移——克服固定 SSM 的局限
03硬件感知的并行扫描算法——比朴素 SSM 快 20-40×
045.5K token 序列上 Mamba 3B 击败同规模 Transformer
05开启 "Attention-free" 架构的再次繁荣（RWKV、RetNet 等）

从 State Space Models 说起

经典线性 SSM：

h_t = A h_{t-1} + B x_t \\ y_t = C h_t

A、B、C 是固定矩阵——状态转移与输入无关。这导致 SSM 在选择性任务（如 "copy 某个特定位置的 token"）上不如 Attention。

`Mamba` 的关键创新：Input-dependent 参数

B_t = \text{Linear}_B(x_t), \quad C_t = \text{Linear}_C(x_t), \quad \Delta_t = \text{Linear}_\Delta(x_t)

让 B、C、Δ 都依赖当前输入——模型可以选择性地记住或遗忘。这打通了 SSM 和 Attention 的核心差距。

硬件效率

不变的 SSM 可以用 FFT 并行化。Mamba 的 B_t、C_t 随输入变化，不能直接 FFT。论文设计了硬件感知的并行扫描（parallel scan）：

在 SRAM 里做 associative scan
避免 HBM 反复读写
达到近 Transformer 的训练吞吐

复杂度对比

| 架构 | 训练复杂度 | 推理复杂度（per token） | 外推 | |---|---|---|---| | Transformer | O(N²d) | O(Nd)（KV Cache） | 差（固定位置） | | Mamba | O(Nd) | O(d)（固定状态大小） | 好 | | RetNet | O(Nd) | O(d) | 中等 |

实测结果

3B Mamba 在 Pile 上 perplexity 匹敌 3B Transformer++（SOTA Transformer 变种）
但仍未在前沿 LLM（7B+）上证明超过 Transformer——Jamba（Mamba + Attention 混合）更有竞争力

局限

In-context learning 能力不如 Transformer（无显式 attention，难做精确 recall）
长上下文推理暗示有衰减（状态大小固定）
小规模研究为主，10B+ 规模工业化案例少

面试视角

面试考点

"Mamba 和 Transformer 的本质区别？" Transformer 的 Attention 是全局的——每个 token 能看到所有历史 token；Mamba 是递归的——所有历史信息压缩到固定大小的隐状态。

"O(N²) vs O(N) 的实际影响？" 短序列（N<2K）差别不大，内存带宽主导。长序列（N>16K）Mamba 快得多。但训练上 Transformer 仍然有 FlashAttention 加持，工程成熟度高。

"Mamba 的 selective 是什么意思？" Input-dependent 的状态转移参数让模型可以"选"记什么——不像固定 SSM 那样对所有输入用同一组 A、B、C。这是 Mamba 超过经典 SSM 的关键。

"Mamba 在推理时的优势？" 状态大小固定——不管序列多长，每生成一个 token 的计算量都一样。Transformer 的 Decoding 要读所有 KV Cache，长序列越来越慢。

"为什么 Mamba 没完全取代 Transformer？"

In-context learning 能力稍弱——固定状态无法精确 retrieve 某个具体 token 的信息
工业 infra 仍围绕 Transformer 优化
Hybrid 架构（Mamba + Attention）可能是务实方向

"Jamba / Zamba 等 Hybrid 架构？" 混用 Mamba 层和 Attention 层——拿 Mamba 的效率 + Attention 的精确 recall。当前研究热点。

"Mamba 是 post-Transformer 的候选吗？" 可能是，但也可能不是。类似问题历史上每隔几年被提出（RNN→LSTM→Transformer），最终胜出的常常是 hybrid 方案或仍是 Transformer + 优化。