Sassha: Sharpness-aware Adaptive Second-order Optimization with Stable Hessian Approximation

We really appreciate the reviewer’s feedback. While we address the reviewer’s specific comments below, we would be keen to engage in any further discussion.

---

“ViTs should be optimized with Adam or Shampoo not with SGD”

There seems to be some confusion. We have already provided results for ViT trained with AdamW in Table 2.

---

“a very baisic resnet can get 94%”

We would appreciate it if you could provide a specific reference for the claimed 94% performance of a basic ResNet. The paper we are referencing [1] reports results similar to ours.

---

“This paper is basically just a different version of Flat Sophia” [2]

Thank you for providing a pointer. With all due respect, however, we disagree with this assessment and find it to be quite stretched, if not unfair.

First, Sassha operates very differently from Flat-Sophia, notably in its sharpness minimization approach [3,4] and second-order techniques (see Appendix B), and such differences are known to yield distinct optimization behavior [5].

More importantly, however, Sassha is supported by a comprehensive empirical and theoretical study that offers insights and presented as a form of research paper, rendering potentially greater value in a significantly more reliable fashion compared to the suggested pointer, which seems to be a code repository in an idea level without providing any numerical result [2].

We sincerely hope that the reviewer sees the value we created in this work to architect a generally well-performing method for standard deep learning settings.

---

“In RL many people have combined SAM and 2nd order opts to boost neuro plasticity. This paper is not novel in this.”

Thank you for your comment. We would appreciate it if you could point us to specific references. This would allow us to more accurately compare under fair settings and discuss their relations to our work.

In the meantime, we hope that the reviewer sees the contributions we made in this work given that the method is rigorously evaluated across diverse standard settings to set a new state of the art, and that leveraging general ideas (i.e., sharpness minimization and second-order optimization) should not necessarily be used as a ground for criticism.

 

Reference

[1] Yao et al., ADAHESSIAN: An Adaptive Second Order Optimization for Machine Learning, AAAI, 2021.
[2] https://github.com/evanatyourservice/flat-sophia.
[3] Wang et al., Improving Generalization and Convergence by Enhancing Implicit Regularization, NeurIPS 2024.
[4] Foret et al., Sharpness-Aware Minimization for Efficiently Improving Generalization, ICLR, 2021.
[5] Dauphin et al., Neglected Hessian component explains mysteries in sharpness regularization., NeurIPS 2024.

We’re sincerely grateful for the reviewer’s thoughtful engagement and recognition of our contribution. It was encouraging and helped us further refine the work. We provide our responses below and welcome further discussion.

---

Justification for design choices

Thank you for your comment. We believe square-rooting can be potentially more effective than clipping or damping for two main reasons. First, square-rooting preserves the relative scale between each element of the Hessian. This property is particularly valuable when the sharpness minimization is underway, where the overall Hessian values tend to be small. In such cases, even small differences between Hessian elements may carry nontrivial curvature information. Square-rooting can help retain this relative structure while also mitigating numerical instability caused by underestimated curvature. In contrast, both clipping and damping operate by abruptly replacing Hessian values based on a predefined and fixed threshold criterion. As a result, when the Hessian is generally small due to sharpness minimization, informative dimensions may fall below the threshold, removing potentially critical variations and hence deteriorating the quality of preconditioning. This behavior can also make the method more sensitive to hyperparameter choices.

Second, square-rooting can further be interpreted as a geometric interpolation between a Hessian-based preconditioner and the identity matrix, which, as theoretically analyzed in [1], provides a natural mechanism for balancing between bias and variance.

Additionally, clipping mechanism in Sophia has been shown to behave like signSGD on certain parameters, which prior studies [2,3] suggest may lead to suboptimal performance depending on the architecture.

While developing a precise account behind the benefit remains a challenge, we hope the reviewer sees our contributions in this work given that the proposed method is rigorously evaluated across diverse under fair settings and sets a new state of the art.

---

Comparison with SAM” section

We apologize for the confusion by the current wording of the section.

First, the main goal of Section 5.3 is to provide a controlled comparison between Sassha and SAM, demonstrating that Sassha is competitive with—or even outperforms—SAM. In line with this objective, we included an experiment in which SAM is given significantly more training iterations (x2) to show that Sassha still maintains superior performance, even under conditions that are favorable to SAM.

The addressing of block Hessian heterogeneity was not intended as a central claim of Section 5.3, but rather as one plausible hypothesis —based on prior literature [4,5]—for why Sassha may perform better than SAM on Transformer-based architectures. We did not mean to position this as the primary claim for the observed performance difference.

We appreciate the valuable feedback and will revise the writing in the final version.

---

On reported performances

“SAM gets 79.1 acc with ResNet50 on ImageNet, while the authors report 76.3”

This discrepancy is primarily due to the significantly larger number of training epochs used in [6] compared to ours (400 epochs in [6] vs. 90 in our setting). Additionally, we find differences in the experimental setup—such as the learning rate scheduler (cosine in [6] vs. multi-step in ours) and the use of label smoothing in [6]—this complicates a direct comparison of final accuracy.

“Sam gets 83.7 acc with WRN-28-10 on CIFAR100 ... However, the authors report 82.9”

First, we would like to clarify that the SAM (with SGD as the base optimizer) result reported in our paper is 83.14%, not 82.9% as mentioned. Additionally, we were unable to find a reference for the 83.7% accuracy cited by the reviewer. If possible, we would appreciate it if you could share the source. In our own investigation, we found that the performance numbers reported in reference [7] are consistent with our results.

Most importantly, we want to emphasize that all methods in our study were evaluated under an identical experimental setup to ensure fair comparison. Furthermore, we have prepared full configurations and code to reproduce all reported results, which will be released alongside the camera-ready version.

 

Reference

[1] Amari et al., When Does Preconditioning Help or Hurt Generalization?
[2] Liu et al., Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training
[3] Karimireddy et al., Error feedback fixes sign sgd and other gradient compression schemes
[4] Zhang et al., Why Transformers Need Adam: A Hessian Perspective
[5] Ormaniec et al., What does it mean to be a transformer? insights from a theoretical Hessian analysis
[6] Foret et al., Sharpness-Aware Minimization for Efficiently Improving Generalization
[7] Wu et al., CR-SAM: Curvature Regularized Sharpness-Aware Minimization

We sincerely thank the reviewer for thoroughly reviewing our work and giving us insightful and constructive feedback. While we address the raised questions below, we would be keen to engage in any further discussion as needed.

---

“How does SASSHA's performance scale to larger models and datasets beyond those presented?”

We appreciate the reviewer’s suggestion. We provide the validation performance of Sassha on larger models (ViT-B-32) below.

We observe that Sassha outperforms both first- and second-order baselines. We will include this in the updated paper. Additionally, we plan to include larger-scale models such as ViT-L and GPT2-small in the camera-ready version.

Note on configuration. For ViT-B, we used the same hyperparameters from ViT-S due to the limited time frame of the rebuttal.

---

“Is there a theoretical explanation for why combining sharpness awareness with second-order information works better than either approach alone?”

We believe that the strong performance of SASSHA stems from the complementary benefits of combining sharpness-awareness and second-order information. Specifically:

• The flatness induced by sharpness minimization has been shown—both theoretically and empirically—to improve generalization [1–5], and
• The effectiveness of preconditioning based on second-order information in adapting to the ill-conditioned geometry of deep learning has been well established in theory [6–9]. More recently, it has also been shown that optimal preconditioners can potentially accelerate the decrease of the population risk [10].

These advantages appear to act synergistically. In contrast, using either technique in isolation may introduce certain limitations (e.g., sharpness minimization alone may struggle with ill-conditioned geometry, while second-order methods alone may converge to sharp minima).

---

“Additional tuning parameter $\rho$ (perturbation radius) compared to pure second-order methods” as a weakness

While it is true that Sassha requires $\rho$, this might not necessarily be a fair criticism since (1) pure second-order methods are not applicable for deep learning, and (2) approximate second-order methods require (i.e., the fair baselines for the purpose of this work) their own hyperparameters to mitigate various issues such as training instability; for instance, AdaHessian requires the spatial averaging block size, Sophia requires a clipping threshold, and Shampoo requires a damping factor. It is also worth noting that Sassha is found to be generally robust to the range of $\rho$ values commonly used for SAM.

 

Reference

[1] Jiang et al., Fantastic generalization measures and where to find them, ICML 2019.
[2] Tsuzuku et al., Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using pac-bayesian analysis, ICML, 2020.
[3] H Petzka et al., Relative Flatness and Generalization, NeurIPS, 2021.
[4] Foret et al., Sharpness-Aware Minimization for Efficiently Improving Generalization, ICLR, 2021.
[5] Orvieto et al., Anticorrelated Noise Injection for Improved Generalization, ICML, 2022.
[6] Boyd et al., Convex Optimization, Cambridge University Press, 2004.
[7] Nocedal et al., Numerical Optimization, Springer, 2006.
[8] Bottou et al., Optimization Methods for Large-scale machine learning, SIAM Review, 2018.
[9] Jiang et al., How Does Adaptive Optimization Impact Local Neural Network Geometry?, NeurIPS, 2023.
[10] Amari et al., When Does Preconditioning Help or Hurt Generalization?, ICLR, 2021.

We really appreciate the reviewer taking the time to engage with our work. We address the specific points below and would be happy to clarify any remaining concerns.

---

Theory for improved generalization

The relationship between flatness and generalization is theoretically well-established, and many prior studies have shown that flat solutions correlate with better generalization performance [1-3]. We are currently developing a theory to prove the improvement of Sassha's generalization. Precisely, we have obtained the result that Sassha finds a flat solution using linear stability analysis [4] (which we will add in the final paper) and plan to use this result and existing flatness-based generalization bounds [3] to establish a generalization theory for Sassha. If you have any other suggestions, please let us know—we will try our best to incorporate them.

---

Convergence analysis for non-convex neural networks

This is indeed a valid limitation. However, convergence analyses under such assumptions are quite standard in optimization literature [6-8]. Moreover, convergence analysis under more realistic conditions, such as non-smoothness, remains a challenge and is being actively studied even for standard deep learning optimizers like SGD and Adam [9].

Nonetheless, we intend to analyze convergence properties of Sassha in non-convex settings. To achieve this, we plan to leverage frameworks for preconditioned optimizers that rely on Hessian eigenstructure analysis [10] and convergence analyses under mild assumptions in stochastic settings [11].

---

Scalability & Large-scale evaluations

The scalability of second-order optimization is a valid concern. However, despite such limitations, the broader community consensus is that the benefits provided by second-order methods—such as preconditioning—are potentially significant, and substantial collective efforts have made its computational complexities comparable to first-order methods[6, 12] ($O(n^3) \rightarrow O(n)$ in time complexity, and $O(n^2) \rightarrow O(n)$ in memory complexity).

Nevertheless, more generally, it is quite natural that leveraging more information for performance improvement can entail increased computations, which is understood as a cost-performance tradeoff. Perhaps more importantly, what matters would be whether or not the tradeoff is reasonable compared to existing alternatives, and precisely in that sense, Sassha sets a solid stand among other second-order baselines.

Also, we hope that the reviewer understands that the reason we have not evaluated billion-parameter scale models is not due to a fundamental limitation of second-order methods, but rather to enormous resource requirements to train models of such scale that we are unable to afford in our environments. In order to train an 8.3B parameter model using Adam, for instance, NVIDIA has employed 512 V100 GPUs (~16 TiB) [14].

---

Hyperparameters

Thank you for sharing your concern. However, we would like to note that the results for Sassha were obtained using values within standard ranges commonly adopted in prior work [4, 13, 17], and Sassha demonstrates strong performance without excessive tuning. We kindly refer you to Appendix D for detailed information on the hyperparameter search spaces across all experimental settings.

Also, regarding the Hessian update interval, we have already shown that Sassha is extremely robust to this hyperparameter, as demonstrated in Fig. 5(a). Please find a more detailed discussion in Section 6.3.

---

Source code

We have prepared a source code that reproduces all results presented in the paper, and we plan to release it alongside the camera-ready version.

 

Reference

[1] Jiang et al., Fantastic generalization measures and where to find them
[2] H Petzka et al., Relative Flatness and Generalization
[3] Foret et al., Sharpness-Aware Minimization for Efficiently Improving Generalization
[4] Wu et al., how sgd selects the global minima in over-parameterized learning: a dynamical stability perspective
[5] Shin et al., Critical Influence of Overparameterization on Sharpness-aware Minimization
[6] Liu et al., Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training
[7] Bottou et al., Optimization Methods for Large-Scale Machine Learning
[8] Reddi et al., On the Convergence of Adam and Beyond
[9] Li, Rakhlin & Jadbabaie, Convergence of Adam Under Relaxed Assumptions
[10] Doikov et al., Spectral Preconditioning for Gradient Methods on Graded Non-convex Functions
[11] He et al., Convergence of Adam for Non-convex Objectives: Relaxed Hyperparameters and Non-ergodic Case
[12] Yao et al., ADAHESSIAN: An adaptive second order optimization for machine learning
[13] Gupta et al., Shampoo: Preconditioned Stochastic Tensor Optimization, ICML, 2018.
[14] Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism