SANIA: Polyak-type Optimization Framework Leads to Scale Invariant Stochastic Algorithms

Authors: Farshed Abdukhakimov, XIANG CHULU, Dmitry Kamzolov, Robert M. Gower, Martin Takáč
Primary area: optimization
Keywords: optimization, learning rate free, adaptive optimizers, polyak step-size

Abstract:
Adaptive optimization methods are widely recognized as among the most popular approaches for training Deep Neural Networks (DNNs). Techniques such as Adam, AdaGrad, and AdaHessian utilize a preconditioner that modifies the search direction by incorporating information about the curvature of the objective function. However, despite their adaptive characteristics, these methods still require manual fine-tuning of the step-size. This, in turn, impacts the time required to solve a particular problem. This paper presents an optimization framework named SANIA to tackle these challenges. Beyond eliminating the need for manual step-size hyperparameter settings, SANIA incorporates techniques to address poorly scaled or ill-conditioned problems. We also explore several preconditioning methods, including Hutchinson's method, which approximates the Hessian diagonal of the loss function. We conclude with an extensive empirical examination of the proposed techniques across classification tasks, covering both convex and non-convex contexts.

Decision: reject
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730399607876

thread_type:
official_review

<thread_object>
creator: reviewer_icar

<content>
summary:
This paper presents SANIA, which is a method that doesn't require manual fine-tuning of the learning rate in commonly used stochastic optimization algorithms, leading to faster optimization. The authors present a framework which generalises common stochastic optimization algorithms. The authors also consider affine and scale invariance which seeks to address poorly scaled or ill-conditioned problems. The authors compare performance of SANIA with commonly used algorithms (which have had fine-tuned parameters) for training classifiers for MNIST, FashionMNIST, CIFAR10 and SVHN.

strengths:
Provides a novel formulation of stochastic optimization algorithms and a novel method that does not require fine-tuning of the learning rate parameter. Investigation into scale invariance is novel. Theorems, equations and ideas are presented clearly. Numerical experiments show that the SANIA methodology achieves similar performance to algorithms with fine-tuned learning rates, and improved robustness as the training curves fluctuate less.

weaknesses:
Numerical experiments lack a comparison to other options for tuning-free methods.

Minor comments:

Occasional incorrect grammar, including after equation (7): "This leads us to Stochastic Polyak step-size method." should read "This leads us to the Stochastic Polyak step-size method." Also bad grammar in the statement of Theorem 1. "Another way to derive this formulation is by solving" (4), I think it may be better to reference appendix B.2 to make it clear why this holds. "interpolation condition" isn't very clearly defined in my opinion, it may be better to be more clear. "in practice it displays better convergence to the true Hessian than other
similar methods like BFGS" - I think this needs a reference.

questions:
In Page 2, SPS has been derived, why not just cite it? The difference between g_t and m_t is unclear to me, these seem to often mean the same thing? Why not keep the notation consistent? In (6) you have written \| w \|_{B_t} is a Euclidean norm, would it be better to say " \| \cdot \|_{B_t} is a Euclidean norm"?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730513305200

thread_type:
official_review

<thread_object>
creator: reviewer_zkyn

<content>
summary:
The work tries to incorporate a Polyak-type optimization framework with existing optimizers for machine learning. In the new method, the optimization direction and step size are spontaneously determined by a simple local optimization for each step. The new method is not difficult in realization and appears to work. From the numerical results, the test accuracy of network trained with the new method is comparable to existing methods (but not better than them).

strengths:
The new method can be naturally incorporated with existing methods and the realization is not difficult.

weaknesses:
From the numerical methods, the new method does not show improvements compared to previous methods.

questions:
1. Instability of the Euler method and noises are believed to be helpful for generalization (see Ref[1]). Relatively large step size is helpful for the training process to jump out of bad local minima. In the new framework, when the steps are spontaneously obtained, the instability are also almost inhibited spontaneously. This may be a severe problem for the new frame work. 

2. From the numerical results, it seems that the test accuracy of Sania is slightly worse than existing methods. The author should also compare the results with SGD, since it usually has competitive generalization ability.

3. In Line 209, what is the meaning of "Otherwise, we replaced step-size parameter γt to parameter fi∗?". The authors should explain more clearly.

4. Due to the anisotropic property of the loss landscape, the estimate f* may not be a good guess for the local optimization. How will this affect the performance of the new method?

[1] The Implicit Regularization of Dynamical Stability in Stochastic Gradient Descent, Lei Wu, Weijie J. Su, ICML 2023;
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706344210

thread_type:
official_review

<thread_object>
creator: reviewer_ztra

<content>
summary:
This paper introduces cubic Newton method with adaptive Polyak step-size, enhancing robustness and efficiency in non-convex tasks. It also proposes scale-invariant variants of AdaGrad and Adam, which improve optimizer performance on poorly scaled data. Extensive experiments validate the effectiveness of these methods across diverse optimization scenarios.

strengths:
1. Thorough theoretical analysis.
2. Clear writing

weaknesses:
1. It would be great to have theoretical/practical comparisons with Sophia [1], which also uses Hutchinson's based method to compute their preconditioner.
2. KATE [2] removes square root to ensure scale invariance of Adagrad. It would be great to have theoretical/empirical comparison with KATE. 

I am willing to improve the score if the above are addressed.

[1] Liu, Hong, et al. "Sophia: A scalable stochastic second-order optimizer for language model pre-training." arXiv preprint arXiv:2305.14342 (2023).
[2] Choudhury, Sayantan, et al. "Remove that Square Root: A New Efficient Scale-Invariant Version of AdaGrad." arXiv preprint arXiv:2403.02648 (2024).

questions:
What are your thoughts on how learning rate schedules such as cosine decay etc compose with $\lambda_t$ schedule defined in (19)?

How does $\lambda_t$ change during training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731425137660

thread_type:
official_review

<thread_object>
creator: reviewer_jmo1

<content>
summary:
This paper proposes a training algorithm generalizing the "Polyak step-size" to second-order algorithms, such as cubic Newton or quasi-Newton. 

Let $f$ be a function to minimize and $w\_t$ the current estimate of the argmin.
For an order-1 method, the Polyak step-size comes down to choosing the step-size such that we would attain the minimum $w^*$ in one step if $f$ was affine between $w\_t$ and $w^*$. For order-2 methods, the authors distinguish the metric used for the parameter space, denoted by $B_t$, and the metric related to the curvature of the local approximation of $f$, represented by a matrix $D_t$.

Then, the authors propose variations of existing algorithms based on their generalization of Polyak step-size.

strengths:
# Originality

As far as I know, the work presented in this paper is original.


# Clarity

Overall, the paper is easy to follow.


# Quality

I am very grateful to the authors for providing additional experiments (Appendix E), which show not only test accuracy, but also test loss. I would have been even better to provide the training loss, since the proposed algorithms have been designed to improve the optimization process (not the generalization).

Overall, the paper seems to be correct.

# Significance

This paper provides practical uses of *SP2: A Second Order Stochastic Polyak*, Li et al., 2023. 
So the present paper is relevant for the community.

weaknesses:
# Originality

This paper can be seen as a practical application of *SP2: A Second Order Stochastic Polyak*, Li et al., 2023.


# Clarity

At the end of Section 2, the authors use the notations: "SANIA $I_d$", "SANIA $(V^{-1})^2$" and "SANIA $\mathrm{diag}(H^{-1})$", while writing that there is no preconditioning. Thus, I understand that "SANIA $A$" refers to Eqn. (6) with $D_t = A$. But there is no formal definition of "SANIA $A$". It should appear somewhere.

Just above Eqn. (9), one can read: $m_t = m_t$, which should be "$m_t = g_t$" (?).

Several mistakes, Figure 1, first plot: 
 * the legend is difficult to understand, since several curves have the same label;
 * some curves do not seem to correspond to their label: apparently, light green should be "SANIA $(V^{-1})^2$";
 * the purple curve is dotted, while it is not dotted in the legend.

questions:
Could the authors provide a formal link between scale invariance and using SANIA? It seems that there is some overlap between the two, but SANIA does not ensure scale invariance by itself. It would be interesting for potential users of SANIA to provide conditions for their algorithm to be scale-invariant.

Addition to Fig. 5: how does SANIA compare to concurrent methods in terms of training loss?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732386659546

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
REBUTTAL TO ALL REVIEWERS

comment:
We would like to sincerely thank all the reviewers for their careful reading of our submission and their thoughtful and constructive feedback. We greatly appreciate the time and effort spent on reviewing our work and for providing suggestions that will help improve our paper.

We are grateful that all reviewers recognized several key strengths of our paper. Reviewer jmo1 appreciated the originality of our work and acknowledged that it is relevant to the community, particularly highlighting the practical application of the generalized Polyak step-size in second-order methods. Reviewer zTra commended the thorough theoretical analysis and clear writing, noting the robustness and efficiency of our proposed method across diverse non-convex optimization scenarios. Reviewer zkYN emphasized the ease of integrating our framework with existing optimizers and mentioned the excellent presentation of our ideas. Finally, Reviewer iCAR highlighted the novelty of our formulation, especially the investigation into scale invariance, and acknowledged the robustness of SANIA in comparison to fine-tuned algorithms. These strengths collectively underscore the significance and potential impact of our contribution to the field.

In response to the issues raised, we have carefully addressed all the points mentioned in the reviews, including clarifications, additional discussions, and improvements to the experimental section. We have uploaded a revised version of the paper, which includes these updates and directly addresses all the comments and suggestions. We hope that these changes address your concerns and improve the overall quality of our submission.

Thank you again for your valuable feedback and suggestions. We are confident that these improvements have strengthened our paper significantly.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733175509974

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
TO ALL REVIEWERS

comment:
Dear Reviewers,

With the discussion period extended, we would like to take this opportunity to further refine and improve our paper based on your insights. We have thoroughly addressed all your questions and concerns but feel the discussion has not yet reached its full potential. 

Could you kindly advise us on any additional steps we could take for you to consider revising your score?

We deeply appreciate the time and effort you have dedicated to reviewing our work and value your feedback greatly.

Best regards,
Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies a general framework for Preconditioned and Second-order Polyak methods. In particular, the paper proposes a first Stochastic Cubic Newton method with polyak step-size and also introduces the new scale invariant versions of AdaGrad and Adam, which make them invariant to some basis transformations. The paper also provides a few experiments to demonstrate the performance of the algorithm. The improvements (if any) are not very significant but the authors highlight that the key benefit of this work is need for  limited tuning of hyperparameters.

The reviews for the paper are very much on the borderline. The reviewers while acknowledging that the work is interesting also highlighted the weakness of lack of comparison with schedule free optimization and somewhat limited empirical evidence. After reading through the paper, I was extremely disappointed with the empirical evidence and I believe these weakness are very important to address before publication. The paper's core claim is about the eliminating the need for manual stepsize hyperparameter setting, however the experiments are done on extremely small settings (where hyperparameter tuning is of not much value) and lack proper comparisons. I think this paper would greatly benefit from thorough empirical analysis on larger settings (e.g. ImageNet, ResNet, VIT). I recommend rejection in the current form.