SSM 家族全景图（下）：Mamba-2 与 Jamba
上一篇分析了 S5、H3 和 RWKV，它们分别代表了“简化数学”“补齐局部记忆”“极简线性 RNN”三条路线。本篇分析家族里理论意义最大的 Mamba-2，以及把 SSM、Attention 与 MoE 结合起来的工业级混合架构 Jamba。最后给出一张覆盖 7 个模型的全景小结。
一、Mamba-1 留下的两个理论短板
Mamba-1 通过选择性机制、并行扫描和硬件感知取得了巨大成功，但它身上仍有两个短板。
第一，并行扫描的计算效率不如矩阵乘法。 树形归约的并行深度是 $O(\log L)$，而 GPU 的矩阵乘法单元（Tensor Core）是为 $O(1)$ 深度的大矩阵乘法优化的。明明硬件最擅长矩阵乘法，Mamba-1 却只能用树形归约，这中间存在效率损失。
第二，状态维度 $N$ 受限。 并行扫描每一步都要搬运并合并整个状态，属于显存带宽受限的算子：状态越大，单位计算量背后的数据搬运越多。因此 Mamba-1 官方模型的状态维度只取很小的值（例如 16），单个状态能承载的信息量有限。
Mamba-1 作者 Tri Dao 与 Albert Gu 在 2024 年提出 Mamba-2，试图回答一个更深层的问题：选择性 SSM 和 Transformer 的 Attention，在数学底层到底是什么关系？
二、核心理论：状态空间对偶（SSD）
1. 核心定理
Mamba-2 的理论突破是一个数学定理：选择性状态空间模型（Mamba-1）的输出，等价于一个半可分矩阵（Semiseparable Matrix）与输入向量的乘积。
具体来说，Mamba-1 的递推
$$h_t = \bar{A}t h{t-1} + B_t x_t, \quad y_t = C_t^\top h_t$$
可以被展开为 $y = M x$，其中 $M \in \mathbb{R}^{L \times L}$ 是一个下三角的半可分矩阵，其元素为
$$M_{ij} = \begin{cases} C_i^\top \left(\prod_{k=j+1}^{i} \bar{A}_k\right) B_j & i \geq j \ 0 & i < j \end{cases}$$
这个矩阵的形状很值得注意：它把整个序列的递推关系一次性写成了一个 $L \times L$ 的矩阵与向量的乘积。
这个框架比“SSM 等价于线性注意力”这句话还要更宽一些：SSM 与带掩码的注意力变体都是结构化半可分矩阵的不同分解形式，因此两种形式可以自由换算——从递推形式出发可以推出注意力式的二次型形式，从二次型形式出发也能重新得到递推形式。这层联系的实际用途在于算法可以互相搬运：注意力那边成熟的分块矩阵乘法可以拿来算 SSM，SSM 这边的递推也可以拿来算线性注意力。Mamba-2 的速度提升正是从这次搬运里来的。
2. “半可分”的含义
这个矩阵虽然大小是 $L \times L$，但它的每一个子块都可以被低秩分解。具体来说，$M$ 可以写成
$$M = L \odot (A B^\top)$$
其中 $L$ 是一个下三角掩码矩阵，$A$、$B$ 是 $L \times r$ 的矩阵（$r$ 是秩），$\odot$ 是逐元素相乘。
“半可分”这个名字描述的就是这个性质：矩阵整体很大，但内部结构高度规律，每个子块都能用低秩表示。
3. 与线性注意力的对偶
线性注意力 $Y = (QK^\top)V$ 也可以被写成类似的形式。这就是“SSM 与线性注意力对偶”的含义：两者本质上是同一个数学结构（半可分矩阵与向量的乘积）的两种不同写法，可以通过简单的变量替换相互转化。
这个结论把 SSM 与 Transformer 家族在理论上连接了起来：Mamba 与线性 Attention 处在同一个数学框架之内。
三、Mamba-2 的架构变化
基于 SSD 理论，Mamba-2 对架构做了四处调整。
1. 多头机制（Multi-Head）
Mamba-1 只有一个“头”，也就是一个大的状态向量。Mamba-2 把扩展后的特征分成 $H$ 个头，每个头独立地运行一个小的 SSD 计算。这与 Transformer 的 Multi-Head Attention 在结构上一一对应。
2. 状态维度增大
Mamba-1 的状态维度 $N$ 受扫描实现的限制，官方模型通常只取很小的值（例如 16）。原因在上一篇已经解释过：Mamba-1 的扫描实现用不到张量核心，状态越大，每一步要搬运的数据就越多，收益反而下降。Mamba-2 改用矩阵乘法作为主要计算之后，这个瓶颈被打开：每个头的状态维度为 $N = P / H$，总状态维度 $P$ 可以取到 128 到 256 这一量级。
3. 状态转移收缩为一个标量
Mamba-1 的 $\mathbf{A}$ 是每个通道一套的实数对角矩阵，对角上保留着 $N$ 个不同的衰减率；Mamba-2 把 $\mathbf{A}$ 收缩成“标量乘单位阵”，也就是同一个头内所有状态维度共享一个衰减标量，状态转移的计算因此变得更简单，个体差异由多头结构去补。
4. 计算方式从并行扫描变为分块矩阵乘法
这是最关键的工程变化。由于 SSD 把问题转化成了矩阵乘法
$$Y = M x = (L \odot (AB^\top)) x$$
这个计算可以通过分块矩阵乘法（Block Matrix Multiplication）来实现：先将序列分成若干个块（例如每块 256 个时间步），块内直接用矩阵乘法计算，充分利用 GPU 的 Tensor Core；块间用一个小的“状态传递”把各块连接起来。
效果很明显：计算被映射到 GPU 最擅长的高效矩阵乘法（GEMM）上；训练速度比 Mamba-1 快 2 到 8 倍；不再需要复杂的树形归约。
这套分块算法值得展开一点，因为它是速度提升的真正来源。SSD 的做法是把序列混合矩阵按 $Q \times Q$ 的块切开（参考实现里 $Q$ 默认取 64，实践中常在 64 到 256 之间），再利用半可分矩阵的性质，把非对角块分解为低秩形式。细看下来，它是四步：
- 计算块内输出：把每个块当成一段短的二次型注意力来计算。
- 计算块内最终状态：同样只用块内的输入。
- 在块之间传递状态：对各块的状态做一次递推扫描，此时序列长度已经缩短到原来的 $1/Q$。
- 用每块真实的初始状态修正输出：把块间状态对输出的贡献补回来。
四步之中有三步完全由矩阵乘法构成，可以充分并行；只有第三步需要扫描，而且作用在缩短后的序列上，占总时间的比例很小。用一句话概括就是：块内走二次型，块间走递推。
随之而来的是系统层面的收益。当矩阵乘法成为主要计算之后，原本只为 Transformer 开发的工程手段（张量并行、序列并行、变长序列处理）都能直接搬到 SSM 上，大规模训练不必再为 SSM 单独造一套并行方案。这样的简化换来的是更简洁的计算，以及更大的状态维度上限。
5. Mamba-2 Block 的数据流
Mamba-2 的一个 Block 内部流程如下：
- 输入 $\mathbf{x} \in \mathbb{R}^{B \times L \times D}$，经过线性层从 $D$ 维扩展到 $2P$ 维。
- 特征被分成两路：主路径形状为 $(B, L, P)$，门控路径形状也是 $(B, L, P)$。
- 主路径先经过一维因果卷积，再经过 SiLU，然后按头拆分成 $(B, L, H, N)$，执行 SSD 计算（也就是分块矩阵乘法）。
- 门控路径经过 SiLU。
- 两路逐元素相乘。
- 经过线性层从 $P$ 维投影回 $D$ 维。
- 加上残差，输出 $(B, L, D)$ 传给下一个 Block；归一化沿用 Pre-Norm，放在每个 Block 的入口，与 Mamba-1 一致。
整体结构与 Mamba-1 的 Block 高度相似，差别集中在核心计算模块：并行扫描换成了分块矩阵乘法，并引入了多头划分。
6. 与 Mamba-1 的对照
理论基础： Mamba-1 是选择性状态空间；Mamba-2 是 SSD（与线性注意力对偶）。
计算方式： Mamba-1 用并行扫描（树形归约）；Mamba-2 用分块矩阵乘法（GEMM）。
多头机制： Mamba-1 没有；Mamba-2 有 $H$ 个头。
状态维度： Mamba-1 常用很小的 $N$（例如 16），受扫描实现的限制；Mamba-2 的总状态维度 $P$ 可以取到 128 到 256 这一量级，明显更大。
状态转移形式： Mamba-1 是每个通道一套的实数对角矩阵；Mamba-2 收缩为每个头一个标量的单位阵。
训练速度： 以 Mamba-1 为基准，Mamba-2 快 2 到 8 倍。
四、Jamba：SSM 与 Transformer 的混合架构
1. 动机：纯 SSM 够用吗
2024 年，AI21 Labs 提出了 Jamba，这是工业界首个大规模验证的 SSM 混合大模型。它的洞察是：SSM 在长上下文和推理速度上优势明显，但 Transformer 的 Attention 在极其精确的信息检索（例如“大海捞针”测试）和复杂逻辑推理上，依然具有不可替代的归纳偏置。
与其争论谁取代谁，不如把两者结合起来。Jamba 的设计哲学就是这种务实主义。
2. 宏观架构：层混合
Jamba 的核心创新是层混合（Layer Mixing）：在深层网络中交替使用 Mamba 层和 Attention 层。具体配置是把网络切成若干个 Jamba block，每个 block 内 7 层是 Mamba 层、1 层是 Attention 层，也就是注意力与 Mamba 的比例为 1:7；官方模型一共 4 个 block，合计 32 层，因此整个大模型里只有 4 层注意力。（容易混淆的一点是：52B 指的是总参数量，32 才是层数。）
采用 1:7 这个比例的理由是：Mamba 层负责高效的序列压缩和长程记忆，推理时是 $O(1)$；少量的 Attention 层负责精确的全局检索，弥补 SSM 在精确定位上的不足；实验表明这个比例在性能和效率之间达到了最佳平衡。
3. 每层的具体设计
两种层的内部结构是统一的：先做一次 RMSNorm，再经过 Mamba 模块或注意力模块，接着再做一次 RMSNorm 并经过 MLP（使用 SwiGLU 激活），外面套残差连接。
Mamba 层 与标准 Mamba-2 类似，包含线性扩展、一维卷积、SSD 核心、门控和线性投影，状态维度 $N$ 取 256。
Attention 层 使用分组查询注意力（GQA）这类多头自注意力结构，并且照常使用 KV Cache。它之所以不构成负担，是因为整个模型只有 4 层注意力，KV Cache 比同规模的纯 Transformer 小得多（官方给出的对比是 256K 上下文下约 4GB，纯 Transformer 约 32GB）。
4. MoE 的集成
Jamba 在通道混合阶段（也就是 MLP 部分）引入了 MoE（混合专家模型），替换方式为隔层进行，也就是每隔一层把 MLP 替换为 MoE：
- 专家数量：每个 MoE 层 16 个专家。
- 激活数量：每个 token 只激活 2 个专家。
- 路由机制：使用 Top-2 路由，也就是选择得分最高的 2 个专家。
由此带来的参数规模是：总参数量 52B，激活参数量 12B。激活参数占比低，是因为每个 token 只激活 $2/16 = 12.5%$ 的专家。
5. 上下文窗口与位置编码
Jamba 的上下文窗口为 256K，也就是 256,000 个 token。
位置编码方面，Jamba 没有使用任何显式的位置编码。原因是每个 Jamba block 都从 Mamba 层开始，token 的顺序信息由 Mamba 层的递推天然携带；注意力层数量很少，而且夹在两个 Mamba 层之间，所以不需要额外注入位置信息。
6. 实验结果
把 Jamba 与两个同期开源模型放在一起比较，可以看到三点差异。
激活参数： Jamba 为 12B，Mistral 7B 为 7B，LLaMA-2 13B 为 13B。
上下文窗口： Jamba 为 256K，Mistral 7B 为 32K，LLaMA-2 13B 为 4K。
MMLU 分数： Jamba 为 62.3，Mistral 7B 为 60.1，LLaMA-2 13B 为 54.8。
Needle-in-Haystack（256K 长度下的“大海捞针”）测试： Jamba 通过，Mistral 7B 与 LLaMA-2 13B 均失败。
也就是说，Jamba 在同等激活参数下性能超过了 LLaMA 2 与 Mistral，同时支持 256K 的超长上下文。
后续版本。 2024 年 8 月，AI21 又发布了 Jamba-1.5 系列：Mini 版本为 12B 激活参数、52B 总参数，Large 版本为 94B 激活参数、398B 总参数，两者都把上下文窗口维持在 256K。Large 版本经过量化之后可以放进单台配备 8 张 80GB 显卡的机器。这个数字值得留意：总参数量已经接近 400B，而推理成本对应的只是 94B 激活参数，这正是 MoE 与混合结构配合之后的效果。
7. Jamba 的启示
Jamba 的工作说明了四件事。
- SSM 和 Attention 是互补关系，两者都有各自的强项。
- 少量的 Attention 层可以弥补 SSM 在精确检索上的不足，并且不会显著增加推理成本。
- MoE 可以进一步扩展模型容量，同时不增加推理时的计算量。
- 未来的大模型架构很可能是混合架构，纯粹的 Transformer 或纯粹的 SSM 都难以覆盖全部需求。
五、全景小结
把 S4、Mamba 与这 5 个变体放在一起，可以得到一张完整的地图。
- S4 是奠基者，解决长序列训练效率问题，关键创新是 HiPPO、DPLR、Cauchy 核与卷积视角。
- S5 是极简派，解决 S4 数学太复杂的痛点，关键创新是 MIMO 架构与纯对角 $\mathbf{A}$ 矩阵。
- H3 是破局者，解决 S4 做不好语言建模的问题，关键创新是 Shift SSM（显式局部记忆）与带衰减对角 SSM 的乘法门控，整体形式与线性注意力同构。
- RWKV 是社区奇迹，解决 Transformer 推理显存爆炸的问题，关键创新是线性 Time-Mixing 与极简公式。
- Mamba-1 是破圈之王，解决 LTI 缺乏内容感知的问题，关键创新是选择性机制（$\Delta$、$B$、$C$ 依赖输入）与并行扫描。
- Mamba-2 是理论巅峰，解决 Mamba-1 训练不够快的问题，关键创新是 SSD 理论（SSM 与线性注意力对偶）与分块 GEMM。
- Jamba 是工业务实派，解决纯 SSM 归纳偏置缺失的问题，关键创新是 Mamba 加少量 Attention 再加 MoE 的混合结构。
一个值得想清楚的问题。 通过 Mamba-2 的 SSD 理论我们知道，选择性 SSM 和线性 Attention 在数学上是等价的。既然数学上等价，为什么在真实的语言建模任务中，Mamba 的表现通常比朴素的线性 Attention 更好？
提示可以从状态衰减的角度入手。线性 Attention 的 KV Cache 是不断累加的，只增不减；而 SSM 的状态带有 $\bar{\mathbf{A}}$ 衰减因子。差异可以这样理解：
第一，衰减因子给模型提供了一个遗忘机制。语言中越近的信息通常越重要，带衰减的状态天然带有近因偏好，而只增不减的累加器会把很久以前的信息一直留在分母里，稀释当前相关的线索。
第二，衰减因子的取值是输入依赖的。Mamba 的 $\bar{\mathbf{A}}_t = \exp(\Delta_t \mathbf{A})$ 由内容决定，模型可以按需选择忘掉多少；朴素线性 Attention 的衰减是固定或缺失的，它没有这种自适应能力。
第三，数值上的差异也来自结构。SSM 的状态被压缩到一个固定大小的 $h \in \mathbb{R}^N$ 中，天然带有正则化效果，而线性 Attention 的记忆随序列增长不断扩大，容易让无关信息抢占权重。
数学等价说明两者表达能力在同一层级，实际效果的区别来自参数化方式、归纳偏置以及数值稳定性。
下一篇，我们会跳出单个模型，把 RNN、Transformer 和 SSM 放在同一个框架下做统一对比，并给出实战选型建议。