MARS: Unleashing the Power of Variance Reduction for Training Large
  Models
Authors: name: Huizhuo Yuan; name: Yifeng Liu; name: Shuang Wu; name: Xun Zhou; name: Quanquan Gu
Published: 2024-11-15 18:57:39

Abstract: 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.