稀疏注意力理论基础
阅读时间: 约 15 分钟
前置要求: Transformer 与注意力机制
本文介绍稀疏注意力的理论基础,包括为什么稀疏注意力有效、常见的稀疏模式以及 UCM 支持的稀疏算法概览。
1. 为什么需要稀疏注意力#
1.1 标准注意力的问题#
graph TB
subgraph problem["标准注意力的挑战"]
P1["O(n²) 计算复杂度"]
P2["O(n) KV Cache 内存"]
P3["长序列带宽瓶颈"]
end
subgraph impact["影响"]
I1["序列长度受限"]
I2["并发数减少"]
I3["推理延迟增加"]
end
P1 --> I1
P2 --> I2
P3 --> I3
具体数据:
- 4K 序列长度: 需要计算 16M 个注意力分数
- 32K 序列长度: 需要计算 1B 个注意力分数
- 128K 序列长度: 需要计算 16B 个注意力分数
1.2 稀疏注意力的核心洞察#
研究发现,注意力权重在实践中通常是高度稀疏的:
graph TB
subgraph observation["观察到的现象"]
O1["大多数注意力权重接近 0"]
O2["少数 token 获得大部分注意力"]
O3["存在可预测的注意力模式"]
end
subgraph conclusion["结论"]
C1["可以只计算重要的注意力"]
C2["可以只加载重要的 KV"]
C3["精度损失可控"]
end
O1 --> C1
O2 --> C2
O3 --> C3
2. 注意力稀疏模式#
2.1 局部注意力(Local Attention)#
相邻的 token 倾向于相互关注:
注意力矩阵示意(对角带状):
T1 T2 T3 T4 T5 T6 T7 T8
T1 [ 1 · · · · · · · ]
T2 [ 1 1 · · · · · · ]
T3 [ · 1 1 · · · · · ]
T4 [ · · 1 1 · · · · ]
T5 [ · · · 1 1 · · · ]
T6 [ · · · · 1 1 · · ]
T7 [ · · · · · 1 1 · ]
T8 [ · · · · · · 1 1 ]
局部窗口大小 = 2
2.2 Sink Token(锚点 Token)#
序列开始的 token 通常获得较高注意力:
T1 T2 T3 T4 T5 T6 T7 T8
T8 [ 1 · · · · · 1 1 ]
↑ ↑ ↑
Sink Token 局部 Token
2.3 语义锚点#
关键词和重要实体获得更多注意力:
输入: "北京是中国的首都,它有着悠久的历史"
注意力分布(对于"它"):
"北京" → 高
"是" → 低
"中国" → 中
"的" → 低
"首都" → 高
...
2.4 稀疏模式分类#
mindmap
root["稀疏注意力模式"]
static["静态模式"]
s1["固定窗口"]
s2["固定步长"]
s3["预定义模式"]
dynamic["动态模式"]
d1["Top-K 选择"]
d2["阈值过滤"]
d3["学习预测"]
hybrid["混合模式"]
h1["局部 + 全局"]
h2["Sink + 局部"]
h3["分层稀疏"]
3. 稀疏注意力的优化原理#
3.1 减少计算#
graph LR
subgraph full["标准注意力"]
F1["计算所有 n² 个分数"]
end
subgraph sparse["稀疏注意力"]
S1["只计算 k 个分数"]
S2["k << n²"]
end
full --> |"O(n²) → O(nk)"| sparse
3.2 减少内存读取#
graph TB
subgraph full["标准 KV 访问"]
F1["读取所有 n 个 KV"]
F2["带宽: O(n)"]
end
subgraph sparse["稀疏 KV 访问"]
S1["只读取 k 个 KV"]
S2["带宽: O(k)"]
end
full --> sparse
3.3 精度与效率权衡#
| 稀疏比例 |
计算节省 |
内存节省 |
精度影响 |
| 10% |
10x |
10x |
可能较大 |
| 30% |
3.3x |
3.3x |
通常可接受 |
| 50% |
2x |
2x |
较小 |
| 70% |
1.4x |
1.4x |
很小 |
4. UCM 稀疏算法概览#
4.1 算法分类#
graph TB
subgraph ucm_sparse["UCM 稀疏算法"]
subgraph retrieval["检索类"]
ESA["ESA
Essential Sparse Attention"]
GSA["GSA
Gather-Scatter Attention"]
end
subgraph ondevice["GPU 端"]
GSAOD["GSA On-Device
GPU 端检索"]
end
subgraph reuse["复用类"]
Blend["Blend
非前缀 KV 复用"]
end
subgraph consistency["一致性"]
KVStar["KVStar
多步一致性"]
end
subgraph position["位置"]
RERoPE["RERoPE
位置外推"]
end
end
4.2 算法对比#
| 算法 |
核心思想 |
适用场景 |
稀疏选择方式 |
| ESA |
检索重要 Block |
长上下文 |
CPU 端预计算 |
| GSA |
收集-分散注意力 |
通用 |
Prefetch 引擎 |
| GSA On-Device |
GPU 端检索 |
低延迟 |
Hash + Top-K |
| Blend |
非前缀复用 |
RAG、多轮对话 |
位置校正 |
| KVStar |
多步一致 |
推测解码 |
状态跟踪 |
| RERoPE |
位置外推 |
超长序列 |
RoPE 修正 |
5. Block 表示法#
5.1 为什么用 Block#
UCM 以 Block(而非单个 Token)为单位管理 KV Cache:
graph TB
subgraph token_level["Token 级别"]
T1["优点: 精细控制"]
T2["缺点: 开销大、碎片化"]
end
subgraph block_level["Block 级别"]
B1["优点: 批量操作、对齐友好"]
B2["缺点: 粒度较粗"]
end
block_level --> |"UCM 选择"| Winner["Block 级别"]
5.2 Block 表示的稀疏选择#
graph TB
subgraph select["Block 选择流程"]
Q["Query Block"]
K1["KV Block 1"]
K2["KV Block 2"]
K3["KV Block 3"]
Kn["KV Block n"]
Q --> |"相似度计算"| Score["相似度分数"]
K1 --> Score
K2 --> Score
K3 --> Score
Kn --> Score
Score --> |"Top-K"| Selected["选中 Block"]
end
5.3 Block 表示方法#
| 方法 |
说明 |
UCM 应用 |
| 均值池化 |
Block 内 KV 平均 |
ESA |
| Hash 编码 |
二进制编码 |
GSA On-Device |
| 首 Token |
使用首个 Token |
简单场景 |
| 加权 |
注意力加权平均 |
高精度场景 |
6. 稀疏注意力的挑战#
6.1 选择开销#
选择哪些 Token/Block 是重要的,本身也有计算开销:
graph TB
subgraph challenge["选择开销"]
C1["需要某种形式的全局信息"]
C2["选择算法本身消耗时间"]
C3["可能需要额外的数据结构"]
end
subgraph solution["UCM 解决方案"]
S1["预计算 Block 表示"]
S2["异步预取"]
S3["GPU 端快速检索"]
end
challenge --> solution
6.2 精度保证#
如何确保稀疏不会损害生成质量:
| 策略 |
说明 |
| 保留 Sink Token |
始终包含序列开始的 token |
| 保留局部窗口 |
始终包含最近的 token |
| 自适应稀疏比例 |
根据任务动态调整 |
| 质量监控 |
监控输出质量指标 |
7. 关键概念总结#
| 概念 |
说明 |
重要性 |
| 注意力稀疏性 |
权重集中在少数 token |
稀疏优化的理论基础 |
| 局部注意力 |
关注相邻 token |
保证局部一致性 |
| Sink Token |
序列开始的锚点 |
避免信息丢失 |
| Block 表示 |
块级别的稀疏选择 |
UCM 的管理单位 |
| Top-K 选择 |
选择最相关的 K 个 |
常用的稀疏方法 |
延伸阅读#