掩蔽Softmax（Masked Softmax）——学习笔记

最近在b站看了沐神的深度学习课程，在手搓注意力机制。偶然习得了Masked Softmax。

在注意力当中，softmax操作用于输出一个概率分布作为注意力权重。在某些情况下，并非所有的值都应该被纳入到注意力汇聚中。例如，为了在机器翻译任务中高效处理小批量数据集，某些文本序列被填充了没有意义的特殊词元。为了仅将有意义的词元作为value来获取注意力汇聚，可以指定一个有效序列长度（即词元的个数），以便在计算softmax时过滤掉超出指定范围的位置。下面的masked_softmax函数实现了这样的掩码softmax操作，其中任何超出有效长度的位置都被掩蔽并置为0。

def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,
                              value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)

上面X代表样本输入，valid_lens代表定义的有效长度。

假设有两个2×4矩阵表示的样本，这两个样本的有效长度分别为2和3。经过掩蔽softmax操作，超出有效长度的值都被掩蔽为0。

masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))

这样，得到的输出为：第一个(2,4)的矩阵的有效长度变为2，第二个(2,4)的矩阵的有效长度变为3

同样，定义有效长度也可以使用二维张量，为矩阵样本中的每一行指定有效长度。

masked_softmax(torch.rand(2, 2, 4), torch.tensor([[1, 3], [2, 4]]))

第一个(2,4)的矩阵第一行有效长度变为1，第二行有效长度变为3；第一个(2,4)的矩阵第一行有效长度变为2，第二行有效长度变为4。