这是我关于 Transformer 系列的第三篇文章。我们将以自上而下的方式介绍其功能。在之前的文章中，我们了解了 Transformer 是什么、它的架构以及它的工作原理。

在本文中，我们将更进一步，深入探讨多头注意力机制 (Multi-head Attention)，即 Transformer 的大脑。

Transformer 中如何使用注意力机制

正如我们在第 2 部分中讨论的那样，Transformer 中的三个地方使用了 Attention：

编码器中的自注意力——输入序列会关注自身
解码器中的自注意力——目标序列会关注自身
解码器中的编码器-解码器注意力——目标序列会关注输入序列

注意输入参数——查询、键和值

注意力层以三种参数的形式获取输入，即查询 (Query)、键 (Key) 和值 (Value)。

这三个参数的结构相似，序列中的每个单词都由一个向量表示。

编码器自注意力机制

输入序列被馈送到输入嵌入和位置编码，后者为输入序列中的每个单词生成一个编码表示，该表示可捕获每个单词的含义和位置。这被馈送到第一个编码器的自注意力中的所有三个参数，即查询、键和值，然后该编码器还会为输入序列中的每个单词生成一个编码表示，现在还合并了每个单词的注意力分数。当它通过堆栈中的所有编码器时，每个自注意力模块也会将其自己的注意力分数添加到每个单词的表示中。

解码器自注意力

进入解码器堆栈后，目标序列被馈送到输出嵌入和位置编码，后者会为目标序列中的每个单词生成一个编码表示，以捕捉每个单词的含义和位置。这会馈送到第一个解码器中自注意力中的所有三个参数（查询、键和值），然后解码器还会为目标序列中的每个单词生成一个编码表示，现在还会合并每个单词的注意力分数。

通过 Layer Norm 后，它被馈送到第一个解码器中的编码器-解码器注意中的查询参数

编码器-解码器注意力机制

与此同时，堆栈中最后一个编码器的输出被传递给编码器-解码器注意力机制中的 Value 和 Key 参数。

因此，编码器-解码器注意力机制既获取目标序列的表示（来自解码器自注意力机制），也获取输入序列的表示（来自编码器堆栈）。因此，它会为每个目标序列单词生成一个带有注意力分数的表示，该表示也会捕获来自输入序列的注意力分数的影响。

当它穿过堆栈中的所有解码器时，每个自我注意力和每个编码器-解码器注意力也会将它们自己的注意力分数添加到每个单词的表示中。

多个注意力头

在 Transformer 中，注意力模块会多次并行重复计算。每个计算都称为注意力头。注意力模块会将其查询、键和值参数拆分成 N 路，并将每个拆分单独传递给单独的头。然后，所有这些类似的注意力计算会组合在一起，以产生最终的注意力分数。这称为多头注意力，它使 Transformer 能够更有效地为每个单词编码多种关系和细微差别。

为了准确了解数据的内部处理方式，让我们在训练 Transformer 解决翻译问题时，了解一下 Attention 模块的工作原理。我们将使用一个训练数据样本，该样本由一个输入序列（英文中的“You are welcome”）和一个目标序列（西班牙语中的“De nada”）组成。

注意超参数

有三个超参数决定数据维度：

嵌入大小 — 嵌入向量的宽度（我们在示例中使用宽度 6）。此维度贯穿整个 Transformer 模型，因此有时被称为其他名称，如“模型大小”等。
查询大小（等于键和值的大小）——三个线性层分别用于生成查询、键和值矩阵的权重的大小（在我们的示例中，我们使用查询大小为 3）
注意力头的数量（我们在示例中使用了 2 个注意力头）

此外，我们还有批次大小，为样本数量提供一个维度。

输入层

输入嵌入和位置编码层产生一个形状矩阵（样本数、序列长度、嵌入大小），该矩阵被输入到堆栈中第一个编码器的查询、键和值。

为了便于可视化，我们将在图片中删除“批次”维度，并专注于剩余维度。

线性层

查询、键和值有三个独立的线性层。每个线性层都有自己的权重。输入通过这些线性层生成 Q、K 和 V 矩阵。

将数据拆分到各个注意力头上

现在数据被分割到多个注意力头上，以便每个注意力头都能独立处理。

但是，需要理解的是，这只是逻辑上的分割。查询、键和值不会在物理上分割成单独的矩阵，每个注意力头一个。查询、键和值分别使用一个数据矩阵，每个注意力头的矩阵部分在逻辑上是分开的。同样，每个注意力头也没有单独的线性层。所有注意力头共享相同的线性层，但只是在数据矩阵的“自己”的逻辑部分上运行。

线性层权重按人头进行逻辑划分

这种逻辑分割是通过在 Attention 头上均匀划分输入数据以及线性层权重来实现的。我们可以通过选择以下查询大小来实现这一点：

查询大小 = 嵌入大小 / 头的数量

在我们的例子中，这就是为什么查询大小 = 6/2 = 3。即使层权重（和输入数据）是一个矩阵，我们也可以将其视为将每个头部的单独层权重“堆叠在一起”。

因此，所有注意力头的计算都可以通过单个矩阵运算来实现，而不需要 N 个单独的运算。这使得计算更加高效，并且由于需要更少的线性层而使模型保持简单，同时仍能实现独立注意力头的功能。

重塑 Q、K 和 V 矩阵

线性层输出的 Q、K 和 V 矩阵被重塑以包含显式头部维度。现在每个“切片”对应每个头部的一个矩阵。

通过交换 Head 和 Sequence 维度，该矩阵再次被重塑。虽然未绘制 Batch 维度，但 Q 的维度现在是 (Batch、Head、Sequence、Query size)。

在下图中，我们可以看到从线性层出来后分割示例 Q 矩阵的完整过程。

最后阶段仅用于可视化——尽管 Q 矩阵是一个单一矩阵，但我们可以将其视为每个头部逻辑上独立的 Q 矩阵。

我们准备计算注意力分数。

计算每个头部的注意力分数

现在，我们有了 3 个矩阵，Q、K 和 V，分布在各个头部。这些矩阵用于计算注意力得分。

我们将仅使用最后两个维度（序列和查询大小）来展示单个头部的计算，并跳过前两个维度（批次和头部）。本质上，我们可以想象，我们正在查看的计算针对每个头部和批次中的每个样本进行“重复”（尽管显然，它们是作为单个矩阵运算而不是循环进行的）。

第一步是对 Q 和 K 进行矩阵乘法。

现在将 Mask 值添加到结果中。在 Encoder Self-attention 中，mask 用于屏蔽 Padding 值，使它们不参与 Attention Score。

解码器自注意力和解码器编码器注意力中应用了不同的掩码，我们将在稍后的流程中讨论。

现在通过除以查询大小的平方根来缩放结果，然后对其应用 Softmax。

Softmax 的输出和 V 矩阵之间进行另一次矩阵乘法。

Encoder Self-attention 中完整的 Attention Score 计算如下：

将每位主管的注意力得分合并在一起

现在，每个头部都有单独的注意力分数，需要将它们组合成一个分数。此合并操作本质上是拆分操作的逆操作。

通过简单地重塑结果矩阵以消除 Head 维度即可完成。步骤如下：

通过交换 Head 和 Sequence 维度来重塑注意力得分矩阵。换句话说，矩阵形状从 (Batch、Head、Sequence、Query size) 变为 (Batch、Sequence、Head、Query size)。
通过重塑为（批次、序列、头部 * 查询大小）来折叠头部维度。这有效地将每个头部的注意力得分向量连接成一个合并的注意力得分。

由于 Embedding size =Head * Query size，合并后的 Score 为 (Batch, Sequence, Embedding size)。在下图中，我们可以看到示例 Score 矩阵合并的完整过程。