BN / LN / RMSNorm
BN / LN / RMSNorm 归一化方法总结
一、背景与动机
深度网络训练中常见问题：
- 梯度消失 / 梯度爆炸
- 不同层输入分布变化（Internal Covariate Shift）
- 收敛慢、训练不稳定
👉 归一化（Normalization）的核心目标：
将特征标准化到稳定分布，加速训练并提升模型稳定性。BN、LN 和 RMSNorm 的本质区别在于归一化维度与是否中心化，其中 BN 依赖 batch，LN 按特征归一化，而 RMSNorm 仅做尺度归一化以提升效率与稳定性。
统一形式：
x ^ = x − μ σ 2 + ϵ \hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}x^=σ2+ϵx−μ
再进行仿射变换：
y = γ x ^ + β y = \gamma \hat{x} + \betay=γx^+β
二、Batch Normalization（BN）
1. 核心思想
对batch 维度 + 空间维度做归一化（常用于 CNN）：
μ c = E B , H , W [ x ] \mu_c = \mathbb{E}_{B,H,W}[x]μc=EB,H,W[x]
σ c 2 = Var B , H , W [ x ] \sigma_c^2 = \text{Var}_{B,H,W}[x]σc2=VarB,H,W[x]
x ^ = x − μ c σ c 2 + ϵ \hat{x} = \frac{x - \mu_c}{\sqrt{\sigma_c^2 + \epsilon}}x^=σc2+ϵx−μc
2. 推理阶段（重要）
使用滑动平均：
μ r u n n i n g = ( 1 − m ) μ + m μ b a t c h \mu_{running} = (1-m)\mu + m\mu_{batch}μrunning=(1−m)μ+mμbatch
σ r u n n i n g 2 = ( 1 − m ) σ 2 + m σ b a t c h 2 \sigma^2_{running} = (1-m)\sigma^2 + m\sigma^2_{batch}σrunning2=(1−m)σ2+mσbatch2
3. 特点
- 依赖 batch size
- 训练 / 推理行为不同
- 适合 CNN
三、Layer Normalization（LN）
1. 核心思想
对单个样本的特征维度做归一化（Transformer 常用）：
μ = E C [ x ] \mu = \mathbb{E}_{C}[x]μ=EC[x]
σ 2 = Var C [ x ] \sigma^2 = \text{Var}_{C}[x]σ2=VarC[x]
x ^ = x − μ σ 2 + ϵ \hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}x^=σ2+ϵx−μ
2. 特点
- 不依赖 batch size
- 训练 / 推理一致
- 适合 NLP / Transformer
四、RMSNorm（Root Mean Square Norm）
1. 核心思想
只做缩放，不减均值：
R M S ( x ) = E [ x 2 ] RMS(x) = \sqrt{\mathbb{E}[x^2]}RMS(x)=E[x2]
x ^ = x E [ x 2 ] + ϵ \hat{x} = \frac{x}{\sqrt{\mathbb{E}[x^2] + \epsilon}}x^=E[x2]+ϵx
y = γ x ^ y = \gamma \hat{x}y=γx^
2. 特点
- 去掉均值中心化（更简单）
- 计算更快
- 在大模型中表现良好（如 LLaMA）
五、三者对比（面试重点）
| 方法 | 归一化维度 | 是否减均值 | 是否依赖Batch | 训练/推理差异 | 典型应用 |
|---|---|---|---|---|---|
| BN | B, H, W | ✅ | ✅ | ✅ | CNN |
| LN | C | ✅ | ❌ | ❌ | Transformer |
| RMSNorm | C | ❌ | ❌ | ❌ | 大模型 |
六、本质区别总结
1. 归一化维度不同
- BN：跨样本
- LN / RMSNorm：单样本
2. 是否中心化（减均值）
- BN / LN：有
- RMSNorm：无
3. 数学表达差异
- BN / LN：
x − μ σ \frac{x - \mu}{\sigma}σx−μ
- RMSNorm：
x E [ x 2 ] \frac{x}{\sqrt{\mathbb{E}[x^2]}}E[x2]x
七、代码实现
# NOTE BN/LN/RMSNormimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFclassBatchNorm(nn.Module):# BN is usually used for CNN, and the input dimensions are B, C, H, W.def__init__(self,channels_dim,eps=1e-5,momentum=0.1):super().__init__()self.eps=eps self.momentum=momentum# NOTE: momentum is the update speed of running_mean and running_varself.register_buffer('running_mean',torch.zeros(1,channels_dim,1,1))self.register_buffer('running_var',torch.ones(1,channels_dim,1,1))self.gamma=nn.Parameter(torch.ones(1,channels_dim,1,1))self.beta=nn.Parameter(torch.zeros(1,channels_dim,1,1))defforward(self,x):ifself.training:mean=x.mean(dim=[0,2,3],keepdim=True)# B,C,H,W -> 1,C,1,1var=x.var(dim=[0,2,3],keepdim=True,unbiased=False)# B,C,H,W -> 1,C,1,1# update running statsself.running_mean=(1-self.momentum)*self.running_mean+self.momentum*mean self.running_var=(1-self.momentum)*self.running_var+self.momentum*varelse:mean=self.running_mean var=self.running_var x_normed=(x-mean)/torch.sqrt(var+self.eps)out=self.gamma*x_normed+self.betareturnoutclassLayerNorm(nn.Module):# LN is usually used for RNN/Transformer, and the input dimensions are B, L, C.def__init__(self,channels_dim,eps=1e-5):super().__init__()self.eps=eps self.gamma=nn.Parameter(torch.ones(1,1,channels_dim))self.beta=nn.Parameter(torch.zeros(1,1,channels_dim))defforward(self,x):mean=x.mean(dim=-1,keepdim=True)# B,L,C -> B,L,1var=x.var(dim=-1,keepdim=True,unbiased=False)# B,L,C -> B,L,1x_normed=(x-mean)/torch.sqrt(var+self.eps)out=self.gamma*x_normed+self.betareturnoutclassRMSNorm(nn.Module):# RMSNorm is a variant of LN, which only normalizes the variance and does not normalize the mean.# It is usually used for RNN/Transformer, and the input dimensions are B, L, C.def__init__(self,channels_dim,eps=1e-5):super().__init__()self.eps=eps self.gamma=nn.Parameter(torch.ones(1,1,channels_dim))defforward(self,x):rms=torch.mean(x**2,dim=-1,keepdim=True)# B,L,C -> B,L,1x_normed=x/torch.sqrt(rms+self.eps)out=self.gamma*x_normedreturnoutif__name__=="__main__":x=torch.rand(10,5,768)LN=LayerNorm(768)x_LN=LN(x)print(x_LN.shape)RMSN=RMSNorm(768)x_RMS=RMSN(x)print(x_RMS.shape)cnn_x=torch.rand(4,12,512,512)BN=BatchNorm(12)x_BN=BN(cnn_x)print(x_BN.shape)