这样图解Transformer应该没人看不懂了吧——多头注意力机制详解
这样图解Transformer应该没人看不懂了吧——多头注意力机制详解
Transformer是自然语言处理领域最具影响力的模型之一，其核心创新之一就是多头注意力机制。本文将深入解析Transformer中的多头注意力机制，帮助读者理解其工作原理和实现细节。
这是关于Transformer系列文章的第三篇部分，我们将用自上而下的方式深入探讨Transformer的功能。在前两篇文章中，我们已经了解了Transformer是什么、它的架构以及工作原理。
Transformer中注意力的应用方式
Transformer在三个地方使用了注意力机制：
- 编码器中的自注意力：输入序列关注自身
- 解码器中的自注意力：目标序列关注自身
- 解码器中的编码器-解码器注意力：目标序列关注输入序列
注意力输入参数——查询、键、值
注意力层接收三个参数作为输入，即查询（Query）、键（Key）和值（Value），它们结构上相似，序列中的每个单词由一个向量表示。
编码器自注意力
输入序列首先经过输入嵌入和位置编码，为每个单词生成包含意义和位置信息的编码表示。这些表示随后作为自注意力中的查询、键和值参数，生成包含注意力分数的编码表示。随着这些表示通过编码器堆栈中的每一层，每个自注意力模块都会将各自的注意力分数加入到每个单词的表示中。
解码器自注意力
在解码器堆栈中，目标序列经过输出嵌入和位置编码，生成包含意义和位置信息的编码表示。这些表示同样作为自注意力中的查询、键和值参数，生成包含注意力分数的编码表示。之后，这些表示通过层归一化，作为编码器-解码器注意力中的查询参数。
编码器-解码器注意力
同时，编码器堆栈的最终输出被传递给编码器-解码器注意力中的键和值参数。因此，编码器-解码器注意力同时获得了目标序列（来自解码器自注意力）和输入序列（来自编码器堆栈）的表示，生成包含输入序列注意力分数影响的目标序列单词表示。随着这些表示通过解码器堆栈中的每一层，每个自注意力和编码器-解码器注意力都会将各自的注意力分数加入到每个单词的表示中。
多个注意力头
在Transformer中，注意力模块并行地重复其计算多次，每次称为一个注意力头。注意力模块将其查询、键和值参数分成N份，每份独立通过一个头进行处理。然后将所有类似的注意力计算结果组合起来，生成最终的注意力分数。这种多头注意力机制使Transformer能够更强大地编码每个单词的多种关系和细微差别。
为了深入理解数据在Transformer内部的处理过程，我们将以一个翻译问题为例，通过训练过程中的一个样本（输入序列为英文“You are welcome”，目标序列为西班牙文“De nada”）来演示注意力模块的工作机制。
注意力超参数
确定数据维度的三个超参数包括：
- 嵌入大小：嵌入向量的宽度（示例中为6）。此维度在整个Transformer模型中保持一致，有时也被称为“模型大小”等。
- 查询大小（等于键和值的大小）：三个线性层分别用于生成查询、键和值矩阵的权重大小（示例中为3）。
- 注意力头数（示例中为2）。
此外，我们还有批处理大小，它代表样本数量的一个维度。
输入层
输入嵌入（Input Embedding）和位置编码（Position Encoding）层生成一个形状为（样本数，序列长度，嵌入大小）的矩阵，该矩阵被送入堆叠中的第一个编码器的查询（Query）、键（Key）和值（Value）部分。为了简化视觉呈现，我们将在图示中省略批次维度，并专注于其余维度。
线性层
查询、键和值分别对应三个独立的线性层，每个线性层都有自己的权重。输入数据通过这些线性层以产生Q、K、V矩阵。
跨注意力头的数据分割
现在，数据被分割到多个注意力头中，以便每个头能够独立处理。但重要的是要理解，这只是一种逻辑上的分割。查询、键和值并没有物理上被分割成每个注意力头一个单独的矩阵。相反，它们分别使用一个数据矩阵，但矩阵中的每个注意力头都有逻辑上独立的部分。同样，也没有为每个注意力头设置单独的线性层。所有注意力头都共享相同的线性层，但仅在其“自己的”逻辑数据部分上操作。
线性层权重按头逻辑划分
这种逻辑分割是通过将输入数据和线性层权重均匀分配到各个注意力头来实现的。我们可以通过选择查询大小（Query Size）来实现这一点，计算公式为：
查询大小 = 嵌入大小 / 注意力头数
在我们的例子中，查询大小 = 6/2 = 3。尽管层权重（和输入数据）是一个单独的矩阵，但我们可以将其视为每个头单独的层权重“堆叠”在一起。因此，所有头的计算都可以通过单个矩阵操作完成，而不需要N个单独的操作。这使得计算更加高效，并且模型更简洁，因为所需的线性层更少，同时仍能实现独立注意力头的强大功能。
重塑Q、K、V矩阵
线性层输出的Q、K、V矩阵被重塑以包含显式的头维度。现在，每个“切片”都对应于每个头的矩阵。接下来，通过交换头和序列维度来再次重塑该矩阵，虽然批次维度没有画出，但Q的维度现在是（批次，头，序列，查询大小）。
在下图中，我们可以看到从线性层出来的Q矩阵被分割的完整过程。最后一步仅用于可视化——尽管Q矩阵是一个单独的矩阵，但我们可以将其视为每个头一个逻辑上独立的Q矩阵。现在我们准备计算注意力分数。
计算每个头的注意力分数
现在我们已经有了按头分割的Q、K、V三个矩阵，这些矩阵用于计算注意力分数。我们将展示使用最后两个维度（序列和查询大小）的单个头的计算过程，并省略前两个维度（批次和头）。基本上，我们可以想象我们正在查看的计算是针对每个头和批次中的每个样本“重复”进行的（尽管显然它们作为单个矩阵操作发生，而不是循环）。
第一步是进行Q和K之间的矩阵乘法。然后，在结果中添加一个掩码值，在编码器自注意力中，掩码用于屏蔽填充值，使其不参与注意力分数的计算。在解码器自注意力和解码器-编码器注意力中，将应用不同的掩码，我们将在稍后的流程中介绍。现在，通过将结果除以查询大小的平方根进行缩放，然后对其应用Softmax函数。最后，在Softmax的输出和V矩阵之间进行另一个矩阵乘法。
编码器自注意力中的完整注意力分数计算过程如下：
将每个头的注意力分数合并
目前，我们为每个头分别计算了注意力分数，这些分数需要被合并成一个单一的分数。这个合并操作本质上是拆分操作的逆过程。合并操作通过重塑结果矩阵来消除头的维度，具体步骤如下：
- 通过交换头和序列维度来重塑注意力分数矩阵。换句话说，矩阵的形状从（批次大小，头数，序列长度，查询大小）变为（批次大小，序列长度，头数，查询大小）。
- 通过重塑为（批次大小，序列长度，头数*查询大小）来合并头维度。这实际上是将每个头的注意力分数向量连接成一个单一的合并后的注意力分数。
由于嵌入大小等于头数乘以查询大小，因此合并后的分数形状为（批次大小，序列长度，嵌入大小）。
在下图中，我们可以看到示例分数矩阵的完整合并过程。
端到端多头注意力
将上述所有步骤整合起来，就是多头注意力的端到端流程。
多头拆分捕捉更丰富的含义
嵌入向量用于捕捉单词的意义。在多头注意力机制中，如我们所见，输入（和目标）序列的嵌入向量在逻辑上被拆分到多个头上。这样做的意义何在？
这意味着嵌入的不同部分可以学习每个单词与其他单词在序列中关系时的不同意义方面。这使得Transformer能够捕捉序列的更丰富解释。虽然这可能不是一个现实的例子，但它有助于我们建立直观理解。例如，一部分可能捕捉名词的“性别”（男性、女性、中性），而另一部分可能捕捉名词的“数”（单数或复数）。这在翻译中可能很重要，因为许多语言的动词使用取决于这些因素。
解码器自注意力与掩码
解码器自注意力的工作原理与编码器自注意力相同，只是它作用于目标序列的每个单词上。同样地，掩码操作会屏蔽目标序列中的填充词。
解码器编码器-解码器注意力与掩码
编码器-解码器注意力从两个来源接收输入。因此，与编码器自注意力（计算每个输入单词与其他输入单词之间的交互）和解码器自注意力（计算每个目标单词与其他目标单词之间的交互）不同，编码器-解码器注意力计算每个目标单词与每个输入单词之间的交互。因此，结果注意力分数中的每个单元格都对应于一个Q（即目标序列单词）与所有其他K（即输入序列）单词和所有V（即输入序列）单词之间的交互，掩码操作会屏蔽目标输出中的后续单词。
结论
本文能让大家对Transformer中的注意力（Attention）模块有了较为清晰的认识。结合我们在之前文章中详细探讨的Transformer的整体端到端流程，现在我们已经覆盖了Transformer架构的详细操作并明确了Transformer的功能。