PyTorch TransformerEncoder中padding机制及src_key_padding_mask失效问题排查

你对src_key_padding_mask的作用理解有偏差，这个掩码的核心作用是在自注意力计算阶段屏蔽padding token的影响，而非让输出的padding位置保持输入时的0值。具体拆解如下：

1. src_key_padding_mask的实际功能 #

src_key_padding_mask（布尔型，True标记padding位置）的作用是：

• 阻止注意力层计算padding token与其他token之间的注意力权重；
• 避免其他token将注意力分配到padding位置上。
  但它不会干预后续的前馈网络（FFN）、残差连接和层归一化操作，这些模块依然会处理padding位置的向量，因此输出的padding位置值会被修改，不会保持输入时的0。

2. 你的代码输出不符合预期的原因 #

你手动将输入的padding位置设为0，但经过TransformerEncoderLayer的处理：

• 注意力层虽然忽略了padding位置的注意力交互，但残差连接会把原始输入（包括0的padding）和注意力输出相加；
• 后续的层归一化、FFN都会对所有位置（包括padding）的向量做运算，最终导致padding位置的输出值偏离0。

3. 实现输出padding位置保持0的方法 #

如果需要让输出的padding位置保持0，可以在得到encoder输出后，手动用掩码将这些位置置0：

output = encoder(input_tensor, src_key_padding_mask=bool_src_key_padding_mask)
# 扩展mask的维度以匹配输出的形状（batch, seq_len, dim），然后将对应位置置0
output = output.masked_fill(bool_src_key_padding_mask.unsqueeze(-1), 0.0)

或者，如果你只是希望后续流程忽略padding位置，也可以直接使用bool_src_key_padding_mask过滤掉这些位置的输出，不一定非要将其置0。

内容的提问来源于stack exchange，提问作者Tom T.