Jamba

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$$ ...