11.6 Метод импульса

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

11.6.1 Основы

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

11.6.1.1 Экспоненциально затухающее среднее

В предыдущем разделе мы рассматривали минибатчевый стохастический градиентный спуск как способ ускорить вычисления. У него есть и приятный побочный эффект: усреднение градиентов уменьшает дисперсию. Минибатчевый стохастический градиентный спуск можно вычислить так:

$${\mathbf{g}}_{t,t-1}={\partial }_{\mathbf{w}}\frac{1}{|{\mathcal{B}}_t|}\sum _{i\in {\mathcal{B}}_t}f({\mathbf{x}}_i,{\mathbf{w}}_{t-1})=\frac{1}{|{\mathcal{B}}_t|}\sum _{i\in {\mathcal{B}}_t}{\mathbf{h}}_{i,t-1}.$$

Чтобы сохранить обозначения простыми, здесь мы используем ${\mathbf{h}}_{i,t-1}={\partial }_{\mathbf{w}}f({\mathbf{x}}_i,{\mathbf{w}}_{t-1})$ как стохастический градиент для примера $i$ при весах ${\mathbf{w}}_{t-1}$, обновлённых на шаге $t-1$. Было бы хорошо, если бы мы могли воспользоваться эффектом снижения дисперсии даже сверх усреднения градиентов по минибатчу. Один из вариантов решения этой задачи — заменить вычисление градиента экспоненциально затухающим средним (leaky average):

$${\mathbf{v}}_t=\beta {\mathbf{v}}_{t-1}+{\mathbf{g}}_{t,t-1},$$

где $\beta \in (0,1)$. Это эффективно заменяет мгновенный градиент усреднённым значением нескольких «прошлых» градиентов. $\mathbf{v}$ называют импульсом (momentum), он накапливает прошлые градиенты. Чтобы объяснить это подробнее, рекурсивно раскроем ${\mathbf{v}}_t$:

$$\begin{array}{r}{\mathbf{v}}_t={\beta }^2{\mathbf{v}}_{t-2}+\beta {\mathbf{g}}_{t-1,t-2}+{\mathbf{g}}_{t,t-1}=\dots ,=\sum _{\tau =0}^{t-1}{\beta }^{\tau }{\mathbf{g}}_{t-\tau ,t-\tau -1}.\end{array}$$

Большое $\beta$ соответствует длинному среднему, тогда как малое $\beta$ лишь немного корректирует обычный градиентный спуск. Новый градиент больше не указывает направление наибольшего спуска для конкретного примера, а указывает направление взвешенного среднего прошлых градиентов. Это позволяет получить большую часть выгоды от усреднения по батчу, не платя за фактическое вычисление градиента для каждого из них.

Описанные рассуждения лежат в основе «ускоренных» методов градиента, например, градиентного спуска с импульсом. Для задач оптимизации с плохой обусловленностью (например, когда одни направления продвигаются гораздо медленнее других, как в узкой долине) ускоренные методы дают дополнительный выигрыш в эффективности. Кроме того, они позволяют усреднять последующие градиенты, чтобы получить более устойчивое направление спуска. Надо признать, что даже для задач выпуклой оптимизации без шума ускорение — одна из ключевых причин, по которым импульс так хорошо работает.

Как и следовало ожидать, благодаря своей эффективности импульс стал предметом активных исследований в глубоком обучении и последующей оптимизации. Например, см. статью (автор — Goh, 2017), в которой приведён подробный анализ и интерактивные анимации. Импульс был предложен в работе (Polyak, 1964). Nesterov (2018) подробно рассматривает его теорию в контексте выпуклой оптимизации. Давно известно, что импульс полезен и в глубоком обучении. Подробности с примерами см. в обсуждении у (Sutskever et al., 2013).

11.6.1.2 Плохо обусловленные задачи

Чтобы лучше понять геометрию метода импульса, вспомним градиентный спуск, хотя его целевая функция явно менее приятна. Напомним, что в разделе 11.3 мы использовали $f(\mathbf{x})=x_1^2+2x_2^2$ — умеренно вытянутый эллипсоид. Мы сделаем эту функцию ещё более вытянутой, растянув её вдоль направления $x_1$:

$$f(\mathbf{x})=0.1x_1^2+2x_2^2.$$

Как и прежде, $f$ имеет минимум в точке $(0,0)$. Эта функция очень плоская в направлении $x_1$. Посмотрим, что происходит при градиентном спуске для этой новой функции.

%matplotlib inline
from d2l import mxnet as d2l
from mxnet import np, npx
npx.set_np()

eta = 0.4
def f_2d(x1, x2):
    return 0.1 * x1 ** 2 + 2 * x2 ** 2
def gd_2d(x1, x2, s1, s2):
    return (x1 - eta * 0.2 * x1, x2 - eta * 4 * x2, 0, 0)

d2l.show_trace_2d(f_2d, d2l.train_2d(gd_2d))

%matplotlib inline
from d2l import torch as d2l
import torch

eta = 0.4
def f_2d(x1, x2):
    return 0.1 * x1 ** 2 + 2 * x2 ** 2
def gd_2d(x1, x2, s1, s2):
    return (x1 - eta * 0.2 * x1, x2 - eta * 4 * x2, 0, 0)

d2l.show_trace_2d(f_2d, d2l.train_2d(gd_2d))

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

eta = 0.4
def f_2d(x1, x2):
    return 0.1 * x1 ** 2 + 2 * x2 ** 2
def gd_2d(x1, x2, s1, s2):
    return (x1 - eta * 0.2 * x1, x2 - eta * 4 * x2, 0, 0)

d2l.show_trace_2d(f_2d, d2l.train_2d(gd_2d))

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

eta = 0.4
def f_2d(x1, x2):
    return 0.1 * x1 ** 2 + 2 * x2 ** 2
def gd_2d(x1, x2, s1, s2):
    return (x1 - eta * 0.2 * x1, x2 - eta * 4 * x2, 0, 0)

d2l.show_trace_2d(f_2d, d2l.train_2d(gd_2d))

По построению градиент в направлении $x_2$ гораздо больше и меняется гораздо быстрее, чем горизонтальный градиент по $x_1$. Поэтому мы оказываемся в затруднительном положении: если выбрать малую скорость обучения, мы гарантируем, что решение не будет расходиться в направлении $x_2$, но миримся с медленной сходимостью по $x_1$. Наоборот, при большой скорости обучения мы быстро продвигаемся по $x_1$, но расходимся по $x_2$. Следующий пример показывает, что даже при небольшом увеличении скорости обучения с $0.4$ до $0.6$ поведение меняется. Сходимость по $x_1$ улучшается, но в целом качество решения становится хуже.

eta = 0.6
d2l.show_trace_2d(f_2d, d2l.train_2d(gd_2d))

11.6.1.3 Метод импульса

Метод импульса (momentum) позволяет решить описанную выше проблему градиентного спуска. Глядя на траекторию оптимизации выше, можно предположить, что вычисление усреднённого прошлого градиента должно хорошо работать. В самом деле, в направлении $x_1$ это позволит суммировать хорошо согласованные градиенты и тем самым увеличит расстояние, которое мы проходим на каждом шаге. Напротив, в направлении $x_2$, где градиент колеблется, усреднённый градиент уменьшит размер шага, поскольку колебания взаимно гасятся. Использование ${\mathbf{v}}_t$ вместо градиента ${\mathbf{g}}_t$ даёт следующие уравнения обновления:

$$\begin{array}{rl}{\mathbf{v}}_t& \leftarrow \beta {\mathbf{v}}_{t-1}+{\mathbf{g}}_{t,t-1},\\ {\mathbf{x}}_t& \leftarrow {\mathbf{x}}_{t-1}-{\eta }_t{\mathbf{v}}_t.\end{array}$$

Заметим, что при $\beta =0$ мы получаем обычный градиентный спуск. Прежде чем углубляться в математические свойства метода, быстро посмотрим, как он работает на практике.

def momentum_2d(x1, x2, v1, v2):
    v1 = beta * v1 + 0.2 * x1
    v2 = beta * v2 + 4 * x2
    return x1 - eta * v1, x2 - eta * v2, v1, v2

eta, beta = 0.6, 0.5
d2l.show_trace_2d(f_2d, d2l.train_2d(momentum_2d))

Как видно, несмотря на ту же скорость обучения, что и раньше, метод импульса хорошо сходится. Посмотрим, что будет, если уменьшить параметр импульса. Если уменьшить его вдвое, до $\beta =0.25$, траектория почти не сходится. Тем не менее это намного лучше, чем расходимость без импульса.

eta, beta = 0.6, 0.25
d2l.show_trace_2d(f_2d, d2l.train_2d(momentum_2d))

Заметим, что метод импульса можно комбинировать со стохастическим градиентным спуском, в частности с минибатчевым стохастическим градиентным спуском. Единственное отличие в том, что в этом случае мы заменяем градиент ${\mathbf{g}}_{t,t-1}$ на ${\mathbf{g}}_t$. Для удобства инициализируем ${\mathbf{v}}_0=0$ в момент времени $t=0$.

11.6.1.4 Эффективные веса выборок

Напомним, что ${\mathbf{v}}_t={\sum }_{\tau =0}^{t-1}{\beta }^{\tau }{\mathbf{g}}_{t-\tau ,t-\tau -1}$. В предельном случае ${\sum }_{\tau =0}^{\infty }{\beta }^{\tau }=\frac{1}{1-\beta }$. Иначе говоря, вместо шага $\eta$ в градиентном спуске или стохастическом градиентном спуске мы используем шаг $\frac{\eta }{1-\beta }$ и одновременно учитываем потенциально лучшее направление спуска. Это сочетает в себе два преимущества. Чтобы показать, как влияют разные значения $\beta$ на веса, посмотрите на график ниже.

d2l.set_figsize()
betas = [0.95, 0.9, 0.6, 0]
for beta in betas:
    x = d2l.numpy(d2l.arange(40))
    d2l.plt.plot(x, beta ** x, label=f'β = {beta:.2f}')
d2l.plt.xlabel('время')
d2l.plt.legend();

11.6.2 Практические эксперименты

Посмотрим, как метод импульса работает на практике. Для этого нам нужна более масштабируемая реализация.

11.6.2.1 Реализация с нуля

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

def init_momentum_states(feature_dim):
    v_w = d2l.zeros((feature_dim, 1))
    v_b = d2l.zeros(1)
    return (v_w, v_b)

def init_momentum_states(features_dim):
    v_w = tf.Variable(d2l.zeros((features_dim, 1)))
    v_b = tf.Variable(d2l.zeros(1))
    return (v_w, v_b)

def init_momentum_states(feature_dim):
    v_w = d2l.zeros((feature_dim, 1))
    v_b = d2l.zeros([1])
    return (v_w, v_b)

def sgd_momentum(params, states, hyperparams):
    for p, v in zip(params, states):
        v[:] = hyperparams['momentum'] * v + p.grad
        p[:] -= hyperparams['lr'] * v

def sgd_momentum(params, states, hyperparams):
    for p, v in zip(params, states):
        with torch.no_grad():
            v[:] = hyperparams['momentum'] * v + p.grad
            p[:] -= hyperparams['lr'] * v
        p.grad.data.zero_()

def sgd_momentum(params, grads, states, hyperparams):
    for p, v, g in zip(params, states, grads):
            v[:].assign(hyperparams['momentum'] * v + g)
            p[:].assign(p - hyperparams['lr'] * v)

def sgd_momentum(params, states, hyperparams):
    a = []
    for p, v in zip(params, states):
        with paddle.no_grad():
            v[:] = hyperparams['momentum'] * v + p.grad
            p[:] -= hyperparams['lr'] * v
        p.grad.zero_()
        a.append(p)
    return a

Посмотрим, как это работает на практике.

def train_momentum(lr, momentum, num_epochs=2):
    d2l.train_ch11(sgd_momentum, init_momentum_states(feature_dim),
                   {'lr': lr, 'momentum': momentum}, data_iter,
                   feature_dim, num_epochs)

data_iter, feature_dim = d2l.get_data_ch11(batch_size=10)
train_momentum(0.02, 0.5)

Когда мы увеличиваем гиперпараметр импульса momentum до 0.9, эффективное число выборок становится равным $\frac{1}{1-0.9}=10$. Чтобы сохранить процесс управляемым, мы немного уменьшаем скорость обучения до $0.01$.

train_momentum(0.01, 0.9)

Дополнительное снижение скорости обучения помогает справиться с трудностями, которые возникают при негладкой оптимизации. Если установить её равной $0.005$, сходимость станет хорошей.

train_momentum(0.005, 0.9)

11.6.2.2 Краткая реализация

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

d2l.train_concise_ch11('sgd', {'learning_rate': 0.005, 'momentum': 0.9},
                       data_iter)

trainer = torch.optim.SGD
d2l.train_concise_ch11(trainer, {'lr': 0.005, 'momentum': 0.9}, data_iter)

trainer = tf.keras.optimizers.SGD
d2l.train_concise_ch11(trainer, {'learning_rate': 0.005, 'momentum': 0.9},
                       data_iter)

trainer = paddle.optimizer.Momentum
d2l.train_concise_ch11(trainer, {'learning_rate': 0.005, 'momentum': 0.9}, data_iter)

11.6.3 Теоретический анализ

Пример с 2D-функцией $f(x)=0.1x_1^2+2x_2^2$ может показаться надуманным. Ниже мы увидим, что он довольно типичен для реальных задач, по крайней мере при минимизации выпуклых квадратичных целевых функций.

11.6.3.1 Выпуклые квадратичные функции

Рассмотрим функцию

$$h(\mathbf{x})=\frac{1}{2}{\mathbf{x}}^{\top }\mathbf{Qx}+{\mathbf{x}}^{\top }\mathbf{c}+b.$$

Это обычная квадратичная функция. Для положительно определённой матрицы $\mathbf{Q}\succ 0$, то есть для матрицы с положительными собственными значениями, минимизатор равен ${\mathbf{x}}^{\ast }=-{\mathbf{Q}}^{-1}\mathbf{c}$, а минимальное значение — $b-\frac{1}{2}{\mathbf{c}}^{\top }{\mathbf{Q}}^{-1}\mathbf{c}$. Поэтому можно переписать $h$ в виде

$$h(\mathbf{x})=\frac{1}{2}(\mathbf{x}-{\mathbf{Q}}^{-1}\mathbf{c}{)}^{\top }\mathbf{Q}(\mathbf{x}-{\mathbf{Q}}^{-1}\mathbf{c})+b-\frac{1}{2}{\mathbf{c}}^{\top }{\mathbf{Q}}^{-1}\mathbf{c}.$$

Градиент задаётся формулой ${\partial }_{\mathbf{x}}f(\mathbf{x})=\mathbf{Q}(\mathbf{x}-{\mathbf{Q}}^{-1}\mathbf{c})$. То есть он равен расстоянию между $\mathbf{x}$ и минимизатором, умноженному на $\mathbf{Q}$. Поэтому и обновление импульса представляет собой линейную комбинацию членов $\mathbf{Q}({\mathbf{x}}_t-{\mathbf{Q}}^{-1}\mathbf{c})$.

Поскольку $\mathbf{Q}$ положительно определена, её можно разложить как $\mathbf{Q}={\mathbf{O}}^{\top }\mathbf{\Lambda }\mathbf{O}$, где $\mathbf{O}$ — ортогональная (поворотная) матрица, а $\mathbf{\Lambda }$ — диагональная матрица с положительными собственными значениями. Это позволяет перейти от переменной $\mathbf{x}$ к переменной $\mathbf{z}:=\mathbf{O}(\mathbf{x}-{\mathbf{Q}}^{-1}\mathbf{c})$ и получить очень простое выражение:

$$h(\mathbf{z})=\frac{1}{2}{\mathbf{z}}^{\top }\mathbf{\Lambda }\mathbf{z}+b^{\prime }.$$

Здесь $b^{\prime }=b-\frac{1}{2}{\mathbf{c}}^{\top }{\mathbf{Q}}^{-1}\mathbf{c}$. Поскольку $\mathbf{O}$ — всего лишь ортогональная матрица, она не искажает градиент по существу. Градиентный спуск в координатах $\mathbf{z}$ принимает вид

$${\mathbf{z}}_t={\mathbf{z}}_{t-1}-\mathbf{\Lambda }{\mathbf{z}}_{t-1}=(\mathbf{I}-\mathbf{\Lambda }){\mathbf{z}}_{t-1}.$$

Важный факт, который следует из этого выражения, состоит в том, что градиентный спуск не смешивает различные собственные подпространства. То есть если работать в собственной системе $\mathbf{Q}$, задача оптимизации решается покоординатно. Это справедливо и для метода импульса:

$$\begin{array}{rl}{\mathbf{v}}_t& =\beta {\mathbf{v}}_{t-1}+\mathbf{\Lambda }{\mathbf{z}}_{t-1}\\ {\mathbf{z}}_t& ={\mathbf{z}}_{t-1}-\eta \left(\beta {\mathbf{v}}_{t-1}+\mathbf{\Lambda }{\mathbf{z}}_{t-1}\right)\\ & =(\mathbf{I}-\eta \mathbf{\Lambda }){\mathbf{z}}_{t-1}-\eta \beta {\mathbf{v}}_{t-1}.\end{array}$$

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

11.6.3.2 Скалярная функция

С учётом полученного результата рассмотрим, что происходит при минимизации функции $f(x)=\frac{\lambda }{2}{x}^2$. Для градиентного спуска имеем

$$x_{t+1}=x_t-\eta \lambda x_t=(1-\eta \lambda )x_t.$$

При $|1-\eta \lambda |<1$ эта оптимизация сходится с экспоненциальной скоростью, поскольку после $t$ шагов получаем $x_t=(1-\eta \lambda {)}^t{x}_0$. Это показывает, как сначала растёт скорость сходимости по мере увеличения скорости обучения $\eta$ до значения, при котором $\eta \lambda =1$. За этим значением градиент начинает расходиться, а при $\eta \lambda >2$ оптимизация расходится.

lambdas = [0.1, 1, 10, 19]
eta = 0.1
d2l.set_figsize((6, 4))
for lam in lambdas:
    t = d2l.numpy(d2l.arange(20))
    d2l.plt.plot(t, (1 - eta * lam) ** t, label=f'λ = {lam:.2f}')
d2l.plt.xlabel('время')
d2l.plt.legend();

Чтобы проанализировать сходимость метода импульса, сначала перепишем уравнения обновления через два скалярных значения: одно для $x$ и одно для импульса $v$. Получаем:

$$\left[\begin{array}{c}{v}_{t+1}\\ x_{t+1}\end{array}\right]=\left[\begin{array}{cc}\beta & \lambda \\ -\eta \beta & (1-\eta \lambda )\end{array}\right]\left[\begin{array}{c}{v}_t\\ x_t\end{array}\right]=\mathbf{R}(\beta ,\eta ,\lambda )\left[\begin{array}{c}{v}_t\\ x_t\end{array}\right].$$

Матрицей $\mathbf{R}$ обозначим $2\times 2$ матрицу, которая определяет поведение сходимости. После $t$ шагов начальные значения $[v_0,x_0]$ становятся равными $\mathbf{R}(\beta ,\eta ,\lambda {)}^t[v_0,x_0]$. Поэтому скорость сходимости определяется собственными значениями $\mathbf{R}$. Прекрасную анимацию можно найти в статье (Goh, 2017). Подробный анализ приведён в работе (Flammarion and Bach, 2015). Вкратце, метод импульса сходится при $0<\eta \lambda <2+2\beta$. Это более широкий диапазон допустимых параметров по сравнению с $0<\eta \lambda <2$ для градиентного спуска. Кроме того, в целом предпочтительны большие значения $\beta$.

Резюме

• Метод импульса заменяет градиент средним значением прошлых градиентов, что существенно ускоряет сходимость.
• Метод импульса полезен как для градиентного спуска без шума, так и для стохастического градиентного спуска с шумом.
• Метод импульса может предотвратить застревание оптимизации в процессе стохастического градиентного спуска.
• Из-за экспоненциального затухания весов прошлых данных эффективное число градиентов равно $\frac{1}{1-\beta }$.
• Для выпуклых квадратичных задач метод импульса допускает явный и подробный анализ.
• Реализация метода импульса очень проста, но требует хранения дополнительного вектора состояния (импульса $\mathbf{v}$).

Упражнения

1. Поэкспериментируйте с другими комбинациями гиперпараметров импульса и скорости обучения, наблюдайте и анализируйте результаты.
2. Попробуйте градиентный спуск и метод импульса на квадратичной задаче с несколькими собственными значениями, то есть $f(x)=\frac{1}{2}{\sum }_i{\lambda }_i{x}_i^2$, например, при ${\lambda }_i=2^{-i}$. Постройте график того, как уменьшается $x$ при инициализации $x_i=1$.
3. Найдите минимум и минимизатор функции $h(\mathbf{x})=\frac{1}{2}{\mathbf{x}}^{\top }\mathbf{Qx}+{\mathbf{x}}^{\top }\mathbf{c}+b$.
4. Что меняется при выполнении стохастического градиентного спуска с импульсом? Что происходит при использовании минибатчевого стохастического градиентного спуска с импульсом? Поэкспериментируйте с параметрами.

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

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

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

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

Литература

• Flammarion, N., & Bach, F. (2015). From averaging to acceleration, there is only a step-size. Conference on Learning Theory, 658–695.
• Goh, G. (2017). Why momentum really works. Distill.
• Nesterov, Y. (2018). Lectures on Convex Optimization. Springer.
• Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics, 4, 1–17.
• Sutskever, I., Martens, J., Dahl, G., & Hinton, G. (2013). On the importance of initialization and momentum in deep learning. International Conference on Machine Learning, 1139–1147.