MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization

2026-06-16

AI Summary

This work addresses the lack of efficient selective parameter update mechanisms in large language model training that simultaneously offer fine-grained control and convergence guarantees. The authors propose MGUP, a method that, at each iteration, applies a larger learning rate to a fixed proportion of parameters while updating the remaining parameters with a smaller yet non-zero learning rate, thereby enabling fine-grained and convergence-controlled selective updates. MGUP provides the first theoretical convergence guarantee for momentum-based optimizers and is designed as a plug-and-play module that seamlessly integrates with mainstream optimizers such as AdamW, Lion, and Muon. Experimental results demonstrate that MGUP consistently achieves superior or more stable performance across diverse tasks, including MAE pretraining, large language model pretraining, and downstream fine-tuning.

Abstract

Efficient optimization is essential for training large language models. Although intra-layer selective updates have been explored, a general mechanism that enables fine-grained control while ensuring convergence guarantees is still lacking. To bridge this gap, we propose \textbf{MGUP}, a novel mechanism for selective updates. \textbf{MGUP} augments standard momentum-based optimizers by applying larger step-sizes to a selected fixed proportion of parameters in each iteration, while applying smaller, non-zero step-sizes to the rest. As a nearly {plug-and-play} module, \textbf{MGUP} seamlessly integrates with optimizers such as AdamW, Lion, and Muon. This yields powerful variants such as \textbf{MGUP-AdamW}, \textbf{MGUP-Lion}, and \textbf{MGUP-Muon}. Under standard assumptions, we provide theoretical convergence guarantees for \textbf{MGUP-AdamW} (without weight decay) in stochastic optimization. Extensive experiments across diverse tasks, including MAE pretraining, LLM pretraining, and downstream fine-tuning, demonstrate that our \textbf{MGUP}-enhanced optimizers achieve superior or more stable performance compared to their original base optimizers. We offer a principled, versatile, and theoretically grounded strategy for efficient intra-layer selective updates, accelerating and stabilizing the training of large-scale models. The code is publicly available at https://github.com/MaeChd/MGUP.

Problem

• selective updates
• stochastic optimization
• convergence guarantees
• large language models
• momentum-based optimizers

Innovation

• selective updates
• momentum-based optimization
• convergence guarantee
• plug-and-play module
• stochastic optimization