超级预训练优化器与寻找之路

AdamW 长期以来是语言模型预训练中占主导地位的优化器，尽管已有许多声称其它优化器可实现 1.4 至 2 倍的加速。我们认为，两个方法论缺陷掩盖了公平比较并阻碍了实际采纳：（i）不平等的超参数调谐，（ii）受限或误导性的评估设置。为解决这两个问题，我们在四个模型规模（0.1B-1.2B 参数）和数据-模型比率（1-8 倍的 Chinchilla 最佳比例）上系统性地研究了十个深度学习优化器。我们发现，进行严格的超参数调谐和在各种模型规模和数据-模型比率上的评估（在训练结束时进行）对于进行公平且有信息的比较至关重要。首先，一个优化器的最优超参数可能对另一个优化器而言是次优的，盲目地转移超参数是不公平的。其次，许多新提出的优化器相对于经过良好调谐的基准的实际加速速度低于声称的，并且随着模型规模的增加，加速速度仅为 1.1 倍（针对 120 亿参数模型）。第三ly，在训练完成目标训练预算之前比较中间检查点可能具有误导性，因为由于学习率衰减，两个优化器之间的排名可能会在训练期间翻转。通过我们的彻底调查，我们发现所有最快的优化器，如 Muon 和 Soap，都使用矩阵作为预条件器——即对梯度与条目级标量相乘，而不是使用标量。然而，矩阵式优化器的加速速度与模型规模成反比，从 0.1B 参数模型相对于 AdamW 的 1.4 倍加速，降至 1.2B 参数模型的仅 1.1 倍。