nn.LayerNorm 实现及原理

原创 python pytorch

1. nn.LayerNorm 函数

nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None)

$$
y = \frac{x-E(x)}{\sqrt{Var(x)+\epsilon}}*\gamma+\beta
$$

• normalized_shape：归一化的维度
• eps：加在方差上的数字，避免分母为0，default=1e-5
• elementwise_affine：bool，True的话会有一个默认的affine参数，即上述公式中的$\gamma$和$\beta$，前者开始为 1，后者为 0，二者均可学习随着训练过程而变化，用来增强表征能力，default=True

2. LayerNorm 在 Transformer 中的应用

在 transformer 中一般采用 LayerNorm，LayerNorm 也是归一化的一种方法，与 BatchNorm 不同的是它是对每单个 batch 进行的归一化，而 batchnorm 是对所有 batch 一起进行归一化的

在 Transformer 中，给定输入 tensor[seq_length, batch_size, d_model]，其中 seq_length 为序列长度，batch_size 为 batch 大小，d_model 为 embedding 的维度，LayerNorm 按指定维度（通常是最后一个维度）归一化，如下所示，我们有多个 batch，其中第一个 batch 的序列长度为 3，包括“I, Love, You”，他们的词向量维度均为 6，在 LayerNorm 的时候，分别对“I”，“Love”，“You”向量进行归一化，即相同颜色的数字

2.1 LayerNorm 的官方实现

下面代码展示了 nn.LayerNorm 的官方使用，假设我们的输入如下

input = torch.arange(1, 19).view(3, 1, 6).type(torch.float32)
'''
tensor([[[ 1.,  2.,  3.,  4.,  5.,  6.]],

        [[ 7.,  8.,  9., 10., 11., 12.]],

        [[13., 14., 15., 16., 17., 18.]]])
'''

官方 nn.LayerNorm 实现的结果如下

norm = nn.LayerNorm(6)
output = norm(input)

'''
tensor([[[-1.4638, -0.8783, -0.2928,  0.2928,  0.8783,  1.4638]],

        [[-1.4638, -0.8783, -0.2928,  0.2928,  0.8783,  1.4638]],

        [[-1.4638, -0.8783, -0.2928,  0.2928,  0.8783,  1.4638]]],
       grad_fn=<NativeLayerNormBackward0>)
'''

手动计算验证一下

mean = torch.mean(input, dim=-1, keepdim=True)
std = torch.std(input, correction=0, dim=-1, keepdim=True) # 这里方差计算是除了N，不是N-1
output = (input - mean) / (std + 1e-5)

'''
tensor([[[-1.4638, -0.8783, -0.2928,  0.2928,  0.8783,  1.4638]],

        [[-1.4638, -0.8783, -0.2928,  0.2928,  0.8783,  1.4638]],

        [[-1.4638, -0.8783, -0.2928,  0.2928,  0.8783,  1.4638]]])
'''

可以发现两者结果是一致的

2.2 LayerNorm 的自定义实现

下面代码展示了 nn.LayerNorm 的自定义实现，方便理解源码

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-5):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, correction=0, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

2.3 LayerNorm 的训练与测试阶段

训练阶段（train mode）

• 对每个样本的当前激活值计算 均值 与 方差（与 batch 其他样本无关）。
• 常见归一化维度为 (hidden_size,)，即对每个 token 的 hidden 向量做标准化。
• 只有 $\gamma$ 和 $\beta$ 参与更新，训练完成后就保持不变了

测试阶段（eval mode）

• 与 BatchNorm 不同，LayerNorm 在训练与推理阶段行为一致。
• 始终使用当前输入的统计量（mean/var）
• 即训练推理的行为完全一致

model.train()  # 与
model.eval()   # 在 LayerNorm 的计算上没有区别

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Oct 09, 2025 11:10 pm