Enhancing Optimizer Stability: Momentum Adaptation of NGN Step-size

W1: We would like to highlight that our convergence guarantees are provided under the assumption that the interpolation error $\sigma^2_{\rm int} := \mathbb{E}[f^* - f_i^*]$ is bounded. This assumption is satisfied if each of the used loss functions is bounded below. This is typically the case for most of the standard losses (MSE, logistic loss, cross-entropy).

Q1: Thank you for this question. Please note that the convergence of NGN-D in the general non-convex regime and under PL is provided in the appendix. The convergence guarantees match those of SGD. However, we highlight that the convergence of NGN-D is a more challenging task since it requires handling the adaptive stepsize. This challenge is overcome in the similar way as it was done for NGN-M: the effective step-size is split into two parts: a deterministic and a stochastic part. In our analysis, this decomposition of the step-size $\gamma_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 a constant $c$ while the error term proportional to $\sigma^2_{\rm int}$ is weighted by $c^2$, which suggests that $c$ can be chosen to trade off the two terms to lead to the exact convergence similarly to the standard analysis of SGD.

Nonetheless, we would like to draw the reviewer's attention to the fact NGN-M reduces to SGDM when the step-size hyper-parameter $c\to 0$ (see discussion in Section 2.2.3 in [1] for more details). Therefore, informally speaking the convergence should be expected for a small enough step-size hyper-parameter but it requires a careful analysis to demonstrate the convergence formally.

[1] Orvieto, Antonio and Xiao, Lin, An adaptive stochastic gradient method with non-negative gauss-newton stepsizes, arXiv preprint arXiv:2407.04358, 2024.

W1: We agree that deriving convergence guarantees in the general non-convex setting is of significant interest. We emphasize that the NGN-D algorithm has been analyzed in the non-convex regime, both in the general case and under the Polyak-Łojasiewicz condition. However, analyzing NGN-M in the non-convex case is considerably more challenging than in the convex setting, so we have decided to leave this analysis for future work. Nonetheless, we would like to draw the reviewer's attention to the fact NGN-M reduces to SGDM when the step-size hyper-parameter $c \to 0$ (see discussion in Section 2.2.3 in [1] for more details). Therefore, informally speaking the convergence should be expected for a small enough step-size hyper-parameter but it requires a careful analysis to demonstrate the convergence formally.

[1] Orvieto, Antonio and Xiao, Lin, An adaptive stochastic gradient method with non-negative gauss-newton stepsizes, arXiv preprint arXiv:2407.04358, 2024.

W2: Implementing NGN-MDv1 in practice might be slightly more computationally expensive. However, we highlight that computing NGN-MDv1 step 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 $\\|\nabla f_{S_k}(x^k)\\|^2_{\mathbf{D}\_k^{-1}}$. This can be done by one pass over the gradient and summing the terms $\frac{1}{(\mathbf{D}\_k)_{j}}(\nabla_j f\_{S_k}(x^k))^2$ for $j\in[d].$ This operation does not require additional matrix multiplication and can be computed while updating $\mathbf{D}\_k$. The rest of the NGN-MDv1 implementation does not add any significant costly operations in comparison with Adam. Besides, the step of NGN-MDv2 does not change much as well since it does not even require the computation of the gradient norm in the weighted norm. We add this discussion about this in the revised version in Section E.2.

W3: Thank you for raising this point. We agree with the reviewer that NLP tasks are of significant interest. We in fact already provided such experiments in the original draft and as we will discuss below, we have added additional experimental results to address this comment.

Before we elaborate further on these experimental results, it is important to highlight that training NLP models presents distinct challenges compared to vision models. Specifically, NLP tasks often require coordinate-wise adaptive step-sizes, making the NGN-M variant (as well as Momo and SGDM) less suitable for these tasks. Developing coordinate-adaptive methods is therefore a critical research direction for improving NLP model training in the future.

Following the previous argument, we demonstrate the performance of only NGN-MD against Momo-Adam and Adam in pretraining large scale language models since NGN-M would not match their best performance. We stress that the pretraining is the most resource intensive, general purpose and challenging NLP task in deep learning. Regarding other experimental results provided in the paper, we refer the reviewer to Figure 9 (additional experiments on LSTM and Transformer models), which might have been overlooked as it is located in the appendix. This figure provides a comparison of NGN-MDv1 and NGN-MDv2 against Momo-Adam and Adam, and demonstrates the good performance of the NGN variants against other methods. In addition, we have extended our experiments in the updated version of the pdf to include comparisons of NGN-M, Momo, SGDM, and NGN in training LSTM and Transformer models. The stability plots are presented in Figure 24 while the detailed discussion is given in Section G.10. Moreover, the best performance of the algorithms is reported in Table 3. From these results, we can draw two conclusions:

• Robustness to Step-Size Selection: NGN-M achieves better resilience to the step-size hyper-parameter choice for these settings than other algorithms which is in line with the results in the main paper.
• Performance of Momentum-Based Algorithms: As expected the performance of momentum-based algorithms does not match that of algorithms that use both momentum and diagonal step-size. This is due to the significantly different conditioning across coordinates in NLP tasks, emphasizing the need for further research into advanced adaptive coordinate-wise step-size methods.

These results underscore the strengths of the NGN variants while highlighting opportunities for future work in the area of NLP models.

Dear reviewer,

We would like to remind you that the discussion period ends soon. Therefore, we would like to know if there are concerns left unaddressed or needed to be clarified. We would be happy to discuss them further.

W1: We believe the reviewer interpreted “stability” differently from how we intended. Specifically, we do not mean that the algorithm employs large step-sizes. Effectively, our method gives an additional adaptivity boost compared to baselines. The method automatically detects if the training goes wrong and implements a lower effective step-size. It is an additional safeguard that is helpful and does not destroy performance.

We emphasize that our focus is on stability from the perspective of step-size hyper-parameter tuning: the performance of our proposed algorithms is less sensitive to the choice of this hyper-parameter compared to baselines such as Adam. This notion of stability is the same as used in prior work (see references cited in the paper). Importantly, the optimal step-size hyper-parameter for Adam varies a lot from one task to another, especially when the domain of data changes (e.g., from vision tasks to NLP). This implies that determining the optimal step size in practice can be cumbersome for users, often requiring significant time and extra computations (which results in significant costs financially). From our experiments, NGN-MDv1 requires less tuning of the step-size hyper-parameter. We will clarify what we exactly mean in the revised version. Moreover, we emphasize that the role of the step-size hyper-parameter $c$ in NGN-MDv1 is the maximum allowed effective step-size. We observe that the peak value of the effective step-size does increase when increasing $c$, and then it gradually decreases at the end of the training. These observations are in line with the learning rate schedulers that are used in practice.

W2: Thank you for pointing this out. We will revise this section to make our explanation clearer. We address your comments in two parts:

Regarding the connection to prior work: first, we agree that the observed phenomenon is strongly related to training at the Edge of Stability (EoS), as explored in [1] and other studies. However, we emphasize that [1] focuses on non-adaptive methods, both with and without momentum. The only work we are aware of that examines EoS behavior in adaptive methods is [2]. According to [2], Adam operates at an adaptive EoS (determined by the eigenvalues of the preconditioned Hessian), even as standard sharpness continues to increase throughout training. Our findings indicate that NGN-M operates at the Edge of Stability (EoS), despite employing adaptive step sizes. This discussion is reported in Section G.8 of the revised version of the paper.

Regarding the current formulation in Section 5.5: We aimed to demonstrate why a large step-size hyperparameter does not negatively impact the performance of NGN-M. Specifically, SGDM fails to converge with a large step-size hyperparameter, diverging when operating beyond the EoS range. In contrast, NGN-M converges under the same step-size hyperparameter due to the adaptive nature of the NGN step size. We believe that the adaptivity of NGN-M enables it to operate effectively within the EoS range, even when the step-size hyperparameter is large. This is what we believe what makes NGN an interesting optimizer to consider in practice.

[1] Cohen et al., Gradient descent on neural networks typically occurs at the edge of stability, arXiv preprint arXiv:2103.00065, 2021.

[2] Cohen et al., Adaptive gradient methods at the edge of stability, arXiv preprint arXiv:2207.14484, 2022.

W3: We thank the reviewer for pointing out the typo. We also changed the iteration counter of the momentum term $m^k$, so that the new iterate $x^{k+1}$ is computed based on the terms from iteration $k$ or before: $x^k, m^{k-1}, \gamma_k, \nabla f_{S_k}(x^k)$. All these changes will be reported in the revised version of the paper.

Q1: We would like to clarify that we only compare NGN-MD against methods with adaptive diagonal step-size and momentum: Momo-Adam and Adam, in exactly the same hyper-parameter scheduling conditions. The only difference is in the particular way how each algorithm computes the adaptive step-size. Therefore, we believe that the provided set of experiments is fair as we compare how the particular choice of adaptive step-size influences the performance. The learning rate scheduler can be incorporated into any of the algorithms we consider. However, doing so may obscure the true differences in adaptive step-size behavior. Introducing a scheduler also adds another layer of complexity, as it requires separate tuning for each task. While there are guidelines available for selecting an appropriate scheduler, these choices may still be suboptimal so we feel this would make the result less trustworthy. We are happy to discuss further if the reviewer disagrees.

Dear reviewer,

We would like to remind you that the discussion period ends soon. Therefore, we would like to know if there are concerns left unaddressed or needed to be clarified. We would be happy to discuss them further.

W3: We will add the names of the algorithms in the caption of each plot in Figure 5 to make the figure more readable. In Figure 5, the two left plots correspond to the comparison of Momo and NGN-M while the two right plots correspond to the comparison of Momo-Adam and NGN-MDv1. We would like to highlight that the performance of the Momo algorithm for the (ViT, CIFAR10) setting is presented in the top row of Fig. 2.

Following the request from the reviewer, we provide the comparison of the adaptive step-size of Adam, Momo-Adam, and NGN-MDv1 in Section G.9 for two different settings: (ResNet20, CIFAR10) and (ViT, CIFAR10); see a more detailed discussion there as well. Note that Adam's adaptive step-size relies solely on normalization using the second-order momentum while both Momo-Adam and NGN-MDv1 have an additional step-size ($\tau_k$ and $\gamma_k$ respectively) factor in addition to the normalization. Therefore, we decided to compare $\frac{\gamma}{(\mathbf{D}\_k)\_{(j)}} \text{ for Adam},\quad \frac{\tau_k}{(\mathbf{D}\_k)\_{(j)}} \text{ for Momo-Adam}, \quad \frac{\gamma_k}{(\mathbf{D}\_k)\_{(j)}} \text{ for NGN-MDv1},$ where $j\in[d]$ corresponds to some specific parameter of the model. We visualize the average quantities defined above for some specific layers of models (i.e., we average the effective step-size of algorithms for all $j$ within some layer). In particular, we present the results for the first convolution layer of each base block of ResNet20, and for the attention layer of $1,3,6$ base blocks of ViT model. We observe that in both experiments the effective coordinate-wise step-size of NGN-MDv1 is smaller than for the other two optimizers. In other words, the adaptive step-size of NGN-MDv1 is more conservative and does not allow the effective step-size to increase too much even when the step-size hyper-parameter is set to be large. This demonstrates that the NGN-MDv1 is less sensitive to the choice of the step-size hyper-parameter while allowing it to reach comparable or superior performance to other optimizers.

Q1: In fact, the main difference between the NGN and SPS${}\_{\max}$ step-sizes is that the NGN step-size is a harmonic mean between the constant step-size of SGD and SPS. On the opposite, SPS${}\_{\max}$ is the minimum between the constant step-size of SGD and SPS. The harmonic averaging can be seen as a soft version of the minimum. Although the difference does not appear to be highly significant at first, it allows NGN to be both theoretically and practically superior. We refer to [1], sec. 2.2.3 for a more detailed discussion on the differences between NGN and SPS step-sizes. We exclude a comparison against MomSPS as its best performance is almost always worse than those of NGN-M and Momo (see Tab 3). Moreover, the Momo framework is based on the SPS step-size, and we found it to be a better alternative to MomSPS across all considered tasks.

[1] Orvieto & Lin, An adaptive stochastic gradient method with non-negative gauss-newton stepsizes, arXiv preprint arXiv:2407.04358, 2024.

W1: Thank you for this insightful comment. We plan to include a more detailed comparison between the two versions of NGN-MD in the revised manuscript. Below is the discussion we incorporated in Section E.1.

Both algorithms use the RMSprop preconditioner $\mathbf{D}\_k$ that performs an exponential moving average of the coordinate-wise squared gradient norm. Moreover, both algorithms use momentum to perform an averaging of the preconditioned updates $\boldsymbol{\Sigma}_k^{-1}\nabla f\_{S_k}(x^k).$ Nonetheless, there are two key differences between the proposed algorithms. First, $\mathbf{D}\_k$ is used as a preconditioner directly in NGN-MDv1 while in NGN-MDv2 it is used to rescale the $c$ constant for each coordinate inside coordinate-wise NGN step-size. Second, NGN-MDv1 uses one global NGN step-size weighted by the $\mathbf{D}_k$ norm while in NGN-MDv2 we use coordinate-wise NGN step-size replacing the full gradient by the corresponding partial derivative. According to the empirical results, both versions with a tuned step-size hyperparameter are competitive with other baselines but NGN-MDv1 demonstrates much better stability performance.

The main difference in comparison with Adam is the order in which the preconditioning and momentum are applied. In both NGN-MDv1 and NGN-MDv2 we average the preconditioned updates $\boldsymbol{\Sigma}\_k^{-1}\nabla f_{S_k}(x^k),$ i.e. we first apply preconditioning and momentum later. In contrast, in Adam and Momo-Adam the stochastic gradients are averaged to construct a 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 hyper-parameter resilience of NGN-MD.

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 hyper-parameter than NGN-MDv2.

W2: We thank the reviewer for this comment. We rewrote this section in the revised version. Nonetheless, we would like to clarify the raised concerns. The derivations in (2) are used to provide an intuition of how one can add a diagonal step-size into NGN by choosing the weight matrix $\boldsymbol{\Sigma}\_k$. These derivations are used exactly in the NGN-MDv1 algorithm with $\boldsymbol{\Sigma}\_k = \mathbf{D}\_k$ where $\mathbf{D}\_k$ is a RMSprop preconditioner. In this case, we have only one global NGN step-size in front of $\mathbf{D}\_k$. For NGN-D the derivations follow a more straightforward intuition. We can update each parameter $j$ using the coordinate-wise NGN step-size where the gradient norm $\\|\nabla f\_{S_k}(x^k)\\|$ is replaced by the corresponding partial derivative $|\nabla\_j f\_{S_k}(x^k)|.$ Namely, each coordinate is updated using $\gamma_k^{(j)} = \frac{c}{1+\frac{c}{2f_{S_k}(x^k)}(\nabla_j f_{S_k}(x^k))^2}$. To derive the NGN-MDv2 algorithm from NGN-D, we observe that each parameter requires a coordinate-wise NGN step-size hyperparameter $c$. To achieve this, we use the RMSprop preconditioner to set the coordinate-wise NGN step-size as $c/(\mathbf{D}\_k)\_{(j)}$. Incorporating a momentum on top leads to the NGN-MDv2 algorithm.

We thank the reviewer for their response. Regarding the second point, we would be grateful if the reviewer could provide further explanation regarding what he expected from the experiments in Section G.9. We clearly observe that NGN stepsize plays the role of the safeguard that is helpful and does not destroy performance. The method automatically detects if the effective step-size is too large and decreases it.