Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training

## Abstract

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 lights 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 1 1 1 To clarify, “late in training” stands for “when applied during the late phase of training”, and is used here for simplicity. .
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 2 2 2 A “valley” refers to the local loss landscape surrounding a set of global minima. .
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.
We provide source code in supplementary materials and will release checkpoints in future.

## 1 Introduction

Understanding the surprising generalization abilities of over-parameterized neural networks yields an important yet open problem in deep learning.
Recently, it has been observed that the generalization of neural networks is closely tied to the sharpness of the loss landscape (Keskar et al., ; Zhang et al., ; Neyshabur et al., ; Jiang et al., ).
This has led to the development of many gradient-based optimization algorithms that explicitly/implicitly regularize the sharpness of solutions.
In particular, Foret et al. () proposed Sharpness-Aware Minimization (SAM), which has substantially improved the generalization and robustness (Zhang et al., ) of neural networks across many tasks, including computer vision (Foret et al., ; Chen et al., ; Kaddour et al., ) and natural language processing (Bahri et al., ).

Despite the empirical success of SAM, its effectiveness is not yet fully understood.
Andriushchenko & Flammarion () has shown that existing theoretical justifications based on PAC-Bayes generalization bounds (Foret et al., ; Wu et al., ) 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 a sufficiently small learning rate (Monzio Compagnoni et al., ), 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-based optimization methods can be attributed to their *implicit bias* toward solutions with favorable properties (Vardi, ).
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., ; Zhang et al., ).
It is known that SAM selects flatter minima over SGD in practice, which represents a form of implicit bias as well.
While the design of SAM is inspired by sharpness regularization, its practical implementation (see Equation 3), which minimizes a 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 and is the primary focus of this work.

Figure: Figure 1: SAM operates efficiently even when applied only during the final few epochs of training. (a) Test error curves for WideResNet-16-8 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 $\widetilde{T}=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 and hyper-parameters are described in Section 3. (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.
Refer to caption: x1.png

In this work, we uncover an intriguing phenomenon in the training dynamics of SAM, which sheds light on 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 1 (a), SAM applied in the last few epochs—after initial training with SGD—can achieve nearly the same test performance as training exclusively with SAM.
Furthermore, our analysis of the solution sharpness reveals that last-few-epochs SAM identifies flatter minima as effectively as full SAM training.
We conclude that *SAM efficiently selects flatter minima over SGD late in training*.
Contrary to the conventional view that early training dynamics are critical (Achille et al., ; Frankle et al., ), 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 I.* SAM initially escapes from the relatively sharp minimum found by SGD exponentially fast; *Phase II.* it then rapidly converges to a flatter minimum over SGD within the same valley.
See Figure 1 (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 are further validated by experiments in both toy models and real-world scenarios.
This two-phase picture explains the 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 effectiveness.

Furthermore, we investigate the necessity of applying SAM during the early phase of training.
Our results show that early-phase SAM has a minimal effect on 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 a more critical role in shaping the final properties of the solution.
Extending this viewpoint to Adversarial Training (AT) (Madry et al., ) 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 during the late phase of training, a topic that has not yet been carefully studied.
Our main discovery unveils a surprising phenomenon in the late dynamics of SAM, and our novel two-phase picture significantly advances our understanding of its effectiveness.
Our results demonstrate the value of analyzing optimization algorithms at the end of training and formulating theories from realistic nontrivial phenomena that provide valuable insights into the underlying mechanisms of deep learning.

## 2 Related Work

Implicit flat-minima bias.
Extensive studies have been dedicated to understand implicit bias in deep learning (Vardi, ).
One notable example is the *flat-minima bias*: SGD and its variants tend to select flat minima (Keskar et al., ; Zhang et al., ), which often generalize better (Hochreiter & Schmidhuber, ; Jiang et al., ).
Many works (Ma & Ying, ; Mulayoff et al., ; Wu & Su, ; Gatmiry et al., ; Wen et al., ) have theoretically justified the superior generalization of flat minima in neural networks.
Wu et al. (; ); Ma & Ying () have taken a 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., ; Li et al., ; Ma et al., ; Damian et al., ).
Further studies (Jastrzębski et al., ; Lyu et al., ) 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. () has proposed Sharpness-Aware Minimization (SAM) algorithm.
SAM and its variants (Kwon et al., ; Zhuang et al., ; Du et al., ; Liu et al., ; Kim et al., ; Mueller et al., ; Becker et al., ) have achieved superior performance across a wide range tasks (Chen et al., ; Kaddour et al., ; Bahri et al., ).
Despite their successes, their effectiveness remains poorly understood.
On the optimization front, Andriushchenko & Flammarion (); Si & Yun (); Shin et al. () have studied the convergence properties of stochastic SAM;
Monzio Compagnoni et al. (); Kim et al. () have demonstrated that SAM requires more time to escape saddle points compared to SGD;
Long & Bartlett () have analyzed the edge of stability phenomenon for SAM.
On the implicit bias side,
Shin et al. (); Behdin et al. () have examined SAM’s flat-minima bias using dynamical stability analysis;
Wen et al. (); Monzio Compagnoni et al. (); Bartlett et al. () have studied how SAM moves towards flat minima;
Andriushchenko & Flammarion () has showed that SAM has the low-rank bias;
and Wen et al. () has revealed that SAM does not only minimize sharpness to achieve better generalization.
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.

## 3 Background and Preliminaries

Notation and Setup.
Let $\mathcal{S}=\{({\bm{x}}_{i},\bm{y}_{i})\}_{i=1}^{n}$ be the training set, where ${\bm{x}}_{i}\in\mathbb{R}^{d}$ is the input and $\bm{y}_{i}\in\mathbb{R}^{k}$ is the target for the $i$-th data point i.i.d. sampled from distribution $\mathcal{S}$.
We consider a model in the form of $f(\bm{x};\bm{\theta})$, where ${\bm{\theta}}\in\mathbb{R}^{p}$ is the model parameters and $f(\bm{x};\bm{\theta})\in\mathbb{R}^{k}$ is the output for the input $\bm{x}$.
The loss of the model at the $i$-th sample $(\bm{x}_{i},\bm{y}_{i})$ is denoted as $\ell(\bm{y}_{i};f(\bm{x}_{i};\bm{\theta}))$, simplified to $\ell_{i}(\bm{\theta})$.
The total loss over the training set is then given by $\mathcal{L}_{\mathcal{S}}(\bm{\theta})=\frac{1}{n}\sum_{i=1}^{n}\ell_{i}(\bm{
\theta})$.
For classification tasks, $\operatorname{Err}_{\mathcal{D}}(\bm{\theta})$ denotes the classification error of the model $f(\cdot;\bm{\theta})$ over dataset $\mathcal{D}$.
For a matrix $A$, let $\left\lVert A\right\rVert_{2}$, $\left\lVert A\right\rVert_{F}$, and $\operatorname{Tr}(A)$ denote its spectral norm, Frobenius norm and trace, respectively.

For our theoretical analysis in Section 5, we always assume $k=1$ for simplicity(^3^33The extension to the case where $k>1$ is straightforward.) and consider the squared loss, i.e., $\ell_{i}(\bm{\theta})=\frac{1}{2}|f(\bm{x}_{i};\bm{\theta})-y_{i}|^{2}$.
We focus on the over-parameterized case in the sense that $\min_{\bm{\theta}}\mathcal{L}(\bm{\theta})=0$.
To characterize the local geometry of the training loss at point $\bm{\theta}$, we consider the Fisher matrix
$G(\bm{\theta})=\frac{1}{n}\sum_{i=1}^{n}\nabla f(\bm{x}_{i};\bm{\theta})\nabla
f
(\bm{x}_{i};\bm{\theta})^{\top}$
and the Hessian matrix $H(\bm{\theta})=G(\bm{\theta})+\frac{1}{n}\sum_{i=1}^{n}\left(f(\bm{x}_{i};\bm{
\theta})-y_{i}\right)\nabla^{2}f(\bm{x}_{i};\bm{\theta})$, where $G(\bm{\theta})\approx H(\bm{\theta})$ 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

$$ $\displaystyle\min_{\bm{\theta}}\mathcal{L}_{\mathcal{S}}^{\rm SAM}(\bm{\theta} ;\rho),\quad{\rm where}\quad\mathcal{L}_{\mathcal{S}}^{\rm SAM}(\bm{\theta}; \rho)=\max_{\left\lVert\bm{\epsilon}\right\rVert_{2}\leq\rho}\mathcal{L}_{ \mathcal{S}}(\bm{\theta}+\bm{\epsilon}).$ (1) $$

Here, $\rho$ is a given neighborhood radius.
However, finding the perturbation in the parameter space that maximizes the loss, i.e., $\bm{\epsilon}$, can be computationally intractable in practice.
Thus Foret et al. () approximated the maximize-loss perturbation using a first order solution:

$$ $\displaystyle\bm{\epsilon}\approx\operatorname*{arg\,max}_{\left\lVert\bm{ \epsilon}\right\rVert_{2}\leq\rho}\left(\mathcal{L}_{\mathcal{S}}\left(\bm{ \theta}\right)+\bm{\epsilon}^{\top}\nabla\mathcal{L}_{\mathcal{S}}\left(\bm{ \theta}\right)\right)=\rho\nabla\mathcal{L}_{\mathcal{S}}\left(\bm{\theta} \right)/\left\lVert\nabla\mathcal{L}_{\mathcal{S}}\left(\bm{\theta}\right) \right\rVert_{2}.$ (2) $$

Consequently, let $\eta$ be the learning rate and $\xi_{t}=\{\xi_{t,1},\cdots,\xi_{t,B}\}$ be the batch indices sampled at iteration $t$, where $B$ is the batch size.
The mini-batch loss is given by $\mathcal{L}_{\xi_{t}}(\bm{\theta})=\frac{1}{B}\sum_{i\in\xi_{t}}\ell_{i}(\bm{
\theta})$.
With a small yet key modification to SGD(^4^44The iterative update rule of SGD is $\bm{\theta}^{t+1}=\bm{\theta}^{t}-\eta\nabla\mathcal{L}_{\xi_{t}}\left(\bm{ \theta}^{t}\right)$), the iterative update rule of SAM with stochastic gradients is:

$$ $\displaystyle\bm{\theta}^{t+1}=\bm{\theta}^{t}-\eta\nabla\mathcal{L}_{\xi_{t}} \left(\bm{\theta}^{t}+\rho_{t}\nabla\mathcal{L}_{\xi_{t}}(\bm{\theta}^{t}) \right),\quad{\rm where}\quad\rho_{t}=\rho/\left\lVert\nabla\mathcal{L}_{\xi_{ t}}\left(\bm{\theta}^{t}\right)\right\rVert_{2}.$ (3) $$

Figure: Figure 2: The illustration of the switching method. Blue dashed lines represent SGD training, while orange dashed lines represent SAM training. $\bm{\theta}^{0}$ denotes the random initialization. $t_{1},t_{2}$ denotes two different switching points.
Refer to caption: x2.png

Our experiments are conducted using the original implementation of SAM, however, for our theoretical analysis in Section 5, we simplify the original SAM by approximating the step size $\rho_{t}$, which is inversely proportional to the gradient norm, with a constant $\rho$ (see Equation 4), as it is easier to analyze.
Andriushchenko & Flammarion () empirically justified this simplification, showing that SAM with a constant $\rho$ achieves generalization performance comparable to the original SAM.
This approximation has been widely adopted in recent theoretical advances on SAM (Andriushchenko & Flammarion, ; Behdin & Mazumder, ; Agarwala & Dauphin, ; Behdin et al., ; Monzio Compagnoni et al., ).

$$ $\displaystyle\bm{\theta}^{t+1}=\bm{\theta}^{t}-\eta\nabla\mathcal{L}_{\xi_{t}} \left(\bm{\theta}^{t}+\rho\nabla\mathcal{L}_{\xi_{t}}(\bm{\theta}^{t})\right).$ (4) $$

Switching Method.
In this paper, we demonstrate our main discovery by switching between SGD and SAM at a specific point during training, referred to as the *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 $t$ epochs to get $\bm{\theta}_{\textrm{SGD}}^{t}$, and then switch to SAM for the remaining $T-t$ epochs to get $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$.
Notably, when $t=T$ or $t=0$, the model is trained exclusively with SGD or SAM, simply denoted as $\bm{\theta}_{\textrm{SGD}}^{T}$ and $\bm{\theta}_{\textrm{{SAM}}}^{T}$, respectively.
See Figure 2 for an illustration.
We fix the total number of training epochs $T$ and vary the switching point $t$ to study the effect of SAM at different training phases.
To ensure a 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. (); Kwon et al. (), we perform experiments on the commonly used image classification datasets CIFAR-10/100 (Krizhevsky et al., ), with standard architectures such as WideResNet (Zagoruyko & Komodakis, ), ResNet (He et al., ), and VGG (Simonyan & Zisserman, ).
We use the standard configurations for the basic training settings shared by SAM and SGD (e.g., learning rate, batch size, and data augmentation) as in the original papers
, and set the SAM-specific perturbation radius $\rho$ to 0.05, as recommended by Foret et al. ().
We also extend our findings to AT (Madry et al., ) and ASAM (Kwon et al., ).
Due to space limits, more experimental details and results can be found in Appendix A.

## 4 SAM Efficiently Selects Flatter Minima Late In Training

Figure: Figure 3: The impact of SAM training proportion on generalization/sharpness when switches from SGD to SAM. The generalization gap between the models $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ (top row) / the sharpness of $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ (bottom row) vs. the SAM training proportion of $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$. $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$ refers to the model obtained by switching from SGD to SAM, while $\bm{\theta}_{\textrm{SAM}}^{T}$ represents the model trained entirely with SAM. Here, $T$ is the total number of training epochs, and $t$ is the switching point. We define SAM training proportion of $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$ as $(T-t)/T$ and adjust it by varying $t$ while keeping $T$ fixed. Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations. Results on more datasets and architectures can be found in Section A.4.
Refer to caption: x3.png

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 to training entirely with SAM.
Taking a step further, we demonstrate that last-few-epochs SAM finds flatter minima as effectively as full SAM training.

Applying SAM in the last few epochs achieves generalization comparable to full SAM training.
We conduct experiments on various datasets and architectures to investigate the effect of applying SAM late in training, after initial training with SGD, on generalization.
Specifically, we use the switching method (detailed in Section 3), which switches from SGD to SAM at epoch $t$, to obtain the model $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$.
For comparison, we train a baseline model entirely with SAM under the same settings, resulting in $\bm{\theta}_{\textrm{SAM}}^{T}$.
Both models are trained for a total of $T$ epochs.
To assess the generalization ability of $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ relative to $\bm{\theta}_{\textrm{SAM}}^{T}$, we calculated the test loss and test error gaps between these two models: $\mathcal{L}_{\mathcal{D}_{\textrm{test}}}\left(\bm{\theta}_{\textrm{SGD}\to
\textrm{SAM},t}^{T}\right)-\mathcal{L}_{\mathcal{D}_{\textrm{test}}}\left(\bm{
\theta}_{\textrm{SAM}}^{T}\right)$ and $\operatorname{Err}_{\mathcal{D}_{\textrm{test}}}\left(\bm{\theta}_{\textrm{SGD
}\to\textrm{SAM},t}^{T}\right)-\operatorname{Err}_{\mathcal{D}_{\textrm{test}}
}\left(\bm{\theta}_{\textrm{SAM}}^{T}\right)$, where $\mathcal{D}_{\textrm{test}}$ denotes the unseen test set.
If the generalization gap between $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ approaches zero, then we say $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ achieve nearly the same generalization performance as $\bm{\theta}_{\textrm{SAM}}^{T}$.

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 of total training epochs $T$ and vary the switching point $t$ to adjust the SAM training proportion of $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$, which is exactly $(T-t)/T$.
We analyze how the generalization gap between $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ changes as the SAM training proportion increases.
In Figure 3, the generalization gap becomes negligible once the SAM training proportion exceeds zero, indicating that even a 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 (), but it 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 sharpness of the solution $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ 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 $\left\lVert H(\bm{\theta})\right\rVert_{2}$.
In Figure 3, the sharpness of the solution $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ 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 of 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 5, we will dig deeper into the underlying mechanism behind this phenomenon and gain new insights into SAM’s implicit bias towards flatter minima.

## 5 How Does SAM Efficiently Select 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 to explore the precise mechanism behind this phenomenon.
Theoretically, we identify a two-phase picture in the training dynamics after switching to SAM in 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 a flatter minimum within the same valley in just a few steps, explaining its efficiency in selecting flatter minima late in training.

**Table 1: Overview of the two-phase picture and corresponding theoretical results.**
| Phase I<br>(escape) | (P1). SAM rapidly escapes from the minimum found by SGD; | Theorem 5.2 |
| --- | --- | --- |
| (P2). However, the iterator remains within the current valley. | Proposition 5.1 |  |
| Phase II<br>(converge) | (P3).<br>SAM converges to a flatter minimum compared to SGD; | Theorem 5.1 |
| (P4).<br>The convergence rate of SAM is extremely fast. | Theorem 5.3 |  |

Figure: Figure 4: Two-phase picture in both toy example and real-world scenarios. (a) Visualization of the two-phase dynamics for Example 5.1. The horizontal gray line represents the set of the global minima $\mathcal{M}=\{(u,v):v=0\}$, where smaller values of $|u|$ correspond to flatter minima. The blue lines trace the SGD iterates, leading to $\bm{\theta}_{\text{SGD}}^{\text{end}}$, while the orange lines show the SAM iterates, which converge to a flatter minimum $\bm{\theta}_{\text{SAM}}^{\text{end}}$, with a smaller value of $|u|$. Notably, $\bm{\theta}_{\text{SGD}}^{\text{end}}$ and $\bm{\theta}_{\text{SAM}}^{\text{end}}$ stay in the same valley around $\mathcal{M}$. (b) The exponentially fast escape from minima found by SGD. Train loss $\mathcal{L}_{\mathcal{D}_{\textrm{train}}}(\bm{\theta}_{\textrm{SGD}\to\textrm {SAM},T}^{T+t})$ v.s. the update step $t$. Here, $T$ is the switching step, chosen to be 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. (c) SAM converges to a flatter minimum within the same valley as the one found by SGD. The loss $\mathcal{L}$ (top row) / the second-order finite difference $\mathcal{L}^{\prime\prime}$ (bottom row) of the interpolated model $\bm{\theta}_{\lambda}$ v.s. the interpolation coefficient $\lambda$. Here, $\bm{\theta}_{\lambda}=(1-\lambda)\bm{\theta}_{\textrm{SGD}\to\textrm{SAM}}^{ \textrm{end}}+\lambda\bm{\theta}_{\textrm{SGD}}^{\textrm{end}}$, and $\mathcal{L}^{\prime\prime}(\bm{\theta}_{\lambda})=\left(\mathcal{L}(\bm{\theta }_{\lambda+h}+\mathcal{L}(\bm{\theta}_{\lambda-h}-2\mathcal{L}(\bm{\theta}_{ \lambda}\right)/(2h)^{2}$, where $h$ is fixed to $0.1$. This is no barrier in the loss curve, indicating $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM}}^{\textrm{end}}$ and $\bm{\theta}_{\textrm{SGD}}^{\textrm{end}}$ stays within the same valley. $\mathcal{L}^{\prime\prime}(\bm{\theta}_{\lambda})$ gradually increases with $\lambda$, implying $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM}}^{\textrm{end}}$ is flatter than $\bm{\theta}_{\textrm{SGD}}^{\textrm{end}}$.
Refer to caption: x4.png

To understand the two-phase picture, let us first consider a toy but representative example.

###### Example 5.1 .

Consider using the shallow neural network $f(u,v;x)=\tanh(v\tanh(ux))$ to fit a single data ($x=1,y=0$) under the squared loss $\ell(y;y^{\prime})=(y-y^{\prime})^{2}/2$, then the loss landscape can be written as $\mathcal{L}(\bm{\theta})=\frac{1}{2}\tanh^{2}(v\tanh(u))$, where $\bm{\theta}=(u,v)$.

Figure 4 (a) visualizes the training dynamics for Example 5.1.
In Figure 4 (a), the dynamics occur within the valley surrounding the set of the minima $\mathcal{M}=\{(u,v):v=0\}$.
On $\mathcal{M}$, the sharpness is given by $\left\lVert H(\bm{\theta})\right\rVert_{2}=\operatorname{Tr}(H(\bm{\theta}))=
\|H(\bm{\theta})\|_{\rm F}=\tanh^{2}(u)$, implying that the smaller the $|u|$, the flatter the minimum.
Initially, SGD converges to a relatively sharp minimum $\bm{\theta}_{\text{SGD}}^{\text{end}}$.
After switching to SAM, the iterator exhibits the two-phase dynamics:

- •
Phase I (escape). SAM first escapes the sharp minimum $\bm{\theta}_{\text{SGD}}^{\text{end}}$ found by SGD within a few iterations (P1), but remains within the same valley around $\mathcal{M}$(^5^55Note that the landscape has another valley around a different set of minima $\mathcal{N}=\{(u,v):u=0\}$) (P2);
- •
Phase II (converge). SAM then rapidly converges to another flatter minimum $\bm{\theta}_{\text{SAM}}^{\text{end}}$ (P3-4), with a smaller $|u|$ compared to $\bm{\theta}_{\text{SGD}}^{\text{end}}$.

This toy example successfully illustrates the two-phase picture after switching to SAM in the late training phase, but it offers only a basic intuition.
In the following, we will formally present our theoretical results, with each claim rigorously supported by corresponding analysis.

###### Example 5.1 .

### 5.1 Theoretical guarantees for P1 and P3 : linear stability analysis

In this subsection, we provide theoretical support for two key claims in our two-phase picture through a linear stability analysis: P1. *SAM escapes from the relatively sharp minima found by SGD exponentially fast*; P3. *SAM converges to a flatter minima compared to SGD*.

During the late phase of training, the model remains close to some global minimum $\bm{\theta}^{\star}$ and can be locally approximated by its linearization:

$$ $f_{\rm lin}(\bm{x};\bm{\theta})=f(\bm{x};\bm{\theta}^{\star})+\left\langle \nabla f(\bm{x};\bm{\theta}^{\star}),\bm{\theta}-\bm{\theta}^{\star}\right\rangle.$ (5) $$

Consequently, we can characterize the local dynamical stability (Strogatz, ) of the SAM iterator near the global minimum on this linearized model.

Similar to the linear stability of SGD (Wu et al., ), we define the linear stability of SAM in Definition 5.1.
Note that for a specific optimizer, only linearly stable minima can be selected.

###### Definition 5.1 (Loss-based linear stability of SAM) .

A global minimum $\bm{\theta}^{\star}$ is said to be linearly stable for SAM if there exists a $C>0$ such that it holds for SAM on the linearized model (Equation 5) that $\mathbb{E}[\mathcal{L}(\bm{\theta}_{t})]\leq C\mathbb{E}[\mathcal{L}(\bm{
\theta}_{0})],\forall t\geq 0$.

Wu et al. () has shown that the linear stability of stochastic algorithms is significantly influenced by the gradient noise $\bm{\xi}(\bm{\theta})=\nabla\ell_{\xi}(\bm{\theta})-\nabla\mathcal{L}(\bm{
\theta})$, where $\xi$ 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., ; Feng & Tu, ; Wu et al., ; ; Mori et al., ; Wojtowytsch, ) have established two key properties of the gradient noise $\bm{\xi}(\bm{\theta})$: (i) its shape is highly anisotropic; (ii) its magnitude is loss dependent.
Specifically, by incorporating both the noise shape and magnitude, Mori et al. () has proposed an approximation for the gradient noise covariance $\Sigma(\bm{\theta})=\frac{1}{n}\sum_{i=1}^{n}\bm{\xi}(\bm{\theta})\bm{\xi}(\bm
{\theta})^{\top}$, suggesting $\Sigma(\bm{\theta})\sim 2\mathcal{L}(\bm{\theta})G(\bm{\theta})$.

Inspired by this approximation, we assume a weaker(^6^66“Weaker” means that the approximation $\Sigma(\bm{\theta})\sim 2\mathcal{L}(\bm{\theta})G(\bm{\theta})$ can imply our Assumption 5.3.) alignment property of the noise covariance $\Sigma(\bm{\theta})$ (see Assumption 5.1) in our analysis, which has been theoretically and empirically validated in recent works (Wu et al., ; Wu & Su, ).

###### Assumption 5.1 (Gradient noise alignment) .

There exists $\gamma>0$ s.t. $\frac{\operatorname{Tr}(\Sigma(\bm{\theta})G(\bm{\theta}))}{2\mathcal{L}(\bm{
\theta})\left\lVert G(\bm{\theta})\right\rVert_{F}^{2}}\geq\gamma$ for all $\bm{\theta}$.

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 5.1 ( P3 . Proof in Section B.1 .) .

Let $\bm{\theta}^{\star}$ be a global minimum that is linearly stable for SAM, and suppose Assumption 5.1 holds.
Then we have
$\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\left(1+\frac{\rho^{2}
\gamma}{B}\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\right)\leq
\frac{B}{\eta^{2}\gamma}$.

Theorem 5.1 characterizes the sharpness of the linearly stable minima for SAM, which is influenced by the learning rate $\eta$, the batch size $B$, and the perturbation radius $\rho$.
Since only linearly stable minima can be selected by optimizer, Theorem 5.1 characterizes the sharpness of global minima selected by SAM.
Notably, taking $\rho=0$ recovers the corresponding result for SGD (see Table 2).

**Table 2: Comparison of the sharpness of global minima selected by SAM and SGD.**
|  | SAM (Theorem 5.1) | SGD (Wu et al., ) |
| --- | --- | --- |
| Sharpness bound | $\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\cdot\left(1+\frac{\rho^<br>{2}\gamma}{B}\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\right)\leq<br>\frac{B}{\eta^{2}\gamma}$ | $\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\leq\frac{B}{\eta^{2}\gamma}$ |

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+\rho^{2}\gamma\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}>1$ in SAM’s bound impose stricter requirement on the sharpness of global minima selected 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., $\bm{\theta}_{t+1/2}=\bm{\theta}_{t}+\rho\nabla\ell_{\xi_{t}}(\bm{\theta}_{t})=
\bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{\theta}_{t})+\rho(\nabla\ell_{\xi_{t
}}(\bm{\theta}_{t})-\nabla\mathcal{L}(\bm{\theta}_{t}))$.
Intuitively, the ascent update $\rho\nabla\mathcal{L}(\bm{\theta}_{t})$ tends to increase the loss, while the noise $\rho(\nabla\ell_{\xi_{t}}(\bm{\theta}_{t})-\nabla\mathcal{L}(\bm{\theta}_{t}))$ introduces extra stochasticity.
These factors together contribute to the flatter-minima selection by SAM.

###### Theorem 5.2 ( P1 . Proof in Section B.1 .) .

Let $\bm{\theta}^{\star}$ be a global minimum, and suppose Assumption 5.1 holds.
If $\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\left(1+\frac{\rho^{2}
\gamma}{B}\left\lVert H(\bm{\theta}^{\star})\right\rVert_{F}^{2}\right)>\frac{
B}{\eta^{2}\gamma}$, then $\bm{\theta}^{\star}$ is linearly non-stable for SAM. Indeed, $\mathbb{E}[\mathcal{L}(\bm{\theta}_{t})]\geq C^{t}\mathbb{E}[\mathcal{L}(\bm{
\theta}_{0})]$ holds for all $t>0$ with $C>1$.

Theorem 5.2, together with Theorem 5.1, characterize the necessary condition for a global minimum to be linearly stable for SAM, namely the stability condition of SAM.
Furthermore, Theorem 5.2 illustrates that if a global minimum fails to meet the stability condition of SAM, the SAM iterator will escapes from the minimum exponentially fast.

Support for P1.
Although the sharpness bound for SGD in Table 2 is an upper bound, Wu et al. () empirically demonstrates that the bound is nearly tight, with SGD solutions approximately satisfying $\left\lVert H(\bm{\theta}_{\text{SGD}}^{\text{end}})\right\rVert_{F}^{2}
\approx\frac{B}{\eta^{2}\gamma}$.
Thus, with the same $\eta$ and $B$, the solution found by SGD cannot 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).

Comparison with other linear stability analyses of SAM (Shin et al., ; Behdin et al., ).
These works have focused on the norm-based linear stability, characterized by $\mathbb{E}[\left\lVert\bm{\theta}_{t}-\bm{\theta}^{\star}\right\rVert^{2}]\leq
C
\left\lVert\bm{\theta}_{0}-\bm{\theta}^{\star}\right\rVert^{2},\forall t\geq 0$, which is stricter than our loss-based stability (see Definition 5.1) due to the presence of degenerate directions in over-parameterized models (see Lemma B.2).
Another key difference is the novel use of the alignment property of gradient noise (see Assumption 5.1) in our analysis, enabling us to derive a concrete bound on the sharpness of SAM’s solutions.

###### Definition 5.1 (Loss-based linear stability of SAM) .

###### Assumption 5.1 (Gradient noise alignment) .

###### Theorem 5.1 ( P3 . Proof in Section B.1 .) .

###### Theorem 5.2 ( P1 . Proof in Section B.1 .) .

### 5.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.*

Sub-quadratic landscape.
In high-loss regions, the local linearization in Equation 5 no 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. () has introduced the concept of the *sub-quadratic landscape*: the loss grows slower than a 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 is able to explain the Edge of Stability phenomenon (Cohen et al., ) (also observed in Wu et al. (); Jastrzebski et al. ()), which the quadratic landscape cannot account for.

While Section 5.1 focuses on a 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., $\theta_{t+1}=\theta_{t}-\eta\nabla\mathcal{L}(\theta_{t}+\rho\nabla\mathcal{L}
(\theta_{t}))$.
This setup is sufficient to illustrate that SAM will remain within the current valley under a sub-quadratic landscape.

###### Definition 5.2 (Sub-quadratic landscape for $p=1$ ) .

Let $\theta^{\star}=0$ be a global minimum of $\mathcal{L}$ with sharpness $\mathcal{L}^{\prime\prime}(0)=a>0$.
The valley $V=[-c,c]$ is termed sub-quadratic if for all $z,z_{1},z_{2}\in V$, it holds that $z\mathcal{L}^{\prime}(z)\geq 0$ and $\mathcal{L}^{\prime\prime}(|z_{1}|)\geq\mathcal{L}^{\prime\prime}(|z_{2}|)$ if $|z_{1}|\leq|z_{2}|$.

Following Theorem 1 in (Ma et al., ), we define the sub-quadratic landscape in Definition 5.2.
Specifically, for any $0<z_{1}<z_{2}<c$, it holds that $\mathcal{L}^{\prime\prime}(z_{1})>\mathcal{L}^{\prime\prime}(z_{2})$ and $\mathcal{L}(z_{1})<\mathcal{L}(z_{2})$, indicating that high-loss regions are flatter than low-loss regions.
The same applies for $0>z_{1}>z_{2}>-c$.

###### Proposition 5.1 ( P2 . Proof in Section B.2 .) .

Under Definition 5.2, assume the landscape is sub-quadratic in the valley $V=[-2b,2b]$.
Then, for all initialization $\theta_{0}\in(-b,b)$, and $\eta,\rho$ s.t. $\eta<\min\limits_{z\in V}b/|\mathcal{L}^{\prime}(z)|$, $\rho\leq\min\{1/a,\eta\cdot\min\limits_{0<|z|<b}|\mathcal{L}^{\prime}(2z)/
\mathcal{L}^{\prime}(z)|\}$, the full-batch SAM will remain within the valley $V$, i.e., $\theta_{t}\in V$ for all $t\in\mathbb{N}$.

Proposition 5.1 shows that under a sub-quadratic landscape, even if SAM escapes the minimum $\theta^{\star}$ when $\eta>{2}/{a(1+a\rho)}$ and $\theta_{0}\approx\theta^{\star}$, it will still remain within the current valley.
This proposition holds when $\eta<\min\limits_{z\in V_{b}}b/|\mathcal{L}^{\prime}(z)|$, where the upper bound $\min\limits_{z\in V_{b}}b/|\mathcal{L}^{\prime}(z)|$ and can be sufficiently large under highly sub-quadratic landscapes (see Example B.1).

Support for P2.
Proposition 5.1 supports our key claim P2 that SAM remains within the current valley during escape. Intuitively, as shown in Section 5.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-quadratic landscape, the high-loss regions are much flatter, which satisfies SAM’s stability condition and prevents it from escaping the current valley.
The experiment in Figure 4 (c) back up this theoretical claim.

###### Definition 5.2 (Sub-quadratic landscape for p = 1 𝑝 1 p=1 italic_p = 1 ) .

###### Proposition 5.1 ( P2 . Proof in Section B.2 .) .

### 5.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 $\bm{\xi}(\bm{\theta})$.

###### Assumption 5.2 (Loss landscape: $L$ -smoothness and Polyak–Łojasiewicz) .

(i) There exists $L>0$ s.t. $\left\lVert H(\bm{\theta})\right\rVert_{2}\leq L$ for all $\bm{\theta}$; (ii) There exists $\mu>0$ s.t. $\left\lVert\nabla\mathcal{L}(\bm{\theta})\right\rVert^{2}\geq 2\mu\mathcal{L}(
\bm{\theta})$ for all $\bm{\theta}$.

Assumption 5.2 are quite standard and is frequently used in the non-convex optimization literature (Karimi et al., ).
Specifically, for neural networks, if $\mathcal{L}(\cdot)$ is four-times continuously differentiable, this assumption holds at least near the set of global minima of $\mathcal{L}$ (Arora et al., ).

###### Assumption 5.3 (Gradient noise magnitude) .

There exists $\sigma>0$ s.t. $\mathbb{E}[\left\lVert\bm{\xi}(\bm{\theta})\right\rVert^{2}]\leq\sigma^{2}
\mathcal{L}(\bm{\theta}),\forall\bm{\theta}$.

The bounded variance assumption $\mathbb{E}[\left\lVert\bm{\xi}(\bm{\theta})\right\rVert^{2}]\leq\sigma^{2}$ is commonly employed in classical optimization theory (Bubeck et al., ).
However, as discussed in Section 5.1, recent studies (Mori et al., ; Feng & Tu, ; Liu et al., ; Wojtowytsch, ) have shown that the magnitude of SGD noise is also bounded by the loss value.
Thus, we introduce Assumption 5.3: $\mathbb{E}[\left\lVert\bm{\xi}(\bm{\theta})\right\rVert^{2}]\leq\sigma^{2}
\mathcal{L}(\bm{\theta})$.
This indicates that the noise magnitude diminishes to zero at 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 in Theorem 5.3.

###### Theorem 5.3 ( P4 . Proof in Section B.3 .) .

Under Assumption 5.2 and 5.3, let $\{\bm{\theta}_{t}\}_{t}$ be the parameters trained by SAM.
If $\eta\leq\min\{1/L,\mu B/2L\sigma^{2}\}$ and $\rho\leq\min\left\{1/L,\mu B/4L\sigma^{2},\eta\mu^{2}/24L^{2}\right\}$, then we have $\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]\leq\left(1-\eta\mu/2\right
)^{t}\mathcal{L}(\bm{\theta}_{0}),\ \forall t\in\mathbb{N}$.

Support for P4.
Theorem 5.3 supports our key claim P4 that the convergence rate of stochastic SAM is significantly fast,
and notably faster than the previous result on SAM’s convergence rate (Andriushchenko & Flammarion, ).

Comparison with other convergence results for stochastic SAM (Andriushchenko & Flammarion, ; Shin et al., ).
These works, along with our result, all study the convergence on stochastic SAM under the standard Assumption 5.2.
Andriushchenko & Flammarion () has shown that stochastic SAM converges under a polynomially fast rate of $\mathcal{O}(1/t)$, but this is much slower than our exponentially fast rate.
The difference lies in the assumptions on gradient noise magnitude: they have used the classical assumption $\mathbb{E}[\left\lVert\bm{\xi}(\bm{\theta})\right\rVert^{2}]\leq\sigma^{2}$, whereas we rely on the more realistic Assumption 5.3.
Shin et al. () has also proved that stochastic SAM converges exponentially fast.
However, their analysis is based on an additional interpolation assumption due to over-parameterization.

###### Assumption 5.2 (Loss landscape: L 𝐿 L italic_L -smoothness and Polyak–Łojasiewicz) .

###### Assumption 5.3 (Gradient noise magnitude) .

###### Theorem 5.3 ( P4 . Proof in Section B.3 .) .

## 6 Is SAM Necessary in the Early Phase of Training?

Since we have both empirically observed and theoretically justified SAM’s efficiency in selecting flatter minima in the late phase of training, 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 at the end of training is more critical in shaping the final solution’s properties.
Based on this viewpoint, we extend our findings from SAM to Adversarial Training.

Figure: Figure 5: The impact of SAM training proportion on generalization/sharpness when switches from SAM to SGD. The generalization gap between the models $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ (top row)/the sharpness of $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$ (bottom row) vs. the SAM training proportion of $\bm{\theta}^{T}_{\textrm{SAM}\to\textrm{SGD},t}$. $\bm{\theta}^{T}_{\textrm{SAM}\to\textrm{SGD},t}$ refers to the model obtained by switching from SAM to SGD, while $\bm{\theta}_{\textrm{SAM}}^{T}$ represents the model trained entirely with SAM. Here, $T$ is the total number of training epochs, and $t$ is the switching point. We define SAM training proportion of $\bm{\theta}^{T}_{\textrm{SAM}\to\textrm{SGD},t}$ as $t/T$ and adjust it by varying $t$ while keeping $T$ fixed. Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations. Results on more datasets and architectures can be found in Section A.4.
Refer to caption: x5.png

Early-phase SAM offers only marginal 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 as Section 4, we use the switching method but switches from SAM to SGD, to obtain the model $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$.
We also train a baseline model entirely with SAM, yielding $\bm{\theta}_{\textrm{SAM}}^{T}$.
By varying $t$ while keep $T$ fixed, we adjust the proportion of SAM training, defined as $t/T$, to examine its impact on the generalization gap between $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$, as well as the sharpness of $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$.
In Figure 5, the generalization gap remains substantial even when trained with SAM for most of the time (e.g., $t/T=0.8$).
A similar trend is observed in the sharpness of $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$.
Only full SAM training (i.e., $t/T=1$) consistently finds the flat minima.
This suggests that applying SAM before SGD yields only limited improvements in generalization and sharpness compared to training entirely with SGD.

Together with findings in Section 4, we conjecture that the optimization algorithm chosen in the late training phase is more critical in determining the final solution’s properties.
However, current findings are based solely on generalization/sharpness—can we extend to other properties?

Figure: 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 model trained with different strategies. (b) Robust/natural error of $\bm{\theta}_{\textrm{SGD}\to\textrm{AT}}^{T}$ vs. the proportion of AT epochs $(T-t)/T$. In both (a) and (b), dashed lines represents the robust error $\operatorname{Err}_{\mathcal{D}_{\textrm{perturb}}}$, and solid lines presents the natural error $\operatorname{Err}_{\mathcal{D}_{\textrm{clean}}}$. Blue and green curves correspond to models trained solely with SGD and AT, i.e., $\bm{\theta}^{T}_{\textrm{SGD}}$ and $\bm{\theta}^{T}_{\textrm{AT}}$, respectively. Orange curves indicate models trained with the switching method, switching from SGD to AT at epoch $t$, i.e., $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{AT},t}$. All models are trained for $T=200$ epochs. In (a), we set the switching epoch $t$ for $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{AT},t}$ to $180$. In (b), we vary $t$ while keep $T$ fixed to adjust the AT training proportion for $\bm{\theta}_{\textrm{SGD}\to\textrm{AT}}^{T}$. All results are averaged over five trials with different random seeds.
Refer to caption: x6.png

Extension: Adversarial Training efficiently finds robust minima late in training.
Adversarial Training (AT) (Madry et al., ), in which networks are trained on adversarial examples (Szegedy et al., ), significantly improves the robustness of neural networks.
Recently, Zhang et al. () highlighted a duality between SAM and AT: while SAM perturbs the model parameters along training (see Equation 1), AT perturbs input samples, with both methods relying on 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 $t$, we obtain the model $\bm{\theta}_{\textrm{SGD}\to\textrm{AT},t}^{T}$.
We evaluate its robustness by classification error on adversarial examples, termed robust error $\operatorname{Err}_{\mathcal{D}_{\textrm{perturb}}}(\bm{\theta})$, and its performance on clean examples, referred to as natural error $\operatorname{Err}_{\mathcal{D}_{\textrm{clean}}}(\bm{\theta})$.
In Figure 6 (a), the robust error of $\bm{\theta}_{\textrm{SGD}\to\textrm{AT},t}^{T}$ drops sharply after switching to AT at $180$ epoch, matching that of full AT training, while the natural error only slightly spikes.
Additionally, we vary the switching point $t$ to examine how the robustness of $\bm{\theta}_{\textrm{SGD}\to\textrm{AT},t}^{T}$ changes with the increasing proportion of AT epochs, i.e., $(T-t)/T$.
In Figure 6 (b), the robust error of $\bm{\theta}_{\textrm{AT}}^{T}$ is comparable to that of models trained entirely with AT (i.e., when $(T-t)/T=1$) once the proportion of AT epochs exceeds 0 0.
These results together confirm that AT efficiently enhances the robustness of neural networks late in training.

## 7 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 a 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 suggests that late training dynamics may play an important role if we want to understand 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 a constant perturbation radius $\rho$ and using squared loss.
We also note that our theory only explains the behaviour of SAM in the late training phase; it cannot account for the phenomena observed for early-phase SAM.

## Appendix A More Experimental Details and Results

### A.1 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 lies in that SAM introduces an extra hyper-parameter, perturbation radius $\rho$.
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-10 dataset.
We train the VGG-16 (Simonyan & Zisserman, ) and the ResNet-20 architecture (He et al., ) on
the CIFAR-10 dataset (Krizhevsky et al., ), using standard data
augmentation and hyperparameters as in the original papers (Simonyan & Zisserman, ; He et al., ).
Data augmentation techniques include random horizontal flips and random
$32\times 32$ pixel crops.
A weight decay of $5\times 10^{-4}$ is applied, and the momentum for gradient update is set to $0.9$.
The learning rate is initialized at $0.1$ and is dropped by $10$ times at epoch $80$(^7^77Note that in the original paper, the learning rate is reduced by a factor of $10$ at epoch $120$ as well. However, based on our theory in Section 5.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.).
The total number of training epochs is $160$.

WideResNet-16-8 on the CIFAR-10 dataset, and WideResNet-28-10 on the CIFAR-100 dataset.
We train the WideResNet-16-8 architecture (Zagoruyko & Komodakis, ) on the CIFAR-10 dataset and the WideResNet-28-10 architecture (He et al., ) on
the CIFAR-100 dataset (Krizhevsky et al., ), using standard data
augmentation and hyperparameters as in the original papers (Zagoruyko & Komodakis, ).
Data augmentation techniques include random horizontal flips, random
$32\times 32$ pixel crop, and cutout regularization (DeVries, ).
A weight decay of $5\times 10^{-4}$ is applied, and the momentum for gradient update is set to $0.9$.
The learning rate is initialized at $0.1$ and is dropped by $5$ times at $60$ and $120$ epochs.
The total number of training epochs is $200$.

Then, we specify the extra hyper-parameter for SAM, i.e., the perturbation radius $\rho$.

Additional setup for SAM and ASAM.
For SAM, the perturbation radius is set to $0.05$ across all architectures and datasets, as recommended by Foret et al. ().
In Section A.4, our findings on SAM generalize to another SAM variant, ASAM (Kwon et al., ).
For ASAM, the perturbation radius is set to $2.0$, as recommended by Kwon et al. ().

In Section 6, 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., ) to generate adversarial examples $\bm{x}^{\textrm{Adv}}$ from the original image $\bm{x}$, with the constraint that $\left\lVert\bm{x}^{\textrm{Adv}}-\bm{x}\right\rVert_{\infty}\leq\epsilon$, where $\left\lVert\cdot\right\rVert_{\infty}$ denotes the $L^{\infty}$ norm.
Following the standard configurations of Madry et al. (), we set the perturbation limit $\epsilon$ to $4/255$ and the step size for PGD attack to $2/255$.

### A.2 More Experiments on Late-Phase SAM

Figure: Figure 7: Additional experiments: the impact of SAM training proportion on generalization/sharpness when switches from SGD to SAM. The generalization gap between the models $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ (top row) / the sharpness of $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ (bottom row) vs. the SAM training proportion of $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$. $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$ refers to the model obtained by switching from SGD to SAM, while $\bm{\theta}_{\textrm{SAM}}^{T}$ represents the model trained entirely with SAM. Here, $T$ is the total number of training epochs, and $t$ is the switching point. We define SAM training proportion of $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$ as $(T-t)/T$ and adjust it by varying $t$ while keeping $T$ fixed. Dots represent the mean over five trials with different random seeds, and error bars indicate standard deviations.
Refer to caption: x7.png

In this subsection, we provide more experimental results to demonstrate the effect of SAM on generalization and sharpness when applied during the late phase of training.
In Figure 7, we include more experiment results for ResNet-20 and VGG-16 on the CIFAR-10 datasets.
Similar to our results in Figure 3, the generalization gap between $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ drops to zero once the SAM training proportion is greater than 0 0.
A similar trend is observed in the sharpness of $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$.
These results further confirms our main discovery that SAM efficiently finds flatter minima over SGD late in training.

### A.3 Extended Experiments on ASAM

In this subsection, we generalize our findings on SAM to another SAM variant, ASAM (Kwon et al., ).
Similar to SAM, we 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, we apply the switching method to investigate the effect of the late-phase ASAM accross various architectures and datasets.
We use the switching method, which switches from SGD to ASAM, to obtain the model $\bm{\theta}_{\textrm{SGD}\to\textrm{ASAM},t}^{T}$.
We also train a baseline model, i.e., $\bm{\theta}_{\textrm{SGD}\to\textrm{SAM},t}^{T}$.
We vary the switching point $t$ while keep $T$ fixed to adjust the proportion of ASAM training and study its impact on the generalization gap between $\bm{\theta}_{\textrm{ASAM}}^{T}$ and $\bm{\theta}_{\textrm{SGD}\to\textrm{ASAM},t}^{T}$, as well as the sharpness of $\bm{\theta}_{\textrm{SGD}\to\textrm{ASAM},t}^{T}$.
As shown in Figure 8, both the generalization gap and the sharpness drops 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.

Figure: Figure 8: The impact of ASAM training proportion on generalization/sharpness when switches from SGD to ASAM. The generalization gap between the models $\bm{\theta}_{\textrm{SGD}\to\textrm{ASAM},t}^{T}$ and $\bm{\theta}_{\textrm{ASAM}}^{T}$ (top row) / the sharpness of $\bm{\theta}_{\textrm{SGD}\to\textrm{ASAM},t}^{T}$ (bottom row) vs. the ASAM training proportion of $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{ASAM},t}$. $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{ASAM},t}$ refers to the model obtained by switching from SGD to ASAM, while $\bm{\theta}_{\textrm{ASAM}}^{T}$ represents the model trained entirely with ASAM. Here, $T$ is the total number of training epochs, and $t$ is the switching point. We define ASAM training proportion of $\bm{\theta}^{T}_{\textrm{SGD}\to\textrm{SAM},t}$ as $(T-t)/T$ and adjust it by varying $t$ while keeping $T$ fixed. Dots represent the mean over three trials with different random seeds, and error bars indicate standard deviations.
Refer to caption: x8.png

### A.4 More Experiments on Early-Phase SAM

In this subsection, we provide more experimental results to demonstrate the effect of SAM on generalization and sharpness when applied during the early phase of training.
In Figure 9, we include more experiment results for ResNet-20 and VGG-16 on the CIFAR-10 datasets.
Similar to our results in Figure 5, the generalization gap between $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ remains substantial until the SAM training proportion exceeds $0.6$.
The same applies for the the sharpness of $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$.
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.

Figure: Figure 9: Additional experiments: the impact of SAM training proportion on generalization/sharpness when switches from SAM to SGD. The generalization gap between the models $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$ and $\bm{\theta}_{\textrm{SAM}}^{T}$ (top row)/the sharpness of $\bm{\theta}_{\textrm{SAM}\to\textrm{SGD},t}^{T}$ (bottom row) vs. the SAM training proportion of $\bm{\theta}^{T}_{\textrm{SAM}\to\textrm{SGD},t}$. $\bm{\theta}^{T}_{\textrm{SAM}\to\textrm{SGD},t}$ refers to the model obtained by switching from SAM to SGD, while $\bm{\theta}_{\textrm{SAM}}^{T}$ represents the model trained entirely with SAM. Here, $T$ is the total number of training epochs, and $t$ is the switching point. We define SAM training proportion of $\bm{\theta}^{T}_{\textrm{SAM}\to\textrm{SGD},t}$ as $t/T$ and adjust it by varying $t$ while keeping $T$ fixed. Dots represent the mean over five trials with different random seeds, and error bars indicate standard deviations.
Refer to caption: x9.png

## Appendix B Proofs in Section 5

In our theoretical analysis of Stochastic SAM, defined in Equation 4, we adopt the technical simplification in Appendix C.2 of Andriushchenko & Flammarion (), which assume that the inner noise and outer noise are independent.
Specifically, we analyze the following variant of Equation 4:

$$ $\displaystyle\bm{\theta}_{t+1}=\bm{\theta}_{t}-\eta\mathcal{L}_{\xi_{t}^{2}}( \bm{\theta}_{t}+\rho\nabla\mathcal{L}_{\xi_{t}^{1}}(\bm{\theta}_{t}))$ (6) $$

where $\xi_{t}^{1}$ and $\xi_{t}^{2}$ are two independent stochastic mini-batches of data of size $B$.

### B.1 Proof of Theorem 5.1 and 5.2

For simplicity, we denote the two-step update rules of SAM defined in Equation 6 as:

$$ $\displaystyle\bm{\theta}_{t+1/2}=\bm{\theta}_{t}+\rho\nabla\mathcal{L}_{\xi_{t }^{1}}(\bm{\theta}_{t}),$ $\displaystyle\bm{\theta}_{t+1}=\bm{\theta}_{t+1/2}-\eta\nabla\mathcal{L}_{\xi_ {t}^{2}}(\bm{\theta}_{t+1/2}).$ $$

For $B$-batch gradient noise, denoted by $\bm{\xi}^{B}(\bm{\theta})$, we have:

$$ $\displaystyle\mathbb{E}[\bm{\xi}^{B}(\bm{\theta})\bm{\xi}^{B}(\bm{\theta})^{ \top}]=\frac{1}{B}\mathbb{E}[\bm{\xi}(\bm{\theta})\bm{\xi}(\bm{\theta})^{\top} ]=\frac{1}{B}\Sigma(\bm{\theta}),$ (7) $$

where $\bm{\xi}(\bm{\theta})$ is the 1-batch gradient noise, used in Assumption 5.1.

For the linearized model ( Equation 5), it holds that:

$$ $\displaystyle\mathcal{L}(\bm{\theta})=\frac{1}{2}(\bm{\theta}-\bm{\theta}^{ \star})^{\top}H(\bm{\theta}^{\star})(\bm{\theta}-\bm{\theta}^{\star}).$ $$

Without loss of generality, we can let $\bm{\theta}^{\star}=0$.
For simplicity, we denote $G:=H(\bm{\theta}^{\star})$.

Next, according to the definition of linear stability (Definition 5.1), we consider the dynamics of stochastic SAM on the linearized model.

First, we have the following estimate for $\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1})\right]$:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1})\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}_{\xi_{t}^{2}}(\bm{\theta}_{t+1/2}))^{\top}G(\bm{\theta}_{t+1/2}- \eta\nabla\mathcal{L}_{\xi_{t}^{2}}(\bm{\theta}_{t+1/2}))\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}(\bm{\theta}_{t+1/2})-\eta\bm{\xi}^{B}(\bm{\theta}_{t+1/2}))^{\top} G(\bm{\theta}_{t+1/2}-\eta\nabla\mathcal{L}(\bm{\theta}_{t+1/2})-\eta\bm{\xi}^ {B}(\bm{\theta}_{t+1/2}))\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}(\bm{\theta}_{t+1/2})^{\top}G(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}(\bm{\theta}_{t+1/2})\right]+0$ $\displaystyle+\frac{1}{2}\eta^{2}\mathbb{E}\left[{\bm{\xi}^{B}(\bm{\theta}_{t+ 1/2})}^{\top}G\bm{\xi}^{B}(\bm{\theta}_{t+1/2})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: linear: batch noise})}}{=}$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}(\bm{\theta}_{t+1/2})^{\top}G(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}(\bm{\theta}_{t+1/2})\right]+\frac{\eta^{2}}{2B}\mathbb{E}\left[ \operatorname{Tr}\left(\Sigma(\bm{\theta}_{t+1/2})G)\right)\right]$ $\displaystyle\overset{\text{Assumption~{}\ref{ass:noise}}}{\geq}0+\frac{\eta^{ 2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\mathbb{E}\left[\mathcal{L}(\bm{ \theta}_{t+1/2})\right]=\frac{\eta^{2}\gamma\left\lVert G\right\rVert_{F}^{2}} {B}\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right].$ $$

Then we estimate the lower bound for $\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]$:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t}-\rho\nabla\mathcal{L} _{\xi_{t}^{1}}(\bm{\theta}_{t}))^{\top}G(\bm{\theta}_{t}+\rho\nabla\mathcal{L} _{\xi_{t}^{1}}(\bm{\theta}_{t}))\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t}+\rho\nabla\mathcal{L} (\bm{\theta}_{t})+\rho\bm{\xi}^{B}(\bm{\theta}_{t}))^{\top}G(\bm{\theta}_{t}+ \rho\nabla\mathcal{L}(\bm{\theta}_{t})+\rho\bm{\xi}^{B}(\bm{\theta}_{t}))\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t}+\rho\nabla\mathcal{L} (\bm{\theta}_{t}))^{\top}G(\bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{\theta}_{ t}))\right]+\frac{1}{2}\rho^{2}\mathbb{E}\left[{\bm{\xi}^{B}(\bm{\theta}_{t})} ^{\top}G\bm{\xi}^{B}(\bm{\theta}_{t})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: linear: batch noise})}}{=}$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t}+\rho\nabla\mathcal{L} (\bm{\theta}_{t}))^{\top}G(\bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{\theta}_{ t}))\right]+\frac{\rho^{2}}{2B}\mathbb{E}\left[\operatorname{Tr}\left(\Sigma( \bm{\theta}_{t})G\right)\right]$ $\displaystyle\overset{\text{Assumption~{}\ref{ass:noise}}}{\geq}$ $\displaystyle\frac{1}{2}\mathbb{E}\left[(\bm{\theta}_{t}+\rho\nabla\mathcal{L} (\bm{\theta}_{t}))^{\top}G(\bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{\theta}_{ t}))\right]+\frac{\rho^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\mathbb{E }\left[\mathcal{L}(\bm{\theta}_{t})\right]$ $\displaystyle=$ $\displaystyle\frac{1}{2}\mathbb{E}\left[\bm{\theta}_{t}^{\top}(I+\rho G)G(I+ \rho G)\bm{\theta}_{t}\right]+\frac{\rho^{2}\gamma\left\lVert G\right\rVert_{F }^{2}}{B}\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]$ $\displaystyle\overset{(I+\rho G)G(I+\rho G)\succeq G}{\geq}$ $\displaystyle\frac{1}{2}\mathbb{E}\left[\bm{\theta}_{t}^{\top}G\bm{\theta}_{t} \right]+\frac{\rho^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\mathbb{E} \left[\mathcal{L}(\bm{\theta}_{t})\right]$ $\displaystyle=$ $\displaystyle\left(1+\frac{\rho^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B} \right)\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right].$ $$

By combining these two estimates, we obtain:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1})\right]\geq\frac{ \eta^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\left(1+\frac{\rho^{2} \gamma\left\lVert G\right\rVert_{F}^{2}}{B}\right)\mathbb{E}\left[\mathcal{L}( \bm{\theta}_{t})\right]$ $$

Hence, if $\bm{\theta}^{\star}$ is linear stable for stochastic SAM, it must hold:

$$ $\displaystyle\frac{\eta^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\left(1+ \frac{\rho^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\right)\leq 1,$ (8) $$

which completes the proof of Theorem 5.1.

Furthermore, if the condition (8) can not be satisfied, i.e., $\frac{\eta^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\left(1+\frac{\rho^{2
}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\right)>1$, then the exponentially fast escape holds:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]\geq C^{t} \mathbb{E}\left[\mathcal{L}(\bm{\theta}_{0})\right],$ $$

where $C=\frac{\eta^{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\left(1+\frac{\rho^
{2}\gamma\left\lVert G\right\rVert_{F}^{2}}{B}\right)>1$. This complete the proof of Theorem 5.2.

### B.2 Proof of Proposition 5.1

As we mentioned in Section 5.2, this proposition focus on a model size of $p=1$ and consider the full-batch SAM, i.e., $\theta_{t+1}=\theta_{t}-\eta\nabla\mathcal{L}(\theta_{t}+\rho\nabla\mathcal{L}
(\theta_{t}))$.

###### Proof of Proposition 5.1 .

First, we define the hitting time:

$$ $\displaystyle T:=\inf\left\{t\in\mathbb{N}:|\theta_{t+1}|\geq b\right\}.$ $$

We aim to prove that $T$ does not exists, which implies that $\theta_{t}\in(-b,b)\in V_{b}$ holds for all $t\in\mathbb{N}$.

Assuming $T$ exists, then one of the following three cases must hold:

- •
Case I: $\theta_{T}=0$. In this case, $\mathcal{L}^{\prime}(\theta_{T})=0$, which implies that $\theta_{T+1}=0$. This contradicts the definition of $T$.
- •
Case II: $0<\theta_{T}<b$.
By the sub-quadratic property and $\rho\leq 1/a$, for the inner update $\theta_{T+1/2}:=\theta_{T}+\rho\mathcal{L}^{\prime}(\theta_{T})$, we have:
$\displaystyle 0<\theta_{T}<\theta_{T+1/2}=\theta_{T}+\rho\int_{0}^{\theta_{T}}
\mathcal{L}^{\prime\prime}(z){\textnormal{d}}z<\theta_{T}+\rho a\theta_{T}\leq
2
\theta_{T}<2b.$
Consequently, we have the two-sided bound for $\theta_{T+1}$:
$\displaystyle\theta_{T+1}=$
$\displaystyle\theta_{T+1/2}-\eta\mathcal{L}^{\prime}(\theta_{T+1/2})>\theta_{T
+1/2}-\eta\sup_{z\in(0,2b)}\mathcal{L}^{\prime}(z)$
$\displaystyle\geq$
$\displaystyle 0-\eta\max_{z\in(-2b,2b)}\mathcal{L}^{\prime}(z)\geq-b;$
$\displaystyle\theta_{T+1}=$
$\displaystyle\theta_{T}+\rho\mathcal{L}^{\prime}(\theta_{T})-\eta\mathcal{L}^{
\prime}(\theta_{T+1/2})$
$\displaystyle<$
$\displaystyle\theta_{T}+\eta\mathcal{L}^{\prime}(\theta_{T+1/2})-\eta\mathcal{
L}^{\prime}(\theta_{T+1/2})=\theta_{T}<b.$
Then we have $|\theta_{T+1}|<b$, which contradicts the definition of $T$.
- •
Case III: $-b<\theta_{T}<0$. The proof for this case is extremely similar to the proof for Case II. This case is also contradicts the definition of $T$.

Thus, we have proved that $T$ does not exist, which implies that $\theta_{T+1}=0$.
∎

The next result provides an example for the highly sub-quadratic landscape.
Intuitively, it means that the decay of $\mathcal{L}^{\prime\prime}(|z|)$ is extremely fast.

###### Example B.1 (highly sub-quadratic landscape) .

Let $\mathcal{L}(0)=\mathcal{L}^{\prime}(0)=0$, $\mathcal{L}^{\prime\prime}(0)=a$, and

$$ $\mathcal{L}^{\prime\prime}(z)=\begin{cases}a-\frac{|z|}{\epsilon},&|z|<a \epsilon\\ 0,&\text{else in $[-2b,2b]$}\end{cases},$ $$

where $0<\epsilon\ll 1$ is a constant to reflect the sub-quadratic degree.

Notice that this example satisfies Definition 5.2. Moreover, it is straightforward that

$$ $\displaystyle\mathcal{L}^{\prime}(z)=\begin{cases}z\left(a-\frac{|z|}{2 \epsilon}\right),&|z|<a\epsilon\\ \frac{a\epsilon}{2},&\text{else in $[-2b,2b]$}\end{cases}.$ $$

Additionally, we can calculate the upper bounds for $\eta$ and $\rho$ in Proposition 5.1:

$$ $\displaystyle\eta\leq\min_{|z|\leq 2b}\frac{b}{|\mathcal{L}^{\prime}(z)|}= \frac{2b}{\epsilon a},$ $\displaystyle\rho\leq\min\left\{\frac{1}{a},\min_{|z|\leq b}\left|\frac{ \mathcal{L}^{\prime}(2z)}{\mathcal{L}^{\prime}(z)}\right|\right\}=\min\left\{ \frac{1}{a},2\right\}.$ $$

One can see that in the highly sub-quadratic case, i.e., $0<\epsilon\ll 1$, the upper bound for $\eta$ goes to $+\infty$.

In addition, we prove Proposition B.1, which illustrates that SAM will escape the current valley under a quadratic landscape.
This result stands in stark contrast to Proposition 5.1 for sub-quadratic landscape, which further confirms our adaptation of the sub-quadratic landscape.

###### Proposition B.1 (Non-local escape behaviour under a quadratic landscape.) .

Assume the landscape in the valley $V_{b}=(-2b,2b)$ be quadratic, i.e., $\mathcal{L}(z)=az^{2}/2$.
Then, for all initialization $\theta_{0}\in V_{b}\backslash\{0\}$, and $\eta,\rho$ s.t. $\eta>\frac{2}{a(1+a\rho)}$, full-batch SAM will escape from the valley $V_{b}$, i.e., there exists $T>0$ s.t. $\theta_{T}\notin V_{b}$.

###### Proof of Proposition B.1 .

By a straightforward calculation, we have

$$ $\displaystyle\theta_{t+1}=\theta_{t}-\eta a\left(\theta_{t}+\rho\cdot a\theta_ {t}\right)=(1-\eta a(1+a\rho))\theta_{t}.$ $$

If $\eta>\frac{2}{a(1+a\rho)}$, then $|\theta_{t}|\geq|1-\eta a(1+a\rho)|^{t}|\theta_{0}|$,
where $|1-\eta a(1+a\rho)|>1$.
This means that there exists $T>0$ s.t. $\theta_{T}\notin V_{b}$.
∎

###### Proof of Proposition 5.1 .

###### Example B.1 (highly sub-quadratic landscape) .

###### Proposition B.1 (Non-local escape behaviour under a quadratic landscape.) .

###### Proof of Proposition B.1 .

### B.3 Proof of Theorem 5.3

For simplicity, we denote the two-step update rules of SAM defined in Equation 6 as:

$$ $\displaystyle\bm{\theta}_{t+1/2}=\bm{\theta}_{t}+\rho\nabla\mathcal{L}_{\xi_{t }^{1}}(\bm{\theta}_{t}),$ $\displaystyle\bm{\theta}_{t+1}=\bm{\theta}_{t+1/2}-\eta\nabla\mathcal{L}_{\xi_ {t}^{2}}(\bm{\theta}_{t+1/2}).$ $$

Then we list a few useful properties under the condition of Theorem 5.3.

###### Lemma B.1 .

Suppose Assumption 5.2. Then we have:

$$ $\displaystyle 2\mu\mathcal{L}(\bm{\theta})\leq\left\lVert\nabla\mathcal{L}(\bm {\theta})\right\rVert^{2}\leq\frac{2L^{2}}{\mu}\mathcal{L}(\bm{\theta}),\ \forall\bm{\theta}.$ (9) $$

###### Proof of Lemma B.1 .

The upper bound is a direct corollary of PL and smoothness. Please refer to Lemma B.6 in Arora et al. () and Lemma 5 in Ahn et al. () for proof details.
∎

The smoothness in Assumption 5.2 also implies the classical quadratic upper bound for the loss (Bubeck et al., ):

$$ $\mathcal{L}(\bm{\theta}+\bm{v})\leq\mathcal{L}(\bm{\theta})+\left\langle\nabla \mathcal{L}(\bm{\theta}),\bm{v}\right\rangle+\frac{L}{2}\left\lVert\bm{v} \right\rVert^{2},\ \forall\bm{\theta},\bm{v}.$ (10) $$

Additionally, Assumption 5.3 holds for the gradient noise $\bm{\xi}(\bm{\theta})$ with batch size 1.
Thus, for $B$-batch gradient noise, denoted by $\bm{\xi}^{B}(\bm{\theta})$, we have $\mathbb{E}[\bm{\xi}^{B}(\bm{\theta})\bm{\xi}^{B}(\bm{\theta})^{\top}]=\frac{1}
{B}\mathbb{E}[\bm{\xi}(\bm{\theta})\bm{\xi}(\bm{\theta})^{\top}]$, implying:

$$ $\mathbb{E}\left[\left\lVert\bm{\xi}^{B}(\bm{\theta})\right\rVert^{2}\right] \leq\frac{\sigma^{2}}{B}\mathcal{L}(\bm{\theta}),\ \forall\bm{\theta}.$ (11) $$

For the learning rate $\eta$ and the permutation radius $\rho$, we choose then as follows:

$$ $\eta\leq\min\left\{\frac{1}{L},\frac{\mu B}{2L\sigma^{2}}\right\},\quad\rho \leq\min\left\{\frac{1}{L},\frac{\mu B}{4L\sigma^{2}},\frac{\eta\mu^{2}}{24L^{ 2}}\right\}.$ (12) $$

###### Lemma B.1 .

###### Proof of Lemma B.1 .

#### B.3.1 Proof of Theorem 5.3

For $\mathcal{L}(\bm{\theta}_{t+1})$, it has the quadratic upper bound:

$$ $\displaystyle\mathcal{L}(\bm{\theta}_{t+1})=\mathcal{L}\left(\bm{\theta}_{t+1/ 2}-\eta\nabla\mathcal{L}_{\xi_{t}^{2}}(\bm{\theta}_{t+1/2})\right)=\mathcal{L} \left(\bm{\theta}_{t+1/2}-\eta\nabla\mathcal{L}(\bm{\theta}_{t+1/2})-\eta\bm{ \xi}^{B}(\bm{\theta}_{t+1/2})\right)$ $\displaystyle\overset{\text{(\ref{proof: equ: quadratic bound})}}{\leq}$ $\displaystyle\mathcal{L}\left(\bm{\theta}_{t+1/2}-\eta\nabla\mathcal{L}(\bm{ \theta}_{t+1/2})\right)-\left\langle\nabla\mathcal{L}\left(\bm{\theta}_{t+1/2} +\eta\nabla\mathcal{L}(\bm{\theta}_{t+1/2})\right),\eta\bm{\xi}(\bm{\theta}_{t +1/2})\right\rangle$ $\displaystyle+\frac{\eta^{2}L}{2}\left\lVert\bm{\xi}^{B}(\bm{\theta}_{t+1/2}) \right\rVert_{2}^{2}.$ $$

Taking the expectation, we obtain the upper bound of $\mathbb{E}[\mathcal{L}(\bm{\theta}_{t+1})]$:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1})\right]\leq\mathbb{ E}\left[\mathcal{L}\left(\bm{\theta}_{t+1/2}-\eta\nabla\mathcal{L}(\bm{\theta} _{t+1/2})\right)\right]+\frac{\eta^{2}L}{2}\mathbb{E}\left[\left\lVert\bm{\xi} ^{B}(\bm{\theta}_{t+1/2})\right\rVert^{2}\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: noise magnitude})}}{\leq}$ $\displaystyle\mathbb{E}\left[\mathcal{L}\left(\bm{\theta}_{t+1/2}-\eta\nabla \mathcal{L}(\bm{\theta}_{t+1/2})\right)\right]+\frac{\eta^{2}L}{2}\frac{\sigma ^{2}}{B}\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: quadratic bound})}}{\leq}$ $\displaystyle\mathbb{E}\left[\mathcal{L}\left(\bm{\theta}_{t+1/2}\right)-\eta \left\lVert\nabla\mathcal{L}(\bm{\theta}_{t+1/2})\right\rVert^{2}+\frac{\eta^{ 2}L}{2}\left\lVert\nabla\mathcal{L}(\bm{\theta}_{t+1/2})\right\rVert^{2}\right ]+\frac{\eta^{2}L}{2}\frac{\sigma^{2}}{B}\mathbb{E}\left[\mathcal{L}(\bm{ \theta}_{t+1/2})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: PL})}}{\leq}$ $\displaystyle\left(1-\eta(2-\eta L)\mu+\frac{\eta^{2}L\sigma^{2}}{2B}\right) \mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: eta})}}{\leq}$ $\displaystyle\left(1-\frac{3\eta\mu}{4}\right)\mathbb{E}\left[\mathcal{L}(\bm{ \theta}_{t+1/2})\right].$ $$

In the similar way, for $\mathcal{L}(\bm{\theta}_{t+1/2})$, it has the quadratic upper bound:

$$ $\displaystyle\mathcal{L}(\bm{\theta}_{t+1/2})=\mathcal{L}\left(\bm{\theta}_{t} +\rho\nabla\mathcal{L}_{\xi_{t}^{1}}(\bm{\theta}_{t})\right)=\mathcal{L}\left( \bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{\theta}_{t})+\rho\bm{\xi}^{B}(\bm{ \theta}_{t})\right)$ $\displaystyle\overset{\text{(\ref{proof: equ: quadratic bound})}}{\leq}$ $\displaystyle\mathcal{L}\left(\bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{\theta }_{t})\right)+\left\langle\nabla\mathcal{L}\left(\bm{\theta}_{t}+\rho\nabla \mathcal{L}(\bm{\theta}_{t})\right),\rho\bm{\xi}(\bm{\theta}_{t})\right\rangle +\frac{\rho^{2}L}{2}\left\lVert\bm{\xi}^{B}(\bm{\theta}_{t})\right\rVert_{2}^{ 2}.$ $$

Taking the expectation, we obtain the upper bound of $\mathbb{E}[\mathcal{L}(\bm{\theta}_{t+1/2})]$:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]\leq \mathbb{E}\left[\mathcal{L}\left(\bm{\theta}_{t}+\rho\nabla\mathcal{L}(\bm{ \theta}_{t})\right)\right]+\frac{\rho^{2}L}{2}\mathbb{E}\left[\left\lVert\bm{ \xi}^{B}(\bm{\theta}_{t})\right\rVert^{2}\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: noise magnitude})}}{\leq}$ $\displaystyle\mathbb{E}\left[\mathcal{L}\left(\bm{\theta}_{t}+\rho\nabla \mathcal{L}(\bm{\theta}_{t})\right)\right]+\frac{\rho^{2}L}{2}\frac{\sigma^{2} }{B}\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: quadratic bound})}}{\leq}$ $\displaystyle\mathbb{E}\left[\mathcal{L}\left(\bm{\theta}_{t}\right)+\rho\left \lVert\nabla\mathcal{L}(\bm{\theta}_{t})\right\rVert^{2}+\frac{\rho^{2}L}{2} \left\lVert\nabla\mathcal{L}(\bm{\theta}_{t})\right\rVert^{2}\right]+\frac{ \rho^{2}L}{2}\frac{\sigma^{2}}{B}\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: PL})}}{\leq}$ $\displaystyle\left(1+\frac{\rho(2+\rho L)L^{2}}{\mu}+\frac{\rho^{2}L\sigma^{2} }{2B}\right)\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]$ $\displaystyle\overset{\text{(\ref{proof: equ: eta})}}{\leq}$ $\displaystyle\left(1+\frac{6\rho L^{2}}{\mu}\right)\mathbb{E}\left[\mathcal{L} (\bm{\theta}_{t})\right].$ $$

Finally, by combining our estimates for $\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1})\right]$ and $\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]$, we obtain:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1})\right]\leq\left(1- \frac{3\eta\mu}{4}\right)\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t+1/2})\right]$ $\displaystyle\leq$ $\displaystyle\left(1-\frac{3\eta\mu}{4}\right)\left(1+\frac{6\rho L^{2}}{\mu} \right)\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]\overset{\text{(\ref {proof: equ: eta})}}{\leq}\left(1-\frac{\eta\mu}{2}\right)\mathbb{E}\left[ \mathcal{L}(\bm{\theta}_{t})\right],$ $$

which implies the convergence rate:

$$ $\displaystyle\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{t})\right]\leq\left(1- \frac{\eta\mu}{2}\right)^{t}\mathbb{E}\left[\mathcal{L}(\bm{\theta}_{0})\right ],\ \forall t\in\mathbb{N}.$ $$

### B.4 Comparison between two Definitions on Linear Stability

###### Definition B.1 (Norm/loss-based linear stability) .

For an update rule $\bm{\theta}_{t+1}=\bm{\theta}_{t}+F(\bm{\theta}_{t};\bm{\xi}_{t})$,

- •
(Norm-based linear stability)
A global minimum $\bm{\theta}^{\star}$ is said to be linearly stable if there exists a $C>0$ such that it holds for SAM on the linearized model (Equation 5) that $\mathbb{E}[\left\lVert\bm{\theta}_{t}-\bm{\theta}^{\star}\right\rVert^{2}]\leq
C
\mathbb{E}[\left\lVert\bm{\theta}_{0}-\bm{\theta}^{\star}\right\rVert^{2}],
\forall t\geq 0$.
- •
(Loss-based linear stability)
A global minimum $\bm{\theta}^{\star}$ is said to be linearly stable if there exists a $C>0$ such that it holds for SAM on the linearized model (Equation 5) that $\mathbb{E}[\mathcal{L}(\bm{\theta}_{t})]\leq C\mathbb{E}[\mathcal{L}(\bm{
\theta}_{0})],\forall t\geq 0$.

###### Lemma B.2 .

The following two-fold results holds:

- •
(L1). If $\bm{\theta}^{\star}$ is norm-based linearly stable, then it must be loss-based linearly stable.
- •
(L2). There exists an update rule and a $G(\bm{\theta}^{\star})$, such that the converse of (L1) does not hold.

###### Proof of Lemma B.2 .

For (L1), it is straightforward by

$$ $\mathcal{L}(\bm{\theta}_{t})=\frac{1}{2}(\bm{\theta}-\bm{\theta}^{\star})^{ \top}G(\bm{\theta}^{\star})(\bm{\theta}-\bm{\theta}^{\star})\leq\left\lVert G( \bm{\theta}^{\star})\right\rVert\left\lVert\bm{\theta}_{t}-\bm{\theta}^{\star} \right\rVert^{2}.$ $$

Thus, the norm-based linear stability can imply the loss-based linear stability.

For (L2), we only need to consider an example containing degenerated directions. Specifically, let

$$ $\displaystyle L(\bm{\theta})=\frac{1}{2}\theta_{1}^{2}+0\cdot\theta_{2}^{2},$ $$

with $\bm{\theta}^{\star}=\mathbf{0}$.

Consider an update rule only occurs on $\theta_{2}$:

$$ $\displaystyle\theta_{t+1,1}=\theta_{t,1};\quad\theta_{t+1,2}=\theta_{t,2}+ \zeta_{t},\ \zeta_{t}\sim\mathcal{N}(0,1),$ $$

starting from $\bm{\theta}_{0}=(1,1)^{\top}$.
Then it is loss-base linearly stable due to

$$ $\displaystyle\mathcal{L}(\bm{\theta}_{t})=\frac{1}{2}\theta_{t,1}^{2}+0\equiv \frac{1}{2}\theta_{t,1}^{2}=\equiv\mathcal{L}(\bm{\theta}_{0}).$ $$

However, it is norm-base linearly nonstable:

$$ $\displaystyle\mathbb{E}\left[\left\lVert\bm{\theta}_{t}\right\rVert^{2}\right] =1+\mathbb{E}\left[|\bm{\theta}_{t,2}|^{2}\right]=1+\mathbb{E}\left[\left|1+ \sum_{s=0}^{t-1}\zeta_{s}\right|^{2}\right]=2+t\to+\infty\text{\ when $t\to+ \infty$}.$ $$

∎

###### Definition B.1 (Norm/loss-based linear stability) .

###### Lemma B.2 .

###### Proof of Lemma B.2 .

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: Could the authors explain how the neighborhood size affects the results when switching from SAM to SGD?
soundness: 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.
The 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.

The 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.
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.
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.

2.In theory, the switch from SGD to SAM is expected to yield certain benefits; however, no experimental evidence is provided to support this.

3.The models discussed are primarily convolution-based, so it would be beneficial to include some transformer-based models as well.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 3
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.
- L081: Efficient sounds like a technical ter
- 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.)
- L115: Please define classification error explicitly.
- L177: "when switches from SGD to SAM"
- 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.
- L286: Why change $T$ and $t$'s usage?
- L291: Typo after $\mathcal{L}''$
- L318: Is there a reason for usage of $\ell$ here?
- 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?
- L354: Please alert the reader that you will introduce Proposition 4.1 to resolve the conflict between the exponential escape and low-loss assumption
- L388: Do these results generalize to cases where $p>1$ and minibatch SAM?
- L406: Can we assume any valley around $\theta^*$ found by SGD to be sub-quadratic?
- 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?
soundness: 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.
- The authors make a clear and convincing case regarding the importance of adopting SAM in late (vs. early training) using empirical and theoretical findings.
- 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.
- The extension of the authors' results to adversarial training both supports their original results and also is an important contribution in and of itself.
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.
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.
- 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can the authors give a clear definition of linear stability? It seems to be a very important definition in the paper.

2. 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.

3. 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.
soundness: 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.

2. The authors identify 2 phases of switching from SGD to SAM. With the theoretical analysis, the authors give 4 properties of SAM.
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
weaknesses: 1. Some definitions are not clear.  
     (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?

    (ii)  What is $a$ in Proposition 4.1, which is used to give an upper bound of $\rho$?

2. 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?

3. 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?

4. 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.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
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.
soundness: 3
strengths: While the empirical observations originate elsewhere, the theoretical analysis appears to be novel and interesting. 
While 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. 
The 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.
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.
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. 
2. 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?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 2
questions: ### Major comments:
- Please address the first three comments in the Weaknesses section.

- I expected to see the loss going down in later steps in Figure 4 (b). Does the figure not show the entire steps?

- "(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?

### Minor comments:
- Are there any experiments using Eq. (4)?

- The paper uses the symbol $\approx$, but its mathematical definition is missing. Could the authors add the precise meaning?

- Proof of Proposition 4.1 could be improved by adding comments on which part the authors used the assumptions.

- Does this refer to Eq. (3)? If so, it should be clearly stated: "Our experiments are conducted using the original implementation of SAM"

---
**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.)
soundness: 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.
- The paper cites relevant references and clearly discuss the differences.
- The paper is well written and easy to follow, except for a few points (please see the Weaknesses section).
- The convergence result in Theorem 4.3 does not rely on convexity.
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)).
In practice, its approximation (Eq. (4)) is implemented for computationally efficiency.
It is not well understood whether this approximated method has the intended effect of finding a flatter solution, and this is what this paper studies.

First, 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.

Next, 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.
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}$.

- The experiments and the theory part use different methods (Eqs. (3) and (4)).

- I could not understand
  - how to obtain the inequality in lines 1137-1138, and
  - the phrase "which implies that $\theta_{T+1} = 0$" in lines 1143-1144.
  - the equality in lines 1085-1086 and the inequality in lines 1087-1088.

- This paper only studies the squared loss.

---
### Relatively minor comments
- 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.

- 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.

- 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.

- 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.

Meta-review: The paper studies the SAM algorithm and analyzes its behavior when applied late in training. It offers valuable insights into the two-phase process where SAM first escapes sharp minima and subsequently converges to flatter ones. The reviewers appreciated the theoretical analysis and empirical results showing that a few epochs of SAM at the end of training can achieve comparable performance to full SAM.

The consensus is that the paper makes a clear contribution towards understanding SAM’s effectiveness. I recommend acceptance, given that the authors implement the changes discussed during the rebuttal and discussion period.

Decision: Accept (Spotlight)