优化器 (1):SGD和Adam
引言
在这一系列的博客中,我将从梯度下降开始介绍,逐步深入现代优化器的设计艺术。
后续介绍的内容主要可以划分为三大类:
- 基础梯度下降方法:BGD、SGD
- 学习率自适应方法:AdaGrad、RMSProp、Adam、AdanW
- 新晋方法:Muon、AdaMuon
在本文中主要包括了前两类,即梯度下降方法和自适应学习率方法。最近的新晋方法Muon则放在在下一篇文章中再详细介绍。
一、基础梯度下降方法
深度学习算法的基石是梯度下降 (Gradient Descent, GD) 算法。
具体来说,给定一个具有可学习参数
在实际问题中,损失函数
1.1. 什么是梯度下降法
一个
梯度下降法是一种迭代求极值的方法。在每一步迭代中,每一个参数都沿着梯度下降的方法前进一小步,直到走到某个极小值:
其中,
1.2. BGD和SGD
BGD (Batch Gradient Descent) 是最简单的梯度下降算法,它完全按照公式
为了解决BGD的效率问题,SGD (Stochastic Gradient Descent) 每次仅使用少量样本(通常称为一个batch)来估计梯度,可以有效兼顾效率和稳定性。
1.3. 梯度下降法的有效性 [2,3]
下面我们将介绍梯度下降法的有效性,即 梯度下降法是欧氏参数空间中的最速下降算法。
这里的最速是指:当更新步长在极小范围内时,选取梯度的负方向能够最大化目标函数下降的程度。用形式化语言表述为:
Theorem 1. (梯度下降是欧氏距离下的最速下降) 给定参数为
定理1的证明在附录中展示。
注意在上面的结论中,对步长的约束条件
- 参数空间
是标准欧几里得空间 - 参数之间构成欧几里得空间的一组标准正交基
当这个假设不能很好满足的时候,梯度下降的最速性质可能会大打折扣。
为此,我们可以考虑使用其他的约束条件,以此摆脱对欧式参数空间的依赖。比如说,在自然梯度下降 (Natural Gradient Descent) 中,使用的是 对数似然函数之间的KL散度 作为约束条件。
Theorem 2. (自然梯度下降是黎曼距离下的最速下降). 定义对数似然函数
其中,
因此,自然梯度下降的参数更新公式为:
1.4. 从动力学视角理解BGD和SGD [4]
这一节的内容来自于苏神的博客。感谢苏神,每次读苏神的博客都是茅塞顿开,建议每个学cs的人都去读一读!
我们重新书写一下BGD的更新公式:
由于学习率一般都有
这正好就是一个ODE动力学系统,而公式
由于BGD需要计算所有样本的梯度,每迭代一次所需的成本太大,因此我们希望随机从训练集
其中,
然而,我们希望最小化的是全样本梯度,即
我们假设二者的梯度之间的差异服从一个方差为
在这种假设下,SGD相当于在动力系统
其中,
这个SDE对应的平衡状态的概率分布为:
从公式
此外,误差的方差
在训练初期,batch size应小一些,那么噪声方差
二、学习率自适应方法
SGD的一大缺陷是对学习率
此外,SGD并非在整个训练集上计算梯度,而是在一个batch上对梯度进行估计。这就导致了梯度的方向不一定正确,有可能让模型沿着错误的方向优化,再也不可能达到全局最小值。
为此,一些现代优化器主要有两大改进点:动态调整学习率、引入动量。
2.1. SGDM (SGD with Momentum)
SGDM在SGD的基础上增加了一阶动量(也被称为一阶矩),考虑了历史梯度下降的方向:
其中
2.1.1. SGDM的动力学表述 [4]
从上面的内容我们可以知道,SGD能够用一个ODE来表述:
这里只考虑了参数
公式
下面我们将展示,公式
首先,我们记
我们对
且注意到,上式在
类似地,从公式
我们记:
分别代入公式
这正好与SGDM的更新公式等价。在数学上,这也被称作“蛙跳积分法” (Leapfrog Intergration),因为速度 (
对比公式
2.1.3. 的选择 [4]
前面我们提到,常数
因此我们需要摩擦力来消耗掉能量,使其能够最终停留在最低点。所以说,引入摩擦力
从公式
在使用SGDM时,如果学习率从
2.2. AdaGrad (Adaptive Gradient Algorithm)
AdaGrad是自适应学习率的先驱,其更新规则如下:
其中,
第一个式子维护了 历史梯度的平方
AdaGrad 开启了自适应学习率的时代,但其过于激进的衰减策略限制了广泛应用。
2.2.1. 为什么用梯度来调节学习率 [5]
在这一小节,我们从动力学的角度来说明为什么我们可以用梯度的大小来调节学习率。
从上面的内容我们知道,梯度下降算法是利用欧拉法求解一个ODE所描述的动力学系统,并最终收敛到最小值的过程:
理论上来说,当
首先我们分析估计的精确程度。由泰勒级数,我们有:
当学习率
这样我们就可以通过梯度来调节学习率。
2.3. RMSProp
AdaGrad过分依赖于历史梯度,在梯度突变时无法迅速响应。
RMSProp在AdaGrad的基础上引入了指数移动平均,可以通过超参数
2.3.1. 为什么要引入指数移动平均 [5]
我们继续 2.2.1 小节中的讨论,来展示为什么RMSProp要引入指数移动平均。
我们将公式
只用到了梯度的符号信息,也就是说不管梯度大小如何,每次迭代中
注意,单纯从求解ODE的角度来看这其实没有问题。因为ODE的解是一条轨迹
所以,为了改善这种情况,同时还能够保留梯度调整学习率的特征,RMSProp就引入了指数移动平均。
此外,滑动平均还有一个重要的原因。在实践中我们每次只能利用一个batch的梯度来估算全样本的梯度,这样一来每次算出来的结果实际上是有偏的,而滑动平均在一定程度上能修正这种偏差。
2.4. Adam (Adaptive Momentum Estimation)
SGDM中考虑了一阶矩,AdaGrad和RMSProp中则考虑了二阶矩。
在Adam中则同时考虑了一阶和二阶矩,使其对于超参数有非常好的鲁棒性。
由于矩估计量的初始值
为了纠正初始偏差,Adam并非直接使用一、二阶矩,而是对其进行了缩放(即
偏差修正的数学原理放在附录中展示。
2.4.1. 从Hessian近似的角度看Adam的优越性 [6]
在这一小节中,我们将从一个新颖的角度来看到Adam这类自适应学习率的优化器:梯度平方的指数滑动平均某种程度上近似于在估计Hessian矩阵的平方,从而Adam、RMSprop等优化器实际上近似于用二阶牛顿法来求解极值问题。
牛顿法通过将损失函数
其中,
我们假设Hessian矩阵
注意上式没有任何额外的参数,因此牛顿法天生就是一种自适应学习率算法。
然而,由于Hessian矩阵的复杂度为参数量的平方,因此在实际应用中,我们往往需要对Hessian矩阵做比较大的简化假设(入对焦矩阵或低秩矩阵)。具体来说,SGD假设了
下面,我们希望证明
证明的要点是考虑梯度
于是我们有:
在训练进入一定阶段后,我们假设
假设 Hessian矩阵是一个对焦矩阵,则上式我们只需要保留对角线元素:
即:
对比Adam的更新公式,可以发现Adam的
这个结论也可以解释为什么Adam中一般都令
2.5. AdamW (Adam with Decoupled Weight Decay)
在实践中,我们通常会在损失函数上加入L2正则化项,以此来防止过拟合:
此时,梯度下降的更新规则可以写为:
这里的
在Adam中使用L2正则化的更新规则为:
在Adam的原始实现中,权重衰减项被包含在了梯度中,权重越大,自适应学习率越小。对于那些需要较大权重的参数,L2正则化带来的惩罚反而因为自适应学习率的变小而被削弱了,使得L2正则化在Adam中效果不佳。
因此,AdamW修正了这一点。它将权重衰减的步骤与梯度更新的步骤分开,权重衰减项
当前,AdamW已经成为了现代大模型训练的标准配置优化器。
Appendix
Proof on Theorem 1
对极限内的目标函数进行一阶泰勒展开得:
因此,最优化任务可以转为如下的带约束的优化问题:
定义拉格朗日函数
其KKT条件为:
- 驻点条件:
; - 互补松弛条件:
; - 原始可行条件:
; - 对偶可行条件:
。
由驻点条件我们有:
代回互补松弛条件得到:
再代回驻点条件得:
因此,代回原极限即得:
Q.E.D
Adam中的初始偏差修正
Adam优化器并非直接使用原始的矩估计
下面我们以一阶矩估计
在Adam中,一阶矩估计的更新规则为指数移动平均:
其中,
我们假设梯度
我们对
对上式取期望得:
可以看到,当
为了消除训练初期的估计偏差,我们可以对