AdEMAMix 优化器：更好、更快、更旧

基于动量的优化器在广泛的机器学习应用中至关重要。这些优化器通常依赖于梯度的指数移动平均 (EMA)，它对较旧梯度的当前贡献进行指数衰减。这是因为梯度是局部线性近似，随着迭代沿损失地形移动会失去其相关性。这项工作质疑了使用单一 EMA 来累积过去梯度的做法，并从经验上证明了这种选择可能是次优的：单一 EMA 无法同时对最近的过去赋予高权重，又对较旧的梯度赋予不可忽略的权重。基于这一观察，我们提出了 AdEMAMix，一种对 Adam 优化器的简单修改，使用两个 EMA 的混合以更好地利用过去的梯度。我们在语言建模和图像分类上的实验表明——相当令人惊讶的是——梯度可以在数万步内保持相关性。它们有助于更快地收敛，并且通常收敛到更低的极小值：例如，在 $101$B tokens 上训练的 $1.3$B 参数 AdEMAMix LLM 的性能与在 $197$B tokens 上训练的 AdamW 模型 ($+95\%$) 相当。此外，我们的方法显著减缓了训练期间的模型遗忘。我们的工作激发了对利用过去梯度的不同类型函数（超越 EMA）的进一步探索。