Transformers are SSMs
Transformers are SSMs 这篇论文的核心贡献，是揭示了一个看似矛盾的数学事实：
结构化半可分矩阵（Structured Semiseparable Matrices）
SSM 的递归产生的矩阵——$K = (C, CA, CA^2, ..., CA^{L-1})$——恰好是下三角半可分矩阵。这类矩阵同时具有两种计算路径：
- SSM 递归路径：$y = ssm(x)$，O(L) 内存，O(L) 计算
- Attention 分解路径：$y = Sx$，O(L²) 内存，O(L²) 计算
两种路径计算的是完全相同的矩阵乘法，只是利用了不同的数学性质。
这意味着 SSM 和 attention 是同一类矩阵的不同计算视角——SSM 利用半可分结构的高效递归计算，attention 则利用矩阵分解的标准二次路径。Mamba-1 的硬件感知并行扫描，本质上是在 SRAM 中计算 SSM 路径，避免 HBM 的 IO 开销。
1.1 从 SSM 矩阵到半可分矩阵
给定 LTI SSM：$h_{t+1} = Ah_t + Bx_t$，$y_t = Ch_t$，展开后得到输出序列：
$y = (C, CA, CA^2, ..., CA^{L-1})x$
这是一个 $L \times L$ 的下三角矩阵乘法。每个列块（对应每个时间步的输入如何影响所有后续输出）由 $CA^k$ 给出，其中 A 是 N×N 对角矩阵（N 是 state dimension）。
关键性质：对于对角 A，这个矩阵恰好是下三角半可分矩阵——每个子矩阵的秩不超过 N。这使得：
- 可以用 SSM 递归路径高效计算（O(L)）
- 可以用 attention 矩阵分解路径精确计算（O(L²)）
1.2 SSD 算法
SSD 算法利用这一对偶性，在不同场景下选择最高效的计算路径：
- 短序列：attention 路径（并行计算优势）
- 长序列：SSM 路径（线性递归优势）
- 混合架构：Mamba-2 的 GVA head 支持两者混合
1.3 与标准 SSM 的关键区别
| 维度 | Mamba-1 SSM | SSD / Mamba-2 |
|---|---|---|
| 理论基础 | 硬件感知并行扫描 | 半可分矩阵对偶性 |
| 并行度 | 序列并行（有限） | Tensor Parallelism |
| State size | 固定小（N=16） | 可达 8×（N=128） |
| 训练速度 | 3× vs S4 | 2-8× vs Mamba-1 |
| 推理 | O(L) 递归，常数 cache | 同左 |
2.1 Grouped Value Attention（GVA）Head
Mamba-2 最重要的架构变化是引入Grouped Value Attention (GVA) head——类比 Transformer 的 multi-head attention：
Transformer：Q, K, V 三个投影，各自 head 数 × head dimension = d_model
Mamba-2 GVA：输入 x → Δ, B, C 三个投影，其中 B, C 输出矩阵（维度 N），然后 group 为 g 组
GVA 的关键设计：把 state dimension N 分成 g 组，每组独立计算后拼接。这使得：
- 并行度更高（组间可以并行）
- Tensor Parallelism 成为可能（每组独立，可跨 GPU 分片）
- 表达能力更强（更多独立的状态向量）
2.2 参数投影的位置变化
Mamba-1：SSM 参数（Δ, B, C）在 block 内部按时间步顺序计算（串行瓶颈）。
Mamba-2：所有投影移到 block 开头并行执行，大幅提升训练效率。
2.3 与 Transformer 的完整对应关系
| Transformer 组件 | Mamba-2 对应组件 | 关键区别 |
|---|---|---|
| Q, K, V 投影 | Δ, B, C 投影 | B, C 是输入依赖的（选择性） |
| Attention score | SSM 递归 | O(L) vs O(L²)，状态压缩 |
| Multi-head attention | GVA head | 每组独立 state，group 拼接 |
| FFN | 主分支 Linear+SiLU | 结构相近 |
| Flash Attention | 硬件感知并行扫描 | 都利用硬件特性优化 IO |
2.4 架构图
3.1 为什么 Mamba-1 不支持 Tensor Parallelism
Mamba-1 的递归形式天然是序列并行的，但无法跨 GPU 分片计算——因为 SSM 参数（Δ, B, C）是时间步依赖的，必须按顺序计算。跨 GPU 分片需要同步中间状态，开销抵消了并行的收益。
3.2 Mamba-2 如何支持 Tensor Parallelism
GVA head 的关键性质：每组的输出可以独立计算。这意味着可以把不同的 GVA head 组分配到不同的 GPU 上，每组独立计算后 All-Gather 聚合。
类比 Megatron 的设计理念：Transformer 的多头机制使得跨 GPU 分片成为可能；Mamba-2 的 GVA head 提供了相同的并行基础。
3.3 实际收益
Mamba-2 在大模型（2.7B+ 参数）场景下支持多 GPU 并行训练，使得模型规模可以继续扩展，而不受单卡显存限制。
4.1 速度与显存基准
Mamba-2 在不同序列长度下的训练速度对比 Mamba-1：
- 短序列（≤ 512）：Mamba-2 接近 Mamba-1
- 中等序列（512-2048）：Mamba-2 快 2-4×
- 长序列（≥ 4096）：Mamba-2 快 4-8×
4.2 State Size 的影响
更大的 state dimension N 带来了更强的表达能力——模型可以在状态向量中存储更多信息：
- N=16（Mamba-1 等效）：标准 SSM 规模
- N=64：2-4× 表达能力提升
- N=128：支持 8× 更大的 state size，Mamba-2 独有
4.3 语言建模基准
Mamba-2（2.7B 参数）在标准语言建模 benchmark 上与 Transformer 持平，推理速度 5× 于 Transformer。
4.4 激活值分析
SSD 的理论意义
Transformers are SSMs 不只是一篇工程论文，更是一篇理论论文。它揭示了：
- SSM 和 attention 是同一类数学对象的不同计算视角
- 半可分矩阵结构是两者统一的理论基础
- 这种统一性为混合架构设计开辟了新路径
对视觉 Mamba 的影响
视觉 Mamba 的后续发展（Vim、VMamba 等）都受益于 Mamba-2 的理论框架。SSD 对偶性让视觉任务可以在 SSM 路径和 attention 路径之间灵活切换，针对不同任务类型选择最优计算路径。