pre-LN 的梯度稳定是用"有效深度"换来的:残差主干上范数膨胀、深层贡献被存量淹没(见本系列 pre-LN 篇)。有没有办法把 post-LN 也训深?DeepNorm(Wang et al., 2022,微软)给出了肯定答案,配方简单到只有两行改动,却让 1000 层的 encoder-decoder 模型稳定收敛。这篇文章拆这两行改动各自在干什么。
配方与改动的位置
DeepNorm 把残差块改成
$$ x_{l+1} = \mathrm{LN}\big(\alpha\, x_l + f_l(x_l)\big), $$
外加一条初始化规则:$f_l$ 内部的所有投影矩阵用 Xavier 初始化后再乘一个小系数 $\beta_l$(论文里对 encoder 与 decoder 的最后一层取 $\beta=0.1$,其余层 $\beta$ 由深度公式给出)。和 post-LN($\alpha=1$)比,只有 $\alpha$ 这一个新东西;和 pre-LN 比,归一化仍在主干上。
$\alpha$ 的取值不是超参搜索出来的,是公式:encoder 有 $N$ 个层时
$$ \alpha = (2N)^{1/2},\qquad \beta_l = (8N)^{-1/2}\ \text{(对大多数层)}, $$
encoder-decoder 里 encoder 取 $(3N)^{1/2}$、decoder 取 $(3M)^{1/2}$ 一类对应形式($M$ 为 decoder 层数)。系数 2 和 3 来自后述推导里"恒等路径与分支路径的数量匹配"。
为什么 α 放大恒等项能稳住训练
回到 post-LN 的病根:$\partial\hat u/\partial u$ 的投影面与主干梯度方向对齐,梯度被逐层投影收缩(见 post-LN 篇)。DeepNorm 的论证换一个角度,从前向扰动有界推出训练稳定。
把 $L$ 层网络看成对恒等映射的扰动:$x_L = x_0 + F(x_0)$,其中 $F$ 是所有分支贡献之和。训练稳定的一个充分条件是 $F$ 有界——深层网络发散的本质是 $F$ 随深度失控增长,参数微小更新被层层放大。对 post-LN 做归纳:设前 $l$ 层复合映射是 $\hat x_l = \mathcal E_l(x_0)$($\mathcal E$ 表示恒等项,$\mathcal F$ 表示分支扰动项,两者都经过 LN),则
$$ \hat x_l = \mathrm{LN}\big(\mathcal E_{l-1}(x_0) + \mathcal F_{l-1}(x_0)\big) \approx \mathrm{LN}\big(\mathcal E_{l-1}(x_0)\big) + \frac{1}{\|\mathcal E\|}\frac{\partial \mathrm{LN}}{\partial \cdot}\mathcal F_{l-1}, $$
当恒等项在 LN 输入中占主导($\|\mathcal E\|\gg\|\mathcal F\|$),LN 对它近似线性,$\hat x_l \approx \hat{\mathcal E}_{l-1} + J\,\mathcal F_{l-1}/\|\mathcal E\|$:LN 的雅可比 $J$ 有 $1/\|\cdot\|$ 因子,归一化的存在天然把每层分支的贡献压掉一个范数因子。问题是 post-LN 里 $\|\mathcal E\|$ 与 $\|\mathcal F\|$ 同量级,这个"恒等主导"近似不成立,归纳链条断掉,这就是 post-LN 只能扛几十层的理论图像。
DeepNorm 乘上 $\alpha=(2N)^{1/2}$ 后,LN 输入变成 $\alpha x_l + f_l$。$\alpha$ 随 $\sqrt{N}$ 增长,恒等项被显式放大,从第 1 层起就保证 $\|\alpha\,\mathcal E\| \gg \|\mathcal F\|$,上面那个"LN 对恒等项近似线性"的归纳假设重新成立;同时 $\beta$ 初始化把分支输出的初始量级压到 $O(1/\sqrt{N})$,让假设在训练初始时刻就成立,之后交给优化器微调。收敛的充分条件被论文归结为:恒等项系数 $\alpha$ 大于扰动项系数的某个随 $N$ 增长的阈值,$\alpha=(2N)^{1/2}$ 是满足条件的最简闭式。
它在 post-LN 与 pre-LN 之间的位置
把三个摆法放在同一坐标系里看。post-LN:归一化在主干,投影收缩梯度,浅层失学,依赖 warmup;pre-LN:归一化进分支,主干无损但范数膨胀、有效深度不足;DeepNorm:归一化在主干(保留 post-LN 的每层输出有界——这点对推理数值稳定是真实收益),用 $\alpha$ 把梯度通路买回来。代价是超参多两个带公式的常数,且 $\alpha$ 依赖总层数 $N$——改深度就要重新算,不像 pre-LN 那样深度无关。
实测层面,DeepNorm 的论文把 encoder-decoder 训到 1000 层(机器翻译任务),同等配置下 post-LN 在几十层就发散,pre-LN 能训但深度利用不足。后续大规模实践里,DeepNorm 没有取代 pre-LN 成为主流——原因更多是生态而非效果:pre-LN 的"深度无关"特性对做模型尺寸扫描更友好,且 μP 等参数化理论给了 pre-LN 更好的超参迁移工具。但 DeepNorm 的思想被大量吸收:残差分支加可学习/随深度的缩放系数(如 GPT-J 之后常见的 $\alpha x + f$ 形式)几乎成了新架构的标配变体。
小结
DeepNorm 的两行改动分工明确:$\alpha=(2N)^{1/2}$ 放大 LN 输入里的恒等项,让"LN 对恒等映射近似线性"的归纳假设在任意深度成立,把 post-LN 断掉的梯度通路接上;$\beta$ 初始化保证训练第 0 步假设就成立。理论要点可以压缩成一句:归一化层的 $1/\|\cdot\|$ 因子天然压制分支扰动,只要让恒等项在输入里占主导,深度本身就不再是稳定性的敌人。