arXiv:2409.03137cs.LGstat.ML2024-09ICLR被引 45

改进优化器,让旧梯度也能发挥作用,训练更快更准

The AdEMAMix Optimizer: Better, Faster, Older

  • 用两个指数移动平均混合,更好利用历史梯度
  • 13亿参数模型仅需1010亿词训练,效果媲美AdamW的1970亿词
  • 显著减缓模型遗忘,适合长周期训练任务

基于动量的优化器广泛应用于各类机器学习任务。这类方法通常依赖梯度的指数移动平均(EMA),对较早梯度按指数衰减。这源于梯度作为局部线性近似,随迭代推进逐渐失效。本文质疑单一EMA累积历史梯度的合理性,实证表明:单一EMA难以同时赋予近期梯度高权重和旧梯度非可忽略权重。为此提出AdEMAMix,是对Adam的简单改进,采用双EMA混合机制以更好利用历史梯度。在语言建模与图像分类任务上的实验显示——令人惊讶地——梯度在数万步内仍具相关性。该方法能加速收敛,并常达到更低极小值:例如,13亿参数的AdEMAMix大模型在1010亿词上训练,表现可比于在1970亿词上训练的AdamW模型(+95%)。此外,该方法显著减缓训练中的模型遗忘。本工作启发未来探索超越EMA的梯度利用方式。

原文摘要 · Abstract (English)

Momentum based optimizers are central to a wide range of machine learning applications. These typically rely on an Exponential Moving Average (EMA) of gradients, which decays exponentially the present contribution of older gradients. This accounts for gradients being local linear approximations which lose their relevance as the iterate moves along the loss landscape. This work questions the use of a single EMA to accumulate past gradients and empirically demonstrates how this choice can be sub-optimal: a single EMA cannot simultaneously give a high weight to the immediate past, and a non-negligible weight to older gradients. Building on this observation, we propose AdEMAMix, a simple modification of the Adam optimizer with a mixture of two EMAs to better take advantage of past gradients. Our experiments on language modeling and image classification show -- quite surprisingly -- that gradients can stay relevant for tens of thousands of steps. They help to converge faster, and often to lower minima: e.g., a $1.3$B parameter AdEMAMix LLM trained on $101$B tokens performs comparably to an AdamW model trained on $197$B tokens ($+95\%$). Moreover, our method significantly slows-down model forgetting during training. Our work motivates further exploration of different types of functions to leverage past gradients, beyond EMAs.

优化器梯度利用大模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。