MLA(Multi-head Latent Attention,多头潜在注意力)
概述
传统 MHA 每个头都要独立缓存 K 和 V,KV Cache 大小是 O(L⋅h⋅dk)O(L \cdot h \cdot d_k)O(L⋅h⋅dk)。
MLA 的核心思想是:不直接缓存每个头的 K、V,而是把所有头的 K、V 联合压缩到一个低维潜在向量 ccc,推理时只缓存 ccc,需要时再上投影还原。
公式:
c=WDKV⋅hc = W_{DKV} \cdot hc=WDKV⋅h
K=WUK⋅c,V=WUV⋅cK = W_{UK} \cdot c, \quad V = W_{UV} \cdot cK=WUK⋅c,V=WUV⋅c
其中:
- hhh 是输入隐藏状态;
- WDKVW_{DKV}WDKV 是下投影矩阵,把高维 hhh 压到低维 rrr;
- WUK,WUVW_{UK}, W_{UV}WUK,WUV 是上投影矩阵,把 ccc 还原成 K、V;
- r≪h⋅dkr \ll h \cdot d_kr≪h⋅dk,所以缓存大幅减少。
缓存复杂度从 O(L⋅h⋅dk)O(L \cdot h \cdot d_k)O(L⋅h⋅dk) 降到 O(L⋅r)O(L \cdot r)O(L⋅r)。
数字例子
设:
text
输入维度 d = 8
头数 h = 2
每个头 d_k = d_v = 4
潜在维度 r = 2
序列长度 L = 3
MHA 的 KV Cache:
每个头独立缓存 K、V,共:
2×L×h×dk=2×3×2×4=48 个数值 2 \times L \times h \times d_k = 2 \times 3 \times 2 \times 4 = 48 \text{ 个数值} 2×L×h×dk=2×3×2×4=48 个数值
MLA 的 KV Cache:
只缓存潜在向量 ccc,共:
L×r=3×2=6 个数值 L \times r = 3 \times 2 = 6 \text{ 个数值} L×r=3×2=6 个数值
缓存减少到原来的 6/48=12.5%6/48 = 12.5\%6/48=12.5%,即减少约 87.5%87.5\%87.5%。
具体计算:
假设输入 h=1,0,1,0,1,0,1,0h = 1, 0, 1, 0, 1, 0, 1, 0h=1,0,1,0,1,0,1,0。
下投影矩阵 WDKVW_{DKV}WDKV 形状 8×28 \times 28×2,设为:
text
W_DKV = [[1, 0],
[0, 1],
[1, 0],
[0, 1],
[1, 0],
[0, 1],
[1, 0],
[0, 1]]
则:
c=WDKVTh=1+1+1+1,0+0+0+0=4,0 c = W_{DKV}^T h = 1+1+1+1, 0+0+0+0 = 4, 0 c=WDKVTh=1+1+1+1,0+0+0+0=4,0
缓存 c=4,0c = 4, 0c=4,0。
上投影还原:
设 WUKW_{UK}WUK 形状 2×42 \times 42×4(还原到每个头的 K 维度),
text
W_UK = [[1, 0, 0, 0],
[0, 1, 0, 0]]
则:
K=WUK⋅c=4,0,0,0K = W_{UK} \cdot c = 4, 0, 0, 0K=WUK⋅c=4,0,0,0
V 同理。这样每个头共享同一个 ccc,但通过不同的上投影矩阵得到不同的 K、V。
权衡
- 优点:KV Cache 大幅减少,显存节省 90%+;
- 缺点:每次注意力计算前要上投影解压缩,计算量增加;
- 结论:显存比计算时间更宝贵,所以 MLA 值得。
DSA(DeepSeek Sparse Attention,稀疏注意力)
概述
传统注意力每个 Query 都要和所有历史 token 计算注意力,复杂度 O(L2)O(L^2)O(L2)。
DSA 的核心思想是:先快速筛选出最重要的 Top-k 个历史 token,只对这些 token 做注意力计算。
两个核心组件:
- Lightning Indexer(闪电索引器):轻量级网络,用少量注意力头、FP8 低精度,快速扫描所有历史 token,计算每个 token 的重要性分数。
- 细粒度 Top-k 选择:根据分数,为每个 Query 动态选出 Top-k 个 token(通常 k=2048),只对这些 token 做注意力。
复杂度从 O(L2)O(L^2)O(L2) 降到 O(L⋅k)O(L \cdot k)O(L⋅k),与序列长度解耦。
数值例子
设序列长度 L=6L = 6L=6,Top-k 中 k=3k = 3k=3。
当前 Query 是位置 5 的 token。历史 token 是位置 0~5。
Lightning Indexer 打分:
假设索引器给出每个历史 token 的重要性分数:
text
位置: 0 1 2 3 4 5
分数: 0.1 0.8 0.3 0.9 0.2 0.7
Top-k 选择:
选分数最高的 3 个:
text
位置 3: 0.9
位置 1: 0.8
位置 5: 0.7
核心注意力:
只对位置 3、1、5 的 K、V 做注意力计算,其他位置忽略。
假设 Q5 与这三个 K 的点积为:
text
Q5·K3 = 2
Q5·K1 = 1
Q5·K5 = 3
缩放后 softmax 得到权重,再对 V3、V1、V5 加权求和。
对比全注意力:
- 全注意力:Query 5 要与 6 个 token 计算,复杂度 O(6);
- DSA:只与 3 个 token 计算,复杂度 O(3。
序列越长,节省越明显。当 L=64kL=64k,k=2048k=2048 时,计算量减少约 32 倍。
特点
- DSA 是插件式的,可以后训练加到已有模型上;
- 长上下文推理成本降低 60-70%,质量几乎不损失。
DeepSeek V4 中的 CSA 和 HCA 混合注意力架构
DeepSeek V4 以 MQA 为基底,用 CSA + HCA 交替,再加滑动窗口捕捉局部信息。
CSA(Compressed Sparse Attention,压缩稀疏注意力)
设计哲学:先压缩,再稀疏选择,保留精细信息。
第一步:KV 缓存压缩
每 mmm 个 token 的 KV 压缩成一个块级表征。不是简单平均,而是可学习的加权压缩:
cj=∑i∈blockjwi⋅KVi c_j = \sum_{i \in \text{block}_j} w_i \cdot KV_i cj=i∈blockj∑wi⋅KVi
其中 wiw_iwi 是模型学习的压缩权重。相邻压缩块共享部分 token,平滑边界。
压缩后 KV 规模降为原来的 1/m1/m1/m。V4 中 m=4m = 4m=4。
数值例子:
设 m=4m = 4m=4,原始 4 个 token 的 KV 值:
text
KV0 = [1, 0]
KV1 = [0, 1]
KV2 = [1, 1]
KV3 = [0, 0]
压缩权重:
text
w0 = 0.4, w1 = 0.3, w2 = 0.2, w3 = 0.1
压缩后:
c=0.4⋅1,0+0.3⋅0,1+0.2⋅1,1+0.1⋅0,0 c = 0.4 \cdot 1, 0 + 0.3 \cdot 0, 1 + 0.2 \cdot 1, 1 + 0.1 \cdot 0, 0 c=0.4⋅1,0+0.3⋅0,1+0.2⋅1,1+0.1⋅0,0
=0.4+0.2,0.3+0.2=0.6,0.5 = 0.4 + 0.2, 0.3 + 0.2 = 0.6, 0.5 =0.4+0.2,0.3+0.2=0.6,0.5
原来 4 个 KV 对,现在只剩 1 个压缩条目。
第二步:DSA 稀疏选择
压缩后,用 Lightning Indexer 快速计算当前 Query 与所有压缩块的相关性分数,只保留 Top-k 个压缩块。
V4 中 Flash 的 k=512k = 512k=512,Pro 的 k=1024k = 1024k=1024。
数值例子:
压缩后有 3 个块,分数:
text
块0: 0.9
块1: 0.3
块2: 0.7
Top-2 选择:块0 和块2。
核心注意力只在块0、块2 上计算,计算量从 O(L)O(L)O(L) 降到 O(k)O(k)O(k)。
CSA 层执行流程
text
KV缓存 → 可学习加权压缩 → Lightning Indexer 低精度打分 → Top-k 选择 → 核心注意力
HCA(Heavily Compressed Attention,重度压缩注意力)
目标:极低成本维护覆盖十万级 token 的全局背景。
HCA 与 CSA 类似,但只做压缩,不做稀疏选择。
- 压缩率 mmm 比 CSA 大得多(比如 m=64m = 64m=64 或更大);
- 压缩后序列极短,可以在短序列上做密集注意力;
- 每个 token 都能无遗漏地看到整个全局背景;
- 因为序列短,计算成本完全可控。
数值例子:
设原始序列 L=64L = 64L=64,压缩率 m=64m = 64m=64。
压缩后只有 1 个全局条目:
c=∑i=063wi⋅KVi c = \sum_{i=0}^{63} w_i \cdot KV_i c=i=0∑63wi⋅KVi
所有 Query 都只与这 1 个全局条目做注意力,计算量从 O(642)O(64^2)O(642) 降到 O(64×1)O(64 \times 1)O(64×1)。
与 CSA 的分工:
| 机制 | 压缩率 | 稀疏选择 | 作用 |
|---|---|---|---|
| CSA | 小 (m=4) | 有 Top-k | 保留精细信息 |
| HCA | 大 (m=64+) | 无 | 极低成本全局背景 |
| 滑动窗口 | - | - | 捕捉最临时的局部信息 |
三者交替配合,兼顾局部细节、精细全局和低成本全局。
总结
| 机制 | 核心思想 | 复杂度 | 数值例子结论 |
|---|---|---|---|
| MLA | 所有头共享低维潜在向量 c,只缓存 c | 缓存 O(L⋅r)O(L \cdot r)O(L⋅r) | 缓存从 48 降到 6,减少 87.5% |
| DSA | Lightning Indexer 打分 + Top-k 选择 | O(L⋅k)O(L \cdot k)O(L⋅k) | 6 个 token 只算 3 个 |
| CSA | 可学习加权压缩 + DSA 稀疏选择 | O(L/m⋅k)O(L/m \cdot k)O(L/m⋅k) | 4 个 KV 压成 1 个,再选 Top-2 |
| HCA | 重度压缩,密集注意力 | O(L/m)O(L/m)O(L/m) | 64 个 token 压成 1 个全局条目 |
一句话: MLA 用低秩压缩减少 KV Cache;DSA 用闪电索引器和 Top-k 选择减少注意力计算;CSA 先压缩再稀疏选择,保留精细信息;HCA 重度压缩,极低成本维护全局背景。三者配合,让 DeepSeek V4 在超长上下文下依然高效。