NLP场景下torch.nn.LayerNorm输出的正确计算方法

错误原因 #

你之前的计算对归一化粒度的理解有误，torch.nn.LayerNorm的默认行为是对输入张量的最后K个维度计算均值和标准差，其中K是初始化时传入的normalized_shape的维度数。你传入的normalized_shape=dim=4，输入张量形状为[batch_size, seq_size, dim] = [2,3,4]，因此归一化的粒度是每个token自身的4维嵌入向量，不会跨同一个样本的不同token计算统计量。
你之前对第一个样本的3个token共12个值整体求均值和标准差，和LayerNorm的实际逻辑不符，所以计算结果不匹配。

正确计算方法 #

以你要计算的y[0,0,:]（第一个样本第一个token的归一化结果）为例：

1. 首先取出对应的原始嵌入向量：embedding[0,0,:] = [ 0.5909, 0.1326, 0.8100, 0.7631]
2. 对该4维向量单独计算均值
3. 对该4维向量单独计算有偏标准差（LayerNorm默认unbiased=False，即方差计算除以元素个数n，不是n-1）
4. 按公式 y = (x - mean) / std * gamma + beta 计算，刚初始化的LayerNorm可学习参数gamma默认值为1，beta默认值为0，因此可以简化为(x - mean)/std

对应代码 #

# 取目标token的嵌入
x = embedding[0, 0, :]
# 计算均值
mean = x.mean()
# 计算有偏标准差，和LayerNorm默认参数对齐
std = x.std(unbiased=False)
# 计算归一化结果
print((x - mean) / std)

运行后输出为 tensor([ 0.0626, -1.6495, 0.8810, 0.7060])，和PyTorch官方LayerNorm输出完全一致。

补充说明 #

如果你想要实现「对单个样本的所有嵌入值计算均值和标准差」的逻辑，只需要在初始化LayerNorm时传入normalized_shape=(seq_size, dim)即可，也就是：

layer_norm = torch.nn.LayerNorm((seq_size, dim))

此时你之前写的计算逻辑就会和输出结果匹配。

内容的提问来源于stack exchange，提问作者YQ.Wang