The AdEMAMix Optimizer: Better, Faster, Older

基于动量的优化器是许多机器学习应用的核心。这些优化器通常依赖于梯度的指数移动平均（EMA），该平均值指数衰减旧梯度的贡献。这解释了为什么梯度是局部线性逼近，随着迭代在损失曲面上移动，它们失去了相关性。本文质疑了使用单个EMA积累过去梯度的做法，并通过实验证明了这种选择可能是次优的：单个EMA无法同时高度重视近期梯度和旧梯度的贡献。基于这个观察，我们提出了AdEMAMix，这是Adam优化器的一个简单修改，采用两个EMA的混合，以更好地利用过去的梯度。我们在语言建模和图像分类上的实验表明，令人惊讶的是，梯度在数万步内仍然保持相关性。它们有助于更快地收敛，并且通常可以到达更低的最小值：例如，一个使用$101$B标记训练的$1.3$B参数的AdEMAMix LLM与使用$197$B标记训练的AdamW模型表现相当（$+95\%$）。此外，我们的方法显著减缓了模型在训练过程中的遗忘。我们的工作激发了对不同类型的函数进行进一步探索，以利用过去的梯度，超越EMA。

本论文旨在解决动量优化器在累积过去梯度时存在的问题，提出了一种新的优化器AdEMAMix，通过混合两个EMA来更好地利用过去的梯度。

论文的关键思路是通过混合两个EMA来更好地利用过去的梯度，以便优化器能够同时高权重地考虑最近的梯度和较老的梯度。

实验表明，过去的梯度在数万步之后仍然有效，有助于更快地收敛到更低的最小值。AdEMAMix LLM在101B个标记上训练的1.3B参数模型与在197B个标记上训练的AdamW模型相比性能相当，而后者训练了更多的标记。此外，AdEMAMix显著减缓了模型在训练过程中的遗忘速度。

最近的相关研究包括使用不同EMA的其他优化器，以及使用其他方法来利用过去的梯度，例如使用非线性滤波器。