Оптимизаторы являются важными инструментами в стеке моделирования. Одним из наиболее широко используемых оптимизаторов является оптимизатор Adam, представленный Kingma и Ba [бумага]. Этот оптимизатор отслеживает скользящие средние градиента (термин импульса) и второй момент (термин энергии) градиента с использованием фильтров экспоненциального скользящего среднего (EMA) и использует квадратный корень из термина энергии для нормализации члена импульса. прежде чем сделать шаг.
Это кажется хорошей идеей, верно? Во-первых, это выглядит как диагональное приближение оптимизации второго порядка, а во-вторых, когда градиенты зашумлены, знаменатель (квадратный корень из энергетического члена) велик по отношению к числителю (импульсный член), а шаги малы. С другой стороны, когда градиенты согласованы, знаменатель приблизительно равен числителю, и мы делаем шаги постоянного размера, равные скорости обучения γ. Вот почему этот оптимизатор де-факто является выбором исследователей и практиков машинного обучения.
Было предложено и изучено множество вариаций оптимизатора Adam (см. AdamW из PyTorch docs и оптимизаторы типа QHAdam из пакета torch_optimizer). Здесь мы предлагаем вариант, который существенно не отклоняется от Адама для зашумленных недавних градиентов, но значительно увеличивает фактический размер шага для параметров с последовательной недавней историей градиента (насколько недавний зависит от параметров EMA β1 и β2).
Наша модификация проста, но эффективна: мы заменяем EMA-фильтр члена энергии градиента на EMA-фильтр члена дисперсии градиента. Это означает окончательное уравнение обновления θ(t) = θ(t-1) - γ * sqrt(SNR) * sign(импульс), где SNR относится к отношению сигнал-шум в «недавней» истории градиента (SNR равно термин, заимствованный из литературы по обработке сигналов и обозначающий отношение средней мощности бесшумного сигнала к дисперсии шума в сигнале). Таким образом, параметры с высоким SNR, то есть с согласованными «недавними» историями градиента, будут иметь гораздо больший размер шага (до бесконечности, если градиент постоянный), чем те, которые имеют зашумленные «недавние» истории градиента (или низкие значения SNR градиента). . Реализация этого оптимизатора проста и приведена ниже для полноты картины.
from typing import Tuple
import torch
from torch.optim.optimizer import Optimizer
class SNRAdam(Optimizer):
r"""Implements the SNRAdam optimization algorithm, which uses std deviation for the denominator rather than
sqrt(energy) term used in conventional Adam. Why is this a good idea? If gradient stddev for a param is small, we
should take larger steps as it means the gradient is consistent over time.
Arguments:
params: iterable of parameters to optimize or dicts defining
parameter groups
lr: learning rate (default: 1e-3)
betas: coefficients used for computing
running averages of gradient and its variance (default: (0.9, 0.999))
eps: term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay: weight decay (L2 penalty) (default: 0)
"""
def __init__(
self,
params,
lr: float = 1e-3,
betas: Tuple[float, float] = (0.9, 0.999),
weight_decay: float = 0.0,
eps: float = 1e-8,
):
if lr <= 0.0:
raise ValueError('Invalid learning rate: {}'.format(lr))
if eps < 0.0:
raise ValueError('Invalid epsilon value: {}'.format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError(
'Invalid beta parameter at index 0: {}'.format(betas[0])
)
if not 0.0 <= betas[1] < 1.0:
raise ValueError(
'Invalid beta parameter at index 1: {}'.format(betas[1])
)
if weight_decay < 0:
raise ValueError(
'Invalid weight_decay value: {}'.format(weight_decay)
)
defaults = {
'lr': lr,
'betas': betas,
'weight_decay': weight_decay,
'eps': eps,
}
super().__init__(params, defaults)
def step(self, closure=None):
"""Performs a single optimization step.
Arguments:
closure: A closure that reevaluates the model and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
lr = group['lr']
beta1, beta2 = group['betas']
weight_decay = group['weight_decay']
eps = group['eps']
for p in group['params']:
if p.grad is None:
continue
d_p = p.grad.data
if d_p.is_sparse:
raise RuntimeError(
'SNRAdam does not support sparse gradients, '
'please consider SparseAdam instead'
)
state = self.state[p]
if weight_decay != 0:
p.data.mul_(1 - lr * weight_decay)
if len(state) == 0:
state['iter_'] = 1
state['exp_avg'] = torch.zeros_like(
p.data, memory_format=torch.preserve_format
)
state['exp_avg_sq'] = torch.zeros_like(
p.data, memory_format=torch.preserve_format
)
iter_ = state['iter_']
exp_avg = state['exp_avg']
if iter_ == 1:
d_sub_p_sq = d_p - exp_avg
else:
d_sub_p_sq = d_p - exp_avg.mul(1.0 / (1 - beta1 ** (iter_ - 1)))
d_sub_p_sq.mul_(d_sub_p_sq)
exp_avg_sq = state['exp_avg_sq']
exp_avg.mul_(beta1).add_(d_p, alpha=1.0 - beta1)
exp_avg_sq.mul_(beta2).add_(d_sub_p_sq, alpha=1.0 - beta2)
p.data.addcdiv_(exp_avg.mul(1.0 / (1 - beta1 ** iter_)),
exp_avg_sq.mul(1.0 / (1 - beta2 ** iter_)).sqrt() + eps, value=-lr)
state['iter_'] += 1
return loss
Мы провели эксперименты, используя простую модель Vision Transformer со 100 000 параметрами (во многом вдохновленную этим средним постом) в наборе данных MNIST с размером партии 4096 в течение 20 эпох со скоростью обучения 1e-3. Результаты указывают на быструю сходимость предложенного оптимизатора по сравнению с оптимизатором Адама (мы построили один из прогонов, но заметили одно и то же поведение последовательно для этой комбинации модели и набора данных):
Чтобы распутать источник выигрыша, мы делаем размер партии = ∞ и сравниваем два алгоритма. Это показывает, происходит ли выигрыш от (i) корректировки стохастики в «стохастическом» градиентном спуске, т. Е. Размер партии меньше, чем полный размер набора данных, или (ii) шума в градиенте, возникающего из траектории оптимизации (часть градиентного спуска ). Мы видим, что выигрыш происходит от компенсации (ii):
