Weight Decay

本文是 LLM Theory 下 SMD 专题的第一篇，是关于 Spherical Motion Dynamics（SMD）的学习笔记，主要参考 Wan et al. 的 Spherical Motion Dynamics: Learning Dynamics of Neural Network with Normalization, Weight Decay, and SGD 以及 NeurIPS 2021 版本 Spherical Motion Dynamics: Learning Dynamics of Normalized Neural Network using SGD and Weight Decay。 朴素的描述模型更新 对于本科一二年级时候的笔者而言，如果想要描述模型的更新量，那么只会考虑这个非常直接的东西：$\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$。 这确实是很直白的，相当于计算了模型权重的欧氏距离。是什么让我们必须放弃这种非常直观的欧氏距离呢？源于常用的归一化。 加上归一化之后，对于模型的输出，$y = \operatorname{BN}(x, k\boldsymbol{W}) = \operatorname{BN}(x, \boldsymbol{W})$，看不出模型尺度对 $y$ 的影响。但是却会从 $\lVert k\boldsymbol{W}_{t+1} - k\boldsymbol{W}_{t} \rVert$ 这一度量手段上产生 $k$ 倍的差距。 这就发现一个很显然的问题了：我们是希望通过观测类似于 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 这样的东西来控制模型的训练，而不是直接看最后的 $y$。因为只通过看 $y$ 的变化，并基于这种变化的规律，去指导如何进行训练前的设置，这件事并不容易。我们是希望找到一个更可以写出明确表达式，更容易观测，意义更明确的指标去观测，而这个指标恰好还要有一些规律和 $y$ 的规律“趋同”，这样我们就得到了一个 $y$ 的近似物，而且这个近似物比 $y$ 更好分析。所以我们会很直接的想到 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$。但是现在，当我们给模型加上BN（为什么加BN不赘述了）以后发现的问题是 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 会随着模型尺度产生成倍的差距，而 $y$ 却不会被尺度影响。这说明 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 不是一个好的指标，因为它错误地预测了 $y$ 随着尺度变化的行为规律。 当然，原则上，我们其实可以发现，上面的地方有一个逻辑漏洞：为什么权重变化一定能预测 $y$ 的变化，这是错误的呀，权重变了 $y$ 不一定变。所以 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 这个东西失效是情理之中的，不能先入为主的认为 $\lVert \boldsymbol{W}_{t+1} - \boldsymbol{W}_{t} \rVert$ 的规律能和 $y$ 趋同，那么出现上面的问题十分合理。我们一个理想中的目标是通过权重变化去预测 $y$ 的变化，这样我们就可以通过设计一套权重变化的方案去精确得到想要的 $y$。理想是很远大的，现实会比较难做，但是可以在最朴素的方案上向前不断推进预测的细化程度。SMD就是这样一种东西。 ...

本文是 LLM Theory 下 SMD 专题的第二篇。上一篇主要把 SGD、SGD + weight decay、Adam、AdamW 的径向/切向更新拆开了；这一篇继续沿着这个拆法，讨论几个笔者读 SMD 时觉得更关键的 insight。 Insight 1：为什么只看 Effective Learning Rate 不够 在讨论 normalization 和 weight decay 的联合作用时，有一句非常关键的话： 对于 scale-invariant weight，任务梯度 $\partial L / \partial \boldsymbol{w}$ 总是与权重 $\boldsymbol{w}$ 垂直。因此，梯度分量总是倾向于增大权重范数，而 weight decay 提供的分量总是倾向于减小权重范数。 这句话其实就是 SMD 的基本物理图像： \[ \text{任务梯度} \Rightarrow \text{切向更新，改变方向，同时二阶增大范数} \] \[ \text{Weight Decay} \Rightarrow \text{径向更新，缩小范数} \] 也就是说，在带 normalization 的网络中，SGD 梯度并不是在普通欧氏空间里随便走，而是在球面的切线方向上推动权重转动；weight decay 则像一个向心力，把权重往原点拉。两者共同作用，最终可能让 weight norm 进入一个稳定状态。 但是这里有一个很容易被忽略的问题： Weight decay 可以对冲掉梯度更新导致的 weight norm 增长，但这并不意味着 gradient norm 本身也被稳定住了。 ...