探索语言模型优化器的新视角:性能与稳定性的平衡

Deconstructing What Makes a Good Optimizer for Language Models

摘要

本文探讨了在语言模型训练中,不同优化算法(如SGD、Adafactor、Adam、Lion和Signum)的性能和稳定性。研究发现,除了SGD之外,其他优化算法在性能和超参数稳定性方面表现相似。文章进一步分析了Adam优化器的简化版本,如Signum和Adalayer,发现这些简化版本能够恢复Adam的性能和稳定性。研究结果表明,优化器的选择可以基于实际考虑,如内存限制和实现简易性,而不是仅仅依赖于性能或超参数稳定性。

原理

本文通过比较多种优化算法在自回归语言模型中的应用,发现除了SGD之外,其他算法如Adam、Adafactor、Lion和Signum在最佳性能和超参数选择稳定性方面表现相似。文章进一步探讨了Adam的简化版本,如Signum(一种带有符号动量的优化器)和Adalayer(一种逐层预处理的Adam变体)。通过实验证明,这些简化版本不仅能够恢复Adam的性能,还能在超参数选择上保持稳定性。特别是Adalayer的实验表明,尽管Adam的预处理对最后一层和LayerNorm参数有较大影响,但网络的其他大部分参数可以通过SGD进行训练。

流程

研究首先通过调整学习率和动量等超参数,对不同优化算法进行性能测试。随后,通过固定最佳学习率和动量,进一步测试其他超参数(如权重衰减、预热持续时间和批次大小)对算法稳定性的影响。实验结果显示,除了SGD外,其他优化算法在不同模型规模和架构变体中表现出相似的性能和稳定性。例如,在150m、300m和600m参数规模的模型中,Adam、Adafactor、Lion和Signum在不同学习率下的验证损失表现出高度一致性。

应用

本文的研究结果对语言模型训练中的优化器选择提供了新的视角。优化器的选择不再仅仅基于性能,而是可以更多地考虑实际应用中的内存限制和实现简易性。此外,研究揭示了简化版本的优化器(如Signum和Adalayer)在保持高性能的同时,还能减少对复杂超参数调整的依赖,这为未来在资源受限环境下的语言模型训练提供了新的可能性。