提出一种动态调整参数更新步长的优化策略,提升大模型训练效率与稳定性。
MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization

- 按比例选择部分参数使用更大步长,其余参数用小步长更新
- 在多个任务中表现优于原始优化器,且保证收敛性
- 可无缝接入AdamW等主流优化器,适合大规模模型训练
高效优化对训练大语言模型至关重要。尽管已有研究探索层内参数选择性更新,但缺乏能实现细粒度控制并保证收敛性的通用机制。为此,我们提出一种新方法MGUP,通过在每轮迭代中为固定比例的参数分配更大步长,其余参数则采用较小但非零步长,增强标准动量优化器的性能。该方法几乎可即插即用,可集成至AdamW、Lion和Muon等优化器,形成如MGUP-AdamW、MGUP-Lion等变体。在标准假设下,我们为无权重衰减的MGUP-AdamW提供了理论收敛性证明。大量实验覆盖MAE预训练、大模型预训练及下游微调任务,结果表明,经MGUP增强的优化器在性能或稳定性上均优于基线。本工作提供了一种有理论支撑、灵活通用的层内选择性更新策略,有效加速并稳定大规模模型训练。代码已开源:https://github.com/MaeChd/MGUP。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。