22 · Mamba 2 — Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (SSD)

SSM Attention Dao & Gu 2024
Tri Dao, Albert Gu · arXiv:2405.21060

Previous: Mamba 1 (S6) — parallel scan cannot use GPU tensor cores, making training slower than equal-FLOPs Transformers. State dimension N (16-64) is severely limited, and model parallelism is difficult due to the recurrent formulation.

This: Structured State Space Duality (SSD) — discovery that SSM = semiseparable matrix = structured masked attention. A is simplified from diagonal to scalar. A dual quadratic form Y = (L ⊙ QKᵀ) · V bridges SSMs and attention. Block decomposition hybrid algorithm combines the best of both worlds.

Effect: 2-8x training speedup via tensor core MatMul, 8x larger state dimension (N=64→256+), full support for Tensor Parallelism and Sequence Parallelism. Pareto-dominates both Transformer++ and Mamba on perplexity vs wall-clock time.

一、核心背景

Mamba 1 (Gu & Dao, 2023) 提出了选择性状态空间模型 (S6)，在长序列建模上取得了重大突破。其核心优势是 O(n) 训练复杂度和 O(1) 推理复杂度（无需 KV 缓存），在推理吞吐上比同规模 Transformer 快 5 倍。

然而 Mamba 1 有两个本质性缺陷，在实际应用中限制了它的竞争力：

IO 瓶颈而非计算瓶颈 — Mamba 1 采用 parallel scan（并行前缀和）实现 O(n) 训练，该算法的计算密度远低于矩阵乘法。在实际 GPU 上，parallel scan 的执行受限于 HBM 带宽（memory-bound），而非计算单元利用率（compute-bound）。这意味着 Mamba 1 无法利用 GPU 的 tensor core（张量核心），而现代 Transformer 的训练效率已经在很大程度上依赖于 tensor core 的密集矩阵乘法。
状态维度 N 严重受限 — Mamba 1 的 SSM 状态维度 N 被限制在 16-64 之间。这是因为 parallel scan 需要为每个时间步处理 N 维状态传递，大 N 会导致 scan 操作的 IO 消耗成倍增长。这种状态容量限制影响了模型记忆长程信息的能力。
模型并行困难 — Mamba 1 的 SSM 层在设计上是一个逐时间步的递归，无法像 Transformer 的注意力机制那样自然地沿序列维度或头维度切分。这使得分布式训练（Tensor Parallelism / Sequence Parallelism）在 Mamba 架构上难以实现。

关键数字：对于 2.8B 参数的模型，Mamba 1 在训练时的 FLOPs 效率低于 Transformer，实际 wall-clock 速度比同 FLOPs 的 Transformer 慢。这是 Mamba 1 未能被大规模采用的根本原因。

在上述背景下，Mamba 2 的目标非常明确：保留 SSM 的线性复杂度优势，同时消除 parallel scan 带来的硬件效率瓶颈，让 SSM 训练能够像 Transformer 那样利用 tensor core 进行高效矩阵乘法。

二、核心洞察：SSM = 结构化矩阵

2.1 矩阵变换形式

论文的核心数学发现始于一个基础观察：任何线性时变 SSM 都可以写成一个矩阵变换。从 SSM 的递推定义出发：

h_t = A_t h_t-1 + B_t x_t

y_t = C_t^T h_t

展开递归，得到每个输出对输入的依赖：

y_t = Σ_s=0^t C_t^T · (A_t · ... · A_s+1) · B_s · x_s

将其写为矩阵形式，定义：

M_j,i = C_j^T · A_j · A_j-1 · ... · A_i+1 · B_i (for j ≥ i)

y = M · x ; M ∈ ℝ^T×T

这个看似简单的改写至关重要。它说明 SSM 的输入-输出映射等价于一个下三角矩阵 M 与输入向量 x 的乘积，其中 M 的每个元素由 A、B、C 参数序列的累积乘积定义。

2.2 N-半可分矩阵（N-Semiseparable Matrix）

上述矩阵 M 在数值线性代数中有一个精确的分类：N-半可分矩阵（N-semiseparable matrix）。其严格定义是：

定义：一个下三角矩阵 M ∈ ℝ^T×T 是 N-半可分的，当且仅当它的任何子矩阵（由连续行和连续列的交集构成）的秩都 ≤ N。直观理解，M 的"偏离对角线的信息"被一个秩-N 的结构所约束。

关键性质是：一个 N-半可分矩阵可以用 O(TN) 个参数完整表示（由 A₁...A_T, B₁...B_T, C₁...C_T 决定），而非 O(T²)。这正是 SSM 效率的数学根源——通过半可分结构隐式编码了 ℝ^T×T 的信息，却不需要显式存储或操作整个矩阵。

SSM 状态大小 N ⇔ 半可分矩阵的边界秩 (rank bound)

这一等价关系建立了 SSM 和结构化矩阵之间的桥梁：

状态维度 N 直接对应于半可分矩阵的秩边界
更大的 N 意味着更多参数来刻画 M，也就意味着更丰富的序列变换能力
Mamba 1 的 N=64 意味着每个下三角子矩阵的秩被限制在 64 以内

洞察：N-半可分矩阵不仅仅是 SSM 的数学等价描述，它还揭示了 SSM 和注意力之间的深层联系。注意力机制的 softmax(QK^T) 也是一个下三角矩阵（因果掩码），但它的每个元素可以独立地取任意值——没有任何秩约束。从这个视角看，SSM 可以理解为一个带有秩约束的结构化注意力。

三、SSD 简化：结构化状态空间对偶

3.1 标量 A 简化

Mamba 1 使用对角线 A 矩阵（N 个可学习参数），每个状态维度有自己的递推衰减率。Mamba 2 将其简化为标量 A：

A_t = a_t · I (标量 × 单位矩阵)

h_t = a_t · h_t-1 + Δ_t · B_t · x_t

y_t = C_t^T · h_t

这一简化虽然在状态维度的表达力上有所牺牲（所有 hidden state 共享同一个衰减率），但带来了计算和推理上的巨大收益：

矩阵乘法兼容 — 标量 A 使得 SSM 递推公式中的状态更新可以拆解为纯矩阵乘法形式
衰减乘积可分解 — a_t 是一个标量，因此 a_t × a_t-1 × ... × a_s 只是一个标量乘积，比矩阵乘积简单得多
状态容量补偿 — 标量 A 的牺牲通过更大的 N（64 → 256+）来弥补，实际表达力反而更强

3.2 对偶（二次）形式

标量 A 简化后，SSM 的矩阵变换形式出现了惊人的结构：它恰好等价于一个结构化的掩码点积注意力。

定义以下变量：

Q_t = C_t — Query 来自 SSM 的 C 参数
K_t = B_t — Key 来自 SSM 的 B 参数
V_t = Δ_t · x_t — Value 来自输入与步长的乘积
L_j,i = a_j × a_j-1 × ... × a_i+1 — 衰减掩码（下三角）

Y = (L ⊙ QK^T) · V

公式 1: SSD Dual Form — 结构化状态空间对偶

展开写：

y_t = Σ_i=1^t (L_t,i · Q_t^T K_i) · V_i

与标准缩放点积注意力对比：

y_t^Attn = Σ_i=1^t softmax(Q_t^T K_i) · V_i

两者惊人地相似，仅有两点区别：

注意力使用 softmax 归一化，将 QK^T 转为概率分布；SSD 使用衰减掩码 L，这是一个随距离指数衰减的确定性权重
注意力的 Q 和 K 来自不同的投影；SSD 的 Q 来自 C、K 来自 B，但都来源于相同的输入 x

核心洞见：SSD 揭示了 SSM 和注意力之间的精确关系。SSM 本质上是将softmax 替换为指数衰减掩码的线性注意力变体。这意味着"SSM vs 注意力"并不是两个截然不同的范式，而是一个连续谱的两端——区别仅仅在于如何构造查询和键之间的相似度矩阵。这也直接呼应了论文"Transformers are SSMs"的标题。

四、SSD 算法：块分解混合计算

4.1 块分解（Block Decomposition）

SSD 的对偶形式本身给出了直接的 O(n²) 二次计算方法（通过矩阵乘法），但这在长序列下不现实。SSD 的核心工程贡献是提出了一种块分解算法，将半可分矩阵 M 沿着主对角线切分为等宽的块：

M = M_intra + M_inter

块内（Intra-block） — 位于同一块内部的元素，使用 SSD 的二次形式（结构化掩码注意力），通过矩阵乘法在 tensor core 上高效计算
块间（Inter-block） — 跨越不同块的元素，使用 SSD 的线性递推形式（类似 Mamba 1 的 scan），但块数量远小于序列长度 T，因此 scan 开销很小

算法直觉：将序列分成 B 个固定大小的块（例如块大小 64-256）。每个块内部执行小规模的注意力计算（O(b²), b = 块大小），块之间通过 SSM 状态传递（O(B)）。当 B << T 时，总复杂度为 O(T·b + T·N)，远优于纯二次形式的 O(T²)。

4.2 计算复杂度对比

算法	训练复杂度	硬件利用率	适用长度
Mamba 1 (Parallel Scan)	O(TN)	低（无 tensor core）	任意
SSD 二次形式	O(T²)	高（tensor core MatMul）	短（< 2K）
SSD 块分解混合	O(T·b + T·N)	高（tensor core MatMul）	任意（最佳）
FlashAttention-2	O(T²)	高（tensor core MatMul）	中（< 8K）

4.3 性能数据

论文报告了 SSD 算法在 A100-80GB GPU 上的实际性能：

与 Mamba 1 对比 — SSD 比 Mamba 1 的 parallel scan 实现快 2-8 倍，序列越长加速比越大
与 FlashAttention-2 对比 — 在序列长度 2K 处两者性能持平；短序列（<2K）SSD 因线性复杂度更具优势；长序列（16K）SSD 达到 6 倍加速
使用 tensor core — SSD 的计算主要由矩阵乘法（MatMul）而非 scan 操作驱动，因此在 A100/H100 上可以充分利用 tensor core 的峰值 FLOPs

五、架构改进

5.1 多头 Mamba 架构

Mamba 1 本质上是单头的——每个 SSM 层处理一个单一的隐藏状态序列。Mamba 2 引入了类似 Transformer 的多头架构：

多头注意力类比（MHA vs MVA） — SSD 的每个"头"有一组独立的 Q、K、V 投影和独立的 SSM 状态（半可分矩阵），允许模型在不同的子空间中捕获不同类型的序列模式
分组值注意力（GVA） — 类似于 GQA（分组查询注意力），Mamba 2 的 SSD 层可以配置多个 SSM 头共享 V 投影，在计算效率和模型容量之间取得折中

5.2 数据依赖投影并行化

Mamba 1 中，B、C、Δ 参数在每个时间步根据当前输入计算，这限制了计算的并行度。Mamba 2 将所有数据依赖的投影集中到块的开始处：

一次性计算：Q, K, V, a = projections(x) (全部矩阵乘法)

这意味着整个序列的 Q、K、V 可以通过一次大型矩阵乘法计算得到，而非逐时间步的小型计算。这种设计：

最大化 tensor core 利用率（大矩阵乘法效率远高于多个小矩阵）
为 Tensor Parallelism 奠定基础（沿头维度切割投影矩阵）
为 Sequence Parallelism 奠定基础（沿序列维度切割后分别计算再同步）

5.3 分布式训练支持

并行策略	Mamba 1	Mamba 2 (SSD)
Tensor Parallelism (TP)	困难（递归状态无法切割）	支持（Megatron 风格，沿头维度切分 SSM 头）
Sequence Parallelism (SP)	不支持	支持（块间 SSM 状态传递可用 all-reduce 同步）
Data Parallelism (DP)	支持	支持
Pipeline Parallelism (PP)	支持	支持

5.4 变长训练

Mamba 1 需要将同批次内的序列填充至相同长度（padding），因为 parallel scan 要求所有序列具有相同的计算图。Mamba 2 的 SSD 算法天然支持变长训练——块分解的灵活性允许不同序列使用不同的块数量，无需填充。这在数据混合训练（不同来源的序列长度差异大）场景中节省了显著的计算量。

六、实验结果

6.1 主实验结果

Mamba 2 在 2.7B 参数规模下（Mamba2-2.7B）进行了基准测试，与同规模和更大规模的模型对比：

Mamba2-2.7B vs Mamba-2.8B — 在几乎相同的参数规模下，Mamba 2 在所有评估任务上优于 Mamba 1，验证了 SSD 架构的有效性和改进的价值
Mamba2-2.7B vs Pythia-2.8B — 显著优于同规模的开源 Transformer 基线
Mamba2-2.7B vs Pythia-6.9B — 在多个基准上以 2.7B 参数超越 6.9B 参数的 Transformer 模型，展示了 SSM 架构的参数效率
Mamba2-2.7B vs Transformer++ — 在困惑度（perplexity）与 wall-clock 时间的 Pareto 曲线上，Mamba 2 明显优于 Transformer++ 和 Mamba 1，即在相同训练时间内取得更低的困惑度

6.2 详细对比表

方面	Mamba 1 (S6)	Mamba 2 (SSD)
A 矩阵结构	对角矩阵（N 参数）	标量 × 单位矩阵（1 参数）
状态维度 N	16-64	64-256+
训练速度	基准（parallel scan）	2-8x 加速（tensor core MatMul）
硬件利用	Scan 操作（无 tensor core）	矩阵乘法（tensor core 高效）
模型并行	困难（递归依赖）	支持 Tensor / Sequence Parallelism
变长训练	不支持（需 padding）	支持（块分解灵活）
对偶形式	无	Y = (L ⊙ QK^T) · V
与 Attention 关系	替代方案	结构化变体（统一理论）

6.3 Pareto 主导

论文最有力的实验结果来自困惑度 vs 训练时间的 Pareto 曲线。在相同计算预算下，Mamba 2 始终取得更低的验证困惑度；在相同的目标困惑度下，Mamba 2 需要的训练时间显著少于 Transformer++ 和 Mamba 1。这种Pareto 主导（即同时在两个指标上优于对手）是架构竞争的黄金标准结果。

七、上下游关联

上一篇：Mamba 1 (S6) (Gu & Dao, 2023) — 提出选择性 SSM 解决时不变 SSM 的内容推理缺陷，但受限于 parallel scan 的硬件效率。Mamba 2 在此基础上解决了工程效率瓶颈。

并行工作：Linear Attention / Katathopoulos Attention — SSD 的对偶形式与线性注意力的"kernel trick"有着深刻的联系。Linear Attention 将 softmax 展开为 φ(Q)φ(K)^T 以获得线性复杂度，SSD 则用指数衰减掩码替代 softmax。两者都通向类似的"线性注意力 = 状态空间模型"结论。

理论前驱：S4 (Gu et al., 2022) — 首个将 SSM 引入深度学习序列建模的工作，提出 HiPPO 初始化。Mamba 2 引用的半可分矩阵理论可以追溯到 S4 的线性代数基础。

下一篇：Mamba 3 (2025) — 进一步扩展到复数状态空间和 MIMO（多输入多输出）SSM，在下游任务上追平了当时的最强 Transformer 模型。Mamba 2 的 SSD 框架为 Mamba 3 提供了数学基础。

混合架构应用：SSD 框架直接促成了Jamba (AI21 Labs)、Samba (Microsoft) 等混合架构——在 Transformer 中交替插入 SSM 层，利用 SSD 的理论统一性实现平滑融合。这些混合模型在多个基准上同时优于纯 Transformer 和纯 SSM 模型。

实际采用：Mamba 2 的 SSD 算法已被集成到多个大模型训练框架中（如 NVIDIA Megatron-LM、Hugging Face Transformers），其 Tensor Parallelism 支持使得 SSM 层可以无缝嵌入现有的大规模训练 pipeline。

八、个人思考

"Transformers are SSMs"是一个精心设计的数学激进主张。标题看起来像是用 SSM 吞噬 Transformer 的挑衅宣言，但论文实际上证明了 SSM 和注意力在数学上共享同一个底层结构——结构化矩阵乘法。这不是一个夸大的断言，而是一个精确的数学等价关系。这种"挑衅性标题 + 精密数学证明"的手法令人印象深刻。

标量 A 的牺牲非常聪明。将 A 从对角矩阵简化为标量，表面上是降低了 SSM 的容量（所有状态维度共享同一个衰减率）。但 SSD 通过允许更大的状态维度 N（64 → 256+）实现了远超 Mamba 1 的总容量。这是一个典型的"trade-off"教科书案例——通过重构计算图来突破硬件瓶颈，表达力的微小损失被规模扩展十倍以上所淹没。

技术和理论的完美配速。Mamba 2 不仅是工程优化（tensor core 加速），还有深刻的理论贡献（SSM-Attention 对偶）。两个作者（Tri Dao, Albert Gu）刚好分别代表了"工程大师"和"理论学者"的角色。Tri Dao 带来了 FlashAttention 中积累的 IO 感知算法经验，Albert Gu 带来了 S4/Mamba 的 SSM 理论深度。这种跨领域协作产生的结果远大于各自领域的简单叠加。

统一框架的战略意义。SSD 的最大贡献可能不是具体的加速比，而是它建立的统一框架。在此之后，SSM 和注意力不再是两个独立的家庭，而是同一个连续谱上的不同点。这意味着过去二十年所有关于注意力机制的理论成果、工程优化（FlashAttention 系列、PageAttention 等）、分布式策略，都可以系统地迁移到 SSM 研究中。反过来，SSM 中的状态空间思想也可能启发注意力机制的改进。这种"大一统"视角才是论文标题的真正含义。

效率与容量权衡的精妙之处。SSD 揭示了一个基本问题：序列建模的效率瓶颈到底在哪里？Mamba 1 认为瓶颈在复杂度（O(n²) vs O(n)），但实践表明硬件利用率（tensor core vs scan）才是实际瓶颈。Mamba 2 证明了：当计算从 scan 换到 MatMul 后，在中等序列长度下，二次形式 + tensor core 比线性形式 + scan 更快。这提醒我们，理论复杂度的降低并不等同于实际性能的提升——硬件特性有时比算法复杂度更重要。

下一步：SSD 将 SSM 推到了可以与 Transformer 同台竞技的位置，但关注点已经转移到混合架构（Hybrid SSM-Attention）。未来的问题可能不是"二选一"，而是"如何最优地混搭"。