MARS: Unleashing the Power of Variance Reduction for Training Large Models

训练深度神经网络——以及最近的大模型——需要高效且可扩展的优化器。自适应梯度算法如Adam、AdamW及其变体在这一任务中发挥了核心作用。尽管在过去十年中开发了多种方差减少算法，旨在加速凸优化和非凸优化中的随机优化，但方差减少在训练深度神经网络或大型语言模型方面并未取得广泛成功。因此，它在现代人工智能中一直是一种不太受欢迎的方法。在本文中，为了释放方差减少在大模型高效训练中的潜力，我们提出了一种统一的优化框架，即MARS（使方差减少发光），该框架通过缩放的随机递归动量技术将预条件梯度方法与方差减少相结合。在我们的框架内，我们介绍了三种基于AdamW、Lion和Shampoo的预条件梯度更新的MARS实例。我们还探讨了我们的算法与现有优化器之间的联系。实验结果表明，在训练GPT-2模型时，MARS在很大程度上始终优于AdamW。

论文试图解决的问题是在训练深度神经网络和大型语言模型时，如何有效地利用方差减少技术来提高优化器的性能。尽管方差减少算法在过去十年中得到了发展，但在深度学习领域并没有广泛成功应用。

论文的关键思路是提出一个统一的优化框架MARS（Make vAriance Reduction Shine），该框架通过引入一种缩放的随机递归动量技术，将预条件梯度方法与方差减少相结合。MARS框架下提出了三种实例，分别基于AdamW、Lion和Shampoo的预条件梯度更新方法。这一思路的新颖之处在于将方差减少技术与现代流行的自适应梯度算法结合起来，以期在大规模模型训练中取得更好的效果。

论文的亮点包括：1) 提出了一个通用的优化框架MARS，能够兼容多种自适应梯度算法；2) 实验结果显示，MARS在训练GPT-2模型时显著优于AdamW；3) 论文不仅介绍了MARS的理论基础，还详细描述了其与现有优化器的关系；4) 作者提供了实验代码，便于其他研究者复现和进一步探索。

最近在这个领域中，相关的研究还包括：1) 《Adaptive Gradient Methods with Dynamic Bound of Learning Rate》（2019年），探讨了自适应梯度方法中的学习率动态调整策略；2) 《Variance Reduced Stochastic Gradient Descent with Neighbors》（2015年），提出了利用邻居信息减少方差的方法；3) 《On the Convergence of Adam and Beyond》（2018年），分析了Adam及其变种的收敛性。