Normalization（归一化）教学文档

理解为什么深层网络需要归一化，以及 RMSNorm 如何工作

---

🤔 1. 为什么（Why）

问题场景：深层网络的数值不稳定

想象你在搭建一个 8 层的神经网络。随着层数加深，你会遇到两个致命问题：

问题 1：梯度消失

Layer 1: 激活标准差 = 1.04
Layer 2: 激活标准差 = 0.85
Layer 3: 激活标准差 = 0.62
Layer 4: 激活标准差 = 0.38
Layer 5: 激活标准差 = 0.21
...
Layer 8: 激活标准差 = 0.016  ← 几乎为 0！

当激活值越来越小，反向传播时梯度也会越来越小，最终接近 0。模型无法学习深层的特征。

问题 2：梯度爆炸

反之，如果数值越来越大，最终会超出浮点数范围，变成 NaN（Not a Number），训练直接崩溃。

---

直觉理解：水压稳定器

🚰 类比：归一化就像给每层网络装一个"水压稳定器"

• 没有稳定器：

  • 第 1 层：水压正常
  • 第 2 层：水压变小
  • 第 3 层：水压更小
  • ...
  • 第 8 层：几乎没水了（梯度消失）

• 有稳定器：

  • 每一层的输出都被"归一化"到标准范围
  • 无论输入如何变化，输出都保持稳定
  • 梯度能顺利传回每一层

---

数学本质

归一化的核心思想：控制数据分布的尺度

输入 x: 可能很大或很小，分布不稳定
       ↓
归一化: 调整到标准范围（均值 ≈ 0，标准差 ≈ 1）
       ↓
输出 x_norm: 分布稳定，便于后续计算

这样做的好处：

1. 梯度稳定：反向传播时梯度不会消失或爆炸
2. 学习率容忍度高：可以用更大的学习率，加速训练
3. 深层网络可训练：可以堆叠更多层而不崩溃

---

📐 2. 是什么（What）

RMSNorm 核心思想

全称：Root Mean Square Normalization（均方根归一化）

一句话总结：只做缩放，不减均值，比 LayerNorm 更简单高效。

---

数学定义

$$\text{RMSNorm}(x)=\frac{x}{\text{RMS}(x)}\odot \gamma$$

其中：

$$\text{RMS}(x)=\sqrt{\frac{1}{d}\sum _{i=1}^d{x}_i^2+ϵ}$$

参数说明：

• $x$：输入向量，形状 [batch_size, seq_len, hidden_dim]
• $d$：隐藏维度（hidden_dim）
• $ϵ$：防止除零的小常数（1e-5）
• $\gamma$：可学习的缩放参数，形状 [hidden_dim]
• $\odot$：逐元素乘法

---

RMSNorm vs LayerNorm 对比

特性	LayerNorm	RMSNorm
公式	$\frac{x-\mu }{\sigma +ϵ}$	$\frac{x}{\text{RMS}(x)+ϵ}$
计算步骤	1. 计算均值 2. 减均值 3. 除以标准差	1. 除以 RMS
是否减均值	✅ 是	❌ 否
参数量	$2d$ (weight + bias)	$d$ (只有 weight)
计算速度	慢	快 7-64%
半精度稳定性	较差（FP16 易下溢）	更好（BF16 友好）
使用场景	BERT, GPT-2	Llama, GPT-3+, MiniMind

---

为什么可以省略"减均值"？

理论依据：

1. 在深层网络中，经过多层变换后，激活值的均值通常已经接近 0
2. 只控制标准差（尺度）就足够稳定训练
3. 省略减均值 → 减少计算 → 更快

实验验证： 在 LLM 训练中，RMSNorm 和 LayerNorm 的最终效果相当，但 RMSNorm 更快。

---

代码实现要点

class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # 可学习参数

    def _norm(self, x):
        # 核心：x / sqrt(mean(x²) + eps)
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        # 先归一化，再乘以可学习的 weight
        return self.weight * self._norm(x.float()).type_as(x)

关键点：

• rsqrt：$1/\sqrt{x}$，比 1/sqrt(x) 更高效
• mean(-1)：在最后一维（hidden_dim）上计算均值
• keepdim=True：保持维度，便于广播
• .float()：确保计算精度（避免 FP16 下溢）
• .type_as(x)：转回原始数据类型

---

Pre-LN vs Post-LN

归一化在 Transformer Block 中有两种放置方式：

Post-LN（旧方案）

# 先计算，后归一化
x = x + Attention(x)
x = LayerNorm(x)          # 归一化在残差之后
x = x + FFN(x)
x = LayerNorm(x)

问题：

• 残差路径上的梯度会被 LayerNorm 打断
• 深层网络（>12 层）训练不稳定
• 需要非常小的学习率

---

Pre-LN（现代方案）

# 先归一化，再计算
x = x + Attention(Norm(x))  # 归一化在子层之前
x = x + FFN(Norm(x))

优势：

• ✅ 残差路径更"干净"（梯度可以直接传播）
• ✅ 每个子层的输入分布稳定
• ✅ 深层网络更容易训练
• ✅ 学习率容忍度更高

MiniMind 使用：Pre-LN + RMSNorm

---

🔬 3. 怎么验证（How to Verify）

实验 1：梯度消失可视化

目的：证明归一化的必要性

方法：

• 创建一个 10 层网络
• 对比有/无归一化的激活标准差变化
• 使用合成数据（随机张量）

运行：

python experiments/exp1_gradient_vanishing.py

预期结果：

• 无归一化：标准差从 1.0 衰减到 0.016（梯度消失）
• 有归一化：标准差保持在 1.0 左右（稳定）

输出图表：results/gradient_vanishing.png

---

实验 2：四种配置对比

目的：对比不同归一化方案的训练效果

方法：

• 训练 4 个小型 Transformer（2 层，256 hidden）
• 数据：TinyShakespeare（1MB）
• 配置：
  1. NoNorm：完全无归一化
  2. Post-LN + LayerNorm
  3. Pre-LN + LayerNorm
  4. Pre-LN + RMSNorm

运行：

python experiments/exp2_norm_comparison.py
# 快速模式（100 步）：
python experiments/exp2_norm_comparison.py --quick

预期结果：

配置	是否收敛	NaN 出现步数	最终 Loss	学习率容忍度
NoNorm	❌	~500 步	NaN	很低
Post-LN	✅	-	~3.5	低（需 LR < 1e-4）
Pre-LN + LayerNorm	✅	-	~2.8	中
Pre-LN + RMSNorm	✅	-	~2.7	高

输出图表：results/norm_comparison.png

---

实验 3：Pre-LN vs Post-LN 深度对比

目的：验证 Pre-LN 在深层网络中的优势

方法：

• 训练两个不同深度的模型（4 层 vs 8 层）
• 对比 Pre-LN 和 Post-LN 的收敛速度

运行：

python experiments/exp3_prenorm_vs_postnorm.py

预期结果：

• 4 层：Pre-LN 和 Post-LN 效果相近
• 8 层：Pre-LN 明显更稳定，Post-LN 需要更小的学习率

输出图表：results/prenorm_vs_postnorm.png

---

💡 4. 关键要点总结

核心结论

1. 为什么需要归一化：

   • 稳定深层网络的激活分布
   • 防止梯度消失/爆炸
   • 提高学习率容忍度

2. 为什么选择 RMSNorm：

   • 比 LayerNorm 更简单（省略减均值）
   • 计算更快（7-64% 提升）
   • 半精度训练更稳定

3. 为什么使用 Pre-LN：

   • 残差路径更干净
   • 深层网络更稳定
   • 训练更容易收敛

---

设计原则

在 MiniMind（以及 Llama、GPT-3 等现代 LLM）中：

# 标准 Transformer Block（Pre-LN）
def forward(self, x):
    # 第一个子层：Attention
    residual = x
    x = self.norm1(x)           # 先归一化
    x = self.attention(x)       # 再计算
    x = residual + x            # 残差连接

    # 第二个子层：FeedForward
    residual = x
    x = self.norm2(x)           # 先归一化
    x = self.feedforward(x)     # 再计算
    x = residual + x            # 残差连接

    return x

记住：Norm → Compute → Residual

---

📚 5. 延伸阅读

必读论文

• RMSNorm: Root Mean Square Layer Normalization (2019)

  • 提出 RMSNorm，证明省略减均值的有效性

• On Layer Normalization in the Transformer Architecture (2020)

  • 系统对比 Pre-LN vs Post-LN

推荐博客

• Layer Normalization 详解
• Why does normalization help?

代码实现

• MiniMind: model/model_minimind.py:95-105 - RMSNorm 实现
• MiniMind: model/model_minimind.py:359-380 - TransformerBlock 中的使用

自测题

• 📝 quiz.md - 完成 5 道选择题巩固理解

---

🎯 完成检查清单

学完本文档后，检查你是否能够：

• [ ] 用自己的话解释梯度消失问题
• [ ] 说出 RMSNorm 的数学公式
• [ ] 解释 RMSNorm 和 LayerNorm 的区别
• [ ] 解释 Pre-LN 比 Post-LN 好在哪里
• [ ] 画出 Pre-LN Transformer Block 的数据流图
• [ ] 从零实现一个 RMSNorm 类

如果还有不清楚的地方，回到实验代码，动手验证！