Transformer模型中的多头注意力机制详解
Transformer模型中的多头注意力机制详解
Transformer模型的核心是多头注意力机制，它通过计算词与词之间的相关性，捕捉序列中的全局依赖关系，允许模型动态关注输入的不同部分，增强上下文理解能力，而不用考虑输入间的举例。多头注意力并行计算多个注意力头，提取多样化的特征，提升模型的表达能力和鲁棒性，广泛应用于机器翻译、文本生成等任务。
0. 前言
按照国际惯例，首先声明：本文只是作者学习的理解，虽然参考了他人的宝贵见解及成果，但是内容可能存在不准确的地方。如果发现文中错误，希望批评指正，共同进步。
本文是Transformer学习系列的第二篇文章：
- 第一篇：提出背景、模型架构及推理过程
- 第二篇：多头注意力机制（本篇）
- 第三篇：位置编码
在前面的Transformer模型架构说明中已经有多处提及（多头）注意力机制，注意力机制是Transformer的核心，本文将详细说明注意力机制的数学原理及作用。
1. 注意力机制（Attention）概述
注意力机制是Transformer模型的核心，因为它赋予了模型识别输入序列中不同元素之间复杂关系的能力，而无需受限于传统的顺序处理方式。通过自注意力（Self-Attention）机制，每个位置的元素都能直接与序列中的所有其他元素交互，从而有效地捕捉长距离依赖关系。
这种机制允许并行化计算，极大地提高了训练效率和模型性能，同时使得模型能够动态地对每个输入元素强调不同的关注点，根据其上下文调整对信息的关注程度。因此，注意力机制不仅加速了模型处理速度，还增强了模型的理解能力和表达能力，成为Transformer架构区别于以往模型的关键创新点。
注意力机制通过计算查询（Query）、键（Key）和值（Value）之间的关系来生成输出。强烈建议先阅读下前文【单点知识】多头注意力机制：torch.nn.MultiheadAttention的原理与实践中的例子。
在Attention is All You Need原文中提出了两种注意力机制：加法注意力和点乘注意力。虽然两者在理论上的复杂度相近，但在实际上点乘注意力性能更好，后面的注意力机制说明都默认是点乘注意力。
注意力机制通过Q（Query，查询）、K（Key，键）和V（Value，值）矩阵工作。Q表示当前词的查询向量，K表示所有词的键向量，V表示所有词的值向量。
首先，计算Q和K的点积并缩放，得到注意力分数，表示词与词之间的相关性。然后，通过 Softmax 将分数转换为概率分布。最后，用这些概率加权V，V矩阵提供了词的实际内容，而注意力机制通过加权V矩阵中的值向量，生成一个上下文相关的表示，帮助模型更好地理解序列中的关系。
2. Q、K、V矩阵是怎么来的？
一说到Transformer的注意力机制，众多文章往往一上来就开始QKV……的解释，令人摸不到头脑。这也是Attention is All You Need原文并没有太交代清楚的地方，我觉得有必要先说明白。
从模型最开始，Transformer输入的n个Token会被转换为( n , d m o d e l )维的矩阵X（回顾下前文【Transformer模型学习】第一篇：提出背景、模型架构及推理过程中的第三章），Transformer 使用三个不同的可学习权重矩阵生成Q、K、V矩阵：
- W Q ：用于生成查询矩阵Q = X ⋅ W Q
- W K ：用于生成键矩阵K = X ⋅ W K
- W V ：用于生成值矩阵V = X ⋅ W V
权重矩阵W Q 和W K 的形状为( d model , d k )，生成的Q和K的形状为( n , d k )。W V 的形状为( d model , d v )，生成的V的形状为( n , d v )。
3. 缩放点积注意力（Scaled Dot-Product Attention）
这是Transformer中使用的注意力机制。它通过计算查询和键的点积，然后进行缩放和Softmax操作来得到权重，最后用这些权重对值进行加权求和。
公式如下：
Attention ( Q , K , V ) = softmax ( Q K T d k ) V
其中，Q、K、V分别表示查询、键和值，d k 是键的维度。
可以更形象一点表示为：
其中，缩放因子d k 是为了防止点积结果过大，导致softmax函数输出接近0或1的极端值，从而减小梯度消失的风险，确保训练过程更加稳定有效。通过缩放点积的结果，可以使softmax函数更有效地分辨不同词之间的关联强度。
最终，如果我们令d v = d m o d e l 就会发现，Attention变成了和输入X的维度一致，这样，输入X就经历了一个“头”的处理，得到了单头注意力。
4. 多头注意力（Multi-Head Attention）
在单头注意力的基础上我们也可以再将查询Q、键K和值V分别投影到多个d v = d m o d e l / h维子空间（在原文中h设定为8，d m o d e l 设定为512），然后在每个子空间中并行计算注意力，最后将结果拼接并再次投影。
整个过程相当于把一个“大头”分解成了多个“小头”，其公式如下：
MultiHead ( Q , K , V ) = Concat ( head 1 , … , head h ) W O
W O 是一个可学习的权重矩阵，用于将多个注意力头的输出拼接后进行线性变换，最终生成多头注意力的输出。其中，每个头head i 的计算方式为：
head i = Attention ( Q W i Q , K W i K , V W i V )
5. 多头注意力的好处
多头注意力的主要好处可以总结为以下几点：
- 捕捉多样化的特征：每个注意力头可以独立学习不同的注意力模式。例如，一个头可能关注局部依赖（如相邻词之间的关系），而另一个头可能关注长距离依赖（如句子中相隔较远的词之间的关系）。
- 支持并行计算，提高效率：多头注意力机制中的每个头的计算是独立的，可以并行进行。这种并行性充分利用了硬件加速器（如GPU或TPU）的计算能力，显著提高了计算效率。
- 增强模型的鲁棒性和表达能力：多个注意力头的组合使得模型能够从不同角度理解输入数据。即使某些头的注意力分配不理想，其他头仍然可以提供有用的信息。
- 更好地建模上下文关系：多头注意力能够同时捕捉输入序列中的多种上下文关系。例如，在自然语言处理任务中，一个词可能与多个其他词有不同的语义关系。
6. 总结
本文详细说明了Transformer的核心——多头注意力机制。它可以通过计算词与词之间的相关性，捕捉序列中的全局依赖关系，允许模型动态关注输入的不同部分，增强上下文理解能力，而不用考虑输入间的举例。多头注意力并行计算多个注意力头，提取多样化的特征，提升模型的表达能力和鲁棒性，广泛应用于机器翻译、文本生成等任务。