Mamba / State Space Models

Mamba 是一类基于状态空间模型（State Space Model, SSM）的序列建模架构，由 CMU 和 Princeton 的研究者提出。它以 O(N) 线性复杂度 挑战 Transformer 的 O(N²) 注意力计算，同时通过选择性注意力（Selective Attention）机制保持了对长上下文的有效建模能力。被视为 Transformer 的潜在竞争者和补充方案。

Definition

状态空间模型是控制理论中的经典概念，描述系统状态随时间的演化：

连续形式: h'(t) = Ah(t) + Bx(t) y(t) = Ch(t) + Dx(t)

其中：

• h(t): 隐藏状态
• x(t): 输入
• y(t): 输出
• A, B, C, D: 系统矩阵

离散化后的形式：

h_t = Āh_{t-1} + B̃x_ty_t = Ch_t + Dx_t

与 RNN 的关系

从形式上看，SSM 类似于 RNN：

• 均为递归状态更新
• 均为 O(1) 的状态大小（与序列长度无关）

但关键区别在于：

• RNN 的状态转移是非线性的（如 tanh）
• SSM 的状态转移是线性的，可以并行化计算

核心机制

线性复杂度的实现

Transformer 的注意力计算是 O(N²)，因为每个 token 都需要与所有其他 token 交互。

Mamba 的线性复杂度来源：

1. 卷积表示: SSM 可以表示为一个卷积核，通过 FFT 在 O(N log N) 时间计算
2. 递推形式: 也可以通过递推方式在 O(N) 时间计算
3. 状态大小固定: 无论序列多长，隐藏状态的维度不变

选择性注意力（Selective Mechanism）

早期 SSM（如 S4, H3）的问题是对输入不敏感——所有输入都被同样处理。Mamba 的核心创新是让 B 和 C 矩阵依赖于输入 x：

Δ_t = σ(Linear(x_t)) → 时间步长 B_t = Linear(x_t) → 输入特异的 B C_t = Linear(x_t) → 输入特异的 C

这意味着：

• 模型可以选择性地关注重要的输入信息
• 忘记不重要的信息
• 这种选择性是 Transformer 注意力的核心优势

硬件感知的实现

Mamba 的实现优化了硬件利用率：

• 使用 Flash Attention 风格的内存访问模式
• 避免在 HBM 和 SRAM 之间频繁传输
• 通过重算（recomputation）减少内存占用

版本演进

S4 (2021)

Gu et al. 提出的结构化状态空间模型：

• 引入结构化矩阵 A（HiPPO 初始化）
• 支持长距离依赖建模
• 但缺乏输入选择性

H3 (2022)

将 SSM 与注意力结合：

• 使用两个 SSM 层模拟 Q 和 K
• 仍然缺乏输入选择性

Mamba (2023)

核心突破：

• 引入输入选择性机制
• 硬件感知的高效实现
• 在 DNA 建模、音频处理等任务上达到 Transformer 水平

Mamba-2 (2024)

进一步优化：

• 统一了 SSM 和注意力的理论框架
• 提出了 SSD (State Space Duality) 架架
• 更好的并行度和扩展性

与 Transformer 的对比

维度	Transformer	Mamba
计算复杂度	O(N²)	O(N)
内存复杂度	O(N²)	O(N)
并行化	容易（独立注意力头）	较复杂
长上下文	需要拨揑或特殊设计	天然支持
训练稳定性	已验证	仍在探索
生态成熟度	极高	中等

各自的优势场景

场景	推荐架构	理由
短文本 (<4K)	Transformer	训练稳定，生态成熟
超长文本 (>100K)	Mamba	线性复杂度，内存效率
实时生成	Mamba	常数级状态，速度快
多模态	Transformer	注意力机制更灵活
代码生成	Transformer	局部模式匹配重要

实际应用与前景

已有应用

• DNA 建模: Mamba 在基因组序列建模上表现突出
• 音频处理: 长序列音频的高效建模
• 时间序列预测: 金融、物流等领域

混合架构趋势

目前的研究趋势不是用 Mamba 完全替代 Transformer，而是：

• Jamba (AI21): Transformer + Mamba 混合架构
• Zamba (Zyphra): 交替使用 Transformer 和 Mamba 层
• Griffin (Google): Gated Linear Recurrent + 局部注意力

Synthesis

Mamba 代表了对 Transformer 架构的重要反思——是否每个 token 都需要与所有其他 token 交互？如果答案是否定的，那么线性复杂度的架构就有显著优势。然而，Mamba 仍面临训练稳定性、生态成熟度和某些任务表现的挑战。最可能的未来是混合架构——在需要全局上下文的层使用 Transformer，在需要高效处理长序列的层使用 Mamba。

Related Concepts

• Transformer Architecture — Mamba 挑战的主流架构
• Attention Mechanism — Transformer 的核心计算单元
• Mamba / State Space Models — 线性注意力的相关研究
• 长上下文技术 — 长上下文建模的关联领域
• Model Merging — 模型融合技术

Sources

• Gu et al. (2021). "Efficiently Modeling Long Sequences with Structured State Spaces." ICLR.
• Dao et al. (2022). "Hungry Hungry Hippos: Towards Language Modeling with State Space Models."
• Gu & Dao (2023). "Mamba: Linear-Time Sequence Modeling with Selective State Spaces."
• Dao & Gu (2024). "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality."
• Jamba 技术报告 (AI21, 2024)