实验 06 · 张量反向与归一化 — shape 契约、LayerNorm 与优化器诊断

本实验将模型训练从标量世界推进到多维张量，聚焦训练稳定性三大支柱：Shape 契约、归一化（LayerNorm vs RMSNorm）机制，以及现代优化器（AdamW vs SGD）在病态损失面上的自适应寻优行为。

网页即时交互实验一：LayerNorm vs RMSNorm 机制对比

现代主流 LLM（如 LLaMA、DeepSeek）普遍用 RMSNorm 替代了传统的 LayerNorm。RMSNorm 抛弃了求均值与去均值平移操作，只保留均方根缩放，在节省显存带宽的同时保持了几乎相同的梯度稳定性。

你可以直接在下方调整均值 $μ$ 与方差 $σ^{2}$ ，观察两者的归一化特征差异：

交互：LayerNorm vs RMSNorm

LLaMA 之后几乎所有模型都换成了 RMSNorm。差别只有一步：要不要减均值。

输入向量 x（逗号分隔）整体平移 c = 0.0（x + c）缩放参数 γ = 1.00 LayerNorm 保留 β 偏置参数

LayerNorm

μ = 1.7000
σ² = 3.1600
σ = 1.7776
y = γ·(x−μ)/σ + β

0.169

-1.519

1.294

-0.675

0.731

参数量：10（γ + β），归约次数 2

RMSNorm

~~μ = 0（跳过）~~
mean(x²) = 6.0500
RMS = 2.4597
y = γ·x / RMS

0.813

-0.407

1.626

0.203

1.220

参数量：5（只有 γ），归约次数 1

$$\text{LN}(x)=\gamma\odot\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta,\qquad \text{RMSNorm}(x)=\gamma\odot\frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2+\epsilon}}$$

教学要点：拖动「整体平移 c」。LayerNorm 输出完全不动（平移不变）， RMSNorm 输出会跟着变 —— 它牺牲了平移不变性。实践证明这个性质没什么用，换来的是省掉一次全量归约（GPU 上约 10~15% 的 norm 层加速）和一半参数。

🐍

实验 06 · RMSNorm 均方根层归一化

抛弃求均值开销，以均方根缩放约束深层激活模长与数值稳定

↗ 独立演练台

# 实验 06 · 动手挑战：手写 RMSNorm 均方根归一化
import math
 
def rms_norm(x, gamma=1.0, eps=1e-6):
    # TODO: 1. 计算均方根 RMS = sqrt(mean(x^2) + eps)
    # TODO: 2. 返回每一个元素经过 (t / RMS) * gamma 缩放后的列表
    pass
 
x = [12.0, -8.0, 3.5, 0.2, 5.5]
out = rms_norm(x, gamma=1.0)
 
assert out is not None, "rms_norm 尚未实现"
norm_val = math.sqrt(sum(t**2 for t in out) / len(out))
print(f"归一化后模长: {norm_val:.6f} (期望接近 1.0)")
assert abs(norm_val - 1.0) < 1e-4
print("🎉 RMSNorm 前向缩放与模长约束断言 100% 通过！")

Python 3行 1, 列 117 行 · 470 字符

运行：⌘ / Ctrl + ↵

就绪 (点击运行)

点击右上角「▶ 运行」或按 ⌘/Ctrl+Enter 在浏览器端直接执行

网页即时交互实验二：病态“峡谷”地形下 AdamW vs SGD 轨迹演进

在深层 Transformer 训练中，损失面往往呈现极端狭长的“峡谷”特征（某个方向极其陡峭，另一方向极度平缓）。传统的 SGD 容易在峭壁间来回剧烈震荡无法前进；而 AdamW 通过一阶动量平滑方向并利用二阶动量自适应调整步长，能迅速沿谷底推进：

交互：AdamW vs SGD

损失面是个「峡谷」（y 方向比 x 方向陡 12 倍）。看两个优化器怎么下山。

learning rate = 0.080weight_decay = 0.020β₁ = 0.90β₂ = 0.999steps = 60

SGD — 终点 loss 0.00054（60 步）

AdamW — 终点 loss 0.01160（60 步）

最优点 (0, 0)

$$m_t=\beta_1 m_{t-1}+(1-\beta_1)g_t,\quad v_t=\beta_2 v_{t-1}+(1-\beta_2)g_t^2$$ $$\theta_t=\theta_{t-1}-\eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}+\lambda\theta_{t-1}\right)$$

教学要点（红旗实验）：把 lr 拉到 0.15 以上，SGD 会在陡峭方向剧烈震荡甚至发散，而 AdamW 因为除以 √v̂ 自动缩小了陡峭方向的步长，依然稳定。再把 β₂ 调到 0.5 —— AdamW 也开始抖了，这说明二阶动量的「记忆长度」才是它稳的原因。

交互实验指南

学习率与震荡：尝试将 SGD 的学习率适当调大，观察红色轨迹是否直接在峡谷两壁发散甚至溢出边界？
权重衰减与正则化：滑动 weight_decay 滑块，观察在解耦权重衰减（Decoupled Weight Decay）下，AdamW 如何在保持优化速度的同时将模型权重约束在紧凑半径内。

本地动手实验与单元测试

在网页端建立训练动力学直觉后，回到本地运行你的张量反传与优化器实现：

bash

cd <仓库根目录>
source .venv/bin/activate
export PYTHONPATH="$PWD/python"

# 运行张量反传与优化器单元测试
pytest python/tests/test_w06_tensor_backward.py -v

下一步

→ 返回第7章正文查看完整理论与故障注入表
→ 进入实验 07：Scaled Dot-Product 与多头注意力

实验 06 · 张量反向与归一化 — shape 契约、LayerNorm 与优化器诊断 ​

网页即时交互实验一：LayerNorm vs RMSNorm 机制对比 ​