NAG-GS: Semi-Implicit, Accelerated and Robust Stochastic Optimizer

Dear reviewer f3XV,

thanks for your comment on our manuscript! Below we have highlighted the main contribution and novelty of our study:

1. We introduced Gauss-Seidel discretization of the second-order ODE related to the accelerated gradient method in the stochastic setup
2. We demonstrated numerically the expected robustness w.r.t. the large step size. Note that at the end of Appendix 1, we detail an additional yet important motivation for proposing a solver capable of handling a broad range of step size values. Specifically, allowing for a larger ratio $\alpha/b$ (b as the mini-batch size) increases the likelihood of converging to wider local minima, ultimately enhancing the generalization performance of the trained model.
3. We proved the convergence theorem in the quadratic case and obtained the optimal learning rate that maximizes the contraction and therefore corresponds to the fastest convergence
4. We presented an extensive benchmarking of NAG-GS in the stochastic setup, where different tasks from computer vision and natural language processing domains are considered. In the presented experimental results, NAG-GS shows superior or on-par performance compared with competitors.

In addition, we have carefully proofread the manuscript and fixed typos.

The second part of the comment

High-Resolution ODE Framework in a nutshell: Introduced by Shi et al. (2019), the framework distinguishes between NAG-SC and Polyak's heavy-ball method, revealing a unique "gradient correction" term for NAG-SC. Incorporating O(√s) terms in the limiting process yields high-resolution ODEs, offering more accurate continuous-time representations.

ODE Formulations: Here-under, we recall two key ODEs which are presented for the heavy-ball method and NAG-SC, highlighting distinctive terms:

The high-resolution ODE for the heavy-ball method: $\ddot{x}(t) + 2 \sqrt{\mu} \dot{x}(t) + (1+\sqrt{\mu s}) \nabla f(x(t)) = 0$

The high-resolution ODE for “NAG-SC”: $\ddot{x}(t) + 2\sqrt{\mu} \dot{x}(t) + \sqrt{s} \nabla^2 f(x(t)) \dot{x}(t) + (1+\sqrt{\mu s}) \nabla f(x(t)) = 0$

According to Shi et al. (2019), these ODEs, with specific initial conditions, provide a better characterization of accelerated methods in continuous time.

Numerical Stability and Step Size: Section 3.3 of Shi et al. (2019) addresses numerical stability, deriving bounds for step size (s) with a forward Euler scheme. A comparison reveals that the Hessian makes the forward Euler scheme for NAG-SC ODE numerically stable with a larger step size (s = O(1/L)). -> Clearly here, there is a strong motivation for a semi-implicit approach to ensure a larger step size than pure explicit discretizations.

However, we would like to raise the following points:

Considerations in Noisy Settings: The previous considerations are relevant in a deterministic setup, but we need for further analysis in noisy gradients. A detailed examination, especially for quadratic forms minimization in noisy settings, could provide insights into step size choices. Furthermore, simple numerical tests as the ones done in Appendix 1.1.5 could be helpful to gain some insights.

Transformation and Condition Number: In our paper, we considered one transformation to derive accelerated ODE, which relies on the embedding of matrix A into a 2 × 2 block matrix G, to reduce the condition number and enable the use of larger step sizes. We do not know if, in the end, using semi-implicit discretization for NAG-SC high-res ODE would offer a similar range for step size choices.

Gradient Corrections and Second-Order Information: In Shi et al. (2019), authors highlight the importance of the good estimation of Hessian-vector products (the gradient-correction terms) in NAG-SC high-res ODE, imposing step size s is not too big, with a claim that even small gradient corrections have a fundamental effect on the scheme's behavior. Moreover, their gradient corrections bring in second-order information from the objective function, which could be an interesting way to develop efficient methods using higher-order methods without explicitly building the Hessian and inverting as it is the case for Newton-steps based methods. Furthermore, in practice, efficient routines exist to accurately compute such Hessian-vector products even in large-scale problems. But again, in the stochastic setting, for instance using stochastic estimates of the Hessian, or stochastic estimates of the Hessian-vector products on a mini-batch, we do not know yet how this approach could behave overall and if stochastic estimates of gradient-corrections terms still bring something. Among many interrogations, we would also be interested to see what would happen close to stationarity. Using additional information on the function to minimize, such as the Hessian, may indeed speed up the converge of the method in the early stages (killing the gradient faster:), that is when the noise is low compared the gradient in terms of “norm”, but what happens when the noise counterpart becomes in the same order of magnitude than the gradient ?

We are eager to further discuss these ideas with the reviewer and appreciate the inspiration provided.

LINKS for the additional experiments http://bitly.ws/wSbz http://bitly.ws/wSbC http://bitly.ws/wSbH

Dear reviewer 13Pv, thanks for your thorough comments and remarks! Below we provide requested clarifications and address the mentioned issues. The main text is also revised respectively.

Weaknesses

• The theoretical analysis is bounded to quadratic case. This is not mentioned accurately in the contribution. Specifically the second main contribution is: We analyze the properties of the proposed method both theoretically and empirically;.

Answer: Thank you for your comment, we have clarified our theoretical contribution for the quadratic case in the main text.

• The introduction is not really an introduction of this work. By just reading the introduction, it is not possible to get an accurate idea of what differences this work has with any other work in "stochastic optimization algorithms".

Answer: We place in the Introduction section the main contributions and the organization of the manuscript that provides step-by-step plan on how we analyze the proposed dmethod theoretically and empirically.

• "The Preliminaries" is a mixture of background, related work and notations. Here, more organization might improve readability.

Answer: thanks for the suggestions! We have slightly revised this section. The main intention here was to briefly introduce the reader to the used concept of ODE-inspired optimizers and clarify the basics necessary for the further development of NAG-GS.

• Gaus-Seidel splitting used here is not a novel idea in discretizing ODEs for acceleration as it was previously discussed in [Luo & Chen (2021)].

Answer: Thank you for your valuable comment. It's acknowledged that deterministic Gauss-Seidel splitting in the context of accelerating SDEs was previously discussed by Luo & Chen (2021). In our work, we introduce novel contributions, particularly in the quadratic case: Extension of Deterministic Results: We provide a detailed analysis of the optimal step size computation, which exceeds the limit proposed by Luo and Chen (refer to Lemma 1 in our paper). Additionally, we identify a scenario (μ > L/2) where we can freely choose the step size, offering a new perspective.

Stochastic Setting under Gaussian Noise: We extend our approach to the stochastic setting under a Gaussian assumption for the gradient noise, introducing a family of accelerated SDEs (Equation 11). We prove the asymptotic convergence of NAG-GS with the same optimal step size as in the deterministic case. We provide optimal choice for the step size leading to the highest possible contraction rate when aiming at minimizing a strongly convex quadratic form. Furthermore, our numerical experiments demonstrate the expected robustness concerning large step sizes across a diverse range of problems.

Additionally, in Appendix 1.2, we explore a fully implicit version. Our findings reveal that, in the stochastic case for quadratic minimization, we can select any step size, mitigating the noise effect. This implies that the spreading of particles can be arbitrarily close to the arg min, showcasing a notable result.

References: Alvarez, F., Attouch, H. An Inertial Proximal Method for Maximal Monotone Operators via Discretization of a Nonlinear Oscillator with Damping. Set-Valued Analysis 9, 3–11 (2001).

Dear Reviewer nCbP,

Thank you for your insightful feedback on our manuscript. We appreciate the opportunity to address the concerns raised and have made several amendments to our paper to reflect these points.

1. Comparison with Other Methods:

   • We recognize the importance of comparing NAG-GS with existing methods like Polyak's momentum method and Triple momentum method. To this end, We have checked the mentioned references and updated Section 4. Both ITEM and Triple momentum do not study the robustness w.r.t, the large step size and consider only deterministic setup. The incorporting these methods in the ODE-inspired setup is the interesting future work. This comparison highlights the unique advantages of NAG-GS, particularly in its original setting. We added additional figures 15 and 16 to the section "Numerical tests for quadratic case" in supplementary materials, which illustrates the fastest asymptotic convergence rate of the NAG-GS for the considered quadratic problems. The figures are also available here 15 and 16. To illustrate the comparison the NAG-GS with Polyak’s momentum metod we also have updated Figure 1.

2. Experimental Improvement and Deviation Statistics:

   • We acknowledge that the improvement in neural network experiments appeared marginal and lacked deviation statistics. To address this, we have now included a set of experiments of NAG-GS and AdamW on the ViT training problem on food-101 dataset with the best hyperparameters found with the hyperparameter search on the portion of data. The number of experiments per method is 5. Mean values and standard errors of the evaluation loss are presented. . The results are included in the section Additional experiments with ViT in Appendix. Link for eval accuracy graph Link for eval loss graph

3. Optimal Learning Rate Discussion:

   • Based on your suggestion, we have included Remark 2 in the main text, which explicitly discusses the derived optimal learning rate for NAG-GS. This addition emphasizes one of the key features of NAG-GS and its relevance in the context of gradient methods.

4. Further Empirical Comparisons and Future Work:

   • We have also updated Figure 1 to illustrate a direct comparison of NAG-GS with Polyak’s momentum method. This visual representation aids in understanding the relative performance of these methods.
   • Recognizing the need for broader empirical evidence, we have conducted additional experiments using Vision Transformers (ViT) and provided results with error bars for a more nuanced evaluation.
   • We see the integration of methods like ITEM and Triple momentum in an ODE-inspired setup as an exciting avenue for future research and have mentioned this as part of our conclusion and future work section.

We believe these revisions and additions significantly strengthen our manuscript and address the concerns you have raised. We are grateful for the opportunity to enhance our work based on your feedback and hope that these changes meet your expectations for a robust and comprehensive analysis.

Dear Reviewer MgiM,

Thank you for your insightful feedback on our paper. Below we provide detailed responses to the mentioned weaknesses and answer the questions.

Convergence Analysis and Empirical Results: We acknowledge that extending analysis to more general functions introduces additional complexities and is the question for the further research. Regarding the requirement of prior knowledge of $\mu$, we approached this by aligning the weight decay parameter used in standard optimizers (like SGD, Adam, etc.) with the strongly convex constant $\mu$. Consequently, in our experiments, we set the constant for weight decay in competitor algorithms and for the $\mu$ in NAG-GS correspondingly.

Algorithm 1 and Non-Convex Losses: In practical applications for non-convex losses, we derived the $\mu$ parameter from the weight decay settings commonly employed in standard optimizers. This approach allowed us to integrate our method into existing frameworks without additional complexity. For some experiments (for example, comparing AdamW and NAG-GS on the ViT experiment on food101) we initially performed hyperparameter search for all hyperparameters for the methods on the portion of the dataset and compared them with the best choice for both.

Clarifications and Typos: We have made several updates to our manuscript to clarify points you've raised:

• Definition of A-stability: We added a high-level definition and reference in the main text and detailed insights in Appendix 1.3.
• Clarification of $\mathcal{R}(\lambda)$: We now specify that $\mathcal{R}(\lambda)$ refers to the real part of $\lambda$.
• SGD-MW: We clarified that this stands for Stochastic Gradient Descent with Momentum and Weight decay, and added this explanation before its first use in the text.
• Baseline Algorithms: The choice of SGD-MW as a baseline for ResNet-20 was due to its comparable memory footprint to our NAG-GS. In contrast, Adam requires two state vectors, instead of one in case of NAG-GS ad SGD-MW. Additionally, the accelerated gradient descent in Figure 1 references the Nesterov accelerated method variant presented in Su et al. (2014).
• Typographical Corrections: We corrected the typo you pointed out and carefully reviewed the manuscript for any other errors.

Additional Insights on Stability and Step Size: $A$-stable methods, exemplified by the implicit Euler scheme, exhibit a notable advantage in that they impose no inherent limitations on the integration step (or step size). On the other hand, fully explicit methods, like the forward Euler scheme (which essentially translates to gradient steps), necessitate certain bounds on integration steps to ensure stability and asymptotic convergence. To strike a balance, we explore a semi-implicit scheme, aiming to permit a broader range of values for the learning rate (or step size) compared to pure explicit discretization. Importantly, this is achieved without incurring the computational cost associated with the fully-implicit approach.

For a comprehensive understanding of this aspect, please refer to the additional insights provided in the new Appendix 1.3 and Appendix 1.2, the latter of which is dedicated to the fully implicit discretization of the proposed accelerated SDE (Equation 11).

In the specific case of the quadratic scenario with $\gamma=\mu$, we can theoretically establish that NAG-GS allows for an optimal step size (($\alpha = \frac{2 \mu + 2 \sqrt{\mu L}}{L-\mu}$)), as detailed in Theorem 1 and Remark 2. This optimal step size is strictly higher than the maximum step size of $\frac{2}{L}$ (itself higher than $\frac{2}{\mu + L}$) permitted by the traditional Gradient Descent (GD) method for convex (or $\mu$-convex) $L$-smooth functions in a deterministic setting, under certain conditions. These conditions include $L > 0$ and ( $\frac{L}{(L+1)^2} < \mu < L$). Empirically, our findings demonstrate that NAG-GS facilitates the utilization of a broad range of step sizes, proving effective across a diverse set of problems when compared to AdamW and SGD with Momentum.

Last but not least, in the end of Appendix 1, we detail an additional yet important motivation of proposing a solver capable of handling a broad range of step size values. Specifically, allowing for a larger ratio $\frac{\alpha}{b}$ ($b$ as the mini-batch size) increases the likelihood of converging to wider local minima, ultimately enhancing the generalization performance of the trained model.

Future Directions: Regarding future work, we have expanded Section 5 of our manuscript to outline potential directions. These include extending our convergence analysis to broader classes of functions and investigating the applicability of our approach to different machine learning models and problem domains, especially for gradient noise with unbounded variance.

Thank you once again for your constructive comments, which have greatly contributed to improving the quality of our paper.