Title: The Crucial Role of Normalization in Sharpness-Aware Minimization
Venue: NeurIPS 2023
Accepted: yes
Keywords: Sharpness-Aware Minimization

Normalization

Deep Learning Theory
Abstract: Sharpness-Aware Minimization (SAM) is a recently proposed gradient-based optimizer (Foret et al., ICLR 2021) that greatly improves the prediction performance of deep neural networks. Consequently, there has been a surge of interest in explaining its empirical success. We focus, in particular, on understanding ***the role played by normalization***, a key component of the SAM updates. We theoretically and empirically study the effect of normalization in SAM for both convex and non-convex functions, revealing two key roles played by normalization: i) it helps in stabilizing the algorithm; and ii) it enables the algorithm to drift along a continuum (manifold) of minima -- a property identified by recent theoretical works that is the key to better performance. We further argue that these two properties of normalization make SAM robust against the choice of hyper-parameters, supporting the practicality of SAM. Our conclusions are backed by various experiments.

The Crucial Role of Normalization in Sharpness-Aware Minimization

Abstract

Sharpness-Aware Minimization (SAM) is a recently proposed gradient-based optimizer (Foret et al., ICLR 2021) that greatly improves the prediction performance of deep neural networks. Consequently, there has been a surge of interest in explaining its empirical success. We focus, in particular, on understanding the role played by normalization , a key component of the SAM updates. We theoretically and empirically study the effect of normalization in SAM for both convex and non-convex functions, revealing two key roles played by normalization: i) it helps in stabilizing the algorithm; and ii) it enables the algorithm to drift along a continuum (manifold) of minima - a property identified by recent theoretical works that is the key to better performance. We further argue that these two properties of normalization make SAM robust against the choice of hyper-parameters, supporting the practicality of SAM. Our conclusions are backed by various experiments.

1 Introduction

Westudy the recently proposed gradient-based optimization algorithm Sharpness-Aware Minimization (SAM) (Foret et al., 2021) that has shown impressive performance in training deep neural networks to generalize well (Foret et al., 2021; Bahri et al., 2022; Mi et al., 2022; Zhong et al., 2022). SAM updates involve an ostensibly small but key modification to Gradient Descent (GD). Specifically, for a loss function L and each iteration t ≥ 0 , instead of updating the parameter w t as w t +1 = w t -η ∇L ( w t ) (where η is called the learning rate ), SAM performs the following update: 1

where ρ is an additional hyper-parameter that we call the perturbation radius . Foret et al. (2021) motivate SAM as an algorithm minimizing the robust loss max ∥ ϵ ∥≤ ρ L ( w + ϵ ) , which is roughly the loss at w (i.e., L ( w ) ) plus the 'sharpness' of the loss landscape around w , hence its name.

The empirical success of SAM has driven a recent surge of interest in characterizing its dynamics and theoretical properties (Bartlett et al., 2022; Wen et al., 2023; Ahn et al., 2023d). However, a major component of SAM remains unexplained in prior work: the role and impact of the normalization factor 1 ∥∇L ( w t ) ∥ used by SAM. In fact, quite a few recent works drop the normalization factor for simplicity when analyzing SAM (Andriushchenko and Flammarion, 2022; Behdin and Mazumder, 2023; Agarwala and Dauphin, 2023; Kim et al., 2023; Compagnoni et al., 2023). Instead of the SAM update (1), these works consider the following update that we call Un-normalized SAM (USAM) :

∗

Suvrit Sra TU Munich / MIT

Apart from experimental justifications in (Andriushchenko and Flammarion, 2022), the effect of this simplification has not yet been carefully investigated, although it is already widely adopted in the community. Thus, is it really the case that such normalization can be omitted 'for simplification' when theoretically analyzing SAM? These observations raise our main question:

What is the role of the normalization factor 1 ∥∇L ( w t ) ∥ in the SAM update (1) ?

1.1 Motivating Experiments and Our Contributions

We present our main findings through two motivating experiments. For the setting, we choose the well-known over-parameterized matrix sensing problem (Li et al., 2018); see Appendix A for details.

Normalization helps with stability. We first pick a learning rate η that allows GD to converge, and we gradually increase ρ from 0 . 001 to 0 . 1 . Considering the early stage of training shown in Figure 1. One finds that SAM has very similar behavior to GD , whereas USAM diverges even with a small ρ - it seems that normalization helps stabilize the algorithm.

0

2

4

6

8

Epoch

0

2

4

6

8

10

Train Loss

=0.1)

=0.01)

=0.005)

=0.001)

0

2

4

6

8

Epoch

0

2

4

6

8

10

USAM (

=0.1)

USAM (

=0.01)

USAM (

=0.005)

USAM (

=0.001)

Normalization permits moving along minima. We reduce the step size by 10 times and consider their performance of reducing test losses in the long run. One may regard Figure 2 as the behavior of SAM, USAM, and GD when close to a 'manifold' of minima (which exists since the problem is over-parametrized) as the training losses are close to zero. The first plot compares SAM and USAMwith the same ρ = 0 . 1 (the largest ρ for which USAM doesn't diverge): notice that USAM and GD both converge to a minimum and do not move further; on the other hand, SAM keeps decreasing the test loss, showing its ability to drift along the manifold. We also vary ρ and compare their behaviors (shown on the right): the ability of SAM to travel along the manifold of minimizers seems to be robust to the size of ρ , while USAM easily gets stuck at a minimum .

0

10000

20000

30000

40000

50000

Epoch

0

1

2

3

4

5

6

7

Test Loss

=0.1)

USAM (

=0.1)

0

10000

20000

30000

40000

50000

Epoch

0

1

2

3

4

5

6

7

=0.1)

=0.25)

=0.5)

=1)

0

10000

20000

30000

40000

50000

Epoch

0

1

2

3

4

5

6

7

USAM (

=0.05)

USAM (

=0.1)

USAM (

=0.01)

Overview of our contributions. In this work, as motivated by Figure 1 and Figure 2, we identify and theoretically explain the two roles of normalization in SAM. The paper is organized as follows.

In Section 2, we study the role of normalization in the algorithm's stability and show that normalization helps stabilize. In particular, we demonstrate that normalization ensures that GD's convergence implies SAM's non-divergence, whereas USAM can start diverging much earlier.

In Section 3, we study the role of normalization near a manifold of minimizers and show that the normalization factor allows iterates to keep drifting along this manifold - giving better performance in many cases. Without normalization, the algorithm easily gets stuck and no longer progresses.

In Section 4, to illustrate our main findings, we adopt the sparse coding example of Ahn et al. (2023a). Their result implies a dilemma in hyper-parameter tuning for GD: a small η gives worse performance, but a large η results in divergence. We show that this dilemma extends to USAM but not SAM. In other words, SAM easily solves the problem where GD and USAM often fail.

These findings also shed new light on why SAM is practical and successful, as we highlight below.

Practical importance of our results. The main findings in this work explain and underscore several practical aspects of SAM that are mainly due to the normalization step. One practical feature of SAM is the way the hyper-parameter ρ is tuned: Foret et al. (2021) suggest that ρ can be tuned independently after tuning the parameters of base optimizers (including learning rate η , momentum β , and so on). In particular, this feature makes SAM a perfect 'add-on' to existing gradient-based optimizers. Our findings precisely support this practical aspect of SAM. Our results suggest that the stability of SAM is less sensitive to the choice of ρ , thanks to the normalization factor.

The same principle applies to the behavior of the algorithm near the minima: Recent theoretical works (Bartlett et al., 2022; Wen et al., 2023; Ahn et al., 2023d) have shown that the drift along the manifold of minimizers is a main feature that enables SAM to reduce the sharpness of the solution (which is believed to boost generalization ability in practice) - our results indicate that the ability of SAM to keep drifting along the manifold is independent of the choice of ρ , again owing to normalization. Hence, our work suggests that the normalization factor plays an important role towards SAM's empirical success.

1.2 Related Work

Sharpness-Aware Optimizers. Inspired by the empirical and theoretical observation that the generalization effect of a deep neural network is correlated with the 'sharpness' of the loss landscape (see (Keskar et al., 2017; Jastrzkebski et al., 2017; Jiang et al., 2020) for empirical observations and (Dziugaite and Roy, 2017; Neyshabur et al., 2017) for theoretical justifications), several recent papers (Foret et al., 2021; Zheng et al., 2021; Wu et al., 2020) propose optimizers that penalize the sharpness for the sake of better generalization. Subsequent efforts were made on making such optimizers scale-invariant (Kwon et al., 2021), more efficient (Liu et al., 2022; Du et al., 2022), and generalize better (Zhuang et al., 2022). This paper focuses on the vanilla version proposed by Foret et al. (2021).

Theoretical Advances on SAM. Despite the success of SAM in practice, theoretical understanding of SAM was absent until two recent works: Bartlett et al. (2022) analyze SAM on locally quadratic losses and identify a component reducing the sharpness λ max ( ∇ 2 L ( w t )) , while Wen et al. (2023) characterize SAM near the manifold Γ of minimizers and show that SAM follows a Riemannian gradient flow reducing λ max ( ∇ 2 L ( w )) when i) initialized near Γ , and ii) η is 'small enough'. Note that while the results of Wen et al. (2023) apply to more general loss functions, our result in Theorem 16 applies to i) any initialization far from the origin, and ii) any η = o (1) and ρ = O (1) . A recent work by Ahn et al. (2023d) formulates the notion of ε -approximate flat minima and analyzed the iteration complexity of practical algorithms like SAM to find such approximate flat minima. A concurrent work by Si and Yun (2023) also analyzes the original version of SAM with the normalization in (1), and makes a case that practical SAM does not converge all the way to optima.

Unnormalized SAM (USAM). USAM was first proposed by Andriushchenko and Flammarion (2022) who observed a similar performance between USAM and SAM when training ResNet over CIFAR-10. This simplification is further accepted by Behdin and Mazumder (2023) who study the regularization effect of USAM over a linear regression model, by Agarwala and Dauphin (2023) who study the initial and final dynamics of USAM over a quadratic regression model, and by Kim et al. (2023) who study the convergence instability of USAM near saddle points. To our knowledge, (Compagnoni et al., 2023) is the only work comparing SAM and USAM dynamics. More preciously, they consider the continuous-time behavior of SGD, SAM, and USAM and find different behaviors of SAM and USAM: USAM attracts local minima while SAM aims at global ones. Still, we remark that as they are considering continuous-time variants of algorithms while we consider discrete (original) versions, our results directly apply to the SAM deployed in practice and the USAM studied in theory.

Edge-of-Stability. In the optimization theory literature, Gradient Descent (GD) was only shown to find minima if the learning rate η is smaller than an 'Edge-of-Stability' threshold, which is related to the sharpness of the nearest minimum. However, people recently observe that when training neural networks, GD with a η much larger than that threshold often finds good minima as well (see (Cohen

et al., 2021) and references therein). Aside from convergence, GD with large η is also shown to find flatter minima (Arora et al., 2022; Ahn et al., 2022; Wang et al., 2022; Damian et al., 2023).

2 Role of Normalization for Stability

In this section, we discuss the role of normalization in the stability of the algorithm. We begin by recalling a well-known fact about the stability of GD: for a convex quadratic cost with the largest eigenvalue of Hessian being β (i.e., β -smooth), GD converges to a minimum iff η < 2 / β . Given this folklore fact, we ask: how do the ascent steps in SAM (1) and USAM (2) affect their stability?

2.1 Strongly Convex and Smooth Losses

Consider an α -strongly-convex and β -smooth loss function L where GD is guaranteed to converge once η < 2 / β . We characterize the stability of SAM and USAM in the following result.

Theorem 1 (Strongly Convex and Smooth Losses) . For any α -strongly-convex and β -smooth loss function L , for any learning rate η < 2 / β and perturbation radius ρ ≥ 0 , the following holds:

SAM. The iterate w t converges to a local neighborhood around the minimizer w ⋆ . Formally,

USAM. In contrast, there exists some α -strongly-convex and β -smooth loss L such that the USAM with η ∈ ( 2 / ( β + ρβ 2 ) , 2 / β ] diverges for all except measure zero initialization w 0 .

As we discussed, it is well-known that GD converges iff η < 2 / β , and Theorem 1 shows that SAM also does not diverge and stays within an O ( √ ηρ ) -neighborhood around the minimum as long as η < 2 / β . However, USAM diverges with an even lower learning rate: η > 2 / ( β + ρβ 2 ) can already make USAM diverge. Intuitively, the larger the value of ρ , the easier it is for USAM to diverge.

One may notice that Equation 3, compared to the standard convergence rate of GD, exhibits an additive bias term of order O ( ηρ 2 ) . This term arises from the unstable nature of SAM: the perturbation in (1) (which always has norm ρ ) prevents SAM from decreasing the loss monotonically. Thus, SAM can only approach a minimum up to a neighborhood. For this reason, in this paper whenever we say SAM 'finds' a minimum, we mean its iterates approach and stay within a neighborhood of that minimum.

Due to space limitations, the full proof is postponed to Appendix C and we only outline it here.

Proof Sketch. For SAM, we show an analog to the descent lemma of GD as follows (see Lemma 9):

By invoking the strong convexity that gives L ( w t ) -L ( w ∗ ) ≤ 1 2 α ∥∇ f ( w t ) ∥ 2 , we obtain

Recursively applying this relation gives the first conclusion. For USAM, we consider the quadratic loss function same as (Bartlett et al., 2022). Formally, suppose that L ( w ) = 1 2 w ⊤ Λ w where Λ = diag( λ 1 , λ 2 , . . . , λ d ) is a PSD matrix such that λ 1 > λ 2 ≥ · · · λ d > 0 . Let the eigenvectors corresponding to λ 1 , λ 2 , . . . , λ d be e 1 , e 2 , . . . , e d , respectively. Then we show the following in Theorem 10: for any η ( λ 1 + ρλ 2 1 ) > 2 and ⟨ w 0 , e 1 ⟩ ̸ = 0 , USAM must diverge. As L ( w ) = 1 2 w ⊤ Λ w is λ 1 -smooth and λ d -strongly-convex, the second conclusion also follows.

Intuitively, the difference in stability can be interpreted as follows: during the early stage of training, w t and ∇L ( w t ) often have large norms. The normalization in SAM then makes the ascent step w t + ρ ∇L ( w t ) ∥∇L ( w t ) ∥ not too far away from w t . Hence, if GD does not diverge for this η , SAM does not either (unless the ρ -perturbation is non-negligible, i.e., ∥ w t ∥ ≫ ρ no longer holds). This is not true for USAM: since the ascent step is un-normalized, it leads to a point far away from w t , making the size of USAM updates much larger. In other words, the removal of normalization leads to much more aggressive steps, resulting in a different behavior than GD and also an easier divergence.

2.2 Generalizing to Non-Convex Cases: Scalar Factorization Problem

Now let us move on to non-convex losses. We consider a scalar version of the matrix factorization problem min U,V 1 2 ∥ UV T -A ∥ 2 2 , whose loss function is defined as L ( x, y ) = 1 2 ( xy ) 2 . Denote the initialization by ( x 0 , y 0 ) , then L ( x, y ) is β ≜ ( x 2 0 + y 2 0 ) -smooth inside the region { ( x, y ) : x 2 + y 2 ≤ β } . Hence, a learning rate η < 2 / β again allows GD to converge due to the well-known descent lemma. The following result compares the behavior of SAM and USAM under this setup.

Theorem 2 (Scalar Factorization Problem; Informal) . For the loss function L ( x, y ) = 1 2 ( xy ) 2 restricted to a β -smooth region, if we set η = 1 / β < 2 / β (so GD finds a minimum), then

SAM. SAM never diverges and approaches a minimum within an O ( ρ ) -neighborhood (in fact, SAM with distinct ρ 's always find the same minimum (0 , 0) ).

USAM. On the other hand, USAM diverges once ρ ≥ 15 η - which could be much smaller than 1.

Thus, our observation in Theorem 1 is not limited to convex losses - for our non-convex scalarfactorization problem, the stability of SAM remains robust to the choice of ρ , while USAM is provably unstable. One may refer to Appendix D for the formal statement and proof of Theorem 2.

2.3 Experiment: Early-Stage Behaviors when Training Neural Networks

0

5

10

15

20

25

30

10 0

4×10 1

6×10 1

2×10 0

3×10 0

CNN with tanh activation

SAM

=0.05

SAM

=0.1

SAM

=0.2

USAM

=0.05

USAM

=0.1

USAM

=0.2

0

5

10

15

20

25

30

10 0

4×10 1

6×10 1

2×10 0

3×10 0

FCN with ReLU activation

SAM

=0.1

SAM

=0.3

SAM

=0.5

USAM

=0.1

USAM

=0.3

USAM

=0.5

As advertised, our result holds not only for convex or toy loss functions but also for practical neural networks. In Figure 3, we plot the early-stage behavior of GD, SAM, and USAM with different ρ values (while fixing η ). We pick two neural networks: a convolutional neural network (CNN) with tanh activation and a fully-connected network (FCN) with ReLU activation. We train them over the CIFAR10 dataset and report the early-stage training losses. Similar to Figure 1, Theorem 1 and Theorem 2, the stability of SAM is not sensitive to the choice of ρ , while USAM diverges easily .

3 Role of Normalization for Drifting Near Minima

Now, we explain the second role of normalization: enabling the algorithm to drift near minima. To convince why this is beneficial, we adopt a loss function recently considered by Ahn et al. (2023a) when understanding the behavior of GD with large learning rates. Their result suggests that GD needs a 'large enough' η for enhanced performance, but this threshold can never be known a-priori in practice. To verify our observations from Figure 2, we study the dynamics of SAM and USAM over the same loss function and find that: i) no careful tuning is needed for SAM ; instead, SAM with any configuration finds the same minimum (which is the 'best' one according to Ahn et al. (2023a)); and ii) such property is only enjoyed by SAM - for USAM, careful tuning remains essential.

3.1 Toy Model: Single-Neuron Linear Network Model

To theoretically study the role of normalization near minima, we consider the simple two-dimensional non-convex loss L ( x, y ) defined over all ( x, y ) ∈ R 2 as

This L was recently studied in (Ahn et al., 2023a) to understand the behavior of GD with large η 's. By direct calculation, the gradient and Hessian of L at a given ( x, y ) can be written as:

Without loss of generality, one may assume ℓ is minimized at 0 (see Appendix E for more details regarding ℓ ). Then, L achieves minimum at the entire x - and y -axes, making it a good toy model for studying the behavior of algorithms near a continuum of minima. Finally, note that the parametrization x × y can be interpreted as a single-neuron linear network model - hence its name.

Before moving on to SAM and USAM we first briefly introduce the behavior of GD on such loss functions characterized in (Ahn et al., 2023a). Since ℓ is even, without loss of generality, we always assume that the initialization w 0 = ( x 0 , y 0 ) satisfies y 0 ≥ x 0 > 0 .

Theorem 3 (Theorems 5 and 6 of (Ahn et al., 2023a); Informal) . For any η = γ/ ( y 2 0 -x 2 0 ) , the GD trajectory over the loss function L ( x, y ) = ℓ ( xy ) has two possible limiting points:

If γ < 2 , then the iterates converge to (0 , y ∞ ) where y 2 ∞ ∈ [ γ / η -O ( γ ) -O ( η / γ ) , γ / η + O ( η / γ )] .

If γ > 2 , then the iterates converge to (0 , y ∞ ) where y 2 ∞ ∈ [ 2 / η -O ( η ) , 2 / η ] .

Intuitively, the limiting point of GD (denoted by (0 , y ∞ ) ) satisfies y 2 ∞ ≈ min { y 2 0 -x 2 0 , 2 / η } . For simplicity, we denote η GD ≈ 2 / y 2 0 -x 2 0 as the threshold of η that distinguishes these two cases.

Interpretation of Ahn et al. (2023a). Fixing the initialization ( x 0 , y 0 ) , it turns out this model has a nice connection to the sparse coding problem, wherein it's desirable to get a smaller y 2 ∞ (which we will briefly discuss in Section 4). According to Theorem 3, to get a smaller y 2 ∞ , one must increase the learning rate η beyond η GD. Hence we mainly focus on the case where η > η GD - in which case we abbreviate y 2 ∞ ≈ 2 / η (see Table 1). However, GD diverges once η is too large - in their language, γ cannot be much larger than 2 . This dilemma of tuning η , as we shall illustrate in Section 4 in more detail, makes GD a brittle choice for obtaining a better y 2 ∞ .

On the other hand, from the numerical illustrations in Figure 4, one can see that SAM keeps moving along the manifold of minimizers (i.e., the y -axis) until the origin. This phenomenon is characterized in Theorem 4; in short, any moderate choice of η and ρ suffices to drive SAM toward the origin - no difficult tuning needed anymore!

In contrast, USAM does not keep moving along the axis. Instead, a lower bound on y 2 ∞ also presents - although smaller than the GD version. As we will justify in Theorem 5, USAM does get trapped at some non-zero y 2 ∞ Thus, a dilemma similar to that of GD shows up: for enhanced performance, an aggressive ( η, ρ ) is needed; however, as we saw from Section 2, this easily results in a divergence.

0.5

0.0

0.5

1.0

1.5

2.0

0

1

2

3

4

5

6

SAM

USAM

Assumptions. To directly compare with (Ahn et al., 2023a), we focus on the cases where y 2 0 -x 2 0 = γ/η and γ ∈ [ 1 2 , 2] is a constant of moderate size; hence, η is not too different from the η GD defined in Theorem 3. In contrast to most prior works which assume a tiny ρ (e.g., (Wen et al., 2023)), we allow ρ to be as large as a constant (i.e., we only require ρ = O (1) in Theorem 4 and Theorem 5).

3.1.1 SAM Keeps Drifting Toward the Origin

We characterize the trajectory of SAM when applied to the loss defined in Equation 5 as follows:

Theorem 4 (SAM over Single-Neuron Networks; Informal) . For any η ∈ [ 1 2 η GD , 2 η GD ] and ρ = O (1) , the SAM trajectory over the loss function L ( x, y ) = ℓ ( xy ) can be divided into three phases:

Initial Phase. x t drops so rapidly that | x t | = O ( √ η ) in O ( 1 / η ) steps. Meanwhile, y t remains large: specifically, y t = Ω( √ 1 / η ) . Thus, SAM approaches the y -axis (the set of global minima).

Middle Phase. x t oscillates closely to the axis such that | x t | = O ( √ η ) always holds. Meanwhile, y t decreases fast until y t ≤ | x t | 2 - that is, | x t | remains small and SAM approaches the origin.

Final Phase. w t = ( x t , y t ) gets close to the origin such that | x t | , | y t | = O ( √ η + ηρ ) . We then show that w t remains in this region for the subsequent iterates.

Informally, SAM first approaches the minimizers on y -axis (which form a manifold) and then keeps moving until a specific minimum. Moreover, SAM always approaches this minimum for almost all ( η, ρ ) 's. This matches our motivating experiment in Figure 2: No matter what hyper-parameters are chosen, SAM always drift along the set of minima, in contrast to the behavior of GD. This property allows SAM always to approach the origin (0 , 0) and remains in its neighborhood, while GD converges to (0 , √ 2 / η ) (see Table 1). The formal version of Theorem 4 is in Appendix F.

3.1.2 USAM Gets Trapped at Different Minima

We move on to characterize the dynamics of USAM near the minima. Like GD or SAM, the first few iterations of USAM drive iterates to the y -axis. However, unlike SAM, USAM does not keep drifting along the y -axis and stops at some threshold - in the result below, we prove a lower bound on y 2 t that depends on both η and ρ . In other words, the lack of normalization factor leads to diminishing drift.

Theorem 5 (USAM over Single-Neuron Networks; Informal) . For any η ∈ [ 1 2 η GD , 2 η GD ] and ρ = O (1) , the USAM trajectory over the loss function L ( x, y ) = ℓ ( xy ) have the following properties:

Initial Phase. Similar to Initial Phase of Theorem 4, | x t | = O ( √ η ) and y t = Ω( √ 1 / η ) hold for the first O ( 1 / η ) steps. That is, USAM also approaches y -axis, the set of global minima.

Final Phase. However, for USAM, once the following condition holds for some round t : 3

| x t | decays exponentially fast, which in turn ensures y 2 ∞ ≜ lim inf t →∞ y 2 t ≳ (1 -η 2 -ρ 2 ) y 2 t .

Remark. Note that USAM becomes GD as we send ρ → 0+ , and our characterized threshold y 2 USAM indeed recovers that of GD (i.e., 2 / η from Theorem 3) because lim ρ → 0+ ( √ 8 ρ / η +1 -1) / 2 ρ = 2 / η .

Compared with SAM, the main difference occurs when close to minima, i.e., | x t | = O ( √ η ) . Consistent with our motivating experiment (Figure 2), the removal of normalization leads to diminishing drift along the minima. Thus, USAM is more like an improved version of GD rather than a simplification of SAM, and the comparison between Theorem 3 and Theorem 5 reveals that USAM only improves over GD if ρ is large enough - in which case USAM is prone to diverges as we discussed in Section 2.

See Appendix G for a formal version of Theorem 5 together with its proof.

3.1.3 Technical Distinctions Between GD, SAM, and USAM

Before moving on, we present a more technical comparison between the results stated in Theorem 3 versus Theorem 4 and Theorem 5. We start with an intuitive explanation of why GD and USAM get stuck near the manifold of minima but SAM does not: when the iterates approach the set of minima, both w t and ∇L ( w t ) become small. Hence the normalization plays an important role: as ∇L ( w t ) are small, w t and w t + ρ ∇L ( w t ) become nearly identical, which leads to a diminishing updates of GD and USAM near the minima. On the other hand, having the normalization term, the SAM update doesn't diminish, which prevents SAM from converging to a minimum and keeps drifting along the manifold.

This high-level idea is supported by the following calculation: recall Equation 6 that ∇L ( x t , y t ) = ℓ ′ ( x t y t ) · [ y t x t ] ⊤ . Hence, when | x t |≪ y t in Final Phase, the 'ascent gradient' direction ∇L ( x t , y t ) (i.e., the ascent steps in Equation 1 and Equation 2) is almost perpendicular to the y -axis. We thus rewrite the update direction (i.e., the difference between w t +1 and w t ) for each algorithm as follows.

For SAM, after normalization, ∇L ( w t ) ∥∇L ( w t ) ∥ is roughly a unit vector along the x -axis. Hence, the update direction is the gradient at w t + 1 / 2 ≈ [ ρ y t ] ⊤ . Once y t is large (making w t + 1 / 2 far from minima), ∇L ( w t + 1 / 2 ) thus have a large component along y t , which leads to drifting near minima.

For GD, by approximating ℓ ′ ( u ) ≈ u , we derive ∇L ( x t , y t ) ≈ [ x t y 2 t x 2 t y t ] ⊤ . When 2 / η > y 2 t , the magnitude of x t is updated as | x t +1 | ≈ | x t -ηx t y 2 t | = | (1 -ηy 2 t ) x t | , which allows an exponential decay. Thus, GD converges to a minimum and stop moving soon after 2 / η > y 2 t .

For USAM, the descent gradient is taken at w t + ρ ∇L ( w t ) ≈ [(1+ ρy 2 t ) x t (1+ ρx 2 t ) y t ] ⊤ . Thus, ∇L ( w t + ρ ∇L ( w t )) ≈ [(1+ ρy 2 t )(1+ ρx 2 t ) 2 x t y 2 t (1+ ρy 2 t ) 2 (1+ ρx 2 t ) x 2 t y t ] ⊤ by writing ℓ ′ ( u ) ≈ u . This makes USAM deviate away from SAM and behave like GD: by the similar argument as GD, USAM stops at a minimum soon after 2 / η > (1 + ρy 2 t )(1 + ρx 2 t ) 2 y 2 t ≈ (1 + ρy 2 t ) y 2 t !

Hence, the normalization factor in the ascent gradient helps maintain a non-diminishing component along the minima, leading SAM to keep drifting. This distinguishes SAM from GD and USAM.

3.2 USAM Gets Trapped Once Close to Minima

In this section, we extend our arguments to nonconvex costs satisfying Polyak-Lojasiewicz (PL) functions (see, e.g., (Karimi et al., 2016)). Recall that f satisfies the µ -PL condition if 1 2 ∥∇L ( w ) ∥ 2 ≥ µ ( L ( w ) -min w L ( w )) for all w . Building upon the analysis of Andriushchenko and Flammarion (2022), we show the following result when applying USAM to β -smooth and µ -PL losses.

Theorem 6 (USAM over PL Losses; Informal) . For β -smooth and µ -PL loss L , for any η < 1 / β and ρ < 1 / β , and for any initialization w 0 , ∥ w t -w 0 ∥ ≤ poly ( η, ρ, β, µ ) · √ L ( w 0 ) -min w L ( w ) , ∀ t .

This theorem has the following consequence: Suppose that USAM encounters a point w 0 that is close to some minimum (i.e., L ( w 0 ) ≈ min w L ( w ) ) during training. Then Theorem 6 implies that the total distance traveled by USAM from w 0 is bounded - in other words, the distance USAM moves along the manifold of minimizers can only be of order O ( √ L ( w 0 ) -min w L ( w )) .

As a remark, we compare Theorem 6 with the recent result by Wen et al. (2023): their result essentially implies that, for small enough η and ρ , SAM iterates initialized close to a manifold of the minimizers approximately track some continuous dynamics (more precisely, a Riemannian gradient flow induced by a 'sharpness' measure they find) and keep drifting along the manifold. This property is indeed in sharp contrast with USAM whose total travel distance is bounded.

The formal statement and proof of Theorem 6 are contained in Appendix H.

3.3 Experiments for Practical Neural Networking Training

0

20

40

60

80

100

Epoch

42

44

46

48

50

52

54

56

58

Test Accuracy

0

20

40

60

80

100

Epoch

30

40

50

60

70

80

90

100

Train Accuracy

SGD

=0.01)

=0.05)

=0.1)

USAM (

=0.01)

USAM (

=0.05)

USAM (

=0.1)

Weclose this section by verifying our claims in practical neural network training. We train a ResNet18 on the CIFAR-10 dataset, initialized from a poor global minimum generated as per Liu et al. (2020)

(we used the 'adversarial checkpoint' released by Damian et al. (2021)). This initialization has 100% training accuracy but only 48% test accuracy - which lets us observe a more pronounced algorithmic behavior near the minima via tracking the test accuracy. From Figure 5, we observe:

GD gets stuck at this adversarial minimum, in the sense that the test accuracy stays ats 48% .

SAM keeps drifting while staying close to the manifold of minimizers (because the training accuracy remains 100% ), which results in better solutions (i.e., the test accuracy keeps increasing).

USAM with small ρ gets stuck like GD, while USAM with larger ρ 's deviate from this manifold.

Hence, USAM faces the dilemma that we describe in Subsection 3.1: a conservative hyper-parameter configuration does not lead to much drift along the minima, while a more aggressive choice easily leads to divergence. However, the stability of SAM is quite robust to the choice of hyper-parameter and they all seem to lead to consistent drift along the minima.

Remark. Apart from the 'adversarial checkpoint' which is unrealistic but can help highlight different algorithms' behavior when they are close to a bad minimum, we also conduct the same experiments but instead initialized from a 'full-batch checkpoint' (Damian et al., 2021), which is the 100% training accuracy point reached by running full-batch GD on the training loss function. The result is plotted as Figure 8 in Subsection B.1. One can observe that USAM still gets stuck at the 'full-batch checkpoint', while SAM keeps increasing its test accuracy via drifting along the minima manifold.

4 Case Study: Learning Threshold Neurons for Sparse Coding Problem

To incorporate our two findings into a single example, we consider training one-hidden-layer ReLU networks for the sparse coding problem, a setup considered in (Ahn et al., 2023a) to study the role of η in GD. Without going into details, the crux of their experiment is to understand how GD with large η finds desired structures of the network - in this specific case, the desired structure is the negative bias in ReLU unit (also widely known as 'thresholding unit/neuron'). In this section, we evaluate SAM and USAM under the same setup, illustrating the importance of normalization.

Main observation of Ahn et al. (2023a). Given this background, the main observation of Ahn et al. (2023a) is that i) when training the ReLU network with GD, different learning rates induce very different trajectories; and ii) the desired structure, namely a negative bias in ReLU, only arises with large 'unstable' learning rates for which GD exhibits unstable behaviors. We reproduce their results in Figure 6, plotting the test accuracy on the left and the bias of ReLU unit on the right. As they claimed, GD with larger η learns more negative bias, which leads to better test accuracy.

0

200

400

600

800

1000

Time (

X Step)

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Test Accuracy

0

200

400

600

800

1000

Time (

X Step)

0.8

0.6

0.4

0.2

0.0

0.2

=0.1

=0.05

=0.025

Their inspiring observation is however a bit discouraging for practitioners. According to their theoretical results, such learning rates have to be quite large - large to the point where GD shows very unstable behavior (à la Edge-of-Stability (Cohen et al., 2021)). In practice, without knowledge of the problem, this requires a careful hyper-parameter search to figure out the correct learning rate. More importantly, such large and unstable learning rates may cause GD to diverge or lead to worse performance. More discussions can be found in the recent paper by Kaur et al. (2022).

In contrast, as we will justify shortly, SAM does not suffer from such a 'dilemma of tuning' -matching with our results in Theorem 4. Moreover, the removal of normalization no longer attains such a property , as we demonstrated in Theorem 5. In particular, for USAM, one also needs to carefully tune η and ρ for better performance - as we inspired in Theorem 5 and Theorem 6, small ( η, ρ ) makes the iterates get stuck early; on the other hand, as we presented in Section 2, an aggressive choice causes USAM to diverge. The following experiments illustrate these claims in more detail.

In Figure 7, we plot the performance of SAM, USAM, and GD with different η 's (while fixing ρ ) -gray lines for GD, solid lines for SAM, and dashed lines for USAM. From the plot, USAM behaves more similarly to GD than SAM: the bias does not decrease sufficiently when the learning rate is not large enough, which consequently to leads poor test accuracy. On the other hand, no matter what

0

200

400

600

800

1000

Time (

X Step)

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Test Accuracy

0

200

400

600

800

1000

Time (

X Step)

1.0

0.8

0.6

0.4

0.2

0.0

0.2

SAM

=0.1

SAM

=0.05

SAM

=0.025

USAM  =0.1

USAM  =0.05

USAM  =0.025

η is chosen for SAM, bias is negative enough and ensures better generalization. Hence, Figure 7 illustrates that compared to SAM, USAM is less robust to the tuning of η .

In Figure 9 (deferred to Subsection B.2), we also compare these three algorithms when varying ρ and fixing η . In addition to what we observe in Figure 7, we show that normalization also helps stability - USAM quickly diverges as we increase ρ , while SAM remains robust to the choice of ρ . Thus, USAM is also less robust to the tuning of ρ . In other words, our observation in Figure 7 extends to ρ .

Hence, putting Figure 7 and Figure 9 together, we conclude that SAM is robust to different configurations of ( η, ρ ) while USAM is robust to neither of them . Hence, the normalization of SAM eases hyper-parameter tuning, which is typically a tough problem for GD and many other algorithms - normalization boosts the success of SAM in practice.

5 Conclusion

In this paper, we investigate the role played by normalization in SAM. By theoretically characterizing the behavior of SAM and USAM on both convex and non-convex losses and empirically verifying our conclusions via real-world neural networks, we found that normalization i) helps stabilize the algorithm iterates, and ii) enables the algorithm to keep moving along the manifold of minimizers, leading to better performance in many cases. Moreover, as we demonstrate via various experiments, these two properties make SAM require less hyper-parameter tuning, supporting its practicality.

In this work, we follow a recent research paradigm of 'physics-style' approaches to understanding deep neural networks based on simplified models (c.f. (Zhang et al., 2022; Garg et al., 2022; von Oswald et al., 2023; Abernethy et al., 2023; Allen-Zhu and Li, 2023; Liu et al., 2023; Li et al., 2023; Ahn et al., 2023a,b,c)). We found such physics-style approaches quite helpful, especially for complex modern neural networks. We hope that our work builds stepping stones for future works on understanding working mechanisms of modern deep neural networks.

Acknowledgments

Kwangjun Ahn was supported by the ONR grant (N00014-20-1-2394) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of the Secretary of Defense. Suvrit Sra acknowledges support from an NSF CAREER grant (1846088), and NSF CCF-2112665 (TILOS AI Research Institute). Kwangjun Ahn also acknowledges support from the Kwanjeong Educational Foundation. We thank Kaiyue Wen, Zhiyuan Li, and Hadi Daneshmand for their insightful discussions.

References

Jacob Abernethy, Alekh Agarwal, Teodor V Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305.17040 , 2023.

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

Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 247-257. PMLR, 2022.

Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang. Learning threshold neurons via the 'edge of stability'. NeurIPS 2023 (arXiv:2212.07469) , 2023a.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023 (arXiv:2306.00297) , 2023b.

Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). arXiv 2310.01082 , 2023c.

Kwangjun Ahn, Ali Jadbabaie, and Suvrit Sra. How to escape sharp minima. arXiv preprint arXiv:2305.15659 , 2023d.

Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673 , 2023.

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

Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In International Conference on Machine Learning , pages 948-1024. PMLR, 2022.

Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 7360-7371, 2022.

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

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

Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference on learning theory , pages 483-513. PMLR, 2020.

Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations , 2021.

Enea Monzio Compagnoni, Antonio Orvieto, Luca Biggio, Hans Kersting, Frank Norbert Proske, and Aurelien Lucchi. An sde for modeling sam: Theory and insights. arXiv preprint arXiv:2301.08203 , 2023.

Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global minimizers. In Advances in Neural Information Processing Systems , 2021.

Alex Damian, Eshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. In International Conference on Learning Representations , 2023.

Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan. Efficient sharpness-aware minimization for improved training of neural networks. In International Conference on Learning Representations , 2022.

Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence , 2017.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations , 2021.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Stanislaw Jastrzkebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint arXiv:1711.04623 , 2017.

Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In 8th International Conference on Learning Representations , 2020.

Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition. In Machine Learning and Knowledge Discovery in Databases: European Conference , pages 795-811. Springer, 2016.

Simran Kaur, Jeremy Cohen, and Zachary C Lipton. On the maximum hessian eigenvalue and generalization. arXiv preprint arXiv:2206.10654 , 2022.

Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang, Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations , 2017.

Hoki Kim, Jinseong Park, Yujin Choi, and Jaewook Lee. Stability analysis of sharpness-aware minimization. arXiv preprint arXiv:2301.06308 , 2023.

Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon 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.

Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory , pages 2-47. PMLR, 2018.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. International Conference on Machine Learning (ICML) (arXiv:2303.04245) , 2023.

Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. ICLR (arXiv:2210.10749) , 2023.

Shengchao Liu, Dimitris Papailiopoulos, and Dimitris Achlioptas. Bad global minima exist and sgd can reach them. Advances in Neural Information Processing Systems , 33:8543-8552, 2020.

Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang 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.

Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao. Make sharpness-aware minimization stronger: A sparsified perturbation approach. In Advances in Neural Information Processing Systems , 2022.

Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems , 30, 2017.

Dongkuk Si and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. NeurIPS 2023 (arXiv:2306.09850) , 2023.

Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems , 35:99839994, 2022.

Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How does sharpness-aware minimization minimize sharpness? In International Conference on Learning Representations , 2023.

Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in Neural Information Processing Systems , 33:2958-2969, 2020.

Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301 , 2022.

Yaowei Zheng, Richong Zhang, and Yongyi Mao. Regularizing neural networks via adversarial model perturbation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 8156-8165, 2021.

Qihuang Zhong, Liang Ding, Li Shen, Peng Mi, Juhua Liu, Bo Du, and Dacheng Tao. Improving sharpness-aware minimization with fisher mask for better generalization on language models. In Findings of the Association for Computational Linguistics: EMNLP 2022 , pages 4064-4085, 2022.

Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, James s Duncan, Ting Liu, et al. Surrogate gap minimization improves sharpness-aware training. In International Conference on Learning Representations , 2022.

Appendix

A Setup of the Motivating Experiment

In the motivating experiments (Figure 1 and Figure 2), we follow the over-parameterized matrix sensing setup as Li et al. (2018) and Blanc et al. (2020). Specifically, we do the following:

Generate the true matrix by sampling each entry of U ⋆ ∈ R d × r independently from a standard Gaussian distribution and let X ⋆ = U ⋆ ( U ⋆ ) ⊤ .

Normalize each column of U ⋆ to unit norm so that the spectral norm of U ⋆ is close to one.

For every sensing matrix A i ( i = 1 , 2 , . . . , m ), sample the entries of A i independently from a standard Gaussian distribution. Then observe b i = ⟨ A i , X ⋆ ⟩ .

In particular, for the experiments, we chose r = 5 , d = 100 , and m = 5 dr .

B Additional Experimental Results

B.1 Running SAM and USAM from Other Initializations

0

20

40

60

80

100

120

140

Epoch

0.750

0.755

0.760

0.765

0.770

0.775

Test Accuracy

0

20

40

60

80

100

120

140

Epoch

0.0

0.2

0.4

0.6

0.8

1.0

Train Accuracy

SGD

SAM

=0.1

SAM

=0.5

SAM  =1

USAM

=0.1

USAM

=0.5

USAM  =1

B.2 Varying ρ While Fixing η in Sparse Coding Example

0

200

400

600

800

1000

Time (

X Step)

0.50

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Test Accuracy

0

200

400

600

800

1000

Time (

X Step)

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

SAM

=0.5

SAM

=0.3

SAM

=0.1

SAM

=0.01

USAM  =0.5

USAM  =0.3

USAM  =0.1

USAM  =0.01

Aside from Figure 7 which varies η while fixing ρ , we perform the same experiment when varying ρ and fixing η . The main observation for SAM is similar to that of Figure 7: different hyper-parameters all keep decreasing the bias and give better test accuracy - even with the tiniest choice ρ = 0 . 01 .

However, for USAM, there are three different types of ρ 's as shown in Figure 9:

For tiny ρ = 0 . 01 , the bias doesn't decrease much. Consequently, the performance of USAM nearly degenerates to that of GD - while SAM with ρ = 0 . 01 still gives outstanding performance.

For moderate ρ = 0 . 1 and ρ = 0 . 3 , USAM manages to decrease the bias and improve its accuracy, though with a slower speed than SAM with the same choices of ρ .

For large ρ = 0 . 5 (where SAM still works well; see the solid curve in blue), USAM diverges.

Thus, the dilemma described in Section 4 indeed applies to not only η but also ρ - matching our main conclusion that normalization helps make hyper-parameter tuning much more manageable.

C Omitted Proof for Smooth and Strongly Convex Losses

We shall first restate Theorem 1 here for the ease of presentation.

Theorem 7 (Restatement of Theorem 1) . For any α -strongly-convex and β -smooth loss function L , for any learning rate η ≤ 2 / β and perturbation radius ρ ≥ 0 , the following holds:

SAM. The iterate w t converges to a local neighborhood around the minimizer w ⋆ . Formally,

USAM. In contrast, there exists some α -strongly-convex and β -smooth loss L such that the USAM with η ∈ ( 2 / ( β + ρβ 2 ) , 2 / β ] diverges for all except measure zero initialization w 0 .

We will show these two conclusions separately. The proof directly follows from Theorem 8 and Theorem 10.

C.1 SAM Allows a Descent Lemma Like GD

Theorem 8 (SAM over Strongly Convex and Smooth Losses) . For any α -strongly-convex and β -smooth loss function L , for any learning rate η ≤ 2 / β and perturbation radius ρ ≥ 0 , the following holds for SAM:

Proof. We first claim the following analog of descent lemma, which we state as Lemma 9:

By definition of strong convexity, we have

where the last inequality uses the fact that ⟨ a, b ⟩ -1 2 ∥ b ∥ 2 ≤ 1 2 ∥ a ∥ 2 . Thus, combining the two inequalities above, we obtain

which after rearrangement becomes

Unrolling this recursion, we obtain

This completes the proof.

Lemma 9 (SAM Descent lemma) . For a convex loss L that is β -smooth, SAM iterates w t satisfy the following when the learning rate η < 2 β and ρ ≥ 0 :

Proof. Let v t ≜ ∇L ( w t ) / ∥∇L ( w t ) ∥ and w t + 1 / 2 ≜ w t + ρv t so we have w t +1 = w t -η ∇L ( w t + 1 / 2 ) . Then the β -smoothness of L yields

We start with upper bounding the norm of ∇L ( w t + 1 / 2 ) :

Hence, as long as η < 2 β , we have the following upper bound on L ( w t +1 ) :

Now we lower bound 〈 ∇L ( w t + 1 / 2 ) , ∇L ( w t ) 〉 . Note that

where the last inequality uses the following standard fact about convex functions: for any w 1 , w 2 , ⟨∇L ( w 1 ) -∇L ( w 2 ) , w 1 -w 2 ⟩ ≥ 0 . Hence, we arrive at

and thus finishing the proof.

C.2 USAM Diverges on Quadratic Losses

Theorem 10 (USAMover Quadratic Losses) . Following (Bartlett et al., 2022), consider the quadratic loss L induced by a PSD matrix Λ . Without loss of generality, let L minimize at the origin. Formally,

where λ max = λ 1 > · · · ≥ λ d > 0 and v max = e 1 , e 2 , . . . , e d are the eigenvectors corresponding to λ 1 , λ 2 , . . . , λ d , respectively. Then the iterates of USAM Equation 2 applied to Equation 8 satisfy:

If η ( λ max + ρλ 2 max ) < 2 , then the iterates converges to the global minima exponentially fast.

If η ( λ max + ρλ 2 max ) > 2 and if ⟨ w 0 , v max ⟩ ̸ = 0 , then the iterates diverge.

Moreover, such a loss function L is λ 1 -smooth and λ d -strongly-convex.

Proof. As ∇ ℓ ( w ) = Λ w , the USAM update Equation 2 reads

Hence, if η ( λ max + ρλ 2 max ) < 2 , then since ∥ ∥ I -η Λ -ηρ Λ 2 ∥ ∥ < 1 , we have that ∥ w t ∥ → 0 exponentially fast. On the other hand, if η ( λ max + ρλ 2 max ) > 2 , then | 1 -ηλ 1 -ηλ 2 1 | > 1 . Since ∥ w t ∥ ≥ | 1 -ηλ 1 -ηλ 2 1 | t |⟨ w 0 , e 1 ⟩| , it follows that the iterates diverge as long as ⟨ w 0 , v max ⟩ ̸ = 0 . Finally, as ∇ 2 L ( w ) = Λ , we directly know that L is λ 1 -smooth and λ d -strongly-convex.

D Omitted Proof for Scalar Factorization Problems

Theorem 11 (Formal Version of Theorem 2) . Consider the scalar factorization loss L ( x, y ) = 1 2 ( xy ) 2 . If ( x 0 , y 0 ) satisfies 2 x 0 > y 0 > x 0 ≫ 0 and η = ( x 2 0 + y 2 0 ) -1 , then i) SAM finds a neighborhood of the origin with radius O ( ρ ) for all ρ , and ii) USAM diverges as long as ρ ≥ 15 η .

This theorem is a combination of Theorem 12 and Theorem 15 that we will show shortly.

D.1 SAM Always Converges on Scalar Factorization Problems

Let w t = [ x t y t ] ⊤ be the t -th iterate. Denote ∇ t = ∇L ( w t ) = [ x t y 2 t x 2 t y t ] ⊤ . For each step t , the actual update of SAM is therefore the gradient taken at

By denoting ˜ ∇ t = ∇L ( ˜ w t ) , the update rule of SAM is

We make the following claim:

Theorem 12 (SAM over Scalar Factorization Problems) . Under the setting of Theorem 11, there exists some threshold T for SAM such that | x t | , | y t | ≤ 5 ρ for all t ≥ T .

Proof. We first observe that SAM over L ( x, y ) always pushes the iterate towards the minima (which are all the points on the x - and y - axes). Formally:

If x t ≥ 0 , then x t +1 ≤ x t . If x t ≤ 0 , then x t +1 ≥ x t .

If y t ≥ 0 , then y t +1 ≤ y t . If y t ≤ 0 , then y t +1 ≥ y t .

This observation can be verified by noticing sign( ρy t z -1 t ) = sign( y t ) sign( x t y t ) = sign( x t ) and similarly sign( ρx t z -1 t ) = sign( y t ) . In other words, sign( x t + ρy t z -1 t ) = sign( x t ) and thus always pushes x t towards the y -axis. The same also holds for y t by symmetry.

In analog to the descent lemma for GD, we can show the following lemma:

Lemma 13. When picking η = ( x 2 0 + y 2 0 ) -1 , we have L ( w t +1 ) -L ( w t ) < 0 as long as x t , y t ≥ 5 ρ .

Proof. As ∇ 2 L ( x, y ) = [ y 2 2 xy 2 xy x 2 ] , we know L is β ≜ ( x 2 0 + y 2 0 ) -smooth inside the region { ( x, y ) : x 2 + y 2 ≤ β } . Then we have (recall that η = ( x 2 0 + y 2 0 ) -1 = 1 / β )

To make sure that it's negative, we simply want

which can be ensured once

If we have x t , y t ≥ 5 ρ , then as z -1 t ≤ min { x -1 t , y -1 t } , we have

which shows the first inequality. The second one follows from symmetry.

Therefore, SAM will progress until x t ≤ 5 ρ or y t ≤ 5 ρ . Without loss of generality, assume that x t ≤ 5 ρ ; we then claim that y t will soon decrease to O ( ρ ) .

Proof. First, show that | x t | remains bounded by 5 ρ . Assume x t ≥ 0 without loss of generality:

where the second last line uses y 2 t ≤ y 2 0 = y 2 0 x 2 0 + y 2 0 η -1 ≤ 4 5 η -1 and the last one uses η = ( x 2 0 + y 2 0 ) -1 ≪ 1 . Meanwhile, we see that y t decreases exponentially fast by observing the following:

where the last line uses z -1 t = ( x 2 t + y 2 t ) -1 / 2 ≥ (2 y 2 t ) -1 / 2 = 2 -1 / 2 y -1 t as y t ≥ 5 ρ ≥ | x t | .

So eventually we have | x t | , | y t | ≤ 5 ρ . Recall that Equation 9 infers | x t +1 | ≤ 5 ρ from | x t | ≤ 5 ρ . Hence, by symmetry, we conclude that | x t +1 | , | y t +1 | ≤ 5 ρ hold as well. Therefore, SAM always finds an O ( ρ ) -neighborhood of the origin, i.e., it is guaranteed to converge regardless of ρ .

D.2 USAM Diverges with Small ρ

For USAM, the dynamics can be written as

We make the following claim which is similar to Theorem 10:

Theorem 15 (USAM over Scalar Factorization Problems) . Under the setup of Theorem 11, for any ρ ≥ 15 η , | x t +1 | ≥ 2 | x t | and | y t +1 | ≥ 2 | y t | for all t ≥ 1 ; in other words, USAM diverges exponentially fast.

Proof. Prove by induction. From Equation 10, our conclusion follows once

According to our setup that η = ( x 2 0 + y 2 0 ) -1 , y 0 ≤ 2 x 0 , and the induction statement that, we have η ≥ (5 x 2 t ) -1 . The second inequality then holds as long as

which is true as x 2 t ≥ x 2 0 , y 4 t ≥ y 4 0 , y 0 ≥ x 0 , and ρ ≥ 15 η ≥ 3 x -2 0 . Note that the bounds on ρ are very loose, and we made no effort to optimize it; instead, we only aimed to show that USAM starts to diverge from a ρ = Θ( η ) ≪ 1 .

E Assumptions in the Single-Neuron Linear Network Model

Assumptions on ℓ . Following Ahn et al. (2023a), we make the following assumptions about ℓ :

(A1) ℓ is a convex, even, 1 -Lipschitz function that is minimized at 0 .

(A2) ℓ is twice continuously differentiable near the origin with ℓ ′′ (0) = 1 , which infers the existence of a constant c > 0 such that | ℓ ′ ( s ) | ≤ | s | for all | s | ≤ c .

(A3) We further assume a 'linear tail' away from the minima, i.e., | ℓ ′ ( s ) | ≥ c/ 2 for all | s | ≥ c and | ℓ ′ ( s ) | ≥ | s | / 2 for | s | ≤ c .

Some concrete example of loss functions satisfying the above assumption include a symmetrized logistic loss 1 2 log(1 + exp( -2 s )) + 1 2 log(1 + exp(2 s )) and a square root loss √ 1 + s 2 One may refer to their paper for more details.

F Omitted Proof of SAM Over Single-Neuron Linear Networks

Theorem 16 (Formal Version of Theorem 4) . For a loss L over ( x, y ) ∈ R 2 defined as L ( x, y ) = ℓ ( xy ) where ℓ satisfies Assumptions (A1), (A2), and (A3), if the initialization ( x 0 , y 0 ) satisfies:

where γ ∈ [ 1 2 , 2] and C ≥ 1 is constant, then for all hyper-parameter configurations ( η, ρ ) such that 4

we can decompose the trajectory of SAM (defined in Equation 1) into three phases, whose main conclusions are stated separately in three theorems and are informally summarized here:

(Theorem 20) Until x t = O ( √ γη ) , we must have y t = Ω( √ γ/η ) , and x t +1 ≤ x t -Ω( √ γη ) .

(Theorem 32) After Middle Phase, we always have | x t | , | y t | = O ( ηρ + √ η ) .

(Theorem 26) After Initial Phase and until y t ≤ | x t | , | x t | = O ( ηρ + √ η ) still holds. Meanwhile, y t +1 ≤ y t -min { Ω( ηρ 2 ) y t , Ω( ηρ ) } (i.e., y t either drops by Ω( ηρ ) or decays by Ω( ηρ 2 ) ).

F.1 Basic Properties and Notations

Recall the update rule of SAM in Equation 1: w t +1 ← w t -η ∇L ( w t + ρ ∇L ( w t ) ∥∇L ( w t ) ∥ ) . By writing w t as [ x t y t ] , substituting L ( x, y ) = ℓ ( xy ) , and utilizing the expressions of ∇L ( x, y ) in Equation 6, we have:

where the second step uses the fact that ℓ is a even function so ℓ ′ ( t ) has the same sign with t . Define z t = √ x 2 t + y 2 t sign( x t y t ) . Then w t + ρ ∇L ( w t ) ∥∇L ( w t ) ∥ = [ x t + ρy t z -1 t y t + ρx t z -1 t ] . Further denoting ℓ ′ ( ( x t + ρy t z -1 t )( y t + ρx t z -1 t ) ) by ℓ ′ t , one can simplify the update rule Equation 1 as follows:

̸

Assumption. Following footnote 1, we assume that x t , y t = 0 for all t .

We are ready to give some basic properties of Equation 12. First, we claim that the sign of ℓ ′ t is the same as the product x t y t . Formally, we have the following lemma:

̸

̸

Lemma 17 (Sign of Gradient in SAM) . If x t = 0 , y t = 0 , then sign( x t ) = sign( x t + ρy t z -1 t ) and sign( y t ) = sign( y t + ρx t z -1 t ) . In particular, if x t = 0 , y t = 0 , we must have sign( ℓ ′ t ) = sign( x t y t ) .

̸

̸

Proof. Note that sign( z t ) = sign( x t y t ) , so sign( y t z -1 t ) = sign( x t ) . Similarly, sign( x t z -1 t ) = sign( y t ) . Thus, we have sign( x t ) = sign( x t + ρy t z -1 t ) and sign( y t ) = sign( y t + ρx t z -1 t ) , which in turn shows that sign( ℓ ′ t ) = sign ( sign( x t + ρy t z -1 t ) sign( y t + ρx t z -1 t ) ) = sign( x t y t ) .

Using Lemma 17, the SAM update can be equivalently written as

We then claim the following property, which can be intuitively interpolated as SAM always goes towards the minimizes (i.e., both axes) for each step, although sometimes it may overshoot.

̸

Lemma 18. Suppose that x t , y t = 0 . Then, the following basic properties hold:

If x t > 0 , then x t +1 < x t . If x t < 0 , then x t +1 > x t .

If y t > 0 , then y t +1 < y t . If y t < 0 , then y t +1 > y t .

Proof. We only show the case that x t > 0 as the other cases are similar. Recall the dynamics of x t :

Since x t = 0 , Lemma 17 implies that sign( ℓ ′ t ) = sign( x t y t ) , and so

̸

where the last line uses the assumption that x t > 0 .

F.2 Initial Phase: x t Decreases Fast while y t Remains Large

As advertised in Theorem 16, we group all rounds until x t ≤ c 4 √ γη as Initial Phase. Formally:

Definition 19 (Initial Phase) . Let t 1 be the largest time such that x t > c 4 √ γη for all t ≤ t 1 . We call the iterations [0 , t 1 ] Initial Phase .

Theorem 20 (Main Conclusion of Initial Phase; SAM Case) . For t 0 = Θ(min { η -1 . 5 ρ -1 γ 1 / 2 , η -2 } ) , we have y t ≥ 1 2 √ γ/η for all t ≤ min { t 0 , t 1 } under the conditions of Theorem 16. Moreover, we have x t +1 ≤ x t -c 4 √ γη for all t ≤ min { t 0 , t 1 } , which consequently infers t 1 ≤ t 0 under the conditions of Theorem 16, i.e., min { t 0 , t 1 } is just t 1 . This shows the first claim of Theorem 16.

Remark. This theorem can be intuitively explained as follows: At initialization, x 0 , y 0 are both Θ( √ 1 /η ) . However, after t 1 iterations, we have | x t 1 +1 | = O ( √ η ) ; meanwhile, y t is still of order Ω( √ 1 /η ) (much larger than O ( √ η ) ). Hence, | x t | and | y t | get widely separated in Initial Phase.

Proof. This theorem is a direct consequence of Lemma 21 and Lemma 24.

Lemma 21. There exists t 0 = Θ(min { η -1 . 5 ρ -1 γ 1 / 2 , η -2 } ) such that | y t | ≥ 1 2 √ γ/η for t ≤ t 0 . If η is sufficiently small s.t. ηρ + √ Cγη ≤ 1 2 √ γ/η , then y t ≥ 1 2 √ γ/η for t ≤ min { t 1 , t 0 } .

Proof. Let t 0 be defined as follows:

Then, it follows that t 0 = Θ(min { η -1 . 5 ρ -1 γ 1 / 2 , η -2 } ) . We first prove the following claim. Claim 22. For all t ≤ t 0 , y 2 t ≥ 1 4 γ/η .

Proof. We prove this claim by induction. We trivially have y 2 0 = Cγ/η ≥ 1 4 Cγ/η . Now suppose that the conclusion holds up to some t < t 0 , i.e., y 2 t ′ ≥ 1 4 γ/η for all t ′ ≤ t . We consider y 2 t +1 -x 2 t +1 :

where (a) used ℓ ′ t ≤ 1 (thanks to Assumption (A1)) and (b) used the induction hypothesis y 2 t ≥ 1 4 Cγ/η . This further implies that

Thus, by the definition of t 0 , we must have y ≥ y -x ≥ γ/η , which proves the claim.

Next, we prove the second conclusion.

Claim 23. If ηρ + √ Cγη ≤ 1 2 √ γ/η , then y t ≥ 1 2 √ γ/η for t ≤ min { t 0 , t 1 } .

Proof. Still show by induction. Let y t ≥ 1 2 √ γ/η > 0 for some t < min { t 0 , t 1 } . Consider y t +1 . By Definition 19, x t is positive for all t ≤ t 1 . Thus, using Lemma 18, we have x t ≤ x 0 ≤ y 0 = √ Cγ/η . Since x t , y t > 0 , we have sign( ℓ ′ t ) = sign( x t y t ) = sign( x t ) > 0 and hence (13) gives

where we used the facts that | ℓ ′ t | ≤ 1 and ( x 2 t + y 2 t ) -1 / 2 ≤ min {| x t | -1 , | y t | -1 } . Hence, we get

since ηρ + √ Cγη ≤ 1 2 √ γ/η ≤ y t . By Claim 22, y t +1 ≥ 0 implies y t +1 ≥ 1 2 √ γ/η as well.

Combining Claim 22 and Claim 23 finishes the proof of Lemma 21.

Lemma 24. For any t ≤ min { t 0 , t 1 } , we have

In particular, if η is sufficiently small s.t. 4 √ C c η -1 ≤ t 0 , then we must have t 1 ≤ 4 c η -1 -1 < t 0 .

2 t +1 2 t +1 2 t +1 1 4

Proof. Since x t > 0 for all t ≤ t 1 , we have sign( ℓ ′ t ) = sign( x t y t ) = sign( y t ) from Lemma 17 and Lemma 21, so the SAM update (13) becomes

Since x t > 2 c √ η Cγ , we have

which implies that ℓ ′ t ≥ c/ 2 from Assumption (A3). Together with Lemma 21, we have

Let t ′ 1 ≜ √ Cγ/η c 4 √ γη = 4 √ C c η -1 . Since x 0 ≤ √ Cγ/η , we have x t ′ 1 < c 4 √ γη as long as t 1 ≤ t 0 . Thus, it follows that t 1 ≤ t ′ 1 -1 < t 0 , as desired.

F.3 Middle Phase: y Keeps Decreasing Until Smaller Than | x t |

Then we move on to the second claim of Theorem 16. We define all rounds until y t < | x t | that are after Initial Phase as Middle Phase. Formally, we have the following definition.

Definition 25 (Middle Phase) . Let t 2 be the first time that y t < | x t | . We call ( t 1 , t 2 ] Middle Phase .

Before presenting the main conclusion of Middle Phase, we first define a threshold B that we measure whether a variable is 'small enough', which we will use throughout this section. Formally,

Theorem 26 (Main Conclusion of Middle Phase; SAM Case) . Under the conditions of Theorem 16, | y t | ≤ √ Cγ/η and | x t | ≤ B throughout Middle Phase. Under the same conditions, we further have y t +1 ≤ y t -min { 1 2 ηρ 2 y t , c 2 √ 2 ηρ } for all t 1 ≤ t ≤ t 2 - showing the second claim of Theorem 16.

Remark. This theorem can be understood as follows. Upon entering Middle Phase, | x t | is bounded by O ( ηρ + √ η ) . This then gets preserved throughout Middle Phase. Meanwhile, in t 2 iterations, y t drops rapidly such that y t 2 +1 ≤ | x t 2 +1 | = O ( ηρ + √ η ) . In other words, x t and y t are both 'close enough' to 0 after Middle Phase; thus, SAM finds the flattest minimum (which is the origin).

Proof. The claims are shown in Lemma 27, Lemma 28 and Lemma 30, respectively.

Lemma 27. If ηρ + √ Cγη < √ Cγ/η , then during the Middle Phase, | y t | ≤ y 0 ≤ √ Cγ/η .

Proof. When sign( y t +1 ) = sign( y t ) , then Lemma 18 implies that | y t +1 | ≤ | y t | . Now consider the case sign( y t +1 ) = -sign( y t ) . For simplicity, say y t > 0 and y t +1 < 0 . Since y t > 0 , we have sign( ℓ ′ t ) = sign( x t y t ) = sign( x t ) and hence Equation 13 gives

which shows that | y | ≤ ηρ + Cγη . This proves the statement.

t +1 √

Lemma 28. Suppose that | x t | ≤ B = max { 2 c √ η Cγ , ηρ + √ Cηγ } . Then we have | x t +1 | ≤ B .

Proof. By symmetry, we only consider the case where x t > 0 . If x t +1 ≥ 0 , then we have 0 ≤ x t +1 ≤ x t due to Lemma 24. If x t +1 < 0 , then since x t > 0 , it follows that

where the last inequality is because | y t | ≤ y 0 ≤ √ Cγ/η . This concludes the proof.

Corollary 29. Note that, by definition of Initial Phase in Definition 19, we already have | x t 1 +1 | ≤ B . Hence, this lemma essentially says | x t | ≤ B = max { 2 c √ η Cγ , ηρ + √ Cηγ } for all t 1 < t ≤ t 2 .

Lemma 30. If y t ≥ | x t | , then we must have

Proof. Since y t ≥ | x t | , we have ( x 2 t + y 2 t ) -1 / 2 ≥ (2 y 2 t ) -1 / 2 = ( √ 2 y t ) -1 . Consider two cases:

if | ( x t + ρy t z -1 t )( y t + ρx t z -1 t ) | ≥ c , then | ℓ ′ t | ∈ [ c 2 , 1] according to Assumption (A3). Moreover, since sign( ℓ ′ t ) = sign( x t y t ) = sign( x t ) , it follows that

ii) otherwise, it follows from (A3) that

As sign( x t ) = sign( ρy t z -1 t ) and sign( y t ) = sign( ρx t z -1 t ) , it follows that

we must have

where we used z -1 t ≥ ( √ 2 y t ) -1 and thus y t z -1 t ≥ 1 / √ 2 .

Combining item 16 and item 17, we obtain the desired conclusion.

F.4 Final Phase: Both x t and y t Oscillates Around the Origin

It only remains to show that the iteration never escapes the origin.

Definition 31 (Final Phase) . We denote by ' Final Phase ' all iterations after t 2 .

Theorem 32 (Main Conclusion of Final Phase; SAM Case) . For all t > t 2 , we have | x t | , | y t | ≤ B where B = max { 2 c √ η Cγ , ηρ + √ Cηγ } .

Remark. As we will see shortly, when entering Final Phase, both | x t | and | y t | are bounded by O ( ηρ + √ η ) . This theorem essentially says that they can never exceed this bound in Final Phase. In other words, in Final Phase, both x t and y t are oscillating around 0 .

Proof. We first show the following lemma, ruling out the possibility of | y t | > B after Middle Phase:

Proof. The proof will be very similar to Lemma 28. Recall that by Definition 25, we must have y t 2 > 0 . If y t 2 +1 > 0 as well, then | y t 2 +1 | = y t 2 +1 ≤ | x t 2 +1 | ≤ B . Otherwise,

where we used | x t 2 | ≤ B . We proved our claim as ηρ + ηB ≤ B .

According to Lemma 28 Lemma 33, we have

Hence, it only remains to do a reduction, stated as follows.

Lemma 34. If ηρ + ηB ≤ B , then | x t +1 | , | y t +1 | ≤ B as long as | x t | , | y t | ≤ B .

Proof. Without loss of generality, let x t > 0 for some t > t 2 and show that | x t +1 | ≤ B - which is identical to the proof of Lemma 27. By symmetry, the same also holds for | x t | .

Our conclusion thus follows from an induction based on Equation 18 and Lemma 34.

G Omitted Proof of USAM Over Single-Neuron Linear Networks

Theorem 35 (Formal Version of Theorem 5) . For a loss L over ( x, y ) ∈ R 2 defined as L ( x, y ) = ℓ ( xy ) where ℓ satisfies Assumptions (A1), (A2), and (A3), if the initialization ( x 0 , y 0 ) satisfies:

where γ ∈ [ 1 2 , 2] and C ≥ 1 is constant, then for all hyper-parameter configurations ( η, ρ ) such that 6

we can characterize the initial and final phases of the trajectory of USAM (defined in Equation 2) by the following two theorems, whose main conclusions are informally summarized below:

(Theorem 39) Until x t = O ( √ γη ) , we must have y t = Ω( √ γ/η ) , and x t +1 ≤ x t -Ω( √ γη ) .

(Theorem 44) Once (1 + ρy 2 t ) y 2 t ≲ 2 / η , | x t | decays exponentially and thus USAM gets stuck.

Different from Theorem 16, there is no characterization of Middle Phase here, which means the iterates can also stop above the threshold ˜ y 2 USAM defined in Equation 7 (the technical reason is sketched in subsubsection 3.1.3, i.e., SAM gradients non non-negligible when y t is large, while USAM gradients vanish once | x t | is small). However, we remark that the main takeaway of Theorem 35 is to contrast SAM (which always attains y 2 ∞ ≈ 0 ) with USAM (which must stop once (1+ ρy 2 t ) y 2 t ≲ 2 / η ).

G.1 Basic Properties and Notations

Recall the update rule of USAM in Equation 2: w t +1 ← w t -η ∇L ( w t + ρ ∇L ( w t )) . Still writing w t as [ x t y t ] ⊤ , we have

̸

Due to the removal of normalization, ℓ ′ are taken twice at different points in Equation 19. For simplicity, we denote ˜ ℓ ′ t = ℓ ′ ( x t y t ) and ℓ ′ t = ℓ ′ (( x t + ρ ˜ ℓ ′ t y t )( y t + ρ ˜ ℓ ′ t x t )) . The update rule can be rewritten as:

Similar to the SAM case, we shall approximate ˜ ℓ ′ t and ℓ ′ t according to the magnitude of x t y t and ( x t + ρ ˜ ℓ ′ t y t )( y t + ρ ˜ ℓ ′ t x t ) . We first have the following lemma similar to Lemma 17:

̸

̸

Lemma 36 (Sign of Gradient in USAM) . If x t = 0 , y t = 0 , then sign( x t ) = sign( x t + ρ ˜ ℓ ′ t y t ) and sign( y t ) = sign( y t + ρ ˜ ℓ ′ t x t ) . In particular, if x t = 0 , y t = 0 , then sign( ℓ ′ t ) = sign( ˜ ℓ ′ t ) = sign( x t y t ) .

̸

̸

Proof. First of all, sign( ˜ ℓ ′ t ) = sign( x t y t ) according to Assumption (A1). Therefore,

giving our first two claims. The last conclusion follows by definition.

Moreover, we also have the following lemma analog to Lemma 18:

Lemma 37. Suppose that x t , y t ≥ 0 . Then, the following basic properties hold:

If x t > 0 , then x t +1 < x t . If x t < 0 , then x t +1 > x t .

If y t > 0 , then y t +1 < y t . If y t < 0 , then y t +1 > y t .

Proof. We only prove the statement for x t > 0 as the proof is similar for other cases. Recall the dynamics of x t :

Since x t = 0 , Lemma 36 implies that sign( ℓ ′ t ) = sign( ˜ ℓ ′ t ) = sign( x t y t ) , and so

where the last line uses the assumption that x t > 0 .

G.2 Initial Phase: x t Decreases Fast while y t Remains Large

The definition of Initial Phase is very similar to the one of SAM - and the conclusion is also analogue, although the proofs are slightly different because of the different update rule.

Definition 38 (Initial Phase) . Let t 1 be the largest time such that x t > 1 2 √ γη for all t ≤ t 1 . We denote by ' Initial Phase ' the iterations [0 , t 1 ] .

Theorem 39 (Main Conclusion of Initial Phase; USAM Case) . For t 0 = Θ(min { η -1 ρ -1 , η -2 } ) , we have y t ≥ 1 2 √ γ/η for all t ≤ min { t 0 , t 1 } under the conditions of Theorem 35. Moreover, we have x t +1 ≤ x t -1 2 √ Cγη for all t ≤ min { t 0 , t 1 } , which consequently infers t 1 ≤ t 0 under the conditions of Theorem 35, i.e., min { t 0 , t 1 } is just t 1 . This shows the first claim of Theorem 35.

Proof. This theorem is a combination of Lemma 40 and Lemma 43.

Similar to Lemma 21, y t in USAM also cannot be too small in the first few iterations:

Lemma 40. There exists t 0 = Θ(min { η -1 ρ -1 , η -2 } ) such that for t ≤ t 0 , we have | y t | ≥ 1 2 √ γ/η . Assuming η ≤ 1 2 , then y t ≥ 1 2 √ γ/η for t ≤ min { t 1 , t 0 } .

Proof. The proof idea follows from Lemma 21. Let t 0 be defined as follows:

Then, it follows that t 0 = Θ(min { η -1 ρ -1 , η -2 } ) . We first show that y 2 t cannot be too small until t 0 :

Claim 41. For all t ≤ t 0 , y 2 t ≥ 1 4 γ/η .

Proof. Prove by induction. Initially, y 2 0 = Cγ/η ≥ 1 4 γ/η . Then consider some t < t 0 such that y 2 t ′ ≥ 1 4 γ/η for all t ′ ≤ t . By Equation 20, we have

Using Assumption (A1), we know that | ℓ ′ t | ≤ 1 and | ˜ ℓ ′ t | ≤ 1 , giving

By definition of t 0 and condition that t < t 0 , we have y 2 t +1 ≥ y 2 t +1 -x 2 t +1 ≥ 1 4 γ/η .

Claim 42. Suppose that √ Cγη ≤ 1 2 √ γ/η , i.e., η ≤ 1 2 √ C . Then y t ≥ 1 2 √ γ/η for all t ≤ t 0 .

Proof. We still consider the maximum single-step difference in y t . Suppose that for some t ′ < min { t 0 , t 1 } , y t ′ ≥ 1 2 √ γ/η for all t ′ ≤ t . According to Equation 20 and Assumption (A1),

where the last inequality used Lemma 37 to conclude that x t ≤ x 0 ≤ y 0 = √ Cγ/η . Hence, as the first term is non-negative and √ Cγη ≤ 1 2 √ γ/η , we must have y t +1 ≥ 0 , which implies y t +1 ≥ 1 2 √ γ/η according to Claim 41.

Putting these two claims together gives our conclusion. Note that the condition η ≤ 1 2 √ C in Claim 42 can be inferred from the assumptions η ≤ 1 2 and C ≥ 1 made in Theorem 35.

After showing that y t never becomes too small, we are ready to show that x t decreases fast enough. Lemma 43. For t ≤ min { t 0 , t 1 } , we have

In particular, if η is sufficiently small s.t. 2 √ Cη -1 ≤ t 0 , then we must have t 1 ≤ 2 √ Cη -1 -1 < t 0 .

Proof. Since x t > 0 for all t ≤ t 1 , we have sign( ℓ ′ t ) = sign( x t y t ) = sign( y t ) from Lemma 36 Lemma 40, so the USAM update (20) simplifies as follows according to Assumption (A1) and Lemma 40

Let t ′ 1 = √ Cγ/η 1 2 √ γη = 2 √ Cη -1 , then t ′ 1 -1 < t 0 and thus x t ′ 1 < 1 2 √ γη , i.e., t 1 < t 0 must holds.

G.3 Final Phase: y t Gets Trapped Above the Origin

Now, we are going to consider the final-stage behavior of USAM.

Theorem 44 (Main Conclusion of Final Phase; USAM Case) . After Initial Phase, we always have | x t | ≤ √ Cγη and | y t | ≤ √ Cγ/η . Moreover, once we have η (1 + ρy 2 t ) y 2 t = 2 -ϵ for some t ≥ t 1 and ϵ > 0 , we must have | x t +1 | ≤ exp( -Ω( ϵ )) | x t | for all t ≥ t . This consequently infers that y 2 ∞ , which is defined as lim inf t →∞ y 2 t , satisfies y 2 ∞ ≥ (1 -4 Cγ ( η + ρCγ ) 2 ϵ -1 ) y 2 t . As ϵ is a constant independent of η and ρ and can be arbitrarily close to 0 , this shows the second claim of Theorem 35.

Proof. The first conclusion is an analog of Lemma 28 and Lemma 27 (which allows a simpler analysis thanks to the removal of normalization), as we will show in Lemma 45. The second part requires a similar (but much more sophisticated) analysis to Lemma 15 of Ahn et al. (2023a), which we will cover in Lemma 47.

Lemma 45. Suppose that | x t | ≤ √ Cγη , | y t | ≤ √ Cγ/η . Assuming ηρ ≤ 1 , then | x t +1 | ≤ √ Cγη and | y t +1 | ≤ √ Cγ/η as well. Furthermore, | x t | ≤ √ Cγη and | y t | ≤ √ Cγ/η hold for all t > t 1 .

Proof. Suppose that x t ≥ 0 without loss of generality. The case where x t +1 ≥ 0 is trivial by Lemma 37. Otherwise, using Lemma 36 and Equation 20, we can write

By Assumption (A1), | ℓ ′ t | and | ˜ ℓ ′ t | are bounded by 1 . By the condition that ηρ ≤ 1 ,

where we used | y t | ≤ sqrtCγ/η . Similarly, for y t +1 , only considering the case where y t ≥ 0 and y t +1 ≤ 0 suffices. We have the following by symmetry

The second part of the conclusion is done by induction. According to the definition of Initial Phase, we have | x t 1 +1 | ≤ 1 2 √ γη . As C ≥ 1 ≥ 1 4 , we consequently have | x t | ≤ √ Cγη for all t > t 1 from the first part of the conclusion. Regarding | y t | , recall Lemma 40 infers y t > 0 for all t ≤ t 1 +1 and Lemma 37 infers the monotonicity of y t , we have y t 1 +1 ≤ y 0 = √ Cγ/η . Hence, | y t | ≤ √ Cγ/η for all t > t 1 as well.

Before showing the ultimate conclusion Lemma 47, we first show the following single-step lemma:

Lemma 46. Suppose that η (1 + ρy 2 t ) y 2 t < 2 and | x t | ≤ √ Cγη for some t . Define ϵ t = 2 -η (1 + ρy 2 t ) y 2 t (then we must have ϵ t ∈ (0 , 2) ). Then we have

Proof. Without loss of generality, assume that x t > 0 . From Equation 20, we can write:

where we used sign( ℓ ′ t ) = sign( ˜ ℓ ′ t ) = sign( x t y t ) = sign( y t ) (Lemma 36). By Assumption (A2),

where we used | x t | ≤ √ Cγη . For the other direction, we have the following by Assumption (A3):

where we used (1 + ρx 2 t ) ≥ 1 and (1 + ρ x 2 t 2 ) ≥ 1 . Equation 22 follows from Equation 23 and Equation 24.

Lemma 47. Let t be such that i) η (1 + ρy 2 t ) y 2 t = 2 -ϵ where ϵ ∈ (0 , 2 9 ) is a constant, and ii) | x t | ≤ √ Cγη . Then we have the following conclusion on lim inf t →∞ y 2 t , denoted by y 2 ∞ in short:

While the ϵ -1 looks enormous, it is a constant independent of η and ρ ; in other words, we are allowed to set ϵ as close to zero as we want. As we only consider the dependency on η and ρ , we can abbreviate this conclusion as y 2 ∞ ≥ (1 -O ( η 2 + ρ 2 )) y 2 t , as we claim in the main text.

Proof. In analog to Equation 23, we derive the following for y t :

where the second inequality uses Lemma 45. Let d t = y 2 t -y 2 t , then we have

where we used the assumption that ρCγη ≤ 1 and the fact that y 2 t is monotonic (thus y 2 t ≤ y 2 t ). According to Equation 22, we have | x t +1 | ≤ | x t | exp( -Ω( ϵ )) for all t ≥ t 2 . Hence, we have

where we used d t = 0 (by definition) and the sum of geometric series. Plugging back x 2 t ≤ Cγη ,

as claimed.

H Omitted Proof of USAM Over General PL functions

Theorem 48 (Formal Version of Theorem 6) . For any µ -PL and β -smooth loss function L , for any learning rate η < 1 / β and ρ < 1 / β , for any initialization w 0 , the following holds for USAM:

where L ∗ is the short-hand notation for min w L ( w ) .

We first state the following useful result by Andriushchenko and Flammarion (2022, Theorem 10).

Lemma 49 (Descent Lemma of USAM over Smooth and PL Losses) . For any β -smooth and µ -PL loss function L , for any learning rate η < 1 / β and ρ < 1 / β , the following holds for USAM:

Proof of Theorem 48. We follow the convention in Karimi et al. (2016): Let X ⋆ be the set of global minima and x p be the projection of x onto the solution set X ⋆ . From β -smoothness, it follows that

Now since L is β -smooth and µ -PL, Theorem 2 from (Karimi et al., 2016) implies that the quadratic growth condition holds, i.e.,

Moreover, from β -smoothness, we have

Thus, by the update rule of USAM (2), it follows that

Now we only to invoke the USAM descent lemma stated before, i.e., Lemma 49, giving

Putting the last two inequalities together then give our conclusion.

Thus, it follows that

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work investigates the role of the normalization factor $1/||\\nabla L(x,w)||$ in the perturbation step of the Sharpness-Aware Minimization (SAM) algorithm. To this end, the authors compare the behavior of gradient descent, SAM and a version of SAM without the normalization (USAM) via experiments and extensive theoretical analysis. The authors find that USAM lacks some of SAMs beneficial properties: On the one hand, it is harder to choose an appropriate perturbation radius for USAM and on the other hand, the ‘drifting along minima’, which is thought to be connected to SAMs empirical success, only occurs in a restricted manner.\n\nSTRENGTHS:\nThe paper is very well written and polished. The motivation (understanding the effects of the normalization factor in the perturbation step of SAM) is outlined clearly. The theoretical analysis is thorough. The toy experiments are clever and well-designed to isolate the effects that are to be investigated, and mostly support and illustrate the conclusions arising from the theoretical analysis.\n\nWEAKNESSES:\nThe authors essentially provide extensive theoretical analysis and toy experiments on why USAM works worse than SAM, which often stems from the fact that the gradient magnitude changes during training. USAM is however not a relevant algorithm in practice - it was only introduced as a simplified version of SAM for easier theoretical analysis. The insights on why SAM works in the first place are thus limited (e.g. that SAM can drift along minima was analyzed in [1], and adding that normalization is important for this to happen adds little to the understanding of SAM).\n\nApart from [2], where USAM was found to work similarly to SAM for a ResNet on CIFAR, it is further unclear how well USAM works in realistic scenarios. While the authors provide some insights (in 2.3 and 3.3), the settings in those experiments are fairly restrictive (one focuses on early-stage training and the other on an adversarial initialisation) and (nicely) highlight isolated effects when dropping the normalization factor. While I think that negative results (like, _In contrast to SAM, USAM does not work in practice, and here we explain why_) can be of value to the community, it is thus unclear to me if the results presented in the paper are indeed relevant for practical neural network training with SAM (or USAM). Presenting results of several models trained with SGD, SAM and USAM on e.g. CIFAR10 & CIFAR100 over a range of $\\rho$ values would bring insights on this and bridge the gap between the analysis and potential practical implications of this paper. In particular, one could e.g. show that for practical deep learning settings picking $\\rho$ is indeed more difficult, or that USAM does not find well-generalizing minima (due to the restricted drifting along minima).\n\n[1] Peter L. Bartlett, Philip M. Long and Olivier Bousque. The Dynamics of Sharpness-Aware Minimization: Bouncing Across Ravines and Drifting Towards Wide Minima. Arxiv preprint\n\n[2] Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware Minimization. ICML 2022\n\nQUESTIONS:\nAs outlined above, my main question is whether the results and claims regarding USAM can indeed be observed in realistic settings for neural network training. This could e.g. be addressed by training a range of models on CIFAR100 for a range of $\\rho$ values and plotting the performance of each algorithm in dependence of $\\rho$ (like e.g. in Figure 4 of [1], just with one such plot per model).\n\nMinor questions:\n- Figure 7: How was $\\rho$ chosen? Is it possible that the optimal $\\rho$ is different for SAM and USAM, i.e. that for a ‘good’ choice of $\\rho$ USAM would also behave like SAM for the selected range of $\\eta$?  \n- Section 2.3: Could avoiding the SAM-perturbation at the early stages of training mitigate the problem of divergence for USAM?\n- Figure 8 in Appendix B: I cannot see the purple line for which USAM diverges ($\\rho=0.5$). Is it not shown in the plot?\n- Appendix D.2: ‘USAM diverges with small $\\rho$’ - shouldn’t this be ‘large $\\rho$’?\n\n[1] Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-aware\nminimization for scale-invariant learning of deep neural networks. ICML 2021\n\nLIMITATIONS:\nI don’t see potential negative societal impact that the authors should have addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

comments: SUMMARY:
The paper investigates the role of the normalization term in the recently proposed optimization algorithm Sharpness-Aware Minimization (SAM). The authors theoretically, and empirically, study the differences among  SAM, its un-normalized counterpart USAM, and gradient descent, under multiple settings (strongly convex and smooth losses, a non-convex case, single-neuron linear networks, sparse coding), and conclude with two key takeaways: normalization in SAM helps (1) stabilize algorithm iterates, and (2) enables the algorithm to drift near minima.

STRENGTHS:
- The paper thoroughly investigates the role of normalization across many different settings, and convincingly show the benefit of the normalization term. 
- The authors provide clear key takeaways for why normalization helps, with both empirical and theoretical support.
- The paper makes a solid contribution to our understanding of why normalization is beneficial
- The authors compared the impact of tuning the learning rate and perturbation radius between SAM and USAM

WEAKNESSES:
- Empirical verification on real-world tasks is limited. For example, generating more results similar to Figure 5 would strengthen the papers conclusions. 
- While having a deeper understanding of a specific term in an optimizer adds value, the overall question of the work does not feel like a pressing issue in ML. The takeaway of the rigorous work is to use the original version of SAM, which most people who use SAM already do.

QUESTIONS:
- Can the takeaways from this work be generalized to the benefits of other optimizers that use normalization (such as normalized gradient descent)?
- Are there any advantages for USAM over SAM?

LIMITATIONS:
No, the authors do not discuss limitations and potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- The paper investigates the difference between SAM and Unclipped SAM (USAM) in terms of stabilizing and drifting along the manifold.\n- The paper reveals a significant disadvantage in USAM compared to SAM, which was previously thought to be similar in various papers.\n- The paper provides mathematical and experimental evidence to support their ideas in diverse settings.\n\nSTRENGTHS:\n- Overall, the idea of investigating the difference between SAM and USAM is clear and interesting.\n- The paper identifies the role of normalization and highlights that USAM behaves more like GD than SAM, potentially getting stuck in local minima and requiring careful tuning.\n- The mathematical proofs are clear and offer unique insights compared to previous papers, such as [1].\n\n[1] An SDE for Modeling SAM: Theory and Insights (ICML’23)\n\nWEAKNESSES:\nPlease refer to questions.\n\nQUESTIONS:\n- I believe the selection of the radius $\\rho$ in USAM and SAM should be different. While I understand that the authors aimed to test various values of $\\rho$, SAM and USAM may require different optimal radius settings to achieve similar results as mentioned in [2]. The chosen radius sets seem to be selected naively without obtaining the best results in each setting and needs to be improved.\n- [2] proposed that SAM and USAM perform similarly in real training problems. However, your experimental results rely on special settings with \"adversarial checkpoints.\" Can you provide insights into the differences between SAM and USAM during \"real training\" from scratch, rather than in the adversarial context?\n- Instead of scheduling based on normalization, have you considered linear scheduling with respect to the learning rate or other scheduling approaches similar to learning rate selection? Additionally, if the early phase of training is crucial, normalizing the radius first and then using a fixed $\\rho$ for that phase may be suitable for your experiments. Could you provide additional experimental results using scheduling methods other than normalization?\n- [3] mentioned that the early stage of training, referred to as the \"Edge of stability,\" differs between SAM and SGD rather than focusing on the property of SAM to be beneficial in later training stages [2]. However, their analysis was conducted using USAM. If USAM is not suitable for sharpness-aware training, how do you interpret their results in relation to your ideas, and how do they compare to your findings?\n\nI’d be happy to raise my score if the authors can address the weaknesses/questions/limitations in the rebuttal.\n\n\n[2] Towards Understanding Sharpness-Aware Minimization (ICML’22)\n\n[3] SAM operates far from home: eigenvalue regularization as a dynamical phenomenon (ICML’23)\n\nLIMITATIONS:\nThe paper has not provided a detailed discussion on limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n**I have read the rebuttal and raised the score to a 7.** This paper studies the importance of the normalization step in SAM (normalizing the gradient used to perturb the weights) for stability and effectiveness at finding a flat region in a manifold of minima (drifting along a continuum of minima). They provide the following theoretical explanations:\n- On a strongly convex, smooth loss and some specific nonconvex settings, for the same range of learning rates, unnormalized SAM may diverge while normalized SAM does not i.e. the SAM perturbation has to be sufficiently small for convergence. Normalized SAM stabilizes within some local neighborhood around the minimizer (this has been shown previously in Bartlett 2022).  \n- Considering the model x*y for some convex loss, where the flattest solution lies at x,y=0, they show that USAM and GD both get stuck after reaching a closer sharper solution (x=0, y > 0), while SAM keeps drifting towards the origin after getting close to a minima (x=0, y > 0). Their contribution here is showing that without the normalization step, the perturbations become too small. \n\nFor these reasons, they empirically show (in a sparse coding problem) that the normalization step of SAM allows SAM to stably achieve flat minima in comparison to large learning rate GD or unnormalized SAM which may both have problems with divergence.\n\nSTRENGTHS:\n**Significance:** The unnormalized SAM has been studied theoretically as a simplification for SAM, and the contributions of the paper emphasize why this simplification may change the dynamics of SAM entirely. \n**Clarity/quality:** The paper is well motivated, easy to read, and thorough. All theoretical conclusions are sound and are supported by clear empirical trends.\n\nWEAKNESSES:\nThe main weakness is significance as unnormalized SAM which the paper focuses on studying is not utilized in practice and the conclusions of the work, though confirms normalization is good, does not provide further prescriptive insight. The characterizations of normalized SAM conducted in the paper have already been shown in previous works, and the new contribution is solely characterizing how unnormalized SAM is different. Yet, unnormalized SAM only appears in a couple works that theoretically analyze SAM but does not appear in practice. As pointed out by the authors, previous works (Wen 2022, Bartlett 2022) have already shown that normalized SAM stabilizes around a local neighborhood around a minima and it achieves flatter minima by continuing to drift along the manifold. The contribution of this work is to formalize how the perturbation needs to be sufficiently small for convergence, but sufficiently large around the manifold of minimas to find the flat solution, which is achieved by normalization. Yet these conclusions, especially the former, seems a bit intuitive/obvious. Additionally, previous theoretical works have indicated that USAM can be tuned to achieve similar performance to SAM (Andriuschenko 2022). Though there exists a difference between USAM and SAM in the toy models in the paper, there aren't any experiments that indicate that not normalizing is a problem on real datasets and models (that USAM requires delicate tuning, or doesn't work at all in terms of achieve similar training curve and test performance as SAM). Perhaps adding such an experiment would help amend the first concern.\n\n-- \nMinor fixes:\nLine 248 diverges --> divergence\n\nQUESTIONS:\n1. The theorem 4, the only place $\\rho$ shows up is in the final phase, where $|x_t|, |y_t|$ are upper bounded by a factor proportional to $\\rho$, as SAM stabilizes around a local neighborhood. Does this imply that SAM converges around the origin regardless of the magnitude of $\\rho$ and $\\rho$ only determines the rate of convergence? \n\n2. In Figure 7, USAM with $\\eta=0.1$ seems to do comparatively to SAM $\\eta=0.1$. From this, it is a bit unclear whether it is truly difficult to tune the hyperparameters for USAM to work. Is it often the case that USAM can be tuned easily to achieve similar performance as SAM? \nFor example in  (Andriuschenko 2022), the training and test curves of SAM and USAM are relatively similar (e.g. Figure 22). \n\n3. What is the dataset/model used in Figure 1?\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
The authors argue that the normalization of the gradient in the SAM model during the ascent step is beneficial. In particular, classic SAM is shown theoretically and empirically in certain scenarios to not-diverge and to be able to drift along the minima manifold. This implies that SAM is more robust in the choice of the step-sizes compared to the unnormalized SAM.

STRENGTHS:
- The paper is in general well-written and accessible despite being a theoretical work.
- The included examples help the understanding of the ideas.
- The theoretical results seem convincing and sensible, but I have not verified the technical details.

WEAKNESSES:
- I think the similarities and differences to Wen 2023 and Compagnoni 2023 could have been discussed more broadly.

QUESTIONS:
Q1. For the non-convex and L-smooth case (Sec 2.2) SAM does not diverge for any function with the same properties? Or is it possible for SAM to diverge in some functions of this class?

Q2. Regarding the function $\mathcal{L}$ in Sec 3.1. Does the $\times$ imply multiplication of $x,y$? The function $\mathcal{L}$ is $L$-smooth?

Q3. Theorem 6 is compared with Wen 2023 where the step-sizes are ``small enough''. Does this result generalize in the setting of the current paper where the argument is that discrete step-size selection is more robust?

LIMITATIONS:
The authors do not discuss any potential limitation of the approach. Implicitly, the limitations are encoded in the assumptions of the theorems. Also, the paper proposes a technical method with no direct negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07