技术博客

归一化的不变性与尺度解耦

第 101 / 147 篇 · 2025-09-26 · ai · 归一化weight decay优化理论

一个常被忽视的观察:现代 Transformer 里几乎所有的矩阵参数(attention 投影、FFN 权重)都"住"在归一化的下游,输出要穿过 RMSNorm 才进残差流。这带来一个参数尺度无关的性质——把 $W$ 整体乘 2,输出几乎不变(归一化会除掉这个因子)。参数化了网络的"方向"而不是"大小"。这篇文章推导这个不变性的几何后果,它和 weight decay 的耦合比直觉深得多。

尺度不变性的定义与直接推论

设参数 $w$(整体看成一个向量),损失 $\mathcal L$ 对它满足

$$ \mathcal L(c\,w) = \mathcal L(w)\quad \forall c>0. $$

对 $c$ 在 $c=1$ 处求导:$\nabla\mathcal L(w)^\top w = 0$。梯度恒与参数正交。这不是近似,是不变性的直接推论。几何图像:损失在 $w$ 的径向(沿 $w$ 本身方向)完全平坦,只在球面上有地形;梯度天然躺在与 $w$ 垂直的切平面上,更新 $w\leftarrow w-\eta\nabla\mathcal L$ 只走切向、不走径向。

哪些参数有这条性质:RMSNorm/LayerNorm 之后的所有线性层($\mathrm{RMS}(Wx)$ 对 $W$ 的逐行缩放不变——行缩放被分母的范数吸收),以及 $g$ 这类增益参数本身(前面 RMSNorm 篇已分析)。 embedding 和 lm_head 不在其列:它们不在归一化下游,径向是有梯度的。

有效学习率:WD 如何变成学习率调度

单独看 SGD:$w_{t+1} = w_t - \eta\nabla\mathcal L$。因为 $\nabla\mathcal L\perp w_t$,参数的范数不变,更新改变的是方向。把 $w$ 写成 $\|w\|\cdot\hat w$,单步的方向改变量(以弧度计)是 $\Delta\theta \approx \eta\|\nabla\mathcal L\|/\|w\|$——范数越大,同样的参数步长对应的角度越小,网络的函数变化越小。$1/\|w\|^2$ 就是"有效学习率"对参数范数的依赖。

现在加 weight decay:$w_{t+1} = (1-\eta\lambda)w_t - \eta\nabla\mathcal L$。分解到径向与切向:切向照旧,径向每一步乘 $(1-\eta\lambda)$,范数稳定值由 WD 衰减与梯度噪声注入的平衡决定(随机梯度噪声有径向分量,WD 持续外拉,稳态 $\|w\|^2 \approx \mathbb E\|\eta g\|^2/(2\eta\lambda)$ 量级,Hoffer et al. 2018 "Norm matters" 的推导)。把稳态范数代回有效学习率:

$$ \eta_{\mathrm{eff}} \approx \frac{\eta}{\|w\|^2} \;\propto\; \frac{\eta\,\lambda}{\text{梯度噪声水平}}. $$

三个参数的定性关系全出来了:调大 WD ⇒ 范数收缩 ⇒ 有效学习率上升;调大 base 学习率 $\eta$ 在 SGD 下反而被范数增长部分抵消(噪声大了范数涨,$\eta_{\mathrm{eff}}$ 只按 $\sqrt{\eta}$ 量级走)。这和无归一化网络里"WD=正则强度"的直觉完全相反:对尺度不变参数,WD 不是正则,是学习率的遥控器。这也解释了带归一化的 Transformer 对 weight decay 系数异常敏感、而且最优 WD 与学习率强耦合的经验现象——你以为在调两个独立旋钮,其实是同一个旋钮的两端。

尺度不变参数空间:损失只在球面上有地形,WD 决定球的大小 w(径向无梯度) ∇L ⊥ w 三个等损失同心球:函数值只依赖 w/‖w‖ WD 的作用链 λ↑ ⇒ ‖w‖ 收缩(向心拉力) ‖w‖↓ ⇒ η_eff = η/‖w‖² 上升 结论:WD 是有效学习率的增函数 (无归一化时恰恰相反:WD=正则)

Adam 下的修正:逐参数缩放的双重作用

上面是 SGD 的账。Adam 给每个参数单独做 $m/\sqrt{v}$ 归一,有两个后果。第一,Adam 的输出步长在参数空间各向同性(每参数每步 $\sim\eta$),$\eta/\|w\|$ 的角度含义更直接。第二,Adam 的 $v$ 估计包含径向分量的统计(虽然真梯度径向分量为零,但随机梯度噪声有),噪声水平影响 $\sqrt{v}$,进而和 WD 的耦合关系里多一个 $v$ 的稳态方程。定性的"WD 调有效学习率"结论在 Adam 下依然成立,文献(如 Loshchilov & Hutter 的 AdamW 论文、Zhang et al. 的 "Which Algorithmic Choices Matter" 消融)里有实测支持:对带归一化的网络,WD 与学习率的比值比各自绝对值更能预测最终效果。

AdamW 的解耦设计在这里显出先见之明。原版 Adam 把 L2 惩罚项也放进 $v$ 的分母归一里——惩罚大的参数自适应步长被压缩,正则失效;AdamW 把 weight decay 挪出梯度、直接乘在参数上,保证"WD 管尺度、梯度管方向"的分工干净。在尺度不变参数占主流的 Transformer 里,这个分工恰好就是"WD 管有效学习率"的配方。

工程含义清单

把上面的理论落成可操作的结论。其一,调参时把 lr 和 wd 当一个二维面上的一条线来扫(固定比值扫),比分别扫高效——大量超参搜索的经验(如 μP 相关工作)支持 lr/wd 比值近似不变的最优线。其二,不同层共享同一个 lr 时,各层的有效学习率天然均衡(范数大的层角度步小),这是归一化带来的隐性好处:不用像 BatchNorm 时代那样担心不同层尺度漂移失衡。其三,监控训练时看参数的 $\|w\|$ 曲线比看梯度范数更有信息量——范数走势直接反映 WD 与学习率的角力结果,范数塌缩往往先于 loss 异常。

小结

归一化把网络参数化从"向量"变成"球面上的点":∇ℓ·w=0,梯度纯切向,WD 因此从正则项变成有效学习率的调节器,η_eff ∝ ηλ/噪声。AdamW 的解耦设计、lr 与 wd 比值扫参、监控范数曲线——这些工程惯例背后是同一条几何:在尺度不变的参数世界里,大小和方向是两个独立的旋钮,weight decay 管的是前者。