为什么梯度裁剪的默认模长是1？
原文链接 | 2025-01-02 | 苏剑林
我们知道，梯度裁剪（Gradient Clipping）是让模型训练更加平稳的常用技巧。常用的梯度裁剪是根据所有参数的梯度总模长来对梯度进行裁剪，其运算可以表示为
\[\begin{equation}\text{clip}(\boldsymbol{g},\tau)=\left\{\begin{aligned}&\boldsymbol{g}, &\Vert\boldsymbol{g}\Vert\leq \tau \\ &\frac{\tau}{\Vert\boldsymbol{g}\Vert}\boldsymbol{g},&\Vert\boldsymbol{g}\Vert > \tau \end{aligned}\right.\end{equation}\]
这样一来，\(\text{clip}(\boldsymbol{g},\tau)\)保持跟\(\boldsymbol{g}\)相同的方向，但模长不超过\(\tau\)。注意这里的\(\Vert\boldsymbol{g}\Vert\)是整个模型所有的参数梯度放在一起视为单个向量所算的模长，也就是所谓的Global Gradient Norm。
不知道大家有没有留意到一个细节：不管是数百万参数还是数百亿参数的模型，\(\tau\)的取值在很多时候都是1。这意味着什么呢？是单纯地复用默认值，还是背后隐含着什么深刻的原理呢？
是什么#
可能有读者觉得，默认值又不一定是最优值，有什么值得纠结的？确实，\(\tau=1\)未必是最优的选择，但它是很多模型的默认选择，并且在这个默认选择下表现尚可，这反过来表明\(\tau=1\)具有普遍的合理性。
这里的“合理性”又指什么呢？让我们回到\(\text{clip}\)运算上。如果\(\Vert\boldsymbol{g}\Vert\)总是小于\(\tau\)，那么\(\text{clip}\)就退化为恒等变换了；如果\(\Vert\boldsymbol{g}\Vert\)总是大于\(\tau\)，那么\(\text{clip}\)就退化成L2归一化。换句话说，\(\text{clip}\)之所以为\(\text{clip}\)，就是因为\(\tau\)产生了适当的区分度，使得大部分的\(\Vert\boldsymbol{g}\Vert\)都是小于\(\tau\)的，只有小部分才是大于\(\tau\)的，这就是\(\tau\)的合理性的含义。
这个当然可以举出反例，而且还不少，这里主要想强调这个现象的普遍性以及这个默认设置的普适性，所以较真的读者大可不必过于执着于个别细节。
因此，我们认为，\(\tau=1\)的普遍合理性的含义，就是不论模型参数量多少、怎么初始化、取何种损失函数，它的梯度总模长都能恰好大致以\(1\)为“异常值”的分界点，这无疑是一个非常不可思议的性质——笔者第一次意识到这个结论时的感受便是如此。
为什么#
为什么会如此“巧合”呢？笔者的答案可能会让人有些意外：因为只有这样，模型才有稳定训练的可能。
让我们考虑损失函数\(\mathcal{L}(\boldsymbol{\theta})\)，优化器更新规则为\(\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta\, \boldsymbol{u}_t\)，那么损失函数的变化近似为
\[\begin{equation}\Delta \mathcal{L} = \mathcal{L}(\boldsymbol{\theta}_{t+1}) - \mathcal{L}(\boldsymbol{\theta}_t) \approx (\boldsymbol{\theta}_{t+1} - \boldsymbol{\theta}_t)\cdot\nabla_{\boldsymbol{\theta}_t}\mathcal{L}(\boldsymbol{\theta}) = -\eta\, \boldsymbol{u}_t\cdot \boldsymbol{g}_t\end{equation}\]
先考虑最简单的SGD，那么\(\boldsymbol{u}_t = \boldsymbol{g}_t\)以及\(\Delta \mathcal{L}=-\eta\Vert\boldsymbol{g}_t\Vert^2\)，即损失函数的变化量正比于梯度模长的平方。我们知道，不管是CV还是NLP，纯粹的SGD（不带动量）都是非常低效的优化器，训练到中后期，平均来说多数任务每步的损失下降量是远不如学习率大小的，也就是\(|\Delta \mathcal{L}| < \eta\)，由此推得\(\Vert\boldsymbol{g}_t\Vert < 1\)。这就表明了\(\Vert\boldsymbol{g}_t\Vert < 1\)是一个能正常收敛的模型的长期表现。
当然，训练初期模型有可能会出现\(\Vert\boldsymbol{g}_t\Vert > 1\)，这是正常的，但很少情况会出现\(\Vert\boldsymbol{g}_t\Vert \gg 1\)，或者说一个优秀的初始化应该避免出现\(\Vert\boldsymbol{g}_t\Vert \gg 1\)，像DeepNorm等的理论依据便是如此。原因是相似的，如果梯度模长太大，那么前期的学习就会过于“激进”，导致提前收敛到不好的局部解。另一个方案是缩小\(\eta\)，这同样能够缩小\(|\Delta \mathcal{L}|\)，这也就是为什么在训练初期我们通常使用Warmup。
顺便说，关于Warmup的理解大家可以参考论文《Optimal Linear Decay Learning Rate Schedules and Further Refinements》，这是笔者认为对Warmup的最合理的分析。
怎么办#
简单来说，就是由于损失函数的变化量正比于梯度模长的平方，所以训练的平稳性决定了梯度模长不能太大，并且长期表现为小于1。而初期如果出现明显大于1的梯度模长，那么通常的策略是Warmup。或者也可以考虑一个更通用的策略：设置另一个阈值\(\mathcal{T}\)，根据\(\boldsymbol{u}_t\cdot \boldsymbol{g}_t\)的值对\(\eta\)进行裁剪
\[\begin{equation}\eta_t = \left\{\begin{aligned}&\eta,& \boldsymbol{u}_t\cdot \boldsymbol{g}_t\leq \mathcal{T} \\ &\frac{\mathcal{T}}{\boldsymbol{u}_t\cdot \boldsymbol{g}_t}\eta,& \boldsymbol{u}_t\cdot \boldsymbol{g}_t > \mathcal{T} \end{aligned}\right.\end{equation}\]
这样就免除了额外的Warmup设置，更加具有自适应性。
对于Adam等优化器，我们可以跟《当Batch Size增大时，学习率该如何随之变化？》一样，通过\(\boldsymbol{u}_t=\text{sign}(\boldsymbol{g}_t)\)来进行近似分析，此时
\[\begin{equation}\Delta \mathcal{L} = -\eta\, \text{sign}(\boldsymbol{g}_t)\cdot \boldsymbol{g}_t = -\eta\, \Vert\boldsymbol{g}_t\Vert_1\end{equation}\]
这里的\(\Vert\Vert_1\)是L1范数，即分量的绝对值之和。由于梯度分量基本都小于1，因此\(\Vert\boldsymbol{g}_t\Vert_1 \gg \Vert\boldsymbol{g}_t\Vert\)，因此同样出于平稳训练的需求，Adam的学习率通常要明显小于SGD的学习率。此外，上式还可以改写成
\[\begin{equation}\Delta \mathcal{L} = -\eta\, \text{sign}(\boldsymbol{g}_t)\cdot \boldsymbol{g}_t = -\eta\, \sqrt{N}\Vert\boldsymbol{g}_t\Vert \cos(\text{sign}(\boldsymbol{g}_t), \boldsymbol{g}_t) \end{equation}\]
这里假设了\(\boldsymbol{g}_t\)没有零分量，因此\(\Vert\text{sign}(\boldsymbol{g}_t)\Vert=\sqrt{N}\)，\(N\)是模型总参数量。实践发现\(\Vert\boldsymbol{g}_t\Vert\)和\(\cos(\text{sign}(\boldsymbol{g}_t), \boldsymbol{g}_t)\)在不同的模型尺度下都大致为常数，因此如果要维持\(\Delta \mathcal{L}\)不变，应该有\(\eta\)反比于\(\sqrt{N}\)，也就是说模型参数量增加到4倍，那么学习率可以考虑减半。
全剧终#
本文对“梯度裁剪的默认模长为1”这一现象给出了自己的一些看法和思考。