外观
Attention 机制面试题
内容整理自学习笔记,仅供面试备考参考;不构成录用、培训或考试承诺。
1. 传统 Attention 存在哪些问题?
- 上下文长度受限:计算复杂度为 O(N²),序列长度增长导致计算量急剧增加
- 推理速度慢:需反复加载巨大的 KV Cache,内存开销大
- 显存占用高:长序列场景下 KV Cache 成为显存瓶颈
2. Attention 的优化方向
| 方向 | 思路 |
|---|---|
| 提升上下文长度 | 稀疏 Attention、线性 Attention、低秩 Self-Attention |
| 加速与减少内存 | FlashAttention(IO 优化)、Multi-Query Attention(KV 共享) |
| 改进多头机制 | Grouped-Query Attention 等变体 |
3. Attention 的主要变体
- 稀疏 Attention:引入稀疏偏置降低复杂度
- 线性化 Attention:解耦 Attention 矩阵与核特征图,以相反顺序计算实现线性复杂度
- 原型与内存压缩:减少 Query/Key-Value 对的数量以缩小 Attention 矩阵
- 低秩 Self-Attention:利用 Self-Attention 的低秩属性进行加速
- 先验 Attention:用先验 Attention 分布补充或替代标准 Attention
- 改进多头机制:探索不同的多头替代方案
4. Multi-Query Attention (MQA)
4.1 Multi-head Attention (MHA) 的问题
- 训练阶段:不影响训练速度
- 推理阶段:每个注意力头有独立的 K、V,需要反复加载巨大的 KV Cache,导致显存开销大,性能受限于内存带宽
4.2 MQA 原理
在所有注意力头上共享 Key 和 Value,仅保留独立的 Query。这大幅减少了 KV Cache 的大小,降低显存占用,提升推理速度。
4.3 MHA vs MQA
| 维度 | Multi-head Attention | Multi-Query Attention |
|---|---|---|
| Q/K/V | 每个头独立的 Q、K、V | 共享 K、V,独立 Q |
| KV Cache | 大 | 小(减少为原来的 1/h) |
| 推理速度 | 受限于内存 | 更快 |
| 训练速度 | 基准 | 相当 |
| 模型效果 | 基准 | 轻微损失 |
4.4 应用模型
Falcon、PaLM、ChatGLM2-6B 均使用 MQA,但实现细节有差异。
- Falcon:为保持参数量,将隐藏维度从 4096 增大至 4544,多余参数分给 Attention 和 FFN
- ChatGLM2:将 FFN 中间维度从 11008 增大至 13696,多余参数分给 FFN
5. Grouped-Query Attention (GQA)
原理:介于 MHA 和 MQA 之间的折中方案,将注意力头分组,每组共享 K 和 V。
应用模型:ChatGLM2、LLaMA2-34B/70B
6. FlashAttention
核心思路:利用分块 Softmax 等价替代传统 Softmax,通过 IO 感知的算法设计实现显存与速度优化。
关键技术:
- HBM vs SRAM:充分利用 GPU 的高速 SRAM,减少对低带宽 HBM 的读写
- 分块 Softmax:将大矩阵分块计算,避免完整 Attention 矩阵的显存占用
- 重计算(Recomputation):反向传播时重新计算中间结果,而非存储
- Kernel 融合:合并多个 CUDA Kernel 减少 IO 开销
优点:节约 HBM 显存,高效利用 SRAM,显著降低显存占用并提升速度。
应用模型:LLaMA、Falcon 等
7. 并行 Transformer Block
将传统串行的 Attention → FFN 结构改为并行执行,即将 Attention 和 FFN 的计算同时进行后再合并:
python
# 传统串行
x = x + Attention(LN(x))
x = x + FFN(LN(x))
# 并行
x = x + Attention(LN(x)) + FFN(LN(x))效果:训练速度提升约 15%。在 8B 参数规模时模型效果有轻微损失,在 62B 规模时无损失。
应用模型:Falcon、PaLM