How to set AdamW's weight decay as you scale model and dataset size

我们展示了AdamW学习的权重可以理解为最近更新的指数移动平均(EMA)。这为如何设置AdamW中的权重衰减以及权重衰减如何随着模型和数据集大小的变化而变化提供了关键见解。特别地，指数移动平均的关键超参数是EMA时间尺度。直观地，EMA时间尺度可以理解为EMA平均数覆盖的最近迭代次数。给定固定的学习率，EMA时间尺度与通常的权重衰减超参数之间存在一一映射。因此，选择EMA时间尺度隐含地设置了权重衰减。重要的是，有自然的指导方针来选择合理的EMA时间尺度：我们需要对所有数据点进行平均，因此EMA时间尺度不应该比一个epoch小(太多)，我们需要忘记早期的更新，因此EMA时间尺度不应该比所有训练epoch的总数大(太多)。在我们的实验中，我们发现最佳的EMA时间尺度与这些指导方针一致，最近大规模LLM预训练运行中选择的超参数也一致(例如Llama 1+2和Stable LM)。关键是，这些指导方针表明，随着模型和数据集的扩大，最佳的EMA时间尺度不应该改变(太多)。这意味着随着数据集大小的增加，最佳的权重衰减应该降低。此外，随着模型大小的增加，如果我们遵循muP建议缩放学习率，最佳的权重衰减也应该增加。

AdamW算法中的权重衰减如何设置才能更好地适应不同的模型和数据集大小？

AdamW算法中的权重衰减可以理解为最近更新的指数移动平均值，因此可以通过EMA时间尺度来设置权重衰减，EMA时间尺度可以理解为EMA平均值覆盖的最近迭代次数。EMA时间尺度的选择应该考虑到需要覆盖所有数据点，同时遗忘早期的更新。根据实验结果，最优的EMA时间尺度符合这些指导方针，并且随着数据集大小的增加，最优的权重衰减应该降低，随着模型大小的增加，最优的权重衰减应该增加。

论文通过理解AdamW算法中的权重衰减为指数移动平均值，提出了EMA时间尺度来设置权重衰减，同时给出了EMA时间尺度的指导方针。实验结果表明这些指导方针是合理的，并且最优的EMA时间尺度随着数据集和模型大小的变化而变化。论文的实验使用了多个数据集和模型，并且给出了开源代码。

在最近的相关研究中，也有一些关于AdamW算法中权重衰减的研究，例如《Fixup Initialization: Residual Learning Without Normalization》和《On Weight Decay and Small Learning Rates》。