当前位置：首页 > news >正文

LayerNorm、BatchNorm、RMSNorm（个人理解）

news 2026/7/9 13:48:56

假设输入x的形状是（batch，seqlen，dim）

1、LayerNorm

含义：对于某一个样本，统计它所有特征的表现。

def layer_norm_custom(x, eps=1e-5): # x shape: (N, L, D) # 对最后一个维度 (Dim) 求均值和方差 mean = x.mean(dim=-1, keepdim=True) # shape: (N, L, 1) var = x.var(dim=-1, unbiased=False, keepdim=True) # 归一化 x_norm = (x - mean) / torch.sqrt(var + eps) # 防止分母为0 return x_norm

注意：在LLM中，每一个token可以被认为是一个单独的样本，所以只在特征dim上进行归一化。

2、BatchNorm

含义：对于某一个特征，统计整个 batch 中所有样本的表现。

def batch_norm_custom(x, eps=1e-5): # x shape: (N, L, D) # 对 N 和 L 维度求均值和方差，保持 D 维度 mean = x.mean(dim=(0, 1), keepdim=True) # shape: (1, 1, D) var = x.var(dim=(0, 1), unbiased=False, keepdim=True) # 归一化 x_norm = (x - mean) / torch.sqrt(var + eps) return x_norm

注意：1、在LLM中，每一个token可以被认为是一个单独的样本，所以是在（batch，seqlen）上做归一化。2、为了防止推理和运行时batch不一致导致性能下降，需要在训练模式时使用批统计量（用 momentum 更新运行统计量）、推理时使用运行统计量对每个特征进行归一化。

3、为什么 Transformer 用 LN 而不用 BN？

序列长度不一：Transformer 处理的序列长度（seqlen）往往是动态的。BN 在测试时需要使用训练集计算出的“滑动平均”均值和方差，如果序列长度变化，BN 的表现会很不稳定。
Batch Size 的限制：在 NLP 任务中，Batch Size 往往受限于显存，有时会很小（甚至为 1）。在这种情况下，BN 的方差计算极不稳定，导致模型性能大幅下降。
序列相关性：LN 在单个样本内部进行归一化，不论 batch size 是多少，模型对每个样本的处理都是独立的，这更符合 Transformer 处理序列特征的需求。

注意：PyTorch 原生的 nn.LayerNorm 和 nn.BatchNorm1d 还包含两个可学习参数：γ (缩放)和 β (平移)，即 y=γ⋅x+β，这样可以保证模型在归一化后仍能恢复原始数据的表达能力。

4、Transformer 中的关键细节：Pre-LN vs Post-LN

在 Transformer 的具体实现中，LN 的位置非常关键：

Post-LN (原始 Transformer 论文)：LayerNorm(x + Sublayer(x))
- 先做残差连接和归一化，再做下一层。
- 问题：靠近输出层的梯度非常大，导致模型训练初期非常不稳定，通常需要极其精细的 Warmup 策略。
Pre-LN (现代主流，如 GPT-2, LLaMA 等)：x + Sublayer(LayerNorm(x))
- 先做归一化，再做残差连接。
- 优点：梯度可以直接通过残差连接流向更浅层，训练极其稳定，收敛更快，不需要太复杂的 Warmup。

5、RMSNorm

含义：是 LayerNorm 的简化替代，跳过均值减法，仅通过激活值的均方根进行归一化。通过可学习参数缩放，以稳定训练、加速收敛，同时保持模型表达能力。

def rms_norm(x, weight, eps=1e-6): # 第1步：计算 RMS（均方根值） rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) # 第2步：归一化并缩放 return x / rms * weight