LayerNorm 前向一句话:减均值、除标准差、乘增益。它的反向却藏着一个投影。这篇把雅可比推出来,看它如何同时压制梯度消失和梯度爆炸。
前向与微分
$$ y = \gamma \odot \hat x, \qquad \hat x = \frac{x - \mu}{\sigma},\;\; \mu = \tfrac{1}{d}\mathbf{1}^\top x,\;\; \sigma = \tfrac{1}{\sqrt{d}}\|x - \mu\mathbf{1}\| . $$
对 $x$ 求微分。均值项 $d\mu = \frac{1}{d}\mathbf{1}^\top dx$;方差项利用 $(x-\mu\mathbf{1})^\top \mathbf{1} = 0$ 化简得 $d\sigma = \frac{1}{d}\,\hat x^\top dx$。代入 $d\hat x = \frac{dx - d\mu \cdot \mathbf{1}}{\sigma} - \frac{\hat x\, d\sigma}{\sigma}$:
$$ d\hat x = \frac{1}{\sigma}\Big(I - \underbrace{\tfrac{1}{d}\mathbf{1}\mathbf{1}^\top}_{\text{剥均值}} - \underbrace{\tfrac{1}{d}\hat x \hat x^\top}_{\text{剥径向}}\Big) dx . $$
记括号内为 $I - P$,其中 $P = \frac{1}{d}(\mathbf{1}\mathbf{1}^\top + \hat x \hat x^\top)$。验证它是投影:$\mathbf{1}^\top \hat x = 0$ 且 $\|\hat x\|^2 = d$,于是 $P^2 = P$——$P$ 是投到 $\{\mathbf{1}, \hat x\}$ 张成平面的秩 2 正交投影。完整雅可比:
$$ \frac{\partial y}{\partial x} = \frac{1}{\sigma}\,\mathrm{diag}(\gamma)\,(I - P). $$
投影的两个后果
径向梯度被剥离。 $I - P$ 杀掉了梯度在 $\hat x$ 方向的分量。这并非巧合而是必然:LN 的输出对输入的径向缩放不变($x \to cx$ 时 $\hat x$ 不变),由链式法则的互易性,损失对 $x$ 的梯度必然落在"缩放不动的方向"上——数学表述就是梯度与产生缩放的方向正交。实践后果重大:参数范数漂移引起的激活整体缩放,不会沿着网络传成梯度幅度的漂移。未归一化的深网络里,各层范数误差逐级累积正是梯度爆炸/消失的主要来源之一;LN 把这条通路从结构上切断了。
1/σ 是负反馈。 雅可比整体乘 $1/\sigma$:前向激活范数萎缩($\sigma$ 变小)时,同量上游梯度被放大传回;激活膨胀则梯度压缩。一个局部稳压器。注意第 6010 篇的 RMSNorm 去掉了剥均值那一项、只保留 $1/\sigma$ 和这个投影的径向半边——现代架构里均值分量本来就接近零,砍掉它换来一次少算的均值。
投影还顺带解释了为什么 LN 后面的线性层初始化可以更随意:进入 $W$ 的梯度已被整形到幅度可控的分布,$W$ 的谱范数扰动不再逐级放大。
放进残差结构:pre-LN 的高速公路
单看 LN 还不够,要放进残差块。pre-LN 下第 $l$ 层的梯度:
$$ \frac{\partial \mathcal{L}}{\partial x_l} = \frac{\partial \mathcal{L}}{\partial x_L}\Big(I + \sum_{k > l}\; \prod_{j}\, J_j\Big), $$
第一项是恒等通路:无论多深,梯度总有一条不穿过任何非线性与归一化的直达路径,范数不随深度衰减——pre-LN 能稳定训练几百层的第一因。第二项是各层雅可比的修正,每层贡献被自己的投影结构限制(投影矩阵谱范数为 1,乘 $\|\gamma\|/\sigma$),增长是加性的而非乘性的。
对照 post-LN:残差出口还要再穿一层归一化,恒等通路被一串 $\frac{1}{\sigma}\mathrm{diag}(\gamma)(I-P)$ 调制——每个 $\sigma$ 都是整条残差流的函数,深层梯度经过层层耦合缩放,没有深度不变性保证。2017 年原版 Transformer 用 post-LN 且必须配 warmup 才不发散,根子在这(pre-LN / post-LN 的全面对比见第 4230 篇)。
一个容易忽略的代价
投影意味着 LN 丢弃梯度在两个方向上的分量。如果某条学习信号恰好需要沿均值方向或径向调整激活,LN 会系统性削弱它。多数任务里这两个方向承载的是漂移而非信息,代价可忽略;但做机理研究(比如直接干预激活范数)时要记得这一刀的存在。
小结
LN 的稳定作用可以精确说成三件事:秩 2 投影剥掉径向与均值分量,$1/\sigma$ 提供局部负反馈,pre-LN 残差结构再给一条深度不变的恒等梯度通路。归一化层从来不"只是让数值好看",它是一个数学结构明确的梯度整形器。