多头注意力是人工神经网络中使用的一种注意力机制，它并行执行多次注意力计算，再将各次计算的输出合并。每次计算称为一个“头”，使用各自独立学习的投影来变换查询、键和值。这使得一个网络层能够通过不同的表示关注不同位置，而不只是生成单一的注意力加权混合表示。多头注意力是Transformer架构的核心组件。(docs.pytorch.org)
起源与架构中的作用
多头注意力由2017年的论文《Attention Is All You Need》提出，该论文介绍了一种用于机器翻译的Transformer。其编码器—解码器架构在三个位置使用了这一机制：编码器的自注意力、解码器的掩码自注意力，以及解码器对编码器输出的注意力。基础模型使用八个头，每个头的查询、键和值均为64维，模型宽度为512。(arxiv.org)
自注意力与交叉注意力的区别在于输入的来源，而不在于包含多少个头。自注意力的三种输入均来自同一个序列；交叉注意力的查询来自一个序列，键和值则来自另一个序列。两者都可以采用多头形式。(docs.pytorch.org)
数学表述
设 (Q)、(K) 和 (V) 是分别包含查询、键和值表示的输入矩阵。对于第 (i) 个头，通过学习得到的投影矩阵构造：
[ Q_i=QW_i^Q,\qquad K_i=KW_i^K,\qquad V_i=VW_i^V. ]
这些投影是线性映射，不过具体实现也可能包含偏置项。每个头计算缩放点积注意力：
[ H_i= \operatorname{softmax} \left(\frac{Q_iK_i^\mathsf{T}}{\sqrt{d_k}}+M\right)V_i. ]
其中，(d_k) 是每个头的查询和键的维度，(K_i^\mathsf{T}) 表示矩阵转置，(M) 是可选的加性掩码。查询与键的内积决定匹配分数。Softmax函数沿键的位置维度运算，将分数转换为非负权重；在应用任何注意力随机失活之前，这些权重之和为一。(docs.pytorch.org)
缩放因子用于抵消点积幅度随维度增加而增大的效应，否则可能使Softmax进入梯度较小的区域。输出的每一行都是投影后的值的加权线性组合。最后：
[ \operatorname{MultiHead}(Q,K,V) =\operatorname{Concat}(H_1,\ldots,H_h)W^O. ]
拼接操作保留各个头的独立输出，再由学习得到的输出投影将它们混合，形成该层的输出表示。(arxiv.org)
维度与实现
当查询位置数为 (n_q)、键位置数为 (n_k) 时，每个头的注意力权重形状为 (n_q\times n_k)。其输出有 (n_q) 行，因此交叉注意力不要求查询序列与键序列长度相等。键和值必须逐位置对应。具体实现通常会增加批次维度，并使用张量组织计算。(docs.pytorch.org)
一种常见配置是将模型宽度 (d) 平均分配给 (h) 个头，使每个头的维度为 (d/h)。因此，在宽度固定时，增加头数会缩小单个头的维度，而不是为每个头都分配一个与模型等宽的表示。头数是架构的一个超参数。软件库可以一次性计算各个投影，再将输出重塑为多个头，从而实现并行计算，无需为每个头单独编写循环。(docs.pytorch.org)
掩码用于限制允许建立的注意力关系。填充掩码排除填充位置；因果掩码阻止解码器中的某个位置访问后续位置。多头注意力通常嵌入更大的模块中，该模块还包含残差连接、层归一化和逐位置前馈网络。位置编码所提供的位置信息需要另行引入，并非基本注意力计算本身固有的信息。(arxiv.org)
计算成本与高效执行
对于覆盖 (n) 个位置的稠密自注意力，若显式存储每个头的注意力矩阵，每个样本需要 (O(hn^2)) 个元素。根据矩阵维度，在所有头的总宽度固定时，分数计算和值聚合操作需要 (O(n^2d)) 的算术运算；输入和输出投影另需约 (O(nd^2)) 的运算。这些计算复杂性估算区分了头数与表示的总宽度。对于长输入，计算量随序列长度呈二次增长这一特点尤为重要。(arxiv.org)
FlashAttention改变的是执行策略，而非注意力机制的数学形式。它采用分块计算，减少图形处理器不同存储层级之间的数据传输，避免在高带宽内存中存储完整的注意力矩阵。它在提高内存效率的同时计算精确注意力，但结果仍可能因浮点运算而存在数值差异。它并未消除稠密注意力的二次算术运算量。(arxiv.org)
注意力头的冗余与键值共享
独立学习的投影使各个头能够表现出不同的行为，但参数不同并不保证每个头都有独特贡献。2019年的研究《Are Sixteen Heads Really Better than One?》发现，在所测试的翻译模型和BERT 语言模型中，许多头可以在推理时移除，而性能损失很小。有些层甚至可以缩减为单个头，但编码器—解码器注意力对剪枝更为敏感。这些发现针对特定模型和任务，并不意味着存在普遍适用的最优头数。(papers.neurips.cc)
多查询注意力保留多个查询头，但共享一个键值头。分组查询注意力则让同一组查询头共享键值头，其形式介于常规多头注意力与多查询注意力之间。这能减少自回归神经网络推理过程中的键值存储需求和内存带宽需求。2023年的GQA论文报告称，在其实验中，分组查询注意力的质量接近常规多头注意力，速度则与多查询注意力相当。(arxiv.org)
文本之外的应用
多头注意力也可用于非语言序列。在视觉Transformer中，图像被划分为多个图像块，并表示为由Transformer模块处理的序列。注意力在图像块的表示之间建立联系，使该架构无需卷积特征提取也能执行图像分类。最初的视觉Transformer研究通过大规模预训练以及在图像识别基准上的评估，展示了这一方法的有效性。(arxiv.org)