PyTorch Transformer Encoder填充掩码失效问题求助

在使用PyTorch的Transformer Encoder时，希望填充序列不影响模型输出，但即使将填充位置的输入值设为0，非填充位置的输出仍存在差异。以下是复现代码：

import torch as th
from torch import nn

# Data
batch_size = 2
seq_len = 5
input_size = 16
src = th.randn(batch_size, seq_len, input_size)

# Set some values to a high value
src[0, 2, :] = 1000.0
src[1, 4, :] = 1000.0

# Generate a padding mask
padding_mask = th.zeros(batch_size, seq_len, dtype=th.bool)
padding_mask[0, 2] = 1
padding_mask[1, 4] = 1

# Pass the data through the encoder of the model
encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(
        d_model=input_size,
        nhead=1,
        batch_first=True,
    ),
    num_layers=1,
    norm=None,
)
out1000 = encoder(src, src_key_padding_mask=padding_mask)

# Modify the input data so that the masked vector does not affect
src[0, 2, :] = 0.0
src[1, 4, :] = 0.0

# Pass the modified data through the model
out0 = encoder(src, src_key_padding_mask=padding_mask)

# Check if the results are the same
assert th.allclose(
    out1000[padding_mask == 0],
    out0[padding_mask == 0],
    atol=1e-5,
)

问题原因 #

1. Dropout的随机干扰
   默认情况下，TransformerEncoderLayer包含概率为0.1的dropout层，两次前向传播会生成不同的随机掩码。即使输入的非填充部分完全一致，随机dropout也会导致输出出现差异，这是断言失败的最直接原因。

2. 填充位置仍参与后续计算
   src_key_padding_mask仅在自注意力层中屏蔽对填充位置的注意力计算——即非填充位置的token不会关注填充位置的token，但填充位置的输入值（无论是1000还是0）仍然会参与残差连接、LayerNorm、前馈网络等后续步骤。虽然这些计算逐token独立进行，不会直接影响非填充位置，但随机组件（如dropout）会间接放大差异。

解决方法 #

1. 关闭dropout以获得确定性结果

在定义TransformerEncoderLayer时显式设置dropout=0.0，消除随机因素的影响：

encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(
        d_model=input_size,
        nhead=1,
        batch_first=True,
        dropout=0.0  # 关闭dropout
    ),
    num_layers=1,
    norm=None,
)

2. 彻底隔离填充位置的影响

如果需要确保填充位置完全不参与任何计算，除了使用src_key_padding_mask，还需要：

• 提前将填充位置的输入置零；
• 在模型输出后，再次将填充位置的输出置零（避免残差连接或前馈网络的残留影响）。

修改后的代码示例：

import torch as th
from torch import nn

batch_size = 2
seq_len = 5
input_size = 16
src = th.randn(batch_size, seq_len, input_size)

# 初始化填充位置为高值
src[0, 2, :] = 1000.0
src[1, 4, :] = 1000.0

padding_mask = th.zeros(batch_size, seq_len, dtype=th.bool)
padding_mask[0, 2] = 1
padding_mask[1, 4] = 1

# 定义无dropout的编码器
encoder = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(
        d_model=input_size,
        nhead=1,
        batch_first=True,
        dropout=0.0
    ),
    num_layers=1,
    norm=None,
)

# 处理输入：填充位置置零
src1000 = src.masked_fill(padding_mask.unsqueeze(-1), 1000.0)
out1000 = encoder(src1000, src_key_padding_mask=padding_mask)
# 输出后再次置零填充位置
out1000 = out1000.masked_fill(padding_mask.unsqueeze(-1), 0.0)

# 修改填充位置为0并处理
src0 = src.masked_fill(padding_mask.unsqueeze(-1), 0.0)
out0 = encoder(src0, src_key_padding_mask=padding_mask)
out0 = out0.masked_fill(padding_mask.unsqueeze(-1), 0.0)

# 验证非填充位置输出一致
assert th.allclose(
    out1000[padding_mask == 0],
    out0[padding_mask == 0],
    atol=1e-5,
)

3. 理解掩码的正确用法

src_key_padding_mask的核心作用是防止注意力机制关注填充位置，而非完全阻断填充位置的计算。如果仅需保证注意力不被填充干扰，关闭dropout后即可让非填充位置的输出一致；如果需要填充位置完全不影响模型状态，则需配合输入和输出的掩码置零操作。

内容的提问来源于stack exchange，提问作者Angelo