Sharpness-Aware Minimization: General Analysis and Improved Rates
Authors: Dimitris Oikonomou
Nicolas Loizou
Venue: ICLR 2025 Poster (2025)
Keywords: Sharpness-Aware Minimization
Convergence Guarantees
Non-Convex Optimization
Generalization in DNNs
Primary area: optimization
TL;DR: Efficient convergence analysis for Sharpness-Aware Minimization under relaxed assumption and practical benefits.

Abstract: Sharpness-Aware Minimization (SAM) has emerged as a powerful method for improving generalization in machine learning models by minimizing the sharpness of the loss landscape. However, despite its success, several important questions regarding the convergence properties of SAM in non-convex settings are still open, including the benefits of using normalization in the update rule, the dependence of the analysis on the restrictive bounded variance assumption, and the convergence guarantees under different sampling strategies. To address these questions, in this paper, we provide a unified analysis of SAM and its unnormalized variant (USAM) under one single flexible update rule (Unified SAM), and we present convergence results of the new algorithm under a relaxed and more natural assumption on the stochastic noise. Our analysis provides convergence guarantees for SAM under different step size selections for non-convex problems and functions that satisfy the Polyak-Lojasiewicz (PL) condition (a non-convex generalization of strongly convex functions). The proposed theory holds under the arbitrary sampling paradigm, which includes importance sampling as special case, allowing us to analyze variants of SAM that were never explicitly considered in the literature. Experiments validate the theoretical findings and further demonstrate the practical effectiveness of Unified SAM in training deep neural networks for image classification tasks.

Published as a conference paper at ICLR 2025
SHARPNESS-AWARE MINIMIZATION: GENERAL ANALYSIS AND IMPROVED RATES
Dimitris Oikonomou
Nicolas Loizou
CS & MINDS Johns Hopkins University
AMS & MINDS
Johns Hopkins University nloi zoul @ jh. edu
doikonol & jh. edu
ABSTRACT
Sharpness-Awand Minimization (SAM) has emerged as powerful method for im- proving generalization in machine learning models by minimizing the sharpness of the loss landscape. However, despite its success, several important questions re- garding the convergence properties SAM in non-convex settings are still open, including the benefits of using normalization in the update rule, the dependence of the analysis on the restrictive bounded variance assumption, and the conver- gence guarantees under different sampling strategies. To address these questions. in this paper, we provide a unified analysis of SAM and its unnormalized vari- ant (USAM) under one single flexible update rule (Unified SAM), and we present convergence results of the new algorithm under a relaxed and more natural as- sumption on the stochastic noise. Our analysis provides convergence guarantees for SAM under different step size selections for non-convex problems and func- tions that satisfy the Polyak-Lojasiewicz (PL) condition (a non-convex generaliza- tion of strongly convex functions). The proposed theory holds under the arbitrary sampling paradigm, which includes importance sampling as special case, allow- ing analyze variants of SAM that were never explicitly considered in the literature. Experiments validate the theoretical findings and further demonstrate the practical effectiveness of Unified SAM in training deep neural networks for image classification tasks.
INTRODUCTION
Consider the classical finite-sum optimization problem
1 n n M M(x)
f(x) min WERd
where each Si is differentiable, Ly-smooth and lower bounded, Let X be the set of minimizers of f, which we assume is non-empty. In practical scenarios, the variable xr represents the model parameters, n is the total number of training instances, and the functions fi are loss functions that measure how close our model is to the i-th training data point. The goal minimize the average loss of all training instances.
Understanding the generalization capabilities of Deep Neural Networks (DNNs) is central concern in machine learning research, (Zhang et al., 2016; Hardt et al., 2016: Neyshabur et al., 2017; Wilson al., 2017: Neyshabur et al., 2018; Zhang al., 2021). The training objective function f has numerous stationary points that perfectly fit the training data (Liu et al., 2020), and each of these points can lead to dramatically varying generalization performances. This suggests that minimizing the training objective using a carefully selected optimization algorithm can lead to convergence at point with improved generalization performance (Foret et al., 2021).
Recent studies have observed that the sharpness of the training loss, that is, how rapidly it changes in a neighborhood around the model parameters, correlates strongly with the generalization error (Keskar et al., 2016; Jiang et al., 2019; Singh et al., 2025; Dziugaite & Roy, 2018). This observation has motivated recent works (Foret et al., 2021; Zheng et al., 2021; Wu et al., 2020; Andriushchenko al., 2023; Xie et al., 2024; Tahmasebi et al., 2024) aiming minimize sharpness to improve

---

Published as a conference paper at ICLR 2025
generalization. More specifically, building on these ideas Foret et al. (2021), proposed reformulating the optimization problem in (I) into a min-max problem of the following form:
min max f(x + +) ZER
where e represents the radius of the desired neighborhood. The merits of such a formulation reside in the fact that essentially we minimize the empirical sharpness measure max]|ells f()+ fxx), which inevitably will lead to flatter minima. The objective now is to find æ that minimizes ff not just at specific point but across the entire e-neighborhood. By taking the first-order Taylor expansion of ff around r and solving for the optimal s", the (Normalized) Sharpness-Awe Aware Minimization (SAM) update rule is obtained:
V/s,(a")) Pt [VSs, (
x x 7t Vss.
(SAM)
where St S (n] is a random subset of data points (mini-batch) with cardinality т sampled in- dependently each iteration L. The normalization of the inner gradient ensures that the point (a)( is good approximation of z", since || x'1 = Pt. This leads to i' = x² + Pt (wf)
more stable optimization process, as explained in Dai et al. (2023).
In an orthogonal direction and building upon SAM, Unnormalized Sharpness- Aware Minimization (USAM) was introduced by Andriushchenko & Flammarion (2022) and further investigated in Shin et al. (2024); Dai et al. (2023). The update rule for USAM is defined as follows:
x'11= 7 Vss. (a' + Pc Ve
(USAM)
In contrast to SAM, USAM omits the normalization thus the point S² can be potentially far from xf making the 2 of USAM updates much larger. This means that the removal of normalization can lead to much more aggressive steps making the USAM potentially more unstable.
Although the two variants appear closely related, the proof techniques and upper bounds used in the convergence analysis of SAM are substantially different from those in USAM. Furthermore, the convergence guarantees of the two variants vary significantly. For example, in the deterministic setting (full-batch), SAM guarantees convergence only to neighborhood of the solution, whereas USAM does not. Additionally, the step sizes "tt and used in the two update rules to guarantee convergence are very different. All of these differences motivate the importance and necessity of novel general analysis of SAM-type algorithms, unifying the two main variants (SAM and USAM) and providing the ability to design and analyze new SAM-like methods filling existing gaps in the theoretical understanding of the update rules.
In this work we develop such unified framework that allows the combination of the two approaches and, at the same time, obtains the best-known convergence guarantees under relaxed assumptions.
Main Contributions. Our main contributions are summarized below.
Unified Framework. We propose the Unified SAM, update rule that is convex combination of SAM and USAM, given by:
++ = x To Ss. ('' +++
1 1-
A VVs, (Tr)
V V.s,, ((ss()))))
(Unified SAM)
where A E 1]. The new formulation captures both USAM and SAM as special cases (A = 0 and À 1, respectively), but more importantly, it opens up a wide range of possible update rules beyond these traditional settings. The unified framework offers the flexibility to adjust the degree of normalization (using different values for X) based on specific model needs, offering a more versatile approach to SAM.
Technical Assumptions on the Stochastic Noise. Existing convergence analyses of stochastic SAM rely heavily on the bounded variance assumption, that is, there exists
0 such that
E f (z) 22 < o². (Andriushchenko & Flammarion, 2022; Si & Yun, 2023; Li & Giannakis, 2023; Harada & liduka, 2024; Mi et al., 2022; Zhuang et al., 2022) sometimes to the much stronger bounded gradient condition E || fs. (x)] < q². where q 2 (Mi et al., 2022;
Zhuang al., 2022). While these assumptions have been crucial in previous analyses, they can be
2

---

Published as a conference paper at ICLR 2025
Work
Assumptions
Arbitrary Sampling?
SAM Variant
PL. Plnnttoss
(Andriushchenko & Flammarion, 2022)
BV
USAM
(Shin et al., 2024)
Interpolation
X
USAM
(Dai al., 2023)
Deterministic
USAM
Theorems 3.2 and 3.5
ER
Unified SAM
General Non Non-convex functions
(Mi al., 2022)
BV, BG
SAM/SSAM
(Zhuang et al.. 2022)
BV. BG
X
GSAM
(Andriushchenko & Flammarion, 2022)
BV
USAM
(Li & Giannakis, 2023)
BV
X
SAM
(Si Yun, 2023)
BV
SAM
Theorem 3.7
ER
Unified SAM
Table 1: Summary of the convergence results in the SAM literature. In all works, smoothness is assumed. The top part of the table is for PL functions and the lower part is for general non-convex functions. Here BV Bounded Variance, BG Bounded Gradients.
overly restrictive. In the literature convergence analysis for stochastic gradient descent (SGD), there have been a lot of efforts recently relaxing such assumptions (Gower et al., 2019; Khaled & Richtárik, 2020: Gower et al., 2021), but, to date, no work has successfully used similar ideas for the analysis of SAM. In our analysis, we relax the bounded pradients/varianc assumptions by utilizing the recently proposed Expected Residual (ER) condition Gower et al. (2021); Khaled & Richtárik (2020). As we explain later, in several scenarios, including smooth non-convex problems, ER holds for free and allows us to provide step sizes for SAM related to the sampling strategies.
Convergence guarantees for Unified SAM. We provide tight convergence guarantees for Uni- fied SAM, for smooth functions satisfying the Polyak-Lojasiewicz (PL) condition (Polyak, 1987; Lojasiewiez, 1963; Karimi et al., 2016) and for general non-convex functions. See also Table 1 for a summary of our results and comparison with closely related works.
PL functions: For constant step-sizes y and p we prove linear convergence for Unified SAM to a neighborhood of the solution. Our theorem holds without requiring the much stronger assumptions interpolation condition or the bounded variance assumption of previous works (Andriushchenko & Flammarion, 2022; Shin al., 2024). Additionally, and Pt. the Unified SAM converges to the exact solution with sublinear O(1/t) rate (under the ER condition). Our theoretical results hold under the arbitrary sampling paradigm and, as such, can capture tight convergence guarantees in the deterministic setting. For PL functions in the deterministic setting (full-batch SAM), we show that USAM converges the exact solution, while SAM does not. This observation was first noted by Si & Yun (2023) for deterministic algorithms. To the best of our knowledge, our work the first that provides tight convergence guarantees, showing this behaviour
prove that for decreasing step sizes "
special case of stochastic algorithms.
Non-convex functions: Under the ER condition, tions Unified SAM with step sizes that depend on T (the total number of iterations) achieves E [[V][(>>)) ||| for given & at sublinear rate. This is the first result that drops the bounded variance assumption for both USAM and SAM, (Andriushchenko & Flammarion, 2022; Li & Giannakis, 2023), and substitutes it with the Expected Residual condition.
show that for general non-convex func-
Finally, corollaries of the main theorems for the above two classes of problems, we obtain the state-of-the--ar convergence guarantees for SGD (a special case of SAM with p = 0), showing the tightness of our analysis.
Arbitrary Sampling. Via a stochastic reformulation of the finite sum problem (1), firstly in- troduced in Gower et al. (2019), we explain how our convergence guarantees of Unified SAM hold under the arbitrary sampling paradigm. This allows us to cover a wide range of samplings for USAM and SAM (and their convex combination via À [o, 1)) that were considered in the literature before, including uniform sampling and importance sampling as special cases. In this sense, our analysis of Unified SAM is unified for different sampling strategies.
3

---

Published as a conference paper at ICLR 2025
Numerical Evaluation. In Section we present extensive experiments validating different pects of our theoretical results (behavior of methods in the deterministic setting, importance sam- pling, and different step-size selections). We also assess the performance of Unified SAM in training deep neural networks for multi-class image classification problems. An open-source implementation of our method is available at https: / / thub. com/ mimi wwwi--iiiiuiii uni fi edsam.
2 UNIFIED SAM WITH ARBITRARY SAMPLING
In this work, we provide theoretical analysis of Unified SAM that allows us to obtain convergence guarantees of any minibatch and reasonable sampling selection.
We are able to do that by leveraging the recently proposed "stochastic reformulation" of the minimization problem (1) from Gower al. (2019; 2021). Following an identical setting to Gower et al. (2021), we assume that we have access to unbiased gradient estimates g(z) E I
such that [[](x)] = Vf(x). For example, we can have g(z) + Lies f((z)) to be mini-batch where S S [n] is chosen uniformly at random with |S|
т. To accommo-
date any form of mini-batching, we utilize the arbitrary sampling notation g(z) = fz(()) := R Di= is [x()), where
Rn is random sampling vector drawn from distribution D such that [D|US] = 1, for i I,...,P. Then the original problem (1) can be reformulated as mingeR Ep [f(x)) == E u.f(()).. Note that it follows immediately from the definition of vz(z). Using this reformulation of the original problem, the update rule of Unified SAM can be rewritten as follows:
sampling vector that E[y(x)] = 1 - Blui/V
x+1 = =x'- - 700 (st (a'+ + Pt (1- (1-
-00
dee [g(2t\||
s(a")).
(Unified SAM)
where gxx²) DD is sampled i.i.d each iteration and P: > 7 0 and X 10. The name unified stems from the fact that the update rule indeed unifies both USAM and SAM, however, we acknowledge that there exist other SAM-like variants that our approach does not cover.
Arbirtary Sampling. Using the stochastic reformulation, the update rule of Unified SAM includes several variants of the algorithm, each related to different sampling, by simply varying the distri- bution 'D (that satisfies ED/u;] 1, Vi [n]). This flexibility implies that different choices of D lead to distinct SAM-type methods (never proposed in the literature before) for solving the original problem (I). In this work we focus on two representative sampling distributions, without aiming to be exhaustive:
Single element sampling: We choose only singleton sets S= {i} for E [m], i.e. P[[S|
1. Each number is sampled with probability P: E [0. 1] or more formally the vector E TR" defined via IP[u- = e,/ ss]p = Pt. where Σi= P: = 1. is clear that E[v;] = 1. For example, when p, = 1/n for all then this reduces to the well-known uniform sampling.
7-nice Sampling: Let e [n). We generate a random subset S [n|] by choosing uniformly from all subsets of size r. More formally, the vector E Rd is defined by PP/r = Eies es]
1/(") z!(n-72 for any subset S [n]] with |S] = 5. Using a double counting argument, one can show that E(e:] (see Gower et al. (2019)).
n!
Importantly, our analysis applies to all forms of mini- mini-batching and supports various choices of sam- pling vectors Later in Section 3.4, we provide additional details on non-uniform single element sampling strategies. addition, it clear that T in the r-nice sampling then we recover the full batch deterministic regime. Later in Section 3.2, we further demonstrate how our analysis encompasses deterministic SAM and USAM as special cases.
3 CONVERGENCE ANALYSIS
In this section, we present our main convergence results. Firstly, we introduce the main assumption for our results, namely the Expected Residual (ER) condition. Then we focus on PL. functions, where we demonstrate linear convergence rate using constant step sizes, and also provide a variant with decreasing step sizes for convergence to the exact solution. Moreover, we extend the analysis general non-convex functions. Lastly, we discuss the use of importance sampling.
4

---

Published as a conference paper at ICLR 2025
3.1 MAIN ASSUMPTION
In all our theoretical results, we rely on the Expected Residual condition.
A ssumption 3.1 (Expected Residual Condition). Let f fiul = inf f(x). We say the Expected
Residual condition holds if there exist parameters A, B, C 2 0 such that for unbiased estimator g(x) of V (x)), we have that for all i: E Rd
Ep |l(()))11 < 2A(S(z) f'r] + [IV C.
(ER)
Most prior works in the SAM literature assume either bounded gradients (e.g., Mi et al. (2022); Zhuang et al. (2022)) or bounded variance (e.g.. Andriushchenko & Flammarion (2022); Harada & liduka (2024)). Both conditions are stronger assumptions than ER. Note that the bounded gradients assumption is captured by ER for A B 0, C' > O, while the bounded variance is obtained for A = 0. B = 1, and C > O. For a detailed analysis of other conditions that automatically satisfy ER, Gower al. (2021) and Khaled & Richtárik (2020). Finally, when each fa is L--smooth under mild assumptions on the distribution D, one can show that ER holds immediately (not an assumption but property of the problem) and has closed-from expressions for the constants A, B, and C. For more details on the expressions A, B, and C in this scenario, please check Appendix B.
3.2 PL FUNCTIONS
One of the popular generalizations of strong convexity in the literature is the Polyak-Lojasieww (PL) condition, (Karimi et al., 2016; Lei al., 2019). We call a function μ-PL if [ > 2x(f(x)) f(=)),, for all x E Rd In the following, we prove linear convergence of Unified SAM for p-PL functions.
Theorem 3.2. Assume that each Sc is L--stotth, J is p-PL. and the ER is satisfied. Set Lmmax max Li. Then the iterates of Unified SAM with
# Lmmo pl ( A =
Pt p= p²:= Smazlu+2[Bu++lip
2b (B++ A)121 02 (1-2)=+1]
satisfy:
B[f(x') f(x 74)² [f(x²) - f(x"){ +N, p(1 Pmxx) [X² + C(1 - c(1-))]].
where N Lmax
1 (C^ + 2
As an immediate corollary of Theorem 3.2, we get the following guarantees for USAM and SAM, when we substitute X = 0 and X = respectively.
Corollary 3.3. Consider the same assumptions in Theorem 3.2.
USAM: For USAM with S Lm and y S L- (pp+))BB++A) it holds:
1
(B A)
S E[f(x') =(x*)((111 (1-1 71)² [f(x)) ((=)) - -[(x*)] Lmax O
(7 (7+ + p(1 2 12 P))
SAM: For SAM with Pt p S T. and = 7 2T. (Bp AB it holds:
plll Limmnn))))
E[f(z') f(x=((()))) s = " [f(x²) f(x")] )] + Lmax
(C^ + p(1 27
To the best of our knowledge, all prior convergence results for (stochastic) SAM have relied on the strong assumption (2022), Si & Yun (2023), Li & Giannakis (2023) and Harada & liduka (2024). In contrast, the above theorem is the first to establish convergence for SAM without relying on this assumption. The closest related works and Dai et al. (2023). The former provides a linear convergence rate for USAM in the interpolated regime, while the latter establishes linear convergence rate for USAM in the deterministic regime.
bounded variance, as seen in works like Andriushchenko Flammarion
the convergence of constant step size SAM for PL functions are Shin et al. (2024)
5

---

Published as a conference paper at ICLR 2025
Our result is the first to demonstrate linear convergence in the fully stochastic regime. Additionally, when p = O, Unified SAM reduces to SGD, and Theorem 3.2 recovers the step sizes and rates (up to constants) of Theorem 4.6 from Gower et al. (2021), demonstrating the tightness of our results.
Recall that with our general analysis via the arbitrary sampling framework, the proposed conver- gence guarantees hold for the 7-nice sampling well (see Section 2). As such, special case of Theorem 3.2 is the case where |S] = n with probability one. That is, we run Unified SAM with full batch, т n. In this scenario, using Proposition B.3, the ER condition holds with A 0,
B = 1 and C = 0 and Theorem 3.2 simplifies as follows.
Corollary 3.4 (Deterministic SAM). Assume that each fi is L,-smooth and f is μ-PL. Then the iterates of the deterministic Unified SAM with S 1max(1+2(1-A77) and y S
1 --mmas 2L 22 0² (11A)+++1 satisfy:
ElS(x') [(x))] s (1 (1-74)² [1(=9) [(x))] Lmaxp(1 + 1-2 +2)
/4
First, observe that the PL parameter j no longer appears in the step sizes p and y. Additionally, when À = O, i.c. USAM, the method converges to the exact solution at linear rate. However, for À 0, and in particular when suggests that even in the deterministic setting, SAM does not fully converge to the minimum. This was first investigated in Si & Yun (2023) and a similar result appear in Dai et al. (2023). We illustrate this phenomenon experimentally in Section 4.1.
1 (SAM), convergence is only up to a neighborhood. This
Finally, as an extension of Theorem 3.2, we also show how to obtain convergence to exact solution with (l/tt) rate for Unified SAM using decreasing step sizes.
Theorem 3.5. Assume that each ft is L,-stooth, f is p-PL and the ER is satisfied. Then the iterates of Unified SAM with pt min 2= and 7t min
{ 2++1 x+ where p" and
x" are defined in Theorem 3.2, satisfy: Elf(x') f(x'))
The detailed expression hidden under the big O notation can be found in Appendix C. A similar result appears in Andriushchenko & Flammarion (2022) where they provide decreasing step size selection for USAM for PL functions and prove a convergence rate of (11/t). However, their result relies on the additional assumption of bounded variance. contrast, our theorem does not require this assumption and is valid for any À [0, 1]. Notably, to the best of our knowledge, this is the first decreasing step size result for SAM.
3.3 GENERAL NON-CONVEX FUNCTIONS
In this section, remove the PL assumption and work with general non-convex functions. First, start with a general proposition that upper bounds the quantity min,= E [[V]((=*) which can serve as an intermediate result for Theorem 3.7. Our approach follows similar derivation to the analysis of SGD in the same setting by Khaled & Richtárik (2020).
Proposition 3.6. Assume that each fo Li-smooth and the ER is satisfied. Then the iterates of
Unified SAM with S min Thmas 3BLmaa (1- and S 8BLmax. sutisfy:
2(1 + 211+2Ay1 2AyL [p 2) 00<11277p2 +7))²
min t=0
E [
[f(=²) pint]]
Ty + 2Laax [Cy + p(1 + m (12 (x² CC(1 ([11))))]-
In order to control the coefficient of the term f(z²) fin' of Proposition 3.6 from exploding we need to carefully select the step sizes and y. This is what the following Theorem 3.7 achieves.

---

Published as a conference paper at ICLR 2025
Theorem 3.7. Let E > 0 and set dg = f(x0) pinl 2 0. For = and 7 = 7," provided that
6aLmax g2
96B, 24(1 - 2) -) /LLmax 5184Lmax A2 2) 8646.A
144C
T T
max
g²
E²
E²
288L 2 max (1 2)21
g2
the iterates of Unified SAM satisfy min VVfx')) Se.
"For the precise expressions of p and 7 we refer to Theorem C.4.
The results in Proposition 3.6 and Theorem 3.7 tight, setting p 0 Unified SAM reduces in SGD and these simplify to the step sizes and rates (up to constants) of Theorem 2 and Corollary from Khaled & Richtárik (2020). Other results for general non-convex functions can be found in Mi et al. (2022), Zhuang et al. (2022) and Li & Giannakis (2023) for SAM and in Andriushchenko & Flammarion (2022) for USAM. However, as mentioned earlier, all these analyses rely on the strong assumption of bounded variance and/or bounded gradients. In contrast, our result uses weaker sumptions and offers guarantees for any E [0, I]. Furthermore, Khanh et al. (2024) have results asymptotic. for general non-convex functions in the deterministic setting though all their results Another closely related work Nam et al. (2023), where they also assume the expected resid- ual condition, however, they additionally assume bounded gradient of f and their results are only asymptotic and hold almost surely.
3.4 BEYOND UNIFORM SAMPLING: IMPORTANCE SAMPLING
In this work, we provide theorems in the PL and non-convex regime through which we can analyze all importance sampling and mini-batch variants of Unified SAM (and as result of USAM and SAM). We achieve that via our general arbitrary sampling analysis. In this section, we focus on the non-convex regime (Theorem 3.7) and provide importance sampling for Unified SAM with single- element sampling. That is, we select probabilities that optimize the iteration complexity (lower bound of T). In particular, to derive the importance sampling probabilities, we substitute the bounds for A, B, and C' from Proposition B.2 into Theorem 3.7, resulting in:
5184 (max, L, 22) 2) n²r2
3 12T
L
T
da Lenax g2
max
24(1 - 1
max
Pi
8646g
max; P:
L.
288
max;
La o
288
2 max
- A)2
g2
22
c²
Having substituted the values of A, B, and C' in the above expression, to obtain the best bound, we require optimizing the quantity max; Ln over the probabilities p;. This results p; Σー r. ap;
Similar probabilities have been proposed for several optimization algorithms, including SGD (Gower et al., 2019; 2021; Khaled & Richtárik, 2020), variance-reduced methods (Khaled et al., 2023), and methods for min-max optimization (Gorbunov al., 2022; Loizou al., 2020; Choud- hury et al., 2024; Beznosikov et al., 2023). To the best of our knowledge, our work is the first to provide importance sampling for SAM algorithms.
4 NUMERICAL EXPERIMENTS
In this section, we evaluate our proposed step sizes for Unified SAM on both deterministic and stochastic PL problems, with experiments designed to illustrate our theoretical findings. Addition- ally, explore different values of À when training deep neural networks to improve accuracy.
4.1 VALIDATION OF THE THEORY
In this part, empirically validate our theoretical results and illustrate the main properties of Unified SAM that our theory suggests in Section 3. In these experiments, we focus on fz- re-reuulri regression problems (problems with strongly convex objective ʃ and components fa and thus PI.)
7

---

Published as a conference paper at ICLR 2025
and we evaluate the performance of Unified SAM synthetic data. The loss function of the Eg- regularized ridge regression is given by
1 - Az b12 + X ||/1
1
n (A 1/= b,)²
Ar
f(x) -
||
2n
2
and the loss function of the fz-regularized logistic regression is given by
f(z) = 2n log (1 exp (-b, (-b; (a)) + 2 |
In both problems A E Raxd b RR are the given data and A is the regularization parameter.
Normalized SAM converges only in neighhorhood. In Corollary 3.4 we highlight that our theoretical results in- dicate that in the deterministic case, USAM achieves lin- ear convergence
the exact solution. In contrast, when > O, and specifically for SAM (A 1), convergence only to a neighborhood of the solution, an unusual out- come for deterministic optimization methods. We vali-
Am02
A-033
A=0.4
date this observation experimentally Figure 1. We run a ridge regression problem with n 100, d 100, X,- O. The matrix A has been generated according to Lenard & Minkoff (1984) such that the condition number of A 10 and the vector b has been sampled from the stan- dard Gaussian distribution. We have used the determin- istic Unified SAM for À = 0.0, 1.0 and we run each algorithm for 50 epochs. Indeed
A-07
A=
A
Iterations
Figure 1: Deterministic Unified SAM for various values of À applied to the ridge regression problem. USAM (A converges the exact solution while the other variants À > 0 converge to neighborhood of the solution.
can see that USAM (A = 0) converges all the way to the exact solution while the other choices of converge to a neighborhood. It is also noteworthy that the neighhorhood increases as À ap- proaches 1.
Constant vs Decreasing Step size. In this part, we compare the performance of Unified SAM under both constant and decreasing step size regimes, as described in Theorem 3.2 and Theorem 3.5, respectively. For this experiment, we consider a logistic regression task with and A, = 3/n. As before the matrix has been generated such that its condition number is 10 and the vector b has been sampled from the standard Gaussian distribution. We run Unified SAM with = 0.0, 0.5, 1.0, using uniform single-element sampling for 10, 000 epochs across trials. The average results of these trials, along with standard deviation, shown in Figure 2. Initially, the trajectory of the decreasing step size follows that the constant step size. However, as the constant step size version of Theorem 3.2 approaches neighborhood near optimality, it stagnates. In contrast, the decreasing step size from Theorem 3.5 allows for continued improvement, leading better overall accuracy.
100, d 100.
Constant Step-size
Constant Step-size
Decreasing Step-size
Decreasing Step-size
Constant Step-size
Decreasing Step-size
Inerations
Iterations
herations
Figure 2:
Comparison between constant and decreasing step size regimes of Unified SAM. From left to right have À 0.0, 0.5, 1.0
Uniform vs Importance Sampling. In this experiment, we demonstrate the benefit of using im- portance sampling compared to uniform sampling. We consider a ridge regression problem with 100, d = 100, and X, 3/n. The cigenvalues of matrix A are sampled uniformly from the interval [1.0, 10.0], and the vector is drawn from a standard Gaussian distribution. We run Unified SAM with À 0.0, 0.5, 1.0, employing single-element uniform sampling for 3, 000 epochs across
8

---

Published as a conference paper at ICLR 2025
Uniform Sampling
Uniform Sampling
Uniform Sampling Importance Sampling
importance Sampling
importance Sampling
cha 100
lterations
Iterations
kerations
Figure 3: Comparison between uniform and importance sampling for Unified SAM. From left to right we have X = 0.0, 0.5, 1.0
five trials. For importance sampling, the probabilities are set as P. = L/ Ly. The averaged results with one standard deviation are presented in Figure 3. The results clearly show that impor- tance sampling enhances the convergence of Unified SAM for all values of X, which is consistent with the discussion in Section 3.4.
4.2 IMAGE CLASSIFICATION
In this section, we evaluate the generalization performance of Unified SAM using the classic image classification task. The models are trained on the CIFAR-10 10 and CIFAR-100 datasets, (Krizhevsky al., 2009). For data augmentation, we apply standard techniques such as random crop, random horizontal flip, and normalization (DeVries, 2017). All experiments are conducted on NVIDIA RTX 6000 Ada GPUs.
Models. We use several models in our experiments, including ResNet- 18 (RN-18) (He et al., 201 6a), and PreActResNett-18 18 (PRN-18) (He et al., 2016b). To demonstrate the scalability of Unified SAM, also include Wider ResNet (WRN-28-DD) 10) (Zagoruyko & Komodakis, 2016).
Different values of X. This subsection explores how varying À in the Unified SAM update rule af- fects generalization performance. We experiment with the PRN-18 and WRN-28- 10 models trained on CIFAR-10 10 and CIFAR-100. The Unified SAM method is trained using p = 0.1, 0.3, 0.4 and A 0.0, 0.5, 1.0, 1/1, 1/1. The last choices of X can be intuitively explained follows: for A = 1 1/1, the algorithm starts as USAM (A = 0) and gradually transitions toward SAM A 1 when L. - Dc, meaning it begins with USAM and approaches SAM. Conversely, for A 1/t, the algorithm behaves the other way around, starting closer from SAM and converging to USAM over time. Following Pang et al. (2021) and Zhang et al. (2024), we set the weight decay to 5 x 10-4, the momentum to 0.9, and train for 100 epochs. The step size 7 is initialized at 0.1 and reduced by a factor of 10 at the 75-th and 90-th epochs.
Each experiment is repeated three times, and we report the averages and standard errors in Tables 2 and and Tables 8 and not always À 0 or À 1. Overall, X 1 1 1/t appears to be a reliable choice. In particular, we observe that USAM (A = 0) does not have the best performance in any of these experiments. On the contrary in the CIFAR10 10 dataset the value A 1 - 1/t is good choice while in the CIFAR100 dataset both A = 1 and A = 1/t offer strong performance.
appendix. The results indicate that, for fixed p. the optimal À value
Table 2: Test accuracy (%) of Unified SAM for WRN-28- 10 on CIFAR10, 10, evaluated across different values of p and X. With bold we highlight the best performance for fixed p.
Unified SAM
A= 0.0
X 0.5
A= 1.0
A 1/6
1 1/6
p = 0.1
95.7
95.68
95.9
95.8410.m
95.81
A = 0.2
95.810.05
05.77.1..0
95.93
95.71 10.13
95.98 10.1
P 0.3
95.35
95.88
95.95
95.68 10.02
95.99
0 = 0.4
95.4640.02
95.7640.1
95.62
95.4640.27
95.79=0.07
SGD
95.35
Unified VaSSO. The Variance-Suppreste Sharpness-Awan Aware Optimization (VaSSO) method, intro- duced in Li & Giannakis (2023), is extension of SAM designed to reduce variance in gradient esti- mates. VaSSO adjusts the direction of gradient updates using a combination of past gradients and the current gradient, controlled by parameter aiming suppress noise during training and enhance

---

Published as a conference paper at ICLR 2025
Table Test accuracy (%) of Unified SAM for WRN-28-10 on CIFARIO0, evaluated across differ- ent values of and X. With bold we highlight the best performance for fixed
Unified SAM
A= 0.0
À = 0.5
1.0
A=
A --1-10 1/t
P 0.1
80.84.0.08
81.01
80.6940.11
S0.81
80.88+0.31
0.2
81.12+0.18
81.45
81.53+0.09
81.31
81.22+0019
P 0.3
80.9440.13
81.6440.16
81.62+0.16
81.0340.14
81.71+00.7
/ 0.4
80.1 10.22
81.70 10.06
80.39
81.59
SGD
79.79
generalization, particularly in overparameterized models. In their work, they prove that VaSSO con- verges the same asymptotic rate as SAM, under the bounded variance assumption. The update rule of VaSSO is defined as follows: d, = (1 - 0)d,-1 + fz'),)),, xt = x² r V (x" + Pt || can incorporate our unification approach, initially designed for SAM, into VaSSO. This leads to the following modified update rule:
dj = (1 0) f)d4-1 + f((''')
r+1 = _. 7h V x' + Pt (1 (1->e
Xe
+
de
(Unified VaSSO)
||
Note that when A = 1 then we recover VaSSO as introduced Li & Giannakis (2023).
In this section, we conduct experiments to evaluate the performance of Unified VaSSO. All models are trained for 200 epochs with batch size of 128. cosine scheduler is employed all cases, with an initial step size of 0.05. The weight decay is set to 0.001. For VaSSO, we use 0 = 0.4, this value provides the best accuracy according to & Giannakis (2023). For the CIFAR-10 dataset, we set p = 0.1, while for CIFAR-100, we use p = 0.2. Each experiment is repeated three times, and we report the average of the maximum test accuracy along with the standard error. The numerical results presented in Tables 4 and 5.
The results show that, for the WideResNet-25 accuracy. For the ResNet-18 model the best values for À appear to be À 0.5 and À = 1/t. Thus, the choice A = 1 - 1/t appears to be a strong overall choice. In any cases, a Å value different from achieves the best performance.
model the value A = 1/t produces the best
Table 4: Test accuracy (%) of Unified VaSSO various neural networks trained CIFAR1O.
Model
A 0.0
À 0.5
A = 1.0
A = 1/t
A 1
R8
96.12=0.02
96. 90.00 96.93
96.22 97.03
96. 00.00
96.34_0.01 97.06
WideResNet-28-10
96.77 =0.07
96.71
Table 5: Test accuracy (%) of Unified VaSSO on various neural networks trained on CIFAR100.
Model
=
À = 0.5
>= 1.0
>= 1/t
X 1- 1/t
ResNet-18
79.82 83.07
80.01
10.07
10.07
79. 79.76.10.u
79.93
79.86 83.6610.11
WideResNee-28-10 10
83.29_0.33
83.51
82.83
5 CONCLUSION
In this work, we introduced Unified SAM, an algorithm that generalizes SAM and USAM via the convex combination parameter À E [0, 1]. The proposed analysis relaxes the common bounded variance assumption used in previous works by employing the ER condition, which enables us to prove convergence under weaker assumptions. In particular, under the ER, we established conver- gence guarantees for Unified SAM for both PL and general non-convex functions, and our analysis encompasses importance sampling and tau-nice sampling. Future work could investigate the opti- mal value of for various classes of problems, extend the unified approach to incorporate adaptive step sizes Pi and 7t. and study Unified SAM in distributed and federated learning settings. Another promising direction to explore how our approach may help minimize the sharpness-awe aware loss by employing the Hessian-regularizeed methods proposed by Wen al. (2023) and Bartlett al. (2023).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awar mini- mization. In ICML,, 2022.
Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-aware minimization leads to low-rank features. In NeurlPS, 2023.
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness- aware minimiza- tion: Bouncing across ravines and drifting towards wide minima. Journal of Machine Learning Research, 24(316):1-36 2023.
Aleksandr Beznosikov, Eduard Gorbunov, Hugo Berard, and Nicolas Loizou. Stochastic gradient descent-ascent: Unified theory and new efficient methods. In AISTATS, 2023.
Sayantan Choudhury, Eduard Gorbunov, and Nicolas Loizou. Single-call stochastic extragradient methods for structured non-monotone variational inequalities: Improved analysis under weaker conditions. NeurIPS, 2024.
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-awan aware minimization. In NeurlPS, 2023.
Terrance DeVries. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: 740.0055,, 2017.
Gintare Karolina Dziugaite and Daniel Roy. Entropy-sgd optimizes the prior of a pac-bayes bound: Generalization properties of entropy-sgd and data-dependent priors. In ICML., 2018.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza-
tion for efficiently improving generalization. In ICLR, 2021.
Eduard Gorbunov, Hugo Berard, Gauthier Gidel, and Nicolas Loizou. Stochastic extragradient:
General analysis and improved rates. In AISTATS, 2022.
Robert Gower, Othmane Sebbouh, and Nicolas Loizou. SGD for structured nonconvex functions: Learning rates, minibatching and interpolation. In AISTATS, 2021.
Robert Mansel Gower, Nicolas Loizou, Xun Qian, Alibek Sailanbayev, Egor Shulgin, and Peter Richtárik. SGD: General analysis and improved rates. In ICML, 2019.
Hinata Harada and Hideaki liduka. Convergence of sharpness-aware minimization algorithms using increasing batch size and decaying learning rate. arXiv preprint 4rXiv:2409.09984, 2024.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In JCML, 2016.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In ECCV, 2016b.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In ICLR, 2019.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-lojasiewicz, condition. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD, pp. 795-811, 2016.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In ICLR, 2016.
Ahmed Khaled and Peter Richtárik. Better theory for sgd in the nonconvex world. arXiv preprint vrXiv:2002.03329, 2020.
Ahmed Khaled, Othmane Sebbouh, Nicolas Loizou, Robert Gower, and Peter Richtárik. Unified analysis of stochastic gradient methods for composite convex and smooth optimization. Journal of Optimization Theory and Applications, 199(2):49--540, 2023.
Pham Duy Khanh, Hoang-Chal Luong, Boris S Mordukhovich, and Dat Ba Tran. Fundamental convergence analysis of sharpness-aware minimization. arXiv preprint 200iv:2401.08000 2024.

---

Published as a conference paper at ICLR 2025
Alex Krizhevsky al. Learning Multiple Layers of Features from Tiny Images. Technical report, University of Toronto, 2009.
Yunwen Lei, Ting Hu, Guiying Li, and Ke Tang. Stochastic gradient descent for nonconvex learn- ing without bounded gradient assumptions. IEEE transactions on neural networks and learning systems, 31(10):4394 4400, 2019.
Melanie L Lenard and Michael Minkoff. Randomly generated test problems for positive definite quadratic programming. ACM Transactions on Mathematical Software (TOMS), 10(1):86-96, 1984.
Bingcong Li and Georgios Giannakis. Enhancing sharpness-aware optimization through variance suppression. In NeurIPS, 2023.
Shengchao Liu, Dimitris Papailiopoulos, and Dimitris Achlioptas. Bad global minima exist and SGD can reach them. In NeurIPS, 2020.
Nicolas Loizou, Hugo Berard, Alexia Seiicoeur-Martiecnn Pascal Vincent, Simon Lacoste-Julies and loannis Mitliagkas. Stochastic hamiltonian gradient methods for smooth games. In JCML, 2020.
Stanislaw Lojasiewicz. A topological property of real analytic subsets. Coll. du CNRS, Les équations aux dérivées partielles, 117(87-89)22, 1963.
Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. In NeurIPS, 2022.
Kyunghun Nam, Jinseok Chung, and Namhoon Lee. Almost sure last iterate convergence of sharpness-aware minimization. 2023.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general-
ization in deep learning. In NewrlPS, 2017.
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. To- wards understanding the role of vvrr-paaametrization in generalization of neural networks. arXiv preprint arXiv: / 2018.
Tianyu Pang. Xiao Yang. Yinpeng Dong, Hang Su, and Jun Zhu. Bag of tricks for adversarial
training. In ICML., 2021.
Boris T Polyak. Introduction to Optimization. New York, Optimization Software, 1987.
Sungbin Shin, Dongyeop Lee, Maksym Andriushchenko, and Namhoon Lee. Analyzing sharpness- aware minimization under overparameteriizttoon arXiv preprini arXiv: 1..1539,,,, 2024.
Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. In NeurlPS, 2023.
Sidak Pal Singh, Hossein Mobahi, Atish Agarwala, and Yann Dauphin. Avoiding spurious sharpness minimization broadens applicability of sam. arXiv preprint arXhv:2502.02400 2025.
Behrooz Tahmasebi, Ashkan Soleymani, Dara Bahri, Stefanie Jegelka, and Patrick Jaillet. A univer- sal class of sharpness-aware minimization algorithms. In ICML, 2024.
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharp-
ness? In ICLR, 2023.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stem, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In NeurlPS, 2017.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust gener- alization. NewrIPS, 2020.
Wanyun Xie, Thomas Pethick, and Volkan Cevher. Sampa: Sharpness-awar aware minimization paral- lelized. In NeurIPS, 2024.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In British Machine Vision
Conference, 2016.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv: /6/ 1611,03 2016. Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
12

---

Published as a conference paper at ICLR 2025
Yihao Zhang, Hangzhou He, Jingyu Zhu, Huanran Chen, Yifei Wang. and Zeming Wei. On the duality between sharpness- -ware aware minimization and adversarial training. ICML, 2024.
Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8156-8165 2021.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvomek, Sekhar Tatikonda, James S. Duncan, and Ting Liu. Surrogate gap minimization improves sharpness- aware training. ICLR, 2022.
13

---

Published as a conference paper at ICLR 2025
Supplementary Material
The Supplementary Material organized as follows: In Appendix A, we give the basic definitions and lemmas that need for the proofs. Appendix C presents the proofs of the theoretical guarantees from the main paper. In Appendix D, we provide additional experiments.
CONTENTS
1 Introduction
2 Unified SAM with Arbitrary Sampling
4
3 Convergence Analysis
4
3.1 Main Assumption
3.2 PL functions
General non-convex functions
Beyond Uniform Sampling: Importance Sampling
Numerical Experiments
4.1 Validation of the theory
4.2 Image Classification
5 Conclusion
10
A Technical Preliminaries
15
A.1 Basic Definitions
15
A.2 Basic Lemmas
15
More on the Expected Residual condition
18
Proofs of the main results
19
Proof of Theorem 3.2
19
C.2 Proof of Theorem 3.5
22
C.3 Proof of Proposition 3.6...
23
C.4 Proof of Theorem 3.7
24
D Additional Experiments
27
14

---

Published as a conference paper at ICLR 2025
A TECHNICAL PRELIMINARIES
A.1 BASIC DEFINITIONS
In this section, we present some basic definitions we use throughout the paper.
Definition A.1 (L-smooth). differentiable function constant L > such that
E IR is L-smooth if there exists
(Vf(z) - - D|| - yll,
(2)
or equivalently
L 2 -11= y112
f(x) f(v) s (Vf(y). 5)
for all z, E Rd.
Lemma A.2. Let f be L-smooth finctionn Then for all R" we have
[V > 2T[f(x)) finf].
Definition A.3 (μ-PL.). We say that a differentiable function f Rdd - satisfies the Polyak- Lojasiewicz condition if
[V/(2)] 2 2n(s(x) - f(x)))
(3)
for all x e Rdd.
Definition A.4 (Interpolation). We say that the interpolation condition holds if there exists x" E X such that
min zER"
fi(x) = fi(x*)
for all i I,...,n.
Proposition A.5 (Young's Inequality). For any a, E R" and 8 + we have
1
3
(a, 6)
|||2
2 1161122
(4)
23
Completing the square we have
||a + b112 < (1 +1+ 8- |a|l? + (1 + 1)1b11
(5)
For 3 1 we gel
|la b112 < 2|aa|1 + 211B112
(6)
A.2 BASIC LEMMAS
Lemma A.6. Assume that each fü is L--smooth and consider the iterates of Unified SAM. Then holds
ED 19 a0' + Pt (1-2) ~x [g(x'(]| m(xz)) )1
A
< 4T.2
maxx Pi 2x
+ 2 222 m Pi (2 )x)2 + 1] Ep 11.
15

---

Published as a conference paper at ICLR 2025
Proof. We have
de
glot)) )}
DD
9
x² Pr
1 1->e+
Xe
||o(
ED
x² + Pt:
1 x (19(
d
((t7)
-x(x9) + g(xx)) 2
(6)
S 2 LDD
9
x² + Pt
1 11- x1 + |[a(zt) s(60))
-g(xt)
2 ||
(2)
ED [1 >e)g(a>))
de
2 2L m BBD
[لa( ((()))
2 ||
_ 4mman p2 (1 21)² ED (1g 112 + 4L2 2 pi X3 + 22 ||
4L2 max 2. m + 2 [2 [2L (1 x₂)² + 1] ED ||9(,))1
This completes the proof.
Lemma A.7. Assume that each Si is Li-smmoth and consider the iterates of Unified SAM. Then for any B > 0 it holds
ED ( 9 x' + ps (1->: [g((a))|| (omm)))ol))))))) s(6t)) vf(at) 2 (1- (1-2) 2
d
L P?
2 | [V
3
I P (1 xx)2
ED ||9(2")113
B
In particular. for B Limaxpe we get
50 9 aa + pr: (1- de + |[((())||
d
oloe)) v)(aa) >
LmaxpPt
ED
[(({)))
2
Pe
Temax (1 >x) ED (1g(x")11
Proof. We have
ED 9 x0 + Pe 1 1- de
de
o(a')) Vf(a')
|19(0()1|
Ep 9 x² + p: (1-
((6)) -((x)), V((())
+ Eo(g(x'), ) vf(x'))
ED 9 x¹ + Pt 1- X +
[19(())|| ((999 -0(0"), ((6))) + +[V
Now
de
)(ot)) - g(x'), Vs(at)
ED
9
xt
+ Pt
1 de
||9(x>)|| \11
(4)
1
ED
9
x¹ + Pr
11 de
((e')) g(
2
3
+
Ep [x(()))
2
(2)
<
I2 P
Ep
(1 -x)e)((x=)= +
|[(()))| s(at)
d,
2
+
\f(''')
233
2
(6)
L ?
3
2(1
2x)²
RD 119(x
2x
+
|(}(}})
23
23
2
16

---

Published as a conference paper at ICLR 2025
I P
I maxx Xx)2
ED [|]((=))1++
3
2
thus
Ep ( gola* + ps (1- de +
||g(z))|| || g(x')), V5(a4) 2
L max pi (1 Xx)2
ED ||
L
m
|V
B
2
+
[V
8
I ? X
[Vf(x')1
2
B
L L
(1 (1
Ep ||
This completes the proof.
Lemma A.8. Let (rx):20 and (ōz)e20 sequences of non-negative real numbers and g >
and N 2 O. Assume that the following recursive relationship holds:
S gō. - sc41 +N
(7)
Then it holds
9" S d0 +
min 0StST
T
Proof.
Set 9t = g for any t Z. Then multiply both sides of (7) with gt. This yields
ge' S It-1d: g:0 Nge-
Summing for f O.... T and telescoping get
Σ geft S 9-100 0r + N gt
Σ
t=0
0- + Σ
gt.
(8)
t=0
Now let G= Et=o gr. Note that the sequence (g:) is decreasing for t 2 0, thus G 2 T9r-1- Now dividing (8) by G we get
T
min 051ST-1
<
G
1=0
9-100
+ N
G
9-100
+ N
Tgr-1
go
T + N.
This completes the proof.
17

---

Published as a conference paper at ICLR 2025
B MORE ON THE EXPECTED RESIDUAL CONDITION
In this section, we introduce several assumptions under which ER automatically satisfied. Similar derivations are provided in Khaled & Richtárik (2020), where the authors focus on the analysis of SGD. Specifically, we consider the following cases:
Assuming bounded gradients, i.e.
Ep 19(2)) 5
Vo E
then the ER holds with = 0, B = -
Assuming bounded variance. i.e.
ED (1)(x)) Vx(i)]] so²
Var E
then the ER holds with = 0, =, = o².
Assuming expected smoothness, i.c.
Ep ll((2) - vf(()) s 2Zl - - Vr ErER" then the ER holds with A = 2C, B 0, ==
Assuming the relaxed strong growth condition, i.e.
ED |=(+)) - /|[ + Vz E
then the ER holds with A 0, B = C.
Assuming the relaxed strong growth condition, i.e.
Ep |=(())] alf(x) - - 0², 4r Ve
then the ER holds with A a, B =
Additionally, if we have more information about the problem and the distribution D, derive stronger bounds on A, B, and C. The more assumptions we make, the tighter these bounds become, demonstrated in the following propositions.
Proposition B.1 (Prop. 2, (Khaled & Richtárik, 2020)). Assume that each fi is Li-smooth and that ED[(?] < for all E [m.. Let o" Σ f1(=") - f3 0, where f inf fs(x). Then ER holds with A max; L, ED (uf). B and C 2Ao
This proposition indicates that each fa smoshh and minimal assumptions hold for the distri- bution D, then ER is satisfied. As an immediate corollary of Proposition B. if the problem further satisfies the interpolation assumption (see Definition A.4), then = The next proposition gives tighter constants for A, B, and C' in the context of the sampling strategies considered in this paper.
Proposition B.2 (Prop. 3, (Khaled Richtárik, 2020)). Assume that each fs T.-sshoth.
For the single element sampling ER holds with A
nT max; Le = 0 and C' -AA0*
2. For the 7-nice sampling ER holds with ((n-1) Tmaxx B where Luuxx max; Li-
a(7 7nn-1)
and C 2AA*
Lastly, if assume x" -onveexity with 7-nice sampling, we obtain the following constants:
Proposition B.3 (Prop. 3.3. (Gower et al., 2021)). Assume that each fo is Ly-ohh and
that there exists x¹ X such that fs is x² convex. In addition, assume that D is the n-F
2(n-T) nice sampling. Then ER holds with A r(n -1n-1) Lmaax B and C r(n-1) T1, where
σ1 EX 1 ET [[+)((
18

---

Published as a conference paper at ICLR 2025
C PROOFS OF THE MAIN RESULTS
In this section we present the proofs of the main theoretical results presented in the main paper, i.c., the convergence guarantees of Unified SAM for PL and smooth functions and general non-convex and smooth functions. We restate the main theorems here for completeness.
In all cases, we use the convention 1/0 oc.
C.I PROOF OF THEOREM 3.2
Theorem C.1. Assume that each fi is Lmmoohhh is p-PL and the ER is satisfied. Set Lmax max Li Then the iterates of Unified SAM with
2 =7 # L-mapp B B A (1-2)² Pt = ps p² = Timas +2 2 1111-133 2)27
2T. BB A) 22 ana 1² AJ
satisfy:
Elf(x') - (((=()-f(xx))) $ (1 711² [f(x')] - f(x))] + N,
where N
Linax #
(Cr + p(1 (11771mmxxx)) m> [X² + C(1 C(1-1)²) 2)²]).
Proof of Theorem 3.2. By combining the smoothness of function f with the update rule of Unified
SAM we obtain:
Lmax | - x" 2¹12
f(x s f(x') + (Vf(x*). x 1+1
=²") +
2
Ae
=(xx)) 71
Vf(x').
x²' + pe
1-
|19(2t)]| a(2²))
Lmxx
9 2² + Pt
1- -e
+
9((2)) -|
2
By taking expectation conditioned on x² we obtain:
() f(x")|=²| - -[f(x') - f(x")]
S - ED (9 (0(2²= - Pr 1 1 de |[{(x!)|| s(x0)) (.(((t)
de
+ Lima ED 9 x² + Pt (1 (1- >e To[ar)]| g(at)) )
de
Lem. A.6 and A.7
((1
LmaxPt
Lmaxped? - Lmaxpe Ax)² ED ||122
5
2
7/ Limx
+
[AT X2 + 22r2 2
(1- - Ax)² + +1]Ra ED (19(2
= =-70 (1- (1- LmmxxPt
|Vf(2')11
2
+ ntt Temax
[px(1 A1) 2 1.2
2: Xc)2 + Ep 1|
+ 9t Tmmxx
(p(1 23 I2 ['
ER
Timaxpt
5
7/
|f(x')11
2
2A
[pc(1 Xx)2 + 2 maxe pi 21 (1 Xx)3 + 7 [s(x') (((=))(((')))
Limax
[pr(1 - 21)2 + 2 12 m pi (1 - Ax) + 7o] [V((2))11
+ Che.
[1t(1 - 21)² - 1 (1 - >x)² + +7] + 71 JenaxaaPPPae P2 + 2 1.2 p
(1 11- LmaxPt BLmaax [pe(1 - 2₂) + 2 12 2 >x) + r8l [[
=-7
2
19

---

Published as a conference paper at ICLR 2025
+ 2 2 Lmax (pr(1 21)22x220200
+ max (1- X1)2 + +7 [f(x") f(xx)]
+ Thmma [Cpc(1 Xx)² 2-000 1 pax Pi Ax)2 + Cre ptd2 27t maxt A
In order to use the fact that is p-PL we need to ensure that the coefficient of |V is non- negative. For this we have
1 ImaxPt
B Lmmx (p - -20² + 2 -\1)² + +% 2 0 +
2
LmaxPt + Bl m1 - -2)² + B Blmaxx maxx Pi Xx) + 1
Solving for "/t we get
LmaxPe Lex 21)2 7t S 2BL (2I maxh X,) + 1)
(10)
The numerator of the upper bound on Tt in (10) must be positive (positive step-size). For this, the following restriction on Pt is needed.
2
2 Lmax Pr- 2BL - 20 Lmux + (1+2B(1 Ax)3))
(11)
Now using the 4-PL condition (3) on (9) we get
B[f(x - -(==")x''-[(f()) f(x")]
< -2/1 (1-
ImnuxPt BLmax [pe(1 (1 - X1)² + 2% I2 2 -x0² twil - [f(x') - f(x*)]
2
+ -2A [pa(1 - >x)²² + 27r 1.2 max (1 - Ax)² + nD [f(x²) - [[xx)]
+ 7/ Temax [Cpe(1 21)² 20 Pa 2x)² C. + p. + 27 27 M
hence
El +1 - -f(x")
< 1 27/e/4 (1- LmaxPt
BLmax [pr(1- 02)2 + 2 (1 (1--))) 2x)² + +mo]
2
+ 22 Lmax [pe - 21)² + 2; L2 2naxx of ~2)² + +72 [f(x*) - f(x')]
+ 7 Lmax [Cpc(1- - -x))² 20 12 2
(1 xx)² OTe + +m + X2]] (12)
The coefficient of [f(x)) f f(x") can be simplified to
1- 2%144 (1
LmaxPt
2 BLmax [pr(1 2,)2 + 2n, 1.2 2 (1 X,)2 70]
+ Lm [p((1 21)2 + 2% L2 Lm
22 (1 2z)² + 7
2T//4 + MuLLmapPt + BB Liax
[pr(1 - 2x) + 2 22
Po (1 - -xx)² + ml
+ 222 [pr - >x + 27, 12 12xx
/2 (1 - >x,)2 + 7r]
- 27p Tm + 27, I Imax Bμ + A) [p:[1 ->,)² 2 1 2x) + 7?.
27444 Li (μ + 2(B A) 2) + + Lmax (Bμ + A) 22 >x)² 1)
20

---

Published as a conference paper at ICLR 2025
Let us bound the above coefficient by 1 1/// (used for the final linear convergence). For doing that, we force the following restrictions for P: and T/::
20 + Tmmaxp (μ + 2( p A) )x)²)
+ +2 Lmax (Bu + A) max (1 - Ax) + +1) S 1 -1-1-4 +
Tt maxPrt (1 2(B + A) - )²)
+ +27 Lmax [Bp + A) (2 2 (1 >x)² +1) S Ti 2
Imaxpt (1z 2(B + A) )1)²
+ +27 Bu + A) (2L P
+ +1)
2/ Lmxx Bu + A) (2L max Pi 21)² +1) S k - -maxpt (p + 2(B A)(1 - -11) +
TemaxPt Pt. (1+ + 2( B/ + A) - >x)3
N1 S 21max (Bu A) (2L m >x)² 1)
(13)
The numerator of the upper bound on in (13) must be positive (as the step-size is positive). Thus:
p - LmaxPr (n + 2(B + A)(1 X₂) N Pt S Lmax (u + 2(Bu + A)(1 Xx)3)
P
(14)
Now by combining Equations (11) and (14) obtain:
Pt p= = min
2
1 Linax (H + 2(B A)(1 x,)22
Lixxx
+ 2B2 X1)2)'
11
max (p + 2(Bu A)(1 -1)3²'
where last equality holds due to:
Lmax (μ + BB - A)(1 - >1)2) Lmax (1 + 2BB1 - >x)2)
T.maxft + 21 mmm By(1 21)² S 2I + 4 4Tmax BBp + A)(1 - Ax) +>
0 < Lmaxti + 2L (Bu + 22) - -x0
Similarly for & by combining Equations (13) and (10) obtain:
2 Lmxx 2B Limax 22)²
-- Lmmxxt (H + 2(B + A)(1 - Ac)3)
S 5-- = min
2B (21 2 pf 22)2 +1) 1)
2L (Bμ A) (2L2 m (1 X1)2 11)
Lmaxpr (μ + BB + A)(1 Xx)²)
2Lmax (Bp + A) m (1 )²)² + 1))
where the last inequality holds due to:
11 Lmaxpr (n+2(B A) )x)3)
2 I Pt 2B max Pr - -x232
mm A) Ax)2 +
2B LBLmas Lmax (2 pf 21)2 1)
J1 LimaxPr (1 + 2(B A)(1 -x)32)
LmaxPt LB (1 A1.2
4
Bμ + A
B
Bp BLmaxPepp 2BLmaxPt (Bμ+ + A)(1 A,)² < 2Bμ 2A
BmaxPrji ALmaxPt
-2BL LBLmaxpr(Bpp A) - xx)
The above derivation simplifies to the below inequality:
Bp + A(2 Immuxpe))
21

---

Published as a conference paper at ICLR 2025
which is true by the restriction on p in (11), that implies the quantity (2 LmaxPt) is non-negative due to:
2
Pt S Lmax (1+2B(11 )x)3) S
Lmax
Now setting constant Pt = p0 =
- :=
and d = E [0. in (12) we have
B[f(x(+1)- f(x")|z²
I (1-7))]See) - [(x')]
+ Xmmax [Cp(1 1 2)² + 2 (1- X)² + Cy + pd² + mnx
(15)
Taking expectation on (15) and using the tower property we get B[f(z'+1) - f(x")]
$ (1 74) [x)( - f(x')1
+ Lmmax [Cp(1 A)² 20' L (1 11- X1)² + +C + pd² + + 2²X²]
(1 7/) [1((t') f(x*)) Iimax (Cr p(1 2- [X² C(1 - -))) (16)
Recursively applying the above and summing up the resulting geometric series yields: |[ f(x'(()))) S (1 - 7M) (f(x²) [(x))]
1 (1 - m4).
+ 7Lmax (Cr + p(1 + 2y L2 p [ + C(1 - 1)7 Σ
j=0
Using the fact that Ej - -11) S 7y completes the proof.
C.2 PROOF OF THEOREM 3.5
Theorem C.2. Assume that each ft is L;-smooth, f is μ-PL and the ER is satisfied. Then the iterates of Unified SAM with pt = min 224 p² and 7/2 min { (++1p 7" where p² and
r are defined in Theorem 3.2, satisfy:
f(f())))((3") 1
Proof of Theorem 3.5. Since % min { 2t+
[(+1)"m 7* 7" and Pt min == S p²
get that inequality (16) holds for any t 2 0, hence we have:
Elf(x'+1] + f(x")] < (1 - mP) [[f((") f(z")]
+ Tamax (Cr, + px + 2hp axx) [X² + C(1 -)33)
(1 mP) (x') f(x"))
Cliax 1 YoP: 1 (1 + Ptppel [X2 + +C11 - 2)a]
(1 xp)) [B[f(x)) f((²)]
+ CLmax' Yope L2 (+ 2L2 Ix+ C(1 - -(1))
(17)
21 where the last inequality follows from 7/ - 1
and P:. S 1
2141 and thus Nopt. 1 ক < #
Now set 6, = B[f(x') - f(x*)) }, R = C' C'Lmax and Q L (1+ 21 [XP + C(1 -)3].
Then the inequality (17) takes the following form:
8.41 (1 7op)o + Qreet + Ry?
(18)
22

---

Published as a conference paper at ICLR 2025
Now since the sequences (x+1)
21+1 and 21+1
decreasing there exists an integer t" N such that 26+11 and
for any & 2 ts we have > (1+1)*M
2F4T. Substituting in (18) we get that for any t 2 t*
we have
12
Q
R(211 + 1)3
5141 S
de + M 1)2
+
(t+ + 1)2
1º + 1)4
t
d + Qu+
< (1 + 1)2
because (2t+1)2 (++1)*
4(t+1)
Multiplying both sides with (t + 1)2 and rearranging we have
[t-1)*
(++1)7
(t + 1) - -Po S Qu + 4R
4²
Summing for t == =t=,,- t",,.... .....: T and telescoping we have
T S (t-)² + Qu 4B
(T -T- -t 1).
4²
Changing notation from T to t get
(2-
Qu 1Rt t t" 1²2
Blf(x*) - [(x)))
+
t²
(1 15-
Qu + 4R1
s
+
4²
0
This completes the proof.
C.3 PROOF OF PROPOSITION 3.6
Proposition C.3. Assume that each fi is Le-smooth and the ER is satisfied. Then the iterates of
Unified SAM with p S min Trmax 8Blmarr(1- and 7 S 8BTmmax satisfy:
2 2(1 + 21+22471 2Ay I.max [p(1 2)2 2 + 7)) [f(x²)- - pint]
min
E [
<
+ 2Lmx [C7 + p(1 + Lmax + ((11-))7 x)2)]
Proof of Proposition 3.6. From Equation (9) we have
Blf(x*+1) -
- pintj
$ 1 ImaxP BL [p(1 - -22 2 L (1 2)² + +.1) |Vf(2')11
2
+ 2myxmax [p(1 - 2)2 + L? p 2)2 + +7] [f(x') fintt]
+ 7 T'max
[Cp(1 - 2)2 + 1-200 max 2 X)² + Ch + p\ - 2 1. 22
< 7 [V)((')) + Lmax [0(1 ))² + L2 max /² ))2 + 7] [ pimt]
2
+ +yLmax [Cp(1 - ^)2 20 2(1 (1- X) + Cr px² 2 27 L 2 222]
(19)
Here inequality (19) true due to:
1 Limaxp BLmax [p(1 - X)2 + L2 max p² (1 X)2 + [7]
1
2 2
+
2
ImmxP + BL - 2)2 + BLmax? + 2B Lame
1
X)2 <
(20)
2
Inequality (20) holds due to the theorem's constraints on y and p. which imply that each factor on the left-hand side of inequality (20) is less than 1/8, as explained below.
23

---

Published as a conference paper at ICLR 2025
Lmoxe S by the condition p.
B (1 - X)² z 5/2 by the condition on p.
BL Y S $ by the condition on ?.
2BBy 2)2 8 because have p S
1
and 7 8.L.max
4lmmox
2ByL 3
max 20 -222 S 2B
I
(1 - 1)2
8
4tsmax
(1
\)
64
1
64 8
Summing up the inequalities in the above bullet points results in inequality (20). By taking expecta- tion and using the tower property in (19), we have
إiuf - -
B[f(x') s
S 7 E
+2 [0(1 - 1)2 2 2 m -1 b [f(x') pine]
2
+ 7 Tmmax [Cp(1 - A)2 + 20yI - A)² + Cy + px² + 2 1 x²]
thus
?
[[f & (1+ 11+220 2 [p(1 1)2 + 12 pa p²(
A)2 + 71) B[f(x*) fiul
2
- -
+ Immax [Cp(1 X)2 +C2CyI2 21 (1 (1- 2)2
Cr + p>² + 2 max
222
(21)
Let us now set:
& 2 - 20
T1 = E| 20 0
(1 + -2A T-max [p(1 - x)2 2- 2 - 2)2 + 7])
(1 2.y [ + p(1 2)2 (1 + 21+2112 2xxx))) )]) >
N 2Tmmax [Cp(1 - 2)2 2. max
r² (1 2)2 + C7 2 pmaxxx X2]
2Lmax
[C7 p(1 + 2ypL )X² + C(1-C(11))) 1)3) 20
and inequality (21) takes the following form:
2 r S gō: det1 + N.
(22)
Applying Lemma A.8 to (22) completes the proof.
C.4 PROOF OF THEOREM 3.7
Theorem C.4. Let and set do f(x0) fiuf 2 0. For
1 1
82
p min 4Tmaxx 8B 8BInax( (1 - 2)2 VT 12L (C(1 - 2)2 + 12)
and
1
1
min
8BLmax 2L(1 - 2) 6Lmax A(1 /T V6ALT
24

---

Published as a conference paper at ICLR 2025
22
24L max (C(11 - A)2 + X 12L C
Then provided that
m
5184L
20
8648pA
144C
T z
max
96B, 24(1 - 2) / m,
g²
8²
8²
g
288L 2 (1 2)2
e²
we have min T-1 E
V (x')] S..
Proof of Theorem 3.7. From
Proposition
C.3 under
the
condition
that
min-
and 7 B.Brmas have
Traas
813lamas (1-x) xy
(1 Lmax [p(1 1)2 2ypl max +7))² [f(x²) - finnt]
E Vf(x')]
+ 2Lmax (Cyy + p(1 + 2xpImmxx C -(1-1)3))
(23)
Using the fact that + x S exp(x), we have that
(1 + 2A [p(1 21 + 2ypp ) +7))²
- exp (2T A [p(1 - 2) + 2yp + 7])
< exp(11 < 3.
(24)
We note that for the second inequality hold, it enough to assume that
27 Lyxa [p(1-A)] x) 2ypl + 7 27 22' A [p(1 X)² 2 27 22 (1 2)²
+ 7]
S 1,
which is true if we have that,
2T AL 1)2 < 1/3
4T'AL X)2 s 1/3
2T ATmax" < 1/3..
Imposing the restriction p S
1 V
the above inequalities can be expressed as (by solving for 7):
VGALT
Using these restrictions on y, able substitute the bound of (24) in (23) to get:
66g 2Lmax [Ch + p(1 + 27 (A² + - -)1)]
min
E
< T
(25)
To make the right hand side of (25) smaller than s², we require that the second term satisfies
Lmax [C++ p(1 + 2yp 2 (A² + C(1 - {)})] < 2
2 [Cp(1 2)² + 20 21 + C7 + px² + L2
maxt
2x21
<
4
2
82
2Lmax [p (C(1 2) + x2) 2L2 max 7p² (C(L 2)2 +12 + C7l
5
For this hold it is enough to have
25

---

Published as a conference paper at ICLR 2025
PC(n A)2 A2)
& G =
2 12 (C(1-A)=+~))
AL (C(1 - 2)² X2) t p<l
21L maa
((11-A)+X3)
2I Cy =² S m
Similarly, for the first term, get that the number of iterations must satisfy:
660
E²
126 yea
T
(26)
Ty
2
Hence so far we need the following restrictions p and
p S min
and y S 8BLmax (by the restrictions of Proposition C.3)
ILmax
8B 3BLmax(1-3)
S and p S 12 12bmax C(1
A(mnsA(1A)j-v and ?
211- A/V3ALma and γ 2 V6ALT
and γ s
2413 (C [C(1-X]=4X]] and
Plugging each of the previous bounds of y into (26) we
T 96 &
T 24 meax 6o(1 Vilmma
T 51841. 2)
86462 Aloaa gd
T
144 Cbu g
T 288 (1
Finally, collecting all terms into a single bound we have:
g2
0 == min
44max
8BLmax (1 ))2 VT 12 (C(1 2)2 + 22)
1
1
min
8BI.max
21(1 - 3)LL
max 6L A(1 - X) V6ALT
6²
2473
(C(1 - [122 + X2) 12
do Lmmxx
A2
A)
8648nA 144C
T >
96B, 24(1
V A,
2
g2
82
288
max 2)2
22
This completes the proof.
26

---

Published as a conference paper at ICLR 2025
D ADDITIONAL EXPERIMENTS
In this section, we present additional experimental evaluations of Unified SAM, following the same setup as in Li & Giannakis (2023). Specifically, we train ResNet-18 and WRN-28-10 10 CIFARIO and CIFARI00 datasets. Standard data augmentation techniques, including random cropping, ran- dom horizontal flipping, and normalization De (2017), are employed. The models are trained for 200 epochs with a batch size of 128, using cosine scheduler starting from 0.05. Weight decay 0.001. For SAM, we use p 0.1 for CIFARIO and p 0.2 for CIFARI00. Each exper- iment is repeated three times, and we report the average of the maximum test accuracy along with the standard error. The numerical results are presented Tables and 7, demonstrate that Unified SAM consistently improves the test accuracy over both USAM and SAM across all tested models. Lastly, we also observe that careful tuning of the parameter essential for achieving optimal performance.
Table Test accuracy (%) of Unified SAM on various neural networks trained on CIFARIO.
Model
= 0.0
== 0.5
1.0
A= 1/L
= 1/L
ResNet- 18
96.13
96. 16.00
96.33
96.3310.03
96.20
96.22
97.26 10.44
96.9810.08
97.05
96.7310.04
96.63
Table 7: Test accuracy (%) of Unified SAM various neural networks trained on CIFARI00.
Model
A 0.0
À = 0.5
==
== 1/t
À = 1 1- 1/t
ResNet-188 WideResNet-28-10
80.21
80.28 10.16
80.14_0.07
80.27+0.09
80. 80. 10.109 83.6220016
83.4940.23
83.7 83.71+0.03
83.55
83.55
Table 8: Test accuracy (%) of Unified SAM for PRN-18 on CIFARIO, evaluated across different values of p and X. With bold we highlight the best performance for fixed p.
Unified SAM
== 0.0
X == 0.5
X = 1.0
A= 1/t
1=1-1/t
p = 0.1
95.2940.09
95.32+0.02
95.59=0.09
95.24+0.08
95.53+0.11
p 0.2
95.25+0.14
95.4840.05
95.540.02
95.38+0.11
95.58 +0.07
= 0.3
95.25
10.24
95.18
95.1210.19
95.2610.1
P 0.4
94.7610.09
94.98
94.7 10.02
94.64
94.61 10.09
SGD
94.821..02
Table 9: Test accuracy (%) of Unified SAM for PRN-18 on CIFAR100, evaluated across different values of and X. With bold we highlight the best performance for fixed p.
Unified SAM
X = 0.5
X >= 1.0
À = 1/t
A= 1=1-1/t
p = 0.1
78.28+0.15
78. 28.28
78.32=0.22
78.33
78.39
P 0.2
78.98
78.68
78.96
78.87
78.7940.1
P = 0.3
79.0
78.95
79.21
78.73
79.27=0.05
P 0.4
78.57
78.7640.3
79.05
78.3640.09
78.7940.13
SGD
76.910.23
27


Reviews:
[{"review_id": "S7X7IbIZcl", "reviewer": "Reviewer_mscb", "summary": "The authors provide a Unified framework (Unified SAM) as a convex combination of SAM ascend and USAM ascend. They provide convergence guarantees for Unified SAM via PL condition. In special cases, the bounds reduce to that of SGD, which is known. The sampling they consider in their method is arbitrary, not restricted, and includes importance sampling. The paper concludes with experiments.", "questions": "This is an interesting paper. I have a question: how does the rate you prove depend on the parameter $\\lambda$? I understand that you derive the results in the paper, but can you explain in words how the convergence rate is changed when $\\lambda$ varies? Another question is, are you making the previous bounds tighter, or do you present a proof that works for new regimes?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- interesting and well-motivated problem\n - very well written\n - clean theoretical contributions and supporting experiments", "weaknesses": "- the paper would benefit from an informal statement of results (in math) in terms of convergence rates at the begining", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fk73595LMM", "reviewer": "Reviewer_Mo81", "summary": "This paper extends and combines previous analyses on the convergence of SAM and unnormalized SAM (U-SAM) by considering a generalized update rule called Unified SAM. Under the Expected Residual Condition, they prove convergence for Unified SAM for loss functions satisfying PL conditions and generalized non-convex loss functions. The convergence bound holds for a wide range of sampling strategies. Intriguingly, they show that importance sampling can minimize this convergence bound. Empirically, they show that Unified SAM matches and sometimes outperforms the original SAM.", "questions": "Questions\n[Relationship of convergence speed and $\\lambda$] In the logistic regression experiments, the convergence is better with larger $\\lambda$, which is coherent with the theory as here $C = 0$. Is there a case that the bound will be minimized at a non-zero $\\lambda$ when $C \\neq 0$ and in general, is there a setting where Unified SAM’s convergence speed will strictly improve over U-SAM?\n[Empirical Studies] In the case of computer vision experiments, is the training loss of Unified SAM or USAM lower than SAM? As the best practice so far seems to be using SAM in the later phase of training, how to disentangle the sharpness reduction benefit of SAM over U-SAM [1,2]?\n\n\n[1] The Crucial Role of Normalization in Sharpness-Aware Minimization, arxiv.org/abs/2305.15287 \n[2] How Does Sharpness-Aware Minimization Minimize Sharpness? arxiv.org/abs/2211.05729", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "[Quality] This paper is well-written and contains solid theoretical and empirical results.\n[Clarity] This paper carefully discusses previous convergence bounds on SAM and U-SAM and shows clearly how the current convergence results generalize previous ones.", "weaknesses": "The connections and distinctions between this work and previous analyses of SGD require clearer articulation. The paper’s central assumption, the Expected Residual Condition, is adopted from [1] and is also employed in [2,3]. In the limiting case where $\\rho = 0$, the results presented here converge to those found in the SGD analyses in [2,3]. While the authors make this reduction explicit, the paper does not address how the implications of the current results and proof techniques differ from those of earlier studies.\nFor instance, the authors suggest that this work \"provides a theoretical justification for applying importance sampling in SAM.\" However, this argument is based on the quantity $\\max_i L_i/p_i$ within the bound, which also appears in previous analyses, such as [1]. Thus, this justification extends to SGD as well and is not specifically related to SAM, a point that should be conveyed more directly.\nThe connection between the current theoretical and empirical results could be further strengthened. The computer vision experiments in this paper demonstrate that Unified SAM can achieve improved validation performance. However, 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. This metric, however, is not included in the paper or its appendix.\nThe presentation can be improved.\nThe current paper contains typos that may obscure understanding. For example, the $g(x)$ in equation (ER) in Assumption 3.1 should be the norm of $g(x)$. Further, given the complexity of the current theoretical bounds, an intuitive interpretation of the current bounds can improve the paper.\n\n[1] SGD: General Analysis and Improved Rates. arxiv.org/abs/1901.09401\n[2] SGD for Structured Nonconvex Functions: Learning Rates, Minibatching and Interpolation arxiv.org/abs/2006.10311 \n[3] Better Theory for SGD in the Nonconvex World. arxiv.org/abs/2002.03329", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2wMASz1rkq", "reviewer": "Reviewer_78w2", "summary": "This paper studies the convergence of sharpness-aware minimization (SAM) for smooth functions. The authors proposed a new unified notion of normalized SAM and unnormalized SAM by linearly interpolating the perturbed point used to take the gradient for the next update. The main contribution of the paper is to establish convergence rates to first-order stationary points for unified SAM under a more general noise assumption called \"Expected Residual (ER) Condition\" and arbitrary sampling methods. For non-convex but PL functions, the authors prove the loss value converges in $O(1/\\epsilon)$ steps.", "questions": "See weakness.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This paper is overall well written (except some typos) and the presentation of the main result is good and easy to understand. The theoretical results for convergence of SAM under the Expected Residual (ER) Condition is new. The authors also provide numerical evaluation for the new algorithm proposed in this paper, that is, unified SAM with $\\lambda_t = 1-1/t$.", "weaknesses": "I have concerns both about the correctness and significance about the main result in this paper. Given these concerns, I do not think the current draft meets the bar of ICLR.\n\n**Correctness**: \n- In Theorem 3.7, the authors wrote \"choose $\\rho\\le\\bar \\rho$ and $\\gamma\\le \\bar\\gamma$...\", the gradient norm will be smaller than $\\epsilon$ in at most $1/\\epsilon^4$ setps. It is not clear if the authors mean there exists $\\rho\\le\\bar \\rho$ and $\\gamma\\le \\bar\\gamma$, or for any $\\rho\\le\\bar \\rho$ and $\\gamma\\le \\bar\\gamma$. The former interpretation make the result trivial because the optimal choice is $\\rho=0$ and unified SAM becomes SGD. Howver, i think the later interpretation makes the current results wrong. As a quick sanity check, learning rate $\\gamma=0$ means the gradient should not change at all. (see elaboration in the next point below)\n\n- In the formal statement of Theorem 3.7, which is Theorem C.4, the authors indeed interpret the condition as \"for any $\\rho\\le\\bar \\rho$ and $\\gamma\\le \\bar\\gamma$\" (line 1100). However, this does not make sense because it does not exclude the case of $\\gamma=0$. A direct cause of this could be that the authors forgot to include Equation (13) into the restrictions that they need to satisfy, which says $T\\ge \\frac{12\\delta_0}{\\gamma \\epsilon^2}$.\n\n- However, I do not think the above issue can be fixed, unless $\\rho(1-\\lambda)=0$. The authors tried to replicate the analysis by Khaled \\&\nRichtarik (2020) for SGD to unified SAM, including how to deal with the seemingly exponential explosion. However, in the case of unified SAM, the term before $[f(x^))-f^{inf}]$ in line 1067 is essentially $\\frac{(1+\\Theta(\\gamma))^T}{T\\gamma}$ when $\\rho(1-\\lambda)\\neq 0$.  Because $\\frac{(1+\\Theta(\\gamma))^T}{T\\gamma} \\ge \\frac{(1+\\Theta(T\\gamma)}{T\\gamma} = \\Theta(1)$, the right-hand side of line 1066 is at least a constant.\n\n**Significance of Theory**: This paper only talks about optimization of SAM, but indeed, SAM is proposed to improve the generalization of SGD. It is not clear to me what the ultimate goal here by studying these convergence bounds of unified SAM. Both from a intuitive sense or from the bounds presented in this paper, SGD (or unified SAM with $\\rho=0$) has the best optimization performance. To me the real problem for SAM is how to show they can efficiently minimize the sharpness-aware loss (either the origial notion of maximal loss after certain perturbation, or the hessian-regularized version proposed in Wen et al. (2023) and Bartlett et al. (2023)), rather than viewing them as a tool to minimize the original training loss and analyze that behavior.\n\n**Significance of Experiments**: The performance of the proposed $1-1/t$ schedule for unified SAM does not beat normalized SAM by a margin which is larger than the standard deviation in most experiments. Sometimes unified SAM even has better generalization. I understand this does not contradict with theory because theory does not try to say anything about generalization. However, it is not clear to me if the main benefit of unified SAM is better optimization or generalization. \n\n**References**:\n\n- Wen, Kaiyue, Tengyu Ma, and Zhiyuan Li. \"How Sharpness-Aware Minimization Minimizes Sharpness?.\" The Eleventh International Conference on Learning Representations. 2023.\n- Bartlett, Peter L., Philip M. Long, and Olivier Bousquet. \"The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima.\" Journal of Machine Learning Research 24.316 (2023): 1-36.\n\n\n# Minor comments:\n1. line 271. \"our results is the first to demonsteate linear convergence in the fully stochastic regime\". I would suggest that the authors do not call this linear converge to avoid confusion, when the final bound for loss is still at least constantly large.\n2. line 334, \"The results in Proposition 3.6 and Theorem 3.7 are tight, as setting ρ= 0 Unified SAM reduces in SGD and these simplify to the step sizes and rates (up to constants) of Theorem 2 and Corollary 1 from Khaled & Richtarik (2020).\" The whole point of the analysis is for the regime $\\rho\\neq 0$. The fact that the analysis in this paper recovers previous result when $\\rho=0$ does not indicate the tightness of the result in the main setting when $\\rho\\neq 0$, especially the dependence on $\\rho$.\n\n# Typos:\n1. line 224, missing l2 norm on $g(x)$\n2. section B in appendix, line 866-890. Missing norm and square over the norm.\n3. line 1121. \"From Theorem 3.7\" should be from \"Proposition C.3\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zuQCHnnRFB", "reviewer": "Reviewer_hnag", "summary": "The paper studies the convergence of SAM and USAM in stochastic settings. It proves these properties for a newly proposed algorithm, Unified-SAM, which includes SAM and USAM as special cases. The analysis relaxes popular assumptions like bounded variance (BV) and gradients (BG), replacing them with expected residual (ER) condition. The proof provides convergence guarantees for SAM under different step sizes in non-convex functions and Polyak-Lojasiewicz (PL) functions. The theory holds under arbitrary sampling paradigms, including importance sampling. The authors also demonstrate Unified-SAM's performance compared to SAM in practical settings.", "questions": "See the Weakness part.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-written and easy to follow.\n- The authors challenge existing assumptions on stochastic noise, such as bounded gradients and bounded variance, in proving SAM's convergence, replacing them with an expected residual condition that encompasses both as special cases.\n- The proof is a slight improvement over existing studies.", "weaknesses": "- The empirical results from the original SAM paper Foret et al. [2020] are established using a constant $\\rho$. This is a crucial point for aligning theoretical and empirical results. To my knowledge, there are no existing works that establish convergence results on the constant $\\rho$ for non-convex functions. However, there are some theoretical papers that use conditions closely related to the constant $\\rho$, but these have not been discussed in the paper under review, for example,  Nam et al. [2023], Khanh et al. [2024], Xie et al. [2024]. The assumption in this paper regarding $\\rho$ is $\\rho = \\min ( \\frac{1}{2t+1}, \\rho^{\\star} )$, which is less general than the assumption in Nam et al. [2023], where $\\rho$ is defined as $\\sum^{\\infty}_{t=0} \\rho_t^2 < \\infty$. Additionally, the assumption on $\\rho$ in Khanh et al. [2024] for the full-batch setting is even more general, as it allows the perturbation radius $\\rho$ to decrease at arbitrary slow rates, which nearly captures the constant $\\rho$.\n- In the asymptotic setting, this paper's result in Theorem 3.7, $\\min_{t=0,...,T-1} \\mathbb{E} || \\nabla f(x^t) || \\leq \\epsilon$, is weaker than the convergence result in the stochastic, non-convex setting in Nam et al. [2023], where the gradient norm approaches zero almost surely. Furthermore, if Theorem 3.7 is considered in the full-batch setting, it is also weaker than the result in Khanh et al. [2024.\n- In Theorem 3.5, why do you write $O\\left( \\frac{1}{t} + \\frac{1}{t^2} \\right)$ instead of $O\\left( \\frac{1}{t} \\right)$, since they are equivalent?\n- Compare your Assumption 3.1 (Expected Residual Condition) with Assumption A.4 in Nam et al. [2023].\n- The name of the algorithm, Unified-SAM, may not be suitable, as it only covers USAM, SAM, and the variant that transitions between SAM and USAM. There are many other SAM-like variants that this algorithm does not cover.\n- As shown in 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.\n- Line 152-153: “This is the first result that drops the bounded variance assumption for both USAM and SAM.\" I suggest rewriting this sentence to clarify that, while the bounded variance assumption is removed, an Expected Residual condition is still required.\n\nReferences:\n- Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization\nfor efficiently improving generalization. ICLR 2021. URL https://arxiv.org/pdf/2010.01412.\n- Pham Duy Khanh, Hoang-Chau Luong, Boris S Mordukhovich, and Dat Ba Tran. Fundamental convergence analysis of sharpness-aware minimization. NeurIPS, 2024. URL https://arxiv.org/pdf/2401.08060.\n- Kyunghun Nam, Jinseok Chung, and Namhoon Lee. Almost sure last iterate convergence of sharpness-aware minimization. Tiny Papers ICLR, 2023. URL https://openreview.net/forum?id=IcDTYTI0Nx.\n- Wanyun Xie, Thomas Pethick, and Volkan Cevher. Sampa: Sharpness-aware minimization parallelized. NeurIPS, 2024. URL https://arxiv.org/pdf/2410.10683v1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)