Mamba-2：把状态空间模型和 Attention 统一到一个数学框架里

一句话结论

Mamba-2 证明结构化状态空间模型（SSM）和注意力机制在数学上是对偶的——同一个序列混合算子既能写成 SSM 的递归形式（推理省显存），也能写成 attention 的矩阵乘形式（训练并行高效），于是 SSM 第一次能像 attention 一样吃满 GPU tensor core，状态维度还能大幅放大。

背景问题

Transformer 的注意力有一个根本代价：对序列长度平方级增长的显存和计算。虽然 Flash Attention 等把常数优化到极致，但 $O(L^2)$ 的本质没变，长上下文依然贵。

更隐蔽的问题在推理：自回归生成时 attention 要保留所有历史 token 的 K、V（KV cache），上下文越长显存越大、每步越慢。这是 MLA、GQA 等压缩方案的动机。

状态空间模型（SSM）走的是另一条路。它源自控制论，把序列建模成一个线性时不变系统：

$$ h_t = \bar{A},h_{t-1} + \bar{B},x_t,\qquad y_t = C,h_t $$

$h_t$ 是隐状态，$\bar{A},\bar{B},C$ 是参数（离散化后）。关键性质：

• 递归形式：生成时每步只依赖 $h_{t-1}$，常数显存，没有 KV cache 爆炸问题。
• 卷积形式：训练时可以展开成卷积，并行计算整个序列。

S4（2022）证明了精心设计的 SSM 在长程依赖任务（Long Range Arena）上能打过 Transformer，但 S4 在语言建模上一直打不过 Transformer——因为它有个致命限制：参数 $\bar{A},\bar{B}$ 是固定的，不随输入变化。

这意味着 SSM 无法「选择性」处理内容：它对所有 token 一视同仁地压缩进状态，既不能像 attention 那样「关注重要 token、忽略无关 token」，也无法做 content-based 检索。这是它在语言上吃亏的根本原因。

核心思路

Mamba 的突破分两步。

第一步：Mamba-1 让 SSM 「选择性」

Mamba-1（Gu & Dao, 2023）的洞察是：让 $\bar{B}$、$C$、以及 $\bar{A}$ 的一个控制衰减的参数 $\Delta$ 都依赖输入 $x_t$。

$$ \Delta_t = \mathrm{softplus}(\text{Linear}(x_t)),\quad B_t = \text{Linear}(x_t),\quad C_t = \text{Linear}(x_t) $$

$\bar{A}$ 通过 $\Delta_t$ 离散化得到，于是 $\bar{A}_t$ 也随输入变化。

「选择性」就这么来的：模型可以学到「看到重要 token 时让 $\Delta_t$ 大，状态快速更新；看到无关 token 时让 $\Delta_t$ 小，状态保持」。这就模拟了 attention 的「该记的记、该忘的忘」，但仍然是常数显存的递归。

代价是：参数依赖输入后，卷积形式不再直接成立（卷积要求参数固定）。Mamba-1 转而用**并行扫描（parallel scan）**算法做训练时并行化，硬件效率比卷积差，但能并行。

第二步：Mamba-2 找到 SSM 与 Attention 的对偶

Mamba-1 的问题：并行扫描不够「GPU 友好」，无法像 attention 那样用矩阵乘吃满 tensor core，训练吞吐偏低；状态维度 $N$（即 $h$ 的维度）也受限于扫描算法，做不大。

Mamba-2 的核心发现：结构化 SSM 可以写成「一个结构化的 attention 矩阵」。

把递归展开：

$$ y = \begin{bmatrix} C_1 \ C_2 \bar{A}_2 \ \vdots \ C_L \bar{A}_L \cdots \bar{A}_2 \end{bmatrix} \begin{bmatrix} \bar{B}_1 \ \bar{B}_2 \ \vdots \ \bar{B}_L \end{bmatrix} \cdot x ;\triangleq; M \cdot x $$

输出 $y$ 可以写成 $y = (L \circ M),x$，其中 $M$ 是一个 $L \times L$ 的矩阵，$L$ 是下三角掩码。这和 attention 的形式 $y = \mathrm{softmax}(QK^\top)V \cdot x$ 几乎一样——只不过 attention 的矩阵是 data-dependent 的全连接 softmax，SSM 的矩阵是 data-dependent 的、由 $\bar{A},\bar{B},C$ 决定的半可分（semiseparable）结构。

这就是论文标题「Transformers are SSMs」的由来：attention 是 SSM 的一个特例（结构无约束的 SSM），而 SSM 是「结构化约束下的 attention」。两者对偶。

对偶的好处立刻显现：

• SSM 可以用矩阵乘实现——既然它就是个结构化矩阵乘 $x$，就能用 attention 的 matmul kernel 跑，吃满 tensor core。
• 状态维度 $N$ 可以大幅放大：Mamba-1 的 $N$ 受限于扫描，典型 16～64；Mamba-2 用矩阵乘后 $N$ 可到 256 甚至更大，模型容量显著提升。
• 可以和 attention 混用：既然同属一个框架，SSM 层和 attention 层可以无缝拼接，催生了 Jamba 等混合架构。

关键技术点

1. 半可分矩阵：SSM 结构的代数刻画

Mamba-2 把 SSM 的混合矩阵刻画为「半可分矩阵（semiseparable matrix）」——它的任意子块都能被低秩分解表示。这个结构既足够丰富（能表达长程依赖），又足够约束（能用 $O(L)$ 显存表示，而非 $O(L^2)$）。这是 SSM 比 vanilla attention 省显存的代数根源。

2. 状态维度放大的连锁收益

Mamba-1 的 $N \approx 16$ 是性能天花板的一部分。Mamba-2 把 $N$ 推到 256 量级，模型每层的「记忆容量」显著扩大，在长上下文检索、长文档理解上质量提升明显。这也是 Mamba-2 能在语言建模质量上更接近同规模 Transformer 的关键。

3. 选择性保留 + 矩阵乘并行

Mamba-2 没有放弃 Mamba-1 的选择性（输入依赖参数），而是找到了「选择性 SSM」的矩阵乘表达。具体说，它把选择性参数化限制在「能让混合矩阵保持半可分结构」的范围内，既保留 content-aware 能力，又能用高效 kernel。这是工程与数学的精巧折中。

4. 混合架构的可行性

既然 SSM 和 attention 对偶，把它们混在一层或交替堆叠都很自然。AI21 的 Jamba、Falcon-Mamba 等模型探索了「大部分 SSM 层 + 少量 attention 层」的配方：SSM 负责长程、attention 负责精准检索。这种混合在长上下文质量与推理成本之间找到了新平衡点。

5. 推理红利不变

Mamba-2 仍然是 SSM，推理时仍是常数显存的递归——没有 KV cache 线性增长问题。在超长上下文生成场景（如长文档续写、长程 agent），这是相对 attention 模型的结构性优势。

效果与局限

效果：

• Mamba-2 在同参数量下，语言建模质量超过 Mamba-1，接近甚至偶有持平同规模 Transformer（取决于任务）。
• 训练吞吐显著高于 Mamba-1（论文报告约 2～8 倍，取决于状态维度），状态维度可放大一个数量级。
• 推理时长上下文成本远低于 Transformer：没有 KV cache 线性膨胀，超长生成场景吞吐优势明显。
• 已被多个开源模型采纳（Jamba、Falcon-Mamba 等），证明工程可用性。

局限：

• 在「精准 in-context 检索」上仍弱于 attention。SSM 把信息压进固定维度状态，对「从超长上下文里精确召回某个具体事实」这类任务，召回精度不如 attention 的显式 KV 匹配。这是 SSM 架构的固有张力。
• 混合架构的配方仍在探索。SSM 与 attention 的比例、放置方式没有定论，不同任务最优配方不同。
• 生态成熟度远不及 Transformer。attention 有 Flash Attention、PagedAttention、GQA 等一整套优化基建，SSM 的 kernel、推理框架、量化工具链都在早期。
• 大规模验证仍不足。目前 Mamba 系列的最大公开模型仍远小于顶级 Transformer 模型的规模，scaling law 在 SSM 上的形态还需更多数据。

谁该关心

• 做长上下文推理 infra 的人：SSM 的常数显存推理对超长生成场景是结构性红利，值得作为 attention 路线的补充评估。
• 架构研究者：SSM-Attention 对偶是近年来最漂亮的架构理论成果之一，无论是否用 Mamba，理解它对设计新混合架构都有启发。
• 追求推理成本极限的团队：在 agent、长文档等场景，SSM 的显存优势可能让单卡跑超长上下文成为可能。
• 关注 scaling law 的人：SSM 在更大规模上是否符合 Transformer 的 scaling 规律、是否有不同形态，是当前开放问题。

一个判断：Mamba 系列不会「取代」Transformer——attention 在精准检索上的优势和生态壁垒太深。但它已经从「有趣但打不过」变成了「在特定场景真的能赢」，混合架构很可能是下一代长上下文模型的主流形态之一。把它当工具箱里的一把新工具，而不是 Transformer 的替代品。