提出MARS框架,让方差减少技术在大模型训练中高效可用。
MARS: Unleashing the Power of Variance Reduction for Training Large Models
- 用缩放随机递归动量统一方差减少与预条件梯度方法。
- 在训练GPT-2时,MARS显著优于AdamW,性能提升明显。
- 适配AdamW、Lion、Shampoo,适合追求高效训练的大模型研究者。
训练深度神经网络,尤其是大模型,需要高效可扩展的优化器。自适应梯度算法如Adam、AdamW及其变体已成为主流。尽管过去十年中发展了众多方差减少算法,旨在加速凸与非凸场景下的随机优化,但这些方法在深度神经网络或大语言模型训练中尚未取得广泛应用,因而长期未被现代AI领域重视。为释放方差减少在大模型训练中的潜力,本文提出统一优化框架MARS(Make vAriance Reduction Shine),通过缩放随机递归动量技术,将预条件梯度方法与方差减少相融合。框架内包含三个实例,分别基于AdamW、Lion和Shampoo的预条件梯度更新。我们还揭示了所提算法与现有优化器之间的联系。在训练GPT-2模型的实验中,MARS始终以显著优势超越AdamW。代码已开源:https://github.com/AGI-Arena/MARS。
原文摘要 · Abstract (English)
Training deep neural networks--and more recently, large models demands efficient and scalable optimizers. Adaptive gradient algorithms like Adam, AdamW, and their variants have been central to this task. Despite the development of numerous variance reduction algorithms in the past decade aimed at accelerating stochastic optimization in both convex and nonconvex settings, variance reduction has not found widespread success in training deep neural networks or large language models. Consequently, it has remained a less favored approach in modern AI. In this paper, to unleash the power of variance reduction for efficient training of large models, we propose a unified optimization framework, MARS (Make vAriance Reduction Shine), which reconciles preconditioned gradient methods with variance reduction via a scaled stochastic recursive momentum technique. Within our framework, we introduce three instances of MARS that leverage preconditioned gradient updates based on AdamW, Lion, and Shampoo, respectively. We also draw a connection between our algorithms and existing optimizers. Experimental results on training GPT-2 models indicate that MARS consistently outperforms AdamW by a large margin. The implementation of MARS is available at https://github.com/AGI-Arena/MARS.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。