Posts

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、Transformer 与 RNN 的统一视角 前面几篇里，我们分别拆解了 S4、Mamba 以及家族中的其他变体。本篇把视角拔高一层，把 RNN、Transformer 和 SSM 这三大门派放在同一张桌子上，先看它们共同的数学本质，再看具体的指标差异，最后给出实战选型建议。 一、统一视角：三者都在对历史做加权求和 如果剥开 RNN、Transformer 和 SSM 复杂的网络结构（FFN、LayerNorm、激活函数），只看它们最核心的序列混合（Sequence Mixing）机制，也就是“当前输出如何依赖过去的输入”，会发现一个统一的事实。 它们在数学本质上都是在对历史输入做加权求和。假设要计算第 $t$ 个时间步的输出 $y_t$，输入序列是 $x_1, x_2, \ldots, x_t$，三种模型都可以被抽象为同一个公式： $$y_t = \sum_{k=1}^{t} \alpha_{t,k} \cdot x_k$$ 这里的 $\alpha_{t,k}$ 是一个权重（可以是标量，也可以是矩阵），它决定了在计算当前输出 $y_t$ 时，历史第 $k$ 个输入 $x_k$ 到底有多重要。 三大模型的根本分歧因此收敛为一个问题：权重 $\alpha_{t,k}$ 是如何计算出来的？ 二、三种计算 $\alpha_{t,k}$ 的思路 1. RNN（LSTM / GRU）：黑盒的隐式压缩 RNN 没有显式计算 $\alpha_{t,k}$ 的公式。它把所有的历史信息一股脑塞进一个固定大小的隐藏状态 $h_t$ 中，通过复杂的非线性门控（遗忘门、输入门）进行压缩。 这个方式有两个特点。 隐式且不可控。 很难解释 RNN 为什么记住了第 3 个词而忘了第 5 个词，注意力分布没有显式表达，只能从状态里反推。 天然的衰减与遗忘。 隐藏状态容量有限，每次更新都有信息损耗，$\alpha_{t,k}$ 会随着时间差 $(t-k)$ 的增大而不可避免地趋向于 0。这就是梯度消失现象在权重视角下的对应表现。 ...

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 优化（中）：DPLR 结构与高效离散化 上一篇介绍了 HiPPO 理论如何推导出理论上最优的 $\mathbf{A}$ 矩阵，同时也指出了两个工程问题：HiPPO-LegS 的时变性和密集矩阵 $O(N^3)$ 的计算代价。 本文将详细介绍 S4 如何通过三个关键技术解决这些问题：用 DPLR 结构将密集矩阵转化为稀疏表示，用双线性变换替代 ZOH 离散化以保持 DPLR 结构，以及用 Z 变换将卷积核序列的计算转化为矩阵求逆问题。 一、DPLR 结构：将密集矩阵变为稀疏表示 1. 问题的起点 HiPPO 推导出的 $\mathbf{A}$ 是一个 $N \times N$ 的密集下三角矩阵。如果直接使用它，离散化后的 $\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$ 和卷积核 $\bar{K}_i = \mathbf{C}\bar{\mathbf{A}}^i\bar{\mathbf{B}}$ 的计算都需要密集矩阵运算，复杂度为 $O(N^3)$。 我们需要一种特殊的矩阵结构，使得幂运算和矩阵求逆都能在 $O(N)$ 内完成。 2. 什么矩阵的幂运算最快 在高等代数中，对角矩阵的幂运算是最简单的。如果 $\mathbf{\Lambda} = \text{diag}(\lambda_0, \lambda_1, \ldots, \lambda_{N-1})$，那么 $\mathbf{\Lambda}^i = \text{diag}(\lambda_0^i, \lambda_1^i, \ldots, \lambda_{N-1}^i)$，复杂度仅为 $O(N)$。 但纯对角矩阵意味着各个状态维度之间完全独立，没有信息交互，无法表达复杂的序列模式。 2.1 S4 内部的状态组织方式：每个特征维度独立维护一个 $N$ 维状态 S4 并非只使用“一个 $N$ 维状态”，这一点在与 Transformer 的对比中非常关键。 ...

SSM 优化（上）：为什么不能直接算，以及 HiPPO 理论 前面几篇文章建立了 SSM 的基础知识和 S4 Block 的架构理解。但我们刻意回避了一个核心问题：矩阵 $\mathbf{A}$ 到底该怎么初始化？ 如果 $\mathbf{A}$ 随机初始化，SSM 在长序列上几乎必然失败——要么数值爆炸，要么瞬间遗忘。HiPPO 理论从数学上证明了什么样的 $\mathbf{A}$ 才是"理论上最优的"。本文将完整介绍 HiPPO 的推导过程和核心结论，为后续的 S4 工程化优化做好铺垫。 一、为什么不能随机初始化 $\mathbf{A}$ 1. 随机初始化的后果 回顾离散化后的 SSM 递推公式：$h_k = \bar{\mathbf{A}}h_{k-1} + \bar{\mathbf{B}}x_k$，其中 $\bar{\mathbf{A}} = e^{\mathbf{A}\Delta}$。 如果 $\mathbf{A}$ 像普通神经网络权重那样随机初始化（例如标准正态分布），会发生三种情况： $\mathbf{A}$ 的特征值实部为正： $\bar{\mathbf{A}}$ 的特征值大于 1。递推 $h_k = \bar{\mathbf{A}}h_{k-1} + \ldots$ 中，矩阵连乘 $\bar{\mathbf{A}}^k$ 会迅速趋于无穷大，状态 $h_k$ 数值爆炸。 $\mathbf{A}$ 的特征值实部为负且绝对值很大： $\bar{\mathbf{A}}$ 的特征值远小于 1。矩阵连乘 $\bar{\mathbf{A}}^k$ 迅速趋于 0，模型瞬间"遗忘"所有历史输入。 $\mathbf{A}$ 是一般随机矩阵： 特征值在复平面上近似均匀分布在某个圆内（Circular Law），无法保证全部实部为负，更无法控制衰减速率的分布。在这种情况下，训练几乎不可能稳定进行。 2. 问题的本质 SSM 需要一个特殊的 $\mathbf{A}$ 矩阵，它必须同时满足两个看似矛盾的要求： ...

数学知识前置讲解：后续优化篇中一定会用到的数学工具 在接下来的 SSM 优化系列文章中，我们会频繁用到一些超出本科线性代数和高等数学范畴的数学知识和算法思想。本文集中讲解这些前置知识，让读者在后续阅读中不会因为数学工具的陌生而阻碍对核心思想的理解。 本文假设读者已具备线性代数基础（矩阵运算、特征值、矩阵指数）、高等数学基础（微分方程、积分）以及生成函数的基本概念。在此基础上，本文重点讲解两个读者可能尚未接触的领域：矩阵运算的计算复杂度分析和快速傅里叶变换（FFT）。同时也会简要介绍 Cauchy 矩阵和递推与卷积的联系。 一、矩阵运算的计算复杂度 SSM 优化的核心目标之一，就是把 $O(N^3)$ 的密集矩阵运算降到 $O(N)$。要理解为什么这是可能的，我们需要先理解：不同矩阵运算的复杂度分别是多少，以及为什么是这个数量级。 1. 为什么需要关心复杂度 在深度学习中，“一个算法能不能用"不仅取决于它在数学上是否正确，还取决于它在工程上是否可行。HiPPO 给出了理论上最优的 $\mathbf{A}$ 矩阵，但由于密集矩阵运算的复杂度过高，直到 S4 通过 DPLR 结构将其优化到可接受的范围，这个理论优势才真正落地。 这里还有一个常被跳过的小票据：S4 在效果上的表现，最早是靠 Long Range Arena（LRA）这类长序列基准里的显著结果让大家信服的。比如在 LRA 六项任务里，S4 的算术平均约为 81.84%，而标准 Transformer 约为 45.32%；到了最长的 Path-X（长度 16000）任务上，Performer、BigBird 这些高效注意力方法也和标准 Transformer 一样崩溃至 0%，S4 仍然达到 86.73%。有的总结材料里会把 S4 的平均写成约 86.1%，这个数和从六项明细算出的 81.84% 并不完全一致，可能来源于不同报告版本或平均方式的差异，所以引用时最好还是以逐任务明细为准。 2. 矩阵-向量乘法：$O(N)$ 与 $O(N^2)$ 的区别 一般矩阵乘向量： 设 $\mathbf{A} \in \mathbb{R}^{N \times N}$，$\mathbf{h} \in \mathbb{R}^N$，计算 $\mathbf{y} = \mathbf{A}\mathbf{h}$。矩阵 $\mathbf{A}$ 有 $N^2$ 个元素，每个输出分量 $y_i = \sum_{j=1}^{N} A_{ij} h_j$ 需要 $N$ 次乘法和 $N-1$ 次加法。总共 $N$ 个分量，总计算量为 $N \times N = N^2$ 次乘法，即 $O(N^2)$。 ...