【模型后训练与深度推理】4 状态空间模型与线性注意力：Mamba架构的原理与实现

状态空间模型与线性注意力：Mamba架构的原理与实现

1.1 为什么Transformer在长序列上遇到瓶颈

1.1.1.1 自注意力的二次复杂度

想象我们正在读一本十万字的小说，Transformer的自注意力机制要求我们在读每一个新词的时候，都要回头和之前所有的词做一次"对视"。读第一个词时回头看零个词，读第二个词回头看一个词，读第一万个词时要回头看九千九百九十九个词。这种全局回望的代价，随着序列长度 $L$ 的增长，呈现出 $O(L^2)$ 的爆炸趋势。

形式化地，给定查询矩阵 $Q\in {\mathbb{R}}^{L\times d}$、键矩阵 $K\in {\mathbb{R}}^{L\times d}$ 和值矩阵 $V\in {\mathbb{R}}^{L\times d}$，自注意力的输出为：

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d}}\right)V$$

其中 $Q{K}^T\in {\mathbb{R}}^{L\times L}$ 这一步，就是二次复杂度的根源。当 $L=2048$ 时，注意力矩阵有大约四百万个元素；当 $L=32768$ 时，这个数字膨胀到十亿级别。如果画成图会是什么样子？想象一个巨大的棋盘，序列中的每个词都要和棋盘上的所有格子打一次招呼，棋盘面积随边长平方增长。

现在我们已经了解了二次复杂度的来源，接下来看看这种复杂度在硬件层面造成的具体困境。

1.1.1.2 KV-Cache的内存墙

在继续之前，我们先看看推理阶段的Transformer在做什么。生成第 $t$ 个token时，模型需要把之前所有 $t-1$ 个token的键和值向量缓存起来，以便计算注意力分数。这个缓存就是KV-Cache。

对于批大小为 $b$、头数为 $h$、每头维度为 $d_h$、层数为 $n$ 的模型，KV-Cache占用的显存为：

$${\text{Memory}}_{\text{KV}}=2\cdot b\cdot n\cdot h\cdot L\cdot d_h$$

系数2来自K和V两个矩阵。当 $L$ 达到十万级别时，KV-Cache可以轻松占满一张高端GPU的全部显存。更麻烦的是，这个内存占用随序列长度线性增长，但GPU的显存是固定的。这意味着Transformer的上下文窗口存在一个硬性的物理上限。

听起来抽象对吧？如果画成图会是什么样子？想象一个仓库，每生成一个新词就要往仓库里塞进两个箱子（K和V）。仓库的门宽度固定，箱子越堆越高，总有一天会顶到天花板。这就是KV-Cache的内存墙。

1.2 状态空间模型的连续时间视角

1.2.1.1 从ODE到离散化：一个弹簧系统的类比

现在我们已经了解了Transformer的瓶颈，接下来看看状态空间模型（State Space Model，SSM）如何从连续时间的视角重新理解序列建模。

想象一个弹簧振子：你用手施加一个外力 $x(t)$，弹簧的状态 $h(t)$（位移和速度）随之变化，而我们观测到的输出 $y(t)$ 是弹簧的位置。这个物理系统的动力学可以用一组线性常微分方程描述：

$$h^{\prime }(t)=Ah(t)+Bx(t)$$
$$y(t)=Ch(t)+Dx(t)$$

这里 $h(t)\in {\mathbb{R}}^N$ 是隐藏状态向量，$x(t)\in \mathbb{R}$ 是标量输入，$y(t)\in \mathbb{R}$ 是标量输出。矩阵 $A\in {\mathbb{R}}^{N\times N}$ 刻画了系统内部的演化规律，$B\in {\mathbb{R}}^{N\times 1}$ 是输入对状态的驱动矩阵，$C\in {\mathbb{R}}^{1\times N}$ 是状态到输出的观测矩阵，$D\in \mathbb{R}$ 是输入到输出的直接通路。

别急，这个连续时间的框架和离散的token序列有什么关系？关系在于：我们可以把文本序列看作一个连续信号在离散时间点上的采样。每个token对应一个脉冲输入，系统的状态在token之间连续演化，而在token到达的瞬间被更新。

1.2.1.2 零阶保持与双线性变换

要把连续系统变成计算机能处理的离散系统，我们需要一个离散化规则。最自然的想法是：在两个token之间，输入保持不变（零阶保持，Zero-Order Hold）。假设token间隔为步长 $\Delta$，那么在区间 $[k\Delta ,(k+1)\Delta )$ 内，$x(t)=x_k$ 为常数。

在这个假设下，状态方程的解为：

$$h_k=\bar{A}{h}_{k-1}+\bar{B}{x}_k$$
$$y_k=\bar{C}{h}_k+\bar{D}{x}_k$$

其中离散化后的参数为：

$$\bar{A}=e^{\Delta A}$$
$$\bar{B}=(\Delta A{)}^{-1}(e^{\Delta A}-I)\cdot \Delta B$$
$$\bar{C}=C,\bar{D}=D$$

推导并不复杂。对 $h^{\prime }(t)=Ah(t)+B{x}_k$ 在区间 $[0,\Delta ]$ 上积分，利用常数变易法即可得到上述结果。如果画成图会是什么样子？想象一条平滑的曲线（连续状态），每隔 $\Delta$ 距离被采样一次，采样点之间曲线按照指数规律 $e^{\Delta A}$ 自然衰减或增长。

在实际使用中，$\Delta$ 本身也可以是一个可学习的参数，它控制着系统对输入的"响应速度"。大步长意味着系统 sluggish，对快速变化的输入反应迟钝；小步长意味着系统敏感，容易受噪声干扰。

1.3 SSM的序列建模形式

1.3.1.1 线性时不变系统的卷积表示

现在我们已经了解了SSM的离散化过程，接下来看看它在序列数据上的计算形式。

回顾离散SSM的递推关系：

$$h_k=\bar{A}{h}_{k-1}+\bar{B}{x}_k$$
$$y_k=\bar{C}{h}_k$$

假设初始状态 $h_0=0$，我们可以把 $h_k$ 完全展开为输入的线性组合：

$$h_k=\sum _{i=1}^k{\bar{A}}^{k-i}\bar{B}{x}_i$$

代入输出方程：

$$y_k=\bar{C}\sum _{i=1}^k{\bar{A}}^{k-i}\bar{B}{x}_i=\sum _{i=1}^k\bar{C}{\bar{A}}^{k-i}\bar{B}{x}_i$$

定义核函数 ${\bar{K}}_j=\bar{C}{\bar{A}}^j\bar{B}$，则输出可以写成输入与核的离散卷积：

$$y_k=\sum _{i=1}^k{\bar{K}}_{k-i}{x}_i=(\bar{K}\ast x{)}_k$$

这个形式极其重要：它意味着整个SSM前向传播可以表示为一次卷积操作。卷积的计算复杂度是多少？使用快速傅里叶变换（FFT）可以在 $O(L\log L)$ 内完成，或者更简单地，我们可以预计算核函数 $\bar{K}$ 并在序列上滑动。

如果画成图会是什么样子？想象一条河流，每个输入 $x_i$ 像一颗石子投入水中，激起的涟漪按照 $\bar{C}{\bar{A}}^{k-i}\bar{B}$ 的权重传播到下游。输出 $y_k$ 就是位置 $k$ 处所有过往涟漪的叠加。

1.3.1.2 结构化状态空间的对角化假设

直接处理一般的矩阵 $\bar{A}$ 仍然计算昂贵。关键观察是：如果 $A$ 是对角矩阵，那么 $e^{\Delta A}$ 也是对角矩阵，且每个对角元独立计算。

设 $A=\text{diag}(a_1,a_2,\dots ,a_N)$，则：

$$\bar{A}=e^{\Delta A}=\text{diag}(e^{\Delta a_1},e^{\Delta a_2},\dots ,e^{\Delta a_N})$$

此时，核函数可以分解为 $N$ 个独立模式的叠加：

$${\bar{K}}_j=\sum _{n=1}^N{C}_n(e^{\Delta a_n}{)}^j{B}_n$$

每个模式 $n$ 都是一个指数衰减（或增长）的序列，衰减速率由 $a_n$ 的实部决定。在实际应用中，$a_n$ 通常被参数化为负实数，确保系统稳定（状态不爆炸）。

这种对角化假设把矩阵运算降维为逐元素运算，是SSM能够高效实现的核心。这在训练中意味着：我们可以并行计算 $N$ 个指数核，然后在通道维度上求和，整个过程不涉及任何矩阵乘法。

1.4 选择性扫描机制

1.4.1.1 为什么LTI不够：内容依赖的困境

现在我们已经了解了线性时不变（LTI）SSM的高效卷积形式，接下来看看它的致命弱点。

LTI系统的核心假设是：核函数 $\bar{K}$ 对所有输入都相同。无论输入是"猫"还是"狗"，系统都用同一套指数衰减模式去处理。这在信号处理中没问题，但在语言建模中却是灾难。语言需要选择性：看到"但是"时应该关注转折，看到"因此"时应该关注因果。LTI SSM就像一个戴着固定滤镜的相机，无法根据画面内容调整焦距。

形式化地说，LTI SSM的 $\bar{A},\bar{B},\bar{C}$ 都是与输入无关的常数矩阵。这意味着模型没有能力执行内容依赖的上下文选择，比如从长文档中精准定位关键信息。

1.4.1.2 选择性参数化的数学形式

Mamba的突破在于：让离散化参数变成输入的函数。具体来说，步长 $\Delta$、输入矩阵 $B$ 和输出矩阵 $C$ 都依赖于当前输入 $x_k$：

$${\Delta }_k=\text{softplus}(W_{\Delta }{x}_k+b_{\Delta })$$
$$B_k=W_B{x}_k$$
$$C_k=W_C{x}_k$$

只有状态转移矩阵 $A$ 保持输入无关（通常是预定义的对角矩阵）。这样，每个token都有自己独特的离散化参数：

$${\bar{A}}_k=e^{{\Delta }_{k}A}$$
$${\bar{B}}_k=({\Delta }_{k}A{)}^{-1}(e^{{\Delta }_{k}A}-I)\cdot {\Delta }_k{B}_k$$

注意下标 $k$ 的出现：${\bar{A}}_k$ 和 ${\bar{B}}_k$ 不再是全局常数，而是每个位置都不相同。这使得系统不再是时不变的，核函数也无法预先计算。

听起来抽象对吧？如果画成图会是什么样子？想象一个流水线，每个工位（token）都有一个可调节的阀门 ${\Delta }_k$ 控制流速，还有一个分拣员 $B_k$ 决定把输入的哪些成分送进状态仓库，以及一个质检员 $C_k$ 决定从仓库中提取哪些成分作为输出。每个工位的阀门开度和分拣策略都取决于当前流经的物品是什么。

1.4.1.3 硬件感知的并行扫描算法

选择性参数化带来了一个工程难题：由于每个位置的 ${\bar{A}}_k$ 和 ${\bar{B}}_k$ 都不同，我们无法再用全局卷积来并行计算整个序列。我们只能回到递推形式：

$$h_k={\bar{A}}_k{h}_{k-1}+{\bar{B}}_k{x}_k$$

这看起来像是必须串行计算。但Mamba的核心工程贡献是：利用并行扫描（Parallel Scan，也称Associative Scan）在GPU上高效实现这个递推。

关键观察是：递推更新可以写成一种结合运算。定义二元运算：

$$(a_1,b_1)\oplus (a_2,b_2)=(a_2{a}_1,a_2{b}_1+b_2)$$

那么递推关系可以表示为：

$$h_k=\text{second}\left(({\bar{A}}_1,{\bar{B}}_1{x}_1)\oplus ({\bar{A}}_2,{\bar{B}}_2{x}_2)\oplus \cdots \oplus ({\bar{A}}_k,{\bar{B}}_k{x}_k)\right)$$

这个运算 $\oplus$ 满足结合律，因此可以使用树形归约（Tree Reduction）在 $O(\log L)$ 并行深度内完成。在继续之前，我们先看看这个结合律为什么成立。展开二元运算：

$$h_2={\bar{A}}_2{h}_1+{\bar{B}}_2{x}_2={\bar{A}}_2({\bar{A}}_1{h}_0+{\bar{B}}_1{x}_1)+{\bar{B}}_2{x}_2={\bar{A}}_2{\bar{A}}_1{h}_0+{\bar{A}}_2{\bar{B}}_1{x}_1+{\bar{B}}_2{x}_2$$

这与 $\oplus$ 的定义完全一致。通过结合律，我们可以把长度为 $L$ 的序列分治为两段，分别扫描后再合并，类似FFT的蝴蝶结构。

在实际使用中，Mamba的实现更进一步：它把矩阵 $A$ 限制为对角形式，这样 ${\bar{A}}_k$ 也是对角矩阵，所有运算都降维为逐元素操作。同时，利用现代GPU的共享内存（SRAM）进行分块计算，避免频繁访问高延迟的全局显存（HBM）。

1.5 Mamba架构完整拆解

1.5.1.1 输入投影与门控机制

现在我们已经了解了选择性扫描的数学原理，接下来看看Mamba在神经网络层面长什么样子。

Mamba块接收输入 $x\in {\mathbb{R}}^{L\times d}$，其中 $L$ 是序列长度，$d$ 是模型维度。第一步是把输入投影到两个分支：一个分支进入SSM核心，另一个分支作为门控信号。

$$x_{\text{ssm}}=W_{\text{in}}x+b_{\text{in}}$$
$$x_{\text{gate}}=W_{\text{gate}}x+b_{\text{gate}}$$

两个分支都通过线性投影得到。SSM分支的维度通常被扩展为 $d_{\text{inner}}$（例如 $2d$），以提供足够的表达能力。

1.5.1.2 SSM核心模块

SSM分支进入核心计算。首先，通过线性投影得到选择性参数：

$${\Delta }_k=\text{softplus}(W_{\Delta }{x}_{\text{ssm},k}+b_{\Delta })$$
$$B_k=W_B{x}_{\text{ssm},k}$$
$$C_k=W_C{x}_{\text{ssm},k}$$

这里 $B_k,C_k\in {\mathbb{R}}^{1\times N}$，${\Delta }_k\in \mathbb{R}$。然后执行选择性扫描：

$$h_k={\bar{A}}_k{h}_{k-1}+{\bar{B}}_k{x}_{\text{ssm},k}$$
$$y_{\text{ssm},k}=C_k{h}_k$$

其中 ${\bar{A}}_k=e^{{\Delta }_{k}A}$，${\bar{B}}_k$ 由双线性变换公式计算。注意 $A$ 是一个预定义的对角矩阵，通常通过HiPPO理论初始化，以确保状态能够记忆历史信息的有效压缩。

如果画成图会是什么样子？想象一个蓝色方框装着输入，里面分出两条绿色管道：一条管道上有三个可调节阀门（Delta、B、C），通向一个紫色仓库（状态），仓库的输出与门控信号相乘，最后汇入主通道。

1.5.1.3 残差连接与输出融合

SSM核心的输出与门控分支进行逐元素乘法（SiLU门控），然后通过输出投影回到原始维度：

$$y_k=W_{\text{out}}(\text{SiLU}(x_{\text{gate},k})\odot y_{\text{ssm},k})$$

最后加上残差连接：

$$\text{output}=y+x$$

整个Mamba块没有注意力层，没有MLP层，核心就是一个带有门控的选择性状态空间模块。这种极简设计使得Mamba在参数量和计算量上都比同等规模的Transformer更高效。

1.6 Mamba与Transformer的复杂度对比

1.6.1.1 计算复杂度：线性vs二次

现在我们已经了解了Mamba的完整前向流程，接下来把它和Transformer放在同一个擂台上比较。

Transformer的自注意力：对于序列长度 $L$ 和模型维度 $d$，计算注意力矩阵 $Q{K}^T$ 需要 $O(L^{2}d)$ 次操作。当 $L$ 翻倍时，计算量变为四倍。

Mamba的选择性扫描：每个token的处理包括：

• 线性投影得到 $\Delta ,B,C$：$O(dN)$
• 计算 ${\bar{A}}_k,{\bar{B}}_k$：$O(N)$（对角矩阵，逐元素）
• 状态更新 $h_k={\bar{A}}_k{h}_{k-1}+{\bar{B}}_k{x}_k$：$O(N)$
• 输出 $y_k=C_k{h}_k$：$O(N)$

每个token的总计算量为 $O(dN)$，整个序列的总计算量为 $O(LdN)$。这里 $N$ 是状态维度，通常远小于 $d$（例如 $N=16$ 而 $d=768$）。关键观察是：计算量随 $L$ 线性增长，而非二次增长。

形式化对比：

架构	计算复杂度	当L翻倍时
Transformer	$O(L^{2}d)$	计算量变为4倍
Mamba	$O(LdN)$	计算量变为2倍

如果画成图会是什么样子？想象两条曲线，一条是开口向上的抛物线（Transformer），一条是直线（Mamba）。在短序列时它们可能相交，但一旦序列长度超过某个阈值，抛物线会迅速甩开直线——只不过这次，直线是更优的那条。

1.6.1.2 内存瓶颈：状态vs KV-Cache

在继续之前，我们先看看推理阶段的内存对比。这是Mamba真正拉开差距的地方。

Transformer的KV-Cache：如前所述，KV-Cache占用 $O(Ld)$ 内存，且随序列长度持续增长。生成长文档时，KV-Cache会挤占模型参数本身的显存空间，最终成为上下文扩展的硬性瓶颈。

Mamba的状态缓存：Mamba在推理时只需要缓存当前状态 $h_k\in {\mathbb{R}}^N$。无论序列多长，内存占用恒定为 $O(N)$。生成第 $k+1$ 个token时，只需读取 $h_k$，计算 $h_{k+1}$，然后丢弃 $h_k$。不需要像Transformer那样保留所有历史KV向量。

形式化对比：

架构	推理内存	当L翻倍时
Transformer	$O(Ld)$	内存翻倍
Mamba	$O(N)$	内存不变

这种常数级内存占用是状态空间模型的本质优势。因为Mamba的下一个输出只依赖于当前状态，而不依赖于完整的历史序列，它天然满足马尔可夫性质。Transformer则因为全局注意力的存在，必须保留全部历史。

1.6.1.3 实际吞吐量的墙在哪里

听起来抽象对吧？如果画成图会是什么样子？想象两个水箱，Transformer的水箱底部有一根越来越粗的管子（KV-Cache），每加一升水就要换一根更粗的管子；Mamba的水箱底部只有一根固定粗细的管子，水通过即走，不留痕迹。

但在实际使用中，Mamba的墙并不在理论上，而在工程实现上。并行扫描虽然把串行深度从 $O(L)$ 降到了 $O(\log L)$，但它引入了额外的同步开销和内存访问模式。在短序列（例如 $L<2048$）时，高度优化的Transformer内核（如FlashAttention）可能仍然更快，因为矩阵乘法的硬件利用率远高于扫描操作。

Mamba的优势区间通常在长序列（$L>8192$）和自回归生成（逐个token解码）场景。在生成阶段，Transformer每步都要加载全部KV-Cache，而Mamba每步只更新一个状态向量。当 $L$ 达到十万级别时，这种差异会变得不可逾越。

1.7 接口与架构总览

1.7.1.1 模块接口对接图

现在我们已经走完了从连续时间ODE到选择性扫描、从卷积表示到复杂度对比的完整路径，接下来看看各个模块之间如何拼接。

Mamba系统对外暴露的接口很简单：输入侧接收"token嵌入序列"，输出侧返回"特征序列"。内部则通过线性投影、选择性参数生成、并行扫描和门控融合四个阶段完成转换。

1.7.1.2 协同设计图：模块交互与数据流

数据在Mamba块内部的流动是单向的，没有循环依赖。输入先分叉为SSM支路和门控支路，SSM支路经过选择性扫描产生内容相关的特征，门控支路提供动态权重，两者相乘后通过残差连接回到主路径。

1.7.1.3 完整三层架构总览

如果把Mamba系统看作一个完整服务，它可以清晰划分为数据层、模型层和计算层。数据层负责词元化和嵌入查找；模型层负责线性投影、选择性扫描和门控融合；计算层负责并行扫描的内核调度、内存分层管理和序列并行策略。

1.7.1.4 全章闭环

现在我们已经走完了从Transformer的二次瓶颈到状态空间模型的连续时间起源、从卷积表示到选择性扫描、从并行扫描算法到复杂度对比的完整路径。让我们在最后把闭环收紧。

Transformer的自注意力通过全局 pairwise 交互获得了强大的表达能力，但代价是 $O(L^2)$ 的计算复杂度和 $O(L)$ 的KV-Cache内存。这种设计在短文本上运行良好，但在长序列和流式生成场景下遇到了物理层面的瓶颈。

状态空间模型提供了一种截然不同的视角：把序列建模看作一个动态系统的演化过程。通过线性ODE的离散化，我们得到了一个可以用卷积高效实现的线性时不变系统。但纯粹的LTI系统缺乏内容选择性，无法根据输入动态调整记忆和关注模式。

Mamba通过让离散化参数 $\Delta ,B,C$ 成为输入的函数，打破了时不变性，同时通过硬件感知的并行扫描算法保持了计算效率。选择性扫描的每个token都有独立的阀门控制状态更新，使得模型既能记住长距离依赖，又能过滤无关噪声。

在复杂度层面，Mamba把计算从 $O(L^{2}d)$ 降到了 $O(LdN)$，把推理内存从 $O(Ld)$ 降到了 $O(N)$。这种线性复杂度和常数状态缓存，使得Mamba在长序列建模上具有根本性的架构优势。

这在训练与实际使用中意味着：当我们需要处理超长文档、实时流式数据或资源受限的边缘设备时，不应该默认使用Transformer。状态空间模型提供了一条绕过二次瓶颈和内存墙的新路径。通过把序列理解为动态系统的轨迹，我们不仅获得了效率，还获得了一种更贴近物理直觉的建模方式——状态在流动中演化，信息在选择中传递，而计算复杂度始终与序列长度保持线性关系。