Pagliardini等人在其9月发布的论文中，提出了一种名为AdEMAMix的新型优化算法。这种算法旨在解决当前广泛使用的Adam及其变体（如AdamW）在利用长期梯度信息方面的局限性。AdEMAMix的核心在于使用两个不同衰减率的指数移动平均（EMA），以更有效地利用历史梯度信息。
研究动机
传统的动量优化器通常使用单一EMA来累积过去的梯度，这种方法面临一个两难困境：较小的衰减率（β）会导致优化器对近期梯度给予较高权重，但快速遗忘旧梯度；较大的衰减率可以保留更多旧梯度信息，但会减慢对近期梯度的响应。研究者们发现，即使在数万步训练之后，梯度信息仍然可能保持有用。这一发现促使他们设计了AdEMAMix，以同时利用近期和远期的梯度信息。
AdEMAMix算法
核心思想
AdEMAMix的核心在于使用两个EMA项：快速EMA（低β值）和慢速EMA（高β值）。这两个EMA项的设计使得AdEMAMix可以同时利用近期的梯度信息和远期的梯度信息，从而提高了优化器的性能。
参数更新规则
AdEMAMix的参数更新规则考虑了两个EMA项和Adam中的二阶矩估计，通过权衡两个EMA项的系数和权重衰减系数，实现了模型的优化。
AdEMAMix的核心思想
稳定性改进
为了提高训练稳定性，研究者们引入了α和β₃的调度器。这些调度器在训练初期逐渐增加α和β₃的值，避免了由于过大的动量值导致的早期训练不稳定。
AdEMAMix的β₃调度器。与线性调度相比，该调度器在β值较小时增长较快，在β值较大时增长较慢，更好地适应了不同β值对优化过程的影响。
实验设置
研究者们在两个主要任务上评估了AdEMAMix的性能：语言建模任务和视觉任务。
语言建模任务
- 模型：Transformer架构，参数规模从110M到1.3B
- 数据集：RedPajama v2
- 评估指标：验证集perplexity、训练速度、模型遗忘程度
视觉任务
- 模型：Vision Transformer (ViT)，24M和86M参数
- 数据集：ImageNet-1k和ImageNet-21k
- 评估指标：测试集损失、Top-1准确率
主要实验结果
语言建模性能
1.3B参数语言模型在不同训练token数下的性能比较。图中显示AdEMAMix仅使用101B tokens就达到了AdamW使用197B tokens的性能，节省了近50%的训练数据。
研究结果表明，AdEMAMix在各种模型规模下均显著优于AdamW：对于110M参数模型，AdEMAMix训练256k步的性能相当于AdamW训练500k步；对于1.3B参数模型，AdEMAMix使用770k步（约101B tokens）即可达到AdamW使用1.5M步（约197B tokens）的性能。
模型遗忘分析
AdEMAMix和AdamW在训练过程中对特定batch的遗忘程度比较。AdEMAMix表现出更慢的遗忘速度，有助于提高学习稳定性。
作者们通过跟踪某个batch在训练过程中的loss变化来衡量遗忘程度。结果表明：AdEMAMix模型遗忘训练数据的速度更慢，使用AdEMAMix时，早期训练的batch对最终模型的影响更大。
视觉任务性能
在视觉任务中，AdEMAMix同样表现出色：在ImageNet-21k上，AdEMAMix consistently优于AdamW，尤其是在数据量较大时；在ImageNet-1k上，当模型容量与数据量比例适当时，AdEMAMix仍能获得性能提升。
计算开销
尽管AdEMAMix引入了额外的计算步骤，但研究者们发现其带来的计算开销可以忽略不计：训练时间仅比AdamW略长（不足2%增加）；在分布式训练环境中，预期额外开销会进一步减少。考虑到AdEMAMix可以显著减少达到同等性能所需的训练步数，作者们认为这微小的额外开销是完全可以接受的。
结论与未来展望
Pagliardini等人通过AdEMAMix成功地在多个任务上展现出显著优于AdamW的性能。这种新的优化器不仅加快了模型收敛速度，还提高了学习稳定性，为大规模神经网络的高效训练提供了新的方法。
研究结果表明，梯度信息可以在数万步训练中保持有效，这一发现为进一步探索非EMA类型的梯度累积方法开辟了新方向。作者们建议未来的研究可能会探索：在更多任务和模型架构上验证AdEMAMix的有效性；研究AdEMAMix对模型泛化能力的影响；探索将AdEMAMix与其他优化技术（如学习率调度、梯度裁剪等）结合的方法。
总的来说，AdEMAMix为深度学习优化领域带来了新的思路，有望在未来的研究和应用中发挥重要作用。研究者们期待看到这种新优化器在更广泛的场景中的应用和进一步的改进。
