The Crucial Role of Normalization in Sharpness-Aware Minimization
Authors: Yan Dai
Kwangjun Ahn
Suvrit Sra
Venue: NeurIPS 2023 poster (2023)
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
Yan Dai IIS, Tsinghua University
Kwangjun Ahn EECS, MIT kjahn@mit edu
Suvrit Sra TU Munich / MIT suvritemit.. edu
yan cun-dai200mails tsinghua. edu. en
Abstract
Sharpness-Awar Aware Minimization (SAM) is recently proposed gradient- based opti- mizer (Foret al., ICLR 2021) that greatly improves the prediction performance of deep neural networks. Consequently, there has been 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 empiri- cally study the effect of normalization in SAM for both convex and non-conved functions, revealing two key roles played by normalization: i) it helps in stabilizing the algorithm; and ii) enables the algorithm to drift along a continuum (manifold) of minima 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 syper-parameters, supporting the practicality of SAM. Our conclusions are backed by various experiments.
Introduction
We study the recently proposed gradient-based optimization algorithm Sharpness-Aware Minimization (SAM) (Foret et al., 2021) that has shown impressive performance 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 wi W1+1 W/ 7 (w.) (where 7) is called the learning rate), SAM performs the following update:'
V(w)) 0++1 = 00 VL W/ PIPPIII [Vw))
where p is an additional hyper- parameter that we call the perturbation radius. Foret al. (2021) motivate SAM as an algorithm minimizing the robust loss max L(w + e), which is roughly the loss at w (i.e., (ww)) plus the "sharpness" of the loss landscape around w, hence its name.
The empirical success of SAM has driven recent surge of interest in characterizing its dynamics and theoretical properties (Bartlett et al., 2022; Wen et al., 2023; Ahn al., 2023d). However, a major component of SAM remains unexplained in prior work: the role and impact of the normalization factor I. 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 al., 2023). Instead of the SAM update (1), these works consider the following update that we call Un-normalized SAM (USAM):
Wh+1 = Wt v L(w, + pV L(w)))
"The first two authors contribute equally. Work done while Yan Dai was visiting MIT.
In principle, the normalization Equation may make SAM ill-defined. However, Wen et al. (2023, Appendix B) showed that except for countably many learning rates, SAM (with any p) is always well-defined for almost all initialization. Hence, throughout the paper, we assume that the SAM iterates are always well-defined.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Apart from experimental justifications in (Andriushchenko and Flammarion, 2022), the effect of this simplification has not yet been carefully investigated, although 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 factor [VC in the SAM update (1)?
1.1 Motivating Experiments and Our Contributions
We present our main findings through two motivating experiments. For the setting, choose the well-known over-parameten wer-paramettrizee matrix sensing problem (Li et al., 2018); see Appendix A for details.
1. Normalization helps with stability. We first pick learning rate r that allows GD to converge, and we gradually increase p from 0.001 to 0.1. Considering the early stage of training shown in Figure 1. One finds that SAM has very similar behavior GD, whereas USAM diverges even with a small - it seems that normalization helps stabilize the algorithm.
GB SAM (p- 0.11
GD
USAM 0.11
0.03)
Train Loss
SAM (p- 0.0011
USWW Ip 00.00001)
Epoch
Epoch
Figure 1: Role of normalization for stabilizing algorithms 0.05).
2. 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 "manifold" of minima (which exists since the problem ) as the training losses are close to zero. The first plot compares SAM and USAM with the same 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 p 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 minimum.
0.1 (the largest p for which USAM doesn't 't diverge); notice that USAM
1p=0.11
(p-00.50
SAK (p55
G0
2- ....
USAM 1p 0.05)
(p- 0:1)
(p-0.3)
(p-
1p=0
10000 20000 10000 40bon 50000
10000 20000
30000
400ns 50020
10000
20000
55000
40005
50000
Epoch
Epoch
Epoch
Figure Role of normalization when close a manifold minimizers (1 0.005).
Overview of our contributions. In this work, motivated by Figure and Figure we identify and theoretically explain the two roles of normalization in SAM. The paper is organized as follows.
1. In Section 2, normalization helps stabilize. particular, we demonstrate that normalization ensures that GD's convergence implies SAM's non-divergence, whereas USAM start diverging much earlier.
study the role of normalization in the algorithm's 's stability and show that
2. In Section 3, study the role of normalization near 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.

---

3. In Section 4, to illustrate our main findings, we adopt the sparse coding example of Ahn et al. (2023a). Their result implies dilemma in myyee-parameter tuning for GD: small 17 gives worse performance, but a large 7 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 p is tuned: Foret et al. (2021) suggest that p can be tuned independently after tuning the parameters of base optimizers (including learning rate 17, momentum 3. and so on). In particular, this feature makes SAM 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 the choice of (t, 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 p. 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 scalc-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 component reducing the sharpness Amax ( ( [(w))). while Wen et al. (2023) characterize SAM near the manifold I of minimizers and show that SAM follows a Riemannian gradient flow reducing Amax (V"((w) when i) initialized near r, and ii) 1/ 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 v) = o(1) and p = 0(1). A recent work by Ahn et al. (2023d) formulates the notion of E-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 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 linear regression model, by Agarwala and Dauphin (2023) who study the initial and final dynamics of USAM over 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-timg variants of algorithms while we consider discrete (original) versions, results directly apply to the SAM deployed in practice and the USAM studied in theory.
Sddge-of-Stability In the optimization theory literature, Gradient Descent (GD) was only shown to find minima if the learning rate is smaller than an "Fudge-o-SStablity"" threshold, which is related to the sharpness of the nearest minimum. However, people recently observe that when training neural networks, GD with 17 much larger than that threshold often finds good minima well (see (Cohen
3

---

et al., 2021) and references therein). Aside from convergence. GD with large is also shown to find flatter minima (Arora et al., 2022; Ahn 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 well-known fact about the stability of GD: for convex quadratic cost with the largest eigenvalue of Hessian being 3 (i.e., 3-smooth), GD converges to a minimum iff vi < 2/8. 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 a-strongly-convex and 3-smooth loss function where GD guaranteed to converge once 17 < 2/8. We characterize the stability of SAM and USAM in the following result.
Theorem (Strongly Convex and Smooth Losses). For any a-strongly-convex and B-smooth loss function C. for any learning rate nl < 2/8 and perturbation radius Z O, the following holds:
1. SAM. The iterate wt converges local neighborhood around the minimizer w². Formally,
n1,8 p²
[(we) [(w") S (1 an m()) (1=))(2-n (C(())) - [(w")) + 2a(2
Vt.
(3)
n)) n3))
USAM. contrast, there exists some o--ttrongly-comvex and 3-smooth loss such that the USAM with rj E (2/(8+062), 2/s] diverges for all except measure zero initialization wo-
As discussed, it is well-known that GD converges iff 2/8, and Theorem 1 shows that SAM also does not diverge and stays within an O neighbooigh around the minimum as long as already
n 2/8. However, USAM diverges with even lower learning rate: ") > 2/(8-p8*) make USAM diverge. Intuitively, the larger the value of p. the easier is for USAM to diverge.
One may notice that Equation 3, compared to the standard convergence rate of GD, exhibits additive bias term of order O(pp²). This term arises from the unstable nature of SAM: the perturbation in (I) (which always has norm p) 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" minimum, mean its iterates approach and stay within neighborhood of that minimum.
Due to space limitations, the full proof is postponed 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):
1 ((w++1) S [(w,) 2 2 - -
322
2
By invoking the strong convexity that gives L(w;) - [(w*) < 1 V (w(w)) obtain
L(w*)) )) + 8728'p²
((w++1) [(w*) S an nn)nnn)))) ())
Recursively applying this relation gives the first conclusion. For USAM, we consider the quadratic
1 w Aw where loss function same as (Bartlett et al., 2022). Formally, suppose that L(w) = 2
A = diag (}1, A2,..., Aa) is PSD matrix such that X X2 2 -d > 0. Let the eigenvectors corresponding to \1, 11.22.... .., Ag be 61, e....., Cd, respectively. Then we show the following in
Theorem 10: for any 17( + pA7) 2 and (Wa e1) # O, USAM must diverge. As [(w) Jw'' Aw is As-smooth and Az-strongly-convex, the second conclusion also follows.
Intuitively, the difference in stability can be interpreted as follows: during the early stage of training, We and V VL(w)) often have large norms. The normalization in SAM then makes the ascent step (((()) Wi + P PTTVE[we]] not too far away from w.. Hence, if GD does not diverge for this 77, SAM does not either (unless the p-perturbation is non-negligible, i.e., ||w:||| > / no longer holds). This not true for USAM: since the ascent step is un-normalized, it leads to a point far away from w,, making the size of USAM updates much larger. In other words, the removal of normalization leads to much more
aggressive steps, resulting in different behavior than GD and also an casier divergence.

---

2.2
Generalizing to Non-Convex Cases: Scalar Factorization Problem
Now let us move on non-convex losses. We consider a scalar version of the matrix factorization
problem ming.v 2 3|VVV - -A whose loss function is defined as [( y) 2/2 (xy)² Denote the initialization by (Io, yo). then [(x, y) is 3 A (r + vd -smooth inside the region (zx y) + + = 33. Hence, a learning rate < 3/3 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(z, v) = 1 (ra)) restricted to B-smooth region, we set 1 1/3 < 3/8 (so GD finds a minimum), then
1. SAM. SAM never diverges and approaches a minimum within an O(p)-neighborhood (in fact, SAM with distinct p's always find the same minimum (0, 0)).
2. USAM. On the other hand, USAM diverges once p 157 which could be much smaller than 1.
Thus, our observation in Theorem is not limited convex losses - for our non-convex scalar- factorization problem, the stability of SAM remains robust to the choice of p. while USAM is provably unstable. One may refer to Appendix for the formal statement and proof of Theorem
2.3 Experiment: Early-Stage Behaviors when Training Neural Networks
CNN with tanh activation
FCN with ReLU activation
SAM p
2x
2 10
SAM p 02
SAM p
USAM 0.0000
USAM p-0.1
USAM p=0.2
USAM p=0.3
207
43.30
4X110
Figure 3: Behaviors of different algorithms when training neural networks (17 = 0.025).
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 7)). We pick two neural networks: a convolutional neural network (CNN) with tanh activation and fully-connected network (FCN) with ReLU activation. We train them over the CIFARIO 10 dataset and report the early-stage training losses. Similar to Figure 1, Theorem 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 drift near minima. To convince why this is beneficial, we adopt 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) 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 C(x, y) defined over all (I, y) E F² as
L: (x. + f(x x y). where é is convex, even, and I-Lipschitz.
(5)
5

---

This L was recently studied in (Ahn et al., 2023a) to understand the behavior of GD with large 17's. By direct calculation, the gradient and Hessian of L at a given (z, y) can be written as:
82
[ 0
Vl(x,)) = (((z)))=((x)) 2
v) = E"
y
+ +((xx)
(6)
Without loss of generality, one may assume is minimized at 0 (see Appendix E for more details regarding f). Then, L achieves minimum at the entire z- and y-axes, making it good toy model for studying the behavior of algorithms near continuum of minima. Finally, note that the parametrization I 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 É even, without loss of generality, we always assume that the initialization wo = (Io. yo) satisfies yo Z0 > 0.
Theorem 3 (Theorems and 6 of (Ahn et al., 2023a); Informal). For any 1 = 7/(0 - xg)), the GD trajectory over the loss function f(x, v) = {(zy) has two possible limiting points:
1. If 7 < 2, then the iterates converge to (0, )oo) where v E ( - 7/4 ((//>)]
2. yf > 2, then the iterates converge (0, yzo) where 1 E 17/n (n)).
Intuitively, the limiting point of GD (denoted by (0, y>o)) satisfies v². R min = 2/1}- For
simplicity, we denote TJD = 2/28 - as the threshold of 17 that distinguishes these two cases.
Interpretation of Ahn et al. (2023a). Fixing the initialization (zo. 30). it turns out this model has a nice connection to the sparse coding problem, wherein it's desirable to get smaller v (which we will briefly discuss in Section 4). According to Theorem 3, to get a smaller v one must increase the learning rate 1) beyond MGD- Hence we mainly focus on the case where 1) > T/D - in which case we abbreviate voc 2/n, (see Table 1). However, GD diverges once y is too large in their language, 7 cannot be much larger than 2. This dilemma of tuning 17, as we shall illustrate in Section in more detail, makes GD brittle choice for obtaining a better y.
SAM
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 r) and p suffices to drive SAM toward the origin no difficult tuning needed anymore!
USAM
GD
Figure 4: Trajectories of dif- ferent algorithms for the ((ry) loss (r) = 0.4 and p = 0.1: ini- tialization (Io,yo) = (2, v 10) is marked by black dot).
In contrast, USAM does not keep moving along the axis. Instead, a lower bound on '². also presents although smaller than the GD version. As we will justify in Theorem 5, USAM does get trapped at some non-zero y² Thus, a dilemma similar to that of GD shows up: for enhanced performance, an aggressive (17, p) is needed; however, as we saw from Section 2, this easily results in a divergence.
Assumptions. To directly compare with (Ahn et al., 2023a), we focus on the cases where vō 2 //n 1/ and 7 E [ 2] is a constant of moderate size; hence, 1) is not too different from the MGD defned in Theorem 3. In contrast to most prior works which assume a tiny P (e.g., (Wen et al., 2023)), we allow p to be as large as a constant (i.e., we only require p (11) 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 follows:
Theorem 4 (SAM over Single-Neuron Networks; Informal), For any E LLTG.. 2nco) and p (L), the SAM trajectory over the loss function [(x, 3) f(ry) can be divided into three phases:
1. Initial Phase. Dt drops so rapidly that |ærl = O(vT) in O())) steps. Meanwhile, ye remains large: specifically, yi = 22( (7)). Thus, SAM approaches the y-axis (the set of global minima).

---

Table 1: Limiting points of GD, SAM, and USAM for the {(xy) loss (assuming v/ > mKD).
Algorithm
GD
SAM
USAM
Limiting Point (0. yec)
yeo 2/1
V
+ pul Du 2/n
Middle Phase. it oscillates closely to the axis such that [axx] = O always holds. Meanwhile, yt decreases fast until yt S | that is, (zze remains small and SAM approaches the origin. Final Phase. we = (It yt) gets close to the origin such that (arel, lyve = O(V)) + np)- We then show that wy 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 specific minimum. Moreover, SAM always approaches this minimum for almost all (n. p)'s. This matches our motivating experiment in Figure 2: No matter what hyper-parameters 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/nn (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 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 yi that depends both 1) and p. In other words, the lack of normalization factor leads to diminishing drift.
Theorem 5 (USAM over Single- Meugle-Neuron Networks; Informal). For any nj nap, 2nau and (1), the USAM trajectory over the loss function ((x, y) = E(xy) have the following properties:
Initial Phase, Similar to Initial Phase of Theorem 4, |xzel = TT) and yt = Q1 /n) hold for the first (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
(VoL+
(1+ (1+ <
i.e., s2 < PUSAM
2p.
(7)
|xt] decays exponentially fast, which in turn ensures 222 P lim vi 2 (1- - p²))
Remark. Note that USAM becomes GD as we send p + 0+, and our characterized threshold viSAM indeed recovers that of GD (i.e., 2/1 from Theorem 3) because limp-0+ B8/n + 1)/2p = 2/1--
Compared with SAM, the main difference occurs when close to minima, i.e., |zz|| = ( Consis- tent 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 formal version of Theorem 5 together with its proof.
3.1.3 Technical Distinctions Between GD, SAM, and USAM
Before moving on, we present more technical comparison between the results stated in Theorem 3 versus Theorem 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. and V((w,) become small. Hence the normalization plays an important role: as V.(w,) small, W, and W, + pV ((w,) 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 it diminish, which prevents SAM from converging to minimum and keeps drifting along the manifold.
2Specifically, either an exponential decay Vt+1 S (1 - np²) )ye or constant drop yt-1 S yt - np appears "The MiSAM in Equation is defined as the solution the equation (1 + py² )y² 2.
7

---

This high-level idea is supported by the following calculation: recall Equation 6 that V VL(zz, Mo) e' (" lyv. xz]"_ Hence, when |,, < in Final Phase, the "ascent gradient" direction ( yr) (i.e., the ascent steps Equation 1 and Equation 2) almost perpendicular to the y-axis. We thus rewrite the update direction (i.e., the difference between w++1 and wy) for each algorithm as follows.
1. For SAM, after normalization.
V[(w)) is roughly a unit vector along the 2-axis. Hence, the [V
update direction the gradient at w/+1/2 Re [p yol". Once ys is large (making 1/+ far from minima), V (2(()1/)))) thus have a large component along yt. which leads to drifting near minima.
2. For GD, by approximating (u) R the magnitude of 2₁ is updated as |æ141| [xt "/ = [11 myr" )=al) which allows an exponential decay. Thus, GD converges to minimum and stop moving soon after 2/6 > vē.
derive (a y1) R [z,v: x.ye] When 2/n > vi.
3. For USAM, the descent gradient is taken w, [V [[1+ [(1+py)) )a,
(+pp))yx] Thus, V = (I++y?)) pzi) ) Pzyy (1+ 21 pzż }æ _ by writing (u) se
u. This makes USAM deviate away from SAM and behave like GD: by the similar argument as
GD, USAM stops minimum soon after 2/n (1 + pyt) + pz 2 uf R (1 + pyi) )y?
Hence, the normalization factor in the ascent gradient helps maintain 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 functions (see, e.g., (Karimi et al., 2016)). Recall that f satisfies the μ-PLL condition if V (2))l u([(w)- - min.. ((w)) for all w. Building upon the analysis of Andriushchenko and Flammarion (2022), we show the following result when applying USAM to 3-smooth and μ-PL losses.
(PL)
Theorem 6 (USAM over PL Losses; Informal). For B-smooth and μ-PL loss for any < 1/8 and < 1/8, and for any initialization wo. ||w woll S poly(T) p. 3, p) VC(w)) min L(w),, VI.
This theorem has the following consequence: Suppose that USAM encounters a point wo that is close some minimum (i.e., ((wa) min L(w)) during training. Then Theorem 6 implies that the total distance traveled by USAM from wi is bounded - in other words, the distance USAM moves along the manifold of minimizers can only be of order Ol O(VL(wo) min,, ((w)).
As remark, we compare Theorem 6 with the recent result by Wen et al. (2023): their result essentially implies that, for small enough 1j and p. SAM iterates initialized close to manifold the minimizers approximately track some continuous dynamics (more precisely, 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 are contained Appendix
3.3 Experiments for Practical Neural Networking Training
160
Test Test.Ace Accurac)
54-
Train Accuracy
52-
(p-0.01)
50.
60
SAM (.05))
50
USAM (p- 0.01)
USAM 0.005
44
USAM (p- 0.1)
42
40
80
100
20
40
69
100
Epoch
Epoch
Figure 5: Training ResNet18 CIFAR-10 from bad global minimum 0.001, batch size 128).
We close this section by verifying our claims in practical neural network training. We train 18 on the CIFAR-10 datasct, initialized from poor global minimum generated per Liu et al. (2020)
8

---

(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 more pronounced algorithmic behavior near the minima via tracking the test accuracy. From Figure 5, we observe:
1. GD gets stuck at this adversarial minimum, in the sense that the test accuracy stays ats 48%.
2. 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).
3. USAM with small P gets stuck like GD, while USAM with larger p's deviate from this manifold.
Hence, USAM faces the dilemma that we describe in Subsection 3.1: a conservative 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 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 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-bay full-bbtch checkpoint", while SAM keeps increasing its test accuracy via drifting along the minima manifold.
Case Study: Learning Threshold Neurons for Sparse Coding Problem
To incorporate our two findings into single example, consider training one- mue-iddee-llaee ReLU networks for the sparse coding problem, a setup considered in (Ahn et al., 2023a) to study the role of n in GD. Without going into details, the crux of their experiment is to understand how GD with large y finds desired structures of the network- in this specific case, the desired structure is the negative bias in ReLU unit (also widely known "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 re- produce their results in Figure 6. plotting the test accuracy on the left and the bias of ReL.U unit on the right. As they claimed, GD with larger ? learns more negative bias, which leads to better test accuracy.
GD -=1
() 10-0)-000 0.05
Test Accuracy
GD 0.025
Bios
Time in X Step)
Time (n X Step))
Figure 6: Behavior of GD for sparse coding problem.
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 Sdge-of-Stabiliyy (Cohen et al., 2021)). In practice, without knowledge of the problem, this requires 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 al. (2022).
In contrast, as we will justify shortly, SAM does not suffer from such a "dilemma of tuning" longer attains
matching with our results in Theorem 4. Moreover, the removal of normalization such a property, as we demonstrated in Theorem 5. In particular, for USAM, one also needs to carefully tune 17 and p for better performance (n, p) 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.
we inspired in Theorem 5 and Theorem 6, small
In Figure 7, we plot the performance of SAM, USAM, and GD with different 1y's (while fixing p) 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 not large enough, which consequently to leads poor test accuracy. On the other hand, no matter what

---

GD
5AM g=0.1
SAM 0=0.05
SAM m=0.025
USAM /=0.1
0=0.05
Test Accuracy
USAM 7=0.025
0.70
Blas
0.63
0.00
-0.
400
600
1002
602
1002
Time Step)
Time Step)
Figure 7: Behaviors of different algorithms for sparse coding problem (p 0.01). The gray curves
(corresponding GD with different learning rates) are taken from Figure 6 with the same set of 1/
?) is chosen for SAM, bias is negative enough and ensures better generalization. Hence, Figure illustrates that compared to SAM, USAM is less robust to the tuning of 17.
In Figure 9 (deferred to Subsection B.2), also compare these three algorithms when varying p and fixing 7). addition to what we observe in Figure 7, we show that normalization also helps stability increase p. while SAM remains robust to the choice of p. Thus, USAM also less robust to the tuning of p. In other words, our observation in Figure 7 extends to
USAM quickly diverges
Hence, putting Figure 7 and Figure 9 together, we conclude that SAM is robust to different con- figurations of (1),P) while USAM is robust to neither of them. Hence, the normalization of SAM eases syper-parameter tuning, which is typically 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 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 pappe--rrraeter tuning, supporting its practicality.
In this work, we follow 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 LN00014-20-1-294) and MIT-IBM Watson as well as a Vannevar Bush fellowship from Office of Secretary of Defense. Suvrit Sra acknowledges support from an NSF CAREER grant (1846088), and NSF CCF-211266 (TILOS Al Research Institute). Kwangjun Ahn also acknowledges support from the K Kwanjeong Educational Foundation. We thank Kaiyue Wen, Zhiyuan Li, and Hadi Daneshmand for their insightful discussions.
10

---

References
Jacob Abernethy, Alckh Agarwal, Teodor Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. arXiv preprint arXiv:2305. 2023.
Atish Agarwala and Yann N Dauphin. Sam operates far from home: eigenvalue regularization as dynamical phenomenon. arXiv preprint arXiv:2302.08869, 2023.
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Proceedings of the 39th International Conference 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.07466), 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.00097), 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 10X55 2023d.
Zeyuan Allen-Zhu and Yuanzhi Li. Physies of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305. /3673, 2023.
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-awarr mini- mization. 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-awann 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-7377 2022.
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness- minimiza- tion: Bouncing across ravines and drifting towards wide minima. arXiv preprint arXiv:2210.01513, 2022.
Kayhan Behdin and Rahul Mazumder. Sharpness-awar minimization: An implicit regularization perspective. arXiv preprint arXiv:2302.//836, 2023.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-ublenbeck like process. Conference on learning theory, pages 483-513 PMLR, 2020.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, 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 Orvicto, 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 Neural Information Processing Systems, 2021.
Alex Damian, Fshaan 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-awee aware minimization for improved training of neural networks, In International Conference 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 Third Conference on Uncertainty 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-30500 2022.
Stanislaw Jastrzkebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint 1711 1,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 proximal- gradient methods under the polyak-tojusiewiezz condition. In Machine Learning and Knowledge Discovery in Databases: European Conference, pages 795-811-8 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 Learning Representations, 2017.
Hoki Kim, Jinseong Park, Yujin Choi, and Jaewook Lee. Stability analysis of sharpness-aware minimization. arXiv preprint arXiv:2301.065 06300, 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 ver-paraaeteriied matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pages 47. PMLR, 2018
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: To- wards mechanistic understanding. International Conference on Machine Learning (ICML) (trXiv:2303.0424555 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 Neural Information Processing Systems, 33:8543-8552, 2020.
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsich, and Yang You. Towards efficient and scalable sharpness-awarr minimization. In Proceedings of the IEEF/CVF Conference 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 Neural Information Processing Systems, 2022.
12

---

Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generaliza- tion in deep learning, Advances neural information processing systems, 30, 2017.
Dongkuk and Chulhee Yun. Practical sharpness-aware minimization cannot converge all the way to optima. NeurlPS 2023 (arXiv:3306098800) 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 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:9983- 9994, 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 general- ization. Advances in Neural Information Processing Systems, 33:2958-222 2969,,9999949 2020.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. arXiv preprint arXiv:2206.044: 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 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 Learning Representations, 2022.
13

---

Appendix
Setup the Motivating Experiment
15
Additional Experimental Results
15
B.1 Running SAM and USAM from Other Initializations
15
B.2 Varying p While Fixing y Sparse Coding Example
15
C Omitted Proof for Smooth and Strongly Convex Losses
16
C.I SAM Allows Descent Lemma Like GD
16
C.2 USAM Diverges on Quadratic Losses
17
D Omitted Proof for Scalar Factorization Problems
18
D.1 SAM Always Converges on Scalar Factorization Problems
Probe.s.........................
18
D.2 USAM Diverges with Small
19
E Assumptions in the Single-Neuron Linear Network Model
20
Omitted Proof of SAM Over Single-Neuron Linear Networks
20
F.1 Basic Properties and Notations
20
F22 Initial Phase: 2t Decreases Fast while y/: Remains Large
21
F3 Middle Phase: y Keeps Decreasing Until Smaller Than :...............................
23
F.4 Final Phase: Both 2, and ye Oscillates Around the Origin
24
Omitted Proof of USAM Over Single-Neuron Linear Networks
25
G.1 Basic Properties and Notations
25
G.2 Initial Phase: 2, Decreases Fast while Remains Large
26
G.3 Final Phase: Yt Gets Trapped Above the Origin
28
H Omitted Proof of USAM Over General PL functions
29
14

---

A Setup of the Motivating Experiment
In the motivating experiments (Figure 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:
1. Generate the matrix by sampling each entry of U R independently from standard Gaussian distribution and let X (**)))
2. Normalize each column of U** to unit norm so that the spectral norm of U is close to one.
3. For every sensing matrix A, (i 1,2....,my m.) sample the entries of A; independently from standard Gaussian distribution. Then observe b, = X*).
In particular, for the experiments, we chose = 5, d = 100, and m = 5dr.
B Additional Experimental Results
B.1 Running SAM and USAM from Other Initializations
1.0
Mry
0.770
n Accurasy
0.705
SGD
0.4-
SAM p-5.
0.755
USAM p=
p-0.5
USMM p=1
0.750
20
60
100
140
100
120
140
Epoch
Epoch
Figure 8: Training R on CIFAR-10 from a "full-batch checkpoint" of Damian et al. (2021)
0.001, batch size 128).
B.2 Varying p While Fixing 7) in Sparse Coding Example
02
GD
SAM p-0.5
SAM p=0.3
5.00
0.-
SAM p-0.1
p-0.01
-02
USAM p-0.5
Toxt Accuracy
1.75
p-0.3
USAM p-0.1
p-0.01
0.00
5.00
0.50
-12
300
X00
620
200
3203
200
400
6:0
1000
Time (1) X Step)
Time (n X Step)
Figure 9: Behaviors of different algorithms for sparse coding problem (n 0.025). Note that USAM with 0.5 (the dashed violet curve) diverges and becomes invisible except for the very beginning.
Aside from Figure which varies y while fixing p. we perform the same experiment when varying and fixing 7). The main observation SAM is similar that of Figure different hyper-parameters all keep decreasing the bias and give better test accuracy even with the tiniest choice p = 0.01.
15

---

However, for USAM, there are three different types of o's as shown in Figure 9:
1. For tiny p 0.01, the bias doesn't t decrease much. Consequently, the performance of USAM nearly degenerates to that of GD while SAM with p = 0.01 still gives outstanding performance.
2. For moderate p = 0.1 and p = 0.3, USAM manages to decrease the bias and improve its accuracy, though with slower speed than SAM with the same choices of p.
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 eyper-parameter tuning much more manageable.
C Omitted Proof for Smooth and Strongly Convex Losses
We shall first restate Theorem here for the of presentation.
Theorem (Restatement of Theorem 1). For any a-strongly-conver and 3-smooth loss function L, for any learning rate rj S 3/8 and perturbation radius Z 0, the following holds:
1. SAM. The iterate converges to local neighborhood around the minimizer wt Formally,
n p²
((w,) L(w*) s (1 (1--amp a (C(())(((()))))))))))))))))) [(w*)) + 2a n3)"
Vt.
USAM. In contrast, there exists some a-stronglyyconvve and B-smooth loss such that the USAM with yj (2/(8+007), 2/3] diverges for all except measure zero initialization wo-
We will show these two conclusions separately. The proof directly follows from Theorem 8 and Theorem 10.
C.1 SAM Allows Descent Lemma Like GD
Theorem 8 (SAM over Strongly Convex and Smooth Losses). For any a-strongly-conee and 3-smooth loss function L, for any learning rate n S 2/8 and perturbation radius 0. the following holds for SAM:
11.3pp
((wi) - L(w") s - (1-an(2--11 (C))))) [(w")) + 2a(2 n18)
Proof. We first claim the following analog of descent lemma, which we state as Lemma 9:
[(w++1) S ((w,)= 2 - - 3233'p2
1
2
By definition of strong convexity, we have
((w.) - L(w*) S (VL(ww)) we - w-) || w² S 1
2a
where the last inequality uses the fact that (a, b) 111b112 s ||a||2 Thus, combining the two inequalities above, we obtain
[(w(+1) S L(w,) (n1)))=o((22))) - [(ww)) 12.33p2
2
which after rearrangement becomes
((we+)) L(w*) s <(1-on(-n3))(((u)) [(ww)) 17233
Unrolling this recursion, we obtain
12 t-1
((w,) L(w*) s (1 an(2 - (B0)) [(w+))
an n²))**
2 k=
16

---

$ (1 an((2 (2 f(w")) 1//33000 0²
Σ
o07(2
-1:3)/
2
k=
3²
(1 (2) [(w")) 20xT(2 1133)
This completes the proof.
Lemma (SAM Descent lemma). For a convex loss L that 3-smooth, SAM iterates wy satisfy the following when the learning rate < and p 2 0:
17238p2
[(w++1) S ((w)) (2 - n5) ( V
Vt.
2
Proof. Let Df (LL(:)/IIVVV and wx+1/2 - wi + pud so we have wM+1 = We 1/ ((weet)).)
Then the B-smoothness of & yields
1²2B
[(w(+1) S [(w)) n(V((w)), V((w.))))
| [[(W4+//111
We start with upper bounding the norm of ((W)+
[V((+++))) VV((w++//) - V[(w))] [F
V((w:))
S p² [V][ 2VV V[(w:)),
Hence, long 1) 2 have the following upper bound [(w)+1))
[(w++1) - [(w,) n, V ((1++1))) V((w))) + 7²3
| ((w))//2)
3²8
s L(w,) 7VV (V v[(w)))
(3p² [[[(w))11
2 V((
1 17 - V[(w)) 12238p2
-(w:)
2
Now lower bound ( (VL(W)+4/2), VL(w:)). Note that
(VL(w, + por), v[(w)) (VL(wr pvz) V((w,), VL(w)) | [[
[V[(w))] VV[(,, + pvr) VL(w;), pu,) + [V:((;)]] 2 [[V[(,))
where the last inequality uses the following standard fact about convex functions: for any Wi W2. (VL(w)) VL(w2), w2) 2 0. Hence, arrive
3 V [V[(w)] - - 12 (VW((())/))) VL(w:))
17233pp
[(w++1) S ((w:))
2
11² 33
1//88
< L(w,)
[V((w))] (17 - 17 VL(w.)]l
2
(2 n VL(w))) 1/238p2
S ((w))
2
and thus finishing the proof.
C.2 USAM Diverges on Quadratic Losses
Theorem 10 (USAM over Quadratic Losses). Following (Bartlett al., 2022), consider the quadratic loss L induced by PSD matrix A. Without loss of generality, let L minimize at the origin. Formally,
C(w) = 2" TAw, where A =
(8)
where Amax ... Aa > 0 and Umax = E1, e2,.... Gd are the eigenvectors corresponding to X1, >2,..., ......,, Ad respectively. Then the iterates of USAM Equation 2 applied to Equation 8 satisfy:
17

---

1. If n(Amax + p.)max) 2, then the iterates converges to the global minima exponentially fast.
2. If (fm((x + p. and if (Wo, Umax) + O, then the iterates diverge. Moreover, such loss function ---smooth and Ag-strongly-convex
Proof. As V((w) Awc the USAM update Equation 2 reads
Wi+1 wt 1/VC(w + (VL(w,)) (I nd m7p7))um we
Hence, if 17(Amas + pA < 2. then since |I- 7A np^lll < 1, we have that ||wl| + 0 exponentially fast. On the other hand, if ()(maxx + pA ||, 2 | nd - e1)l. follows that the iterates diverge as long as (wo, Umax) *
2. then | .n niii > 1. Since
Finally, as V²L(w) A, we directly know that L is Ss-smooth and x4-strongly-convex.
D Omitted Proof for Scalar Factorization Problems
Theorem 11 (Formal Version of Theorem 2). Consider the scalar factorization loss L(x, y) (xy))² If (zo; yo) satisfies 2no yo > 3o &> O and n (a + vū) then i) SAM finds neighborhood of the origin with radius (p) for all p. and ii) USAM diverges as long as p 2 157.
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 Wt = [x Ye T be the t-th iterate. Denote V, = V ((ws) = [xeee
x&ue]. For each step t, the
actual update of SAM is therefore the gradient taken at
24 + ptr 1/ + p
where 27 sign (x4
W wi +
Vo vē
By denoting V. = V V(+.), the update rule of SAM is
St 11 + 2y 1) (yl + prrez 2)
x++1
We v. =
Vt++
U++1
y/: (7 py1a 1) (ye + parezi
We make the following claim: Theorem 12 (SAM over Scalar Factorization Problems). Under the setting of Theorem 11, there exists some threshold for SAM such that S for all t 2
Proof. We first observe that SAM over L(x, y) always pushes the iterate towards the minima (which are all the points on the z- and y- axes). Formally:
If 2, 2 0, then 1141 T₁- If F4 < O, then Tz41 Zz-
If yı 2 O. then yt+1 S yt. yt S O, then 36++1 2 yc-
This observation can be verified by noticing sigu(pytz; similarly sign(pu:= pushes 3y towards the y-axis. The same also holds for ye by symmetry.
sign(y:) sign( = sign(x:) and sign(1:). In other words, sign(x; + pu: = sign(x.) and thus always
In analog to the descent lemma for GD, we can show the following lemma:
Lemma 13. When picking 1j = (xg + vz) -1 we have L(w++1) - L(we) < as long as Tt, Bt 2 5p.
y² 2xy x²
Proof. As V²[(x,y) {(x, v) : 2 + x² 33.. Then we have (recall that = (w) + = = 1/2)
we know is 3 A (ag + -mooth inside the region
2xy
8
((w41) [(w,) I (VL(w,), we+1 - w.) + 2
||a+1 - w|ll
18

---

V. v. + - 2V, Vi)).
To make sure that it's negative, we simply want
0 S (V. 2. Vi) = ziii aa + pVL 1) ")i + poizz
=xxce_12) ))) (Is + py ())==(((x " + p
zi.i (o) + +y122 )) (yr+ +pxrz 1) (xs + py1z) ))² (vo + pzizz
which can be ensured once
(ze + pyr2;" (ye + pa 1)² 2 20,1/ (x) + pulz) 1)² (ye + partz s 207 yr.
If we have It, 2 5p, then as 2, I min( Ue J) we have
+ (01=5) + p))===)) & (Ia + p)(u p)² S 1..29 < 2n.vi
which shows the first inequality. The second one follows from symmetry.
Therefore, SAM will progress until ie S 5p or ys I 5p. Without loss of generality, assume that It S 5p; we then claim that will soon decrease to (pp).
Lemma 14. Suppose that | & 5p but y 2 5p. Then |1441l s 5p but yt+1 S (1 2mpp yn-
Proof. First, show that |a, remains bounded by 5p. Assume F₁ without loss of generality:
T++1 It 11( p3gt= 1) (y+ 2)
2 xz.
m5pp + (y) + p)²
4
4
2 766
+
p² + 2pp
2 -5p,
(9)
51
5n
where the second last line uses oi S v/ vū)-z < Meanwhile, we that yt decreases exponentially fast by observing the following:
+ot " S 11 and the last one uses rj
oh
(x
Ut+1 = y: 710 + pxx,===)000 + pb
S y/: MB8e (19111))
1 _110² y/+
< 1
where the last line uses 2, (x? + -1/2 (2y2) 2-1/2 as yr 2 5p |.2//-
1/2
So eventually we have |rr|, Wel < 5p. Recall that Equation 9 infers | 5p from [zzel < 5p. Hence, by symmetry, we conclude that |x |y++1| S 5p hold as well. Therefore, SAM always finds an -eighheigbooooood of the origin, i.e., it is guaranteed to converge regardless of
D.2
USAM Diverges with Small p
For USAM, the dynamics can be written as
xe+1
It + pr,y? D1
[xt n7xx + pxxy)) (2) :: + pa,y, y1
ye+1
1
17
+ pazy:)
I/ n(x; + prr ) (2(1y, + palyi))
(10)
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 p 15n, | 2|ae| and |ye+1| 2 2lye| for all 2 1; other words, USAM diverges exponentially fast.
19

---

Proof. Prove by induction. From Equation 10, our conclusion follows once
(x + parv (ve + p=lv=)² 2 3) 3/
n (ar +
Uv+ + p=iv)) 2 3/yel-
According to our setup that 7 = (xî yτ) - yo I 2rg. and the induction statement that, we have 2 (5x))-1. The second inequality then holds as long as
(px (p= 15x2lyel. i.c., x² 2 15,
which is true zi 2 z.g y/² 2 vo- yo To. and p 2 15n- 3zd Note that the bounds p very loose, and we made no effort to optimize it; instead, we only aimed to show that USAM starts to diverge from p ((n) << 1.
E Assumptions in the Single-Neurol Linear Network Model
Assumptions é. Following Ahn al. (2023a), we make the following assumptions about &
(A1) l is a convex, even, 1-Lipschitz function that minimized at O.
(A2) l' is twice continuously differentiable near the origin with ("(0) = 1, which infers the existence of constant > 0 such that [e'(s) S |s| for all |s| S
(A3) We further assume "linear tail"' away from the minima, i.e., [e()s) c/2 for all |s| 2 and
[*))) 2 /2 for [s] S c.
Some concrete example of loss functions satisfying the above assumption include symmetrized logistic loss 1 log(1 exp(-2x) + 1 log(1 exp(2s)) and square root loss vI + 8 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 over (x. E R2 defined as L(x,y)
(xry where l satisfies Assumptions (Al), (A2), and (A3), if the initialization (zo, yo) satisfies:
v/ -3 = 2 "
vo c2. 1
yo
where y E 2] and C' 2 is constant, then for all hyper-parameter configuratinoo (n, p) such thar
+ VCm S min 2 1..0 T.. V7.
4VC n - 100 2)))
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 sunmarized here:
1. (Theorem 20) Until 2; = =(VTT), must have yt ((V and x++1 S It ((v
(Theorem 26) After Initial Phase and until ye |æ,). |æe| = O( +T) still holds. Meanwhile, Vt+1 yt min (ain((((pp))))))) )v, ((((gpp))e;;((pp)))))) [(p))} (i.e., Ys either drops by ((pp)) decays by ((pp)))
3. (Theorem 32) After Middle Phase, we always have (rel. Iycl = O(1p vn).
F.1 Basic Properties and Notations
C(ws) Recall the update rule of SAM in Equation 1: wi+1 + VLLW + + P[vz(())) ). By writing
Zt ya
as have:
substituting L(x, y) = E(ry). and utilizing the expressions of C(r, y) Equation we
[((w)
[xe] +
((x
x7 yt
sign
[1
23
yt
W. +
+
p
(11)
[
y/
x
1/2
2t
As yo 0, must have 7) < thus, these conditions automatically hold when 7 o(1) and p O(1).
20

---

where the second step uses the fact that e is a even function so (t) has the same sign with t. Define Further denoting e' ((xre + Pyr2 ()=
v=]+u? Then We + []]L( ))()) = It + pyp
2, =1gn(2(84))
Ut + p
))) by fy one can simplify the update rule Equation as follows:
Ft
n7 ((xx + -):22 + pr
[1 + p
W1+1
Vt
Ft + PV++
yt + p It Pit+
n
(12)
Assumption. Following footnote 1, we assume that Fz, y1 # 0 for all 4.
We are ready to give some basic properties of Equation 12. First, we claim that the sign of Ej the same as the product Tzzr- Formally, we have the following lemma:
Lemma 17 (Sign of Gradient in SAM). If + O, yt F O, then sign(x=) = sign(2; ; + p and
sign( = sign(y: ye + p=;25") In particular, if x4 + O, Ve # O, we must have sign (ft) =
Proof. Note that sign(=:) = sign(a)))t), so sign(ve = sign(a.). Similarly, sign( (tg) Thus, we have (x;) = sign(z: + 23:2 and sign(w)) = sign(v: + p ) which in turn shows that sign(() = sign (x; + p sign(ye + 2 )) = sign(r13)).
Using Lemma 17, the SAM update can be equivalently written as
Tz+1 Mt+1
xx ye
V1 Zt
paz (x] + v/² PUl (1r + 1//a
n
ml
(13)
We then claim the following property, which be intuitively interpolated as SAM always goes towards the minimizes (i.e., both axes) for each step, although sometimes may overshoot.
Lemma 18. Suppose that Tts St + 0. Then, the following basic properties hold:
ot > 0, then T++1 < Tt. If x4 < 0, then x++1 > Et.
yf > 0, then yt++1 < yt- If yt < 0, then Vt++1 >
Proof. We only show the case that x; > 0 as the other cases are similar. Recall the dynamics of xy;
Tt+1 xt nlt pzz
p xt. 176 Ve
Since x, # O, Lemma 17 implies that sign(E)) sign (zeye). and
Ft+1 Tt nC, p2ī xe np, V/
It - -isiin(ne. pssgu(zr::)))))) + v²) 1/=,,//,---- sigu((r:))))) sign(ve) (y:) Iyel
= Iz --141 plz? + 4/2 n |val S Iz.
where the last line uses the assumption that v, > 0.
F22 Initial Phase: x; Decreases Fast while y, Remains Large
As advertised in Theorem 16, we group all rounds until x, S ft/m as Initial Phase. Formally: Definition 19 (Initial Phase), Let t be the largest time such that > 1 all t f1. We call the iterations [0, f1] Initial Phase.
Theorem 20 (Main Conclusion of Initial Phase; SAM Case). For to
we have Yt > 2 / h////fffee for all t < min(to. t under the conditions of Theorem 16. Moreover, we have I++1 I Zt /mm for all S min 1n), which consequently infers t S to under the conditions of Theorem 16, i.e., min(to, [1) is just f1. This shows the first claim of Theorem 16.
Remark. This theorem can be intuitively explained as follows: At initialization, Zo. yo are both @ ((11) However, after & iterations, we have |xt = O(vT); meanwhile, ys is still of order Ω( 1/m) (much larger than ((v7)). Hence, [xr) and |ye| get widely separated in Initial Phase.
21

---

Proof. This theorem is direct consequence of Lemma 21 and Lemma 24.
Lemma 21. There exists to (1 "y2 rf 2}) such that lval 3V7/nforr i S to. n is sufficiently small s.I. np + vCm - 7/1, then yt 2 2 for t & min(41, to).
Proof. Let la be defined as follows:
a max { t
(:
(1- 1
8
n/²
>
1
(14)
Then, follows that tu = (1min 1.6p 1 ,n 2)). We first prove the following claim.
Claim 22. For all t S to. oi t//n.
Proof. We prove this claim by induction. We trivially have vn Cy/n 2 1///n. Now suppose that the conclusion holds up to some to, i... si 2 /// for all S t. We consider vi+1 - =2+1
v 18+1 (1 np41 pz 1) - x²) (n)()) - x;)
(1 27f/e= pz: + (Z, pzi 1² n/² ((,))) (vz -x7)
(a)
I (1 2mp2 172) (2) xi)
2 (1 21py." -1 112) (v? x!t)
W6 (1 VC
8 11 , 1/2 112) (1² 2²),
where (a) used & < (thanks to Assumption (Al)) and (b) used the induction hypothesis ył 2 Cr/n. This further implies that
1/) H+1 (1b - -x²) = 0(11 8
1+]
vit1 - rit1 2 (1-
8 1.
1/2
5
1/2
n²
VC"
vC
Thus, by the definition of to. must have v241
v2+ x241 2 tr/n, which proves the claim.
Next, we prove the second conclusion.
Claim 23, If up + VCTi S $ V7Tn, then ye 2 1/2 /7/ for S min(!o, 4
Proof. Still show by induction. Let yı 2 1/2 v 7/1
for some min(to. f1). Consider yo41-
By Definition 19, xe is positive for all t < f1. Thus, using Lemma 18, we have xt S 30 S yo VCr/n. Since Zt, Ve > O, we have sign( (f) = sign(x = sign(x:) > 0 and hence (13) gives
+ + 1/2 ///l - 1/ [/
We+1 y/
2 y: 7p 7 1
where we used the facts that |el < and (x? ツو ) 1/2 S min 1ya1-1). Hence, we get
Vi+1 2 Vt 1p v C-m 2 0
since np + VOm S 1/// Vr. By Claim 22, M/++1 2 implies y1+1 2 1 Vn/n as well.
Combining Claim 22 and Claim 23 finishes the proof of Lemma 21.
Lemma 24. For any t < min } we have
It-1 S 8t 1 //m.
In particular, if is sufficiently small s.t. 1VC n < to. then must have t & 2 - 1 < to
22

---

Proof. Since > 0 for all f Σ f1, we have sign(()) = sign (z:ye) = sign(V:) from Lemma 17 and
Lemma 21, so the SAM update (13) becomes
p(xt + =²²)=//2= 1/2xxx -7/fll vel-
11+1
Since 2, 2 C
have
7 2c, 2V
(at + pu + paizz Jzlli 2 20
which implies that f c/2 from Assumption (A3). Together with Lemma 21, have
It-1 It - -7/eal p( + 1/
1/2 xtt -n/e/ l
cl 7 = v $ xt - -/leal lyel - xt 22
11
VCr/n 1vC 17 -1, Since 20 S VO/nn have Iÿ < as long as t S fo. Thus, i it follows that 1 & t 1 < to. as desired.
Let i4
E3
Middle Phase: y Keeps Decreasing Until Smaller Than [20]]
Then we move on to the second claim of Theorem 16. We define all rounds until y/ < [ayg that are after Initial Phase Middle Phase. Formally, we have the following definition.
Definition 25 (Middle Phase). Let to be the first time that yt |zz].. We call (t1, t2] Middle Phase.
Before presenting the main conclusion of Middle Phase, we first define a threshold B that we measure whether variable is "small enough", which we will use throughout this section. Formally.
2 7 B A max (oV Cy
up np+ Vom
Theorem 26 (Main Conclusion of Middle Phase; SAM Case), Under the conditions of Theorem 16, lyel I CCr/n and |xa] S B throughout Middle Phase. Under the same conditions, we further have } for all t, < t2 showing the second claim of Theorem 16.
V641 $ yt minf hn yt: nP)
Remark. This theorem can understood as follows. Upon entering Middle Phase, |az] is bounded by O(np + /T). This then gets preserved throughout Middle Phase. Mcanwhile, in t2 iterations, yt drops rapidly such that U1,+1 S | = O(pp + /n). In other words, 2; and yr 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 hp + VOm /Ch/n, then during the Middle Phase, 1yo S yo S VCr/n-
Proof. When sign(w:+1) 1 sign(y)). then Lemma 18 implies that |B++ < /yel. Now consider the case sign(1:++) (y:). For simplicity, say y/ 0 and yr+1 < Since ye 0, have (ig)(E)) = sign()))))) = sign(2r) and hence Equation 13 gives
yo+1 y: 7/61ppa0 y2) 1/2 /pel Gl1el
2 - xel 2 -pp 778
Cj mp C
- m
"
which shows that [y++1l S up + VCmi. This proves the statement.
Lemma 28. Suppose that | < B max {
2. no no+ + vCmr.. Then we have |14411 B.
23

---

Proof. By symmetry, we only consider the case where > 0. If 24+1 2 0, then we have 0 - It-1 I zr due Lemma 24. If 2 0, then since 2, 0, it follows that
2++ = =ign(x+) | 7) (.
y²) - 1/2 | nllllll
2 |. - ny C//h 7p V Vmm,
(15)
where the last inequality is because 1yel - yo VC7/n. This concludes the proof.
Corollary 29. Note that, by definition of Initial Phase in Definition 19, already have | S B.
Hence, this lemma essentially says |0, S B = max 2 V 11 + vCn for all 1 < 1 S l2.
Lemma 30. If yt 2 [otb, then must have
Det1 S yr min (2 2000ur. yt,
c
2v2" 1p
Proof. Since Vv 2 /2te/. we have (r? + -
1/3 2 (2y?) = (V2y)) Consider two cases:
i) (3x + p392, )y: + p=r; 2 then |e] E 1s. 1] according to Assumption (A3). Moreover, since sign(() (C, = (z:))))))))))))))))))))) = sign (2), it follows that
1/2y yv nl Dt Vt+1 = - 7 p(x? + v
= 1 nlelpp(2 + v²,) )ul n) crll
< 1 - neep v /2yr)) /22)))
1
S yt
np
2 /2 y/ (((p))
(16)
ii) otherwise, follows from (A3) that
|e 2 2 (xx 0312, 1)) + pmrzz
As sign(x;) sign and sign sign(pxx12 21)) ), it follows that
14 p2j yt 2 (xx+ + pyl (8) + px(25)) p2r y
2
1
(pye V)2
2
we must have
Vet1 V/- 2 (0p 11)())) S V: 2 npp v = (1 Q((pp)))
(17)
where we used 2, 2 (V2y)) and thus yル 2 1/v2.
Combining item 16 and item 17, we obtain the desired concl fonclusion.
F.4 Final Phase: Both xa and y 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 l2- Theorem 32 (Main Conclusion of Final Phase; SAM Case). For all >> t2, we have 1rl1/96 S B where B max 2 V + VCmn).
Remark. we will see shortly, when entering Final Phase, both |ata and |yt are bounded by O(np + TT) This theorem essentially says that they can never exceed this bound in Final Phase. In other words, Final Phase, both I, and yt are oscillating around O.
Proof. We first show the following lemma, ruling out the possibility of lyal > B after Middle Phase:
24

---

n Lemma 33. If no + nB < B then 1y,+1l S B max <o Cr up vOm
Proof. The proof will be very similar to Lemma 28. Recall that by Definition 25, we must have yt, 0. If Mtz+1
well, then [ Mt++1 S [Tt++1 I B. Otherwise,
Ut++1 = sign Vt, 1b/++1
1e,1 7/ 1p(=2 + v 1/2 yt2l - n/c,, 1 2 -1p - nBB
where we used [ S
B. We proved our claim as mp + nB - B.
According to Lemma 28 Lemma 33, we have
2 CVV Cry up vCm].
Ft2 11 /1g++| S B = max
(18)
Hence, it only remains to do reduction, stated as follows.
Lemma 34. If mp + BB < B, then |2141 1o+11 & as long as |2t lyel S B.
Proof. Without loss of generality, let xy > for some l2 and show that |2441 S B which identical to the proof of Lemma 27. By symmetry, the same also holds for (xz]]-
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) E R2 defined as L(x, v)
((rr) where I satisfies Assumptions (AI), (A2), and (A3), if the initialization (æo Vo) satisfies:
ve = 2
c2 Tj
yo
2 x
>
v
where y 2] 2] and C 2 is constant, then for all hyper-parameter configuuatosss (n) p) such that
1 77 2' uo S min {1 f b. 2VOn O(min(n p
,n "), CT(1+ 2 16,
can characterize the initial and final phuses of the trajectory of USAM (defined in Equation 2) by the following two theorems, whose main conclusions are informally summarized below:
(Theorem 39) Until xt = (TTT), we must have y: St (VV7/n), and T++1 S It (2(YTT)-
2. (Theorem 44) Once (1 + pu? )u? $ 2/1, |xa 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 M defined in Equation 7 (the technical reason is sketched in subsubsection 3.1.3, i.c., SAM gradients non non-negligible when 3/s large, while USAM gradients vanish once |2t is small). However, remark that the main takeaway of Theorem 35 is to contrast SAM (which always attains v2. x 0) with USAM (which must stop once (1 + pyf yo? S ?/1).
G.1 Basic Properties and Notations
Recall the update rule of USAM in Equation 2: w2+1 <- We V/WL + /VL(w))). Still writing Wt as [ee ye] T we have
w++1
xx Vl
VL
Zt Vt
pl (x+ye)
1/1
It
C5 up + vOm) we only need to ensure that nB & VCry,, which can be done by 1 S min(cCTy, C 2, }.. As the RHS order 9(1), = O(1) again suffices.
Recall that max 2
"Smillar to Theorem 16, these conditions satisfied once o(1) and p O(1).
25

---

xt yt
yt + pl (1 ((01 + + (219)))) + pe (z13))=1)) [t + pll (213))))))
(19)
Due to the removal of normalization, l' are taken twice at different points in Equation 19, For simplicity, we denote & = l' (1133) and Ey ((x + (ln + pllxz)). The update rule can be rewritten as:
J++1
npl, (ye + plax)), M4+1 1/ ne, (Ir + pl(y)).
(20)
Similar to the SAM case, shall approximate & and É according the magnitude of ztye and (2a, + pe (1+ x.) We first have the following lemma similar Lemma 17:
Lemma (Sign of Gradient in USAM). If 2t + 0, y + O, then sign( = sign(a, + pl(y)) and sign(ye) = sign + pl(a)). In particular, if # O, yt # 0, then sign(() = sign(E)) = sign(
Proof. First of all, sign ( = sign(x:3)) according to Assumption (A1), Therefore,
sign(x: + pε sign( + psigu(a..) )be) sign(ra),
sign(y: + +(( sign(y: + + (ign(
sign(Ve),
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 Ft, Yt 2 0. Then, the following basic properties hold:
σt > 0, then T < Tt- If Jt < 0, then Fe+1 > Tt.
If y > 0, then yt+1 < yt- If yb < 0, then yt+1 > yt-
Proof. We only prove the statement for n, 0 as the proof is similar for other cases. Recall the dynamics of vt:
Te+1 = mpl (ye + vl,x)).
Since 2, O, Lemma 36 implies that (ig) sign sign(r,3t). and
It+1 Tt n (ye + pl(x)e)
It 77 sign((r 1441 (ye + (s)g((2r0)))()))))
It - (lvi (()x))) I 21,
where the last line uses the assumption that x; > 0.
G.2 Initial Phase: He Decreases Fast while y: 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 t1 be the largest time such that 24 > for all S t1. We denote by "Initial Phase" the iterations t1).
Theorem 39 (Main Conclusion of Initial Phase; USAM Case). For to = 2m 2)). we have yt 2 ////n for for all t I min(to. ti} under the conditions of Theorem 35. Moreover, we have I1+1 S 2t 1 WCmm for all t < min(to, }1 which consequently infers t < to under the conditions of Theorem 35, i.e. min(to, {13 just t1. This shows the first claim of Thhoorem 35.
Proof. This theorem is combination of Lemma 40 and Lemma 43.
Similar to Lemma 21, ye in USAM also cannot be too small the first few iterations:
Lemma 40. There exists to = ((min(m -1p-,, "..." -2}) Assuming 1. then yı / t min
such that for S to. we have 136 2 4/7
26

---

Proof. The proof idea follows from Lemma 21. Let to be defined as follows:
to max { t (1 (1-1)² 2mp)
1
(21)
4
Then, it follows that to (min 10 p-= ,,, 2)) We first show that ył cannot be too small until to:
Claim 41. For all E los vi 2 4//n.
Proof. Prove by induction. Initially, V6 = Oy/n 2 ///m. Then consider some < to such that ve 2 77/n for all l' < t. By Equation 20, we have
Vi+1 r²t11 -V: n1, (ot + pl:v.)) (Ix- npii (yr + plaxze)
(vž x;)) (n²r)² (1s + +E (ye + pl)xz))
anen (ye + x², - 2y: . np,i (0r + pl(ve)
(0? - =²)) + (x) - vřt (n 2
=i?) + 20p2' pÊ,(2? - v²²)
(1 (1- (net)² + (nl))((p)))) 2 - -2ne pf: ) =x?)-
(vž
Using Assumption (A1), we know that | and | S 1, giving
v 2,41 2 (1 122 2np) (vi x/) 2 2 (1 -1- 1+] '1/000 -20).
By definition to and condition that t < to, we have v221 2 v2²11 - -12+1 2 2/
Claim 42. Suppose that VOm S V7/n, ns JVC. Then y: 2 1 \/ for all S to-
Proof. We still consider the maximum single-step difference in yt- Suppose that for some < min(to, h.) ye 2 2 7//n for all t' S t. According to Equation 20 and Assumption (A1),
V++1 = y/ n24 (xz + pk 2 yt - (({; + py)) 2 (1 πp)nr - -1 Cr/n.
where the last inequality used Lemma 37 to conclude that S xo yo VCy/n. Hence, as the first term non-nee negativt and vOm S Vn/n. we must have ye+1 O, which implies
gt+1 2 1 V/nn according to Claim 41.
Putting these two claims together gives our conclusion. Note that the condition v) S V in Claim 42 can be inferred from the assumptions 1) < and C 2 1 made Theorem 35.
After showing that Ye never becomes too small, we are ready to show that Fe decreases fast enough. Lemma 43. For S min(lo, J. we have
It-1 S
2 vm.
In particular, if 1 is sufficiently small s.t. 2 s to, then we must have t 22 to
Proof. Since x, > 0 for S l1 we have (e) sign( sign(v)) from Lemma 36 Lemma 40, the USAM update (20) simplifies follows according to Assumption (Al) and Lemma 40
1 -
F++1 xt -/lll (lval + el s x, n/ya S xr
Let tj
V07a \/79
2/C On then ti - to and thus Izς < 2 i.., t < to must holds.
27

---

G.3 Final Phase: yI 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
|zyl| S VCri and lvol= S VCT/n. Moreover, once we have sy(1 + put)ut =2 for some t 2 1 and 6 > 0, we must have [2++1| < exp(- for all E. 2 t. This consequently infers that v which is defined as lim inf,- Vi. satisfies y 2 (1 4C +
pC-)² w)? As is constant
independent of 1) and p 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 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 |aaa S VCm, laal S VCylm Assuming no S 1. then ret1ll S VCy and (y--1 S /Cy/n as well. Furthermore, | S VCyn and luel S /Cr/n hold for all t > t1-
Proof. Suppose that 2 2 0 without loss of generality. The case where 2z+1 2 0 is trivial by Lemma 37. Otherwise, using Lemma 36 and Equation 20, we can write
Fe+1 It n7 (Ve + pl, It 17 eil (lual 11))
By Assumption (AI). (e] and |& are bounded by 1. By the condition that mp S 1,
x(+1 2 (1 np)z, nb 2 C'm,
where we used |yrl S sqrtOr/nn Similarly, for vi+1, only considering the case where 2 0 and Ve+1 I suffices. We have the following by symmetry
yt+1 = yt n-, (2r + p( 2 y - u<(x x pyc) = (1-mp)y) 7/ 2 -n On 2 -V -Cy/n
where we used | S VChm-
The second part of the conclusion is done by induction. According to the definition of Initial Phase, we have [1t S / As C 1 1 we consequently have |aza S /Cym for all t > & from the first part of the conclusion. Regarding [wel, recall Lemma 40 infers y/t > 0 for all L S 4 + and
Lemma 37 infers the monotonicity of ye, we have 4t S yo = v CC//n Hence, 1yel S VCh/n n for all t >> & as well.
Before showing the ultimate conclusion Lemma 47, we first show the following single-step lemma:
Lemma 46. Suppose that 17(1 + pyf y9i < 2 and |arl S vorn some Define e, n/1 put ) (then we must have 6 E 2)). Then we have
|x++||= S [ exp ( (-min/ ( (1 pCrm)ce,
2 2- fz }).
(22)
8
Proof. Without loss of generality, assume that x2 > 0. From Equation 20, can write:
le+1 It. - 11/ (ye + p(( x. - 1/ Cl (luel + p1()) where we used (fg) = sign ( sign( sig) (Lemma 36). By Assumption (A2),
Ft41 Zt 17 (x+ por (ve pL ) (leel + p1
St 17 (1 + / el@illvel) (1vrl + pl(a)are (laal + paills 1) 2 zt 17 + p (19a + poI (e + paflval)
(1 111 + + + 2
2 (1 (1 (1
+
2 - -(1 - + p1+pCmm))))
(23)
28

---

where we used | - yn For the other direction, we have the following by Assumption (A3):
I(+1 S Zt (2x pl, pl)
(19c +
21
(It + p pleelleele ((vel + pl@g)aze)
131ll + P 2i
It
Wel
2
2
(oe + pi (1vel + pluelaz))
+ p 1t
S It 11
1yv
8
2
(1 pu?) pat)
1
24
Vi
[
!
(1 + puf yo 22
2
It+
(24)
8
where we used (1 + pxt) 2 and (1 + / 2 1. Equation 22 follows from Equation 23 and Equation 24.
Lemma 47. Ler t be such that i) n7(1 + put yo? 2 where e E (0, ( is a constant, and ii) [xd S VOm. Then we have the following conclusion on lim in yi, denoted by y in short:
you = lim inf vt 2 4C= pCg)
12200
While the e-1 looks enormous, it is constant independent of and in other words, we are allowed as close to zero want. As we only consider the dependency on r and p. we can abbreviate this conclusion as yoo 2 (1 - O(1² + p²)) as we claim in the main text.
Proof. In analog to Equation 23, we derive the following for yr:
M 2 1 pu)) +1) pat) vi 2 11 (1-
2m 1-2n(1+
pO1//n + vi:
where the second inequality uses Lemma 45. Let de = y2 v2. then we have
d++1 de + 2n/(1 + pCT/n) + pCmm)= v S de + 4n 2 pC 22 where we used the assumption that pChn I 1 and the fact that y2 is monotonic (thus yf S v2) According to Equation 22, we have |x++a|= S |x.] exp(eepp(()))) for all t 2 f2. Hence, we have
20
do lim sup d, s de + 1-000
4n (1 (C 4n ' pC 2 Σ
where we used d, = 0 (by definition) and the sum of geometric series. Plugging back xf < Cym.
22c = lim of 2 de - +
1+000
as claimed.
H Omitted Proof of USAM Over General PL functions
Theorem 48 (Formal Version of Theorem 6). For any μ-PL and B-smooth loss function for any learning rate 7 < 1/3 and < 1/2,for any initialization Wy, the following holds for USAM:
2,32 f
(1 (1-2 2nnn(1 /13)
(22 113 (2211 (1 ((--8)))
nB))) -1/2
||w woll S 1)(1 + 8p)
V ((wo)
Lr
41 2 0,
2
where L the short-hand notation for min, L(w).
We first state the following useful result by Andriushchenko and Flammarion (2022, Theorem 10).
29

---

Lemma 49 (Descent Lemma of USAM over Smooth and PL Losses), For any 8-smooth and p-PL loss function L. for any learning rate dj < 1/8 and p 1/3, the following holds for USAM:
[(wi) - (1-2um(1 p8) (I n (1 (1--83)) po1)) (C(wo) - L*), A2 0.
Proof of Theorem 48. We follow the convention in Karimi et al. (2016): Let X' the set of global minima and Vp be the projection of onto the solution X From B-smoothness, follows that
[V = [z) V (x,,)) - 3 2 xpll
Vr.
Now since L 8-smooth and prP Theorem 2 from (Karimi et al., 2016) implies that the quadratic growth condition holds, i.e.,
(c(x) - L*) 2 1 xpl??
Vr.
Thus, follows that
202 |VL(())] (C(r)) -(');
Vo.
1
Moreover, from 3-smoothness, we have
VV + PV & [V + + +
= (1+ ())) Vz.
Thus, by the update rule of USAM (2), it follows that
|| - woll S
VL(u: + +VL((w)))
$ (1 + 3p) Σ VL(w))]
2,32
2n(1 + 3p))
(E(w:) L*)
V
P
2,32 t-1
n1 + 8p)
V /((w) - L-.
1=0
Now we only to invoke the USAM descent lemma stated before, i.e., Lemma 49, giving
[-1
1 - 2um(11 - (B3) (1- 1/3 (1 - -03)) 1/2
[(w;)
L <
Σ
(1-n88)
v ((wn) L
2
0
2=0
2jtn(1 p33) (n,
1/3 (1-p8))
2 (1 - ns))) 1/2
((wo) L-.
Putting the last two inequalities together then give our conclusion.
30


Reviews:
[{"review_id": "tFTA0StFw3", "reviewer": "Reviewer_ymqN", "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.", "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?\n\nQ2. Regarding the function $\\mathcal{L}$ in Sec 3.1. Does the $\\times$ imply multiplication of $x,y$? The function $\\mathcal{L}$ is $L$-smooth?\n\nQ3. 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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is in general well-written and accessible despite being a theoretical work.\n- The included examples help the understanding of the ideas.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments! Here are our responses:\n\n---\n\n### Comparison Between Our Paper and Wen et al. (2023); Compagnoni et al. (2023)\n\nThanks for mentioning these two seminal papers on SAM! We will definitely discuss more in our revision. Here is a draft of our discussion:\n\n1. Wen et al. (2023) characterized the behavior of SAM when *1)* initialized near the minima manifold, and *2)* the step size $\\eta$ is small enough, by showing that SAM roughly follows a Riemannian gradient flow on $\\lambda_{\\max}(\\nabla^2 \\mathcal L(w))$ (i.e., their \"sharpness\" measure maximum eigenvalue of Hessian) along the minima manifold. **This result is closely related to ours in Section 3** as we also study the behavior of SAM when close to a global minimum.  \nHowever, there are also several differences:\n    * Admittedly, they consider more general loss functions than ours. However, as a side-effect, they require \"small enough\" $\\eta$ which they do not specify. They also require the initialization to be sufficiently close to the minima manifold.\n    * In contrast, we allow any $\\eta=o(1)$ and $\\rho=\\mathcal O(1)$ (we include detail requirements in Theorem 16) and focus on a specific loss function. **This gives a much more accurate characterization of the algorithmic behavior of SAM**, even when the initialization is far from the origin.\n\n2. Compagnoni et al. (2023) consider the continuous-time behavior of SGD, SAM, and USAM. As we mentioned in the related work section, this is the only work (to our knowledge) that distinguishes USAM from SAM. **Indeed, they also show different behaviors of SAM and USAM** (i.e., USAM is attracted to local minima while SAM aims at global ones). However, as they are considering continuous-time variants of algorithms while we consider discrete (original) versions, **our results directly apply to the SAM in practice and the USAM in theory**.\n\n---\n\n### Does SAM Diverge in Section 2.2 (Non-Convex but Smooth Loss Functions)?\n\nYes, it is possible for SAM to diverge when facing non-convex loss functions, similar to GD which diverges when $\\eta$ is too large.\n\nHowever, we would like to emphasize that the purpose of Section 2 is to illustrate that in the early stage of training, *1)* SAM is **more robust than USAM** as USAM diverges even with a rather small $\\eta$ or $\\rho$, and *2)* SAM has **comparable robustness to GD** as an $\\eta$ ensuring convergence of GD can also make SAM converge.\n\nIn other words, one of our main messages in that section is **\"SAM is better than USAM\", not \"SAM is better than GD\"**. The superiority of SAM over GD is left for Section 3, where SAM exhibits the ability to drift along the minima manifold (which is beneficial as suggested by various previous works, both theoretically and empirically) while GD easily gets trapped at \"bad\" minima.\n\n---\n\n### Function $\\mathcal L$ in Section 3.1\n\nYes, $\\times$ means the multiplication. When restricting $\\mathcal L$ in a compact region, it becomes $L$-smooth.\n\n---\n\n### \"Small Enough\" Step Sizes in (Wen et al., 2023)\n\nThanks for referring to this subtle difference. Indeed, the conclusion by Wen et al. (2023) requires a \"small enough\" step size for the dynamics of SAM to track a Riemannian gradient flow induced by $\\lambda_{\\max}(\\nabla^2 \\mathcal L(w))$ and thus decrease the \"sharpness\".\n\nNevertheless, from our Theorem 4 (behavior of SAM on single-neuron networks) and Figure 5 (behavior of SAM on ResNet18 over CIFAR-10), **we conjecture that when the step sizes are not \"small enough\", SAM can still drift along the manifold** and thus contrasts USAM's behavior characterized in Theorem 6. Proving or disproving this would be an interesting future direction.\n\n---\n\nThank you once again for taking the time to review our paper. Given that you have not raised any concerns regarding the correctness or novelty of our work, if you believe that we have addressed your points, we would appreciate it if you would consider raising your score."}, {"review_id": "8XbxDOeaG7", "reviewer": "Reviewer_pAM7", "summary": "**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.", "questions": "1. 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?", "limitations": "None", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 1, "strengths": "**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.", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments! Here are our responses:\n\n---\n\n### USAM is Not Practical\n\nWe agree that USAM is by no means practical. However, the reason for comparison is solely to understand SAM better and doesn't have to do with USAM's impracticality. In particular, we are trying to **understand the success of SAM in practice rigorously and theoretically**: Undoubtably, SAM exhibits ground-breaking success in practice (and likely thus amassing ~700 citations in ~2 years); thus, better understanding SAM is critical for designing better algorithms in the future.\n\nIndeed, extensive effort has been made study SAM in the theoretical ML literature. However, we notice that **most of such works assume that the normalization of SAM can be omitted \"for simplicity\"** (e.g., Andriushchenko and Flammarion (2022); Behdin and Mazumder (2023); Agarwala and Dauphin (2023); Kim et al. (2023)). However, as far as we know, **the only rationale behind such a simplification is just one experiment** on CIFAR10 conducted by Andriushchenko and Flammarion (2022).\n\nIn our paper, we take a step back and **challenge this common \"simplification\", asking whether normalization plays a critical role in SAM**. To make our results scientifically rigorous, our paper considers the variant of SAM that **only** removes normalization, namely USAM. This provides a precise and scientific ablation study. Hence, **we did not consider USAM for its practicality** but **for providing a precise and scientific understanding of SAM**. By considering USAM, we successfully and rigorously identify (one of) the reasons behind SAM's success, namely the normalization step.\n\nFinally, we would like to remark that **previous papers do not yet achieve such an accurate identification**. For example, as you mentioned, Wen et al. (2023) show that SAM can drift along the minima manifold for a \"flatter\" minimum. However, the crucial role played by the normalization step is beyond the scope of their paper. In our work, by comparing USAM and SAM, we identify its crucial role.\n\n> ### References\n> * Maksym Andriushchenko and Nicolas Flammarion. Towards Understanding Sharpness-Aware Minimization. *Proceedings of the 39th International Conference on Machine Learning*, PMLR 162:639-668, 2022.\n> * Kayhan Behdin and Rahul Mazumder. On Statistical Properties of Sharpness-Aware Minimization: Provable Guarantees. *arXiv preprint arXiv:2302.11836*, 2023.\n> * Atish Agarwala, Yann Dauphin. SAM operates far from home: eigenvalue regularization as a dynamical phenomenon. *Proceedings of the 40th International Conference on Machine Learning*, PMLR 202:152-168, 2023.\n> * Hoki Kim, Jinseong Park, Yujin Choi, and Jaewook Lee. Stability Analysis of Sharpness-Aware Minimization. *arXiv preprint arXiv:2301.06308*, 2023.\n> * Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How Sharpness-Aware Minimization Minimizes Sharpness? In *The Eleventh International Conference on Learning Representations*, 2023.\n\n---\n\n### Non-Toy Experiments\n\nThanks for mentioning this! We want to point out that **our experiments in Figure 5 are done on practical ResNet18 architecture on the CIFAR-10 dataset**. In Figure 5, we initialized SGD, SAM, and USAM at the same checkpoint, which is a poor global minimum of the loss function (which we call \"adversarial checkpoint\"; the same experiment is also conducted from a more realistic checkpoint yielded by full-batch GD, which is included in the attached PDF file).\n\n**From Figure 5, we observe that USAM is indeed harder to tune in real-world settings**: A conservative configuration of $(\\eta,\\rho)$ makes it unable to drift along the manifold (and is thus trapped at the bad minimum); an aggressive configuration drives it away from the manifold and results in worse performance. In sharp contrast, SAM keeps drifting along the manifold towards better-generalizing minima and achieves much better performance.\n\n---\n\n### Role of $\\rho$ in Theorem 4\n\nThanks for raising such a good question! In fact, $\\rho$ also determines the convergence speed in the \"Middle Phase\". In Footnote 2, we mention that the dynamics of $y_t$ can be upper bounded by either $y_{t+1}\\lesssim (1-\\eta \\rho^2)y_t$ or $y_{t+1}\\lesssim y_t-\\eta \\rho$. Thus, a **larger $\\rho$ may drive $y_t$ faster towards the origin**. We will include this as a remark after Theorem 4 to avoid possible confusion.\n\nAnd as a final remark, we would like to mention that **we made no effort to optimize the convergence rate bound**. The objective of Theorem 4 is to show that the trajectory will approach a neighborhood of the origin \"pretty fast\". It remains open how many iterations are needed for SAM to find such a neighborhood.\n\n---\n\n### USAM versus SAM in Figure 7 when $\\eta=0.1$\n\nIndeed, when picking $\\eta=0.1$ in Figure 7, the performance of USAM and SAM are similar. Meanwhile, for other choices of $\\eta$ (e.g., $0.025$), USAM performs much worse while SAM remains good.\n\nHowever, this tendency is not valid in general. As presented in Figure 8 of Appendix B, when setting the algorithmic parameters $(\\eta,\\rho)$ too aggressively, USAM diverges (the $\\rho=0.5$ curve becomes invisible in the plot because it diverges). Thus, **tuning the parameters of $\\eta$ and $\\rho$ is indeed a dilemma** in USAM, just as in GD.\n\n---\n\n### Dataset in Figure 1\n\nIn Figure 1, we adopt the over-parametrized matrix sensing problem setup introduced by Li et al. (2018). The exact setup can be found in Appendix A.\n\n> ### References\n> * 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.\n\n---\n\nWe thank you once again for such a valuable review of our paper. We are more than happy to answer any further questions!\n"}, {"review_id": "6aygjSP5V7", "reviewer": "Reviewer_chBL", "summary": "- 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.", "questions": "- 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\n", "limitations": "The paper has not provided a detailed discussion on limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- 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)", "weaknesses": "Please refer to questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments! Here are our responses:\n\n---\n\n### Should the Choice of $\\rho$ Be Different in SAM and USAM?\n\nYes, the optimal tuning of algorithmic parameters $(\\eta,\\rho)$ likely differs between SAM and USAM. However, we would like to emphasize that **we are not comparing the performance of SAM and USAM** when fixing the same $(\\eta,\\rho)$. Instead, **we are comparing the robustness of SAM and USAM** to parameter tuning (i.e., $\\eta$ and $\\rho$).\n\nIn other words, our experiments like Figures 7 and 8 are **not** for \"fixing a same $(\\rho,\\eta)$ and comparing the behavior of SAM and USAM\"; instead, we consider \"**how much does change $\\eta$ or $\\rho$ affects their algorithmic behaviors**\". Specifically,\n1. In Figure 7, we fix $\\rho$ and use various different $\\eta$'s. We discover that fixing the same $\\rho$, SAM with different $\\eta$'s finds the same optimum; however, USAM with different $\\eta$ behaves pretty differently. Thus, **Figure 7 illustrates that USAM is less robust to the tuning of $\\eta$**.\n\n3. In Figure 8, we instead fix $\\eta$ and try different $\\rho$'s. It turns out that SAM is still robust to different $\\rho$'s while USAM is not. Again, we are not comparing the performance of SAM and USAM with the same $(\\eta,\\rho)$; instead, we identify the effect of varying $\\rho$ on the algorithmic behaviors. **Our conclusion from Figure 7 generalizes to $\\rho$ as well**.\n\nHence, putting Figures 7 and 8 together, we conclude that **SAM is robust to different configurations of $(\\eta,\\rho)$ while USAM is robust to neither of them**. We will make our conclusion more transparent in the final version.\n\n------\n\n### More Practical Setting than \"Adversarial Checkpoint\"?\n\nThanks for noticing this! We used the \"Adversarial Checkpoint\" by Damian et al. (2021) as a \"poor global minimum\" which has 100% training accuracy but only 48% test accuracy. This initialization **helps highlight different algorithm's behavior when they are close to \"bad minimum\"**. If we use a \"good minimum\" instead, we would not see the difference in behaviors as clear. \n\nFor better understanding, we also conducted 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. Such initialization is thus more realistic compared to the \"adversarial checkpoint\".\n\nThe training trajectory initialized from this \"full-batch checkpoint\" is included in the attachment. In short, **USAM still gets stuck, while SAM attains better test accuracy via drifting along the minima manifold**. We will include this new experimental result in our final version.\n\n> ### References\n> * Alex Damian, Tengyu Ma, and Jason D. Lee. Label Noise SGD Provably Prefers Flat Global Minimizers. In *Advances in Neural Information Processing Systems*, 2021.\n\n---\n\n### Better Learning Rate Scheduling?\n\nIn practice, there are various ways of tuning learning rates, e.g., linear or cosine scheduling. However, such tunings are a bit tricky to analyze in theory -- none of the previous papers has succeeded. \n\nRegarding the \"no perturbation in the early stage\" variant of USAM, we believe it probably avoids the divergence issue. However, this only reflects our results in Section 2. In Section 3 (or more specifically, Figure 5), we illustrate that **\"USAM without perturbation in its early phase\" cannot guarantee a good performance, either**.\n\nFor more context, in Figure 5, we initialize all algorithms from the same \"poor global minimum\" derived by SGD and study which algorithms can escape it. Essentially, to arrive at this \"poor global minimum\", all algorithms are not adopting perturbation in their early stages. However, as suggested by our plot, USAM loses the ability to find better global minima, matching our conclusions in Theorems 5 and 6. This means that **normalization is not only important in the initial stages, but also critical in the final phase for better minima**. Consequently, removing perturbation in the early phase of USAM only resolves the divergence issue but still cannot ensure better performance.\n\n---\n\n### Does the 'USAM versus SGD' Result in [3] Contradit Ours?\n\nThanks for mentioning [3]! In [3], the authors show that while USAM (which they call SAM) exhibits a similar EoS phenomenon as GD, **USAM has a smaller EoS threshold**: `We will show that SAM also induces an EOS stabilization effect, but at a smaller eigenvalue than GD` (from their Section 2.2.2). Hence, their result also precisely shows that USAM is less stable than GD: When the learning rate $\\eta$ exceeds the EoS threshold of USAM but is below the one of GD (i.e., the condition $\\eta \\in (2/(\\beta+\\rho\\beta^2),2/\\beta]$ in our Theorem 1), USAM diverges, but GD does not. As a result, **their observation essentially supports our claim**.\n\nMeanwhile, we would like to mention that our early-stage analyses additionally reveal that SAM does not diverge for the range of $\\eta$ where USAM diverges and is thus more robust to the choice of $\\eta$, which is beyond the scope of [3] as [3] only studied USAM instead of SAM.\n\n---\n\nThank you once again for taking the time to review our paper. Given that you have not raised any concerns regarding the correctness or novelty of our work, if you believe that we have addressed your points, we would appreciate it if you would consider raising your score.\n"}, {"review_id": "vCDNUvmnXi", "reviewer": "Reviewer_FoGd", "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. ", "questions": "- Can the takeaways from this work be generalized to the benefits of other optimizers that use normalization (such as normalized gradient descent)?\n- Are there any advantages for USAM over SAM?", "limitations": "No, the authors do not discuss limitations and potential negative societal impact.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper thoroughly investigates the role of normalization across many different settings, and convincingly show the benefit of the normalization term. \n- The authors provide clear key takeaways for why normalization helps, with both empirical and theoretical support.\n- The paper makes a solid contribution to our understanding of why normalization is beneficial\n- 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. \n- 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments! Here are our responses:\n\n---\n\n### Is the Main Takeaway \"Using SAM instead of USAM\"?\n\nThe takeaway is not this one if interpreted literally, but slightly more nuanced. Let us summarize the **two main takeaways** of our paper, one for the theoretical ML community and one for the empirical ML community:\n\n1. For the theoretical ML community, our main takeaway is **normalization plays a critical role and cannot be neglected when analyzing SAM** -- when investigating the success of SAM (which is important for future algorithm designs), normalization plays a non-negligible role; removing normalization can drastically change the algorithmic behaviors.  \n\n2. For the empirical ML community, the main takeaway is **the practicality of SAM is (in part) backed up by the normalization**. Previously, empirical researchers successfully applied SAM in various applications, and theoretical researchers characterized SAM's ability to perform well. However, the reason behind such a success is still unclear until our work. Our work suggests that **when designing other algorithms in the future, whether using normalization should be taken seriously**.\n\n---\n\n### Advantages of USAM over SAM?\n\nThanks for mentioning this. The only advantage of USAM over SAM that we are aware of is its simplicity in analyses -- and this is probably the main reason why various previous theoretical works focused on USAM instead of SAM (e.g., Andriushchenko and Flammarion (2022); Behdin and Mazumder (2023); Agarwala and Dauphin (2023); Kim et al. (2023)).\n\nStill, we would like to mention that **the main reason for considering USAM is not its practicability or any other advantage** (more details can be found in our first response to Reviewer a57m or Reviewer pAM7). For context, we aim to understand the considerable success of SAM in practice. However, we noted that existing theoretical works mostly ignore the normalization in SAM \"for simplicity\". Our work challenges this \"simplification\" and **justifies that normalization plays a critical role**. We only remove the normalization in SAM and keep all other parts unchanged for a scientifically rigorous conclusion. This gives USAM -- hence, **the reason for considering USAM is for a precise ablation study**, not its practicability (or any other advantage).\n\n\n> ### References\n> * Maksym Andriushchenko and Nicolas Flammarion. Towards Understanding Sharpness-Aware Minimization. *Proceedings of the 39th International Conference on Machine Learning*, PMLR 162:639-668, 2022.\n> * Kayhan Behdin and Rahul Mazumder. On Statistical Properties of Sharpness-Aware Minimization: Provable Guarantees. *arXiv preprint arXiv:2302.11836*, 2023.\n> * Atish Agarwala, Yann Dauphin. SAM operates far from home: eigenvalue regularization as a dynamical phenomenon. *Proceedings of the 40th International Conference on Machine Learning*, PMLR 202:152-168, 2023.\n> * Hoki Kim, Jinseong Park, Yujin Choi, and Jaewook Lee. Stability Analysis of Sharpness-Aware Minimization. *arXiv preprint arXiv:2301.06308*, 2023.\n\n---\n\n### More Experiments\n\nThanks for mentioning this! We conducted another experiment similar to Figure 5, where we used a more realistic \"full-batch checkpoint\" instead of an \"adversarial checkpoint\" (see our second response to Reviewer chBL for more details); indeed, our conclusion still holds. The plot is contained in our attachment.\n\nIt would definitely be great to conduct large-scale experiments on larger datasets like CIFAR-100 and larger neural networks, as it can further verify our theoretical conclusions. Although we won't be able to share the full results this time, we will make attempts to add some experimental results in our final version.\n\nWe also would like to highlight that **our paper is a theoretical paper aiming at identifying the reason behind SAM's success**; and, we did succeed in justifying the crucial role of normalization by *i)* rigorously comparing SAM's behavior with the un-normalized variant USAM in theory, and *ii)* carefully designing experiments inspired by our theory to support our claim via real-world models (ResNet18 on CIFAR-10).\n\n---\n\nThank you once again for taking the time to review our paper. Given that you have not raised any concerns regarding the correctness or novelty of our work, if you believe that we have addressed your points, we would appreciate it if you would consider raising your score.\n"}, {"review_id": "qXDBEQ6eTp", "reviewer": "Reviewer_a57m", "summary": "This 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.", "questions": "As 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", "limitations": "I don’t see potential negative societal impact that the authors should have addressed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "The 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful comments! Here are our responses:\n\n---\n### Why Compare SAM to USAM?\n\nWe agree that USAM is by no means practical. However, the reason for comparison is solely to understand SAM better and doesn't have to do with USAM's impracticality. In particular, we are trying to **understand the success of SAM in practice rigorously and theoretically**: It is undoubtedly that SAM exhibits a ground-breaking success in practice (getting ~700 citations in ~2 years); thus, better understanding the success of SAM is critical for designing better algorithms in the future.\n\nIndeed, extensive effort has been made to answer this question in the theoretical ML literature. However, we notice that **most of such works assume that the normalization of SAM can be omitted \"for simplicity\"** (e.g., Andriushchenko and Flammarion (2022); Behdin and Mazumder (2023); Agarwala and Dauphin (2023); Kim et al. (2023)). However, as far as we know, **the only rationale behind such a simplification is just one experiment** on CIFAR10 conducted by Andriushchenko and Flammarion (2022).\n\nIn our paper, we take a step back and **challenge this common \"simplification\", asking whether normalization plays a critical role in SAM**. To make our results scientifically rigorous, our paper considers the variant of SAM by **only** removing the normalization, namely USAM. This provides a precise and scientific ablation study. Hence, **we did not consider USAM for its practicality** but **for providing a precise and scientific understanding of SAM**. By considering USAM, we successfully and rigorously identify (one of) the reasons behind SAM's success, namely the normalization step.\n\nFinally, we would like to remark that **previous papers do not yet achieve such an accurate identification**. For example, as you mentioned, Wen et al. (2023) show that SAM can drift along the minima manifold for a \"flatter\" minimum. However, the crucial role played by the normalization step is beyond the scope of their paper. In our work, by comparing USAM and SAM, we identify its crucial role.\n\n---\n### More Experiments on Larger Datasets\n\nThanks for mentioning this. It would be great to conduct large-scale experiments on larger datasets like CIFAR-100 and larger neural networks, as it can further verify our theoretical conclusions. Although we won't be able to share the full results this time, we will attempt to add some experimental results in our final version.\n\nWe also would like to highlight that **our paper is a theoretical paper aiming at identifying the reason behind SAM's success**; and we did succeed in justifying the crucial role of normalization by *i)* rigorously comparing SAM's behavior with the un-normalized variant USAM in theory, and *ii)* carefully designing experiments inspired by our theory to support our claim via real-world models (ResNet18 on CIFAR-10).\n\n---\n### Should the Choice of $\\rho$ Be Different between SAM and USAM?\n\nYes, the optimal tuning of algorithmic parameters $(\\eta,\\rho)$ likely differs between SAM and USAM. However, we emphasize that **we are not comparing the performance of SAM and USAM** when fixing the same $(\\eta,\\rho)$. Instead, **we are comparing the robustness of SAM and USAM** to parameter tuning (i.e., $\\eta$ and $\\rho$).\n\nIn other words, our experiments like Figures 7 and 8 are **not** for \"fixing a same $(\\rho,\\eta)$ and comparing the behavior of SAM and USAM\"; instead, we consider \"**how much does change $\\eta$ or $\\rho$ affects their algorithmic behaviors**\". Specifically,\n1. In Figure 7, we fix $\\rho$ and use various different $\\eta$'s. We discover that fixing the same $\\rho$, SAM with different $\\eta$'s finds the same optimum; however, USAM with different $\\eta$ behaves pretty differently. Thus, **Figure 7 illustrates that USAM is less robust to the tuning of $\\eta$**.\n2. In Figure 8, we instead fix $\\eta$ and try different $\\rho$'s. It turns out that SAM is still robust to different $\\rho$'s while USAM is not. Again, we are not comparing the performance of SAM and USAM with the same $(\\eta,\\rho)$; instead, we identify the effect of varying $\\rho$ on the algorithmic behaviors. **Our conclusion from Figure 7 generalizes to $\\rho$ as well**.\n\nHence, putting Figures 7 and 8 together, we conclude that **SAM is robust to different configurations of $(\\eta,\\rho)$ while USAM is robust to neither of them**. We will make our conclusion more transparent in the final version.\n\n---\n### Does Removing Perturbation in the Early Phase of USAM Resolves Divergence?\n\nWe believe it is possible as removing normalization in the early stage essentially makes USAM same as GD. This is also our motivation for Section 3, where we illustrate that **\"USAM without perturbation in its early phase\" cannot guarantee a good performance, either**.\n\nAs our Figure 5 suggests, USAM loses the ability to find a better global minimum, matching our conclusions in Theorems 5 and 6. This indicates that **normalization is not only important in the initial stages, but also critical in the final phase in order to escape bad minima**. Consequently, removing perturbation in the early phase of USAM only resolves the divergence issue but still cannot ensure better performance.\n\n ----\n### USAM Diverges with \"Small\" or \"Large\" $\\rho$?\n\nThanks for mentioning this! We indeed meant \"small $\\rho$\". According to Theorem 15, USAM diverges when $\\rho\\approx 15\\eta=o(1)$; however, by Theorem 12, SAM doesn't diverge even when $\\rho=\\Omega(1)$. Hence, **USAM starts to diverge with a \"small\" $\\rho$ compared to SAM**. We will change the titles of Appendices D.1 and D.2 to \"SAM Converges with a Large $\\rho$\" and \"USAM Diverges from a Small $\\rho$\", respectively.\n\n---\nThank you once again for taking the time to review our paper. Given that you have not raised any concerns regarding the correctness or novelty of our work, if you believe that we have addressed your points, we would appreciate it if you would consider raising your score."}]

Author rebuttal:
Dear Reviewers,

We thank all the reviewers for the valuable feedback. Specifically, thank you for appreciating that our motivation is "**clear and interesting**" (Reviewers  pAM7, chBL, and a57m), our paper is "**well-written, easy-to-read, and thorough**" (Reviewers ymqN, pAM7, and a57m), the theoretical analysis is "**convincing and sound**" (Reviewers ymqN, pAM7, chBL, and a57m), the empirical results "**clearly support our key takeaways**" (Reviewers ymqN, pAM7, FoGd, and a57m), our paper "**identifies the role of normalization and highlights its importance**" (Reviewers pAM7, chBL, and FoGd).

Here, we first dicuss the two main concerns shared by the reviewers, namely *i) why consider USAM (which seems to be impractical)*, and *ii) what are the main takeaways of our paper*. More detailed responses are included in our responses to Reviewer a57m and Reviewer FoGd, respectively.

1. *Why consider USAM?*  

We agree that USAM is not practical, but **we consider it in order to understand SAM better**. To provide brief context, despite the success of SAM, its theory is at its nascent stage, and many existing theoretical works _ignore the normalization step_ in SAM "for simplicity."  
Our work is motivated precisely by this state of affairs in the theory of SAM, and our work reveals that such "for simplicity" choice is unwarranted,  as **normalization plays an important role** in SAM. 
  
Moreover, in order to conduct a principled ablation based investigation, **we only remove normalization in SAM, and keep all other parts unchanged**. This reason underlies **the reason for considering USAM;** we are not claiming USAM's practicality (or any other advantage).

2. *What are the main takeaways of our paper?*  

There are **two main takeaways,** one for the theoretical ML community, and another one for the empirical ML community.  

For the theoretical community, our takeaway is that **normalization plays a critical role and cannot be neglected when analyzing SAM**. We believe this message should be valuable to the community, as various previous theoretical papers have omitted normalization -- we demonstrate that such an omission drastically changes the algorithmic behavior of SAM.
  
For the empirical community, we find **the success of SAM is (at least in part) backed up by the normalization** -- although SAM exhibits  great success in practice, and the success is supported by various theoretical works, its working mechanism is far from understood. We start this investigation by specifically focusing on the importance of normalization.

---

We also thank the reviewers for catching the confusing arguments in our draft -- we will update them in the final version as promised. Many thanks to all the reviewers for your efforts.

Finally, we also included an additional plot as an attachment. Its explanation can be found in our second response to Reviewer chBL.

We hope our response resolves your main concerns. If you have any further questions, please do not hesitate to contact us. We are more than happy to answer them!

Thank you and Best Regards,  
Authors


Meta-review: The authors investigate the importance of normalization in the Sharpness-Aware Minimization (SAM) algorithm. They do this by comparing how gradient descent, SAM, and unnormalized SAM behave through experiments and theory. The main finding is that normalization is crucial and must not be overlooked when analyzing SAM. Reviewers consider the work relevant, with valuable theoretical outcomes. They've also provided suggestions to enhance the paper's presentation. We recommend authors to include these in the final version. In summary, I recommend accepting this paper.
Decision: Accept (poster)