12.10. Adam
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 Colab 中打开 Notebook
在 SageMaker Studio Lab 中打开 Notebook

在本节之前，我们已经详细讨论了许多用于高效优化的技术。下面我们来回顾一下：

  • 我们看到，在解决优化问题时， 12.4节中的随机梯度下降比梯度下降更有效，例如，由于它对数据冗余的内在恢复能力。

  • 我们看到， 12.5节中的小批量随机梯度下降可以通过矢量化，在一个小批量中使用更大的观测集，从而显著提高效率。这是高效的多机、多GPU和整体并行处理的关键。

  • 12.6节中添加了一种机制，用于聚合过去梯度的历史以加速收敛。

  • 12.7节中的Adagrad通过按坐标缩放实现了计算上高效的预处理器。

  • 12.8节中的RMSProp将按坐标缩放与学习率调整解耦。

Adam （Kingma and Ba, 2014）算法将所有这些技术汇总到一个高效的学习算法中。不出所料，作为深度学习中更强大和有效的优化算法之一，它非常受欢迎。不过，它并非没有问题。特别是，（Reddi et al., 2019）表明，在某些情况下，由于方差控制不佳，Adam可能会发散。在后续工作中，Zaheer et al.（2018）提出了一个名为Yogi的Adam的热修复，解决了这些问题。稍后我们将详细介绍。现在，让我们先回顾一下Adam算法。

12.10.1. 算法

Adam算法的关键组成部分之一，是它使用指数加权移动平均值（也称为泄漏平均值）来获得动量和梯度二阶矩的估计。也就是说，它使用状态变量

(12.10.1)\[\begin{split}\begin{aligned} \mathbf{v}_t & \leftarrow \beta_1 \mathbf{v}_{t-1} + (1 - \beta_1) \mathbf{g}_t, \\ \mathbf{s}_t & \leftarrow \beta_2 \mathbf{s}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2. \end{aligned}\end{split}\]

这里\(\beta_1\)和\(\beta_2\)是非负的权重参数。它们的通常选择为\(\beta_1 = 0.9\)和\(\beta_2 = 0.999\)。也就是说，方差估计值的移动速度比动量项的移动速度要慢得多。请注意，如果初始化\(\mathbf{v}_0 = \mathbf{s}_0 = 0\)，我们有相当大的初始偏差朝向较小的值。这可以通过使用\(\sum_{i=0}^{t-1} \beta^i = \frac{1 - \beta^t}{1 - \beta}\)来重新规范化项来解决。相应地，规范化的状态变量由下式给出

(12.10.2)\[\hat{\mathbf{v}}_t = \frac{\mathbf{v}_t}{1 - \beta_1^t} \textrm{ and } \hat{\mathbf{s}}_t = \frac{\mathbf{s}_t}{1 - \beta_2^t}.\]

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

(12.10.3)\[\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}\)，以便在数值稳定性和保真度之间取得良好的平衡。

现在我们已经准备好计算更新所需的所有部分。这有点平淡无奇，我们有一个简单的更新形式

(12.10.4)\[\mathbf{x}_t \leftarrow \mathbf{x}_{t-1} - \mathbf{g}_t'.\]

回顾Adam的设计，它的灵感很明显。动量和尺度在状态变量中清晰可见。它们相当独特的定义迫使我们对项进行去偏（这可以通过稍微不同的初始化和更新条件来修正）。其次，考虑到RMSProp，这两个项的组合非常直接。最后，显式学习率\(\eta\)允许我们控制步长以解决收敛问题。

12.10.2. 实现

从零开始实现Adam并不令人生畏。为方便起见，我们将时间步计数器\(t\)存储在hyperparams字典中。除此之外，一切都很简单。

pytorchmxnettensorflow
%matplotlib inline
import torch
from d2l import torch as d2l


def init_adam_states(feature_dim):
    v_w, v_b = torch.zeros((feature_dim, 1)), torch.zeros(1)
    s_w, s_b = torch.zeros((feature_dim, 1)), torch.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):
        with torch.no_grad():
            v[:] = beta1 * v + (1 - beta1) * p.grad
            s[:] = beta2 * s + (1 - beta2) * 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)
                                                       + eps)
        p.grad.data.zero_()
    hyperparams['t'] += 1

%matplotlib inline
from mxnet import np, npx
from d2l import mxnet as d2l

npx.set_np()

def init_adam_states(feature_dim):
    v_w, v_b = np.zeros((feature_dim, 1)), np.zeros(1)
    s_w, s_b = np.zeros((feature_dim, 1)), np.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

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


def init_adam_states(feature_dim):
    v_w = tf.Variable(tf.zeros((feature_dim, 1)))
    v_b = tf.Variable(tf.zeros(1))
    s_w = tf.Variable(tf.zeros((feature_dim, 1)))
    s_b = tf.Variable(tf.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)

我们准备好使用Adam来训练模型。我们使用学习率为\(\eta = 0.01\)。

pytorchmxnettensorflow
data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
d2l.train_ch11(adam, init_adam_states(feature_dim),
               {'lr': 0.01, 't': 1}, data_iter, feature_dim);

loss: 0.243, 0.193 sec/epoch
../_images/output_adam_f5876e_15_1.svg
data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
d2l.train_ch11(adam, init_adam_states(feature_dim),
               {'lr': 0.01, 't': 1}, data_iter, feature_dim);

loss: 0.243, 1.878 sec/epoch
../_images/output_adam_f5876e_18_1.svg
data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
d2l.train_ch11(adam, init_adam_states(feature_dim),
               {'lr': 0.01, 't': 1}, data_iter, feature_dim);

loss: 0.242, 1.616 sec/epoch
../_images/output_adam_f5876e_21_1.svg

更简洁的实现是直接的，因为adam是Gluon trainer优化库中提供的算法之一。因此，我们只需要为Gluon中的实现传递配置参数。

pytorchmxnettensorflow
trainer = torch.optim.Adam
d2l.train_concise_ch11(trainer, {'lr': 0.01}, data_iter)

loss: 0.243, 0.152 sec/epoch
../_images/output_adam_f5876e_27_1.svg
d2l.train_concise_ch11('adam', {'learning_rate': 0.01}, data_iter)

loss: 0.244, 0.816 sec/epoch
../_images/output_adam_f5876e_30_1.svg
trainer = tf.keras.optimizers.Adam
d2l.train_concise_ch11(trainer, {'learning_rate': 0.01}, data_iter)

loss: 0.246, 1.637 sec/epoch
../_images/output_adam_f5876e_33_1.svg

12.10.3. Yogi

Adam的问题之一是，当\(\mathbf{s}_t\)中的二阶矩估计爆炸时，即使在凸设置中它也可能无法收敛。作为一种修复，Zaheer et al.（2018）为\(\mathbf{s}_t\)提出了一个精炼的更新（和初始化）。为了理解发生了什么，让我们将Adam的更新重写如下：

(12.10.5)\[\mathbf{s}_t \leftarrow \mathbf{s}_{t-1} + (1 - \beta_2) \left(\mathbf{g}_t^2 - \mathbf{s}_{t-1}\right).\]

每当\(\mathbf{g}_t^2\)具有高方差或更新稀疏时，\(\mathbf{s}_t\)可能会过快地忘记过去的值。一个可能的解决方法是将\(\mathbf{g}_t^2 - \mathbf{s}_{t-1}\)替换为\(\mathbf{g}_t^2 \odot \mathop{\textrm{sgn}}(\mathbf{g}_t^2 - \mathbf{s}_{t-1})\)。现在更新的大小不再取决于偏差的大小。这就产生了Yogi的更新：

(12.10.6)\[\mathbf{s}_t \leftarrow \mathbf{s}_{t-1} + (1 - \beta_2) \mathbf{g}_t^2 \odot \mathop{\textrm{sgn}}(\mathbf{g}_t^2 - \mathbf{s}_{t-1}).\]

作者还建议在较大的初始批次上初始化动量，而不仅仅是初始的逐点估计。我们省略了细节，因为它们对讨论不重要，而且即使没有这些，收敛性也仍然很好。

pytorchmxnettensorflow
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)
                                                       + eps)
        p.grad.data.zero_()
    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);

loss: 0.243, 0.165 sec/epoch
../_images/output_adam_f5876e_39_1.svg
def yogi(params, states, hyperparams):
    beta1, beta2, eps = 0.9, 0.999, 1e-3
    for p, (v, s) in zip(params, states):
        v[:] = beta1 * v + (1 - beta1) * p.grad
        s[:] = s + (1 - beta2) * np.sign(
            np.square(p.grad) - s) * 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

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);

loss: 0.248, 1.783 sec/epoch
../_images/output_adam_f5876e_42_1.svg
def yogi(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(s + (1 - beta2) * tf.math.sign(
                   tf.math.square(grad) - s) * 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)
    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);

loss: 0.244, 1.680 sec/epoch
../_images/output_adam_f5876e_45_1.svg

12.10.4. 小结

  • Adam将许多优化算法的特性结合到一个相当鲁棒的更新规则中。

  • Adam在RMSProp的基础上创建，也对小批量随机梯度使用指数加权移动平均。

  • Adam使用偏差校正来调整在估计动量和二阶矩时的缓慢启动。

  • 对于梯度具有显著方差的情况，我们可能会遇到收敛问题。可以通过使用更大的小批量或切换到\(\mathbf{s}_t\)的改进估计来修正它们。Yogi提供了这样一种替代方案。

12.10.5. 练习

  1. 调整学习率并观察和分析实验结果。

  2. 你能重写动量和二阶矩更新，使其不需要偏差校正吗？

  3. 为什么在收敛时需要降低学习率\(\eta\)？

  4. 尝试构造一个Adam发散而Yogi收敛的案例？

pytorchmxnettensorflow

讨论

讨论

讨论