Title: Practical Sharpness-Aware Minimization Cannot Converge All the Way to Optima
Venue: NeurIPS 2023
Accepted: yes
Keywords: Sharpness-Aware Minimization

convex optimization
Abstract: Sharpness-Aware Minimization (SAM) is an optimizer that takes a descent step based on the gradient at a perturbation $y_t = x_t + \rho \frac{\nabla f(x_t)}{\lVert \nabla f(x_t) \rVert}$ of the current point $x_t$. 
Existing studies prove convergence of SAM for smooth functions, but they do so by assuming decaying perturbation size $\rho$ and/or no gradient normalization in $y_t$, which is detached from practice. To address this gap, we study deterministic/stochastic versions of SAM with practical configurations (i.e., constant $\rho$ and gradient normalization in $y_t$) and explore their convergence properties on smooth functions with (non)convexity assumptions.
Perhaps surprisingly, in many scenarios, we find out that SAM has limited capability to converge to global minima or stationary points.
For smooth strongly convex functions, we show that while deterministic SAM enjoys tight global convergence rates of $\tilde \Theta(\frac{1}{T^2})$, the convergence bound of stochastic SAM suffers an inevitable additive term $\mathcal O(\rho^2)$, indicating convergence only up to neighborhoods of optima.
In fact, such $\mathcal O(\rho^2)$ factors arise for stochastic SAM in all the settings we consider, and also for deterministic SAM in nonconvex cases; importantly, we prove by examples that such terms are unavoidable.
Our results highlight vastly different characteristics of SAM with vs. without decaying perturbation size or gradient normalization, and suggest that the intuitions gained from one version may not apply to the other.

Practical Sharpness-Aware Minimization Cannot Converge All the Way to Optima

Abstract

Sharpness-Aware Minimization (SAM) is an optimizer that takes a descent step based on the gradient at a perturbation y t = x t + ρ ∇ f ( x t ) ∥∇ f ( x t ) ∥ of the current point x t . Existing studies prove convergence of SAM for smooth functions, but they do so by assuming decaying perturbation size ρ and/or no gradient normalization in y t , which is detached from practice. To address this gap, we study deterministic/stochastic versions of SAM with practical configurations (i.e., constant ρ and gradient normalization in y t ) and explore their convergence properties on smooth functions with (non)convexity assumptions. Perhaps surprisingly, in many scenarios, we find out that SAM has limited capability to converge to global minima or stationary points. For smooth strongly convex functions, we show that while deterministic SAM enjoys tight global convergence rates of ˜ Θ( 1 T 2 ) , the convergence bound of stochastic SAM suffers an inevitable additive term O ( ρ 2 ) , indicating convergence only up to neighborhoods of optima. In fact, such O ( ρ 2 ) factors arise for stochastic SAM in all the settings we consider, and also for deterministic SAM in nonconvex cases; importantly, we prove by examples that such terms are unavoidable . Our results highlight vastly different characteristics of SAM with vs. without decaying perturbation size or gradient normalization, and suggest that the intuitions gained from one version may not apply to the other.

1 Introduction

Modern neural networks are armed with a large number of layers and parameters, having a risk to overfit to training data. In order to make accurate predictions on unseen data, generalization performance has been considered as the most important factor in deep learning models. Based on the widely accepted belief that geometric properties of loss landscape are correlated with generalization performance, studies have proved theoretical and empirical results regarding the relation between sharpness measures and generalization [14, 17, 19, 21, 29]. Here, sharpness of a loss at a point generally refers to the degree to which the loss varies in the small neighborhood of the point.

Motivated by prior studies that show flat minima have better generalization performance [19, 21], Foret et al. [16] propose an optimization method referred to as Sharpness-Aware Minimization (SAM) . A single iteration of SAM consists of one ascent step (perturbation) and one descent step. Starting from the current iterate x t , SAM first takes an ascent step y t = x t + ρ ∇ f ( x t ) ∥∇ f ( x t ) ∥ to (approximately) maximize the loss value in the ρ -neighborhood of x t , and then uses the gradient at y t to update x t to x t +1 . This two-step procedure gives SAM a special characteristic: the tendency to search for a flat minimum, i.e., a minimum whose neighboring points also return low loss value. Empirical results [3, 9, 16, 20, 25] show that SAM demonstrates an exceptional ability to perform well on different models and tasks with high generalization performance. Following the success of SAM,

numerous extensions of SAM have been proposed in the literature [6, 15, 18, 23, 24, 26-28, 3032, 34, 35].

On the theoretical side, various studies have demonstrated different characteristics of SAM [1, 2, 4, 5, 12, 13, 22, 33]. However, comprehending the global convergence properties of practical SAM on a theoretical level still remains elusive. In fact, several recent results [2, 18, 27, 31, 35] attempt to prove the convergence guarantees for SAM and its variants. While these results provide convergence guarantees of SAM on smooth functions, they are somewhat detached to the practical implementations of SAM and its variants. They either (1) assume decaying or sufficiently small perturbation size ρ [18, 27, 31, 35], whereas ρ is set to constant in practice; or (2) assume ascent steps without gradient normalization [2], whereas practical implementations of SAM use normalization when calculating the ascent step y t .

1.1 Summary of Our Contributions

To address the aforementioned limitations of existing results, we investigate convergence properties of SAM using gradient normalization in ascent steps and arbitrary constant perturbation size ρ . We note that to the best of our knowledge, the convergence analysis of SAM have not been carried out under the two practical implementation choices. Our analyses mainly focus on smooth functions, with different levels of convexity assumptions ranging from strong convexity to nonconvexity. We summarize our contributions below; a summary of our convergence results (upper and lower bounds) can also be found in Table 1.

For deterministic SAM, we prove convergence to global minima of smooth strongly convex functions, and show the tightness of convergence rate in terms of T . Furthermore, we establish the convergence of SAM to stationary points of smooth convex functions. For smooth nonconvex functions, we prove that SAM guarantees convergence to stationary points up to an additive factor O ( ρ 2 ) . We provide a worst-case example that always suffers a matching squared gradient norm Ω( ρ 2 ) , showing that the additive factor is unavoidable and tight in terms of ρ .

For stochastic settings, we analyze two versions of stochastic SAM ( n -SAM and m -SAM) on smooth strongly convex, smooth convex, and smooth (Lipschitz) nonconvex functions. We provide convergence guarantees to global minima or stationary points up to additive factors O ( ρ 2 ) . In case of m -SAM, we demonstrate that these factors are inevitable and tight in terms of ρ , by providing worst-case examples where SAM fails to converge properly.

1.2 Related Works: Convergence Analysis of SAM

Recent results [18, 27, 31, 35] prove convergence of stochastic SAM and its variants to stationary points for smooth nonconvex functions. However, these convergence analyses are limited to only smooth nonconvex functions and do not address convex functions. Also, as already pointed out, all the proofs require one crucial assumption detached from practice: decaying (or sufficiently small ) perturbation size ρ . If ρ becomes sufficiently small, then the difference between f SAM ( x ) = max ∥ ϵ ∥≤ ρ f ( x + ϵ ) and f ( x ) becomes negligible, which means that they undergo approximately the same updates. Due to this reason, especially in later iterates, SAM with negligible perturbation size

would become almost identical to gradient descent (GD), resulting in a significant difference in the convergence behavior compared to constant ρ .

As for existing proofs that do not require decaying ρ , Andriushchenko and Flammarion [2] prove convergence guarantees for deterministic SAM with constant ρ on smooth nonconvex, smooth Polyak-Łojasiewicz, smooth convex, and smooth strongly convex functions. Moreover, the authors prove convergence guarantees of stochastic SAM, this time with decaying ρ , on smooth nonconvex and smooth Polyak-Łojasiewicz functions. However, their analyses are also crucially different from the practical implementations of SAM because their proofs are for a variant of SAM without gradient normalization in ascent steps, with updates in the form of x t +1 = x t -η ∇ f ( x t + ρ ∇ f ( x t )) . This form can be considered as vanilla SAM with an 'effective' perturbation size ρ ∥∇ f ( x t ) ∥ , which indicates that even with constant ρ , the effective perturbation size will become smaller and smaller as the algorithm converges towards a stationary point. As in the case of decaying ρ discussed above, this can make a huge difference in the convergence behavior.

Figure 1 illustrates a comparison between SAM with and without gradient normalization, highlighting the disparity between their convergence points. As we predicted above, it is evident that they indeed reach

USAM

SAM

entirely distinct global minima, exhibiting different convergence characteristics.

Discussions on the related works on the relationship between sharpness and generalization, as well as other theoretical properties of SAM, are provided in Appendix A.

1.3 Notation

The euclidean norm of a vector v is denoted as ∥ v ∥ . We let R ∗ + ≜ { x ∈ R | x > 0 } , and use [ · ] + to define max {· , 0 } . We use O ( · ) and Ω( · ) to hide numerical constants in our upper and lower bounds, respectively. We use ˜ O ( · ) to additionally hide logarithmic factors. When discussing rates, we sometimes use these symbols to hide everything except T , the number of SAM iterations. When discussing additive factors that arise in our convergence bounds, we use these symbols to hide all other variables except the perturbation size ρ . For an objective function f and initialization x 0 of SAM, we define f ∗ ≜ inf x f ( x ) and ∆ ≜ f ( x 0 ) -f ∗ .

2 Sharpness-Aware Minimization: Preliminaries and Intuitions

Before presenting the main results, we first introduce deterministic and stochastic SAM, and present definitions of function classes considered in this paper. We next introduce virtual gradient map and virtual loss that shed light on our intuitive understanding of SAM's convergence behavior.

2.1 Sharpness-Aware Minimization (SAM)

We focus on minimizing a function f : R d → R , where the optimization variable is represented by x ∈ R d , using the Sharpness-Aware Minimization (SAM) optimizer. Instead of minimizing the vanilla objective function f ( x ) , SAM [16] aims to minimize the SAM objective f SAM ( x ) , where

For ρ ∈ R ∗ + , the SAM loss f SAM ( x ) outputs the worst function value in a ρ -ball { w ∈ R d | ∥ x -w ∥ ≤ ρ } around x . Assuming 'sufficiently small' ρ , the inner maximization of (1) can be solved by Taylor-approximating the objective:

Using this approximate solution ˆ ϵ ( x ) , we can define the approximate SAM objective ˆ f SAM as ˆ f SAM ( x ) ≜ f ( x +ˆ ϵ ( x )) . In order to run gradient descent (GD) on ˆ f SAM ( x ) , one needs to calculate its gradient; however, from the definition of ˆ ϵ ( x ) , we can realize that ∇ ˆ f SAM ( x ) has terms that involve the Hessian of f . Here, SAM makes another approximation, by ignoring the Hessian term:

From these approximations, one iteration of SAM is defined as a set of two-step update equations:

As seen in (2), we use x t to denote the iterate of SAM at the t -th step. We use T to denote the number of SAM iterations. We refer to the hyperparameter ρ ∈ R ∗ + as the perturbation size and η ∈ R ∗ + as the step size . Note that in (2), the perturbation size ρ is a time-invariant constant; in practice, it is common to fix ρ as a constant throughout training [3, 9, 16].

According to the SAM update in (2), the update cannot be defined when ∥∇ f ( x ) ∥ = 0 . In practice, we add a small numerical constant (e.g., 10 -12 ) to the denominator in order to prevent numerical instability. In this paper, we ignore this constant and treat ∇ f ( x t ) ∥∇ f ( x t ) ∥ as 0 whenever ∥∇ f ( x t ) ∥ = 0 .

2.2 SAM under Stochastic Settings

To analyze stochastic SAM, we suppose that the objective is given as f ( x ) = E ξ [ l ( x ; ξ )] , where ξ is a stochastic parameter (e.g., data sample) and l ( x ; ξ ) indicates the loss at point x with a random sample ξ . Based on the SAM update in (2), we can define stochastic SAM under this setting:

We define g ( x ) = ∇ x l ( x ; ξ ) and ˜ g ( x ) = ∇ x l ( x ; ˜ ξ ) . Here, ξ and ˜ ξ are stochastic parameters, queried from any distribution(s) which satisfies E ˜ ξ l ( x ; ˜ ξ ) = E ξ l ( x ; ξ ) = f ( x ) .

There are two popular variants of stochastic SAM, introduced in Andriushchenko and Flammarion [2]. Stochastic n -SAM algorithm refers to the update equation (3) when ξ and ˜ ξ are independent. In contrast, practical SAM algorithm in Foret et al. [16] employs stochastic m -SAM algorithm, which follows the update equation (3) where ξ is equal to ˜ ξ . We will consider both versions in our theorems.

2.3 Function Classes

We state definitions and assumptions of function classes of interest, which are fairly standard.

Definition 2.1 (Convexity) . A function f : R d → R is convex if, for any x , y ∈ R d and λ ∈ [0 , 1] , it satisfies f ( λ x +(1 -λ ) y ) ≤ λf ( x ) + (1 -λ ) f ( y ) .

Definition 2.2 ( µ -Strong Convexity) . A differentiable function f : R d → R is µ -strongly convex if there exists µ > 0 such that f ( x ) ≥ f ( y )+ ⟨∇ f ( y ) , x -y ⟩ + µ 2 ∥ x -y ∥ 2 , for all x , y ∈ R d .

Definition 2.3 ( L -Lipschitz Continuity) . A function f : R d → R is L -Lipschitz continuous if there exists L ≥ 0 such that ∥ f ( x ) -f ( y ) ∥ ≤ L ∥ x -y ∥ , for all x , y ∈ R d .

Definition 2.4 ( β -Smoothness) . A differentiable function f : R d → R is β -smooth if there exists β ≥ 0 such that ∥∇ f ( x ) -∇ f ( y ) ∥ ≤ β ∥ x -y ∥ , for all x , y ∈ R d .

Next we define an assumption considered in the analysis of stochastic SAM (3).

Assumption 2.5 (Bounded Variance) . The gradient oracle of a differentiable function f : R d → R has bounded variance if there exists σ ≥ 0 such that

2.4 SAM as GD on Virtual Loss

Convergence analysis of SAM is challenging due to the presence of normalized ascent steps. In order to provide intuitive explanations of SAM's convergence properties, we develop tools referred to as

virtual gradient map and virtual loss in this section. These tools can provide useful intuitions when we discuss our main results.

In order to define the new tools, we first need to define Clarke subdifferential [10], whose definition below is from Theorem 6.2.5 of Borwein and Lewis [7].

Definition 2.6 (Clarke Subdifferential) . Suppose that a function f : R d → R is locally Lipschitz around a point x and differentiable on R d \ W where W is a set of measure zero. Then, the Clarke subdifferential of f at x is ∂f ( x ) ≜ cvxhull { lim t →∞ ∇ f ( x t ) | x t → x , x t / ∈ W } , where cvxhull denotes the convex hull of a set.

Clarke subdifferential ∂f ( x ) is a convex hull of all possible limits of ∇ f ( x t ) as x t approaches x . It can be thought of as an extension of gradient to a nonsmooth function. It is also known that for a convex function f , the Clarke subdifferential of f is equal to the subgradient of f .

Using the definition of Clarke differential, we now define virtual gradient map and virtual loss of f , which can be employed for understanding the convergence of SAM.

Definition 2.7 (Virtual Gradient Map/Loss) . For a differentiable function f : R d → R , we define the virtual gradient map G f : R d → R d of f to be G f ( x ) ≜ ∇ f ( x + ρ ∇ f ( x ) ∥∇ f ( x ) ∥ ) . Additionally, if there exists a function J f : R d → R whose Clarke subdifferential is well-defined and ∂J f ( x ) ∋ G f ( x ) for all x , then we call J f a virtual loss of f .

If a virtual loss J f is well-defined for f , the update of SAM (2) on f is equivalent to a (sub)gradient descent update on the virtual loss J f ( x ) , which means, x t +1 = x t -ηG f ( x t ) . The reason why we have to use Clarke subdifferential to define the virtual loss is because even for a differentiable and smooth function f , there are cases where the virtual gradient map G f is discontinuous and the virtual loss J f (if exists) is nonsmooth; see the discussion below Theorem 3.1.

Note that if the differentiable function f is one-dimensional ( d = 1 ), the virtual loss J f ( x ) is always well-defined because it can be obtained by simply (Lebesgue) integrating G f ( x ) . However, in case of multi-dimensional functions ( d > 1 ), there is no such guarantee, although G f is always well-defined. We emphasize that the virtual gradient map G f and virtual loss J f are mainly used for a better intuitive understanding of our (non-)convergence results. In formal proofs, we use them for analysis only if J f is well-defined.

Lastly, we note that Bartlett et al. [4] also employ a similar idea of virtual loss. In case of convex quadratic objective function f ( x ) , the authors define u t = ( -1) t x t and formulate a (different) virtual loss ˜ J f ( u ) such that a SAM update on f ( x t ) is equivalent to a GD update on ˜ J f ( u t ) .

3 Convergence Analysis Under Deterministic Settings

In this section, we present the main results on the (non-)convergence of deterministic SAM with constant perturbation size ρ and gradient normalization. We study four function classes: smooth strongly convex, smooth convex, smooth nonconvex, and nonsmooth Lipschitz convex functions.

3.1 Smooth and Strongly Convex Functions

For smooth strongly convex functions, we prove a global convergence guarantee for the best iterate.

Theorem 3.1. Consider a β -smooth and µ -strongly convex function f . If we run deterministic SAM starting at x 0 with any perturbation size ρ > 0 and step size η = min { 1 µT max { 1 , log ( µ 5 ∆ T 2 β 6 ρ 2 )} , 1 2 β } to minimize f , we have

The proof of Theorem 3.1 can be found in Appendix B.2. As for relevent existing studies, Theorem 1 of Dai et al. [13] can be adapted to establish the convergence guarantee f ( x T ) -f ∗ = ˜ O ( 1 T ) , by selecting the step size η = min { 1 µT max { 1 , log ( µ 2 ∆ T β 3 ρ 2 )} , 1 β } and employing a similar proof technique as ours. We highlight that Theorem 3.1 achieves a faster convergence rate compared to the concurrent bound by Dai et al. [13]. Moreover, Theorem 11 of Andriushchenko and Flammarion [2] proves the convergence guarantee for this function class: ∥ x T -x ∗ ∥ 2 = O (exp( -T )) , but assuming

SAM without gradient normalization , and the boundedness of ρ . Here, we get a slower convergence rate of ˜ O ( exp( -T ) + 1 T 2 ) , but with any ρ > 0 and with normalization.

By viewing SAM as GD on virtual loss, we can get an intuition why SAM cannot achieve exponential convergence. Consider a smooth and strongly convex function: f ( x ) = 1 2 x 2 . One possible virtual loss of f is J f ( x ) = 1 2 ( x + ρ sign( x )) 2 , which is a nonsmooth and strongly convex function, for which the exponential convergence of GD is impossible [8].

Indeed, we can show that the sublinear convergence rate in Theorem 3.1 is not an artifact of our analysis. Interestingly, the ˜ O ( exp( -T ) + 1 T 2 ) rate given in Theorem 3.1 is tight in terms of T , up to logarithmic factors. Our next theorem provides a lower bound for smooth strongly convex functions.

Theorem 3.2. Suppose β µ ≥ 2 . For any choice of perturbation size ρ , step size η , and initialization x 0 , there exists a differentiable, β -smooth, and µ -strongly convex function f such that

holds for deterministic SAM iterates.

For the proof of Theorem 3.2, refer to Appendix B.3. By comparing the rates in Theorems 3.1 and 3.2, we can see that the two bounds are tight in terms of T and ρ . The bounds are a bit loose in terms of β and µ , but we believe that this may be partly due to our construction; in proving lower bounds, we used one-dimensional quadratic functions as worst-case examples. A more sophisticated construction may improve the tightness of the lower bound, which is left for future work.

3.2 Smooth and Convex Functions

For smooth and convex functions, proving convergence of the function value to global optimum becomes more challenging, due to the absence of strong convexity. In fact, we can instead prove that the gradient norm converges to zero.

Theorem 3.3. Consider a β -smooth and convex function f . If we run deterministic SAM starting at x 0 with any perturbation size ρ > 0 and step size η = min { √ 2∆ √ β 3 ρ 2 T , 1 2 β } to minimize f , we have

The proof of Theorem 3.3 is given in Appendix B.4. As for relevant existing studies, Theorem 11 of Andriushchenko and Flammarion [2] proves the convergence guarantee for this function class: f (¯ x ) -f ∗ = O ( 1 T ) , where ¯ x indicates the averaged x over T iterates, while assuming SAM without gradient normalization and bounded ρ . Here, we prove a weaker result: a convergence rate of O ( 1 √ T ) to stationary points , albeit for any ρ > 0 and with normalization.

One might expect that a reasonably good optimizer should converge to global minima of smooth convex functions. However, it turns out that both showing convergence to global minima and finding a non-convergence example are quite challenging in this function class.

Indeed, we later show non-convergence examples for other relevant settings. For stochastic SAM, we provide a non-convergence example (Theorem 4.4) in the same smooth and convex function class, showing that the suboptimality gap in terms of function values cannot have an upper bound, and therefore rendering convergence to global minima impossible. For deterministic SAM in nonsmooth Lipschitz convex functions, we show an example (Theorem 3.6) where convergence to the global minimum is possible only up to a certain distance proportional to ρ .

Given these examples, we suspect that there may also exist a non-convergence example for the determinstic SAM in this smooth convex setting. We leave settling this puzzle to future work.

3.3 Smooth and Nonconvex Functions

Existing studies prove that SAM (and its variants) with decaying or sufficiently small perturbation size ρ converges to stationary points for smooth nonconvex functions [2, 18, 27, 31, 35]. Unfortunately, with constant perturbation size ρ , SAM exhibits a different convergence behavior: it does not converge all the way to stationary points.

Theorem 3.4. Consider a β -smooth function f satisfying f ∗ = inf x f ( x ) > -∞ . If we run deterministic SAM starting at x 0 with any perturbation size ρ > 0 and step size η = 1 β to minimize f , we have

Refer to the Appendix B.5 for the proof of Theorem 3.4. For a comparison, Theorem 9 of Andriushchenko and Flammarion [2] proves the convergence for this function class: 1 T ∑ T t =0 ∥∇ f ( x t ) ∥ 2 = O ( 1 T ) , but again assuming SAM without gradient normalization , and boundedness of ρ . Our Theorem 3.4 guarantees O ( 1 T ) convergence up to an additive factor β 2 ρ 2 .

One might speculate that the undesirable additive factor β 2 ρ 2 is an artifact of our analysis. The next theorem presents an example which proves that this extra term is in fact unavoidable.

Theorem 3.5. For any ρ > 0 and η ≤ 1 β , there exists a β -smooth and Θ( βρ ) -Lipschitz continuous function such that, if deterministic SAM is initialized at a point x 0 sampled from a continuous probability distribution, then deterministic SAM converges to a nonstationary point, located at a distance of Ω( ρ ) from a stationary point, with probability 1 .

Here we present a brief outline for the proof of Theorem 3.5. For a given ρ , consider a one-dimensional function f ( x ) = 9 βρ 2 25 π 2 sin ( 5 π 3 ρ x ) . Figure 2(a) demonstrates the virtual loss for this example. By examining the virtual loss J f and its stationary points, we observe that SAM iterates x t converge to a non-stationary point, located at a distance of Ω( ρ ) from the stationary point. This means that the limit point of SAM will have gradient norm of order Ω( βρ ) , thereby proving that the additive factor in Theorem 3.4 is tight in terms of ρ . A detailed analysis of Theorem 3.5 is provided in Appendix B.6.

Remark: why do we ignore η = Ω( 1 β ) ? As the reader may have noticed, Theorem 3.5 only considers the case η = O ( 1 β ) , and hence does not show that the additive factor is inevitable for any choice of η > 0 . However, one can notice that if η = Ω( 1 β ) , then we can consider a one-dimensional Ω( β ) -strongly convex quadratic function and show that the SAM iterates blow up to infinity. For the same reason, in our other non-convergence results (Theorems 4.2 and 4.4), we only focus on η = O ( 1 β ) .

3.4 Nonsmooth Lipschitz Convex Functions

Previous theorems study convergence of SAM assuming smoothness. The next theorem shows an example where SAM on a nonsmooth convex function converges only up to Ω( ρ ) distance from the global minimum x ∗ . This means that for constant perturbation size ρ , there exist convex functions that prevent SAM from converging to global minima. In Appendix B.7, we prove the following:

Theorem 3.6. For any ρ > 0 and η < 7 ρ 4 , there exists a nonsmooth Lipschitz convex function f such that for some initialization, deterministic SAM converges to suboptimal points located at a distance of Ω( ρ ) from the global minimum.

4 Convergence Analysis Under Stochastic Settings

In this section, we present the main results on the convergence analysis of stochastic SAM, again with time-invariant (constant) perturbation size ρ and gradient normalization . We consider both types of stochastic SAM: n -SAM and m -SAM, defined in Section 2.2. We study four types of function classes: smooth strongly convex, smooth convex, smooth nonconvex, and smooth Lipschitz nonconvex functions, under the assumption that the gradient oracle has bounded variance of σ 2 (Assumption 2.5). Our results in this section reveal that stochastic SAM exhibits different convergence properties compared to deterministic SAM.

4.1 Smooth and Strongly Convex Functions

Theorem 3.1 shows the convergence of deterministic SAM to global optima, for smooth strongly convex functions. Unlike this result, we find that under stochasticity and constant perturbation size ρ , both n -SAM and m -SAM ensure convergence only up to an additive factor O ( ρ 2 ) .

Theorem 4.1. Consider a β -smooth, µ -strongly convex function f , and assume Assumption 2.5. Under n -SAM, starting at x 0 with any perturbation size ρ > 0 and step size η =

(a)

(b)

(c)

(d)

(e)

f

f(1)

Jf（1)

Jf

0.015

f(2)

f

Ef

0.010

0.005

1.5

.0

0.5

0.5

1.0

1.5

0.005

.010

1

2

3

0.015

∩nn:

20

f(1)

0

Jr（1)

f(2)

Jr(2)

5

f

Ef

.0

0

1

2

3

5

6

8

10

Under m -SAM, additionally assuming l ( · , ξ ) is β -smooth for any ξ , the inequality continues to hold.

For the proof, please refer to Appendix C.2. Theorem 4.1 provides a convergence rate of ˜ O ( 1 T ) to global minima, but only up to suboptimality gap 2 β 2 ρ 2 µ . If σ ≤ βρ , then Theorem 4.1 becomes:

thereby showing a convergence rate of O (exp( -T )) modulo the additive factor.

For relevant existing studies, Theorem 2 of Andriushchenko and Flammarion [2] proves the convergence guarantee for smooth Polyak-Łojasiewicz functions: E f ( x T ) -f ∗ = O ( 1 T ) , but assuming stochastic SAM without gradient normalization , and perturbation size ρ decaying with t . In contrast, our analysis shows that the convergence property can be different with normalization and constant ρ .

The reader might be curious if the additional O ( ρ 2 ) term can be removed. Our next theorem proves that in the case of high gradient noise ( σ > βρ ), m -SAM with constant perturbation size ρ cannot converge to global optima beyond a suboptimality gap Ω( ρ 2 ) , when the component function l ( x ; ξ ) is smooth for any ξ . Hence, the additional term in Theorem 4.1 is unavoidable , at least for the more practical version m -SAM.

Theorem 4.2. For any ρ > 0 , β > 0 , σ > βρ and η ≤ 3 10 β , there exists a β -smooth and β 5 -strongly convex function f satisfying the following. (1) The function f satisfies Assumption 2.5. (2) The component functions l ( · ; ξ ) of f are β -smooth for any ξ . (3) If we run m -SAM on f initialized inside

a certain interval, then any arbitrary weighted average ¯ x of the iterates x 0 , x 1 , . . . must satisfy E [ f (¯ x ) -f ∗ ] ≥ Ω( ρ 2 ) .

Here we provide a brief outline for Theorem 4.2. The in-depth analysis is provided in Appendix C.3. Given ρ > 0 , β > 0 , σ > βρ , we consider a one-dimensional quadratic function f ( x ) whose component function l ( x ; ξ ) is carefully chosen to satisfy the following for x ∈ [ ρ 6 , 13 ρ 6 ] :

For values of x outside this interval [ ρ 6 , 13 ρ 6 ] , each component function l ( x ; ξ ) takes the form of a strongly convex quadratic function. Figures 2(b) and 2(c) illustrate the original and virtual loss function plots of l ( x ; ξ ) . The local concavity of component function plays a crucial role in making an attracting basin in the virtual loss, and this leads to an interval [ ρ 6 , 13 ρ 6 ] bounded away from the global minimum 0 , from which m -SAM iterates cannot escape.

For this scenario, we obtain f ( x t ) -f ∗ = Ω( ρ 2 ) and ∥∇ f ( x t ) ∥ 2 = Ω( ρ 2 ) for all iterates. From this, we can realize that the additive factor in Theorem 4.1 for m -SAM is unavoidable in the σ > βρ regime, and tight in terms of ρ . Moreover, the proof of Theorem 4.2 in Appendix C.3 reveals that even in the σ ≤ βρ case, an analogous example gives ∥ x t -x ∗ ∥ = Ω( ρ ) for all iterates; hence, m -SAM fails to converge all the way to the global minimum in the small σ regime as well.

4.2 Smooth and Convex Functions

We now move on to smooth convex functions and investigate the convergence guarantees of stochastic SAM for this function class. As can be guessed from Theorem 3.3, our convergence analysis in this section focuses on finding stationary points. Below, we provide a bound that ensures convergence to stationary points up to an additive factor.

Theorem 4.3. Consider a β -smooth, convex function f , and assume Assumption 2.5. Under n -SAM, starting at x 0 with any perturbation size ρ > 0 and step size η = min { √ ∆ √ β [ σ 2 -β 2 ρ 2 ] + T , 1 2 β } to minimize f , we have

Under m -SAM, additionally assuming l ( · , ξ ) is β -smooth for any ξ , the inequality continues to hold.

The proof of Theorem 4.3 can be found in Appendix C.4. Theorem 4.3 obtains a bound of O ( 1 √ T ) modulo an additive factor 4 β 2 ρ 2 . Similar to Theorem 4.1, if σ ≤ βρ , then Theorem 4.3 reads

hence showing a convergence rate of O ( 1 T ) modulo the additive factor. Since the non-convergence example in Theorem 4.2 provides a scenario that E ∥∇ f ( x t ) ∥ 2 = Ω( ρ 2 ) for all t , we can see that the extra term is inevitable and also tight in terms of ρ .

Theorem 4.3 sounds quite weak, as it only proves convergence to a stationary point only up to an extra term. One could anticipate that stochastic SAM may actually converge to global minima of smooth convex functions modulo the unavoidable additive factor. However, as briefly mentioned in Section 3.2, the next theorem presents a counterexample illustrating that ensuring convergence to global minima, even up to an additive factor, is impossible for m -SAM.

Theorem 4.4. For any ρ > 0 , β > 0 , σ > 0 , and η ≤ 1 β , there exists a β -smooth and convex function f satisfying the following. (1) The function f satisfies Assumption 2.5. (2) The component functions l ( · ; ξ ) of f are β -smooth for any ξ . (3) If we run m -SAM on f initialized inside a certain interval, then any arbitrary weighted average ¯ x of the iterates x 0 , x 1 , . . . must satisfy E [ f (¯ x ) -f ∗ ] ≥ C , and the suboptimality gap C can be made arbitrarily large and independent of the parameter ρ .

Here, we present the intuitions of the proof for Theorem 4.4. As demonstrated in Figures 2(d) and 2(e), the local concavity of the component function significantly influences the formation of attracting

basins in its virtual loss, thereby creating a region from which the m -SAM updates get stuck inside the basin forever.

Also note that we can construct the function to form the basin at any point with arbitrary large function value (and hence large suboptimality gap). Therefore, establishing an upper bound on the convergence of function value becomes impossible in smooth convex functions. A detailed analysis for the non-convergence example is presented in Appendix C.5.

4.3 Smooth and Nonconvex Functions

We now study smooth nonconvex functions. Extending Theorem 3.4, we can show the following bound for stochastic n -SAM.

Theorem 4.5. Consider a β -smooth function f satisfying f ∗ = inf x f ( x ) > -∞ , and assume Assumption 2.5. Under n -SAM, starting at x 0 with any perturbation size ρ > 0 and step size

The proof of Theorem 4.5 is provided in Appendix C.6. Notice that the non-convergence example presented in Theorem 3.5 (already) illustrates a scenario where E ∥∇ f ( x t ) ∥ 2 = Ω( ρ 2 ) , thereby confirming the tightness of additive factor in terms of ρ .

The scope of applicability for Theorem 4.5 is limited to n -SAM. Compared to n -SAM, m -SAM employs a stronger assumption where ξ = ˜ ξ . By imposing an additional Lipschitzness condition on the function f , m -SAM leads to a similar but different convergence result.

Theorem 4.6. Consider a β -smooth, L -Lipschitz continuous function f satisfying f ∗ = inf x f ( x ) > -∞ , and assume Assumption 2.5. Additionally assume l ( · , ξ ) is β -smooth for any ξ . Under m -SAM, starting at x 0 with any perturbation size ρ > 0 and step size η = √ ∆ √ β ( σ 2 + L 2 ) T to minimize f , we have

Corollary 4.7. Under the setting of Theorem 4.6, we get

The proofs for Theorem 4.6 and Corollary 4.7 are given in Appendix C.7. Since Theorem 3.5 presents an example where E ∥∇ f ( x t ) ∥ = Ω( ρ ) and E ( ∥∇ f ( x t ) ∥ -βρ ) 2 = Ω( ρ 2 ) , we can verify that the additive factors in Theorem 4.6 and Corollary 4.7 are tight in terms of ρ .

As for previous studies, Theorems 2, 12, and 18 of Andriushchenko and Flammarion [2] prove the convergence of n , m -SAM for smooth nonconvex functions: 1 T ∑ T t =0 E ∥∇ f ( x t ) ∥ 2 = O ( 1 √ T ) , but assuming SAM without gradient normalization , and sufficiently small ρ . For m -SAM on smooth Lipschitz nonconvex functions, Theorem 1 of Mi et al. [27] proves 1 T ∑ T t =0 E ∥∇ f ( x t ) ∥ 2 = ˜ O ( 1 √ T ) , while assuming decaying ρ . From our results, we demonstrate that such full convergence results are impossible for practical versions of stochastic SAM.

5 Conclusions

This paper studies the convergence properties of SAM, under constant ρ and with gradient normalization. We establish convergence guarantees of deterministic SAM for smooth and (strongly) convex functions. To our surprise, we discover scenarios in which deterministic SAM (for smooth nonconvex functions) and stochastic m -SAM (for all function class considered) converge only up to unavoidable additive factors proportional to ρ 2 . Our findings emphasize the drastically different characteristics of SAM with vs. without decaying perturbation size. Establishing tighter bounds in terms of β and µ , or searching for a non-convergence example that applies to n -SAM might be interesting future research directions.

Acknowledgments

This paper was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant (No. 2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)) funded by the Korea government (MSIT), two National Research Foundation of Korea (NRF) grants (No. NRF-2019R1A5A1028324, RS-2023-00211352) funded by the Korea government (MSIT), and a grant funded by Samsung Electronics Co., Ltd.

References

A. Agarwala and Y. N. Dauphin. Sam operates far from home: eigenvalue regularization as a dynamical phenomenon. arXiv preprint arXiv:2302.08692 , 2023.

M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning , pages 639-668. PMLR, 2022.

D. Bahri, H. Mobahi, and Y. Tay. Sharpness-aware minimization improves language model generalization. arXiv preprint arXiv:2110.08529 , 2021.

P. L. Bartlett, P. M. Long, and O. Bousquet. The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513 , 2022.

K. Behdin and R. Mazumder. Sharpness-aware minimization: An implicit regularization perspective. arXiv preprint arXiv:2302.11836 , 2023.

K. Behdin, Q. Song, A. Gupta, A. Acharya, D. Durfee, B. Ocejo, S. Keerthi, and R. Mazumder. msam: Micro-batch-averaged sharpness-aware minimization. arXiv preprint arXiv:2302.09693 , 2023.

J. M. Borwein and A. S. Lewis. Convex analysis and nonlinear optimization: Theory and examples, 2010.

S. Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning , 8(3-4):231-357, 2015.

X. Chen, C.-J. Hsieh, and B. Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. arXiv preprint arXiv:2106.01548 , 2021.

F. H. Clarke, Y. S. Ledyaev, R. J. Stern, and P. R. Wolenski. Nonsmooth analysis and control theory , volume 178. Springer Science & Business Media, 2008.

J. M. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065 , 2021.

E. M. Compagnoni, A. Orvieto, L. Biggio, H. Kersting, F. N. Proske, and A. Lucchi. An sde for modeling sam: Theory and insights. arXiv preprint arXiv:2301.08203 , 2023.

Y. Dai, K. Ahn, and S. Sra. The crucial role of normalization in sharpness-aware minimization. arXiv preprint arXiv:2305.15287 , 2023.

L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning , pages 1019-1028. PMLR, 2017.

J. Du, H. Yan, J. Feng, J. T. Zhou, L. Zhen, R. S. M. Goh, and V. Y. Tan. Efficient sharpnessaware minimization for improved training of neural networks. arXiv preprint arXiv:2110.03141 , 2021.

P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur. Sharpness-aware minimization for efficiently improving generalization. arXiv preprint arXiv:2010.01412 , 2020.

S. Hochreiter and J. Schmidhuber. Simplifying neural nets by discovering flat minima. Advances in neural information processing systems , 7, 1994.

W. Jiang, H. Yang, Y. Zhang, and J. Kwok. An adaptive policy to employ sharpness-aware minimization. In The Eleventh International Conference on Learning Representations , 2022.

Y. Jiang, B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio. Fantastic generalization measures and where to find them. arXiv preprint arXiv:1912.02178 , 2019.

J. Kaddour, L. Liu, R. Silva, and M. Kusner. When do flat minima optimizers work? In Advances in Neural Information Processing Systems , 2022.

N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04836 , 2016.

H. Kim, J. Park, Y. Choi, and J. Lee. Stability analysis of sharpness-aware minimization, 2023. URL https://arxiv.org/abs/2301.06308 .

H. Kim, J. Park, Y. Choi, W. Lee, and J. Lee. Exploring the effect of multi-step ascent in sharpness-aware minimization. arXiv preprint arXiv:2302.10181 , 2023.

J. Kwon, J. Kim, H. Park, and I. K. Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning , pages 5905-5914. PMLR, 2021.

H. Lee, H. Cho, H. Kim, D. Gwak, J. Kim, J. Choo, S.-Y. Yun, and C. Yun. Enhancing generalization and plasticity for sample efficient reinforcement learning. arXiv preprint arXiv:2306.10711 , 2023.

Y. Liu, S. Mai, X. Chen, C.-J. Hsieh, and Y. You. Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12360-12370, 2022.

P. Mi, L. Shen, T. Ren, Y. Zhou, X. Sun, R. Ji, and D. Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. arXiv preprint arXiv:2210.05177 , 2022.

T. Möllenhoff and M. E. Khan. Sam as an optimal relaxation of bayes. arXiv preprint arXiv:2210.01620 , 2022.

B. Neyshabur, S. Bhojanapalli, D. McAllester, and N. Srebro. Exploring generalization in deep learning. Advances in neural information processing systems , 30, 2017.

Y. Shi, Y. Liu, K. Wei, L. Shen, X. Wang, and D. Tao. Make landscape flatter in differentially private federated learning. arXiv preprint arXiv:2303.11242 , 2023.

H. Sun, L. Shen, Q. Zhong, L. Ding, S. Chen, J. Sun, J. Li, G. Sun, and D. Tao. Adasam: Boosting sharpness-aware minimization with adaptive learning rate and momentum for training deep neural networks. arXiv preprint arXiv:2303.00565 , 2023.

P. Wang, Z. Zhang, Z. Lei, and L. Zhang. Sharpness-aware gradient matching for domain generalization. arXiv preprint arXiv:2303.10353 , 2023.

K. Wen, T. Ma, and Z. Li. How does sharpness-aware minimization minimize sharpness? arXiv preprint arXiv:2211.05729 , 2022.

X. Zhang, R. Xu, H. Yu, H. Zou, and P. Cui. Gradient norm aware minimization seeks first-order flatness and improves generalization. arXiv preprint arXiv:2303.03108 , 2023.

J. Zhuang, B. Gong, L. Yuan, Y. Cui, H. Adam, N. Dvornek, S. Tatikonda, J. Duncan, and T. Liu. Surrogate gap minimization improves sharpness-aware training. arXiv preprint arXiv:2203.08065 , 2022.

Contents

A More Details on the Related Works of SAM

Sharpness and Generalization Hochreiter and Schmidhuber [17] propose an algorithm for finding low-complexity models with high generalization performance by finding flat minima. Keskar et al. [21] show that the poor generalization performance of large-batch training is due to the fact that large-batch training tends to converge to sharp minima. Jiang et al. [19] empirically show correlation between sharpness and generalization performance.

Motivated by prior studies, Foret et al. [16] propose Sharpness-Aware Minimization (SAM) that focuses on finding flat minima by aiming to (approximately) minimize f SAM ( x ) = max ∥ ϵ ∥≤ ρ f ( x + ϵ ) instead of f . Empirical results [3, 9, 16, 20, 25] show outstanding generalization performance of SAMonvarious tasks and models, including recent state-of-the-art models such as ViTs, MLP-Mixers, and T5.

Other Theoretical Properties of SAM. Bartlett et al. [4] prove that for quadratic functions, SAM dynamics oscillate and align to the eigenvector corresponding to the largest eigenvalue of Hessian, and show interpretation of SAM dynamics as GD on a 'virtual loss'. Wen et al. [33] prove that for sufficiently small perturbation size ρ of SAM, its trajectory follows a sharpness reduction flow which minimizes the maximum eigenvalue of Hessian. Dai et al. [13] examine the properties of SAM with and without normalization. They specifically focus on the stabilization and 'drift-alongminima' effects observed in SAM with normalization, which are not present in SAM without normalization. Compagnoni et al. [12] derive a continuous time SDE model for SAM with and without normalization, and employ the model to explain why SAM has a preference for flat minima. Furthermore, they demonstrate that SAM with and without normalization exhibit distinct implicit regularization properties, especially for ρ = O ( √ η ) .

Agarwala and Dauphin [1] prove that SAM produces an Edge Of Stability [11] stabilization effect, at a lower eigenvalue than gradient descent, under quadratic function settings. Kim et al. [22] show that saddle point acts as an attractor in SAM dynamics, and stochastic SAM takes more time to escape saddle point than stochastic gradient descent. Behdin and Mazumder [5] study the implicit regularization perspective of SAM, and provide a theoretical explanation of high generalization performance of SAM.

While Agarwala and Dauphin [1], Behdin and Mazumder [5], Kim et al. [22] provide useful insights on theoretical aspects of SAM, we point out that all proofs are restricted to SAM without gradient normalization in ascent steps. Considering the vastly different behaviors of SAM with and without normalization (Figure 1), it is not immediately clear whether these insights carry over to the practical version.

B Proofs for (Non-)Convergence of full-batch SAM

In this section, we show detailed proofs and explanations regarding convergence of SAM with constant ρ . The SAM algorithm we consider is defined as a two-step method:

B.1 Additional Function Class and Important Lemmas

In this section, we define an additional function class to use in convergence proofs. We also state and prove a few lemmas that are used in our theorem proofs.

Definition B.1 (Polyak-Lojasiewicz) . A function f : R d → R satisfies µ -PL condition if there exists µ > 0 such that,

for all x ∈ R d .

It is well-known that µ -strongly convex functions are µ -PL, but not vice versa.

Lemma B.2. For a differentiable and µ -strongly convex function f , we have

For a differentiable and convex function f , the inequality continues to hold with µ = 0 .

Proof. If a differentiable function f is µ -strongly convex, then its gradient map x ↦→ ∇ f ( x ) is µ -strongly monotone, i.e.,

Using this fact, we have

and this finishes the proof.

Lemma B.3. For a β -smooth and µ -strongly convex function, with step size η ≤ 1 2 β , we have

For a β -smooth and convex function, the inequality continues to hold with µ = 0 .

Proof. Starting from the definition of β -smoothness, we have

Since we assumed η ≤ 1 2 β , both η ( 1 -ηβ 2 ) ≥ η 2 and η (1 -ηβ ) ≥ η 2 hold. Applying Lemma B.2 to the above, we get

and this finishes the proof.

B.2 Convergence Proof for Smooth and Strongly Convex Functions (Proof of Theorem 3.1)

In this section, we prove Theorem 3.1, restated below for the sake of convenience.

Theorem 3.1. Consider a β -smooth and µ -strongly convex function f . If we run deterministic SAM starting at x 0 with any perturbation size ρ > 0 and step size η = min { 1 µT max { 1 , log ( µ 5 ∆ T 2 β 6 ρ 2 )} , 1 2 β } to minimize f , we have

Proof. The proof is divided into two cases, based on the ∥∇ f ( x t ) ∥ observed throughout the entire optimization process. In the first case, the gradient norm ∥∇ f ( x t ) ∥ at x t is sufficiently large for all t = 0 , . . . , T -1 ; in this case, we show that the SAM algorithm converges linearly. The other scenario corresponds to the case where there exists an iteration index t ∈ { 0 , . . . , T -1 } such that ∥∇ f ( x t ) ∥ is smaller than a certain threshold. In this case, we can show from the PL inequality (4) that we are already close to global optimality.

From Lemma B.3, we have

Now, recall from Definition B.1 that µ -strongly convex functions satisfy µ -PL inequaltiy (4). From this, we get

Case 1: when ∥∇ f ( x t ) ∥ remains large. Let us now consider the first case, where ∥∇ f ( x t ) ∥ ≥ ηβ 3 ρ µ holds for all t = 0 , . . . , T -1 . In this case, (5) becomes

which holds for all t = 0 , . . . , T -1 . Unrolling the inequality, we obtain

Case 2: when some ∥∇ f ( x t ) ∥ is small. In the other case where there exists t such that ∥∇ f ( x t ) ∥ ≤ ηβ 3 ρ µ , notice from µ -PL inequality (4) that

Therefore, combining the two cases, it is guaranteed that

We now elaborate how the choice of step size η = min { 1 µT max { 1 , log ( µ 5 ∆ T 2 β 6 ρ 2 )} , 1 2 β } results in the convergence rate in the theorem statement. Naturally, there are four cases, depending on the outcomes of the min and max operations.

Case B: log ( µ 5 ∆ T 2 β 6 ρ 2 ) ≥ 1 and 1 µT log ( µ 5 ∆ T 2 β 6 ρ 2 ) ≥ 1 2 β . Substituting η = 1 2 β ≤ 1 µT log ( µ 5 ∆ T 2 β 6 ρ 2 ) to (6),

Case C: log ( µ 5 ∆ T 2 β 6 ρ 2 ) ≤ 1 and 1 µT ≤ 1 2 β . Putting η = 1 µT ≥ 1 µT log ( µ 5 ∆ T 2 β 6 ρ 2 ) into (6),

Case D: log ( µ 5 ∆ T 2 β 6 ρ 2 ) ≤ 1 and 1 µT ≥ 1 2 β . By substituting η = 1 2 β ≤ 1 µT to (6) we obtain

Combining the four cases, we conclude

thereby completing the proof.

B.3 Lower Bound Proof for Smooth and Strongly Convex Functions (Proof of Theorem 3.2)

In this section, we prove Theorem 3.2, restated below for the sake of convenience.

Theorem 3.2. Suppose β µ ≥ 2 . For any choice of perturbation size ρ , step size η , and initialization x 0 , there exists a differentiable, β -smooth, and µ -strongly convex function f such that

holds for deterministic SAM iterates.

The proof is divided into three cases, depending on the step size. For each case, we will define and analyze a one-dimensional function to show the lower bound. In doing so, without loss of generality we will fix an initialization x 0 because we can appropriately shift the function for different choices of x 0 .

For η ≤ 1 2 µT , we show that there exists a function f such that

For 1 2 µT ≤ η ≤ 2 β , we show existence of a function f that satisfies

For η ≥ 2 β , we show that there exists a function f such that

Combining the three results, we can see that the suboptimality gap of the best iterate is at least Ω ( β 3 ρ 2 µ 2 T 2 ) , hence proving the theorem. Below, we prove the statements for the three intervals.

Case 1: η ≤ 1 2 µT . Consider a function

Suppose we start at initialization x 0 = 2 β 2 ρ µ 2 ≥ 8 ρ . Note from the definition of ∇ f that for any x t ≥ ρ , we have y t = x t + ρ ∇ f ( x t ) |∇ f ( x t ) | = x t + ρ and ∇ f ( y t ) = µ ( x t + ρ -ρ ) = µx t . Therefore, the first SAM update can be written as

Since 1 -ηµ ≥ 1 -1 2 T ≥ 1 2 , the inequality x 1 ≥ ρ still holds and the same argument can be repeated for the second update, yielding

In fact, (1 -ηµ ) t ≥ (1 -1 2 T ) T ≥ 1 2 for all t ≤ T , so all the iterates up to x T stay above ρ . Thus, for any t = 0 , . . . , T , the iterate x t can be written and bounded from below as

Therefore, in this case, the suboptimality gap of the best iterate is at least

Case 2: 1 2 µT ≤ η ≤ 2 β . Consider a function

Suppose we start at initialization x 0 = ηβρ 4 -ηβ . We are going to show that the SAM iterates oscillate between ± ηβρ 4 -ηβ . Indeed, if x t = ηβρ 4 -ηβ , then ∇ f ( x t ) > 0 and

Then, after SAM update, we get

The same argument can be repeated to show that x t +2 = -x t +1 = x t . As a result, the iterates oscillate between ± ηβρ 4 -ηβ forever. In this case, the suboptimality gap is bounded from below by

Applying η ≥ 1 2 µT yields

Case 3: η ≥ 2 β . Consider a function

Let the initialization be x 0 = βρ µ . For any x t ≥ 0 , we have y t = x t + ρ and ∇ f ( y t ) = β ( x t + ρ ) . Then, the resulting SAM update becomes

Since we have η ≥ 2 β , we have 1 -ηβ ≤ -1 and x t +1 has the opposite sign as x t and its absolute value is at least as large as | x t | . We can similarly check that any further SAM update changes the sign and does not decrease the absolute value. Therefore, for all t = 0 , . . . , T , we have | x t | ≥ | x 0 | .

Consequently, the suboptimality gap of the best iterate is at least

Remarks on validity of lower bound. Lastly, we comment on why we choose different functions and initialization for different choices of η and why it suffices to provide a matching lower bound for Theorem 3.1. In convergence upper bounds in the form of Theorem 3.1, we aim to prove an upper bound on the following minimax risk :

where A denotes the class of algorithms , F denotes the class of functions , and ζ ( A ( f )) is the suboptimality measure for algorithm's output A ( f ) for function f . In the context of Theorem 3.1, our algorithm class A corresponds to the choices of the 'hyperparameters' η , ρ , and x 0 of SAM, and the function class F here is the class analyzed by the theorem: the collection of β -smooth and µ -strongly convex functions. From this viewpoint, Theorem 3.1 can be thought of as an upper bound on (7), with the choice of ζ ( A ( f )) ≜ min t ∈{ 0 ,...,T } f ( x t ) -f ∗ .

Hence, showing a matching lower bound for Theorem 3.1 amounts to showing a matching lower bound for the minimax risk (7). For this purpose, it suffices to show that for each A ∈ A , there exists a choice of f ∈ F such that a certain lower bound holds. Therefore, we are allowed to choose different choices of f for each different choice of hyperparameters η , ρ , and x 0 . In fact, in our proof, we choose different choices of f and x 0 for each different choice of η ; however, this is without loss of generality once we notice here that starting an algorithm at x 0 to minimize f ( x ) is equivalent to starting at 0 to minimize f ( x -x 0 ) . Hence, even though we choose different f 's and x 0 's for different choices of η in our proof of Theorem 3.2, this is sufficient for providing a matching upper bound for Theorem 3.1.

Admittedly, some authors show stronger versions than what we show, where a single function takes care of all possible hyperparameters. Indeed, such results provide lower bounds for sup f ∈F inf A ∈A ζ ( A ( f )) . Recalling that sup inf ≤ inf sup , one can notice that these sup inf lower bounds are in fact much stronger than what suffices.

B.4 Convergence Proof for Smooth and Convex Functions (Proof of Theorem 3.3)

For smooth and convex function, we prove the convergence of gradient norm to zero. The theorem is restated for convenience.

Theorem 3.3. Consider a β -smooth and convex function f . If we run deterministic SAM starting at x 0 with any perturbation size ρ > 0 and step size η = min { √ 2∆ √ β 3 ρ 2 T , 1 2 β } to minimize f , we have

Proof. Using Lemma B.3, for smooth and convex function f , we have

Adding up the inequality for t = 0 , . . . , T -1 , and the dividing both sides by T , we get

We now spell out how the choice of step size η = min {√ 2∆ β 3 ρ 2 T , 1 2 β } results in the convergence rate in the theorem statement. There are two cases to be considered, depending on the outcome of the min operation.

Case B: √ 2∆ β 3 ρ 2 T ≥ 1 2 β . Substituting η = 1 2 β ≤ √ 2∆ β 3 ρ 2 T to (8) yields

Merging the two cases, we conclude

which can be rewritten as completing the proof.

B.5 Convergence Proof for Smooth and Nonconvex Functions (Proof of Theorem 3.4)

In this section, we prove the convergence up to an additive factor β 2 ρ 2 for smooth and nonconvex functions. The theorem is restated for convenience.

Theorem 3.4. Consider a β -smooth function f satisfying f ∗ = inf x f ( x ) > -∞ . If we run deterministic SAM starting at x 0 with any perturbation size ρ > 0 and step size η = 1 β to minimize f , we have

Proof. Starting from the definition of β -smoothness, we have

The inequality can be rearranged as

Adding up the inequality for t = 0 , . . . , T -1 , and the dividing both sides by T , we get

Substituting η = 1 β to (9) yields

B.6 Non-convergence for a Smooth and Nonconvex Function (Proof of Theorem 3.5)

In this section, we spell out the proof of our counterexample that SAM with constant perturbation size ρ provably fails to converge to a stationary point. Here we restate the theorem.

Theorem 3.5. For any ρ > 0 and η ≤ 1 β , there exists a β -smooth and Θ( βρ ) -Lipschitz continuous function such that, if deterministic SAM is initialized at a point x 0 sampled from a continuous probability distribution, then deterministic SAM converges to a nonstationary point, located at a distance of Ω( ρ ) from a stationary point, with probability 1 .

Proof. For x ∈ R , consider the following one-dimensional function: given ρ and β ,

It is easy to check that this function is β -smooth and 3 βρ 5 π -Lipschitz continuous. Also, let

which is the set of points x where ∇ f ( x ) = 0 . For SAM with perturbation size ρ , given the current iterate x t , the corresponding y t = x t + ρ ∇ f ( x t ) |∇ f ( x t ) | is given by

This leads to the following a virtual gradient map G f :

From this virtual gradient map, we can define

̸

and Y is the set of all points x where G f ( x ) = 0 and ∇ f ( x ) = 0 .

Since f is a one-dimensional function, a virtual loss J f can be obtained by integrating G f . One possible example is

which is a piecewise β -smooth function that is minimized at points in Y and locally maximized at (non-differentiable) points in X . Since J f is well-defined, we can view SAM as GD on J f .

For sufficiently small η ≤ 1 β and any initialization x 0 ∈ R \ X , the initialization belongs to one of the intervals (( -0 . 3 + 0 . 6 k ) ρ, (0 . 3 + 0 . 6 k ) ρ ) . For such small enough η , we can guarantee that the SAM iterates x t will stay in the interval (( -0 . 3 + 0 . 6 k ) ρ, (0 . 3 + 0 . 6 k ) ρ ) for all t ≥ 0 . For example, suppose that x t ∈ ( -0 . 3 ρ, 0 . 3 ρ ) . Then, the next iterate will stay in the same interval if and only if the image of the interval ( -0 . 3 ρ, 0 . 3 ρ ) under a map x ↦→ x -3 ηβρ 5 π cos ( 5 π 3 ρ x + 5 π 3 ) is a subset of ( -0 . 3 ρ, 0 . 3 ρ ) .

and the containment is true if ηβ ≤ 1 . By symmetry, the same argument can be applied to any intervals of the form (( -0 . 3 + 0 . 6 k ) ρ, (0 . 3 + 0 . 6 k ) ρ ) .

Thus, for any initialization x 0 ∈ R \ X , all SAM iterates stay inside the interval which x 0 belongs to. Inside the interval, by β -smoothness of J f , the following descent lemma always holds:

Therefore, if we add the inequalities up for t = 0 , . . . , T -1 , we get

where J ∗ f ≜ inf x J f ( x ) > -∞ . Since the inequality holds for all T , the series is summable, which in turn implies that ∥ G f ( x t ) ∥ → 0 as t → ∞ . However, the only point in the interval (( -0 . 3 + 0 . 6 k ) ρ, (0 . 3 + 0 . 6 k ) ρ ) satisfying ∥ G f ( x ) ∥ = 0 is the one in Y ∩ (( -0 . 3 + 0 . 6 k ) ρ, (0 . 3 + 0 . 6 k ) ρ ) . Hence, if x 0 ∈ R \ X and η ≤ 1 β , SAM must converge to a point in Y . Since x 0 is drawn from a continuous probability distribution, x 0 ∈ R \ X holds almost surely. This finishes the proof.

B.7 Non-convergence for a Nonsmooth and Convex Functions (Proof of Theorem 3.6)

For nonsmooth convex function, depending on the initialization, SAM can converge to a suboptimal point with distance Ω( ρ ) from the global minimum. The theorem is restated for convenience.

Theorem 3.6. For any ρ > 0 and η < 7 ρ 4 , there exists a nonsmooth Lipschitz convex function f such that for some initialization, deterministic SAM converges to suboptimal points located at a distance of Ω( ρ ) from the global minimum.

(a)

(b)

SAM

4

V2

V1

-2

0

1

2

E

4

-3

2

-1

0

2

3

Proof. For x = ( x (1) , x (2) ) ∈ R 2 , consider a 2-dimensional function,

A set { v 1 , v 2 } can be established as a basis for R 2 , with v 1 = e 1 , and v 2 = 2 √ 5 e 1 + 1 √ 5 e 2 . Any vector x ∈ R 2 can be uniquely expressed by v 1 and v 2 . Consider the region: { x = b (1) v 1 + b (2) v 2 | b (1) < 0 } .

This area can be partitioned into four separate regions.

Figure 3(a) demonstrates the regions A , B , C , D , as well as the vectors v 1 and v 2 . When x belongs to the set x = b (1) v 1 + b (2) v 2 | b (1) < 0 , we can examine four different scenarios.

Case A: x = b (1) v 1 + b (2) v 2 ∈ A . x ∈ A , so y = x + ρ ∇ f ( x ) ∥∇ f ( x ) ∥ and -∇ f ( y ) are as follows.

As a result, the SAM updates ∇ f ( y ) only affect v 2 , and does not affect on v 1 . Thus, if x t ∈ A , the next SAM iterate x t +1 remains in A .

Case B: x = b (1) v 1 + b (2) v 2 ∈ B . We can verify that y = x + ρ ∇ f ( x ) ∥∇ f ( x ) ∥ = x + ρ v 2 . Therefore, -∇ f ( y ) = - √ 5 v 2 . Therefore, the SAM update shifts in the direction of -v 2 . Hence, if x t ∈ B , the next SAM iterate will be x t +1 ∈ ( B ∪ C ∪ D ) .

For x t = b (1) t v 1 + b (2) t v 2 ∈ C , if b (1) t < √ 5 b (2) t < b (1) t + ρ , then SAM update shifts in the direction of v 2 , resulting in x t +1 ∈ ( B ∪ C ∪ D ) . Otherwise, the next SAM iterate x t +1 shifts in the direction of + v 1 . Consequently, x t +1 can either remain in B ∪ C ∪ D , or move to A .

Assume that x t +1 moves to A . Given that b (1) t < -7 ρ 2 for x t ∈ C , we can verify that b (1) t +1 < -7 ρ 2 + η .Based on our previous observation in Case A where SAM updates only affect v 2 when x ∈ A , we can conclude that for all subsequent iterates x i with i > t +1 , the condition b (1) i < -7 ρ 2 + η continues to hold.

We can check that -∇ f ( y ) = √ 5 v 2 for all cases. So the SAM update shifts in the direction of + v 2 . As a result, if x t ∈ D , the next SAM iterate x t +1 will fall into x t +1 ∈ ( B ∪ C ∪ D ) .

Furthermore, we can verify that for x ∈ ( B ∪ C ∪ D ) , b (1) < -7 ρ 2 holds. Therefore, summing up all the cases, we can conclude that if the initial iterate x 0 is chosen from x 0 ∈ ( B ∪ C ∪ D ) , then all subsequent SAM iterates x t will satisfy b (1) t < -7 ρ 2 + η . Since the global minimum is located at x ∗ = (0 , 0) , it follows that ∥ x t -x ∗ ∥ > | 7 ρ/ 2 -η | √ 5 for every t > 0 . Thus, for η < 7 ρ 4 , we can ascertain that in this particular function, ∥ x t -x ∗ ∥ > 7 ρ 4 √ 5 for every t > 0 , thereby proving that the distance from the global minimum is at least Ω( ρ ) . The trajectory plot of SAM on this example function is illustrated in Figure 3(b).

C Proofs for (Non-)Convergence of Stochastic SAM

In this section, we provide in-depth demonstrations and proofs regarding convergence of stochastic SAM with constant perturbation size ρ . The objective is defined as f ( x ) = E ξ [ l ( x ; ξ )] , where ξ represents a stochastic parameter (e.g., data sample) and l ( x ; ξ ) represents the loss at point x with a random sample ξ . The update iteration for stochastic SAM is specified as follows:

We define g ( x ) = ∇ x l ( x ; ξ ) and ˜ g ( x ) = ∇ x l ( x ; ˜ ξ ) , where ξ and ˜ ξ are stochastic parameters. Stochastic SAM comes in two variations: n -SAM, where ξ and ˜ ξ are independent, and m -SAM, where ξ is equal to ˜ ξ .

C.1 Important Lemmas Regarding Stochastic SAM

In this section, we present a number of lemmas that are utilized in our theorem proofs regarding stochastic SAM. In order to do this, we also introduce extra notation, ˆ y t = x t + ρ ∇ f ( x t ) ∥∇ f ( x t ) ∥ , as a deterministically ascended parameter.

Lemma C.1. For a differentiable and µ -strongly convex function f , we have

For a differentiable and convex function f , the inequality continues to hold with µ = 0 .

Proof.

where we use Lemma B.2 in the last inequality, thereby completing the proof.

Lemma C.2. Under n -SAM, for a β -smooth and µ -strongly convex function f , we have

Under m -SAM, additionally assuming l ( · , ξ ) is β -smooth for any ξ , the inequality continues to hold.

Proof. First we consider n -SAM. Starting from Lemma C.1, we have

Next we consider m -SAM. additionally assuming l ( · , ξ ) is β -smooth for any ξ , starting from Lemma C.1, we have

completing the proof.

Lemma C.3. Consider a β -smooth, µ -strongly convex function f , and assume Assumption 2.5. Under n -SAM, with step size η ≤ 1 2 β , we have

Under m -SAM, additionally assuming l ( · , ξ ) is β -smooth for any ξ , the inequality continues to hold.

Proof. Starting from the definition of β -smoothness, we have

Since we assumed η ≤ 1 2 β , η (1 -ηβ ) ≥ η 2 hold. Applying Lemma C.2, we get

completing the proof.

C.2 Convergence Proof for Smooth and Strongly Convex Functions (Proof of Theorem 4.1)

In this section, we demonstrate the convergence result of stochastic SAM for smooth and strongly convex functions. For convenience, we restate the theorem here.

Theorem 4.1. Consider a β -smooth, µ -strongly convex function f , and assume Assumption 2.5. Under n -SAM, starting at x 0 with any perturbation size ρ > 0 and step size η = min { 1 µT · max { 1 , log ( µ 2 ∆ T β [ σ 2 -β 2 ρ 2 ] + )} , 1 2 β } to minimize f , we have

Under m -SAM, additionally assuming l ( · , ξ ) is β -smooth for any ξ , the inequality continues to hold.

Proof. We start the proof from Lemma C.3; in order to apply the lemma, additionally assuming β -smoothness for component functions l ( · , ξ ) is necessary for m -SAM.

Since µ -strongly convex functions satisfy µ -PL inequality (B.1), we get

Depending on the value of σ , there are two cases in which the convergence rate varies.

Case A: σ ≤ βρ . In this case, we have

and our choice of η must be 1 2 β . Unrolling the inequality and substituting η = 1 2 β draws out

Case B: σ > βρ . In this case, we have

Again, unrolling the inequality draws out

Similar to Section B.2, substituting η = min { 1 µT · max { 1 , log ( µ 2 ∆ T β ( σ 2 -β 2 ρ 2 ) )} , 1 2 β } can result in four cases.

Case B-1: log ( µ 2 ∆ T β ( σ 2 -β 2 ρ 2 ) ) ≥ 1 , and 1 2 β ≥ 1 µT log ( µ 2 ∆ T β ( σ 2 -β 2 ρ 2 ) ) . Setting η = 1 µT log ( µ 2 ∆ T β ( σ 2 -β 2 ρ 2 ) ) ,

Merging all four cases, we get

thereby finishing the proof.

C.3 Non-Convergence of m -SAM for a Smooth and Strongly Convex Function (Proof of Theorem 4.2)

In this section, we present a formal analysis of our counterexample, which shows that stochastic m -SAM with constant perturbation size ρ may fail to fully converge to a global minimum, and can get close to the global minimum by only Ω( ρ ) . This counterexample indicates that the O ( ρ 2 ) term in Theorem 4.1 is unavoidable. For readers' convenience, we restate the theorem.

Theorem 4.2. For any ρ > 0 , β > 0 , σ > βρ and η ≤ 3 10 β , there exists a β -smooth and β 5 -strongly convex function f satisfying the following. (1) The function f satisfies Assumption 2.5. (2) The component functions l ( · ; ξ ) of f are β -smooth for any ξ . (3) If we run m -SAM on f initialized inside a certain interval, then any arbitrary weighted average ¯ x of the iterates x 0 , x 1 , . . . must satisfy E [ f (¯ x ) -f ∗ ] ≥ Ω( ρ 2 ) .

Proof. Given ρ > 0 , β > 0 , σ ≥ 0 , and η ≤ 3 10 β , we choose p ≜ 2 3 , c ≜ ( 1 + p 4 ) ρ . For x ∈ R , and a constant a > 0 which will be chosen later. Using these p , c , and a , we construct the counterexample. We consider a one-dimensional smooth strongly convex function

and the stochastic function l ( x ; ξ ) can be given as

where each component functions are as described below:

First, it is easy to verify that E l ( x ; ξ ) = f ( x ) . We can also confirm that f , f (1) are a -smooth, and f (2) is ( 1+ p 1 -p a ) -smooth. Furthermore, ∥∇ f ( x ) - ∇ f (1) ( x ) ∥ 2 ≤ a 2 (2 ρ + c ) 2 holds, along with

Now choose a = min { β 5 , σ 5 ρ } ≤ β 5 . This choice ensures that f , f (1) , and f (2) are all β -smooth, as desired. Also notice that 2 a 2 · ( 19 ρ 6 ) 2 ≤ 2 · ( 19 5 · 6 ) 2 σ 2 < σ 2 , so the function satisfies Assumption 2.5.

For the remaining of the proof, we investigate the virtual gradient maps G f (1) and G f (2) within the specified region of interest: c -ρ ≤ x ≤ c + ρ . We will then demonstrate that if we initialize m -SAM in this interval, all subsequent iterations of m -SAM will remain inside [ c -ρ, c + ρ ] , as required above.

Within this interval, the perturbed iterate y = ρ ∇ f (1) ( x ) ∥∇ f (1) ( x ) ∥ and the virtual gradient map G f (1) of f (1) can be described as follows.

Additionally defining c ′ ≜ p 1+ p c , we now calculate y = ρ ∇ f (2) ( x ) ∥∇ f (2) ( x ) ∥ and the virtual gradient map G f (2) of f (2) .

Here, given c = ( 1 + p 4 ) ρ , we can verify that

where 0 , c ′ , and c are the global minimum (or maximum) of f , f (2) , and f (1) , respectively.

Recall that a ≤ β 5 . Since η ≤ 3 10 β , we have ηa ≤ 3 50 , which will be useful for the rest of the proof. Now, we analyze x t +1 , the next iterate of m -SAM. For c -ρ ≤ x t ≤ c + ρ , the next iterate x t +1 of m -SAM using G f (1) ( x t ) is

We will show that x t +1 must stay within the interval [ c -ρ, c + ρ ] . For the case x t = c , the inclusion is trivial. To analyze the rest, we divide into two cases.

Case A-1: c -ρ ≤ x t < c . Since ηa ≤ 3 50 , we have

Case A-2: c < x t ≤ c + ρ . Since ηa ≤ 3 50 , we have

Cases A-1 and A-2 show that x t +1 also remains in [ c -ρ, c + ρ ] , when we update m -SAM using G f (1) ( x t ) .

(a)

(b)

f(1)

Jf(1)

f(2)

Jf2)

f

Ef

3

3

2

3

-2

1

2

3

4

1

We next examine the case involving G f (2) . The next iterate x t +1 of m -SAM using G f (2) ( x t ) is

Again, we divide it into two cases, since the x t = c ′ case is obvious.

Since ηa ≤ 3 50 , we have 1 -η a (1+ p ) 1 -p = 1 -5 ηa ≥ 0 , so

where the last inequality used ηa ≤ 3 50 .

Cases B-1 and B-2 demonstrate that x t +1 also remains within the interval [ c -ρ, c + ρ ] when we update m -SAM using G f (2) ( x t ) .

To gain a better intuitive understanding, we provide Figures 4(a), 4(b), demonstrating the original and virtual loss functions of f (1) and f (2) . In Figure 4(b), we can examine that m -SAM updates generate an attraction basin in J f (1) , thereby making a region that m -SAM cannot escape.

The aforementioned case analyses indicate that when the initial point x 0 falls within the interval [ c -ρ, c + ρ ] , all subsequent iterations of m -SAM will also remain within this interval, regardless of the selected component function for updating. Given c = ( 1 + p 4 ) ρ = 7 6 ρ , we can conclude that all subsequent iterations of m -SAM are points located at a distance of at least 1 6 ρ from the global optimum, for any β > 0 , ρ > 0 , and σ > 0 .

Moreover, in case of σ > βρ , it follows that f ( x ) = a 2 x 2 = β 10 x 2 . Consequently, the suboptimality gap at any timestep is at least

thereby finishing the proof.

C.4 Convergence Proof for Smooth and Convex Functions (Proof of Theorem 4.3)

In this section, we establish the convergence of stochastic SAM for smooth and convex functions to near-stationary points. For ease of understanding, we provide the theorem statement here.

Theorem 4.3. Consider a β -smooth, convex function f , and assume Assumption 2.5. Under n -SAM, starting at x 0 with any perturbation size ρ > 0 and step size η = min { √ ∆ √ β [ σ 2 -β 2 ρ 2 ] + T , 1 2 β } to minimize f , we have

Under m -SAM, additionally assuming l ( · , ξ ) is β -smooth for any ξ , the inequality continues to hold.

Proof. Westart from Lemma C.3 with µ = 0 . In order to do this, additionally assuming β -smoothness of component functions l ( · , ξ ) is necessary for m -SAM.

which can be rewritten as

Adding up the inequality for t = 0 , · · · , T -1 , and dividing both sides by T , we get

The convergence rate varies depending on two different cases determined by the value of σ .

Case A: σ ≤ βρ . In this case, it must hold that η = 1 2 β . By this choice of η , we have

Case B: σ > βρ . Setting η = min {√ ∆ β ( σ 2 -β 2 ρ 2 ) T , 1 2 β } , we consider two cases.

Merging the two cases, we conclude

thereby completing the proof.

C.5 Non-Convergence of m -SAM for a Smooth and Convex Function (Proof of Theorem 4.4)

In this section, we provide an in-depth analysis of our counterexample, which demonstrates that stochastic m -SAM with constant perturbation size ρ can provably converge to a point that is not a global minimum. Furthermore, the suboptimality gap in terms of the function value can be made arbitrarily large, hence indicating that proving convergence of m -SAM to global minima (modulo some additive factors O ( ρ 2 ) ) is impossible. The theorem is restated for convenience.

Theorem 4.4. For any ρ > 0 , β > 0 , σ > 0 , and η ≤ 1 β , there exists a β -smooth and convex function f satisfying the following. (1) The function f satisfies Assumption 2.5. (2) The component functions l ( · ; ξ ) of f are β -smooth for any ξ . (3) If we run m -SAM on f initialized inside a certain interval, then any arbitrary weighted average ¯ x of the iterates x 0 , x 1 , . . . must satisfy E [ f (¯ x ) -f ∗ ] ≥ C , and the suboptimality gap C can be made arbitrarily large and independent of the parameter ρ .

Proof. Given ρ > 0 , β > 0 , σ > 0 , we set an arbitrary constant c > 5 ρ 4 , and a parameter a > 0 which will be chosen later. For x ∈ R , consider a one-dimensional smooth convex function,

where the functions f (1) and f (2) are given by the following definitions.

and l ( x ; ξ ) can be given as

It can be verified that f (1) is ( max { 8 a ρ , β }) -smooth, and f (2) is ( max { 8 a ρ · p 1 -p , β }) -smooth. Moreover, ∥∇ f ( x ) -∇ f (1) ( x ) ∥ 2 ≤ 9 a 2 holds, as well as ∥∇ f ( x ) -∇ f (2) ( x ) ∥ 2 ≤ 9 a 2 · p 2 (1 -p ) 2 . Consequently, E ∥∇ f ( x ) -g ( x ) ∥ 2 ≤ 9 a 2 · p 1 -p .

By selecting p > 1 2 , and setting a = min { βρ (1 -p ) 8 p , σ √ 1 -p 3 √ p } , we can check that E l ( x ; ξ ) = f ( x ) , and the component functions are β -smooth. Additionally, f satisfies Assumption 2.5.

In the following analysis, we examine the virtual gradient maps G f (1) and G f (2) in the specified region of interest: c -ρ ≤ x ≤ c + ρ . For this specific interval, we are going to show that if we start m -SAM inside this interval, then all the subsequent iterates of m -SAM must stay inside the same interval [ c -ρ, c + ρ ] .

In this region, y = x + ρ ∇ f (1) ( x ) ∥∇ f (1) ( x ) ∥ and G f (1) are as follows.

Additionally defining c ′ ≜ c -ρ 8 p , we can compute y = x + ρ ∇ f (2) ( x ) ∥∇ f (2) ( x ) ∥ and G f (2) as follows.

Now recall that η ≤ 1 β . Given that a ≤ βρ (1 -p ) 8 p , we can derive η ≤ 1 β ≤ ρ a · 1 -p 8 p . Furthermore, the next iterate x t +1 of m -SAM using G f (1) ( x t ) for c -ρ ≤ x t ≤ c + ρ is

We will now show that x t +1 must remain in the interval [ c -ρ, c + ρ ] . For the case x t = c , it is obvious. Dividing the rest into three cases,

Case A-1: c -ρ ≤ x t ≤ c -3 ρ 4 . Since 0 ≤ η · 8 a ρ ( x t + ρ -c ) ≤ 2 ηa ≤ ρ · 1 -p 4 p ≤ ρ 4 , we have

which proves that x t +1 also remains in [ c -ρ, c + ρ ] .

(a)

(b)

:0

f(1)

Jf(1)

0

f(2)

Jf(2)

f

Ef

.0

1

2

3

5

8

4

b

8

10

5

Case A-2: c -3 ρ 4 ≤ x t < c , c < x t ≤ c + 3 ρ 4 . Since 0 ≤ 2 ηa ≤ ρ · 1 -p 4 p ≤ ρ 4 , we have

thereby proving that x t +1 also remains in [ c -ρ, c + ρ ] .

Case A-3: c + 3 ρ 4 ≤ x t ≤ c + ρ . Since 0 ≥ η · 8 a ρ ( x t -ρ -c ) ≥ -2 ηa ≥ -ρ · 1 -p 4 p ≥ -ρ 4 , we have

which shows that x t +1 also remains in [ c -ρ, c + ρ ] .

We next consider the case of G f (2) . The next iterate x t +1 of m -SAM using G f (2) ( x t ) for c -ρ ≤ x t ≤ c + ρ is given by

Again, the x t = c ′ case is trivial. Considering the remaining two cases,

Case B-1: c -ρ ≤ x t < c ′ . Since 0 ≥ ηa · 1 -2 p 1 -p ≥ ρ · 1 -2 p 8 p ≥ -( c + ρ -c ′ ) , we get

which indicates that x t +1 also remains in [ c -ρ, c + ρ ] .

Case B-2: c ′ < x t ≤ c + ρ . Since 0 ≤ ηa · 1+2 p 1 -p ≤ ρ · 1+2 p 8 p ≤ c ′ -( c -ρ ) , we have

which shows that x t +1 also remains in [ c -ρ, c + ρ ] .

The case analyses above indicate that if x 0 is initialized in [ c -ρ, c + ρ ] , the subsequent iterates of m -SAM also remain in [ c -ρ, c + ρ ] , regardless of the component function chosen by m -SAM for update. Figures 5(a) and 5(b) demonstrate the original loss function and virtual loss function of component functions f (1) and f (2) . Due to the concavity of f (1) , a basin of attraction is formed in J f (1) , thereby creating a region where the iterations of m -SAM cannot escape.

Therefore, the suboptimality gap at any timestep is at least

Therefore, regardless of iterate averaging scheme, the suboptimality gap in terms of function value will stay above a ( c -ρ ) + a 2 2 β .

Moreover, this suboptimality gap can be made arbitrarily large if we choose larger values of c ; notice that c > 5 ρ 4 is the only requirement on c . Consequently, it is impossible to guarantee convergence to global minima, even up to an additive factor, for m -SAM. This finishes the proof.

C.6 Convergence Proof for Smooth and Nonconvex Function under n -SAM (Proof of Theorem 4.5)

In this section, we establish the convergence result of stochastic n -SAM for smooth and nonconvex functions. For convenience, we restate the theorem.

Theorem 4.5. Consider a β -smooth function f satisfying f ∗ = inf x f ( x ) > -∞ , and assume Assumption 2.5. Under n -SAM, starting at x 0 with any perturbation size ρ > 0 and step size

Proof. Starting from the definition of β -smoothness, we have

Rearranging the inequality, we get

Adding up the inequality for t = 0 , · · · , T -1 , and dividing both sides by T , we get

Substituting η = min { 1 2 β , √ ∆ √ Tβσ 2 } to (12), we get two cases.

√

Merging two cases, we can conclude that

thereby completing the proof.

C.7 Convergence Proof for Smooth Lipschitz Nonconvex Function under m -SAM (Proof of Theorem 4.6 and Corollary 4.7)

In this section, the convergence proof for stochastic m -SAM for smooth, Lipschitz, and nonconvex functions is presented. The notation ˆ y t = x t + ρ ∇ f ( x t ) ∥∇ f ( x t ) ∥ is used here once again. The theorem is restated for convenience.

Theorem 4.6. Consider a β -smooth, L -Lipschitz continuous function f satisfying f ∗ = inf x f ( x ) > -∞ , and assume Assumption 2.5. Additionally assume l ( · , ξ ) is β -smooth for any ξ . Under m -SAM, starting at x 0 with any perturbation size ρ > 0 and step size η = √ ∆ √ β ( σ 2 + L 2 ) T to minimize f , we have

Proof. The proof technique resembles Mi et al. [27]. Starting from the definition of β -smoothness, we have

Rearranging the inequality, we get

Adding up the inequality for t = 0 , · · · , T -1 , and dividing both sides by T , we get

Substituting η = √ ∆ √ β ( σ 2 + L 2 ) T to (13) yields

thereby completing the proof.

From the result of Theorem 4.6, Corollary 4.7 can be derived. The corollary is restated for the ease of reference.

Corollary 4.7. Under the setting of Theorem 4.6, we get

Proof. Starting from the result of Theorem 4.6, we have

Taking the square root on both sides,

Rearranging the inequality, we obtain

thereby completing the proof.

(a)

(b)

(c)

(d)

(e)

(f)

f

f(1)

f(1)

SAMiterates

0.015

f(2)

f(2)

f

f

0.010

SAM iterates

SAM iterates

0.005

2

0.000

10

:.0

1.5

1.0

0.5

0

0.5

1.0

1.5

2.

0.005

5

-0.010

1

2

3

0.015

Z

8

1

0.70

smooth deterministic SAM

4.5

smooth strongly convex m-SAM

10

smooth convex m-SAM

4.0

0.65

6

3.5

0.60

3.0

8

×0.55

2.5

7

0.50

2.0

6

1.5

0.45

1.0

4.86

0.40

0

250

500

750

1000

1250150017502000

0.69

0.5

0

250

500

750

1000

1250150017502000

4

0

250

00S

750

1000

1250150017502000

D Simulation Results of SAM on Example Functions

Figure 6 shows the results of SAM simulations on the example functions considered in Sections 3 and 4. It demonstrates that SAM indeed does not converge in our worst-case example constructions.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper analyzes the convergence properties of Sharpness-Aware Minimization (SAM) with **constant** perturbation size $\\rho$ and gradient normalization applied to the updates; this has not been done in prior works. Both deterministic and stochastic settings are considered. The authors show the following:\n\n1. For the strongly convex + smooth case in the deterministic setting, SAM converges to the global minimum. A matching lower bound for the rate is also presented although I'm not sure if it's correct; please see Weakness #2. For the general convex + smooth case in the deterministic setting, the authors are only able to show convergence to a stationary point (in terms of gradient norm).\n\n2. More importantly, the authors show that in all other settings considered in the paper, SAM cannot converge to the global/local minimum. Specifically, this includes the smooth non-convex and non-smooth Lipschitz convex cases in the deterministic setting and smooth + (strongly) convex and smooth non-convex cases in the stochastic setting. Some lower bounds are also presented to complement the upper bounds.\n\nIn summary, compared to the case of decaying perturbation size $\\rho$ analyzed in prior theoretical works on SAM, this work shows that using a constant $\\rho$ can inhibit convergence.\n\nSTRENGTHS:\n**1.** Unlike prior theoretical works on SAM, this paper considers the use of constant perturbation size and gradient normalization in the SAM update which is more aligned with practice. \n\n**2.** The theory is comprehensive in the sense that the common function classes are covered and there is analysis for both the deterministic and stochastic settings.\n\n**3.** In some cases, matching lower bounds are also presented to support the upper bounds. Some of the constructions and techniques for the lower bounds seemed novel to me (for e.g., Theorem 3.6) but I'm not an expert on lower bounds. \n\n**4.** The paper is also written more or less clearly.\n\nWEAKNESSES:\n**1.** In the abstract, the authors write: \"*Perhaps surprisingly, in many scenarios, we find out that SAM has limited capability to converge to global minima or stationary points*\". But I don't find the non-convergence of SAM **with constant $\\rho$** (and constant step-size $\\eta$) very surprising; in fact, the asymptotic convergence in Theorems 3.1 and 3.3 surprises me. I say this because of the following reason. Suppose deterministic SAM converges to some point $\\tilde{x}$. Then, we must have $\\nabla f\\Big(\\tilde{x} + \\rho \\frac{\\nabla f(\\tilde{x})}{||\\nabla f(\\tilde{x})||}\\Big) = 0$. This means that we must have $\\tilde{x} + \\rho \\frac{\\nabla f(\\tilde{x})}{||\\nabla f(\\tilde{x})||} = x^{\\ast}$, where $x^{\\ast}$ is some stationary point of $f$. So, $|| \\tilde{x} -  x^{\\ast}|| = \\rho$, i.e., if SAM with constant $\\rho$ converges to some point, then that point must be $\\rho$ distance away from a stationary point of $f$. Notice that if $\\rho$ is a decreasing function of $t$, this (apparent) issue of non-convergence won't arise. \n\n**2.** There seems to be an issue in Theorem 3.2 which states that $\\frac{\\beta}{\\mu} \\geq 2$. But in the proof of Theorem 3.2, one-dimensional quadratics are used to obtain lower bounds; for these one-dimensional quadratics the **tightest possible** smoothness constant = **tightest possible** strong convexity parameter. So actually, $\\beta = \\mu$ in all 3 cases. Now in Case 1, $x_0 = 2\\rho$. But then, we get $x_t \\geq \\frac{x_0}{2} = \\rho$. Unfortunately, this yields a vacuous lower bound of 0. Can the authors comment on this? It seems that if the current proof strategy is to be used, we would need at least 2-d quadratics where $\\beta > \\mu$.\n\n**3.** In Theorem 3.4, the dependence of the convergence bound w.r.t. $T$ seems loose and I don’t think the Lipschitz assumption is required. In the second equation of the proof, there is a $-\\frac{\\eta}{2}||\\nabla f(y_t)||^2$ term and a $\\frac{\\eta^2 \\beta}{2}||\\nabla f(y_t)||^2$ term; these two terms can be combined together and as long as $\\eta \\beta < 1$, the coefficient of $||\\nabla f(y_t)||^2$ is negative. So I think we can get $O(1/T)$ convergence to a $O(\\beta^2 \\rho^2)$ stationary point with a fixed step-size independent of $T$. Also, if we follow this approach, we don’t even Lipschitzness.\n\n**4.** It is unfortunate that the lower bounds for the stochastic case needed to be corrected in the supplementary material.\n\n**5.** (Minor weakness but a limitation nevertheless) SAM was proposed to improve generalization. But this paper has no results illustrating how SAM improves generalization such as by converging to points that are *approximately* flat minima (or because it inhibits convergence preventing over-fitting).\n\nA missed reference (just for the author's information): https://openreview.net/pdf?id=IcDTYTI0Nx shows the convergence of SAM with non-constant perturbation sizes.\n\nQUESTIONS:\nPlease address Weaknesses 1, 2 and 3. My current rating is mainly due to Weakness 1; if the authors can point out some fallacy in my understanding or convince me that the *non-convergence* of SAM isn't all that obvious, I can increase my score.\n\nLIMITATIONS:\nSome minor limitations have been discussed such as the loose dependence on $\\beta$ and $\\mu$ in Theorem 3.1. But important limitations such as Weakness #5 have not been discussed. No foreseeable negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSAM is a very practical algorithm for improving generalization in deep learning, however, even the convergence properties of SAM are not well understood. The paper studies the convergence/non-convergence of SAM under various standard setups in optimization, including smooth/nonsmooth, convex/strongly convex/nonconvex, deterministic/stochastic optimization problems. The results of this paper provide a detailed and complete description of the convergence of SAM for these problems.\n\nSTRENGTHS:\n1. Compared with existing works, the paper analyzes SAM with fixed $\\rho$ and normalization steps, which is the algorithm that people use in practice. Under these practical configurations, this paper provides different characteristics of SAM.\n2. The upper bounds give sufficient conditions for the convergence of SAM and the convergence rate of the algorithm under these conditions. Additionally, this paper also provides lower bounds to show the tightness of their upper bounds in the dependency of certain important problem parameters such as $\\rho$ and $T$.\n3. I found the lower bounds very interesting and they provide many insights about SAM. In contrast to gradient descent (GD), the result in this paper shows that deterministic SAM (full-batch SAM) fails to converge to a  stationary point even for smooth nonconvex objectives. In contrast to stochastic gradient descent (SGD), stochastic SAM (m-SAM) even fails to converge to a global minimum for smooth strongly convex objectives. These results justify the limited capability of SAM in optimization and suggest the differences between SAM and other common optimizers such as SGD/Adam/RMSProp.\n\nWEAKNESSES:\nSee the questions part.\n\nQUESTIONS:\n1. It seems that we only have the non-convergence result of stochastic SAM for smooth and strongly convex functions for m-SAM. Can we prove similar lower bounds for n-SAM? If not, I hope the authors can explain the underlying intuition. By the way, it is proper to say \"stochastic SAM for every function classes we consider, fail to converge properly\" in Line 359?\n2. I can not fully understand the purpose of Sec. 3.4 as well as Thm. 3.6. It seems that for merely convex objectives, gradient descent also fails to converge to a global minima? See the second inequality in Thm. 2.1.7 in Nesterov's celebrated book Lectures \non Convex Optimization\", Second Edition.\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the convergence properties of sharpness-aware minimization in a specific setting with the use of gradient normalization and arbitrary constant perturbation. The paper established convergence rates for both deterministic and stochastic SAM with various assumptions on the convexity of the objective function. The authors argue the term $\\mathcal{O}(\\rho^2)$ appearing multiple times in their bounds is in fact unavoidable by establishing corresponding lower bounds.\n\nSTRENGTHS:\n- The paper is well-written and easy to follow.\n- The analysis covers different levels of convexity.\n- The author proves lower bounds to justify their proposed convergence rate.\n\nWEAKNESSES:\n- The proof of the upper bounds seems to have followed that of [1].\n- Several cases are left out, including the lower bound for the smooth and convex case of deterministic SAM, and the lower bound for the smooth and strongly-convex case of stochastic SAM with small variance.\n- Generalization is not considered.\n\nQUESTIONS:\n- General questions:\n  -  The paper considers SAM with constant stepsizes and gradient normalization and this setting seems to be a common choice in practice. However, I do think more justification for this setting is needed, especially when [1] derived faster convergence rates for SAM without gradient normalization.\n  - It seems no methodological conclusions for the implementation of SAM are provided in the paper. Based on the results in this paper, (how) should we decrease the perturbation $\\rho$ in the algorithm so that the additive term $\\mathcal{O}(\\rho^2)$ would be resolved and thus the convergence be faster?\n\n- Small issues:\n  - There is only one lower bound result in Table 1, which looks weird to me. \n  - In line 830, does \"Assumption 4.2\" refer to Theorem 4.2?\n\n[1] M. Andriushchenko and N. Flammarion. Towards understanding sharpness-aware minimization. In International Conference on Machine Learning, pages 639–668. PMLR, 2022.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper examines convergence properties of sharpness-aware minimization, which is an empirically-popular algorithm for training deep neural networks, yet whose theoretical properties are poorly understood. This paper makes several contributions to the analysis of sharpness-aware minimization, which are novel in this growing literature.\n\nFor smooth and convex functions, this paper proves that the best iterate of SAM converges to the global minimizer. This is expected as, for such functions, there exists a global minimizer; thus, even after perturbations, the convergent point should still be the global minimizer.\n\nFor smooth and nonconvex functions, this paper proves that there is an additional bias term in the convergent point compared with a local minimizer. Furthermore, a lower bound example is constructed, showing that this bias term is necessary.\n\nSTRENGTHS:\n- The paper is nicely written and provides enough background knowledge for readers to understand the topic. For each theorem, the authors accompany their results with a sketch of the high-level proof arguments.\n\n- The lower-bound examples are especially interesting in terms of complementing the convergence rates.\n\nWEAKNESSES:\n- The results in the main text are a bit dense. It would be easier to read if the paper is formatted in a more friendly manner.\n\n- There are no experiments/simulations to complement the theoretical results. I consider this a limitation since SAM is primarily motivated as an empirically-successful algorithm, so a connection between theoretical results to their practical implications would be important for this audience.\n\nQUESTIONS:\n- I wonder if the authors could discuss the practical implications of their findings. For example, how does the extra bias term affect the empirical performance of SAM; Does it inject some kind of regularizer that is actually helpful for the empirical performance, or is it possible to design a better algorithm that eliminates this bias term altogether?\n\n- There is a concurrent paper that studies a very related question:\n\nDai, Y., Ahn, K., & Sra, S. (2023). The Crucial Role of Normalization in Sharpness-Aware Minimization. arXiv preprint arXiv:2305.15287.\n\nIt would be helpful if the authors could discuss the relation to this paper.\n\nLIMITATIONS:\nThe authors discuss the limitations of their results along with presentations in the main text. The potential societal impact of their work should be minimal, given the technical nature of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}