MHA、MQA、GQA 都是深度学习中注意力机制的相关概念,其中 MHA 是标准的多头注意力机制,MQA 和 GQA 则是其优化变体,以下是它们的区别、优缺点介绍:
区别
- MHA(多头注意力):是 Transformer 架构的核心组件,会将输入的查询(Query)、键(Key)和值(Value)分别投影到多个不同的头中,每个头都拥有专属的 Q、K、V 映射参数,它们独立计算注意力,然后将多个头的结果拼接后投影回原始维度。
- MQA(多查询注意力):属于 MHA 的优化方案,其所有的查询头共享同一组 K、V 参数,每个头依然保持独立的 Q,以此在保持查询头多样性的同时,大幅减少内存占用。
- GQA(分组查询注意力):是 MHA 和 MQA 的折中方案,它先把查询头划分成若干个组,每组配备独立的 K、V 参数,而组内所有查询头共享该组的 K、V 来计算注意力,最后将头部输出拼接后经线性层整合。
优缺点
- MHA
- 优点:能让模型在不同的子空间中学习不同的注意力模式,提高了模型的表达能力,可使模型关注输入序列的不同部分,并且支持并行计算,能在 GPU 上高效执行。
- 缺点:每个查询头都需要与多个键计算注意力,计算量较大,并且还需要存储多个 Q、K、V 头,大模型中显存占用较多。
- MQA
- 优点:由于所有查询头共享 K、V,极大减少了计算成本,尤其适用于推理阶段,也降低了内存占用,对处理超长文本较为友好。
- 缺点:仅有一个 K、V 组,可能会让模型的表达能力受到影响,难以很好区分不同查询头的语义信息,上下文信息捕捉能力有所下降。
- GQA
- 优点:相比 MHA,减少了键值计算的开销,计算效率有所提升,同时在视觉 Transformer 等计算机视觉任务里也能表现不错,可降低视觉 Transformer 在图像数据上的计算复杂度。
- 缺点:因查询头共享 K、V,会让模型在灵活性上有一定损失,相对标准 MHA,其表达能力可能会降低。
备注:
在 MQA(多查询注意力)中,K、V 的维度不会除以头数 H 。
在标准的 MHA(多头注意力)里,假设模型隐藏层维度是 dmodel,注意力头数为 H,通常会将 dmodel 均匀划分到各个头中,所以每个头的 K、V 维度是 dmodel/H,K、V 的张量形状是 (batch_size, seq_len, H, dmodel/H) 。但 MQA 中所有查询头共享同一组 K、V 参数,其 Key 与 Value 线性层输出维度为 (batch_size, seq_len, dmodel),也就是不再需要按头数对维度进行切分,K、V 本身的维度就是原始的模型维度,后续会通过广播机制将它们扩展到和查询头数量相匹配的形状以完成注意力计算 。
MQA减少计算量体现在两方面:
1、减少多余投影的计算次数。
2、 标准的MHA每个头都拥有专属的 Q、K、V 映射参数 每个头的 Ki 是独立存储的 MQA 通过 K 的共享和广播,减少了内存访问的离散性,提高了硬件利用率。
3、减少对额外的每个头的K V投影矩阵的梯度计算。