LayerNorm 的位置决定初始化时的梯度尺度

BLEU 8.45 与 34 之间的一个 warm-up

在 IWSLT14 德译英上训练一个 6 层编码器、6 层解码器的标准 Transformer，优化器是 Adam，峰值学习率 $10^{-3}$ 或 $5\times10^{-4}$。去掉学习率预热（warm-up），BLEU 只有 8.45；加上 4000 步预热，BLEU 到 34 左右。模型、数据、峰值学习率都没动，差别只在前几千步学习率怎么走。【文中有据，第 4 页及图 2(a)】

预热的写法很简单，前 $T_{\text{warmup}}$ 步学习率从 0 线性升到峰值：

$$\mathrm{lr}(t)=\frac{t}{T_{\text{warmup}}}\,\mathrm{lr}_{\max},\qquad t\le T_{\text{warmup}}$$

代入数字看它在做什么。$T_{\text{warmup}}=4000$、$\mathrm{lr}_{\max}=10^{-3}$ 时，第 1 步的学习率是 $2.5\times10^{-7}$，第 400 步是 $10^{-4}$，到第 4000 步才用上峰值。也就是说，前几百步模型几乎是原地不动地在看梯度。把预热缩到 500 步，$\mathrm{lr}_{\max}=5\times10^{-4}$ 时 BLEU 是 31.16，$10^{-3}$ 时掉到 2.77。预热长度和峰值学习率绑在一起，单调一个都不行。

论文图 2（第 4 页）左两幅是 Adam，右两幅是 SGD。该看的是每幅图里那几条贴着底部的 BLEU 曲线，它们全是去掉预热或预热太短的配置。原文没有说 8.45 对应哪一个峰值学习率。

这个现象在当时有两种流行解释。一种来自大批量训练（Goyal et al. 2017）：批量很大时，一开始用大学习率会把模型推到坏区域，所以要预热。但这里每批只有 4096 个 token，谈不上大批量。另一种来自 RAdam（Liu et al. 2019a，On the Variance of the Adaptive Learning Rate and Beyond）：Adam 的二阶矩估计 $v_t$ 在最初几步只见过很少几个梯度，方差很大，导致自适应学习率忽大忽小，预热是在等 $v_t$ 估准。作者用一个对照把第二种解释推开了一半：换成没有二阶矩的普通 SGD，不预热同样失败，15 个 epoch 里 BLEU 刚过 0。既然不用 Adam 也要预热，原因就不只在 Adam 身上。【文中有据，第 4 至 5 页】

这个 SGD 对照的说服力要打个折扣。图 2(b) 里带预热的 SGD 也远低于 Adam，SGD 训 Transformer 本来就训不好，它失败的原因可能不止一个。后面讲到 Adam 时还会回到这里。

作者押注的是另一个方向：预热只发生在前几千步，那就去看第 0 步，也就是初始化那一刻，梯度长什么样。

这背后有三张牌。第一张是一个几乎无法反驳的观察，一个调度超参数让 BLEU 在 8.45 和 34 之间来回。第二张是工具，同一批作者（Di He、Liwei Wang、Tie-Yan Liu）前一年刚用动力系统的视角分析过 Transformer（Lu et al. 2019），平均场理论（mean field theory）分析初始化的路数他们是熟的。第三张最硬：要推荐的改法 Pre-LN 早就存在，Tensor2Tensor 里有实现，fairseq 里是一个命令行开关，任何人不改一行代码就能试。所以这篇论文的贡献不是发明 Pre-LN，而是解释它为什么好训，并据此把预热拿掉。同期 Nguyen & Salazar（2019，Transformers without Tears）也从实验上发现 Pre-LN 不预热能训，这篇论文多的是理论。【推断，依据：引言第 2 页、相关工作第 2 页】

两种摆法

Transformer 的每一层有两个子层，自注意力和前馈网络（FFN），每个子层外面都有一条残差连接和一个 LayerNorm（下文简称 LN）。论文要比较的只是 LN 放在哪里。写成两行：

$$\text{Post-LN:}\quad x\leftarrow \role{norm}{\mathrm{LN}}\big(\role{in}{x}+\role{op}{F}(\role{in}{x})\big)\qquad\qquad \text{Pre-LN:}\quad x\leftarrow \role{in}{x}+\role{op}{F}\big(\role{norm}{\mathrm{LN}}(\role{in}{x})\big)$$

$\role{op}{F}$ 依次是注意力和 FFN。Post-LN 是 2017 年原版 Transformer 和 BERT 的写法，先相加再归一化。Pre-LN 把 LN 挪到子层入口，先归一化再进子层，子层输出直接加回主干；因为主干从头到尾没有被归一化过，Pre-LN 在整个网络的最后还要再加一个 LN，再接输出层。

图 1上一行 Post-LN，下一行 Pre-LN，都只画一层，$N\times$ 表示重复 $L$ 次。点下方的箭头一步步看从损失回到最后一层 $W^2$ 的路径；点 Pre-LN 那个紫色的 Final LN，会收起或展开它的内部。

把这张图当成一条主干线来看，差别一下就出来了。Post-LN 的主干上每隔一个子层就坐着一个 LN，信号和梯度每过一层都要穿过两次归一化。Pre-LN 的主干上只有加号，从词嵌入一路加到最后的 LN，中间没有任何东西改变主干的尺度；子层像挂在主干旁边的支路，只往主干上加东西。

后面的推导都在这两种主干上做。为了能算，论文做了三处简化（第 5 页）：只看单头注意力；所有权重矩阵都取 $d\times d$；把注意力里的 $W^Q$、$W^K$ 初始化为零，这样初始时每个位置对所有位置的注意力都是均匀的 $1/n$，注意力子层就退化成对整个序列的 $\role{in}{x}W^V$ 取平均。前两处只影响常数，第三处是为了让引理 2 能算出来，它把注意力在初始化时可能带来的数据相关的混合全部去掉了。真实训练里 $W^Q$、$W^K$ 用的是 Xavier 初始化，不是零。

顺带纠正论文引言里一处引用。第 1 页把 GPT-2（Radford et al. 2019）列为基于 Post-LN 的预训练模型，但 GPT-2 论文第 2.3 节写明它把 LN 挪到了每个子层的输入处，并在最后一个注意力块后加了一个 LN，这正是 Pre-LN。换句话说，到 2019 年最大的生成式模型已经在用 Pre-LN 了。【文中有据，本文第 1 页与 GPT-2 论文第 2.3 节】

LayerNorm 的雅可比

要知道梯度穿过 LN 会被放大还是缩小，先看 LN 在几何上做了什么。取 $\gamma=1$、$\beta=0$（初始化时就是这样）：

$$\role{norm}{\mathrm{LN}}(x)=\frac{x-\mu\mathbf 1}{\sigma},\qquad \mu=\frac1d\sum_k x_k,\qquad \sigma^2=\frac1d\sum_k (x_k-\mu)^2$$

分两步看。减去均值，是把 $x$ 投影到和全 1 向量垂直的那个超平面上，记投影后的向量为 $y$。除以 $\sigma$，等于除以 $\lVert y\rVert/\sqrt d$，于是输出的长度永远是 $\sqrt d$。所以 LN 的输出总落在一个半径为 $\sqrt d$ 的球面上，输入只决定落在球面的哪个方向。

这带出一个直接的推论：把输入整体放大 10 倍，输出一点不变，$\mathrm{LN}(10x)=\mathrm{LN}(x)$。现在请读者先想一下，同一个小扰动 $\delta$ 加在一个短的 $x$ 上和加在一个长的 $x$ 上，哪一次输出变得多？

答案是短的那次。$x$ 越长，同样大小的 $\delta$ 让它的方向转过的角度越小，而 LN 的输出只关心方向。转过的角度约等于 $\lVert\delta\rVert/\lVert y\rVert$，再乘上球的半径 $\sqrt d$，就是输出的变化量。这就是论文引理 3 的内容：

$$J_{\mathrm{LN}}(x)=\frac{\sqrt d}{\lVert y\rVert_2}\Big(I-\frac{y y^\top}{\lVert y\rVert_2^2}\Big)\Big(I-\frac1d\mathbf 1\mathbf 1^\top\Big),\qquad \lVert J_{\mathrm{LN}}(x)\rVert_2=O\Big(\frac{\sqrt d}{\lVert x\rVert_2}\Big)$$

后面两个括号都是投影矩阵，特征值只有 0 和 1：一个去掉沿 $y$ 方向的分量（沿半径方向推，输出不变），一个去掉沿全 1 方向的分量（所有坐标一起加同一个数，减均值时被消掉）。真正决定大小的只有前面那个标量 $\sqrt d/\lVert y\rVert$。反向传播时梯度乘的是 $J_{\mathrm{LN}}^\top$，同一个标量也乘在梯度上。

所以只要知道每个 LN 的输入有多长，就知道梯度穿过它时被缩放了多少。下一节就去算这个长度，两种摆法在这里分道扬镳。

残差流的粗细

先准备一个小工具，论文的引理 1：一个各坐标独立、服从 $N(0,\sigma^2)$ 的 $d$ 维向量过一次 ReLU，长度平方的期望从 $\sigma^2 d$ 变成 $\tfrac12\sigma^2 d$。理由一句话：每个坐标有一半概率被置零，另一半原样保留，而正半边的二阶矩正好是整体的一半。

再看一个 FFN 子层对长度的作用。输入长度平方为 $d$，权重按 $N(0,1/d)$ 初始化，过 $W^1$ 之后每个坐标方差为 1，过 ReLU 剩一半，过 $W^2$ 长度平方期望为 $d/2$。支路输出和主干输入在期望上不相关（交叉项里有一个均值为零的随机矩阵），所以相加时长度平方直接相加。

有了这两步，引理 2 就是记账。

Post-LN。每一层的输入刚被 LN 过，长度平方恰好是 $d$。注意力子层在均匀注意力的简化下加上一个均值向量，再过 LN 又回到 $d$；FFN 子层加上 $d/2$，得到 $\mathbb E\lVert x^{\text{post},5}_{l}\rVert^2=\tfrac32 d$，然后被最后那次 LN 拉回 $d$。每一层都一样，和 $l$ 无关。

Pre-LN。主干只做加法，没人把它拉回去。每一层往主干上加两个支路：注意力支路贡献 0 到 $d$ 之间（取决于序列里各向量平均之后还剩多长），FFN 支路贡献 $d/2$。从 $\lVert x_0\rVert^2=d$ 起步，一层一层累加：

$$\Big(1+\frac l2\Big)d\ \le\ \mathbb E\lVert \role{in}{x^{\text{pre}}_{l}}\rVert^2\ \le\ \Big(1+\frac{3l}{2}\Big)d$$

第 6 层时在 $4d$ 到 $10d$ 之间，第 24 层时在 $13d$ 到 $37d$ 之间。残差流随深度线性变粗。

我用文末的最小实现把这两条直接量了一遍（$d=512$，6 层，单头，Xavier 初始化，$d_{ff}=2d$，5 个随机种子）。Post-LN 的 $\lVert x^{\text{post},5}_l\rVert^2/d$ 六层依次是 1.45、1.45、1.47、1.42、1.46、1.50，种子间标准差不超过 0.04。理论值 1.5 是在 $d\times d$、方差 $1/d$ 的设定下算的；换成我用的 $d_{ff}=2d$ 加 Xavier，同样的记账给出 $1+4/9\approx1.44$，和实测对得上。Pre-LN 的 $\lVert x_l\rVert^2/d$ 依次是 1.52、2.15、2.89、3.74、4.69、5.75，每层大约加 0.85，落在引理 2 的区间里，靠近下沿。

图 2Pre-LN 的六层加最后的 LN，从下往上叠。指针在图上上下移动，选中的一层亮起轮廓，它上面的层抬起来；读数行给出这一层输出的残差流长度（实测）和梯度穿过最后那个 LN 时要乘的倍数。

现在把这一节和上一节接起来。Pre-LN 里，从损失回到任何一个参数的梯度，都要先穿过最后那个 LN，而它的入口是全网络最粗的那段残差流。按引理 3，梯度在这里被乘上 $\sqrt d/\lVert x_L\rVert\approx 1/\sqrt{1+cL}$。6 层时实测是 $1/\sqrt{5.75}=0.42$。层数越多，这个倍数越小，而且它一次性乘在所有参数的梯度上。Post-LN 的最后一个 LN 的入口永远是 $\tfrac32 d$ 左右，倍数固定在 $\sqrt{2/3}\approx0.82$，不随层数变。

这就是整篇论文的机制。它的漂亮之处在于两个引理都只是记账，没有任何一步依赖注意力的具体行为；它的局限也在这里，后面会看到，这笔账只记到了最后一层。

定理 1 里的两个上界

有了引理 2 和引理 3，论文的主定理就是把它们代进链式法则。定理 1 只管最后一层 FFN 的第二个矩阵 $\role{param}{W^{2,L}}$：

$$\text{Post-LN:}\ \Big\lVert \frac{\partial\role{loss}{\tilde{\mathcal L}}}{\partial \role{param}{W^{2,L}}}\Big\rVert_F\le O\big(d\sqrt{\ln d}\big)\qquad\qquad \text{Pre-LN:}\ \Big\lVert \frac{\partial\role{loss}{\tilde{\mathcal L}}}{\partial \role{param}{W^{2,L}}}\Big\rVert_F\le O\Big(d\sqrt{\frac{\ln d}{L}}\Big)$$

成立的概率至少是 $0.99-\delta-\epsilon/(0.9+\epsilon)$，前提是隐状态的长度满足一个集中性条件。

证明的骨架

拿 $W^{2,L}$ 的第 $p$ 行第 $q$ 列这一个元素来看。它对损失的影响只走一条路：改变 FFN 输出的第 $q$ 个坐标，穿过一个 LN，到达损失。于是这一个元素的梯度是三个因子相乘（附录 E 式 34 到 38）：

三个因子相乘，再对 $d^2$ 个元素求平方和开根号。Post-LN 得到 $\sqrt{d^2\cdot \tfrac{d}{(3/2)d}\cdot 2\ln 100d}=O(d\sqrt{\ln d})$；Pre-LN 的第二个因子多了 $1/(1+L/2)$，于是多出 $1/\sqrt L$。

代入 $d=512$ 看看这个界松到什么程度。第三个因子的上界是 $\sqrt{2\ln(51200)}\approx4.66$，而一个 ReLU 单元的典型大小是 $\sqrt{\mathbb E[\mathrm{ReLU}(z)^2]}=\sqrt{0.5}\approx0.71$。证明给每个元素都按最坏的那个隐藏单元算，单这一项就松了 6 倍多。再看概率：论文实测的集中性是 $(0.1,0.125)$-bounded（第 6 页），代进去 $0.99-0.125-0.1/1.0=0.765$。所以定理 1 在作者自己的实验设定下，是一个以约 77% 的概率成立、常数没有给出的上界。

$(\epsilon,\delta)$-bounded 本身的意思不复杂：随机变量以至少 $1-\delta$ 的概率落在期望的 $\pm\epsilon$ 倍之内。它在证明里只起一个作用，把引理 2 的期望长度换成这一次实际的长度，好让引理 3 能用。

上界和上界相比

论文从定理 1 得出的结论是 Post-LN 输出层附近的梯度大，Pre-LN 的小得多（第 6 页）。这一步在逻辑上有个缺口：两个都是上界，上界大不代表实际值大。要说 Post-LN 的梯度大，需要一个下界，论文没有给。定理 1 能严格支持的只有一半：Pre-LN 最后一层的梯度不会超过一个随 $L$ 变小的数。【文中有据，定理 1 只有不等号】

所以实验验证在这里承担了比通常更重的分量。论文图 3(c)(d) 量了 6 到 14 层的模型，Post-LN 最后一层 $W^2$ 的梯度一直在 1.6 附近，Pre-LN 从 0.31 降到 0.125。这两条趋势我在最小实现里都复现了（图 3 右），而且两种统计量都一样，比较可信。

论文图 3（第 7 页）(a)(b) 是 6-6 模型各层 FFN 两个矩阵的梯度，橙色 Post-LN、蓝色 Pre-LN、绿色是 Post-LN 预热之后；(c)(d) 是最后一层 $W^2$ 随模型层数的变化。注意 (c)(d) 的纵轴量级差了约 4 倍，而且两幅图都没有误差棒，虽然正文说重复了 10 个种子。

但有两处细节和定理对不上。

第一处是下降的速度。$1/\sqrt L$ 从 6 层到 14 层应该降到原来的 0.65，论文图 3(c) 降到了 0.40，我量到的是 0.23。实际下降比定理快得多。这和上界不矛盾，但说明除了最后那个 LN，还有别的东西随层数在变。一个候选是输出层看到的向量里，词嵌入本身所占的份额：Pre-LN 的主干直接带着 $x_0$ 走到输出，层数越多，它在最后被 LN 归一化之后剩下的比例越小。【猜测】

第二处是固定层数时谁大谁小。6 层时我量到的 Pre-LN 最后一层梯度反而比 Post-LN 大（每批范数 1.63 对 1.09），8 层以后才反过来。论文图 3 在 6 层时是 Post-LN 1.6 对 Pre-LN 0.31。两边设定不同（我是单头、只有编码器、$d_{ff}=2d$、目标是 $t\mapsto t+1$ 的映射），所以这不是对论文的否定，而是提醒：$O(\cdot)$ 藏起来的常数，在 6 层这种浅模型上足以决定谁大。论文的结论在更深的模型上才稳。

图 3最小实现的实测，5 个种子的均值，(a) 带标准差。(a) 6 层模型各层 $W^2$ 的梯度，统计量和论文图 3(b) 相同（先对 20 批梯度取平均，再取 F 范数）。(b) 残差流长度，阴影是引理 2 的区间。(c) 最后一层 $W^2$ 的梯度随总层数的变化；实线是每批梯度范数的平均，也就是定理 1 管的量，虚线是论文的统计量；点线是从 Pre-LN 6 层的值按 $1/\sqrt L$ 外推。

梯度沿深度的分布

定理 1 只覆盖最后一层。论文在 3.3 节末尾说，对其他层做了一些初步分析，主要结论是 Post-LN 的梯度在输出附近大、随层号 $l$ 变小很可能衰减，Pre-LN 的很可能各层差不多。措辞从定理 1 的 we prove 降到了 preliminary 和 likely，作者自己对这部分的把握明显弱一档。【文中有据，第 6 页】

附录 F 的论证是这样的：梯度从输出回到第 $l$ 层，要穿过上面 $L-l$ 层，每层都有两个 LN；FFN 后面那个 LN 的入口长度平方是 $\tfrac32 d$，按引理 3 每过一次乘 $\sqrt{2/3}$，于是得到一个 $(2/3)^{(L-l)/2}$ 的指数衰减项。

这里漏了一个因子。梯度往下走时，穿过的不只是 LN，还有残差相加 $x+F(x)$ 的雅可比 $I+\partial F/\partial x$。前向时 FFN 支路让长度平方从 $d$ 涨到 $\tfrac32 d$，用的是引理 1 的记账：随机矩阵乘一次、ReLU 砍一半。反向时梯度穿过同一条支路，乘的是同样分布的随机矩阵的转置、同样一半的 ReLU 掩码，长度平方也涨到原来的 $\tfrac32$ 倍。前向多出来的这一份，正是 LN 要除掉的那一份；反向多出来的这一份，正好补上 LN 除掉的。一层下来是 $\tfrac32\times\tfrac23=1$，不衰减。

论文的式 (62) 其实把这个因子写进去了，只是估计时只留下 LN 那一项。所以在论文自己的理想化设定下，初始化时的 Post-LN 梯度沿深度应该是平的。我的实测也是平的：6 层的 $W^2$ 梯度依次是 0.54、0.51、0.49、0.50、0.52、0.52（论文统计量），每批范数在 1.05 到 1.17 之间（图 3 左）。

这和论文图 3(a)(b) 不一致，那里 Post-LN 从第 1 层的 0.13 一路涨到第 6 层的 1.65。差别从哪来，我没法从论文里判断。论文的模型是 6 层编码器加 6 层解码器，图里的第 1 到 6 层没说是编码器还是解码器；编码器的梯度要先穿过整个解码器的交叉注意力才能回来；真实数据、dropout 0.1、标签平滑也都和我的玩具设定不同。【猜测】我能确定的只有一点：附录 F 的指数衰减论证本身不成立，图 3(a)(b) 的递增需要别的解释。

Pre-LN 这边，我量到的是从底层到顶层下降 2.5 倍（3.77 到 1.49），论文图 3(b) 是 0.57 到 0.30，下降约 1.9 倍。方向一致，但正文说各层几乎一样，图上是差了将近一倍。原因也来自残差相加的雅可比：Pre-LN 的梯度往下走时每层都加上一份支路贡献，而支路入口的 LN 在底层残差流细的地方增益更大，所以越往下累积得越多。

下面这段代码就是上面所有实测数字的来源。它对应正文的表 1 和 3.3 节的初始化设定，图 3 的三幅子图都由它产生。

# 按论文第 3 节和表 1 实现，非官方代码。只看初始化那一刻：Post-LN 和 Pre-LN 各层 W2 的梯度有多大。
import sys, json, torch, torch.nn as nn, torch.nn.functional as F
d, n, V, B = 512, 32, 1000, 8          # 宽度、序列长、词表、每批句数

class Layer(nn.Module):
    def __init__(s, pre):
        super().__init__()
        s.pre = pre
        s.Wq, s.Wk, s.Wv = (nn.Linear(d, d) for _ in range(3))
        s.W1, s.W2 = nn.Linear(d, 2 * d), nn.Linear(2 * d, d)   # IWSLT 配置：512 -> 1024 -> 512
        s.ln1, s.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        for m in (s.Wq, s.Wk, s.Wv, s.W1, s.W2):
            nn.init.xavier_normal_(m.weight); nn.init.zeros_(m.bias)   # 论文 3.3 节：Xavier，偏置为 0

    def attn(s, x):                                     # x: (B, n, d)
        a = s.Wq(x) @ s.Wk(x).transpose(1, 2) / d ** .5  # (B, n, n)；单头，和定理的设定一致
        return a.softmax(-1) @ s.Wv(x)

    def forward(s, x):
        ffn = lambda h: s.W2(F.relu(s.W1(h)))
        if s.pre:                                       # 表 1 右栏：LN 在残差支路里面，主干上只有加法
            x = x + s.attn(s.ln1(x))
            return x + ffn(s.ln2(x))
        x = s.ln1(x + s.attn(x))                        # 表 1 左栏：每次相加之后都过一次 LN
        s.x5 = x + ffn(x)                               # 引理 2 里的 x^{post,5}，留着看范数
        return s.ln2(s.x5)

class Net(nn.Module):
    def __init__(s, L, pre):
        super().__init__()
        s.E = nn.Parameter(torch.randn(V, d) / d ** .5)  # fairseq 的做法：嵌入按 N(0, 1/d) 初始化
        s.layers = nn.ModuleList(Layer(pre) for _ in range(L))
        s.final = nn.LayerNorm(d) if pre else nn.Identity()   # Pre-LN 多出来的最后一个 LN，梯度缩放全靠它

    def forward(s, tok, tgt):
        x, s.norms = s.E[tok] * d ** .5, []             # 输入再乘 sqrt(d)，和论文假设一致：|x0|^2 约等于 d
        for layer in s.layers:
            x = layer(x)
            h = x if layer.pre else layer.x5            # Post-LN 的 x_l 被 LN 钉在 d 上，看 x5 才有信息
            s.norms.append((h.pow(2).sum(-1).mean() / d).item())    # 引理 2 预测：Post 约 1.5，Pre 随层数线性涨
        logits = s.final(x) @ s.E.t()                   # 输出层和词嵌入共享；LN 之后 |h|^2 = d，logits 标准差约 1
        return F.cross_entropy(logits.reshape(-1, V), tgt.reshape(-1))

def grads(L, pre, batches=20):
    net = Net(L, pre)
    acc = [torch.zeros_like(l.W2.weight) for l in net.layers]; each = [0.] * len(net.layers)
    for _ in range(batches):                            # 论文 3.4 节的做法：先对各批的梯度取平均，再取 F 范数
        net.zero_grad()
        tok = torch.randint(V, (B, n)); tgt = (tok + 1) % V   # 目标有规律，平均梯度才不是纯噪声
        net(tok, tgt).backward()
        for k, l in enumerate(net.layers):
            acc[k] += l.W2.weight.grad / batches
            each[k] += l.W2.weight.grad.norm().item() / batches   # 定理 1 管的是这个：单批梯度的范数
    return [a.norm().item() for a in acc], net.norms, each

if __name__ == '__main__':                              # 每个配置 5 个种子，报均值和标准差
    ms = lambda xs: [round(float(torch.tensor(c).mean()), 3) for c in zip(*xs)] + [round(float(torch.tensor(c).std()), 3) for c in zip(*xs)]
    out = {}
    for pre in (False, True):
        k = 'pre' if pre else 'post'
        runs = [(torch.manual_seed(s), grads(6, pre))[1] for s in range(5)]
        out[k] = {'W2_per_layer': ms([r[0] for r in runs]), 'W2_per_batch': ms([r[2] for r in runs]), 'norm2/d': ms([r[1] for r in runs])}
        lastW2 = {L: [(torch.manual_seed(s), grads(L, pre, 10))[1] for s in range(5)] for L in (6, 8, 10, 12, 14)}
        out[k + '_lastW2_vs_L'] = {L: ms([[r[0][-1], r[2][-1]] for r in rs]) for L, rs in lastW2.items()}
    json.dump(out, sys.stdout, indent=1)

梯度尺度和 Adam 之间隔着一层

到这里为止，整篇论文的论证都是在比较梯度的大小。可是所有实验用的都是 Adam，而 Adam 的更新量几乎不看梯度的大小。

算一下 Adam 的第一步。$m_1=(1-\beta_1)g$，$v_1=(1-\beta_2)g^2$，做完偏差校正 $\hat m_1=g$，$\hat v_1=g^2$，于是

$$\Delta\theta_1=-\mathrm{lr}\cdot\frac{\hat m_1}{\sqrt{\hat v_1}+\epsilon}=-\mathrm{lr}\cdot\frac{g}{|g|+\epsilon}\approx-\mathrm{lr}\cdot\mathrm{sign}(g)$$

每个坐标都恰好移动 $\mathrm{lr}$。梯度是 1.6 还是 0.016，第一步走的距离一模一样（论文里 $\epsilon=10^{-9}$，完全可以忽略）。之后几步 $\hat m/\sqrt{\hat v}$ 也一直在 $O(1)$ 的量级，除非梯度的符号来回翻转。所以定理 1 里差出来的那个 $\sqrt L$，在 Adam 下不会直接变成更新步长的差别。梯度大加上学习率大，所以不稳定这句话，严格说只对 SGD 成立。论文正文没有处理这一层。【推断，依据：Adam 的更新式和论文第 4 页的优化器设定】

那 Pre-LN 的好处在 Adam 下靠什么传过去？换一个空间看。Adam 让每个参数都走大约 $\mathrm{lr}$ 这么远，那么真正决定稳定性的是参数走这么远之后，网络的输出变了多少。在 Pre-LN 里，任何一层的参数变化，要影响输出都得穿过最后那个 LN，而引理 3 说，这个 LN 对前向扰动和反向梯度乘的是同一个 $\sqrt d/\lVert x_L\rVert$。也就是说，最后那个 LN 把梯度压小 $\sqrt L$ 倍的同时，也把每一步更新在输出上造成的变化压小了同样的倍数。论文的机制在 Adam 下依然成立，只是应该在函数空间里讲，而不是在梯度空间里讲。【推断】

Post-LN 这边，Liu et al. 2020（Understanding the Difficulty of Training Transformers）量的正是这个量：同样大小的参数更新之后，输出变化在 Post-LN 里随层数线性增长，在 Pre-LN 里只按对数增长。他们把 Post-LN 的问题归结为输出过度依赖各层的残差支路，小的参数扰动被逐层放大。这个解释和本文的结论兼容，而且不需要在 Adam 下绕弯。

论文自己的 RAdam 对照也该放在这个框架里重读。图 4(a)(b) 里，RAdam 不预热就能把 Post-LN 训好（第 8 页原文：although RAdam trains the Post-LN Transformer well without the warm-up stage）。这说明修好 Adam 早期二阶矩估计的方差，单独就足以救回 Post-LN，它是对 RAdam 那派解释的支持，而作者把这一组结果读成 LN 位置的影响盖过了优化器的影响，依据是 Pre-LN 换成 RAdam 几乎没有差别。两种读法合在一起，更可能的图景是：失败需要两个条件同时成立，早期更新又大又乱（Adam 的二阶矩还没估准），并且网络输出对参数扰动敏感（Post-LN）。拿掉任何一个条件，训练都能走下去。预热、RAdam、Pre-LN 是拿掉条件的三种办法。【推断，依据：图 2(a)、图 4(a)(b)】

玩具规模上的同一个现象

把同样的对照缩到 CPU 上几分钟能跑完的规模，现象依然在。设定是 8 层、$d=128$ 的因果语言模型，数据由一条人造的马尔可夫链生成，词表 256，每个 token 只有 4 个可能的后继，所以最优损失是 $\ln 4\approx1.39$。优化器和论文一样是 Adam，$\beta=(0.9,0.98)$，训练 500 步，每个配置 2 个种子。

图 4玩具规模的训练损失，2 个种子取平均。实线不预热，虚线预热 300 步；灰色 Post-LN，紫色 Pre-LN；点线是最优损失 $\ln 4$。数字只作示意，脚本是 train_toy.py。

峰值 $10^{-3}$ 时，Post-LN 不预热的那条线 500 步里一直停在 5.38，加上 300 步预热后，它在同一个平台上待到第 250 步左右才突然往下走，第 500 步到 1.77。Pre-LN 不预热，第 100 步已经到 2.3，第 500 步到 1.57。峰值提到 $3\times10^{-3}$，Post-LN 预热 300 步也不够了，两条线都停在平台上，这和论文里预热 500 步配 $10^{-3}$ 只剩 2.77 BLEU 是同一种情况；Pre-LN 预热和不预热都照常下降。【实测，玩具设定】

有两点比曲线本身更值得记住。第一，Post-LN 的失败方式不是发散。没有 NaN，损失也没有爆掉，它只是停在一个平台上：5.38 比均匀猜测的 $\ln 256=5.55$ 略低，大约是只学会了各个 token 出现频率、完全不看上下文的水平。论文里 BLEU 8.45 而不是 0，也更像这种卡住，而不是崩掉。第二，这个小实验顺带补了下一节会说到的那个缺失对照的玩具版本：Pre-LN 加预热，终点和不加差不多（1.62 对 1.57），只是晚了大约 100 步。在这个规模上，预热对 Pre-LN 没有好处，只有时间成本。

拿掉预热的三组实验

第 4 节把结论落到三个任务上。设定差别比正文读起来要多，先列成表：

出自第 7 页 4.1 节和附录 A。所有实验都是 Adam，$\beta=(0.9,0.98)$。

结果分三条。IWSLT 上两者最终都在 BLEU 34、验证损失 4 左右，Pre-LN 第 9 个 epoch 的检查点和 Post-LN 第 15 个差不多。WMT 上趋势相同。BERT 上 Pre-LN 的验证损失更快下降，MRPC 和 RTE 两个下游任务在各个检查点上也都更高。【文中有据，图 4、图 5】

论文图 5（第 8 页）(a) BERT 预训练的验证损失，(b)(c) 各检查点微调后的下游准确率，误差棒来自 6 个种子。注意 (a) 里 Pre-LN 的曲线在约 80 万步就停了，Post-LN 一直画到 100 万步。

正文里一处写反的数字

第 8 页的原话是 Post-LN 在 50 万步达到验证损失 1.69，Pre-LN 在 70 万步达到相近的值，由此得出 40% 的加速。按这句话，Pre-LN 反而更慢。对照图 5(a)：50 万步处 Pre-LN 约 1.70，Post-LN 约 1.75，Post-LN 要到 70 万步左右才降到 1.69。两个模型名写反了，应该是 Pre-LN 50 万步、Post-LN 70 万步，$700/500=1.4$，这才是 40% 的来历。【文中有据，第 8 页与图 5(a)】作者接着指出预热只有 1 万步，而省下的是 20 万步，所以加速不只是省掉了预热那段时间，这个论证是对的，前提是数字读对。

混在一起的变量

调度不同。IWSLT 上 Pre-LN 用的是第 8 个 epoch 乘 0.1 的阶梯衰减，Post-LN 用的是平方根倒数衰减。学习率一刀切下去之后验证损失通常会突然掉一截，而正文拿来比较的恰好是 Pre-LN 第 9 个 epoch，切完之后的第一个检查点。Pre-LN 收敛更快这个结论里，有多少来自 LN 的位置，有多少来自那一刀，从论文给的实验里分不开。缺的对照是 Post-LN 预热之后接同样的阶梯衰减。【推断】

学习率不同。BERT 上 Pre-LN 的峰值学习率是 Post-LN 的 3 倍。作者说试过让 Post-LN 也用 $3\times10^{-4}$，训练发散了。这正是论文想说明的事，Pre-LN 能吃更大的学习率，所以这里的不公平是有意为之的，可以接受。只是没交代 Post-LN 在 $3\times10^{-4}$ 时试过多长的预热。

少了一个最关键的对照。全文没有一处是 Pre-LN 加预热。论文的标题级结论是 Pre-LN 可以安全地去掉预热，要支持这句话，需要看到去掉预热之后 Pre-LN 没有损失什么，也就是 Pre-LN 加预热和不加预热的对比。没有这一组，读者只知道 Pre-LN 不预热能训，不知道预热对它是不是依然有帮助。上一节的玩具实验里，预热只让 Pre-LN 晚了约 100 步、终点不变；后来大模型的实践给了另一个间接回答，最后一节再说。【文中有据，第 7 页实验设定中没有这一配置】

种子。翻译和预训练的每条曲线都是单次运行。微调做得好，每个检查点 6 个种子、报 95% 置信区间（附录 A.3），这一点值得肯定。IWSLT 上 BLEU 的种子间波动通常在零点几，所以两者相当的结论站得住，谁略高一点则读不出来。

附录 A 的脚注里还有一个细节：Pre-LN 在 IWSLT 上拿到的 35.5 BLEU 用的是另一套配置，峰值 $7.5\times10^{-4}$，第 8000 步衰减，dropout 从 0.1 改成 0.3。这个数字和正文图 4 的比较不在同一个设定下。

从第一版到定稿的一条时间线

把前面各节的线索按可能的推进顺序串起来。这一节整体是推断，每一步标出依据。

Pre-LN 的代价和后来的走向

Pre-LN 好训，但它不是免费的。残差流随层数线性变粗，意味着第 $l$ 层往主干上加的那一份，在主干里所占的比例大约是 $1/\sqrt l$。层数很深时，后面几层每一层对主干的相对改动都很小，整个网络的行为更像一个更宽、更浅的网络。苏剑林在为什么 Pre Norm 的效果不如 Post Norm（科学空间，2022）里用这个角度解释了一个被反复观察到的现象：Post-LN 一旦训得起来，最终效果常常比 Pre-LN 好。Liu et al. 2020 的实验里也是这个结论，他们的 Admin 初始化就是想让 Post-LN 既保留效果又能稳定训练。

2022 年的 DeepNet（Wang et al.，DeepNet: Scaling Transformers to 1,000 Layers）走得更远：回到 Post-LN，在残差相加时给主干乘一个大于 1 的系数，同时把支路的初始化缩小，用这种 DeepNorm 训到了 1000 层。从这篇论文的角度看，这是在 Post-LN 里人为地让 LN 入口的主干占比变大，从而把每层的更新在输出上的影响压住，和 Pre-LN 最后那个 LN 在做的事情是同一个方向。

另一条线是大语言模型的实践。GPT-2、GPT-3、LLaMA 都用 Pre-LN（LLaMA 换成了 RMSNorm），但没有一个去掉预热：GPT-3 用前 3.75 亿个 token 线性预热，LLaMA 用 2000 步。原因不难理解：在几十万步的训练里，两千步的预热几乎没有成本，却能同时挡住 Adam 早期二阶矩不准、超大批量、早期损失尖峰这些风险。这篇论文让预热从必需变成了保险，大模型的训练者选择继续买这份保险。这也间接回答了上面那个缺失的对照：在大规模上，Pre-LN 加预热至少不比不加差。【推断】

归一化层是一个尺度旋钮

这篇论文最值得带走的不是 Pre-LN 本身，而是引理 3 背后的那个视角：一个归一化层会把穿过它的梯度和穿过它的扰动按入口长度的倒数一起缩放。所以看一个网络好不好训，可以先去量每个归一化层入口的长度。哪里长度在涨、哪里被反复拉回，梯度和更新在函数空间里的尺度就在哪里被改写。这个视角对 BatchNorm、RMSNorm、QK-Norm 一样适用。

第二件事是方法上的。论文的理论只证明了上界，附录里的推广漏了一个因子，图 3 的统计量和定理不完全对应。这些问题都可以用一个六十行的脚本在一分钟内看出来。读一篇讲初始化或训练稳定性的论文时，先把它的设定写成最小实现、打出它说的那个量，往往比反复读推导更快找到它真正成立的范围。

留下来值得做的问题有三个。一是给定理 1 配一个下界，把Post-LN 的梯度大从上界的比较变成真正的结论。二是在 Adam 的更新规则下重写整个论证，直接分析一步更新在输出上造成的变化，而不是梯度的大小。三是找一个摆法，同时拿到 Post-LN 的最终效果和 Pre-LN 的稳定性；DeepNorm 是一个答案，但它要按层数调系数，离换一个开关就能用还有距离。