Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training
Authors: Zhanpeng Zhou
Mingze Wang
Yuchen Mao
Bingrui Li
Junchi Yan
Venue: ICLR 2025 Spotlight (2025)
Keywords: Sharpness-Aware Minimization
Implicit Bias
Training Dynamics
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We study the implicit bias of SAM during the late phase of training, revealing that SAM efficiently selects flatter minima over SGD even when applied in the last few epochs.

Abstract: Sharpness-Aware Minimization (SAM) has substantially improved the generalization of neural networks under various settings. Despite the success, its effectiveness remains poorly understood. In this work, we discover an intriguing phenomenon in the training dynamics of SAM, shedding light on understanding its implicit bias towards flatter minima over Stochastic Gradient Descent (SGD). Specifically, we find that *SAM efficiently selects flatter minima late in training*. Remarkably, even a few epochs of SAM applied at the end of training yield nearly the same generalization and solution sharpness as full SAM training. Subsequently, we delve deeper into the underlying mechanism behind this phenomenon. Theoretically, we identify two phases in the learning dynamics after applying SAM late in training: i) SAM first escapes the minimum found by SGD exponentially fast; and ii) then rapidly converges to a flatter minimum within the same valley. Furthermore, we empirically investigate the role of SAM during the early training phase. We conjecture that the optimization method chosen in the late phase is more crucial in shaping the final solution's properties. Based on this viewpoint, we extend our findings from SAM to Adversarial Training.

Published as a conference paper at ICLR 2025
SHARPNESS- MINIMIZATION EFFICIENTLY SE- LECTS FLATTER MINIMA LATE IN TRAINING
Zhanpeng Zhou -i, Mingze Wang"; Yuchen Mao', Bingrui Li", Junchi Yan 'Shh of Computer Science & Sch. of Artificial Intelligence, Shanghai Jiao Tong University PPeking University 'Tsinghua University "Shanghai Artificial Intelligence Laboratory 1zzp1012, janjunchi wasjtu.dd....
ABSTRACT
Sharpness-Aware Minimization (SAM) has substantially improved the generaliza- tion of neural networks under various settings. Despite the success, its effectiveness remains poorly understood. In this work, we discover an intriguing phenomenon in the training dynamics of SAM, shedding light on understanding its implicit bias towards flatter minima over Stochastic Gradient Descent (SGD). Specifically, we find that SAM efficiently selects flatter minima late in training', Remarkably, even a few epochs of SAM applied at the end of training yield nearly the same generalization and solution sharpness as full SAM training. Subsequently, delve deeper into the underlying mechanism behind this phenomenon. Theoretically, we identify two phases in the learning dynamics after applying SAM late in train- ing: i) SAM first escapes the minimum found by SGD exponentially fast; and ii) then rapidly converges
flatter minimum within the same valley. Furthermore, we empirically investigate the role of SAM during the early training phase. We conjecture that the optimization method chosen in the late phase is more crucial in shaping the final solution's 's properties. Based on this viewpoint, we extend our findings from SAM to Adversarial Training. We released our source code https: //git thub. com zzp1 012/ / SAM-in-MMaaa -atn--aate.ae..
INTRODUCTION
Understanding the surprising generalization abilities of over- peeepaaramterized neural networks yields an important yet open problem deep learning. Recently, has been observed that the generalization of neural networks is closely tied to the sharpness of the loss landscape (Keskar et al., 2017: Zhang et al., 2017; Neyshabur et al., 2017; Jiang al., 2020). This has led to the development of many gradient-based optimization algorithms that explicitly/impliiity regularize the sharpness of solutions. In particular, Foret et al. (2021) proposed Sharpness-Awe has substantially improved the generalization and robustness (Zhang et al., 2024) of neural networks across many tasks, including computer vision (Foret et al., 2021; 1; Chen et al., 2022; Kaddour et al., 2022) and natural language processing (Bahri et al., 2022).
Minimization (SAM), which
Despite the empirical success of SAM, its effectiveness is not yet fully understood. Andriushchenko & Flammarion (2022) has shown that existing theoretical justifications based PAC-Bayes gen- eralization bounds (Foret et al., 2021; Wu et al., 2020a) are incomplete in explaining the superior performance of SAM. Additionally, recent theoretical analyses of SAM's dynamics and properties often rely on unrealistic assumptions, such as sufficiently small learning rate or perturbation radius, which undermines the validity of their conclusions. Understanding the hidden mechanisms behind SAM remains an active area of research.
Recent works show that the effectiveness of gradient- bassed optimization methods can be attributed to their implicit bias toward solutions with favorable properties (Vardi, 2023). One notable example of such implicit bias is that Stochastic Gradient Descent (SGD) and its variants tend to find flat minima, which often leads to better generalization (Keskar et al., 2017; Zhang et al., 2017). It is known that
Equal contribution.
Corresponding author. The work in part supported by NSFC (62222607) and Shanghai Municipal Science and Technology Major Project (2021SHZDZZ0012). by ate training" refers to the phase when the optimization process nearing the late stages of training.

---

Published as a conference paper at ICLR 2025
SGD SAM
SAM
0:08
Test Eiroo
0.07
0.06
0.05
MAMaolttion gpodia)
SQD sahaoon
0.04
spode
0.03
100
125
150
175 200
Epoch
(a)
(b)
Figure 1: SAM operates efficiently even when applied only during the final few epochs of training. (a) Test error curves for WideResNeet-6-22 on CIFAR- 10 trained with different strategies. The blue and green baseline curves represent training solely with SGD and SAM, respectively. The orange curve shows the test error of a neural network initially trained with SGD up to epoch = 175, followed by SAM training up to epoch T = 200. The first 100 epochs are omitted for clarity, Curves are means and standard deviations over five trials with different random seeds. The detailed settings described in Section 2. (b) A schematic picture of the training trajectory after applying SAM late in training. The contour plot represents the loss landscape, with darker regions indicating lower loss. The orange dotted lines depict the path of the SAM iterator.
and syper-parameters
SAM selects flatter minima over SGD in practice, which represents form of implicit bias as well. While the design of SAM inspired by sharpness regularization, its practical implementation (see Equation (3)), which minimizes first-order approximation of the original objective (see Equations (1) and (2)), does not explicitly achieve this. Understanding SAM's implicit bias towards flatter minima, especially in comparison to SGD, is of paramount importance for explaining its effectiveness.
In this work, we uncover an intriguing phenomenon the training dynamics of SAM, which sheds light understanding its implicit bias toward flatter minima. Specifically, we find that SAM operates efficiently even when applied only during the last few epochs of training". As shown in Figure I (a). SAM applied in the last few epochs prochs-atter initial training with SGD-car can achieve nearly the same test performance as training exclusively with SAM. Furthermore, our analysis of the solution sharpness reveals that last-fewsepochs SAM identifies flatter minima effectively
full SAM training. We conclude that SAM efficiently selects flatter minima over SGD late in training. Contrary the conventional view that early training dynamics are critical (Achille al., 2019; Frankle et al., 2020), our discovery underscores the importance of the late training phase for generalization.
Subsequently, we delve deeper into the mechanism behind the efficiency of SAM in selecting flatter minima over SGD during the late phase of training. Theoretically, we identify a two-phase picture in the learning dynamics after applying SAM late in training: Phase SAM initially escapes from the relatively sharp minimum found by SGD exponentially fast; Phase II. it then rapidly converges to flatter minimum over SGD within the same valley. See Figure (b) for an illustration. As outlined in Table 1, we characterize the two-phase picture into four key claims (P1-4), each supported by rigorous analyses of SAM's convergence, escape behavior, and bias toward flatter minima. These claims further validated by experiments in both toy models and real-worke scenarios. This two-phase picture explains SAM's ability to efficiently find flatter minima, even when applied only in the last few epochs, yielding novel insights into its implicit bias and effectivenes..
Furthermore, we investigate the necessity of applying SAM during the early phase of training. Our results show that early-phase SAM has minimal effect both generalization and sharpness of the final solution. Based on these findings, we conjecture that the choice of the optimization algorithm during the late phase of training plays more critical role in shaping the final properties of the solution. Extending this viewpoint to Adversarial Training (AT) (Madry et al., 2018) and robustness, we find that similar to SAM, AT efficiently selects robust minima when applied late in training.
In summary, our work focuses on the implicit bias of SAM in the late training phase, which has not yet been carefully studied. Our two-phase picture advances the understanding of its effectiveness.
2 AAnrrishhheenoo & Flammarion (2022) observed similar finding but did not explore it in depth. Surprising. we first show that even few epochs of SAM the end of training yield generalization comparable to full SAM training. We study this further via extensive experiments under various setup and rigorous theoretical analysis.
2

---

Published as a conference paper at ICLR 2025
2 BACKGROUND AND PRELIMINARIES
Notation and Setup. Let Derain = (ar, y:) and Drest = ((; 9.) be the training set
and the testing set, where 2; E R is the input and yi E IR is the corresponding target. We consider model in the form of fax; 0), where O E IR' are the model parameters and f(x;; 0) R is the output of the input æ. The loss of the model at the i-th sample (æt, yi) is denoted as E (a;;; 0)), simplified to (0). The loss over the training set and the testing loss then given by LDreain (0)
n >== (0) and D (0) = m Si=n+1 (0), respectively. Unless otherwise specified, we denote L(0)= = CD..ai (0) for simplicity. For classification tasks, we define the classification error on the test set Dusst as EIr.. (f). For a matrix A, let ||A|||2 2. || ps and Tr(A) denote its spectral norm, Frobenius norm and trace, respectively.
For our theoretical analysis in Section 4, we always assume k: = 1 for simplicity and consider the squared loss, i.e., ()) (2;; 0) y: 2 We focus on the over-parameterized case in the sense that ming ((0) = 0. To characterize the local geometry of the training loss at point 0, we consider the Fisher matrix G()) = - >' Vf(z;; O)V 0) T and the Hessian matrix ((0) = G(0) Σi= (f( y/) V (ax; O), where G(O) Rs (O)) in the low-loss region.
Sharpness-Aware Minimization. The central idea behind SAM is to minimize the worst-case loss within a neighborhood of the current parameters. Specifically, SAM problem can be formulated as
min C SAM p):
where [SAM (0: p) = max L(O + c).
(I)
Here, p a given neighborhood radius. However, finding the perturbation in the parameter space that maximizes the loss, i.e. can be computationally intractable in practice. Thus Foret et al. (2021) approximated the maximize-loss perturbation using a first order solution:
arg max (C + (0) = PVC (0) VVV (0)ll2
(2)
Consequently, let 17 be the learning rate and &, (f........ E...) [n]] be the batch indices sampled
at iteration t, where B is the batch size. The mini-batci batch loss is given by Cz. (0)= 1/2 Exes, (()).
With a small yet key modification to SGD, the update rule of SAM with stochastic gradients is:
o 0² - Lg. (0' + P (V))((0)) where P = p/ Vc. (0')ll-
(3)
Our experiments are conducted using SAM in Equation (3), whereas our theoretical analyses in Section apply the simpli- fied SAM in Equation (4). Specifically, we approximate the step size in Equation (3) with a constant p. We also assume the independence between the randomness in the inner and outer updates, i.e., <), <2. These simplifications facilitate math- ematical tractability while capturing the central behavior of SAM in training dynamics and generalization, which has been supported by empirical evidence (Andriushchenko marion, 2022). They have also been widely adopted in recent theoretical advances on SAM (Andriushchenko Flammarion, 2022; Behdin& & Mazumder, 2023; Agarwala & Dauphin, 2023; Behdin et al., 2023; Monzio Compagnoni et al., 2023).
8"
e SGD
Flam-
0 0 SAM., 0o
0 6Sn-
Figure 2: The illustration of the switching method. Blue dashed lines represent SGD training. while orange dashed lines represent SAM training. 0 denotes the random ini- tialization. f1, t2 denotes two dif- ferent switching points.
0++1 0- - V (0' -0'+ +
(4)
Switching Method. In this paper, we demonstrate our main dis- covery by switching between SGD and SAM at a specific point during training, referred to as switching method. Specifically,
let T be the total number of training epochs and t be the switching epoch, considering the switch from SGD to SAM without loss of generality. We initially train the model from scratch with SGD for epochs to get 0aa and then switch to SAM for the remaining T t epochs to get - Notably, when t T or t 0, the model is trained exclusively with SGD or SAM, simply denoted
"The extension the case where k > is straightforward.
3

---

Published as a conference paper at ICLR 2025
WideRResNet 28
CIRAR-100
WiteReeNet 1e-ist------- 16-8 on CIBAR-10
Genpral gRD betwoon
0.20
0.64
0.040
0.015
0.032
0.01
0.35
0.63
0.30
0.02
0.06
0.006
0.00
-0.05
-0.1
-2.000
-0.003
0.8
Sheqpeeess of
50
0.0
6.2
2.4
0.5
0.8
1.0
0.0
0.2
0.4
0.5
0.8
1.0
SAMbaining proportion when obtaining SSD-SSAM
T-t
T
Di L (SSDD-SAM)))
Lore (OSAM)
Err (SAn))
ErT
(0%a)))
[0 SGD- 2ll
Figure 3: The impact of SAM training proportion on generalizatioo/sharpees when switching from SGD SAM. The generalization gap between the models 0% -Scw-S and 0 (top row) the sharpness of -SAD- (bottom row) vs. the SAM training proportion of OSAD-SAMM Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations. Results more datasets and architectures can be found in Appendix B.2.
6. and AM respectively. See Figure 2 for illustration. We fix the total number of training epochs T and vary the switching point f to study the effect of SAM at different training phases. To ensure fair comparison, we ensure the same initialization, mini-batch orders, and hyper-parameters across all models, regardless of the switching point. Consequently, the training trajectories of models with different switching points align precisely before the switch.
Main Experimental Setup. Following Foret et al. (2021): ): Kwon et al. (2021), we perform experi- ments on the commonly used image classification datasets CIFAR- 10/100 (Krizhevsky et al., 2009), with standard architectures such as WideResNet (Zagoruyko & Komodakis, 2016), ResNet (He et al., 2016), and VGG (Simonyan Zisserman, 2015). We use the standard configurations for the basic training settings shared by SAM and SGD (c.g.. learning rate, batch size, and data augmentation) as in the original papers and set the SAM-specific perturbation radius p to 0.05, as recommended by Foret et al. (2021). We also extend our findings to AT (Madry al., 2018) and ASAM (Kwon et al., 2021). Due to space limits, more experimental details and results can be found in Appendix B.
3 SAM EFFICIENTLY SELECTS FLATTER MINIMA LATE IN TRAINING
In this section, we provide empirical evidence indicating that SAM efficiently selects flatter minima over SGD late in training. We first show that applying SAM only during the last few epochs of training achieves generalization comparable training entirely with SAM. Taking a step further, we demonstrate that festsss SAM finds flatter minima as effectively as full SAM training.
Applying SAM in the last few epochs achieves generalization comparable full SAM training. We conduct experiments on various datasets and architectures to investigate the effect of applying SAM late training, after initial training with SGD, on generalization. Specifically, we use the switching method (detailed in Section 2), which switches from SGD to SAM at epoch t, to obtain the model Buw---AM.. For comparison, we train a baseline model entirely with SAM under the same settings, resulting 0a Both models are trained for total of I epochs. To assess the generalization ability of d.M.- relative to oaam: we calculated the test loss and test error gaps between these two models: LD... (0 (Sam.)) Lo (0%Am) and Erro,a (OMMD))) Errr (0%a)), where Drest denotes the unseen test set. If the generalization gap between 0 ISAM, and 0 approaches zero, then we say OSaD- achieve nearly the same generalization performance as 8
Surprisingly, we find that applying SAM only in the final few epochs can achieve nearly the same generalization performance as training entirely with SAM. In particular, we fix the number total training epochs T and vary the switching point t adjust the SAM training proportion of 0% -SAM.s which is exactly (T - t)/T. We analyze how the generalization gap between 0% -SMM, and 0 changes the SAM training proportion increases. In Figure 3, the generalization gap becomes
4

---

Published as a conference paper at ICLR 2025
negligible once the SAM training proportion exceeds zero, indicating that even few epochs of SAM applied at the end of training yield generalization performance comparable to full SAM training. This phenomenon is consistently observed across various datasets and architectures.
A similar phenomenon was noted by Andriushchenko & Flammarion (2022), but was not given significant attention. We isolate and extend this phenomenon by showing that even fewer epochs of SAM (e.g. 1-5 epochs for WideResNet- 16-8 on CIFAR-10) applied in the late phase of training yield test error and loss comparable to full SAM training.
Applying SAM in the last few epochs finds flatter minima as effectively as full SAM training. Furthermore, we investigate the effect of applying SAM in the late phase of training on the sharpness of the final solution. Similarly as before, we analyze how the sharpness of the solution s (AM, changes w.r.t the SAM training proportion. We measure the sharpness along the most curved direction of the loss surface, i.e., the spectral norm of the Hessian matrix | (0)l2- In Figure 3, the sharpness of the solution -SMM, dramatically drops once the SAM training proportion is greater than zero, reflecting a similar trend to the generalization gap. This indicates that applying SAM in the last few epochs training is as effective as training entirely with SAM in achieving flatter minima.
In summary, we see that SAM efficiently selects flatter minima over SGD when applied late in training. This suggests a fine-tuning scheme for SAM: start from a pretrained checkpoint and continue training with SAM for several epochs, reducing computational costs while maintaining similar performance. In Section 4, we will dig deeper into the underlying mechanism behind this phenomenon and gain new insights into SAM's implicit bias towards flatter minima.
How SAM EFFICIENTLY SELECTS FLATTER MINIMA LATE IN TRAINING
We have seen that SAM efficiently selects flatter minima over SGD when applied late in training. In this section, we aim explore the precise mechanism behind this phenomenon. Theoretically, identify a two-phase picture in training dynamics after switching to SAM the late phase, which is characterized by four key claims (P1-4), as outlined in Table 1. This picture demonstrates that SAM provably escapes the minimum found by SGD and converges to flatter minimum within the same valley in just few steps, explaining its efficiency in selecting flatter minima late in training.
Phase I
(P1). SAM rapidly escapes from the minimum found by SGD; Corollary 4.2
(escape)
(P2). However, the iterator remains within the current valley.
Proposition 4.1
Phase
(P3). SAM selects a flatter minimum compared to SGD;
Theorem 4.1
(converge)
(P4). The convergence rate of SAM is extremely fast.
Theorem 4.3
Table 1: Overview of the two-phase picture and corresponding theoretical results.
To understand the two-phase picture, let us first consider a toy but representative example.
Example 4.1. Consider using the shallow neural network f(u, v; tanh( y tanh(uz)) fir a single data point (x = 1, = 0) under the squared loss E(y: y/) = (y - 2/) then the loss landscape can be written as (()) = tanh? tanh(u)). where f = (u, v).
Figure 4 (a) visualizes the training dynamics for Example 4.1. In Figure 4 (a), the dynamics occur within the valley surrounding the set of the minima M = {(u, v)
0). On M, the sharpness [r) = tanh tanh²(w), implying that the smaller the /ul, the flatter the minimum. Initially, SGD converges relatively sharp minimum 0 After switching to SAM, the iterator exhibits the two-phase dynamics:
given by [0)]2 = Tr(H(())
Phase I (escape). SAM first escapes the sharp minimum Oend found by SGD within few iterations (P1), but remains within the same valley around M (P2);
Phase II (converge). SAM then rapidly converges to another flatter minimum Oand with a smaller [u] compared to SGD
SAM (P3-4),
NNote that the landscape has another valley around different set of minima V ((,,)) v) : a==00
5

---

Published as a conference paper at ICLR 2025
Toy model on 2 tanh (v tanh(ux))
WideResNet: on C-FAR-10
10-1
Phoce Phapet
560
Eiteats (o (co0o)
L' L(O,)
0.1-
10
0.0
012
-0.2
0.30
-0.2
0.05
-03
0.4
L.0
1.2
1.4
19
0.00
0.0
0.2
0.6
Strp
[b)
(c)
Figure 4: (a) Visualization of the two-phase dynamics for Example 4.1. The horizontal gray line represents the set of the global minima M {(u, v) v 0}, where smaller values [u] correspond to flatter minima. The blue lines trace the SGD iterates, leading to S while the orange lines show the SAM iterates, which converge to flatter minimum SAM Notably, SGD and godd stay the same valley around M. (b) The exponentially fast escape from minima found by SGD.
SAM
Train loss L (0SD SCD-SAM. v.s. the update step t'. Here, t is the switching step, chosen to be sufficiently large for SGD to converge, and E' is the number of updates after switching to SAM. The red line represents the mean over 100 trials with different randomness. (c) SAM converges to a flatter minimum within the same valley as the one found by SGD. The loss L (top row) / the second-orde finite difference L" (bottom row) of the interpolated model 0g V.s. the interpolation coefficient 2C(0x)) /(2h)?, SGD-SAM and SGD stay within the same valley. L" (²)0)) gradually increases with A, implying 0gad SGM is flatter than SGD
Here, 0^ (1 x)0 SCD--SM A SGD- and (²)0)) (C(0x+n) + L(0x-
where h is fixed to 0.1. This is no barrier in the loss curve, indicating ggnd
This toy example successfully illustrates the two-phase picture after switching to SAM in the late training phase, but it offers only basic intuition. In the following, will formally present our theoretical results, with each claim rigorously supported by corresponding analysis.
4.1 THEORETICAL GUARANTEES FOR P1 AND P3: LINEAR STABILITY ANALYSIS
In this subsection, we provide theoretical support for two key claims our two-phase picture through a linear stability analysis: P1. SAM escapes from the relatively sharp minima found by SGD exponentially fast; P3. SAM selects a flatter minimum compared to SGD.
During the late phase of training, the model remains close to some global minimum O* and can be locally approximated by its linearization:
(i; 0) [ax; 0*) + (V 0-,,0 0*)
Consequently, we can characterize the local dynamical stability (Strogatz, 2000) of the SAM iterator near the global minimum on this linearized model.
Similar to the linear stability of SGD (Wu et al., 2022), we define the linear stability of SAM in Definition 4.1. Note that for specific optimizer, only linearly stable minima can be selected.
Definition 4.1 (Loss-based linear stability of SAM). A global minimum 0* is said be lin- early stable for SAM if there exists a C' > 0 such that, when SAM optimizes the loss (()) =
2n Liz1 (fin (x;;; )2 on the linearized model (Equation (5)), the following condition holds: E ] & (E))( V' 2 O. Here, 0 could be any point near 0 where the linearization is valid.
Wu et al. (2022) has shown that the linear stability of stochastic algorithms is significantly influenced by the gradient noise (()) (l) (())). where denotes the index of a random sample. Thus, it is crucial to consider the structural properties of the gradient noise in our analysis.
Gradient noise structure. A series of studies (Zhu et al., 2018; Feng & Tu, 2021; Wu et al., 2020b; 2022; Mori et al., 2022; Wojtowytsch, 2024) have established two key properties of the gradient noise E(0): (i) its shape is highly anisotropic; (ii) its magnitude is loss dependent. Specifically, by incorporating both the noise shape and magnitude, Mori et al. (2022) has proposed an approximation for the gradient noise covariance E(0) > (()) T. suggesting Σ(() (()))

---

Published as a conference paper at ICLR 2025
Inspired by this approximation, we assume a weaker alignment property of the noise covariance Σ(O) (see Assumption 4.1) in our analysis, which has been theoretically and empirically validated in recent works (Wu et al., 2022; Wang & Wu, 2023).
Assumption 4.1 (Gradient noise alignment). There exists 7 > 0 s.t.
THOGOO 22 GOOII
2 y for all O.
By analyzing the linear stability of SAM under this gradient noise alignment assumption, we establish two theorems supporting the two key claims P1 and P3.
Theorem 4.1 (P3. Proof in Appendix C.I.). Let 0* be a global minimum that is lin- early stable for SAM in Equation (4), and suppose Assumption 4.1 holds. Then we have (0)))) (1+ ( -
(1(())))))) S ."
Theorem 4.1 characterizes the sharpness of the linearly stable minima for SAM, which is influenced by the learning rate 17. the batch size B, and the perturbation radius p. Since only linearly stable minima can be selected by optimizer, Theorem 4.1 characterizes the sharpness of global minima selected by SAM. Notably, taking p = recovers the corresponding result for SGD (see Table 2).
SAM (Theorem 4.1)
SGD (Wu al., 2022)
Sharpness bound
|10 (1+ (1+ (H
|| [(()))
(H(0)) S /
Table 2: Comparison of the sharpness of global minima selected by SAM and SGD.
Support for P3. We compare the bounds on the sharpness of the global minima selected by SAM and SGD in Table 2. The additional term 1 + p²y [IH((')) >|2 > 1 in SAM's bound impose stricter requirement on the sharpness for by SAM compared to SGD, which supports our key claim P3 that SAM tends to select flatter global minima over SGD. This term arises from SAM's inner update, i.e. O OF + pV (0*) 0 + ((0) + p p(Veg (0') VL(O)). Intuitively, the ascent update pv (O*) tends to increase the loss, while the noise P(Vee (0') VC(O')) introduces extra stochasticity. These factors together contribute to the flatter-minima selection by SAM.
Corollary 4.2 (P1. Proof in Appendix C.I.). Let 0" be a global minimum, and suppose Assump- tion 4.1 holds. If (0*) (1 +
B [[H(0))|?)) > T. then 0 is linearly non-stable for SAM in Equation (4). Indeed, [F((())) 2 C'B(C(0") holds for all t with C >>
p²x
Corollary 4.2, together with Theorem 4, characterizes the necessary condition for a global minimum to be linearly stable for SAM, namely the stability condition of SAM. Furthermore, Corollary 4.2 illustrates that if a global minimum fails to meet the stability condition of SAM, the SAM iterator will escape from the minimum exponentially fast. Notably, such escape occurs within the low-loss regions, where the linear approximation in Equation (5) holds. Complementary analysis of dynamics in the high-loss regions is provided in Proposition 4.1.
Support for P1. Although the sharpness bound for SGD in Table 2 is an upper bound, Wu et al. (2022) empirically demonstrates that the bound is nearly tight, with SGD solutions approximately satisfying || H (Oge ~ . satisfy the stability condition of SAM. This leads to our key claim P1 that SAM escapes from the sharp minima found by SGD exponentially fast, as validated by our experiments in Figure 4 (b).
B_ Thus, with the same 7 and B, the solution found by SGD cannot
Comparison with other linear stability analyses of SAM (Shin et al., 2023; Behdin al., 2023). These works have focused on the norm-based linear stability, characterized by Bi|l0' 1- <
C' C||00 0- 11 Vt 2 O, which is stricter than our loss-based stability (see Definition 4. 1) due to the presence of degenerate directions in over- perr-aaamterizeed models (see Lemma C.2). Another key difference is the novel use of the alignment property of gradient noise (see Assumption 4.1) in our analysis, enabling us to derive a concrete bound on the sharpness of SAM's solutions.
4.2 THEORETICAL GUARANTEE FOR P2: BEYOND LOCAL ANALYSIS
In this subsection, we provide theoretical support for another key claim in our two-phase picture: P2. SAM remains within the current valley during escape.
Rs is not used in any formal proofs. Here, indicates empirical finding that has not been rigorously proven.
7

---

Published as a conference paper at ICLR 2025
Sub-quadratic landscape. In high-loss regions, the local linearization in Equation (5) longer provides a valid approximation for neural networks, and consequently, the quadratic approximation of the neural network loss landscape fails to capture its true structure. Beyond the quadratic landscape, Ma et al. (2022) has introduced the concept of the sub-quadratic landscape: the loss grows slower than quadratic function, and the landscape in high-loss regions is significantly flatter than in low-loss regions. The sub-quadratic landscape has been empirically verified in high-loss regions, and able to explain the Edge of Stability phenomenon (Cohen al., 2021) (also observed in Wu et al. 20; Jastrzebski et al. (2020)), which the quadratic landscape cannot account for.
While Section 4. focuses on local analysis near global minima, the analysis here is necessarily non-local, as SAM might iterate into high-loss regions during escape. In this subsection, we need to show that SAM remains within the current valley during escape under a sub-quadratic landscape.
Additional setup. For simplicity, we assume a model size of p 1 and consider the full-batch SAM, i.e. g++1 l v/ L( + PVL((0)) This setup is sufficient to illustrate that SAM will remain within the current valley under sub-quadratic landscape.
Definition 4.2 (Sub-quadratic landscape for 1). Let = 0 be a global minimum of L with > 0. The valley V = [-c, is termed sub-quadratic if for all 2, 21, 22 V, it holds that =L (=) 2 0 and ("(==l) 2 ("( if |21 S |=2).
sharpness (0)
Following Theorem I in (Ma al., 2022), we define the sub-quadratic landscape in Definition 4.2.
Specifically, for any 0 < 21 < 2g < c. it holds that (21)) > (²22) and ((21) < (22)), indicating that high-loss regions are flatter than low-loss regions. The same applies for 0 21 > 2/ > -c
Proposition 4.1 (P2. Proof in Appendix C.2.). Under Definition 4.2, assume the landscape is sub-quadratic in the valley V = [-2b, 26). Then, for all initialization @ E (-b. b). and n) p s.t. 1) < min (/ (z)), P S min(1/ 1/a, 7, 7 min [2'(22) ('()), the full-batch SAM (see Equation (10))
=EV
0<|=/<<
will remain within the valley V, ie.. f V for all t N.
Proposition 4.1 shows that under sub-quadratic landscape, even if SAM escapes the minimum #* when 1j > 2/a(1 + ap) and 6 = f², it will still remain within the current valley This proposition holds when 17 < min 6/), where the upper bound min b/l(((zz) can be sufficiently large under =eV
=EV highly sub-quadratic landscapes (see Example C.I). Furthermore, this proposition can be extended to almost all initializations within the valley V. See Proposition C.2.
Support for P2. Proposition 4.1 supports our key claim P2 that SAM remains within the current valley during escape. Intuitively, as shown in Section 4.1, the low-loss region around the sharp minimum found by SGD cannot satisfy SAM's stability condition, causing SAM to escape toward higher-loss regions; however, under the sub-quadratie landscape, the high-loss regions are much flatter, which satisfies SAM's stability condition and prevents from escaping the current valley. The experiment in Figure 4 (c) backs up this theoretical claim.
4.3 THEORETICAL GUARANTEE FOR P4: CONVERGENCE ANALYSIS
In this subsection, we provide theoretical support for the final claim in our two-phase picture: P4. The convergence rate of SAM is significantly fast.
To establish the convergence results for stochastic optimization algorithms, it is often necessary to make assumptions about the loss landscape and the magnitude of gradient noise {()).
Assumption 4.2 (Loss landscape: I--smoothness and
(i) There exists >>
s.t. ||H(0))2 < for all 0: (ii) There exists μ > 0 s.t. ||V[())] 2 20C(() for all 0.
Assumption 4.2 quite standard and is frequently used in the non-convex optimization litera- ture (Karimi et al., 2016). Specifically, for neural networks, L(-) is four-times continuously differentiable, this assumption holds at least near the set of global minima of L (Arora al., 2022).
Assumption 4.3 (Gradient noise magnitude). There exists > s.t. (E] < o L(()) V0.
The bounded variance assumption ([] o² is commonly employed in classical optimization theory (Bubeck et al., 2015). However, discussed in Section 4.1, recent studies (Mori et al., 2022;
Multiple valleys exist in the settings of Proposition 4.1 (see Remark C.3 for more details).
8

---

Published as a conference paper at ICLR 2025
WidkResN 28- on CIFAR-100
WideResNNes 16-8 on CIFAR-10
Sherpiiss Gezzralzzaion gop between
0.30
0.06
0.06
05 10M--500.. and SAM
0.25
0.05
0.05
0.15
0.15
0.03
0.03
SAM-SSDD.
0.30
0.02
0.02
0.006
0.05
0.01
0.00
0.00
0.00
0.000
-0.01
-0.003
0.0
0.2
0.4
0.6
0.0
1.0
150
30
20
0.0
0.2
0.4
0.6
0.0
1.0
0.0
0.2
0.4
0.6
0.8
1.0
SAM training proportion when obtaining SGn
Diuest (SAM-SGD.))
SAM
ErTo
(SAAx-SGD2)
Br (AAM)
I|M Sax- )ll
Figure 5: The impact of SAM training proportion on generaliaation/harpnes. when switching from SAM SGD. The generalization gap between the models 05AM and O (top row)/the sharpness of - (bottom row) vs. the SAM training proportion of 8SAM-SGD, Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations. Results more datasets and architectures can be found in Appendix B.5.
Feng & Tu, 2021; Liu et al., 2021; Wojtowytsch, 2024) have shown that the magnitude of SGD noise is also bounded by the loss value. Thus, we introduce Assumption 4.3: E (()) This indicates that the noise magnitude diminishes to zero global minima, contrasting sharply with the classical bounded variance assumption and closely aligning with real-world scenarios.
Under these assumptions, we present the convergence result for stochastic SAM Theorem 4.3.
Theorem 4.3 (P4. Proof in Appendix C.3.). Under Assumption 4.2 and 4.3, let {0'): }: be the parameters trained by SAM in Equation (4). If < min(1/21, HB/220' and p S min (1/4L, pB/ pB/4Lo², n////2AL" /2122), then have E [C(()) S (1 mu/2)" L()"), VI NEE N.
Support for P4. Theorem 4.3 supports our key claim P4 that the convergence rate of stochas- tic SAM is significantly fast, and notably faster than the previous result on SAM's convergence rate (Andriushchenko & Flammarion, 2022).
Comparison with other convergence results for stochastic SAM (Andriushchenko & Flammarion, 2022; Shin et al., 2023). These works, along with our result, all study the convergence stochastic SAM under the standard Assumption 4.2. Andriushchenko & Flammarion (2022) has shown that stochastic SAM converges under polynomially fast rate of (1/tt, but this is much slower than exponentially fast rate. The difference lies in the assumptions on gradient noise magnitude: they have used the classical assumption E[ll{(0) (0) < o², whereas we rely on the more realistic Assumption 4.3. Shin al. (2023) has also proved that stochastic SAM converges exponentially fast. However, their analysis is based on an additional interpolation assumption due to over-parameterization.
5 Is SAM NECESSARY IN THE EARLY PHASE OF TRAINING?
Since our main discovery in Section 3, one might question the necessity of using SAM in the early phase. In this section, we demonstrate that applying SAM during the early phase offers only marginal benefits for generalization and sharpness compared to full SGD training. Thus, we conjecture that the optimization algorithm chosen the end of training more critical in shaping the final solution's properties. Based this viewpoint, extend our findings from SAM to Adversarial Training.
Early-phase SAM may offer limited improvements over SGD in generalization and sharpness. We conduct experiments under various settings to investigate the effect of applying SAM during the early phase of training, followed by SGD, on generalization and sharpness. Similarly to Section 3, we the switching method but switch from SAM to SGD, to obtain the model 05a--sOD.. In Figure 5, the generalization gap remains substantial even when trained with SAM for most the time (e.g.. t/T 0.8). A similar trend is observed in the sharpness of 0 0 Only full SAM training (i.e.,

---

Published as a conference paper at ICLR 2025
WideResNN 16-8 CIFRR 10 with wwwhhhhehhhhssee Adversarial Traning
1.0
Classficaion Enror Brrg(
Classfiction Broor
0.4
0.2
0.2
0.0
0.0
160
170
180 Bpoch
190
200
0.0
0.2
0.6
0.8
1.0
AT braining proportion when obtaining OS.
T-t
(a)
(b)
0cO
Ar
8on AT.
perturbt (0)
ErTo (0)
Figure 6: AT improves robustness efficiently even when applied only during the final few epochs of training. (a) Robust/natural error vs. training epochs for models trained with different strategies. (b) Robust/natural error of OKGD-AAT vs. the AT training proportion (T t)/T. Dashed lines denotes the robust error Err-D and solid lines denotes the natural error ErtD.... Blue and green curves represent models trained solely with SGD and AT, i.e., foaD and 0 respectively. Orange curves indicate models which switching from SGD to AT at epoch f, i.c. O.CD-ATT.. All results are averaged over five trials with different randomness.
I/T consistently finds the flat minima. This suggests that applying SAM before SGD could yield only limited improvements in generalization and sharpness compared to full SGD training.
Together with findings Section 3, we conjecture that the optimization algorithm chosen in the late training phase is more critical in determining the final solution's 's properties. However, current findings are based solely on generalizatioo/sharpeess can we extend to other properties?
Extension: Adversarial Training efficiently finds robust minima late in training, Adversarial Training (AT) (Madry al., 2018), which networks are trained on adversarial examples (Szegedy et al., 2014), significantly improves the robustness of neural networks. Recently, Zhang et al. (2024) highlighted duality between SAM and AT: while SAM perturbs the model parameters during training (see Equation (I)), AT perturbs input samples, with both methods relying gradients for perturbations. Here, we pose the question: can our findings on SAM be extended to AT?
Similar to SAM, we discover that applying AT only during the final few epochs yields robustness comparable to training entirely with AT. With the switching method, which switches from SGD to AT at epoch on adversarial examples, termed robust error E (0), and its performance on clean examples, referred to as natural error Err..... (0). In Figure 6 (a), the robust error of 0 drops sharply after switching to AT 180 epoch, matching that of full AT training, while the natural error only slightly spikes. Additionally, we vary the switching point & to examine how the robustness of 0 changes with the increasing proportion of AT epochs, i.e., (T-t)/T. In Figure (b), the robust error of 0%r is comparable to that of models trained entirely with AT (i.e., when (T /T/T = 1) once the proportion of AT epochs exceeds O. These results together confirm that AT efficiently enhances the robustness of neural networks late in training.
obtain the model 0% AAT, We evaluate its robustness by classification error
6 CONCLUSION AND LIMITATIONS
In conclusion, we discovered an intriguing phenomenon that SAM efficiently selects flatter minima over SGD when applied late in training. We provided theoretical explanation for this by conducting a two-phase analysis of the learning dynamics after switching to SAM in the late training phase. Our results suggest that late training dynamics may play an important role in understanding the implicit bias of optimization algorithms. Since we extended our findings to AT with the switching method, a natural future direction is to generalize this novel approach to other optimization algorithms.
Limitations. We note that even if our theoretical setup closely reflects real settings, it still relies on several simplifications, such as assuming constant perturbation radius p and using squared loss. We also note that our theory only explains the behavior of SAM in the late training phase; it cannot account for the phenomena observed for early-phase SAM.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Alessandro Achille, Matteo Rovere, and Stefano Soatto. Critical learning periods in deep networks. In International Conference on Learning Representations, 2019. URL https: /ttp:///ppeeeeeeeeeeeeeee net/forumt 1d=3keSt 2
Atish Agarwala and Yann Dauphin. SAM operates far from home: eigenvalue regularization as a dynamical phenomenon. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 152-168. PMLR, 23-29 Jul 2023. URL http::ttt//.occ........tt /ttp://ooceceeeeeeoo ngs ml www.spressvv.. 02/ agarwala23a. 3
Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 247-257. PMLR, 17-23 Jul 2022. URL https: / pttpp//occeigg...e////p mlr. press / v1 62 ahn22a html. 22
Kwangjun Ahn, Ali Jadbabaic, and Suvrit Sra. How to escape sharp minima with random perturba- tions. arXiv preprint arXiv:2305. 15659, 2023. 30
Maksym Andriushchenko and Nicolas Flammarion. Towards understanding sharpness-aware mini- mization. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 639-668. 668. PMLR, 17-23 Jul 2022. URL https: /pppoceedinggss.e ngs..ngs.... ml www.sprrssvlllv.... 62 http:///..//ddd..ussseeeeeeeee/ 3. 9, 17, 18, 21, 25
html. 1, 2,
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. International Conference on Machine Learning, pp. 948-1024. PMLR, 2022. 8, 30
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Pa- pers), pp. 7360-7371, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10.18653/v1/20222 www.:9222...lloo..5.00....0 URL https: : hatpa/aaaanntocccoocceee org 20222222222222 acl-long 508. 1, 17
Peter L Bartlett, Philip M Long, and Olivier Bousquet. The dynamics of sharpness-aware minimiza- tion: Bouncing ravines and drifting towards wide minima. Journal of Machine Learning Research, 24(316):1-36, 2023. 17
Marlon Becker, Frederick Altrock, and Benjamin Risse. Momentum-sam: Sharpness aware mini- mization without computational overhead. arXiv preprint arXiv:2401.12085 2024. 17
Kayhan Behdin and Rahul Mazumder. On statistical properties of sharpness-awarr minimization: Provable guarantees. arXiv preprint /30..2332...88366, 2023. 3
Kayhan Behdin, Qingquan Song, Aman Gupta, Sathiya Keerthi, Ayan Acharya, Borja Ocejo, Gregory Dexter, Rajiv Khanna, David Durfee, and Rahul Mazumder. msam: sharpness-aware minimization. arXiv preprint XXiv:2302.09699 2023. 3, 7,
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by 483-513. PMLR, 2020. 17
Ubleebbek like process. In Conference learning theory, pp.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):2113357, 2015. 8, 30

---

Published as a conference paper at ICLR 2025
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong. When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning 1,
Representations, 2022. URL
17
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations, 2021. URL https: /opoereeeeeeeeeeeeeeeeeeee wewiew.rreee........ net mortmform?? ii jo- 22
Yan Dai, Kwangjun Ahn, and Suvrit Sra. The crucial role of normalization in sharpness-awart minimization. In shirty-seventh Conference on Neural Information Processing Systems, 2023.
URL htttps://oppnreeeeee wwt.:wwwwweeeeeeeeeee.e..eeeeeee n forum? id=zigg (40 FeRiA.. 17
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 0-27749-2274611 2021. 17
Terrance DeVries. Improved regularization convolutional neural networks with cutout, arXiv
preprint arXiv:
2017. 18
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference Learning Representations, 2021. URL https::/ /pw.r/eeee.ee.eeeeeeeee net/ / Eorum? idd iiddidiibbbbbie dNTTy. 18
Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan. Efficient sharpness- aware minimization for improved training of neural networks. In International Conference on Learning Representations, 2022, URL / /
net / forum? CeL=nOOCTTNNRRGN 17
Yu Feng and Yuhai Tu. The inverse variance-flatie relation in stochastic gradient descent is critical for finding flat minima. Proceedings of the National Academy of Sciences, 1119))22015117118. 2021. 6,
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL https: : / opwnreviewwww.eooooo.. net/ forum? d= 6tt.mmmmmmmmommmmmmmmmmm rM. 1, 3, 4, 17,
Jonathan Frankle, David J. Schwab, and Ari S. Morcos. The early phase of neural network training. In International Conference on Learning Representations, 2020. URL URLhttpp://pppeeeeviee net / Eorum? ww1.k1k11.P.........
Khashayar Gatmiry, Zhiyuan Li, Ching-Yar Gao Chuang, Sashank Reddi, Tengyu Ma, and Stefanie Jegelka. The inductive bias of flatness regularization for deep matrix factorization. arXiv preprint arXiv:2306.13253 2023. 17
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
doi: 10. 10.1109/CVPR.2016.0000 4. 18
Sepp Hochreiter and Jürgen Schmidhuber. Flat Minima. Neural Computation, 9(1):1442, 01 1997. ISSN 0899-7667. doi: 10. 10.116660 162/neco. 1997.9...19. 1997. 9. 1.
paattps://doi.. http:///dooo.ogoo/./eooo wtt.g//.....././/...
Stanistaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint 1rXiv:1711.004 2017. 17
Stanislaw Jastrzebski, Maciej Szymezak. Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories deep neural networks. arXiv preprint 2rX1v:2000..95522, 2020. 8
12

---

Published as a conference paper at ICLR 2025
Zhiwei Jia and Hao Su. Information-thert theoretic local minima characterization and regularization. In International Conference on Machine Learning. pp. 4773-4783. PMLR, 2020. 17
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020. URL
net wtte:/oormnnndddieeee
1,
17
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimizers work? In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: /openeeveeeeeeeeeeeeeeeeeee net/ forum? id=vDeh2yxIvui. 1,
Ryo Karakida, Tomoumi Takase, Tomohiro Hayase, and Kazuki Osawa. Understanding gradient regularization in deep learning: Efficient finite-difference computation and implicit bias. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 15809-158277 PMLR, 23-29 Jul 2023. URL https: http:///ooeeeeeeeeeeeeeeeeeecece mlr press / ka rak wetraakida222. da23a. html. 17
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasewwizz condition. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part 16. pp. 795-811. Springer, 2016. 8
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. URL https: : /toppppeeett
net pert/formm? i d= oyRR1 Ygg. 1, 17
Hoki Kim, Jinseong Park, Yujin Choi, and Jaewook Lee. Stability analysis of sharpness-aware minimization. arXiv preprint arXiv:2301.0633 2023. 17
Minyoung Kim. Da Li, Shell X Hu, and Timothy Hospedales. Fisher SAM: Information geom- etry and sharpness aware minimisation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 11148-11161. PMLR, 17-23 Jul 2022. URL https: pttpp//eeeeeee.cce/ee///oo www.presss..v1222... kim22f. html. 17
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6800
2014. 18
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
4, 18
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- minimization for invariinv learning of deep neural networks. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 5905-5914. PMLR, 18-24 Jul 2021. URL ht https: http:////cceeeeeeeeeeeeee mlr presss/ / v139 kwon2lb html. 17, 18, 19
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss land- scape of neural nets. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Asso- ciates, Inc., 2018 URL hwtph:tts////tttt httptt/ss///pccceeeiieiiooeeeeee neurips. oc wwwp//aapeeaeeeeeeeeeeeeeeeeeeee 23
ss// /paper/
2018 / 4 http:///222200000000066666666663
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss?-a mathematical framework. International Conference on Learning Representations, 2022. 17
Kangqiao Liu, Liu Ziyin, and Masahito Ueda. Noise and fluctuation of finite learning rate stochastic gradient descent. International Conference on Machine Learning, pp. 7045-7056. PMLR, 2021.
13

---

Published as a conference paper at ICLR 2025
Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You. Towards efficient and scalable sharpness-awart minimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12360-12370 2022. 17
Philip M. Long and Peter L. Bartlett Sharpness-aware minimization and the edge of stability. Journal of Machine Learning Research, 25(179):1-20, 2024. URL http: /tmpr//mmrr../.m// w.w..orrrr....o..oooooo.oooooo.l patpg//pps////////
v25/23-1222 2255/23-185... html. . 17, 22
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35:
34689-34708, 2022. 17
Chao Ma and Lexing Ying. On linear stability of SGD and input-smoothness of neural networks.
Advances in Neural Information Processing Systems, 34: 34:16805- 18155--6617,,,, 2021. 17
Chao Ma, Daniel Kunin, Lei Wu, and Lexing Ying. Beyond the quadratic approximation: the multiscale structure of neural network loss landscapes. arXiv preprint arXiv:2204.11322 2022. 8, 17
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. To- wards deep learning models resistant adversarial attacks. In International Conference on Learn- ing Representations, 2018. URL http:/::////peeeeeeeeeee www.neeeoooomndddaa.e.e.....
2,
4, 10, 18
Enea Monzio Compagnoni, Luca Biggio, Antonio Orvieto, Frank Norbert Proske, Hans Kersting, and Aurelien Lucchi. An SDE for modeling SAM: Theory and insights. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato. and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 352009252253. PMLR, 23-29 Jul 2023. URL https: /apps://proee http:///poceeeeeeee/e/ee ml www.mmlrrprrs... w2t./moo22iii htt//moniiiiccccc www.acoomaacccoooo i23a html. 17
Takashi Mori, Liu Ziyin, Kangqiao Liu, and Masahito Ueda. Power-law escape rate of sgd. In International Conference Machine Learning, pp. 15959-15975. PMLR, 2022. 6, 8
Maximilian Mueller, Tiffany Joyce Vlaar, David Rolnick, and Matthias Hein. Normalization layers are all that sharpness-aware minimization needs. In seheenh-seventh Conference on Neural Information Processing Systems, 2023. URL https:: /ttp:://///peeeeieeeee / forum? dd1AArl13yyy9 17, 18
Rotem Mulayoff, Tomer Michacli, and Daniel Soudry. The implicit bias of minima stability: view from function space. Advances in Neural Information Processing Systems, 34:1 14117499-1760 17761, 2021.
17
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, and Nati Srebro. Exploring generaliza- tion in deep learning. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, Vish- wanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, Curran Associates, Inc., 2017. URL https://prroeedddsss..uuu neurssss..eeuuuspo. www.eeewwis.s...peeeeeeeeeeeeeee ccppaper_ https//aaeeee/oo///ee////eaaeaas
Sungbin Shin, Dongyeop Lee, Maksym Andriushchenko, and Namhoon Lee. The effects of overpa- rameterization on sharpness-aware minimization: An empirical and theoretical analysis. arXiv preprint arXiv:2311.175399 2023. 9, 17
Dongkuk Si and Chulhee Yun. Practical sharpness- aware minimization cannot converge all the way to optima. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL ftttps://openrevee
net forum? innmmiddiii j JNOLHqM. 17
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015. Conference Track Proceedings, 2015. URL http: http:///wwwxvv..//b/eeesse/ee www//aa/////ooooooooooooo/.e.eee 409 1556. 4, 18
14

---

Published as a conference paper at ICLR 2025
Steven H. Strogatz. Nonlinear Dynamics and Chaos: With Applicctin Physics, Biology. Chemistry and Engineering. Westview Press, 2000. 6
Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus. Intriguing properties of neural networks. International Conference on Learning Representations, 2014. 10
Gal Vardi. On the implicit bias deep-learning algorithms. Commun, ACM, 66(6)) may 2023. ISSN 0001-0782. doi: 10. 10.1145/3570070 17
URL
1tt.://1..111115111/1 1,
Mingze Wang and Lei Wu. A theoretical analysis of noise geometry in stochastic gradient descent. arXiv preprint 2rXXiv2310000000 2023. 7
Mingze Wang. Haotian He, Jinbo Wang, Zilin Wang, Guanhua Huang, Feiyu Xiong, Zhiyu Li, Weinan E, and Lei Wu. Improving generalization and convergence by enhancing implicit regularization. arXiv preprint arXiv:2405. 2024. 17
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. How sharpness-aware minimization minimizes sharpness? In The Eleventh International Conference on Learning Representations, 2023a. URL https: // /pperrevieww net/ forum? -ieedddd5ppppp SpDgWmpp 6x. 17
Kaiyue Wen, Tengyu Ma, and Zhiyuan Li. Sharpness minimization algorithms do not only minimize sharpness achieve better generalization. Advances in Neural Information Processing Systems, 2023b. 17
Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type part it: Continuous time analysis. Journal of Nonlinear Science, 34(1): 166, 2024. 6,
Dongxian Wu, Shu- Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust general- ization. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 2958-2969. Curran Associates, Inc., 2020a. URI. https: http:///ooococcococccooccooccccc ngs. neurips wwtpcccppppceeeeeeeeeeeee www.el//////eeeeee...ee...eeeeee lellee/ppaa pessppppp//220/ Ptwp:/e2200-PPap..Pppppp fiiee//eeeccce22eeeeel1ee2222223
Jingfeng Wu, Wenqing Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes sgd. In International Conference on Machine Learning. pp. 10367-10376. PMLR, 2020b. 6
Lei Wu and Weijie J Su. The implicit regularization of dynamical stability in stochastic gradient descent. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 37656-37684, PMLR, 23-29 Jul 2023. URL https:t///// http///pooceeeeeeeee/ee/eeeeesee mwr. wttppppssss//////eess w202/w223222222 html. 17
Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018. 8, 17
Lei Wu, Mingze Wang, and Weijie Su. The alignment property of sgd noise and how it helps select flat minima: A stability analysis. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 4680-4693. Curran Associates, Inc., 2022. URL https://proceede/ neurips wwwprips...ppppprrr.e..... httpp/ppprfills/ppppeereee file /il1:/e5ccsarrd55c22aa2eeeeeeee Paaper Poperseeeeeeeeeeeeeeeeeeeeeeeee pdf.. 17
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint 160iv: 146. 2016. 4, 18
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017. URL https: : http///ppeeeeeeeeeeeeeeeee wwt.:/..eee..rreeeeeeeeeooooo htt.:oormmmmmmsyyddoooa d=Sy8gda9xxx 17
httts::/openrrvieee
15

---

Published as a conference paper at ICLR 2025
Yihao Zhang, Hangzhou He, Jingyu Zhu, Huanran Chen, Yifei Wang, and Zeming Wei. On the duality between sharpness-awen aware minimization and adversarial training. arXiv preprint arXiv:2402.15152. 2024. 1, 10, 17
Tongtian Zhu, Fengxiang He, Kaixuan Chen, Mingli Song, and Dacheng Tao. Decentralized sgd and average-direction sam are asymptotically equivalent. In International Conference on Machine Learning, pp. -3006. PMLR, 2023. 17
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. arXiv preprint arXiv:1803. 103.000803 2018.
Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, sekhar tatikonda, James Duncan, and Ting Liu. Surrogate gap minimization improves sharpness- aware training. In International Conference on Learning Representations, 2022. URL https: //openreview..o. net/ htt.:/oormmmmddddddamaaaaaaa
17
16

---

Published as a conference paper at ICLR 2025
A RELATED WORK
Implicit flat-minima bias. Extensive studies have been dedicated to understanding the implicit bias in deep learning (Vardi, 2023). One notable example the flat-minima bias: SGD and its variants tend to select flat minima (Keskar et al., 2017: Zhang et al., 2017), which often generalize better (Hochreiter & Schmidhuber, 1997: Jiang al.. 2020). Many works (Ma Ying. 2021; Mulayoff et al., 2021; Wu & Su, 2023: Gaimiry et al., 2023; Wen et al., 2023b) have theoretically justified the superior generalization of flat minima neural networks. Wu et al. (2018; 2022); Ma Ying (2021) have taken dynamical stability perspective to explain why SGD favors flat minima; in-depth analysis of SGD dynamics near global minima shows that the SGD noise drives SGD) towards flatter minima (Blanc et al., 2020: Li et al., 2022: Ma et al., 2022: Damian et al., 2021). Further studies (Jastrzębski et al., 2017: Lyu et al., 2022) have explored how training components, like learning rate and normalization, influence the flat-minima bias of SGD.
Understandings towards SAM. Inspired by the implicit flat-minima bias of SGD, Foret et al. (2021) has proposed Sharpness- Aware Minimization (SAM) algorithm. SAM and its variants (Kwon et al., 2021; Zhuang et al., 2022; Du et al., 2022; Liu et al., 2022; Kim et al., 2022; Mueller et al., 2023; Becker et al., 2024; Wang al., 2024) have achieved superior performance across a wide range tasks (Chen et al., 2022; Kaddour et al., 2022; Bahri et al., 2022). Despite their successes, their effectiveness remains poorly understood. On the optimization front, Andriushchenko & Flammarion (2022); Si & Yun (2023); Shin et al. (2023) have studied the convergence properties of stochastic SAM; Monzio Compagnoni et al. (2023); Kim al. (2023) have demonstrated that SAM requires more time to escape saddle points compared to SGD; Long & Bartlett (2024) have analyzed the edge of stability phenomenon for SAM. On the implicit bias side, Shin et al. (2023); Behdin et al. (2023) have examined SAM's flat-minima bias using dynamical stability analysis; Wen et al. (2023a); Monzio Compagnoni et al. (2023); Bartlett et al. (2023) have studied how SAM moves towards flat minima; Andriushchenko & Flammarion (2022) has showed that SAM has the low-rank bias; and Wen et al. (2023b) has revealed that SAM does not only minimize sharpness to achieve better generalization. Dai et al. (2023) has characterized the role of Normalization in SAM. In comparison, our work focuses on the dynamics of SAM in the late training phase, a topic largely unattended in the literature, offering both theoretical and empirical analyses of SAM's convergence, escape behavior, and bias toward flatter minima. Additionally, several studies have explored the connection between SAM and other algorithms: Zhang et al. (2024) explored the duality between SAM and AT; Zhu et al. (2023) demonstrated the equivalence between average- -vrrage-direction SAM and decentralized SGD. Another related work Jia Su (2020), which introduces a sharpness-regularized training algorithm with formulation similar to SAM. Notably, its ablation study demonstrates the algorithm's 's effectiveness when applied in the mid or late training stages. Karakida et (2023) has analyzed the implicit bias of Gradient Regularization (GR), which shares similar form with SAM, and showed that GR tends to select solutions in rich regimes.
17

---

Published as a conference paper at ICLR 2025
B MORE EXPERIMENTAL DETAILS AND RESULTS B. DETAILED EXPERIMENTAL SETTINGS
In this subsection, we introduce the detailed experimental settings.
In this work, we mainly compare the training dynamics of SAM and SGD. In fact, SAM and SGD share almost identical training settings, with the only difference lying in that SAM introduces an extra
perturbation radius p. Therefore, we first outline the basic training settings for training with SGD across different network architectures and datasets.
VGG-16 and ResNet-20 on the CIFAR R-10 10 dataset. We train the VGG-16 (Simonyan & Zisserman, 2015) and the ResNet-20 architecture (He et al., 2016) on the CIFAR-10 10 dataset (Krizhevsky et al., 2009), using standard data augmentation and hyperparameters as in the original papers (Simonyan & Zisserman, 2015; He et al., 2016). Data augmentation techniques include random horizontal flips and random 32 x 32 pixel crops. A weight decay of 5 x 10-4 is applied, and the momentum for gradient update is set to 0.9. The learning rate is initialized 0.1 and is dropped by 10 times at epoch 80?. The total number of training epochs is 160.
WideRes on the CIFAR-10 dataset, and Wide 10 the CIFAR-ION dataset. We train the WideResNet- 16-8 architecture (Zagoruyko & Komodakis, 2016) on the CIFAR-10 dataset and the WideResNet-28-10 architecture on the CIFAR- 100 dataset (Krizhevsky et al., 2009), following the standard settings as in the original papers (Zagoruyko & Komodakis, 2016). Data augmentation techniques include random horizontal flips. random 32 32 pixel crop, and cutout regularization Vrevs,es, 2017). A weight decay of 5 x 10-4 is applied, and the momentum for gradient update is set to 0.9. The learning rate initialized at 0.1 and is dropped by times at and 120 epochs. The total number of training epochs is 200.
Then, we specify the extra hyper-paramete for SAM, i.e., the perturbation radius p.
Additional setup for SAM, ASAM and USAM. For SAM, the perturbation radius is set to 0.05 across all architectures and datasets, as recommended by Foret et al. (2021). In Appendix B.3, our findings SAM generalize to another SAM variant, ASAM (K (Kwon et al., 2021). For ASAM, the perturbation radius is set to 2.0, as recommended by Kwon et al. (2021). In Appendix B.4, we further generalize our findings to USAM (Andriushchenko & Flammarion, 2022) verify the simplified SAM update rule used in our theoretical analyses. For USAM, the perturbation radius is set to 0.1.
In Section 5, we also extend our findings to Adversarial Training (AT). AT shares almost the same settings as SGD as well, except that AT trains neural networks on adversarial examples generated from the original images. Here, we specify the additional settings for creating adversarial examples.
Additional setup for AT. We considered untargeted PGD attack (Madry et al., 2018) to generate adversarial examples x Adv from the original image æ, with the constraint that ||a,Ady
co S e, where ||- denotes the _ norm. Following the standard configurations of Madry et al. (2018), set the perturbation limit to 4/255 and the step size for PGD attack to 2/255.
Additional setup for experiments on VITs. In Figure 8, we train the VIT-T/S architecture (Doso- vitskiy et al., 2021) ) on the CIFAR-I00 dataset, following the settings in Mueller et al. (2023). Data augmentation techniques include random horizontal flips, random 32 x 32 pixel crops and Auto ment. The early phase optimization is done with Adam Kingma (2014), with a constant learning rate 1 10 A weight decay of 5 x 10-4 is applied, and the batch size set to 64. The late phase optimization is done with SAM, with an extra hyperparameter, the perturbation radius, set to 0.1. The total number of training epochs is 200.
B.2 MORE EXPERIMENTS ON LATE- PHASE SAM
In this subsection, we provide more experimental results under various settings to demonstrate the effect of SAM on generalization and sharpness when applied during the late phase of training.
NNtte that in the original paper, the learning rate reduced by factor of at epoch 120 well. However. based on our theory in Section 4.1, the learning rate affects the solution sharpness (see Table 2). To fairly compare the sharpness of solutions found by the switching method, which switches at different epochs, we remove the learning rate decay during the late phase of training. The same applies for other architectures and datasets.
18

---

Published as a conference paper at ICLR 2025
ResNet:20 CTRAR-10
VGG-16 CIFXR-10
Gevenlization gap between
0.015
8amn SAM, and AAM
0.035
0.12
0.012
0.08
0.032
0.009
0.04
0.006
0.02
0.003
0.003
0.00
0.000
-3.02
-0.003
-0.03
-0.003
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.8
1.0
Shappeess of -
SAM braining proportion when obtaining 0
T-t
Digs (0500-SAME)) Ligest (AM))
Erries (SDD-SAM:
Rer puss (o
[0GGo--s Ill
Figure 7: Additional experiments for VGG-16 and ResNet-20 CIFAR-10: the impact of SAM training proportion on generalization/haarpnes when switching from SGD to SAM. The generalization gap between the models 0 -SM, and 0 (top row) / the sharpness of OSGD-SAM, (bottom row) vs. the SAM training proportion of OSAD- Dots represent the mean over five trials with different random seeds, and bars indicate standard deviations.
VII-T on CIFAR-100 100
VIT-S on CIFAR-100 100
Generalization gyip betveen
0.20
0.05
0.30
0.040
0 Adam--Shhh and 0SAM
0.36
0.04
0.18
0.074
0.04
0.01
0.008
0.00
0.00
-0.04
-0.01
0.00
0.000
0.2
0.4
0.6
1.0
0.0
0.2
0.4
0.5
SAM taimingg proportion when obtaining 0/
Goneralization gap. D Te)
Adam -SAM.)
Lorere (0)AM)
ErT (o -AMM.)
ErT (0%AM.)
Figure 8: Additional experiments for ViTs on CIFAR-100: the impact of SAM training pro- portion on generalization when switching from Adam to SAM. The generalization gap between the models /aam-sSAM and ObAM vs. the SAM training proportion of AAdm--SAM.. Dots represent the mean over five trials with different random seeds, and error bars indicate standard deviations.
In Figure 7, we include more experimental results for ResNet-20 and VGG-16 on the CIFAR-10 datasets.
In Figure 8, we include experiment results for ViT-T/S on the CIFAR-100 datasets. Notice that due an implementation limitation in Py PoTrch, computing the gradient twice consecutively for ViTs is not feasible. Thus, only generalization results are presented. Additionally, as Adam a more common practice for training ViTs, we replace SGD with Adam before the switch.
In Figure 9, we include further experimental results using the squared loss. Figure 9, both the generalization gap and sharpness follow trends consistent with our previous findings in Figure 3, thereby bridging the gap between our theoretical and experimental settings.
B.3 EXTENDED EXPERIMENTS ON ASAM
In this subsection, we generalize our findings on SAM to another SAM variant, ASAM (Kwon et al., 2021). Similar to SAM, find that applying ASAM during the late training phase, after initial training with SGD, achieves generalization ability and solution sharpness comparable to full ASAM training. Still, apply the switching method to investigate the effect of the late-phase ASAM
19

---

Published as a conference paper at ICLR 2025
VG CIFAR- (Squared Loss)
VG CIFAR- 10 (Squared Loss)
Generalizztion gep betwoen
0 3500- and 6a0s
1.6
0.04
1.2
0.20
1.2
0.018
0.9
0.055
0.6
0.000
0.4
0.006
0.3
0.005
0.0
0.000
0.4
0.006
-0.3
0.005
-0.8
-0.012
-0.100
Shrppmess of
500
100
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
SAM branning propartion when obtaining -Sn-
T-t
Enuse (OSDD-SMM))
(aase)
ErT (OSCD-SAXL.)
ErTn
(Aa)))
|| (ODD--sax
Figure 9: Additional experiments using squared loss: the impact of SAM training proportion
generalization when switching from SGD to SAM. The generalization gap between the models i and 0 the USAM training proportion of 0gmp -SMM.- Dots represent the mean three trials with different random seeds, and error bars indicate standard deviations.
ResNet on CIRAR-10
VGG- 16 CIRAR-10
Generaiizaiin gap between
ASAM, ASAM
0.15
0.042
0.20
0.035
0.08
0.32
0.028
3c0 03DD-ASAM:
0.09
0.021
0.08
0.007
0.04
0.007
0.00
0.00
0.000
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
2.1
0.5
0.
1.0
Sharpeess
050D ASAME
triiig proportion when obtaining 05cD SSDM-
T-E
WideResNet 28- 10 CIBAR 100
WideResNett 16-8 on CIFAR-10
Genesalization gep between
Sco ASAA, and ASAM
0.30
0.00
0.060
0.025
0.25
0.05
0.032
0.020
0.15
0.03
0.016
0.010
0.000
0.05
0.61
0.00
-3.008
00.00
00000
0.0
0.5
0.0
0.6
Sharpness of
S ASAM.
78
30
20
0.0
0.2
0.5
1.0
0.0
0.2
0.4
0.6
0.8
1.0
ASAM proportion obtaiinn 0eop SAM
Lyrem (Raam))
Erip (0 (0am--ssAM)
Ert (8)
|| (P Ill
Figure 10: The impact ASAM training proportion on generalizattinnsharpees when switching from SGD to ASAM. The generalization gap between the models 03m--A and T (top row) / sharpness of 05D--ASAM. (bottom row) vs. the ASAM training proportion of 05aD-ASSMM.. Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations.
20

---

Published as a conference paper at ICLR 2025
WideResNNtt 28 10 CTBAR- 100
WideReeNee 16-8 CTRAR-10 10
Geveralizztion gAp betwenn
0500-118A3 and 0fisAM
0.25
0.05
0.000
0.15
0.20
0.04
0.032
0.02
0.20
0.02
0.16
0.006
0.05
0.01
0.000
0.003
-0.05
-0.01
-0.008
0.003
0.0
0.0
0.0
0.0
Sharpeess of
0sen -15--
25
0.0
USAMuaining proportion when obtainingi SCD-
T-t
Li (OSCD-USAMz) Aaus
(o(s))
BrTn (ODD-
Erun ( (Sssa)))
I|R ( Sooco--ss.x.
The impact of USAM training proportion on generalization when switching from SGD to USAM. The generalization gap between the models -UGD- and o vs. the USAM training proportion of US Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations.
Figure 11:
across various architectures and datasets. We use the switching method, which switches from SGD ASAM, obtain the model 6 A.a.-A We also train baseline model, i.e., 0% SSM.- We vary the switching point t while keeping 7 fixed to adjust the proportion of ASAM training and study its impact on the generalization gap between ooA and - as well as the sharpness of ASAMA.. As shown in Figure 10, both the generalization gap and the sharpness drop once the ASAM training proportion exceeds zero, implying that ASAM, even when applied only during the last few epochs, achieves comparable generalization performance and solution sharpness as training entirely with ASAM. These results further highlight the importance of studying the implicit bias of different optimization algorithms towards the end of training.
B.4 EXTENDED EXPERIMENTS ON USAM
In this subsection, we generalize our findings on SAM to another SAM variant, USAM (An- driushchenko & Flammarion, 2022), further justify the simplification made on the SAM update rule in our theoretical analysis. With small modification to Equation (3), the update rule of USAM with stochastic gradients
0 = 1 (Ø' Cg.
(6)
Consistently with the main paper, we measure how the generalization gap between O SSDD-USAM,
and USAAM as well as the sharpness SSDD- changes as the USAM training proportion
- increases. In Figure 11, both the generalization gap and sharpness drop sharply once the USAM training proportion exceeds zero. This finding validates the effectiveness of USAM in capturing the main behavior of SAM, further bridging the gap between theoretical results and empirical
findings.
B.5 MORE EXPERIMENTS ON EARLY-PHASE SAM
In this subsection, we provide more experimental results demonstrate the effect of SAM on generalization and sharpness when applied during the carly phase of training.
The effect of early-phase SAM on the generalization ability. In Figure 12, we include more experimental results for ResNet-20 and VGG-16 on the CIFAR-10 datasets. Similar our results in Figure 5, the generalization gap between 0AMM- and 0 remains substantial until the SAM training proportion exceeds 0.6. The same applies for the sharpness of 0f These results further validate that applying SAM only during the early phase offers limited improvements on generalization and sharpness of the final solution compared to full SAM training.
21

---

Published as a conference paper at ICLR 2025
ResNet on CHARR-1P
VGG- 16 CIARR 10
Geveralizaion gep bebowenn
0.010
0 Saat--oohy B
0.16
0.000
0.008
0.12
0.006
0.12
0.006
0.04
0.002
0.04
0.002
0.00
0.000
0.00
0.000
-0.08
0..04
-0.05
0.004
0.0
0.2
0.0
1.0
Sherpesss of 0SAM--ADDE
0.0
SAM training propartion when
whopattion.hem.cccaiiiit obbtaiiig -SAM-S
Dreest (AMM-)
(OSA))
Ert (5Ax--SDD..)
ErTree (ase.)
|| (Sxxt-aa )l22
Figure 12: Additional experiments for VGG-16 and ResNet-20 on CIFAR-10: the impact of SAM training proportion on generalizationShaarpees when switching from SAM to SGD. The generalization gap between the models 0'AM--GDD. and 0 (top row //th sharpness of 65AM (bottom row) vs. the SAM training proportion of 0SMM- Dots represent the mean over five trials with different random seeds, and error bars indicate standard deviations.
Fully Mally--oonecteelnne net on CIFAR-10 5k subset
200
Sharmees ||/(0)]12
150
SAM-GD. 1- 10000
100
SAM-GD, 5000
50-
2000
5000
Iteraiio
Figure 13:
The impact of (early-phase) SAM on the sharpness evolution. The sharpness evolution curves for fully-connected nets on CIFAR-10 5k subset trained with different strategies. The blue and purple baseline curves represents training entirely with GD and (full-batch) SAM, respectively. The red, green and orange curves represent the switching method, which switches from GD to (full-batch) SAM at different switching point I, with / E (2000, 5000, 10000]. A constant learning rate vj = 0.01 (2/n 200) is used for both GD and (full-batch) SAM.
The effect of carly-phase SAM the sharpness evolution. Other than focusing on the sharpness of the final solution, we are also interested in the effect of early-phase SAM on the evolution process of the sharpness.
Furthermore, we explore how sharpness at the iterator evolves when switching from SAM to (S)GD, Recent works (Cohen al., 2021: Ahn al., 2022) have shown that when training neural networks with Gradient Descent (GD), the sharpness H (0)ll, tends to increase monotonically until it arrives 2/n. where r the learning rate. This is commonly known as the Edge of Stability (EoS) phenomenon". Specifically, we follow the setup of Cohen et al. (2021), measuring sharpness evolution for models trained with the switching method, which switches from (full-batch) SAM' to GD, as well as models trained entirely with GD and SAM. In Figure 13, we observe the EoS phenomenon for models trained entirely with GD. However, for models trained entirely with SAM, sharpness initially decreases, then slightly increases before converging to value much smaller than 2/7. As for models trained using the switching method, sharpness evolution mirrors that of full SAM training up to the switching point t. After f, the sharpness rapidly increases and then oscillates around
Despite Long Bartlett (2024) noted an "edge of stability" for SAM, here, the Edge of Stability phenomenon refers specifically to GD. In this experiment, we adopt the SAM with full-hatch gradient for fair comparison with GD.
22

---

Published as a conference paper at ICLR 2025
WideResNet- 16-8 on CIFAR-I0: f), Dreat ( (OSDD-SAM. + ad + f [6
1 Random vector yandom..
2 Random vector gandoma.
3 Random vector random.
-1.00
0.2 0.4 0.6
0.2
0.4
0.6
0.8
0.0
0.6
0.8 1.
a
0,8 0): SAM
1, 0): 0
Figure 14: 2D visualization of the loss landscape around SGD and SGD-to-SAM solutions. The orange dot (a 1, 3 0) denotes the SGD solution, and the orange triangle denotes (a 0, 0) denotes the solution obtained by switching from SGD to SAM. The contour plot represents the loss landscape, with darker regions indicating lower loss. The random vectors sampled three times.
WideRResNet 10 on CIFAR-100
1.00
Daosst (ASAM-SDD.)
Err (ASAM-SGDA
0.95
0.90
0.29
0.85
0-0.02
0.19
p=
0.00
0.05
0.18
0:00.00
0.2
P=02
SAM training proportion when obtairing SaM--son
Figure 15: The impact of perturbation radius generalization when switching from SAM to SGD. The test loss (left) / test error (right) of the model 0 AMM- vs. the SAM training proportion. Different perturbation radius E {0.02, 0.05, 0.1, 0.15, 0.2] are chosen for SAM. Dots represent the mean over five trials with different random seeds, and error bars indicate standard deviations.
2/1, regardless of how late the switch occurs. This implies that early-phase SAM cannot prevent the emergence of the EoS phenomenon. Once switching to GD, the iterator immediately escapes the flat minima found by SAM, resulting in relatively sharper minima. Together with our findings in Figures 5 and 12, we firmly say that early-phase SAM has a limited impact on both generalization and sharpness of the final solution.
B.6 OTHER EXPERIMENTS
A 2D visualization of loss landscape Other than the 1D interpolation in Figure 4 (c), we provide 2D visualization of the loss landscape, using the approach proposed by Li et al. (2018). In this setup, we center the visualization at o SGD-+SAM. and select two directions: òtarget and 8aadom- Specifically, Starget 0 S SGD random direction sampled from (O. I,)) with layer-wise normalization (Li et al., 2018) applied. We then plot the loss landscape as function f(a, 8) L(O 1.G.--SAM.. + @ (.58aadom) ). In Figure 14, it is clear that 0LCD- and 0gGD stay within the same valley. Moreover, the contours wider compared to 0g indicating 0g SSDD-SA.. corresponds flatter
pointing towards the SGD solution, and 8random is
around T BGDDSSAM, minimum. This two-dimensional approach provides additional insights into the structure the loss landscape and further supports two key claims P2 and P3.
The effect of perturbation radius on generalization when switching from SAM to SGD In our theory (see Table 2), the perturbation radius p influences the sharpness of the global minima selected by SAM. Specifically, larger perturbation radius leads to the selection of flatter global minima. In contrast, SGD's behavior is independent of the perturbation radius. However, discussed in
23

---

Published as a conference paper at ICLR 2025
16d CIFAR-10
10"4
116,,
0.10
Catasin (0 (SAAA)
520-EAM
0.05
0.00
10-3
1000
2000
3000
4000
2000
3000
Stept
Step
(a)
(b)
Extending Figure (c) to include more update steps. Train/test loss LD SODD v.s. the update step t'. Here, t is the switching step, chosen to sufficiently large for SGD to converge, and t' is the number of updates after switching to SAM. The red line represents the mean over 100 trials with different randomness.
Figure 16:
conjecture in Section 5, the properties of the final solutions are primarily shaped by the optimization method chosen the late training phase. Consequently, when switching from SAM to SGD, the generalization properties of the final solution are expected to be independent of the perturbation radius, as SGD dominate in the final stage of training.
To further explore this, we have conducted experiments to investigate the effect of perturbation radius on generalization when switching from SAM to SGD. Specifically, we varied the perturbation radius E (0.02, 0.05, 0.1, 0.15, 0.2) for SAM and measured how the test loss and error of SAM-SGD, change the SAM training proportion increases. In Figure 15, observe significant differences in test loss еπтоr across different values, even when training with SAM for most of the time (e.g., 0.8). Notably, only full SAM training (i.e., = shows a pronounced dependence on where a larger perturbation radius leads to better generalization. These findings align with our expectation that the perturbation radius has minimal impact on generalization when transitioning from SAM to SGD, further validating our theoretical results in Section 4 and conjecture in Section 5,
Extending Figure (c) to include more update steps. Figure 16, we observe that the training loss initially increases exponentially fast and then gradually decreases, consistent with our theoretical predictions.
24

---

Published as a conference paper at ICLR 2025
C PROOFS IN SECTION 4
In our theoretical analysis of Stochastic SAM, defined in Equation (4), we adopt the technical simplification in Appendix C.2 of Andriushchenko Flammarion (2022), which assumes that the inner noise and outer noise are independent. Specifically, we analyze the following variant of Equation (4):
0 = 0² nL: Lz; /V (0'))
where E and & are two independent stochastic mini-batches of data of size B.
C.I PROOF OF THEOREM 4.1 AND COROLLARY 4.2
For simplicity, we denote the two-step update rules of SAM defined in Equation (7) as:
22+1/2 0º + pV Cg)
0 = L (0
For B-batch gradient noise, denoted by (0) = (g VC(0), we have:
1 FlE(O)E(0)T]
T (0)5 (0)
B =()))
(8)
B
where & are the sample indices in a mini-batch and {()) is the 1-batch gradient noise, used in Assumption 4.1.
For the linearized model EEquation (5)), holds that:
1 (0- (0)=)(0))(-- 07).
(())
2
Without loss of generality, we can let 0 0. For simplicity, denote G ((*)).
Next, according to the definition of linear stability (Definition 4), we consider the dynamics of stochastic SAM the linearized model.
First, we have the following estimate for [C(+++)]]:
EE [[]
E (0' 1/V (0++1/2) - (12
2
(00 -V2((++1/2))) - T (0++1/2) (
G(0 (12 - n
2
(0'
VVC (VC(0)1))))))) +0
2 n (0 T G (0+4/1))
(0' - VV G(O' (V((+111))))
1/2 +
Tr (5
E
2B
Assumption
127 7 [((041/3)]
172y || ((0))1/2)]
0+
Remark C.1 (Clarifications E). The expectation E is taken with respect the random variables (el,52), (t 2 0), which are the indices of independently random examples sampled from Dursin-
Remark C.2 (0++1))) and 0* L are uncorrelated). Observing that 0' (( @(g), e? { and g'+1/2 (EE). ...... ..!), we the law of total
expectation to derive
B (0 - ((
25

---

Published as a conference paper at ICLR 2025
=B [3 [e (0+;1/1)) - - - si ... ...]]
-Ree.st.-...
c (V st. ,)))
Ve( (....................)...
,)),,,,,,,,,,,,...... Mt s..... .tt) ...
-Bej.ef....
[0]
Then estimate the lower bound for EE [c(O
[(0)+113)]
E (0' -VVL:: G(O* +VLL
E [00 + + (E G(0 + (E)))
1 , (0') GE
E [00 +VVC(0')) "VVL(O'))GG(O G( (VL((')) +
2
E [00* + +VVC(0')) TG((') +V + /² E
IT (E())G))
2B
Assumption
B [10' +VVL(O)) G(O' +VV(()))) + p²
|GI|
E [c(0']]
2
B
+ GG [C(0')]
- E o (I+ (I- + PG)
B
(4)
2 2 (0r G0r)
p²y IGG [C(0']]
+
B
p²y
|GG|F
= 1+ B
E [[(0')]
where holds due to (I PG) + pC) G 2p0F p3 GG
By combining these two estimates, we obtain:
[(()))) 2 y7²y ||
p2xx G
1+
E [C()')
B
B
Hence, 0* is linear stable for stochastic SAM, it must hold:
2 ||G
(1 1+ p²y |IGI|
< 1,
B
which completes the proof of Theorem 4.1.
Furthermore, if the condition (9) can not be satisfied, i.c.. exponentially fast escape holds:
(1+
p'>|GG
> 1, then the
B
[(0') 2 C [c())]]
where C - 72y 1Gll
(1
? B
1. This complete the proof of Corollary 4.2.
C.2 PROOF OF PROPOSITION 4.1
As we mentioned in Section 4.2, this proposition focus on model size of p 1 and consider the full-batch SAM, i.e., g++1= v PVL(('))
(10)
26

---

Published as a conference paper at ICLR 2025
Proof of Proposition 4.1. First, we define the hitting time:
T: inf i N: | bj.
We aim to prove that 7 does not exists, which implies that 0c E (-b,b) C V holds for all & E N. Assuming T exists, then one the following three cases must hold:
Case I: 0² = 0. In this case, L' O, which implies that 0++1 = 0. This contradicts the definition T.
Case II: < 0T gT+1/2 = HT + pL"(t"), we have:
By the sub-quadratic property and p < 1/a, for the inner update
0 0² < 0 o +0/ o" o= + paol S 20 2b.
=00 (=)
Consequently, we estimate the two-sided bounds for 0=+1
The lower bound for @T+1;
0T -0" - -C(((0+1/3) > 00-17 sup C'(z)
=E(U.26)
20 max
f'(z) -b.
=E(-2b.20
The upper bound for 7T
@T+1
LT 11' T/c' (0T+1/22 TT ( ()
TT + nn nn''(67 (0T(((1112)) (
(8T+1/2) TT <
To justify we prove p < 17 (2160-11/3 when L' L'(0T) > 0 (noticing 8T C'(0T) 2 by sub-quadratic property) Since 0² 0111/2 20 we aconsider two scenarios:
C'(6")
If fort1/2 (z)dx S 0, then
L'(0T+1/2) L'(20T) L'(OT)
JgT41/2 (2) L'(TT)
C2(20T)
(ATT)
(2(2)) ('(x)
we obtain p < ('(0211/2)
Using the condition p < min
0<|=/<>
(((07)
If 22 1/2 ("(z)d= > 0. then since the sub-quaratic property L"((=z)) > ((2(=
for |æa| I [22), we have ("(0T+1/2) > O. Thus, (²) 2 L" (0T++1/22 > for
gr 11/2 implying:
f gT+1/2
=
= L(
(=) > C'(0T).
50T
C100++112 (((0))
Using the condition < 1), we obtain
('(87+1/2)
Therefore, have shown p 1) (((6T) implying
Thus, we obtain | < b, which contradicts the definition of
Case III: - 0² < 0. The proof for this case extremely similar to the proof for Case This case is also contradicts the definition of T.
Thus, we have proved that 7' does not exist, which implies i.e., 0¹ E (-b.b) C...IEN
27

---

Published as a conference paper at ICLR 2025
The next result provides example for the highly sub-quadratic landscape. Intuitively, it means that the decay of C ((2)) is extremely fast. Example C.I (highly sub-quadratic landscape). Let L(() (C) 0, (") 0, and
E= |z| < ac
(2)
else in [-2b, 2b]
where 0 < 1 is constant to reflect the sub-quadratic degree.
Notice that this example satisfies Definition 4.2. Moreover, is straightforward that
C'(2)
2
2c
1=1
|=| ae
else in [-2b, 26)
Additionally, we can calculate the upper bounds for and p in Proposition 4.1:
b
2b
1 S min
24
[=)=
EG
L' (22)
P 2
min bs
min
a 2
(z)
One can see that in the highly sub-quadratic case, i.e., 0 < e & 1, the upper bound for goes to 400.
In addition, we prove Proposition C.1, which illustrates that SAM will escape the current valley under quadratic landscape. This result stands stark contrast to Proposition 4.1 for sub sub-quadratic landscape, which further confirms our adaptation of the sub-quadratic landscape.
Proposition C.1 (Non-local escape behaviour under a quadratic landscape.). Assume the landscape in the valley V,, (-2b, 26) be quadratic, i.e., L(2) az? /2. Then, for all initialization 6 Vo VoY(0), and Th 8.1. v > a[I+apy. full-batch SAM will escape from the valley Vn, i.e., there exists > 0 s.t. 8T & Vb.
Proof of Proposition C.I. By straightforward calculation, we have
a++1 = 0² na (0' + -p al*) = (1 - 1jarjaa1-0 + ap))
17
2 a[I+ap] then |0|| 2 1 na(1 + 100) where [1 na(1 ap) > 1. This means that
there exists T S.t. 0" $ V..
Remark C.3 (multiple valleys in Proposition 4.1). Proposition 4.1 assumes the landscape is sub- quadratic within the valley = 26, 26j, but it does not impose any assumptions about the landscape outside V. Therefore, additional valleys can exist in (-o0, 22b) U (2b, +20). For example, consider
m0
(()) = 1 - cos
2b
which satisfies the sub-quadratic property within V many valleys:
[-2b, 26). This loss function has infinitely
[2(2]: 1)6, 2(2k + 1)b), k Z.
By Proposition 4.1, SAM remains in the current valley V [b 2b} and cannot enter into other valleys.
Proposition C.2 (Extension of Proposition 4.1). Under Definition 4.2, assume the landscape is
sub-quadratic in the valley V = (-, 2b). Then, for any E (0, 1) and all initialization / E {--(2
C' C'
e)b, (2-e)6), and n, s.t. 7 min(2-e)6/l( p S min zeV
al2 17, min 0<|=|<((-<66
(7
the full-batch SAM will remain within the valley V, iee,, 0f V for all i E N.
Notice that Proposition C.2 generalizes Proposition 4.1 almost all initializations within the valley
(-2b, 2b). By choosing sufficiently small e, the result effectively applies to all initializations within 2h, 2b).
28

---

Published as a conference paper at ICLR 2025
Proof of Proposition C.2. This proof is highly similar to that of Proposition 4.1. First, we define the hitting time:
T inf { N 2 -)bb
We aim to prove that T' does not exists, which implies that 0 E (-(2-e)o (2 (2-)) C V holds for all N.
Assuming I exists, then one of the following three cases must hold:
Case I: 0 = 0. In this case, L'( = O, which implies that 0T+1 0. This contradicts the definition of T.
Case II: 0 < 0r < (2 - -)).. By the sub-quadratic property and S al2-c}' for the inner update 0T+1/2 = TT + pL'(6T). we have:
0 0t < 0" 0² +0l (2)2) < 0" + pao0 S
2 0² < 2b.
Consequently, we estimate the two-sided bounds for +T
The lower bound for 0z+1,
8 -0"- (f'(00+1/2) 0- sup C'(2)
=E(0.26)
20 - max
C2(2) 2 - (2 e}b.
=E( =E(-26,20)
The upper bound for 02++1
gT+1 0" - C1C' (0T+1/2) < < + +P'((0") nn
( - + n (0T+1/2) -n (0T+122)))) = -0 < <22 - c]]b
To justify (#), we prove p < ('(07+112) when C'(0T) > (noticing 6T C'(OT) 2 by sub-quadratic property) Since 0e 01 < --c
2 0 we consider two scenarios:
0 If Jarin/2 1/2 (z)dd < 0, then
L'(0T+1/2) L'(TT)
C'(2-00) fgr-1/2 (z(2d -'(2-00')
L'(0T)
L'(0T)
('(0T12/2)
Using the condition p < < min
we obtain p < x1 ('(0T)
(<)=)<(2-e)b
(1)
If 00
1/2 (2(2)d2>> > 0, then since the sub- qubrauaratic property ("((z))) 2 L' (122))
for |li I |22), we have (²(0++1/2) > 0. Thus, L² 2 L" L"(( (0T+1/2) > for
gT < z < 11/2 implying:
f 17t1//2
= - +
(² > C'(0²).
J0T
(((0T41/2)
Using the condition p 17, obtain p < < (((TT) Therefore, we have shown p r) (()
('(6T+1/2)
implying (t).
Thus, we obtain |0 < b, which contradicts the definition of
Case III: b 8H < 0. The proof for this case extremely similar to the proof for Case II. This case is also contradicts the definition of T.
Thus, we have proved that 7' does not exist, which implies i.e., 0² e (-b,b) C...EEN
29

---

Published as a conference paper at ICLR 2025
C.3 PROOF OF THEOREM 4.3
For simplicity, we denote the two-step update rules of SAM defined in Equation (7) as:
g = 0 + [V
0++ 0 VV
Then we list a few useful properties under the condition of Theorem 4.3. Lemma C.I. Suppose Assumption 4.2. Then we have:
21. ()) = |V v L()), A0.
(11)
p
Proof of Lemma C.I. The upper bound is a direct corollary of PL. and smoothness. Please refer to Lemma B.6 in Arora et al. (2022) and Lemma in Ahn et al. (2023) for proof details.
The smoothness in Assumption 4.2 also implies the classical quadratic upper bound for the loss (Bubeck et al., 2015):
I. L(O u) & ((0) + (VL(0), v) + 2 ||012 v0, v.
(12)
Additionally, Assumption 4.3 holds for the gradient noise {(0) with batch size 1. Thus, for B-batch gradient noise, denoted by - (0). we have (0) (00) ] (()(())(()))) implying:
S o² [()), 4O.
E
(13)
For the learning rate and the permutation radius we choose then as follows:
μB
HB
n S min
p P min
(14)
2L 2L2
4L 4Lo² 241
C.3.1 PROOF OF THEOREM 4.3
For 1/0)+1/2) it has the quadratic upper bound:
((0 = L (0 + (V (0)) C (0' + (12) (0' +VV((()) + (e' +VV(())) PE(O') + D²I
+ +
|| " (1
2
Taking the expectation, we obtain the upper bound of [[C(0++/33)]
E B[c(0+1/2)] = E [c (0'
+VV(()))]
p²L B
||5 (00) [²]
2
(1))
S L (0' + +VV((')) 0²L
[C(O')]
B
E
C (0') + =0
+
L²L
||
+
0 LL
E [(()))
2
2 B
1+ 0(2 + PL)1?
12102
E [[(()')
2B
(4))
1+
E [[((0'] < (1 11+ T E [c('))
5 E [c(0')]
u
In the similar way, for L(0++1), it has the quadratic upper bound:
((1'+2) (0' - LeV (0(+1/2)) = L (o' - -
30

---

Published as a conference paper at ICLR 2025
(12) S (0' - -VC(041133)
(VC (V0(00-1V20000 (0'-n1((0+112) (100000)))))0)0)11))))
L | ( + 2
2
Taking the expectation, we obtain the upper bound of [E[((O++)))
E [[((t(t)] & E ( (- -V +
72L
<
[8(0+11/))
E
2
c (0' (  0²
(1)
[(((+113)]
S
B
B
(19)
K ((0') 1, (VC(O"), VL(0') 112)
L/L
VC(0++122)
2
n2 o²
E [(0++1/3]
+
2
2
< (C()) (4 V2( -210011121)
(
17 TTTT + 12
[ [(0)+11/2)
+
IR
(C(0) 7 (00 VC 1
(4)
4 T [c(Bt)1/2) ]
3111 B [
<E
[ (0) 7 / VC(0))
4
(1)
[ 7/ -/0== TdL |V
3np
c(0++1/3)
SE
4
(1)
<
(00)
31
|V((0)) ||]
+
3T14 E [(0++1/3)]
4
4
(4)
37p
[C(0']]
37/
[(0++1/2)
Finally, by combining our estimates for E [C( and E [<(00-1/2)], we obtain:
E [c()++1)] s (1- 37p
E
[[((0)]]
+ 37/11
4 E [(()+1/2)]
(1 (1-
311 2
E [[(')]
+
3mp² 5
E [[((')] s (1- n
E [(0")
4 4
2
which implies the convergence rate:
E [<(O']] S (1- 2 E [C(O")] ViEN.
n
C.4 COMPARISON BETWEEN TWO DEFINITIONS ON LINEAR STABILITY
Definition C.1 (Norm/loss- based linear stability). For an update rule g++1 0¹ + F(O':E))
(Norm-based linear stability) A global minimum 0* is said to be linearly stable if there exists
>> such that it holds for SAM the linearized model (Equation (5)) that E| 0" 1 < 0-
CE| 0
|| Vt 2 0.
(Loss-based linear stability) A global minimum 0* is said to be linearly stable if there exists a C > such that it holds for SAM on the linearized model (Equation (5)) that ([C((')] < CE[[(0"), Vt 2 0.
Lemma C.2. The following two-fold results holds:
(LI). Of is norm-hased linearly stable, then it must be loss-based linearly stable.
(I.2). There exists an update rule and a G(0*), such that the converse of (L.I) does not hold.
31

---

Published as a conference paper at ICLR 2025
Proof of Lemma C.2. For (L1), it is straightforward by
L(0') - (0-00
00) s |[G 0 - 0'||²
Thus, the norm-based linear stability imply the loss-based linear stability.
For (L2), we only need to consider an example containing degenerated directions. Specifically, let
L( 00 +0-0 02.
with 0 0.
Consider an update rule only occurs on 82:
0t+1 0::
0,+1 0 Gt Ge ~ 1),
starting from = T Then it is loss-base linearly stable due to
(()') 2 (0)))
1
1 (01) C()
However, it is norm-base linearly nonstable:
E [1] 1 + E 110
=1
1 Σ Se
+ t + to0 when t +00.
32


Reviews:
[{"review_id": "MwhfPXVYez", "reviewer": "Reviewer_6K1k", "summary": "The paper delves into the behavior of Sharpness-Aware Minimization (SAM) in the later stages of neural network training, demonstrating that SAM can effectively evade the sharp minima typically identified by Stochastic Gradient Descent (SGD) and instead settle into flatter minima. It outlines a two-stage learning process in which SAM initially hastens the departure from sharp minima and subsequently swiftly reaches flatter minima, a pattern observed even when SAM is employed only during the final training epochs. The research indicates that the selection of the optimization algorithm during the late training phase has a substantial impact on the ultimate attributes of the solution, underscoring the significance of scrutinizing optimization tactics at this critical juncture.", "questions": "Could the authors explain how the neighborhood size affects the results when switching from SAM to SGD?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This paper offers an in-depth examination of Sharpness-Aware Minimization (SAM) throughout the training process, with a particular emphasis on its late-phase behavior, an area that has been understudied in prior research.\nThe paper elucidates a distinct two-phase learning pattern of SAM, illustrating its ability to swiftly exit sharp minima and subsequently settle into flatter minima, thereby enriching the theoretical grasp of SAM's optimization trajectory.\n\nThe study underscores the pivotal impact of the optimization algorithm selected for the late training phase, proposing a potential optimization strategy that could reduce computational demands while sustaining optimal performance, thus offering practical relevance for deep learning applications in real-world contexts.", "weaknesses": "1.The theoretical results have limited applicability to real training scenarios. They do not sufficiently explain why transitioning the training method from SGD to SAM results in a large generalization gap compared to using the full SAM training method. The experiments in the paper demonstrate that switching from SAM to SGD often leads to poorer results.\n\n2.In theory, the switch from SGD to SAM is expected to yield certain benefits; however, no experimental evidence is provided to support this.\n\n3.The models discussed are primarily convolution-based, so it would be beneficial to include some transformer-based models as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XC3CFI7fYH", "reviewer": "Reviewer_WFBm", "summary": "The authors demonstrate that effects of sharpness-aware minimization (SAM) are independent of whether it is applied in early training. They show that when applied late, SAM-driven iterates first escape the narrow minimum found by SGD, and then find a wider minimum. The authors demonstrate that the earlier application of SAM has no discernible effect on the solution. Inspired by these results, the authors find a similar pattern in adversarial training (AT) as well.", "questions": "- L000: The concepts of early vs. late training are used throughout the training, but is not given clear definitions. Please commit to an explicit definition (in accordance with L306), or discuss why this is not necessary/possible. Doing so earlier in the paper is important.\n- L081: Efficient sounds like a technical ter\n- L110: The notation for dataset and distribution are the same, $\\mathcal{S}$. Please differentiate, and differentiate their use as subscripts to $\\mathcal{L}$ as well (as this is almost exclusively used with distributions.)\n- L115: Please define classification error explicitly.\n- L177: \"when switches from SGD to SAM\"\n- L177: The ordering of the x-axes in Figure 3 is confusing. I understand that the authors aimed for consistency between Figures 3, 5, and 6; but I think their choice makes the visualization less legible. I believe having x axes stand for changing time will make the graph easier to read. I ultimately leave this choice to authors' discretion.\n- L286: Why change $T$ and $t$'s usage?\n- L291: Typo after $\\mathcal{L}''$\n- L318: Is there a reason for usage of $\\ell$ here?\n- L352: After changing the use of $T$ and $t$ at L286, Theorems 4.1 and 4.2 changes this use one more time (from $T$ to $T + t$ to $0$ to $t$). Can these be made more consistent?\n- L354: Please alert the reader that you will introduce Proposition 4.1 to resolve the conflict between the exponential escape and low-loss assumption\n- L388: Do these results generalize to cases where $p>1$ and minibatch SAM?\n- L406: Can we assume any valley around $\\theta^*$ found by SGD to be sub-quadratic?\n- L524: Are you making a claim about robustness-accuracy trade-off? I.e. do you expect robustness to increase without sacrificing vanilla accuracy if the AT is started late?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper's findings are novel, important, while also counterintuitive given previous research's emphasis on early training. The findings are likely to be of interest to research communities that focus on SAM in specific, and training dynamics and generalization in general.\n- The authors make a clear and convincing case regarding the importance of adopting SAM in late (vs. early training) using empirical and theoretical findings.\n- The paper is written well and is easy to follow, with the authors' visualizations being especially instructive. Theoretical results are promptly compared with previous results in the literature.\n- The extension of the authors' results to adversarial training both supports their original results and also is an important contribution in and of itself.", "weaknesses": "- The divergence between paper's experimental and theoretical setting is notable especially in the choice of the loss function. This is acknowledged in the paper and is acceptable given the novelty of their results.\n- Although the paper is written well overall, certain unclearly defined concepts, choice of notation, and design choices make it hard to follow at times. I provide more detailed feedback regarding these in the section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NmkELMhuAt", "reviewer": "Reviewer_cw28", "summary": "The authors explore how algorithm SAM converges to flat local minima and find a phase transition when switching SGD to SAM. By identifying the quick convergence property of SAM, the authors propose to use SAM in the later few epochs, which significantly reduces the computation of using SAM for flat local optima. The experimental results further verify the effectiveness of the proposed algorithm", "questions": "1. Can the authors give a clear definition of linear stability? It seems to be a very important definition in the paper.\n\n2. Can the authors provide proof when the function has two valleys, SAM will not go from one valley to another? In my view, Proposition 4.1 can not support iterates generated by SAM stay in the same valley.\n\n3. What are the advantages of SAM shown in Theorem 4.3? SGD converges at the same rate under the Assumptions 4.2 and 4.3.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors propose to use SAM in the last few epochs instead of starting from the beginning. The new method can reduce almost half of the computation of SAM while retaining the generalization ability of SAM.\n\n2. The authors identify 2 phases of switching from SGD to SAM. With the theoretical analysis, the authors give 4 properties of SAM.", "weaknesses": "1. Some definitions are not clear.  \n     (i) Definition 4.1:  Does it mean that SAM is used to optimize linearized function but is evaluated on the real loss or that SAM both trains and tests on the linearized loss?\n\n    (ii)  What is $a$ in Proposition 4.1, which is used to give an upper bound of $\\rho$?\n\n2. Is the first claim in Theorem 4.2 the contrapositive of Theorem 4.1? If not, can you explain the difference between these two claims?\n\n3. After Theorem 4.1 and Theorem 4.2, the authors claim that SAM escapes from the sharp minima exponentially fast, which sounds like a good result. However, in Theorem 4.2, the result says that the final loss is $C^t$ larger than the initial one, which seems to be a bad result.  What is the relationship of these two results or the exponentially fast claim should be derived from some other theorem?\n\n4. When the authors want to prove that SAM can stay in the same valley as SGD, there is only one valley in Proposition 4.2. Thus, to show SAM stays in the same valley, the function should have at least two valleys.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xRgas2hmKE", "reviewer": "Reviewer_bzn5", "summary": "The paper studies effects of SAM at the end of training upon an SGD baseline, and claims that using SAM for the final stages of training is sufficient to improve generalization. Considering a trajectory where SGD is switched to SAM at epoch $t$, the paper's theoretical analyses show that SAM rapidly escapes the SGD minima, while staying in the same valley. SAM then proceeds to converge rapidly to a new minimum that is flatter. These claims are supported by empirical evidence, and the authors posit a general conjecture that these loss objectives are only necessary toward the end of training.", "questions": "1. In 'Support for P3', can you support why increasing the loss might lead to flatter minima and better generalization? The empirical evidence suggests that SAM decreases the training loss, in addition to test loss.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "While the empirical observations originate elsewhere, the theoretical analysis appears to be novel and interesting. \nWhile the analyses are limited to the interplay between SGD and SAM during the later phases of training, it is sound and, this line of reasoning can conceivably generalize to other learning algorithms of interest, allowing the potential significance to be large. \nThe paper is clearly structured and easy to follow. The authors introduce a well-defined scope with thorough justification, both theoretically and empirically, and clearly discuss the motivation and role of assumptions. They then propose and support empirically a few general impacts of these insights to broader aspects of optimization, such as the role of SAM early-in-training and adversarial training late-in-training.", "weaknesses": "1. For some experiments (Fig 9) using SGD toward the end of training outperformed SAM in terms of accuracy, which is mildly contradicting toward an implicit claim that at the end-of-training, SAM outperforms SGD. Some commentary in the main section on these results would be appreciated. \n2. The size of the valley in Proposition 4.1 for P2, $(-2b,2b)$, appear extremely wide when the parameter is initialized between $(-b, b)$, which weakens the significance of this proof. The empirical evidence (Fig 4c) for this claim also appears limited. Can a more general visualization of the loss landscape be considered?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "27alY9DIgO", "reviewer": "Reviewer_S3eC", "summary": "This paper studies Sharpness Aware Minimization (SAM), an optimization method that modifies the Stochastic Gradient Descent (SGD) in a way that the solution will prefer a flatter region (Eq. (1)).\nIn practice, its approximation (Eq. (4)) is implemented for computationally efficiency.\nIt is not well understood whether this approximated method has the intended effect of finding a flatter solution, and this is what this paper studies.\n\nFirst, the paper empirically shows that when we switch from the SGD to the SAM during the updates, it increasingly improves the solution as we make the switching time earlier, both in terms of the loss and the sharpness (Figure 3). However, using SAM for the last few updates already accounts for the majority of the improvement, and the performance increment becomes very small after that. This suggests that applying SAM in a final few epochs may be nearly as good as doing so for the whole training.\n\nNext, the authors theoretically study several properties of SAM. Theorems 4.1 and 4.2 show that SAM escapes from a solution with large sharpness exponentially fast. Furthermore, with a few additional assumptions, Theorem 4.3 shows SAM converges to the global minimum exponentially fast.", "questions": "### Major comments:\n- Please address the first three comments in the Weaknesses section.\n\n- I expected to see the loss going down in later steps in Figure 4 (b). Does the figure not show the entire steps?\n\n- \"(P3). SAM converges to a flatter minimum compared to SGD; Theorem 4.1\" The result of Theorem 4.1 seems to be a little overstated here. If I understand correctly, Theorem 4.1 does not show SAM converges although Theorem 4.3 does. Note that the latter assumes additional conditions (Assumptions 4.2 and 4.3). Could the authors revise the sentence to a more accurate one?\n\n### Minor comments:\n- Are there any experiments using Eq. (4)?\n\n- The paper uses the symbol $\\approx$, but its mathematical definition is missing. Could the authors add the precise meaning?\n\n- Proof of Proposition 4.1 could be improved by adding comments on which part the authors used the assumptions.\n\n- Does this refer to Eq. (3)? If so, it should be clearly stated: \"Our experiments are conducted using the original implementation of SAM\"\n\n---\n**Edit**: During the discussion period, the authors addressed all of my concerns. I updated the Rating score from 6 to 8. (I was about to set it to 7, but there wasn't this option.)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The theorems are interesting and nicely presented with interpretations. They successfully explain how SAM escapes a sharp region and converges to a flat region.\n- The paper cites relevant references and clearly discuss the differences.\n- The paper is well written and easy to follow, except for a few points (please see the Weaknesses section).\n- The convergence result in Theorem 4.3 does not rely on convexity.", "weaknesses": "- Eq. (4) and the update rule that the authors use in theory (the two-step update rule defined at the beginning of C.1) are not equivalent after eliminating $\\theta_{t+1/2}$.\n\n- The experiments and the theory part use different methods (Eqs. (3) and (4)).\n\n- I could not understand\n  - how to obtain the inequality in lines 1137-1138, and\n  - the phrase \"which implies that $\\theta_{T+1} = 0$\" in lines 1143-1144.\n  - the equality in lines 1085-1086 and the inequality in lines 1087-1088.\n\n- This paper only studies the squared loss.\n\n---\n### Relatively minor comments\n- Theorems 4.1 and 4.2 rely on the linear approximation (5), so it does not explain the whole dynamics after escaping from a sharp minimum and before getting close to a flat minimum. Furthermore, combined with the squared loss, the analysis is in a classical strongly convex setup.\n\n- The single-data setup of Example 4.1 is unrealistic, and I do not know how relevant it is. Moreover, Proposition 4.2 considers the case in which there is only one parameter and has almost no practical usefullness.\n\n- Theorem 4.1 and 4.2 seem to be almost equivalent (contrapositive) to each other. Theorem 4.2 is slightly more detailed in that it provides the explicit inequality showing the exponential bound. Theorem 4.2 directly implies Theorem 4.1, so the latter should be a corollary.\n\n- I am not totally convinced that Assumptions 4.1-4.3 will be actually satisfied in practical situations. To be fair, many supporting references are provided in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)