Hybrid momentum methods improve deep learning training performance
Hybrid Momentum Methods for Optimizing Deep Learning Models
Machine Learning
Summary
Training deep learning models involves adjusting them based on gradients, which show how to improve. Traditional momentum methods focus mostly on recent gradients, sometimes ignoring older useful information. The authors introduce a hybrid approach that combines fast attention to recent changes with a longer view of past gradients, aiming to get the best of both. Their method improves training results on tasks like image recognition, language translation, and large language model training.
What this means in practice
- •For deep learning engineers: Optimize training of neural networks for tasks like image and language processing by integrating hybrid momentum methods to improve convergence and final accuracy.
- •For machine translation developers: Enhance neural machine translation model training efficiency and accuracy using hybrid momentum algorithms combined with orthogonalization techniques.
Authors
Haifeng Xu, Yijin Ren, Qi Deng, Chenyu Zhou, Dongdong Ge, Yinyu Ye
Abstract
Momentum is central to stochastic optimization in deep learning, including the training of large language models. Classical momentum methods typically rely on a fast exponential moving average (EMA) scheme, which discounts past gradients on a fixed timescale and may underweight older gradients that remain informative. This creates a trade-off between retaining long-term history and adapting to recent changes. We introduce Hybrid Momentum (HyM), which combines a linearly weighted average of past gradients with a geometrically discounted fast momentum. The former assigns weights proportional to the gradients' iteration indices and retains a growing span of history, while the latter emphasizes recent gradients. This hybrid construction is modular and can be combined with adaptive preconditioning and other update transformations. In particular, we propose HyMUON, which applies Newton--Schulz orthogonalization to the hybrid direction. Illustrative examples and controlled experiments demonstrate the benefits of combining gradient information across these two timescales. Across image classification, machine translation, and language modeling, HyM and HyMUON match or improve upon strong baselines. HyMUON improves over Muon on IWSLT14, and in both the 100k- and 600k-step GPT-2 124M OpenWebText regimes. For LLaMA-1.3B pretraining on C4, HyMUON achieves a final validation loss of 2.777, compared with 2.903 for Muon.