- 优化器 (1) SGD和Adam
- 优化器 (2) Muon
引言
在这一系列的博客中，我将从梯度下降开始介绍，逐步深入现代优化器的设计艺术。
后续介绍的内容主要可以划分为三大类：
- 基础梯度下降方法：BGD、SGD
- 学习率自适应方法：AdaGrad、RMSProp、Adam、AdanW
- 新晋方法：Muon、AdaMuon
在本文中主要包括了前两类，即梯度下降方法和自适应学习率方法。最近的新晋方法Muon则放在在下一篇文章中再详细介绍。
一、基础梯度下降方法
深度学习算法的基石是梯度下降 (Gradient Descent, GD) 算法。
具体来说，给定一个具有可学习参数 $ \theta\in\mathbb{R}^{n} $ 的参数化模型，我们希望找到最优的参数 $ \hat{\theta} $ 使得某个损失函数 $ \mathcal{J}(\theta) $ 最小，即：
\[\begin{equation} \hat{\theta}=\arg\min_{\theta\in\mathbb{R}^{n}}\mathcal{J}(\theta) \end{equation}\]在实际问题中，损失函数 $\mathcal{J}$ 的形式往往非常复杂。为了找到损失函数的最小值对应的参数 $\theta$，我们使用梯度下降法，每次沿着损失函数梯度的方向（也就是损失函数变化最大的方向）移动一小步，直到走到某个极小值。
1.1. 什么是梯度下降法
一个 $n$ 元函数 $f$ 在某个点 $x=(x_1,x_2,\dots,x_n)$ 的梯度是在该点上该函数变化率最大的向量，即：
\[\begin{equation} \nabla_xf:=\left( \frac{\partial f}{\partial x_1},\frac{\partial f}{\partial x_2},\dots,\frac{\partial f}{\partial x_n} \right)\in\mathbb{R}^{n} \end{equation}\]梯度下降法是一种迭代求极值的方法。在每一步迭代中，每一个参数都沿着梯度下降的方法前进一小步，直到走到某个极小值：
\[\begin{equation} \theta_{t+1}=\theta_t-\eta\cdot\nabla_\theta\mathcal{J}(\theta_t) \end{equation}\]其中，$\eta$ 称为学习率 (learning rate)，是梯度下降法最重要的超参数之一。注意这里要把学习率和更新步长区分开来，在后续我们会介绍这两者的区别。
1.2. BGD和SGD
BGD (Batch Gradient Descent) 是最简单的梯度下降算法，它完全按照公式 $(3)$ 的更新方式来训练参数。在每次更新参数时，都使用整个训练集来计算损失函数的梯度。
为了解决BGD的效率问题，SGD (Stochastic Gradient Descent) 每次仅使用少量样本（通常称为一个batch）来估计梯度，可以有效兼顾效率和稳定性。
\[\begin{equation} \theta\leftarrow\theta-\eta\cdot\nabla_\theta \mathcal{J}(\theta;x^{(i:i+n)},y^{(i:i+n)}) \end{equation}\]1.3. 梯度下降法的有效性 [2,3]
下面我们将介绍梯度下降法的有效性，即 梯度下降法是欧氏参数空间中的最速下降算法。
这里的最速是指：当更新步长在极小范围内时，选取梯度的负方向能够最大化目标函数下降的程度。用形式化语言表述为：
Theorem 1. (梯度下降是欧氏距离下的最速下降) 给定参数为 $\theta$ 的函数 $f:\Theta\mapsto\mathbb{R}$，我们有：
\[\begin{equation} \lim_{\varepsilon\rightarrow 0} \frac{1}{\varepsilon}\left( \arg\min_{\|\delta\|\le\varepsilon}f(\theta+\delta) \right) =-\frac{\nabla_\theta f}{\|\nabla_\theta f\|} \end{equation}\]定理1的证明在附录中展示。
注意在上面的结论中，对步长的约束条件 $|\delta|\le\varepsilon$ 是居于欧几里得距离定义的。因此这里隐含了如下两个假设：
- 参数空间 $\Theta$ 是标准欧几里得空间
- 参数之间构成欧几里得空间的一组标准正交基
当这个假设不能很好满足的时候，梯度下降的最速性质可能会大打折扣。
为此，我们可以考虑使用其他的约束条件，以此摆脱对欧式参数空间的依赖。比如说，在自然梯度下降 (Natural Gradient Descent) 中，使用的是 对数似然函数之间的KL散度 作为约束条件。
Theorem 2. (自然梯度下降是黎曼距离下的最速下降). 定义对数似然函数 $l(\theta):=\log p(x\mid \theta)$，我们有：
\[\begin{equation} \begin{aligned} \lim_{\varepsilon\rightarrow 0} \frac{1}{\varepsilon}\left( \arg\min_{D_{KL}(l(\theta)\|l(\theta+\delta))\le\varepsilon^2}l(\theta+\delta) \right) &\approx \lim_{\varepsilon\rightarrow 0} \frac{1}{\varepsilon}\left( \arg\min_{\delta^TF(\theta)\delta\le\varepsilon^2}l(\theta+\delta) \right)\\ &=-CF^{-1}(\theta)\nabla_\theta l(\theta) \end{aligned} \end{equation}\]其中，$F(\theta):=\mathbb{E}\left[ \nabla_\theta l(\theta)\nabla_\theta l(\theta)^T \right]$ 是对数似然函数的Fisher信息矩阵。
因此，自然梯度下降的参数更新公式为：
\[\begin{equation} \theta\leftarrow\theta-\eta\cdot F^{-1}(\theta)\nabla_\theta l(\theta) \end{equation}\]1.4. 从动力学视角理解BGD和SGD [4]
这一节的内容来自于苏神的博客。感谢苏神，每次读苏神的博客都是茅塞顿开，建议每个学cs的人都去读一读！
我们重新书写一下BGD的更新公式：
\[\begin{equation} \theta_{n+1}=\theta_n-\eta\cdot\nabla_\theta\mathcal{J}(\theta_n) \end{equation}\]由于学习率一般都有 $0\lt\eta\ll1$，因此我们可以变形为：
\[\begin{equation} \begin{aligned} \dot{\theta} &\approx \frac{\theta_{n+1}-\theta_{n}}{\eta}\\ &=-\nabla_\theta\mathcal{J}(\theta) \end{aligned} \end{equation}\]这正好就是一个ODE动力学系统，而公式 $(8)$ 则是动力系统 $(9)$ 的一个欧拉解法。由于 $(9)$ 是一个保守动力系统（即总能量保持不变），因此它最终可以收敛到一个不动点，并这个不动点是一个极小值点（但未必是全局最小的）。
由于BGD需要计算所有样本的梯度，每迭代一次所需的成本太大，因此我们希望随机从训练集 $S$ 中抽取一个子集 $R\subseteq S$，然后每次只计算 $R$ 中样本的梯度。这就是SGD的迭代方法：
\[\begin{equation} \theta_{n+1}=\theta_n-\eta\cdot\nabla_\theta\mathcal{J}_R(\theta_n) \end{equation}\]其中，
\[\begin{equation} \mathcal{J}_R(\theta_n)=\frac{1}{|R|}\sum_{x\in R}\mathcal{J}(x;\theta_n) \end{equation}\]然而，我们希望最小化的是全样本梯度，即 $\nabla_\theta\mathcal{J}(\theta_n)$，我们希望通过随机变量 $\nabla_\theta\mathcal{J}_R(\theta_n)$ 来估计全样本梯度。
我们假设二者的梯度之间的差异服从一个方差为 $\sigma^2$ 的正态分布，即：
\[\begin{equation} \nabla_\theta\mathcal{J}(\theta_n)-\nabla_\theta\mathcal{J}_R(\theta_n)=\xi_n\sim\mathcal{N}(\mu,\sigma^2) \end{equation}\]在这种假设下，SGD相当于在动力系统 $(9)$ 中引入了一个高斯噪声：
\[\begin{equation} \begin{aligned} \dot{\theta} &=-\nabla_\theta\mathcal{J}(\theta)+\sigma\xi \end{aligned} \end{equation}\]其中，$\xi\sim\mathcal{N}(0,1)$。也就是说，SGD对应的动力系统从ODE变为了SDE。我们一般称这个动力学方程为朗之万方程。
这个SDE对应的平衡状态的概率分布为：
\[\begin{equation} P(\theta)\sim\exp\left(-\frac{\mathcal{J}(\theta)}{\sigma^2} \right) \end{equation}\]从公式 $(14)$ 我们可以看到，在SDE的平衡状态中，参数 $\theta$ 出现概率越高的地方对应的损失函数 $\mathcal{J}(\theta)$ 越小。也即是说，如果我们执行SGD足够多次，理论上 $\theta$ 能够大概率收敛到全剧最优解。
此外，误差的方差 $\sigma^2$ 是取决于子集 $R$ 的大小（即batch size）的。batch size越大，梯度估计的方差越小，反之亦然。因此，在使用SGD作为优化器时，我们可以有以下的经验法则：
在训练初期，batch size应小一些，那么噪声方差 $\sigma^2$ 就会大一些，越接近均匀分布，算法就会遍历更多的区域。随着迭代次数的增加，慢慢地就会越来越接近最优区域，这时候方差应该要下降，使得极值点更为突出。也就是说，batch size应该要随着迭代次数而缓慢增加。同样地，学习率则应该缓慢减小。
二、学习率自适应方法
SGD的一大缺陷是对学习率 $\eta$ 的设置比较敏感，特别是在一些非凸优化问题中，学习率太大或太小都容易导致无法收敛，或收敛到局部最小而非全局最小。
此外，SGD并非在整个训练集上计算梯度，而是在一个batch上对梯度进行估计。这就导致了梯度的方向不一定正确，有可能让模型沿着错误的方向优化，再也不可能达到全局最小值。
为此，一些现代优化器主要有两大改进点：动态调整学习率、引入动量。
2.1. SGDM (SGD with Momentum)
SGDM在SGD的基础上增加了一阶动量（也被称为一阶矩），考虑了历史梯度下降的方向：
\[\begin{equation} \begin{aligned} v_t&=\gamma v_{t-1}+\eta\nabla_\theta \mathcal{J}(\theta_t)\\ \theta_{t+1}&=\theta_t-v_t \end{aligned} \end{equation}\]其中 $\gamma$ 为衰减因子，一般取 $\gamma=0.9$，这就意味着下降方向主要是动量方向，并略微偏向当前时刻的梯度方向。
2.1.1. SGDM的动力学表述 [4]
从上面的内容我们可以知道，SGD能够用一个ODE来表述：
\[\begin{equation} \dot{\theta}=-\nabla_\theta\mathcal{J}(\theta) \end{equation}\]这里只考虑了参数 $\theta$ 的一阶导。一个直接的问题是我们能不能引入二阶导数呢？为此，我们考虑一般的情况：
\[\begin{equation} \ddot{\theta}+\lambda\dot{\theta}=-\nabla_\theta\mathcal{J}(\theta) \end{equation}\]公式 $(17)$ 表示了一个标准的牛顿力学系统，其中 $\lambda$ 在系统中起到类似摩擦力的作用。
下面我们将展示，公式 $(17)$ 描述的动力系统的一个欧拉解法就是公式 $(15)$ 所描述的SGDM算法，且SGDM算法的收敛速度要快于SGD。
首先，我们记 $\gamma:=\dot{\theta}$，则公式 $(17)$ 可变形为
\[\begin{equation} \dot{\gamma}=-\lambda\gamma-\nabla_\theta\mathcal{J}(\theta) \end{equation}\]我们对 $\gamma$ 进行中心差分离散化得：
\[\begin{equation} \begin{aligned} \frac{\theta_{n+1}-\theta_n}{\eta} &\approx\dot{\theta}_{n+\frac{1}{2}}\\ &=\gamma_{n+\frac{1}{2}} \end{aligned} \end{equation}\]且注意到，上式在 $t_{n+\frac{1}{2}}$ 处具有 $O(\eta^2)$ 的二阶精度，公式 $(9)$ 中的前向查分则只有 $O(\eta)$ 的一阶精度。
类似地，从公式 $(18)$ 我们可以得到以下的二阶近似结果：
\[\begin{equation} \begin{aligned} \frac{\gamma_{n+\frac{1}{2}}-\gamma_{n-\frac{1}{2}}}{\eta} &\approx\dot{\gamma}_{n}\\ &=-\lambda\left(\frac{\gamma_{n+\frac{1}{2}}+\gamma_{n-\frac{1}{2}}}{2}\right)-\nabla_\theta\mathcal{J}(\theta) \end{aligned} \end{equation}\]我们记：
\[\begin{equation} \begin{aligned} v_{n+1}&=\eta\cdot\gamma_{n+\frac{1}{2}}\\ \beta&=\frac{1-\lambda\cdot\frac{\eta}{2}}{1+\lambda\cdot\frac{\eta}{2}}\\ \alpha&=\frac{\eta^2}{1+\lambda\cdot\frac{\eta}{2}}\\ \end{aligned} \end{equation}\]分别代入公式 $(19)$、$(20)$ 得：
\[\begin{equation} \begin{aligned} v_{n+1}&=\beta v_n-\alpha\nabla_\theta\mathcal{J}(\theta)\\ \theta_{n+1}&=\theta_n+v_{n+1} \end{aligned} \end{equation}\]这正好与SGDM的更新公式等价。在数学上，这也被称作“蛙跳积分法” (Leapfrog Intergration)，因为速度 ($v_{n+1}$) 和位置 ($\theta_{n+1}$) 是交替更新的。
对比公式 $(8)$ 和 $(22)$ 我们可以看到，SGD的迭代步长为 $O(\eta)$，而SGDM的迭代步长实际上为 $O(\alpha)=O(\eta^2)$。因此，动量可以在同等学习率、不损失精度的情况下，为越过不那么好的极小值点提供了来自动力学的可能性，使得整个算法以更大步长前进。
2.1.3. $\lambda$ 的选择 [4]
前面我们提到，常数 $\lambda>0$ 这一项相当于摩擦力，用来消耗系统中的能量。如果没有摩擦力，那么不管学习率多小，只要不为0，则SGDM算法都永远不会收敛到极小值点。这是能量守恒定律告诉我们的，因为在最小值点的动能是最大的，也就是速度是最快的。
因此我们需要摩擦力来消耗掉能量，使其能够最终停留在最低点。所以说，引入摩擦力 $\lambda$ 对于算法的收敛性是必要的。但是摩擦力也不能太大，太大的摩擦力会导致算法过早收敛。我们是否有一些经验准则，来指导 $\lambda$ 的选择呢？
从公式 $(21)$ 可以看到，当 $\lambda$ 固定时，如果学习率 $\alpha$ 降低，此时 $\eta$ 降低，因此 $\beta$ 将随之升高。其中的比例可以估算得到如下结论：
在使用SGDM时，如果学习率从 $\alpha$ 降低到 $r\alpha$，则 $\beta$ 应该升高到 $1-(1-\beta)\sqrt{r}$。
2.2. AdaGrad (Adaptive Gradient Algorithm)
AdaGrad是自适应学习率的先驱，其更新规则如下：
\[\begin{equation} \begin{aligned} G_t&=G_{t-1}+g_t\odot g_t\\ \theta_{t+1}&=\theta_t-\frac{\eta}{\sqrt{G_t+\varepsilon}}\odot g_t \end{aligned} \end{equation}\]其中，$g_t=\nabla_\theta\mathcal{J}(\theta_t)$ 表示当前的梯度，$\odot$ 表示向量逐元素相乘。
第一个式子维护了 历史梯度的平方 $G_t$ （也被称为二阶矩）。每个参数实际的学习率为 $\frac{\eta}{\sqrt{G_t+\varepsilon}}$ 而非固定的常数 $\eta$，这意味着对于历史梯度较大的参数，其学习率较小，而没有被频繁更新的参数则获得较大的学习率。
AdaGrad 开启了自适应学习率的时代，但其过于激进的衰减策略限制了广泛应用。
2.2.1. 为什么用梯度来调节学习率 [5]
在这一小节，我们从动力学的角度来说明为什么我们可以用梯度的大小来调节学习率。
从上面的内容我们知道，梯度下降算法是利用欧拉法求解一个ODE所描述的动力学系统，并最终收敛到最小值的过程：
\[\begin{equation} \begin{aligned} \dot{\theta}_t &\approx \frac{\theta_{t+\eta}-\theta_{t}}{\eta}\\ &=-\nabla_\theta\mathcal{J}(\theta) \end{aligned} \end{equation}\]理论上来说，当 $\eta$ 越小，估计的精度就越高，但收敛所需的迭代次数也就越多。因此最恰当的方案是使得每一步的学习率够用就好。
首先我们分析估计的精确程度。由泰勒级数，我们有：
\[\begin{equation} \begin{aligned} \theta_{t+\eta} &=\theta_t+\eta\dot{\theta}_t+O(\eta^2)\\ &=\theta_t-\eta\nabla_\theta\mathcal{J}(\theta_t)+O(\eta^2)\\ \end{aligned} \end{equation}\]当学习率 $\eta$ 足够小时，我们可以期望平方误差项有上界：$O(\eta^2)\lt\eta\lvert\nabla_\theta\mathcal{J}(\theta_t)\rvert$。也就是说，如果我们将 $\eta\lvert\nabla_\theta\mathcal{J}(\theta_t)\rvert$ 控制在一定范围内，那么整个估计的误差界也就得到了控制。因此，我们可以令其为某一个常数 $\bar{\eta}=\eta\lvert\nabla_\theta\mathcal{J}(\theta_t)\rvert$。于是，学习率就可以写为：
\[\begin{equation} \eta=\frac{\bar{\eta}}{\lvert\nabla_\theta\mathcal{J}(\theta_t)\rvert} \end{equation}\]这样我们就可以通过梯度来调节学习率。
2.3. RMSProp
AdaGrad过分依赖于历史梯度，在梯度突变时无法迅速响应。
RMSProp在AdaGrad的基础上引入了指数移动平均，可以通过超参数 $\beta$ 来控制动量的权重。
\[\begin{equation} \begin{aligned} v_t&=\beta v_{t-1}+(1-\beta)g_t\odot g_t\\ \theta_{t+1}&=\theta_t-\frac{\eta}{\sqrt{v_t+\varepsilon}}\odot g_t \end{aligned} \end{equation}\]2.3.1. 为什么要引入指数移动平均 [5]
我们继续 2.2.1 小节中的讨论，来展示为什么RMSProp要引入指数移动平均。
我们将公式 $(26)$ 代入梯度下降的迭代公式中得：
\[\begin{equation} \begin{aligned} \theta_{t+\eta} &=\theta_t-\eta\cdot\nabla_\theta\mathcal{J}(\theta_t)\\ &=\theta_t-\frac{\bar{\eta}}{|\nabla_\theta\mathcal{J}(\theta_t)|}\cdot\nabla_\theta\mathcal{J}(\theta_t)\\ &=\theta_t-\bar{\eta}\cdot\text{sign}[\nabla_\theta\mathcal{J}(\theta_t)] \end{aligned} \end{equation}\]只用到了梯度的符号信息，也就是说不管梯度大小如何，每次迭代中 $\theta$ 都只是移动固定的长度，变成像是网格搜索一样的策略。
注意，单纯从求解ODE的角度来看这其实没有问题。因为ODE的解是一条轨迹 $(t,\theta_t)$，动态学习率虽然使得 $\theta$ 的步长固定，但时间 $t$ 的步长却变得不固定了。但我们关心的是优化问题，即求 $\mathcal{J}(\theta)$ 的极小值点，那么我们更希望 $\theta$ 的更新步长不固定。
所以，为了改善这种情况，同时还能够保留梯度调整学习率的特征，RMSProp就引入了指数移动平均。
此外，滑动平均还有一个重要的原因。在实践中我们每次只能利用一个batch的梯度来估算全样本的梯度，这样一来每次算出来的结果实际上是有偏的，而滑动平均在一定程度上能修正这种偏差。
2.4. Adam (Adaptive Momentum Estimation)
SGDM中考虑了一阶矩，AdaGrad和RMSProp中则考虑了二阶矩。
在Adam中则同时考虑了一阶和二阶矩，使其对于超参数有非常好的鲁棒性。
\[\begin{equation} \begin{aligned} m_t&=\beta_1m_{t-1}+(1-\beta_1)g_t\\ v_t&=\beta_2v_{t-1}+(1-\beta_2)g_t\odot g_t\\ \hat{m}_t&=\frac{m_t}{1-\beta_1^t}\\ \hat{v}_t&=\frac{v_t}{1-\beta_2^t}\\ \theta_{t+1}&=\theta_t-\frac{\eta}{\sqrt{\hat{v}_t}+\varepsilon}\odot\hat{m}_t \end{aligned} \end{equation}\]由于矩估计量的初始值 $m_0=v_0=0$，在训练初期，$m_t$ 和 $v_t$ 会较为偏向0，这会导致其低估了真实的矩。
为了纠正初始偏差，Adam并非直接使用一、二阶矩，而是对其进行了缩放（即 $\hat{m}_t$ 和 $\hat{v}_t$），使其期望值更接近真实的矩。
偏差修正的数学原理放在附录中展示。
2.4.1. 从Hessian近似的角度看Adam的优越性 [6]
在这一小节中，我们将从一个新颖的角度来看到Adam这类自适应学习率的优化器：梯度平方的指数滑动平均某种程度上近似于在估计Hessian矩阵的平方，从而Adam、RMSprop等优化器实际上近似于用二阶牛顿法来求解极值问题。
牛顿法通过将损失函数 $\mathcal{J}(\theta)$ 展开到二阶来更新参数 $\theta_{t+1}$：
\[\begin{equation} \mathcal{J}(\theta) \approx \mathcal{J}(\theta_t)+g_t^T(\theta-\theta_t)+\frac{1}{2}(\theta-\theta_t)^T\mathcal{H}_t(\theta-\theta_t) \end{equation}\]其中，$g_t=\nabla_{\theta}\mathcal{J}(\theta_{t})$ 表示梯度，$\mathcal{H}{t}=\nabla{\theta}^2\mathcal{J}(\theta_{t})$ 表示Hessian矩阵。
我们假设Hessian矩阵 $\mathcal{H}_{t}$ 是正定的，则上式存在唯一最小值：
\[\begin{equation} \theta_{t+1}=\theta_t-\mathcal{H}_t^{-1}g_t \end{equation}\]注意上式没有任何额外的参数，因此牛顿法天生就是一种自适应学习率算法。
然而，由于Hessian矩阵的复杂度为参数量的平方，因此在实际应用中，我们往往需要对Hessian矩阵做比较大的简化假设（入对焦矩阵或低秩矩阵）。具体来说，SGD假设了 $\mathcal{H}_t=\eta_t^{-1}I$，而Adam则是假设了 $\mathcal{H}_t=\eta_t^{-1}\text{diag}(\sqrt{\hat{v}_t}+\varepsilon)$。
下面，我们希望证明 $\eta_t^{-1}\text{diag}(\sqrt{\hat{v}_t}+\varepsilon)$ 是Hessian矩阵 $\mathcal{H}_t$ 的一个更好的近似。
证明的要点是考虑梯度 $g$ 在损失函数极小值点 $\theta^*$ 上的一阶泰勒展开：
\[\begin{equation} \begin{aligned} g_\theta &\approx g_{\theta^*}+\mathcal{H}_{\theta^*}(\theta-\theta^*)\\ &=\mathcal{H}_{\theta^*}(\theta-\theta^*) \end{aligned} \end{equation}\]于是我们有：
\[\begin{equation} g_\theta g_\theta^T\approx \mathcal{H}_{\theta^*}(\theta-\theta^*)(\theta-\theta^*)^T\mathcal{H}_{\theta^*}^T \end{equation}\]在训练进入一定阶段后，我们假设 $\theta-\theta^*$ 服从正态分布 $\mathcal{N}(0,\sigma^2)$，则：
\[\begin{equation} \begin{aligned} \mathbb{E}\left[g_\theta g_\theta^T\right] &\approx \mathcal{H}_{\theta^*}\mathbb{E}\left[(\theta-\theta^*)(\theta-\theta^*)^T\right]\mathcal{H}_{\theta^*}^T\\ &=\sigma^2\mathcal{H}_{\theta^*}\mathcal{H}_{\theta^*}^T \end{aligned} \end{equation}\]假设 Hessian矩阵是一个对焦矩阵，则上式我们只需要保留对角线元素：
\[\begin{equation} \text{diag}\left(\mathbb{E}\left[g_\theta\odot g_\theta\right]\right) \approx\sigma^2\mathcal{H}_{\theta^*}^2 \end{equation}\]即：
\[\begin{equation} \mathcal{H}_{\theta^*}\approx\frac{1}{\sigma}\text{diag}\left(\sqrt{\mathbb{E}\left[g_\theta\odot g_\theta\right]}\right) \end{equation}\]对比Adam的更新公式，可以发现Adam的 $\hat{v}t$ 是对梯度平方的滑动平均，这一步可以看做是在近似 $\mathbb{E}\left[g\theta\odot g_\theta\right]$，也就是Hessian矩阵 $\mathcal{H}_{\theta^{*}}$ 的一个估计。
这个结论也可以解释为什么Adam中一般都令 $\beta_{1}\lt\beta_{2}$。为了更加准确的估计Hessian，$\hat{v}t$ 的滑动平均应该尽可能长期，因此 $\beta{2}$ 应该很接近1。而动量 $\hat{m}t$ 是梯度的滑动平均，如果太过于长期的话，结果会将接近 $g{\theta^{*}}=0$，这反而不好，因此动量的滑动平均应该更加局部一些。
2.5. AdamW (Adam with Decoupled Weight Decay)
在实践中，我们通常会在损失函数上加入L2正则化项，以此来防止过拟合：
\[\begin{equation} \mathcal{J}(\theta)=\mathcal{J}_{ori}(\theta)+\frac{\lambda}{2}\|\theta\|^2 \end{equation}\]此时，梯度下降的更新规则可以写为：
\[\begin{equation} \theta_{t+1}=\theta_t-\eta\cdot\left(\nabla_\theta\mathcal{J}_{ori}(\theta_t)+\lambda\theta_t\right) \end{equation}\]这里的 $\lambda\theta_t$ 就是所谓的权重衰减项 (weight decay)。
在Adam中使用L2正则化的更新规则为：
\[\begin{equation} \theta_{t+1}=\theta_t-\frac{\eta}{\sqrt{\hat{v}_t}+\varepsilon}\left(\hat{m}_t+\lambda\theta_t\right) \end{equation}\]在Adam的原始实现中，权重衰减项被包含在了梯度中，权重越大，自适应学习率越小。对于那些需要较大权重的参数，L2正则化带来的惩罚反而因为自适应学习率的变小而被削弱了，使得L2正则化在Adam中效果不佳。
因此，AdamW修正了这一点。它将权重衰减的步骤与梯度更新的步骤分开,权重衰减项 $\lambda\theta_t$ 不参与梯度计算，而是直接作用于参数更新本身。AdamW的更新规则为：
\[\begin{equation} \theta_{t+1}=\theta_t-\eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\varepsilon}+\lambda\theta_t\right) \end{equation}\]当前，AdamW已经成为了现代大模型训练的标准配置优化器。
Appendix
Proof on Theorem 1
对极限内的目标函数进行一阶泰勒展开得：
\[\begin{equation} \begin{aligned} \arg\min_{\|\delta\|\le\varepsilon}f(\theta+\delta) &\approx \arg\min_{\|\delta\|\le\varepsilon}f(\theta)+(\nabla_\theta f)^T\delta\\ &=\arg\min_{\|\delta\|\le\varepsilon}(\nabla_\theta f)^T\delta \end{aligned} \end{equation}\]因此，最优化任务可以转为如下的带约束的优化问题：
\[\begin{equation} \min_\delta(\nabla_\theta f)^T\delta\qquad\text{s.t.}\|\delta\|^2\le\varepsilon^2 \end{equation}\]定义拉格朗日函数
\[\begin{equation} \mathcal{L}(\delta,\lambda):=(\nabla_\theta f)^T\delta+\lambda(\|\delta\|^2-\varepsilon^2) \end{equation}\]其KKT条件为：
- 驻点条件：$\nabla_\delta\mathcal{L}(\delta,\lambda)=0$；
- 互补松弛条件：$\lambda(|\delta|^2-\varepsilon^2)=0$；
- 原始可行条件：$|\delta|^2-\varepsilon^2\le 0$；
- 对偶可行条件：$\lambda\ge0$。
由驻点条件我们有：
\[\begin{equation} \begin{aligned} &\nabla_\theta f+2\lambda\delta=0\\ &\delta=-\frac{1}{2\lambda}\nabla_\theta f \end{aligned} \end{equation}\]代回互补松弛条件得到：
\[\begin{equation} \begin{aligned} &\lambda \left( \left\|-\frac{1}{2\lambda}\nabla_\theta f\right\|^2-\varepsilon^2 \right)=0\\ &\lambda=\frac{1}{2\varepsilon}\|\nabla_\theta f\| \end{aligned} \end{equation}\]再代回驻点条件得：
\[\begin{equation} \delta^*=-\varepsilon\frac{\nabla_\theta f}{\|\nabla_\theta f\|} \end{equation}\]因此，代回原极限即得：
\[\begin{equation} \lim_{\varepsilon\rightarrow 0} \frac{1}{\varepsilon}\left( \arg\min_{\|\delta\|\le\varepsilon}f(\theta+\delta) \right) =-\frac{\nabla_\theta f}{\|\nabla_\theta f\|} \end{equation}\]Q.E.D
Adam中的初始偏差修正
Adam优化器并非直接使用原始的矩估计 $m_t$ 和 $v_t$，而是对其进行了偏差修正：
\[\begin{equation} \begin{aligned} \hat{m}_t&=\frac{m_t}{1-\beta_1^t}\\ \hat{v}_t&=\frac{v_t}{1-\beta_2^t} \end{aligned} \end{equation}\]下面我们以一阶矩估计 $m_t$ 为例，说明偏差修正的数学原理。二阶矩估计 $v_t$ 可以完全类比。
在Adam中，一阶矩估计的更新规则为指数移动平均：
\[\begin{equation} m_t=\beta_1m_{t-1}+(1-\beta_1)g_t \end{equation}\]其中，$\beta_1$ 是衰减率，$g_t=\nabla_\theta\mathcal{J}(\theta_{t-1})$ 是在第 $t$ 步获得第梯度向量，$m_0$ 的初始值为0。
我们假设梯度 $g_t$ 是一个平稳的随机过程，即其期望不随时间变化：
\[\begin{equation} \mathbb{E}[g_t]=\mu \end{equation}\]我们对 $m_t$ 的递推式进行展开得：
\[\begin{equation} \begin{aligned} m_t &=\beta_1m_{t-1}+(1-\beta_1)g_t\\ &=\beta_1\left(\beta_1m_{t-2}+(1-\beta_1)g_{t-1}\right)+(1-\beta_1)g_t\\ &=\beta_1^2m_{t-2}+\beta_1(1-\beta_1)g_{t-1}++(1-\beta_1)g_t\\ &=\cdots\\ &=(1-\beta_1)\sum_{i=1}^t\beta_1^{t-i}g_i \end{aligned} \end{equation}\]对上式取期望得：
\[\begin{equation} \begin{aligned} \mathbb{E}[m_t] &=\mathbb{E}\left[(1-\beta_1)\sum_{i=1}^t\beta_1^{t-i}g_i\right]\\ &=(1-\beta_1)\sum_{i=1}^t\beta_1^{t-i}\mathbb{E}[g_i]\\ &=\mu(1-\beta_1)\sum_{i=1}^t\beta_1^{t-i}\\ &=\mu(1-\beta_1)\frac{1-\beta_1^t}{1-\beta_1}\\ &=\mu(1-\beta_1^t)\\ &=(1-\beta_1^t)\cdot \mathbb{E}[g_t] \end{aligned} \end{equation}\]可以看到，当 $t$ 较小时（训练初期），$1-\beta_1^t\approx 0$，此时 $m_t$ 严重低估了真实梯度 $g_t$ 的梯度。在训练后期，$1-\beta_1^t\approx 1$，偏差逐渐消失。
为了消除训练初期的估计偏差，我们可以对 $m_t$ 进行修正，使其变成对 $g_t$ 的无偏估计量，即：
\[\begin{equation} \begin{aligned} \hat{m}_t&:=\frac{m_t}{1-\beta_1^t}\\ \mathbb{E}[\hat{m}_t] &=\frac{1}{1-\beta_1^t}\mathbb{E}[m_t]\\ &=\mathbb{E}[g_t] \end{aligned} \end{equation}\]