Title: MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization
Venue: NeurIPS 2025
Accepted: yes
Keywords: Stochastic Optimization; Convergence Analysis; Adaptive Optimizers ; Nonconvex Optimization; Large Language Models; Momentum-Gradient Alignment
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.


Peer reviews:
comments: SUMMARY:
This paper proposes a new monentum-gradient alignment policy for stochastic optimization. 
It provides theoretical guarantees of the proposed methods, and evaluates the performance on a few benchmark NN training tasks.

STRENGTHS AND WEAKNESSES:
Strength: This paper have a clean introduction of the main idea -- I enjoy reading the paper a lot. It has sound theoretical guarantees for the proposed algorithm, and demonstrates the efficiency with extensive numerical experiments.

Weakness: As mentioned by the authors, the proposed algorithm combines the idea of a few existing methods. The contribution is a bit incremental.

QUESTIONS:
1. In equation (3), the scaling factor is 1/tau if signs are aligned and tau if signs are not aligned. What's the intuition here to choose these scaling factors as the reverse of each other, instead of choosing in a more empirical/data-driven way?

2. In Assumption 4.2, it is assumed that function f is L-smooth. Is it possible to show a convergence result for with relaxed (non-smooth) conditions? For example for relu activation. 

3. In Theorem 4.2, it is better to give a specific probability upper bound (at least mention it in a remark) instead of just stating it is 'with high probability'. It can help the reader understand how the probability bound depends on the parameters.

LIMITATIONS:
See weakness and questions.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-21

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a new plug-in method, MGUP, for improving momentum updates. This new strategy can be viewed as a smoother version of the cautious optimizer. Theoretical convergence guarantees and empirical results are provided to show the effectiveness of the algorithm.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:\n1. The paper explains the idea well, and to smooth the cautious optimizer sounds to be a reasonable idea.\n\n2. The paper provides theoretical proof for the convergence of the algorithm and illustrates the importance of employing a nonzero update even when the gradient and momentum are not well aligned.\n\n**Weakness**:\n1. The theoretical illustration for using nonzero $\\gamma$ doesn't seem to be convincing enough. We basically need a lower bound, or maybe even a simple concrete case study for why setting $\\gamma=0$ can lead to unstable performance of the algorithm. Also, I think it would be better to also talk about work related to the convergence theory.\n\n2. The empirical improvement doesn't seem to be significant enough, especially for the language model tasks. The improvement upon the baseline is not that significant while the experiment scale is not large, at least insufficient to justify the effectiveness of the approach. This is important because we need strong empirical evidence to support MGUP, since the motivation of the MGUP idea isn't that convincing theoretically. I would be happy to raise the score if you could explain the benefits of doing MGUP compared to the cautious optimizer and the original Adam.\n\nQUESTIONS:\n1. I am concerned with the computation burden of MGUP, since it requires sorting during the update, which basically has $O(d\\log d)$ complexity. \n\n2. MGUP-Muon seems to have a different implementation compared to MGUP-Lion and MGUP-Adam, since it does MGUP to compute $\\phi_t$ before computing the update $u_t$. Could you explain this difference?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an adaptive mechanism that adjusts the step size based on the alignment of the momentum and the gradient. The paper provides a convergence analysis of the proposed method, MGUP, with Adam under mild conditions. The theoretical analysis also shows the necessity of keeping all step sizes non-zero for convergence. Numerical experiments demonstrate that the proposed method outperforms its baseline algorithms without MGUP.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\n1. The paper provides a rigorous convergence analysis for the Adam-type algorithm with different step sizes.\n2. Numerical results on ViT training and LLaMa fine-tuning show that the proposed algorithm outperforms its baselines without MGUP.\n\nWeakness:\n1. On the necessity of using $\\gamma >0$. Remark 4.1 and 4.3 only show that under the current analysis, $\\gamma>0$ is a necessary condition. However, it is not sufficient to show whether $\\gamma>0$ is necessary due to the analysis or the algorithm design. If the authors can provide an example that demonstrates that choosing $\\gamma=0$ leads to non-convergence, the claim would be significantly strengthened.\n2. Regarding remark 4.2, the authors failed to provide further discussion on the existence of such an $M_T < \\infty$. Since the convergence results only show that with high probability, or in expectation, the gradient norm is diminishing. Proving that the per-sample gradient is bounded along the trajectory is also a non-trivial extension.\n3. The experiments failed to report the standard deviation. For LLaMA2 and Qwen2.5 training tasks, the accuracy and loss are relatively close to each other. It is challenging to determine whether the performance improvement is due to the random seed or if there is a statistically significant difference.\n4. The analysis uses constant $\\alpha$ and $\\gamma$, independent of $s=u_t g_t$. It does not explain why MGUP has better performance. Also, it does not provide insight on how to choose $\\phi_i$ to accelerate the convergence of MGUP.\n\nMinor point: \n4. Why choose $\\alpha = 1/\\tau$ and $\\gamma = \\tau$? Is there any specific reason? It is also a bit confusing since $\\gamma, \\alpha$ do not appear in Algorithms 1 and 2.\n\nQUESTIONS:\nPlease address the weakness.\n\nIt is unclear to me how the theoretical analysis reflects the unique properties of the proposed algorithm.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MGUP, a novel mechanism designed to enhance momentum-based optimizers like AdamW, Lion, and Muon. The core idea is to selectively and differentially adjust learning rates within each layer. Specifically, for a fixed proportion of parameters with the highest momentum-gradient alignment score, MGUP applies an amplified step size. For the remaining parameters, it applies a smaller, nonzero step size. The authors provide theoretical convergence guarantees for a variant, MGUP-AdamW, under standard stochastic non-convex optimization assumptions. Through extensive experiments on tasks including image pretraining, language model pretraining, and fine-tuning, the paper demonstrates that MGUP-enhanced optimizers consistently achieve faster convergence and/or better generalization compared to their original counterparts and other state-of-the-art methods.\n\nSTRENGTHS AND WEAKNESSES:\nPros:\n\n- The paper is well-motivated by the limitations of existing selective update methods. It identifies a key problem with methods like Cautious Optimizers, which can discard too much gradient information by nullifying updates (setting step size to zero) for misaligned parameters. MGUP's approach of applying a decayed, non-zero step size (γ>0) is a simple yet principled solution that is well-justified by the theoretical analysis.\n\n- The empirical evaluation is thorough and provides strong evidence for MGUP's effectiveness. \n\n- The design of MGUP as a lightweight, nearly plug-and-play module that can be integrated with various optimizers is a major practical advantage. The paper successfully demonstrates this by creating and evaluating MGUP-AdamW, MGUP-Lion, and MGUP-Muon.\n\nCons:\n\n- MGUP introduces a key hyperparameter, τ, which defines the proportion of parameters to receive amplified updates. While the authors provide a sensitivity analysis (Figure 3b), the choice of τ appears to be data and model-dependent, which could add to the tuning burden for practitioners.\n\n- he convergence guarantees are formally provided for MGUP-AdamW (without weight decay). While the empirical results for MGUP-Lion and MGUP-Muon are strong, the paper notes that their theoretical properties require further study. This limits the theoretical scope of the paper's claims, though the empirical success across different optimizers is a mitigating factor.\n\nQUESTIONS:\nRegarding the hyperparameter τ, have you identified any general heuristics for its selection? For instance, does its optimal value correlate with the task (e.g., pretraining vs. fine-tuning), model size, or the base optimizer being used?\n\nLIMITATIONS:\nI do not see any potential negative societal impact in addition to the development of LLMs themselves.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}