Saved in:
Bibliographic Details
Main Authors: Yuan, Huizhuo, Liu, Yifeng, Wu, Shuang, Zhou, Xun, Gu, Quanquan
Format: Preprint
Published: 2024
Subjects:
Online Access:https://arxiv.org/abs/2411.10438
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915478169976832
author Yuan, Huizhuo
Liu, Yifeng
Wu, Shuang
Zhou, Xun
Gu, Quanquan
author_facet Yuan, Huizhuo
Liu, Yifeng
Wu, Shuang
Zhou, Xun
Gu, Quanquan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2411_10438
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle MARS: Unleashing the Power of Variance Reduction for Training Large Models
Yuan, Huizhuo
Liu, Yifeng
Wu, Shuang
Zhou, Xun
Gu, Quanquan
Machine Learning
Optimization and Control
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.
title MARS: Unleashing the Power of Variance Reduction for Training Large Models
topic Machine Learning
Optimization and Control
url https://arxiv.org/abs/2411.10438