归一化：BatchNorm、LayerNorm、RMSNorm

最新推荐文章于 2026-05-15 14:31:19 发布

原创最新推荐文章于 2026-05-15 14:31:19 发布 · 2.1k 阅读

18 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

标签

#深度学习 #人工智能 #自然语言处理

LLM核心组件专栏收录该内容

7 篇文章

订阅专栏

前言

在神经网络中，前一层网络参数变化会导致后一层输入分布发生变化，这种现象被称为“内部协变量偏移”。这会影响模型的学习效率，导致训练不稳定甚至梯度消失/爆炸。

解决的办法就是：将每一层的输入数据做归一化，使其具有一致的分布（如均值为0，方差为1），从而让网络更容易收敛。

一、归一化的动机

动机：不同尺度的输入导致训练困难

神经网络的每一层都会学习到某种特征表示，如果前一层输出的数值跨度很大，下一层就很难学习到规律，容易造成：

梯度消失/爆炸

收敛速度慢

模型训练不稳定

参数更新效率低

想象一下，假设有一群选手（特征）在同一条笔直的跑道上比赛（参与训练），有的起点在起跑线，有的却提前了10米，有的落后了10米。很明显，这对比赛是几部公平的。这就好比网络中某些特征分布在0附近，有的却分布在100附近，学习起来就很困难。

归一化就像把所有选手都重新调整到同一条起跑线（均值为0，方差为1），减少内部协变量偏移，从相同的起点开始训练。

归一化操作可以使得各层输入的数据分布保持相对稳定。

二、BatchNorm（批量归一化）

BatchNorm对同一个通道中的所有batch计算一次均值和方差，使其均值为0，方差为1，并引入课学习的缩放和偏移参数。具体公式为：

$\mu =\frac {1}{B}\sum\limits_{i=1}^Bx_i,\sigma ^2=\frac {1}{B}\sum\limits_{i=1}^B(x_i-\mu)^2$

标准化：

$\hat x=\frac {x_i-\mu}{\sqrt{\sigma^2+\epsilon}}$

可学习的缩放 $\alpha$ 与偏移 $\beta$ ：

$y_i=\alpha \hat x_i+\beta$

假设数据输入形状为 $[2,3,2,2]$ ，其中batch_size为2，通道数为3，高和宽都为2，则做BN时，分别对每个通道里batch计算均值和方差，因为此时的通道数为3，则会有三个均值和方差，也会计算三个 $\alpha$ 和 $\beta$ 。

具体代码实现如下：

import torch
import torch.nn as nn

# 支持图像，shape为[B, C, H, W]
def batchnorm2d(x, eps=1e-05):
    alpha = nn.Parameter(torch.ones(x.shape[1]))
    beta = nn.Parameter(torch.zeros(x.shape[1]))
    mean = torch.mean(x, dim=(0, 2, 3), keepdim=True)
    var = torch.var(x, dim=(0, 2, 3), keepdim=True, unbiased=False)
    out = (x-mean)/torch.sqrt(var+eps)
    output = out * alpha.view(1, -1, 1, 1) + beta.view(1, -1, 1, 1)
    print(mean, var)
    return output
bn2d_input = torch.arange(1, 25).view(2, 3, 2, 2).to(torch.float32)
output_2d = batchnorm2d(bn2d_input)
#ptroch内置的2dBN：
bn2d = nn.BatchNorm2d(3)
output_bn2d = bn2d(x)
print(output_2d, output_bn2d)

# 支持序列，shape为[B, C, D] C这里为句子长度，也可以当做通道数，句子里的每一个单词当做一个通道。
def batchnorm1d(x, eps=1e-05):
    alpha = nn.Parameter(torch.ones(x.shape[1]))
    beta = nn.Parameter(torch.zeros(x.shape[1]))
    mean = torch.mean(x, dim=(0, 2), keepdim=True)
    var = torch.var(x, dim=(0, 2), keepdim=True, unbiased=False)
    out = (x-mean)/torch.sqrt(var+eps)
    output = out * alpha.view(1, -1, 1) + beta.view(1, -1, 1)
    return output
bn1d_input = torch.arange(1, 13).view(2, 3, 2).to(torch.float32)
output_1d = batchnorm1d(bn1d_input)
# Pytorch内置1dBN
bn1d = nn.BatchNorm1d(3)
output_bn1d = bn1d(x)
print(output_1d, output_bn1d)

BN一般不支持处理序列数据，因为BN是根据通道处理归一化的，而在序列数据中，只能把句子中的每一个单词当做一个通道。但是由于每个句子中的单词长度都不一致，因此，对于短的句子，只能填充0。这样的话，对那些填充0的通道做均值和方差就毫无意义。而且每个batch中的多个句子来说，它们同一通道里的每个单词之间并无关联，不像图片（RGB三通道）的每个通道都代表一种像素关系。

另外，BN需要额外保存均值和方差，用于推理阶段进行归一化操作。因为BN在训练和推理阶段是不一样的：训练阶段，计算当前batch的均值和方差，也就是每个batch实时计算，用于当前batch的归一化；而在推理阶段，推理时没有batch，因此要使用训练时累积的全局均值和方差来计算。即：

BN层内部会维护两个缓冲变量：

running_mean：训练过程中累计计算的全局均值

running_var：训练过程中累计计算的全局方差

它们在训练过程中更新：

$\begin{aligned} running\_mean&=(1-momentum)*running_mean + momentum*\mu_{batch} \\ running_var&=(1-momentum)*running_var+momentum*\sigma^2_{batch} \end{aligned}$

在 model.eval() 模式下，BN使用这两个变量进行归一化。

三、LayerNorm（层归一化）

LN用于处理序列数据，它是对序列数据中的单个样本的所有特征维度做归一化。具体公式为：

$\mu = \frac {1}{D}\sum\limits_{j=1}^Dx_j,\sigma^2=\frac {1}{D}\sum\limits_{j=1}^D(x_i-\mu)^2$

$D$ 为特征维度，标准化+缩放偏移：

$y_i=\alpha_j*\frac {x_j-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_j$

具体代码实现如下：

# 用于处理序列数据，shape为[B, seq_len, hidden_dim]
def layernorm(x, eps=1e-05):
    alpha = nn.Parameter(torch.ones(x.shape[-1]))
    beta = nn.Parameter(torch.zeros(x.shape[-1]))
    mean = torch.mean(x, dim=-1, keepdim=True)
    var = torch.var(x, dim=-1, keepdim=True, unbiased=False)
    print(mean, var)
    out = (x-mean)/torch.sqrt(var + eps)
    output = out * alpha + beta
    return output

x = torch.arange(1, 13).view(2, 3, 2).to(torch.float32)

output_ln = layernorm(x)
print(output_ln)
# pytroch内置LN
ln_py = nn.LayerNorm(2)
output_py = ly_ln(x)
print(output_py)

LN与batch_size无关，更适合序列数据。并且与BN不同，LN的训练和推理一致，因此推理时不需要额外保存均值和方差。

四、RMSNorm

RMSNorm 沿用了LayerNorm，相比LayerNorm来说，去掉了均值计算，减少了计算开销，但是效果基本不变。LLaMA架构就是使用了RMSNorm。

原理：

$RMS_j=\sqrt{\frac {1}{D}\sum\limits_{j=1}^Dx^2_j+\epsilon}$

$y_i=\alpha_j*\frac {x_j}{RMS_j}$

具体实现：

import torch
import torch.nn as nn
class RMSNorm(nn.Module):
    def __init__(self, hidden_dim, eps=1e-05):
        super(RMSNorm, self).__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(hidden_dim))
    def forward(self, x):
        input_dtype = x.dtype
        x = x.to(torch.float32)
        var = x.pow(2).mean(dim=-1, keepdim=True)
        x = x * torch.rsqrt(var + self.eps)
        return self.weight * x.to(input_dtype)
x = torch.arange(1, 13).view(2, 3, 2).to(torch.float32)
rms = RMSNorm(hidden_dim=2)
output = rms(x)
print(output.shape, output)