MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization
Authors: Da Chang
Ganzhao Yuan
Venue: NeurIPS 2025 spotlight ()
Keywords: Stochastic Optimization; Convergence Analysis; Adaptive Optimizers ; Nonconvex Optimization; Large Language Models; Momentum-Gradient Alignment
Primary area: optimization

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.

MGUP: A Momentum-Gradient Alignment Update Policy for Stochastic Optimization
Da Chang 134 Ganzhao Yuan
Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences Shenzhen University of Advanced Technology. "Pengcheng Laboratory University of Chinese Academy of Sciences
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 MGUP, a novel mechanism for selective updates. MGUP augments standard momentum-basee 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 nearly plug-and-pla module, MGUP seamlessly integrates with optimizers such as AdamW, Lion, and Muon. This yields powerful variants such MGUP-AdamYs
MGUP-Lion, and MGUP-Muon. Under standard assumptions, we provide theoretical convergence guarantees for MGUP-Adam W (without weight decay) in stochastic optimization. Extensive experiments across diverse tasks, including MAE pretraining, LLM pretraining, and downstream fine-tuning, demonstrate that our MGUP-enhande optimizers achieve superior or more stable performance compared to their original base optimizers. We offer 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 wtttpps: :// //gi thub. Com/MaeChd/MGUS
Introduction
Recent studies reveal that the learning matrix during Large Language Model (LLM) training exhibits low-rank properties, suggesting that learning predominantly occurs in a low-dimensional space [ This observation has catalyzed the development of methods such Galore 3 and LDAdam 4
which use gradient low-rank decomposition to achieve performance comparable to full-rank updates while reducing memory consumption. Although low-rank properties do not directly imply sparsity, the insight that optimization occurs in a low www-diensional space provides a crucial foundation for selective parameter updates. This principle is exemplified by SIFT (5), which achieves efficient adaptation through gradient- based sparse parameter updates, leveraging the low intrinsic dimensionality and sparse gradient characteristies inherent in I.MMs. Building this foundation, several innovative layer-wise selective update methods have emerged, including AutoFreeze [6], LOMO [7]. LISA (8), and BAdam (9. By strategically freezing certain layers while updating others, these methods achieve performance comparable to, or even surpassing, that of full full-parameter updates.
While layer-wise selective updates show promise, finer mrainer parameter selection remains under- explored. Although SIFT [5] investigates sparse intra-layer updates, systematic methodology for identifying the most critical parameters within each layer is still lacking. This research gap motivates the development of novel intra-layer sparse update strategies.
*Corresponding author: yuanganchao@formail.co
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Recently, Liang et al. [T0] have proposed Cautious Optimizers, a novel intra-layer sparse update strategy. This approach selectively updates only parameters where momentum and gradient are aligned (i.e., m, O g: > 0), enabling larger updates for aligned directions while skipping misaligned ones. Conceptually, extends earlier adaptive optimizers like AdaBelief [1], which adjusts step sizes using (m, gx)², 2 but introduces parameter selection based on momentum-gradient alignment.
However, both methods have notable limitations. AdaBelief's update mechanism relies heavily on Adam's second-moment estimatiion, which restricts its applicability to optimizers that do not compute second moments (e.g., Lion [2] Muon [3] Furthermore, Cautious Optimizers lack rigorous theoretical convergence guarantees in the stochastic setting. Although the strategy offers theoretical insights in the deterministic case, its convergence properties remain unverified under stochastic conditions. This raises a crucial question:
Within the stochastic optimization setting, can the concept of intra-layer sparsity in updates, based on momentum- gradient
-errennnnn-rrdiennt direction consistency, truly serve plug-and-play mechanism?
If so, what are the boundaries of its effectiveness? If not, what are the underlying reasons?
We explore this issue in detail in the theoretical analysis presented in Section + Specifically, we demonstrate that for Adam variants incorporating a mask, simply setting the update step to zero for parameters where momentum and gradient directions misaligned significantly impacts the convergence properties of stochastic optimization. This motivates rethinking how to perform selective parameter updates more effectively in stochastic optimization settings to maintain favorable convergence properties. For example, without guided parameter selection, certain extreme cases occur: (i) only small fraction of parameters receive substantial updates (potentially leading to unstable training), or (ii the updates for the vast majority of parameters are overly suppressed (potentially resulting in slow training). Therefore, we propose that a promising policy involves not only considering the alignment between momentum and gradient direction but also regulating the proportion of parameters receiving substantial versus minor updates, strike balance between training efficiency and stability.
Motivated by theoretical analysis and resulting design considerations, introduce novel selective update method: MGUP (Momentum- Gradient alignment Update Policy). MGUP updates parameters selectively and differentially by sorting the values of the element-wise product m, @ ge- Specifically, the top parameters ranked by m, @ ge receive a scaled step size a (a > 1), while the rest receive Y T/ (y < 1), where 7). is the base step size from the original optimizer. MGUP is inspired by the cautious update strategy, refining it in line with the principles of AdaBelief and Cautious Optimizers by dynamically adjusting update strength based on momentum-gradien alignment.
Our contributions summarized follows:
We develop a novel selective parameter update mechanism, MGUP, which assigns larger step sizes to a subset of parameters and smaller ones to the rest. As a allg-and-play mechanism, MGUP can be integrated into momentum-bases optimizers such as AdamW, Lion, and Muon, yielding variants we refer to as MGUP-Adamt MGUP-Lion, and MGUP-Muon.
We establish the convergence of the Adam optimizer with the MGUP mechanism in the stochastic setting, providing theoretical guarantees for its reliability.
We validate the proposed MGUP optimizers through
key experiments, including MAE pretraining of VYT-27M on CIFAR-10: autoregressive pretraining of LLaMMA-71M and Qwen2.5-150M on Wikitext- 1033; and fine-tuning of RoBERTa-base on GLUE and LLaMA2-7B for GSM-8K. These results show the robustness and versatility of MGUP across diverse models and tasks.
2
Related Work
In this section, we review the basic principles of stochastic optimization methods relevant to the momentum-gradden approach. We consider minimizing the objective function as follows:
min fxx) where f(x) Eg [x Rd
(I)

---

Here, f Rd + R is differentiable and possibly nonconvex function, & represents a random vector, such as training data point, sampled from an unknown data distribution D.
In the context of solving problem ① momentum- based methods are foundational in large-scale machine learning optimization, accumulating past gradient information to accelerate convergence and navigate complex loss landscapes. The standard momentum update, an exponentially weighted moving average (EWMA) of gradients, is given by
m, 3 m (1 B1)g
where 31 is the decay factor, m, denotes the momentum vector, and ge denotes the gradient the t-th iteration. This technique smooths gradient estimates, empirically and theoretically accelerating convergence and enhancing training stability [4
T3 16 17].
While standard momentum is a robust baseline, research has sought to improve it, primarily through: (i) reducing stochastic gradient estimate variance and (ii) adapting learning based on momentum and gradient characteristics.
Variance reduction techniques, such as SPIDER 188. STORM [9), SUPER-ADAM [20], and MARS [21]. operate by substituting the original stochastic gradient gr with a gradient estimator gi that exhibits lower variance. This refined estimator is then used in the momentum update: m, = -m;-1 + (1 - g) While these methods theoretically accelerate convergence, they often necessitate additional computation storage (e.g., storing past gradients). In contrast, MGUP adopts a distinct strategy, focusing on adaptively adjusting the update magnitude based on the characteristics of momentum and the current stochastic gradient, rather than directly altering the variance of the gradient estimation.
Another significant method involves adapting the optimization step based on the perceived reliability or characteristics of the momentum estimate. The intuition guiding this class of methods can be summarized as:
Increase step size for trustworthy momentum; Decrease step size for untrustworthy momentum.
This adaptation is often implemented by modulating the momentum vector, which can be represented generally as:
X X, m) C Φi.
where de is a scaling factor, often applied element-wise, determined by gradient statistics.
Early adaptive methods, like Adagrad (22), introduced per- per-paraeeter learning rates by accumulating squared gradients. The widely adopted Adam optimizer (23) builds on this by using EWMAs for both the first moment m, and the second moment v. of the gradients:
V2 =32 + (1 - 82)g2
The update step is then element wise scaled by 1/ 1/vV V.V.+c, f with v, being a bias-corrected v, and é >> 0 is small constant. This enables Adam to adapt the learning rate per parameter based on historical gradient magnitudes. Subsequent research delved into various scaling factors, frequently investigating the interplay between the current gradient gi and the accumulated momentum m,, The AdaBelief optimizer [II] modifies Adam's 's second moment by using the squared difference between momentum and the current gradient, (m, g²)² instead of the raw squared gradient g?. The update rule for the second moment V2 is as follows, with the initial condition Vo = 0:
Vt = -2V.-1 + (1 32) - gz) = (1 - B2) Ei f -
- g
The term (m, )² measures "belief" in the current gradient by its consistency with momentum. Significant deviation increases the corresponding element in V,, reducing that parameter's 's effective step size. This mechanism aims to merge Adam's 's rapid convergence with SGD's generalization. Denote m,,; and gt, as the i-th elements of the momentum vector m, and gradient vector gi- respectively. m,,; and gt, have different signs, (m, )2 is typically larger than g², (for similar magnitudes), increasing Vr, and adaptively decreasing the step size. Meanwhile, a more direct approach to leveraging the sign consistency between momentum and gradient is taken by the Cautious Optimizers [TO]. It employs an element-wise mask y to selectively apply momentum updates:
PP I[m, C ge 0),
Xx+1 = Xx - 7, O Pr.
3

---

Here, II(-) is the indicator function, which equals when its argument is positive and otherwise. If the signs m,, and gr,; are align, the momentum component my,i is scaled by a otherwise, the update for that component is nullified. This "Cautious Updating" strategy aims to prevent updates from potentially conflicting gradient information.
However, these advanced adaptive methods have notable limitations. AdaBelief's reliance on second-moment estimation restricts its applicability primarily Adam-style optimizers, thereby rendering it incompatible with newer methods like Lion [2] and Muon [13] that perform well without this component. The Cautious Optimizer, while more broadly applicable, lacks formal stochastic convergence guarantees. As analyzed in Section 4 its binary masking mechanism can aggressively discard gradient information. This behavior may slow convergence, especially in scenarios where the signs of the momentum and gradient align infrequently. Furthermore, Cautious Adam exhibits non- convergent behavior, highlighting critical flaw its design (see the counterexample Appendix
f(x)
f(x)
I -/ di 1/3"
D= V/le' (5)) mi s0 m/ %%
Increuse stop for trussworthy momentum
Decrense step untrustwortyy momention
m² -7/
m1
m' N
x
x2
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.
3 The Proposed Method
This section introduces the MGUP Mlomentum-Grrdicnn alignment Update Policy) mechanism for solving Problem 1I. Our motivation is to address limitations observed in methods such as AdaBelief and Cautious Optimizers. Figuve[Iprovides a conceptual illustration of the MGUP idea. The pseudocode for specific implementation variant, MGUP-A -daum. is detailed in Algorithm for the implementation.
The core steps of MGUP are as follows; see Algorithm 2
Step Compute Alignment Scores. For each parameter i, calculate its alignment score St. m -g.i-
Step 2: Top K Selection. Sort all parameters based on their alignment scores s,, in decreasing
order, and identify the index Zoop of the top K entries, where K -[ d] with 7 E (0, 1).
Step 3: Differentiated Update. Adjust the step size T)t, computed by the original optimizer as
follows: (i) If parameter E Ioopx. its effective step size is set to 77,.- (ii) If parameter i Topk. its effective step size is set to y 7t, Here, a represents the amplification factor, while y denotes the decay factor. In practice, and y can be set to 1/T and т, respectively, where (0, 1).
adjustment based on sign judgment a concept from prior work (e.g., Cautious Optimizers 10 Similarly, we can define the Cautious-MGUn mechanism as:
Φ1. = L/T if m, gr,
mm gr, S 0.
In contrast, MGUP offers more flexible and robust adjustment strategy by introducing top-K selection and sorting mechanism based on the element-wis. wise product between the momentum and gradient.
We clarify the selection criterion associated with the momentum. While intuitively related to momentum-gradien consistency, MGUP-AdamW's implementation in Algorithm can use the

---

product of the final update vector (typically mi/lyV: + c)) and the gradient gt, not just mo- mentum m, and gradient gr. This because, in specific contexts, especially when training large language models, the difference between the selections based on Ut, gr, and ma,i Bt, may
be negligible. Research [24 25 26 27 suggests that within certain model layers, the second moment v; v adaptive scaling might be relatively uniform. This implies an approximation where (m1/y /V1//vi,....... ma/ /va) R3 (m,/c, /un//c,.... ma/c) for some constant c. Consequently, the sign and relative magnitude ordering from u g would closely mirror that from m gr,i. Thus, MGUP- Adam W can be intuitively seen as a selection strategy guided by momentum-gradien alignment.
Remark 3.1. For optimizers with simpler update structures, such as Lion, Muon, or standard Momenntn, M gi, can be directly used the alignment score.
Remark 3.2. We explain why MGUP is expected accelerate convergence. (i) MGUP's mechanism uses a greedy strategy (due to sorting). Greedy strategies play a key role in accelerating heuristic algorithms. When the stochastic gradient and the update share the same sign, their positive product favors larger step sizes; when their signs differ; the negative product favors smaller ones. Large steps drive acceleration, while small nonzero steps ensure convergence. The necessity of small yet nonzero step sizes is analyzed in Section 4. with a counterexample in AnpendixxAjillusstatiii step sizes. (ūi MGUP increases the average update magnitude. In MGUP, a fraction of parameters update with an increased learning rate of (1/T)Ir, while the remaining 1 T use a reduced rate of TIr. The resulting average learning rate is - ( individual step sizes are roughly uniform-for example, when Adam's 's early updates approximate sign(g:) ) the overall update magnitude of MGUP increases by a factor of (1+T- 7²) compared to Adam. This provides an intuitive explanation for its acceleration effect. We recommend = 2 the
the failure of zero
Ir- Ir (1+T T T²) > Ir. When
default, since arg max+ T
Algorithm 1 MGUP-AdamW
Algorithm 2 MGUP
Input: Learning rate 1j > O, initial solu- tion X E Rd, momentum factors $1 31 E [0, 1), weight decay coefficient X, stability term c 0, ratio E 1).
Input:
Alignment score vector St = U2 gr Rd ratio E (0, 1).
(S1) Let Tupk be the index set of the largest
K elements of 8, E Rdd with K = [T- [d]
Set mg- = 0, Vo 0.
(S2) Set Φr,i = { T, else.
1/T, i E Tupk:
for / 1 to T do
Compute the stochastic gradient g: = (x)
return Φi
m: = fimm-- + (1 -1- 31 31)g
V. = 12V1-1 + (1 (1-33)) (g)
g:)
Remark 3.3, The MGUP method can be eas- ily plugged into existing momentum- based optimization algorithums in a satg-andd-lay manner. Examples include Lion fI2] and Muon [13] (see Appendix H for the pseu- docode of MGUP-Lion and MGUP-Muon).
VI- g
u, Vv.+e' ?1- 7
T-
Φ MGUP(u, (u, O 0t)
X; (1 - m)))x;
X = Xxt Te O u,
end for
4 Convergence Analysis
In this section, we rigorously establish both the expected convergence and high- high-probability conver- gence guarantees for Algorithm | in the stochastic setting.
For the convergence analysis of Algorithm !
we make the following assumptions:
Assumption 4.1. The function f is bounded below. There exists f co such that f(x) 2 f",
for all x E Rd. Assumption 4.2. The function f is L-smooth: [V(()) - fx(x) < L - x|/-
Assumptions 4.1 and 4.2 are standard in the analysis of nonconvex optimization
28
20
30
We first show that, under the following standard assumption, the MGUP-Adam (aaamm((wtheee Wwithouut weight decay) can achieve the expected convergence rate of Olog(T))/VT).
5

---

Assumption 4.3. The stochastic gradient is unbiased with bounded variance. That is, there exists > such that for all E R E [xx Vx()), and E|| fxx:()
f [ S o².
Additionally, we assume that [) is M-Lipschitz, Le., [V)( & M for all x and
Assumption 4.3 is very common in the literature
31
Theorem
states our general
non-convex convergence result.
Theorem 4.1. Let B1, = -1 1/2, a² I. 8.
32 S 1, and it n1 /1111//pp. We define the following: E1 M
Umin
5L. pu
and 83
1. Here, Umin 5 and Umax .. for some constant
P
2u
2L'
7
learning rate 7n and any 0. Let p 10L that E2 O, and define Emi min(s1, Ea). Under Assumptions
Whin
and 4.2 for Algorithm [ weight decay), it holds that:
+ E V
where G
31.2 of pPc21
[(x)) f(x + (T+( VT - -2(2V--T))
Remark 4.1. Convergence relies the condition p 10La Notably, if y can be O, Umax
Moon
approaches infinity, making 10L unbounded. This renders the condition >> 10L ill-defined, Winnn
as p would need to be infinitely large, which is unattainable and adversely affects convergence.
Remark 4.2. While our analysis assumes global Lipschitz continuity, the algorithm can be im- plemented using MT jetr [xx((x El instead of a global bound M. This approach only requires bounded gradients along the optimization trajectory, typically yields tighter bounds, and remains fully compatible with our theoretical guarantees. Furthermore, setting
1-1/2//
V--8 t-1/2/p /p is justified since V-88 is bounded and can be absorbed into the 1-3
instead of 1
constant 7 without loss of generality. See Appendix forr details.
Next, under the assumption of coordinate-wise random noise, we show that the MGUP- AdthW(without weight decay) also achieves rate of (ppoly(oog(T))/ with high probability. Assumption 4.4. The stochastic gradient is unbiased with coordinate-wise bounded variance. That is, there exists a; 0 such that for all E Rd, fx: {)] fxx),, and fx; fxx: V f((x)) S of for all i.
Note that the coordinate- wise noise bound in Assumption 4.4 stronger than the standard bound E V fx; {) V f(x)ll < a², as the latter can be readily derived from the former. This relaxed choice made facilitate the application of probabilistic inequalities, thereby achieving improved convergence properties.
Assumption 4.4 commonly used in the literature 34 35 36 37
38
Theorem 4.2. Let 0 S 31 82 1, 32 1 1/T, 7 Cov1 82
10_ 10-- for some con- 1) max 11 7, 1/7, 7 12, and 8s max 1-8, (1+8
VT- - -
stants Co > 0,3 > 2. Under Assumptions 4. 4..2 and|4.4 for Algorithum 7 wiithout weight decay),
then for any given d E (0, 1/2), it holds that with probability at least 20,
+ (Vf(x)) S 1/T-)
Remark 4.3. Setting 0 is crucial for ensuring the stable convergence of the algorithm. The convergence proof relies on surrogate stepsizes (defined in equations (10) and (12y) to manage the complex interplay between stochastic gradients and adaptive stepsizes. The theoretical framework for employing these surrogate slepsizes within the proof is informed by the methodologies presented
in 35 39 40), as follows:
Yt+1 = y: mesh C
B1 (T.bb-iOdr
la 2 (x: x,-
b.
I-27
where the precise definitions of all terms are provided Appendix L Notably, if were set to 0, the St., could approach infinity for some component i when òi a and ΦC-1, 7 -. Such
ratio 0--1.
occarrences might prevent parameter updates in certain iterations, thereby hindering convergence. Consequently, 7 is set to a positive value instead of O. For more detailed discussion, please refer to Appendix E

---

Remark 4,4. Theorem 4.1 and Theorem 4.2 are independent of the specific mask selection mechanism.
Combining Remark 4.1 and Remark 4.3 for Adam variants employing mask, simply nullifying the update step when momentum and gradient directions misalign (tantamount to setting 7 0) markedly alters the convergence properties of stochastic optimization.
5 Experiments
In this section, we evaluate the performance of the proposed MGUP optimizers both pretraining and supervised fine-tuning (SFT) tasks. All experiments are conducted using two NVIDIA V100 (32GB) GPUs and four NVIDIA RTX 4090 (24GB) GPUs. Detailed experimental settings provided in Appendix
Datasets. We use the image dataset CIFAR-1O, the text dataset Wikitext-103, and the language model fine-tuning benchmarks GLUE and GSM-8K.
Compared Methods. We compare MGUP-AdamW, MGUP-Lion, MGUP-Muon with (7)
AdamW [1]) (ii) Cautious Oppimizers(C-Adaane
C-Lion, C-Muon) (10), (iii) Lion [[2], (iv)
Muon 13 42 as well as other state-of-the art memory-efficien optimization methods such as (v)
GaLore
3 (v)) LDAdam 4, (vīi) adam-mini [27 and (vit) Adam-8Bit
43
Unless stated otherwise, the default setting for MGUP-enhanced optimizers is r = = 1/a
5.1 Pretraining
Image MAE Pretraing We pre-train a straightforward VIT model 44 with the Masked Autoen- coder (MAE) framework 455 on the CIFAR-10 dataset. For this experiment, we set the learning rate to 1.5 x 10-4, the MAE mask rate to 75%, and train for 200 epochs. We compare MGUP-AdamW with the standard AdamW and C-AdamW optimizers by evaluating their training and validation losses. The results, presented in Figure 2 show that MGUP-Adane W consistently achieves lower training and validation loss throughout the training process. In contrast, the performance of C-AdamW gradually falls behind that of AdamW.
C-AcamW
C-AdamWW
Adami
AdamW
0.16
MGUP AdamiV
0.16
MGUP AdamW
Valication Loss
0.14
0.14
Traiiing
0.12
0.12
0.10
0.10
0.08
0.08
0.06
0.06
100
200
100
150
Epochi
Epachs
(a) VIT training curve
(b) ViT validation curve
Figure 2: VIT MAE training and validation curves on CIFAR-10
Language Modeling We employ straightforward LLaMA2-7IM model 46 and Qwen2.5- 50M model (77. both of which are pretrained on the WikiText- 103 dataset.
LLaMA2-7IM on WikiText-103. To assess optimizer performance on smaller language model, we train LLaMA2-71M on WikiText- 103, evaluating validation loss. We compare Adam Lion, and Muon variants using a learning rate of 3c-4, batch size of 480, and 2000 training steps. As shown in Figure 3a| the results highlight several key differences. Among the Adam-type optimizers, MGUP-Adam's W achieves 1.6x speedup over standard AdamW and exhibits superior generalization compared to C-Adam W. For the Lion-type optimizers, MGUP-Lion demonstrates a 2.5x speedup standard Lion; unlike the unstable C-Lion which shows early loss spikes, MGUP-Lion maintains training stability. With the Muon-type optimizers, MGUP-Muon yields a ~ 1.2x speedup relative to Muon and delivers better generalization than C-Muon.
While serves as the primary hyperparameter our approach, it essential to examine how variations in influence the performance of
We conduct experiments with 7 E {0.3, 0.5, 0.7)
7

---

and {0, 0.5, 0.9} to evaluate this relationship across different hyperparameter configurations. The comparative results are presented in Figure 3bj The analysis indicates the following: (i) with fixed, increasing beyond a certain threshold degraded performance; (ii) with fixed, larger y generally improved performance. The findings in (ii) precisely corroborate the discussion on the setting of y in Section
1wen2.5-150MM on WikfText-103. We also evaluate optimizers on a larger Qwen2.5-150M model using WikiText- 103 (Figure 4). For these experiments, we use a learning rate of le-3, a batch size of 160, and 1500 training steps. For Adam-type optimizers, MGUP-AdamW demonstrates higher speedup than standard AdamW W and better generalization than C-AdamW. For Muon-type optimizers, MGUP-Muon achieves 1.1x speedup over standard Muon and superior generalization compared to C-Muon.
55-
5445
101
17V
030
(a) LLaMA2 validation curve
(b) Sensitivity analysis
Figure 3: LLaMA2-71M validation curve and sensitivity analysis on WikiText-103
5.0
5.0
C.AdamWW
C. AcammW
AdamW
AdemtW
C-Muon
Muon
Muon
Muon
MGUP-MA WMMP-
MGUP-AdamMW
Training
Validaatoon
3.5
3.0
3.0
2.5
2.5
1200
1200
1500
Iterationn
Iteration
(a) Qwen2.5 training curve
(b) Qwen2.5 validation curve
Figure 4: Qwen2.5-1500 training and validation curves WikiText- 103
5.2
Fine- MunniTning
/13mm
Lea
Liodlig
(a) Adamm-type
(b) Lion-type
(c) Muon-type
Figure Adamw-type,
optimizers average performance across GLUE tasks
We conduct comprehensive experiments on downstream tasks, with particular emphasis on supervised fine-tuning (SFT) scenarios. Our evaluation encompassed two representative tasks: fine-tuning

---

the RoBERTa-besd base model 48 on the GLUE benchmark and the LLaMA2-7B model [6 on the GSM-8K.
GLUE Benchmark Evaluation. To evaluate performance and generalization on diverse Natural Language Understanding (NLU) tasks, we experiment the GLUE benchmark, which comprises tasks varying in dataset size and complexity. We perform a learning rate search within the range of le-5 to 5e-5 for most optimizers, and within the range of le-6 to 5e-6 for Lion-type optimizers. The best performance for each task is reported in Table [ On most tasks, MGUP-AdamY W and MGUP-Muon achieve sate-of-the--as art results. Notably, MGUP-Adam reaches an average optimal performance of 85.15 across all GLUE tasks.
Figure 5 shows how the average GLUE score changes across the tested learning rates. MGUP- Adam W. MGUP-Lion, and MGUP-Muon consistently outperform their standard counterparts AdamW, Lion, and Muon across this range. Additionally, all MGUP-enhanced optimizers demon- strate greater robustness compared to the cautious variants: C-AdamW, C-Lion, and C-Muon.
GSM-8K Fine-tuning. We further evaluate MGUP-AdamW by fine-tuning LLaMA2-7B on the challenging GSM-8K dataset, a critical indicator of fine-tuning effectiveness due to typically low
zero-shot accuracy [49] We conduct learning rate grid search (from le-5 to 5e-5), consistent with
[]. As shown in Table 2 MGUP-Adamm W achieves lower training loss per epoch and the highest validation accuracy 34.96% %, outperforming baseline optimizers.
Table I: Comparison of best results of fine-tuning RoBERTa-base model on GLUE benchmark.
RTE 2.5k
MRPC 3.7k
STS-B 7k
CoLA 8.5k
SST-2 67k
QNLI 105k
QQP 364k
Method
Avg.
AdamW 41
72.93
90.44
90.55
60.32
94.84
92.79
91.34
84.74
Lion 12
67.15
87.50
89.39
60.57
94.84
93.00
91.32
83.39
Muon
64.62
81.13
87.33
59.34
94.27
93.11
91.72
81.65
Adam-mini 27
56.32
87.01
89.49
56.32
93.35
92.02
89.58
80.44
GaLore(r=8)
3
69.45
86.19
88.97
55.12
94.15
92.01
89.86
82.25
LDAdamW(r=3)
4
67.58
88.32
90.03
60.60
94.49
92.82
91.23
83.58
C-AdamW W 100
71.12
89.22
90.25
57.29
93.92
92.62
91.39
83.69
C-Lion
10
67.87
88.73
89.58
57.78
94.50
92.81
91.41
83.23
C-Muon 10
70.04
88,24
90.04
59.81
94.84
93.19
91.75
83.98
MGUP-Lion
71.12
88.24
90.07
61.23
94.27
93.04
91.33
84.18
MGUP-Muon
70.40
88.24
89.84
61.07
94.61
93.24
91.78
84.17
MGUP-AdamW
75.81
90.44
90.54
59.83
94.95
93.08
91.43
85.15
Table 2: Fine-tuning results for LLaMA-2 GSM-8K.
Adam
AdamW-8b
LDAdamW (rank 512)
GALore (rank 512)
C-AdamW
MicroAdamm
MGLP-Adam
Model
Metric
Accuracy
34.53
34.42
34.88
34,62 0.070
34.68 0.081
34.58
34,96
Train loss
0.064
0.069
0.073
0.057
0.056
6 Conclusion
We introduce MGUP, a novel intra-layer parameter selection mechanism based momentum- gradient alignment, and integrated into AdamW, Lion, and Muon yields MGUP-AdamW, MGUP. Lion, and MGUP-Muon. Empirically, MGUP Optimizers demonstrate competitive convergence speeds and superior generalization over their base versions across diverse tasks, including large language model training. Theoretically, we establish stochastic convergence guarantees for MGUP- Adam Wwithout weight decay) under standard non-convex assumptions, achieving a rate near the known optimum. Limitations include the pre-selection of J, inviting future work on adaptive methods. Our theoretical analysis also primarily covers MGUP-Adam W (without weight decay). Thus, while empirically effective with optimizers like Lion and Muon, MGUP's theoretical properties (e.g., the necessity of
0) in these diverse frameworks require further study.

---

Acknowledgments
This work was supported by NSFC (61772570), and Guangdong Natural Science Funds for Distin- guished Young Scholar (2018B030306025),
References
[1] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens tiny subspace. arXiv preprint arXiv: /8/2.044/ 2018.
[2] Brett W. Larsen, Stanislav Fort, Nic Becker, and Surya Ganguli. How many degrees of freedom do we need train deep networks: loss landscape perspective. In International Conference on Learning Representations (ICLR), 2022.
[3] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficiees LLM training by gradient low-rank projection. In International Conference on Machine Learning (ICML), 2024.
[4] Thomas Robert, Mher Safaryan, Ionut- Vlad Modoranu, and Dan Alistarh. Ldadam: Adaptive optimization from cowddimensional gradient statistics. In International Conference Learning Representations (ICLR), 2025.
[5] Weixi Song, Zuchao Li, Lefei Zhang, Hai Zhao, and Bo Du Sparse is enough in fine-tuning pre-trained large language models. In International Conference on Machine Learning (ICML), 2024.
[6] Yuhan Liu, Saurabh Agarwal, and Shivaram Venkataraman. Autofreeze: Automatically freezing model blocks to accelerate fine-tuning. ArXiv, abs/2102.01386, 2021.
[7] Kai Lv, Yuqing Yang, Tengxiao Liu, Qinghui Gao, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources. In Annual Meeting of the Association for Computational Linguistics, 2023.
[8] Rui Pan, Xiang Liu, Shizhe Diao, Renjie Pi, Jipeng Zhang, Chi Han, and Tong Zhang. LISA: layerwise importance sampling for memory-efficient large language model fine-tuning. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
[9] Qijun Luo, Hengxu Yu, and Xiao Li. Badam: A memory efficient full parameter optimization method for large language models. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
[10] Kaizhao Liang, Lizhang Chen, Bo Liu, and Qiang Liu. Cautious optimizers: Improving training with one line of code. ArXiv, abs/2411. 2024.
[11] Juntang Zhuang. Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in Neural Information Processing Systems (NeurIPS), 33:18795- 18806, 2020.
[12] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V. Le. Symbolic discovery of optimization algorithms. In Advances in Neural Information Processing Systems (NeurIPS). 2023.
[13] Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024.
[14] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning (ICML), ICML' 13. JMLR.org, 2013.
[15] John Chen, Cameron R. Wolfe, Zhaoqi Li, and Anastasios Kyrillidis. Demon: Improved neural network training with momentum decay. ICASSP 2022 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 3958-3962, 2019.
10

---

[16] Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning (ICML), 2022.
[17] Jingwen Fu, Bohan Wang, Huishuai Zhang, Zhizheng Zhang, Wei Chen, and Na Zheng. When and why momentum accelerates sgd: An empirical study. ArXiv, abs/2306.09000, 2023.
[18] Cong Fang, Chris Junchi Li, Zhouchen Lin, and Tong Zhang. SPIDER: near-optimal non- convex optimization via stochastic path-integrated differential estimator. In Advances in Neural Information Processing Systems (NeurIPS), pages 687-697, 2018.
[19] Ashok Cutkosky and Francesco Orabona. Momentum maased variance reduction in non-convex SGD. In Advances in Neural Information Processing Systems (NeurIPS), pages 15210-15219, 2019.
[20] Feihu Huang, Junyi Li, and Heng Huang. SUPER-ADAM: faster and universal framework of adaptive gradients. In Advances in Neural Information Processing Systems (NeurIPS), pages 0074-9085, 2021.
[21] Huizhuo Yuan, Yifeng Liu, Shuang Wu, Xun Zhou, and Quanquan Gu. Mars: Unleashing the power of variance reduction for training large models. ArXiv, abs/2411.100336 2024.
[22] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[23] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations (ICLR), San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[24] Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham M. Kakade. De- constructing what makes good optimizer for autoregressive language models. In International Conference on Learning Representations (ICLR), 2025.
[25] Zhiming Zhou, Qingru Zhang, Guansong Lu, Hongwei Wang. Weinan Zhang, and Yong Yu. Adashift: Decorrelation and convergence of adaptive learning rate methods. In International Conference Learning Representations (ICLR), 2019.
[26] Jiaxuan Wang and Jenna Wiens. Adasgd: Bridging the gap between sgd and adam. ArXiv, abs/2006.16511, 2020.
[27] Yushun Zhang, Congliang Chen, Ziniu Li. Tian Ding, Chenwei Wu, Diederik P. Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. In International Conference on Learning Representations (ICLR), 2025.
[28] Dongruo Zhou, Yiqi Tang, Ziyan Yang, Yuan Cao, and Quanquan Gu. On the convergence of adaptive gradient methods for nonconvex optimization. ArXiv, abs/1 808.00678,, 2018.
[29] Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of A class of adam-type algorithms for non-convex optimization. In International Conference on Learning Representations (ICLR), 2019.
[30] Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang. A novel convergence analysis for algorithms of the adam family. ArXiv, abs/21 12.03459. 2021.
[31] Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of adam under relaxed assumptions. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
[32] Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of adam's 's iteration complexity. In Advances in Neural Information Processing Systems (NewrIPS), 2023.
[33] Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng Yan. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEFE Transactions on Pattern Analysis and Machine Intelligence, 46(12):9508-955 99200, 2024.

---

[34] Qi Zhang, Yi Zhou, and Shaofeng Zou. Convergence guarantees for rmsprop and adam in generalized-smen smoohh non-convex optimization with affine noise variance. Trans. Mach. Learn. Res., 2025, 2025.
[35] Yusu Hong and Junhong Lin. On convergence of adam for stochastic optimization under relaxed assumptions. In Advances in Neural Information Processing Systems (NeurIPS), 2024.
[36] Xiaoyu Li and Francesco Orabona. A high probability analysis of adaptive sgd with momentum. arXiv preprint trXiv:2007.14294 2020.
[37] Alexandre Défossez, Léon Bottou, Francis R. Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research, 2022.
[38] Yusu Hong and Junhong Lin. High probability convergence of adam under unbounded gradients and affine variance noise. ArXiv, abs/231 1. 1.02000, 2023.
[39] Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson. Global convergence of the heavy-ball method for convex optimization. 2015 European Control Conference (ECC), pages 310-315, 2014.
[40] Yan Yan, Tianbao Yang, Zhe Li, Qihang Lin, and Yi Yang. A unified analysis of stochastic momentum methods for deep learning. In Proceedings of the Eienty-Sevents International Joint Conference on Artificial Intelligence, IJCAI 2018, July 13-19, 2018, Stockholm, Sweden, pages 2955-2961. ijeai.org, 2018.
[41] Uya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), 2019.
[42] Jingyuan Liu, Jianling Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng. Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu Yuzhi Wang, Jianzhou Wang, Meng Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for IIm training. ArXiv, abs/2502. 2025.
[43] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations (ICLR), 2022.
[44] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.
[45] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Doll'an ar, and Ross B. Girshick. Masked autoencoders are scalable vision learners. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 15989- 2021.
[46] Hugo Touvron, Louis Martin, Kevin R. Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Daniel M. Bikel. Lukas Blecher, Cristian Cantón Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony S. Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel M. Kloumann, A. V. Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog. Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, R. Subramanian, Xia Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu. Zhengxu Yan. Iliyan Zarov, Yuchen Zhang, Angela Fan, Mclissa Hall Melanic Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. ArXiv, abs/2307 099888, 2023.
[47] Qwen An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Guanting Dong. Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxin Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu,
12

---

Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin. Tianhao Li, Tingyu Xia, Xingzhang Ren, Xuancheng Ren. Yang Fan, Yang Su, Yi-Chao Zhang, Yunyang Wan, Yuqi Liu, Zeyu Cui, Zhenru Zhang, Zihan Qiu, Shanghaoran Quan, and Zekun Wang. Qwen2.5 technical report. ArXiv, abs/2412.151 2024.
[48] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. ArXiv, .bwp/190.111699 1692, 2019.
[49] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110. 2021.
[50] Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo. Adam can con- verge without any modification on update rules. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[51] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning (ICML), pages 4596-4604. PMLR, 2018.
[52) Zhenyu Zhang, Ajay Kumar Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang- Q-galore: Quantized galore with INT4 projection and layer-adaptive low- rank gradients. In Conference on Parsimony and Learning, Stanford University, USA, 24-27 March 2025, volume 280 of Proceedings of Machine Learning Research, pages 1035-1050. PMLR, 2025.
[53] Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle. Greedy layer-wise training of deep networks. In Advances in Neural Information Processing Systems (NeurIPS), pages 153-160. MIT Press, 2006.
[54] Geoffrey F. Hinton, Simon Osindero, and Yee-Whye Teh. A fast learning algorithm for deep helief nets. Neural Computation, 18(7):127-1554, 2006.
[55] Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimiza- tion with adaptive learning rate. In Findings of the Association for Computational Linguistics (ACL), pages 2488-12502,, 2024.
[56] Yurii Nesterov. A method for solving the convex programming problem with convergence rate o11/22). Proceedings of the USSR Academy of Sciences, 269:543-547, 1983.
[57] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Lecture 6a overview of mini- batch gradient descent. Coursera Lecture slides https://class. coursera. org/neuralnets-2012 N01/tecture,/Onling 2012.
[58] Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence adam and beyond. In International Conference on Learning Representations (ICLR), 2018.
[59] Timothy Dozat. Incorporating nesterov momentum into adam. 2016.
[60] Liangehen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. In International Conference on Learning Representations (ICLR), 2019.
[61] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations (ICLR), 2020.
[62] Zeke Xie, Xinrui Wang, Huishuai Zhang, Issei Sato, and Masashi Sugiyama. Adaptive inertia: Disentangling the effects of adaptive learning rate and momentum. In International Conference on Machine Learning (ICML), 2020.
[63] Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In International Conference Learning Representations (ICLR), 2024.
13

---

[64] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning (ICML). 2018.
[65] Fangyu Zou, Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pages 1 11127-11135, 2019.
[66] Yuqing Liang, Meixuan He, Jinlan Liu, and Dongpo Xu. Convergence of adam for non-convex objectives: relaxed hyperparameters and non-ergodic case. Mach. Learn., 114(3):75, 2025.
14

---

Appendix
The appendices are structured as follows:
Appendix
A
gives counterexample showing that Cautious Adam may diverge.
Appendix
summarizes additional related work.
Appendix
provides the definitions and lemmas related Theorem
Appendix
offers the formal proof of Theorem 4.1
Appendix presents the definitions and lemmas related to Theorem 4.2
Appendix includes the formal proof of Theorem 4.2
Appendix supplies additional details regarding the experimental setup.
Appendix contains the pseudocode for other MGUP-type algorithms.
Appendix| presents more experimental results.
A Motivating Counterexample: The Necessity of y > 0
To intuitively demonstrate the necessity of non-zero decayed step size (? > 0) for misaligned updates, we present a counterexample where optimizers that nullify updates (i.e., 7 = 0), such as
Cautious Adam (C-Adam) [T0], fail to converge. We adapt classic construction from [50].
Consider the one-dimensional objective function f(x) Ei= [a()), where the stochastic compo- nents are defined as:
nz,
-1
f((x)
for i 0.
+22 (1 +
3n
I < -1
2
- E,
x ΣΣ
s((z)
for i 0.
2 2)²
I x<
The full objective is f(x) = Σn- fi()), which simplifies to:
I, ( [ 2)
x ス 3 x h<
f(x)
We analyze the behavior of C-Adam and MGUP-Adam on a counterexample with its global minimum v= 22 starting from an initial point xo -0.5. In this environment, the optimizer encounters frequent, small negative gradients St -1 and rare, large positive gradients = n.
Analysis of C-Adam's Failure. The stochastic nature of the gradients induces "pulse-decay" dynamic in the momentum term m, A rare positive gradient pulse pushes m, to a high value, after which the frequent negative gradients cause it decay. This leads persistent oscillations of the momentum around zero, a behavior empirically confirmed in Figure [b)
This momentum instability is detrimental to C-Adam (9 0). When my >> 0, the frequent, correctly- signed gradients 9: are misaligned with the momentum m&9t < causing the optimizer to skip the update. When m, < 0, these same gradients are aligned meg: 0, but they produce an incorrect update, pushing the parameter x away from the optimum x" = -2. Figure 6e provides clear evidence for this dysfunction: C-Adam's updates are either null (Ar: 0) strictly positive
(Ax, > 0), moving in the wrong direction. As a result, the iterates not only stagnate but actively diverge from the minimum, as illustrated by the trajectory in Figure 6a
MGUP's Advantage. In stark contrast, MGUP-Adam leverages its safeguarding mechanism (Y 0) to overcome this failure mode. The critical scenario is when > 0 and the gradient is It -1. Instead of inaction, MGUP performs a small, corrective update of size YT), in the proper negative direction. This ensures a persistent, albeit small, push towards the optimum. The scatter plot Figure (c demonstrates that MGUP-Adam consistently performs updates the correct direction (Arr < 0). These small but steady corrective steps enable the optimizer to escape the challenging region and successfully converge the true minimum x" 2, as shown by its trajectory in
Figure 6a
15

---

Coldan
Coldim Updile CARS
MGJP-Adam
MOUP--dam Updala (A6)) - (No Updzn)
(Q)
Optiumis
Itcraton
Tonadon
Paration
(a) Parameter Trajectory
(h) Momentum Dynamics of C- Adam
(c) Analysis of Update Steps
Figure 6: Analysis of C-Adam's failure and MGUP-Adam's success in the counterexample. (a) MGUP-Adam converges the optimum x² 2 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 (Ar, 0) or updates in the wrong direction (Ar, > 0), whereas MGUP-Adam consistently updates in the correct direction (Az, < 0).
B
More Related Work
B.1
Efficient Training and Parameter Update Strategies
Large Language Model (LLM) training often exhibits gradient updates with inherent low-rank characteristics 2 This observation has spurred the development of methods aiming to enhance training efficiency and reduce memory footprint. Techniques like Adafactor significantly
memory needs by applying low-rank decomposition to Adam's second-order moments [5T]. while Adam- mini achieves further optimization using Transformer-specifie Hessian-based storage strategies [27]. More directly exploiting gradient structure, GaLore enhances memory and computational efficiency through low-rank projection of gradient matrices B, and LDAdam complements this by optimizing within low -iw-imensiooal gradient subspaces [4. Furthermore, Q-Galore integrates quantization with low-rank projections, boosting efficiency, particularly in snource-constrained scenarios [52].
While the aforementioned methods focus on compressing or projecting the update information, an- other significant avenue for efficiency involves selectively updating only subset of model parameters. This strategy operates on the principle that not all parameters contribute equally to learning at every stage. Such selective approaches have been particularly explored for accelerating the fine-tuning or adaptation phase of LLMs, although their applicability to large-scale pre-training is less established compared full parameter updates.
For instance, SIFT [S] achieves efficient adaptation through gradient- sparse parameter updates, leveraging the low intrinsic dimensionality and sparse gradient characteristics observed in LLMs during min-ttnning. Applying the selective principle at a coarser granularity, layer-wise and block- wise strategies have also proven effective, primarily in fine-tuning contexts. Building upon early the LOMO method enables efficient gradient calculation and grouped updates, allowing for farl-parameter fine-tuning with reduced memory [7]. later enhanced by AdaL.OMO with adaptive learning rates [55). The LISA method introduces an innovative layer selection strategy based on parameter norms further optimize the update process during fine- tuning [8) Similarly, BAdam implemented a block coordinate optimization framework, selecting parameter blocks for Adam updates, thereby reducing memory and computation specifically for adaptation tasks [9].
unsupervised pre-training concepts 53 54
These diverse approaches highlight prominent pathways towards more efficient training and adapta- tion: leveraging low-rank approximations, or strategically selecting which parameters or parameter groups receive updates, with many selective methods currently specialized for post-pre-training stages.
B.2 Evolution of Adaptive Optimization Methods
First-order optimization methods play a critical role in deep learning. Building on early foundational work, the Momentum method accelerates the optimization process by accumulating historical gradi-
16

---

ents [56] Subsequently, the RMSprop algorithm introduced the concept of adaptive learning rates, which enables distinct update steps for different parameters [57]] The Adam algorithm merges the benefits of Momentum and RMSprop, adaptively adjusting both first and second moments, and is widely used for its effective adaptive moment estimation [23].
Building upon Adam, researchers have proposed several enhanced variants. The AMSGrad method [58] aims to strengthen optimization stability by utilizing the maximum of historical second- order moments, while the NAdam approach [59] incorporates Nesterov momentum to enhance performance. To address challenges with weight decay and learning rate adjustment in Adam, the AdamW method [41] enhances L2 regularization through decoupled weight decay, and the AdaBound method [60] introduces learning rate bounds prevent excessively large or small update steps. The RAdam method [6]] is designed to enhance convergence stability by rectifying variance estimation during the early stages of training. Finally, the AdaBelief algorithm [IT] refines the computation of second-order moments by employing the exponential moving average of gradient deviations from their mean, design choice intended improve generalization performance.
Recent advances have further expanded the capabilities of adaptive optimization methods. Xie al. [62] presents the Adai framework from the perspective of dynamical systems, accelerating training and improving minima selection by decoupling the effects of adaptive learning rates and momentum. In [33], the Adan method is introduced. This method incorporates a novel Nesterov momentum estimation approach designed to accelerate convergence without incurring additional gradient computation overhead. More recently, the C-Adam W method is proposed in [10], which employs a masking strategy aimed at enhancing optimization efficiency.
Beyond Adam variants, other optimizers also demonstrate considerable advantages. The Sophia method, detailed in [63]. enhances Adam's 's second-moment estimation through efficient diagonal Hessian approximation combined with coordinate-wise clipping. This approach has demonstrated superior performance in language model pretraining. In contrast, the Lion optimizer, presented in [12), designed optimize memory efficiency and computational speed. It achieves this by tracking momentum exclusively and employing a sign- based operator to standardize update magnitudes. Notably, the Muon method, proposed in [13] and originating from the framework of Shampoo [64), incorporates Newton-Schulz-ine Schulz-itraated orthogonalization of gradient momentum. This technique is intended enhance convergence dynamics through adaptation to parameter curvature.
B.3 A Brief Review on the Convergence of Adam
The convergence theory for the Adam optimizer has evolved from early uncertainty to rigorous proof. Initially, Reddi al. (58) revealed its risk of non-convergence with a convex counterexample. Early work to address this established conditional guarantees; for instance, Chen et al. [29] provided the first convergence rate in non-convex settings, while Zou al. [65] identified necessary hyperparameter coupling conditions to ensure stability. A key turning point was the work of Zhang et al. [50], who first proved that the unmodified Adam algorithm is convergent, attributing prior failures to a mismatch between hyperparameters and the specific problem rather than an inherent algorithmic flaw. Building on this, He al. [66] strengthened the guarantee from ergodic to the more practical last-iterate convergence. Finally, Wang et al. (32 resolved the debate by proving that Adam achieves optimal iteration complexity of ((e-1) matching the theoretical lower bound and providing a firm theoretical foundation for its excellent empirical performance.
C Expectation Convergence Lemmas
C.1 Definition
Recall the form of Algorithm Let g, denote the stochastic gradient. Let's consider 31,5 = 1 1-1 1-11 Th 1/2 / Therefore, we rewrite the formal definition of the algorithm,
17

---

m. = BI, 1, + (1 B1,
Ve = B2 + - -32)ge;
V = pVt max(e, vvr,, 0,
T) T/
h. V
Tp
X++1
Here, we incorporate = nyt 1/111 into v and set To de (ar, 7), we set = and y (0, 1). Then, we define that
Without loss of generality, for
M
Umax
Mmin
Umax
1 17y
Umin
(5)
he 2 (x++1 -xt),
S; m.
vf(xx)).
C.2 Lemma C1
Lemma C.1. Suppose that {F4, Ar} are two nonnegative sequences. Assume Eat1 & (1 (t
1)-1/2) + A-11 and 2 1/2. Then we have:
1-1/2 E, 1 20, - Ec+1 + Alt1).
Proof. We derive:
1/2E, c c(E: Est1 + Ast1)
11/2 E, - (Ee + Act1) c. Er (1+ ()+1)-1/22) + Act1)
Es (1-1/2 dt + 1)-1/2
-1/2
= E. . (t+ (++1) 11/2 ((+11)
(0)
E, ( 11 (2 c)
2 0.
where ()) follows from E.+1 5 (1 (1 1)-1/2) + Aetsi (0) due (+4 < 21/2; (x) due to our choice e = 25
C.3 Lemma
Lemma C.2. We have the following results for all t 1, pytumin I min(h:) s pVtumax.
18

---

Proof. We obtain:
Ve, (1- 32) Σ
82 gg gj.
$ (1 82) (max 87,3)
Σ 3/2
jE(t]
j=1
(o)
max
g
%E
S M
where (o) is due to E/
1-
(+) is due to we assume that f(x; is M-Lipschitz
1-
82
T- 82
Ve,i 2 0. for all x. Additionally, we note that
Thus, we conclude:
V1, M2
This implies:
Ves E
Next, according to the definition of Umin and Ummax:
M '9"
Umin
Umax
Therefore, have:
pytee
pVtM
h
pV tymini PV bumax
11
777
C.4 Lemma
Lemma C.3. Let Umin: Menaxx Fe, St be given in (5) We have the following inequality:
L
Bl|ra
El|s. 2
B[[(x++1)] {
2x pumax
p Prrni
vt
2
Proof. First, have the following inequalities:
max(h;)² 1
Xt+1 Xx 12.
||x+11 xe
min(h;)² x²
2 |x41 x min(h.) max )2,x), 1
min(h;)² x²
max(())) 1
||x+1 x
(6)
min (h,) R²
h C (Xe+1 x:)
x² min(h:)
||rell
R² min(h,)
19

---

Applying the descent Lemma to the algorithm, we have
[(x++1) S [(x:) + (V](x)), Xx+1 -x-x, xx) + L
- 2
Xx+11 xx112
= (x,, + m, X6+1 x1) - -mm+- - (x(x)), X - -x1) + | - Xxo12
2 (xx)
1
L
2 |xr+1 xrllir (m: v.f(x,),x++ X++1 xx) + 2
||xx+1 x1112
2 [(x)) 2 ||x++1 - xell +
m - (xxx))]13 + (3+ 1)4
| - Xxe12
2,3..
1
(3+ 1)L
$ ((x)) 2 | xtlli.
m [x(x))1
1h ② - (x)11
2BL
Emin(h,)²
2 f(x)) 2x² min(he)
(8+ + L)L
|1r112 +
||re112
||
2 miin(h.))
2BL
1
(8 1)L
||s 12
f(xi) 2x Pelpytumman
||re112
2p
23L
where (o) follows from the equality (m., X(+1 xx) + || X 1l = 00; (o) is due to Young's
inequality; (+) results from Inequality ( and (+) is derived from Lemma C.2
Then, by setting 3 = vt and taking the expectation of both sides, we obtain:
(Vē + 1)L
ElS S s(x.)
E Fr
E
[sr11
2x² putumax
2p2tu22 Lu L p² D2vtul
2v
S f(x))
E
s
(7)
2n² Dyttmax
2
E
E Es
f(x,)
2x² pumax
p² prumin
VE
21L
C.5 Lemma C.4
Lemma C.4. Let St be given in 5 We define S, El|s R, E|[rell Pt = f(x)) f(xxx(xxx)))))))r))) 5t El 7. 82
and E3 ZL. Assume that p is sufficiently large such min(e1, Ea). The following result holds:
51.
2x Manax
pr
that > 0. Let Emin
T
Pi + 20 Logg + VT -T - 1)..
Σ
R S. <
Emin
Proof. First, derive the following equalities:
S; m (x(e))
-1,rm,-1 + (1 1-,11)) - vx(x))
-1,,(m--1 v((x(-1)) + - vf((,)
= -1,18--1 + /L.((V - vf((,)) + B1, Vf(x.)).
20

---

Then, we have:
24 2 (1 Vs(x,))]
1+ 3 1 82,1 (ialIl/(((x V
S (1+ (1+,))8,,lss+--111
1
(2 81.))8 +
1+ (1+ 1 1 B1,3
89.11 ( - v
2 339
$ (231, - L1es
31,LL' 1| - -x112
1 B1,
I,2
-1, + 1 31,t ||x-1 - -xe12
where (o) is due to Young's inequality for any B > 0; (o) follows from setting / = - B1,1 ; (+)
due to Assumption 4.2 and (+) results from fact that 23 B. -(31, 1)2 S 0,
which implies 2/1,+ ff. I 1.
Therefore, have:
(0) BI[(1 - (g. ff(x))) + E|
E
5
5 31,0)² + 31,1|s: 11. + - B1,t
Xt-1 x
2
o² 31,) + 11,111s--111 +
① (x; x,)1)
(1 31,t) min (h,
where (o) is due to Assumption 4.3 E E [x; V fz)). Then,we we define
1
A -= 02 10-21111111 31,1²² + (1 - B1,1) min 1h--))
R²
Therefore, we have
8:¹ < S1 + Ac-
Then, using Lemma C.I and B1, = 1 1/2 we obtain:
t- 1/1 2(S - 5/4 + Ac+1)
212
-(S S/41) +222 +1+
R
(I+ (1+1)-11/ 1)11/2 min(h.)
(0)
20
[t+1)
1/2
I.2R?
(8)
< 2(S) 1
+2
t+1
n2u2 min Vi
(+)
20
4T.2
2 2(S - S4n1
[+1
Vi
where (o) is due to Lemma C.2
relies on (111 < 1/2 22
Then, from Lemma C.3 it follows that:
L
E
ElIra11
E
E (xx)
2x² pumax
puuz
Vt
2VLL
L
R2
S?
2x pulmax
puuz
Vi
2/LL
21

---

Adding both sides by sit + 1-12
St yields:
t/2
Bff [x(+)(+1]] - f(x)) + t1tt-
+
S
2L
R.
1-1
+
E1t
E1t
St
2k² pumax
p²u
V
L
L
R
t
I 7 Ss?
2
- + [((+1)
20²
+ 4LR
2x pulmx
P
Vt
pru
5L
R2
+ - ( 1 + L(t +
20²
t1t
2x pumax
p² Daumin
Vt
A0
where (o) is due to Inequality Using the definition of Pt, 82, we further derive:
20² e + + 1 s P: P41
(1 ( 1)-1
L
This leads
202 Emin (1/2 + R2 + 1 Pi+1 +
L (t+1))
(9)
Summing Inequality (9) over from to T yields:
T
1/2 (t (t-1/2 + R² - S?)
Emin
[=1
T
$
Σ
(P. P.+1 + 20 I, (t 1)-1)
1=1
(0)
Pi+ 11220211 log(TTT + +1,,
where (o) is due Σ This further leads to
+ S log(T + 1).
T
+ 20
log + 11
T
Σ
R + S.
VT Σ 11/2
Emin
1=1
Pi + 2 log (T +1)
VT - 2(VT = (log((T) (T)
Emin
where (o) is due to Σ 11/2 2 2(VT-1)
22

---

D Proofs of Theorem 4.1
Proof. First, have the following inequalities:
(x(x++))
[V[(x++1) - m. h, @ (X++1 - xx)112
Σ
[IV[(x++1) + m. he (xe+1 xx)112
[V]((++1) v.f(x,) 5, h, 3 (xx41-xx)1
T
IIIV V fxx )13 + ||sc| + ||ا 2 (xe+1 - x:)1 + 2(s,. by 2 (x(+1 - -x))
2VV((x fxxx)) 2 f(x+4) Vf(()),, h, (x1+1 - xx))
T
3(1 [ (x(+11) fx( 11 + 1 + ||ra 11:)
1=1
3L ||x+11111 - x,113 + 3|| 11 + 3
1=1
31.²
|re12 + 3| 12 + 3|1re
[min mmin(h,))²
where is due Young's inequality; (o) due to Assumption 4.2
Let's take the expectation both sides,
31,2
E /
R² + +5, + 3R;
(min h, ))2
3L22
R2 3R2 + +5)
Pu min min
3L2
R2 + 3 + St)
P
3
T
T Σ
Σ
R +3
R + S?)
p²22
1=1
t=1
3122p²
T
1²2 +3)
Σ
R² + S?
31 + +3022
f(x1)) -f(x)) 2022 log(T +1) VT 2 11)
12e2
Emin
(T)oT( = Ö(VT),
where results from Lemma C.2
(o) is due 1 1: (+) follows from Lemma C.4
Thus, we have
T
((T
min t=,..
V f (x++1)
E
T
23

---

E Probability Convergence Lemmas
This proof refers to the following literature 35
37
36 38
E.1
Definition
Let ga denote the stochastic gradient. We define the noise E. g. Vx(xx, and the coordinate-
wise noise as Sa,i = ga,i (x(x)).. Furthermore, we define two auxiliary sequences {Ps and { $211.
31
P1 O1 X1 1 B1
(x. X x
(10)
yz = + x,, Vs 2 2, (x(x)), Sa,i = ga,; (x()))
6. = g.
Then by the definition of Assumption
4.4
we continue define useful notations:
el'
Dr V log
G
max fxxx)]] jE
(s) Dr / + 2G²
G1 Drv
m,
V
(11)
m, 1- 31 V. 1 33
b. Vv + == VB +
32 -2) tc.
a = /v. V. - v 32V2-11 (1) (Gr(s)la)* +e.
It is noted that yt can also be expressed in the following form:
gx 81
7.b, @ Φ. 1s O (xx xx-1).
Ys+1 n...
(12)
b, 1 81 n 3 0
Without loss of generality, for Φe E fa, 7). we set a = 1 and y E (0.1).
Then, define A,
11)
E.2
Lemma
F.
Lemma E.1. Suppose that {as}s21 a real number sequence. Given S 31 < 32 S 1,e
we
define c₂- = Σi= B aaj,0 ds = T- Eÿ= 34 Ja and es = =ド 132
85 0 then
(1-B1)(1- - /12 (log (1+ - +log =log33).
1
Vt 2
c Cg
(log (1+ =t t tlog 1log
Vt 1.
Es (1- 31) -31/32)
Proof.
See the proof of 35 Lemma A.2].
E,3 Lemma
Lemma E.2. Suppose [Ze) sE(r) is martingale difference sequence with respect G.. Sr-
Assume that for each E [T],o, is random variable only dependent by Si, Cr and satisfies that
E
Texp(ZZ/03) (1, ... S.-11 <e.
then any > O, and for any (0, 1).it holds that
T
T
Z
A
log
5
0
28.
24

---

Proof. See the proof of [[36], Lemma 1].
E,4 Lemma
Lemma E.3. Let Ba My m, be given in Algorithm and b,, be defined in 17 1f0 and F,(t) =1 E s==1 g then Vt 1, have,
31
B2 <1
g+
d
Fi(1)
log
bg
32
55
m.
1 31
d
Fi(t)
log
b
(1
B2)
B1/32)
135
2=1
m,
1 1-
Fi(t)
log
b
32
B2)
B11B2)
B5
i=1
Fi(()
log
b
32)(1 - 3r/32)
35
Proof. First, we have:
2 2 (3) (1 32)
Next, the following inequalities and equality hold:
b" 2 s. + 2 2 (1 (1-B2)
32 gj, + +22 ) Σ
m. = (1 - -11) Σ
35-1 Bi
j=l
5=1
For the first expression, follows that:
Ba
B2 8=1
2
Σ; 1=
08-1
8=1
Bj
(0)
1
1
32
log
1+
2
at -
2 8
t tog 32
3=1
Fi(())
log
1 B2
13
where (o) is due using Lemma
E.1
For the second expression, we have:
(1 3,)
Σ
88
81,
Σ
1 - B2
22
05-1 gj
s=1
(0) (1 B1)²
t
log
1+
g
-loog 32
1 - 32
(1 - 31)(1 31/B2)
62
#=1
1- 31
F((t)
log
(1 (1 - 31/32)
3
where (o) is due using Lemma E.1 and setting 32 1.
25

---

For the third inequality, we derive:
2
(1 - (1-3 11 8/
m
2 32) + (1 32) Σ}=1
18+1-
2
8j
2
(1- (1-31)² [2)= 89- 81
2 - 82) + (1 32) Ei=1 1=
02
35-1
8),
==1
2
Ej= 81 81
(1- -,)²
(1 12)
==1 32
2===1 1 P2
Bjd
( (1 81)²
1
32
t
log
1
B8,
log 32
(1 82)
(1- 31)(1 - -11B2)
1- 31
(F((t)
81 (1) /r/32) log
where (o) is due using Lemma E.1 and setting Bg 1.
For the fourth inequality, we derive:
m 5,i
(1 (1-3,)
P j=1
3
88
1 1 B2
22 + Ei=1
B2
2
s=1
Bjji
(1 - 31)²
t
log
1+
8/ 8 8.,
- -1 loo 32
1 - B2
(1- - B1) - 32/32)
2
s=1
1
log (F((1)
< (1 31 B2/B2)
185
where (o) is due to using Lemma E.1 and setting 32 S 1.
E.5 Lemma E.4
Lemma E.4. Let qu 7--1, y, by, b,- be given in Algorithm 7 and 70
then we have
1s b.b 1,8
S w, Vt 2 2,
7 00
1+ih where w cov D2 + 1, max (1,7, 1/7]
Proof. To proceed with the proof, we first establish the following. For all 2 2, given the conditions
0 5 1
95-1
3i and
S
8, it follows that:
1 1-83
1-39 03
ns
35
1-
3
1/ 1
1 32 3--1
1
Bi
32-1(1 (1 32)
82
1 +
v
(1
12)
/ Bz.
1- B2
Next, We have:
b
++ Vo-1,
E Vo-la 1
b c+ V B2V 111 + (1- - 3g
e+ v B3V
V
26

---

For 0 Pa, {1, 71, the ratio satisfies it follows that:
د.. 0.-1.
< max 1, 7... Defining c max 117.77)
The nrbb- n
1bbb-11
1+ 82
S
+15 q/
+ 1.
1/ 10
7/- bay
82
E.6 Lemma MS
Lemma E.5. Let m, b, be given in Algorithm
[and| 17 with B1 < B2 < 1, respectively, Then,
(1 B1)(1 - (3)
VS 1.
b
B2)(1 B1/B2)²
Consequently if is I-smooth and we set v) Cov for some constant Co > then have
[Vf((,)] < [V(() + LCosy +Cos
Vt 121. 1.
1- 1/32/32
Proof. First,we derive:
(1 (1-B3) 31)
88-1-j
m,
j=
1,2
B2)
2j=1 82
Bj
(o) 31
e-1
g8 1
8
3
V1 B2
02
1-1
2j=1
fj
1 8
s-1-j
28-
2
1-j
2
-1- - 3
B2
1-
Σ
8-1
3 38-1-
8-1
s- s-1-j
VI 3
32
1=1
j=1
9
8-1
1 31
1 28-1
8g
V1-
31
91
3a
(1-31)(1 - 31
<
(1 B2) [
81 Dg
where (o) follows from applying Cauchy-Schward inequality, which gives us (Ej=1 1.
8 84-1-5
g2,,)² 5
25 Σj- -- --ا
-g g3 (+) is due
For the second conclusion, have:
Vf((x- [|vf(x))| = [V [V((x)) (f(xx))--((xx-1)))) [xxx)) L||x,
X--1ll-
Furthermore, it follows that:
- m.
|x,, x,-1/1 S vdlxx, x,-1 = 7---1VV Vd
b
VVd Vd
m
Σ1 \
d
Co
d
b
(1 - B2) (1-
0
1
91
3g
27

---

We obtain:
d
d
[V (IV)(xi) + LCD,
S [[f((1)) )l + LCBs
1 8, 82
1 8.
E.7
Lemma E.6
Lemma E.6. Suppose that f is L-smooth and Assumption 4.1 holds, then for any x E Rd, we have
[V & 2x( f")
If' = Co < 32 < 1, let any xx. Yt defined in 10
then we have
2L 2ld
|V S 2ll))(y))l (1 (1) B1/B2)
Ys
Proof. For the first conclusion, we define x x 1. fx().. According to the descent Lemma for
L-smooth functions, have:
L || x1| S f(x)
1
f(x) < f(x) + (Vf()), xx) +
2L V
2
Rearranging the terms and noting that f(x) 2 f where f denotes the optimal value of f. it follows that:
V fxx] < 2L (xx f(x)) - 2L ff).
For the second conclusion, utilizing the norm inequality and the L-smoothness property, we obtain:
[V / + (y:))l
$ v 21 - xx
2L2,37
21IV/(>))11++
||++ xt-1
(1
B1)2
LB 1Bivddn-11
m
21[ +2
1-81 31
bi-1
2r.Cdd
21 +
where (o) is due to Young's inequality; (+) relies on the inequality ?) S - S
CovI-8 -1
followed
by Lemma E.5
E.8
Lemma E.7
Lemma E.7. Given T suppose that for any s E
[T], coordinate-wise Esri S..i V)((,);
satisfies Assumption 4.4
Then for any given 8 E (0, 1). it holds that with probability at least 1 8,
e < D
Vs [T].
Furthermore, we have
max JE(a]
||y < Gs))),
max SE(a]
|g, s Gr(s),
max Vlyllo 5 (Gr(s))3, e jEls]
Proof. First, we define Wa,1 expectation yields:
for all E [T]. According to Assumption 4.4 taking the full
E (exp(wa.)) < exp(1).
28

---

By applying the Markov inequality, for any & E (0, 1), we have:
max Wa,i 2 = P eE(T]
exp
max Way sE[2]
2 exp(d)
< exp( (
exp
max Wa. $E
<exp(-00
exp(ws,)
S <exp(-5)T exp(1).
This implies that, with probability at least - 6,
ET
2², S log
as E [ of
Consequently, it follows that:
115.11 < D2 (1lo|2 + 2G?) < (s),
where Dr and Gr(s) are appropriately defined constants or functions
Next, applying Young's ' inequality and given Dr 2 1, for any E [s]. we obtain:
g < (xx))( + 2 < 2D} (1o=||2 + |V s (Gr(s))?
Finally, we employ mathematical induction to establish the concluding result. For any E [d], note that the base case holds since:
V1, (1- - 32)gi, -1-
(Gr(s))²
Assume that for some s' [s). the inequality v,, S (Gr holds for all j e [s']. Then, for the inductive step at j = s + 1,
V = 32,. + (1 32) S ((b((G(()) + (1 (x)( (Gr(s))
Thus, by induction, it follows that:
Vy S (Gr(s))² Vj [s].
Since this inequality holds for all e [d], we conclude that the desired result obtained.
E.9
Lemma E.g
Lemma E.8. Given T 2 1. bf = (b..)s and (a the definitions in 11 and Lemma
E.7 holds, then for all s E [T] i [d].c c 11,711/7, given by Algorithm
1
Gr(s) B2
ag.
b
1
(Gr(s) + 32
b
by-
c
(G((s) Br
ag.i
bg-
ag be-
16292-1|+ o--1
where Bg
oVI-12
Proof. First, prove the first inequality:
1
n.
1- B2 1 (GT(s))²
agá
b
a,,by,s /Vaii + Vūni
(0)
32
(GT(s))2
a, D
VVs + 32V.-15 1, +(1 (T)
Gr(s)vI 32
ba.i
29

---

where (o) applies the result from Lemma E.7
ga |
v (Gr(s))? and (+) is due to
Veii 2 /1 (s).
Next, we prove the second incquality using vā - vb S va for - b S a:
1
[Va.,
b
as,i
b, b,
(1 32) (GT V.-LA
1,
Vv,1 + VV-1.
1
(1- (Gr())
V + V
Gr(s)\T-13
where (o) is due to the fact that 1,-1, S (Gr(s))?, as stated in Lemma E.7 and (+) follows from the inequality VT- (s) I 0,.;-
Finally, we prove the third inequality, similar to the proof the second inequality:
Vo-11
a.,i
b iss
1
1e2vos Vo-Lul
b b,is, CV Vs,i + Vv
1
(2B2(V G7(s)) + (c ( v=-1,))
1,i
/VVo + VV--11
1
(122 11 + 162 (1)
1,18,
cvI (3)((s)
Let B3
then we have:
evI-18
BrG((s)
b 1,4 a,,i
b1-1,
This completes the proof of all three inequalities.
E.10
Lemma E.9
Lemma E.9. Given 2 and 8 (0, 1). If Assumptions 4.4 holds, then for any > 0,
2(1 -1----V- with probability at least 6,
1
V (xxx)
d T
A.1.1 <
7
+ log(
23
a
==1
1
V fxxx)
2
(1) 82
g+
A.1.2 <
7
23
va,
2(1 31)
Proof. Recalling the definitions of a, and Gr(s), for any s E [T] and i E [d]. we have:
1
agd
(Gr(s) + +0VT
Cr(s)VI-B3 (11s)iI-B, 5
0.vI- - 32
Then, we define:
x_ nod. V (x(x))(far
no V ).E
Ne
X,,
Wa,
04
ag
ag,i
a,,i
30

---

where fx,)) and a,, are measurable with respect to Fa-1, = (f1,s- (5--1,) and Ex, is the noise at step s. Thus:
F.f (Li, ... (a.)).
Next, applying the Bauchy-Schwarc inequality, we obtain:
V (xxx): Øe E.
V [x(x)
(x(x)) as
11
02
Given Assumption 4.4 Eg ( 3) = Vx(x), and:
(Vf(x;E); v)((x),) - o
almost surely, it follows that:
X²
E
exp
s,
Fa-1.
S E
exp
F Fo.1i
11 V fxx))
E
exp
17 f(xx)
1Fo-1a S exp((1)
Then, by invoking Lemma E.2
we derive:
31
T
S
w 5,
+
log
4
X
8=1
8=1
31
17
V
f
(x+)²
T
of +
log
4
a²
8
8=1
3
V
fxx))
T
0f +
log
4(1
311v1
B2 s=1
as
GT
Vx(x))
(5
T
log
B2
a,
8
where (o) follows from
0.11- and 1/4
follows from oi S GT.
T
Setting&
2(1 2(1-81)v1-84 92
3rGrB we obtain:
A.1.1 Σ V((x))
<
1/o
(x.))
log
T
ag
28
a,,i
s=1
1=1
1
fx(x,,
T g
1/s
+
log
2,3
Va
Next, we bound A.1.2 as follows:
A.1.2
Σ
7
(v((x)),
1
d. 3 ga
b,
8=1
1
ns
[Vf(x))g...
a,,i
b
1=1
(6)
Gr(s)/I-
Br
n
(f(x)).g..)
(xx))
(1
32)
d
(Gr(s))
ng.2
23
ag,
2
agd
ba
im s=1
ngGT (t)
32
g+
23
a.
2(1 B1)
b
1=1
#=1
where (o) follows from the result of Lemma E.8
(+) due to Young's 's inequality, and relies on
aud 0vv1- and 77 < 1-
31

---

E.11
Lemma E.10
Lemma E.10. Given T 2 1, if Lemma E.7
holds, then for all t E [T]
1/4
(x.)
3(GT(1)) T- :32
m
B.1
B
a
2(1 31)3
b
BGT (t)m
m
2V
2(1 - 31)3
b,
v(1 (1) :3) - 32)
Proof. Decompose A, ① (x, xx-1) follows:
1/ b. O O
m O Px-1
A 2 Xx xx-11 =
l8
1/--1
118-1 bg
@ P--1
b bg-1
1
- @
3 mm-1
bg
ag
na
m, ② Φx-1
3 "1,-1- - -17-- - 11
a
bg-1
bg-1
Then, we have
31
ne- 11-11
3 P0-1 V/((x))
B.1 -
As
1- B1
ba-1
31
75-19.- b
M v((x.)
B1
bg
B1
17s
178
① 0 2 1 vf(x.)
B1
b,
as
B.1.1
B1
1-P8-1
vf(xs))
+
m,-1,
1
31
as
b,-1
B.1.2
B1
Φ--1
② m vs(xs))
+
(11-1
110)
1
31
bs-1
1.1.3
For B.1.1:
81
1s
n
1, Vf(xs))
mm-11
1-
B
b,
a
81
1
1
11
m. (xss)
1- 31
bagi
agi
8
1
1
Σ
n
M ((xs))
1 1- 3
baib
as,1
1=1
d
31
Gr(s)
B2 a,.bb.i
VV((x) m ). 1.
B
vf(x)}
31. - 32)
(GT (s))2
m 1.1 b
2,9
as.i
2(1 - -3
as,i
7/
V
(x.)
B(GT (1) + +e))+(/1 / B2
ms-1
2/3
a,
2(1 B1) 3
b
where (o) follows from Lemma
E.8 and (+) follows from Young's inequality.
32

---

For B.1.2, let Oui = and for any c₁ E f1,7, 1/1), let
00-1,
1 82 2 2 72 (32)
182 -- - 72)
33
max
vI 32 7V 32
7V $2
Then, we derive:
3
770s-1
m3-1. Vf(x.))
- 31
a,
d
3
1
m Vf(x.)
1-
B1
ag.i
1=1
B1
Gr(s)MaB3
[V[(x))
31
D,, 1,
1=1
(+)
7/8
vf(x.)²
)
B
d
(GT(s))²
m 1.i
+ 2(1 - (11-
2/3
as,i
as,i
by
i=1
1/8
(x(x))
(I) + e)m183
m
23
a
2(1 - 31)
b,
where (o) follows from Lemma
E.8 and the condition di S 1, and (+) is due to Young's inequality.
For B.1.3, we derive:
31
④
m
vf(x))
(1s-1 17,)
1- B
b
B1
/1
33
33
1
1
32
② ② m,
(((x.))
17
V
1- 31
1
Bi
82
32
30-1
b
Thus,
n811 / - B2
1
Φ--1 ② m,
B.1.3 <
v.((x)),
1- B1
3:-
1
3
ba-1
B.1.3.1
(
Os-1 4 m
B2
v((x),
31)(1 35-11
b
B.1.3.2
Note that V (x,)]] S G, < G. for all s S t. Then, by applying the Cauchy-Schwarr inequality, and the condition di S 1, we have:
Lemma B.5
Φ--1
V B
(x.),
b
(1 (1- (1))1)(1- - 15-1)
$ /I 351 fx(x))) ||
m,-1
VdGuy.
b
(1 - 31) - -31/32)
Therefore, summing B.1.3.1 over E [t]. using 31 E [0, 1}. and noting that B. 1.3. vanishes when
5 ==
VdnG.
1- 81
B.1.3.15 <
I- 3
V (1 - 32)(1 - B1/B2)
1 89
1 1 39
VdniG.
(1-8881 (1-88)(11 - (1=))11 - -Br132)
33

---

Similarly, since [V,(x,)]
I G. S G2 for all s S and 1 - -1-1 2 1 B1, and di S 1, we have:
1
Φ @ m
v)(xs):
81
81-1
b
1
[V(())]|
m
S VVdGr
39
b,-1
-1- (1-32)(1- -1:133))))3
Thus, have:
VdnGe
1
B.1.3.2 -
V B2 11 32
1-B1
V (1 (1-B2) B2) B1/32)
==2
VdngGr
VdnG.
S
(1-33)0 B2)(1 B1/B2)
v(1 31 (3
B2)(1 - 101--11885 31/
Therefore, combining all terms, we obtain:
Vx(x))
3GT(1) B2
m
1/8
v
B.1 <
2
Vag
2(1 B1)3
b
23
a
(GT(())Be
mg
2 BddnG
2(1 B1)3
v(1 - -1) B2) - 91/32)
Vs(x.)
[GT(1) 1- 132 m.-1 2
3 Va,
2(1 31)3
b,
Σ
BGT(()nB3
Ma-1
2 2vdrnGr
2(1 B1)3
b,
V(1 31)
E.12
Lemma E.11
Lemma E.11. Let 32 1 1-//T, and Fi(1) be given in Lemma E.3
We have log
F((1)
(llo(T)))
Proof.
First, we have
1
1
32
1/T
- log B2 = =log
<
1/T <
(13)
82
82
1
T
where we apply log(1/a) (1- a)/a, Va E (0. It follows that
F(T)
log
S log(F(T)) + +2 log(e? (TT())
(14)
According to the definition of F(()) in Lemma
E.3 we have
Σ
3 80 2 Σ
(x.)).) 1 22 < (o ()(x))))
(15)
2 [1allae1
Σ
| |V
34

---

Thus, we obtain
=1(1=1+
Σ
B2
8
Ba.
2 #=1
( [ + (1V7(x))) + +1ICO B1/B2
(0)
11+-
(16)
2
(1011 +2/V/(xii| 1+ 1- /1/12
21'00d (3)
where (o) follows from using Lemma E.5
(+) is due (a + b12 v 2a² 2b
Therefore combining (13), (14) and (6). we arrive at log
[((T)
0 (log(T))
8
Proofs of Theorem 4.2
Proof. Applying the descent Lemma to the algorithm, we have
L
f(y++1) S f(y.) + (Vf(y.). -y)) +
- Ilyy+11 -y
2
= =(ys))
()((s), -ndes
gr
31
A
@ (x.
xx-1)
by
1 31
L
8. +
3
+ 2
A.
(x, -x,-1)
b
1
B1
f(y.)
V.((y)),
g.
+ 1 31
B1
(VS(ys), A. (x,
by
L
g+
31
+ 2
7.00 -
+
Ag ① (x, - -x--1)
b.
1- 31
$
f(xi) +
Σ 7
Vf(ys),
og
2
g+
Σ
t
31
(Vf(y.), A. (x, - -x--1))
b,
1 - 31
sm1
8=1
+
n.o
g+
A 2 (x. xx-1)
2
b,
B1
Then, we define
1]8
Vf(y.).
0
g+
b.
#=1
B1
Σ
(Vf(y.). A 2 (x, - x,-1)),
1- B1
==1
C= ΣΣΣ
- nad. ④
g+
B1
A 2 (x, - xx-_1) 2
2
b,
31
==1
Now, decomposing A and B,
Vi(v.).
g+
=
1/8
b,
1=1
-78 (V((x)))
ga
+
Σ
(f(x.)
Vf(y.),
o. 3
g.
b
b,
1=1
==1
A11
A.2
35

---

31
B
(Vf(y.). (x, x,-1))
Vx(x)); As (x, x,-1))
1- B
B1
B.1
B1
Vf(y,) VI(x.), (x, -x.-
B.2
Subsequently, using the conclusions of Lemma E.9 and Lemma
E.8
we have
v V o
V fxx))
A.1
n
Va
Φ ④ E.
1/
(xx(x)
1 [f(xx))
os
g
a
a
b,
A.1.1
A.1.2
V (x,)
fxx,,)
2
T
A.1 <
Σ N8?
Va.
+
23 ==1
Σ
7
Va.
+
log
8
8=1
(x,,
n/ (1)TT- 32
g
ns
23
va,
2(1 31)
b
==1
fxx))
T
B (t) 32
Σ
ga
1/8
+ log
+
3
0
2(1 - 31)
b
#=1
For A.2, applying Young's ' inequality, we have
11. (V/(x)) vs(ys). &.
g
g+
7s [)(x)) - vy
b
b
LA 2
g,
V fx(x)) ff(y,))l
+
2L
2
bg
L
Xx,
X 11 +
LT)²
g+
2(1 - 31)
2(1 - 81)
b,
Tnt²
m
Tnj
g+
+
2(1 )
b
2(1 )
where is based on dy S 1 and applying the Cauchy-Schwarr Schwarz inequality; (o) follows from applying Young's ss inequality; (x) is due to
||x x 2 $ n
m
S 17
m
b
b.
Thus, summing over E [t]. we obtain:
Ln
m
Lnd
g
A.2 < 2(1 31
+
b
2(1 B1
b,
For (B.1). using Lemma
E.
have
B.1 <
Σ
t
1/4
(x(x)
B(GT(1)))/I - :3
m
3
a
2(1 - -3133
b,
(G)(8)
m
2
2Vdn
2(1 - 31
b
v(1 3, B2)(1 33
36

---

For B.2, applying vector inequalities and Lemma E.4 we have
81
B.2
A |0010 X, x, Vlyr)) (x.)))
1 1 31
LS5WW
||x; - X3-11 11?
(1 - B1) s=1
1
2
m 1
1 (1 - 31)²
b
==1
Finally, for the upper bound of C, we have
C CSL
L
g.
L
A 1
Ix
xx-1 11
(1 31)2
Lap
g+
2
m 2
+
31)2
b,
(1 - B1)2
ba-1
Therefore, we define
C1 1 B2
+ 3Ld²
2(1 B1)
2(1 31)2²
Lnj
2L
C2 2(1 - 81)2
(1 31)2"
BGT I Br 3GT1133
C3 2(1 31)3 2(1 - 31)33
2/dn
C = V(1 (1 31) - B2) 31/32)
T.C.dd
C5 (1 - 31) - -B1/32)
Then, we have
f(y++1) & f(x1) +
3
1
?
V((,))
va.
2
+
log
T
+
(t) 32
2(1 - -31)
1
b.
g+
2
Ln7
m
LIF
t
g
2
V((xx)
+ 2(1 11)3 1=1
+
+
ns
b
2(1
B1)
b
B
yas
s=1
(GT(1)/V - B3
m,-1
@Gr(t)n33
Σ
m
+ 2(1 - -3133
+
b
2(1 3,)3
b
2/dnGr
V(1- - 31) - 32)(1 - B2)
Ln7²
Σ
g,
12
20772
Σ
m
+
(1 B1)
b
(1 B1)2 s=1
b,-1
s=l
2
N
Vx(x.)
d
log
T
+C1
Σ
g.
+C2
+
m
Va.
b,
b
C3
Σ
m
C.Gr.
b
Then, according to Lemma
E.6
we have
fxx++11111
< f +C2c
< <4L(f(y++1) - - 2C,.
37

---

Thus, we have
(2
V vx(x,)
4Ld
(T)
V <4L(f((()))) - f )++ + 41
B
-
1/8
va.
+
A
log
+ 20g
==1
2
1 Σ
g+ b
+ 2
m.g
+ 4LC
Σ
ma-1
+
4LCCGG.
b
b
Recall Lemma
we have:
g
Fi(t)
log
bg
32
83
mg
1 81
Fi(t)
log
bg
(1 - B2) - 31132)
m,
1 B1
d
(F((l)
log
bs+1
B2
32) B1/B2)
$5
2=1
m
Σ
log
Fi(1)
b
32) 31/32)
.35
Let
2L B2
6Lr/²
F(T))
D1
31)2 +
dlog
(1-31)) 31)3
2121)2(1 + 4w²)
F((T)
D2 (1 1-)311111111 -B22 821332)
dlog
31
2L33Gr(VI 32 (33)
(((T)
D3 B2 - (1) - 81) - 31/32)
dlo
8LGTvdn
D4 V(1 - 31) B2)(1 - 31/32)
4Ld
D3 = log
5
x == 2(1 31)v1-1 B2
3nGrB8
Finally, given 0 < 31 < B2 < 1, 11 Covl - -2, E (3:1) and 3 > 2, we define
83 max 1-9 VT- -1-F
(1+60)
102->2 102- and
(VI+ 1/32 +) max 7, 1/7) With
TVI-B
these definitions, derive that 2² satisfies
-
D2 + D3 + D4 D5 + 2Cg.
(17)
Next, we proceed with a proof by mathematical induction. First, we assume that G. GTT Gr, VI. E [t] Thus.
2
V f((x))
V G² 4L
G²
a
Thus, G:41 max(Gr, | < G. which confirms the validity of the initial hypothesis.
38

---

Since Lemma E.7 and Lemma
E.9 each hold with probability at least 1-8, they hold simultaneously with probability at least 1-28. We have.
|as
max sE[d]
B2V +11 (1- +c
max re[d]
(1- B2)
Σ 32-8) + +GT(s))) te
(1 32)
Σ
85 Gz c GT tc, As E [T].
j=1
Then, combining 17 - Cov I - B2, E (1 - B2) B2), for any 8 E [T], we can obtain
1/8
Cov(1 (1 35)(1 - 32)
> Cov B2
las 0c
GrvI- - 52 + c/(1 (1 - 82)(1 52) 1 1- 3
GT + E 32
Thus,
|
Gr + Gr- -
n
Cov1 B2
Therefore, letting y 2.
n
[f(x,)11 < L
T Σ
178
vf(x,)
G V [f(x++1)12
G²
ag
10000
va,
4(7 2/83
4(7 - 2/3)
==1
V2G²
|10122 G2
cT
[V((x))11
<
+e |
log
4(7 2/B)
1 1- 82
V2C²
|0112 + G T
V 1ogg
eTl'
(o)
1/2),
4(7 - -2
T
where (o) is due Lemma E.1 have G² (pp)))(ogg(T))
G Experimental Details
G.1 Pretraining on CIFAR10
The VIT-27M model 43
undergoes pretraining the CIFAR-10 10 dataset with comprehensive hyperparameter specifications provided in Table Our training protocol employs a base learning rate of 1.5 10 4 coupled with a cosine decay schedule over 200 training epochs. The optimization configuration utilizes Adam parameters with weight decay coefficient = 0.05, numerical stability constants c 1 10-8 and momentum terms $1 0.9, 82 0.95. To maintain training stability while processing large-scale inputs, we implement a batch size of 4096 through gradient accumulation with step size of 20, ensuring memory efficiency without compromising convergence dynamics.
Table 3: Hyperparameters used for training ViT
# Params
31 B2
Learning Rate Weight Decay Batch Size
Warmup Epochs 20
27.6M
0.9
0.95
1.5c-4
0.05
4096
G.2 Pretraining on WikiText-103
The LLaMA2-71M model 466 and Qwen2.5-150M model [7 were pre-trained on the Wikitext-100 103 dataset. Identical learning rates and scheduling protocols were systematically implemented across all
optimizers during the training process. Comprehensive experimental specifications are tabulated in
Table and T
39

---

Table 4: Hyperparameters used for training LLaMA2-7IM on WikiText-103
Adam-Type
Lion-Type
Muon-Type
Model Size
71M
Hidden Size
512
Head
8
Depth
12
Training Steps
2034
Warmup Steps
203
Maximum Length
1024
Batch Size
480
Learning Rate
3e-4
Warmup Scheduling
linear from 3e-5
Learning Rate Scheduling
cosine to 10%
Numerical precision
b
Weight Decay
0.01
31
0.9
0.9
0.95
B2
0.999
0.99
0.95
Momentum
X
X 0.95
Table 5: Hyperparameters used for training Qwen2.5-150M on WikiText- 103
Adam- Muon-Type
Model Size
150M
Hidden Size
640
Head
10
Depth
12
Training Steps
1525
Warmup Steps
154
Maximum Length
1024
Batch Size
160
Learning Rate
le-3
Warmup Scheduling
linear from 6e-5
Learning Rate Scheduling
cosine to 10%
Numerical precision
6
Weight Decay
0.01
31
0.95
32
0.95
0.95
Momentum
X
0.95
G.3 Fine-Tuning on GLUE
We fine-tune the pre-trained RoBERTa-Bass model 48 on the GLUE benchmark using the Hug- ging Face implementationfl For all tasks except QQP, we employ a batch size of 32, while QQP uses larger batch size of 128 due to its dataset characteristics. The model is trained uni- formly for epochs across all tasks with maximum sequence length of 512. For each task, we perform a grid search over learning rates. For most optimizers, the learning rate range is {1c-5, 2c-5, 3c-5, 4c-5, 5e-5} and the weight decay is set to 0,01. For Lion-type optimizers, the learning rate range (1e-6, 2e-6, 3e-6, 4e-6, 5e-6) and the weight decay set to 0.1. The complete hyperparameter configurations are summarized in Table 6 For MGUP-AGE Addm W, A, Adam-mini, and C-Adam W, we use 31 0.9 and B2 0.999. For LDAdam and Galore, we use 0.9 and 32 = 0.99. For Lion, C-Lion, and MGUP-Lion, we use B = 0.95 and 32 0.98.
html
40

---

Table Hyperparameters used for fine-tuning on GLUE.
Hyperparameter
MRPC
STS-B
CoLA
RTE
SST-2
QNLI
QQP
Batch Size
32
32
32
32
32
32
128
Weight Decay (Most)
0.01
Weight Decay (Lion-type)
0.1
Epochs
3
Max Seq Len
512
G.4
Fine-Tuiig on GSM8K
We fine-tune the pre-trained LLaMA2-7B model [66 using the IIm-foundry codebase" with evalu- ation via standardized am-evaluation-hantn on the GSM8K benchmark with the Hugging Face implementatior" [6] The fine-tuning process employs consistent hyperparameters across all optimizers, including MGUP-Adam with total of 702 training steps, including 20 warm-up steps. The batch size is set to 32, and the maximum sequence length 512. We use learning rate of 5e-5 and optimizer parameters 31 0.9 and B2 0.999. The complete hyperparameter configurations are summarized in Table
Adam-8bit, Adam W, and C-Adam Specifically, we train for 3 epochs
Table 7: Hyperparameter configurations for fine- fine-tuning LLaMA2-7B on GSM8K.
Hyperparameter
Value
Epochs
3
Training Steps
702
Warm-up Steps
20
Batch Size
32
Maximum Length
512
Learning Rate
5e-5
31
0.9
B2
0.999
H Other MGUP-type Algorithms
Algorithm 3 MGUP-Lion
Input: Learning rate 7) > 0. initial parameters X E Rd loss function f(x). 31. 32 [0, 10, weight decay coefficient stability term >> 0, ratio 7 (0, 1). for t = 1 to T" do
momentum factors
Compute the stochastic gradient g: (x(;())
5,g m + (1)g:)
m; 32m; 1 + 82)gr
Φ: MGUP(u; g:)
x; = (1 A) O Xx
X++1 Xx th Ur
end for
41

---

Algorithm 4 MGUP-Muon
Input: Learning rate 7n 0, initial parameters X. E R=*x, loss function f(X), momentum factors 3 E [0, 1), weight decay coefficient X, ratio 7 E (0,
for t 1 to I' do
Compute the stochastic gradient G. = vf((;;) Et)
M, MM,-1 G2
Φ = MGUP(M, O G
U, (Nwton-Schulz(MM,)
X, (1 n)) @ X
X++1 X. ne: 3 U.
end for
I
More Results
As shown in Figure 7 which depicts the training curves under learning rate of 5e-5, MGUP-AdamW achieves lower training loss per epoch, outperforming baseline optimizers.
1.0
C-AdamW
AdamW-8Bit
AdamW
0.8
MGUP-AdamW
Training Loss
0.6
0.4
0.2
100
200
300
400
500
600
700
Iteration
Figure 7: Training curves of LLaMA2-7B on GSM-8K.
L.1 Memory cost
Table 8 reports memory for fine-tuning LLaMA-7B on two NVIDIA V100 32GB GPUs. We use miero batch size of for GSM8K fine-tuning. Our algorithm introduces transient elevation in peak memory overhead while maintaining unaltered static memory allocation throughout the computational process.
Table 8: Memory for FineTuing LLaMA2-7B GSM8K
Adam-8bit
AdamW
C-AdamW
MGUP-AdamW
Peak Reserved Memory
25.38GB
32.09GB
32.98GB
33.82 GB
42

---

1.2
Time cost
Table 9 documents the runtime measurements for both fine-tuning and pre- pre-training processes con- ducted primarily on single NVIDIA RTX 4090 24GB GPU, with one exceptional case: LLaMA2-7B fine-tuning utilizing two NVIDIA V100-32GB GPUs. For GSM8K fine-tuning experiments, maintained micro-batch size of throughout the process. In pre-training configurations, gradient accumulation strategy was implemented to optimize memory utilization.
Remark L.I. In all experiments, documented the duration from initiation completion rather than the algorithm's 's execution time. Discrepancies may arise due to variations in GPU operational states.
Table 9: Runtime for FineTTuing((I)) and PreTraininggPP) tasks
Model
Task
AdamW
C-AdamW
MGUP-AdamW
ViT-28M
CIFAR IO(PT)
Ih5m
Ih 6m
Ih 6m
LLaMA2-7IM
SoeTText-133PP)
5h 36m
5h 37m
5h 37m
1-
WikiText- (03(TT)
54m
1h55m
1h56m
RoBERTa- Base
QQP(FT)
25m
30m
34m
LLaMA2-7B
GSM8K(FT)
15h 53m
21h 37m
22h 58m
43

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
In the introduction, we highlighted our contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
In the conclusion section, we discuss the limitations our work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
44

---

Each theoretical result is accompanied by hyperlink that links directly to its complete proof in the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide all necessary information for reproducing the main experimental results of this paper in the Experimental Details section of the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
45

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification:
We provide code and data links in the supplemental materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ priips.cc/ubbliccguu.eee CoodeSumeeieeeeee ssi iessonPolll mollicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all necessary information for the experimental setting in the Experimental Details section of the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: The statistical significance of the experimental results beyond the scope of this work.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
46

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide the necessary information on computer resources in the More Results section of the Appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't ' make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics Ehhnshhrrrr/////rrrr........./op
Answer: [Yes]
Justification: Our research aligns with the ethical guidelines outlined by NeurIPS.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: This work is related to optimization techniques, however, both positive and negative social impacts will stem from the specific models and applications developed using this optimization technique, rather than the optimization technique itself. For these reasons, this basic research paper does not contain a detailed discussion of specific social impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
47

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This work poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We cite the original papers that generated the code packages or datasets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
48

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta
latasets.papersstt
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper introduces novel optimization algorithm and provides its pseu- docode within the manuscript. No new datasets, standalone software packages, or pre-trained models are released as separate assets alongside this work.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooell as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This work does not include Crowdsourcing and research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification:
This work does not include Crowdsourcing and research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
49

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: This work does not involve L.M as any significant, original or non-standard component.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https://neurd) htt.ep///eerrrpp.cc///oofeeeeeee /tt//////ooeeeeeee/222////.... for what should or should not be described.
50


Reviews:
[{"review_id": "bDATRWP9Wt", "reviewer": "Reviewer_xAKK", "summary": "This 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.", "questions": "Regarding 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?", "limitations": "I do not see any potential negative societal impact in addition to the development of LLMs themselves.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Pros:\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.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s thoughtful comments and suggestions.\n\n---\n\n**Q1.** MGUP introduces a key hyperparameter, $\\tau$, which defines the proportion of parameters to receive amplified updates. While the authors provide a sensitivity analysis (Figure 3b), the choice of $\\tau$ appears to be data and model-dependent, which could add to the tuning burden for practitioners.\n\n**A1.** We would like to highlight MGUP's robustness to the hyperparameter $\\tau$. As depicted in the sensitivity analysis in Figure 3b, simply setting $\\tau=0.5$ provides strong performance across all tasks.\n\n---\n\n**Q2.** The 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\n**A2.** Our theoretical analysis focuses on MGUP-Adam, as Adam remains one of the most popular optimizers, making the convergence guarantees for MGUP-Adam meaningful and timely. \n\nWe believe the analysis may serve as a foundation for future theoretical investigations into MGUP-Lion and MGUP-Muon. \n\nAt the same time, this focus helps preserve the clarity and conciseness of the theoretical exposition in the paper.\n\n---\n\n**Q3.** Regarding the hyperparameter $\\tau$, 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\n**A3.** In some tasks, we observed that MGUP is not highly sensitive to the choice of $\\tau$. A fixed value of $\\tau=0.5$ generally performs well across tasks. \n\nWhile the optimal value may vary slightly depending on the specific setting (e.g., task type, model size, or base optimizer), strong performance is typically maintained within the range $\\tau \\in [0.3,0.7]$, suggesting that extensive tuning is unnecessary in practice."}, {"review_id": "wXyJiAW4cn", "reviewer": "Reviewer_Vt5g", "summary": "This 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.", "questions": "Please address the weakness.\n\nIt is unclear to me how the theoretical analysis reflects the unique properties of the proposed algorithm.", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\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.", "was_revised": true, "final_justification": "The response of the authors addressed my initial concern. Therefore, I would raise my socre.", "rebuttal": "Thank you for the reviewer’s valuable comments. Below we reply point‑by‑point, using the reviewer’s wording for clarity.\n\n---\n\n**Q1.** 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.\n\n**A1.** First, we provide below a counterexample, based on Section 3.2 of (Zhang, Y., Chen, C., Shi, N., Sun, R., & Luo, Z. Adam Can Converge Without Any Modification on Update Rules. NeurIPS 2022), which demonstrates that setting $\\gamma=0$ can lead to non-convergence in practice.\n\nConsider $f(x) = \\sum_{i=0}^{n-1} f_{i}(x)$ for $x\\in \\mathbb{R}$, and then we define $f_i(x)$ as \n\n$\\begin{align*}\nf_i(x) &=\n\\begin{cases}\nnx, & x\\geq-1 ;\\\\\n\\frac{n}{2}(x+2)^2-\\frac{3n}{2}, & x<-1\n\\end{cases}\n& \\text{for } i=0,\\end{align*}\n$\n\n\n$\\begin{align*}\nf_i(x) &=\n\\begin{cases}\n-x, & x\\geq-1 ;\\\\\n-\\frac{1}{2}(x+2)^2+\\frac{3}{2}, & x<-1\n\\end{cases}\n& \\text{for } i>0.\n\\end{align*}$\n\nso\n\n$\nf(x)=\n\\begin{cases}\nx, & x\\geq-1 ;\\\\\n\\frac{1}{2}(x+2)^2-\\frac{3}{2}, & x<-1\n\\end{cases}\n$\n\n\nThe optimum is $x^*=-2$.\n\nWe construct a counterexample to demonstrate the non-convergence of Cautious-Adam in a specific stochastic environment. We initialize at $x = -0.5$. Let's analyze what happens when $x\\ge-1$.\n\n1.Dynamic Analysis of Momentum ($m_t$)\n\nIn this counterexample, the behavior of the momentum $m_t$ follows a \"pulse-decay\" model:\n\n（1）Pulse:A positive gradient $g_t = n$ with an extremely low probability pushes $m_t$ to a large positive value.\n\n（2）Decay: A negative gradient $g_t = -1$ with a high probability causes $m_t$ to continuously decay towards its steady state of -1. This process inevitably leads to $m_t$ crossing zero into the negative region.\n\nThis oscillatory cycle results in the momentum spending roughly equal amounts of time on either side of zero. Therefore, we can reasonably approximate that **$P(m_t > 0) \\approx 0.5$.**\n\n2.Decomposing the Behavior of Cautious-Adam in This Scenario\n\nBased on the assumption of momentum oscillation, we analyze the four core behaviors of Cautious-Adam:\n\n（1）Case A: A positive gradient $g_t = n$ is sampled (low probability of $1/n$)\n\n（1.1）A.1 ($m_t > 0$, probability $\\approx 0.5$): The directions align, and a correct update is performed. (Total probability $\\approx 0.5/n$)\n\n（1.2）A.2 ($m_t < 0$, probability $\\approx 0.5$): The directions are opposed, and the update is skipped. (Total probability $\\approx 0.5/n$)\n\n（2）Case B: A negative gradient $g_t = -1$ is sampled (high probability of $(n-1)/n \\approx 1$)\n\n（2.1）B.1 ($m_t > 0$, probability $\\approx 0.5$): The directions are opposed, and the update is skipped. (Total probability $\\approx 0.5$)\n\n（2.2）B.2 ($m_t < 0$, probability $\\approx 0.5$): The directions align, and an incorrect update (moving away from the optimum) is performed. (Total probability $\\approx 0.5$)\n\n3.Conclusion\n\nThe behavior of Cautious-Adam is dominated by high-probability events. Ultimately, the dynamics of the algorithm are as follows:\n\n(1)With a probability of $\\approx 50\\%$, the update is skipped.\n\n(2)With a probability of $\\approx 50\\%$, the algorithm updates in the wrong direction.\n\n(3)Only with an extremely low probability of $\\approx 0.5/n$ does the algorithm update in the correct direction.\n\nTherefore, not only does Cautious-Adam's convergence stagnate, but it also actively moves away from the optimal point with high probability. This provides strong evidence for its non-convergence. In contrast, MGUP avoids this issue by setting $\\gamma > 0$, which ensures that it still progresses in the correct direction in case B.1.\n\nWe set n=10 and conducted an empirical study. The data we recorded are shown below (due to character limits, we only display some of the updates, up to a total of 900 steps). As can be seen, Cautious-Adam gradually diverges from the optimal point of -2, whereas MGUP-Adam is able to converge to -2. This is consistent with our analysis. \n\n| Step | Cautious-AdamW x           | MGUP-AdamW x |\n| ---- | ------------------- | ------------ |\n| 0    | -0.5                | -0.5         |\n| 50   | -0.178418443        | -0.710841    |\n| 100  | -0.005119875        | -0.767780244 |\n| 150  | 0.35917407274246216 | -0.48992154  |\n| 200  | 0.583101749420166   | -0.627167225 |\n| …    |                     |              |\n| 700  | 2.395268678665161   | -1.837388873 |\n| 750  | 2.733309268951416   | -1.848099828 |\n| 800  | 2.945746898651123   | -1.914497256 |\n| 850  | 3.069368839263916   | -1.959119558 |\n| 900  | 3.2514820098876953  | -1.992537737 |\n\n---\n\n**Q2.** 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.\n\n**A2.** Thank you for pointing out the imprecision. This assumption is mild as $|| \\nabla f(x;\\xi)|| \\le || \\nabla f(x;\\xi) - \\nabla f(x) || + || \\nabla f(x)||  < +\\infty$ and $f(x)$ is a continous function.\n\n---\n\n**Q3.** 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.\n\n**A3.** Thank you for your suggestion. We conducted an additional four sets of experiments with different random seeds for AdamW, Cautious-AdamW, and MGUP-AdamW on the Qwen2.5 training task, doubling the number of iterations and calculating their standard deviations. Please refer to the table below. Our setup utilizes a single NVIDIA RTX 4090 GPU.\n\n|            | Avg loss | Std         |\n| ---------- | -------- | ----------- |\n| MGUP-AdamW | 2.748    | 0.005223205 |\n| Cautious-AdamW    | 2.787    | 0.003388492 |\n| AdamW      | 2.778    | 0.006125102 |\n\n---\n\n**Q4.** The analysis uses constant and , 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. Minor point: 4. Why choose $\\alpha = 1/\\tau$ and $\\gamma=\\tau$? Is there any specific reason? It is also a bit confusing since $\\alpha,\\gamma$ do not appear in Algorithms 1 and 2.\n\n**A4.** Thank you for raising these important questions.\n\n1.We will first address the first question: **\"It does not explain why MGUP has better performance. Also, it does not provide insight on how to choose to accelerate the convergence of MGUP.\"**\n\n(1)**MGUP's mechanism uses a greedy strategy (due to sorting).** Greedy approaches play a crucial accelerating role in heuristic algorithms. When the stochastic gradient and the current update have the same sign, their product is positive, making them more likely to be selected for a large step-size update. Conversely, when they have opposite signs, their product is negative, making them more likely to be selected for a small step-size update. Large step-size updates drive acceleration, while small step-size updates ensure convergence.\n\n(2)**MGUP increases the average update magnitude.** In MGUP, a proportion of $\\tau$ parameters execute updates with an increased learning rate of $1/\\tau \\cdot \\text{lr}$, while a proportion of $1-\\tau$ parameters execute updates with a reduced learning rate of $\\tau \\cdot \\text{lr}$. This means that when $0 < \\tau < 1$, the average learning rate is $\\tau \\cdot \\text{lr} \\cdot (1/\\tau) + (1-\\tau) \\cdot \\text{lr} \\cdot \\tau = \\text{lr} \\cdot (1+\\tau-\\tau^2) > \\text{lr}$. If the step size for each parameter is roughly equivalent, for example, when Adam's updates are approximately $\\text{sign}(g_t)$ in the initial phase, this implies that the overall magnitude of MGUP-Adam's updates is increased by a factor of $(1+\\tau-\\tau^2)$ compared to Adam. This is a potential underlying intuition behind MGUP-Adam's acceleration: MGUP can accelerate updates by an approximate factor of $(1+\\tau-\\tau^2)$.\n\n2.Next, we address the second minor point: **\"Why choose $\\alpha = 1/\\tau$ and $\\gamma=\\tau$? Is there any specific reason? It is also a bit confusing since $\\alpha,\\gamma$ do not appear in Algorithms 1 and 2.\"**\n\n$\\alpha$ represents the factor for increasing the step size, and $\\gamma$ represents the factor for decreasing the step size. In experiments, once the threshold is determined, the choice of $\\gamma$ is robust, as shown in **Figure 3(b)** of the paper. Designing $\\alpha$ and $\\gamma$ to be related to $\\tau$ aims to reduce hyperparameter tuning and simplify the design, resulting in the final algorithm having only one hyperparameter.\n\n---\n\n**Q5.** It is unclear to me how the theoretical analysis reflects the unique properties of the proposed algorithm.\n\n**A5.** Thank you for raising this important question. \n\n(1) MGUP incorporates a protective mechanism where $\\gamma>0$ ensures convergence, a feature absent in existing algorithms like Cautious AdamW. \n\n(2) Furthermore, our approach to expected convergence and the underlying proof methodology differ from current convergence proofs, representing an additional contribution. Unlike existing work that commonly relies on specific treatments of $\\beta_2$ (e.g., decay), our Theorem 4.1 establishes convergence by decaying the first moment parameter $\\beta_1$, which constitutes a unique aspect of our proof."}, {"review_id": "zq5liQNjMn", "reviewer": "Reviewer_kuKg", "summary": "This 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.", "questions": "1. 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?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I am generally satisfied with the rebuttal, thus raise my score to 4.", "rebuttal": "We sincerely thank the reviewer for his/her constructive feedback.\n\n---\n\n**Q1.** The theoretical illustration for using nonzero 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\n**A1.** Thank you for raising this important point.\nTo show non-convergence when $\\gamma=0$, we'll use a counterexample from the **literature [1]** section 3.2.\n\nConsider $f(x) = \\sum_{i=0}^{n-1} f_{i}(x)$ for $x\\in \\mathbb{R}$, and then we define $f_i(x)$ as \n\n$\\begin{align*}\nf_i(x) &=\n\\begin{cases}\nnx, & x\\geq-1 ;\\\\\n\\frac{n}{2}(x+2)^2-\\frac{3n}{2}, & x<-1\n\\end{cases}\n& \\text{for } i=0,\\end{align*}\n$\n\n\n$\\begin{align*}\nf_i(x) &=\n\\begin{cases}\n-x, & x\\geq-1 ;\\\\\n-\\frac{1}{2}(x+2)^2+\\frac{3}{2}, & x<-1\n\\end{cases}\n& \\text{for } i>0.\n\\end{align*}$\n\nso\n\n$\nf(x)=\n\\begin{cases}\nx, & x\\geq-1 ;\\\\\n\\frac{1}{2}(x+2)^2-\\frac{3}{2}, & x<-1\n\\end{cases}\n$\n\nThe optimum is $x^*=-2$.\nWe initialize at $x = -0.5$. Let's analyze what happens when $x\\ge-1$ for Cautious-Adam.\n\n1.Dynamic Analysis of Momentum ($m_t$)\nIn this counterexample, momentum ($m_t$) follows a \"pulse-decay\" model: rare, large positive gradients push $m_t$ high, followed by frequent negative gradients causing it to decay, inevitably crossing zero. This oscillation leads to $P(m_t>0)\\approx 0.5$.\n\n2.Decomposing the Behavior of Cautious-Adam in This Scenario\nBased on the assumption of momentum oscillation, we analyze the four core behaviors of Cautious-Adam:\n\n(1)Case A: A positive gradient $g_t = n$ is sampled (low probability of $1/n$)\n\n(1.1)A.1 ($m_t > 0$, probability $\\approx 0.5$):The directions align, and a correct update is performed. (Total probability $\\approx 0.5/n$)\n\n(1.2)A.2 ($m_t < 0$, probability $\\approx 0.5$): The directions are opposed, and the update is skipped. (Total probability $\\approx 0.5/n$)\n\n(2)Case B: A negative gradient $g_t = -1$ is sampled (high probability of $(n-1)/n \\approx 1$)\n\n(2.1)B.1 ($m_t > 0$, probability $\\approx 0.5$):The directions are opposed, and the update is skipped. (Total probability $\\approx 0.5$)\n\n(2.2)B.2 ($m_t < 0$, probability $\\approx 0.5$):The directions align, and an incorrect update (moving away from the optimum) is performed. (Total probability $\\approx 0.5$)\n\n3.Conclusion\n\nThe behavior of Cautious-Adam is dominated by high-probability events. Ultimately, the dynamics of the algorithm are as follows:\n\n(1)With a probability of $\\approx 50\\%$, the update is skipped.\n\n(2)With a probability of $\\approx 50\\%$, the algorithm updates in the wrong direction.\n\n(3)Only with an extremely low probability of $\\approx 0.5/n$ does the algorithm update in the correct direction.\n\nTherefore, Cautious-Adam not only stagnates but also diverges from the optimal point with high probability, demonstrating its non-convergence. In contrast, MGUP avoids this by setting $\\gamma>0$, ensuring continued progress towards the optimum. Our empirical study with n=10 shows Cautious-Adam diverging from the optimal point of -2, while MGUP-Adam converges to -2, consistent with our analysis. (Due to character limits, you can view partial results in our response to the second Reviewer Vt5g.)\n\nThen, regarding your suggestion :**'We basically need a lower bound'** .\n\nOur theoretical results indicate that $\\gamma$ only needs to be a strictly positive constant, thus making this infimum unnecessary. We believe that this concrete counterexample more clearly reveals the inherent instability of setting $\\gamma=0$ than a generalized theoretical lower bound would. \n\nRegarding your suggestion: **'Also, I think it would be better to also talk about work related to the convergence theory’**.\n\nThank you for your suggestion. We'll summarize the development of Adam's convergence theory and additionally cite articles [1], [2], [3], [4], [5], and others. Unlike existing work that commonly relies on specific treatments of $\\beta_2$ (e.g., decay), our Theorem 4.1 establishes convergence by decaying the first moment parameter $\\beta_1$, which constitutes a unique aspect of our proof.\n\n[1] Zhang, Y., Chen, C., Shi, N., Sun, R., & Luo, Z. (2022). Adam Can Converge Without Any Modification on Update Rules. \n\n[2] Zou, F., Shen, L., Jie, Z., Zhang, W., & Liu, W. (2019). A sufficient condition for convergences of adam and rmsprop.\n\n[3] Wang, B., Fu, J., Zhang, H., Zheng, N., & Chen, W. (2023). Closing the Gap Between the Upper Bound and the Lower Bound of Adam's Iteration Complexity. \n\n[4] Chen, X., Liu, S., Sun, R., & Hong, M. (2018). On the Convergence of A Class of Adam-Type Algorithms for Non-Convex Optimization. \n\n[5] He, M., Liang, Y., Liu, J., & Xu, D. (2023). Convergence of adam for non-convex objectives: Relaxed hyperparameters and non-ergodic case. \n\n---\n\n**Q2.** 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\n**A2.** We emphasize MGUP's empirical superiority through its efficiency, robustness, and extensive validation. MGUP significantly enhances robustness and stability, reducing learning rate sensitivity and maintaining training stability, outperforming methods like C-Lion. Our findings are broadly and diversely validated across diverse model architectures (ViT, LLaMA, MoE), scales (27M-7B parameters), and tasks (pre-training, NLU, mathematical reasoning fine-tuning). This multifaceted improvement and consistent validation offer compelling evidence for our method.\n\nNext, we will focus on elaborating the core advantages of MGUP over the Cautious Optimizer and Adam in response to your follow-up questions.\n\n**1.MGUP-Adam vs. Cautious Adam.**\n\n(1) Cautious Adam lacks theoretical results for stochastic optimization, whereas we ensure global convergence by employing a safeguard parameter.\n\n(2) The inability of cautious optimizers' methods to guarantee convergence under extreme conditions accounts for their lack of theoretical results, as illustrated by the counterexample in **A1**.\n\n(3) MGUP-Adam sets a threshold $\\tau$, to ensure that the number of parameters being updated does not reach extreme cases. For instance, an extreme case could involve updates to merely 1% of the parameters.\n\n**2.MGUP-Adam vs. Adam.**\n\n(1)MGUP's mechanism uses a greedy strategy (due to sorting). Greedy approaches play a crucial accelerating role in heuristic algorithms. When the stochastic gradient and the current update have the same sign, their product is positive, making them more likely to be selected for a large step-size update. Conversely, when they have opposite signs, their product is negative, making them more likely to be selected for a small step-size update. Large step-size updates drive acceleration, while small step-size updates ensure convergence.\n\n(2)MGUP increases the average update magnitude. In MGUP, a proportion of $\\tau$ parameters execute updates with an increased learning rate of $1/\\tau \\cdot \\text{lr}$, while a proportion of $1-\\tau$ parameters execute updates with a reduced learning rate of $\\tau \\cdot \\text{lr}$. This means that when $0 < \\tau < 1$, the average learning rate is $\\tau \\cdot \\text{lr} \\cdot (1/\\tau) + (1-\\tau) \\cdot \\text{lr} \\cdot \\tau = \\text{lr} \\cdot (1+\\tau-\\tau^2) > \\text{lr}$. If the step size for each parameter is roughly equivalent, for example, when Adam's updates are approximately $\\text{sign}(g_t)$ in the initial phase, this implies that the overall magnitude of MGUP-Adam's updates is increased by a factor of $(1+\\tau-\\tau^2)$ compared to Adam. This is a potential underlying intuition behind MGUP-Adam's acceleration: MGUP can accelerate updates by an approximate factor of $(1+\\tau-\\tau^2)$.\n\n---\n\n**Q3.** 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\n**A3.** We understand your concern. However:\n\n(1) The complexity is approximately linear, and the sorting for TopK is O(dlogK). We can set K to an acceptably small value, in which case the complexity remains nearly unchanged.\n\n(2) Our experiments in Qwen2.5 show that despite 4 random trials with the same iterations, sorting has no significant impact on total execution time. Crucially, our approach yields a superior objective function result. Our setup utilizes a single NVIDIA RTX 4090 GPU.\n\n|            | average time   |\n| ---------- | ----------- |\n| MGUP-AdamW | 4h20m49s |\n| Cautious-AdamW    | 4h20m40s |\n| AdamW      | 4h20m03s |\n\n---\n\n**Q4.** MGUP-Muon seems to have a different implementation compared to MGUP-Lion and MGUP-Adam, since it does MGUP to compute  before computing the update . Could you explain this difference?\n\n**A4.** All three of our algorithms adhere to the same framework: parameter selection is based on the dot product of the mini-batch gradient and the EMA  of the gradient. In Adam and Lion, the element-wise direction of the momentum and the update vector are consistent. However, this is not the case in Muon, because Muon introduces a gradient whitening step, $(G G^\\top)^{-1/2}$. Specifically, the process $M_t \\stackrel{SVD}{=}V_1SV_2^\\top \\Rightarrow U_t = V_1V_2^\\top$ is likely to change the element-wise direction of the update vector, which is inconsistent with our momentum-gradient alignment concept. Therefore, we use $s=M_t \\odot G_t$ instead of $s=U_t \\odot G_t$."}, {"review_id": "U7aLPoqFHa", "reviewer": "Reviewer_go7i", "summary": "This paper proposes a new monentum-gradient alignment policy for stochastic optimization. \nIt provides theoretical guarantees of the proposed methods, and evaluates the performance on a few benchmark NN training tasks.", "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?\n\n2. 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. \n\n3. 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "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.\n\nWeakness: As mentioned by the authors, the proposed algorithm combines the idea of a few existing methods. The contribution is a bit incremental.", "was_revised": true, "final_justification": "Thank you for the response. I will keep my rating.", "rebuttal": "We thank the reviewer for his/her thoughtful comments and suggestions.\n\n---\n\n**Q1.** 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?\n\n**A1.** We will now elaborate on the intuition for choosing these scale factors.\n\n(1) MGUP increases the average update magnitude. In MGUP, a proportion of $\\tau$ parameters execute updates with an increased learning rate of $1/\\tau \\cdot \\text{lr}$, while a proportion of $1-\\tau$ parameters execute updates with a reduced learning rate of $\\tau \\cdot \\text{lr}$. This means that when $0 < \\tau < 1$, the average learning rate is $\\tau \\cdot \\text{lr} \\cdot (1/\\tau) + (1-\\tau) \\cdot \\text{lr} \\cdot \\tau = \\text{lr} \\cdot (1+\\tau-\\tau^2) > \\text{lr}$. If the step size for each parameter is roughly equivalent, for example, when Adam's updates are approximately $\\text{sign}(g_t)$ in the initial phase, this implies that the overall magnitude of MGUP-Adam's updates is increased by a factor of $(1+\\tau-\\tau^2)$ compared to Adam. This is a potential underlying intuition behind MGUP-Adam's acceleration: MGUP can accelerate updates by an approximate factor of $(1+\\tau-\\tau^2)$.\n\n(2) With the goal of reducing hyperparameter tuning, we recognize that introducing two separate hyperparameters for adjustment $\\alpha$ for increasing the step size and $\\gamma$ for decreasing it—could add to the burden of practical application. In our experiments, however, we found that simply choosing a $\\tau$ in the range of 0.3 to 0.7 is sufficient to achieve good results.\n\n---\n\n**Q2.** 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.\n\n**A2.** Extending our convergence results to non-smooth scenarios is a highly valuable but also challenging theoretical task. We are confident that this approach can be extended to our method, though other issues may arise in the process.\n\n---\n\n**Q3.** 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.\n\n**A3.** Thank you for your suggestion, which helps to improve the clarity of our theoretical results. However, we would like to clarify that we have already stated in the **Theorem 4.2** that the convergence holds with a probability of at least $1−2\\delta$, where $\\delta\\in(0,1/2)$ is a user-selected confidence parameter. We use the notation $\\tilde{O}$ to hide polylogarithmic factors. For the specific formula, please refer to the result above line 626 in the appendix."}]

Author rebuttal:

Meta-review: This paper introduces MGUP, a mechanism for intra-layer selective updates which comes with strong theoretical guarantees as well. The proposed method is a plug-and-play module, and the authors combine it with AdamW, Muon, and Lion and demonstrate its effectiveness in various experiments.

Reviewers all agree that the theory is very strong. Even though there are some concerns about the results not being much better than AdamW, I do not think that is a major issue. There was a concern about the work being incremental (reviewer g07i), but I do not think it is a valid reason to reject this paper. It is an interesting method with strong theoretical guarantees. I recommend accepting this paper.
Decision: Accept (spotlight)