PyTorch Transformer使用src_key_padding_mask输出不符合预期如何解决

结论先行 #

• 你的src_key_padding_mask使用方式完全符合PyTorch的API要求，没有参数错误
• 你预期掩码位置输出为0/负无穷是对该参数作用的常见误解

原理说明 #

src_key_padding_mask的作用范围仅限制在注意力权重计算环节：
计算自注意力得分时，框架会自动将padding位置对应的注意力得分置为负无穷，经过Softmax运算后该位置的注意力权重为0，保证padding token不会对其他有效token的表征计算产生干扰。
但TransformerEncoder层除了注意力模块外，还包含残差连接、层归一化、FeedForward全连接层，这些网络层会正常处理padding位置的原始输入值，所以最终输出中padding位置的数值不会自动变为0或负无穷。

验证方法 #

你可以对比out1和out2中有效非padding位置的输出值，二者会存在明显差异，这就说明mask已经起到了过滤padding干扰的作用，符合预期功能。

实现padding位置输出为0的方案 #

如果业务上需要padding位置的输出为0，可以在编码器输出后手动乘以有效位置掩码，示例代码如下：

# 转换为float类型并扩展到和输出相同的特征维度
valid_mask = (~mask.T).unsqueeze(-1).float()
out1 = out1 * valid_mask

注：你刻意不添加位置编码的操作不影响mask的作用逻辑，该部分设计没有问题。

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