We would like to thank the reviewer for their time. Below, we address questions and concerns raised by the reviewer.

• We agree with the reviewer that Foret et al. [2020] are established using a constant $\rho$. However, we politely disagree with the comment that “no existing works that establish convergence results on the constant $\rho$ for non-convex functions.” As we mentioned in our paper, there are some existing results about constant step-size SAM in the general non-convex setting. For example, in [Andriushchenko, 2022] and [Li & Giannakis, 2023] the authors prove convergence of USAM and SAM respectively, with constant stepsizes that depend on $T$ (total number of iterations) with rate $O(1/\sqrt{T})$. Both papers make the strong assumption of bounded variance. All of the above papers include limitations in their convergence guarantees and this is precisely what our work fixes. We provide tight convergence analysis of both USAM and SAM under relaxed assumptions (and we did that under the unified SAM framework we proposed).

  Regarding the mentioned papers, let us provide more details below. We have already cited them in the updated version of our work (see pdf).

  Khanh et al. [2024]: This work only considers the deterministic setting and all the results are asymptotic. In our paper, we focus on the stochastic regime (which has the deterministic as a special case) and we provide convergence rates in our results.

  Nam [2023]: We were not aware of this work. Thank you for bringing it to our attention. As the reviewer mentioned, we make the same main assumption with this work, namely the expected residual assumption. However, there the authors additionally assume bounded gradient of f and the results are only asymptotic and hold almost surely while they do not prove any convergence rates (do not provide how fast the proposed method is). Furthermore, they have no theory for PL functions.

  Xie [2024]: We were not aware of this work. Actually, it would be impossible to compare with it as it was published on arxiv on Oct 14, while the ICLR deadline was Oct 1. However, thank you for bringing it to our attention. Compared to our work, this paper only focuses on the deterministic regime, arguably a much simpler setting. This is really concurrent work, and our submission should not be judged because it is missing such (impossible to have) comparisons.

• We politely disagree with the reviewer's comment on the prior works and their relation to our paper. In particular, the reviewer's second statement related to prior work is misleading. None of the previous results mentioned provide convergence rates.

  In our work, we provide convergence rates, and thus, the two results are not directly comparable. In Thm 3.7 we show that Unified SAM yields the rate $O(\epsilon^{-4})$ rate for finding a stationary point of nonconvex smooth functions. Our results are clearly not weaker than Nam [2023] or Khanh [2024] as the reviewer pointed out. These papers only provide asymptotic convergence. Furthermore, as mentioned in the previous point, in Nam [2023] also assumes a bounded gradient of the function on top of the ER condition, making the results impractical as this condition is very restrictive in practical scenarios.

• It is updated (please see line 296).

• Our assumption 3.1 and assumption A.4 in Nam et al. [2023] are equivalent.

• The name Unified SAM is justified since it indeed unifies the analysis of both USAM and SAM. As we never claim that we have a unified theory for all different variants of SAM we believe that the name of the method is okay. We agree with the reviewer that there are many other SAM-like variants that our approach does not cover and we add a sentence related to this in the updated version.

• The reviewer mentioned: “Tables 2 and 3, the proposed Unified-SAM does not show significant improvement over SAM. This is a point that diminishes the importance and contribution of the paper.” We politely disagree with this statement; even if minor improvement in general our results in terms of generalization aligned well with existing state-of-the-art approaches in the literature on improving SAM. Please check the following papers with similar improvements over SAM: [Li, 2023] Table 1,2; and [Kwon, 2021] Table 6. Even the papers the reviewer mentioned have similar improvements to ours (see Xie [2024] Table 1,2).

• We have updated this sentence to the new PDF file.

References (not mentioned in our paper):

[Kwon, 2021]: ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks

Dear Reviewer hnag,

In our rebuttal, we explained in more detail that none of the papers mentioned in your review captures the convergence guarantees of our work and their importance in understanding the behavior of stochastic SAM-type methods. They either focus on the arguably much simpler deterministic case or if they have stochastic results, they only have asymptotic convergence (convergence to a minimum but not convergence rates).

One of the papers mentioned became available online only after the ICLR submission deadline (it was impossible to be aware of this).

Please let us know if you have any other concerns.

In our opinion, based on the original points of the reviewer, none of the issues raised justify a score of rejection/borderline rejection. If you agree that we managed to address all issues, please consider raising your mark. If you believe this is not the case, please let us know so we can respond.

Thanks, The authors

Thank you for investing time in the review process and for the detailed review. Below, we address questions and concerns raised by the reviewer.

The reviewer has concerns about the correctness and significance of our results and claims that because of these issues, the current draft does not meet the bar for ICLR. We disagree with this judgment, and we provide more details below. In our opinion, all points mentioned as significant limitations of our work are simply clarification points.

Correctness: The reviewer raised three points related to the correctness of our results.

The first two points are related to the choice of stepsize gamma. That is, the reviewer argues that having $\gamma=0$ makes the theorem useless. We agree with this statement, and we point out here that $\gamma$ as a step-size is always positive by default. This applies to any practical optimization algorithm (including gradient descent and its different variants). If the step-size is zero, then there is no progress of the method (the point $x^k$ of the method remains fixed), and there is no convergence by default.

As such, the two first correctness points the reviewer mentioned are not really issues related to the correctness of our approach. $\gamma$ is not allowed to be zero in any of our convergence guarantees.

The third point related to our analysis of SAM in the non-convex setting was indeed a minor issue, and we appreciate that the reviewer pointed it out. As we explained below and in our updated PDF, this has an easy fix as well. In particular, in our previous draft, in Theorem 3.7 we had step-sizes $\rho=O(1)$ and $\gamma=O(1/T)$, which indeed takes care of the exponential explosion term but might not guarantee convergence. However, the fix was quite easy (see updated file, we have used red fonts highlighting the changes in the stepsizes). We just needed to choose $\rho=O(1/\sqrt{T})$ and $\gamma=O(1/\sqrt{T})$. We have updated all the statements and proofs in the new draft.

Thank you for catching this. We greatly appreciate the feedback.

Significance of Theory:

We agree with the statement of the reviewer that for SAM, one of the main challenges is how to show the method can efficiently minimize the sharpness-aware loss (either the original notion of maximal loss after a certain perturbation or the hessian-regularized version). However, in our opinion, for any optimization/training algorithm, the generalization and optimization aspects are of the same importance. Theory in terms of (i) generalization performance and minimizing sharpness-aware loss aims to understand which solution the method converges and explain how this leads to better performance in DNNs and (ii) optimization focus on the algorithm's speed to reach this stationary point.

Both are valuable aspects of any training algorithm. In this work, even if we present some generalization results in experiments we primarily focus on the optimization aspect of SAM variants and explain how fast the different variants of SAM are under realistic assumptions (no strong assumptions made in prior works).

There has been plenty of research output in the last couple of years on papers analyzing the convergence guarantees (no generalization) of the SAM-type method. See, for example, [Khanh, 2024], [Li, 2023], [Andriushchenko, 2022]. Our work belongs in this category of papers. Having said that, the two papers mentioned by the reviewer are very interesting, and results about efficiently minimizing sharpness-aware loss are orthogonal directions to our paper. We suspect combining these papers and our work would be an interesting future direction. We will cite them in the camera-ready version and clarify further the difference between optimization and minimization of the sharpness-aware loss.

We thank the reviewer for the review and positive evaluation.

Below, we address the questions raised by the reviewer.

Thanks for the suggestion. Let us highlight that we do have a table (Table 1) that has informal statements of our theorems and other related works.

How the rates change as $\lambda$ varies:

The choice of $\lambda$ indeed affects the performance of the method. To understand the connections of $\lambda$ and the two step-sizes of the algorithm $\rho$ and $\gamma$, let us focus on the PL problems, one class of problems we focus on in this work (similar statements can be obtained for nonconvex as well - but the connection are arguably more complicate to present).

We have the following dependence between the step-sizes $\rho$ and $\gamma$ of the proposed method, and the parameter $\lambda$. In particular, we have that $\rho=O(1/(1-\lambda)^2)$ and $\gamma=O((1-\lambda)^2)$, so when $\lambda$ increases from 0 to 1, $\rho$ increases while $\gamma$ decreases. As a result, with the increase of $\lambda$, the convergence of the method is slower. However, as we highlighted in Thm 3.2, any $\lambda \in [0,1]$, we always have linear convergence.

About the bounds and proofs:

Our work relaxes conditions used in previous analyses and, at the same time, provides new and improved convergence rates for both USAM and SAM. Furthermore, we extend the flexibility of sampling selection in SAM-type methods (via the unified SAM approach) as we provide an analysis under the arbitrary sampling framework that includes important sampling as a special case (sampling strategy that improves the theoretical complexity of our theorems over the more classical uniform sampling).

We would like to thank the reviewer for their time and the positive evaluation. We appreciate the comments on the strengths of our work and the characterization of having solid theoretical and empirical results.

Below, we address the concerns raised by the reviewer.

This work and previous analyses of SGD:

As we mentioned in our paper (lines 156-158) as corollaries of the main theorems on the analysis of SAM for the two classes of problems we focus on (PL and non-convex problems), we obtain the state-of-the-art convergence guarantees for SGD (for ρ = 0), showing the tightness of our analysis.

The two methods (SGD and SAM) are conceptually different. As we explained in our paper, SAM is proposed as a method for direct sharp minima during the training process, while SGD does not necessarily possess such property. This is possible as we allow having positive \rho in the update of SAM (that itself can be interpreted as a solver of the min-max problem - see line 056).

The proofs for the convergence guarantees of the two methods are substantially different from each other. For example, one important difference is that for the analysis of SAM, one needs to handle gradient norms and inner products that do not appear in the analysis of SGD (see lemmas A6 and A7 in the appendix for precise statements).

On Importance sampling:

The quantity $\max_i L_i / np_i$ naturally arises in almost all stochastic methods for solving smooth optimization problems. See for example [Gower, 2019; Khaled, 2020] for SGD and [Choudhury, 2024] for variational inequalities. So, it’s not surprising (it is actually expected) that it appears in the analysis of SAM as well. In most papers focusing on the analysis of stochastic methods, the quantity $\max_i L_i / np_i$ as a lower bound on the number of iterations required to achieve specific accuracy. As we explained in Section 3.4 this is the reason that the probability $p_i = L_i / \sum_{j=1}^n L_j$ should be used instead of the more classical uniform sampling.

On Loss Plots in experiments:

The reviewer mentioned, “since the primary focus of this paper is on the convergence properties of these methods, training loss would serve as a more relevant metric for linking the theory with empirical findings.”

For this, let us note that the first half of our experiments, in Section 4.1, Figs 1, 2, 3, are all training loss plots (exactly related to what the reviewer requested). In these plots, we verify exactly the theoretical findings of our work. Other works on the analysis of SAM-type methods have often focused on theoretical analyses and empirical results without the use of explicit loss plots; see [Li & Giannakis, 2023; Zhuang et al, 2022; Mi et al, 2022] (only present tables with generalization performance). The focus of our work is providing formal proofs and empirical benchmarks, which, in our opinion, offer a clearer assessment of SAM's performance across theoretical and practical scenarios. We agree with the reviewer, and this is precisely the reason why, in our submission, we included loss plots as well.

Typo: Thank you for catching the typo. It is fixed.

Theoretical Benefit of Unified SAM’s over U-SAM:

This is indeed an interesting question. In our work, we focus primarily on exploring the convergence of SAM, and U-SAM and their potentially interesting in-between variants ($\lambda \in (0,1)$ under relaxed assumptions and meaningful step-size selections. However, based on our current theoretical results, providing a closed-form expression of the best $\lambda$ is a challenge by itself (a different optimization problem). In this work, our primary goal is to show that more variants beyond the two most popular SAM and USAM can be analyzed and have beneficial practical performance (see experiments on Unified SAM). A possible future direction will be to describe when Unified SAM is theoretically advantageous against USAM or SAM (with strong theoretical guarantees).

Empirical Studies: In the computer vision experiments, there are cases where the training loss of Unified SAM is lower than SAM, and vice-versa. We can easily include such plots in the camera-ready version (experiments have already been run). Regarding the final question, “How to disentangle the sharpness reduction benefit,” we are unsure what the reviewer meant. Can you please clarify?

Thanks again for the review and the positive evaluation of our work. Reading your comments, we believe that all pointed weaknesses are simple clarifications. If you agree that we managed to address all issues, please consider raising your mark to support our work. If you believe this is not the case, please let us know so that we have a chance to respond.