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

Abstract: Efficient optimization is essential for training LLMs. 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.

• The paper introduces MGUP, a policy that modulates parameter updates via momentum-gradient alignment to ensure stable convergence in stochastic non-convex settings.
• MGUP seamlessly integrates with AdamW, Lion, and Muon by selectively amplifying updates, backed by rigorous theoretical convergence guarantees.
• Empirical results show significant speedups and state-of-the-art performance across tasks from image pretraining to fine-tuning benchmarks like GLUE.

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

Motivation and Context

Optimization in deep neural networks, specifically for LLMs and large-scale vision and LLMs, has driven the pursuit of methods that enable both memory efficiency and fast, stable convergence. Prevailing approaches leverage low-rank structure and selective parameter updates, including gradient projection (e.g., GaLore [DBLP:conf/icml/Zhao0CWAT24]) and sparse update strategies (e.g., SIFT [DBLP:conf/icml/SongLZ0024]), but most operate at a coarse (layer or block) granularity. Fine-grained, theoretically justified, and plug-and-play intra-layer selective update policies remain absent, especially for optimizers outside Adam-style methods.

The MGUP framework addresses these deficiencies by introducing a selection methodology based on momentum-gradient alignment, generalizable to Adam, Lion, and Muon optimizers, with rigorous convergence properties in stochastic non-convex settings. Prior efforts such as AdaBelief and Cautious Optimizer are either limited by their dependence on second-moment estimation or lack theoretical convergence guarantees under stochasticity, and can fail in practice when misaligned updates are naively zeroed.

The MGUP Policy

The central innovation of MGUP lies in its parameter-wise modulation of step sizes based on the instantaneous elementwise alignment between momentum and stochastic gradient.

Figure 1: The key idea of MGUP involves adaptively adjusting the learning rate by leveraging the element-wise product of the stochastic gradient and momentum.

Specifically, given alignments $s_{t,i} = m_{t,i} g_{t,i}$ (where $m_{t,i}$ and $g_{t,i}$ are the $i$th element of the momentum and gradient, respectively), MGUP sorts these scores and selects a proportion $\tau$ of coordinates (top-K) for an amplified step ($\alpha \cdot \eta_t$), with the rest updated with a nonzero decayed step ($\gamma \cdot \eta_t$), typically $\alpha = 1/\tau$ and $\gamma = \tau$. This policy prevents the over-suppression or instability emergent in zero-masked strategies (i.e., $\gamma = 0$ as used by Cautious Optimizer [Liang2024CautiousOI]).

Algorithmically, MGUP is seamlessly inserted into AdamW, Lion, and Muon, resulting in MGUP-AdamW, MGUP-Lion, and MGUP-Muon, without requiring modification to core adaptive logic.

Theoretical Analysis

MGUP-AdamW is analyzed under standard non-convex stochastic optimization assumptions (objective bounded below, $m_{t,i}$0-smoothness, unbiased stochastic gradients with bounded variance). Importantly, contrary to prior approaches, the analysis explicitly demonstrates that nonzero $m_{t,i}$1 for non-selected parameters is essential to guarantee convergence:

• MGUP-AdamW achieves a convergence rate of $m_{t,i}$2 in expectation, matching the best-known for adaptive optimizers.
• High-probability guarantees are also attained under coordinate-wise variance control, with $m_{t,i}$3 rates.

This analysis reveals that mask-based optimizers with strict hard masking (i.e., zero updates for misalignment) fundamentally degrade stochastic optimization guarantees; small but nonzero updates ensure retained progress towards optimum (see Appendix and counterexample discussion).

Figure 2: Analysis of C-Adam's failure and MGUP-Adam's success in the counterexample. (a) MGUP-Adam converges to the optimum $m_{t,i}$4 while C-Adam diverges. (b) The momentum in C-Adam oscillates around zero, leading to unstable update decisions. (c) A scatter plot of update steps shows C-Adam either skips updates ($m_{t,i}$5) or updates in the wrong direction ($m_{t,i}$6), whereas MGUP-Adam consistently updates in the correct direction ($m_{t,i}$7).

Empirical Results

MGUP-enhanced optimizers are validated over both pretraining and supervised fine-tuning regimes spanning image and text domains.

• ViT MAE Pretraining (CIFAR-10): MGUP-AdamW delivers consistently lower training and validation loss vs AdamW and C-AdamW.

Figure 3: ViT MAE training and validation curves on CIFAR-10.

• Autoregressive pretraining (WikiText-103): For LLaMA2-71M, MGUP-AdamW achieves 1.6x speedup in convergence compared to AdamW and shows improved generalization over C-AdamW. MGUP-Lion achieves 2.5x speedup versus Lion, and MGUP-Muon modestly outperforms standard Muon, notably stabilizing the training curve.

Figure 4: LLaMA2-71M validation curve and MGUP-AdamW sensitivity analysis on WikiText-103.

• Scaling to Qwen2.5-150M: MGUP-AdamW and MGUP-Muon again manifest accelerated and more stable convergence on both training and validation loss curves versus baselines.

Figure 5: Qwen2.5-150M training and validation curves on WikiText-103.

• Fine-tuning on GLUE benchmark (with RoBERTa-base): MGUP-AdamW achieves a new state-of-the-art average GLUE score (85.15), outperforming AdamW, Lion, Muon, and competitive memory-efficient optimizers for both average and most individual tasks.

Figure 6: Adamw-type, Lion-type, Muon-type optimizers average performance across GLUE tasks.

• GSM8K with LLaMA2-7B: MGUP-AdamW delivers the lowest fine-tuning loss and highest valid accuracy, confirming robustness in complex reasoning tasks.

Robustness and Sensitivity

Comprehensive sensitivity analysis demonstrates that, for MGUP-AdamW, larger $m_{t,i}$8 generally improves performance at fixed $m_{t,i}$9, and excessive $g_{t,i}$0 degrades accuracy. The recommended default ($g_{t,i}$1) maximizes the average effective learning rate and empirical convergence, consistent with the theoretical analysis regarding the average step increase (see Section 3 of the paper).

Implementation and Memory/Runtime Profile

While MGUP introduces transient peak memory overhead—by requiring auxiliary storage for selection and alignment scores—static allocation is unchanged. Training runtime overhead is negligible for small- and medium-scale models, with increments for very large models dominated by top-K selection cost; these are offset by the improved convergence speed and generalization, especially for large-scale LLM pretraining and adaptation.

Broader Implications and Future Directions

Practically, MGUP provides an optimizer-agnostic, easily integrable mechanism for dynamic parameter selection, with implications for scaling LLMs and vision transformers under resource-constrained settings. Theoretically, it resolves key limitations of prior selective update schemes and broadens the applicability of momentum-gradient alignment in stochastic non-convex regimes. The key insight—that small, nonzero updates on non-selected parameters preserve convergence—is relevant for the future development of sparse, modular, or distributed optimization strategies.

Key open questions include the development of adaptive policies for $g_{t,i}$2, understanding MGUP's properties for more complex optimizers (especially those lacking second-moment estimation), and further characterizing the landscape of sparse selection—potentially blending rank-adaptive or quantized updates with MGUP-like mechanisms.

Conclusion

MGUP provides a justified, general, and empirically validated plug-and-play intra-layer selective update policy for momentum-based optimizers. It succeeds in accelerating convergence and improving generalization—especially in large-scale, non-convex stochastic optimization settings—while preserving theoretical soundness. The framework merits consideration for both fundamental optimization research and practical deployment in memory- and compute-constrained deep learning.

For full details and code: (2606.17526).