Enhancing Optimizer Stability: Momentum Adaptation of The NGN Step-size
Authors: Rustem Islamov
Niccolò Ajroldi
Antonio Orvieto
Aurelien Lucchi
Venue: NeurIPS 2025 poster ()
Keywords: Adaptive Algorithms
Theoretical Analysis of Adaptive Methods
Polyak Stepsize
Stability
Primary area: optimization
TL;DR: We propose a novel algorithms, called NGN-M with a strong theoretical convergence analysis and extensive numerical evaluations showing the robustness of our algorithm to the choice of the learning rate hyperparameter.

Abstract: Modern optimization algorithms that incorporate momentum and adaptive step-size offer improved performance in numerous challenging deep learning tasks. However, their effectiveness is often highly sensitive to the choice of hyperparameters, especially the step-size.  Tuning these parameters is often difficult, resource-intensive, and time-consuming. Therefore, recent efforts have been directed toward enhancing the stability of optimizers across a wide range of hyperparameter choices [Schaipp et al., 2024]. In this paper, we introduce an algorithm that matches the performance of state-of-the-art optimizers while improving stability to the choice of the step-size hyperparameter through a novel adaptation of the ${\sf NGN}$ step-size method  [Orvieto and Xiao, 2024]. Specifically, we propose a momentum-based version ${\sf NGN}$-${\sf M}$ that attains the standard convergence rate of $\mathcal{O}(1/\sqrt{K})$ under less restrictive assumptions, without the need for interpolation condition or assumptions of bounded stochastic gradients or iterates, in contrast to previous approaches. Additionally, we empirically demonstrate that the combination of the ${\sf NGN}$ step-size with momentum results in enhanced robustness to the choice of the step-size hyperparameter while delivering performance that is comparable to or surpasses other state-of-the-art optimizers.

Enhancing Optimizer Stability: Momentum Adaptation of The NGN Step-size
Rustem Islamov! Niccoló Ajroldi² Antonio Orvieto Aurelien Lucchi
University of Basel Max Planck Institute for Intelligent Systems
3 ELLIS Institute Tübingen 'Tübingen Al Center
Abstract
Modern optimization algorithms that incorporate momentum and adaptive step- size offer improved performance in numerous challenging deep learning tasks. However, their effectiveness is often highly sensitive to the choice of hyperpa- rameters, especially the learning rate (LR). Tuning these parameters is often dif- ficult, resource-intene been directed toward enhancing the stability of optimizers across wide range hyper-parameter choices [79]. this paper, we introduce an algorithm that matches the performance of bate-of-thhear optimizers while improving stability through a novel adaptation of the NGN step-size method [66]. Specifically, we propose momentum-base, version (NGN-M) that attains the standard convergence rate of O(1/V/) under common assumptions, without the need for interpolation condition or assumptions of bounded stochastic gradients or iterates, in contrast previous approaches. Additionally, bination of the NGN step-size with momentum results in high robustness while delivering performance that is comparable or surpasses other state-of-the-a optimizers.
and time-comee consuminn.ing. Therefore, recent efforts have
empirically demonstrate that the com-
1 Introduction
Adaptive methods such as Adam [44] and RMSprop [30] are widely used in machine learning due to their established advantages over (momentum) SGD, particularly in tasks such as training Transformers [8, 88, 89]. These methods adaptively scale the step-size across different dimensions (parameters) based on their respective statistics, effectively acting as a diagonal preconditioning.
Although these methods perform well in practice, existing theoretical analyses typically require strin- gent assumptions on the noise structure of the stochastic gradients, such as sub-Gaussian noise [49] or affine noise models [90, 106]: Relaxing these assumptions remains an open challenge. Another well- known issue of Adam is its performance sensitivity to the LR hyperparameter [96, 10], particularly when training Transformers, where loss spikes are commonly observed [60, 97]. This often necessi- tates careful adjustments of the hyperparameters throughout the training process [107, 11], which can be costly in terms of computational resources [64]. Consequently, there has been growing interest in developing optimization methods that are more robust to hyperparameter selection [79]. In addition
to adapting LR, Adam and other oatt--of-the-ar optimizers also rely on momentum [71], a broadly
used technique that has been shown to enhance performance both theoretically [14, 19, 35, 36) and
practically [10, 21, 38]. Besides speeding up convergence, momentum is known as technique to reduce the variance of stochastic algorithms [57, 15], improving stability as well as generalization in some settings [38].
In this work, we address the aforementioned drawbacks of Adam by developing new algorithm based on the recently proposed NGN step-size [66]. an improved variant of the Stochastic Polyak Step-size [54] that has demonstrated strong resilience to LR hyperparameter tuning. In particular,
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

NGN was shown never to diverge for any choice of the LR hyperparameter in the convex setting, and exhibit strong curvature adaptation properties strengthened by theoretical guarantees. However, the step-size of Orvieto and Xiao [66] simply adapts the LR through a scalar multiplier, leaving to future work the incorporation of momentum and coordinate-wis variants - needed in complex problems such optimizing transformers, as motivated above. Here, we develop momentum and step-size adaptive version of NGN designed to enhance robustness' in terms of hyperparameter selection. We also present theoretical analysis alongside a practical evaluation of this approach, showcasing its improvements over current state-of-the-art methods.
In summary, our contributions are as follows:
1. We introduce a new algorithm named NGN-M that combines the NGN step-size with momen- um. We theoretically show that NGN-M - achieves convergence rate (11/RR) in the convex regime without the typical requirements of interpolation or bounded gradient assumptions found in earlier works on Polyak step-size;
2. We focus on the problem of adapting the step-size rule towards a coordinate wise diagonal preconditioning. By integrating this diagonal step-size strategy with momentum, we develop a new variant of NGN, called NGN-MD;
3. The theoretical results are supported by extensive empirical validation in various deep learning settings where we demonstrate that NGN-M and NGN-MD not only preserve the robustness property of the NGN step-size, but improve it further in many cases. LR hyperpa- rameter resilience comes together with better or comparable performance state-of-the-art algorithms.
2
Related Works
Polyak Step-size. When training deep network with standard optimizers, a tuned LR is crucial but time-consuming and resource-intensive [26]. This issue is at the root of recent research focusing on transferring hyperparameters across architectures at different scales, therefore avoiding expensive tuning pipelines [99, 100, 7J. Yet, in the convex setting, choosing LR can already be difficult an issue that was studied already in Polyak [72] and gave rise the first adaptive method: the Polyak Stepsize (PS). Recently, there has been a renewed interest in adapting PS to modern settings [54, 67, 391. delivering theoretically principled way to scale the gradient magnitude during training adaptively. PS-inspired methods have gained increasing interest for their simplicity and adaptability, as they utilize local curvature and smoothness information to accelerate algorithms and facilitate faster convergence. Orvieto and Xiao [66] recently introduced a variant of the Stochastic Polyak step-size, called NGN, which further enhances the robustness to LR hyperparameter and solidifies the link to Gauss-Newton preconditioning. The theoretical analysis in Orvieto and Xiao [66] demonstrated that NGN does not diverge regardless of the choice of LR hyperparameter, and converges fast when the LR is appropriately tuned. In contrast, the current theory of the SPS step-size with fixed LR hyperparameters [54] proves convergence the exact solution only if the interpolation condition holds".
Polyak Step-size and Heavy-ball Momentum. Heavy- ball momentum methods, stemming from the work of Polyak [71], have gained significant attention the years due to their benefits, including acceleration on convex quadratics [37, 48, 6). convex-like [92], and non-convex problems [14), as well their variance reduction abilities [57, 15]. This has led to growing interest in the combination of Polyak step-size and heavy-ball momentum, which is an active area of research [3, 77, 3, 93, 63, 28). Recently, Schaipp al. [79] demonstrated that geometrically principled combination of SPS and momentum leads to lower sensitivity to LR hyperparameter, although they did not provide strong theoretical convergence guarantees.
Diagonal Polyak Step-size. Coordinate- wiseiintte-wieeeeee adaptive step-sizes are essential in training Trans- former architectures due to the varying parameter-we wise scaling and conditioning of the problem
'It is worth emphasizing that the terms robustness and stability have been used a different sense in the literature [12, 103]. In this work, we focus on the stability the choice of the LR hyperparameter. In our notation, this means that of 0.
2

---

Table 1: Summary of existing methods exploiting Polyak-type adaptive step-sizes and their conver- gence guarantees. Mom.-Supports momentum; SupporS diagonal step-sizes. a is defined in Section 4. O notation hides absolute, problem-dependent constants and logarithmic factors.
Method
Rate (m)
Mom.
Diag.
Comments
SPS [54]
O + dia
x
X
Cony, non-vanishing neighbourhood
Strong convexity Conv. to non-vanishing neighbourhood
ALR-SMAG [93]
O((1 n)* + f
/
X
Bounded stoch, gradients Interpolation
Momo [79]
O((vK)
/
X
Momo-Adam [79]
x
/
V
Momo framework for Adam
MomSPSMax [63]
olyK offe)
/
x
Conv, non-vanishing neighbourhood
NGN [66]
O(/VR)
x
x
IAM [28]
0(/VK)
/
x
Knowledge of f((")
NGN-M (Alg.
[This work]
0(/vx)
/
x
NGN-MDv1 and NGN-MDv2 (Alg. 2)
[This work]
x
/
/
Combination of NGN-M and NGN-D
NGN-D (Alg. 3)
O(VvR)
x /
[TThi work]
[62, 108]. Algorithms employing parameter-wise LR, such as Adam and SignSGD [4], typically outperform non-diagonal methods in language modeling tasks by addressing issues such class imbalance (where certain words appear more frequently than others) [46, 47] and heavy-tailed noise [104, 105, 13]. It is, therefore, paramount in current setups to deliver adaptive LR improve- ments targeted to the coordinate- wise (diagonal) regime. However, most Polyak-step-sizze satep-size-basee based algorithms only focus propose a diagonal-wise modification of Polyak-step-size by either using Adam preconditioner [79] weight matrix incorporating second-order information from the objective function [51, 76].
single LR for all parameters [54, 93, 27, 63, 66). Only few works
Quantitative Measure of Robustness. To quantify robustness, we adopt the learning-rate (LR) sensitivity metric of Wortsman et al. [97]. Let l denote the final performance metric (e.g., negative test accuracy) when training with LR 7, and by the value at initialization. We define e = min by the best achievable metric within the LR range [a, bj. The LR sensitivity is then Eye(a,b) (m. e } (*). We estimate this expectation by averaging over the LR values in sweep grid for each algorithm and task.
Comparison to prior work. Table I provides theoretical comparison of various Polyak step-size- based algorithms that incorporate momentum and/or diagonal step-size, highlighting the differences between the theoretical results presented in this work and those from prior works.
3 Algorithm design of NGN-M and NGN-D
In Orvieto and Xiao [66], the NGN step-size is derived by applying a Gauss-Newton update on a regularized first-order expansion of r(x) == vf(x). At the current point wt, they linearized
r(a* + p) se r(**) + Vr(**) p. Thus the next iterate is given as r = x + p² where
_ = argmin, (r(a*) Vr(a*)
1
||
(1)
2c
It turns out that the problem above has closed-form solution
n = 74 fx
where
Y
ff(
2)(25)
with " representing the NGN step-size. In Orvieto and Xiao [66], convergence guarantees were established for both convex and general non-convex settings. Importantly, the convex analysis shows
3

---

that NGN exhibits non-divergence property, regardless of the step-size hyperparameter (see Theorem 4.5 in [66]). Due this property, the NGN step-size is a strong candidate achieve better robustness w.r.t. the choice of LR hyperparameter.
3.1 How to Add Momentum and What to Expect?
There are several approaches to combining the adaptive Polyak-type step-size with heavy-ball momentum. Broadly, existing algorithms can be divided into two categories: the first category involves computing the Polyak step-size in the usual manner and incorporating it into the standard heavy-ball update [63]. In contrast, algorithms from the second category first determine an update direction using exponential weighted averaging of the stochastic gradient and momentum variable, and then compute the Polyak-type step-size based on the computed direction [93, 79]. Following this principled approach, we test two possible versions for combining the NGN step-size and momentum:
m k Bm (1 - 3VV (x*)
Tkk 1+ 2/ (at) (V (1*))*
Ver.1
k = 3m k- + (1 Vss,
Ver.2:
Yk
1t 255y (utt) 1m=1
xh+1 = x² - m
x 7km
Before we proceed, we should answer the question: "What do we expect from the combination of NGN step-size and momentum?". First, we aim to preserve, and ideally enhance, NGN's robustness the LR hyperparameter. To this end, we propose incorporating (heavy-ball) momentum, which is known to increase the range of LR that leads to convergence [71]. Additionally, we seek improved performance, achieving accelerated convergence akin to the advantage of SGD with momentum (SGDM) over standard SGD in convex settings. With these goals in mind, we now show that version 1 meets all of these criteria, while version is less suitable. To gain some intuition regarding the performance of these two variants, we start by conducting a simple experiment on a quadratic function (=) = 2 Ar b1[2 where is data matrix from the normalized Diabetes dataset [85] and is vector of labels. Based on the results from Figure 1 (first), we observe that variant achieves accelerated convergence as SGDM for middle-range step-size hyperparameters (c E {10*, 102)) and does not diverge for large LR hyperparameter (c E (10*)). Conversely, version 2 has a worse convergence rate than version 1 for middle-range LR hyperparameters and diverges for large ones (c {10*)): see Figure 1 (second). Therefore, theoretically analyze and practically test version 1, which we call NGN-M.
3.2 Evidence of Robustness of NGN-M
To illustrate the advantages of the design choice of NGN-M, we first consider the Rosenbrock function
f(z y) = (x 1)² + 100 z² z²² whose minimizer is (1, 1). Starting from (-1.2, 1.2.11) 1). we run both NGN-M and SGDM over wide range of LR hyperparameters {10 {10-3 ..... 10?). As shown in Figure 1, we observe that (i) for small LR hyperparameter both methods successfully converge to (1;1); (ii) SGDM already diverges for LR hyperparameter 10-2; By contrast, NGN-M remains stable even up to e = 102, thanks to its adaptive I.R that automatically adjusts with the local curvature. Figure H.3 further traces the optimization trajectories: NGN-M converges reliably for every tested value of c, whereas SGDM fails outside its narrow stability window. Finally, in Appendix H. we repeat these experiments on synthetic multi-modal function and find that NGN-M consistently finds the global minimum, while SGDM typically becomes trapped in a nearby suboptimal local minimum.
3.3 Diagonal Step-size for NGN
We propose two alternatives to make NGN step-size parameter-wiss adaptive. In the first approach, we modify an approach of (1): The next iterate A++ obtained by minimizing an approximation of the regularized first-order Taylor expansion of r(x) = / (x(x around z² namely, 2+1 x + p*
where for preconditioning matrix Σg
1 = argmin, (((at) Vr(x p)² +
(2)
2c
The intuition that Σg E Rdxd can penalize each parameter with its own weight while in vanilla NGN the penalization is the same for all parameters, and ff is an objective function we aim to minimize.

---

Tein LO3S
Train LOGS
107
107
80
120
200
Ieertions
Iterations
SCEM of
Ven
We
Mow
Figure 1: First: Comparison of SGDM, NGN. NGN-M for linear regression normalized Diabetes dataset varying a step-size hyperparameter. Second: Comparison of two options on how momentum be used in combination with NGN step-size. Third and fourth: Comparison of SGDM and NGN-M on the Rosenbrock function.
Performing simple derivations (see Appendix F), we obtain the following update rule
1+ 2f(2)) [VS( Ex v(("")
(3)
Note that by choosing Σg to be an identity matrix, the step-size "Yk in (3) reduces to the vanilla NGN step-size.
Alternatively, we can adopt simpler, parameter-wise rule: For each parameter j. we replace the full gradient norm in the NGN step-size with its own partial derivative V V,ʃs, (a*)). Both of the described per-coordinate variants can be further adjusted by an RMSprop-style preconditioner D diag( (algg(DD)oa (D))())..... (D))(a) and lead to the following update rule (see Alg. 2 for full description)
(.)1) (5)
c/(De)o c/DDE (V, fs, (at))²
7k 14 2f(un
1V, (w*)
NGN-MDv1
D
NGN-MDv2
1+
25(45)
E 1 7 D
Σ-1 diag(7/
7k (d0)
I. A+1 -1- (a*) + B1 Bxx x 1
We highlight that both versions have the same number of hyperparameters as Adam. From an empirical evaluation of the two versions of NGN-MD in Figure 2, observe that the first choice improves the performance of NGN-M while maintaining robustness to the LR hyperparameter. A more detailed discussion on the two versions of NGN-MD algorithms deferred to Appendix F.1. However, the robustness of NGN-MD also depends on the choice of preconditioner. When the preconditioner is sensitive to variations the loss landscape or hyperparameters, NGN-MDv may be less robust than NGN-M. In our experiments, we find that the RMSprop preconditioner performs well in practice. Other preconditioners, such as AdaFisher [25], could integrated into NGN-MD, potentially providing more accurate curvature approximations and improved stability with respect to the LR. We leave systematic study of these alternatives to future work.
In the special case 31 = 0 and Σg I, NGN-MDv2 reduces to NGN-D (Algorithm 3). To the best of our knowledge, NGN-D is the first algorithm that uses a per-parrae per-parameter Polyak- type step-size while achieving the standard (11 / rate under smoothness and bounded noise variance assumptions; see detailed discussion in Appendix C.
4 Theoretical Analysis of NGN-M
4.1 Problem Formulation and Notation
We consider the classic Empirical Risk Minimization (ERM) problem that typically appears when training machine learning models, namely,
min DER
f(x)
n -E(0)) f((x)
5

---

Algorithm NGN-M
1: Input: x- r0 Rd step-size hyperparameter
momentum parameter 8 E [0,
2: for k ===00 0.1..... K do
3: Sample Sk [n]]
4:
218, (a [[V
5: x k+1 x B) (x + B(a 2k-1)
6: end for
Algorithm 2 NGN-MD
1: Imput: x" E IR", step-size hyperparameter e > O, momentum parameters 31,82 E [0, 1),
stabilization parameter 0, second-order momentum e do
2: for k 0.=01 1.1...... K
3: Sample S S [n]
4: o* 320
(1 (2) s (a*)
V.s (x*))
5: D; = diag(el II + V /(1 - 35)
6: For NGN-MDv1: 7/k 250, (ct))
D/
7: For NGN-MDv1: Σx D
8: For NGN-MDv2: Σx diag((?)
7k
(d) where (2)
2484 (=)))))))))) (V,
e((D))
(2
9: xkt1 = x* - 101-(11-- (1-31) E V (x*) + 11* xb-1)
10:
end for
where r are the parameters of a model we aim to train, is the number of data points in the dataset, is the number of parameters, denotes the solution (4), and fi represents the loss associated with the i-th data point/batch. We assume that each Sa is differentiable and non-negative' and that the global optimal value is bounded, i.e. f argming (1) E R.. Moreover, we assume that we have access to mini-batch stochastic losses fs during training such that f3 = argmin, fs(x) < co for any S _ [n] picked uniformly at random.
We analyze the convergence of NGN-M under assumptions that are often used in the analysis of the Polyak step-size [54, 67, 66, 63, 79].
Assumption 4.1. Each fi is convex and L-smooth, i.e., for all x, y Rd and i E [n] we have (V)(()) x) 2 fz( - fy) and [[(r( - ff(l))l L| - all.
Assumption 4.2. The interpolation of = Es(f* - fål and positive opos == Es[fš] errors are bounded. We say that the interpolation holds if of = O, where S is sampled mini-batch.
4.2 Convergence Guarantees
Theorem 4.3. Let Assumptions 4.1, 4.2 hold. Let the step-size hyperparameter c > and the X momentum parameter 3 = 1-x be constants where X S min(cL, - cI+ Then
the iterates of NGN-M (Algorithm 1) satisfy
([(( < 1| + 8cL(1 + 2c + 2cL max (2eL - 1,03 of
ck
where K- chosen uniformly at random from (20) K- Moreover, we set e O(I/R) then obtain [[ f(=*)] < ((1/R)).
The convergence of NGN-M is provided the convex setting, which is motivated by recent works that observe convex-like structures in the landscape of networks [34, 31] and agreement between convex theory and practice [80]. Importantly, we show that (i) when the constant e is sufficiently small, NGN-M attains the same convergence rate SGDM [22]. Moreover, for any choice of we demonstrate that the NGN-M iterates provably converge to a neighborhood of the optimum and
CCommon losses, e.g. entssyy satisfy this condition.

---

Houl Ias PLISTIAY
Frou Tost Auturacy
NGN
SGDM
Momo
NGN-M
Resnet20, CITARIO
20
scon Nomo
50100
Resnet 10, CIFARIO
Nomr
Pimm
Stopsiee
Stesdizo
Stepsze
VI.. CIFARIO
Flodl lost Recureey
Fial Tost McDiaay
Adam
Momo Adam
NGN-MDv1
Resnet20, CIFARIO
37
27
13
Adar Mare
Homeodder trau
Resnet110, CIFARIO
Walk MON
32
32
27
Stepaze
I0"
Stoodize
Stepsize
VIT, CIFARIO
41
Resne120 for CIFAR 10
Resnet 10 for CIFAR 100
VIT for CIFAR 10
Figure 2: Stability performance of algorithms varying LR hyperpa-
Table 2: Test accuracy LR sen- sitivity of the different opti-
rameter (e for NGN-M, NGN-MDv1 and NGN-MDv2, for Momo and
Momo-Adaa Adam, and LR for SGDM and Adam). We refer to Figures I.1 to 1.3, 1.5, and 1.8 for the results on additional workloads.
mizers shown in Figure 2.
thereafter remain within it; (it) Unlike prior works on Polyak step-size, our analysis does not rely on strong assumptions such as bounded gradients, interpolation, or bounded domain; (tii) For small values of e, NGN-M converges to the exact solution while algorithms such as MomSPS and ALR-SMAG were shown to converge up to non- vanishing neighborhood of the solution only, due to an inherent limitation of the stochastic Polyak step-size [65]. Regarding the momentum parameter B, the typical (large) value 3 0.9 performs well in our own experiments. Theoretically, however, B is recommended to be chosen sufficiently small to ensure convergence with the NGN step-size [63]. This discrepancy between theoretical guidance and practical implementation has also been observed in prior works on momentum [23, 53, 93, 92, 63]. Interestingly, under the additional interpolation condition oi commonly used choice 3 = 0.9 (see Appendix D.2). This suggests that the small-3 requirement may reflect limitations of current proof techniques rather than intrinsic restriction of NGN-M. Extending the analysis to arbitrary , in the stochastic regime without interpolation remains an open problem. Our intuition, however, that simultancously choosing both c and 33 without restriction is not feasible: fixing one hyperparameter arbitrarily necessitates imposing constraints on the other to prevent divergence; (v) Theorem 4.3 requires the total iteration count K to be specified in advance; this assumption is standard in the complexity analysis of optimization algorithms [24, 61]. Since this can be impractical, we also establish convergence under diminishing step-size of order l/V in Appendix D.3, which removes the need to preset K; (vi) Finally, we corroborate our analysis as we run NGN-M with the theory-derived values of to a quadratic problem that satisfies all our assumptions: We observe NGN-M's M's rapid convergence with theoretical step-size hyperparameters in seeceee-eeeeeee
we can establish convergence even for large momentum values, including the
practice- see Appendix H.3 and Figure H.4 therein.
Key Ingredients of the Proof. We discuss the key steps the proof to highlight the main challenges in the analysis. First, we make use of the Iterative Moving Average (IMA) formulation of momentum [81]. Specifically, we define a sequence of virtual iterates {=*] whose update rule of the form
1
A
h+l x 7A (ss,
1+X & +
A+1 where 20 = x and 3 1+
I+X
Next, one of the key technical strategies we follow is splitting the step-size Yh into two parts: fixed 3 term p [1+2L) Occ) and changing term 5h S 1+2el O This decomposition of the step-size ? enables us to regulate the balance between the descent term, which drives improvement in the objective, and the error term, which reflects possible inaccuracies. More precisely, the descent term weighted by e while the error term proportional to of weighted by c², which suggests that has to be chosen to tradeoff the two terms to lead to the exact convergence similarly to the standard analysis of SGD [22]. contrast, MomSPS and Momo algorithms achieve the exact convergence only under the interpolation regime.
7

---

SGD
NGN
SGOM
Momo
NGN-
2101
Resneti8, 18, ImageNetlk
Roal That
Aam
Adam
Momo Adam NGNHDVi
Home
Resnet ImageNelk
NOVHDD1
12
Steosice
Stopsi0e
Stopsze
VIT Tiny, ImageNetik
Resnet18 for ImageNetIk
Resnet18 for ImageNetlk
VIT-tiny for ImageNet1k
29
26
18
Figure Stability performance on ImageNetIk varying the LR hyperparameter. NGN-M and NGN-MDv1 achieve higher accuracy for wider range of the LR hyperparameters, We refer Figure I.4 for results on train loss stability and additional results on ImageNet32.
Table 3:
Test accuracy LR sen- sitivity of the different opti-
mizers shown in Figure 3.
Rral Tost Purp Muyy
Adam
REUUXDV
Frall Test Perp
SCN-MEN22
mno Test
MGG HDx
Stepsize
Steps 20
Scepsize
Stepsize
70M Transformer++
160M Transformer++
410M Transformer++
Transformer++
Figure 4: Comparison of stability the LR hyperparameter across model sizes and optimizers in language modeling. We refer to Figures 1.11 to I. 14 for the results that report update magnitude when training 160M model and training dynamics across all model sizes.
Experiments
We now turn to the empirical evaluation of the proposed algorithms against several benchmarks. The detailed experiment setup, including the choice of hyperparameters well additional experimental results and details, can found in Appendix I. The best performance of algorithms is reported in Tables (momentum-basee algorithms), 8 (algorithms with momentum and component-wise step-size), and 9 (algorithms with component-wise step-size). For clarity and quick reference, all links the paper's empirical results summarized in Table while Appendix provides additional details about the training and tokenization pipeline.
Comparison on Standard Benchmarks. First, test the performance of NGN-M against other methods that use momentum, such as SGDM, Momo, MomSPS, ALR-SMAG, and NGN. The tests include the training of Resnet20 [29] and VIT [17] the CIFARIO dataset [45], and Resnetl 10 on CIFAR 100. Second, we test the performance of NGN-MD against Adam and Momo-Adam that contrary NGN-M both use component- wise preconditioning. All experiments in this section do not use LR schedulers or weight decay.
From Tables 7 and we observe that NGN-M and NGN-MDv1 exhibit competitive performance across all settings we tested, matching the best performance of other algorithms. Importantly, NGN-M and NGN-MDv1 demonstrate significantly greater robustness to the choice of the LR hyperparameter. Indeed, Figure 2 shows that the range of LR hyperparameter that allows NGN-M and NGN-MDv1 to perform optimally is much wider: We can, for instance, step-sizes that are 1-2 orders of magnitude larger than the optimal one without a significant drop in performance. This is particularly evident when training ResNet20 and ViT models. Besides, improves the stability of NGN across all settings. A similar trend can be observed when considering the LR sensitivity metric: see Table 2. We refer to Appendix I for additional ablation studies against other optimizers and results when training small-scale NLP models.
clearly observe that momentum consistently
Vision Experiments on ImageNet. Now we switch to larger tasks and datasets. We train ResNet18 on ImageNetlk [16]. This represents the task which we pair our proposed algorithms with LR schedule. As illustrated in Figure 3, NGN-M and NGN-MDv1 achieve the highest test accuracy, while exhibiting higher robustness across larger LR, improving over both NGN and Momo. Among adaptive methods, NGN-MDv1 compares favorably against Adam and Momo-Adam, while once again

---

00000
1991
006333
00014
COKG
CORR16
0.80
C00000
102
102
Eftectvc Stepsies
Effective Stepsize
Effectiv Steasize
10
10%
10
10
10"2
10
10
15
10
Iterationn, x101
Rerations, x110
Bteratoon,
Iterationn, x10
Momo
NGN-M
-omo-Adam
NGN-MDv1
Figure 5: The step-size of Momo, NGN-M (two left), Momo-Adam and NGN-MDv1 (two right) during the training of ViT on CIFARTO. 10. We demonstrate the step-sizes тą for Momo and Momo-Adee and " for NGN-M and NGN-MDv1 varying step-size hyperparameters dgg and of the algorithms (indicated in the legend). We refer to Figures 1.9 and L.10 for the results in training Resnet20.
achieving higher performance on a wider range of LR (Table 8). Appendix 1.4 reports additional ablations ImageNet32 and train loss stability results.
Next, we test the effectiveness of the algorithms when training ViT-Tiny on ImageNetlk. This model is trained for longer horizon, making it notoriously sensitive to the initial LR and requiring an adaptive step-size. We follow the protocol of Schaipp et al. [79]. As highlighted in Figure 3 and Table 8. NGN-MDv1 achieves the highest test accuracy across adaptive methods. Moreover, at larger LR, Adam diverges, whereas both Momo-Adam and NGN-MDv1 maintain more stable training dynamics. The LR sensitivity metric reported in Table 3 supports observations.
Language Modeling. Pre-training Large Language Models represents a challenging optimization task. To achieve competitive performance, optimizers with adaptive step-size are needed, and preventing instabilities in low-precision training often requires careful hyperparameter tuning. To evaluate the capability of NGN-MDv1 in this setting, we train decoder-only transformers [74] with 70M, 160M, 410M, and 1B parameters around Chinchilla optimum [33] on SlimPajama-66778 [86]. For each model, we retune the LR, using seeds for the first three models and seed for the 1B.
As reported in Figure and Table 8, we note that NGN-MDv1 matches the performance of Adam across all model sizes. However, NGN-MDv1 achieves competitive performance even for LR hyperparameter = 1 while Adam's 's performance drops significantly. This phenomenon is consistent across all scales we tested, suggesting that the optimal LR of NGN-MDv1 is shifted towards larger values, but also that the algorithm is less sensitive to such hyperparameter. We additionally discuss how to introduce weight decay in NGN-MDv1 and report additional ablations on its role in this training task in Appendix G. Moreover, we report the ablation studies when varying the momentum parameter in Appendix 1.12, demonstrating the improved stability to momentum parameters.
Effective Step-size of NGN-M and NGN-MDv1. As shown in Figure 5, the effective LR of NGN- M and NGN-MDv1 is inherently adaptive: it rises sharply at the start, and then gradually decreases, resembling annealing schedules commonly used in practice. By contrast, the effective step-size of Momo and Momo-Adam is largely fixed for large on, effectively reducing them SGDM and Adam and limiting their resilience. Evidence across ResNet20 training (Figures 1.6, I.7, 1.9 and 1.) and large-scale language modeling (Figures I.11 to L.13) shows that the NGN step-size is more conservative, automatically decreasing when needed to stabilize training-even for large c. This adaptivity underlies the LR robustness of NGN-M and NGN-MDv1.
Table 4: Train time of Adam and NGN-MDv1 when training language models.
Model
Method
Time per Iteration (sec)
Time per Optimizer Update (sec)
AdamW
1. 1.63
0.0048 0.0130=0.0m02
70M
NGN-MDv1
1. 1.65#0..01
AdamW
3.33
0.0088
160M
NGN-MDv1
3..7
(1.023920.000
AdamW
8..08
0.0838400 10.0009
410M
NGN-MDv1
8. 8.68+0.06
1.215440.000

---

Implementing NGN-MDv1 can be slightly more expensive, but the overall cost modest as we show next. First, we emphasize that the implementation does not even require additional matrix-vector products since the preconditioner is diagonal; the only extra work is computing [V)()))) D. which amounts to an additional pass over the gradient. This can
Computation Cost of NGN-MD.
also be incorporated into the update of Dg, avoiding extra matrix operations. In practice, implementation is about 2.5x slower per update than Py PyToocch's 's AdamW (see Table 4) due to the need for two passes over the gradient. However, since forward/backward computations dominate runtime, the overall training speed remains largely comparable. Since our focus is on demonstrating the stability benefits of the NGN step-size, we leave efficiency improvements future work. The extended discussion is reported in Appendix F.2.
naive
6 Conclusion and Future Work
This work introduced several novel adaptations of the NGN step-size method, incorporating support for momentum and/or diagonal step-size. We provided a theoretical analysis of the convergence rates for these algorithms and conducted an extensive empirical evaluation of their performance. The experimental results show that combining momentum with the NGN step-size yields high robustness to step-size hyperparameter choices and performs competitively with state-of-the-art algorithms across various settings.
Given the significant complexity of the task, we defer the theoretical explanation of the step-size resilience properties NGN-M for large values of 3 and analysis in the non-convex setting to future work, including classes of structured non-convex functions such as PL [73], a-B--ondition [34], or Aiming [52]. would also be worthwhile to study NGN-M under weaker smoothness assumptions [104, 1j. Furthermore, while the two proposed methods for incorporating weight decay into NGN- MDv1 outperform AdamW in training language models, they still exhibit some sensitivity to the step-size hyperparameter. This may, in part, be due to the limited understanding of the expected effects of the weight decay technique, topic that requires further investigation. We acknowledge that computing NGN step-size at large scale may cause runtime overhead, and discuss this limitation in Appendix F22 We also recognize that integrating NGN-MDv1 with advanced parallelism schemes, such Tensor Parallelism [83] or ZeRO-2 [75], introduces additional compute and communication overhead, and will require further adaptation of the algorithm. Nevertheless, our results provide valuable guidance for developing inherently more stable optimizers. As next step, it would be fascinating to investigate whether the resilience of emerging methods like Muon [40] can be further improved by incorporating the NGN step-size.
Acknowledgement
Rustem Islamov and Aurelien Lucchi acknowledge the financial support of the Swiss National Foundation, SNF grant No 207392. Antonio Orvieto acknowledges the financial support of the Hector Foundation.
References
[1] Foivos Alimisis, Rustem Islamov, and Aurelien Lucchi. Why do we need warm-up? a theoretical perspective. arXiv preprint arXhv:2510.03164 2025. (Cited on page 10)
[2] Maksym Andriushchenko, Francesco D' Angelo, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? arXiv preprint 10.04423 20233. (Cited on pages 45 and 48)
[3] Mathieu Barré, Adrien Taylor, and Alexandre d' Aspremont. Complexity guarantees for polyak steps with momentum. In Proceedings of Thirty Third Conference on Learning Theory, 2020. (Cited on page 2)
[4] Jeremy Bernstein, Yu-Xiang Wang, Kamyar Arizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, 2018. (Cited on pages and 25)
10

---

[5] Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Ho- race He, Connor Leaby, Kyle McDonell, Jason Phang. Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Gpt- neox-20b: An open-source autoregressive language model. arXiv preprint arXiv: 2204.06745, 2022. (Cited on page 51)
[6] Raghu Bollapragada, Tyler Chen, and Rachel Ward. On the fast convergence of minibatch heavy ball momentum, arXiv preprint (rXiv:2206.07555 2022. (Cited on page 2)
[7] Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309. /6620, 2023. (Cited on page 2)
[8] Tom B Brown. Language models are few-shot learners. arXiv preprint 200i::0005.11666 2020. (Cited on page 1)
[9] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong. Thang Luong, Cho-Jui Hsieh, Yifeng Lu, al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems, 2024. (Cited on page 53)
[10] Dami Choi, Christopher Shallue, Zachary Nado, Jaehoon Lee, Chris Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv:1910005446, 2019. (Cited on page 1)
[11] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 2023. (Cited on page 1)
[12] Jeremy M Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gra- dient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv: 10t::1203300005 2021. (Cited on page
[13] Enea Monzio Compagnoni, Tianlin Liu, Rustem Islamov, Frank Norbert Proske, Antonio Orvieto, and Aurelien Lucchi. Adaptive methods through the lens of SDEs: Theoretical insights on the role of noise. In The Thirteenth International Conference on Learning Representations 2025, URL https:tt///// http://///peeeeiee...o.m
(Cited on page 3)
[14] Ashok Cutkosky and Harsh Mehta. Momentum improves normalized sgd. In International conference machine learning. PMLR, 2020. (Cited on pages and 2)
[15] Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. Advances in neural information processing systems, 2019. (Cited on pages and 2)
[16] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, leee, 2009. (Cited on page 8)
[17] 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 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. (Cited on page
[18] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 2011. (Cited on page 55)
[19] Ilyas Fatkhullin, Alexander Tyurin, and Peter Richtárik. Momentum provably improves error feedback! Advances in Neural Information Processing Systems, 2024. (Cited on page
[20] Simon Foucart. Lecture 6: Matrix norms and spectral radii. lecture notes for the course NSTP187 at Drexel University, Philadelphia, PA, Fall, 2012, 2012. (Cited page 43).
11

---

[21] Jingwen Fu, Bohan Wang, Huishuai Zhang. Zhizheng Zhang, Wei Chen, and Nanning Zheng. When and why momentum accelerates sgd: An empirical study. arXiv preprint arXiv:2306.00000 09000, 2023. (Cited on page 1)
[22] Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301. /1235, 2023. (Cited pages 6, 7, 25, 27, 30, and 36)
[23] Euhanna Ghadimi, Hamid Reza Feyzmahdavian, and Mikael Johansson. Global convergence of the heavy-ball method for convex optimization. In 2015 European control conference (ECC), 2015. (Cited page 7)
[24] Saeed Ghadimi and Guanghui Lan, Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM journal on optimization, 2013. (Cited on page
[25] Damien Martins Gomes, Yanlei Zhang, Eugene Belilovsky, Guy Wolf, and Mahdi Hos- seini. Adafisher: Adaptive second order optimization via fisher information. arXiv preprint arXiv:2405. /6397, 2024. (Cited on page 5)
[26] lan J. Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. (Cited on page 2)
[27] Robert Gower, Othmane Sebbouh, and Nicolas Loizou. Sgd for structured nonconvex functions: Learning rates, minibatching and interpolation. International Conference on Artificial Intelligence and Statistics, 2021. (Cited on page 3)
[28] Robert Gower, Guillaume Garrigos, Nicolas Loizou, Dimitris Oikonomou, Konstantin Mishchenko, and Fabian Schaipp. Analysis of an idealized stochastic polyak method and its application to black-box model distillation. arXiv preprint arXiv:2504.01 2025. (Cited on pages 2 and 3)
[29] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, 2016. (Cited on page 8)
[30] Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-bater gradient descent. Lecture notes, 2012. (Cited on page 1)
[31] Tran Hoang, Qinzi Zhang, and Ashok Cutkosky, Empirical tests of optimization assumptions in deep learning. arXiv preprint rxXiv22407.01825, 2024. (Cited on page 6)
[32] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997. (Cited on page 51)
[33] Jordan Hoffmannn, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. arXiv preprini arXiv:2203. 15556. 2022. (Cited on page 9)
[34] Rustem Islamov, Niccoló Ajroldi, Antonio Orvieto, and Aurelien Lucchi. Loss landscape characterization neural networks without over-parametrization. Advances in Neural Information Processing Systems, 2024. (Cited on pages 6 and 10)
[35] Rustem Islamov, Yuan Gao, and Sebastian Stich. Near optimal decentralized optimization with compression and momentum tracking. arXiv preprint arXiv:2405.2011, 2024. (Cited on page
[36] Rustem Islamov, Samuel Horvath, Aurelien Lucchi, Peter Richtarik, and Eduard Gorbunov. Double momentum and error feedback for clipping with fast rates and differential privacy. arXiv preprint arXiv: 2502.11682 2025. (Cited on page 1)
12

---

[37] Prateek Jain, Sham M Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Accelerating stochastic gradient descent for least squares regression. In Conference On Learning Theory, 2018. (Cited on page 2)
[38] Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, 2022. (Cited on page I)
[39] Xiaowen Jiang and Sebastian U Stich. Adaptive sgd with polyak stepsize and line-search: Robust convergence and variance reduction. Advances in Neural Information Processing Systems, 2024. (Cited on page 2)
[40] 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. URL
https://kellerjoone
github. htt.:iiopppps//iie mww./. (Cited on page 10)
[41] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv: 2001.0836/ 2020. (Cited on page 51)
[42] Andrej Karpathy. char-rnn. httpps//ggttuub...m/aaaaaaattt/t com
2015. (Cited on
page 51)
[43] Andrej Karpathy. Nanogpt. https://github. cotplau...ook//aaaa/aeaa com./aarppaaayy/aaoPPPP,, 2022. (Cited on page 51)
[44] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015. (Cited on pages and 55)
[45] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Learning multiple layers of features from tiny images. Scientific Report, 2009. (Cited on page 8)
[46] Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In The Eleventh International Conference on Learning Representations, 2023. (Cited on page 3)
[47] Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy- tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv: 2402.19449, 2024. (Cited on page 3)
[48] Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of mini-batch momentum: batch size saturation and convergence in high dimensions. Advances in Neural Information Processing Systems, 2022. (Cited on page 2)
[49] Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of adam under relaxed assumptions. Advances in Neural Information Processing Systems, 2024. (Cited on page I)
[50] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006. 15704, 2020. (Cited on page 45)
[51] Shuang Li, William Swartworth, Martin Takáč, Deanna Needell, and Robert M Gower. Sp2: A second order stochastic polyak method. arXiv preprint arXiv:2207.08177, 2022. (Cited on page 3)
[52] Chaoyue Liu, Dmitriy Drusvyatskiy, Misha Belkin, Damek Davis, and Yian Ma. Aiming towards the minimizers: fast convergence of sgd for overparametrized problems. Advances in neural information processing systems, 2023. (Cited on page 10)
[53] Yanli Liu, Yuan Gao, and Wotao Yin. An improved analysis of stochastic gradient descent with momentum. Advances in Neural Information Processing Systems, 2020. (Cited on page 7)
[54] Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak siep-size for sgd: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics, 1.2 (Cited pages 1, 2, 3, and 6)
13

---

[55] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 1711.05101, 2019. (Cited on page 46)
[56] Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv: wtt../902..9.... 2019. (Cited page 53)
[57] Jerry Ma and Denis Yarats. Quasi-hyperholic momentum and adam for deep learning. arXiv preprint
2018. (Cited on pages and 2)
[58] Stephen Merity, Caiming Xiong, James Bradbury. and Richard Socher. Pointer sentinel mixture models. In Proceedings of 4th International Conference on Learning Representations (ICLR 2016), 2016. (Cited on page 52)
[59] Tomas Mikolov, Martin Karafíát, Lukas Burget, Jan Cernocký, and Sanjeev Khudanpur. Recurrent neural network based language model. Proceedings of the 11th Annual Conference of the International Speech Communication Association, INTERSPEECH 2010, 2010. (Cited on page 51)
[60] Igor Molybog, Peter Albert, Moya Chen, Zachary Vetioo, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory adam instability in large-scale machine learning. arXiv preprint arXiv:2304.09871, 2023. (Cited on page
[61] A. Nemirovski, A. Juditsky, G. Lan, and A. Shapiro. Robust stochastic approximation approach to stochastic programming. SIAM Journal on Optimization, 2009. (Cited on page 7)
[62] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 2022. (Cited on page 3)
[63] Dimitris Oikonomou and Nicolas Loizou. Stochastic polyak step-sizes and momentum: Convergence guarantees and practical performance. arXiv preprint MxXiv:2406.0114: 2024. (Cited on pages 2, 3, 4, 6, 7, and 25)
[64] Sharir Or, Barak Peleg, and Yoav Shoham. The cost of training nlp models: A concise overview. arXiv preprint arXiv:2004.0890 2020. (Cited on page 1)
[65] Francesco Orabona and Ryan D'Orazio. New perspectives on the polyak stepsize: Surrogate functions and negative results, arXiv preprint 250.::29900.2299, 2025. (Cited on page 7)
[66] Antonio Orvieto and Lin Xiao. An adaptive stochastic gradient method with non-negative gauss-newton stepsizes. arXiv preprint arXiv: 2407.04355 2024. (Cited on pages 1, 2, 3, 4, 6,
and 26)
[67] Antonio Orvieto, Simon Lacoste-Julien, and Nicolas Loizou. Dynamics of sgd with stochastic polyak stepsizes: Truly adaptive variants and convergence to exact solution. Advances in Neural Information Processing Systems, 2022. (Cited on pages and 6)
[68] Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categoriza- tion with respect to rating scales. In Proceedings of the ACI., 2005. (Cited on page 52)
[69] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. NIPS 2017 Workshop Autodiff, 2017. (Cited on page 51)
[70] Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 2024. (Cited on page 58)
[71] Boris T. Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathematics and mathematical physics, 1964. (Cited on pages 1, 2, 4, and 42)
[72] Boris T Polyak. Introduction to optimization. New York, Optimization Software, 1987. (Cited page 2)
14

---

[73] Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal vychisli- tel' noi matematiki matematticessko fiziki, 1963. (Cited on page 10)
[74] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. Technical report, OpenAI, 2019. (Cited on page 9)
[75] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory opti- mizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16. IEEE, 2020. (Cited pages 10 and 45)
[76] Peter Richtárik, Simone Maria Giancola, Dymitr Lubezyk, and Robin Yadav. Local curvature descent: Squeezing more curvature out of standard and polyak gradient descent. arXiv preprint arXiv:2405. 16574, 2024. (Cited on page 3)
[77] Samer Saab, Shashi Phoha, Minghui Zhu, and Asok Ray. An adaptive polyak heavy-bal method. Machine Learning, 2022. (Cited on page
[78] Mher Safaryan and Peter Richtárik. Stochastic sign descent methods: New algorithms and better theory. In International Conference Machine Learning, 2021. (Cited on page 25)
[79] Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, and Robert M. Gower. MoMo: Momentum models for adaptive learning rates. In Proceedings of the s1 International Conference on Machine Learning, 2024. (Cited on pages 2, 3, 4. 6, and 51)
[80] Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, and Francis Bach. The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. arXiv preprint MXXi:2550111966 2025. (Cited on page 6)
[81] Othmane Sebbouh, Robert M Gower, and Aaron Defazio. Almost sure convergence rates for stochastic gradient descent and stochastic heavy ball. In Conference on Learning Theory, 2021. (Cited pages 7 and 25)
[82] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv: 2002.05202, 2020. (Cited on page 51)
[83] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-biliion parameter language models using model parallelism. arXiv preprint (tt::1909.000000 2019. (Cited pages 10 and 45)
[84] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409. 1556, 2014. (Cited on page 51)
[85] J. W. Smith, J. E. Everhart, W. C. Dickson, W. C. Knowler, and R. S. Johannes. Using the adap learning algorithm to forecast the onset of diabetes mellitus. In Symposium on Computer Applications and Medical Care, 1988. (Cited page 4)
[86] Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023. (Cited on pages 9 and 51)
[87] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104. 09864, 2023. (Cited on page 51)
[88] Hugo Touvron, Matthieu Cord. Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training eefta-efficintt image transformers distillation through attention. In International conference on machine learning, 2021. (Cited page 1)
[89] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. arXiv preprint arXiv: 2302.13977 2023. (Cited on pages I and 51)
15

---

[90] 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. Advances Neural Information Processing Systems, 2024. (Cited on page 1)
[91] Jun-Kun Wang, Chi-Heng Lin, and Jacob D Abernethy. A modular analysis of provable acceleration via polyak's s momentum: Training a wide relu network and deep linear network. In International Conference on Machine Learning, pages 10816-10827. PMLR, 2021. (Cited
on page 43)
[92] Kun Wang, Chi-Heng Lin, Andre Wibisono, and Bin Hu. Provable acceleration of heavy ball beyond quadratics for class of colyak-lojasiewict functions when the non-convexity is averaged- out. In International conference on machine learning, 2022. (Cited on pages 2 and 7)
[93] Xiaoyu Wang, Mikael Johansson, and Tong Zhang. Generalized polyak step size for first order optimization with momentum. In International Conference on Machine Learning, 2023, (Cited pages 2, 3, 4, and 7)
[94] Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. Journal of Machine Learning Research, 2020. (Cited on page 27)
[95] Ross Wightman. Pytorch image models. saarrhh-imagemmdeel 2019. (Cited on page 51)
httpe:://gtthbb..co//rrgggg
[96] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural informa- tion processing systems, 2017. (Cited on page 1)
[97] Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, et al. Small-scale proxies for large-scale transformer training instabilities. arXiv preprint 23Xv:2309911322, 2023. (Cited on pages and 3)
[98] Lechao Xiao. Rethinking conventional wisdom machine learning: From generalization to scaling. arXiv preprint arXiv: 15156, 2024. (Cited on pages 45 and 48)
[99] Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022. (Cited on page 2)
[100] Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310. 178/33 2023. (Cited on page 2)
[101] Biao Zhang and Rico Sennrich. Root mean square layer normalization. arXiv preprint 11Xiv:1910007467, 2019. (Cited on page 51)
[102] Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. arXiv preprint 181iv:1810.1228/ 2018. (Cited on page 45)
[103] Hongyang R Zhang, Dongyue Li, and Haotian Ju. Noise stability optimization for finding flat minima: A hessian-based regularization approach. arXiv preprint arXiv:2306.088. 2023. (Cited on page 2)
[104] Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint 1rXiv:195..1188/ 2019. (Cited on pages and 10)
[105] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 2020. (Cited on page 3)
[106] Qi Zhang, Yi Zhou, and Shaofeng Zou. Convergence guarantees for rmsprop and adam in geweralized-smonnee non-convex optimization with affine noise variance. arXiv preprint arXiv:2404.0144 2024. (Cited on page 1)
16

---

[107] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01060 2022. (Cited on page 1)
[108] Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why transformers need adam: A hessian perspective. arXiv preprint arXiv:2402. 16788, 2024. (Cited
page 3)
[109] Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikkonda, Nicha Dvoreek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. In Advances in Neural Information Processing Systems, 2020. (Cited on page 53)
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose a novel scheme how to combine NGN step-size with momentum and component- wise step-size. We provide extensive theoretical and numerical analysis of NGN-M and NG to support the claims made the abstract and introduction.
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 the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as 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]
We include a discussion on the limitations of our work in the conclusion and appendix.
Justification:
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 on 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 few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
17

---

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 to 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, 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 a complete (and correct) proof?
Answer: [Yes]
Justification: We provide all assumptions used in the analysis in Section 4 and Appendix C.
The proofs of convergence and stability are deferred Appendix C. D.I, D.3, and E.
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 in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but 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 the Pytorch- implementation in the supplementary. We use an open-source implementation of models and public datasets in experiments. All training details are reported in Appendix I.
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 provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to 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 to either make 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 are appropriate to the research performed.
While NeurIPS 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
18

---

(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should descrihe 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- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing verifying the results. 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We use only publicly available code and datasets and provide links them in Appendix I.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/ htt////dd//Coddeiiie
Pompcolly for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not 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: / //nips. /t/pcc/ppl/co public/guider Seee//CooSuuiiieec asssooooicyy for more details.
The authors should provide instructions 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. 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 as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is 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 understand the results?
Answer: [Yes]
Justification: We provide all training details in Appendix and Table 6.
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 to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
19

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For each experiment, we plot the mean and a standard deviation across at least random seeds (if the opposite is not stated). do not include error bars when training IB language model due to limited resource availability which is necessary to run this experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at 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 run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It 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 figures symmetric error bars that would yield results that are out range (c.g. negative error 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide description of compute resource we used for each experiment in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well 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 't make 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 https://neuripp: sttges///orrrp..cccpbbl///////i httpp//ect//c.cccccccceeeeeeeeee ines?
Answer: [Yes]
Justification: The research is conducted following the NeurlPS instructions, including the regulations regarding anonymity.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
20

---

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: [Yes]
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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 negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to 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:
Guidelines:
The answer NA means that the paper poses 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 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 is 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
21

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
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 or dataset.
The authors should state which version of the asset is used and, if 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.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcom thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is available online, the authors are encouraged to reach out the asset's ss 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:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of 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 anonymize your assets (if applicable). You can either create an anonymized URL or include an 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:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
22

---

Question: 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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
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 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
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
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 L.LM policy (https:/// httppss////urpp..... http:////eerrpp..cc/oofeeeeeeeee /ttp:///ooeeeeecc///2/////1/ for what should or should not be described.
23

---

Appendix
Contents
Equivalent Formulations of NGN-M
25
Technical Lemmas and Definitions
25
Convergence of NGN-D
26
C.1 Convergence in General Non-convex Setting
28
C.2 Convergence under Pł.-condition
D Convergence of NGN-M
30
D.I Convergence NGN-M Stochastic Setting
Stetg...........................
D.2 Convergence of NGN-M under Interpolation.
34
D.3 Convergence of NGN-M with Decaying Step-size
36
E
Stability of NGN-M on Simple Problem
40
How to Derive Diagonal NGN-based Step-size?
44
F.1 Design Comparison of NGN-MDv1 and NGN-MDv2
44
F.2 Computation Cost of NGN-MD
45
G How to add weight decay to NGN-MDv1?
45
G.1
Combining NGN-MDv1 and Weight Decay Regularization
46
G.2 Empirical Validation of the Proposed Combinations
48
Additional Experiments on Toy Problems
H.1
Additional Experiments on the Problem with Many Minima
49
H.2
Comparison Rosenbrock Function
49
H.3
Comparison on Quadratic Function with Theoretical Step-size
49
Additional Experiments and Training Details
51
LI
Training Details
1.2
Comparison Algorithms that Support Momentum
51
L.3
Comparison of Algorithms that Support Momentum and Diagonal Step-size Additional ImageNet Experiments
51
L.4
52
I.5
Additional Comparison against Lion, Adabelief, Adabound
53
16
Comparison of Adaptive Step-sizes Adam, Momo-Adam and NGN-MDV
53
1.7
Extended Comparison of Momentum-based Algorithms on NL. Tasks
54
I.8
Comparison of Algorithms with Diagonal Step-size
55

---

1,9
Effective Step-size of NGN-M, Momo, NGN-MDvI, and Momo-Adam
56
L,10
Effective Updates in Training Language Models
58
1.11
Training Dynamics Training Language Models
58
I.12
Ablation Study of Momentum Parameters
58
A
Equivalent Formulations of NGN-M
We remind that the iterates of NGN-M are the following
x k+1 = x (1 3) Ssa (x*) + B(x xk-1)
= x
3)
25 ( fs. (2*))]2 Vfsx(=)) + 3(x* - 2h-1).
We can rewrite the update rule using Iterative- Average (IMA) approach presented in Proposi- tion 1.6, Sebbouh et [81]. Lemma A.1 (Proposition C.8 [63], Lemma 7.3 in [22]). The iterates (x*)
generated by NGN-M are
equivalent to the sequence {x*} generated by IMA update
A
8+1
t 7he V fs, (x*),
x+11
k+1
(5)
1+X
1+
where
A
ktl = A(xkt1 n²),
x- and
2 x
1+ X
Proof. Let the sequences fx*} and {2r} be defined according to Equation (5), Let 3 be defined as Then we have
14X
k+1 1 + A
1+1
-7 VSS. (x*))
1+X
+ X
X
((1 + A)xt - AI -m Vis. (a²1)
1+ X
1 X
x
1
7k fs. (x+*) +
1 (x -I*-1).
1+1
remains use (6) have 3
and
1+X
1+1 1+X
B Technical Lemmas and Definitions
Definition B.1. We say that the function é admits L-smooth with parameters (Ia...,, Ta), by 2 0 V.j E [d], if the following inequality holds for all x, h e Rd
(sr h) < ((x) + (Vo(x), h) + t Lh.
Remark B.2. If we set for all j E [d] L, == then Definition B.1 reduces to standard L-smoothness.
This assumption typically used in the context of coordinate adaptive algorithms such as SignSGD [4, 78).
Definition B.3. The function d: Rd + satisfies Pl-condition with constant > 0 if for all x, Rd we have
|f 2x(((z)) -
(8)
Assumption B.4. We assume that the coordinate-wiss variance of the stochastic estimator is bounded, ie. for x E Rd and j E [d] we have
Es -
(9)
25

---

Lemma B.5 (Lemma 4.9 from [66]). Let each ft be L-smooth for all i, then the step-size of NGN satisfies
E
[
c
(10)
Lemma B.6 (Lemma 4.2 from [66]). Let each fo be L-smooth for all i, then the iterates of NGN satisfy
4cL
2c² L
2cL 1 0 2cL +1" S's
N V (x
(k(Ss (1*)) [(²)
max
(11)
+ 2cL
+ cL
Lemma B.7 (Gradient Upper Bound), Let d: R - R satisfy Definition B.1. Then, for all Rd and all j E [d] we have
2L,(f(x) - /r) z (
(12)
Proof. From Definition B.1 we have
min f(x) S min f(x hycz) S f(x) min AyER
V)f(2)hy +
Tuj 2
n3
ZERd
RyeR
Now we can explicitly compute the minimum in the right-hand side. The optimal value is achieved at
Kr -ー= Lr
{z(())
therefore,
fr S f(x) + vzf(x) + 1 (hi)²
= f(x) (V,[(x))) + 2 (V,[(z))²
Lj"
2Ij
1
f(x) (V,f(x))²
2L
which is equivalent to the statement of the lemma.
C
Convergence of NGN-D
First, we provide NGN-D pseudocode and the main convergence results.
Algorithm 3 NGN-D
Input: _ E Rd step-size parameter
2: for k == 0, 0,.... K 1 do
3: Sample a batch S₁ C [ni] and compute Ss. and VSs, (x*)
(1) 4: Compute 7
215, (ak) [yfs(("*))
5: Update
k+1 ()
= -
(1) 11 Vifs. (x*)
6: end for
Theorem C.1. Let each fi satisfies Definition B.I. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters {c, )d such that Cj S 1/2L, satisfy
12( 1*)
min E oskek
18L,
(13)
Cmin K
Cmin
j=1
where Cmin == minjeld Cj. Moreover, if = O(=2) for all j E [d] then aftrr K = O(E we obtain min IE [[]]f(**]]3) < O(E²). 0SkeK
26

---

NGN-D converges with classic rate (1/VK) similar to Adagrad [94]. We highlight that, to the best of our knowledge, NGN-D is the first algorithm that uses diagonal Polyak-type stepsize and converges under standard smoothness and bounded variance assumptions without requirements of bounded gradients and interpolation.
Theorem C.2. Let f satisfies PL-condition and each fo satisfies Definition B.1. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters {c,}
such that Cj z min(2/22, 0/u} satisfy
9
E [f((x*) - 1 S HCmin/6) (f(2))
+
Σ
L
(14)
puCmin j=1
where Cmin = minjeld Cj. Moreover, if oj O(s) for all j E [d] then after K
max ()({=))() loge-1 iterations obtain F. [](*)) lr] s O(e).
To the best of our knowledge, this is the first result of the convergence of the Polyak-like step-size algorithm under the PÉ-condition. The convergence guarantees are similar to that SGD [22].
Now we are ready to derive the step-size bounds.
Lemma C.3 (Step-size Bounds). Let fs. (x)): Rd + R be a stochastic loss of batch Sy at iteration k. Let Js, (x) satisfy Definition (B.1). Consider 2/ as in NGN-D (Algorithm 3), then we have
1 03 c)
75 E
(15)
Proof. From Lemma B.7 we have 2fs - fs.) 2 [Vffs("*))) Since we assume that each f3, 2 0, then fs, (x*) 2 (v, (x ))² or equivalently,
0 (V,ss, < T.j.
2fs, (z)
Therefore, for all j E [d] we have
$ Cj cc₁
+
62
V
Js )x²
21. (w
and
y/
Cj
Cj
iss. (x*))
1 L
2fs, (x*)
which concludes the proof.
Lemma C.4 (Fundamental Equality). Consider 74 as in NGN-D (Algorithm 3). Then the following equality holds
2 (V,ss = Cj 13-11
fsa (2
(16)
Cj
Proof. From NGN-D (Algorithm 3) we have
(1+
Cj
(V)5s((="))))
2/t =c,-
2fs, (x+)
which one can rewrite as
Cj
-
(V -
2S5, (x*)
2f5 (x
It is left to divide both sides by
27

---

C.1 Convergence in General Non-convex Setting
Theorem C.1. Let each fi satisfies Definition B.I. Assume that Assumption B.4 holds. Then the
iterates of NGN-D (Algorithum 3) with step-size parameters {ep}z=1 such that Cj S 1/2,, satisfy
12(f(x)) -
min E [ OSkcK
Σ j=1
183,
(13)
Cmin K
Cmin
where Cmin = minjeld Gj. Moreover if O(=2) for all E [d] then after K = O(E obtain
min E [[ O(
OSkcK
Proof.
First, we write separable Definition B.1
d
f(xht1)- - f(x²)) = f f x
23 v Vifss -f(x
v fz((x*) -7 vffs Σ Vifs, (xt))²
d
- Σ (,f(a*) V,ss +
(17)
Note that both o and V (x*) depend on the realization SR, thus we can not directly apply conditional expectation with respect to x² as in this case we would have to analyze the product 2 Vfs (z*). Given bounds of the step-size 7/ from Lemma C.3, we can write the step-size as follows
GL,
9 +U I+ It where w/ E [0, 1] is a random variable. Varying the value of oy from 0 to we cover the whole range of 75. Thus, we continue as follows
W (iff(("") (2*)
G4 V (x*)
9 V,[(=*)V,/ss
1+
Lg,,
$ 1 1+ , v,f(z (V (xx) + 1
1/ V (x*)
L
S G v,J(2* V4, /s,, (x*)
1+ VV (x+))
1+
Now we use the inequality |ab| S / + 112 1 a b12, and derive
2E, [IV Nyfs < [,f(2** ) | fs. (I + [V V (f
2v. |2 2 | (x**) (1
< (){
Therefore, we get
GL ((V)S(*'))PO))
-Ex V,VV" 5
9 _- (55((*))
1+ct, Gl, of. ,
( (1-
=
(18)
28

---

We plug in (18) into (17) and get
d
+ 4,0 Lx 1V,
Ex [f(x*t))]
f(xx)
Σ
Ex
5=1
d
+ 4
Σ
[bf(E*)
-j 1 c,ر,
j=1
L
+ L 1
+
og
1
2
If Cj S 2z,.
1, we get
3lzr
Ex [f(
- f(x))
$ Σ
cj
V
2 0z
12
j=1
We continue follows
3l
Ex [[(x+tt]] -
Cmin
| [V
||
oj.
(19)
12
Taking full expectation and unrolling the recursion above for all iterations {0, ...., K 1}. Thus, obtain
K-1
min E [[] OSRCK
Σ
[[IV)(=)) <
12
(f(x") f*) 18
Cmin
Cmin jml
k=
we choose each Cj
Cn vk
such that co. S
22, we ensure that og
22g as well. Plugging this
step-size into the bound we get
min F [[IV](>*) osk<K
12
(S(x²) 1*)
18
Lyjo
U,J K
min Vk
min
K
j=1
12
18
(5(x9)- (
T T, 00.5
C, v K
Cu,min V KK
j=1
where C0.ii min Co.-- If we choose K= 0(=-) we get that .E[d]
min E OSkcK
O(e).
C.2 Convergence under PE-condition
Theorem C.2. Let f satisfies PL-condition and each fo satisfies Definition B.l. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters {cgb such that ej I min(1/22,, 6/u} satisfy
9
[[((*k)- - 1 (1- MCmin /6) ([(=))(())))))))))))))))())))))) -5)
L
(14)
p&miin j=1
where Cmin == minjeld Cj. Moreover, if Cj
O(e) for all j E [d] then after K =
max(O(==1), log iterations we obtain E [f(**)] f*1 < O
29

---

Proof. We obtain (19) and use Definition B.3
Eg [f(x - f(x*)
Cmin
Σ
oj
12
-=1
ptmmin (f(z)) +
Σ
3 5
6
j=1
Subtracting f from both sides of the inequality above and taking full expectation we obtain
d
31
3l
E [[((a [*] S (1 /cmin/6)) [
+
Unrolling the recursion above for fo..... ........ K 13 iterations we derive
9 9L
[[f(**)~ f*] (1- HCmmn6)) K (f(x²) - f*) Gmin ==1
A
Now we follow the proof of Lemma A.3 in Garrigos and Gower [22]. Let us choose Cj min (/2,, s/2dA;3- Together with the choice of 2 max max jEld] # f log -
12A, 12L,
we get
- /Cmin// x (f(x*) f³) s
Now we have two cases:
1. Gmin does depend on e, then we have
1
A,e O(E)).
Cmia
2. Cmin does depend on E, i.e. Cmin (E), then we have
1 A - O(s).
Cmin
Therefore, combining all together we get
E -
{ # log 2(((a9)-8)) iterations.
[1 12A,
12L,
after K 2 max max
jEld]
D Convergence of NGN-M
D.1
Convergence of NGN-M in Stochastic Setting
Theorem 4.3. Let Assumptions 4.1, 4.2 hold. Let the step-size hyperparameter > and the momentum parameter 3 X T be constants where X I (.5 2cL) 1). Then
the iterates of NGN-M (Algorithm 1) satisfy
([((==1)-f(<"" < 12"=2" + + 2cL 2cL max (2eL 1,03 op ck
where HK- is chosen uniformly at random from {I".... { K-1), Moreover, if set O(//R) then obtain E [](K=-1) f(x")) < 00/VK).
30

---

Remark D.1. In fact, if x I (L+cD)(122oL) then it implies that X for any >> 0.
1 c
because
(1+2)
Proof. To prove the convergence of NGN-M we consider IMA formulation Equation (5):
1
x =1 = x., 2 x = 2h - 7/ Vis
x ++1
x
A+1
1+ x
1+1
where B 1+x gtz xlt1 + N(xht1 - x²))
At iteration &s 0 we have
2 = 20 -000000 V.so (x.9) x0 V..
Therefore, we get
|| x²1
11 2²1 -
20 =x")
2
V (x,")
70
Lem. B.6 S
4cL
11 2-112-2701 V. (z ),20 xx') +
70 (fs, (x²) f50)
+ 2cL
L max
2ch -
2cL +1 0 ) 1300
(20)
1 + cL
Let 70 p 70 where p = (I+cL)(+2eL)
Then we have
70 = -P
Lem, B.5 S
(1 + cl. + 2cI.)
1 + 3cL + 2c -1
(1 + cb)(1 [1|11220 (12L))
+ 3cL
L (1 + cL)(1 + 2cL)
L
1cL
Using the above we continue from (20)
conv.
4cL
|| - 11
S
11
11100220(sso (fso)(()))) +
1 2el 70 (x9) 13))
T.
( 22cI. -
22L +1' 0 0)
+ 1 1 cL
max
130
< 1B xx 122l22([s0
Sso (x*)) - 27o (x²) 150 + (z=))=((==) ( 130
4cL
2
{ 12cL - 1 0 ]
+ 1+ 2cL
r0(Ss0 (x)) -15)) + 1 cL
max
2el +
2cL
12 11112-2p(fsn ffs((x)) - 7 p 1 2cl 70) (fso(a") - -2%)
2c L
12cL -
+ 25o (f%, f2,)) +
max
2cl + 1 f3..
(21)
1 cL
Here we have
2cL
70- p + 2cL
70-
1 + 2cL 70- p
1 2ch 70-
(1+ cL) + 2cL)
Lem.B.5
1
1 2cL 1 cL (1+ (1 2cL)
= 0,
31

---

70 S
I. I+2eL.
and (1=) f3 2 0. Hence, we get
6
1| - x²1 1 1 |120 - x²112-2p(1 (ppsss (20) - {z=))")) +
(fso (x") -3%)
1 - dcl
2
max 22c -1 2cl. +1 0) S/00
+ cL
Rearranging terms and taking expectation we get
6c L
2pE (f(x')
- [xx)]] S E [||=1 -11='=x""0 1x'1=111'"
11 -x"11"
2
+ 2cT. iit
2c²L
2cL 1
2c.1. + 1 0
+ ct.
max
poss
(22)
Next, for 1: > 0 we can use the relation 2² = x² + 1(0* - x=-1). We expand || - x2
||
x²1
1| - 212 2 V fse (I*) ) - x²) + V (x²
Lem. A.1
||* - 20112 1 V Ss. ), - -x)) 27x V Ss, (at)), 126-11
+ ~Rl V fsa (12
conv. S
1| - x²1 -2r (fse (z²)) - (s)( - 2Tx A(fs - fsa (2h-1))
+ v2l V fsx fs )1
Lem, B.6
|| - x' x²112 2 -
/ (J (x (x))) 2Ts (2*) Ssy (2=-11)
4cL
2 (2cL -
+ + 2cL
(fs, (x*) - fs.)) +
1 1 cL 12cL + 1 0 S.r.
max
Let Thk p + Nr. where 2h 2 0, and p is constant step-size independent of Ss which will be defined later. Therefore, we have
|| x11
11a x 1-1 p (x*) (x))x))) (fs (x*) fs²
ln (x*) f3,) ) + 27x 2fs (x
fb.))
4cL
L
12cL - 1
+ 1 2cL 7k (fse (2 - -f3,) 1 cL
max
10cL + -1:0}
0 f3L.
= 1I - x211 -pp S (x*) - fs(((") (x")) - 27 (fse (x*)) - .(S.) + (x=)
- -55.)
2nxA(f. (2.)) -5%,) + +2 (fse (xM-1) f8,.)
4cL
L
12cL - 1 )
+ 1 2cL
78 (k -33,) 1 cch
max
22l + 0 fzr
2cT.
11ah - xx111-20( 1p-sss (x*tss -fs (x*)) - (T4 + +nk
(fsa (3a
+ 2el 7k
+ 2fs (2")
-f3)) 2yn (aa-1)
2ca L
12cL fS..
+ 1 + cL
max
(23)
2cL + 1
We need find p such that
2cL
R + n
1 2cL 7/ 2 0
Since 7n 7 p. then we have
2cL
7k- p 7k7
7k 2 0
+ 2cL
2cI.
+ 7k
1+
2 p.
+ 2cL
32

---

The inequality above is satisfied if it is satisfied for the lower bound on 7 (which is 9/1+cl), i..
( + 2cL + +2) 2 p.
c
1 cL
We can take p ([+cL](1+2oL) since À 2 0.
m = 1 - p
c
S c (1 + c1) + 2cL)
1 + 222 1,2
(1+ cL) + 2cL)
3 333
$ c²
(1+ cL) + 2cL)
3
2cL
Using the above, we get from (23)
|| x=1 112
1 xx x=1 12112-2p(s5s
2p([ss (x*) Sss (x")) + (S - Sss (x*))
6c
+ 2cM(fs (x*) - ((,)
(fse (x*) -z3)
1 dcl
+ 2
(2eT. -
max
2cL +1 1: 0 S'L..
1 + cL
Taking expectations we get
E [[] ==]
S
E [1| - [3] 2pl [f(x*) ((x)) + 2cAAE
[f(x f(x))]
6 L
26 L
max (2cL- -1
2eT. + 1 0
02 pos (24)
+ 12c + 1 1
Jit +
1 cT.
Rearranging terms we get
2pE [f(=*) - f(x")] - 2cAE ((aa-1)-f((x)) S E [112* -x-1]]] - [[| x x21
L
2cd +
iint
2
1 2cL
L max
2cc 1
2cL 11.1111) 1
ppns (25)
cl
Combining Equation (22) and Equation (25) for iterations
(1...., K - 1} we get
K-1 1
K-1
2pE [f(z²) - f(xx)(]]] []==")+++0) + 2p
Σ
E [f((at) - f(x*)] - -2cA
E [f(( f(x²)]
R=1
k=
K-1 E [f(xt) -(xx)] - 2cA
K-2
2p
Σ
[[((tt)- - -[(x)]
k=U
k=0
K-1
$ (2p 2cA)
Σ k=0
E [ -s(x))
2.21.
(2eT 0
$ 11 x= 212
1 + 2cL oi +
1 cl
max
12cL +
pos
6c L
2c max
2cL -
+ 2cA +
(K-
(K - iu
+ +KK- - 1).
2cL + 1 0
pps
1 2cL
1 cL
GG2L.
2. L
2cl.
$ 11 x²12² 2cd +
Ko + K.
1 cL max
0 pos.
(26)
1 2cL
1cc.
33

---

We need to ensure that cλ > 0 which is satisfied for À such that
2(1 + cL)(1 2cL) cd
+1 > 2X(1 cl. + 2cI.).
Note that we also assume that S cL. Therefore, from (26) get
K-1 Σ -E(
~ 1200 - x²1
6c
+ 2(p- cA)
2cA +
In
K 0
2(p cX)K
1 2cT.
L max
2cL -
+ 2(p cλ)
2cL +1:
Opos
1 cL
11 x²1
8c
o
2(p c\)K
2(p ct)
2c L
2cL - 0) 2
2(p cλ)
cl.
max
2CT. 0)
pos
(27)
Since p- c² and setting x. be uniformly random chosen from (I0 3........ ,....,,.K. -k11 get
[((xt))f(x*") f(x*)] S || -x112
8c p p1+al cL
2 max
2cl. - )
pK +
of
2cL + 1 } post
(28)
P
where we use the convexity f and Jensen's inequality. Plugging the value of p (I+cL](122cL) inside we get
- [(x')] s 1120 120--011 cl.) 2cI.) 8cT(1 + +cl((1 2cL.)o
F -
cK
+ 2eL max {2cL 1, 03 oposs
(29)
Choosing == O(1/VK) we get
|1:0 |
[[ -(x") 00
Jmt
pas
max (2cL - 1,0}
(30)
VK
VK
VK
Therefore, if K 3 then E [f(**) -0 f(x*)]
O().. It remains to notice that 20 x0 to
derive the statement of the theorem.
D.2 Convergence of NGN-M under Interpolation
In this section, we show that NGN-M provably converges for large momentum values 3-including the default /3 0.9 provided the LR hyperparameter is chosen sufficiently small. This natural: the step-size is too large, the momentum term accumulates excessive past error, which the algorithm may be unable to correct, potentially causing divergence. short, convergence can be ensured in two complementary ways: (i) use small momentum parameter while allowing any LR hyperparameter (Theorem 4.3); or (ii) restrict the LR hyperparameter e while permitting large (near- momentum parameter /3 (Theorem D.2).
Theorem D.2. Let Assumption 4.1 hold, Assume that there exists minimizer x² shared across all functions fú, i.e. f = f((x*) all i E [n]. Let NGN-M is run with momentum parameter B T+X X (0, A), A 2 9, step-size S min { 4AL 2 Then the iterates of NGN-M satisfy 1
E -(x))))() < 3(1x> x² 11/////2(((2)) -
K
ck
Remark D.3. Note that the requirement A 2 allows setting the momentum parameter 3 0.9, which is default choice in practice.
34

---

Proof. The proof is similar to Theorem 4.3, but we take into account the additional interpolation assumption. We start with
||
1
11" x² 01 -y2 (VIs, (x*) x - x²) + V (1²)
Lem. A.1
12ah - x²11" 21 V fsa (x (x=)), 1. + 11x 1²-11- x²)
7 V fs. (z* (z²))1
112h x² 2 27k Vfs (z*), xthh - -x²) 2 2%x \ fs. (x=*),²² 1n-1)
+78 V.s. (zxh )112
Asm. 4.1 S
(Ts. 11 - x² x²112 -11 (x - -ss (x')) - 2nxAss (x²²) )- Sss, (3 (xx- +2I (3)²) - 13s)
Interp.
12 - x² 20 227 (1²) - -ss( - -2yx (xx Ss, (x
+2 (x") - fsx(x²))
11 x¹1 11122-27c + A - LT)( - f(xx))
+2 --(xx))) (31)
Since [s - (x*) 2 O, with the choice c S 1
2L, we have
1 - Lin 2 1 + A - Lc X+ 1 2
Moreover, by Lemma B.5 we have yA 2 _ and 9h Thus, continue the bound (31) on the distance || x² as follows
2c
||2 x²1 < 11 x²11
I+ +L (1/2+ A)(fs, (x)))
- =2cd(fse (2)ss(//11) (2h-1)- (z')).
Taking the expectation, we obtain the final bound on the distance
[[] x= [13] s E [| - x² [P + (1/2 + A)E A)E [f(ah) - f(x))]
+ 2cAEE - (32)
Note that for 0, the bound simpler, since the momentum term is zero
E [1Ia* x= II]] S || x011²
+ (f(=²) - f(x'))
+
= |1 x² 1 + cL
C
(f(x9) - f(x'))
(33)
where we use 20 2. Summing the bounds (32) and (33) for iterations E (o...., K 1], we obtain
1 + (f(x²) - f(xx))) +
C
20(1/2+ + X)
Σ
K-1
E[f(x)) - f(x")]
K-1
Σ
2cAE [f(x - -(x=)))))(())]
1 cL
k=1
1 T+ (f(x²) - -((x)) + 20(1/2 + A)
c
[(( [[x')]
1l
1 1+cL cL -x)
1/2 1/2+1 x
K-2
Σ E[f(2t) f(x'))
2c
K-1
<
E [||**
x0] - [[]=At]- 11]] + (c)( - s(x)))
K=0
||
- -f(z)))
35

---

Note that
1/++
1
AcL
cL A 4(1+ cLu
>
8 AcL,
- cL
where the last inequality holds by the choice of c
1
S 4XI.
Therefore, we obtain
INI
c/2
K-1
20(1/2 + 1) E [f f(x')
Σ E[f(=*)
(=((")))("='] $ 1 (f(2") f(x")) +
cL K=
1+ cL
+ 2c 1/2
K-2
(f(((t)) -
1 cL
k=1
< c
c (f(x9) - [(x))) + 20(1/2 +1)
[[(xK-1) [(x))]
+
2 11/2+> ->)
K-2
Σ
E[f(at) - -(())))((++))
1
R=1
S 11 - 22\(2+2c\(((²)) -(x)))
Thus, obtain the rate
1 K-1
S 2(1 2(1 cc) - x 1''1--111 (2=)) - f(x)))
Σ
- -
K K=
ck
31110- x0 2=21)=/22 f(x*))
ck
where the last inequality holds by the choice of c.
D.3 Convergence of NGN-M with Decaying Step-size
Lemma D.4. We have
Σ
log(K K + 2),
4 VK
+ 1.
(34)
vk +1
5
km
k=
Proof. We refer to Lemma A.8 from Garrigos and Gower [22].
To prove the convergence of NGN-M with decaying O₁ consider IMA formulation (see Section A in the paper):
1
== 20, get1 -ttl=2 -
Vs. (2*), 7k
Ck
x
1 1+
cb
V fs,
)1
Bjs. (2t)
AR+1
1
k+1
kt1
1 A6t1
1 A
where Cc
Co Ao 0.
VKt+1'
Theorem D.5. Assume that each fs is convex and L-smooth, and that Assumption 3.2 holds. Let the step-size hyperparameter is set .R. momentum parameter x S min + +c L.) -. Then the iterates of NGN-M satisfy
S 5(1+ + co -10 112
E [[((3- -((=))
4coVK
log (K + + VK
+ 10Lco(1 + L) + 20 L)
+ 5c + L) log( K + 2)
max {2cpL 1, 02
(35)
2
pas
where jk-1=
K-1 =Kk-
Ek=0 x
t Pk
L10+2,00+20,)) L.)
C
36

---

Proof. At iteration A = O we have
zt -0 = x² - -70 V.s
Therefore, we get
| x²') 2=1
11 x=112-2 (V S.0 )=
20 - =²² + 70 V Sso (12
Lem. B.6
LCgL
S
1|10 x²1 1=112-2%0 / fs. (z ), - -x") +
1 + 200 70 r0 (20) - f50
L0LL
20p L + 1 .0 S..
200
+ 1 + 1o
max
(36)
Let 7o - Po + To where Po = [1+o2][14200)
Then we have
70 70- - -Po
Lem. B.5 S
Co
C0 (1 L1 + L)
Co + 3c + 2.2 -1
(1+ (1+l))
30yL
aL (1+ (L + 200L)
3e
1 2coL
Using the above we continue from (36)
conv.
- S5o (x))) + 4coL
||
x²11
S 11 1=112-220006
1 2coL (o(Fs((=)) - {80)
20
1200L 1 1
+ 1 + Co..
max
2e0 +1 0 S3.
< 21 x= - (ppo(fss(1)) (20) - (x")) - -7 (fs,(2)) (2 ) f3) + 27o (fsu((()²") 5%%.
fss
AcoL
20²L
220 - 0) sso
r0 (I¹) (10)- -5%)) + 1 c
max
200 + :0)
1 200L
2c
11 x² 111--20o(fs (x.) fso (x))) -2 70 Pu
1+ 200L 70) (fso f5o
2
12o -
(Joo (x²) - (8S))
20yL 1 0 S'3.
(37)
1 CoL
Here we have
200L
2o Pa- 1 200L
110
70 Po
1 + 20,
1
Co
70
1 2cL
(1 + cI + 200
Lem.B.5 B.5
1
Co
CD
1 200
1+ coL
(1 + c1 + 2cL)
=0,
L 1+2ooL and fs. (-²)- f3, 2 0. Hence, we get
50 S
GoğL
|| x 1|2 < 1|:0 11 18-2ppo (po (x0) - (so()²)) 1 20 I (Sso (x)) -55)
2.
22 I -
+
max
20yL + 1' 0
1 GL
37

---

Rearranging terms and taking the expectation we get
6c L
2pdE ( -
S F [||= - =x11? - - -x=11++
1 + 200 o
20L
{ 220L -1
200 + +1.0) ) pos
+1.
max
(38)
+ CoL
Next, for 1: we can use the relation
2/ xh >x(2 tk -xx-1). We expand || - x01
||2 1+2
x=' x²11
1|* x²12 221k (V Ssn (x 2t x²) + V Ss. (x²))²
||* 2012--214 VV Ssn (I. ), x²) - 26 V Ss, (1 1 -08-11
+ >2l Vss, (x²))]
conv.
||a 2--12=212(2(( (z*) - -fse(=')) 2nx (x+) fSx (2th-11)
+ 7 fss (12
Lem. B.6 S
||* x²1 112 1-2222202 2 (Js. (x²)) - -SS. -ss(((')) -2TAAx (x*) )- fsa. (2x-11)
Lc&L
207L
12xk - 1
20 Li + 1 0 ))
+ 20x L. Th (Ss, (x*)) - -fS,) + 1 Ca
.0) SoL..
max
Let 7h = Ph + 7ik where p, Tt 2 0, and is constant step-size independent of Sy. which will be defined later. Therefore, we have
|| x11 <
1F* - 2p (x*) Ss. (x")) -x -xx
2n (x*) 13.) + 27x 2TsA(Ss (1
1(5x) -
4c&L
20
max { 12c - 1 220 L + 0 f3x
+ 1 20 Tk (fs, (fs)((") - -f3,) + 1 Cr L
1 x² 112-2p( fsx (x*) )- - fs. (x²)) 2 (fse (x²)) - -sS.) + -2 (fs (x²) - -53.)
2nx(fs, (x. - f8,) + 27s (x^-11 {(s))
+ 1 2c4L The (f. (x*) -fs)) + 1 C+L
4c
L0
2c L + 0 fsr
2c L -
max
+ 20 7K ) (fs. (x*) -f²))
2c
11ah -2p
(x*)- fsr (x*)) - 2 Fis + 7xd +oo
+ 25M(Ss (x))(("))
-(8)) 27 (IOh-1)
-f's))
L
120 L 1 20 To + 1 0 Sāc.
+ + Ct
max
(39)
We need to find Pk such that
2c
m + 7nd
?/ 2 0
+ 20
Since 7/2 = - - Pk. then we have
20x
Fik - Ph: + 7b>x
7/k 20
1 2c L
2c
7k 11+1p Av
Pk.
1 2cx
The inequality above is satisfied if it is satisfied for the lower bound on 7h (which is 9/1+cl), i.c.
Ck
(1+2
1 L + +2) 2p
1
+ L
38

---

We can take Pk = (1-22) since X 2 0.
% 7/-- Pk
Ck
S Ch
(1 + ca [.) + 20x I)
1 30g I. + 20 -1
= -ck (1 + L) L)
+ 3cx
S Cg (1 + ca 1. + L)
36 L
1 20
Using the above, we get from (39)
|| x²11 - 11 - x= 2pp - (x")) + 20x (x fsx (x*))
+ (fs (x*)
-(%,) + 1 20x L (fsx(x*) - (sa
20
120 L - 20 1 +
+ 1 C.
max
0 15xx
Taking expectations, we get
E [[|a [|| x=]
S E[1l** - x² x²1]²] - E [f(xt) - -(ax')] + 2c [f(x f(x))]
+ 200do= + 6e
267.
12c4L -
22 + 0 02 o
i I+ 1
max
1 ++-20x11
Rearranging terms, we get
2p&E [5(**) - [(x))] - Ax [ - -s(x)) < E [1I=* - =x=11]]- - E [1] 11133]]
T.
+ 20xxx +
iin
1 2cL
20
12 L
+
1 + CkL
max
20 + 1 {
0
x2
pas
(41)
Combining (38) and (41) for iterations {l...., K 13 we get
K
K-1 PRE [f(x*) - [(x)) - 2
EpdE [5(x9) (((=)) +2
Σ
GoARE [
- [3x)))
1m
km
K-1 2
K-2
Σ
PkE [f(tt) - -
2 Σ
Cxd+EE (f((t) - f(xx)]
k=
k=
K-1
< 2
Σ (pk - cxddEE [f(aa) [f(x))
k=0
20
200L - 1 (
< ||
11+
1 20
22 +
1 LoL
max
200 +
_ opos
K-1
K
207
12cx -
2crda= + 1 20 Ib
Jit
2
1 1 C,
max
2cx Iu +1
0
2
pos
(42)
k=1
Note that choosing A = min min farl, 0.5(1 + 4 2 + 20 () ensures that 2 2 cx In- deed, we have
Pk
> ckds
2 2(1 ct L1 + 2c
+ 1 > 2 (1 + c [1 + L).
39

---

Therefore, from (42) and the facts that λo = 0 and d S CaL we get
60
Σ pxl -((o))(((>)))
s 11 -x-11 Σ
2ck-x + 20x 1
int
K-1
20 L
2ct.
+ Σ
op
max
1 c
2ch +
0=
K-1
< 11 x11 =11 +12+8To0
2
int
Ck
k=
K-1
22 L
Σ k=0
2c L max
2ck + 1.0]
0
ppos
(43)
We have by Lemma D.4
Σ Pk
K-1
K-
Ch
2
K-1
Ck
2
4co 4coV
(1 + c& - 2cxL)
(1 + co + 200L)
5 co + 2c0L)*
k=
k=0
Lem D.4
of log K +22,
(44)
K-1
2c
K-1
220 1 20 + 1' 00 s c log(K + 2) max 20g I.
Ck max
0 1
c
max
200 + 0
20, L + 1
0=
k=
Therefore, using (44), =0 = xx" in (43) and dividing both sides in (43) by Σk= Pk we derive
K-1
Σ
ス k- kk=
Pk
E [f(x*)
f(x)) S 1:00 101
+ 8Lago
log K +
EKJ Pk
int
K-1 =k= Pk
k=
+ 20 L
log( K
Σk-1 Pk
+)
max
220L -
200 + 1 0
pos
(45)
With an lower bound on Ek=1 and Jensen's inequality we conclude that
5(1+ 5(L+co coL) + 1|x 212
E [[((a - f(x")] -ff"))=))))]
4coVK K
log (K + + VK
+ 10Lcp(1 + (1 + 200 L)
log( K + +))
+ 5c + coL)
max {2cgL 1, o2
(46)
2VK
where K1 K-
k
Et=o Pk
E Stability of NGN-M on Simple Problem
We consider 1D convex functions of the form f(x) LL (x)) that satisfy the following assumption.
Assumption E.1. There exists a constant C' such that C(1 2 xx)x)x(xx).
Note that + px() 2 and deg(1+p"(z)) deg(ap(z)p/(z). Therefore, this assumption mild. Remark F.2, For example, the function f(xx) = 2² + x²) (i.e. p(x) x) convex and satisfies Assumption E.1 with C = 1.
40

---

Remark E.3. Let p(x) amx² (x) (x) mamm"" Therefore, the constant C' should be expected of order C m, where m = deg(p())).
Ej=0 azz'. Then for large values of z in magnitude, p(a) ~
The function f(x) is non-negative for any æ E RR and its minimum f = 0 is attained at x = 0 by design. Let us compute a step of NGN-M on this problem
xkt1 xt -(1 -11- -)
(1-3)
(f(()))²2 (xx* B(x xh-1)
2Lc(1 + +
xb 18-(( -(1- (1) (3)
k yp' ( x² + B(at - xh-1)
21 [1t+""ppp'p
2Lc(1 + + (x*
p
2. - (1 11-1) () 2he
I² 3(0 xb-1).. (47)
(1 + p² (x²) + x p( p(=x))((<>)))
Note that the convexity of f implies that
f(0)= f(x) f())+
0 Lx -=
02 Lx p 2(x))())(())
xp(x)p' (x) 2 - (1+
(48)
In particular, (48) implies that 1 + p²(x) + xp(x)/((x) 2 (1+p²(x)) > 0. Therefore, we can
obtain lower and upper bounds on 7/.-
Lemma E.4. Let Assumption E.1 hold with a constant C > 0 and f(x) = x² + p² (x)) be convex.
Let t.
Then we have 2 E 2 2
Proof. Indeed, the upper bound on ny follows from the following inequality
2Le(1 + (x + p(a""""
4 = 1+ (+p7(zz) (1 (x*
2he
2Lc(1 + p² (x
2he (x (1 + p2 (x*) + p(a (
1 1+ + x² )p' $2
(49)
due to (48). The lower bound can be obtained as follows
2Lo(1 + (x*)) + x p'
tx 1 14p (1 p2 "*
2he
2L + 12 (x
(1+ p2 +c (1))())))))))))))))
2Le
2 2(1 + (x+) + x x<* (1 p
Lc1+p' p2+))(x²))
= 1 Lc(L + (x+) + 2'p(==
Le(1
1+ Le(1 +
Lc(1
(50)
2Lc(1 + (1
C)
41

---

The update rule of NGN-M can be rewritten as
k+1
z (1 3) + B(x* x 11-1).
(51)
Let us consider the joint dynamics of w (1=4)T (1T=-1T)T E Rad, We have that
z xb-1
xx (1 3l + 3(00 - )
w
xk-1
(1 - Bi + +1
31 ( 1/b
I
(xh-1)
Gw
(52)
where
G I
B 31
-3I
(53)
0
Now we are ready to prove the convergence of NGN-M on this simple problem for any value 2 2
Theorem E.5. Let ((x) =² (1+ be convex and Assumption E.1 holds. Let B (2(1+))-11 and c 2 zI. Then the iterates of NGN-M on fx) converge to the minimum f 0. (11+C)+11
Proof. We follow the standard proof of SGD with Polyak momentum [71]. At this stage, we need to estimate the eigenvalues of G. To do so, will proceed with permutation matrix III which transforms the matrix to the block-diagonal matrix as
G
0
0
G=
(54)
0
G
where
G 1 1 31+3-(11 00==1+---00 (1- 8)r
)
(55)
Since the matrix G is block-diagonal matrix, have ||||S max, | Therefore, the problem now simplified to bounding the spectral radii of the individual blocks G, for = 1,2,.... d. The two eigenvalues i1 and u2 of G; are the roots of the quadratic
quu) == 2 (1 3 (1 3)%))u + =
(56)
which take different values depending on the discriminant A = (1 + B (1 3) 4,3. Let us find the values of when the discriminant is negative. We need to satisfy the inequality
(1 + (1-3)))) 45
-)1-8) 2(1 B)(1 8)9x 48 0
+ (1 - + B)(1 3)5x 50
+ (1- s 2(1 8)hu
1 1+ 1 1+3
(57)
2
1
Since the function 1+y2 for y E we satisfy the last inequality, and consequently the discriminant non-positive, if we choose
2I+CJ 2
attains the maximum 40+003+1 y पा. then
4(I+C)
1+3 4(1 002++ < 1 1-3
(58)
4(1 C
1 i odd .j =
"The permutation matrix is defined as Ig
1 even . = + 4. Note that permutation matrices
0 clse
preserve eigenvalues.
42

---

The above inequality satisfied for B E of [8 have A; s 0 for all [d]. Therefore, the zeros of the quadratic ((u) complex, and equal in absolute value
(2
1).
Therefore, we obtain that for such choice
[2(++0)+1)3
|u |22 <
(59)
This gives us that ||G:|| vB 1. Therefore, the algorithm converges for any value of 3 this range.
remains to use Lemma 11 from Foucart [20] which says that for a given matrix E Rdxd and e > 0, there exists matrix norm ||-|| such that
||| p(A) +
(60)
where p(A) max [AA: eigenvalue of A} (spectral radius of A).
Asymptotically (as + co, one can show (see Theorem 12 in (20)) that
|| 2 ((p(G)))
(61)
where p(G) S vB < 1 in our analysis. Therefore, NGN-M with hyperparameters c 2 and / 2 converges.
Remark E.6. For example, NGN-M converges on ((x) x²² +²) for any and , 2
9
25
Theorem E.5 shows that NGN-M remains stable even with an arbitrarily large step-size hyperparameter Thanks the adaptive nature of NGN step-size, the actual update scale is automatically shrunk when necessary, preserving convergence. Importantly, this is possible with a choice of momentum parameter close 1, which extends the results of Section 4. We acknowledge that our current analysis is restricted to the special convex class of 1D functions [(x) = x² + p² (ar) satisfying Assumption F. Extending such stability guarantees wider function classes with large momentum 3 remains a significant open challenge.
To support the theoretical result, we test the performance of NGN-M and GDM (Gradient Descent with Momentum) on the problem f(x) = (1 + x²), which is convex and satisfies Assumption E.I; see Figure I. We run both algorithms, varying the step-size hyperparameter in {10 We run algorithms for 10 iterations. We stop training if the loss reaches threshold 10 15 or exceeds 10"h for the first time. We observe that (t) for small step-size hyperparameters, both methods converge but do not reach the threshold 10 15, (2i) NGN-M reaches the threshold even for extremely large values of the step-size hyperparameter while GDM diverges. (iii) the fastest convergence of GDM is achieved with the step-size hyperparameter 10-2 after 691 iterations while the fastest convergence of NGN-M is achieved with 10 after 269 iterations. other details, NGN-M achieves faster convergence and much more stable to the choice of the step-size hyperparameter. These results align well with our theoretical analysis.
10a).
20
Final Train Loss
10
105
NGN-M GDM
10 5.
10 15
10
410
100
100
10
10
10
103
10*
Stepsize
Figure E.1: Comparison of SGDM and NGN-M when minimizing function f(x) x² +
'A non-asymptotic version of the analysis can be derived using Theorem 5 by [91]
43

---

F How to Derive Diagonal NGN-based Step-size?
Here we provide derivations of how combine NGN and diagonal step-size following Section 3.3 for completeness.
We consider the following model
p² argmin [Sanc +p) (r(x*) Vr(x*)Tp))
1
(62)
2c
where r(x) = Vf(z). We compute the gradient of RHS of (62) w.r.t and equal it to zero:
Vpf +p) = -r(x(²)) + Vn(a*) Vr(a*) + ExP
= (2Vr(=t)Vr(a))) + - p+ vr(aa)
Therefore, we have
(2Vr(=t)Vr(tt)) --
-1
p² -
)r((*))
Using Shermann- -Morriso formula (A uv -)-1= A
A-' A-'
1-u l-u A- with A 1/2E we derive
2, Vr(x*) (r(2)))) 1
n² +24
2r(x")rr(x)
1 2c(r(x*) E/ Vr(ett
2cVr(x*) 2/1 (r(z*)
)
ET'vr(**)
+ 20Vr (2
E5 Vr((*))
)cr(a*)
Dx 1 Vr(o").
1 2c EE Vr(wk)
Now we plug-in in r(æ*) = v ((*t) and Vr(**)
2/5( ff(**) and obtain
2cv f(**)
1 + 2c [y(
(x*k)
Σ Vf(x*) 2 f(*h))
Σx
(z**)
2f(a"
F.1 Design Comparison of NGN-MDv1 and NGN-MDv2
The derivations in (3) are used to provide an intuition of how one can add diagonal step-size into NGN by choosing the regularization matrix Σg. By choosing Σ, = D3 we recover the update direction of NGN-MDv1. In this case, we have only one global NGN step-size in front of D The design of NGN-MDv2 follows a more straightforward intuition. In particular, can be seen as a direct extension of NGN diagonal case by replacing the squared gradient norm [[,,,, (a*) [[ by the squared partial derivative (V fss. (x²))² for each parameter j [d].
The main difference in comparison with Adam is the order in which the preconditioning and mo- mentum is applied. In both NGN-MDV1 and NGN-MDv2 we average the preconditioned updates ' (x*), i.c. we first apply preconditioning and momentum later. In contrast, in Adam the stochastic gradients averaged to construct new momentum term, and then the momentum is pre- conditioned. other words, the momentum is applied first and then it followed by preconditioning. We believe this change might be one of the reasons behind the step-size hyperparameter resilience as well.
In practice, we found out that the tuned performance of NGN-MDv1 is slightly better than that of NGN-MDv2, Moreover, NGN-MDv1 demonstrates higher resilience to the choice of the step-size hyperparameter than NGN-MDv2.
44

---

Table 5: Train time of Adam and NGN-MDv1 when training language models.
Model
Method
Time per Iteration (sec)
Time per Optimizer Update (sec)
AdamW NGN-MDv1
1.63 + 0.01 1.65 + 0.01
0.0048 + 0.0002 0.0130 + 0.0002
70M
AdamW NGN-MDv1
3.33 + 0.03
0.0088 0.0003
160M
3.37 + 0.02
0.0239 0.0003
AdamW
8.41 + 0.06
0.0838 + 0.0009
410M
NGN-MDv1
8.68 + 0.06
0.2154 = 0.0007
F.2
Computation Cost of NGN-MD
Implementing any version of NGN-MD in practice might be slightly more computationally expensive. However, highlight that computing step of NGN-MD does not involve matrix- operations since the preconditioner a diagonal matrix, and the matrix notation is used only for the convenience of presentation. The additional computation cost that we have in NGN-MDv1 is the computation of [[V (x I This can naïvely be done by one additional pass over the gradient and summing
the terms [D Vv (x*))² for j E [d]. This operation does not require additional matrix multipli- cation. However, it be computed more efficiently while updating Dx. The of the NGN-MDv1 implementation does not add any significantly costly operations in comparison with Adam.
We compare in Table 5 the time per iteration and optimizer update when training language models from Section 5 using AdamW and NGN-MDv1. We notice that our naive implementation of NGN- MDv1 is about 2.5 times slower than PyTorch's 's AdamW. This is expected since algorithm requires two passes over the gradient. Nevertheless, in this setting training time is dominated by forward and backward computations, keeping NGN-MDv1 competitive with AdamW. Moreover, noted above, this overhead can be largely eliminated by computing the weighted gradient concurrently with the second -omentum e update. We do not aim provide the most efficient implementation of NGN-MDv1 as the primary goal of our work is to highlight the stability advantages that NGN step-size brings in the training of neural networks.
F221 Distributed Training
In a vanilla DDP implementation [50], computing the weighted gradient norm | (VV[s((*))) |1= is
straightforward since gradients are replicated across devices. We only require an additional all-reduce to synchronize fs, (x²) across devices, which is, however, lightweight communication (just single float) and, principle, can even be overlapped with the backward pass.
However, with more sophisticated types of parallelism, like Tensor Parallel [83] or ZeRO-2 [75]. computing the weighted gradient norm introduces additional communication, as gradients are sharded across devices. This could still be implemented efficiently by accumulating squared gradient entries in each device and all-reducing only a single float, but will, nevertheless, result in computation and communication overhead for NGN-MDv1. We acknowledge that our methods might not be scalable to large distributed training, and adjustments are needed to make NGN-MDv1 work this case. Nonetheless, we believe that our findings offer useful insights toward designing more stable optimization algorithms.
G
How to add weight decay to NGN-MDv1?
Regularization techniques serve a fundamental purpose in minimizing generalization error. Or- thogonal to their role for generalization, modern deep learning tasks often benefit from the use of weight decay [98]. Despite its widespread application, the role of weight decay is poorly understood. Andriushchenko et al. [2] suggested that it might provide implicit regularization by stabilizing the loss in over-parameterized neural networks and helping to balance the bias-variance tradeoff that leads networks. However, even in the case of SGD, there still uncertainty regarding how the weight decay mechanism should be incorporated, as various implementations may exist [102].
to lower training loss in
45

---

We propose two ways of adding weight decay to NGN-MDv1. The first variant follows the approach of [55], adding decoupled weight decay A:
r x² -ex -1 - V (x* + B1 - 11-1).
(63)
In this update rule, the weight is added separately from the update direction Σg "Ss, (x*). We call the resulting algorithm (63) Dec-NGN-MDv1, that stands for decoupled NGN-MDv1.
G.1 Combining NGN-MDyl and Weight Decay Regularization
We now discuss how to combine NGN-MDv1 and weight decay, following the idea that weight decay should perform weight regularization.
We consider the following model
1 2c
A -1 +
(x +p) = (r(**) + vr((") 2)² +
By taking the gradient of fEL,A w.I.t p we get
(r(a"))))\((( 0 = 2(r(=*))))))))))))))(())))) + Vr(x*) vr(++)++++ + ExP 2Ex (x* +p)
(2Vr(=+)V7(a7) + - + +1Ex p p+ -r (r(=t) x
Therefore, we get
p² = - (2Vr(at)Vrr(a))
+ - + +2Exs) (2r(**)Vr(xt) + \Ena*).
Using Sherman Sherman-Morrison formula (A + uv Tu-)-1 A-1 (2Vr(x*) we get that
I4a A- with = (1+1/c)22 and
1 En A
2Vr(I*)Vr
c
20a
D
Vr Vr(zz Vr EE'
Σ
1+
vr(x)))) 'v((*t)
1+
Therefore, we have
2ca
Dj Vr (r***) Ex
p² C Ex
(1+Ac)2
(r(())) + ))
1 Ac
1+ he
Vr Ex Vr(ar))
2cr )z
2e Vr
Σ vr(x*)
I+he
Ea
1 Xc
1 I+Ac
2c
Vr(zt)
20
Ae k
It 21
1+ Ac
2c Vr Vr(*))
I+Ac
(x)(x*)
Ex 'r(2a)
1+ AC
1+ Ttxc
Vr(a*)
2r² -
Ac
I+
1 Ac
1 2c Vr(2" (zx E3 Vr(xt))
46

---

Algorithm 4 NGN-MDV1W
1: Input: x.9 Rd step-size parameter O, momentum parameters ,1, [0, 1), weight decay
parameter 2 0, stabilization parameter & > 0
2: for k= 0, L...., K do
3: Sample a batch S& S [n]] and compute fs. and V (x*)
4: Compute y = 32 k + (1 (2) fs. (I* O Ssc (z(1))
5: Compute D diag(a cI + v w/ /11 ()))
6: Compute
V S3p (3x Taat
(1+Xc]
2/ (a+
7k- = 1+ 2/3, (w {c) fse (z*) D
7: Update 2h+1
1 x - (1 BI V fsc (2²) + 31(20 - 1 1+Ac
8: end for
[]+ denotes max(0,.),
Using the connection
Vr(**)
Vs(x*) and r(x*)( = (x**) we get
22 2(1())
p* 20/ 20//fxx
Vf(x*)] E 'V((tt) Dx
2 f(at)) vf(x*)
1 Ac
2c 3f(**][1+X)
Σ Vff("* Vf(x)) x
1 + A
4[ (arty Vs(att E (V((x*))
c/(122))
ed
En V)(**)
x
1 2 (1+Ac)
1+ Ac
(((**)
CA
>-'(((t
1 Ac [1+Aa]
Σ
To summarize, the update of NGN-Dv1W is the following
x x + p
T Tz
1 I Ae
[a
Ex VV((tt
1+tcl+
2)(2 111 11=
(((+Ac)
Σx (('
1
V
* Σx P (x*))
1
It
(64)
+ AC
1 11
To prevent the step- size next to Σj
V f((*)) from being negative, the final update has the form
ch
A+1
Σ 'V(('),
(65)
+
Ac
[1>x]
where [1 = max(:,0). Now we can add momentum on top and obtain the following update of NGN-MDV1W
1
I+>c
"E'(((+)) + B(a x(-1).
(66)
1 Ac
1
Vf(r*) 111
This combination of NGN-MDv1 and weight decay is summarized in Algorithm 4. We highlight that now the weight decay is incorporated inside the adaptive step-size as well as regularizing the coefficient next to *
47

---

AdamW
Momo-Adamw
NGN-MDV1W
Dec-NGN-MDV.
30idaston Peplestty
vg dastopnnn Peppeett
VO cosponn Peploxtt
VO GOadSpnnn Poplostt
3
8
8
S0ops 22
Stepsize
Stopsize
Stapaze
10
10
10
Figure G.I: Adding weight decay when pretraining 70M Transformer When properly tuned, a value of weight decay > enhances the performance of all algorithms. NGN-MDv1 retains his characteristic stability, and achieves smaller perplexity in all scenarios.
G.2 Empirical Validation of the Proposed Combinations
Having two possible ways of adding weight decay to NGN-MDV1, we test them on pretraining a 70M transformer on language modeling. The validation perplexity at the end of training is reported in Figure G. 1. We note that when weight decay is turned off, both NGN-MDv1W and Dec-NGN-MDv1 reduce to NGN-MDv1.
First, observe that when weight decay is properly tuned, all algorithms improve the baseline case with no weight decay, which consistent with the observation of Xiao [98] and Andriushchenko al. [2] on AdamW. We also note that Dec-NGN-MDv1 and NGN-MDv1W require smaller weight decay value compared to the other algorithms, Finally, the stability and performance of NGNMDv1 preserved by both variations, allowing training with larger LR, and significantly improving over AdamW and Momo-Adam.
We do not observe substantial difference between the two proposed modifications of NGN-MDv1 for this task. We remark however that these two versions serve substantially different purposes, and pretraining language models might not be the most representative task evaluate the effect of adding regularization.
48

---

H Additional Experiments on Toy Problems
H.1 Additional Experiments on the Problem with Many Minima
Now, we provide a simple example of minimizing a function
f(x) - (sin(1 + cos(-- + x)) 0.222)2+ + (sin(1 + cos(IT - x)) + 0.2r)*
(67)
that has many sharp sub-optimal local and flat global minima. We compare the performance of NGN-M and SGDM varying the step-size hyperparameter in {10', 10', 100 100} and the starting point in [-20, 20] with step 4/30". Based on the results in Figure H.1 (right). we conclude that (i) for small step-sizes, both methods likely get stuck sub-optimal local minima and reach the global minima only if they are initialized close enough to it; (ii) for large step-sizes, we observe less runs of SGDM reaching the global minima; (1it) in contrast, for NGN-M with large step-sizes, we observe more runs reaching the global minima. This is possible due to the adaptive nature of the NGN step-size that forces NGN-M to converge to the flatness of the global minima.
100
300
3G2M
Functton Valuee
MOR-M 1e
Frequency
2004
200-
Function
100
10B
-00
10
-20
100
Parameter
Parameter
Figure H.1: Comparison of SGDM and NGN-M when minimizing function in (67).
H.2 Comparison on Rosenbrock Function
Now we present the results where we compare NGN-M and SGDM when minimizing the Rosenbrock function. We report the trajectories of optimizers and training dynamics in Figure H.2 and Figure H.3.
We observe that NGN-M converges for all values of c, indicating its high resilience the choice of step-size hyperparameter. In contrast, SGDM already diverges for the step-size hyperparameter 10 This can be explained by the adaptive nature of NGN step-size, which decreases the effective step-size of NGN-M for a more stable convergence. This is especially evident from the trajectories of algorithms. Indeed, NGN-M effectively moves in the complex valley of the Rosenbrock function, adapting to the local curvature.
NGN-M c- 10 SGDM Ir-
NGN-M 1-
NGN-M c- 10
SGDM 10
0-
-2
2
2
-1
0
1
2
2
-1
Figure H.2: Trajectories of NGN-M and SGDM when minimizing the Rosenbrock function and varying the step-size hyperparameter.
H.3 Comparison on Quadratic Function with Theoretical Step-size
Next, we run NGN-M with theoretical choice of step-size hyperparameter 1/vK and e 1/vk (see Theorem 4.3 and Theorem D.5 for more details) against fixed choices c {10-3 10 10-4). The
TThis step chosen small enough that the initial point can be close to any local minima within [-20,
49

---

Stepsize
25-
107
100
Reration MP
Keration uP
Reration 107
10
10
Training dynamics of NGN-M and SGDM when minimizing the Rosenbrock function and varying the step-size hyperparameter.
Figure H.3:
comparison is made on quadratic function f(z) = (A + rI)x where A E R400x400 and R100 are sampled from standard normal distribution. The constant controls the condition number the problem.
We test the performance of NGN-M varying the condition number of the problem and the number of it- erations; see Figure H.4. We observe that in all the cases, the choice 1 1/vk T leads to faster convergence, supporting our theoretical claims. The choice 1/v K demonstrates competitive performance as well, but is slightly pessimistic the beginning of training. In contrast, the choice E (10-3, which is default value in practice, is too small and does not lead to fast convergence.
-1003, 10-1),
These experiments demonstrate that when the problem satisfies all assumptions needed in the analysis, the choice of the step-size hyperparameter given by the convergence theorems is good starting point in practice and can serve baseline when tuning
Cond. number
Cond. number 99
Cond. number 1000
10
10
10'
[ini
100
10"
10"
10a
0
50
100
50
100
0
50
100
Iteration
Iteration Cond. number
Iteration
Cond. number
99
Cond. number
1000
10a
10
10
10a
10%
NGN -
10
10
10
NGN - 1/VK
NGN c/VE
10
13
10
2
10
0
500
1000
0
500
1000
500
1000
Iteration
Iteration
Iteration
Cond. number =
Cond. number
99
Cond. number
1000
10
10
10
Tead At*)
10
10
10
10
10
2
10-15
10
5000
10000
5000
10000
5000
10000
Iteration
Iteration
Iteration
Figure H.4: Training dynamics of NGN-M with several choices of the step-size hyperparameter varying the condition number of the quadratic problem.
50

---

I Additional Experiments and Training Details
L.1 Training Details
The detailed experiment setup with hyperparameters and training details is presented in Table 6. We provide links to the exact model architectures used in our experiments (the links are clickable) as well as links to the tables and figures for each workload. We demonstrate the results averaged across 3 different random seeds for small and middle-range size experiments. We use standard values of momentum parameters (31, (B2) = (0.9, 0.999) if the opposite is not specified. The step- size hyperparameter is tuned across powers of 10 (for some workloads we add additional values of the step-size hyperparameter shown in the step-size resilience plots). We use PyTorch [69] implementation of Adam. The implementation of MomSPS, Momo, Momo-Adam are provided in the corresponding papers. Finally, when employing SGD-M, we set dampening equal to 0.9.
For vision transformers experiments, we follow the setup of Schaipp et al. [79], and use Pytorch Image Models codebase [95]. We train vit_tiny._ patchi6tchhhh _24 for 200 epochs Imagenet using a cosine LR schedule with a linear warmup of 5 epochs. Differently than Schaipp et al. [79], we train in bfloat1, instead of float16, and do employ weight decay regularization.
For pre-training Transformers on Causal Language Modeling, we build upon the nanoGPT [43] implementation, augmenting it with Rotational Positional Embedding [87], RMSNorm [101], and SwiGLU [82]. We call this enhanced version Transformer++. Models are trained with a batch size of 256, context length of 2048 tokens, vocabolary size of 50280 and make use of GPT-Neox tokenizer [5]. We adopt an enhanced training recipe, made popular by large language models such as LLaMa [89]. These modifications include: training in bfloat16; employing a linear LR warm-up for 10% of the training steps, followed by cosine annealing to 10-5; omitting biases from linear layers; using (8, 32) (0.9, 0.95) for all algorithms; clipping gradient norms above 1; no weight tying between embedding and last linear layer. All models are trained on SlimPajama-6277 [86]. a cleaned and deduplicated version of RedPajama We report validation perplexity separate subset of Slim-Pajama consisting of 10M tokens. The total compute is estimated following Kaplan et al. [41], where the estimated number of floating-point operations (FLOPs) is 6 Number of Parameters X Number of Tokens.
Experiments of small and middle size are performed on 1xRTX 4090. We perform ImageNet32 experiments on 2xA 100-40GB, and ImageNetlk k experiments on 4xA100-SXM4440GB For pretrain- ing Transformers Language Modeling, employ LLH100-HBM3-800 GPUs. With multiple devices in use, we employ Distributed Data Parallel to parallelize the training process.
I.2 Comparison Algorithms that Support Momentum
In the main paper, we provided the test performance only. Now we additionally illustrate the performance of algorithms w.r.t. training loss convergence. Figure 1.1 demonstrates that NGN-M is the most robust algorithm for the choice of the step-size hyperparameter from this perspective as well. In Figure I.1. we additionally demonstrate the performance of the algorithms (VGG16 [84], CIFAR10) and (MLP, MNIST) workloads where NGN-M matches the performance of the state-of-the. art algorithms in this setting and archives higher resilience to the step-size hyperparameter choice. The best performance results are reported in Table 7 and showcase that NGN-M always matches the performance of other optimizers improves it.
I.3 Comparison of Algorithms that Support Momentum and Diagonal Step-size
Next, illustrate the performance of the algorithms that support both momentum and diagonal step-size. According to the results in Figures 1.2 and I.3, NGN-MDv1 achieves the best resilience to the step-size hyperparameter choice among all considered algorithms. Again, NGN-MDv1 is the most stable algorithm to the choice of step-size hyperparameter w.r.t. training loss convergence. Its best performance is competitive to that of other algorithms but the step-size hyperparameter range that gives such performance is wider.
Moreover, we support our claims about stability on additional workloads such as (VGG16, CIFAR10) (in Figure I.1), (MLP, MNIST), (LSTM [32], PTB [59]), and (Transformer [43], Tiny Shakespeare [42]) workloads. We observe that NGN-MDv1 attains higher robustness to the choice of the step-size
51

---

Table 6: Summary of experiment setup with all the details on hyperparameters used in each case.
Model
Phitaat
Perforaunce
Stabiiity
Effective Results
Heeentions
Sine
Cammants
Resne20
CIEAR10
12, Fig. 2, 12,
1.0,
50
125
Roaeillo
CIFAR100
Tol 7,8
100
128
Vodis
CIFAR10
Tob 7,8
Fig. 11,12
90
128
MNIST
Fig.
128
stat 160
VIT
CHFARIO
166 7,8
Fig. 13,
Fig 5.1.1.001
200
512
ISEM
150
layeess
LSTM
Wilitesn
Tob.
150
flayers
Transforree
Rouin
Tub. 8.9
'lbb. 18
2000
16
heads g
Traisforee
Shatesspeze
Tiny
Tub 3.9
Fig. 13,
2000
16
hcaaih &
layers
R8
InzgeNett2
Til 7,8,
Fig. 14
43
128
comsitun loarsing ra12 schiedile; Wegilt decay
kraming nie clacee
Resnetts
InsgeNetik
Tub. 7.8
Pg. 2,14
90
256
30 epochs 1.1.
wizght
cosine lcaring nuce worrs-up
VIT Tie
ImsgeNetik
Tab,
Ri. 3
200
512
schedale
for spechs
weight decicy bfloo
dirr A heade
1 layers 6. costeat leegth 2018
TIM Transfurmer-+
ShinPajano-67788
Tol 8,5
Fip. 4, 1.1
3400
256
(.5 10..59] 0,000 beloasi6
olpping tarm linear for 10K of lizneiuss
d12n A hesdl
* layess 6, costest longth 2018 (2) 36) (0.8 0.90) brloa116 clppin fix L0KS of iltenionss
TIM Transforme +-+
FieeWeb
Tul 10, 12, 13
4600
128
layers 12, contest longh 2018
16004 Transfomerr
StiniPajaaa 627B
Tub 8,5
Fi.. 4,1
Fig. LIE,
4500
256
(Pr.5) (0.9) 0,03) hel
Cliping soom for leerutioss
up
dimu 1004, e needs 10 (19 3.1 (0.0.3.00) 3,06), boioatld etgping zorm 1. lineer worm up
4100 Transformert++
StimPaamm--7778
Tah
Pig 4,134
13500
256
dim==2448 hracs 16, comeest leagth
IB Tradormmne++
SlinPajano-62718
Tak
Fig.
13500
256
(0.9, 0,90), biloat16 oliging norm 1. lince worm-rp fir 10F Heneiims
Table 7: The best validation score (with one standard deviation across 3 runs; accuracy for computer vision tasks; perplexity for NLP tasks) for the best learning rate choice for each method that supports momentum.
Model
Dataset
NGN
SGDM
NGN-M
MomSPS
Momo
ALR-SMAG
Resnet20
CIFARIO
88.30
85.4220.70
88, 88.7620.05
87.2 87.20+0.39
88. 88.8620.14
88.88+0.19
Resnetl10
CIFARI00
64. 24.7640.2
57. 57.1642.00
24.98_0.22
63.37
64.81
64. 64.7341.81
VGG16
CIFARIO
90.2140.1
89.6 89.67
00.4220.06
87 87.2640.21
90.43g0.17
00.49
MLP
MNIST
98.0440.07
97.65 97.6322011
97, 20.97
97. 73..03
97.97
97 74.0440.00
VIT
CIFARIO
83.344024
83.7420.11
84.95=0.29
83.7740.27
85.47
85.5440.30
Resnet18
ImageNet32
48.63
48.56
48.29
N/A
48.68
N/A
Resnet18
ImageNetlk
67.00
66.73
67.12
N/A
67.09
N/A
Tiny
Transformer
Shakespeare
9.27+0.19
8. 8.73
7.67
N/A
8. 8.80
N/A
Rotten
Transformer
Tomatoes
9.01
8.75
7.12
N/A
8. 8.65
N/A
LSTM
Wikitext-2
75.33
82.07
75.51
75.51
20.22
N/A
76.09
N/A
hyperparameter. Finally, the performance results on (LSTM, Wikitext-2 [58)) and (Transformer, Rotten Tomatoes [68]) are reported in Table 8. The results demonstrate competitive performance of NGN-MDv1 against other benchmarks across all considered workloads.
1.4 Additional ImageNet Experiments
Now we turn to the experiments involving training R on ImageNer and ImageNet32. In Figure 1.4 we provide the train loss curves and results on (Resnet18, 8,, ImageNet32) workload that demonstrate that NGN-M and NDN-MDv1 attain better resilience to the step-size hyperparameter
52

---

Table 8: The best validation score (with standard deviation; accuracy for computer vision tasks; perplexity for NLP tasks) for the best learning rate choice for each method that supports diagonal step-sizes and momentum.
Model
Dataset
Adam
Momo Adam
NGN/MDe1
NGN-MDD2
Lion
Adabaliat
Adabound
CIFARIO
20.9
89.41
g9
87.80go
88.09g007
87.47
g5.00
Reanet110
CIFARI00
j4.12
67. 10.5a
66. 10
61.3310
61 11.851777
61.28
VGG16
CIFARIO
90.28 10.20
90.95 1o.3s
10
90.07
MLP
MNIST
97. 17.44
97. 17.96
98. 10 1o
97.07
NIA
NIA
N/A
VIT
CIFARIO
p0.96
15.74
85 15.0510.10
86.561os1
86.89 La.t
85.06
10.32
Transformer
Rotien Tomatoes
6.80
20.0
6.90
6.83
N/A
N/A
N/A
Tiny Shakespeare
Transformer
6,80
6.80
6.82
NIA
NIA
N/A
LSTM
PTB
70.95
71.09
70.84
71.37
NIA
N/A
N/A
LSTM
Wikitext-2
81.49 31.49
82.2330.61
75.24
21.90
N/A
N/A
N/A
Resnet18
ImageNet32
48.11
48.09
48.06
47.55
NIA
NIA
N/A
Resnet18
ImageNetlk
67.17
67.06
67.15
67.32
N/A
N/A
N/A
ViT-Tay
ImageNet3k
71.05
71.22
71.345g099
N/A
NIA
NIA
N/A
Transformer+
70M
SlimPajama-627B
34.38
34.96
33 33.84
N/A
N/A
N/A
Transformer-+
160M
SlimPajuma 6270
24.03
24.29
23.322000
NIA
N/A
NIA
N/A
Transformer
410M
SlimPajama-6278 16.65 uo.r
1o.0
18
NA
Transformer-+
IB
SlimPajama-633 627B
13.09
N/A
13.11
N/A
NIA
NIA
N/A
choice than competitors not only from the train loss point of view well. The best performance of algorithms is provided in Table 7 and According to them, both NGN-M and NGN-M achieve competitive performance against considered benchmarks.
1.5
Additional Comparison against Lion, Adabelief, Adabound
This section compares algorithms from Section 5. Moreover, we include the comparison against Lion [9], Adabound [56], and Adabelief [109]. The results are presented in Table 8.
We observe that NGN-MDv1 and NGN-MDv2 both achieve competitive performance across various Deep Learning workloads. In Figure I.5, we observe that Lion, Adabound and Adabelief algorithms do not match always the performance of NGN-MDv1 and Adam: Adabelief has worse performance on (Resnet20, CIFAR10) workload; Adabound has worse performance on (Resnet20, CIFAR 10), (Resnetl 10, CIFAR 100), and (ViT, CIFAR 10) workloads; Lion has worse performance (Resnet 10, CIFAR 100) workload. Moreover, their resilience to the step-size hyperparameter choice is lower than that of NGN-MDv1. To summarize, NGN-M and NGN-MDv1 are the most robust algorithms the choice of step-size hyperparameter.
L6 Comparison of Adaptive Step-sizes of Adam, Momo-Adam, and NGN-MDv1
Next, conduct experiments compare the adaptive step-size of Adam, Momo-Adam, and NGN-MDv1. Note that ResNet20 model consists of base blocks, and each block has convolution layers. In Figure 1.6 [layer1.0 coo.col...oo.... 1, layer2.0.come layer3.0.con/ of ResNet20 that corresponds to the first convolu- tion layer within each base block. Similarly, in Figure 1.7 plot the average adaptive step-size of the layers j E (layer0.0.fn.tt to layer3.0.fn..r wtt...ffttttoooo. _kv, aayer5.0.fn.to.gee that corresponds to the attention layers of the first, fourth, and sixth base blocks.
plot the average adaptive step-size of the layers j
Since the adaptivity Adam is only the second-order momentum applied as a normalization, in our experiment we compare the following quantities
Tk (Dx)g)
7kk
(D))g) for Adam,
for Momo-Adam,
for NGN- MDv1
(68)
(D+)
where y is the step-size hyperparameter of Adam.
53

---

NGN
SGDM
Momo
NGN-M
10P
10a
101
Final Final Loss
Final Final Train Loss
10
10
10
Final
10"
10
10
10-11
10-11
10"2
10
10
10'-
10'
10-2
10
Stepsize
Stepsize
Stepsize
Resnet20 for CIFAR 10
Resnet1 for CIFAR 100
VIT for CIFAR 10
Final Trxt Arcuracy
8
10
Final Loss
10
10
10
10
10
10
10
10P
Stepsize
Stepsize
VGG16 for CIFAR 10
VGG16 for CIFAR 10
10"
Final ledt Accuracy
10
80
Final Train
40
10-2
20
10
10"
10
10
10
10
Stepsize
Stepsize
MLP for MNIST
MLP for MNIST
Figure L1: Stability performance of algorithms supporting momentum varying step-size hyper- arameter (e for NGN and NGN-M, as for Momo, and step-size for SGDM). We observe that NGN-M achieves the training loss close to the best possible for a wider range of the step-size hyperparameter.
Let us first describe the results for ResNet20 in Figure 1.6. We observe that NGN-MDv1 tends to set smaller effective step-size compared two other algorithms. This especially visible for the large step-size hyperparameter values where the adaptive step-size of NGN-MDv1 is by several orders in magnitude smaller than that Adam and Momo-Adam. In contrast, the coordinate-wis adaptive step-size of Momo-Adam is mostly follow that of Adam. Considering that the stability performance of NGN-MDv1 is much higher for this task, this happens mainly due to the fact that the adaptation conservative than that of Momo-Adam.
mechanism of NGN-MDv1 step-size is
Now switch to the results on ViT model in Figure 1.7. Here both Momo-Adam and NGN-MDv1 tend to utilize smaller effective coordinate step-size, by several orders in magnitude smaller than that of Adam. However, the adaptation mechanism of NGN-MDv1 still more conservative than that of Momo-Adam, especially for large siep hyperparameters. We also highlight that in this experiment the best performance of NGN-MDv1 is achieved with 10-3 When we vary the step-size hyperparameter the effective coordinate-wise step-size does not change dramatically, especially for layers.O. w.w.fin.to_gkk layer.
I.7 Extended Comparison of Momentum-based Algorithms NLP Tasks
We switch to comparison of NGN-M, Momo, NGN, and SGDM on NLP tasks. In particular, we consider the training of Transformer (based on NanoGPT) on the Tiny Shakespeare and Rotten Tomatoes datasets and LSTM on the Wikitext-2 dataset from Appendix I.3. We report the results in Figure while the best performance is shown in Table 7. First, note that all algorithms do not match the best performance of those that incorporate diagonal step-size and momentum (see Table 8). Such results are expected since the training NI.P models has significantly different coordinate- wise conditioning. Nonetheless, NGN-M algorithm achieves better resilience to the step-
54

---

Adam
Momo-Adam
NGN-MDv1
NGN-MDv2
10'
101
Final Train Loss
Final Train Loss
20
10a
101
10a
10
10-1
10--1
10-3 10-110
10-"10-1100
1000011000-1000
102
10
10-5
10-3
16
10
110-2
10
Stepsize
Stepsize
Stepsize
Resnet20 for CIFAR 10
Resnet 10 for CIFAR 100
VIT for CIFAR 10
Final Test Accuracy
10
Final Train
10h
Stepsize
Stepsize
VGG16 for CIFAR 10
VGG16 for CIFAR 10
Figure I.2: Stability performance of algorithms supporting momentum and diagonal step-size varying step-size hyperparameter (c for NGN-MDv1 and NGN-MDv2, ag for Momo-Ades and step-size for Adam). We observe that NGN-MDv1 achieves the training loss close to the best possible for wider range of the step-size hyperparameter.
Adam
Momo-Adam
NGN-MDv1
NGN-MDv2
Rinal Test Accuracy
100
Final Test Perplexity
10
Final Perplexity
10a 1o
80
10
60-
10>
40
10a
20
10
10
110
10
10P
10
10a
10
10
10a
10-5
10
10
10
10
10"
10
10
Stepsize
Stepsize
Stepsize
Transformer for Tiny Shakespeare
MLP for MNIST
LSTM for PTB
10a
10a
107
Final FinalTrahe Train Loss
30'
Final Train Loss
102
10P
10"
10-1-
Final
10-3
10
10
30
10a
10
10a
10
Stepsize
Stepsize
Stepsize
Transformer for Tiny Shakespeare
ML.P for MNIST
L.STM for PTB
Figure 1.3: Stability performance of algorithms supporting momentum and diagonal step-size varying step-size hyperparameter (c for NGN-MDv1 and Adam). We observe that NGN-MDv1 achieves the training loss close to the best possible for wider range of the step-size hyperparameter.
ag for Momo-Adam, and step-size for
size hyperparameter choice, especially in the training of Transformer models. Therefore, NGN-M across various model architectures and task domains.
1.8 Comparison of Algorithms with Diagonal Step-size
Now we compare algorithms with diagonal step-size such as NGN-D, Adagrad [18], and RMSprop [44]. Since NGN-D requires to find constants {cy]5_1 where d is the size of the model. Finding sufficiently good constants yj might be challenging task since d is large number. Therefore,
55

---

Final Test Accuracy
60+
NGN
NGN
SCDM
SGOM
40-
5G0
Rinal Train
Mome
Nomo
NGN
20
5COM
11
10
10
10
10
10
10
10"
10
30
10
10"
10'
10
Stepsize
Stepsize
Stepsize
Resnet]8 for ImageNer32
Resnet18 for ImageNet32
Resnet18 for ImageNetIk
50-
Final Test Accuracy
Adam
Adam Momo Adam
40-
Mome-Adam
Final Train Loss
Homo Adam
NGN--MDV1
NGN-MDV1
NGN-MDva
80-
NGN-MDV2
20-
6
10-
10
10
10
10
10"
10
10
10
10
Stepsize
Stepsize
Stepsize
Resnet18 for ImageNet32
Resnet18 for ImageNet32
Resnet18 for ImageNetIk
Figure 1.:: Stability performance of algorithms supporting momentum (first row), and momentum with diagonal step-size (second row) varying step-size hyperparameter (e for NGN, NGN-M. NGN- MDv1, and NGN-MD NGN-MDv2, og for Momo and Momo-Adam, and step-size for SGD, SGDM, and Adam).
NGN
Momo
Adam
NGN-MDv1
Lion
Adabound
SGDM
NGN-M
Momo-Adam
NGN-MDv2
Adabelief
Final Test Accuracy
Final Test Accuracy
Ai
Accuracy
100
80
60
60-
40
40
201
10
10
10
10
103
07
10
10-4
10
10"
10
100310
10
10
10 10
10"
10
10
Stepsize
Stepsize
Stepsize
Rital Train Loss
Final Train Loss
20*
10
10
11
10
10'
10"
10
10
10
30
110
10" Stepsize
'10
16
Stepsize
Stepsize
Resnet20 for CIFARIO
Resnetl for CIFARIO0
ViT for CIFARIO
Figure I.5 Stability performance of various optimizers for (Resnet20, CIFARIO), (Resnetl 10,
CIFAR100), (VIT, CIFAR10) workloads.
RMSprop preconditioner D them ej = cD())) We leave the exploration of how set constants Cj properly for future research.
For each method, we tune its learning rate hyperparameter over the powers of 10: (10-1,.. 102) and present the best performance averaged across 3 random seeds in Table 9. We observe that NGN-D performs similarly to RMSprop. NGN-D has slightly worse performance on (LSTM, PTB) dataset but significantly better (LSTM, Wikitext-2) workload. Besides, Adagrad always has the worst performance. Moreover, these algorithms do not have high resilience choice hyperparameter. Therefore, we omit their comparison from this perspective.
1.9 Effective Step-size of NGN-M, Momo, NGN-MDv1, and Momo-Adam
Next, we compare the effective step-size applied throughout the training with NGN-M. Momo, NGN- MDv1, and Momo-Adam in Figures 1.9 and 1.10. First, both NGN-M and Momo perform warm-up
56

---

0.0001
0.001
0.01
0.1
1.0
10.0
100.0
10
10"h
10
Effective Stepsize
Effective Stepsize
Effective Stepsize
10'
207
20"
20*
20*
101
201
201
10
20
10
15
20
10
15
20
Iterations, x10"
Iterations, X10
Iterations, 110
layer1.0 c.0.o
layerI.0 10.conv
layerI.0.com 0.coonv
Adam
Momo-Adam
NGN-MDv1
1010
10
10
Effective Stepsize
Effective Stepsize
Effective Stepsize
107
100
307
1044
101
10'
301
10-2
10-2
Iterations, x10
Iteratioon,
-0
Iterations, 1100"
A.ocr2.0.comvl
1ayer2.0.comv
layer2.0.co 1onvv
Adam
Aoomo-Adam
NGN-MDv1
Effective Stepsize
Efrective Stepsize
Effective Stepsize
107
100
107
204-
20*
30"
30*
10-3
10-2
10
10
Iterations,
Iterations,
Iterations, 110
layer3.0.co 1cer3.
layer3.0.com/
layer3.0.com 1ver3.0.conv
Adam
Momo-Adam
NGN-MDv1
Figure L6: The adaptive stepsize of Adam (first column), Momo-Adam (second column), and NGN-MDv1 (third column) algorithms in training ResNet20 model on CIFARIO dataset. We plot the average stepsize [De)a] - (for Adam), (De) (for Momo-Adam), and 72 (Da)us (for NGN-MDv1) for the
first convolution layer within each of 3 base blocks of ResNer20 architecture varying the step-size hyperparameter the algorithms (c for NGN-M and NGN, og for Momo, and learning rate parameter for Adam).
Table 9: The best validation (with standard deviation; accuracy for image classification; perplexity for language modeling) for the best learning rate choice for each method that supports diagonal step-sizes.
Model
Dataset
Adagrad
RMSprop
NGN-D
Resnet20
CIFARIO
85.90
86.71 $6.7140.64
86.9820.15
Transformer
Rotten Tomatoes
7.77
6.87
6.92+0.03
Transformer
Tiny Sheaksper
7.77=0.05
7.00
6.90
LSTM
PTB
99.24_2.13
69.00+o.17
71.5420.11
LSTM
Wikitext-2
113.19
79.4840.45
75.44
in the beginning: the effective step-size increases at the beginning of the training. Then we observe the main difference between the two algorithms above: effective step-size of Momo for sufficiently large step-size hyperparameter is not adaptive within some part of the training, it always hits the upper bound. Consequently, during that part of the training Momo reduces to SGDM. In contrast, the effective step-size of NGN-M is always adaptive: it gradually decreases after short warm-up. This trend is similar to the sfate-of-the-art learning rate schedulers used in practice. Similar observations can be made in comparison of NGN-MDv1 and Momo-Adam.
57

---

0.0001
0.000316
0.001
0.00316
0.01
0.0316
109
30
Effectivve Stepsize
Effective Stepsize
20"
20
10"
10P
10P
10a
107
Effective
20"
10"
10"
10
10-2
10-3
10-3
10
20
111
Iterations, x10"
Iterations
Iterations, x 10
layers.0.00 0.t.s....ffffffffff.. _ok
layers.0.0.fl wwo.rs.....fn.to.oo. gkv
layers.0.0k 0.0.fn to gkv
Adam
Momo-Adam
WGG--DDV
10"
10"
30"-
Effective Stepsize
10"
Effective Stepsize
30
Effective Stepsize
30*
20"
30*
10
10a
10*
20"
20"
10-1
10-1
10-2
10-2
10
10
20
10
20
10
15
20
Iterations, X10"
Iteeations, 1
Iterations, =
syverrs..0.fm.tooog
3.0ers...0.ff..t.
layors.3 3.0.fn to qkv NGN-MDv1
Adam
Momo-Ade Adamo-Adamm
10
10
10a
Effective Stepsize
20*
Effective Stepsize
30*
Efrective Stepsize
307
10a
10a
108
20
10
10a
103
10P
10"
30"
10-2
10-1
10
20
10
Iterations, X10
Iterations, X
Iterations, 20
layers.5.0.fa.to qkv Adam
layers.5 Uovers.5.0.fii. to qkv
1oyers.5.0.fmm f...ff.to__ggv NGN-MDv1
Momo-Adam
Figure L.7: The adaptive stepsize of Adam (first column). Momo-Adam (second column), and NGN-MDv1 (third column) algorithms in training ViT model on CIFARIO dataset. We plot the average stepsize (D)))g) attention layer within each of the first, fourth, and sixth base blocks of ViT architecture varying the step-size hyperparameter of the algorithms (c for NGN-M and NGN, ag for Momo, and learning rate parameter for Adam).
? (for Adam), (Da)g) (for Momo-Adamm and (Da) (for NGN-MDv1) for the
I.10 Effective Updates in Training Language Models
In this section, we demonstrate the magnitude of updates when training 160M language model with Adam and NGN-MDv1 and varying the step-size hyperparameter across different layers of the model: see the results in Figures I.11 to L.13. We demonstrate that NGN-MDv1 is a more conservative algorithm: the effective update is smaller than that of Adam due the adaptive nature of the step-size. This is especially evident when training 160M language model with a step-size hyperparameter 0.03: The updates of Adam become considerably larger than the update of NGN-MDv1. This property is a key factor behind the difference in training dynamics: NGN-MDv1 can stabilize at a significantly lower training loss.
L11 Training Dynamics in Training Language Models
Now we report the training dynamics in the training language across all tested sizes.
I.12 Ablation Study of Momentum Parameters
In this section, study the sensitivity of NGN-MDv1 and Adam to the choice of the learning rate and momentum hyperparameters, when training 70M language model on Web dataset [70]. To do that, fix 3 = 0.9 (or Bg 0.95) and make sweep over the learning rate hyperparameter and
58

---

NGN
NGN
Final Train LoSs
Final Train Loss
SGON
SGON
Nomo
NGN
SCOM
10
10
10
10
10
10
111 Stepsize
107
10
10
10
10
10'
10"
Stepsize
Stepsizee
10º
Final Test Perplexity
NGN
Final Test Perplexity
NGN
10":
102
Moma
NGN--M
NGN M
NGN-M
100
100
10a
10
10 Stepsize
10
10 Stepsize
10"
10
Stepsize
LSTM for Wikitext-2
Transformer for Tiny Shakespeare
Transformer for Rotten Tomatoes
Figure 18:: Stability performance of algorithms supporting momentum and diagonal step-size varying step-size hyperparameter (e for NGN-M and NGN, og for Momo, and step-size for SGDM). We observe that NGN-M achieves the training loss close to the best possible for wider range of the step-size hyperparameter.
0.0001
0.001
0.01
0.1
1.0
10.0
100.0
1000.0
10a
10
Effectiv Stspsse
Effective Stepsize
Stepsine
10
10
10
20
10
100
200
200
400
50K
10
100
250
300
400
500
Rerations, =100
Hert ations
Iterations,
Ieeraiinns
Momo for Resnet20 for CIFAR 10
NGN-M for Resnet20 for CIFAR 10
0.316
10.0
31.6
1.7
100
107
Effective Stepsiee
Effectve Stepcize
10*
Effective
20
100
.300
-000
000
500
10
100
233
3001
10
300
Rerations,
Iteraiion
Beeations,
Momo for ViT for CIFAR 10
NGN-M for VIT for CIFAR 10
Figure 1.9: The step-size of Momo and NGN-M during the training. We demonstrate the step-sizes Tk for Momo and Tyu for NGN-M varying step-size parameters ogg for Momo and e for NGN-M.
32 (or learning rate hyperparameter and (1). We report the final test perplexity averaged over 3 runs for each of hyperparameters.
We summarize our findings from Table 10, Table 11, Table 12, and Table 13 as follows:
Low Ir (3e-3): NGN-MDv1 and Adam show similar sensitivity changes in both 31 and 82.
Moderate Ir (le-2): NGN-MDv1 is noticeably more robust than Adam to extremes of B1, while both optimizers perform similarly across 32 (though Adam's performance degrades slightly at B2 = 0.999).
High Ir (3e-2): Both methods suffer when 81 is small (or 32 is large), but NGN-MDv1 recovers lower perplexity at larger 31 values (smaller (3 values), whereas Adam fails to reach comparable performance.
59

---

0.0001
0.001
0.01
0.1
1.0
10.0
100.0
1000.0
10
10a
Effectie Steppie
Siapiie
Effective Stepsze
Eflective Siepazee
10
10"
201
600
500
1900
1000
20
400
1300
1003
2000
herations, =10"
Ber ations
herations,
Momo-Adam for Resnet20 for CIFAR 10
NGN-MDv1 for Resnet20 for CIFAR 10
0.0001
0.0316
Effective Stepoto
Effective Steppize
Electtvv Stepsze
10
10
100
300
320
10
350
-30
500
Iterations XII
fer attions
Iterations, 107
Ber ationss
Momo-Adam for VYT for CIFAR 10
NGN-MDv1 for ViT for CIFAR 10
Figure 1.10: The step-size of Momo-Adam and NGN-MDv1 during the training. We demonstrate the step-sizes т. for Momo-Adea Aamm--Adam& and "/k for NGN-MDv1 varying step-size parameters og for Momo and for NGN-MDv1.
Training Loss
Embedding Weight
Layer Attention QRV Proj
Layer 1 MLP to Proj
Effetto Sp Neo
3.3
3.00
Toining
6:75
1830
2303
6206
2000
3200
1660
1906
2000
1038
1600
5936
Layer MLP Out
Attention
MLP
Layer MLP Out
Efterieve Uplate Norza
Eftertrve Updare Norza
0.5
0.4
1
Stop
Slap
Stup
Step 1M Hooad
Layer 12 Atbeation QKV Fro]
Layer 12 MLP to Proj
Layer 12 MLP Out Proj
Ellodive Upcate Noen
0.c
12
[itecton Uppar
0.4
Effestre
1820
6260
3200
5600
1.00
2000
2030
1600
5036
Figure 1.11: Magnitude of updates when training 160M language model with Adam and NGN-MDv1 and step-size hyperparameter 0.003.
To conclude, NGN-MDv1 demonstrates greater robustness to changes in momentum parameters at high Ir, and consistently attains lower perplexity than Adam, even when both methods' performance deteriorates (we refer to the cases when both algorithms cannot achieve perplexity around 50).
60

---

Optimizot
Adon
MON
umbedding Weight
Attention
MLP Proj
EEecito Updatee Nomm
120 Update Narm
Nom
Troning
1850
3000
4006
2000 3300
2006
1000
2000
2002
6000
5830
1000
2000
4036
S0ep
5lap
Step
Step
Layer Dut Proj
Layar Proj
Layer Proj
Layer
Out Proj
Ellcure Update Nont
Updatt Nomm
Eilctire Updatt Noms
Effction Updatt
Layer 12 Atbention OKV Proj
12 M.P Proj
Layer 12 M.P Out Proj
1 Head
Eifeetto tpttes Norm
Effrctoo Updarr Nava
22.5
18.0
800
0.0
6.0
1850
2303
3000
4200
2000
2000
5006
1.000
2000
3002
6000
5000
1000
2007
2000
4030
Step
5tap
Step
Step
Figure I.12: Magnitude of updates when training 160M language model with Adam and NGN-MDv1 and step-size hyperparameter 0.01.
Table 10: Test perplexity of NGN-MDv1 when varying the learning rate and 31 hyperparame- ters when training 70M language model on the Fine dataset.
Table 11: Test perplexity of Adam when varying the learning rate and B1 hyperparameters when training 70M language model on the Feb dataset.
A
0.6
81 = 0.8
81 = 0.9
A = 0.99
Ir
8
= 0.6
8
0.8
3 0.9
81 0.99
3e-4
49.9 + 0.2
47.4 0.2
47.0 + 0.2
49.7 0.3
3e-4
49.4 + 0.2
46.5 0.1
46.2 0.3
57
le-3
41.5
39.9 0.2
38.6 0.1
40.2 0.3
le-3
41.4 0.2
39.6 6 0.1
38.5 28.5::0 0.1
45.0 # 0.2
3e-3
40
36.9 :: 0.3
35.9 + 0.1
37.2 + 0.1
3e-3
40.7 2 0.1
37.0 4: 0.1
36.0 ::: 0.1
220 :: 70
1e-2
54 +: 16
37 2:2
34.7 + 0.3
35.9 0.1
le-2
160 +: 60
4:2
36 4
210 + 110
3e-2
278 + 6
129 +:2
34.6 + 0.1
35.6 0.1
3e-2
420 +: 20
340 50
320 60
330 + 130
Table 12: Test perplexity of NGN-MDv1 when varying the learning rate and 3n hyperparame- ters when training 70M language model on the Web dataset.
Table 13: Test perplexity of Adam when varying the learning rate and B2 hyperparameters when training 70M language model on the Feb dataset.
1r 3:
0.6
8,
0.
- 0.95
& 0,509
82
0.6
3g 0.8
3: - 0.9
8: 0195
3g - 0,569
Re-4
51.8 + 0.6
4924 0.4
478=033 0.3
47.0+02
47.0 0.2
36-4
46.1
0.2
466+ 0.1
465 00 0.2
46.2+00.
46.5 + 0.1
1e:3
1..6103.
40.51..1
39.3 0.
38.6.10.1
38.9 0.1.
le3
38.8 0.1
39.0...2
38.9.0.1
38.5
39.5 1.0
3633
00.4002
37.5+ 0.2
95.3-0 0.1
3..9++0.1
36.5 + 0.4
36-3
38.8 0.3
36.3+
36.1
36.0++0.1
36.7 0.8
1o.2
3.2
0.19
35.1
31.7 0.3
36.0
35.4
3.10
31.9 0:3
2:2
0..8+00.
36.8+
3..6++0.1
250
550
210-4:00
500
61

---

Optimiont
Adun
MIN
Embedding Weight
Attention Pool
Elletto Update Namm
Ellectie Update Namm
7.5
Touning
to
3.5
S0mp
Step
Step Layer MLP In Pruj
Step
Laque MLP Out Pooj
Layer Atlantion GKV Pruj
Layar 7 MLP Out Proj
Effodioo Sislos Nara
Elleetive Updane Norz
Effeetoo Updaew Norm
Effedtoo Updare Nirn
1000
2000
6180
5500
1993
2000
3000
6000
5938
1983
2301
0000
1006
1000
2065
3000
2030
5000
Attention
MLL Froj
MLP Out
LM3 Head
Ellicttv Upaars Nara
Elficto Uppde
Eltvcton Uppoe
Step
Step
Strpp
Step
Figure 1.1: Magnitude of updates when training 160M language model with Adam and NGN-MDv1 and step-size hyperparameter 0.03.
Optimizur NEN-MDel
Optimteer
Adam NGN MDD1
Teaining Lows
Training Loss
0.003
0,01
0.03
0.0
0.2
0.4
0.6
0.8
1.0
1.2
0.0
0.5
1.0
1.5
2.0
2.5
Tokans
Tokens
70M Transformer
160M Transformer+
Optimirre
Opttmiteer
Adam
MGN-MDD
5.5
NGN MDv1
Trainiog Loss
LR
0.01
Training
0.03
4.0
0.00G
2.5
0.0
1.0
1.5
2.0
Toleens
Tokeens
1w10
410M Transformer
1B Transformer
Figure I.14: Training dynamics when training language model at different sizes.
62


Reviews:
[{"review_id": "dOfW4Pvwne", "reviewer": "Reviewer_ST2q", "summary": "Summary: \n\nThis paper introduces {NGN-M}, a momentum-enhanced variant of the {Non-negative Gauss-Newton step-size} (NGN) method by Orvieto & Xiao. It also proposes {NGN-MD}, which incorporates coordinate-wise diagonal step-size adaptation similar to the classic Adam. The main goal is to improve optimizer {stability} with minimal hyperparameter tuning, especially the learning rate. Key contributions include: (1) a theoretical convergence guarantee for NGN-M at the standard $\\mathcal{O}(1/\\sqrt{K})$ rate in convex stochastic optimization {without} requiring restrictive assumptions like interpolation or bounded gradients (which prior adaptive methods often needed); (2) extension to a diagonal-momentum version (NGN-MD) to address per-parameter conditioning in deep networks; and (3) experiments demonstrating that NGN-M/D maintain high performance while being significantly less sensitive to learning-rate choices than baselines (SGD with momentum, Adam, and recent adaptive methods).", "questions": "Questions: \n\nPlease see the weaknesses for question on analysis  — it seems we either require pre-knowledge of k or we need the c to go to 0 which means the method will align to NGN. Did I miss something? \n\nIs there an intuitive way to understand of why momentum improves stability? The authors have cited some prior references too that have made a similar claim, but I am finding it hard to intuitively understand this — if it is just because of variance reduction, surely there would be more direct approaches that achieve that goal ? \n\nThis is not super relevant to the paper’s acceptance but it would be interesting to know the authors’ perspective on this — Would it be interesting to explicitly tie in the step size behavior to the warmup kind of schedules? For example — see Small-scale proxies for large-scale Transformer training instabilities. Iclr 2024\n\nGiven the lack of theoretical arguments, have the authors given thought to if NGN-MDv1 is always stable, or can the coupling of heavy-ball momentum and an RMSprop preconditioner ever possibly cause divergence ? I understand such analyses are notoriously hard.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strenghts: \n\nThe theoretical contribution is one of the paper’s strongest aspects. The authors build on the NGN step-size proposed by Orvieto & Xiao, which was shown to never diverge for \\emph{any} learning-rate choice in convex problems and to automatically adapt to curvature. By incorporating heavy-ball momentum (Polyak’s momentum), the paper advances prior theory on combining adaptive step-sizes with momentum. Earlier attempts in this direction often required strong assumptions or achieved weaker guarantees. For example, the original {Stochastic Polyak Step-size} (SPS) method by Loizou et al guarantees convergence to the exact solution only under the interpolation assumption (informally, zero gradient noise at the optimum) or else converges to a neighborhood. Recent work by Oikonomou/Loizou introduced momentum into SPS but similarly could only ensure convergence within a non-vanishing error ball or required decaying step-sizes. In contrast, this paper’s Theorem 4.3 proves that {NGN-M converges at $\\mathcal{O}(1/\\sqrt{K})$ in convex stochastic settings without assuming interpolation, bounded gradients, or even bounded iterates}. This is a significant theoretical improvement over prior adaptive momentum methods. The proofs appear mathematically rigorous and leverage techniques from both momentum analysis and adaptive step-size methods. The derivation of the NGN-M update is grounded in a Gauss–Newton approximation to the squared-root loss (following Orvieto et al), but adding the momentum  term raises non-trivial analysis challenges. This paper handles this by using an “Iterative Moving Average” view of momentum and carefully coupling the momentum term with the adaptive step-size in their Lyapunov analysis. They highlight that unlike prior works, they do not require a bounded domain or bounded gradients assumption for stability. This is a notable strength: the method is provably stable for arbitrarily large chosen $c$ (learning-rate hyperparameter), just like the non-divergence property of the NGN. The proofs (mostly deferred to appendices) appear to be sound; they build on techniques from convex optimization and stochastic process bounds. \n\nThe paper compliments its theoretical contribution with extensive set of experiments covering both simpler settings to illustrate stability as well as large scale experiments to demonstrate effectiveness of the proposed technique. NGN-M converges reliably for {every} tested $c$, whereas SGDM diverges outside a narrow band of tuned values (even at moderate mis-tuning). This  supports the claim that NGN’s adaptive step-size “never diverges” and that momentum further helps by accelerating convergence for intermediate $c$ values without sacrificing stability. The evaluation is done vs several important baselines. And across many different settings, the proposed method-variants seem to outperform these baselines on validation accuracy. The experimental methodology is also sound. \n\nThe paper is also well-written with the contributions, theorem statements, contexts and existing works as well as assumptions and proofs clearly laid out. \n\nWeaknesses: \n\nFor analysis theorem 4.3 requires pre-knowing total number of iterations (so c can be set accordingly to obtain the bound to be the same as SGDM), or as done in the appendix, follow a diminishing schedule on c, where c goes to 0 as k increases — in this case, does the method not “eventually” boil down to NGN without the momentum or is close to it? \n\nThe runtimes can be possibly concerning — the major contribution of the paper seems to be removing the need for tuning the hyper-parameter, but if the range of stable SGDM is largely known from previous experiences on similar datasets, I am a bit worried that the increased runtimes can possibly offset this advantage. For the reported generalization performance, were the models trained for the same runtime? Or same number of epochs? If it is the latter, and the fact that NGN does “more work” per iteration, is it possible that running SGDM (within the acceptable stable set of parameters ) can lead to better performance? \n\nA brief ablation on $\\beta$ would strengthen the empirical section.\n\n\nMinor:\n“[2, 64, 2, 79, 53]” (line 80) contain a duplicated “[2]” entry", "was_revised": true, "final_justification": "I appreciate the authors providing more references for clarity on the \"fixed horizon\" sort of norms that are accepted in optimization. While I have worked in optimization for several years, I learnt something new from this paper and the discussions. I hope the authors will provide more clarity on these norms in the community so that an uninformed reader can place the results in context. i will raise my score.", "rebuttal": "We thank the reviewer for valuable comments\n\n***Answer to W1:*** The reviewer is right, we acknowledged this concern in the discussion after Theorem 4.3. Nonetheless, in Section F, we demonstrate that for simple functions like convex polynomials, ***the restriction on $\\beta$ being small can be avoided***. This is substantiated by experimental evidence showing that ***NGN-M and NGN-MD perform well***. This makes us believe that the restriction on $\\beta$ can be avoided but it requires a more careful analysis, which we were not able to obtain at the moment. Moreover, we highlight that other momentum-based algorithms, including ***MomSPS$_{\\max}$, have similar limitations***. We also note that we can remove the restriction on $\\beta$ by making the stepsize decreasing or under, additional assumptions like bounded iterates similar to [53]. While these modifications improve the theoretical analysis, we believe they would make the algorithm less favorable in practice.\n\n***Answer to W2:*** We directly addressed the concern regarding runtime in ***section G.2*** (also mentioned in the conclusion). Table 2 reports the time needed for performing one step of training and optimizer.step() when training language models at different scale. We observe that the time needed for one training step of NGN-MDv1 ***does not increase significantly and is comparable to that of AdamW*** (see middle column of Table 2). This happens because *** backpropagation is the most time-consuming part*** of training. Therefore, the final training time of AdamW and NGN-MDv1 is comparable: for all tested scales, it ***increases by a maximum of 3%***. Therefore, running the algorithms under a fixed time budget will not affect the results significantly, and the improved robustness of NGN-M/NGN-MDv1 to the choice of lr hyperparameter will remain.\n\n***Answer to W3:*** Following the reviewer's comment, we provide sweeps varying not only the lr hyperparameter but also the momentum parameters $\\beta_1,\\beta_2$ when training 70M language model on FineWeb dataset. To do that, we fix $\\beta_1=0.9$ (or $\\beta_2=0.95$) and make a sweep over lr and $\\beta_2$ (or lr and $\\beta_1$). We report the final test perplexity averaged over 3 runs for each set of hyperparameters. We thank the reviewer for this valuable suggestion, which has enabled us to strengthen our empirical evaluation. We will include these new results in the revised manuscript.\n\n|NGN-MDv1 $(\\beta_2=0.95$)|$\\beta_1=0.6$|$\\beta_1=0.8$|$\\beta_1=0.9$| $\\beta_1=0.99$|Adam ($\\beta_2=0.95$)|$\\beta_1=0.6$|$\\beta_1=0.8$|$\\beta_1=0.9$| $\\beta_1=0.99$|\n|-|-|-|-|-|-|-|-|-|-|  \n|lr 3e-4|$49.9\\pm0.2$|$47.4\\pm0.2$|$47.0\\pm0.2$|$49.7\\pm0.3$||$49.4\\pm0.2$|$46.5\\pm0.1$|$46.2\\pm0.3$|$57\\pm1$|\n|lr 1e-3|$41.5\\pm0.2$|$39.9\\pm0.2$|$38.6\\pm0.1$|$40.2\\pm0.3$||$41.4\\pm0.2$|$39.6\\pm0.1$|$38.5\\pm0.1$|$45.0\\pm0.2$|\n|lr 3e-3|$40\\pm1$|$36.9\\pm0.3$|$35.9\\pm0.1$|$37.2\\pm0.1$||$40.7\\pm0.1$|$37.0\\pm0.1$|$36.0\\pm0.1$|$220\\pm70$|\n|lr 1e-2|$54\\pm16$|$37\\pm2$|$34.7\\pm0.3$|$35.9\\pm0.1$||$160\\pm60$|$41\\pm2$|$36\\pm2$|$210\\pm110$|\n|lr 3e-2|$278\\pm6$|$129\\pm2$|$34.6\\pm0.1$|$35.6\\pm0.1$||$420\\pm20$|$340\\pm50$|$320\\pm60$|$330\\pm130$|\n\n|NGN-MDv1 ($\\beta_1=0.9$)|$\\beta_2=0.6$|$\\beta_2=0.8$|$\\beta_2=0.9$| $\\beta_2=0.95$|$\\beta_2=0.999$|Adam ($\\beta_1=0.9$)|$\\beta_2=0.6$|$\\beta_2=0.8$|$\\beta_2=0.9$| $\\beta_2=0.95$|$\\beta_2=0.999$|\n|-|-|-|-|-|-|-|-|-|-|-|-|  \n|lr 3e-4|$51.8\\pm0.6$|$49.2\\pm0.4$|$47.8\\pm0.3$|$47.0\\pm0.2$|$47.0\\pm$0.2||$46.1\\pm0.2$|$46.6\\pm0.1$|$46.5\\pm0.2$|$46.2\\pm0.3$|$46.5\\pm0.1$|\n|lr 1e-3|$42.6\\pm0.3$|$40.5\\pm0.1$|$39.3\\pm0.2$|$38.6\\pm0.1$|$38.9\\pm0.1$||$38.8\\pm0.1$|$39.0\\pm0.2$|$38.9\\pm0.1$|$38.5\\pm0.1$|$39.5\\pm0.6$|\n|lr 3e-3|$39.4\\pm0.2$|$37.5\\pm0.2$|$36.3\\pm0.1$|$35.9\\pm0.1$|$36.5\\pm0.4$||$38.8\\pm0.3$|$36.3\\pm0.1$|$36.1\\pm0.2$|$36.0\\pm0.1$|$36.7\\pm0.8$|\n|lr 1e-2|$37.8\\pm0.2$|$35.9\\pm0.1$|$35.1\\pm0.3$|$34.7\\pm0.3$|$35.0\\pm0.3$||$35.4\\pm0.2$|$35.0\\pm0.1$|$34.9\\pm0.3$|$36\\pm2$|$41\\pm3$|\n|lr 3e-2|$37.8\\pm0.3$|$35.8\\pm0.1$|$34.9\\pm0.1$|$34.6\\pm0.1$|$250\\pm50$||$550\\pm250$|$120\\pm80$|$160\\pm5$|$210\\pm60$|$500\\pm20$|\n\nWe summarize our findings as follows:\n\n- Low lr ($\\le$ 3e-3): NGN-MDv1 and Adam show similar sensitivity to changes in both $\\beta_1$ and $\\beta_2$.\n\n- Moderate lr (1e-2): NGN-MDv1 is noticeably more robust than Adam to extremes of $\\beta_1$, while both optimizers perform similarly across $\\beta_2$ (though Adam’s performance degrades slightly at $\\beta_2 = 0.999$).\n\n- High lr (3e-2): Both methods suffer when $\\beta_1$ is small (or $\\beta_2$ is large), but NGN-MDv1 recovers lower perplexity at larger $\\beta_1$ values (smaller $\\beta_2$ values), whereas Adam fails to reach comparable performance.\n\nTo conclude, NGN-MDv1 demonstrates greater robustness to changes in momentum parameters at high lr, and consistently attains lower perplexity than Adam, even when both methods’ performance deteriorates (we refer to the cases when both algorithms cannot achieve perplexity $\\lesssim$ 50).\n\n***Answer to Q1:*** Demonstrating the improved stability from using momentum is a ***challenging task***, especially for an adaptive algorithm like NGN-M. However, to gain intuition, we can consider minimizing a convex quadratic function $\\min_{x}\\frac{1}{2}x^\\top Q x + b^\\top x + c.$ The ***maximum allowed stepsize of GD*** that leads to convergence is $\\boldsymbol{\\gamma}<\\mathbf{\\frac{2}{\\boldsymbol{\\lambda}\\_{\\max}(Q)}}$, where $\\lambda_{\\max}(Q)$ is the maximum eigenvalue of $Q$. When using Heavy-ball algorithm instead, the maximum allowed stepsize for convergence is $\\mathbf{\\frac{2(1+\\boldsymbol{\\beta})}{\\boldsymbol{\\lambda}\\_{\\max}(Q)}}$ [1], which is about ***two times larger than for GD*** for a standard choice of $\\beta=0.9$. We acknowledge that training neural networks is far from quadratic,s but it gives an intuition why we should expect an improved robustness to the choice of the lr hyperparameter. This is also supported by recent works that observe ***convex-like structures in the loss landscape*** of neural networks [2-4] and ***agreement between convex theory and practice*** [5].\n\n[1] Polyak, Some methods of speeding up the convergence of iteration methods. Ussr comp. math. and math. phys., 1964\n\n[2] Tran et al., Empirical tests of optimization assumptions in deep learning. arXiv preprint arXiv:2407.01825.\n\n[3] Islamov, et al. Loss landscape characterization of neural networks without over-parametrization. NeurIPS 2025\n\n[4] Guille-Escuret et al. No wrong turns: The simple geometry of neural networks optimization paths. arXiv preprint arXiv:2306.11922, 2023.\n\n[5] Schaipp et al. The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. ICML 2025\n\n***Answer to Q2:*** We believe there is a strong connection between our adaptive approach of setting the lr hyperparameter and lr scheduling, including warmup. According to Figure 5 in the main paper and Figures J.9-J.12 in the appendix, our proposed adaptive scheme ***serves as a safeguard*** during training. If the lr hyperparameter is too large, which might lead to unstable convergence/divergence, our scheme makes the effective lr smaller and stabilizes the training. This mechanism is similar in spirit to warm-up, but with two key differences. First, unlike traditional warm-up, which is typically applied only at the beginning of training, our approach can activate dynamically at any point during training. Second, while warm-up schedules require manual specification of the warm-up duration, our method automatically adapts based on the local curvature. We thank the reviewer for the reference and will incorporate this discussion into the revised version.\n\n***Answer to Q3:*** We expect NGN-MDv1 to retain the robustness properties of NGN-M in most cases. However, its robustness also depends on the choice of preconditioning. If the chosen preconditioner is highly sensitive to variations in the loss landscape or hyperparameters, NGN-MDv1's resilience to learning rate selection may degrade. In such cases, alternative preconditioning strategies, such as those based on Shampoo [6] or the Fisher information matrix [7], might be worth exploring to improve stability.\n\n[6] Gupta et al., Shampoo: Preconditioned stochastic tensor optimization, ICML 2018\n\n[7] Gomes et al., Adafisher: Adaptive second order optimization via fisher information. arXiv preprint arXiv:2405.16397, 2024."}, {"review_id": "QHAYpeXHHG", "reviewer": "Reviewer_3pMo", "summary": "It extends the NGN step-size to consider momentum. Besides proposing the momentum-based version, NGN-M, it proves the standard convergence under less restrictive assumptions. Additionally, it conducts experiments to show the enhanced robustness to the choice of the step-size hyperparameter while delivering performance that is comparable to or surpasses other state-of-the-art optimizers.", "questions": "1. Can you conduct experiments on larger models and datasets? e.g.， pretrained MoE LLMs\n2. Can you add LR scheduler and weight decay in experiments?\n3. Can you give the variances of different step sizes for all compared methods?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1. It provides a theoretical analysis and proof for the explanation.\n2. It conducts comprehensive experiments on different tasks and models.\n3. Good introduction on the background and related work to show the position in the literature.\n\nWeakness:\n1. The proposed algorithm is the extension of NGN, so the solution itself seems to be incremental. \n2. The robustness is not well quantified by measuring some metrics like performance variance.\n3. Since learning rate schedulers and weight decay are widely adopted in practice, it is a pity that the experiment settings ignore such practical factors. Therefore, it leaves room for the claim to be questioned.", "was_revised": true, "final_justification": "Since my main concerns have been addressed, I plan to maintain my positive score.", "rebuttal": "***Answer to W1:*** We agree that NGN-M builds upon the NGN method, but this extension is far from trivial. In Section 3.1, we compare two variants of coupling momentum and the NGN stepsize. We demonstrate that ***the second variant*** (Adam-type momentum) ***does not preserve the non-divergence feature of NGN*** in the convex case. Moreover, in the convex setting, the coupling should lead to acceleration if $c$ is small enough. Hence, the integration of momentum into NGN is not straightforward and requires a careful analysis, which we provide in the paper.\n\n***Answer to W2:*** We thank the reviewer for this insightful comment, as it allowed us to rethink how to better demonstrate the robustness of the proposed method, and eventually improve our work further. We can quantify the robustness in a similar way to how it was done in [1]. We define $\\ell\\_{\\gamma}$ to be the final metric (e.g., train loss or negative test accuracy that we need to minimize) achieved when training a model with the lr hyperparameter $\\gamma\\in[a,b]$, while $\\ell(x_0)$ is the metric at initialization. Moreover, we define the best possible metric $\\ell^\\*\\_{\\gamma_*}$ that can be achieved by training a model with $\\gamma\\in[a,b],$ i.e., $\\ell^\\*\\_{\\gamma_*} = \\min\\_{\\gamma\\in[a,b]}\\ell\\_{\\gamma}.$ The lr sensitivity is then defined as $\\mathbb{E}\\_{\\gamma\\in[a,b]}[\\min\\\\{\\ell\\_{\\gamma}, \\ell(x_0)\\\\} - \\ell^\\*\\_{\\gamma_*}]$. We measure this quantity for various tested algorithms and tasks, approximating the expectation by averaging $\\min\\\\{\\ell\\_{\\gamma}, \\ell(x_0)\\\\} - \\ell^\\*\\_{\\gamma_*}$ over the grid of values for $\\gamma$. This metric ***captures the maximum deviation from optimal performance***. Across all experiments, both NGN‑M and NGN‑MDv1 show consistently lower sensitivity to the lr hyperparameter than the other baselines, quantitatively confirming our claims.\n\n\n\n| Resnet20, CIFAR10 | SGDM | NGN | Momo | NGN-M | Adam | Momo-Adam | NGN-MDv1 | NGN-MDv2 |\n| -------- | -------- | -------- | -------- | -------- |-------- |-------- | -------- |-------- |\n| Test accuracy LR sensitivity     | 0.22     | 0.10     | 0.13 | ***0.03*** | 0.18 | 0.08 | ***0.04*** | 0.15|\n\nHere, $\\gamma \\in \\{10^{-2}, 3\\cdot 10^{-2}, \\dots, 3\\cdot 10^0, 10, 10^1, 10^2\\}$ for the SGDM, NGN, Momo, NGN-M and $\\gamma \\in \\{10^{-4}, \\dots, 10^0\\}$ for Adam, Momo-Adam, and NGN-MD.\n\n\n| Resnet110, CIFAR100 | SGDM | NGN | Momo | NGN-M | Adam | Momo-Adam | NGN-MDv1 | NGN-MDv2 |\n| -------- | -------- | -------- | -------- | -------- |-------- |-------- | -------- |-------- |\n| Test accuracy LR sensitivity     | 0.25 | 0.17 | 0.16 | ***0.05*** | 0.19 | 0.18 | ***0.10*** | 0.18 |\n\nHere, $\\gamma \\in \\{10^{-2}, \\dots, 10^2\\}$ for the SGDM, NGN, Momo, NGN-M and $\\gamma \\in \\{10^{-4}, \\dots, 10^0\\}$ for Adam, Momo-Adam, and NGN-MD.\n\n| ViT, CIFAR10 | SGDM | NGN | Momo | NGN-M | Adam | Momo-Adam | NGN-MDv1 | NGN-MDv2 |\n| -------- | -------- | -------- | -------- | -------- |-------- |-------- | -------- |-------- |\n| Test accuracy LR sensitivity     | 0.28 | 0.19 | 0.24 | ***0.01*** | 0.41 | 0.21 | ***0.06*** | 0.33 |\n\nHere, $\\gamma \\in \\{10^{-1}, 3\\cdot 10^{-1}, 10^0, 3\\cdot 10^0, 10^1, 3\\cdot 10^1\\}$ for the SGDM, NGN, Momo, NGN-M and $\\gamma \\in \\{10^{-4}, 3\\cdot 10^{-4}, \\dots, 10^{-2}, 3\\cdot 10^{-2}\\}$ for Adam, Momo-Adam, and NGN-MD.\n\n\n| Test perplexity LR sensitivity | 70M | 160M | 410M | 1B |\n| -------- | -------- | -------- | -------- | -------- |\n| Adam | 78.6 | 85.2 | 78.6 | 96.3 |\n| NGN-MDv1 | 25.4 | 30.9 | 25.4 | 15.5 |\n\nHere, for 70M, 160M, 410M models, $\\gamma \\in \\{10^{-3}, 1.7\\cdot 10^{-3}, 3\\cdot10^{-3}, 5.5\\cdot 10^{-3}, 10^{-2}, 3\\cdot10^{-2}\\}$. For 1B model, $\\gamma \\in \\{10^{-3}, 3\\cdot10^{-3}, 5.5\\cdot 10^{-3}, 10^{-2}, 3\\cdot10^{-2}\\}$.\n\n[1] Wortsman, Mitchell, et al. \"Small-scale proxies for large-scale transformer training instabilities.\" arXiv preprint arXiv:2309.14322 (2023).\n\n***Answer to W3 and Q2:*** We respectfully disagree with this claim. Learning rate scheduling and weight decay are not required to achieve strong performance in our core experiments, as shown in Figure 2. When training Resnet18 on Imagenet1k, we use learning rate decay every 30 epochs by 0.1 without weight decay, while when training ViT on Imagenet1k, we use cosine annealing with warm-up and weight decay. For language modeling, we use linear warm-up followed by cosine annealing. To summarize, for each training task we ***employ standard training procedures*** that are summarized in Section J.1 and Table 3. Moreover, we provide a discussion on how weight decay can be added to NGN-MD in Section H and provide additional ablations by varying the weight decay parameter $\\lambda$ in Figure H.1.\n\n***Answer to Q1:*** The empirical part of the work involves testing the proposed algorithms on ***10 datasets***, including both image and text source of data; ***6 types of architectures*** (MLP, VGG, Resnet, LSTM, ViT, Transformer) with the largest ***1B model***. In all the cases, the proposed algorithms achieve competitive performance with an increased robustness to the lr hyperparameter, which is the main focus of the work. In our view, this ***comprehensive set of experiments sufficiently supports the validity of our claims***. While we agree that scaling to even larger models and datasets would strengthen the empirical evidence, such experiments are not feasible within the rebuttal period (a single run of a 1B model with our compute resources takes 11 hours on 8xH100-80Gb). Demanding larger-scale experiments would exceed our available compute resources, and, more broadly, raises concerns about fairness since such a requirement tends to favor institutions with disproportionate access to large-scale GPU infrastructure, rather than evaluating contributions on the basis of methodological innovation.\n\n***Answer to Q3:*** We thank the reviewer for this comment. Note that Figures 5 and J.9 do show the standard deviation. However, since the stepsize changes rapidly, this cannot be observed in Figure 5. The reviewer still can observe the variance in the 2nd and 4th columns of Figure J.9. The standard deviation is relatively small. We will update the plots by using exponential smoothing so that the variance can be better observed in all figures."}, {"review_id": "p5W9E7v9qo", "reviewer": "Reviewer_HqrG", "summary": "The paper addresses the sensitivity of momentum-based optimizers to hyperparameter variance to improve the stability of neural network training. Specifically, the authors adapt the Non-negative Gauss-Newton (NGN) stepsize method to the optimizers incorporating momentum. Empirical experiments show the proposed optimizers, NGN-M and NGN-MD, achieve more stable performance when facing parameter variance, and comparable performance to traditional optimizers like Adam.", "questions": "Apart from the weaknesses noted above, I would appreciate if the authors could answer the following questions:\n\n1. When comparing stability across optimizers (NGN-M, NGN-MDv1, NGN-MDv2, Momo, momo-Adam, SGDM and Adam), is it fair to adjust $c$ and $\\alpha_0$ for the comparison?\n\n2. The discussion of NGN-MD is relatively shorter than that of NGN-M. Are there any drawbacks to using NGN-MD, such as increased training time or other trade-offs?\n\n3. In Related Work, the authors focus on the Polyak step-size framework, but there are other methods aiming to improve training stability (e.g., [1–3]). Could you briefly compare those approaches with the PS-based method and highlight the advantages of adopting the Polyak framework?\n\n[1] Zhang, H. R., Li, D., & Ju, H. (2023). Noise stability optimization for finding flat minima: A hessian-based regularization approach. arXiv preprint arXiv:2306.08553.\n\n[2] Gomes, D. M., Zhang, Y., Belilovsky, E., Wolf, G., & Hosseini, M. S. (2024). AdaFisher: Adaptive Second Order Optimization via Fisher Information. arXiv preprint arXiv:2405.16397.\n\n[3] Li, X., Luo, J., Zheng, Z., Wang, H., Luo, L., Wen, L., ... & Xu, S. (2024). On the Performance Analysis of Momentum Method: A Frequency Domain Perspective. arXiv preprint arXiv:2411.19671.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**  \n- The authors naturally extends the NGN to momentum-based optimizers to address the hyperparameter sensitivity issue, where the motivation is well justified.\n- The authors conduct extensive experiments and theory analysis to back up and prove the effectiveness of their method.\n\n**Weaknesses**  \n-  **Captions, notation, and equations lack clarity**  \n   - The caption of Table 1 defines $x^*$ but never uses it.  \n   - $S_k$ appears in the NGN-M equations without explanation (though can be inferred from later texts).  \n   - Figure 1 labels subfigures as “left”, “second left”, “third and fourth”, yet line 123 refers to “left” when it seems to mean “second left.” Also, NGN-M isn’t identified as Version 1 until the paragraph’s final sentence, causing confusion.  \n   - Equations such as $f^* = \\arg\\min_x f(x)\\in\\mathbb{R}$ are nonstandard.  \n\n   I would encourage the authors to review all notations and equations for consistency and clarity or it would be difficult to follow.\n\n-  **Lack of Reinforcement Learning(RL) experiments**  \n   RL training is well-known for its high sensitivity to hyperparameters, so it would be valuable to know whether NGN-M/NGN-MD deliver similar stability gains in that domain, which is of great importance. Considering the limited time, even a brief comparison on one or two RL tasks would make the evaluation more convincing.", "was_revised": true, "final_justification": "My concerns raised during review were (1) clarity issues, (2) missing experiments, (3) comparison fairness, (4) broader related work needed. In the rebuttal, the authors address (1), (3), (4). While I still think that RL would be helpful for a more comprehensive evaluation of their method. However, after the rebuttal, I think it is an optional choice considering the response given by the authors and the large experiments the authors have conducted. My opinion is that the strengths of the paper outweigh its weaknesses.", "rebuttal": "We thank the reviewer for valuable comments. \n\n***Answer to W1:*** We thank the reviewer for a careful read of our paper. We will incorporate all proposed changes in the revised version of the work. The notation for $f^* \\in \\mathbb{R}$ means that the objective is bounded from below, i.e., $f^* > -\\infty$.\n\n***Answer to W2:*** Thank you for the suggestion to apply our results to reinforcement learning (RL). However, we would like to emphasize that our current experimental focus on the standard classification setting is deliberate: sensitivity to hyperparameters is already a well-documented and significant challenge in this domain. By demonstrating the robustness of our method in such a widely studied and sensitive setting, we believe our work already makes a significant contribution. But we agree that RL is an exciting future direction due to its high sensitivity, and we will mention this in the paper and consider it as follow-up work, thank you again for pointing this out.\n\n***Answer to Q1:*** In our view, it is a fair comparison. Momo and NGN-M (Momo-Adam and NGN-MDv1) are based on the SPS (Stochastic Polyak Stepsize) and the NGN stepsizes, respectively, and they have hyperparameters $\\alpha_0$ and $c$. These hyperparameters play a similar role: ***they adjust the effective lr hyperparameter*** during training, not allowing it to be larger than $\\alpha_0$ or $c$, respectively. We therefore believe it is fair to adjust these parameters and would be happy to discuss this further if the reviewer thinks otherwise.\n\n***Answer to Q2:*** NGN-MD is an extension of the NGN-M algorithm. Therefore, all the discussion on the lr hyperparameter sensitivity of NGN-M applies to NGN-MD as well. Moreover, the derivation of the NGN-M update rule can be seen as a special case of NGN-MDv1 with $\\mathbf{\\Sigma}_k=\\mathbf{I}.$ From an implementation point of view, using a general diagonal preconditioner $\\mathbf{\\Sigma}_k$ leads to per-parameter adaptive lr. In the case of RMSprop preconditioning, performing one step of training increases since we need to compute the second-order momentum $v^k$. To link, the difference in training time of NGN-M and NGN-MDv1 is similar to that between SGD and Adam. However, the robustness of NGN-MDv1 also depends on the preconditioning. In the cases when the preconditioning is sensitive to the changes in the loss landscape or hyperparameters, NGN-MDv1 might be less robust to the choice of the lr hyperparameter than NGN-M.\n\n***Answer to Q3:*** [1] This paper studies how to set momentum parameters $\\beta_1$ and $\\beta_2$ based on frequency domain analysis. However, we did not read any claim about the robustness of their approach to changes of the hyperparameters. Based on Figure 4 in [1], the ***hyperparameters should be chosen in a specific way***, which may limit their robustness in practice. Therefore, ***we ask the reviewer to clarify which part of the work is related to ours***. Nonetheless, their idea on changing the momentum parameter is orthogonal to ours, which makes the comparison difficult. But their momentum scheduling can be combined with our NGN-M algorithm, which would potentially lead to improved results. \n\n[2] Thank you for the reference, and we will cite this work in the revision. [2] introduces AdaFisher, a ***novel way to choose diagonal preconditioning*** by leveraging the composite structure of neural networks. Specifically, the authors derive a diagonal ***approximation to the Fisher Information Matrix*** (FIM) and use it to rescale each parameter update. This curvature‐aware scaling yields more accurate local stepsizes and, in their experiments, substantially reduces sensitivity to the lr choice when training ResNet‑50 on CIFAR‑100. It therefore remains an open question whether the same stability gains carry over to other architectures and datasets. Therefore, the robustness of their approach is underexplored. However, their new ***AdaFisher preconditioning can be incorporated into NGN-MDv1***, which potentially leads to ***better curvature approximation of AdaFisher and improved stability of NGN-MDv1*** to the choice of the lr hyperparameter. \n\n[3] The approach is based on computing the update direction as $\\frac{1}{2}(\\nabla f(x+u) + \\nabla f(x-u)),$ where $u$ is sampled from some distribution. Such an approach leads to ***convergence to flatter minima*** where the trace of the Hessian is minimized, i.e., they study another. This is an alternative to the SAM algorithm, which minimizes the maximum eigenvalue of the Hessian. However, this approach is again orthogonal to ours. While the method potentially might lead to flatter minima, the efficiency of the algorithm is tested; it is still sensitive to the choice of the lr hyperparameter: see Figure 4 in [3], which is the main focus of our work. However, the main disadvantage of this approach is the necessity of computing at least two backpropagations per step, which increases both time and memory consumption during the training when the model is sufficiently large. Thus, ***we ask the reviewer to clarify which part of the work is related to ours*** since this work seems not to study robustness to the choice of hyperparameters."}, {"review_id": "D1lddzKkeT", "reviewer": "Reviewer_R39v", "summary": "This paper introduces new optimization algorithms, NGN-M and NGN-MD, which enhance the recently proposed NGN (Non-negative Gauss-Newton) step-size method by incorporating momentum and coordinate-wise (diagonal) updates. The primary goal is to improve the stability and robustness of the optimizer with respect to the step-size hyperparameter, a common issue with adaptive optimizers like Adam, especially in training large models like Transformers.\n\nThe paper introduces NGN-M: A momentum-based version of the NGN step-size. The authors provide a theoretical analysis showing that NGN-M achieves a convergence rate of O(1/ K) in the convex setting without requiring restrictive assumptions like bounded gradients or the interpolation condition.\n\nThe paper also introduces NGN-MD: A variant that integrates a coordinate-wise diagonal preconditioning strategy with momentum, making it suitable for training large-scale models like Transformers. Two versions, NGN-MDv1 and NGN-MDv2, are proposed.\n\nEmpirical results demonstrate that NGN-M and NGN-MD are significantly more robust to the choice of the step-size hyperparameter compared to state-of-the-art optimizers like Adam, SGDM, and Momo. They achieve comparable or better performance while maintaining stability over a much wider range of learning rates.", "questions": "1.  Many recent studies have used the affine gradient noise assumption. The paper claims existing literature relying on bounded gradient is not true.\n\n2. Minor typos.\n\n2.1 Page 6, Line 190: \"MomSPS and ALR-SMAG were shown to converge up to a non-vanishing neighborhood of the solution only  \" - The word \"only\" seems misplaced and the footnote reference is superscripted. Suggest rephrasing for clarity, e.g., \"...were shown to converge only to a non-vanishing neighborhood...\".", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strength**\n\n1. Well-Motivated Algorithms: The paper extends the NGN step-size to incorporate momentum and diagonal adaptive learning, which is important for NGN able to apply to state-of-the-art deep learning models.  The proposed combination of the NGN step-size with momentum and diagonal preconditioning is a logical and well-justified extension of prior work.\n\n\n2. Comprehensive Empirical Evaluation: The experiments are extensive, covering a wide range of models (CNNs, ViTs, LLMs), datasets (from small to very large-scale), and comparisons against relevant popular optimizers. This provides evidence for the practical utility of the proposed methods. The stability plots (Figures 2, 3, and 4) are particularly effective at illustrating the main claim of the paper—that NGN-M and NGN-MD are significantly more robust to the step-size hyperparameter than their counterparts.\n\n**Cons**\n\nTheoretically, the analysis does not cover nonconvex objectives, which has been widely studied in recent convergence analysis of Adam. Moreover, their assumption on gradient noise is affine, rather than bounded.\n\nIt is of benefit to study the hyperparameter sensitivity under a clean setup: no learning rate scheduling, no weight decay and no extremely large batch. However, the paper lacks at least a strong baseline to demonstrate the efficacy of the proposed method in training nowadays LLMs. For example, the Llama 2, with standard weight decay, 4M batch size and standard learning rate scheduling, the default setting in available torchtitan codebase. This is a significant gap, as NGN-MD is positioned as the more practical algorithm for large models.", "was_revised": true, "final_justification": "The authors' response resolved my comments on the lack of non-convex analysis and clarified the related hyperparameter choices.\n\nHowever, the paper lacks at least a strong baseline to demonstrate the efficacy of the proposed method in training nowadays LLMs. This point is quite important to me given that this is a paper positioning as a fundamental optimizer design, which should demonstrate its efficacy in most challenging and real optimization tasks.", "rebuttal": "We thank the reviewer for valuable comments.\n\n***Answer to W1:*** In Section C, ***we provide a nonconvex convergence analysis for NGN‑D***, which may have been overlooked by the reviewer—we kindly ask them to revisit this section. Moreover, that analysis readily extends to the ***more general ABC noise model*** [1], which covers the affine noise case mentioned by the reviewer as a special instance. We will include this broader analysis in the revised manuscript.\n\nWe also agree that developing a nonconvex theory for NGN‑M would strengthen the analysis, and we mention this explicitly in the conclusion as an important direction for future study. While convexity is typically considered a strong assumption, especially for neural networks training, a growing body of empirical evidence shows that ***convex convergence bounds often predict the behavior of large‑scale language models quite accurately*** [1], and that neural loss landscapes frequently ***exhibit convex‑like regions around the minimizer*** [3-5]. Finally, our theoretical contributions significantly advance prior Polyak‐stepsize adaptive methods by ***relaxing assumptions*** and ***tightening convergence guarantees***, and we validate these findings through a comprehensive empirical evaluation.\n\n[1] Khaled and Richtarik. Better theory for SGD in the nonconvex world. TMLR, 2023.\n\n[2] Schaipp et al., The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. ICML 2025\n\n[3] Tran et al., Empirical tests of optimization assumptions in deep learning. arXiv preprint arXiv: 2407.01825.\n\n[4] Islamov, et al. Loss landscape characterization of neural networks without over-parametrization. NeurIPS 2024\n\n[5] Guille-Escuret et al. No wrong turns: The simple geometry of neural networks optimization paths. arXiv preprint arXiv: 2306.11922, 2023.\n\n***Answer to W2:*** We highlight that the results demonstrated in Figure 2 ***do not use any lr scheduling and weight decay***, as the reviewer suggests, since they are not required to achieve strong performance in those settings. We demonstrate an improved stability of NGN-M and NGN-MDv1 to the lr hyperparameter. Moreover, we do not claim that the proposed adaptive scheme of setting lr hyperparameter replaces other techniques such as warm-up or weight decay. In contrast, ***all of them can be combined together for a further improvement in robustness*** of the algorithms since warm-up and weight decay are known to improve sensitivity of algorithms to the choice of the lr hyperparameter [3]. Therefore, for training tasks where warm-up and weight decay are necessary, we ***employ standard training procedures*** that are summarized in Section J.1 and Table 3. \n\nThe empirical part of the work involves testing the proposed algorithms on ***10 datasets***, including both image and text source of data; ***6 types of architectures*** (MLP, VGG, Resnet, LSTM, ViT, Transformer) with the largest ***1B model***. In all the cases, the proposed algorithms achieve competitive performance with an increased robustness to the lr hyperparameter, which is the main focus of the work. In our view, this ***comprehensive set of experiments sufficiently supports the validity of our claims***. While we agree that scaling to even larger models and datasets would strengthen the empirical evidence, such experiments are not feasible within the rebuttal period (a single run of the 1B model with our compute resources takes 11 hours on 8xH100-80Gb). Demanding larger-scale experiments would exceed our available compute resources and, more broadly, raises concerns about fairness since such a requirement tends to favor institutions with disproportionate access to large-scale GPU infrastructure, rather than evaluating contributions on the basis of methodological innovation.\n\n[3] Wortsman, Mitchell, et al. \"Small-scale proxies for large-scale transformer training instabilities.\" arXiv preprint arXiv:2309.14322 (2023).\n\n***Answer to Q1:*** We would like to clarify this point. In point 1 of the contributions, we primarily focus on the literature of adaptive methods with a Polyak stepsize, where zero interpolation (see [45, 53, 79]) and bounded gradient (see [66]) assumptions were used to demonstrate convergence. We are aware of other adaptive algorithms (e.g., Adam) that have been analyzed under the affine gradient noise assumption; however, they are outside the scope of that comparison. We will clarify this inaccuracy in the revised version of the paper. Moreover, as it was stated above NGN-D algorithm can be analyzed under a more relaxed ABC assumption on the gradient noise: see the answer to W1.\n\n***Answer to Q2:*** We thank the reviewer for pointing this out. We will revise the writing accordingly."}]

Author rebuttal:

Meta-review: The paper introduces a momentum enhanced version of the NGN method developed by Orvieto and Xiao. The paper also introduced a coordinate-wise diagonal step-size adaptation similar to classic Adam. Like the Orvieto, Xiao paper, the core idea here is to show optimizer stability with minimum parameter tuning required and to do it with step size that can incorporate the elements that have shown practical advantages. The authors prove this theoretically by providing non-asymptotic rates and show some experiments demonstrating their claims. 

Overall 3/4 reviewers appreciated the paper's contributions with the rebuttals resolving most of the concerns raised by the reviewers. The 4th reviewer had primary one concern remaining which was experiments on LLMs. I believe the authors have performed a wide-spate of experiments across tasks and settings within reach of an academic setup. While larger scale experiments embellish the paper, from the reviews and discussion it can be concluded that even without that the paper represents a solid contribution to optimization literature and can be published.
Decision: Accept (poster)