Title: Enhancing Optimizer Stability: Momentum Adaptation of The NGN Step-size
Venue: NeurIPS 2025
Accepted: yes
Keywords: Adaptive Algorithms

Theoretical Analysis of Adaptive Methods

Polyak Stepsize

Stability
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

Abstract

Modern optimization algorithms that incorporate momentum and adaptive stepsize offer improved performance in numerous challenging deep learning tasks. However, their effectiveness is often highly sensitive to the choice of hyperparameters, especially the learning rate (LR). 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 hyper-parameter choices [79]. In this paper, we introduce an algorithm that matches the performance of state-of-the-art optimizers while improving stability through a novel adaptation of the NGN step-size method [66]. Specifically, we propose a momentum-based version ( NGN-M ) that attains the standard convergence rate of O (1 / √ K ) under common 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 NGN step-size with momentum results in high robustness while delivering performance that is comparable to or surpasses other state-of-the-art optimizers.

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 stringent 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 wellknown 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 necessitates 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 state-of-the-art 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 a 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 a 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,

NGN was shown never to diverge for any choice of the LR hyperparameter in the convex setting, and to 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-wise variants - needed in complex problems such as optimizing transformers, as motivated above. Here, we develop a momentum and step-size adaptive version of NGN designed to enhance robustness 1 in terms of hyperparameter selection. We also present a 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:

We introduce a new algorithm named NGN-M that combines the NGN step-size with momentum. We theoretically show that NGN-M achieves a convergence rate O ( 1 / √ K ) in the convex regime without the typical requirements of interpolation or bounded gradient assumptions found in earlier works on Polyak step-size;

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 ;

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 hyperparameter resilience comes together with better or comparable performance to state-of-the-art algorithms.

2 Related Works

Polyak Step-size. When training a 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, 7]. Yet, in the convex setting, choosing LR can already be difficult - an issue that was studied already in Polyak [72] and gave rise to the first adaptive method: the Polyak Stepsize ( PS ). Recently, there has been a renewed interest in adapting PS to modern settings [54, 67, 39], delivering a 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 to the exact solution only if the interpolation condition holds 2 .

Polyak Step-size and Heavy-ball Momentum. Heavy-ball momentum methods, stemming from the work of Polyak [71], have gained significant attention over the years due to their benefits, including acceleration on convex quadratics [37, 48, 6], convex-like [92], and non-convex problems [14], as well as 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 et al. [79] demonstrated that a 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-wise adaptive step-sizes are essential in training Transformer architectures due to the varying parameter-wise scaling and conditioning of the problem

[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 as 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 improvements targeted to the coordinate-wise (diagonal) regime. However, most Polyak-step-size-based algorithms only focus on a single LR for all parameters [54, 93, 27, 63, 66]. Only a few works propose a diagonal-wise modification of Polyak-step-size by either using Adam preconditioner [79] as a weight matrix or incorporating second-order information from the objective function [51, 76].

Quantitative Measure of Robustness. To quantify robustness, we adopt the learning-rate (LR) sensitivity metric of Wortsman et al. [97]. Let ℓ γ denote the final performance metric (e.g., negative test accuracy) when training with LR γ , and ℓ 0 the value at initialization. We define ℓ ∗ := min γ ∈ [ a,b ] ℓ γ as the best achievable metric within the LR range [ a, b ] . The LR sensitivity is then E γ ∈ [ a,b ] [min { ℓ 0 , ℓ γ } -ℓ ∗ ] . We estimate this expectation by averaging over the LR values in our sweep grid for each algorithm and task.

Comparison to prior work. Table 1 provides a theoretical comparison of various Polyak step-sizebased 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 ) := √ f ( x ) . At the current point x k , they linearized r ( x k + p ) ≈ r ( x k ) + ∇ r ( x k ) ⊤ p. Thus the next iterate is given as x k +1 = x k + p k where

It turns out that the problem above has a closed-form solution

with γ k 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

that NGN exhibits a non-divergence property, regardless of the step-size hyperparameter c (see Theorem 4 . 5 in [66]). Due to this property, the NGN step-size is a strong candidate to 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:

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 to 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 2 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 f ( x ) = 1 2 ∥ Ax -b ∥ 2 where A is a data matrix from the normalized Diabetes dataset [85] and b is a vector of labels. Based on the results from Figure 1 (first), we observe that variant 1 achieves accelerated convergence as SGDM for middle-range step-size hyperparameters ( c ∈ { 10 1 , 10 2 } ) and does not diverge for large LR hyperparameter ( c ∈ { 10 3 } ). Conversely, version 2 has a worse convergence rate than version 1 for middle-range LR hyperparameters and diverges for large ones ( c ∈ { 10 3 } ) : see Figure 1 (second). Therefore, we 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 ( x, y ) = ( x -1) 2 +100( y -x 2 ) 2 , whose minimizer is at (1 , 1) . Starting from ( -1 . 2 , 1) , we run both NGN-M and SGDM over a wide range of LR hyperparameters { 10 -3 , . . . , 10 2 } . 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 c = 10 2 , thanks to its adaptive LR 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.1 we repeat these experiments on a 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-wise adaptive. In the first approach, we modify an approach of (1): The next iterate x k +1 is obtained by minimizing an approximation of the regularized first-order Taylor expansion of r ( x ) := √ f ( x ) around x k , namely, x k +1 = x k + p k where for a preconditioning matrix Σ k

The intuition is that Σ k ∈ R d × d can penalize each parameter with its own weight while in vanilla NGN the penalization is the same for all parameters, and f is an objective function we aim to minimize.

.

0

40

80

120

160

200

Iterations

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

10 4

Train Loss

0

40

80

120

160

200

Iterations

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

10 4

10 6

Train Loss

SGDM lr

10 3

SGDM lr

10 2

SGDM lr

10 1

SGDM lr

10 0

NGN lr

10 3

NGN lr

10 2

NGN lr

10 1

NGN lr

10 0

NGN-M lr

10 3

NGN-M lr

10 2

NGN-M lr

10 1

NGN-M lr

10 0

Ver. 1 lr

10 3

Ver. 1 lr

10 2

Ver. 1 lr

10 1

Ver. 1 lr

10 0

Ver. 2 lr

10 3

Ver. 2 lr

10 2

Ver. 2 lr

10 1

Ver. 2 lr

10 0

¡

2

¡

1

0

1

2

x

0

1

2

y

NGN-M

c

=10

¡

3

NGN-M

c

=10

¡

2

NGN-M

c

=10 2

¡

2

¡

1

0

1

2

x

0

1

2

y

SGDM lr

=10

¡

3

SGDM lr

=10

¡

2

Performing simple derivations (see Appendix F), we obtain the following update rule

Note that by choosing Σ k to be an identity matrix, the step-size γ k in (3) reduces to the vanilla NGN step-size.

Alternatively, we can adopt a simpler, parameter-wise rule: For each parameter j , we replace the full gradient norm in the NGN step-size with its own partial derivative ∇ j f S k ( x k ) . Both of the described per-coordinate variants can be further adjusted by an RMSprop -style preconditioner D k = diag (( D k ) (1) , . . . , ( D k ) ( d ) ) and lead to the following update rule (see Alg. 2 for a full description)

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, we 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 is 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 in 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 be integrated into NGN-MD , potentially providing more accurate curvature approximations and improved stability with respect to the LR. We leave a systematic study of these alternatives to future work.

In the special case β 1 = 0 and Σ k = 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-parameter Polyak-type step-size while achieving the standard O (1 / √ K ) 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,

Algorithm 1 NGN-M

1: Input: x -1 = x 0 ∈ R d , step-size hyperparameter c > 0 , momentum parameter β ∈ [0 , 1)

2: for k = 0 , 1 , . . . , K -1 do

3: Sample S k ⊆ [ n ]

4: γ k = c 1+ c 2 f S k ( x k ) ∥∇ f S k ( x k ) ∥ 2

6: end for

Algorithm 2 NGN-MD

1: Input: x 0 ∈ R d , step-size hyperparameter c > 0 , momentum parameters β 1 , β 2 ∈ [0 , 1) , stabilization parameter ε > 0 , second-order momentum v 0 = 0

2: for k = 0 , 1 , . . . , K -1 do

3: Sample S k ⊆ [ n ]

4: v k = β 2 v k -1 +(1 -β 2 )( ∇ f S k ( x k ) ⊙∇ f S k ( x k ))

5: D k = diag( ε I + √ v k / (1 -β k 2 ))

7: For NGN-MDv1: Σ -1 k = γ k D -1 k

6: For NGN-MDv1: γ k = c 1+ c 2 f S k ( x k ) ∥∇ f S k ( x k ) ∥ 2 D -1 k

10: end for

where x are the parameters of a model we aim to train, n is the number of data points in the dataset, d is the number of parameters, x ∗ denotes the solution to (4), and f i represents the loss associated with the i -th data point/batch. We assume that each f i is differentiable and non-negative 3 and that the global optimal value is bounded, i.e. f ∗ = argmin x f ( x ) ∈ R . Moreover, we assume that we have access to mini-batch stochastic losses f S during training such that f ∗ S := argmin x f S ( x ) < ∞ 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 f i is convex and L -smooth, i.e., for all x, y ∈ R d and i ∈ [ n ] we have ⟨∇ f i ( x ) , y -x ⟩ ≥ f i ( x ) -f i ( y ) and ∥∇ f i ( x ) -∇ f i ( y ) ∥≤ L ∥ x -y ∥ .

Assumption 4.2. The interpolation σ 2 int := E S [ f ∗ -f ∗ S ] and positive σ 2 pos := E S [ f ∗ S ] errors are bounded. We say that the interpolation holds if σ 2 int = 0 , where S is a sampled mini-batch.

4.2 Convergence Guarantees

Theorem 4.3. Let Assumptions 4.1, 4.2 hold. Let the step-size hyperparameter c > 0 and the momentum parameter β = λ 1+ λ be constants where λ ≤ min { cL, 0 . 5(1+ cL ) -1 (1+2 cL ) -1 } . Then the iterates of NGN-M (Algorithm 1) satisfy

where x K -1 is chosen uniformly at random from { x 0 , . . . , x K -1 } . Moreover, if we set c = O ( 1 / √ K ) then we obtain E [ f ( x K -1 ) -f ( x ∗ ) ] ≤ O ( 1 / √ K ) .

The convergence of NGN-M is provided in 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 c is sufficiently small, NGN-M attains the same convergence rate as SGDM [22]. Moreover, for any choice of c , we demonstrate that the NGN-M iterates provably converge to a neighborhood of the optimum and

10

¡

4

10

¡

2

10 0

10 2

Stepsize

0

20

40

60

80

100

SGDM

Momo

NGN-M

10

¡

4

10

¡

2

10 0

10 2

Stepsize

0

20

40

60

80

SGDM

Momo

NGN-M

10

¡

1

10 0

10 1

10 2

Stepsize

0

20

40

60

80

100

SGDM

Momo

NGN-M

10

¡

5

10

¡

4

10

¡

3

10

¡

2

10

¡

1

10 0

10 1

10 2

Stepsize

0

20

40

60

80

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

5

10

¡

3

10

¡

1

10 1

Stepsize

0

20

40

60

80

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

4

10

¡

3

10

¡

2

Stepsize

0

20

40

60

80

100

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

SGDM

Momo

NGN-M

Resnet20, CIFAR10

18

26

20

12

Resnet110, CIFAR10

31

36

31

24

ViT, CIFAR10

27

28

31

1

Adam

Momo-Adam

NGN-MDv1

Resnet20, CIFAR10

37

27

13

Resnet110, CIFAR10

32

32

27

ViT, CIFAR10

41

21

6

Resnet20 for CIFAR 10

Resnet110 for CIFAR 100

ViT for CIFAR 10

thereafter remain within it; ( ii ) Unlike prior works on Polyak step-size, our analysis does not rely on strong assumptions such as bounded gradients, interpolation, or a bounded domain; ( iii ) For small values of c , NGN-M converges to the exact solution while algorithms such as MomSPS and ALR-SMAG were shown to converge up to a non-vanishing neighborhood of the solution only, due to an inherent limitation of the stochastic Polyak step-size [65]. Regarding the momentum parameter β , the typical (large) value β = 0 . 9 performs well in our own experiments. Theoretically, however, β 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 σ 2 int = 0 , we can establish convergence even for large momentum values, including the commonly used choice β = 0 . 9 (see Appendix D.2). This suggests that the smallβ requirement may reflect limitations of current proof techniques rather than an intrinsic restriction of NGN-M . Extending the analysis to arbitrary β in the stochastic regime without interpolation remains an open problem. Our intuition, however, is that simultaneously choosing both c and β 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 a diminishing step-size of order 1 / √ k 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 c to a quadratic problem that satisfies all our assumptions: We observe NGN-M 's rapid convergence with theoretical step-size hyperparameters in practice-see Appendix H.3 and Figure H.4 therein.

Key Ingredients of the Proof. Wediscuss the key steps of 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 { z k } whose update rule is of the form

Next, one of the key technical strategies we follow is splitting the step-size γ k into two parts: a fixed term ρ = c (1+ cL )(1+2 cL ) = O ( c ) and a changing term ˜ γ k ≤ 3 c 2 L 1+2 cL = O ( c 2 ) . This decomposition of the step-size γ k 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 is weighted by c while the error term proportional to σ 2 int is weighted by c 2 , which suggests that c has to be chosen to tradeoff the two terms to lead to the exact convergence similarly to the standard analysis of SGD [22]. In contrast, MomSPS and Momo algorithms achieve the exact convergence only under the interpolation regime.

10

¡

2

10

¡

1

10 0

10 1

10 2

Stepsize

0

20

40

60

SGD

SGDM

Momo

NGN-M

10

¡

5

10

¡

4

10

¡

3

10

¡

2

10

¡

1

Stepsize

0

20

40

60

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

4

10

¡

3

10

¡

2

Stepsize

0

10

20

30

40

50

60

70

Adam

Momo-Adam

NGN-MDv1

Resnet18 for

ImageNet1k

Resnet18 for

ImageNet1k

ViT-tiny for

ImageNet1k

SGD

SGDM

Momo

NGN-M

Resnet18, ImageNet1k

18

5

6

6

5

Adam

Momo-Adam

NGN-MDv1

Resnet18, ImageNet1k

16

16

12

ViT-Tiny, ImageNet1k

29

26

18

10

¡

3

10

¡

2

Stepsize

33

34

35

36

37

38

39

Final T est Perplexity

Adam

NGN-MDv1

10

¡

3

10

¡

2

Stepsize

23

24

25

26

27

28

29

Final T est Perplexity

Adam

NGN-MDv1

10

¡

3

10

¡

2

Stepsize

16

17

18

19

20

21

22

Final T est Perplexity

Adam

NGN-MDv1

10

¡

3

10

¡

2

Stepsize

12

13

14

15

16

17

18

Final T est Perplexity

Adam

NGN-MDv1

70

MTransformer++

160

MTransformer++

410

MTransformer++

1

B Transformer++

5 Experiments

We now turn to the empirical evaluation of the proposed algorithms against several benchmarks. The detailed experiment setup, including the choice of hyperparameters as well as additional experimental results and details, can be found in Appendix I. The best performance of algorithms is reported in Tables 7 (momentum-based 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 to the paper's empirical results are summarized in Table 6, while Appendix I provides additional details about the training and tokenization pipeline.

Comparison on Standard Benchmarks. First, we 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] on the CIFAR10 dataset [45], and Resnet110 on CIFAR100. Second, we test the performance of NGN-MD against Adam and Momo-Adam that contrary to NGN-M - both use component-wise preconditioning. All experiments in this section do not use LR schedulers or weight decay.

From Tables 7 and 8 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, use 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, we clearly observe that momentum consistently 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.

Vision Experiments on ImageNet. Now we switch to larger tasks and datasets. We train a ResNet18 on ImageNet1k [16]. This represents the first task in 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

Momo

NGN-M

Momo-Adam

NGN-MDv1

0.0001

0.000316

0.001

0.00316

0.01

0.0316

0.0001

0.000316

0.001

0.00316

0.01

0.0316

10²

10²

10-1

10-1

Stepsize

Stepsize

Stepsize

Stepsize

101

101

10

-2

10

2

Effective

100

Effective

100

Effective

10

-3

Effective

10

-3

10

-1

10

~1

10

10

1

E

10-2

0

5

10

15

21

10-2

0

5

10

15

21

10-5

0

5

10

15

21

10-

0

5

10

15

20

Iterations，x103

Iterations,x103

Iterations，x103

Iterations,x103

achieving higher performance on a wider range of LR (Table 8). Appendix I.4 reports additional ablations on ImageNet32 and train loss stability results.

Next, we test the effectiveness of the algorithms when training ViT-Tiny on ImageNet1k. This model is trained for a 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 a 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 our 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 70 M, 160 M, 410 M, and 1 B parameters around Chinchilla optimum [33] on SlimPajama-627B [86]. For each model, we retune the LR, using 3 seeds for the first three models and 1 seed for the 1 B.

As reported in Figure 4 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 c = 10 -2 while Adam '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 a 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 I.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 NGNM 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 α 0 , effectively reducing them to SGDM and Adam and limiting their resilience. Evidence across ResNet20 training (Figures I.6, I.7, I.9 and I.10) and large-scale language modeling (Figures I.11 to I.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 .

Computation Cost of NGN-MD. Implementing NGN-MDv1 can be slightly more expensive, but the overall cost is 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 ∥∇ f ( x k ) ∥ 2 D -1 k , which amounts to an additional pass over the gradient. This can also be incorporated into the update of D k , avoiding extra matrix operations. In practice, our naive implementation is about 2 . 5 × slower per update than PyTorch'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 to future work. The extended discussion is reported in Appendix F.2.

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 of NGN-M for large values of β and analysis in the non-convex setting to future work, including classes of structured non-convex functions such as PL [73], α -β -condition [34], or Aiming [52]. It would also be worthwhile to study NGN-M under weaker smoothness assumptions [104, 1]. Furthermore, while the two proposed methods for incorporating weight decay into NGNMDv1 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, a topic that requires further investigation. We acknowledge that computing NGN step-size at a large scale may cause runtime overhead, and discuss this limitation in Appendix F.2. We also recognize that integrating NGN-MDv1 with advanced parallelism schemes, such as 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 a 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.

References

Foivos Alimisis, Rustem Islamov, and Aurelien Lucchi. Why do we need warm-up? a theoretical perspective. arXiv preprint arXiv:2510.03164 , 2025. (Cited on page 10)

Maksym Andriushchenko, Francesco D'Angelo, Aditya Varre, and Nicolas Flammarion. Why do we need weight decay in modern deep learning? arXiv preprint arXiv:2310.04415 , 2023. (Cited on pages 45 and 48)

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)

Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning , 2018. (Cited on pages 3 and 25)

Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. Gptneox-20b: An open-source autoregressive language model. arXiv preprint arXiv: 2204.06745 , 2022. (Cited on page 51)

Raghu Bollapragada, Tyler Chen, and Rachel Ward. On the fast convergence of minibatch heavy ball momentum. arXiv preprint arXiv:2206.07553 , 2022. (Cited on page 2)

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.16620 , 2023. (Cited on page 2)

Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005.14165 , 2020. (Cited on page 1)

Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in neural information processing systems , 2024. (Cited on page 53)

Dami Choi, Christopher J Shallue, Zachary Nado, Jaehoon Lee, Chris J Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv:1910.05446 , 2019. (Cited on page 1)

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)

Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065 , 2021. (Cited on page 2)

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://openreview.net/forum?id=ww3CLRhF1v . (Cited on page 3)

Ashok Cutkosky and Harsh Mehta. Momentum improves normalized sgd. In International conference on machine learning . PMLR, 2020. (Cited on pages 1 and 2)

Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. Advances in neural information processing systems , 2019. (Cited on pages 1 and 2)

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition . Ieee, 2009. (Cited on page 8)

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021. (Cited on page 8)

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)

Ilyas Fatkhullin, Alexander Tyurin, and Peter Richtárik. Momentum provably improves error feedback! Advances in Neural Information Processing Systems , 2024. (Cited on page 1)

Simon Foucart. Lecture 6: Matrix norms and spectral radii. lecture notes for the course NSTP187 at Drexel University, Philadelphia, PA, Fall , 2012, 2012. (Cited on page 43)

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.09000 , 2023. (Cited on page 1)

Guillaume Garrigos and Robert M Gower. Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235 , 2023. (Cited on pages 6, 7, 25, 27, 30, and 36)

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 on page 7)

Saeed Ghadimi and Guanghui Lan. Stochastic first-and zeroth-order methods for nonconvex stochastic programming. SIAM journal on optimization , 2013. (Cited on page 7)

Damien Martins Gomes, Yanlei Zhang, Eugene Belilovsky, Guy Wolf, and Mahdi S Hosseini. Adafisher: Adaptive second order optimization via fisher information. arXiv preprint arXiv:2405.16397 , 2024. (Cited on page 5)

Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning . MIT Press, 2016. (Cited on page 2)

Robert Gower, Othmane Sebbouh, and Nicolas Loizou. Sgd for structured nonconvex functions: Learning rates, minibatching and interpolation. In International Conference on Artificial Intelligence and Statistics , 2021. (Cited on page 3)

Robert M 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.01898 , 2025. (Cited on pages 2 and 3)

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , 2016. (Cited on page 8)

Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Lecture notes , 2012. (Cited on page 1)

Tran Hoang, Qinzi Zhang, and Ashok Cutkosky. Empirical tests of optimization assumptions in deep learning. arXiv preprint arXiv:2407.01825 , 2024. (Cited on page 6)

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation , 1997. (Cited on page 51)

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, 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 preprint arXiv:2203.15556 , 2022. (Cited on page 9)

Rustem Islamov, Niccoló Ajroldi, Antonio Orvieto, and Aurelien Lucchi. Loss landscape characterization of neural networks without over-parametrization. In Advances in Neural Information Processing Systems , 2024. (Cited on pages 6 and 10)

Rustem Islamov, Yuan Gao, and Sebastian U Stich. Near optimal decentralized optimization with compression and momentum tracking. arXiv preprint arXiv:2405.2011 , 2024. (Cited on page 1)

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)

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)

Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning , 2022. (Cited on page 1)

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)

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://kellerjordan.github.io/posts/muon/ . (Cited on page 10)

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.08361 , 2020. (Cited on page 51)

Andrej Karpathy. char-rnn. https://github.com/karpathy/char-rnn , 2015. (Cited on page 51)

Andrej Karpathy. Nanogpt. https://github.com/karpathy/nanoGPT , 2022. (Cited on page 51)

Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR) , 2015. (Cited on pages 1 and 55)

Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Learning multiple layers of features from tiny images. Scientific Report , 2009. (Cited on page 8)

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)

Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavytailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv: 2402.19449 , 2024. (Cited on page 3)

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)

Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of adam under relaxed assumptions. Advances in Neural Information Processing Systems , 2024. (Cited on page 1)

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)

Shuang Li, William J Swartworth, Martin Takáˇ c, Deanna Needell, and Robert M Gower. Sp2: A second order stochastic polyak method. arXiv preprint arXiv:2207.08171 , 2022. (Cited on page 3)

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)

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)

Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In International Conference on Artificial Intelligence and Statistics , 2021. (Cited on pages 1, 2, 3, and 6)

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 1711.05101 , 2019. (Cited on page 46)

Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv:1902.09843 , 2019. (Cited on page 53)

Jerry Ma and Denis Yarats. Quasi-hyperbolic momentum and adam for deep learning. arXiv preprint arXiv:1810.06801 , 2018. (Cited on pages 1 and 2)

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)

Tomas Mikolov, Martin Karafiá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)

Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on adam instability in large-scale machine learning. arXiv preprint arXiv:2304.09871 , 2023. (Cited on page 1)

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)

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)

Dimitris Oikonomou and Nicolas Loizou. Stochastic polyak step-sizes and momentum: Convergence guarantees and practical performance. arXiv preprint arXiv:2406.04142 , 2024. (Cited on pages 2, 3, 4, 6, 7, and 25)

Sharir Or, Barak Peleg, and Yoav Shoham. The cost of training nlp models: A concise overview. arXiv preprint arXiv:2004.08900 , 2020. (Cited on page 1)

Francesco Orabona and Ryan D'Orazio. New perspectives on the polyak stepsize: Surrogate functions and negative results. arXiv preprint arXiv:2505.20219 , 2025. (Cited on page 7)

Antonio Orvieto and Lin Xiao. An adaptive stochastic gradient method with non-negative gauss-newton stepsizes. arXiv preprint arXiv: 2407.04358 , 2024. (Cited on pages 1, 2, 3, 4, 6, and 26)

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 2 and 6)

Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. In Proceedings of the ACL , 2005. (Cited on page 52)

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. NIPS 2017 Workshop Autodiff , 2017. (Cited on page 51)

Guilherme Penedo, Hynek Kydlíˇ cek, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et 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)

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)

Boris T Polyak. Introduction to optimization. New York, Optimization Software , 1987. (Cited on page 2)

Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal vychislitel'noi matematiki i matematicheskoi fiziki , 1963. (Cited on page 10)

Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training. Technical report, OpenAI, 2019. (Cited on page 9)

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-16. IEEE, 2020. (Cited on pages 10 and 45)

Peter Richtárik, Simone Maria Giancola, Dymitr Lubczyk, 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)

Samer Saab, Shashi Phoha, Minghui Zhu, and Asok Ray. An adaptive polyak heavy-ball method. Machine Learning , 2022. (Cited on page 2)

Mher Safaryan and Peter Richtárik. Stochastic sign descent methods: New algorithms and better theory. In International Conference on Machine Learning , 2021. (Cited on page 25)

Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, and Robert M. Gower. MoMo: Momentum models for adaptive learning rates. In Proceedings of the 41st International Conference on Machine Learning , 2024. (Cited on pages 1, 2, 3, 4, 6, 9, and 51)

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 arXiv:2501.18965 , 2025. (Cited on page 6)

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 on pages 7 and 25)

Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv: 2002.05202 , 2020. (Cited on page 51)

Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 , 2019. (Cited on pages 10 and 45)

Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556 , 2014. (Cited on page 51)

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 on page 4)

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)

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)

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning , 2021. (Cited on page 1)

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothé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.13971 , 2023. (Cited on pages 1 and 51)

Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of adam's iteration complexity. Advances in Neural Information Processing Systems , 2024. (Cited on page 1)

Jun-Kun Wang, Chi-Heng Lin, and Jacob D Abernethy. A modular analysis of provable acceleration via polyak's momentum: Training a wide relu network and a deep linear network. In International Conference on Machine Learning , pages 10816-10827. PMLR, 2021. (Cited on page 43)

Jun-Kun Wang, Chi-Heng Lin, Andre Wibisono, and Bin Hu. Provable acceleration of heavy ball beyond quadratics for a class of polyak-lojasiewicz functions when the non-convexity is averaged-out. In International conference on machine learning , 2022. (Cited on pages 2 and 7)

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 on pages 2, 3, 4, and 7)

Rachel Ward, Xiaoxia Wu, and Leon Bottou. Adagrad stepsizes: Sharp convergence over nonconvex landscapes. Journal of Machine Learning Research , 2020. (Cited on page 27)

Ross Wightman. Pytorch image models. https://github.com/rwightman/ pytorch-image-models , 2019. (Cited on page 51)

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 information processing systems , 2017. (Cited on page 1)

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 arXiv:2309.14322 , 2023. (Cited on pages 1 and 3)

Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling. arXiv preprint arXiv: 2409.15156 , 2024. (Cited on pages 45 and 48)

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)

Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813 , 2023. (Cited on page 2)

Biao Zhang and Rico Sennrich. Root mean square layer normalization. arXiv preprint arXiv:1910.07467 , 2019. (Cited on page 51)

Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse. Three mechanisms of weight decay regularization. arXiv preprint arXiv:1810.12281 , 2018. (Cited on page 45)

Hongyang R Zhang, Dongyue Li, and Haotian Ju. Noise stability optimization for finding flat minima: A hessian-based regularization approach. arXiv preprint arXiv:2306.08553 , 2023. (Cited on page 2)

Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint arXiv:1905.11881 , 2019. (Cited on pages 3 and 10)

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)

Qi Zhang, Yi Zhou, and Shaofeng Zou. Convergence guarantees for rmsprop and adam in generalized-smooth non-convex optimization with affine noise variance. arXiv preprint arXiv:2404.01436 , 2024. (Cited on page 1)

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068 , 2022. (Cited on page 1)

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 on page 3)

Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. 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 NGN-MD to support the claims made in 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 in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We include a discussion on the limitations of our work in the conclusion and appendix.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a 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 on a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 to Appendix C, D.1, 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 crossreferenced.

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 if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: We provide the Pytorch-based implementation in the supplementary. We use an open-source implementation of models and public datasets in our 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 are provided or not.

If the contribution is a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We use only publicly available code and datasets and provide links to them in Appendix I.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a 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) is 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide all training details in Appendix I 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

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 3 random seeds (if the opposite is not stated). We do not include error bars when training 1B language model due to a 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 are accompanied by error bars, confidence 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).

It should be clear whether the error bar is the standard deviation or 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide a description of compute resource we used for each experiment in Appendix I.

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 of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research is conducted following the NeurIPS instructions, including the regulations regarding anonymity.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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]

Justification: 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.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is 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 are 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

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 a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of 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 is used.

At submission time, remember to 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 as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

Institutional review board (IRB) approvals or equivalent for research with human subjects

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 an 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, 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, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Contents

Appendix

A Equivalent Formulations of NGN-M

We remind that the iterates of NGN-M are the following

We can rewrite the update rule using Iterative-Moving Average (IMA) approach presented in Proposition 1.6, Sebbouh et al. [81].

Lemma A.1 (Proposition C.8 [63], Lemma 7.3 in [22]) . The iterates { x k } generated by NGN-M are equivalent to the sequence { x k } generated by IMA update

where

Proof. Let the sequences { x k } and { z k } be defined according to Equation (5). Let β be defined as λ 1+ λ . Then we have

It remains to use (6) as we have β = λ 1+ λ and 1 -β = 1 -λ 1+ λ = 1 1+ λ .

B Technical Lemmas and Definitions

Definition B.1. We say that the function ϕ admits L -smooth with parameters L := ( L 1 , . . . , L d ) , L j ≥ 0 ∀ j ∈ [ d ] , if the following inequality holds for all x, h ∈ R d

Remark B.2 . If we set for all j ∈ [ d ] L j := L then Definition B.1 reduces to standard L -smoothness.

This assumption is typically used in the context of coordinate adaptive algorithms such as SignSGD [4, 78].

Definition B.3. The function ϕ : R d → R satisfies PŁ-condition with constant µ > 0 if for all x, y ∈ R d we have

Assumption B.4. We assume that the coordinate-wise variance of the stochastic estimator is bounded, i.e. for all x ∈ R d and j ∈ [ d ] we have

Lemma B.5 (Lemma 4.9 from [66]) . Let each f i be L -smooth for all i , then the step-size of NGN satisfies

Lemma B.6 (Lemma 4.2 from [66]) . Let each f i be L -smooth for all i , then the iterates of NGN satisfy

Lemma B.7 (Gradient Upper Bound) . Let ϕ : R d → R satisfy Definition B.1. Then, for all x ∈ R d and all j ∈ [ d ] we have

Proof. From Definition B.1 we have

Now we can explicitly compute the minimum in the right-hand side. The optimal value is achieved at

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

1: Input: x 0 ∈ R d , step-size parameter c > 0

2: for k = 0 , 1 , . . . , K -1 do

3: Sample a batch S k ⊆ [ n ] and compute f S k and ∇ f S k ( x k )

4: Compute γ ( j ) k = c 1+ c 2 f S k ( x k ) ( ∇ j f S k ( x k )) 2

5: Update

6: end for

Theorem C.1. Let each f i satisfies Definition B.1. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters { c j } d j =1 such that c j ≤ 1 / 2 L j satisfy

where c min := min j ∈ [ d ] c j . Moreover, if c j = O ( ε 2 ) for all j ∈ [ d ] then after K = O ( ε -4 ) we obtain min 0 ≤ k<K E [ ∥∇ f ( x k ) ∥ 2 ] ≤ O ( ε 2 ) .

therefore,

NGN-D converges with classic rate O ( 1 / √ K ) 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 PŁ-condition and each f i satisfies Definition B.1. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters { c j } d j =1 such that c j ≤ min { 1 / 2 L j , 6 / µ } satisfy

where c min := min j ∈ [ d ] c j . Moreover, if c j = O ( ε ) for all j ∈ [ d ] then after K = max {O ( ε -1 ) , O (1) } log ε -1 iterations we obtain E [ f ( x K ) -f ∗ ] ≤ O ( ε ) .

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 of SGD [22].

Now we are ready to derive the step-size bounds.

Lemma C.3 (Step-size Bounds) . Let f S k ( x ): R d → R be a stochastic loss of batch S k at iteration k. Let f S k ( x ) satisfy Definition (B.1) . Consider γ k j as in NGN-D (Algorithm 3), then we have

Proof. From Lemma B.7 we have 2 L j ( f S k ( x k ) -f ∗ S k ) ≥ ( ∇ j f S k ( x k )) 2 . Since we assume that each f ∗ S k ≥ 0 , then 2 L j f S k ( x k ) ≥ ( ∇ j f S k ( x k )) 2 , or equivalently,

Therefore, for all j ∈ [ d ] we have

and which concludes the proof.

Lemma C.4 (Fundamental Equality) . Consider γ k j as in NGN-D (Algorithm 3). Then the following equality holds

Proof. From NGN-D (Algorithm 3) we have

which one can rewrite as

It is left to divide both sides by 2 f S k ( x k ) c j

.

C.1 Convergence in General Non-convex Setting

Theorem C.1. Let each f i satisfies Definition B.1. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters { c j } d j =1 such that c j ≤ 1 / 2 L j satisfy

where c min := min j ∈ [ d ] c j . Moreover, if c j = O ( ε 2 ) for all j ∈ [ d ] then after K = O ( ε -4 ) we obtain min 0 ≤ k<K E [ ∥∇ f ( x k ) ∥ 2 ] ≤ O ( ε 2 ) .

Proof. First, we write separable Definition B.1

Note that both γ k j and ∇ j f S k ( x k ) depend on the realization S k , thus we can not directly apply conditional expectation with respect to x k , as in this case we would have to analyze the product γ k j ∇ j f S k ( x k ) . Given bounds of the step-size γ k j from Lemma C.3, we can write the step-size as follows

where ν k j ∈ [0 , 1] is a random variable. Varying the value of ν k j from 0 to 1 we cover the whole range of γ k j . Thus, we continue as follows

Now we use the inequality | ab |≤ 1 2 a 2 + 1 2 b 2 + 1 2 | a -b | 2 , and derive

Therefore, we get

We plug in (18) into (17) and get

If c j ≤ 1 2 L j , we get

We continue as follows

Taking full expectation and unrolling the recursion above for all iterations { 0 , . . . , K -1 } . Thus, we obtain

If we choose each c j = c 0 ,j √ K such that c 0 ,j ≤ 1 2 L j we ensure that c j ≤ 1 2 L j as well. Plugging this step-size into the bound we get

where c 0 , min := min j ∈ [ d ] c 0 ,j . If we choose K = O ( ε -4 ) we get that

C.2 Convergence under PŁ-condition

Theorem C.2. Let f satisfies PŁ-condition and each f i satisfies Definition B.1. Assume that Assumption B.4 holds. Then the iterates of NGN-D (Algorithm 3) with step-size parameters { c j } d j =1 such that c j ≤ min { 1 / 2 L j , 6 / µ } satisfy

where c min := min j ∈ [ d ] c j . Moreover, if c j = O ( ε ) for all j ∈ [ d ] then after K = max {O ( ε -1 ) , O (1) } log ε -1 iterations we obtain E [ f ( x K ) -f ∗ ] ≤ O ( ε ) .

Proof. We obtain (19) and use Definition B.3

Subtracting f ∗ from both sides of the inequality above and taking full expectation we obtain

Unrolling the recursion above for { 0 , . . . , K -1 } iterations we derive

Now we follow the proof of Lemma A.3 in Garrigos and Gower [22]. Let us choose c j = min { 1 / 2 L j , ε / 2 dA j } . Together with the choice of K ≥ max j ∈ [ d ] max { 1 ε 12 A j µ , 12 L j µ } log 2( f ( x 0 ) -f ∗ ) ε we get

Now we have two cases:

1. c min does not depend on ε, then we have

2. c min does depend on ε, i.e. c min = O ( ε ) , then we have

Therefore, combining all together we get

after K ≥ max j ∈ [ d ] max { 1 ε 12 A j µ , 12 L j µ } log 2( f ( x 0 ) -f ∗ ) ε iterations.

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 c > 0 and the momentum parameter β = λ 1+ λ be constants where λ ≤ min { cL, 0 . 5(1+ cL ) -1 (1+2 cL ) -1 } . Then the iterates of NGN-M (Algorithm 1) satisfy

where x K -1 is chosen uniformly at random from { x 0 , . . . , x K -1 } . Moreover, if we set c = O ( 1 / √ K ) then we obtain E [ f ( x K -1 ) -f ( x ∗ ) ] ≤ O ( 1 / √ K ) .

Remark D.1 . In fact, if λ ≤ 1 (1+ cL )(1+2 cL ) , then it implies that λ ≤ 1 cL because 1 x > 1 (1+ x )(1+2 x ) for any x > 0 .

Proof. To prove the convergence of NGN-M we consider IMA formulation Equation (5):

At iteration k = 0 we have

Therefore, we get

Let γ 0 = ρ + ˜ γ 0 where ρ = c (1+ cL )(1+2 cL ) . Then we have

Using the above we continue from (20)

Here we have

˜ γ 0 ≤ 3 c 2 L 1+2 cL , and f S 0 ( x 0 ) -f ∗ S 0 ≥ 0 . Hence, we get

Rearranging terms and taking expectation we get

Let γ k = ρ + ˜ γ k , where ρ, ˜ γ k ≥ 0 , and ρ is a constant step-size independent of S k which will be defined later. Therefore, we have

We need to find ρ such that

Since ˜ γ k = γ k -ρ , then we have

The inequality above is satisfied if it is satisfied for the lower bound on γ k (which is c / 1+ cL ), i.e.

We can take ρ = c (1+ cL )(1+2 cL ) since λ ≥ 0 .

Using the above, we get from (23)

Taking expectations we get

Rearranging terms we get

Combining Equation (22) and Equation (25) for iterations { 1 , . . . , K -1 } we get

We need to ensure that ρ -cλ > 0 which is satisfied for λ such that

Note that we also assume that λ ≤ cL . Therefore, from (26) we get

Since ρ -cλ ≥ ρ 2 and setting x k be uniformly at random chosen from { x 0 , . . . , x K -1 } we get

where we use the convexity of f and Jensen's inequality. Plugging the value of ρ = c (1+ cL )(1+2 cL ) inside we get

Choosing c = O (1 / √ K ) we get

Therefore, if K ≥ O ( ε -2 ) then E [ f ( x k ) -f ( x ∗ ) ] ≤ O ( ε ) . It remains to notice that z 0 = x 0 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 β -including the default β = 0 . 9 -provided the LR hyperparameter c is chosen sufficiently small. This is natural: if the step-size is too large, the momentum term accumulates excessive past error, which the algorithm may be unable to correct, potentially causing divergence. In short, convergence can be ensured in two complementary ways: ( i ) use a small momentum parameter while allowing any LR hyperparameter c (Theorem 4.3); or ( ii ) restrict the LR hyperparameter c while permitting a large (near-1) momentum parameter β (Theorem D.2).

Theorem D.2. Let Assumption 4.1 hold. Assume that there exists a minimizer x ∗ shared across all functions f i , i.e. f ∗ i = f ( x ∗ ) for all i ∈ [ n ] . Let NGN-M is run with momentum parameter β = λ 1+ λ , λ ∈ (0 , Λ] , Λ ≥ 9 , step-size c ≤ min { 1 4Λ L , 1 2 L } . Then the iterates of NGN-M satisfy

Remark D.3 . Note that the requirement Λ ≥ 9 allows setting the momentum parameter β = 0 . 9 , which is a default choice in practice.

Proof. The proof is similar to Theorem 4.3, but we take into account the additional interpolation assumption. We start with

Since f S k ( x k ) -f S k ( x ∗ ) ≥ 0 , with the choice c ≤ 1 2 L , we have

Moreover, by Lemma B.5 we have γ k ≥ c 1+ cL and γ k ≤ c . Thus, we continue the bound (31) on the distance ∥ z k +1 -x ∗ ∥ 2 as follows

Taking the expectation, we obtain the final bound on the distance

Note that for k = 0 , the bound is simpler, since the momentum term is zero

where we use z 0 = x 0 . Summing the bounds (32) and (33) for iterations k ∈ { 0 , . . . , K -1 } , we obtain

Note that

where the last inequality holds by the choice of c ≤ 1 4Λ L ≤ 1 4 λL . Therefore, we obtain

Thus, we obtain the rate

where the last inequality holds by the choice of c .

D.3 Convergence of NGN-M with Decaying Step-size

Proof. We refer to Lemma A.8 from Garrigos and Gower [22].

To prove the convergence of NGN-M with decaying c k we consider IMA formulation (see Section A in the paper):

where c k = c 0 √ k +1 , λ 0 = 0 .

Theorem D.5. Assume that each f i is convex and L -smooth, and that Assumption 3.2 holds. Let the step-size hyperparameter is set c k = c 0 √ k , momentum parameter λ k ≤ min { c k L, 0 . 5(1+ c k L ) -1 (1+ 2 c k L ) -1 } . Then the iterates of NGN-M satisfy

Proof. At iteration k = 0 we have

Therefore, we get

Let γ 0 = ρ 0 + ˜ γ 0 where ρ 0 = 0 (1+ c 0 L )(1+2 c 0 L ) . Then we have

Using the above we continue from (36)

Here we have

˜ γ 0 ≤ 3 c 2 0 L 1+2 c 0 L , and f S 0 ( x 0 ) -f ∗ S 0 ≥ 0 . Hence, we get

Rearranging terms and taking the expectation we get

Let γ k = ρ k + ˜ γ k , where ρ, ˜ γ k ≥ 0 , and ρ is a constant step-size independent of S k which will be defined later. Therefore, we have

We need to find ρ k such that

Since ˜ γ k = γ k -ρ k , then we have

The inequality above is satisfied if it is satisfied for the lower bound on γ k (which is c / 1+ cL ), i.e.

Using the above, we get from (39)

Taking expectations, we get

Rearranging terms, we get

Combining (38) and (41) for iterations { 1 , . . . , K -1 } we get

Note that choosing λ k = min { c k L, 0 . 5(1 + c k L ) -1 (1 + 2 c k L ) -1 } ensures that ρ k 2 ≥ c k λ k . Indeed, we have

Therefore, from (42) and the facts that λ 0 = 0 and λ k ≤ c k L we get

Therefore, using (44), z 0 = x 0 in (43) and dividing both sides in (43) by ∑ K -1 k =0 ρ k we derive

With an lower bound on ∑ K -1 k =0 and Jensen's inequality we conclude that

where ˆ x K -1 = ∑ K -1 k =0 ρ k ∑ K -1 k =0 ρ k x k .

E Stability of NGN-M on a Simple Problem

We consider 1D convex functions of the form f ( x ) = Lx 2 (1 + p 2 ( x )) that satisfy the following assumption.

Assumption E.1. There exists a constant C such that C (1 + p 2 ( x )) ≥ xp ( x ) p ′ ( x )) .

Note that 1+ p 2 ( x ) ≥ 1 and deg(1+ p 2 ( x )) = deg( xp ( x ) p ′ ( x )) . Therefore, this assumption is mild. Remark E.2 . For example, the function f ( x ) = x 2 (1 + x 2 ) (i.e., p ( x ) = x ) is convex and satisfies Assumption E.1 with C = 1 .

Remark E.3 . Let p ( x ) = ∑ m j =0 a j x j . Then for large values of x in magnitude, p ( x ) ∼ a m x m , p ′ ( x ) ∼ ma m x m -1 . Therefore, the constant C should be expected of order C ≈ m, where m = deg( p ( x )) .

The function f ( x ) is non-negative for any x ∈ R and its minimum f ∗ = 0 is attained at x = 0 by design. Let us compute a step of NGN-M on this problem

Note that the convexity of f implies that

In particular, (48) implies that 1 + p 2 ( x ) + xp ( x ) p ′ ( x ) ≥ 1 2 (1 + p 2 ( x )) > 0 . Therefore, we can obtain lower and upper bounds on ˆ γ k .

Lemma E.4. Let Assumption E.1 hold with a constant C > 0 and f ( x ) = x 2 (1 + p 2 ( x )) be convex. Let c ≥ 1 2 L . Then we have ˆ γ k ∈ [ 1 2(1+ C ) , 2 ] .

Proof. Indeed, the upper bound on ˆ γ k follows from the following inequality

due to (48). The lower bound can be obtained as follows

The update rule of NGN-M can be rewritten as

Let us consider the joint dynamics of w k := ([ x k ] ⊤ , [ x k -1 ] ⊤ ) ⊤ ∈ R 2 d . We have that

where

Now we are ready to prove the convergence of NGN-M on this simple problem for any value c ≥ 1 2 . Theorem E.5. Let f ( x ) = x 2 (1+ p 2 ( x )) be convex and Assumption E.1 holds. Let β ≥ (2(1+ C ) -1) 2 (2(1+ C )+1) 2 and c ≥ 1 . Then the iterates of NGN-M on f ( x ) converge to the minimum f ∗ = 0 .

2 L

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, we will proceed with a permutation matrix Π 4 which transforms the matrix G to the block-diagonal matrix as

Since the matrix G is a block-diagonal matrix, we have ∥ G ∥≤ max i ∥ G i ∥ . Therefore, the problem is now simplified to bounding the spectral radii of the individual blocks G i , for i = 1 , 2 , . . . , d . The two eigenvalues u 1 and u 2 of G i are the roots of the quadratic

which take different values depending on the discriminant ∆ := (1 + β -(1 -β )ˆ γ k ) 2 -4 β . Let us find the values of β when the discriminant is negative. We need to satisfy the inequality

Since the function 1+ y 2 2 y for y ∈ [ 1 2(1+ C ) , 2 ] attains the maximum 4(1+ C ) 2 +1 4(1+ C ) at y = 1 2(1+ C ) , then we satisfy the last inequality, and consequently the discriminant is non-positive, if we choose

4 The permutation matrix Π is defined as Π ij =      1 i odd , j = i 1 i even , j = 2 n + i 0 else . Note that permutation matrices preserve eigenvalues.

where

The above inequality is satisfied for β ∈ [ (2(1+ C ) -1) 2 (2(1+ C )+1) 2 , 1 ) . Therefore, we obtain that for such choice of β we have ∆ i ≤ 0 for all i ∈ [ d ] . Therefore, the zeros of the quadratic q ( u ) are complex, and are equal in absolute value

This gives us that ∥ G i ∥≤ √ β < 1 . Therefore, the algorithm converges for any value of β in this range.

It remains to use Lemma 11 from Foucart [20] which says that for a given matrix A ∈ R d × d , and ϵ > 0 , there exists a matrix norm ∥·∥ such that

where ρ ( A ) = max {| λ | : λ eigenvalue of A } (spectral radius of A ).

Asymptotically 5 (as k →∞ , one can show (see Theorem 12 in [20]) that

where ρ ( G ) ≤ √ β < 1 in our analysis. Therefore, NGN-M with hyperparameters c ≥ 1 2 and β ≥ 1 0 converges.

Remark E.6 . For example, NGN-M converges on f ( x ) = x 2 (1 + x 2 ) for any c ≥ 1 2 and β ≥ 9 25 .

Theorem E.5 shows that NGN-M remains stable even with an arbitrarily large step-size hyperparameter c . Thanks to 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 to 1 , which extends the results of Section 4. We acknowledge that our current analysis is restricted to the special convex class of 1D functions f ( x ) = x 2 (1 + p 2 ( x )) satisfying Assumption E.1. Extending such stability guarantees to wider function classes with large momentum β 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 ) = x 2 (1 + x 2 ) , which is convex and satisfies Assumption E.1; see Figure E.1. We run both algorithms, varying the step-size hyperparameter in { 10 -4 , . . . , 10 4 } . We run algorithms for 10 5 iterations. We stop training if the loss reaches a threshold 10 -15 or exceeds 10 10 for the first time. We observe that ( i ) for small step-size hyperparameters, both methods converge but do not reach the threshold 10 -15 ; ( ii ) 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 c = 10 1 after 269 iterations. In 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.

10

¡

4 10

¡

3 10

¡

2 10

¡

1

10 0

10 1

10 2

10 3

10 4

Stepsize

10

¡

15

10

¡

5

10 5

10 15

10 25

Final Train Loss

NGN-M

GDM

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

where r ( x ) = √ f ( x ) . We compute the gradient of RHS of (62) w.r.t. p and equal it to zero:

Therefore, we have

Using Shermann-Morrison formula ( A + uv ⊤ ) -1 = A -1 -A -1 uv ⊤ A -1 1+ u ⊤ A -1 v with A = 1 / c Σ k we derive

Now we plug-in r ( x k ) = √ f ( x k ) and ∇ r ( x k ) = 1 2 √ f ( x k ) ∇ f ( x k ) and obtain

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 a diagonal step-size into NGN by choosing the regularization matrix Σ k . By choosing Σ k = D k we recover the update direction of NGN-MDv1 . In this case, we have only one global NGN step-size in front of D k . The design of NGN-MDv2 follows a more straightforward intuition. In particular, it can be seen as a direct extension of NGN to diagonal case by replacing the squared gradient norm ∥∇ f S k ( x k ) ∥ 2 by the squared partial derivative ( ∇ j f S k ( x k )) 2 for each parameter j ∈ [ d ] .

The main difference in comparison with Adam is the order in which the preconditioning and momentum is applied. In both NGN-MDv1 and NGN-MDv2 we average the preconditioned updates Σ -1 k ∇ f S k ( x k ) , i.e. we first apply preconditioning and momentum later. In contrast, in Adam the stochastic gradients are averaged to construct new momentum term, and then the momentum is preconditioned. In other words, the momentum is applied first and then it is 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 .

F.2 Computation Cost of NGN-MD

Implementing any version of NGN-MD in practice might be slightly more computationally expensive. However, we highlight that computing a step of NGN-MD does not involve matrix-vector operations since the preconditioner is 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 ∥∇ f S k ( x k ) ∥ 2 D -1 k . This can naïvely be done by one additional pass over the gradient and summing the terms 1 ( D k ) j ( ∇ j f S k ( x k )) 2 for j ∈ [ d ] . This operation does not require additional matrix multiplication. However, it can be computed more efficiently while updating D k . The rest 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 NGNMDv1 is about 2 . 5 times slower than PyTorch's AdamW . This is expected since our 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, as noted above, this overhead can be largely eliminated by computing the weighted gradient concurrently with the second-momentum v k update. We do not aim to 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.

F.2.1 Distributed Training

In a vanilla DDP implementation [50], computing the weighted gradient norm ∥∇ f S k ( x k ) ∥ 2 D -1 k is straightforward since gradients are replicated across devices. We only require an additional all-reduce to synchronize f S k ( x k ) across devices, which is, however, a lightweight communication (just a single float) and, in 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 it will, nevertheless, result in a 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 in 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. Orthogonal 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 to lower training loss in under-parameterized networks. However, even in the case of SGD , there is still uncertainty regarding how the weight decay mechanism should be incorporated, as various implementations may exist [102].

We propose two ways of adding weight decay to NGN-MDv1 . The first variant follows the approach of [55], adding decoupled weight decay λ :

In this update rule, the weight is added separately from the update direction Σ -1 k ∇ f S k ( x k ) . We call the resulting algorithm (63) Dec-NGN-MDv1 , that stands for decoupled NGN-MDv1 .

G.1 Combining NGN-MDv1 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

By taking the gradient of f Σ k ,λ w.r.t. p we get

Therefore, we get

Using Sherman-Morrison formula ( A + uv ⊤ ) -1 = A -1 -A -1 uv ⊤ A -1 1+ u ⊤ A -1 v with A = ( λ + 1 / c ) Σ k and u = v = √ 2 ∇ r ( x k ) we get that

Therefore, we have

Algorithm 4 NGN-MDv1W

1: Input: x 0 ∈ R d , step-size parameter c > 0 , momentum parameters β 1 , β 2 ∈ [0 , 1) , weight decay parameter λ ≥ 0 , stabilization parameter ε > 0

2: for k = 0 , 1 , . . . , K -1 do

3: Sample a batch S k ⊆ [ n ] and compute f S k and ∇ f S k ( x k )

4: Compute v k = β 2 v k -1 +(1 -β 2 )( ∇ f S k ( x k ) ⊙∇ f S k ( x k ))

5: Compute D k = diag( ε I + √ v k / (1 -β k 2 ))

6: Compute

8: end for

[ · ] + denotes max { 0 , ·} .

Using the connection

To summarize, the update of NGN-Dv1W is the following

To prevent the step-size next to Σ -1 k ∇ f ( x k ) from being negative, the final update has the form

where [ · ] + := max {· , 0 } . Now we can add momentum on top and obtain the following update of NGN-MDv1W

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 x k .

AdamW

Momo-AdamW

NGN-MDv1W

Dec-NGN-MDv1

10

¡

3

10

¡

2

Stepsize

33

35

37

39

Validation Perplexity

10

¡

3

10

¡

2

Stepsize

33

35

37

39

Validation Perplexity

10

¡

3

10

¡

2

Stepsize

33

35

37

39

Validation Perplexity

10

¡

3

10

¡

2

Stepsize

33

35

37

39

Validation Perplexity

λ

= 0

λ

= 10

-

3

λ

= 10

-

2

λ

= 10

-

1

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, we observe that when weight decay is properly tuned, all algorithms improve over the baseline case with no weight decay, which is consistent with the observation of Xiao [98] and Andriushchenko et al. [2] on AdamW . We also note that Dec-NGN-MDv1 and NGN-MDv1W require a smaller weight decay value compared to the other algorithms. Finally, the stability and performance of NGNMDv1 are preserved by both variations, allowing training with larger LR, and significantly improving over AdamW and Momo-Adam .

We do not observe a 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 to evaluate the effect of adding regularization.

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

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 0 , 10 1 , 10 2 , 10 3 } and the starting point in [ -20 , 20] with a step 4 / 30 6 . Based on the results in Figure H.1 (right), we conclude that ( i ) for small step-sizes, both methods likely get stuck at 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; ( iii ) 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.

¡

20

¡

10

0

10

20

Parameter

x

0

100

200

300

Frequency

SGDM

10

¡

1

SGDM

10 0

SGDM

10 1

SGDM

10 2

0

100

200

Function Value

¡

20

¡

10

0

10

20

Parameter

x

0

100

200

300

Frequency

NGN-M

10

¡

1

NGN-M

10 0

NGN-M

10 1

NGN-M

10 2

0

100

200

Function Value

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 to the choice of step-size hyperparameter. In contrast, SGDM already diverges for the step-size hyperparameter 10 -2 . 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.

¡

2

¡

1

0

1

2

x

¡

1

0

1

2

¡

3

y

NGN-M

c

=10

¡

3

SGDM lr

=10

¡

3

¡

2

¡

1

0

1

2

x

¡

1

0

1

2

¡

3

y

NGN-M

c

=10

¡

2

SGDM lr

=10

¡

2

¡

2

¡

1

0

1

2

x

¡

1

0

1

2

¡

3

y

NGN-M

c

=10

¡

1

H.3 Comparison on Quadratic Function with Theoretical Step-size

Next, we run NGN-M with theoretical choice of step-size hyperparameter c = 1 / √ K and c k = 1 / √ k (see Theorem 4.3 and Theorem D.5 for more details) against fixed choices c ∈ { 10 -3 , 10 -4 } . The

0

0

:

5

1

1

:

5

2

Iteration

k;

£

10 3

10

¡

8

10

¡

5

10

¡

3

10 1

f

(

xk

)

¡

f

(

x

¤

)

NGN-M

0

0

:

5

1

1

:

5

2

Iteration

k;

£

10 3

10

¡

8

10

¡

5

10

¡

3

10 1

f

(

xk

)

¡

f

(

x

¤

)

SGDM

0

0

:

5

1

1

:

5

2

Iteration

k;

£

10 3

10

¡

4

10

¡

3

10

¡

2

Eff. Stepsize of NGN-M

10

¡

3

10

¡

2

10

¡

1

10 0

10 1

10 2

comparison is made on quadratic function f ( x ) = 1 2 ∥ ( A + r I ) x -y ∥ 2 , where A ∈ R 400 × 400 and y ∈ R 400 are sampled from standard normal distribution. The constant r controls the condition number of the problem.

We test the performance of NGN-M varying the condition number of the problem and the number of iterations; see Figure H.4. We observe that in all the cases, the choice 1 / √ k leads to faster convergence, supporting our theoretical claims. The choice 1 / √ K demonstrates competitive performance as well, but it is slightly pessimistic at the beginning of training. In contrast, the choice c ∈ { 10 -3 , 10 -4 } , which is a default value in practice, is too small and does not lead to fast convergence.

These experiments demonstrate that when the problem satisfies all assumptions needed in the analysis, the choice of the step-size hyperparameter c given by the convergence theorems is a good starting point in practice and can serve as a baseline when tuning c.

0

50

100

Iteration

k

10 3

10 2

10 1

10 0

f

(

xk

)

¡

f

(

x

¤

)

f

(

x

0)

¡

f

(

x

¤

)

Cond. number

∙

=3

0

50

100

Iteration

k

10 1

10 0

Cond. number

∙

=99

0

50

100

Iteration

k

10 1

10 0

Cond. number

∙

=1000

0

500

1000

Iteration

k

10 2

10

¡

3

10

¡

8

10

¡

13

f

(

xk

)

¡

f

(

x

¤

)

f

(

x

0)

¡

f

(

x

¤

)

Cond. number

∙

=3

c

=10

¡

3

c

=10

¡

4

c

=1

=

p

K

c

=1

=

p

k

0

500

1000

Iteration

k

10

¡

2

10

¡

1

10 0

10 1

Cond. number

∙

=99

0

500

1000

Iteration

k

10

¡

2

10

¡

1

10 0

10 1

Cond. number

∙

=1000

0

5000

10000

Iteration

k

10 0

10

¡

5

10

¡

10

10

¡

15

f

(

xk

)

¡

f

(

x

¤

)

f

(

x

0)

¡

f

(

x

¤

)

Cond. number

∙

=3

0

5000

10000

Iteration

k

10 1

10

¡

1

10

¡

3

Cond. number

∙

=99

0

5000

10000

Iteration

k

10 1

10 0

10

¡

1

10

¡

2

Cond. number

∙

=1000

I Additional Experiments and Training Details

I.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 ( β 1 , β 2 ) = (0 . 9 , 0 . 999) if the opposite is not specified. The stepsize 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 a vit_tiny_patch16_224 for 200 epochs on Imagenet1k, using a cosine LR schedule with a linear warmup of 5 epochs. Differently than Schaipp et al. [79], we train in bfloat16 , instead of float16 , and do not 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 ( β 1 , β 2 ) = (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-627B [86], a cleaned and deduplicated version of RedPajama We report validation perplexity on a separate subset of Slim-Pajama consisting of 10 Mtokens. The total compute is estimated following Kaplan et al. [41], where the estimated number of floating-point operations (FLOPs) is 6 × Number of Parameters × Number of Tokens.

Experiments of small and middle size are performed on 1xRTX 4090. We perform ImageNet32 experiments on 2xA100-40GB, and ImageNet1k experiments on 4xA100-SXM4-40GB. For pretraining Transformers on Language Modeling, we employ 8xH100-HBM3-80GB 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 I.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 on (VGG16 [84], CIFAR10) and (MLP, MNIST) workloads where NGN-M matches the performance of the state-of-theart 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 or improves it.

I.3 Comparison of Algorithms that Support Momentum and Diagonal Step-size

Next, we illustrate the performance of the algorithms that support both momentum and diagonal step-size. According to the results in Figures I.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

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.

I.4 Additional ImageNet Experiments

Now we turn to the experiments involving training Resnet18 on ImageNet1k and ImageNet32. In Figure I.4 we provide the train loss curves and results on (Resnet18, ImageNet32) workload that demonstrate that NGN-M and NDN-MDv1 attain better resilience to the step-size hyperparameter

choice than competitors not only from the train loss point of view as well. The best performance of algorithms is provided in Table 7 and 8. According to them, both NGN-M and NGN-M achieve competitive performance against considered benchmarks.

I.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, CIFAR10), (Resnet110, CIFAR100), and (ViT, CIFAR10) workloads; Lion has worse performance on (Resnet110, CIFAR100) 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 to the choice of step-size hyperparameter.

I.6 Comparison of Adaptive Step-sizes of Adam, Momo-Adam, and NGN-MDv1

Next, we conduct experiments to compare the adaptive step-size of Adam , Momo-Adam , and NGN-MDv1 . Note that ResNet20 model consists of 3 base blocks, and each block has 3 convolution layers. In Figure I.6 we plot the average adaptive step-size of the layers j ∈ { layer1.0.conv1, layer2.0.conv1, layer3.0.conv1 } of ResNet20 that corresponds to the first convolution layer within each base block. Similarly, in Figure I.7 we plot the average adaptive step-size of the layers j ∈ { layer0.0.fn.to_qkv, layer3.0.fn.to_qkv, layer5.0.fn.to_qkv } that corresponds to the attention layers of the first, fourth, and sixth base blocks.

Since the adaptivity of Adam is only in the second-order momentum applied as a normalization, in our experiment we compare the following quantities

where γ is the step-size hyperparameter of Adam .

SGDM

Momo

NGN-M

10

¡

4

10

¡

2

10 0

10 2

Stepsize

10 2

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

10

¡

4

10

¡

2

10 0

10 2

Stepsize

10 2

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

10

¡

1

10 0

10 1

10 2

Stepsize

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

Resnet20 for CIFAR 10

Resnet110 for CIFAR 100

ViT for CIFAR 10

10

¡

4

10

¡

2

10 0

10 2

Stepsize

0

20

40

60

80

10

¡

4

10

¡

2

10 0

10 2

Stepsize

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

VGG16 for CIFAR 10

VGG16 for CIFAR 10

10

¡

3

10

¡

1

10 1

10 3

Stepsize

0

20

40

60

80

100

10

¡

3

10

¡

1

10 1

10 3

Stepsize

10 0

10

¡

1

10

¡

2

10

¡

3

Final Train Loss

MLP for MNIST

MLP for MNIST

Let us first describe the results for ResNet20 in Figure I.6. We observe that NGN-MDv1 tends to set smaller effective step-size compared to two other algorithms. This is 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 of Adam and Momo-Adam . In contrast, the coordinate-wise 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 mechanism of NGN-MDv1 step-size is more conservative than that of Momo-Adam .

Now we switch to the results on ViT model in Figure I.7. Here both Momo-Adam and NGN-MDv1 tend to utilize smaller effective coordinate-wise step-size, by several orders in magnitude smaller than that of Adam . However, the adaptation mechanism of NGN-MDv1 is still more conservative than that of Momo-Adam , especially for large step-size hyperparameters. We also highlight that in this experiment the best performance of NGN-MDv1 is achieved with c = 10 -3 . When we vary the step-size hyperparameter c , the effective coordinate-wise step-size does not change dramatically, especially for layers.0.0.fin.to_qkv layer.

I.7 Extended Comparison of Momentum-based Algorithms on 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 I.8 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 of NLP models has significantly different coordinate-wise conditioning. Nonetheless, NGN-M algorithm achieves better resilience to the step-

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

5 10

¡

4 10

¡

3 10

¡

2 10

¡

1

10 0

10 1

10 2

Stepsize

10 2

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

10

¡

5

10

¡

3

10

¡

1

10 1

Stepsize

10 3

10 2

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

10

¡

4

10

¡

3

10

¡

2

10

¡

1

Stepsize

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

Resnet20 for CIFAR 10

Resnet110 for CIFAR 100

ViT for CIFAR 10

10

¡

6

10

¡

4

10

¡

2

10 0

Stepsize

0

20

40

60

80

10

¡

6

10

¡

4

10

¡

2

10 0

Stepsize

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

VGG16 for CIFAR 10

VGG16 for CIFAR 10

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

4

10

¡

2

10 0

10 2

Stepsize

0

20

40

60

80

100

10

¡

4

10

¡

3

10

¡

2

10

¡

1

Stepsize

10 1

10 2

10 3

10 4

10 5

Final T est Perplexity

10

¡

5

10

¡

4

10

¡

3

10

¡

2

10

¡

1

10 0

10 1

10 2

Stepsize

10 0

10 1

10 2

10 3

10 4

10 5

Final T est Perplexity

MLP for MNIST

LSTM for PTB

Transformer for

10

¡

4

10

¡

2

10 0

10 2

Stepsize

10 2

10 1

10 0

10

¡

1

10

¡

2

10

¡

3

Final Train Loss

10

¡

4

10

¡

3

10

¡

2

10

¡

1

Stepsize

10 2

10 1

10 0

Final Train Loss

10

¡

5

10

¡

3

10

¡

1

10 1

Stepsize

10 3

10 2

10 1

10 0

Final Train Loss

MLP for MNIST

LSTM for PTB

Transformer for

size hyperparameter choice, especially in the training of Transformer models. Therefore, NGN-M across various model architectures and task domains.

I.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 { c j } d j =1 where d is the size of the model. Finding sufficiently good constants c j might be a challenging task since d is a large number. Therefore, we

10

¡

2

10

¡

1

10 0

10 1

10 2

Stepsize

0

20

40

60

SGD

SGDM

Momo

NGN-M

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

Stepsize

2

4

6

8

10

12

Final Train Loss

SGD

SGDM

Momo

NGN-M

10

¡

2

10

¡

1

10 0

10 1

10 2

Stepsize

1

2

3

4

5

Final Train Loss

SGD

SGDM

Momo

NGN-M

Resnet18 for ImageNet32

Resnet18 for ImageNet32

Resnet18 for ImageNet1k

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

Stepsize

0

10

20

30

40

50

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

Stepsize

2

4

6

8

10

12

Final Train Loss

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

10

¡

5

10

¡

4

10

¡

3

10

¡

2

10

¡

1

Stepsize

1

2

3

4

5

Final Train Loss

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

Resnet18 for ImageNet32

Resnet18 for ImageNet32

Resnet18 for ImageNet1k

SGDM

Momo

NGN-M

Adam

Momo-Adam

NGN-MDv1

NGN-MDv2

Lion

Adabelief

Adabound

10

¡

5

10

¡

3

10

¡

1

10 1

10 3

Stepsize

0

20

40

60

80

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

Stepsize

0

20

40

60

80

10

¡

5

10

¡

4

10

¡

3

10

¡

2

10

¡

1

10 0

10 1

10 2

Stepsize

0

20

40

60

80

100

10

¡

5

10

¡

3

10

¡

1

10 1

10 3

Stepsize

10 2

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

10

¡

6

10

¡

4

10

¡

2

10 0

10 2

Stepsize

10 2

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

10

¡

5 10

¡

4 10

¡

3 10

¡

2 10

¡

1

10 0

10 1

10 2

Stepsize

10 1

10 0

10

¡

1

10

¡

2

Final Train Loss

Resnet20 for CIFAR10

Resnet110 for CIFAR100

ViT for CIFAR10

use RMSprop preconditioner D k to set them as c j = c/ ( D k ) ( j ) . We leave the exploration of how to set constants c j properly for future research.

For each method, we tune its learning rate hyperparameter over the powers of 10 : { 10 -4 , . . . , 10 2 } 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 on (LSTM, Wikitext-2) workload. Besides, Adagrad always has the worst performance. Moreover, these algorithms do not have high resilience to the choice of hyperparameter. Therefore, we omit their comparison from this perspective.

I.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 , NGNMDv1 , and Momo-Adam in Figures I.9 and I.10. First, both NGN-M and Momo perform a warm-up

0

:

0001

0

:

001

0

:

01

0

:

1

1

:

0

10

:

0

100

:

0

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

layer1.0.conv1

Adam

layer1.0.conv1

Momo-Adam

layer1.0.conv1

NGN-MDv1

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

layer2.0.conv1

Adam

layer2.0.conv1

Momo-Adam

layer2.0.conv1

NGN-MDv1

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10 1

10 4

10 7

10 10

Effective Stepsize

layer3.0.conv1

Adam

layer3.0.conv1

Momo-Adam

layer3.0.conv1

NGN-MDv1

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 a short warm-up. This trend is similar to the state-of-the-art learning rate schedulers used in practice. Similar observations can be made in comparison of NGN-MDv1 and Momo-Adam .

0

:

0001

0

:

000316

0

:

001

0

:

00316

0

:

01

0

:

0316

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

layers.0.0.fn.to_qkv

Adam

layers.0.0.fn.to_qkv

Momo-Adam

layers.0.0.fn.to_qkv

NGN-MDv1

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

layers.3.0.fn.to_qkv

Adam

layers.3.0.fn.to_qkv

Momo-Adam

layers.3.0.fn.to_qkv

NGN-MDv1

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

0

5

10

15

20

Iterations,

£

10 3

10

¡

2

10

¡

1

10 0

10 1

10 2

10 3

10 4

10 5

Effective Stepsize

layers.5.0.fn.to_qkv

Adam

layers.5.0.fn.to_qkv

Momo-Adam

layers.5.0.fn.to_qkv

NGN-MDv1

I.10 Effective Updates in Training Language Models

In this section, we demonstrate the magnitude of updates when training 160 M 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 I.13. We demonstrate that NGN-MDv1 is a more conservative algorithm: the effective update is smaller than that of Adam due to the adaptive nature of the step-size. This is especially evident when training 160 Mlanguage 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.

I.11 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, we study the sensitivity of NGN-MDv1 and Adam to the choice of the learning rate and momentum hyperparameters, when training 70M language model on FineWeb dataset [70]. To do that, we fix β 1 = 0 . 9 (or β 2 = 0 . 95 ) and make a sweep over the learning rate hyperparameter and

10

¡

4

10

¡

2

10 0

10 2

10 1

Stepsize

1

3

5

7

9

Final Train Loss

SGDM

Momo

NGN-M

10

¡

5

10

¡

3

10

¡

1

10 1

10 3

Stepsize

1

3

5

7

9

Final Train Loss

SGDM

Momo

NGN-M

10

¡

5

10

¡

3

10

¡

1

10 1

10 3

Stepsize

1

3

5

7

9

Final Train Loss

SGDM

Momo

NGN-M

10

¡

4

10

¡

2

10 0

10 2

10 1

Stepsize

10 1

10 2

10 3

10 4

10 5

Final T est Perplexity

SGDM

Momo

NGN-M

10

¡

5

10

¡

3

10

¡

1

10 1

10 3

Stepsize

10 3

10 2

10 1

10 0

Final T est Perplexity

SGDM

Momo

NGN-M

10

¡

5

10

¡

3

10

¡

1

10 1

10 3

Stepsize

10 3

10 2

10 1

10 0

Final T est Loss

SGDM

Momo

NGN-M

LSTM for

Wikitext-2

Transformer for

Transformer for

Momo

for Resnet20 for CIFAR 10

NGN-M

for Resnet20 for CIFAR 10

0

:

1

0

:

316

1

:

0

3

:

16

10

:

0

31

:

6

Momo

for ViT for CIFAR 10

NGN-M

for ViT for CIFAR 10

0.0001

0.001

0.01

0.1

1.0

10.0

100.0

1000.0

103

103,

103

103

Stepsize

Stepsize

EffectiveStepsize

Stepsize

101

101

101

101

Effective

10

1

Effective

10

1

10

-1

Effectve

10

~1

10

10

3

10

3

10

3

0

5

10

15

21

0

100

200

300

400

500

0

10

15

20

0

100

200

300

400

500

Iterations,x103

Iterations

Iterations,x103

Iterations

10²

10²

102

10²

Stepsize

Stepsize

Stepsize

Stepsize

10'

10

10'

101

Effective

100

Effective

10°

Effective

100

Effectve

10°

10

I~

10

10

~1

10

10-2

0

5

10

15

21

102

0

100

200

300

400

500

10-2,

0

5

10

15

20

102

0

100

200

300

400

500

Iterations,x103

Iterations

Iterations,×103

Iterations

β 2 (or learning rate hyperparameter and β 1 ). We report the final test perplexity averaged over 3 runs for each set of hyperparameters.

We summarize our findings from Table 10, Table 11, Table 12, and Table 13 as follows:

Low lr (3e-3): NGN-MDv1 and Adam show similar sensitivity to changes in both β 1 and β 2 .

Moderate lr (1e-2): NGN-MDv1 is noticeably more robust than Adam to extremes of β 1 , while both optimizers perform similarly across β 2 (though Adam 's performance degrades slightly at β 2 = 0 . 999 ).

High lr (3e-2): Both methods suffer when β 1 is small (or β 2 is large), but NGN-MDv1 recovers lower perplexity at larger β 1 values (smaller β 2 values), whereas Adam fails to reach comparable performance.

Momo-Adam

for Resnet20 for CIFAR 10

NGN-MDv1

for Resnet20 for CIFAR 10

Momo-Adam

for ViT for CIFAR 10

NGN-MDv1

for ViT for CIFAR 10

0.0001

0.001

0.01

0.1

1.0

10.0

100.0

1000.0

103

103

103

103

Effective Stepsize

EffectiveStepsize

eStepsize

EffectiveStepsize

101

101.

101

T

10-

10-1

Effective

10-1.

10-1

10-3

10-3

10-3

0

5

10

15

0

400

8001200

16002

20

0

5

10

15

0

400

8001200

1600

2000

Iterations,×103

Iterations

Iterations,×103

Iterations

0.0001

0.000316

0.001

0.00316

0.01

0.0316

10-

10-1

Effective Stepsize

EffectiveStepsize

Effective Stepsize

Effective Stepsize

10-2

10

10-2

10

10-4

10-4

10-4

10-4.

10-5

0

5

10

15

10-5

0

100

200300

400

10-5

0

5

10

15

10-5

0

100

200300

400

500

Iterations,×103

Iterations

Iterations,×103

Iterations

0

1000

2000

3000

4000

5000

Step

3

4

5

6

7

8

Training Loss

Training Loss

0

1000

2000

3000

4000

5000

Step

0

1

2

3

Effective Update Norm

Embedding Weight

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

0.8

Effective Update Norm

Layer 1 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0.00

0.25

0.50

0.75

1.00

1.25

Effective Update Norm

Layer 1 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

0.8

Effective Update Norm

Layer 1 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

Effective Update Norm

Layer 7 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

0.8

1.0

Effective Update Norm

Layer 7 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

0.8

Effective Update Norm

Layer 7 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

Effective Update Norm

Layer 12 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

0.8

1.0

Effective Update Norm

Layer 12 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.2

0.4

0.6

Effective Update Norm

Layer 12 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0

1

2

3

Effective Update Norm

LM Head

Optimizer

Adam

NGN-MDv1

To 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 around 50 ).

0

1000

2000

3000

4000

5000

Step

3

4

5

6

7

8

Training Loss

Training Loss

0

1000

2000

3000

4000

5000

Step

0.0

2.5

5.0

7.5

10.0

12.5

Effective Update Norm

Embedding Weight

0

1000

2000

3000

4000

5000

Step

0

1

2

3

Effective Update Norm

Layer 1 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0

1

2

3

4

Effective Update Norm

Layer 1 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0

1

2

3

Effective Update Norm

Layer 1 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Effective Update Norm

Layer 7 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0

1

2

3

4

Effective Update Norm

Layer 7 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.5

1.0

1.5

2.0

2.5

Effective Update Norm

Layer 7 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Effective Update Norm

Layer 12 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0

1

2

3

4

Effective Update Norm

Layer 12 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Effective Update Norm

Layer 12 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0.0

2.5

5.0

7.5

10.0

12.5

Effective Update Norm

LM Head

Optimizer

Adam

NGN-MDv1

0

1000

2000

3000

4000

5000

Step

5.5

6.0

6.5

7.0

7.5

8.0

Training Loss

Training Loss

0

1000

2000

3000

4000

5000

Step

0

5

10

15

20

Effective Update Norm

Embedding Weight

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

Effective Update Norm

Layer 1 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0.0

2.5

5.0

7.5

10.0

12.5

15.0

Effective Update Norm

Layer 1 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

10

Effective Update Norm

Layer 1 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

Effective Update Norm

Layer 7 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

10

12

Effective Update Norm

Layer 7 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

Effective Update Norm

Layer 7 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

Effective Update Norm

Layer 12 Attention QKV Proj

0

1000

2000

3000

4000

5000

Step

0

2

4

6

8

Effective Update Norm

Layer 12 MLP In Proj

0

1000

2000

3000

4000

5000

Step

0

1

2

3

4

5

6

Effective Update Norm

Layer 12 MLP Out Proj

0

1000

2000

3000

4000

5000

Step

0

10

20

30

Effective Update Norm

LM Head

Optimizer

Adam

NGN-MDv1

0.0

0.2

0.4

0.6

0.8

1.0

1.2

Tokens

1e9

4

5

6

7

8

9

10

Training Loss

Optimizer

Adam

NGN-MDv1

LR

0.003

0.01

0.03

3.5

3.6

0.0

0.5

1.0

1.5

2.0

2.5

Tokens

1e9

3

4

5

6

7

Training Loss

Optimizer

Adam

NGN-MDv1

LR

0.003

0.01

0.03

3.1

3.2

70M Transformer

++

160M Transformer

++

0

1

2

3

4

5

6

7

Tokens

1e9

3

4

5

6

7

8

Training Loss

Optimizer

Adam

NGN-MDv1

LR

0.003

0.01

0.03

2.80

2.85

0.0

0.5

1.0

1.5

2.0

Tokens

1e10

2.5

3.0

3.5

4.0

4.5

5.0

5.5

6.0

6.5

Training Loss

Optimizer

Adam

NGN-MDv1

LR

0.001

0.003

0.006

2.54

2.56

++

++

410M Transformer

1B Transformer

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\n1.  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...\".\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-16"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nApart 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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-16"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
1. It provides a theoretical analysis and proof for the explanation.
2. It conducts comprehensive experiments on different tasks and models.
3. Good introduction on the background and related work to show the position in the literature.

Weakness:
1. The proposed algorithm is the extension of NGN, so the solution itself seems to be incremental. 
2. The robustness is not well quantified by measuring some metrics like performance variance.
3. 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.

QUESTIONS:
1. Can you conduct experiments on larger models and datasets? e.g.， pretrained MoE LLMs
2. Can you add LR scheduler and weight decay in experiments?
3. Can you give the variances of different step sizes for all compared methods?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSummary: \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).\n\nSTRENGTHS AND WEAKNESSES:\nStrenghts: \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\n\nQUESTIONS:\nQuestions: \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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}