从SGD到AdamW:优化算法的演化

想象一下,训练一个神经网络,就像把你蒙上眼睛,丢进一片连绵不绝的山脉里,你的任务是找到最低的山谷。

这个"最低山谷",就是我们模型的最佳状态------损失函数的最小值。

问题来了:你看不见全局地图,只能靠脚下的坡度来判断方向。下一步该往哪走?步子迈多大? 这就是优化器 (Optimizer) 要解决的核心问题。它就像你的智能登山杖,指导你如何高效地"下山"。

一个好的优化器,能带你快速、稳定地到达目的地;而一个差的,可能会让你在某个小山坡上来回打转,甚至迷失方向。

在这篇文章里,我们就来一场优化器的进化之旅。从最简单直接的"凭感觉走"------随机梯度下降(SGD)出发,一步步看它如何变得更聪明、更强大,最终进化成当今大模型训练普遍采用的优化算法------Adam 和 AdamW。

SGD:随机梯度下降

最初的梯度下降法需要一次性把所有数据都扔给模型,通过遍历所有训练数据,计算模型在整个数据集上的总损失(计算出所有样本的损失再计算平均值)然后再计算出梯度,这意味着每一次更新参数都需要遍历整个数据集,虽然说这样计算出来的梯度方向非常准确,但当今数据集普遍都非常庞大,这种方法对算力和内存都有非常高的需求。

SGD的改进方向是每次只遍历一部分的数据,每次只根据这一部分数据计算出来的梯度来更新参数:

\[\theta_{t+1} = \theta_t - \gamma g_t \]

其中\(\theta\)是参数,\(\gamma\)是学习率,\(g\)是梯度。

因为每次只遍历一部分数据(这一部分数据的大小我们称为batch_size),所以梯度的计算会非常快,并且对内存也会更加友好。这个遍历的过程会一直重复,直到遍历完所有数据(这里称作一个Epoch)。

但这种方法也会有一定的缺点,每次只遍历一部分数据意味着充满随机性和噪音,损失函数会有一定的上下波动,收敛过程较慢。

话又说回来,噪音和随机性有时候反而会跳出局部最优解。

Momentum

上面说到在使用SGD算法时,我们有事会遇到震荡的问题,导致模型收敛较慢,这种现象是由于每次迭代后梯度变化较大导致的。

想像一下我们的损失函数像是一个被拉得很长的、像峡谷或椭圆形山谷一样的地形,在这个峡谷里两侧的峭壁非常陡峭,梯度很大。而沿着峡谷底部走向最低点的坡度却非常平缓,梯度很小。我们的目标,是让小球(代表模型参数)从峡谷的一端,平稳地滚动到谷底的最低点。

现在,我们把 SGD 控制的小球放在峡谷的一侧峭壁上。会发生什么?

  1. 第一步: 小球只看脚下,它发现通往对面峭壁的方向是最陡的下坡路(梯度最大)。于是,砰 的一下,它朝着对面冲了过去。
  2. 用力过猛: 由于峭壁方向的梯度很大,这一步的步长可能也很大,导致小球不仅冲到了对面,甚至可能冲到了比起始点更高的位置。
  3. 第二步: 在新的位置,小球再次环顾四周,发现"回头"的方向又是最陡的。于是,砰 的一下,它又朝着原来的方向冲了回去。

这个过程会不断重复。小球就在峡谷的两侧峭壁之间来回"之"字形地反弹、震荡。虽然它整体上确实在向着谷底移动,但这种移动是极其低效和缓慢的。大部分能量都浪费在了这种无意义的左右摇摆上。这就是 SGD 的核心困境:

  • 在梯度变化剧烈的方向上(峭壁方向),它会剧烈震荡,难以稳定。
  • 在梯度变化平缓的方向上(谷底方向),它又因为梯度太小而前进缓慢。
    我们陷入了一个两难境地:如果调大学习率想让它在谷底走快点,它在峭壁上的震荡就会失控;如果调小学习率来抑制震荡,它在谷底的前进速度又会变得遥遥无期。

如何解决?引入"惯性"的力量 ------ 动量 (Momentum)

我们该如何让这个小球变得更"聪明"一点呢?问一个简单的问题:一个真实世界里的重铁球 会这样来回反弹吗?不会。一个有质量的铁球会带有惯性 ,或者说动量

  • 在峭壁方向(垂直于峡谷): 铁球冲向对面峭壁时,峭壁会给它一个反作用力。这样一来一回,它在这个方向上的速度会因为反复的碰撞和转向而被抵消掉。
  • 在谷底方向(沿着峡谷): 在这个方向上,重力始终在稳定地拉着它前进。这个方向上的力是持续不断的,所以铁球的速度会不断累积,越滚越快。
    这就是动量法 (Momentum) 的核心思想
    我们给小球增加一个"速度"变量 v,它会累积过去的梯度信息:
  1. 抵消震荡: 在来回震荡的方向上,梯度方向一正一负,反复变化。当把这些梯度加权平均后,它们会相互抵消,使得这个方向上的更新幅度变小。
  2. 加速前进: 在方向一致的梯度上(如峡谷底部),梯度方向始终不变。当把它们加权平均后,会不断累积,使得这个方向上的更新速度越来越快。
    看看在数学上如何实现动量法:

\[\theta_{t+1} = \theta_t - \gamma v_t \]

其中,\(v_t = \mu v_{t-1} + g_t\) 代表当前时刻的梯度;\(\mu\) 参数用于指数加权移动平均,该方法可以减小更早时刻梯度对当前梯度的影响,通常 \(\mu\) 取值为 0.9。

如果说上一时刻的梯度比较陡,也就是梯度是一个较大的负数,那么当前时刻的梯度会被减去一个较大的值,导致参数不会发生太大的变化,就好像小球在冲向谷底时会慢慢刹车,从而减小震荡的程度。

RMSProp

前面提到的优化器中的学习率均是一个固定的参数,但在复杂的、高维的损失函数空间中,不同参数对最终损失的"敏感度"和"重要性"是截然不同的,一个"一刀切"的更新步长会对优化过程造成极大的阻碍。

依旧把损失函数空间比做一个山谷,这个山谷的宽度方向非常陡峭,但长度方向非常平缓,理想的策略是在宽度方向用小步子,在长度方向上用大步子,对应来说就是在损失函数空间在梯度比较大的参数上的学习率比较小,在梯度比较小的参数上的学习率比较大。

RMSProp最核心的优点就是自适应学习率:

\[\theta_{t+1} = \theta_t - \frac{\gamma}{\sqrt{s_t + \epsilon}} g_t \]

其中,\(s_t = \alpha s_{t-1} + (1 - \alpha)g_t^2\) 代表着历史中所有梯度的平方和,\(\alpha\) 代表指数加权移动平均法的参数,它控制了历史梯度信息被遗忘的速度,\(\gamma\)是全局学习率,\(\epsilon\)是一个非常小的数,防止分母为0。

RMSprop 通过\(\sqrt{s_t + \epsilon}\)对每个参数的更新进行了归一化。这使得每个参数都有一个量身定制的学习率,非常适合处理特征稀疏或者不同参数梯度尺度差异巨大的情况(例如在自然语言处理中,某些词的词向量很少被更新)。

Adam

动量法是在梯度上做文章,RMSProp是在学习率上做文章,那能不能把这两种方法结合起来呢?Adam就是二者结合起来的一种优化器:

\[\theta_{t+1} = \theta_t - \frac{\gamma}{\sqrt{s_t + \epsilon}} v_t \]

其中,\(s_t = \beta_2 s_{t-1} + (1 - \beta_2)g_t^2\);\(v_t = \beta_1 v_{t-1} + (1 - \beta_1)g_t\)。

AdamW

虽然理论上Adam算法的性能更优,但人民发现Adam有时的表现并不如动量法,尤其是在模型泛化能力上。

我们知道,L2 正则化的目标是防止模型过拟合,通过在损失函数中增加一个惩罚项,惩罚过大的权重。

\[L_{\text{total}} = L_{\text{original}} + \frac{\lambda}{2} \|w\|^2 \]

其中 \(\lambda\) 是正则化强度,\(w\) 是模型的权重。

L2 正则化如何变成"权重衰减":

当计算总损失的梯度时,这个惩罚项会引入一个新的梯度项:

\[\nabla L_{\text{total}} = \nabla L_{\text{original}} + \lambda w \]

在更新权重时(以 SGD 为例):

\[w_{t+1} = w_t - \eta (\nabla L_{\text{original}} + \lambda w_t) \]

\[w_{t+1} = (1 - \eta \lambda) w_t - \eta \nabla L_{\text{original}} \]

看到 \((1 - \eta \lambda)\) 这一项了吗?它意味着在每次应用梯度更新之前,权重 \(w_t\) 自身都会被乘以一个小于 1 的系数进行"衰减"。这就是权重衰减的由来。对于 SGD 来说,L2 正则化和权重衰减是等价的。

再看看L2正则化在Adam中是怎么样的:

\[\theta_{t+1} \leftarrow \theta_t - \gamma \frac{\beta_1 v_{t-1} + (1 - \beta_1)(g_t + \lambda \theta_t)}{\sqrt{s_t + \epsilon}} \]

我们可以看到\(\lambda\theta_t\)这个旨在让权重衰减的项,本身也被Adam的自适应分母\(\sqrt{s_t + \epsilon}\)给缩放了,这样的后果是对于那些历史梯度比较大的权重,它们的有效权重衰减强度会变得非常小。通常,梯度较大的权重往往是模型中比较重要的、需要被好好正则化的权重。但 Adam 的实现方式反而减小了对这些权重的正则化力度。它将正则化强度和梯度的历史大小错误地耦合在了一起。

AdamW的作者在论文中指出了上述问题,并提出了一种简单而优雅的解决方案:将权重衰减与梯度更新解耦:

\[Adam_{update} =\frac{\gamma}{\sqrt{s_t + \epsilon}} v_t \]

\[\theta_{t+1} = \theta_t - Adam_{update} - \gamma \lambda\theta_t \]

在实际实现过程中还需要进行偏置矫正: