Adam算法

:label:sec_adam

本章我们已经学习了许多有效优化的技术。 在本节讨论之前，我们先详细回顾一下这些技术：

  • 在 :numref:sec_sgd中，我们学习了：随机梯度下降在解决优化问题时比梯度下降更有效。
  • 在 :numref:sec_minibatch_sgd中，我们学习了：在一个小批量中使用更大的观测值集，可以通过向量化提供额外效率。这是高效的多机、多GPU和整体并行处理的关键。
  • 在 :numref:sec_momentum中我们添加了一种机制，用于汇总过去梯度的历史以加速收敛。
  • 在 :numref:sec_adagrad中，我们通过对每个坐标缩放来实现高效计算的预处理器。
  • 在 :numref:sec_rmsprop中，我们通过学习率的调整来分离每个坐标的缩放。

Adam算法 :cite:Kingma.Ba.2014将所有这些技术汇总到一个高效的学习算法中。 不出预料，作为深度学习中使用的更强大和有效的优化算法之一，它非常受欢迎。 但是它并非没有问题，尤其是 :cite:Reddi.Kale.Kumar.2019表明，有时Adam算法可能由于方差控制不良而发散。 在完善工作中， :cite:Zaheer.Reddi.Sachan.ea.2018给Adam算法提供了一个称为Yogi的热补丁来解决这些问题。 下面我们了解一下Adam算法。

算法

Adam算法的关键组成部分之一是：它使用指数加权移动平均值来估算梯度的动量和二次矩，即它使用状态变量

$$\begin{aligned} \mathbf{v}t & \leftarrow \beta_1 \mathbf{v}{t-1} + (1 - \beta1) \mathbf{g}_t, \ \mathbf{s}_t & \leftarrow \beta_2 \mathbf{s}{t-1} + (1 - \beta_2) \mathbf{g}_t^2. \end{aligned}$$

这里$\beta1$和$\beta_2$是非负加权参数。 常将它们设置为$\beta_1 = 0.9$和$\beta_2 = 0.999$。 也就是说，方差估计的移动远远慢于动量估计的移动。 注意，如果我们初始化$\mathbf{v}_0 = \mathbf{s}_0 = 0$，就会获得一个相当大的初始偏差。 我们可以通过使用$\sum{i=0}^t \beta^i = \frac{1 - \beta^t}{1 - \beta}$来解决这个问题。 相应地，标准化状态变量由下式获得

\hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_1^t} \text{ and } \hat{\mathbf{s}}_t = \frac{\mathbf{s}_t}{1 - \beta_2^t}.

有了正确的估计，我们现在可以写出更新方程。 首先，我们以非常类似于RMSProp算法的方式重新缩放梯度以获得

\mathbf{g}_t’ = \frac{\eta \hat{\mathbf{v}}_t}{\sqrt{\hat{\mathbf{s}}_t} + \epsilon}.

与RMSProp不同，我们的更新使用动量$\hat{\mathbf{v}}_t$而不是梯度本身。 此外，由于使用$\frac{1}{\sqrt{\hat{\mathbf{s}}_t} + \epsilon}$而不是$\frac{1}{\sqrt{\hat{\mathbf{s}}_t + \epsilon}}$进行缩放，两者会略有差异。 前者在实践中效果略好一些，因此与RMSProp算法有所区分。 通常，我们选择$\epsilon = 10^{-6}$，这是为了在数值稳定性和逼真度之间取得良好的平衡。

最后，我们简单更新：

\mathbf{x}t \leftarrow \mathbf{x}{t-1} - \mathbf{g}_t’.

回顾Adam算法，它的设计灵感很清楚： 首先，动量和规模在状态变量中清晰可见， 它们相当独特的定义使我们移除偏项（这可以通过稍微不同的初始化和更新条件来修正）。 其次，RMSProp算法中两项的组合都非常简单。 最后，明确的学习率$\eta$使我们能够控制步长来解决收敛问题。

实现

从头开始实现Adam算法并不难。 为方便起见，我们将时间步$t$存储在hyperparams字典中。 除此之外，一切都很简单。

```{.python .input} %matplotlib inline from d2l import mxnet as d2l from mxnet import np, npx npx.set_np()

def init_adam_states(feature_dim): v_w, v_b = d2l.zeros((feature_dim, 1)), d2l.zeros(1) s_w, s_b = d2l.zeros((feature_dim, 1)), d2l.zeros(1) return ((v_w, s_w), (v_b, s_b))

def adam(params, states, hyperparams): beta1, beta2, eps = 0.9, 0.999, 1e-6 for p, (v, s) in zip(params, states): v[:] = beta1 v + (1 - beta1) p.grad s[:] = beta2 s + (1 - beta2) np.square(p.grad) v_bias_corr = v / (1 - beta1 hyperparams[‘t’]) s_bias_corr = s / (1 - beta2 hyperparams[‘t’]) p[:] -= hyperparams[‘lr’] * v_bias_corr / (np.sqrt(s_bias_corr) + eps) hyperparams[‘t’] += 1

  2. ```{.python .input}
  3. #@tab pytorch
  4. %matplotlib inline
  5. from d2l import torch as d2l
  6. import torch
  8. def init_adam_states(feature_dim):
  9.     v_w, v_b = d2l.zeros((feature_dim, 1)), d2l.zeros(1)
  10.     s_w, s_b = d2l.zeros((feature_dim, 1)), d2l.zeros(1)
  11.     return ((v_w, s_w), (v_b, s_b))
  13. def adam(params, states, hyperparams):
  14.     beta1, beta2, eps = 0.9, 0.999, 1e-6
  15.     for p, (v, s) in zip(params, states):
  16.         with torch.no_grad():
  17.             v[:] = beta1 * v + (1 - beta1) * p.grad
  18.             s[:] = beta2 * s + (1 - beta2) * torch.square(p.grad)
  19.             v_bias_corr = v / (1 - beta1 ** hyperparams['t'])
  20.             s_bias_corr = s / (1 - beta2 ** hyperparams['t'])
  21.             p[:] -= hyperparams['lr'] * v_bias_corr / (torch.sqrt(s_bias_corr)
  22.                                                        + eps)
  23.         p.grad.data.zero_()
  24.     hyperparams['t'] += 1

```{.python .input}

@tab tensorflow

%matplotlib inline from d2l import tensorflow as d2l import tensorflow as tf

def init_adam_states(feature_dim): v_w = tf.Variable(d2l.zeros((feature_dim, 1))) v_b = tf.Variable(d2l.zeros(1)) s_w = tf.Variable(d2l.zeros((feature_dim, 1))) s_b = tf.Variable(d2l.zeros(1)) return ((v_w, s_w), (v_b, s_b))

def adam(params, grads, states, hyperparams): beta1, beta2, eps = 0.9, 0.999, 1e-6 for p, (v, s), grad in zip(params, states, grads): v[:].assign(beta1 v + (1 - beta1) grad) s[:].assign(beta2 s + (1 - beta2) tf.math.square(grad)) v_bias_corr = v / (1 - beta1 hyperparams[‘t’]) s_bias_corr = s / (1 - beta2 hyperparams[‘t’]) p[:].assign(p - hyperparams[‘lr’] * v_bias_corr / tf.math.sqrt(s_bias_corr) + eps)

  2. 现在，我们用以上Adam算法来训练模型，这里我们使用$\eta = 0.01$的学习率。
  4. ```{.python .input}
  5. #@tab all
  6. data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
  7. d2l.train_ch11(adam, init_adam_states(feature_dim),
  8.                {'lr': 0.01, 't': 1}, data_iter, feature_dim);

此外，我们可以用深度学习框架自带算法应用Adam算法，这里我们只需要传递配置参数。

```{.python .input} d2l.train_concise_ch11(‘adam’, {‘learning_rate’: 0.01}, data_iter)

  2. ```{.python .input}
  3. #@tab pytorch
  4. trainer = torch.optim.Adam
  5. d2l.train_concise_ch11(trainer, {'lr': 0.01}, data_iter)

```{.python .input}

@tab tensorflow

trainer = tf.keras.optimizers.Adam d2l.train_concise_ch11(trainer, {‘learning_rate’: 0.01}, data_iter)

  2. ## Yogi
  4. Adam算法也存在一些问题：
  5. 即使在凸环境下，当$\mathbf{s}_t$的二次矩估计值爆炸时，它可能无法收敛。
  6.  :cite:`Zaheer.Reddi.Sachan.ea.2018`为$\mathbf{s}_t$提出了的改进更新和参数初始化。
  7. 论文中建议我们重写Adam算法更新如下：
  9. $$\mathbf{s}_t \leftarrow \mathbf{s}_{t-1} + (1 - \beta_2) \left(\mathbf{g}_t^2 - \mathbf{s}_{t-1}\right).$$
  11. 每当$\mathbf{g}_t^2$具有值很大的变量或更新很稀疏时，$\mathbf{s}_t$可能会太快地“忘记”过去的值。
  12. 一个有效的解决方法是将$\mathbf{g}_t^2 - \mathbf{s}_{t-1}$替换为$\mathbf{g}_t^2 \odot \mathop{\mathrm{sgn}}(\mathbf{g}_t^2 - \mathbf{s}_{t-1})$。
  13. 这就是Yogi更新，现在更新的规模不再取决于偏差的量。
  15. $$\mathbf{s}_t \leftarrow \mathbf{s}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2 \odot \mathop{\mathrm{sgn}}(\mathbf{g}_t^2 - \mathbf{s}_{t-1}).$$
  17. 论文中，作者还进一步建议用更大的初始批量来初始化动量，而不仅仅是初始的逐点估计。
  19. ```{.python .input}
  20. def yogi(params, states, hyperparams):
  21.     beta1, beta2, eps = 0.9, 0.999, 1e-3
  22.     for p, (v, s) in zip(params, states):
  23.         v[:] = beta1 * v + (1 - beta1) * p.grad
  24.         s[:] = s + (1 - beta2) * np.sign(
  25.             np.square(p.grad) - s) * np.square(p.grad)
  26.         v_bias_corr = v / (1 - beta1 ** hyperparams['t'])
  27.         s_bias_corr = s / (1 - beta2 ** hyperparams['t'])
  28.         p[:] -= hyperparams['lr'] * v_bias_corr / (np.sqrt(s_bias_corr) + eps)
  29.     hyperparams['t'] += 1
  31. data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
  32. d2l.train_ch11(yogi, init_adam_states(feature_dim),
  33.                {'lr': 0.01, 't': 1}, data_iter, feature_dim);

```{.python .input}

@tab pytorch

def yogi(params, states, hyperparams): beta1, beta2, eps = 0.9, 0.999, 1e-3 for p, (v, s) in zip(params, states): with torch.no_grad(): v[:] = beta1 v + (1 - beta1) p.grad s[:] = s + (1 - beta2) torch.sign( torch.square(p.grad) - s) torch.square(p.grad) v_bias_corr = v / (1 - beta1 hyperparams[‘t’]) s_bias_corr = s / (1 - beta2 hyperparams[‘t’]) p[:] -= hyperparams[‘lr’] * v_bias_corr / (torch.sqrt(s_bias_corr)

  1.                                                    + eps)
  2.     p.grad.data.zero_()
  3. hyperparams['t'] += 1

data_iter, feature_dim = d2l.get_data_ch11(batch_size=10) d2l.train_ch11(yogi, init_adam_states(feature_dim), {‘lr’: 0.01, ‘t’: 1}, data_iter, feature_dim);

  2. ```{.python .input}
  3. #@tab tensorflow
  4. def yogi(params, grads, states, hyperparams):
  5.     beta1, beta2, eps = 0.9, 0.999, 1e-6
  6.     for p, (v, s), grad in zip(params, states, grads):
  7.         v[:].assign(beta1 * v  + (1 - beta1) * grad)
  8.         s[:].assign(s + (1 - beta2) * tf.math.sign(
  9.                    tf.math.square(grad) - s) * tf.math.square(grad))
  10.         v_bias_corr = v / (1 - beta1 ** hyperparams['t'])
  11.         s_bias_corr = s / (1 - beta2 ** hyperparams['t'])
  12.         p[:].assign(p - hyperparams['lr'] * v_bias_corr
  13.                     / tf.math.sqrt(s_bias_corr) + eps)
  14.     hyperparams['t'] += 1
  16. data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
  17. d2l.train_ch11(yogi, init_adam_states(feature_dim),
  18.                {'lr': 0.01, 't': 1}, data_iter, feature_dim);

小结

  • Adam算法将许多优化算法的功能结合到了相当强大的更新规则中。
  • Adam算法在RMSProp算法基础上创建的，还在小批量的随机梯度上使用EWMA。
  • 在估计动量和二次矩时，Adam算法使用偏差校正来调整缓慢的启动速度。
  • 对于具有显著差异的梯度，我们可能会遇到收敛性问题。我们可以通过使用更大的小批量或者切换到改进的估计值$\mathbf{s}_t$来修正它们。Yogi提供了这样的替代方案。

练习

  1. 调节学习率，观察并分析实验结果。
  2. 你能重写动量和二次矩更新，从而使其不需要偏差校正吗？
  3. 当我们收敛时，为什么你需要降低学习率$\eta$？
  4. 尝试构造一个使用Adam算法会发散而Yogi会收敛的例子。

:begin_tab:mxnet Discussions :end_tab:

:begin_tab:pytorch Discussions :end_tab:

:begin_tab:tensorflow Discussions :end_tab: