带掩码的多头自注意力

在我们考察解码器堆栈时，请回想它的主要作用：一次生成一个元素地输出序列。在机器翻译等任务中，这意味着逐词生成翻译后的句子。这种按序生成过程对解码器内部使用的自注意力机制提出了特定要求。与编码器可以同时处理整个输入序列不同，解码器必须仅根据已生成的令牌来预测下一个令牌。它不能预见当前正在构建的输出序列中的未来令牌。

这就是带掩码的多头自注意力发挥作用的地方。它是我们在第2章中讨论的多头自注意力机制的一种修改，专门设计用于在训练和推理期间强制执行这种“不可预见”的约束。

为什么需要掩码

考虑序列到序列模型的训练过程。我们通常使用“教师强制”机制，即解码器被输入正确的目标序列（向右移位，通常带有序列开始令牌）作为输入，以预测每个位置的下一个令牌。例如，为了预测目标翻译的第三个词，解码器会接收到前两个正确的词。

如果在此处使用标准自注意力，位置i的注意力机制将能够整合目标序列中所有位置的信息，包括位置i+1、i+2等。这就像作弊；模型可以简单地从输入中复制下一个词，而不是学习根据前面的词和编码器的输出来预测它。模型需要学习条件概率分布$P(\text{输出}_t | \text{输出}_{<t}, \text{编码器\_输出})$。允许对未来令牌的注意力会破坏这种条件依赖性。

在推理期间（生成新序列时），未来令牌本来就是未知的。因此，注意力机制在训练和推理之间必须保持一致，只关注先前生成的令牌。

掩码机制

带掩码的多头自注意力通过修改每个注意力头内部的缩放点积注意力计算来实现这一约束。在对注意力分数应用softmax函数之前，会添加一个掩码。

回顾缩放点积注意力公式：

$$\text{分数} = \frac{QK^T}{\sqrt{d_k}}$$

在带掩码的自注意力中，我们修改如下：

$$\text{带掩码的分数} = \frac{QK^T}{\sqrt{d_k}} + M$$ $$\text{注意力权重} = \text{softmax}(\text{带掩码的分数})$$ $$\text{注意力}(Q, K, V) = \text{注意力权重} V$$

掩码$M$通常是一个矩阵，其中模型不允许关注的位置的元素被设置为一个非常大的负数（实际上是负无穷大），而允许关注的位置的元素被设置为零。

对于长度为$L$的目标序列，掩码$M$将是一个$L \times L$矩阵。对于第$i$个令牌（$Q$矩阵的第$i$行），掩码确保它只关注令牌$j$，其中$j \le i$（$K$矩阵的第$0$到$i$列）。其中$j > i$的元素$M_{ij}$被设置为$-\infty$，而$j \le i$的元素被设置为$0$。

当一个大的负数被添加到未来位置的注意力分数时，随后的softmax函数会给这些位置分配一个非常接近零的概率。这有效阻止了任何信息从未来令牌流入当前令牌的表示中。

让我们可视化长度为4的序列的掩码：

   关注目标 -->   位置 1  位置 2  位置 3  位置 4
查询自位置 1: [  0      -inf   -inf   -inf ]
查询自位置 2: [  0       0     -inf   -inf ]
查询自位置 3: [  0       0      0     -inf ]
查询自位置 4: [  0       0      0      0   ]

这里，0代表一个允许的注意力连接（添加此掩码后，原始分数不变），-inf代表一个被掩码的连接（添加此掩码后，分数实际上变为负无穷大）。

下图说明了输出序列“The model predicts well”中位置3（“predicts”）的令牌允许的注意力连接，假设它是逐步生成的。

该图显示，来自位置3（“predicts”）的查询可以关注位置1、2和3的键/值，但对位置4（“well”）的注意力被掩盖掉了。

与多头注意力的结合

这种掩码过程在多头结构的每个注意力头内部独立应用。整体过程与标准多头注意力相同：

1. 将输入嵌入投影到多组Q、K、V向量中（每头一组）。
2. 对于每个头，计算缩放点积注意力分数。
3. 应用注意力掩码，将掩码矩阵$M$添加到分数中。
4. 应用softmax函数以获得注意力权重。
5. 计算V向量的加权和。
6. 连接所有头的输出。
7. 应用最终的线性投影。

与编码器的自注意力相比，唯一的区别是步骤3，即掩码的应用。

总结

带掩码的多头自注意力是标准自注意力机制的一种修改，对于Transformer的解码器部分不可或缺。通过阻止位置关注输出序列中的后续位置，它确保了模型的预测是自回归的，这意味着当前步骤的预测仅依赖于先前生成的步骤和输入序列。这是通过在softmax计算之前，向注意力分数中添加一个掩码矩阵（包含零和负无穷大）来实现的，从而有效地将未来令牌的权重归零。这种机制应用于解码器自注意力层中多头注意力结构的每个头部。这与编码器的自注意力形成对比，编码器自注意力允许每个位置关注输入序列中的所有位置。