Transformer -2：Mask self-attention 掩码机制-CSDN博客

本文链接：https://blog.csdn.net/m0_75108877/article/details/147933970

在自然语言处理的 Transformer 模型中，掩码 (Mask) 机制是一个关键概念。它用于处理序列中的填充位置、防止信息泄露，以及控制注意力的流动方向。本文将通过一段具体的 PyTorch 代码，深入浅出地解释掩码机制的实现原理。

一、为什么需要掩码？

在处理序列数据时，我们经常会遇到两个问题：

变长序列：不同句子的长度可能不同，但神经网络通常需要固定大小的输入。
信息泄露：在解码器中，我们不希望模型看到未来的信息（例如，在预测第 i 个词时，不能使用 i+1 之后的词）。

掩码就是解决这些问题的关键技术。它通过在计算过程中 "遮蔽" 某些位置，使模型忽略这些位置的信息。

二、代码实现：构建编码器自注意力掩码

让我们从一段代码开始，逐步解析掩码的构建过程：

import torch
import torch.nn.functional as F
import numpy as np

# 假设我们有以下参数
batch_size = 2  # 批次大小
src_len = torch.tensor([2, 4])  # 两个序列的实际长度
max_src_seq_len = 5  # 最大序列长度

# 步骤1: 构建有效位置的掩码
vaild_encoder_pos = torch.unsqueeze(
    torch.cat([
        torch.unsqueeze(
            F.pad(torch.ones(L), (0, max_src_seq_len - L)), 0
        ) for L in src_len
    ]), 2
)
# 原始序列：src_len = [2, 4]
# padding 后：[tensor([1., 1., 0., 0., 0.]), tensor([1., 1., 1., 1., 0.])]（每个元素 shape [5]）
# 第一次 unsqueeze：[tensor([[1., 1., 0., 0., 0.]]), tensor([[1., 1., 1., 1., 0.]])]（每个元素 shape [1, 5]）
# cat 拼接：tensor([[1., 1., 0., 0., 0.], [1., 1., 1., 1., 0.]])（shape [2, 5]）
# 第二次 unsqueeze：tensor([[[1., 1., 0., 0., 0.]], [[1., 1., 1., 1., 0.]]])（shape [2, 5, 1]）

为每个序列创建一个长度为max_src_seq_len的向量
实际词的位置填充为 1，填充位置填充为 0
最终得到形状为[batch_size, max_src_seq_len, 1]的张量

例如，对于序列长度[2, 4]，我们会得到：

tensor([[[1.], [1.], [0.], [0.], [0.]],  # 第一个序列：前2个位置有效
        [[1.], [1.], [1.], [1.], [0.]]])  # 第二个序列：前4个位置有效

三、构建掩码矩阵

下一步，我们需要将这个一维掩码扩展为二维矩阵：

# 步骤2: 构建掩码矩阵
vaild_encoder_pos_matrix = torch.bmm(
    vaild_encoder_pos, 
    vaild_encoder_pos.transpose(2, 1)
)
invaild_encoder_pos_matrix = 1 - vaild_encoder_pos_matrix
mask_encoder_self_attention = invaild_encoder_pos_matrix.to(torch.bool)

这里的关键操作是torch.bmm，它执行批次矩阵乘法：

vaild_encoder_pos形状：[batch_size, max_src_seq_len, 1]

vaild_encoder_pos.transpose(2, 1)形状：[batch_size, 1, max_src_seq_len]
矩阵乘法结果形状：[batch_size, max_src_seq_len, max_src_seq_len]

vaild_encoder_pos_matrix这个矩阵的含义是：

有效位置：如果两个位置 i 和 j 都有效，则matrix[i][j] = 1
无效位置：如果 i 或 j 中有一个是填充位置，则matrix[i][j] = 0

例如，对于第一个序列（长度 2），vaild_encoder_pos_matrix矩阵为：

tensor([[[1, 1, 0, 0, 0],
         [1, 1, 0, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0],
         [0, 0, 0, 0, 0]]])

这表示：只有前两个位置之间的注意力计算是有效的，其他位置都被掩码。

invaild_encoder_pos_matrix这个矩阵的含义是：

有效位置：如果两个位置 i 和 j 都有效，则matrix[i][j] = 0
无效位置：如果 i 或 j 中有一个是填充位置，则matrix[i][j] = 1

tensor([[0., 0., 1., 1., 1.],
         [0., 0., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.],
         [1., 1., 1., 1., 1.]])

mask_encoder_self_attention这个矩阵的含义是：

有效位置：如果两个位置 i 和 j 都有效，则matrix[i][j] = False
无效位置：如果 i 或 j 中有一个是填充位置，则matrix[i][j] = True

tensor([[[False, False,  True,  True,  True],
         [False, False,  True,  True,  True],
         [ True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True],
         [ True,  True,  True,  True,  True]])

四、应用掩码到注意力分数

最后，我们将掩码应用到注意力分数上：

# 步骤3: 应用掩码到注意力分数
score = torch.randn(batch_size, max_src_seq_len, max_src_seq_len)# 随机初始化一个QK矩阵，这里只是考虑了这个形状，并没有用到真正的Q\K
mask_score = score.masked_fill(mask_encoder_self_attention, -np.inf) #True就是无效，置为-∞
prob = F.softmax(mask_score, dim=-1)

1这里发生了什么？

随机生成注意力分数：score是一个随机矩阵，表示每个位置对其他位置的 "注意力程度"

应用掩码：使用masked_fill函数，将掩码矩阵中为True的位置（无效位置）填充为-∞
Softmax 计算：填充后的分数经过 softmax 函数，无效位置的概率接近 0，有效位置的概率被保留

例如，对于第一个序列，掩码后的分数和概率可能是：

掩码前的分数:
tensor([[[-0.82,  1.23,  0.45, -0.12,  0.78],
         [ 0.34, -0.91,  1.56, -0.54,  0.23],
         [ 0.67, -0.23, -0.89,  1.34, -0.45],
         [ 0.98, -0.45,  0.12, -0.76,  1.56],
         [ 0.23, -0.89,  0.56, -0.12, -0.78]]])

掩码后的分数:
tensor([[[-0.82,  1.23, -inf, -inf, -inf],
         [ 0.34, -0.91, -inf, -inf, -inf],
         [-inf, -inf, -inf, -inf, -inf],
         [-inf, -inf, -inf, -inf, -inf],
         [-inf, -inf, -inf, -inf, -inf]]])

Softmax后的概率:
tensor([[[0.11, 0.89, 0.00, 0.00, 0.00],
         [0.78, 0.22, 0.00, 0.00, 0.00],
         [0.00, 0.00, 0.00, 0.00, 0.00],
         [0.00, 0.00, 0.00, 0.00, 0.00],
         [0.00, 0.00, 0.00, 0.00, 0.00]]])