Flash-atten/Cross-atten/SSM/Mamba/Jamba技术演进
这部分知识的了解缘起于和同行业对大模型趋势的闲聊,引出了目前网络结构的进化和迭代。自己后续如果还吃算法的饭,那么对最新技术趋势的跟踪是必须的。
另外最近对于resnet代表的shortcut,感觉就是人类眼睛看东西的天然性,我们只是关注了动态的部分,对于保持不动的画面其实是不予以关注的;这对应网络要学习变动的量;对于动态的部分,动态的物体不同也会对我们本身大脑有不同的刺激,这其实就对应着attention。
对于transformer代表的大语言模型在处理文字时是一个字一个字的输入,其实也是对应着人类阅读或者说话,本质也是一个一个往外蹦;有时一句连续的蹦,也可以很好的对应目前的MTP技术。
本质上大模型和神经脑科学还是存在非常强劲的联系,这是一门类似医学的科学,有学界的知识和经验,也有know-how的驱动。
1、Sparse attention和Dense attention
1.1、 什么是 Attention(注意力)机制?
在传统序列模型(如 RNN 或 LSTM)中,模型处理文本时必须像人类看书一样,从左到右逐字读取。由于隐藏状态(Hidden State)的容量有限,随着句子变长,前面的信息会被后面的信息逐渐“冲淡”,这就是经典的长距离依赖问题(Long-Range Dependency)。
2017 年,Google 在论文 《Attention Is All You Need》 中彻底颠覆了这一模式,提出了 Self-Attention(自注意力) 机制。其核心思想是:消除时序的先后顺序,让序列中的每一个 Token(词/符号)在一步之内,都能直接与序列中的其他所有 Token 发生计算。
我们可以用一个通俗的“检索系统”比喻来理解它的数学抽象:
当一个词需要计算自己的上下文表示时,它会发出一个询问(Query),去匹配整个序列中所有词的特征标签(Key),计算出相关性权重(注意力分数),最后根据这个权重对所有词的实际内容(Value)进行加权求和。
1. 数学图解与计算流水线
Self-Attention 的核心计算公式非常简洁:
Attention(Q,K,V)=softmax(QKTdk)VAttention(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)=softmax(dkQKT)V
整个计算过程可以拆解为以下四个步骤:
- 线性映射(线性变换):
输入矩阵 XXX(大小为 N×dmodelN \times d_{model}N×dmodel,其中 NNN 为序列长度)通过三个不同的权重矩阵 WQ,WK,WVW_Q, W_K, W_VWQ,WK,WV 进行线性变换,得到三个矩阵:
- Query (QQQ): Q=XWQ(N×dk)Q = XW_Q \quad (N \times d_k)Q=XWQ(N×dk) —— “我要寻找什么”
- Key (KKK): K=XWK(N×dk)K = XW_K \quad (N \times d_k)K=XWK(N×dk) —— “我能提供什么以供匹配”
- Value (VVV): V=XWV(N×dv)V = XW_V \quad (N \times d_v)V=XWV(N×dv) —— “我的实际业务内容是什么”
- 计算点积注意力分数(Similarity):
计算 QQQ 和 KTK^TKT 的乘积,得到一个 N×NN \times NN×N 的矩阵。矩阵中的每一个元素 Sij=qi⋅kjS_{ij} = q_i \cdot k_jSij=qi⋅kj 代表了第 iii 个词和第 jjj 个词之间的原始相关性得分。 - 缩放与归一化(Scaled & Softmax):
- 缩放(Scale):将得分除以 dk\sqrt{d_k}dk(dkd_kdk 为向量维度)。这是为了防止在维度很高时,点积结果过大,导致进入 Softmax 后的梯度极其微弱(梯度消失)。
- 归一化(Softmax):在每一行上执行 Softmax 函数,将原始得分转化为 0∼10 \sim 10∼1 之间、且每行总和为 1 的概率分布。这个 N×NN \times NN×N 的矩阵就是注意力权重矩阵(Attention Map)。
- 加权求和(Aggregation):
将归一化后的注意力矩阵乘以 VVV。最终输出的矩阵(N×dvN \times d_vN×dv)中,每一个词的向量都融入了整个序列中与其相关的其他词的信息。
1.2、 Dense Attention(稠密注意力)的瓶颈
上述标准的 Self-Attention 就是 Dense Attention(稠密注意力)。之所以叫“稠密”,是因为每一个 Token 都要和序列中的其他所有 Token 计算一次点积。
1. 空间与时间复杂度分析
如果一个文本序列的长度是 NNN:
- QQQ 的大小是 N×dkN \times d_kN×dk,
- KTK^TKT 的大小是 dk×Nd_k \times Ndk×N,
- 它们相乘得到的注意力矩阵大小是 N×NN \times NN×N。
这意味着,无论是在计算量(FLOPs)还是显存占用上,Dense Attention 的复杂度都是 O(N2)\mathcal{O}(N^2)O(N2)(平方级复杂度)。
2. 灾难性的长文本瓶颈
- 当 N=1,024N = 1,024N=1,024(约1千字)时,N2≈106N^2 \approx 10^6N2≈106(百万级计算)。
- 当 N=128,000N = 128,000N=128,000(大模型长文本,如 128k 上下文)时,N2≈1.6×1010N^2 \approx 1.6 \times 10^{10}N2≈1.6×1010(百亿级计算与显存占用)。
由于显存随上下文长度呈平方级暴涨,标准的 Dense Attention 无法直接扩展到极长的文本(如整本书、长视频分析),这就逼迫学术界和工业界寻找替代方案,从而诞生了 Sparse Attention(稀疏注意力)。
1.3、 Sparse Attention(稀疏注意力)的崛起
Sparse Attention 的核心逻辑是:“我不需要看所有人,我只想看最重要的人。” 它通过人工设计或动态学习的方式,将原先 N×NN \times NN×N 稠密矩阵中大量的计算点“剪枝”掉(将其强制设为 0 或不计算),从而将计算复杂度从 O(N2)\mathcal{O}(N^2)O(N2) 降低到 O(NlogN)\mathcal{O}(N \log N)O(NlogN) 甚至 O(N)\mathcal{O}(N)O(N)(线性级)。
在发展过程中,主要演进出了以下几种经典的稀疏模式:
1.3.1. 局部/滑动窗口注意力 (Local / Sliding Window Attention)
- 做法:每个 Token 只关注自己左边和右边固定窗口大小(如 www 个)的邻居词。
- 效果:复杂度直接降为 O(N×w)\mathcal{O}(N \times w)O(N×w)。因为在语言中,靠得近的词通常关系最紧密。
1.3.2. 空洞/步长注意力 (Dilated / Strided Window Attention)
- 做法:类似于图像中的空洞卷积,Token 每隔 sss 个词看一个(例如只看第 2, 4, 6, 8 个词)。
- 效果:在不增加计算量的前提下,通过“跳跃式”查看,扩大了模型单层的全局视野。
1.3.3. 全局锚点注意力 (Global Attention / Global Tokens)
- 做法:由于前两种方法丢失了远距离的全局信息,稀疏模型通常会人为设定几个“特殊特权词”(如
[CLS]符号或序列开头的几个词)。所有普通 Token 都要和这些全局 Token 计算注意力,全局 Token 之间也互相计算。 - 效果:搭建了一条长距离信息传递的“高速公路”。
- (注:将滑动窗口、空洞和全局锚点三者结合的典型代表就是著名的 BigBird 和 Longformer 模型。)
1.3.4. 基于聚类/哈希的动态稀疏 (Routing / Reformer)
- 做法:不使用固定的几何窗口,而是使用局部敏感哈希(LSH)或聚类算法,在运行时动态找出向量空间中最接近的(即可能最相关的) KKK 个词进行注意力计算。
1.4、 Dense Attention vs Sparse Attention 深度对比
| 特性维度 | Dense Attention (稠密) | Sparse Attention (稀疏) |
|---|---|---|
| 计算与显存复杂度 | O(N2)\mathcal{O}(N^2)O(N2) (随长度呈平方暴涨) | O(N)\mathcal{O}(N)O(N) 或 O(NlogN)\mathcal{O}(N \log N)O(NlogN) (近线性) |
| 长文本支持能力 | 极差(受限于显存,通常在 4k-8k 遭遇瓶颈) | 极强(可以轻松扩展到 32k、128k 甚至更高) |
| 信息无损度 | 100% 无损(保留完整的全局交互) | 有损(存在信息漏看、长距离捕捉能力退化) |
| 硬件计算效率 | 极高(矩阵极其规整,完美契合 GPU 的 Tensor Core 密集乘法) | 较差 / 需定制优化(稀疏矩阵导致 GPU 算力利用率低,往往需要定制 Triton 算子) |
| 经典应用代表 | 标准 Transformer, GPT-4 基础版, LLaMA | Longformer, BigBird, Reformer, 以及混合架构中的 Mamba/Jamba |
1.5、 现代大模型的演进趋势
在工业界落地时,虽然 Sparse Attention 理论上复杂度低,但由于它破坏了矩阵的连续性,导致 GPU 矩阵乘法器(Tensor Core)经常跑不满,出现“理论上飞快,实际上卡顿”的尴尬局面。
为了解决这个问题,现代大模型(2024-2026年期间)走向了以下两个分化与融合的路线:
- FlashAttention 家族(Dense 的极致硬件优化):
大家发现 Dense 没问题,问题出在显存读写太慢。FlashAttention 通过在 GPU 的 SRAM(高速缓存)和 HBM(显存)之间进行分块计算,不改变 Dense 的数学原理,通过算法工程将计算速度提升数倍。这使得很多大模型在长文本时依然坚守 Dense Attention。 - 混合架构(Hybrid Systems,如 Jamba / MoE+SSM):
最新的长文本架构(如基于 Mamba 演进的 Hybrid 架构)选择将 Dense Attention 层 和 线性复杂度的选择性状态空间模型(Mamba / SSM 层) 或者 Sparse Attention 层 交替堆叠。通过这种方式,既用少量 Dense 层死死锁住完美的全局关联记忆,又用线性/稀疏层极大地拉长了上下文窗口,实现了性能与效率的终极平衡。 - 动态稀疏(DSA)的核心思想:
“先筛选,再计算(Select-then-Compute)”。在 Autoregressive 解码的每一步中,它会实时根据当前生成的 Query,去动态计算和评估谁才是最相关的历史 KV 块,只把最关键的前 KKK 个块载入 SRAM 参与 Attention 计算,其余的直接剪枝、不读不算。
1.6、 题外话:Cross Attention(交叉注意力)
Cross-Attention 是多模态模型和 Encoder-Decoder 架构的灵魂,用于让一个序列去“查阅”另一个序列的内容。
从数学机理与矩阵维度来说,假设有两个完全不同的序列:
- 序列 A(Query 来源):长度为 LAL_ALA,维度为 din_Ad_{in\_A}din_A。例如:Decoder 正在生成的文本、或者 Stable Diffusion 中的图像特征 latent。
- 序列 B(Key/Value 来源):长度为 LBL_BLB,维度为 din_Bd_{in\_B}din_B。例如:Encoder 提取的源语言特征、或者 CLIP 提取的文本 Prompt 嵌入。
其计算流程和严格的维度变化如下:
- 投影(Projection):
通过三组不同的可学习权重矩阵,将输入投影到相同的内部维度 dkd_kdk 和 dvd_vdv:
- Q=A⋅Wq(维度: [LA,din_A]×[din_A,dk]→[LA,dk])Q = A \cdot W_q \quad (\text{维度: } [L_A, d_{in\_A}] \times [d_{in\_A}, d_k] \rightarrow [L_A, d_k])Q=A⋅Wq(维度: [LA,din_A]×[din_A,dk]→[LA,dk])
- K=B⋅Wk(维度: [BB,din_B]×[din_B,dk]→[LB,dk])K = B \cdot W_k \quad (\text{维度: } [B_B, d_{in\_B}] \times [d_{in\_B}, d_k] \rightarrow [L_B, d_k])K=B⋅Wk(维度: [BB,din_B]×[din_B,dk]→[LB,dk])
- V=B⋅Wv(维度: [BB,din_B]×[din_B,dv]→[LB,dv])V = B \cdot W_v \quad (\text{维度: } [B_B, d_{in\_B}] \times [d_{in\_B}, d_v] \rightarrow [L_B, d_v])V=B⋅Wv(维度: [BB,din_B]×[din_B,dv]→[LB,dv])
- 注意力权重计算(Attention Map):
Score=Q⋅KTdk(维度: [LA,dk]×[dk,LB]→[LA,LB])\text{Score} = \frac{Q \cdot K^T}{\sqrt{d_k}} \quad (\text{维度: } [L_A, d_k] \times [d_k, L_B] \rightarrow [L_A, L_B])Score=dkQ⋅KT(维度: [LA,dk]×[dk,LB]→[LA,LB])
这个 Score\text{Score}Score 矩阵的第 (i,j)(i, j)(i,j) 个元素,代表了序列 A 的第 iii 个 Token 对序列 B 的第 jjj 个 Token 的关注程度。
3. Softmax 与加权输出:
Attention(Q,K,V)=softmax(Score)⋅V(维度: [LA,LB]×[LB,dv]→[LA,dv])\text{Attention}(Q,K,V) = \text{softmax}(\text{Score}) \cdot V \quad (\text{维度: } [L_A, L_B] \times [L_B, d_v] \rightarrow [L_A, d_v])Attention(Q,K,V)=softmax(Score)⋅V(维度: [LA,LB]×[LB,dv]→[LA,dv])
最终输出的序列长度与 序列 A 一致(LAL_ALA),但每个位置的特征已经融合了 序列 B 的信息。
- 典型应用:
Transformer Decoder: 在文本翻译时,用目标语言的 QQQ 去注意源语言的 K,VK, VK,V。
多模态融合(如 Stable Diffusion): 用文生图模型中图像特征的 QQQ,去注意文本提示词(Prompt)提取出的 K,VK, VK,V,从而让文本指引图像的生成。
2、linear attention的出师未捷
工业界在 2020 年到 2023 年期间对 Linear Attention 进行了疯狂的压榨与尝试,最终撞上了无法解决的性能南墙,才逼得研究人员另辟蹊径,最终分化出了 FlashAttention(既然改算法不行,那就死磕硬件优化)和 Mamba/Jamba(既然线性注意力记不住,那就改用控制论状态空间)这两条统治级的路线。
2.1、 Linear Attention 的提出背景:Transformer 的“阿喀琉斯之踵”
要理解 Linear Attention 的提出,必须回到 2020 年前后。当时 Transformer 已经在 NLP 领域确立了绝对的霸主地位,但所有科学家和工程师头上都悬着一把达摩克利斯之剑:O(N2)\mathcal{O}(N^2)O(N2) 的平方级计算与显存复杂度。
当时业界面临着极其残酷的现实背景:
2.1.1. 硬件算力的无形屏障
当年的显卡主要是 V100 和刚上市的 A100(80GB 显存版在当时是奢侈品)。
- 在标准 Transformer 中,序列长度只要超过 4K(约 3000 字),显存就会被爆发式增长的注意力矩阵直接撑爆(OOM)。
- 这导致当时所有的大模型(如早期的 GPT-3 175B、BERT)其上下文长度都极其短小(通常只有 512 到 2048 个 Token)。模型成了“高度近视眼”,根本看不了长文档,更别提分析代码库或书籍了。
2.1.2. 传统 RNN 并行化的彻底绝望
有人提出:“既然 Transformer 费显存,那退回经典的 LSTM/RNN 架构行不行?RNN 的显存是恒定的 O(1)\mathcal{O}(1)O(1)。”
- 答案是不行。因为 RNN 必须算完字 AAA 才能算字 BBB(前后死锁),完全无法利用 GPU 几千个核心进行并行训练。在动辄千亿参数的大模型时代,用 RNN 训练无异于用马车拉火箭,时间成本不可接受。
2.1.3. 唯一的数学突破口:矩阵乘法结合律
在这样的绝望背景下,2020 年,学术界(以 Linear Transformers, Performer, Linformer 等论文为代表)突然发现了一个数学上的“后门”:
“如果我们能想办法把标准 Attention 外层的 softmax\text{softmax}softmax 函数摘掉,整个公式就会变成纯粹的矩阵连乘。这样根据矩阵结合律,我们先算后两个矩阵相乘,就能把复杂度从 O(N2)\mathcal{O}(N^2)O(N2) 瞬间降到 O(N)\mathcal{O}(N)O(N)(线性级)!”
这就是 Linear Attention 诞生的历史背景:它背负着在“保持 Transformer 可并行训练特性的同时,把计算成本打到线性级”的救世主使命。
2.2、 为什么现在的演进趋势(2024-2026)抛弃了它?
既然 Linear Attention 理论上这么完美,为什么在如今的工业界落地趋势中,大家反而不提它了?因为它的数学假设在实际落地时遭遇了惨败。
2.2.1. 致命缺陷:摘掉 Softmax 导致“长文本健忘症”
标准 Attention 的 softmax\text{softmax}softmax 包含指数运算(exe^xex),它能起到超级放大器的作用——让相关的词权重无限接近 1,不相关的无限接近 0。
- Linear Attention 为了追求线性复杂度,用核函数(如 ϕ(x)=elu(x)+1\phi(x) = \text{elu}(x) + 1ϕ(x)=elu(x)+1)替代了 softmax\text{softmax}softmax。
- 结果导致模型的注意力变得极其平滑和模糊。随着文本拉长,历史信息在一个固定大小的中间矩阵(d×dd \times dd×d)里不断叠加、覆盖。
- 最终,模型在长文本下出现了严重的“健忘”,根本无法实现像“在 10 万字里精准找到某一个数字(大海捞针)”这样的高难度长文本任务。
2.2.2. 硬件尴尬:“理论上飞快,实际上卡顿”
Linear Attention 改变了矩阵相乘的顺序。虽然它的总计算量(FLOPs)变小了,但代价是矩阵形状变得极不规整。
英伟达的 Tensor Core 喜欢计算超大的、规整的方形矩阵乘法。Linear Attention 产生的大量瘦长条形状的向量乘法,导致 GPU 的硬件流水线经常跑不满(利用率极低)。在实际测试中,它往往比标准的 Transformer 还要慢。(主要还是效果的原因导致英伟达不愿意在硬件上适配)
2.3. Linear Attention 是如何逼出当前两大路线的?
工业界在 Linear Attention 上撞得头破血流之后,终于痛定思痛,在 2024-2026 年期间分化出了你在前文看到的两个终极方向:
[ Transformer 遇到 N² 瓶颈 ]
│
▼
[ 诞生先驱: Linear Attention ]
│ (遭遇长文本健忘、硬件不适配惨败)
┌───────────────────┴───────────────────┐
▼ ▼
【工程派: FlashAttention 路线】 【算法派: Mamba / Jamba 路线】
既然改算法(丢Softmax)会变傻, 既然线性注意力记不住,干脆彻底抛弃
那就死守 Dense Attention,纯粹从 Attention,用更高级的“选择性状态空间”
硬件 I/O 层面优化内存读写加速。 来实现真正的、有动态过滤的线性 RNN。
3、FlashAttention
在现代大语言模型(LLM)的演进过程中,随着上下文长度(Context Window)迈向 100K 甚至百万级别,标准的 Transformer 遭遇了巨大的计算瓶颈。为了打破这个瓶颈,业界诞生了两个极其重要的代表性技术:FlashAttention 和 Mamba。
有趣的是,它们解决长文本问题的思路完全不同:
- FlashAttention 属于工程派:它没有改变任何数学原理,纯粹通过极致的硬件加速,让标准的 Attention 跑得飞快。
- Mamba 属于算法派:它彻底改变了数学结构,用全新的状态空间模型(SSM)取代了 Attention,把计算复杂度降到了线性级。
一、 FlashAttention:极致的硬件“时间管理大师”
1. 核心解决的痛点:I/O 速度瓶颈(Memory-Bound)
在标准 Attention 的计算中,Softmax 需要产生一个 N×NN \times NN×N 的大矩阵(如前所述,100k 上下文会产生百亿级元素)。
在 GPU 芯片内部,计算核心(Tensor Core)的速度极快,但显存(HBM,高带宽显存)的读写速度相对很慢。
标准的计算流程是:
- 从 HBM 读取 Q,KQ, KQ,K,计算出 QKTQK^TQKT,写回 HBM;
- 从 HBM 读取 QKTQK^TQKT,计算 Softmax,写回 HBM;
- 从 HBM 读取 Softmax 结果和 VVV,计算最终输出,写回 HBM。
大量的中间结果在显存里反复“读·写·读·写”,导致 GPU 的计算核心大部分时间都在干等数据传输。这就是所谓的 I/O 瓶颈。
2. FlashAttention 是怎么做的?
FlashAttention 的核心思想是:“分块计算(Tiling),利用 GPU 内部的高速缓存(SRAM)。”
GPU 内部有一种容量极小(只有几百 KB)但速度快到飞起的缓存叫 SRAM。FlashAttention 重新设计了算法底层算子(利用 CUDA 编写):
- 不一次性算完:它把大矩阵切成一个一个的小方块(Tiles)。
- 在缓存里闭环:把一个小方块的 QQQ 和 KKK 加载到超快的 SRAM 里,在 SRAM 里直接把点积、Softmax 的局部缩放、以及与 VVV 的乘法全部算完。
- 在线更新(Online Softmax):Softmax 原本需要全局最大值和分母和。FlashAttention 引入了一个数学技巧,可以在不看到完整序列的情况下,通过动态维护局部统计量,边分块计算边修正 Softmax 的值。
3. 最终效果
由于中间那张庞大的 N×NN \times NN×N 注意力矩阵自始至终都没有被完整写入慢速的 HBM 显存,只是在超快的 SRAM 内部闪烁了一下就输出了最终结果。核心的提升点也在于该巨型矩阵只在SRAM内部计算,不需要到HBM反复读写。
- 显存占用:从 O(N2)\mathcal{O}(N^2)O(N2) 骤降到 O(N)\mathcal{O}(N)O(N)。
- 速度:速度提升了 2 到 4 倍。
- 无损:由于只是改变了计算的硬件顺序,其数学结果与标准 Attention 100% 完全一致。目前(2026年)几乎所有主流 Transformer 模型(如 LLaMA、GPT-4)的底层都在默认使用 FlashAttention 家族。
4、 Mamba:打破 Transformer 统治的选择性状态空间模型
如果说 FlashAttention 是在旧铁路上把火车开到了极致,那么 Mamba 就是直接换了一条高铁轨道。
4.1. 核心突破:从 O(N2)\mathcal{O}(N^2)O(N2) 到 O(N)\mathcal{O}(N)O(N) 的跨越
传统 Attention 无论怎么工程优化,每个词看全场的数学本质没变。而 Mamba 引入了选择性状态空间模型(Selective SSM),它的计算复杂度天生就是 O(N)\mathcal{O}(N)O(N)(线性级)。无论文本多长,每增加一个词,计算量只线性增加一点点。
4.2. Mamba 的底层原理:从数学到硬件的飞跃
Mamba 的核心思想源于状态空间模型(State Space Model, SSM)。传统的 SSM 结合了 RNN 的低推理开销(O(1)\mathcal{O}(1)O(1))与 Transformer 的高效并行训练(O(N)\mathcal{O}(N)O(N))优势。而 Mamba 在此基础上,通过引入“选择性”和“硬件感知算法”,彻底解决了传统 SSM 的性能瓶颈。
我们可以从数学演进到硬件实现来逐层拆解:
4.2.1. 什么是“连续系统”?(物理直觉)
在没有大模型之前,状态空间模型(SSM)是自动化控制工程(比如阿波罗登月飞船的导航系统)的核心算法。它的本质是:如何用一盘不断变化的“录像带”(连续隐藏状态 h(t)h(t)h(t)),去实时记录和预测世界的变化。
我们来看它的连续状态方程:
h′(t)=Ah(t)+Bx(t)h'(t) = Ah(t) + Bx(t)h′(t)=Ah(t)+Bx(t)
- x(t)x(t)x(t) 是当前的外界输入(比如飞船传感器的实时气流)。
- h(t)h(t)h(t) 是飞船当前的内部记忆状态(蓄水池)。
- h′(t)h'(t)h′(t) 是状态的变化率(导数)。
- 矩阵 AAA 的微观含义:它是整个系统的“遗忘因子”。它决定了如果没有新输入时,历史记忆以多快的速度衰减。
- 矩阵 BBB 的微观含义:它是系统的“吸纳因子”。它决定了当前的新输入 x(t)x(t)x(t) 有多大比例能转换成记忆。
为什么要“离散化”?
飞船在大自然中面对的是连续的时间轴(0.0001秒、0.0002秒……),但计算机处理文本时是一字一字的(Token 1, Token 2, Token 3……),时间是断开的。
所以,我们必须把连续的微分方程,切成一步一步的离散递推公式。
离散化的微观细节(零阶保持 ZOH)
为了把连续变成离散,模型引入了一个至关重要的参数:步长 Δ\DeltaΔ(Delta)。
你可以把 Δ\DeltaΔ 理解为相机的快门曝光时间。
通过数学转化(零阶保持法),连续的 AAA 和 BBB 被公式重塑成了计算机认识的 Aˉ\bar{A}Aˉ 和 Bˉ\bar{B}Bˉ:
Aˉ=exp(ΔA)\bar{A} = \exp(\Delta A)Aˉ=exp(ΔA)
Bˉ=(ΔA)−1(exp(ΔA)−I)⋅ΔB\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta BBˉ=(ΔA)−1(exp(ΔA)−I)⋅ΔB
变成了我们最终看到的类似 RNN 的循环公式:
ht=Aˉht−1+Bˉxth_t = \bar{A}h_{t-1} + \bar{B}x_tht=Aˉht−1+Bˉxt
- 对齐直觉:如果快门时间 Δ\DeltaΔ 很大,意味着时间过去了很久,旧的记忆 Aˉ=exp(ΔA)\bar{A} = \exp(\Delta A)Aˉ=exp(ΔA) 就会衰减得非常厉害;新输入 Bˉ\bar{B}Bˉ 占的比重就会变大。
4.2,2. Mamba 的核心创新:选择性机制(Selective SSM)
上一节讲的经典 S4 模型有一个致命缺陷:它是“时间不变”的(Time-Invariant)。
这意味着,在整个模型的训练和推理过程中,矩阵 AAA、BBB 以及步长 Δ\DeltaΔ 都是固定死的参数。无论进来的词是什么,模型的快门时间和记忆衰减率都一模一样。
为什么固定参数会变傻?
假设我们给模型输入两句话:
- 句子甲:“Joshua 今天去买了一杯咖啡,然后他……”
- 句子乙:“Joshua 的银行卡密码是 883921,然后他……”
在经典 S4 模型里,因为 Δ\DeltaΔ 和 BBB 是固定的,模型处理“然后”这个废话,和处理“密码是 883921”这个核心信息的态度完全一样,这就会导致有价值的信息在传递几百个词后被彻底稀释。
Mamba 的选择性(Selectivity)改造:让参数变成输入的函数
Mamba 的突破就在于:它让 BBB、CCC 和 Δ\DeltaΔ 变成了由当前输入的词 xtx_txt 实时算出来的动态变量。
当每一个 Token xtx_txt 进入 Mamba 层时,内部会发生如下微观演变:
- 动态生成 BtB_tBt 和 Δt\Delta_tΔt:
Bt=LinearB(xt),Δt=LinearΔ(xt)B_t = \text{Linear}_B(x_t), \quad \Delta_t = \text{Linear}_{\Delta}(x_t)Bt=LinearB(xt),Δt=LinearΔ(xt)
每个词进来,先过一个极其轻量的线性层,为自己量身定制一套当下的系统参数。
2. 微观博弈(记还是忘?):
- 当读到废话“然后”时,线性层故意输出一个极大的 Δt\Delta_tΔt。带入离散化公式,Aˉt=exp(极大×A)\bar{A}_t = \exp(\text{极大} \times A)Aˉt=exp(极大×A),因为 AAA 是负数,Aˉt\bar{A}_tAˉt 瞬间坍塌接近于 0。
公式变成:ht≈0⋅ht−1+Bˉtxth_t \approx 0 \cdot h_{t-1} + \bar{B}_t x_tht≈0⋅ht−1+Bˉtxt。前面的历史记忆直接被清空,腾出空间! - 当读到关键隐私“密码是 883921”时,线性层输出一个极小的 Δt\Delta_tΔt。此时 Aˉt=exp(极小×A)≈1\bar{A}_t = \exp(\text{极小} \times A) \approx 1Aˉt=exp(极小×A)≈1。
公式变成:ht≈1⋅ht−1+Bˉtxth_t \approx 1 \cdot h_{t-1} + \bar{B}_t x_tht≈1⋅ht−1+Bˉtxt。历史记忆被牢牢锁死,完美继承!
总结:Mamba 的选择性机制,让模型拥有了像人类一样的“专注力”,能够根据当下的文本内容,实时调节记忆蓄水池的闸门大小。
4.2.3. 破局工程:硬件感知并行扫描(Hardware-aware Parallel Scan)
有了选择性机制后,Mamba 遇到了一个极其痛苦的工程灾难:无法并行训练。
在传统 RNN 里,因为 ht=Aˉtht−1+Bˉtxth_t = \bar{A}_t h_{t-1} + \bar{B}_t x_tht=Aˉtht−1+Bˉtxt,算第 3 个词必须等第 2 个词的 h2h_2h2 算完,算第 2 个词必须等 h1h_1h1。这种前后死锁的依赖,导致 GPU 的几千个计算核心无法同时开工,训练速度慢到无法接受。
(注:经典的 S4 能并行是因为它的参数是固定的,可以用数学转换成全局卷积;而 Mamba 引入了选择性,参数每个时间步都在变,卷积这条路被彻底堵死了。)
为了拯救训练速度,Mamba 启动了硬件感知并行扫描(Parallel Scan)。
工程细节:内存阶梯与内核融合(Kernel Fusion)
在英伟达 GPU 内部,存在严重的“贫富差距”:
- HBM(普通显存):容量很大(如 80GB),但读写速度极其缓慢。
- SRAM(高速缓存):容量极小(只有几百KB),但读写速度快到飞起。
Mamba 团队意识到,RNN 并行慢,很大程度上是因为每个时间步计算出的隐藏状态 hhh(大小为:Batch_Size × 序列长度 × 通道维度 × 状态维度)太庞大了。如果每算一个词,就把 hhh 写入 HBM,下一个词再从 HBM 读出来,GPU 核心绝大部分时间都在等显存传输数据(Memory-Bound)。
Mamba 的 CUDA 底层优化做了三件事:
- 绝不离开高速缓存:
在训练开始时,把整个序列的长文本 XXX 以及基础参数,一次性加载到超快的 SRAM 缓存中。 - 利用“关联扫描(Parallel Scan)”在 SRAM 内部强行并行:
虽然公式看似是串行的,但在数学上,状态的传递满足结合律。
例如:h4=Aˉ4(Aˉ3h2+Bˉ3x3)+Bˉ4x4h_4 = \bar{A}_4(\bar{A}_3 h_2 + \bar{B}_3 x_3) + \bar{B}_4 x_4h4=Aˉ4(Aˉ3h2+Bˉ3x3)+Bˉ4x4。
GPU 可以启动树状并行算法(像淘汰赛一样两两配对合并),在 SRAM 内部,将原本需要跑 64 步的串行循环,在 log2(64)=6\log_2(64) = 6log2(64)=6 步内,强行并行并行算完。 - 内存降维打击(不保存隐状态):
在扫描计算的过程中,中间产生的那个遮天蔽日的超大状态矩阵 hhh,自始至终在 SRAM 内部快速闪烁、迭代,绝对不往慢速的 HBM 显存里写一个字节。当一整个块算完,拿到最终的输出 YYY 后,把 hhh 直接在缓存里抹去(反向传播时再重新在线计算)。
对齐直觉:硬件感知并行扫描,就是把所有干活的工人和原材料关在一个极其狭小但速度极快的全封闭车间(SRAM)里,用最高效的并行算法把活干完,只把成品运出工厂。 从而彻底甩掉了显存读写的拖累,让流水线速度暴涨。
4.3. 关联扫描怎么做的
关联扫描(Parallel Scan,也称 Prefixes Sum 或 Scan)是 Mamba 能够既像 RNN 一样进行“选择性记忆”,又像 Transformer 一样进行“大规模并行训练”的幕后功臣。
很多人感到困惑:RNN 后一个词的计算明明依赖前一个词的结果(ht=f(ht−1)h_t = f(h_{t-1})ht=f(ht−1)),这怎么可能像矩阵乘法一样同时并行呢?
秘诀在于数学上的结合律(Associativity)。Mamba 巧妙地将“前后死锁的递归”转化成了“可以像打淘汰赛一样同时计算的二叉树”。
下面我们彻底拆解它的微观运转机制。
4.3.1、 核心数学前提:打通“结合律”
我们先来看 Mamba 的离散递归公式(这里为了方便直观理解,我们简化掉 BBB 和 xxx,只看状态 hhh 和衰减因子 AAA 的传递):
h1=A1h0+x1h_1 = A_1 h_0 + x_1h1=A1h0+x1
h2=A2h1+x2h_2 = A_2 h_1 + x_2h2=A2h1+x2
h3=A3h2+x3h_3 = A_3 h_2 + x_3h3=A3h2+x3
h4=A4h3+x4h_4 = A_4 h_3 + x_4h4=A4h3+x4
如果我们要算 h4h_4h4,按照传统的串行逻辑,必须先算 h1→h2→h3→h4h_1 \rightarrow h_2 \rightarrow h_3 \rightarrow h_4h1→h2→h3→h4。
但如果我们把前面的公式强行展开、代入呢?
h4=A4(A3(A2(A1h0+x1)+x2)+x3)+x4h_4 = A_4 \Big( A_3 \big( A_2(A_1 h_0 + x_1) + x_2 \big) + x_3 \Big) + x_4h4=A4(A3(A2(A1h0+x1)+x2)+x3)+x4
把括号全部拆开,重新组合:
h4=(A4A3A2A1)h0+(A4A3A2)x1+(A4A3)x2+A4x3+x4h_4 = (A_4 A_3 A_2 A_1) h_0 + (A_4 A_3 A_2) x_1 + (A_4 A_3) x_2 + A_4 x_3 + x_4h4=(A4A3A2A1)h0+(A4A3A2)x1+(A4A3)x2+A4x3+x4
注意看!在这个展开的公式中,所有的 AAA 之间都是纯粹的乘法。而在数学中,乘法是满足结合律的——也就是说,计算 (A4×A3)×A2(A_4 \times A_3) \times A_2(A4×A3)×A2 和计算 A4×(A3×A2)A_4 \times (A_3 \times A_2)A4×(A3×A2) 的结果完全一样!
既然计算顺序可以被任意调换,我们就不必傻傻地从左到右挨个算了,我们可以利用 GPU 的几千个核心“两两成对”地同时算。
4.3.2、 算法图解:两阶段的“二叉树”并行
为了利用结合律,关联扫描算法(最经典的是 Blelloch Scan)在 GPU 内部的 SRAM(高速缓存) 里将计算分为了两个阶段:上推阶段(Up-Sweep)和下归阶段(Down-Sweep)。
假设我们现在有一个长度为 8 的文本序列,我们要并行算出所有位置的隐藏状态。
第一阶段:上推(Up-Sweep)构建二叉树
这一阶段的目的是由底向上,把局部的乘积算出来,就像世界杯淘汰赛一样。
- 第 1 步(并行):
核心 1 计算:A1×A2A_1 \times A_2A1×A2 的中间结果
核心 2 计算:A3×A4A_3 \times A_4A3×A4 的中间结果
核心 3 计算:A5×A6A_5 \times A_6A5×A6 的中间结果
核心 4 计算:A7×A8A_7 \times A_8A7×A8 的中间结果
*(这 4 个计算是完全同时(并行)*发生的!) - 第 2 步(并行):
核心 1 把上一轮算好的 (A1A2)(A_1 A_2)(A1A2) 和 (A3A4)(A_3 A_4)(A3A4) 相乘,得到 (A1A2A3A4)(A_1 A_2 A_3 A_4)(A1A2A3A4)。
核心 2 把上一轮算好的 (A5A6)(A_5 A_6)(A5A6) 和 (A7A8)(A_7 A_8)(A7A8) 相乘,得到 (A5A6A7A8)(A_5 A_6 A_7 A_8)(A5A6A7A8)。 - 第 3 步:
把两个大块相乘,最终得到整句话的全局总乘积。
到此为止,树顶拿到了全局信息,但我们的目标是拿到每一个时间步的独立状态(比如 h1,h2,h3h_1, h_2, h_3h1,h2,h3 各自的值)。于是进入第二阶段。
第二阶段:下归(Down-Sweep)分发状态
这一阶段就像是“族谱向下分发家产”,从树顶向下,利用前面留在二叉树各个节点上的中间结果,并行地把每一个位置的最终状态“拼凑”出来。
由于树的高度只有 log2(N)\log_2(N)log2(N),对于一个长度为 N=8N=8N=8 的序列:
- 传统 RNN 串行需要走 8 步。
- 关联扫描并行只需要走 $\log_2(8) \times 2 = $ 6 步。
如果序列长度 N=1024N = 1024N=1024:
- 传统 RNN 串行需要走 1024 步。
- 关联扫描并行只需要走 $\log_2(1024) \times 2 = $ 20 步!
文本越长,这种成百上千倍的并行加速优势就越恐怖。
4.3.3、 Mamba 算子层面的终极魔鬼细节
如果你去看 Mamba 的论文或 CUDA 源代码,你会发现它比上述通用的二叉树扫描还要更极致,因为它必须要解决“硬件开销”和“选择性(数据在变)”的问题。
1. 块状扫描(Chunk-based Scan)与寄存器白嫖
在实际的 GPU 运算中,如果真的对 100k 长度建立一个完整的超大二叉树,树节点的临时内存开销会把超快的 SRAM 缓存挤爆。
所以 Mamba 实际运行的是混合扫描:
- 它把 100k 长的序列切成若干个大小为 64 或 128 的 Chunk(小块)。
- 块内串行:在每个线程内部,利用 GPU 的 Registers(寄存器,比 SRAM 还要快一个数量级的终极芯片核心),把这 64 个词用极快的速度串行扫完,算出这个块的“局部总和”和“最终状态”。
- 块间并行扫描:把每个块的最终状态当成基础元素,在整个 GPU 的数百个线程块(Thread Blocks)之间,跑上面讲的二叉树关联扫描。
这样既利用了寄存器的极致速度,又甩掉了大树的内存开销。
2. 算子融合(Kernel Fusion)不留痕迹
我们之前提到,Mamba 的衰减参数 AtA_tAt 是通过当前的词 xtx_txt 动态算出来的(Aˉt=exp(ΔtA)\bar{A}_t = \exp(\Delta_t A)Aˉt=exp(ΔtA))。
在普通的深度学习代码(如 PyTorch)中,如果你写:Delta = linear(X) →\rightarrow→ A_bar = exp(Delta * A) →\rightarrow→ h = scan(A_bar)
每一次箭头的转化,GPU 都需要把数据从 HBM(大显存)倒腾到计算核心,再写回 HBM,这会产生灾难性的读写延迟。
Mamba 的做法是:将“动态生成参数”和“二叉树关联扫描”全部写进同一个 CUDA 内核(Kernel)里。
当数据 XXX 被载入 SRAM 的那一刻起,它在当下就原地完成了“线性投影 →\rightarrow→ 算出 Δt\Delta_tΔt →\rightarrow→ 离散化出 Aˉt\bar{A}_tAˉt →\rightarrow→ 跑二叉树扫描”的全部动作。在整个并行的流水线中,中间变量自始至终没有在显存里留下过任何蛛丝马迹,算完即焚,只把最终的 YYY 输出来。
这就是 Mamba 能够在训练速度上和 FlashAttention 掰手腕的底层工程底气。
4.4、 最新研究进展与演进路线
大模型架构经历了激烈的迭代,针对纯 Mamba 架构的局限性,学术界和工业界取得了突破性的最新进展:
1. 混合架构(Hybrid Models)成为绝对主流
纯 Mamba 架构在进行“精确信息检索(如 Needle In A Haystack 大海捞针测试)”和“多步程序化逻辑推理”时,能力略逊于标准的 Attention。因此,最新的研究彻底倒向了 Hybrid 架构:
- Jamba 1.5 & Samba:AI21 Labs 等机构将 Mamba 层、Attention 层与 MoE(混合专家) 进行了深度融合。通常采用 8:1 或 4:1 的比例,即每隔几层 Mamba 才插入一层 Attention。这样既保留了 Attention 100% 的硬核逻辑和位置记忆能力,又将整体 KV Cache 空间压缩了 80% 以上。
- 注意力与状态空间的等价性研究:最新理论证明,经过特定设计的选择性 SSM 在数学上可以被视为一种带有特定衰减因子的线性注意力(Linear Attention)。这为两种架构的相互融合提供了坚实的理论支撑。
2. Mamba-2 的架构升级
Mamba 的原作者(Albert Gu 和 Tri Dao)发布了 Mamba-2。这一重大升级核心解决了 Mamba-1 依然存在的硬件加速痛点:
- 结构化状态空间双重性(SSD):Mamba-2 提出了一个全新的理论框架,将 SSM 和注意力机制紧密结合。
- 对 Tensor Core 的完美适配:Mamba-1 的并行扫描算子虽然快,但它利用的是 GPU 的通用计算单元(ALU),没有完全白嫖到英伟达专为矩阵乘法设计的 Tensor Core(张量核心)。Mamba-2 重新设计了状态矩阵结构(将其从原本的 1 维通道阻塞改为了块状矩阵乘法形式),使得模型训练可以调用 Tensor Core。
- 速度再次飞跃:在相同参数规模下,Mamba-2 的训练速度比 Mamba-1 快了 2 到 8 倍,彻底抹平了与 FlashAttention-2 在大型集群上的训练效率差距。
3. 硬件与基础设施生态的繁荣
- Trident & Custom Operators:主流推理框架(如 vLLM、TensorRT-LLM)已经开始原生支持 Mamba 及混合架构的算子优化,打通了 Mamba 在云端大规模上线的最后一步。
- 非 Transformer 原生芯片的崛起:一些非 GPU 的新型 AI 芯片(如神经拟态芯片、数据流架构芯片)发现,Mamba 这种免去全局 KV Cache 缓存的算法更加契合它们的硬件设计,正在引发新一轮的软硬件协同创新。
5. 什么是线形注意力
- 从“广义线性注意力”的数学定义来看,你说得完全正确:RNN 确实在数学形式上与线性注意力高度统一。 2. “线性注意力”在今天(2026年)的学术界,绝对不仅仅指当年那个死板的、摘掉 Softmax 的
Linear Attention论文算法,它已经演变成了一个庞大的“线性时序建模”广义家族(Generalized Linear Attention)。
5.1、 为什么说 RNN 属于广义的线性注意力?(数学对偶性)
在 2020 年之前,大家觉得 RNN 和 Transformer 是两个完全对立的门派:
- RNN 是串行的、通过一个固定大小的隐状态 hth_tht 像滚雪球一样往前递推。
- Transformer 是并行的、通过一张大表(KV Cache)让每个词看全场。
但是,随着 Katharopoulos 等人在 2020 年提出线性注意力,以及后来 Mamba-2 证明的 SSD(结构化状态空间双重性) 理论,学术界达成了一个终极共识:RNN 和 线性注意力,其实是同一个数学公式在不同计算顺序下的“一体两面”(对偶性,Duality)。
我们用最直观的矩阵维度和乘法顺序来证明这一点:
设序列长度为 NNN,特征维度为 ddd。我们有三个矩阵 Q,K,VQ, K, VQ,K,V(它们的维度都是 N×dN \times dN×d)。
视角 A:从上往下算(广义线性注意力视角)
如果我们去掉标准 Attention 的 Softmax,利用矩阵乘法的结合律,先计算 KT×VK^T \times VKT×V:
M=KT×V(维度是 d×d)M = K^T \times V \quad (\text{维度是 } d \times d)M=KT×V(维度是 d×d)
然后再用 QQQ 去乘这个中间矩阵 MMM:
Output=Q×M(维度是 N×d)\text{Output} = Q \times M \quad (\text{维度是 } N \times d)Output=Q×M(维度是 N×d)
因为计算 MMM 时,KTK^TKT 是 d×Nd \times Nd×N,VVV 是 N×dN \times dN×d,相乘的复杂度是 O(N⋅d2)\mathcal{O}(N \cdot d^2)O(N⋅d2)。随着文本长度 NNN 变长,计算量呈线性级 O(N)\mathcal{O}(N)O(N) 增长。这就是线性注意力。
视角 B:从左往右看(RNN 视角)
现在我们把视角切换到 RNN。RNN 的核心是:当前隐状态 = 旧隐状态 + 新输入。
我们把上面线性注意力的中间矩阵 MMM(大小为 d×dd \times dd×d)拆开,看它随着时间步 ttt 是如何演进的。你会惊奇地发现:
Mt=Mt−1+ktT⋅vtM_t = M_{t-1} + k_t^T \cdot v_tMt=Mt−1+ktT⋅vt
这不就是一个标准的 RNN 隐状态更新方程吗?
- MtM_tMt 就是 RNN 的隐藏状态(Hidden State)。
- ktT⋅vtk_t^T \cdot v_tktT⋅vt 就是当前时间步的新输入(Rank-1 更新)。
- 当前步的输出则是:yt=qt⋅Mty_t = q_t \cdot M_tyt=qt⋅Mt。
结论
- 当你选择一次性把所有序列的 KT×VK^T \times VKT×V 算完,再乘以 QQQ 时,它就是线性注意力(可并行训练)。
- 当你选择一字一字地把 ktT⋅vtk_t^T \cdot v_tktT⋅vt 累加进 MMM 矩阵,每一步用 qtq_tqt 去读取时,它就是 RNN(低能耗流式推理)。
因此,在广义上,所有线性级复杂度的时序模型(包括传统 RNN、LSTM、RWKV、Mamba),在数学本质上都可以被归为“广义线性注意力(Generalized Linear Attention)”这个大族谱里。
5.2. 线性注意力绝不仅仅是“Linear Attention”
你提到的第二个观点非常高明:线性注意力不是单单指那一种特定算法。 在 2020 年 《Transformers are RNNs》 论文刚提出时,大家管那种“直接摘掉 Softmax、用普通激活函数 ϕ(Q)ϕ(K)T\phi(Q)\phi(K)^Tϕ(Q)ϕ(K)T 替代”的特定算法叫 Linear Attention。那个初版算法由于对历史信息的分配过于平滑(没有指数放大),导致大模型得了“严重健忘症”,在工业界流产了。
但是,“线性注意力”作为一种底层设计思想(即:利用结合律把复杂度降为 O(N)\mathcal{O}(N)O(N)),在随后的几年里经历了极其疯狂的打补丁与进化演学。目前的广义线性注意力家族包括以下几大流派:
5.2.1. 动态衰减派(以 RWKV 为代表)
既然初版线性注意力容易遗忘,RWKV(Receptive Weighted Key Value)引入了基于时间的指数衰减因子。
它在 Mt=Mt−1+ktTvtM_t = M_{t-1} + k_t^T v_tMt=Mt−1+ktTvt 的基础上,强行让旧记忆乘以一个不断衰减的权重 www。这就使得模型能够分清“远近亲疏”,距离现在越近的词权重越高,极大地缓解了健忘症。
5.2.2. 状态空间派(以 Mamba-2 / SSD 为代表)
Mamba-2 提出的 Structured State Space Duality (SSD),是广义线性注意力演进的巅峰。
它不再像 RWKV 那样人工设计死板的衰减因子,而是将控制论里的“选择性时序控制”融合了进来。Mamba-2 在块(Chunk)内部的计算,数学形式上就是一种高度结构化、带有输入相关动态衰减掩码(Input-dependent Decay Mask)的强力线性注意力机制。
5.2.3. 硬件感知派(以 FlashLinear / GLA 为代表)
最新的研究(如 Segment-Based Linear Attention)不再纠结于完全干掉 Transformer,而是借鉴了 FlashAttention 的分块思想,在硬件寄存器层面重新设计线性注意力的形状,让其既有 RNN 恒定显存的优势,又能吃满英伟达的 Tensor Core。
5.3. 终极对齐:为什么我们之前说 Mamba“不属于”线性注意力?
既然它们在广义上是一家亲,为什么在上一轮回答或者很多技术博客里,大家依然习惯性地说“Mamba 不属于线性注意力”呢?
这其实是学术界语境(数学本质)与工业界语境(技术起源)的划分问题:
- 工业界/技术起源语境:当我们说“Mamba 不属于线性注意力”时,是在强调 Mamba 的亲生父亲是控制论中的“微分方程/状态空间模型(SSM)”,它不是靠魔改 Transformer 的代码、抠掉 Softmax 诞生出来的。它们的出发点、参数设计和物理意义完全不同。
- 学术界/数学本质语境:当我们把视距拉到最终的矩阵计算时,Mamba-2 的作者自己扯下了这层外衣,证明了 SSM 经过特定约束后,就是线性注意力。
5.4. 总结
传统的、广义的 RNN 确实是线性注意力在时序递推状态下的另一种写法。
- 线性注意力(Linear Attention) 已经从 2020 年的一个特定算法名字,演变成了今天一整个涵盖了 RWKV、Mamba-2、GLA 等众多前沿模型的“广义线性时序架构家族”。
- 它们共同的魂魄,就是那条完美的、甩掉了 Softmax 霸权的数学公式:O(N)\mathcal{O}(N)O(N) 线性复杂度。
6、 Softmax 的“全局分母”和矩阵结合律
在标准的 Attention 计算中,公式如下:
Attention(Q,K,V)=softmax(QKTdk)VAttention(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)=softmax(dkQKT)V
我们假设输入文本的序列长度是 NNN(例如 10,000 个词),特征维度是 ddd(例如 128)。
- Q∈RN×dQ \in \mathbb{R}^{N \times d}Q∈RN×d
- K∈RN×d ⟹ KT∈Rd×NK \in \mathbb{R}^{N \times d} \implies K^T \in \mathbb{R}^{d \times N}K∈RN×d⟹KT∈Rd×N
- V∈RN×dV \in \mathbb{R}^{N \times d}V∈RN×d
6.1. 纯矩阵乘法的结合律
如果没有外层的 softmax\text{softmax}softmax,公式只是三个矩阵单纯相连:Q×KT×VQ \times K^T \times VQ×KT×V。
根据矩阵乘法的结合律,我们有两种计算顺序:
-
顺序甲(先算左边):(Q×KT)×V(Q \times K^T) \times V(Q×KT)×V
-
第一步:(N×d)×(d×N)→(N \times d) \times (d \times N) \rightarrow(N×d)×(d×N)→ 产生 N×NN \times NN×N 的巨型中间矩阵。
-
计算复杂度:O(N2⋅d)\mathcal{O}(N^2 \cdot d)O(N2⋅d)。
-
顺序乙(先算右边):Q×(KT×V)Q \times (K^T \times V)Q×(KT×V)
-
第一步:(d×N)×(N×d)→(d \times N) \times (N \times d) \rightarrow(d×N)×(N×d)→ 产生 d×dd \times dd×d 的极小中间矩阵(与文本长度 NNN 无关)。
-
第二步:(N×d)×(d×d)→(N \times d) \times (d \times d) \rightarrow(N×d)×(d×d)→ 最终输出。
-
计算复杂度:O(N⋅d2)\mathcal{O}(N \cdot d^2)O(N⋅d2)。
6.2. Softmax 带来的“物理锁死”
然而,softmax\text{softmax}softmax 函数直接把“顺序乙”这条完美的线性逃生通道给炸毁了。
我们来看 Softmax 的数学公式。对于 QKTQ K^TQKT 产生的矩阵中第 iii 行、第 jjj 列的某一个原始得分 SijS_{ij}Sij,经过 Softmax 映射后的最终注意力权重 AijA_{ij}Aij 是:
Aij=eSij∑k=1NeSikA_{ij} = \frac{e^{S_{ij}}}{\sum_{k=1}^{N} e^{S_{ik}}}Aij=∑k=1NeSikeSij
注意看分母的那个求和符号 ∑k=1N\sum_{k=1}^{N}∑k=1N。
这个分母意味着:第 iii 行的某一个元素要想完成归一化,必须知道这一行所有其他元素的指数总和。
这就形成了一个极其残酷的计算依赖:
- 你不能只算一个局部。
- 你不能把 QQQ 和 KKK 拆开。
- 你必须强行把 Q×KTQ \times K^TQ×KT 的完整矩阵(大小为 N×NN \times NN×N)全部显式地计算出来。 只有让每一个元素都暴露出自己的原始得分,这一行才能闭环算出那个全局分母。
因为 softmax(QKT)V≠Q×(softmax(KT)V)\text{softmax}(Q K^T) V \neq Q \times (\text{softmax}(K^T) V)softmax(QKT)V=Q×(softmax(KT)V),非线性的 Softmax 就像一把挂在括号上的大锁,迫使大模型必须先执行 (Q×KT)(Q \times K^T)(Q×KT),从而将计算和显存复杂度死死钉在了 O(N2)\mathcal{O}(N^2)O(N2)。
6.3、 空间的无情吞噬:为什么显存也会发生平方爆炸?
Softmax 不仅带来了计算量(FLOPs)的平方增长,更致命的是它带来了硬件显存(Memory)的平方暴涨。
在传统 Transformer 计算中:
- 你的显卡计算核心(Tensor Core)吐出了 Q×KTQ \times K^TQ×KT 的原始结果(N×NN \times NN×N 矩阵)。
- 为了执行接下来的 Softmax 操作,这个包含了 N2N^2N2 个浮点数的巨型矩阵,必须被实体化(Instantiated)并写入 GPU 的普通显存(HBM)中。
当 N=1,024N=1,024N=1,024 时,1024×1024≈1001024 \times 1024 \approx 1001024×1024≈100 万个元素,显存毫无压力。
当 N=128,000N=128,000N=128,000(现代大模型长文本的标配)时,128,000×128,000≈163128,000 \times 128,000 \approx \mathbf{163}128,000×128,000≈163 亿个元素!
如果使用 FP16(半精度浮点数,每个元素占 2 字节),单单存储这一个注意力矩阵,就需要消耗:
16.384×109×2 Bytes≈32.7 GB16.384 \times 10^9 \times 2 \text{ Bytes} \approx \mathbf{32.7 \text{ GB}}16.384×109×2 Bytes≈32.7 GB
这还没算上模型的参数、激活值以及多头注意力(Multi-Head)的倍数!一张大显存的 A100 显卡(80GB)瞬间就会被这一个 Softmax 算子彻底挤爆(Out of Memory)。
6.4、 Softmax 霸权
- 因为有 Softmax:大模型拥有了无与伦比的“超级过滤器”,它通过 exe^xex 的强力指数放大,能够让模型在 10 万字里精准捕捉到两三个关键词之间的硬核逻辑关联。
- 代价则是:你必须接受硬件层面上无法逃避的 O(N2)\mathcal{O}(N^2)O(N2) 算力与显存惩罚。
后来诞生的一切技术(如 FlashAttention 是通过硬件缓存分块数学修正来隐式化这个 N×NN \times NN×N 矩阵以节省显存;而 Mamba/线性注意力 是直接用别的非线性函数平替掉 Softmax 以解锁矩阵结合律),全都是为了造反这个“平方复杂度”的物理枷锁。
7、DSA算法层的极致优化
Dynamic Sparse Attention(DSA,动态稀疏注意力) 是当前大模型在突破超长上下文(Long-Context)和端侧推理瓶颈时,技术演进最激进、最前沿的绝对硬核方向。
7.1. 为什么必须要从“静态稀疏”走向“动态稀疏”?
- 静态稀疏(如滑动窗口、固定 Block 稀疏)的痛点:它们假设模型对历史 Token 的关注度是随距离衰减的,或者只关注固定位置(比如开头几个 Token)。但实际大模型在处理复杂推理、大海捞针(Needle in a Haystack)等长文本任务时,注意力的分布是完全由当前的 Query 决定的(Content-Dependent)。
- 动态稀疏(DSA)的核心思想:“先筛选,再计算(Select-then-Compute)”。在 Autoregressive 解码的每一步中,它会实时根据当前生成的 Query,去动态计算和评估谁才是最相关的历史 KV 块,只把最关键的前 KKK 个块载入 SRAM 参与 Attention 计算,其余的直接剪枝、不读不算。
7.2. 动态稀疏注意力(DSA)的通用工作三阶段
为了在硬件上跑得通、跑得快,DSA 并没有傻傻地去算完 O(N2)O(N^2)O(N2) 再做 Top-K,而是采用了一种分层的巧妙设计:
阶段一:分块与元数据抽象 (Blocking & Metadata)
- 物理显存中的 KV Cache 被划分为固定大小的物理页或块(例如每 32 个 Token 为一个 Block)。
- 针对每个 Block,实时维护或预先算好一个极轻量级的元数据(Metadata)结构(通常不到 full KV 尺寸的 1%)。这个元数据可以是 Key 向量的均值/方差、高低维投影,或是每个维度的物理边界(Bounding Box)。
阶段二:Query 驱动的动态筛选 (Select)
- 当前输入一个新 Query 时,模型首先让这个 Query 仅与这些轻量级的元数据进行快速低成本的匹配计算,以此来估计每个 Block 的潜在最高注意力得分。
- 过滤出得分最高的 Top-KKK 个关键物理块(Blocks)。
阶段三:精确注意力计算 (Compute)
- 只把这 Top-KKK 个被选中的 KV 块从显存(HBM)中捞到 GPU 的片上缓存(SRAM)中。
- 配合定制的硬件算子(如基于 FlashAttention 改造的动态稀疏算子),仅对这些块进行精准的 Softmax 和 Value 聚合计算。
7.3. 目前主流的 Dynamic Sparse Attention 算法路线
根据“如何高效挑出那 Top-KKK 个块”,业界衍生出了三大流派:
- 启发式几何边界流派:Quest
- 原理:由 MIT 实验室等提出的 Quest 算法。它在 KV Cache 页面中实时跟踪 Key 值的最大值(Max)和最小值(Min)。当 Query 到来时,利用 Query 向量与 Key 向量的 Bounding Box 几何关系,快速推导出一个注意力分数的数学上限(Upper Bound)。
- 特点:完全不需要训练(Training-free),纯靠数学逻辑在线筛选。在 128K 甚至更长上下文中,能只捞取极少比例的 KV Cache 页就保持近乎无损的准确度。
- 聚类检索流派:ClusterKV
- 原理:在 Prefill(预填充)阶段,利用 K-means 等聚类算法对长文本中所有的 Key 向量进行在线聚类,抽象出若干个中心点(Centroids)。在 Decode(解码)阶段,Query 只需要和这几个中心点做比对,就能瞬间定位到关联的 Token 簇。
- 模型原生/轻量索引流派:DeepSeek DSA & RTPurbo
- 原理:直接在模型架构中引入一个原生的、极轻量级的 Indexer(索引器)模块(例如 DeepSeek 或者是近期研究如 RTPurbo 的做法)。这个 Indexer 通常只有很低的维度(如 16 维),专门在运行时通过小规模参数动态预测 Token 的重要性,直接驱动硬件的 Top-K 硬件路由。
- 特点:属于模型原生支持的稀疏化,软硬件协同极高,但需要少量的微调或联合训练。
7.4. 动态稀疏的系统级“死穴”与最新突破
DSA 虽然在数学上完美把计算复杂度砍了下来,但在系统底层落地时曾遭遇过致命的系统瓶颈:
瓶颈:序列化依赖与显存碎片化
- 选择-计算依赖锁死:以前的 Attention 是直接读内存开算,而 DSA 必须要“先算完 Selection,知道要哪个块,才能去读 KV Cache 计算 Attention”。这个动态查找过程直接变成了推理关键路径上的延迟(Latency)瓶颈。
- 硬件不友好:动态挑出来的块在物理显存里往往是高度离散、碎片化、随机分布的。这会导致 GPU 发生严重的非连续内存读取(Non-coalesced Memory Access),直接把显存带宽卡死。
目前像 vLLM、SGLang 等前沿推理框架正在通过软硬件深度共轨来消解这些瓶颈:
- 投机注意力(Speculative Attention / PRR 机制):利用注意力选择在时间上的稳定特性(即这一步关注的块,下一步大概率也关注),在 Selection 还没算完时,直接投机性预取(Speculative Prefetch)上一轮的块进入缓存,实现计算与访存的完美流水线重叠(Overlap)。
- 层次化稀疏(DHSA / Dynamic Hierarchical Sparse Attention):引入二级架构,先在 Chunk 级别做粗粒度动态预测,再在 Token 级别做精细过滤,大幅缓解硬件层面的碎片化读取痛点,将端侧大模型(如 Gemma-2B/3B 在手机端)的 100K 长文本处理能力直接推向实用。
7.5. 核心生产指标表现
根据最新生产级测试,在硬件(如 H100 / A100 架构)上运行长文本任务时,DSA 结合软硬件协同优化带来的提升非常可观:
- 吞吐量(Throughput)提升:在处理超过 64K 到 128K 的长文本检索(如长文 RAG、多轮对话代码生成)时,DSA 相比于 vanilla vLLM 可以提升 1.4倍 至 3.1倍 的系统吞吐量。
- 首字延迟(TTFT - Time to First Token)暴降:在并发请求率高且 Prompt 极长的复杂场景下,由于 DSA 的“Select-then-Compute”机制在 Prefill(预填充)阶段大幅削减了不必要的 KV 块读写,其平均首字延迟最多可降低 9倍(源自 SparseServe 2025/2026 最新压测数据)。
- 显存占用(KV Cache Memory)显著缩减:配合 Progressive(渐进式)或阈值过滤策略,DSA 可以安全地剪枝掉高达 60% - 80% 的非必要注意力计算,让单卡能够容纳更大的 Batch Size。
7.6. 生产落地的“隐性代价”(什么时候效果不明显?)
尽管指标惊人,但在真实的生产线部署(Production Deployment)中,DSA 存在几个不容忽视的局限性:
① 短文本场景:纯粹的负优化
如果你的线上业务 90% 以上的 Prompt 都在 2K 或 4K Token 以内,请直接关闭 DSA。
因为在短文本下,标准 Attention 的 KV Cache 读写本来就很小,GPU 根本不卡带宽。此时开启 DSA,反而需要让 Query 额外跑一遍 Indexer(索引器)或聚类算法,这增加了额外的计算延迟,会导致 TTFT 和每秒生成 Token 数(TPOT)不降反升。
② 精度损失的“长尾效应”
动态稀疏本质上是一种近似计算(Approximation)。
- 在普通的总结、润色、闲聊任务中,由于注意力集中在头部和尾部,DSA 几乎能做到 0 精度损失。
- 但在极端复杂的长文本代码 Debug、密集的数学推导、或者超长文本中多处关键信息的“多跳推理”(Multi-hop Reasoning)场景下,DSA 的 Top-K 偶尔会漏掉某些隐藏在角落里的关键 Token,从而引发模型胡言乱语(幻觉率上升)。生产中通常需要动态调大 KKK 值来对齐精度,而这会消耗一部分性能。
③ 硬件不友好与显存碎片化
动态挑出来的 KV 块在物理显存里通常是高度离散、随机分布的。如果底层的推理引擎(比如早期的 PagedAttention 变体)没有专门针对非连续内存读取做内核级优化(Non-coalesced Memory Access),那么“动态查找”所带来的开销甚至会直接抵消掉稀疏化带来的红利。
7.7. 生产落地的最佳实践建议
如果你们团队正在考虑在生产网关中引入 DSA,可以参考目前大厂的主流配置策略:
- 动态路由机制(Dynamic Routing):
在 API Gateway 层做冷热分流。判断输入的 Prompt 长度,如果< 16K,走标准高效的 GQA + FlashAttention 路由;如果> 32K,动态切换到 GQA + Dynamic Sparse Attention + FlashAttention。 - 软硬件深度绑定:
不要尝试自己裸写 PyTorch 实现。生产中必须依赖如 vLLM(其已正式对 DeepSeek DSA 提供原生支持及优化)、SGLang 或 SparseServe 这种针对 CUDA Kernel、Tiling 机制和 HBM-SRAM 数据移动做过深度融合重写的专有引擎。 - 首选模型原生 DSA:
像原生就带有低维 Indexer 设计的模型,在 DSA 下的精度和吞吐表现,远远好于在标准 MHA 模型上生搬硬套 Quest 等“外挂式”算法。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐



所有评论(0)