原始 Transformer(Vaswani et al., 2017)是 post-LN:每个子层的输出先加回残差,再过 LN。后来的 GPT 系、LLaMA 几乎全换成 pre-LN:先 LN 再进子层。这个看似微小的搬动,是训练几十层模型能不发散的关键,但它不是没有代价的。这篇用上一篇的残差流语言把两边都分析掉。
post-LN 为什么深不了
post-LN 的第 $l$ 层是 $x_{l+1} = \mathrm{LN}(x_l + f_l(x_l))$。问题出在最后一层:梯度从 loss 回传到 $x_{L-1}$ 要穿过一次 LN。LN 的雅可比(本系列 RMSNorm vs LayerNorm 那篇推过)是
$$ \frac{\partial\,\mathrm{LN}(x)}{\partial x} = \frac{1}{\|x\|}\Big(I - \frac{xx^\top}{\|x\|^2}\Big)\cdot\sqrt d $$
(按单位化后的 $\hat x = x/\|x\|$ 写法),其核心是把梯度投影到垂直于 $x$ 的子空间,并整体乘以 $1/\|x\|$。于是回传到残差分支的梯度被这个投影矩阵缩放:如果子层的输出 $f_l$ 相对于 $x_l$ 很小(初始化时正是如此——$f_l$ 的参数从零附近初始化,$x_l$ 却是前面所有层累加的产物),$\|x_l\|$ 远大于 $\|f_l\|$,那么 $\mathrm{LN}$ 的缩放项近似 $1/\|x_l\|$,残差分支收到的梯度被压小一个量级,而子层分支 $f_l$ 收到的梯度是正常尺度。逐层累积这个失衡,靠近输入的层拿到的梯度指数级衰减——这就是"warmup 必配 post-LN、层数一深就不收敛"的定量来源。
一个佐证:Xiong et al.(2020,On Layer Normalization in the Transformer Architecture)实测了两种摆放下各层梯度范数随深度的分布,post-LN 底层梯度比顶层小好几个数量级;并给出了一个事后补救——把残差分支乘一个 $\alpha$ 常数(Liu et al. 的 DeepNorm 前身思路),让两条分支的梯度尺度匹配。
pre-LN 的高速公路与它的债
pre-LN 把 LN 挪进子层内部:$x_{l+1} = x_l + f_l(\mathrm{LN}(x_l))$。这样 LN 只影响子层的输入,残差分支从头到尾没有任何非线性、没有任何归一化,梯度 $\partial \mathcal L/\partial x_l$ 沿总线直通——几十层照样训得动,连 warmup 都可以缩短甚至取消(小学习率下)。
欠的债在输出端:$x_L = x_0 + \sum_l f_l$ 里,末层的 $\|x_L\|$ 随 $\sqrt L$ 膨胀,而最终预测是 $\mathrm{LN}(x_L)$ 乘 unembedding 矩阵——$x_L$ 里真正起作用的"信号"只是它在这个被归一化后的方向。更麻烦的是各子层的写入之间没有任何归一化来约束相对尺度:每一层的 $f_l$ 都直接面对一个越来越大、内容越来越杂的总线,深层子层的输入分布漂移剧烈,末几层实际上承担了"把失控的总线压回可预测分布"的活。实践中它表现为:pre-LN 的下游迁移和超参敏感性略逊于 post-LN,以及必须额外小心残差输出的初始化尺度(常见做法是给 $W_O$ 和 FFN 输出层乘 $1/\sqrt{2L}$ 之类的小因子)。
两边合起来就得到现代架构的标准答案:pre-LN 为主,叠一个残差缩放补丁——DeepNorm 把 post-LN 的两分支梯度匹配思想反过来用在 pre-LN 上:残差乘 $\alpha > 1$、子层乘 $\beta < 1$,把 $\sqrt L$ 增长线性化,160 层也训得稳。代价是它重新引入了一点 warmup 需求。
小结
推导上的分水岭在 LN 雅可比那项 $1/\|x\|\cdot(I - \hat x\hat x^\top)$:post-LN 让梯度在每层出口都过这道闸门,深层被压消失;pre-LN 把闸门挪进子层内部,残差分支全程无阻,换来的是 $\sqrt L$ 范数膨胀与输出分布漂移。现代栈的选择是 pre-LN 加上残差缩放(DeepNorm 或输出层小初始化),把两边的债都还到最小。下两篇顺着这个工程化思路走:怎么在不牺牲质量的前提下,把层内的串行依赖也砍掉一部分。