优化器十一：Adam

前言

  本文介绍Adam算法。

回顾AdaGrad与RMSProp

  在之前的文章中我们介绍了AdaGrad，它提出的动机是考虑到这样一个简单事实：缩小历史改变量较大的参数的梯度；放大历史改变量较小的参数的梯度。直观而言，这种方法很适合稀疏梯度的情况（如Embedding层）。其参数更新式为：

$$ \theta^{(\tau+1)}_i=\theta^{(\tau)}_i-\frac{\eta}{\sqrt{\sum\limits_{k=1}^{\tau}\left(g^{(k)}_i \right)^2+\varepsilon}}\cdot g^{(\tau)}_i\tag{1} $$

  而后我们从赋权的角度引出了RMSProp，其让梯度的权重动态地进行指数衰减从而避免了分母逐渐减小的问题。这不仅使得RMSProp具有AdaGrad适用于稀疏梯度的优点，其还具有适合非平稳目标的特点（参见AdaDelta）。RMSProp的参数更新式为：

$$ \theta^{(\tau+1)}_i=\theta^{(\tau)}_i-\frac{\eta}{\sqrt{\sum\limits_{k=1}^{\tau}\beta^{\tau-k}\left(g^{(k)}_i \right)^2+\varepsilon}}\cdot g^{(\tau)}_i\tag{2} $$

矩估计

1.用矩估计的观点理解自适应策略

  接下来我们从矩估计的观点重新理解自适应策略。首先来看AdaGrad，在学习率衰减的情况下：$\eta^{(\tau)}=\eta/\sqrt{\tau}$，我们有：

$$ \begin{align} \theta^{(\tau+1)}_i&=\theta^{(\tau)}_i-\frac{\eta/\sqrt \tau}{\sqrt{\frac {1}{\tau}\sum\limits_{k=1}^{\tau}\left(g^{(k)}_i \right)^2}}\cdot g^{(\tau)}_i\tag{3} \end{align} $$

  对于自适应学习率的分母根号下的值，我们有：

$$ \begin{align} \mathbb{E}[\frac {1}{\tau}\sum\limits_{k=1}^{\tau}\left(g^{(k)}_i \right)^2]=\mathbb{E}[g_i^2]\tag{4} \end{align} $$

所以此时分母根号下的值是 $g^2_i$的一个无偏估计！
  对于RMSProp，在实践中我们通常采用以下的版本：

$$ \begin{align} v^{(\tau)}_i&=\beta v_i^{(\tau-1)}+(1-\beta)\left(g_i^{(\tau)} \right)^2\\ \theta^{(\tau+1)}&=\theta^{(\tau)}_i -\frac{\eta}{\sqrt{v_i^{(\tau)}}}\cdot g^{(\tau)}_i \end{align}\tag{5} $$

  对于$v_i^{(\tau)}$我们将其展开：

$$ v^{(\tau)}_i=(1-\beta)\sum_{k=1}^\tau \beta^{n-k}\left(g_i^{(\tau)} \right)^2+\beta^\tau v^{(0)}_i\tag{6} $$

  在两边取期望：

$$ \begin{align} \mathbb{E}[v_i^{(\tau)}]&=\mathbb{E}[g_i^2]\cdot (1-\beta)\cdot \frac{1-\beta^\tau}{1-\beta}+\beta^\tau v_i^{(0)}\\ &=(1-\beta^\tau)\mathbb{E}[g^2_i]+\beta^\tau v_i^{(0)} \end{align} $$

  所以当$v_i^{(0)}$初始化为 $\mathbb{E}[g_i^2]$时，$v_i^{(\tau)}$同样可以认为是 $g^2_i$的一个无偏估计！所以从矩估计的观点而言，自适应学习率的框架为：

$$ \theta^{(\tau+1)}_i=\theta^{(\tau)}_i-\frac{\eta}{\sqrt{v_i^{(\tau)}+\varepsilon}}\cdot g^{(\tau)}_i\tag{7} $$

其中 $v_i^{(\tau)}$是随机梯度二阶原点距（$\mathbb{E}[g_i^2]$）的一个估计。而$\mathbb{E}[g_i^2]$可以认为是未中心化的方差：$\mathbb{E}[g_i^2]=\mathbb{D}[g_i]+\mathbb{E}^2[g_i]$。形如（7）式的参数更新策略对于历史改变量较小（此时方差较小，如平原、鞍点等情况）的参数将放大步长；对于历史改变量较大（此时方差较大，如峡谷等情况）的参数将缩小步长。由矩估计来理解自适应策略的视角，相比我们直接得出参数更新式要更加自然。

2.用矩估计的观点理解动量

  回顾我们介绍的动量方法，对于HB而言，其参数更新式为：

$$ \begin{cases} \boldsymbol{m}^{(\tau)}=\beta^{(\tau)}\boldsymbol{m}^{(\tau-1)}+(1-\beta^{(\tau)})(-\boldsymbol{g}^{(\tau)}) \\ \boldsymbol{\theta}^{(\tau+1)} = \boldsymbol{\theta}^{(\tau)} + \eta\cdot \boldsymbol{m}^{(\tau)} \end{cases}\tag{8} $$

把它展开：$\theta^{(\tau+1)}_i=\theta^{(\tau)}_i-\eta\cdot\bigg(\underbrace{ (1-\beta)\sum\limits_{k=1}^\tau \beta^{\tau-k}g^{(k)}_i+\beta^\tau m^{(0)}_i }_{m^{(\tau)}_i}\bigg)$，类似于RMSProp的考虑，我们可以将$m^{(\tau)}_i$视作随机梯度一阶矩（$\mathbb{E}[g_i]$）的一个估计。而且这种指数移动平均的方式我们也证明了比单纯使用 $g_i$也即SGD更有效。

3.结合动量与自适应学习率

  Kingma与Ba于2014年结合了自适应学习率与动量提出了Adam算法，并为这两者给出了矩估计的同一解释，这正是Adam名称的来源：Adaptive Moment Estimation。Adam的参数更新框架为：

$$ \theta^{(\tau+1)}_i=\theta^{(\tau)}_i-\frac{\eta}{\sqrt{v^{(\tau)}_i}+\varepsilon}\cdot m^{(\tau)}_i\tag{9} $$

其中$m^{(\tau)}_i,v^{(\tau)}_i$分别是随机梯度 $g_i$的一阶矩与二阶原点距的估计。我们分别采用EMA的方式，则有：

$$ \begin{align} m^{(\tau)}_i&=\beta_1\cdot m_i^{(\tau-1)}+(1-\beta_1)\cdot g^{(\tau)}_i\\ v^{(\tau)}_i&=\beta_2\cdot v_i^{(\tau-1)}+(1-\beta_1)\cdot \left(g_i^{(\tau)} \right)^2\\ \theta^{(\tau+1)}_i&=\theta^{(\tau)}_i-\frac{\eta}{\sqrt{v^{(\tau)}_i}+\varepsilon}\cdot m^{(\tau)}_i \end{align}\tag{10} $$

  在实际情况中，我们通常将$m,v$初始化为0，而$v$初始化为0的影响往往大于$m$（这是由于 $\beta_2$往往大于$\beta_1$），所以在初始阶段我们将会得到过大的步长，往往会影响最终结果甚至导致不收敛。鉴于此，我们将渐进无偏的$m,v$进行初始偏差修正：

$$ \begin{align} \hat m^{(\tau)}_i &=m^{(\tau)}_i/(1-\beta_1^\tau)\\ \hat v^{(\tau)}_i &=v^{(\tau)}_i/(1-\beta_2^\tau) \end{align}\tag{11} $$

  综合（10）（11），我们便得到了Adam的参数更新式：

$$ \begin{align} m^{(\tau)}_i&=\beta_1\cdot m_i^{(\tau-1)}+(1-\beta_1)\cdot g^{(\tau)}_i\\ v^{(\tau)}_i&=\beta_2\cdot v_i^{(\tau-1)}+(1-\beta_2)\cdot \left(g_i^{(\tau)} \right)^2\\ \hat m^{(\tau)}_i &=m^{(\tau)}_i/(1-\beta_1^\tau)\\ \hat v^{(\tau)}_i &=v^{(\tau)}_i/(1-\beta_2^\tau)\\ \theta^{(\tau+1)}_i&=\theta^{(\tau)}_i-\frac{\eta}{\sqrt{\hat v^{(\tau)}_i}+\varepsilon}\cdot \hat m^{(\tau)}_i \end{align}\tag{12} $$

收敛性

  原文使用了在线凸优化的分析框架：设$\mathcal L\in \mathcal F_0^{1,1}$，考察$R(\tau)=\sum\limits_{k=1}^\tau \Big[\mathcal L(\boldsymbol{\theta}^{(k)})-\mathcal L(\boldsymbol{\theta}^*)\Big]$。由于原文的证明比较混乱而且Lemma 10.3与Lemma 10.4与主定理的证明有一些问题，这里我们给出一个清晰一点的证明过程。类似于AdaGrad的证明，由凸性：

$$ \mathcal L(\boldsymbol{\theta}^{(\tau)})-\mathcal L(\boldsymbol{\theta}^*)\leq \left\langle \boldsymbol{g}^{(\tau)} ,\boldsymbol{\theta}^{(\tau)}-\boldsymbol{\theta}^* \right\rangle=\sum_{i=1}^P g^{(\tau)}_i \left(\theta^{(\tau)}_i-\theta^*_i \right)\tag{13} $$

为了得出$g^{(\tau)}_i \left(\theta^{(\tau)}_i-\theta^*_i \right)$的表达式，我们从Adam的参数迭代式出发：

$$ \theta^{(\tau+1)}_i=\theta^{(\tau)}_i-\underbrace{\frac{\eta^{(\tau)}}{1-\prod\limits_{k=1}^\tau\beta_1^{(\tau)}}}_{\gamma^{(\tau)}}\cdot \frac{m_i^{(\tau)}}{\sqrt{\hat v^{(\tau)}_i}}\tag{14} $$

  这里我们沿用了Adam原文的假设，令动量衰减因子$\beta_1$为$\tau$的一个函数：$\beta_1^{(\tau)}$，进而在最后时获得最优上界。其中：$\beta_1^{(\tau)}\leq \cdots\leq \beta_1^{(1)}\leq 1$。所以：

$$ 2\gamma^{(\tau)}\frac{m_i^{(\tau)}}{\sqrt{\hat v_i^{(\tau)}}}(\theta^{(\tau)}_i-\theta^*_i)=(\theta^{(\tau)}_i-\theta_i^*)^2-(\theta^{(\tau+1)}_i-\theta_i^*)^2+\frac{\left(\gamma_i^{(\tau)} m_i^{(\tau)} \right)^2}{\hat v_i^{(\tau)}}\tag{15} $$

  我们代入$m^{(\tau)}_i=\beta_1\cdot m_i^{(\tau-1)}+(1-\beta_1^{(\tau)})\cdot g^{(\tau)}_i$则有：

$$ g^{(\tau)}_i \left(\theta^{(\tau)}_i-\theta^*_i \right)=\underbrace{\frac{\sqrt{\hat v_i^{(\tau)}}\left[(\theta^{(\tau)}_i-\theta_i^*)^2-(\theta^{(\tau+1)}_i-\theta_i^*)^2\right]}{2(1-\beta_1^{(\tau)})\gamma^{(\tau)}}}_{A}+\underbrace{\frac{\gamma^{(\tau)}\left(m_i^{(\tau)} \right)^2}{2(1-\beta_1^{(\tau)})\sqrt{\hat v_i^{(\tau)}}}}_{B}+\underbrace{(-1)\cdot\frac{\beta_1^{(\tau)}}{1-\beta_1^{(\tau)}}m_i^{(\tau-1)}\left(\theta^{(\tau)}_i-\theta^*_i \right)}_C\tag{16} $$

  我们需要对A、B、C进行放缩从而得到$g^{(\tau)}_i \left(\theta^{(\tau)}_i-\theta^*_i \right)$上界。首先我们来看C。假设参数与最优解之间间隔不会太远：$\left(\theta_i^{(\tau)}-\theta^*_i \right)^2\leq D^2_i$，则：

$$ \sum_{k=1}^\tau C\leq \sum_{k=1}^\tau \frac{\beta_1^{(k)}}{1-\beta_1^{(k)}}m_i^{(k-1)}\cdot D_i\tag{17} $$

  对于$m_i^{(k-1)}$，由于$m_i^{(0)}=0,m_i^{(\tau)}=\sum_{k=1}^\tau (1-\beta_1^{(k)})\left(\prod\limits_{s=k+1}^\tau\beta_1^{(s)} \right)\cdot g_i^{(k)}$，我们进一步假设梯度有界：$\left(g_i^{(\tau)} \right)^2\leq G_i^2$，所以：

$$ m_i^{(\tau)}\leq\sum_{k=1}^\tau (1-\beta_1^{(k)})\left(\prod\limits_{s=k+1}^\tau\beta_1^{(s)} \right)\cdot G_i=\left(1-\prod_{k=1}^\tau \beta_1^{(k)} \right)G_i\leq G_i\tag{18} $$

  所以：$\sum_{k=1}^\tau C\leq G_iD_i\sum_{k=1}^\tau\frac{\beta_1^{(k)}}{1-\beta_1^{(k)}}$。对于B，由于：$\frac{\left(m_{i}^{(\tau)}\right)^{2}}{\sqrt{\hat{v}_{i}^{(\tau)}}}=\sqrt{1-\beta_{2}^{(\tau)}}\frac{\left(m_{i}^{(\tau)}\right)^{2}}{\sqrt{v_{i}^{(\tau)}}}\leq\frac{\left(m_{i}^{(\tau)}\right)^{2}}{\sqrt{v_{i}^{(\tau)}}}$。又：

$$ \begin{align} m_i^{(\tau)}&=\sum_{k=1}^\tau (1-\beta_1^{(k)})\left(\prod\limits_{s=k+1}^\tau\beta_1^{(s)} \right)\cdot g_i^{(k)}\\ v_{i}^{(\tau)}&= \left(1-\beta_{2}\right)\sum_{k=1}^{\tau}\beta_{2}^{\tau-k}\left(g_i^{(k)} \right)^2 \end{align}\tag{19} $$

  我们做变换：

$$ \begin{align} \left(m_{i}^{(\tau)}\right)^{2}&=\left(\sum_{k=1}^{\tau}\frac{(1-\beta_1^{(k)})\left(\prod_{s=k+1}^\tau\beta_1^{(s)} \right)}{\sqrt{\left(1-\beta_{2}\right)\beta_{2}^{\tau-k}}}\cdot\sqrt{\left(1-\beta_{2}\right)\beta_{2}^{\tau-k}}\cdot g_i^{(k)} \right)^{2}\\ &\leq\sum_{k=1}^{\tau}\left(\frac{(1-\beta_1^{(k)})\left(\prod_{s=k+1}^\tau\beta_1^{(s)} \right)}{\sqrt{\left(1-\beta_{2}\right)\beta_{2}^{\tau-k}}}\right)^{2}\cdot\sum_{k=1}^{\tau}\left(\sqrt{\left(1-\beta_{2}\right)\beta_{2}^{\tau-k}}\cdot g_i^{(k)} \right)^{2} &Cauchy-Schwarz\\ &=\sum_{k=1}^{\tau}\frac{(1-\beta_1^{(k)})^2\left(\prod_{s=k+1}^\tau\beta_1^{(s)} \right)^2}{\left(1-\beta_{2}\right)\beta_{2}^{\tau-k}}\cdot \underbrace{\sum_{k=1}^{\tau}(1-\beta_2)\beta_2^{\tau-k} \left(g_i^{(\tau)} \right)^2}_{v_i^{(\tau)}}\tag{20} \end{align} $$

  所以对于B我们有：

$$ \sum_{k=1}^\tau B= \sum_{k=1}^\tau\frac{\gamma^{(k)}\left(m_{i}^{(k)}\right)^{2}}{2(1-\beta_{1}^{(k)} )\sqrt{\hat{v}_{i}^{(k)}}} \leq \sum_{k=1}^\tau \frac{\gamma^{(k)}}{2(1-\beta_1^{(k)})}\cdot \sum_{s=1}^{k}\frac{(1-\beta_1^{(s)})^2\left(\prod_{r=s+1}^k\beta_1^{(r)} \right)^2}{\left(1-\beta_{2}\right)\beta_{2}^{k-s}}\cdot G_i\tag{21} $$

  其中，我们由（18）可以类似得到$v_i^{(\tau)}\leq G_i$，从而可得（21）式。最后，我们来看最繁琐的A。首先我们代入$\gamma^{(\tau)}\geq \eta^{(\tau)}$：

$$ \begin{align} \sum_{k=1}^\tau A&=\sum_{k=1}^\tau \frac{\sqrt{\hat v_i^{(k)}}\left[(\theta^{(k)}_i-\theta_i^*)^2-(\theta^{(k+1)}_i-\theta_i^*)^2\right]}{2(1-\beta_1^{(k)})\gamma^{(k)}}\\ &\leq \sum_{k=1}^\tau \frac{\sqrt{\hat v_i^{(k)}}\left[(\theta^{(k)}_i-\theta_i^*)^2-(\theta^{(k+1)}_i-\theta_i^*)^2\right]}{2(1-\beta_1^{\color{Red}(1)})\eta^{(k)}}\\ &=\frac{\sqrt{\hat v_i^{(1)}}\left(\theta_i^{(1)}-\theta^*_i \right)^2}{2(1-\beta_1^{(1)})\eta^{(1)}}-\frac{\sqrt{\hat v_i^{(\tau)}}\left(\theta_i^{(\tau+1)}-\theta^*_i \right)^2}{2(1-\beta_1^{(1)})\eta^{(\tau)}}\\ &\quad+\sum_{k=2}^\tau \left(\theta_i^{(\tau)}-\theta_i^* \right)^2\cdot\left[\frac{\sqrt{\hat v_i^{(k)}}}{2(1-\beta_1^{(1)})\eta^{(k)}}-\frac{\sqrt{\hat v_i^{(k-1)}}}{2(1-\beta_1^{(1)})\eta^{(k-1)}} \right]\tag{22} \end{align} $$

  在$\Gamma_{k}=\frac{\sqrt{\hat{v}_{i}^{(k)}}}{\eta^{(k-1)}}-\frac{\sqrt{\hat{v}_{i}^{(k-1)}}}{\eta^{(k-1)}}\geq 0$恒成立时，有：

$$ \begin{align} \sum_{k=1}^\tau A&\leq \frac{D_i^2\sqrt{\hat v_i^{(\tau)}}}{2\eta^{(\tau)}(1-\beta_1^{(1)})}\leq \frac{D_i^2G_i}{2\eta^{(\tau)}(1-\beta_1^{(1)})}\tag{23} \end{align} $$

  综合（17）（18）（21）（23）：

$$ \begin{align} R(\tau)\leq& \frac{\sum_{i=1}^d D_i^2G_i}{2\eta^{(\tau)}(1-\beta_1^{(1)})}+\left(\sum_{i=1}^d G_iD_i\right)\sum_{k=1}^\tau \frac{\beta_1^{(k)}}{1-\beta_1^{k}}\\ &+\left(\sum_{i=1}^dG_i\right)\cdot\sum_{k=1}^\tau \frac{\gamma^{(k)}}{2(1-\beta_1^{(k)})}\cdot \sum_{s=1}^{k}\frac{(1-\beta_1^{(s)})^2\left(\prod_{r=s+1}^k\beta_1^{(r)} \right)^2}{\left(1-\beta_{2}\right)\beta_{2}^{k-s}}\tag{24} \end{align} $$

  对于第一项而言无法放缩了，对于第二项：

$$ \left(\sum_{i=1}^d G_iD_i\right)\sum_{k=1}^\tau \frac{\beta_1^{(k)}}{1-\beta_1^{k}}\leq \left(\sum_{i=1}^d G_iD_i\right)\frac{1}{1-\beta_1^{(1)}}\sum_{k=1}^\tau \beta_1^{(k)}\tag{25} $$

  对于第三项：

$$ \begin{align} &\left(\sum_{i=1}^dG_i\right) \cdot\sum_{k=1}^\tau \frac{\gamma^{(k)}}{2(1-\beta_1^{(k)})}\cdot \sum_{s=1}^{k}\frac{(1-\beta_1^{(s)})^2\left(\prod_{r=s+1}^k\beta_1^{(r)} \right)^2}{\left(1-\beta_{2}\right)\beta_{2}^{k-s}}\\ =&\left(\sum_{i=1}^dG_i\right) \cdot\sum_{k=1}^\tau \frac{\eta^{(k)}}{2(1-\beta_1^{(k)})(1-\prod_{s=1}^k \beta_1^{(s)})}\cdot \sum_{s=1}^k \Bigg[\left(\frac{1-\beta_1^{(s)}}{\sqrt{1-\beta_2}} \right)^2\cdot\prod_{r=s+1}^k \left(\frac{\beta_1^{(r)}}{\sqrt{\beta_2}} \right)^2 \Bigg]\\ \leq &\left(\sum_{i=1}^dG_i\right)\cdot \frac{\sum_{k=1}^\tau \eta^{(k)}}{2(1-\beta_1^{(1)})^2(1-\beta_2)(1-c)}\tag{26} \end{align} $$

其中最后一步我们设：$\frac{\beta_1^{(\tau)}}{\sqrt{\beta_2}}\leq \sqrt c<1$。所以：

$$ R(\tau)\leq \frac{C_1}{\eta^{(\tau)}}+C_2\cdot\sum_{k=1}^\tau \beta_1^{(\tau)}+C_3\cdot \sum_{k=1}^\tau \eta^{(k)}\tag{27} $$

令：$\eta^{(\tau)}=\eta_0/\sqrt \tau$，$\beta_1^{(\tau)}=\beta_1/\sqrt{\tau}$，则：$\frac{R(\tau)}{\tau}$有最优上界$O(\frac{1}{\sqrt \tau})$。综合上述所有假设与结论，我们有定理：

定理1

  设有：$\mathcal L\in\mathcal F_{0}^{1,1}$，$\mathbb{E}\left[\left(\theta_i^{(\tau)}-\theta^*_i \right)^2\right]\leq D^2_i,\mathbb{E}\left[\left(g_i \right)^2\right]\leq G^2_i$，$\frac{\beta_1^{(\tau)}}{\sqrt{\beta_2}}\leq \sqrt c<1$，且$\Gamma_{k}=\frac{\sqrt{\hat{v}_{i}^{(k)}}}{\eta^{(k-1)}}-\frac{\sqrt{\hat{v}_{i}^{(k-1)}}}{\eta^{(k-1)}}\geq 0$恒成立，则若：$\eta^{(\tau)}=\eta_0/\sqrt \tau$，$\beta_1^{(\tau)}=\beta_1/\sqrt{\tau}$，有：

$$ \frac{R(\tau)}{\tau}\leq O(\frac1{\sqrt \tau}) $$

  这个阶数与SGD和AdaGrad是一致的。在现有证明框架下，Adam算法的动量项必须衰减，并且至少以平方根的速率衰减；然而在工程实践中，动量项的超参数$\beta_1$一般设置为常数。Adam的相关实验结果参见：

  由定理1我们可以看出Adam证明中最大的Bug在哪了，即Adam无法保证：$\Gamma_{k}=\frac{\sqrt{\hat{v}_{i}^{(k)}}}{\eta^{(k-1)}}-\frac{\sqrt{\hat{v}_{i}^{(k-1)}}}{\eta^{(k-1)}}\geq 0$对于所有$k=1,\cdots,\tau$恒成立！那么在在线凸优化的框架下真的无法证明Adam的收敛性吗？Adam非凸情况又有什么表现呢？我们接下来介绍一些关于Adam收敛性的工作。

References

1. Adam: A method for stochastic optimization，Diederik P. Kingma, Jimmy Lei Ba，ICLR，2015
2. 【科研喂饭】深度学习算法收敛性证明之Adam