如何随模型和数据集规模缩放 AdamW 的权重衰减

随着我们试图构建更大的模型，最优 AdamW 权重衰减超参数随模型和数据集规模的缩放至关重要，但对此理解甚少。我们表明，AdamW 学习到的权重可以理解为近期更新的指数移动平均（EMA）。这为如何设置 AdamW 中的权重衰减以及权重衰减应如何随模型和数据集规模缩放提供了关键见解。特别地，指数移动平均的关键超参数是 EMA 时间尺度。直观上，EMA 时间尺度可以理解为 EMA 平均的近期迭代次数。我们发现，以 epoch 为单位的最优时间尺度在改变模型和数据集规模时大致恒定。此外，给定学习率，EMA 时间尺度与权重衰减超参数之间存在一一对应关系。因此，如果最优 EMA 时间尺度恒定，则意味着随着数据集规模增大，最优权重衰减应下降；随着模型规模增大，最优权重衰减应增加（如果我们遵循 muP 关于学习率缩放的推荐）。我们在 CIFAR-10 和 ImageNet 上训练的 ResNet-18 和 Vision Transformer 上，以及在 OpenWebText 上预训练的 NanoGPT 上验证了这些缩放规则。最后，我们发现随着训练进行，除非适当缩放权重衰减，否则 muP 的学习率缩放对于 AdamW 会失效。