Enhancing Optimizer Stability: Momentum Adaptation of The NGN Step-size

We thank the reviewer for valuable comments.

Answer to W1: In Section C, we provide a nonconvex convergence analysis for NGN‑D, which may have been overlooked by the reviewer—we kindly ask them to revisit this section. Moreover, that analysis readily extends to the more general ABC noise model [1], which covers the affine noise case mentioned by the reviewer as a special instance. We will include this broader analysis in the revised manuscript.

We also agree that developing a nonconvex theory for NGN‑M would strengthen the analysis, and we mention this explicitly in the conclusion as an important direction for future study. While convexity is typically considered a strong assumption, especially for neural networks training, a growing body of empirical evidence shows that convex convergence bounds often predict the behavior of large‑scale language models quite accurately [1], and that neural loss landscapes frequently exhibit convex‑like regions around the minimizer [3-5]. Finally, our theoretical contributions significantly advance prior Polyak‐stepsize adaptive methods by relaxing assumptions and tightening convergence guarantees, and we validate these findings through a comprehensive empirical evaluation.

[1] Khaled and Richtarik. Better theory for SGD in the nonconvex world. TMLR, 2023.

[2] Schaipp et al., The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. ICML 2025

[3] Tran et al., Empirical tests of optimization assumptions in deep learning. arXiv preprint arXiv: 2407.01825.

[4] Islamov, et al. Loss landscape characterization of neural networks without over-parametrization. NeurIPS 2024

[5] Guille-Escuret et al. No wrong turns: The simple geometry of neural networks optimization paths. arXiv preprint arXiv: 2306.11922, 2023.

Answer to W2: We highlight that the results demonstrated in Figure 2 do not use any lr scheduling and weight decay, as the reviewer suggests, since they are not required to achieve strong performance in those settings. We demonstrate an improved stability of NGN-M and NGN-MDv1 to the lr hyperparameter. Moreover, we do not claim that the proposed adaptive scheme of setting lr hyperparameter replaces other techniques such as warm-up or weight decay. In contrast, all of them can be combined together for a further improvement in robustness of the algorithms since warm-up and weight decay are known to improve sensitivity of algorithms to the choice of the lr hyperparameter [3]. Therefore, for training tasks where warm-up and weight decay are necessary, we employ standard training procedures that are summarized in Section J.1 and Table 3.

The empirical part of the work involves testing the proposed algorithms on 10 datasets, including both image and text source of data; 6 types of architectures (MLP, VGG, Resnet, LSTM, ViT, Transformer) with the largest 1B model. In all the cases, the proposed algorithms achieve competitive performance with an increased robustness to the lr hyperparameter, which is the main focus of the work. In our view, this comprehensive set of experiments sufficiently supports the validity of our claims. While we agree that scaling to even larger models and datasets would strengthen the empirical evidence, such experiments are not feasible within the rebuttal period (a single run of the 1B model with our compute resources takes 11 hours on 8xH100-80Gb). Demanding larger-scale experiments would exceed our available compute resources and, more broadly, raises concerns about fairness since such a requirement tends to favor institutions with disproportionate access to large-scale GPU infrastructure, rather than evaluating contributions on the basis of methodological innovation.

[3] Wortsman, Mitchell, et al. "Small-scale proxies for large-scale transformer training instabilities." arXiv preprint arXiv:2309.14322 (2023).

Answer to Q1: We would like to clarify this point. In point 1 of the contributions, we primarily focus on the literature of adaptive methods with a Polyak stepsize, where zero interpolation (see [45, 53, 79]) and bounded gradient (see [66]) assumptions were used to demonstrate convergence. We are aware of other adaptive algorithms (e.g., Adam) that have been analyzed under the affine gradient noise assumption; however, they are outside the scope of that comparison. We will clarify this inaccuracy in the revised version of the paper. Moreover, as it was stated above NGN-D algorithm can be analyzed under a more relaxed ABC assumption on the gradient noise: see the answer to W1.

Answer to Q2: We thank the reviewer for pointing this out. We will revise the writing accordingly.

We thank the reviewer for valuable comments.

Answer to W1: We thank the reviewer for a careful read of our paper. We will incorporate all proposed changes in the revised version of the work. The notation for $f^* \in \mathbb{R}$ means that the objective is bounded from below, i.e., $f^* > -\infty$.

Answer to W2: Thank you for the suggestion to apply our results to reinforcement learning (RL). However, we would like to emphasize that our current experimental focus on the standard classification setting is deliberate: sensitivity to hyperparameters is already a well-documented and significant challenge in this domain. By demonstrating the robustness of our method in such a widely studied and sensitive setting, we believe our work already makes a significant contribution. But we agree that RL is an exciting future direction due to its high sensitivity, and we will mention this in the paper and consider it as follow-up work, thank you again for pointing this out.

Answer to Q1: In our view, it is a fair comparison. Momo and NGN-M (Momo-Adam and NGN-MDv1) are based on the SPS (Stochastic Polyak Stepsize) and the NGN stepsizes, respectively, and they have hyperparameters $\alpha_0$ and $c$. These hyperparameters play a similar role: they adjust the effective lr hyperparameter during training, not allowing it to be larger than $\alpha_0$ or $c$, respectively. We therefore believe it is fair to adjust these parameters and would be happy to discuss this further if the reviewer thinks otherwise.

Answer to Q2: NGN-MD is an extension of the NGN-M algorithm. Therefore, all the discussion on the lr hyperparameter sensitivity of NGN-M applies to NGN-MD as well. Moreover, the derivation of the NGN-M update rule can be seen as a special case of NGN-MDv1 with $\mathbf{\Sigma}_k=\mathbf{I}.$ From an implementation point of view, using a general diagonal preconditioner $\mathbf{\Sigma}_k$ leads to per-parameter adaptive lr. In the case of RMSprop preconditioning, performing one step of training increases since we need to compute the second-order momentum $v^k$. To link, the difference in training time of NGN-M and NGN-MDv1 is similar to that between SGD and Adam. However, the robustness of NGN-MDv1 also depends on the preconditioning. In the cases when the preconditioning is sensitive to the changes in the loss landscape or hyperparameters, NGN-MDv1 might be less robust to the choice of the lr hyperparameter than NGN-M.

Answer to Q3: [1] This paper studies how to set momentum parameters $\beta_1$ and $\beta_2$ based on frequency domain analysis. However, we did not read any claim about the robustness of their approach to changes of the hyperparameters. Based on Figure 4 in [1], the hyperparameters should be chosen in a specific way, which may limit their robustness in practice. Therefore, we ask the reviewer to clarify which part of the work is related to ours. Nonetheless, their idea on changing the momentum parameter is orthogonal to ours, which makes the comparison difficult. But their momentum scheduling can be combined with our NGN-M algorithm, which would potentially lead to improved results.

[2] Thank you for the reference, and we will cite this work in the revision. [2] introduces AdaFisher, a novel way to choose diagonal preconditioning by leveraging the composite structure of neural networks. Specifically, the authors derive a diagonal approximation to the Fisher Information Matrix (FIM) and use it to rescale each parameter update. This curvature‐aware scaling yields more accurate local stepsizes and, in their experiments, substantially reduces sensitivity to the lr choice when training ResNet‑50 on CIFAR‑100. It therefore remains an open question whether the same stability gains carry over to other architectures and datasets. Therefore, the robustness of their approach is underexplored. However, their new AdaFisher preconditioning can be incorporated into NGN-MDv1, which potentially leads to better curvature approximation of AdaFisher and improved stability of NGN-MDv1 to the choice of the lr hyperparameter.

[3] The approach is based on computing the update direction as $\frac{1}{2}(\nabla f(x+u) + \nabla f(x-u)),$ where $u$ is sampled from some distribution. Such an approach leads to convergence to flatter minima where the trace of the Hessian is minimized, i.e., they study another. This is an alternative to the SAM algorithm, which minimizes the maximum eigenvalue of the Hessian. However, this approach is again orthogonal to ours. While the method potentially might lead to flatter minima, the efficiency of the algorithm is tested; it is still sensitive to the choice of the lr hyperparameter: see Figure 4 in [3], which is the main focus of our work. However, the main disadvantage of this approach is the necessity of computing at least two backpropagations per step, which increases both time and memory consumption during the training when the model is sufficiently large. Thus, we ask the reviewer to clarify which part of the work is related to ours since this work seems not to study robustness to the choice of hyperparameters.

Answer to W1: We agree that NGN-M builds upon the NGN method, but this extension is far from trivial. In Section 3.1, we compare two variants of coupling momentum and the NGN stepsize. We demonstrate that the second variant (Adam-type momentum) does not preserve the non-divergence feature of NGN in the convex case. Moreover, in the convex setting, the coupling should lead to acceleration if $c$ is small enough. Hence, the integration of momentum into NGN is not straightforward and requires a careful analysis, which we provide in the paper.

Answer to W2: We thank the reviewer for this insightful comment, as it allowed us to rethink how to better demonstrate the robustness of the proposed method, and eventually improve our work further. We can quantify the robustness in a similar way to how it was done in [1]. We define $\ell\_{\gamma}$ to be the final metric (e.g., train loss or negative test accuracy that we need to minimize) achieved when training a model with the lr hyperparameter $\gamma\in[a,b]$, while $\ell(x_0)$ is the metric at initialization. Moreover, we define the best possible metric $\ell^\*\_{\gamma_*}$ that can be achieved by training a model with $\gamma\in[a,b],$ i.e., $\ell^\*\_{\gamma_*} = \min\_{\gamma\in[a,b]}\ell\_{\gamma}.$ The lr sensitivity is then defined as $\mathbb{E}\_{\gamma\in[a,b]}[\min\\{\ell\_{\gamma}, \ell(x_0)\\} - \ell^\*\_{\gamma_*}]$. We measure this quantity for various tested algorithms and tasks, approximating the expectation by averaging $\min\\{\ell\_{\gamma}, \ell(x_0)\\} - \ell^\*\_{\gamma_*}$ over the grid of values for $\gamma$. This metric captures the maximum deviation from optimal performance. Across all experiments, both NGN‑M and NGN‑MDv1 show consistently lower sensitivity to the lr hyperparameter than the other baselines, quantitatively confirming our claims.

Here, $\gamma \in \{10^{-2}, 3\cdot 10^{-2}, \dots, 3\cdot 10^0, 10, 10^1, 10^2\}$ for the SGDM, NGN, Momo, NGN-M and $\gamma \in \{10^{-4}, \dots, 10^0\}$ for Adam, Momo-Adam, and NGN-MD.

Here, $\gamma \in \{10^{-2}, \dots, 10^2\}$ for the SGDM, NGN, Momo, NGN-M and $\gamma \in \{10^{-4}, \dots, 10^0\}$ for Adam, Momo-Adam, and NGN-MD.

Here, $\gamma \in \{10^{-1}, 3\cdot 10^{-1}, 10^0, 3\cdot 10^0, 10^1, 3\cdot 10^1\}$ for the SGDM, NGN, Momo, NGN-M and $\gamma \in \{10^{-4}, 3\cdot 10^{-4}, \dots, 10^{-2}, 3\cdot 10^{-2}\}$ for Adam, Momo-Adam, and NGN-MD.

Here, for 70M, 160M, 410M models, $\gamma \in \{10^{-3}, 1.7\cdot 10^{-3}, 3\cdot10^{-3}, 5.5\cdot 10^{-3}, 10^{-2}, 3\cdot10^{-2}\}$. For 1B model, $\gamma \in \{10^{-3}, 3\cdot10^{-3}, 5.5\cdot 10^{-3}, 10^{-2}, 3\cdot10^{-2}\}$.

[1] Wortsman, Mitchell, et al. "Small-scale proxies for large-scale transformer training instabilities." arXiv preprint arXiv:2309.14322 (2023).

Answer to W3 and Q2: We respectfully disagree with this claim. Learning rate scheduling and weight decay are not required to achieve strong performance in our core experiments, as shown in Figure 2. When training Resnet18 on Imagenet1k, we use learning rate decay every 30 epochs by 0.1 without weight decay, while when training ViT on Imagenet1k, we use cosine annealing with warm-up and weight decay. For language modeling, we use linear warm-up followed by cosine annealing. To summarize, for each training task we employ standard training procedures that are summarized in Section J.1 and Table 3. Moreover, we provide a discussion on how weight decay can be added to NGN-MD in Section H and provide additional ablations by varying the weight decay parameter $\lambda$ in Figure H.1.

Answer to Q1: The empirical part of the work involves testing the proposed algorithms on 10 datasets, including both image and text source of data; 6 types of architectures (MLP, VGG, Resnet, LSTM, ViT, Transformer) with the largest 1B model. In all the cases, the proposed algorithms achieve competitive performance with an increased robustness to the lr hyperparameter, which is the main focus of the work. In our view, this comprehensive set of experiments sufficiently supports the validity of our claims. While we agree that scaling to even larger models and datasets would strengthen the empirical evidence, such experiments are not feasible within the rebuttal period (a single run of a 1B model with our compute resources takes 11 hours on 8xH100-80Gb). Demanding larger-scale experiments would exceed our available compute resources, and, more broadly, raises concerns about fairness since such a requirement tends to favor institutions with disproportionate access to large-scale GPU infrastructure, rather than evaluating contributions on the basis of methodological innovation.

Answer to Q3: We thank the reviewer for this comment. Note that Figures 5 and J.9 do show the standard deviation. However, since the stepsize changes rapidly, this cannot be observed in Figure 5. The reviewer still can observe the variance in the 2nd and 4th columns of Figure J.9. The standard deviation is relatively small. We will update the plots by using exponential smoothing so that the variance can be better observed in all figures.

We thank the reviewer for valuable comments

Answer to W1: The reviewer is right, we acknowledged this concern in the discussion after Theorem 4.3. Nonetheless, in Section F, we demonstrate that for simple functions like convex polynomials, the restriction on $\beta$ being small can be avoided. This is substantiated by experimental evidence showing that NGN-M and NGN-MD perform well. This makes us believe that the restriction on $\beta$ can be avoided but it requires a more careful analysis, which we were not able to obtain at the moment. Moreover, we highlight that other momentum-based algorithms, including MomSPS$_{\max}$, have similar limitations. We also note that we can remove the restriction on $\beta$ by making the stepsize decreasing or under, additional assumptions like bounded iterates similar to [53]. While these modifications improve the theoretical analysis, we believe they would make the algorithm less favorable in practice.

Answer to W2: We directly addressed the concern regarding runtime in section G.2 (also mentioned in the conclusion). Table 2 reports the time needed for performing one step of training and optimizer.step() when training language models at different scale. We observe that the time needed for one training step of NGN-MDv1 does not increase significantly and is comparable to that of AdamW (see middle column of Table 2). This happens because *** backpropagation is the most time-consuming part*** of training. Therefore, the final training time of AdamW and NGN-MDv1 is comparable: for all tested scales, it increases by a maximum of 3%. Therefore, running the algorithms under a fixed time budget will not affect the results significantly, and the improved robustness of NGN-M/NGN-MDv1 to the choice of lr hyperparameter will remain.

Answer to W3: Following the reviewer's comment, we provide sweeps varying not only the lr hyperparameter but also the momentum parameters $\beta_1,\beta_2$ when training 70M language model on FineWeb dataset. To do that, we fix $\beta_1=0.9$ (or $\beta_2=0.95$) and make a sweep over lr and $\beta_2$ (or lr and $\beta_1$). We report the final test perplexity averaged over 3 runs for each set of hyperparameters. We thank the reviewer for this valuable suggestion, which has enabled us to strengthen our empirical evaluation. We will include these new results in the revised manuscript.

We summarize our findings as follows:

• Low lr ($\le$ 3e-3): NGN-MDv1 and Adam show similar sensitivity to changes in both $\beta_1$ and $\beta_2$.

• Moderate lr (1e-2): NGN-MDv1 is noticeably more robust than Adam to extremes of $\beta_1$, while both optimizers perform similarly across $\beta_2$ (though Adam’s performance degrades slightly at $\beta_2 = 0.999$).

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

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 $\lesssim$ 50).

Answer to Q1: Demonstrating the improved stability from using momentum is a challenging task, especially for an adaptive algorithm like NGN-M. However, to gain intuition, we can consider minimizing a convex quadratic function $\min_{x}\frac{1}{2}x^\top Q x + b^\top x + c.$ The maximum allowed stepsize of GD that leads to convergence is $\boldsymbol{\gamma}<\mathbf{\frac{2}{\boldsymbol{\lambda}\_{\max}(Q)}}$, where $\lambda_{\max}(Q)$ is the maximum eigenvalue of $Q$. When using Heavy-ball algorithm instead, the maximum allowed stepsize for convergence is $\mathbf{\frac{2(1+\boldsymbol{\beta})}{\boldsymbol{\lambda}\_{\max}(Q)}}$ [1], which is about two times larger than for GD for a standard choice of $\beta=0.9$. We acknowledge that training neural networks is far from quadratic,s but it gives an intuition why we should expect an improved robustness to the choice of the lr hyperparameter. This is also supported by recent works that observe convex-like structures in the loss landscape of neural networks [2-4] and agreement between convex theory and practice [5].

[1] Polyak, Some methods of speeding up the convergence of iteration methods. Ussr comp. math. and math. phys., 1964

[2] Tran et al., Empirical tests of optimization assumptions in deep learning. arXiv preprint arXiv:2407.01825.

[3] Islamov, et al. Loss landscape characterization of neural networks without over-parametrization. NeurIPS 2025

[4] Guille-Escuret et al. No wrong turns: The simple geometry of neural networks optimization paths. arXiv preprint arXiv:2306.11922, 2023.

[5] Schaipp et al. The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. ICML 2025

Answer to Q2: We believe there is a strong connection between our adaptive approach of setting the lr hyperparameter and lr scheduling, including warmup. According to Figure 5 in the main paper and Figures J.9-J.12 in the appendix, our proposed adaptive scheme serves as a safeguard during training. If the lr hyperparameter is too large, which might lead to unstable convergence/divergence, our scheme makes the effective lr smaller and stabilizes the training. This mechanism is similar in spirit to warm-up, but with two key differences. First, unlike traditional warm-up, which is typically applied only at the beginning of training, our approach can activate dynamically at any point during training. Second, while warm-up schedules require manual specification of the warm-up duration, our method automatically adapts based on the local curvature. We thank the reviewer for the reference and will incorporate this discussion into the revised version.

Answer to Q3: We expect NGN-MDv1 to retain the robustness properties of NGN-M in most cases. However, its robustness also depends on the choice of preconditioning. If the chosen preconditioner is highly sensitive to variations in the loss landscape or hyperparameters, NGN-MDv1's resilience to learning rate selection may degrade. In such cases, alternative preconditioning strategies, such as those based on Shampoo [6] or the Fisher information matrix [7], might be worth exploring to improve stability.

[6] Gupta et al., Shampoo: Preconditioned stochastic tensor optimization, ICML 2018

[7] Gomes et al., Adafisher: Adaptive second order optimization via fisher information. arXiv preprint arXiv:2405.16397, 2024.