Mamba

SSM 的应用、工程实践与前沿方向 前几篇我们把 SSM 的数学原理（HiPPO、DPLR）、架构演进（S4 到 Mamba）以及底层优化（并行扫描、硬件感知）都拆解了一遍。本篇从理论走向真实业务：先看 SSM 如何跨越文本的边界，在视觉、语音、生物医疗等领域落地；再给出一份实操层面的工程指南；然后客观审视它目前的局限，最后梳理这个方向正在往哪里走。 一、跨模态应用：SSM 在哪里发挥作用 SSM 的核心优势是“在极长序列下保持线性复杂度和常数级推理显存”。只要数据满足“序列极长”或“流式到达”这两个特点之一，SSM 就比 Transformer 更有优势。下面分四个领域来看。 1. 计算机视觉：Vision Mamba（Vim）与 VMamba 要解决的问题。 Transformer（ViT）处理图像时，需要把图像切成 Patch 并展平为 1D 序列。但图像的 Patch 数量很大，例如 $224 \times 224$ 的图像切成 196 个 Patch，Attention 的 $O(L^2)$ 复杂度让高分辨率图像的处理变得极其昂贵。 架构改造的矛盾。 SSM 是 1D 因果模型，只能看过去不能看未来；而图像是 2D 的，没有天然的“因果方向”，因为一个像素的上下文分布在它的上下左右。 真实应用中的解法。 Vim（Vision Mamba）采用双向扫描（Bidirectional Scan）：把图像展平后正向扫一遍，再反向扫一遍，把两个方向的状态拼接起来，从而获得全局感受野。VMamba（VSS Block）采用更进一步的交叉扫描（Cross-Scan）：将图像分别按“行、列、反向行、反向列”展平成 4 个 1D 序列，分别跑 4 个 SSM，最后把 4 个结果合并。这弥补了 SSM 缺乏 2D 空间感知能力的缺陷。 结果。 在 ImageNet 分类和 COCO 目标检测上，VMamba 以显著更低的 FLOPs（计算量）超过了 Swin Transformer。 ...

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

SSM 家族全景图（上）：S5、H3 与 RWKV 前面的文章里，我们深入拆解了 SSM 家族的两座大山：S4 与 Mamba。S4 用 HiPPO、DPLR 和 FFT 把长序列建模做到了当时的最好水平，Mamba 用选择性机制和并行扫描把它推向了语言建模的舞台中央。 本篇跳出单个模型，俯瞰整个家族在 S4 奠基、Mamba 破圈之后长出的其他分支。我们选 5 个最具代表性的变体，每个都按“动机、核心创新、架构与数学”三个部分来拆解。S5、H3、RWKV 放在本篇，Mamba-2 与 Jamba 放在下一篇。 一、S5（Simplified S4）：用最简数学达到同等效果 1. 动机：S4 的工程实现太复杂 S4 虽然强大，但工程实现极其繁琐。要处理复数域、DPLR 结构、Woodbury 恒等式、Cauchy 矩阵和 FFT。研究者自然会问：能不能去掉这些复杂的数学技巧，用最简单的矩阵实现同等的性能？ 2. SISO 的含义：$d$ 个独立的小 SSM 要理解 S5 的改进，先要弄清 S4 的 SISO（Single-Input Single-Output，单输入单输出）策略到底指什么。 假设输入为 $\mathbf{X} \in \mathbb{R}^{B \times L \times d}$，其中 $d = 768$ 是特征维度。S4 的处理方式是把 $d$ 维特征向量拆成 $d$ 个独立的标量通道，每个通道独立地跑一个一维的 SISO SSM。 以通道 $i$ 为例，它的小 SSM 结构是这样的：输入是标量 $x_t^{(i)} \in \mathbb{R}$；状态是 $h_t^{(i)} \in \mathbb{R}^N$（$N$ 维，例如 64）；参数是 $\mathbf{A}^{(i)} \in \mathbb{R}^{N \times N}$、$\mathbf{B}^{(i)} \in \mathbb{R}^{N \times 1}$、$\mathbf{C}^{(i)} \in \mathbb{R}^{1 \times N}$；输出是标量 $y_t^{(i)} \in \mathbb{R}$。 ...

Mamba（下）：并行扫描与硬件感知优化 上一篇我们给 SSM 装上了“选择性”，让它变成线性时变（LTV）系统，代价是卷积视角失效、S4 赖以生存的 FFT 加速无法使用。如果退回最朴素的 RNN 循环视角，用 for 循环一步步算 $h_t$，GPU 上训练 Mamba 会比 Transformer 慢得多。 本篇从两个层面拆解 Mamba 的应对办法：数学优化，把串行递推变成可并行的计算；工程优化，消除 GPU 上的内存带宽瓶颈。最后给出 Mamba 与 Transformer、S4 的整体对比。 一、问题的起点：串行递推 无法使用 FFT 之后，我们只能回到递推公式： $$h_t = \bar{\mathbf{A}}t h{t-1} + \bar{\mathbf{B}}_t x_t$$ 这个式子存在严格的串行依赖：算 $h_2$ 必须等 $h_1$ 算完，算 $h_3$ 必须等 $h_2$ 算完。在拥有上万个计算核心的 GPU 上，串行执行等于浪费了绝大部分算力。 破局的思路是：寻找一种数学方法，把看似串行的递推转化为可以并行计算的形式。答案就是并行扫描（Parallel Scan），而它成立的前提是操作必须满足结合律（Associativity）。 二、并行扫描（Parallel Scan） 1. 把递推写成仿射变换 为了符号简洁，令 $u_t = \bar{\mathbf{B}}_t x_t$，它是一个已知向量。递推式变为： $$h_t = \bar{\mathbf{A}}t h{t-1} + u_t$$ 这在数学上是一个仿射变换（Affine Transformation）。我们可以把每一步的“转移”抽象为一个二元组 $(\mathbf{A}_t, u_t)$，它的作用是把前一个状态 $h$ 映射到新状态： ...

Mamba（上）：动机、架构与选择性机制 前六篇我们从连续时间的状态空间方程出发，走过了离散化、HiPPO 理论、S4 的 Block 结构，以及 S4 赖以高效的 DPLR、双线性变换、Cauchy 核与 FFT 加速。S4 已经能在长序列任务上全面超过 Transformer，但它在语言建模上始终差一口气。根本原因在于它是线性时不变（LTI）系统：训练完成后，矩阵 $\mathbf{B}$、$\mathbf{C}$、$\mathbf{A}$ 就固定了，对所有输入一视同仁。 本篇进入 SSM 发展史上最关键的一步——Mamba。我们仍然按“自顶向下”的顺序展开：先看它要解决什么问题，再看它的整体架构，最后放大到一个 Mamba Block 的内部，逐个拆解五个组件，并把选择性机制背后的数学讲透。并行扫描与硬件感知优化这两块工程内容留到下一篇。 一、Mamba 要解决什么问题 1. 序列建模需要的两种能力 如果把人类阅读长文的过程抽象出来，一个优秀的序列模型需要同时具备两种能力。 第一是压缩能力（Compression）。 模型要能把长达几十万字的背景信息压缩并维持在一个固定大小的“脑容量”（隐藏状态）里，而且不能出现“脑容量爆炸”，也就是显存 OOM。 第二是内容感知与选择性（Content-Awareness / Selectivity）。 面对海量信息，模型要能根据当前看到的内容，动态决定什么该牢牢记住，什么可以左耳进右耳出。 举个具体的例子。当读到“请注意，接下来的密码是 8492”时，模型应该高度专注，把“8492”刻进记忆；而当读到“今天天气不错，顺便说一句”这类废话时，模型应该降低关注度，保住之前记下的核心内容不被覆盖。 2. 现有模型各自的困境 Transformer 通过自注意力机制（$QK^\top$）天然具备很强的内容感知能力，注意力权重完全由当前 Query 和 Key 的内容决定。它的代价是灾难性的 $O(L^2)$ 复杂度，长文本的压缩成本极高。 经典 SSM（如 S4） 具备很强的压缩能力，推理显存只有 $O(1)$，但它是线性时不变的。矩阵 $\mathbf{B}$、$\mathbf{C}$、$\mathbf{A}$ 在训练好之后就固定了，对所有输入一视同仁。它像一台没有感情的录音机，匀速记录一切，做不到“遇到关键词就竖起耳朵”。 3. Mamba 的目标 Mamba 要打破 LTI 的枷锁，在保留 SSM 线性复杂度 $O(L)$ 与常数级推理显存的前提下，赋予 SSM 内容感知（也就是选择性）的能力，使它在语言建模等需要精确信息检索的任务上能够与 Transformer 正面竞争。 二、Mamba 的整体架构 1. 宏观结构：Decoder-only 堆叠 Mamba 的宏观架构很纯粹，是一个典型的 Decoder-only 堆叠结构，与 GPT 等 Transformer 模型高度一致。数据流如下： ...

SSM 优化（下）：Cauchy 核、FFT 加速与工程优化 上一篇介绍了 S4 通过 DPLR 结构、双线性变换和 Z 变换将卷积核计算的复杂度从 $O(LN^2)$ 降到 $O(LN)$。但 $O(LN)$ 在超长序列（$L = 16384$）和大状态维度（$N = 256$）下仍然有约 $4 \times 10^6$ 的运算量。本文将介绍 S4 如何利用 Cauchy 矩阵的特殊结构和 FFT 进一步加速，以及训练时的反向传播策略和显存优化。最后讨论 S4 的根本局限性，以及它如何引向 Mamba 的诞生。 读者在阅读本篇时，可能会对一个结构细节有疑问：S4 中的状态到底是一个总的 $N$ 维向量，还是别的什么东西？这一点在跟 Transformer 对比时特别容易混淆，所以本文先把它澄清掉，避免后续关于“选择性机制”和 Mamba 的讨论建立在错误的前提上。 一、Cauchy 矩阵与快速求值 1. 一个先决问题：S4 真的只有一个 $N$ 维状态吗？ 在很多地方，S4 被简略地说成是“一个 $N$ 维状态空间”。这其实是对整个模型结构的过度简化，容易让人误以为 S4 和 Transformer 的内部机制很像。 更准确的说法是： S4 层通常作用于 $d$ 个特征维度（例如 $d = 768$ 或更大）。 对每一个特征维度 $j$，都会独立地维护一个 $N$ 维状态空间。换句话说，每个特征维度有自己的 $\mathbf{A}_j$、$\mathbf{B}_j$、$\mathbf{C}_j$ 体系（实际参数化后可能会共享结构，但变量是独立可学的）。 所以总的状态规模大约是 $d \times N$，但在每一个特征维度内部，状态维度 $N$ 仍然很小，典型值大概是 $16 \sim 128$。 这意味着什么？ ...

SSM 入门：从连续微分方程到离散递推 上一篇文章我们从宏观角度介绍了 SSM 的定位和优势。本文将进入 SSM 的数学基础：连续时间的线性状态空间方程如何定义、如何离散化为计算机可执行的递推公式、离散化后如何获得"循环"和"卷积"两种等价视角，以及一个完整的 SSM Block 由哪些组件构成。 本文的目标是让读者建立起对 SSM 的初步知识体系，不涉及后续的数学优化和工程优化。 一、连续时间的 SSM 方程 1. 两个核心方程 在连续时间 $t$ 下，一个线性状态空间模型由两个方程定义。 状态方程： $$h’(t) = \mathbf{A}h(t) + \mathbf{B}x(t)$$ 其中 $h’(t)$ 是 $h(t)$ 对时间 $t$ 的导数，即 $\frac{dh(t)}{dt}$。这个方程描述的是：在当前时刻，内部记忆 $h(t)$ 的变化率由两部分决定——自身当前状态 $\mathbf{A}h(t)$ 和新输入 $\mathbf{B}x(t)$。 观测方程： $$y(t) = \mathbf{C}h(t) + \mathbf{D}x(t)$$ 这个方程描述的是：如何从当前的内部记忆 $h(t)$ 中提取出我们关心的输出 $y(t)$。 2. 各变量的含义 $x(t) \in \mathbb{R}^d$：当前时刻的输入信号。在 NLP 中，这是某个 token 的 Embedding 向量。 $h(t) \in \mathbb{R}^N$：当前时刻的隐藏状态，即模型对历史信息的"记忆快照"。 $y(t) \in \mathbb{R}^d$：当前时刻的输出信号，将传递给下一层或用于最终预测。 $\mathbf{A} \in \mathbb{R}^{N \times N}$：状态转移矩阵。它控制记忆如何随时间自然演化。如果 $\mathbf{A}$ 的特征值为负，记忆会随时间衰减（遗忘旧信息）；如果为正，记忆会指数增长（数值爆炸）。 $\mathbf{B} \in \mathbb{R}^{N \times d}$：输入投影矩阵。它控制新输入以多大的强度和方式写入状态。可以理解为一个"写入阀门"。 $\mathbf{C} \in \mathbb{R}^{d \times N}$：输出投影矩阵。它控制如何从状态中提取输出。可以理解为一个"读取探针"。 $\mathbf{D} \in \mathbb{R}^{d \times d}$：直接穿透矩阵，代表输入绕过记忆直接到达输出的路径。在实际的深度学习 SSM 中，$\mathbf{D}$ 通常被设为 0 或极小常数，因为它不经过记忆，对序列建模没有实质贡献。 三个核心参数的分工可以概括为：$\mathbf{B}$ 负责写入，$\mathbf{A}$ 负责维持和遗忘，$\mathbf{C}$ 负责读出。 ...

状态空间模型（SSM）：从 RNN 与 Transformer 的困境出发 在序列建模的发展历程中，RNN 和 Transformer 是两座绕不过去的丰碑。RNN 以递推的方式逐步处理序列，结构简洁但受困于串行计算和长程依赖问题；Transformer 用自注意力机制彻底解决了长程依赖和并行化训练的难题，却引入了序列长度平方级的计算和显存开销。 这两种架构各有不可忽视的缺陷。有没有一种可能，让序列模型同时拥有 RNN 推理时的高效和 Transformer 训练时的并行？ 这就是 SSM（State Space Model，状态空间模型）试图回答的问题。本文将从 RNN 和 Transformer 的痛点出发，解释 SSM 是什么、它为什么能做到"鱼与熊掌兼得"，以及它从 2020 年到 2024 年的演进路线。 一、RNN 与 Transformer 的痛点 1. RNN 的核心问题 RNN（包括 LSTM、GRU）的核心是隐藏状态的递推更新：$h_t = f(h_{t-1}, x_t)$。这种递推结构带来了两个根本性问题。 第一，训练无法并行。 由于 $h_t$ 强依赖于 $h_{t-1}$，序列中不同时间步的计算必须严格串行进行。在 GPU 这种为大规模矩阵并行运算设计的硬件上，RNN 的训练效率极低。假设序列长度为 $L$，RNN 的训练需要 $L$ 个串行时间步，无法像矩阵乘法那样在 GPU 上一次算完。 第二，长程依赖衰减。 即使有 LSTM 的门控机制缓解梯度消失问题，当序列长度达到数千甚至数万时，信息从第一个 token 传递到最后一个 token 仍然需要经过漫长的递推链路，每一次传递都会有信息损耗。就像传话游戏，传的人越多，最后听到的内容和最初说的偏差越大。 2. Transformer 的核心问题 Transformer 用自注意力机制解决了 RNN 的两个痛点：通过计算所有 token 之间的两两相关性，实现了全局信息的直接交互和完全并行的训练。但这个解决方案引入了新的代价。 ...