Research map / Optimization / Optimizers
MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MGUP, a plug-and-play selective-update mechanism that scales step sizes up for the top-K momentum-gradient aligned coordinates and down (but nonzero) for the rest, integrated into AdamW, Lion, and Muon, with convergence analysis for MGUP-AdamW.
Problem
Cautious optimizers zero out updates for momentum-gradient misaligned parameters, which lacks stochastic convergence guarantees and can discard gradient information; AdaBelief relies on second-moment estimates and does not transfer to Lion or Muon.
Approach
Rank coordinates by the element-wise product of momentum (or update) and gradient, give the top-tau fraction a step multiplier 1/tau and the rest a multiplier tau (nonzero), so the mechanism is a fixed-proportion, sign-agnostic selection that keeps small updates for all parameters.
Key results
- MGUP-AdamW reaches an average GLUE score of 85.15 when fine-tuning RoBERTa-base, versus 84.74 for AdamW.
- MGUP-AdamW achieves 34.96% validation accuracy on GSM-8K fine-tuning of LLaMA2-7B, versus 34.53% for AdamW.
- On LLaMA2-71M / WikiText-103, MGUP-AdamW reports a 1.6x speedup over AdamW, and MGUP-Lion a 2.5x speedup over Lion.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add MGUP to AdamW | AdamW · GLUE (Avg.) | score: 84.74 -> 85.15 | Helped |
| Add MGUP to AdamW | AdamW · GSM-8K | accuracy: 34.88 -> 34.96 | Helped |
| Add MGUP to AdamW | AdamW · GSM-8K | train loss: 0.057 -> 0.056 | Helped |
| Increase τ beyond a threshold (ablation) | MGUP-AdamW · WikiText-103 (LLaMA2-71M) | validation loss | Hurt |
| Increase γ while holding τ fixed (ablation) | MGUP-AdamW · WikiText-103 (LLaMA2-71M) | validation loss | Helped |
Surprising findings
- Setting gamma=0 (zeroing misaligned updates, as in Cautious Adam) can cause divergence on a constructed one-dimensional counterexample.
Limitations
- Convergence theory covers only MGUP-AdamW without weight decay; Lion and Muon variants lack theoretical guarantees.
- The fraction tau is a pre-selected hyperparameter requiring tuning.
- Theorem bounds rely on global gradient bound M and gamma>0 conditions that are sensitive to parameter choices.
Experimental setup
Datasets: CIFAR-10 WikiText-103 GLUE GSM-8K
Benchmarks: CIFAR-10 WikiText-103 GLUE GSM-8K
Metrics: validation loss training loss GLUE accuracy/correlation averages GSM-8K accuracy speedup
Baselines: AdamW Lion Muon C-AdamW C-Lion C-Muon GaLore LDAdam Adam-mini Adam-8bit
Builds on: Cautious Optimizers AdaBelief Adam Lion Muon
Closest papers by meaning
- ZO-AdaMU Optimizer: Adapting Perturbation by the Momentum and Uncertainty in Zeroth-Order Optimization (AAAI 2024 · similarity 0.89)
- Unleashing the Potential of Large Language Models as Prompt Optimizers: Analogical Analysis with Gradient-based Model Optimizers (AAAI 2025 · similarity 0.89)
- Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training (ICLR 2024 · similarity 0.88)
- Align-Pro: A Principled Approach to Prompt Optimization for LLM Alignment (AAAI 2025 · similarity 0.88)
- APT: Adaptive Pruning and Tuning Pretrained Language Models for Efficient Training and Inference (ICML 2024 · similarity 0.88)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.