Adam: инерция, масштаб и поправка смещения

Зачем нужен Adam

У нейросети миллионы весов, и их градиенты живут в очень разных масштабах: у эмбеддингов редких токенов — редкие крупные толчки, у последних слоёв — постоянные мелкие, у слоёв нормализации — третьи. Обычный SGD делает шаг $\eta g$, и один темп $\eta$ для всех одновременно слишком велик для одних весов и слишком мал для других. К тому же мини-батчевый градиент шумит, а поверхность потерь полна вытянутых оврагов и плато.

Adam (adaptive moment estimation, Кингма и Ба, 2014; ICLR 2015) решает обе задачи одной формулой:

• первый момент $m_t$ — экспоненциальное среднее градиентов: инерция, которая гасит шум и зигзаги (как Momentum);
• второй момент $v_t$ — экспоненциальное среднее квадратов градиента: типичный размер градиента каждого веса, на который делится шаг (как RMSProp);
• поправка смещения — деление на $1-\beta^t$, которое убирает систематическую недооценку моментов на первых шагах.

Результат: шаг каждого веса порядка $\eta$ независимо от масштаба его градиента, без ручной подстройки под каждый слой. Поэтому Adam — оптимизатор по умолчанию в большинстве задач, а его вариант AdamW обучает трансформеры и большие языковые модели.

Интуиция: компас и спидометр

Представьте, что вы ведёте машину по бездорожью по показаниям дрожащего компаса. Глупо поворачивать руль при каждом его рывке — разумнее следить за усреднённым направлением. Это $m_t$. А скорость стоит выбирать по тому, насколько обычно сильны рывки: на ровной дороге можно ехать уверенно, на ухабах — осторожно. Это $\sqrt{v_t}$: шаг Adam — направление, делённое на типичную силу толчков.

На рисунке — поток градиентов одного веса: среднее $+0{,}5$ до шага 100 и $-0{,}5$ после, плюс шум. Оранжевая линия — сглаженный градиент $\hat m_t$, бирюзовые — коридор $\pm\sqrt{\hat v_t}$, внизу — итоговый шаг в единицах $\eta$.

Формула по косточкам

$g_t$

градиент на мини-батче шага $t$ (шаги нумеруются с 1);

$m_t$

первый момент: экспоненциальное скользящее среднее (EMA) градиента, $m_0=0$;

$v_t$

второй момент: EMA поэлементного квадрата градиента, $v_0=0$; у каждого веса свои $m$ и $v$;

$\beta_1,\ \beta_2$

коэффициенты затухания, по умолчанию $0{,}9$ и $0{,}999$: память примерно $1/(1-\beta)$ шагов — 10 и 1000;

$1-\beta^t$

поправка смещения: компенсирует старт моментов с нуля;

$\eta$

темп обучения (learning rate), по умолчанию $10^{-3}$; здесь он задаёт типичную длину шага в единицах весов;

$\epsilon$

маленькая добавка ($10^{-8}$) против деления на ноль.

Пройдём три шага руками для градиентов $0{,}5$, $-0{,}2$, $0{,}4$ при стандартных $\beta_1=0{,}9$, $\beta_2=0{,}999$. Последний столбец — каким был бы шаг без поправки смещения.

Обратите внимание на третий шаг: градиенты меняли знак, поэтому $\hat m_3=0{,}23$ заметно меньше типичного размера градиента $\sqrt{\hat v_3}=0{,}39$, и Adam шагает лишь на $0{,}6\eta$. Чем согласованнее градиенты, тем ближе шаг к $\eta$.

Поправка смещения: откуда 1 − βᵗ

Моменты стартуют с нуля и первые шаги «помнят» этот ноль. Раскроем рекурсию для $m_t$:

Это взвешенная сумма градиентов, но сумма весов равна не 1, а $(1-\beta_1)(1+\beta_1+\dots+\beta_1^{t-1})=1-\beta_1^t$. Недостающая доля $\beta_1^t$ досталась начальному нулю. Отсюда и поправка.

Свойства шага

Масштабная инвариантность

Умножим все градиенты веса на $c>0$. Тогда $m_t$ умножится на $c$, $v_t$ — на $c^2$, $\sqrt{v_t}$ — на $c$, и отношение $\hat m_t/\sqrt{\hat v_t}$ не изменится. Шаг Adam — однородная функция нулевой степени от градиентов (пока $\sqrt{\hat v_t}\gg\epsilon$). Эффективный темп каждого веса $\eta/(\sqrt{\hat v_t}+\epsilon)$ подстраивается сам: большой у весов с мелкими градиентами, маленький — у весов с крупными.

Шаг и отношение сигнал/шум

Пусть градиент веса стационарен со средним $\mu$ и стандартным отклонением $\sigma$. Тогда $\hat m_t\approx\mu$, $\hat v_t\approx\mu^2+\sigma^2$ (второй момент — это квадрат среднего плюс дисперсия), и средний шаг

При чистом сигнале ($\sigma\ll|\mu|$) шаг $\approx\eta$, в сильном шуме ($\sigma\gg|\mu|$) он пропорционален отношению сигнал/шум. Это автоматическое «торможение»: у минимума средний градиент стремится к нулю, а шум остаётся — и Adam сам укорачивает шаги.

Особые случаи

• $\beta_1=\beta_2=0$: $\hat m_t=g_t$, $\sqrt{\hat v_t}=|g_t|$, шаг $\eta\operatorname{sign}(g_t)$ — знаковый градиентный спуск (signSGD).
• $\beta_1=0$: RMSProp с поправкой смещения.
• $\epsilon\gg\sqrt{\hat v_t}$: шаг $\approx(\eta/\epsilon)\,\hat m_t$ — SGD с инерцией и темпом $\eta/\epsilon$.
• Adam — диагональный предобусловливатель: он масштабирует каждую координату отдельно и не умеет поворачивать оси (см. рисунок 6 на странице SGD).

Математика за формулой

Соберём вместе то, что прячется за четырьмя строками.

• EMA как взвешенное среднее. $\hat m_t=\sum_i w_i g_i$ с весами $w_i=(1-\beta_1)\beta_1^{t-i}/(1-\beta_1^t)$, которые положительны и в сумме дают 1. Вес убывает вдвое примерно каждые $\ln2/(1-\beta)$ шагов: 7 для $\beta_1=0{,}9$ и 693 для $\beta_2=0{,}999$.
• Эффективное окно $1/(1-\beta)$: $m$ усредняет примерно 10 последних градиентов, $v$ — примерно 1000. Направление реагирует быстро, масштаб — медленно и потому устойчиво.
• Второй момент — не дисперсия. $v$ оценивает $\mathbb E[g^2]=\mu^2+\sigma^2$, то есть «нецентрированный» размер градиента. Поэтому при чистом сигнале шаг $\approx\eta$, а не бесконечен.
• Связь с методом Ньютона. Иногда говорят, что $\sqrt v$ аппроксимирует кривизну. Это верно лишь отчасти: $v$ — оценка диагонали эмпирической матрицы Фишера (по градиентам мини-батчей), которая для некоторых потерь у минимума близка к гессиану, но Adam берёт из неё квадратный корень. Надёжнее понимать Adam как нормировку шага, а не как приближение Ньютона.

Гонка оптимизаторов в 3D

Сравним пять методов на двух поверхностях. Первая — плато с ямой: $f=1-e^{-q(x,y)}$, $q=0{,}2(x-1)^2+1{,}2(y-0{,}5)^2$. Вдали от ямы $f\approx1$, градиент крошечный, а сама яма вытянута. Вторая — банан Розенброка.

AdamW: правильное затухание весов

Классическая $L_2$-регуляризация добавляет к потерям $\tfrac\lambda2\|\theta\|^2$, то есть к градиенту — $\lambda\theta$. Для SGD это то же самое, что на каждом шаге умножать веса на $1-\eta\lambda$ (затухание весов, weight decay). Для Adam — нет: слагаемое $\lambda\theta$ попадает в $m$ и $v$ и делится на $\sqrt{\hat v_t}$. Лощилов и Хаттер (2017; ICLR 2019) предложили отделить затухание от градиента:

В AdamW моменты считаются по градиенту одних лишь потерь, а все веса затухают с одной и той же скоростью $\eta\lambda$ за шаг. В Adam с $L_2$ затухание каждого веса делится на его $\sqrt{\hat v_t}$: веса с крупными градиентами почти не регуляризуются, а веса без градиента потерь сжимаются с постоянной скоростью $\approx\eta$ за шаг, независимо от $\lambda$.

Adam в обучении: разогрев и расписание

Даже адаптивному методу нужен график темпа. Самая распространённая схема для трансформеров — линейный разогрев (warmup) от нуля до $\eta_{\max}$ за первые сотни или тысячи шагов, затем косинусное снижение до $\eta_{\min}$:

Почему снижение помогает: из раздела о свойствах шаг Adam у минимума пропорционален $\eta$, умноженному на отношение сигнал/шум. Чем меньше $\eta$, тем меньше облако, в котором веса дрожат вокруг минимума.

Почему разогрев нужен: на первых шагах $\hat v_t$ оценён по нескольким градиентам и сильно шумит, а веса только что инициализированы, и градиенты особенно велики. Лю и соавторы (RAdam, 2019) показали, что дисперсия эффективного темпа $\eta/\sqrt{\hat v_t}$ в начале обучения велика, и предложили выпрямлять её аналитически; разогрев решает ту же проблему проще.

В коде: с нуля и в PyTorch

Реализация с нуля — двадцать строк. Флаг decoupled переключает затухание весов между AdamW и $L_2$ внутри градиента.

pythonimport numpy as np

class Adam:
    """Adam; при weight_decay > 0 и decoupled=True — AdamW. Как в torch.optim."""
    def __init__(self, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, decoupled=True):
        self.lr, (self.b1, self.b2), self.eps = lr, betas, eps
        self.wd, self.decoupled = weight_decay, decoupled
        self.m = self.v = None
        self.t = 0

    def step(self, theta, g):
        if self.m is None:
            self.m, self.v = np.zeros_like(theta), np.zeros_like(theta)
        if self.wd and not self.decoupled:
            g = g + self.wd * theta                        # L2 внутри градиента (Adam)
        self.t += 1
        self.m = self.b1 * self.m + (1 - self.b1) * g      # инерция
        self.v = self.b2 * self.v + (1 - self.b2) * g * g  # масштаб
        m_hat = self.m / (1 - self.b1 ** self.t)           # поправка смещения
        v_hat = self.v / (1 - self.b2 ** self.t)
        if self.wd and self.decoupled:
            theta = theta - self.lr * self.wd * theta      # AdamW: затухание отдельно
        return theta - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

# Три шага из таблицы: g = 0.5, −0.2, 0.4 (lr = 1, чтобы шаг был в единицах η)
opt, th = Adam(lr=1.0), np.zeros(1)
for g in [0.5, -0.2, 0.4]:
    new = opt.step(th, np.array([g]))
    print(f"t={opt.t}: m̂={opt.m[0] / (1 - 0.9 ** opt.t):.4f}  шаг/η={th[0] - new[0]:.4f}")
    th = new
# t=1: m̂=0.5000  шаг/η=1.0000
# t=2: m̂=0.1316  шаг/η=0.3456
# t=3: m̂=0.2306  шаг/η=0.5955

# Масштабная инвариантность: градиенты ×1000 дают те же шаги
rng = np.random.default_rng(0)
a, b, x, y = Adam(lr=0.01), Adam(lr=0.01), np.zeros(1), np.zeros(1)
for t in range(100):
    g = rng.normal(1.0, 0.5, size=1)
    x, y = a.step(x, g), b.step(y, 1000 * g)
print(x, y, abs(x - y))   # [-0.95081564] [-0.95081565] [8.64890637e-09] — разница только из-за ε

Выход совпадает с таблицей из раздела «Формула по косточкам», а torch.optim.Adam на том же первом шаге сдвигает вес ровно на $\eta$. Типичный цикл обучения с AdamW, раздельными группами затухания, разогревом и косинусом:

pythonimport math
import torch
import torch.nn.functional as F

decay = [p for p in model.parameters() if p.ndim >= 2]     # матрицы весов
no_decay = [p for p in model.parameters() if p.ndim < 2]   # смещения и нормализации
opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
                         {"params": no_decay, "weight_decay": 0.0}],
                        lr=3e-4, betas=(0.9, 0.95), eps=1e-8)

total, warmup = 10_000, 500
def lr_factor(step):                       # множитель к lr: разогрев, затем косинус до 10%
    if step < warmup:
        return (step + 1) / warmup
    p = (step - warmup) / (total - warmup)
    return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * p))
sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_factor)

for step, (x, y) in zip(range(total), loader):
    loss = F.mse_loss(model(x), y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # обрезка градиента
    opt.step()
    sched.step()                          # расписание — после каждого шага

Гиперпараметры и практика

• Начните с умолчаний. $\eta=10^{-3}$ (или $3\cdot10^{-4}$ для трансформеров), $\beta=(0{,}9;\,0{,}999)$, AdamW с $\lambda=0{,}01$. Подбирать в первую очередь стоит $\eta$, затем $\lambda$ и расписание.
• Обрезка градиента по норме ($1{,}0$) почти бесплатна и спасает от редких выбросов.
• Сохраняйте состояние оптимизатора (opt.state_dict()) вместе с весами: при продолжении обучения с нулевыми $m$ и $v$ первые шаги снова будут «холодными».
• Когда SGD лучше. На задачах компьютерного зрения хорошо настроенный SGD с инерцией часто даёт лучшее обобщение (Уилсон и др., 2017). Для трансформеров, рекомендательных систем с эмбеддингами и задач с разреженными градиентами выигрывает Adam.

Частые ошибки

1. Adam(weight_decay=…) вместо AdamW. Это $L_2$ внутри градиента: затухание делится на $\sqrt{\hat v_t}$ и работает неравномерно (рисунок 6).
2. Темп от SGD. $\eta=0{,}1$ хорош для SGD с инерцией, но для Adam это шаг около $0{,}1$ в единицах весов на каждом шаге — обучение почти наверняка разойдётся.
3. Затухание для всех параметров. Смещения и веса LayerNorm/BatchNorm обычно исключают из weight decay: их сжатие к нулю вредит, а не регуляризует.
4. Неудачный $\epsilon$. Слишком маленький в float16 округляется до нуля, слишком большой ($\gtrsim\sqrt{\hat v}$) молча отключает адаптацию — шаг становится $(\eta/\epsilon)\hat m$.
5. Без разогрева на трансформере. Большой $\eta$ с первых шагов при шумной оценке $\hat v_t$ часто приводит к взрыву потерь в начале обучения.
6. Расписание не на том шаге. scheduler.step() раз за эпоху там, где расписание рассчитано на шаги (или наоборот) растягивает или сжимает разогрев в сотни раз.

Связи с другими формулами

• SGD, Momentum и RMSProp: числитель Adam — EMA-инерция, знаменатель — $\sqrt v$ из RMSProp; там же условие устойчивости и число обусловленности.
• Кросс-энтропия и softmax: потеря, которую AdamW минимизирует в языковых моделях; её градиент по логитам $s-y$ и есть $g_t$ последнего слоя.
• MSE и Huber: от масштаба целевой переменной зависит масштаб градиента — Adam к нему нечувствителен, SGD — нет.
• GELU и SiLU: активации трансформеров, которые обучаются AdamW.

Итоги

Проверь себя