AMSGrad:改进Adam算法的收敛性.

Adam等优化方法使用平方梯度的指数滑动平均值的平方根对梯度进行缩放,这些算法在较大的输出空间中经常无法收敛到最优解。作者证明其原因是算法中使用了指数滑动平均,通过赋予这些算法对过去梯度的“长期记忆”可以解决收敛问题,基于此提出了AMSGrad算法。

1. 自适应优化算法

自适应优化方法通常使用指数滑动平均累计梯度的低阶矩估计,如Adam算法在更新过程中计算梯度和平方梯度的指数滑动平均值:

\[m_t = \beta_1 \cdot m_{t-1} + (1-\beta_1) \cdot g_t\] \[v_t = \beta_2 \cdot v_{t-1} + (1-\beta_2) \cdot g_t^2\]

其中$m_t$是一阶矩(均值)的估计,$v_t$是二阶矩(未中心化的方差)的估计。若忽略偏差修正和分母的小值,则参数的更新步骤为:

\[\theta_t = \theta_{t-1} - \alpha \cdot \frac{m_t}{\sqrt{v_t}}\]

注意到指数滑动平均值也可以表示为:

\[m_t = \beta_1 \cdot m_{t-1} + (1-\beta_1) \cdot g_t \\ = \beta_1 \cdot (\beta_1 \cdot m_{t-2} + (1-\beta_1) \cdot g_{t-1}) + (1-\beta_1) \cdot g_t \\ = \beta_1^tm_0 + \beta_1^{t-1} (1-\beta_1) g_{1} + ... + \beta_1 (1-\beta_1) g_{t-1} + (1-\beta_1) g_t \\ = (1-\beta_1)\sum_{i=1}^{t}\beta_1^{t-i}g_{i}\] \[v_t = \beta_2 \cdot v_{t-1} + (1-\beta_2) \cdot g_t^2 \\ = (1-\beta_2)\sum_{i=1}^{t}\beta_2^{t-i}g_{i}^2\]

2. 收敛性问题

Adam论文中,原作者声称Adam在凸优化问题上是收敛的。然而本文作者指出,即使在一些简单的一维凸环境中,Adam也不能收敛到最优解,造成该问题的原因是在计算自适应学习率时使用了指数滑动平均值。

为便于讨论,假设$\beta_1=0$(即不使用动量),则自适应优化公式为:

\[v_t = (1-\beta_2)\sum_{i=1}^{t}\beta_2^{t-i}g_{i}^2\] \[\theta_t = \theta_{t-1} - \frac{\alpha_t}{\sqrt{v_t}} \cdot g_t\]

参数更新中的自适应学习率为$\frac{\alpha_t}{\sqrt{v_t}}$。观察学习率的倒数的变化:

\[\Gamma_{t+1} = \frac{\sqrt{v_{t+1}}}{\alpha_{t+1}}-\frac{\sqrt{v_t}}{\alpha_t}\]

通常随着更新轮数$t$增大,自适应学习率应该减小。假设基本学习率$\alpha$不变,对于SGDAdaGrad等算法,有$v_{t+1} \geq v_t$,因此$\Gamma_{t} \geq 0, \forall t$,算法是收敛的。

当引入指数滑动平均时:

\[v_{t+1} = \beta_2 \cdot v_{t} + (1-\beta_2) \cdot g_{t+1}^2\]

此时无法保证$v_{t+1} \geq v_t$,因此可能出现$\Gamma_{t}<0$的情况,使得算法不能收敛到最优解。

3. 一个使Adam不收敛的例子

假设特征$x$的取值为$[-1,1]$,损失函数为:

\[f_t(x) = \begin{cases} Cx, & \text{for }t\text{ mod }3=1 \\ -x, & \text{otherwise} \end{cases}\]

其中$C>2$。注意到该损失函数是一个凸函数,且极小值点为$x=-1$。在更新中,每三轮中的一轮算法会获得一次较大的梯度$C$,使得参数向最优值($-1$)更新;其余两轮算法的梯度是$-1$,使得参数向错误的方向(增大)更新。

不失一般性地假设$\beta_2 = 1/(1+C^2)$,则算法在前几轮的更新如下:

\[v_{1} = \frac{C^4}{1+C^2} \\ x_1 = x_{0} - \frac{\sqrt{1+C^2}}{C} \cdot \alpha\] \[v_{2} = \frac{C^2(1+2C^2)}{(1+C^2)^2} \\ x_2 = x_{0} - (\frac{\sqrt{1+C^2}}{C}-\frac{1+C^2}{C\sqrt{1+2C^2}}) \cdot \alpha\] \[v_{3} = \frac{C^2(2+4C^2+C^4)}{(1+C^2)^3} \\ x_3 = x_{0} - (\frac{\sqrt{1+C^2}}{C}-\frac{1+C^2}{C\sqrt{1+2C^2}}-\frac{(1+C^2)^{3/2}}{C\sqrt{2+4C^2+C^4}}) \cdot \alpha\]

通常$\alpha>0$,若假设$x_0=0$,则经过三次更新后$x_3>0$,这意味着第一次更新中获得的较大的梯度$C$经过指数滑动平均后的影响已经减小,后两轮更新提供的错误梯度使得算法没有收敛到最优解。

4. AMSGrad算法

作者指出Adam算法中无法保证$v_{t+1} \geq v_t$,因此它具有较差的收敛性,并提出了改进的AMSGrad算法。

AMSGrad算法使用截止到当前轮数的平方梯度的滑动平均$v_t$的历史最大值来代替该平均值,使得梯度更新的有效步长$\frac{\alpha_t}{\sqrt{v_t}}$至少不会增加。

\[\hat{v}_t = \max(\hat{v}_{t-1},v_t)\] \[\theta_t = \theta_{t-1} - \alpha \cdot \frac{m_t}{\sqrt{\hat{v}_t}}\]

5. 实验结果

AMSGrad算法在逻辑回归和神经网络模型等非凸优化问题上取得比Adam算法更好的收敛性: