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.

---

More comparisons?

It appears that there may be a misunderstanding. All approximate second-order algorithms cited in Sections 4.1 and 4.2 (i.e. AdaHessian [1], Sophia [2], and Shampoo [3]) are compared in Section 5. Other algorithms cited in Sections 4.1 and 4.2 [5-11] are not approximate second-order methods and are thus excluded from the comparison (except SAM [4]). We note that for SAM [4] in Section 5.4, additional comparisons are presented in Section 6.3 and Appendix G.2. If this does not address your concerns or if there are any other comparisons you would like us to make, please let us know. We would be keen to provide them during the discussion period.

---

Errata

Some incomplete references in line 113.

Thank you for bringing this to our attention. We will have them fixed in the revised version.

---

Remark

We thank the reviewer for taking the time to review our work. Please let us know if there are any remaining questions or suggestions that could help improve the quality of our paper.

 

Reference
[1] Yao et al., Adahessian: An adaptive second order optimizer for machine learning. AAAI, 2021.
[2] Liu et al., Sophia: A scalable stochastic second-order optimizer for language model pre-training. ICLR, 2024.
[3] Gupta et al., Shampoo: Preconditioned stochastic tensor optimization. ICLR, 2018.
[4] Foret et al., Sharpness-aware minimization for efficiently improving generalization. ICLR, 2021.
[5] Chaudhari et al., Entropy-SGD: Biasing gradient descent into wide valleys. ICLR, 2017.
[6] Izmailov et al., 2018 Averaging weights leads to wider optima and better generalization. UAI, 2018.
[7] Orvieto et al., Anticorrelated noise injection for improved generalization. ICML, 2022.
[8] Levenberg et al., A method for the solution of certain non-linear problems in least squares. QAM, 1944.
[9] Marquardt et al., An algorithm for least-squares estimation of nonlinear parameters. SIAM, 1963.
[10] Amari et al., adaptive method of realizing natural gradient learning for multilayer perceptrons. Neural computation, 2000.
[11] Nesterov & Polyak et al., cubic regularization of newton method and its global performance. Math. prog., 2006.

whether it is possible to compare the performance of ‘sassha’ with those algorithms in Section 4.1&4.2 that are related to the design of ‘sassha’ and its idea to some extent.

First, we list up all the referenced works in Sections 4.1 and 4.2 here and describe in detail why they are excluded for comparisons:

Section 4.1 [1-3] : other sharpness minimization algorithms

• Entropy-SGD [1] demands several additional backpropagations per update to minimize local entropy, which incurs impractical costs in deep learning while offering only marginal improvements over SGD.
• SWA [2] is an algorithm that averages weight points to find flat solutions, requiring keeping several parameters and computing multiple gradients in each iteration. It is also known to be less effective compared to SAM [4].
• Anti-PGD [3] proposes to inject a special type of noise in the weights to artificially mimic the noise of SGD, which is not as effective as SAM.

Section 4.2 [5-8] : classical second-order techniques

• Levenberg-Marquardt [5-6] requires resource-intensive computing and storing the inverse of the Gauss-Newton matrix, tracking changes in the loss function, and demanding additional matrix calculations to adaptively adjust its damping factor.
• Natural Gradient Descent [7] similarly demands the computation and storage of the inverse of the Fisher information matrix for preconditioning.
• Cubic Regularization [8] involves computing the Hessian matrix during optimization, repeatedly calculating Hessian-vector products, and solving complex subproblems.

In short, we chose not to compare the sharpness-minimization algorithms [1-3] for two key reasons: (i) their mechanisms differ fundamentally from Sassha's and are not directly related; and (ii) they are more resource-intensive and less competitive than SAM, making SAM a sufficient baseline for demonstrating Sassha's advantages. Similarly, the classical second-order algorithms [5-8] were excluded due to their computational infeasibility in deep learning, requiring $\mathcal{O(d^3)}$ for computation and $\mathcal{O(d^2)}$ for memory.

We emphasize that the algorithms selected for comparison in this paper were carefully chosen to ensure relevance, scalability, and objectivity. The exclusion of certain algorithms from our comparison was not simply because they are not approximate second-order algorithms, but was aimed at providing a focused and cohesive evaluation.

---

More comparison

Nevertheless, we additionally conducted further experiments on algorithms related to Sassha:

The square root technique used in Sassha for stability improvement is related to damping and clipping. We also conducted experiments comparing the square root with these two techniques.

We provide additional experiments comparing Sassha with SAM on Language tasks.

Pretraining

Finetuning

In short, Sassha was carefully designed to become a practical approximate second-order method and strong alternative to existing methods as demonstrated by these comparisons.

If there is any other comparison that the reviewer believes should be included, we kindly request the reviewer for specific pointers to it along with rationale. We would be keen to provide them in the final version, if not during the discussion period. Otherwise, we would greatly appreciate it if the reviewer could reconsider the earlier critique and kindly reflect that on the rating of our work.

 

Reference
[1] Chaudhari et al., Entropy-SGD: Biasing gradient descent into wide valleys. ICLR, 2017.
[2] Izmailov et al., 2018 Averaging weights leads to wider optima and better generalization. UAI, 2018.
[3] Orvieto et al., Anticorrelated noise injection for improved generalization. ICML, 2022.
[4] Jean Kaddour et al., When Do Flat Minima Optimizers Work?. Neurips, 2022
[5] Levenberg et al., A method for the solution of certain non-linear problems in least squares. QAM, 1944.
[6] Marquardt et al., An algorithm for least-squares estimation of nonlinear parameters. SIAM, 1963.
[7] Amari et al., adaptive method of realizing natural gradient learning for multilayer perceptrons. Neural computation, 2000.
[8] Nesterov & Polyak et al., cubic regularization of newton method and its global performance. Math. prog., 2006.

We are greatly encouraged by the reviewer’s positive and constructive feedback, which we believe will tremendously improve our work. While we address specific comments below, we would be keen to engage in any further discussion.

---

W1: Technical originality?

“combining sharpness-awareness with second-order information is relatively straightforward and could be considered incremental”

With all due respect, we disagree with the reviewer’s criticism. Simply put, this “combining” process was far from trivial. Achieving seamless integration of sharpness minimization into second-order optimization, while retaining their benefits and avoiding instability and inefficiency, indeed required extra careful engineering. This process involved extensive verifications of a diverse set of algorithmic design choices.

It is also worth noting that such engineering is a standard procedure in optimization. Seemingly simple but well-thought-through ideas often result in significant performance improvement and have been appreciated by the community. For example, Clipping is the core component that distinguishes Sophia from other approximate second-order optimizers [1]. Similarly, AdamW’s decoupling of weight decay from L2 regularization is a widely appreciated innovation despite its simplicity [2].

“The empirical investigation of sharpness measures largely confirms known intuitions about the relationship between curvature and generalization.”

We also find this criticism hard to accept, as to the best of our knowledge, no prior work has explicitly investigated nor confirmed what we show in this work (i.e., the correlation between the poor generalization of approximate second order methods and the sharpness of their solution) to any degree matching ours.

Based on these points, we hope the reviewer could reconsider the value of our contributions.

---

W2: Critical experimental gaps?

The authors avoid testing on large-scale models (>100M parameters), raising questions about the scalability to modern deep neural networks.

We test Sassha on the fine-tuning task using the BERT Base model (110M parameters). The results are shown below:

This indicates that Sassha achieves better performance than others in most tasks. We will include this in the updated paper.

the absence of results on fundamental computer vision tasks, such as object detection and semantic segmentation, suggests potential limitations in the performance consistency of the method

While it is understandable to be curious about the full potential, we find this request quite out of context and beyond the scope of this work.

More importantly, the paper lacks comparison with recent sharpness-aware variants like ASAM [1] and GSAM [2], making it impossible to assess whether the proposed method represents genuine progress in the field.

We perform experiments to compare Sassha against advanced SAM variants [4-5]. For a fair comparison, we also evaluate G-Sassha (Sassha with surrogate gap-guided sharpness from GSAM[5]). The results are below.

We find that the performance of Sassha is comparable to, and mostly better than, these SAM variants. In addition, we note clearly that GSAM requires tuning additional hyperparameters (which are known to be quite important but sensitive). ASAM, according to Kwon et al., may need a wider $\rho$ search range than the original SAM, which implies a more involved search process.

W3: Technical concerns?

The square root pre-conditioner, while empirically shown to provide stability benefits, appears to be an arbitrary choice without a rigorous mathematical foundation

We thank the reviewer’s constructive criticism. Using square-rooted preconditioner can be viewed as geometric interpolation between the first order method $H^{-\alpha}$, which has been demonstrated to enable selecting an optimal preconditioner that balances the bias and the variance of the population risk, thereby minimizing generalization error [6]. In general, $\alpha = 1/2$ (i.e., square root) has shown moderate performance across various scenarios [6, 7, 8]. Additionally, square rooting can alleviate specific instabilities that approximate second-order optimizers suffer near minima. Specifically, since we posit that the commonly used diagonal Hessian estimators can scale inversely quadratic to the gradient, the preconditioned gradient update can scale inverse to the gradient [9, 10]. This behavior can lead to instability as the gradient scale diminishes as the model reaches near minima. Square-rooting the preconditioner can help mitigate such instabilities. We will add this discussion to improve our manuscript.

The hyper-parameter k for Hessian update intervals introduces additional complexity

Sassha remains robust even with a very large $k$ as demonstrated in Fig 5 (a) in Section 6.2.. This is potentially due to the sharpness minimization process contributing to Hessians less frequently changing, thereby enabling reuse. We refer the reviewer to Section 5.4 where we provide an ablation analysis on this aspect. While we may not be able to precisely point to an “optimal” $k$, we can at least say that a reasonable value seems to suffice.

the memory requirements of storing and updating the diagonal Hessian approximation, though lower than full second-order optimizers, could still be prohibitive for large-scale applications

The memory requirement is not prohibitive. Storing the diagonal Hessian requires the same amount of memory as storing second moments in Adam. Also, while updating the diagonal Hessian using Hessian-vector products (HVP) incurs additional memory overhead compared to gradient computation, [11] indicates that this is acceptable for most large-scale applications.

Nonetheless, we also plan to provide efficient implementations that leverage HVP-free Hessian approximations such as the diagonal of the Gauss-Newton matrix, which only introduces a 5% increase in memory cost compared to Adam while preserving similar performance [1].

We deeply appreciate the reviewer for recognizing our efforts and leaving us a list of valuable suggestions and feedback, which will undoubtedly improve our work tremendously. We will ensure to incorporate all of these in the final version of the paper. Below, we address the reviewer’s last comment in detail.

---

A deeper empirical analysis of these variations would provide valuable insights for practitioners in the community.

We sincerely thank the reviewer for the thoughtful suggestion, and we agree that this would be valuable to practitioners. Here, we present further empirical analysis and practical considerations for interpreting these variations in our results.

GSAM vs. Sassha. While the accuracy difference between Sassha and GSAM may seem marginal in some cases (e.g., ImageNet), it is crucial to note that achieving comparable results with GSAM required us to perform significantly more hyperparameter tuning. In contrast, Sassha consistently achieved strong performance without these additional efforts.

Specifically, GSAM introduces an additional hyperparameter $\alpha$ to penalize dominant eigenvalues, demanding as much tuning effort as tuning $\rho$. For instance, in the ResNet-50 setup, tuning GSAM involved $8.75 \times$ larger search grids compared to Sassha. GSAM (mostly follows the same search ranges presented in the GSAM paper)

• ResNet on ImageNet:
  • $\rho: \\{\textcolor{red}{0.01, 0.05}, 0.1, 0.15, 0.2, 0.25, \textcolor{red}{0.3}\\}$
  • $\textcolor{red}{\alpha: \\{0.01, 0.05, 0.1, 0.15, 0.2\\}}$

Additionally, we would like to provide our interpretation of how GSAM might achieve good performance in ImageNet. As discussed in [1], there could exist a discrepancy between the actual flatness (dominant Hessian eigenvalue) of the loss and the “perturbed loss” definition used by Sassha, which can worsen under a more complex and intricate loss landscape of ImageNet. Consequently, GSAM, designed to minimize dominant eigenvalues explicitly via a surrogate gap, may be better positioned to find flatter solutions than Sassha, which employs a sharpness minimization mechanism similar to SAM. Nevertheless, we note once again that GSAM requires a much larger tuning budget than Sassha to achieve such results.

BERT finetuning. Here, while Sassha outperforms on subsets of tasks, it outperforms by larger margins compared to marginal differences in underperforming cases. Specifically, on tasks where Sassha outperforms, the performance gaps are more pronounced compared to that of underperforming cases ({+2.17%, +1.675%, +0.815% +0.23%} vs. {-0.11%, -0.035%, -0.11%}). Thus, we can say that on average, Sassha outperforms in the BERT fintuning by +0.53%. We additionally note that it is common practice to evaluate GLUE finetuning tasks through their average score [2, 3].

Update on STSB. The initial numbers we reported for the STSB results of Sassha were mistakenly put 84.86/84.76. This was a typo, and now we fixed it with correct results of 89.29/89.71.

 

Reference
[1] Zhuang et al., Surrogate Gap Minimization Improves Sharpness-Aware Training. ICLR, 2022.
[2] Clark et al., ELECTRA: pre-training text encoders as discriminators rather than generator. ICLR, 2020.
[3] Hu et al., LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS. ICLR, 2022.

We appreciate the reviewer for taking the time to review our work. While we respond to the reviewer’s specific comments as below, please do let us know if there is anything else we can address further.

---

Sassha just combines SAM and common second-order optimization techniques?

We respectfully disagree with the reviewer. While the underlying idea behind Sassha may seem straightforward, it should not be judged without accounting for the technical and engineering challenges required to fully realize it. Extensive efforts were devoted to carefully engineering and examining various techniques, which all play a crucial role in mitigating specific issues arising from naively applying sharpness minimization to second order methods:

• Absolute Hessian enables Sassha to avoid saddle points while preserving optimal rescaling of negative Hessians, unlike clipping or damping.
• Square root is introduced to prevent divergent behavior from curvature under-estimation from introducing sharpness minimization.
• Lazy Hessian updating is intentionally set to default to mitigate computation overhead based on investigations revealing that sharpness minimization can help enable this without severe performance degradation.

We also note that this way of engineering is standard in optimization since ideas can achieve remarkable effectiveness when supported by seemingly simple but carefully considered techniques. For instance, Sophia has demonstrated that a simple clipping mechanism was sufficient to enhance second-order optimizers in training LLMs, gaining wide recognition by the community. Thus, reducing Sassha as merely combining SAM with common second-order optimizers neglects the complexities and efforts required to realize this concept, which is an unjust representation.

To more directly support our claims, we compare Sassha to simple combinations of SAM + second-order optimizer (Sophia-H):

We observe that it indeed underperforms Sassha given similar search space. This indicates that the techniques we designed for Sassha are non-trivial, and should be taken account of when evaluating Sassha.

---

Sassha vs. SAM & other first-order baselines under a fair setting

We appreciate the reviewer for the suggestion. We provide further comparison with SAM, please refer to Appendix G.2. For other first-order baselines, here we provide results for SGD and AdamW with twice the epoch budget than Sassha.

We observe that despite the first-order baselines being given more training budgets, they are still unable to outperform Sassha. We believe that this is due to the synergy between preconditioning and the bias towards flat solutions.

---

Lack of theoretical analysis on why SOO finds sharp solutions

We thank the reviewer for the insightful comment. Here we provide a theoretical account of why second-order methods might converge to sharp solutions as described in [2].

According to the linear stability analysis of Wu et al. [2], SGD selects minima with the maximum Hessian eigenvalue ($\lambda_{\max}(H^\star)$, i.e., sharpness) satisfying the following condition: $\lambda_{\max}(H^\star)\leq2/\eta$ where the $\eta$ denotes the step size. This indicates that SGD escapes from sharp minima and converges to flat minima.

However, approximate second-order optimizers (with preconditioner $D^{-1}\approx H^{-1}$) will converge to minima of the following condition: $\lambda_{\max}(D^{-1}H^\star)\leq2/\eta$ From this, we observe that when $D^{-1}H^\star$ is close to the identity matrix, this condition is fulfilled as long as $\eta\leq2$. This implies that, unlike SGD, approximate second-order optimizers can converge to virtually any minima, including sharper ones that SGD might escape from.

We will refer to these in the revised version.

---

Questions

In Table 7, why the average sharpness of SGD and SAM is negative values?

The formula we used for computing the average sharpness [3] is as follows: $\mathbb{E}_{z\sim\mathcal{N}(0,1)}[L(x^{\star}+\rho z/\|z\|)-L{x^\star}],$ which can yield negative values if the solution $x^\star$ from the optimizer isn’t the exact minima of $L$, which can often be the case for deep learning.

 

Reference
[1] Liu et al., Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training, ICLR, 2024.
[2] Wu et al., How SGD Selects the Global Minima in Over-parameterized Learning: A Dynamical Stability Perspective. NeurIPS, 2018.
[3] Wen et al., How Sharpness-Aware Minimization Minimizes Sharpness? ICLR, 2023.

We sincerely appreciate the reviewer for recognizing our contribution and raising the score. We are also grateful for the reviewer’s insightful feedback, which we believe will significantly enhance our paper. We will put our best efforts to reflect the reviewer's suggestions in the updated version.

---

it seems to indicate that the model has not been well-optimized

Firstly, it is not quite likely that our models are not well-optimized, since their performance aligns well with previously reported results in prior works [1, 2].

Even so, negative average sharpness values can naturally be measured for non-convex $L$ in deep learning (depending on the selected $\rho$). This is because the sharpness, by definition, simply measures the loss difference within a $\rho$-radius. This is slightly different from our earlier explanation (which was rooted from assuming a convex scenario), which we now see could be misleading. We apologize for the confusion and would like to clarify our stance. Please note that we have included this average sharpness measure alongside other metrics to provide a comprehensive reference as in prior work [3].

Before we wrap up, in response to the reviewer’s comment on “minor improvements", we would like to gently highlight that the performance gains achieved by Sassha are in fact quite prominent. As shown in Table 2, Sassha consistently outperforms existing approximate second-order methods by at least 1%, and up to around 5%. Seeing Sassha exhibits improved performance across various tasks (even compared to SAM), we have a strong belief in the generality and robustness of Sassha.

 

Reference
[1] Zagoruyko et al. Wide Residual Networks. CoRR, 2016.
[2] Kwon et al., ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. ICML, 2021.
[3] Wen et al., How Sharpness-Aware Minimization Minimizes Sharpness? ICLR, 2023.

SAM for language tasks

We provide the validation performance of SAM on various language model tasks below.

Pretraining:

Finetuning:

We observe that Sassha outperforms SAM across most tasks. We will include this in the updated paper.

---

Comparisons to SAM variants

We conduct experiments to compare Sassha with advanced SAM variants [4-5]. We also evaluate G-Sassha (Sassha with surrogate gap guided sharpness from GSAM[5]) for fair comparison. The results are below.

We find that Sassha is competitive with these advanced SAM variants. In addition, we note clearly that ASAM, according to its authors, may require a broader $\rho$ search range compared to the original SAM, leading to a more complex tuning process. Similarly, GSAM involves tuning additional hyperparameters that are critical yet quite sensitive.

Search space. Here, we provide detailed setup and hyperparameter search space. For ResNet, we use SGD as the base optimizer for ASAM and GSAM, while for ViT, AdamW with gradient clipping set to 1.0 serves as the base optimizer. For all models, cross entropy loss is used, and the best learning rate and weight decay of the base optimizer are selected in experiments with ASAM and GSAM. All algorithms are evaluated with fixed rho (not using rho scheduling). For ResNet on CIFAR, we apply multi-step decay with a decay rate of 0.1, and for ViT, we use cosine learning rate decay with 8 warm-up epochs.
For GSAM and G-Sassha:

• ResNet on CIFAR:
  • $\rho$: {0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4}
  • $\alpha$: {0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3}
• ResNet on ImageNet:
  • $\rho$: {0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3}
  • $\alpha$: {0.01, 0.05, 0.1, 0.15, 0.2}
• ViT on ImageNet:
  • $\rho$: {0.1, 0.2, 0.3, 0.4, 0.5, 0.6}
  • $\alpha$: {0.1, 0.2, 0.3}

For ASAM:

• $\rho$: {0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.5, 2}

We really appreciate the reviewer’s positive and constructive feedback. We believe this has led us to tremendously improve our work. While we make clarifications to specific comments below, please let us know if there is anything we need to address further. We would be keen to engage in any further discussion.

---

Isn’t it just additional tweaks to second-order methods (or Sophia-H) and applying SAM?

We think it is reasonable to project that adding SAM can generally enhance the performance of the respective base optimizers. However, the process of implementing this effectively for second-order optimizers is far from trivial. Achieving seamless integration of sharpness minimization and second-order optimization to preserve their benefits while avoiding instability required careful engineering that involves exploration of various algorithmic design choices grounded in established prior studies (e.g., square root, absolute, see Section 4) and validating them through extensive testing (see Section 6).

It is also worth noting that such engineering is a standard procedure in the literature since seemingly simple but well-chosen “tweaks” have been demonstrated to significantly improve performance, which are often appreciated by the community. Clipping, for instance, is a fundamental component in Sophia’s design and is, in fact, its sole feature distinguishing it from basic approximate second-order optimizers (just to be clear, bias correction has nothing to do with clipping, it is for correcting biases from zero-initializing first and second-moment variables).

Nonetheless, we provide results (val accuracy) for simple Sophia-H + SAM on various tasks.

These results indicate that Sassha indeed outperforms the combination of Sophia-H + SAM. We attribute this to (i) the limitations of clipping and (ii) its reduced compatibility with SAM compared to the square root method.

(i) The use of clipping results in Sophia partially performing signSGD over a subset of parameters [1], which may lead to suboptimal convergence in typical situations (on non-transformer architectures) [2, 3].

(ii) When curvature becomes flat due to SAM, clipping may lead to a significant increase in the number of Hessian entries replaced by a very small constant $\epsilon$. This situation raises the sensitivity to hyperparameters like the clipping threshold and makes the optimization process more dependent on careful tuning.

Conversely, square root does not enforce hard adjustments to the magnitude or direction of the update vector or to Hessian entries. Instead, it smoothly interpolate Hessian toward the identity matrix and performs Newton updates. A more detail explanation of the square root approach are provided in on sqrt below.

For these reasons, we believe it would be unreasonable to consider Sassha as merely SAM with a second-order method as the base optimizer or Sophia with a few tweaks.

---

300 epochs for ImageNet

We appreciate the reviewer’s suggestion. The experiment is conducted as shown below,

The result shows that Sassha remains consistently better, even with extended training epochs.

---

Large search space of rho

$\rho$ was searched over $\{0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3\}$, which is slightly larger than what was used in the label noise experiments. As the reviewer suggested, we extended this range to $0.8$ in increments of $0.05$. The results are as follows.

RN20 CIFAR10 (Sassha 92.983):

RN32 CIFAR10 (Sassha 94.093):

RN32 CIFAR100 (Sassha 72.143):

We find that SAM still underperforms Sassha despite it being given a much larger search space (and hence increased search budget). We will include the search space and these results in the revised version.

Questions

Can the authors explain why Sophia-H and AdaHessian are so slow in terms of wall clock time (G.3). Also, part of the speedup compared to Sophia-H might vanish when Sophia-H is used with a k-interval like Sassha. I think this should be stated more clearly.

We thank the reviewer for the observation regarding the wall clock time differences.

The longer wall clock time of Sophia-H and AdaHessian compared to Sassha arises from our decision not to apply the lazy Hessian update to both algorithms in the experiments. This choice was made to compare Sassha against the two algorithms under their best performance configurations. Specifically, Liu et al. [1] demonstrated that Sophia achieves its peak performance at k=1, and AdaHessian [8] does not support lazy Hessian updates in its original design.

We believe this to be fairer as our claim for Sassha focuses on enhancing the generalization performance of second-order optimizers, not their efficiency. Introducing lazy Hessian updates for Sophia in our main experiments would improve speed but at the cost of reduced performance compared to what is currently reported in our paper, which still underperforms ours. Nonetheless, we will make this point clearer in the revision.

Is there a reason that Sophia-H is investigated, but Sophia-G is omitted?

We would first like to clarify that ‘H’ and ‘G’ refer to Hutchinson’s unbiased diagonal Hessian estimator and the Gauss-Newton-Bartlett biased diagonal Hessian estimator, respectively, for computing the preconditioner. However, as noted by Liu et al. [1], there is no discernable performance difference between the two.

In our experiments, we used Sophia-H to ensure a consistent comparison of the effects of sharpness minimization while keeping the diagonal Hessian estimator identical.

Did the authors use Shampoo or distributed Shampoo?

We used the original Shampoo [9] for our experiments.

---

Minor remarks

We appreciate the reviewer’s detailed feedback and suggestions. We will make the necessary adjustments in the revised paper.

 

Reference
[1] Liu et al., Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training, ICLR, 2024.
[2] Kunstner et al., Noise Is Not the Main Factor Behind the Gap Between SGD and Adam on Transformers, but Sign Descent Might Be. ICLR, 2023.
[3] Karimireddy et al., Error Feedback Fixes SignSGD and other Gradient Compression Schemes. ICML, 2019.
[4] Kwon et al., ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. ICML, 2021.
[5] Zhuang et al., Surrogate Gap Minimization Improves Sharpness-Aware Training. ICLR, 2022.
[6] Amari et al., When does Preconditioning Help or Hurt Generalization? ICLR, 2021.
[7] Kunstner et al., Limitations of the Empirical Fisher Approximation for Natural Gradient Descent. NeurIPS, 2019.
[8] Yao et al., ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning. AAAI, 2020.
[9] Gupta et al., Shampoo: Preconditioned stochastic tensor optimization. ICLR, 2018.

We really appreciate the reviewer’s continuing engagement in our work. We address the specific points below.

---

Novelty

In essence, it underlines that for SASSHA to be a valid contribution, it needs to improve over SAM-baselines, as improvements over second-order optimizers could have been achieved by a naive baseline.

First, we would like to note that the performance improvements of Sassha over naive baselines (Sophia-H + SAM) should not be overlooked, as they are noticeable and consistent. For instance, Sassha outperforms Sophia-H + SAM by +0.453%p on ResNet20/CIFAR-10, which is more noticeable compared to widely recognized ASAM which improves over its baseline (SAM) by +0.26%p on the same task [2]. Also note that Sassha shows mostly competitive performance or even surpasses SAM in various cases, whereas Sophia-H + SAM underperforms SAM.

More importantly, achieving these results with Sophia-H + SAM required extensive hyperparameter search to tune its additional clipping threshold and $\epsilon$ ($28\times$ larger search grids in our setup), as the value suggested in the original Sophia paper [3] does not generalize well beyond language tasks. In contrast, Sassha consistently achieves robust performance across all tested tasks without any additional tuning, as it shares the same set of hyperparameters as SAM.

Thus, with respect, we find the critique (“the improvements can be achieved by a naive baseline”) too stretched to accept, since Sassha (1) consistently outperforms Sophia-H + SAM, and (2) it does so without requiring extra extensive hyperparameter tuning.

---

Comparison to SAM

While I can see improvements of SASSHA for certain scenarios (e.g. language pretraining, ViT on ImageNet), differences are often small (CIFAR, esp. with the newer SAM variants) or inconsistent (Bert-Base ft, RN50 ImageNet).

We relate to the reviewer’s concern. However, we find it a bit unfairly harsh, by being only focused on “performance” without concerning budgets (for newer SAM variants) or partly incorrect (for Bert-Base ft).

First, SAM variants require considerably more hyperparameter tuning to achieve generalization performance comparable to Sassha. For example, GSAM introduces an additional hyperparameter $\alpha$, demanding as much tuning effort as tuning $\rho$. Similarly, ASAM, as noted by its authors, typically necessitates exploring a broader $\rho$ range, as its appropriate value is approximately 10 times larger than that of SAM. In our setup, tuning GSAM and ASAM involved $4.5\times\sim15.75\times$ and $3\times$ larger search grids compared to Sassha, respectively.

GSAM (mostly followed the same search ranges presented in the GSAM paper)

• ResNet on CIFAR:
  • $\rho: \\{\textcolor{red}{0.01, 0.05}, 0.1, 0.15, 0.2, 0.25, \textcolor{red}{0.3, 0.35, 0.4}\\}$
  • $\textcolor{red}{\alpha: \\{0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3\\}}$
• ResNet on ImageNet:
  • $\rho: \\{\textcolor{red}{0.01, 0.05}, 0.1, 0.15, 0.2, 0.25, \textcolor{red}{0.3}\\}$
  • $\textcolor{red}{\alpha: \\{0.01, 0.05, 0.1, 0.15, 0.2\\}}$
• ViT on ImageNet:
  • $\rho: \\{0.1, 0.2, 0.3, 0.4, \textcolor{red}{0.5, 0.6}\\}$
  • $\textcolor{red}{\alpha: \\{0.1, 0.2, 0.3\\}}$

ASAM

• $\rho: \\{0.1, 0.15, 0.2, 0.3, \textcolor{red}{0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.5}\\}$

Next, it appears that the reviewer may have misinterpreted the fine-tuning results (on squeezeBERT) in our first response since Sassha consistently outperforms SAM across almost all tasks. Notably, Sassha achieves average improvements of approximately 3% in MRPC and 1% in QQP.

ViT results

For ViT-S32, I’m surprised by the low numbers of SAM+AdamW. For almost identical settings (300 epochs, cosine lr decay, lr 0.001, 1024 batch size, 8 epochs warmup), numbers that are >2.5% larger are achieved (e.g., in [1]). I’m unsure about the reason for this, but I encourage the authors to double-check their implementation carefully.

Upon investigation, we found key differences in the setup that may have resulted in this performance difference. Specifically, the official implementation for MSAM [4] revealed that they set the ViT depth to 16 and employed advanced methods such as cutout regularization and rho scheduling. By contrast, we configured the ViT depth with 12 following the original ViT paper and its variants [5-7] and did not use the aforementioned additional techniques. We also note that our results are the average of three random seeds, whereas their results are based on a single run. Nevertheless, we believe this does not affect the conclusion of our results, since all algorithms were compared fairly in the identical experiment settings.

---

Final remark

We sincerely appreciate the time and effort the reviewer dedicated to providing such thoughtful feedback. With respect, however, we find that some comments are quite focused on “performance perspective”, thereby potentially overlooking other significant contributions presented in our work, most notably for example:

• We conducted extensive experiments to investigate the correlation between the poor generalization of approximate second-order methods and the sharpness of their solution. To the best of our knowledge, no prior work has explicitly investigated this to any degree matching ours.
• We developed a new practical and robust second-order method that produces strong and consistent results across diverse sets of deep learning experiments, which is validated through a series of ablation analyses, and also, theoretically proven with a convergence analysis.
• We analyzed the positive impact of sharpness minimization on Hessian reuse. Our empirical investigations suggest that this effect may stem from sharpness minimization slowing the rate of change in the Hessian between iterations. This observation highlights sharpness minimization as a promising approach to addressing the inefficiencies associated with Hessian computation in second-order methods.

We hope these contributions are appropriately recognized and considered in the evaluation of our work.

 

Reference
[1] Becker et al., “Momentum-sam: Sharpness aware minimization without computational overhead,” arXiv, 2024.
[2] Kwon et al., ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks. ICML, 2021.
[3] Liu et al., Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training, ICLR, 2024.
[4] MarlonBecker, Official implementation of MSAM, 2024.
[5] Dosovitskiy et al., An image is worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021.
[6] Xu et al., Supplementary to GroupViT: Semantic Segmentation Emerges from Text Supervision. CVPR, 2022.
[7] Touvron et al., Training data-efficient image transformers & distillation through attention. ICML, 2021.

We sincerely thank the AC for their time and efforts to actively participate in reviewing our paper, encouraging discussions, and sharing feedback that would enrich the quality of our work. We address your question below.

---

Compare to Li et al. [1]

With respect, we believe that there are major differences between our work and [1], which we detail below.

First, the purpose of [1] differs significantly from Sassha. Precisely, Sassha aims to mitigate the poor generalization issue of second-order optimizers by incorporating sharpness minimization. In contrast, [1] focuses on improving the adversarial step in SAM by addressing the issue of an “overly-friendly adversary” via variance suppression.

Consequently, their resulting algorithms share little similarity. Specifically, Sassha, designed around approximate second-order optimization, employs diagonal Hessian preconditioning in its descent step. Furthermore, as its name suggests, Sassha introduces techniques (e.g., square root and absolute function) to ensure stable Hessian approximation. Meanwhile, Vasso [1], built around SAM, introduces an exponential moving average (EMA) to its adversarial step for efficient variance suppression.

Thus, the adversarial direction used by Sassha differs from that of [1], since Sassha utilizes the non-EMA perturbation step as in the original SAM.

We summarize these differences in the following table:

---

Discrepancy in ViT-S/32 performance

We’re really grateful for your observation on this matter.

Upon closer inspection, we find this discrepancy to arise primarily from differences in experimental configurations. Specifically, $\textcolor{red}{\text{[1]}}$ reports results using the more advanced $\textcolor{red}{**AdamW**}$ optimizer with a total of $\textcolor{red}{**300 epochs**}$ for optimization. In contrast, $\textcolor{blue}{\text{our vanilla training}}$ result of 62.94% was achieved using $\textcolor{blue}{**SGD**}$ with a significantly shorter training duration of $\textcolor{blue}{**90 epochs**}$.

For reference, we also provide results using AdamW. With 90 epochs, our method achieves 66.46% accuracy, and we anticipate that with 300 epochs—similar to [1]—our results would align closely with theirs.

However, more importantly, we believe this would not undermine the validity of our conclusion, since all comparisons between Sassha and other methods are conducted under identical experiment settings.

---

Final remark

Once again, we really appreciate the AC for sparing their valuable time to review our work. We hope that our findings and discussions prove insightful and that our contributions can be acknowledged as recognized by most of the reviewers.

 

Reference
[1] Li et al., Enhancing Sharpness-Aware Optimization Through Variance Suppression, Neurips, 2023.