MARS：释放方差缩减在大模型训练中的威力

训练深度神经网络——以及最近的大模型——需要高效且可扩展的优化器。自适应梯度算法如Adam、AdamW及其变体一直是这一任务的核心。尽管过去十年开发了许多旨在加速凸和非凸设置下随机优化的方差缩减算法，但方差缩减在训练深度神经网络或大语言模型中未获广泛成功。因此，它在现代AI中一直是一种不受青睐的方法。在本文中，为释放方差缩减在大模型高效训练中的威力，我们提出了一个统一优化框架MARS（Make vAriance Reduction Shine），该框架通过缩放随机递归动量技术，将预条件梯度方法与方差缩减相协调。在我们的框架内，我们引入了三个MARS实例，分别基于AdamW、Lion和Shampoo利用预条件梯度更新。我们还揭示了我们的算法与现有优化器之间的联系。在GPT-2模型训练上的实验结果表明，MARS始终以较大优势优于AdamW。MARS的实现可在https://github.com/AGI-Arena/MARS获取。