AdEMAMix: 通过双重EMA优化神经网络训练

介绍

背景

在机器学习中，AdamW是用于大型神经网络训练的最常用优化器之一。尽管其广泛使用，但AdamW依赖于单一的指数移动平均（EMA）来对梯度进行加权，这在处理新旧梯度信息时常常表现不足。这种局限性导致在优化过程中难以平衡对较新和较旧梯度的加权，从而影响模型的收敛速度和优化效果。

目标

AdEMAMix旨在实现与AdamW相同的优化效果，但通过减少所需的训练令牌数量来提高效率。通过有效地整合最近和较旧的梯度信息，AdEMAMix可以加快模型的训练过程，从而在相同的时间内完成更多的任务。

创新点

AdEMAMix的创新之处在于其采用了双重EMA的方法，这使得优化器能够同时捕捉不同时间段的梯度信息。通过结合两个不同衰减率的EMA，AdEMAMix能够改善梯度信息的平衡，增强模型的收敛性和稳定性。

优化器

AdamW

AdamW是应用广泛的优化器，适合大型神经网络训练。它使用单一EMA来平衡梯度更新，这虽然在某些情况下有效，但在处理长时间序列数据时可能会导致信息遗失。

单一EMA

在AdamW中，使用的是单一的指数移动平均方法。这种方法在很大程度上依赖于最近的梯度变化，可能忽略了较旧但仍具参考价值的信息。

优缺点

AdamW的优点包括易于实现和良好的收敛性；但其缺点则在于过度依赖单一EMA，导致在长时间训练中可能错失重要的梯度信息。

AdEMAMix

AdEMAMix是一种新的优化器，通过结合两个EMA以提高优化效率。它不仅克服了AdamW在长时间序列数据上的不足，还提升了模型的整体表现。

双重EMA

双重EMA是AdEMAMix的核心，与传统的单一EMA不同，它同时使用快速EMA和缓慢EMA来平衡不同时间段的梯度信息。

快速EMA

快速EMA运用较低的衰减率，能够即时响应最近的变化。这对于快速适应梯度的剧烈波动至关重要。

缓慢EMA

缓慢EMA则使用较高的衰减率，利用较旧的梯度信息。这有助于增强模型的长期记忆，减缓忘记训练数据的速度。

更新规则

AdEMAMix在更新过程中对Adam的更新规则进行了改进，使得能够同时考虑两个EMA。这种结合不仅提高了梯度信息的有效性，还为从AdamW中切换到AdEMAMix提供了平滑的过渡，通过将第二个动量项初始化为零来实现。

实验结果

性能比较

实验结果表明，AdEMAMix在语言建模和视觉任务上都优于AdamW。其优化稳定性提升，收敛速度加快，并且在较长的训练过程中减缓了模型对训练数据的遗忘。

应用领域

AdEMAMix的设计使其在多个应用领域具有良好的适应性和影响力，特别是在需要长时间训练的任务中，能够利用减半的训练令牌达到与AdamW相同的效果。这一突破性进展为深度学习领域提供了更多选择和可能性。

总之，AdEMAMix的推出为优化器的设计带来了新的思路，特别是在如何平衡和利用梯度信息方面，为未来的研究和应用奠定了基础。