11.10 Алгоритм Adam

В этой главе мы уже изучили множество эффективных методов оптимизации. Прежде чем переходить к новому, вспомним их подробно:

В разделе 11.4 мы узнали, что стохастический градиентный спуск эффективнее градиентного спуска при решении задач оптимизации.
В разделе 11.5 мы узнали, что использование большего набора наблюдений в минибатче даёт дополнительную эффективность за счёт векторизации. Это ключ к эффективному параллелизму на нескольких машинах, нескольких GPU и в целом.
В разделе 11.6 мы добавили механизм, который суммирует историю прошлых градиентов, чтобы ускорить сходимость.
В разделе 11.7 мы использовали предобусловливатель, который эффективно вычисляется за счёт покоординатного масштабирования.
В разделе 11.8 мы отделили масштабирование каждой координаты от управления скоростью обучения.

Алгоритм Adam (Kingma and Ba, 2014) объединяет все эти приёмы в один эффективный алгоритм обучения. Неудивительно, что он стал очень популярным: это один из самых мощных и эффективных оптимизаторов, используемых в глубоком обучении. Однако у него есть и проблемы. В частности, Редди и соавторы (Reddi et al., 2019) показали, что Adam иногда может расходиться из-за плохого контроля дисперсии. В качестве доработки Захир и соавторы (Zaheer et al., 2018) предложили «быстрое исправление» для Adam под названием Yogi. Ниже мы рассмотрим алгоритм Adam.

11.10.1 Алгоритм

Одна из ключевых составляющих алгоритма Adam — использование экспоненциально взвешенных скользящих средних для оценки первого момента (импульса) и второго момента градиента. То есть используются переменные состояния

$v_{t} s_{t} \leftarrow β_{1} v_{t - 1} + (1 - β_{1}) g_{t}, \leftarrow β_{2} s_{t - 1} + (1 - β_{2}) g_{t}^{2} .$

Здесь $β_{1}$ и $β_{2}$ — неотрицательные весовые параметры. Обычно их выбирают равными $β_{1} = 0.9$ и $β_{2} = 0.999$ . То есть оценка дисперсии смещается гораздо медленнее, чем оценка импульса. Заметим, что если инициализировать $v_{0} = s_{0} = 0$ , возникает заметное начальное смещение. Его можно устранить, воспользовавшись равенством $\sum_{i = 0}^{t} β^{i} = \frac{1 - β ^{t}}{1 - β}$ . Соответственно, нормированные переменные состояния получаются по формулам

$\hat{v}_{t} = \frac{v _{t}}{1 - β _{1}^{t}} и \hat{s}_{t} = \frac{s _{t}}{1 - β _{2}^{t}} .$

Теперь, имея корректные оценки, можно записать уравнения обновления. Сначала мы перемасштабируем градиент очень похожим на RMSProp образом:

$g_{t}^{'} = \frac{η v ^ _{t}}{s ^ _{t} + ϵ} .$

В отличие от RMSProp, наше обновление использует импульс $\hat{v}_{t}$ , а не сам градиент. Кроме того, из-за масштабирования с помощью $\frac{1}{s ^ _{t} + ϵ}$ , а не $\frac{1}{s ^ _{t} + ϵ}$ , результаты немного различаются. На практике первый вариант работает чуть лучше, поэтому он и отличает алгоритм от RMSProp. Обычно мы выбираем $ϵ = 1 0^{- 6}$ , что обеспечивает хороший баланс между численной устойчивостью и точностью.

Наконец, выполняем простое обновление:

$x_{t} \leftarrow x_{t - 1} - g_{t}^{'} .$

Если взглянуть на алгоритм Adam, его замысел становится очевидным. Во-первых, импульс и масштаб явно выражены в переменных состояния. Их довольно своеобразное определение позволяет убрать смещение (это можно исправить с помощью немного другой инициализации и условий обновления). Во-вторых, комбинация двух слагаемых из RMSProp очень проста. Наконец, явная скорость обучения $η$ позволяет управлять размером шага для решения проблем сходимости.

11.10.2 Реализация

Реализовать алгоритм Adam с нуля нетрудно. Для удобства мы храним номер временного шага $t$ в словаре hyperparams. Во всём остальном всё просто.

MXNet:

%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

PyTorch:

%matplotlib inline
from d2l import torch as d2l
import torch

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

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)

PaddlePaddle:

%matplotlib inline
from d2l import paddle as d2l
import warnings
warnings.filterwarnings("ignore")
import paddle

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
    a = []
    for p, (v, s) in zip(params, states):
        with paddle.no_grad():
            v[:] = beta1 * v + (1 - beta1) * p.grad
            s[:] = beta2 * s + (1 - beta2) * paddle.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 / (paddle.sqrt(s_bias_corr)
                                                       + eps)
        p.grad.zero_()
        a.append(p)
    hyperparams['t'] += 1
    return a

Теперь обучим модель с помощью описанного алгоритма Adam, используя скорость обучения $η = 0.01$ .

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

Кроме того, Adam можно использовать из фреймворка глубокого обучения: достаточно передать параметры конфигурации.

MXNet:

d2l.train_concise_ch11('adam', {'learning_rate': 0.01}, data_iter)

PyTorch:

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

TensorFlow:

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

PaddlePaddle:

trainer = paddle.optimizer.Adam
d2l.train_concise_ch11(trainer, {'learning_rate': 0.01}, data_iter)

11.10.3 Yogi

У Adam есть и свои проблемы: даже в выпуклой постановке он может не сходиться, когда оценка второго момента $s_{t}$ «взрывается». Захир и соавторы (Zaheer et al., 2018) предложили улучшенное обновление и инициализацию для $s_{t}$ . В статье предлагается переписать обновление Adam следующим образом:

$s_{t} \leftarrow s_{t - 1} + (1 - β_{2}) (g_{t}^{2} - s_{t - 1}) .$

Всякий раз, когда $g_{t}^{2}$ принимает большие значения или обновления очень разрежены, $s_{t}$ может слишком быстро «забыть» прошлые значения. Эффективное решение — заменить $g_{t}^{2} - s_{t - 1}$ на $g_{t}^{2} ⊙ sgn (g_{t}^{2} - s_{t - 1})$ . Это и есть обновление Yogi: величина шага теперь не зависит от размера отклонения.

$s_{t} \leftarrow s_{t - 1} + (1 - β_{2}) g_{t}^{2} ⊙ sgn (g_{t}^{2} - s_{t - 1}) .$

В статье авторы также предлагают инициализировать импульс не поточечной оценкой, а большим начальным минибатчем.

MXNet:

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

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

TensorFlow:

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

PaddlePaddle:

def yogi(params, states, hyperparams):
    beta1, beta2, eps = 0.9, 0.999, 1e-3
    a=[]
    for p, (v, s) in zip(params, states):
        with paddle.no_grad():
            v[:] = beta1 * v + (1 - beta1) * p.grad
            s[:] = s + (1 - beta2) * paddle.sign(
                paddle.square(p.grad) - s) * paddle.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 / (paddle.sqrt(s_bias_corr)
                                                       + eps)
        p.grad.zero_()
        a.append(p)
    hyperparams['t'] += 1
    return a

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

Резюме

Adam объединяет множество приёмов оптимизации в довольно мощное правило обновления.
Adam создан на основе RMSProp и дополнительно использует экспоненциально взвешенные скользящие средние (EWMA) для минибатчевых стохастических градиентов.
При оценке импульса и второго момента Adam использует коррекцию смещения, чтобы компенсировать медленный старт.
Для градиентов с заметно различающейся дисперсией могут возникать проблемы со сходимостью. Их можно смягчить, используя больший минибатч или переключившись на улучшенную оценку $s_{t}$ . Альтернативу предлагает Yogi.

Упражнения

Подберите скорость обучения и понаблюдайте и проанализируйте результаты экспериментов.
Попробуйте переписать обновления импульса и второго момента так, чтобы коррекция смещения не требовалась.
Почему при сходимости нужно уменьшать скорость обучения $η$ ?
Постройте пример, в котором Adam расходится, а Yogi сходится.

MXNet: Обсуждение

PyTorch: Обсуждение

TensorFlow: Обсуждение

PaddlePaddle: Обсуждение

Литература

Kingma, D. P., & Ba, J. (2014). Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.
Reddi, S. J., Kale, S., & Kumar, S. (2019). On the convergence of Adam and beyond. arXiv preprint arXiv:1904.09237.
Zaheer, M., Reddi, S., Sachan, D., Kale, S., & Kumar, S. (2018). Adaptive methods for nonconvex optimization. Advances in Neural Information Processing Systems, 9793–9803.

Keyboard shortcuts

Погружение в глубокое обучение

11.10 Алгоритм Adam

11.10.1 Алгоритм

11.10.2 Реализация

11.10.3 Yogi

Резюме

Упражнения

Литература