如何在计算 Attention Score 时对 Padding 做 Mask 操作

在使用神经网络进行自然语言处理任务时，输入序列通常会有不同的长度。为了使得这些序列能够批处理输入，通常需要对较短的序列进行填充（Padding）操作。然而，在计算 Attention Score 时，这些填充部分不应参与运算。本文将详细讲解如何在计算 Attention Score 时对 Padding 做 Mask 操作。

什么是 Attention Score

在 Transformer 模型中，Attention Score 是一种用于表示序列中不同位置的词对目标词的重要性程度的机制。它的计算通常依赖于 Query 和 Key 的内积，并经过缩放和 softmax 归一化。Attention Score 的公式如下：

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

其中， $Q$ 是 Query 矩阵， $K$ 是 Key 矩阵， $V$ 是 Value 矩阵， $d_k$ 是 Key 的维度。

为什么需要对 Padding 做 Mask 操作

Padding 的部分在 Attention Score 计算中不应该被考虑，否则会引入无关的噪声。我们可以使用 Mask 操作来避免这种情况发生。Mask 操作通过将 Padding 的部分设为负无穷大，从而在 softmax 步骤中将其归零，使得模型不会关注到填充的部分。

Mask 操作的公式

假设 $S$ 为未进行 Mask 操作的 Attention Score 矩阵，而 $M$ 为 Mask 矩阵，其中 Padding 部分为 0，有效部分为 1。那么进行 Mask 操作后的 Attention Score 可以表示为：

S_{\text{masked}} = S + M \times (-\infty)

在经过 softmax 操作后，Padding 部分将变为 0。更为具体地，假设 $S_{ij}$ 表示第 $i$ 个 Query 对第 $j$ 个 Key 的得分，那么 Mask 操作可以表示为：

S'_{ij} = \begin{cases} S_{ij}, & \text{if } M_{ij} = 1 \\ -\infty, & \text{if } M_{ij} = 0 \end{cases}

最终的 Attention Score 计算过程为：

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{S'}{\sqrt{d_k}}\right)V

代码实现

在实现中，我们通常使用一个布尔 Mask 矩阵来表示填充位置。例如，假设 mask 是一个形状为 (batch_size, seq_len) 的布尔矩阵，其中填充部分为 False，非填充部分为 True。以下是一个典型的 Mask 应用代码：

python
import torch
import torch.nn.functional as F

# 假设 Q, K, V 为 batch_size x seq_len x d_k
Q = torch.rand(2, 5, 64)
K = torch.rand(2, 5, 64)
V = torch.rand(2, 5, 64)

# 假设 mask 为 batch_size x seq_len，填充位置为 0，有效位置为 1
mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 0]]).bool()

# 计算未进行 Mask 操作的注意力分数 S
d_k = Q.size(-1)
S = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))

# 对 Mask 矩阵扩展形状，使其与 S 维度一致
mask = mask.unsqueeze(1).expand(-1, S.size(1), -1)

# 使用 -inf 填充 Padding 部分
S = S.masked_fill(~mask, float('-inf'))

# 计算 Attention Score 并应用到 V
attention = F.softmax(S, dim=-1)
output = torch.matmul(attention, V)

print(output)

代码详解

创建 Query、Key、Value 矩阵：我们假设 $Q$ 、 $K$ 和 $V$ 的维度为 (batch_size, seq_len, d_k)。
生成 Mask 矩阵：mask 为布尔矩阵，表示填充和有效部分。
计算 Attention Score 矩阵 $S$ ：通过 torch.matmul(Q, K.transpose(-2, -1)) 计算内积，再缩放除以 $\sqrt{d_k}$ 。
应用 Mask 操作：使用 masked_fill 函数将填充部分设为 -inf，有效部分保持不变。
计算 Softmax 并得到最终输出：使用 softmax 对注意力分数进行归一化，最后通过 torch.matmul 计算得到注意力加权后的输出。

总结

在 Attention Score 计算过程中，对 Padding 部分进行 Mask 操作是非常必要的，以确保填充部分不影响结果。通过将填充部分设为负无穷大，模型可以在 softmax 归一化时忽略这些部分，从而专注于真正的有效信息。通过本文的讲解和代码示例，相信大家能够清晰理解如何在计算 Attention Score 时对 Padding 做 Mask 操作，从而在实践中有效应用。