UNDERSTANDING SHARPNESS-AWARE MINIMIZATION
Year: 2021

Abstract:
Sharpness-Aware Minimization (SAM) is a recent training method that relies on worst-case weight perturbations. SAM significantly improves generalization in various settings, however, existing justifications for its success do not seem conclusive. First, we analyze the implicit bias of SAM over diagonal linear networks, and prove that it always chooses a solution that enjoys better generalisation properties than standard gradient descent. We also provide a convergence proof of SAM for non-convex objectives when used with stochastic gradients and empirically discuss the convergence and generalization behavior of SAM for deep networks. Next, we discuss why SAM can be helpful in the noisy label setting where we first show that it can help to improve generalization even for linear classifiers. Then we discuss a gradient reweighting interpretation of SAM and show a further beneficial effect of combining SAM with a robust loss. Finally, we draw parallels between overfitting observed in learning with noisy labels and in adversarial training where SAM also improves generalization. This connection suggests that, more generally, techniques from the noisy label literature can be useful to improve robust generalization.

heading: 1 INTRODUCTION
text: Understanding generalization of overparametrized deep neural networks is a central topic of the current machine learning research. Their training objective has many global optima where the training data are perfectly fitted (Zhang et al., 2016), but different global optima lead to dramatically different generalization performance (Liu et al., 2019). However, it has been observed that stochastic gradient descent (SGD) tends to converge to well-generalizing solutions, even without any explicit regularization methods (Zhang et al., 2016). This suggests that the leading role is played by the implicit bias of the optimization algorithms used (Neyshabur et al., 2014): when the training objective is minimized using a particular algorithm and initialization method, it converges to a specific solution with favorable generalization properties. However, even though SGD has a very beneficial implicit bias, significant overfitting can still occur, particularly in the presence of label noise (Nakkiran et al., 2019) and adversarial perturbations (Rice et al., 2020).
Recently it has been observed that the sharpness of the training loss, i.e., how quickly it changes in some neighborhood around the parameters of the model, correlates well with the generalization error (Keskar et al., 2016; Jiang et al., 2019), and generalization bounds related to the sharpness have been derived (Dziugaite & Roy, 2018). The idea of minimizing the sharpness to improve generalization has motivated recent works of Foret et al. (2021) and Wu et al. (2020) which propose to use worst-case perturbations of the weights on every iteration of training in order to improve generalization. We refer to this method as Sharpness-Aware Minimization (SAM) and focus mainly on the version proposed in Foret et al. (2021) that performs only one step of gradient ascent to approximately solve the weight perturbation problem before updating the weights.
Despite the fact that SAM significantly improves generalization in various settings, the existing justifications based on the generalization bounds provided by Foret et al. (2021) and Wu et al. (2020) do not seem conclusive. The main reason is that their generalization bounds do not distinguish the worst-case robustness to worst-case weight perturbation from average-case robustness to Gaussian noise. However the latter does not sufficiently improve generalization as both Foret et al. (2021) and Wu et al. (2020) report. Moreover, their analysis does not distinguish whether the weight perturbation is applied based on some or on all training examples which, as we will discuss, has an important impact on generalization.
We aim at further investigating the reasons for SAM’s success and make the following contributions:
• We first discuss that the SAM formulation can lead to two different objectives. We show that the gradient flows on these objectives over diagonal linear networks are implicitly biased towards solutions which enjoy better generalization properties compared to the standard gradient flow.
• We provide convergence results for the SAM algorithm used with stochastic gradients for non-convex objectives and discuss the convergence and generalization behavior of SAM for deep networks.
• We discuss why SAM can prevent overfitting in the noisy label setting by interpreting SAM as a gradient reweighting scheme and showing the further benefits of combining it with a robust loss.
• Finally, we draw parallels between overfitting in learning with noisy labels and in adversarial training where SAM also improves generalization. This connection suggests that techniques from the noisy label literature can be more generally useful to improve robust generalization.

heading: 2 RELATED WORK
text: Here we discuss relevant works on robustness in the weight space and cover the main references on overfitting in the noisy label setting and in adversarial training, two settings where weight-space robustness also improves generalization.
Weight-space robustness. Works on weight-space robustness of neural networks date back at least to 1990s (Murray & Edwards, 1993; Hochreiter & Schmidhuber, 1995). Random perturbations of the weights are used extensively in deep learning (Jim et al., 1996; Graves et al., 2013), and most prominently in approaches such as dropout (Srivastava et al., 2014). Many practitioners have observed that using SGD with larger batches for training leads to worse generalization (LeCun et al., 2012), and Keskar et al. (2016) have shown that this degradation of performance is correlated with the sharpness of the found parameters. This observation has motivated many further works which focus on closing the generalization gap between small-batch and large-batch SGD (Wen et al., 2018; Haruki et al., 2019; Lin et al., 2020). More recently, Jiang et al. (2019) have shown a strong correlation between the sharpness and the generalization error on a large set of models under a variety of different settings hyperparameters, beyond the batch size. This has motivated the idea of minimizing the sharpness during training to improve standard generalization, leading to Sharpness-Aware Minimization (SAM) (Foret et al., 2021). SAM modifies SGD such that on every iteration of training, the gradient is taken not at the current iterate but rather at a worst-case point in its vicinity. Zheng et al. (2021) concurrently propose a very similar weight perturbation method which also successfully improves standard generalization on multiple deep learning benchmarks. Wu et al. (2020) have also proposed a similar algorithm with the same motivation, although Wu et al. (2020) focuses on improving robust generalization of adversarial training.
Overfitting in learning with noisy labels. Arpit et al. (2017) have observed that deep networks noticeably overfit in the presence of mislabeled samples but early stopping can mitigate the problem. There are multiple other approaches that can mitigate this overfitting behavior: removing (Song et al., 2020) or downweighting noisy points (Jiang et al., 2018; Huang et al., 2020), using robust losses (Ghosh et al., 2017; Zhang & Sabuncu, 2018; Menon et al., 2019). Most of these approaches explicitly or implicitly leverage the early-learning phenomenon (Liu et al., 2020) where the network tends to fit first correctly labeled points and noisy points later in training. In the context of adversarial robustness, Sanyal et al. (2020) have discussed that fitting label noise can aggravate adversarial vulnerability, although removing label noise is not sufficient to achieve adversarial robustness.
Overfitting in adversarial training. Adversarial training in deep learning has been formulated as a robust optimization problem by Madry et al. (2018) where the worst-case perturbations are typically found via projected gradient descent. Rice et al. (2020) have described the robust overfitting phenomenon in adversarial training suggesting that early stopping is highly beneficial and provides a competitive baseline. However, they do not provide explanations about the reasons behind the overfitting phenomenon. The state-of-the-art approaches propose different ways to improve robust generalization and mitigate robust overfitting such as using additional training data (Carmon et al., 2019), weight averaging (Gowal et al., 2020; Chen et al., 2021), or advanced data augmentations (Rebuffi et al., 2021). Additionally, Wu et al. (2020) have shown that combining adversarial training with a method similar to sharpness-aware minimization also improves robust generalization.

heading: 3 IMPLICIT BIAS AND CONVERGENCE OF SAM
text: In this section, we discuss different objectives related to Sharpness-Aware Minimization (SAM) and their generalization properties, then derive convergence results for the stochastic SAM algorithm, and further discuss the convergence and generalization of SAM for deep networks.

heading: 3.1 IMPLICIT BIAS OF SAM
text: SAM objectives. Let {xi, yi}ni=1 be the training data points and `i(w) be the loss of a classifier parametrized by weights w ∈ R|w| and evaluated at point (xi, yi). Foret et al. (2021) theoretically base their approach on the following objective which we denote as MaxSum SAM:
MaxSum SAM: min w∈R|w| max ‖δ‖2≤ρ
1
n n∑ i=1 `i(w + δ). (1)
They justify this objective via a PAC-Bayesian generalization bound, although they show empirically (see Fig. 3 therein) that the following objective, denoted as SumMax SAM, leads to better generalization:
SumMax SAM: min w∈R|w|
1
n n∑ i=1 max ‖δ‖2≤ρ `i(w + δ). (2)
The update rule of SAM for these objectives amounts then to a variation of gradient descent with step size γt where the gradients are taken at intermediate points wit+1/2, i.e., wt+1 = wt − γtn ∑n i=1∇`i(wit+1/2). The two objectives, however, differ in how the points w i t+1/2 are computed since they approximately maximize different losses with inner step sizes ρt:
MaxSum: wit+1/2 = wt + ρt n n∑ j=1 ∇`j(wt) vs. SumMax: wit+1/2 = wt + ρt∇`i(wt). (3)
We next show formally that SumMax SAM has a better implicit bias than ERM and MaxSum SAM.
Implicit bias of SumMax and MaxSum SAM. The implicit bias of gradient methods is well understood for linear models where all gradient-based algorithms enjoy the same implicit bias. For diagonal linear neural networks, first-order algorithms have a richer implicit bias. We consider here a sparse regression problem and a predictor parametrized as fu,v = u u−v v. Before understanding how SAM induces a preferable bias, we first recall the seminal result of Woodworth et al. (2020): assuming global convergence, the solution selected by the gradient flow initialised at α ∈ Rd and denoted βα∞ solves the constrained optimisation problem:
βα∞ = arg min β∈Rd s.t. Xβ=y φα(β). (4)
The potential φα whose precise expression is given Eq. (12) in App. A interpolates nicely between the `1 and the `2 norms according to the initialization scale α: Small initialisations lead to low `1-type solutions which are known to induce good generalisation properties. Large initialisations lead to low `2-type solutions. Our main result is that both MaxSum and SumMax dynamics bias the flow towards solutions which still minimise the potential φα. However it does so with effective parameter αMaxSum and αSumMax which are strictly smaller than α for suitable inner step size ρ. Thus the chosen solution has better sparsity-inducing properties than the solution of the gradient flow.
Proposition 1 (Informal). Assuming global convergence, the solutions selected by the MaxSum and the SumMax algorithms defined Eq. (3), taken in the infinitesimally-small-γt limit and initialised at α, solve the optimisation problem (4) with effective parameters αSumMax and αSumMax which satisfy:
αSumMax = αe −ρ∆SumMax+O(ρ2) and αMaxSum = αe−ρ∆MaxSum+O(ρ 2),
where ∆SumMax and ∆MaxSum are two entrywise positive vectors for which typically: ‖∆SumMax‖1 ≈ d ∫ ∞
0
L(w(s))ds and ‖∆MaxSum‖1 ≈ d
n ∫ ∞ 0 L(w(s))ds.
The results are formally stated in Proposition 4 and 5 in App. A. The SumMax implementations has better bias properties since its effective scale of α is considerably smaller than the one of MaxSum. It is worth noting that the vectors ∆SumMax and ∆MaxSum are linked with the integral of the loss function along the flow. Thereby, the speed of convergence of the training loss impacts the magnitude of the biasing effect: the slower the convergence, the better the bias, similarly to what observed for SGD by Pesme et al. (2021).
Empirical evidence for the implicit bias. We compare the performance of different methods (ERM, MaxSum and SumMax SAM) on the training and test losses in Fig. 1. As predicted, the methods enjoy various generalization abilities: ERM and MaxSum SAM enjoy the same performance whereas SumMax benefits from a better implicit bias. We also note that the training loss of all the variants is converging to zero but as alluded before the convergence of SumMax SAM is slower. We show a similar experiments in App. A with stochastic variants of the algorithms. As expected, their performances are better than their deterministic counterparts (Keskar et al., 2016; Pesme et al., 2021).
The results on the implicit bias presented above require that the algorithm converges to zero training error. Therefore we analyze next the convergence of the SAM algorithm for general nonconvex functions in the practically relevant stochastic case. Note that we cannot expect rates as fast as the one observed for the previous simple model. They were due to its special structure for which fast convergence rates have been proven (Yun et al., 2021).

heading: 3.2 CONVERGENCE OF SAM ALGORITHM
text: To make the SAM algorithm practical, Foret et al. (2021) propose to combine SAM with stochastic gradients. By denoting the batch indices at time t by It, this leads to the following update rule:
wt+1 = wt − γt |It| ∑ i∈It ∇`i ( wt + ρt |It| ∑ j∈It ∇`j(wt) ) . (5)
Importantly, the same batch It is used for the inner and outer gradient steps which allows to interpolate between the MaxSum (for |It| = n) and SumMax (for |It| = 1) update rules defined in Eq. (3). However, using batch size |It| = 1 is inefficient in practice since it does not fully utilize the computational accelerators such as GPUs so for most experiments, Foret et al. (2021) use a higher |It| to balance the generalization improvement with the computational efficiency. In the following, we analyze the convergence of the update rule from Eq. (5). We make the following assumptions on the training loss L(w) = 1n ∑n i=1 `i(w):
(A1) (Bounded variance). It exists σ ≥ 0 s.t. E[‖∇`i(w)−∇L(w)‖2] ≤ σ2, i ∼ U(J1, nK), w ∈ Rd. (A2) (Individual β-smoothness). It exists β > 0 s.t. ‖∇`i(w) − ∇`i(v)‖ ≤ β‖w − v‖ for all w, v ∈ Rd and i ∈ J1, nK. (A3) (Polyak-Lojasiewicz). It exists µ > 0 s.t. 12‖∇L(w)‖ 2 ≥ µ(L(w)− L∗) for all w, v ∈ Rd.
We have the following convergence result for the stochastic SAM algorithm (see proof in App B.2). Proposition 2. Assume (A1-2) for the iterates (5). For any T ≥ 0 and for step-sizes γt = 1√Tβ and ρt =
1 T 1/4β , we have:
1 T E [ T−1∑ t=0 ‖∇L(wt)‖2 ] ≤ 4 β √ T (L(w0)− L∗) + 8σ2 b √ T ,
In addition, under (A3), with step-sizes γt = min{ 8t+43µ(t+1)2 , 1 2β } and ρt = √ γt/β:
E[L(wT )]− L∗ ≤ 3β2(L(w0)− L∗)
µ2T 2 +
22βσ2
µ2bT
We make several remarks regarding this theorem:
• We recover the rates of SGD with the usual condition on the step size γt (Ghadimi & Lan, 2013; Karimi et al., 2016). The ascent step-size ρt has to be O( √ γt) to ensure convergence.
Therefore it tolerates a slower decrease than γt. This finding is aligned with the fact that the ascent step-size of SAM should not be decreased as drastically as the descent step-size when training neural networks.
• The proof relies on the bound 〈∇L(wt + η∇L(wt)),∇L(wt)〉 ≥ (1− ηβ)‖∇L(wt)‖2 which shows that SAM-step is well aligned with the gradient step (seem Lemma 16 in App B.2).
• We can obtain the same convergence result for the different variants of SAM described in the previous section. We also provide in App B.1 tighter convergence rates for the full-batch case.
SAM with normalized gradient. When SAM algorithm was first introduced, the gradient was normalized in the inner step, in order to have an ascent-step of size exactly ρ. This normalization is not needed to get similar performance in applications and is not favorable for optimization purposes. Indeed when used with constant step-size normalized SAM is oscillating, and the averaged iterate converges to a different point which is not necessary flatter than the solution. When decreasing step-sizes are used, then the behaviors are comparable to the ones in the unnormalized case. In App. B.3, we show experiments that suggest that SAM with and without normalization steps achieves similar improvements in terms of generalization. However, for all other experiments in the paper, we rely on the normalized SAM as introduced in Foret et al. (2021) which is the main focus of our study.

heading: 3.3 CONVERGENCE AND GENERALIZATION OF SAM FOR DEEP NETWORKS
text: Here we relate the theoretical results from the previous section with the behavior of SAM for deep networks on a standard image classification dataset.
ERM and SAM both converge but generalize differently. We compare the behavior of ERM and SAM by training a ResNet-18 on CIFAR-10 for 1000 epochs using mini-batch SGD with momentum and piece-wise constant learning rates (see App. C for experimental details) and plot the results over epochs in Fig. 2. We observe that not only the ERM model but also the model trained with SAM fits all the training points and converges to a nearly zero training loss: 0.0012 for ERM vs 0.0009 for SAM. However, the SAM model has significantly better generalization performance: 3.76% vs 5.03% test error. Moreover, we observe no noticeable overfitting on standard CIFAR-10: the best and last model differ at most by 0.1% test error for both methods.
When is implicit bias of SAM beneficial during training? To better understand the implicit bias of SAM over training iterations, we perform the following experiment: we train with ERM for 900 epochs and then use SAM for the remaining 100 epochs (ERM→ SAM), and vice versa (SAM → ERM). We can see from Fig. 2 (right) that for SAM→ ERM once SAM converges to a well-generalizing minimum thanks to its implicit bias, then it is not important whether we continue optimization with SAM or with ERM.
In particular, we do not observe overfitting when switching to ERM. We note that it is still important here that SAM converges since SAM has to minimize the training loss below some threshold after which ERM can be potentially applied.
At the same time, for ERM→ SAM we observe a different behavior: the test error clearly improves when switching from ERM to SAM. This suggests that SAM (using a higher ρ than the standard value, see App. C) can help to escape a suboptimal minimum where ERM converges to. This phenomenon is interesting since it suggests that such a fine-tuning scheme can save computations as we can potentially start from any pre-trained model and improve its generalization on the same dataset which is different from fine-tuning with SAM in the transfer learning setup covered in Foret et al. (2021).

heading: 4 WHY CAN SAM HELP AGAINST OVERFITTING OF NOISY LABELS?
text: In this section, we discuss another important aspect of SAM: its beneficial effect for the noisy label setting. We first analyze the development of the gradient direction over training and then interpret SAM as a gradient reweighting schemes where the gradients are reweighted according to the sharpness of the loss.
Overfitting on noisy labels and the effect of SAM. By the noisy label setting we refer to a situation where some fraction of the training labels is changed to uniformly random labels and kept fixed throughout the training.1 This setting is challenging since standard training with ERM leads to overfitting as illustrated in Fig. 3a. In particular, it is commonly observed (Liu et al., 2020) that at the beginning of the training ERM fits the correctly labeled training points but, after some point, starts to also fit the incorrectly labeled points, which coincides with overfitting in terms of the test error. Thereby, the noisy label setting is different from the standard setting from Sec. 3.3: training to zero loss is harmful and leads to overfitting, and thus early stopping is needed (Li et al., 2020b). Fig. 3a illustrates that early stopping indeed leads to a significantly better test error for the ERM model compared to the model taken at the last epoch. At the same time, SAM noticeably improves generalization over ERM, although later in training SAM also starts to fit the noisy points and thus also requires early stopping, either explicitly via a (potentially noisy) validation set or implicitly via restricting the number of training epochs as in Foret et al. (2021).
Benefits of following the gradient direction given by SAM. Consider a batch of points Sall := {(x, y)}ni=1 where some labels y are noisy. We define Sclean and Snoisy to be the subsets of Sall with clean and noisy labels respectively. We also define Strue to be the same set of points as Sclean but where all the labels are correct. We denote by LS(w) = 1|S| ∑ (x,y)∈S `x,y(w) the loss on a set of points S where w denotes the model parameters. The gradient update then can be decomposed as:
∇LSall(w) = |Sclean| |Sall| ∇LSclean(w) + |Snoisy| |Sall| ∇LSnoisy (w).
Here we want to analyze the contribution of each part of the gradient for the ERM and SAM weight updates. For this, in Fig. 3 (b) and (c), we compare the following directions over training epochs of an ERM model: ∇LSall(w) vs. ∇LStrue(w) and ∇LSall(w) vs. ∇LSnoisy (w) for the current point w (i.e., the ERM update) and the point w + SAM (i.e., the SAM update). To be robust to the label noise, we would like the update direction of an algorithm to be maximally aligned with the ideal direction∇LStrue(w) and minimally aligned with the noisy direction∇LSnoisy (w). The key observation of Fig. 3b and Fig. 3c is that SAM is changing the direction of the gradient by making it less aligned with the noisy direction at the later stage of training when the model starts to fit the label noise. Therefore, the SAM update is less dominated by the noisy labels and can resist overfitting longer. To better understand this effect, we discuss next the behavior of SAM when training a linear classifier in a binary classification setting.
SAM improves generalization even for linear models. Since using SumMax SAM (Eq. (2)) leads to better generalization than MaxSum SAM, we consider this technique for an overparametrized
1Thus, it is different from the setting where new label noise is added on each iteration (HaoChen et al., 2020).
linear model fx(w) = 〈w, x〉 trained with the cross-entropy loss ` using gradient descent. The training objective of SumMax SAM then becomes:
min w∈Rd
1
n n∑ i=1 max ‖δ‖2≤ρ ` (yi 〈w + δ, xi〉) = min w∈Rd 1 n n∑ i=1 ` (yi 〈w, xi〉 − ρ ‖xi‖2) . (6)
We consider the following simple dataset: x ∼ N (yµ, σ2Id), where y ∈ {−1, 1} is the label, and µ is some unit vector. We use σ = 0.1, input dimension d = 100, n = 50 training samples, and we randomize 90% of the training labels. We select the best step sizes for ERM and SAM via a grid search. We can see from Fig. 4 that SAM with early stopping consistently (over ten random seeds) improves the test error over ERM. This illustrates that benefits of SAM extend beyond non-convex learning tasks with multiple minima and suggests that the linear setting can be a good starting point for understanding the benefit of SAM.
SAM as a gradient reweighting scheme. For a general non-linear binary classifier fx(w), on every iteration of gradient descent, the update direction is given by the gradient on all points:
∇LSall(w) = 1
n n∑ i=1 `′(yi · fxi(w))︸ ︷︷ ︸ loss derivative · ∇fxi(w)︸ ︷︷ ︸ direction . (7)
When we perturb the weights with SAM, we change in general both the loss derivative `′(y · fx(w)) and the direction given by the classifier’s gradient ∇fx(w). However, for a linear classifier fx(w), only `′(y ·fx(w)) is changed while the direction∇fx(w) = x is independent of the weights. The fact that SAM helps even in this simple setting suggests that its crucial effect can lie only in modifying the loss derivative `′(y · fx(w)) which leads to beneficial reweighting of the per-example gradients of noisy and clean inputs. We note that gradient reweighting is also the idea behind robust losses (Ghosh et al., 2017; Zhang & Sabuncu, 2018) as they also affect only the loss derivative while keeping the direction∇fxi(w) the same. This suggests the following interpretation of SumMax SAM:
Gradient reweighting interpretation of SAM: SumMax SAM reweights per-example gradients amplifying the derivative of clean points more than the derivative of noisy points so that the overall gradient direction∇LSall(w) is more aligned with∇LSclean(w).
We observe from Eq. (6) that SAM offsets the margin yi 〈w, xi〉 of each example by ρ ‖xi‖2 More generally, using SumMax SAM for non-linear models guarantees that the offset is non-negative for each example as maximization of a monotonic loss ` implies minimization of the margin. Note that this would not be the case for MaxSum SAM or average-case perturbations of the weights. We discuss next why offsetting the loss can be valuable to reweight the gradients of different examples.
We first note that the early-learning phenomenon which is proven to hold under some assumptions in Liu et al. (2020) implies that the correctly labeled points are fitted first during training so that at some early stage they tend to have a higher margin than the noisy points. We illustrate this in Fig. 5 for a linear model. We then consider the following two losses: cross-entropy loss (CE) and generalized cross-entropy loss (GCE) (Zhang & Sabuncu, 2018) (we hide the dependency of f on w for brevity)
`CE(y · fx) = log(1 + exp(−y · fx)), `GCE(y · fx) = 1/q ( 1− (1 + exp(−y · fx))−q ) .
We plot these losses in Fig. 6a (using q = 0.7 for GCE following Zhang & Sabuncu (2018)) together with the derivative change from adding an offset c in Fig. 6b and 6c. We observe based on Fig. 6b that the offset affects more significantly the points with small positive margins (i.e., mostly correctly labeled points due to the early-learning phenomenon) whose derivative is amplified more making it closer to its minimal value of -1. Fig. 6c suggests that GCE has a better effect than CE since GCE does not only decrease the derivative of points with small positive margin but it also increases the derivative of the points with small negative margin (i.e., mostly noisy points) making it closer to zero. This suggests that combining GCE with SAM can lead to further benefits which is indeed confirmed experimentally in App. D.
Does the loss derivative change explain the behavior of a non-linear model? To further support our gradient reweighting interpretation for non-linear models, we show that changing only the loss derivative `′(yi · fxi(wSAM )) (see Eq. (7)) and discarding the change of the direction leads to the main benefit in SAM. For this, we train a ResNet-18 on a two-class CIFAR-10 under 80% label noise with early stopping and study different ways to modify the gradient ∇LS(w) used on each iteration of training. We report results in Table 1 where we observe that SAM DIRECTION leads only to a small improvement compared to ERM (13.47% to 11.30%) while SAM DERIVATIVE is helping most leading to 7.45% test error. This gives more evidence that the main improvement from using SAM comes from changing the derivative of the loss `′(yi · fxi(wSAM )) and not the direction ∇fxi(wSAM ) which is coherent with our gradient reweighting interpretation of SAM. We further report extra experimental details and the plots over epochs in App. D.

heading: 5 WHY CAN SAM HELP AGAINST ROBUST OVERFITTING?
text: Here we draw parallels between the overfitting observed in the noisy label setting and in robust training (Rice et al., 2020) where SAM also improves generalization (Wu et al., 2020).
The effect of SAM is similar for robust overfitting. In Fig. 7, we plot the error under adversarial perturbations obtained via the PGD attack (we use 10 iterations, see details in App. D) for an `∞ adversarially trained model (Madry et al., 2018) with = 8/255 trained with ERM and SAM for 1000
epochs. For this experiment, we use the SAM algorithm as in Foret et al. (2021) which also leads to similar improvements of robust error as the method introduced in Wu et al. (2020) which relies on a different scaling of the layerwise perturbation bounds. We observe that, when trained for sufficiently long (contrary to Wu et al. (2020) that only train for 200 epochs), SAM also leads to overfitting, showing high similarity to the noisy label setting (Fig. 3a).
Implicit label noise in adversarial training. Related to the point above, a recent work of Dong et al. (2021) argues about the presence of inputs which are “noisy” for adversarial training, and whose removal mitigates robust overfitting. They further discuss that such inputs become ambiguous for a human observer after adding `∞ adversarial perturbations. The problem stems from the formulation of adversarial examples which is widely used: max‖δ‖∞≤ `(x+ δ) instead of the label-preserving formulation max‖δ‖∞≤ , y(x)=y(x+δ) `(x+ δ) which is obviously impractical since we cannot query the true label y(x+ δ) at an arbitrary point during training. Thus, the commonly used adversarial perturbations can change the effective label y(x + δ) leading to implicit label noise. This can be confirmed by implementing labelpreserving adversarial training for a linear model where we know the ground truth model. We can use the same two Gaussian example as in Fig. 4 but without explicit label noise and perform adversarial training in the subspace orthogonal to the discriminative direction µ which ensures that the label of x+ δ is not affected. We show experimentally in Fig. 8 that (1) robust overfitting can be observed also in simple linear models and (2) using the label-preserving formulation of adversarial training prevents it. This suggests that robust overfitting is strongly connected to the labels of adversarial examples.
Mitigating implicit label noise by SAM and other
methods. Performing label-preserving adversarial training requires the knowledge of the ground truth labeling function which is impractical. Thus, in practice one has to rely on algorithms that prevent fitting the implicit label noise in other ways. We hypothesize that an early-learning phenomenon similar to the one observed in training with noisy labels (Liu et al., 2020) should also hold for adversarial training, and our gradient reweighting explanation of SAM from Sec. 4 should also extend to this setting. We also note that other solutions against robust overfitting that have been recently proposed in the literature—such as early stopping (Rice et al., 2020), weight averaging (Gowal et al., 2020; Rebuffi et al., 2021), bootstrap (Chen et al., 2021), data removal (Dong et al., 2021)—also rely on the early-learning phenomenon and have been previously explored in the noisy label literature. We confirm this further on GCE (Zhang & Sabuncu, 2018) and semi-supervised pairing terms (Luo et al., 2019) in App. D. This further suggests that both settings are closely related, and improvements in one setting can also translate to the other one which can explain the improvement from SAM against robust overfitting.

heading: 6 CONCLUSIONS AND OUTLOOK
text: We have discussed two different objectives motivated by SAM and analyzed their implicit bias for a diagonal linear network. Then we have provided a convergence proof for the variant of SAM used in practice and confirmed empirically that the best generalization performance is achieved at a solution with nearly zero training loss on a standard image classification tasks. At the same time, in the noisy label setting, it is harmful to achieve low training loss even with SAM since this would involve fitting the noisy labels. We have discussed why SAM can be helpful in this case: we observed that SAM amplifies the contribution of the correctly labeled points compared to the noisy ones in the gradient update. Finally, we argue that the overfitting observed in noisy label setting shares a lot of similarities with overfitting observed in adversarial training which can explain why SAM is also helpful there.

heading: REPRODUCIBILITY STATEMENT
text: Theoretical results. We provide proofs of our theoretical results in the Appendix in Sec. A for the implicit bias and in Sec. B for the convergence of different variants of SAM. The assumptions A1-A3 for the convergence proofs are mentioned at the beginning of Sec. 3.2. An additional assumption A2’ is used only for the statements in the appendix and thus is introduced only there.
Empirical results. To facilitate reproducibility of our experiments, in Sec. C we specify all the hyperparameters of our experiments on deep networks and linear models. Moreover, in Sec. D, we provide more details about some of the experiments from the main part.

heading: Appendix

heading: ORGANIZATION OF THE APPENDIX
text: The appendix contains proofs related to sharpness and convergence of different variants of SAM that we omitted from the main part of the paper due to the space constraints. Moreover, the appendix also contains additional implementation details and extended experimental results. The appendix is organized as follows:
• Sec. A: proofs related to the implicit bias of SumMax and MaxSum SAM. • Sec. B: proofs related to the convergence of different variants of SAM. • Sec. C: experimental details for the epxeriments with deep networks and linear models. • Sec. D: additional experiments about the importance of the change in the loss derivative
for deep networks, improvement of standard generalization under label noise by adversarial training, and the experiment with the GCE loss and semi-supervised pairing term that mitigate robust overfitting.

heading: A THEORETICAL ANALYSIS OF THE IMPLICIT BIAS
text: To understand why SAM is generalizing better than ERM we consider the simpler problem of noiseless regression with 2-layer diagonal linear network for which we can precisely characterize the implicit bias of the optimizations algorithms in hand.
Overparametrised noiseless regression. We consider linear regression with inputs, outputs (x1, yi, . . . , xn, yn) ∈ (Rd+1)n. We consider the overparametrised case and assume that there exists at least one interpolating parameter β∗ ∈ Rd which fits the training set: yi = 〈β∗, xi〉 for all 1 ≤ i ≤ n. We study the quadratic loss for which the training loss is written as:
L(β) := 1
4n n∑ i=1 (〈β, xi〉 − yi)2 = 1 4n n∑ i=1 〈β − β∗, xi〉2.
2-layer diagonal linear network. The simplest parametrisation of β is to consider β = w which corresponds to the classical least-squares framework. In this case, first order methods (GD, with and without momentum) converge towards the same solution: they have the same implicit bias. The same holds for the variants of SAM we consider and the linear case is not helping us to distinguish them and to show their primacy on ERM. To theoretically resolve this question, we study a simple non-linear parametrisation: the quadratic parametrisation β = w2+ − w2− with w = [w+, w−]> ∈ R2d. The training loss in w is therefore written:
L(w) := 1
4n n∑ i=1 (〈w2+ − w2−, xi〉 − yi)2 1 4n n∑ i=1 (〈β, xi〉 − yi)2 with β = w2+ − w2−. (8)
This parametrisation can be viewed as a simplified linear network of depth 2. Additionally to being a toy neural model, this model has received recently a lot of attention since it is possible to precisely characterize the implicit bias of descent algorithm on it.
Optimization algorithms. We consider minimizing the loss L(w) using three different algorithms:
• Gradient descent with infinitesimally small stepsize (in the gradient flow limit): ẇt = −∇L(wt). (9)
• The MaxSum SAM algorithm with infinitesimally small outer-stepsize and inner step size ρ ≥ 0 :
ẇt = −∇L(wt + ρ∇L(wt)). (10) • The SumMax SAM with infinitesimally small outer-stepsize and inner step size ρ ≥ 0:
ẇt = − 1
n n∑ i=1 ∇`i(wt + ρ∇`i(wt)) (11)
Implicit bias of the Gradient flow. Let us define the function φα for α ∈ Rd as
φα(β) = d∑ i=1 α2i q(βi/α 2 i ) where q(z) = ∫ z 0 arcsinh(u/2)du = 2− √ 4 + z2 + z arcsinh(z/2).
(12) Note that if we have α ∈ R then we consider φα = φα1. We will make great uses of this function to precisely characterize the implicit bias of the different algorithms we consider. Indeed, following Woodworth et al. (2020) we can show the following result for the gradient flow dynamics (9)
Proposition 3 (Theorem 1 of Woodworth et al. (2020)). For any 0 < α <∞, if the solution β∞ of the gradient flow (9) started from w0 = α1 for the squared parameter problem in Eq.(8) satisfies Xβ∞ = y, then
β∞ = arg min β∈Rd φα(β) s.t. Xβ = y, (13)
where φα is defined in Eq.(12).
It worth noting that the implicit regularizer φα interpolates between the `1 and `2 norms (see Woodworth et al., 2020, Theorem 2). Therefore an implicit bias depending on the scale of the initialization is observed. The algorithm, started from α1, converges to the minimum `1-norm interpolator for α small and to the minimum `2-norm interpolator for α large. The proof follows from (a) the KKT condition for the optimization problem (13): ∇φα(w) = X>ν for a Lagrange multiplier ν and (b) the closed form solution obtained by integrating the gradient flow, w = b(X>ν) for some function b and some vector ν. Identifying∇φα(w) = b−1(w) leads to the solution. Considering the same proof technique, we now derive the implicit bias of both variants of SAM algorithms.
Implicit bias of MaxSum SAM algorithm. We characterize the implicit bias of the MaxSum dynamics (10) in the following proposition. Recall that the function φα is defined in Eq. (12).
Proposition 4. For any 0 < α < ∞, if the solution β∞ of the gradient flow (10) started from w0 = α1 for the squared parameter problem in Eq.(8) satisfies Xβ∞ = y, then
β∞ = arg min β φαMaxSum(β) s.t. Xβ = y,
where αMaxSum = αe − ρ 4n2
∫∞ 0 (X>rs) 2ds+O(ρ2).
We note that for ρ small enough αMaxSum < α. The scale of the vector 1n2 ∫∞ 0 (X>rs)
2ds (which influences the bias effect) is related to the loss integral dn ∫∞ 0 L(w(s))ds since ‖rs‖2 = nL(w(s)) (see intuition in Eq.(16)). Thereby the speed of convergence of the loss controls the magnitude of the biasing effect. However in the case of MaxSum SAM, as explained below, this effect is typically negligible because of the extra prefactor dn and this implementation behaves similarly as ERM (as shown in the experiments in Sec. 3.1).
Proof. We follow the proof technique of Woodworth et al. (2020). We first derive the equation satisfies by the flow
ẇ(t) = − 1 2n X̃>r+(t) ◦ w+(t)
= − 1 2n X̃>r+(t) ◦ (w(t) + ρ 2n X̃>r(t) ◦ w(t)),
where w+(t) = w(t) + ρ∇L(w(t) is the intermediate SAM step, r(t) = X̃w(t)2 − y and r+(t) = X̃w+(t)
2 − y are the residual of w(t) and w+(t) and we denote by X̃ = [X −X]. We can directly integrate this ODE to obtain a closed form expression for w(t):
w(t) = w(0) ◦ exp(− 1 2n X̃> ∫ t 0 r+(s)ds) ◦ exp(− ρ 4n2 ∫ t 0 (X̃>r+(s)) ◦ (X̃>r(s))ds).
Using that the flow is initialized at w(0) = α1 and the definition of β(t) yields to
β(t) = u(t)2 − v(t)2
= α2 exp(− 1 n X> ∫ t 0 r+(s)ds) ◦ exp(− ρ 2n2 ∫ t 0 (X>r+(s)) ◦ (X>r(s))ds)
− α2 exp( 1 n X> ∫ t 0 r+(s)ds) ◦ exp(− ρ 2n2 ∫ t 0 (X>r+(s)) ◦ (X>r(s))ds)
= 2α2 exp(− ρ 2n2 ∫ t 0 (X>r+(s)) ◦ (X>r(s))ds) sinh(− 1 n X> ∫ t 0 r+(s)ds).
Recall we are assuming that β∞ is global minimum of the loss, i.e., Xβ∞ = y. Thus
Xβ∞ = y
β∞ = bαMaxSum(X >ν), with and bα(z) = 2α2 ◦ sinh(z), ν = − 1n ∫∞ 0 r+(s)ds and
αMaxSum = α exp(− ρ
4n2 ∫ ∞ 0 (X>r+(s)) ◦ (X>r(s))ds). (14)
Denoting by
∇φα(β) = b−1α arcsinh( 1
2α2 ◦ β), and integrating this equation, we obtain the expression of φα = ∑d i=1 α 2 i q(βi/α
2 i ) where q(z) =∫ z
0 arcsinh(u/2)du = 2−
√ 4 + z2 + z arcsinh(z/2).
Therefore β∞ satisfies the KKT conditions for the minimum norm interpolator problem arg minβ φα(β) s.t. Xβ = y, i.e.
Xβ∞ = y
∇φα(β∞) = X>ν. This proves the first part of the result. Now using the definition of r+(s) we obtain
r+(t) = X̃w+(t) 2 − y
= X̃(w(t) + ρX̃>r(t) ◦ w(t))2 − y = r(t) + 2ρX̃(X̃>r(t)) ◦ w(t)2 + ρ2X̃(X̃>r(t))2 ◦ w(t)2
= r(t) + 2ρX(X>r(t)) ◦ (u(t)2 + v(t)2) + ρ2X(X>r(t))2 ◦ β(t).
And
X>r+(t) = X >r(t) +O(ρ),
which concludes the second part of the result.
Implicit bias of SumMax SAM algorithm We characterize similarly the implicit bias of the SumMax dynamics (11) in the following proposition. Recall that the function φα is defined in Eq. (12). Proposition 5. For any 0 < α < ∞, if the solution β∞ of the gradient flow (11) started from w0 = α1 for the squared parameter problem in Eq.(8) satisfies Xβ∞ = y, then
β∞ = arg min β φαSumMax(β) s.t. Xβ = y,
where αSumMax = αe− ρ 4n
∫∞ 0 ∑n i=1 x 2 i (x > i β(s)−yi) 2ds+O(ρ2).
In addition, if we assume that there exist R,B ≥ 0 such that the features are bounded ‖x‖ ≤ R almost surely and that the trajectory of the flow is bounded ||β(t)‖ ≤ B for all t ≥ 0, we have for all ρ ≤ 1
R2 √ B(B+‖β∗‖) :
αSumMax ≤ α
We note that for ρ smaller than an explicit constant depending on the problem quantities, we can assert that αMaxSum < α. Furthermore the scale of the vector 1n ∫∞ 0 ∑n i=1 x 2 i (x > i β(s) − yi)2ds (which
influences the bias effect) is related to the loss integral d ∫∞
0 L(w(s))ds, through ‖rs‖2 = nL(w(s))
(see intuition Eq. (17)). Thus the speed of convergence of the loss controls the magnitude of the biasing effect. This effect is typically important in the case of SumMax SAM (see experiments Sec.3.1).
Proof. The proof follows the same lines as the proof of Proposition 4 Let us denote by x̃i = [xi−xi]. We then have that the dynamics of the flow (11) satisfies
ẇ(t) = − 1 2n n∑ i=1 x̃iri,+(t) ◦ w+(t)
= − 1 2n n∑ i=1 x̃iri,+(t)(1 + ρ 2 x̃iri(t)) ◦ w(t),
where w+(t) is the intermediate SAM step, ri(t) = x̃>i w(t) 2 − yi and ri,+(t) = x̃iw+(t)2 − yi are the residual of w(t) and w+(t) at the observation (xi, yi). Integrating we obtain:
w(t) = w(0) ◦ exp(− 1 n X̃> ∫ t 0 r+(s)ds) ◦ exp(− ρ 2n n∑ i=1 x̃2i ∫ t 0 ri,+(s)ri(s)ds).
The remainder of the proof is similar to the one of Proposition 4 and we directly obtain that
αSumMax = α exp(− ρ
4n n∑ i=1 x̃2i ∫ t 0 ri,+(s)ri(s)ds). (15)
Using the definition of ri,+(t) we have
ri,+(t) = x̃ > i w+(t) 2 − yi = x̃>i w(t)
2 ◦ (1 + ρri,tx̃i)2 − yi = x̃>i w(t)
2 ◦ (1 + 2ρri,t>xi + ρr2i,t>x2i )− yi = ri,t + 2ρri,t(u(t) 2 + v(t)2)>x2i + ρ 2r2i,tβ(t) >x3i .
And therefore
x2i ri,+(t)ri(t) = x 2 i ri(t) 2[1 + 2ρ(u(t)2 + v(t)2)>x2i + ρ 2ri(t)β(t) >x3i ])
So if ρ is such that 1 + ρ2ri(t)β(t)>x3i ≥ 0 then αSumMax < α. Using Cauchy-Schwarz inequality |ri(t)β(t)>x3i | ≤ ||xi||(||β|| + ||β∗||)||β||||x3i || and Holder inequality ||x3i || ≤ ||xi||3. Thus |ri(t)β(t)>x3i | ≤ ||xi||4(||β|| + ||β∗||)||β||. And we have the result for ρ ≤ R−2(B(B + ‖β∗‖))−1/2.
Comparison between SumMax and MaxSum implicit bias and connection with the loss integral We wish to compare the two leading terms IMaxSum(t) = 1n2 (X >r(t))2 = 1n2 ( ∑n i=1 xiri) 2
and ISumMax(t) = 1n ∑n i=1 x 2 i ri(t)
2 which comprise the first order in the expansion of the α and relate them to the loss value at w(t).
We first note using Cauchy-Schwarz inequality that ISumMax(t) ≥ IMaxSum(t) but we aim at obtaining a more quantitative result, even though the derivation will be informal. Comparing their `1-norms amounts to compare the empirical fourth moment ‖IMaxSum(t)‖1 = (w − w∗)> 1n ∑n i=1 ‖xi‖2xix>i (w − w∗) and second moment squared ‖ISumMax(t)‖1 = (w − w∗) >[ 1n ∑n i=1 xix > i ]
2(w − w∗). We can compare the typical operator norm of these two random matrices. Following the Bai-Yin’s law, the operator norm of a Wishart matrix is with high probability ‖ 1n ∑n i=1 xix > i ‖ ≈ dn and that
with high-probability, the squared norm of a Gaussian vector is ‖xi‖ ≈ d. Therefore we obtain that
‖ 1 n n∑ i=1 ‖xi‖2xix>i ‖ ≈ d‖ 1 n n∑ i=1 xix > i ‖ ≈ d2 n
‖[ 1 n n∑ i=1 xix > i ] 2‖ ≤ ‖[ 1 n n∑ i=1 xix > i ]‖2 ≈ ( d n )2 Therefore in the overparametrized regime (d >>> n) we typically have that ISumMax(t)IMaxSum(t) ≈ n and the biasing effect of the SumMax implementation would tend to be O(n) time better than the biasing effect of the MaxSum implementation.
However this first insight only enables to compare IMaxSum(t) and ISumMax(t). It is not informative on the intrinsic biasing effect of MaxSum and SumMax SAM. With this aim, we would like to relate the quantities IMaxSum(t) and ISumMax(t) to the loss function evaluated in w(t). With high probability (still using the concentration of a Gaussian vector) we have that
‖IMaxSum(t)‖1 = (w(t)− w∗)> 1
n n∑ i=1 ‖xi‖2xix>i (w(t)− w∗)
≈ d(w(t)− w∗)> 1
n n∑ i=1 xix > i (w(t)− w∗)
= dL(w(t)). (16)
And using the concentration of Wishart matrices, i.e., 1d [XX >] ≈ I for large dimension d, we also have
‖ISumMax(t)‖1 = 1
n2 (w(t)− w∗)>X>XX>X(w(t)− w∗)
d
n2 (w(t)− w∗)>X>
1 d [XX>]X(w(t)− w∗)
≈ d n (w(t)− w∗)> 1 n [X>X](w(t)− w∗)
= d
n L(w(t)). (17)
These approximations provide some intuitions on why the biasing effect of both SAM implementation can be related to the integral of the loss. We let a formal derivation of these results as future work.
Additional experiments on implicit bias of stochastic methods. We provide an additional experiment to investigate the performance of stochastic implementation of the ERM, MaxSum and SumMax SAM. As explained by Pesme et al. (2021), we observe in Fig.9 that the stochastic implementations enjoy a better implicit bias than their deterministic counterparts. We note that the fact that small batch versions generalize better than full batch version is commonly observed in practice for deep networks Keskar et al. (2016). We let the characterization of the implicit bias of these stochastic implementations as future works.

heading: B CONVERGENCE OF SAM ALGORITHM
text: In this section we prove the convergence of the SAM algorithm to a stationary point of the function L or to a minimum when the function is convex.

heading: B.1 CONVERGENCE OF DETERMINISTIC SAM ALGORITHM
text: Here we aim to understand if the one-step approximation of SAM makes sense from an optimization perspective, i.e., whether the algorithm proposed in (Foret et al., 2021) can successfully minimize the
training objective. This is an important aspect since if an optimization algorithm cannot sufficiently minimize the training loss, the resulting classifier typically will not lead to good generalization performance, even if it has some useful implicit bias (Neyshabur et al., 2014). We first consider the following first-order optimization algorithm on the objective L parameterized by w:
wt+1 = wt − γt∇L(wt + ρt∇L(wt)), (18)
where ρt and γt are the step sizes of the inner and outer steps, respectively. This update rule is reminiscent of the extra-gradient algorithm (Korpelevich, 1977) but with an ascent in the inner step, i.e., the gradient is evaluated not at the current iterate but at the point which maximises locally the function. Moreover, this update rule can also be seen as a realization of the general extrapolated gradient descent framework suggested in Lin et al. (2020). However, taking an ascent step for extrapolation is not discussed there, and the convergence properties of the update rule from equation 18, to the best of our knowledge, have not been proven.
Summary of the convergence results. Let us first recall the definition of smoothness.
(A2’) (β-smoothness). It exists β>0 such that ‖∇L(w)−∇L(v)‖ ≤ β‖w − v‖ for all w, v ∈ Rd.
When the function L is β-smooth, convergence to stationary points can be obtained. Proposition 6. Assume (A2’). For any γ < 1/β and ρ < 1/β, the iterates equation 18 satisfy for all T ≥ 0:
1
T T−1∑ t=0 ‖∇L(wt)‖2 ≤ 2 γ(1− ρβ)T (L(w0)− L∗),
If, in addition, the function L satisfies (A3), then:
L(wT )− L∗ ≤ ( 1− γ(1− ρβ)µ
2
)T (L(w0)− L∗).
We can make the following remarks:
• We recover the rates of gradient descent but with constants increasing with the ascent step-size ρ.
• The condition ρ < 1/β is necessary since the point w + 1/β∇L(w) can be a local maximum of L. Such w would be a fixed point of the algorithm without being a stationary point of L.
• The proof crucially relies on the bound 〈∇L(wt+ρ∇L(wt)),∇L(wt)〉 ≥ (1−ρβ)‖∇L(wt)‖2 which shows that SAM-step is well aligned with the gradient step (see Lemma 7) and on a descent inequality similar to the classical one for gradient descent (see Lemma 8).
• For non-convex functions, full details are provided in Proposition 9. When the function satifies an addition Polyak-Lojasiewicz inequality, details are in Proposition 10.
• For convex functions, 〈∇L(wt + ρ∇L(wt)),∇L(wt)〉 ≥ ‖∇L(wt)‖2 and convergence holds for any step-size ρ given that γρ is small enough. Details are provided in Proposition 11 .
Auxillary Lemmas. The following lemma shows that the SAM update is well correlated with the gradient∇L(w) and will be cornerstone to our proof. Lemma 7. Let L be a differentiable function and w ∈ Rd. We have the following bound for any ρ ≥ 0:
〈∇L(w + ρ∇L(w)),∇L(w)〉 ≥ (1 + αρ)‖∇L(w)‖2 where α =  −β if L is β-smooth, 0 if L is convex µ if L is µ-strongly convex.
Proof. We simply add and subtract a term ‖∇L(w)‖2 in order to make use of classical inequalities bounding 〈∇L(w1) − ∇L(w2), w1 − w2〉 by ‖w1 − w2‖2 for smooth or convex functions and w1, w2 ∈ Rd.
〈∇L(w + ρ∇L(w)),∇L(w)〉 = 〈∇L(w + ρ∇L(w))−∇L(w),∇L(w)〉) + ‖∇L(w)‖2
= 1/ρ〈∇L(w + ρ∇L(w))−∇L(w), ρ∇L(w)〉+ ‖∇L(w)‖2
≥ (1 + αρ)‖∇L(w)‖2,
where the last inequality is using that
〈∇L(w1)−∇L(w2), w1 − w2〉 ≥ α‖w2 − w1‖2, where α =  −β if L is β-smooth, 0 if L is convex µ if L is µ-strongly convex.
The next lemma shows that the decrease of function values of the SAM algorithm defined in Eq.(18) can be controlled similarly as in the case of gradient descent Nesterov (2004).
Lemma 8. Assume (A2’). For any γ ≤ 1/β, the iterates (18) satisfies for all t ≥ 0:
L(wt+1) ≤ L(wt)− γ(1− ρβ) (
1− γβ 2
(1− ρβ) ) ‖∇L(wt)‖2.
If in addition the function L satisfies (A3) with potentially µ = 0, then for all γ, ρ ≥ 0 such that γβ(2− ρβ) ≤ 2, we have
L(wt+1) ≤ L(wt)− γ (
1− γβ 2
+ ρµ ( 1− γβ − γρβ 2
2
)) ‖∇L(wt)‖2.
We note that the constraints on the step-size are different depending on the assumptions on the function L. In the non-convex case, ρ has to be smaller than 1/β, whereas in the convex case, it has to be smaller than 2/β.
Proof. Let us define by wt+1/2 = wt + ρ∇L(wt) the SAM ascent step. Using the smoothness of the function L (Assumption (A2’)) (see, e.g., Nesterov, 2004), we obtain
L(wt+1) ≤ L(wt)− γ〈∇L(wt+1/2),∇L(wt)〉+ γ2β
2 ‖∇L(wt+1/2)‖2.
The main trick is to use the binomial squares
‖∇L(wt+1/2)‖2 = −‖∇L(wt)‖2 + ‖∇L(wt+1/2)−∇L(wt)‖2 + 2〈∇L(wt+1/2),∇L(wt)〉,
to bound
L(wt+1) ≤ L(wt)− γ〈∇L(wt+1/2),∇L(wt)〉+ γ2β
2 ‖∇L(wt+1/2)‖2
= L(wt)− γ2β
2 ‖∇L(wt)‖2 +
γ2β
2 ‖∇L(wt+1/2)−∇L(wt)‖2 − γ(1− γβ)〈∇L(wt+1/2),∇L(wt)〉
≤ L(wt)− γ[1− ρβ − γβ
2 (1− ρβ)2]‖∇L(wt)‖2,
where we have used Lemma 7 and that ‖∇L(wt+1/2) − ∇L(wt)‖2 ≤ β2‖wt+1/2 − wt‖2 ≤ β2ρ2‖∇L(wt)‖2. If in addition the function L is convex then we can use its co-coercivity (see, e.g., Nesterov, 2004) to bound ‖∇L(wt+1/2)−∇L(wt)‖2 ≤ β〈∇L(wt+1/2)−∇L(wt), wt+1/2 −wt〉 and obtain a tighter bound:
L(wt+1) ≤ L(wt)− γ〈∇L(wt+1/2),∇L(wt)〉+ γ2β
2 ‖∇L(wt+1/2)‖2
= L(wt)− γ2β
2 ‖∇L(wt)‖2 +
γ2β
2 ‖∇L(wt+1/2)−∇L(wt)‖2 − γ(1− γβ)〈∇L(wt+1/2),∇L(wt)〉
≤ L(wt)− γ(1− γβ
2 )‖∇L(wt)‖2 − γ(1− γβ −
γρβ2
2 )〈∇L(wt+1/2)−∇L(wt),∇L(wt)〉
≤ L(wt)− γ(1− γβ 2 + ρµ(1− γβ − γρβ
2
2 ))‖∇L(wt)‖2,
where we have used Lemma 7.
Convergence proofs. Using the previous Lemma 8 recursively, we can bound the average gradient value of the iterates (18) of SAM algorithm and ensure convergence to stationary points. Proposition 9. Assume (A2’). For any γ < 1/β and ρ < 1/β, the iterates equation 18 satisfies for all T ≥ 0:
1
T T∑ t=0 ‖∇L(wt)‖2 ≤ L(w0)− L(wT ) Tγ(1− ρβ)[1− γβ2 (1− ρβ)] .
Proof. Using the Lemma 8 we obtain γ(1− ρβ) (
1− γβ 2
(1− ρβ) ) ‖∇L(wt)‖2 ≤ L(wt)− L(wt+1).
And summing these inequalities for t = 0, · · · , T − 1 yields
1
T T−1∑ t=0 ‖∇L(wt)‖2 ≤ L(w0)− L(wT ) Tγ(1− ρβ)[1− γβ2 (1− ρβ)] .
When the function L additionally satisfies a Polyak-Lojasiewicz condition, i.e., the Assumption (A3), linear convergence of the function value to the minimum function value can be obtained. This is the object of the following proposition: Proposition 10. Assume (A2’-3). For any γ < 1/β and ρ < 1/β, the iterates (18) satisfies for all T ≥ 0:
L(wt)− L∗ ≤ ( 1− 2γµ(1− ρβ) (
1− γβ 2
(1− ρβ) ))t (L(w0)− L∗).
Proof. Using the Lemma 8 and that the function L is µ Polyak-Lojasiewicz (Assumption (A3)) we obtain
L(wt+1) ≤ L(wt)− 2µγ(1− ρL) (
1− γβ 2
(1− ρL) ) (L(wt)− L∗).
And subtracting the optimal value L∗ we get L(wt)− L∗ ≤ ( 1− 2γµ(1− ρβ) (
1− γβ 2
(1− ρβ) )) (L(wt−1)− L∗)
≤ ( 1− 2γµ(1− ρβ) (
1− γβ 2
(1− ρβ) ))t (L(w0)− L∗).
When the function L is convex, convergence of the average of the iterates can be proved. Proposition 11. Assume (A2’) and L convex. For any step-sizes γ and ρ such that γβ(1 + ρβ) < 2, then the averaged w̄T = 1T ∑T−1 t=0 wt of the iterates (18) satisfies for all T ≥ 0:
L(w̄T )− L∗ ≤ 2ρβ + 1
γ(2− γβ(1 + ρβ))T ‖w0 − w∗‖2,
If, in addition, the function L is µ-strongly convex, then: ‖wT − w∗‖2 ≤ ( 1− γµ(2− γβ(1 + ρβ)) )T (2ρ+ 1)‖w0 − w∗‖2.
The proof is using a different astute Lyapunov function (which works for the non-strongly convex case).
Proof. Let us define by Vt = [L(wt)− L(w∗)] + 12ρ‖wt − w∗‖ 2 and by wt+1/2 = wt + ρ∇L(wt) the SAM ascent step.
Vt+1 − Vt ≤ − γ
ρ 〈∇L(wt+1/2), wt − w∗〉 − γ〈∇L(wt+1/2),∇L(wt)〉+
γ2 2ρ (1 + ρβ)‖∇L(wt+1/2)‖2
= −γ ρ 〈∇L(wt+1/2), wt + ρ∇L(wt)− w∗〉+
γ2 2ρ (1 + ρβ)‖∇L(wt+1/2)‖2
= −γ ρ 〈∇L(wt+1/2), wt+1/2 − w∗〉+
γ2 2ρ (1 + ρβ)‖∇L(wt+1/2)‖2
≤ −γ ρ (1− γβ 2 (1 + ρβ))〈∇L(wt+1/2), wt+1/2 − w∗〉.
If L is convex then L(wt+1/2)− L(w∗) ≤ 〈∇L(wt+1/2), wt+1/2 − w∗〉 and therefore we obtain γ
ρ (1− γβ 2 (1 + ρβ))L(wt+1/2)− L(w∗) ≤ Vt − Vt+1.
Using the definition of wt+1/2 we always have that L(wt+1/2) ≥ L(wt) + ρ‖∇L(wt)‖2 therefore γ
ρ (1− γβ 2 (1 + ρβ))L(wt)− L(w∗) ≤ Vt − Vt+1.
And taking the sum and using Jensen inequality we finally obtain:
L(1/T T∑ t=0 wt)− L(w∗) ≤ V0 − VT+1 T γρ (1− γβ 2 (1 + ρβ)) .
If L is µ-strongly convex, we use that 〈∇L(wt+1/2), wt+1/2 − w∗〉 ≥ µ‖wt+1/2 − w∗‖2 to obtain
‖wt+1/2 − w∗‖2 = ‖wt + ρ∇L(wt)− w∗‖2 = ‖wt − w∗‖2 + 2ρ〈∇L(wt), wt − w∗〉+ ρ2‖∇L(wt)‖2
≥ ‖wt − w∗‖2 + 2ρ〈∇L(wt), wt − w∗〉 ≥ ‖wt − w∗‖2 + 2ρ[L(wt)− L(w∗)] ≥ 2ρVt.
Therefore we have
Vt+1 ≤ (1− γµ(2− γβ(1 + ρβ))Vt ≤ (1− γµ(2− γβ(1 + ρβ))t+1V0.

heading: B.2 CONVERGENCE OF STOCHASTIC IMPLEMENTATIONS OF SAM ALGORITHM

heading: B.2.1 CONVERGENCE OF MAXSUM GRADIENT ALGORITHM
text: When the SAM algorithm is implemented with the MaxSum objective as optimization objective, two different batches are using in the ascent and descent steps. We obtain the MaxSum algorithm defined as
wt+1 = wt − γt b ∑ i∈It ∇`i ( wt + ρt b ∑ i∈Jt ∇`i(wt) ) , (19)
where It and Jt are two different mini-batches of data of size b. For this variant of the algorithm we obtain a similar convergence result that the one obtain with the same batch for the two steps. Proposition 12. Assume (A1), (A2’) for the iterates (19). For any T ≥ 0 and for step-sizes γt = 1√Tβ and ρt = 1T 1/4β , we have:
1 T E T−1∑ t=0 ‖∇L(wt)‖2] ≤ 4 β √ T (L(w0)− L∗) + 8σ2 b √ T ,
In addition, under (A2), with step-sizes γt = min{ 8t+43µ(t+1)2 , 1 2β } and ρt = √ γt/β:
E[L(wT )]− L∗ ≤ 3β2(L(w0)− L∗)
µ2T 2 +
22βσ2
bµ2T
We obtain the same convergence result that in Proposition 2, but under the relaxed smoothness assumption (A2’).
As in the deterministic case, the proof relies on two lemmas which shows that the SAM update is well correlated with the gradient and that the decrease of function values can be controlled.
Auxilliary Lemmas. The following lemma shows that the SAM update is well correlated with the gradient ∇L(wt). Let us denote by ∇Lt+1(w) = 1b ∑ i∈It ∇`i(w), ∇Lt+1/2(w) = 1 b ∑ i∈Jt ∇`i(w) and wt+1/2 = wt + ρ∇Lt+1/2(wt) the SAM ascent step.
Lemma 13. Assume (A1) and (A2). Then for all ρ ≥ 0, t ≥ 0 and w ∈ Rd,
E〈∇Lt+1(w + ρ∇Lt+1/2(w)),∇L(w)〉 ≥ (1/2− βρ)‖∇L(w)‖2 − β2ρ2σ2
2 .
The proof is similar to the proof of Lemma 7. Only the stochasticity of the noisy gradients has to be taken into account. In this aim, we consider instead the update which would have been obtained without noise, and bound the remainder using the bounded variance assumption (A1).
Proof. Let us denote by ŵ = w + ρ∇L(w), the true gradient step. We first add and subtract ∇Lt+1/2(ŵ)
〈∇Lt+1(w + ρ∇Lt+1/2(w)),∇L(w)〉 = 〈∇Lt+1(w + ρ∇Lt+1/2(w))−∇Lt+1(ŵ),∇L(w)〉 − 〈∇Lt+1(ŵ),∇L(w)〉.
We bound the two terms separately. We use the smoothness of L (Assumption (A2’)) to bound the first term:
−E〈∇Lt+1(w + ρ∇Lt+1/2(w))−∇Lt+1(ŵ),∇L(w)〉 = −E〈∇L(w + ρ∇Lt+1/2(w))−∇L(ŵ),∇L(w)〉
≤ 1 2 E ‖∇L(w + ρ∇Lt+1/2(w))−∇L(ŵ)‖2 + 1 2 ‖∇L(w)‖2 ≤ β 2
2 E ‖w + ρ∇Lt+1/2(w)− ŵ‖2 +
1 2 ‖∇L(w)‖2
≤ β 2ρ2
2 E ‖∇Lt+1/2(w)−∇L(w)‖2 +
1 2 ‖∇L(wt)‖2
≤ β 2ρ2σ2
2b +
1 2 ‖∇L(w)‖2,
where we have used that the variance of a mini-batch of size b is bounded by σ2/b. Note that this term can be equivalently bounded by βρσ/ √ b‖∇L(w)‖ if needed. For the second term, we directly apply Lemma 7 to obtain
E〈∇Lt+1(ŵ),∇L(w)〉 = E〈∇L(ŵ),∇L(w)〉 ≥ (1− βρ)‖∇L(w)‖2.
The next lemma shows that the decrease of function values of SN-SAM algorithm can be controlled similarly as in the case of stochastic gradient descent.
Lemma 14. Let us assume (A1, A2’) then for all γ ≤ 1/(2β) and ρ ≤ 1/(2β), the iterates equation 19 satisfies
EL(wt+1) ≤ EL(wt)− γ/4E ‖∇L(wt)‖2 + γβσ2(γ + ρ2beta).
This lemma is analogous to Lemma 8 in the stochastic case. The proof is very similar, with the slight difference that Lemma 13 is used instead of Lemma 7.
Proof. Let us define by wt+1/2 = wt + ρ∇Lt+1/2(wt). Using the smoothness of the function L (A2), we obtain
L(wt+1) ≤ L(wt)− γ〈∇Lt+1(wt+1/2),∇L(wt)〉+ γ2β
2 ‖∇Lt+1(wt+1/2)‖2.
Taking the expectation and using that the variance is bounded (A1) yields to
EL(wt+1) ≤ EL(wt)− γ E〈∇L(wt+1/2),∇L(wt)〉+ γ2β
2 E ‖∇Lt+1(wt+1/2)‖2
≤ EL(wt)− γ E〈∇L(wt+1/2),∇L(wt)〉+ γ2β E ‖∇Lt+1(wt+1/2)−∇L(wt+1/2)‖2 + γ2β E ‖∇L(wt+1/2)‖2
≤ EL(wt)− γ E〈∇L(wt+1/2),∇L(wt)〉+ γ2βσ2/b+ γ2β E ‖∇L(wt+1/2)‖2.
The main trick is still to use the binomial squares
‖∇L(wt+1/2)‖2 = −‖∇L(wt)‖2 + ‖∇L(wt+1/2)−∇L(wt)‖2 + 2〈∇L(wt+1/2),∇L(wt)〉
to bound
EL(wt+1) ≤ EL(wt)− γ E〈∇L(wt+1/2),∇L(wt)〉+ γ2β
2 E ‖∇L(wt+1/2)‖2 + γ2σ2β/b
= EL(wt)− γ2LE ‖∇L(wt)‖2 + γ2β E ‖∇L(wt+1/2)−∇L(wt)‖2
− γ(1− 2γβ)E〈∇L(wt+1/2),∇L(wt)〉+ γ2σ2β/b = EL(wt)− γ2β E ‖∇L(wt)‖2 + γ2L3 E ‖wt+1/2 − wt‖2
− γ(1− 2γβ)(1/2 + αρ)E ‖∇L(wt)‖2 + γ(1− 2γL)σ2ρ2β2/2 + γ2σ2β/b = EL(wt)− γ2β E ‖∇L(wt)‖2 + γ2β3ρ2 E ‖∇Lt+1/2(wt)‖2
− γ(1− 2γβ)(1/2 + αρ)E ‖∇L(wt)‖2 + γ(1− 2γβ)σ2/bρ2β2/2 + γ2σ2β/b = EL(wt)− γ2β E ‖∇L(wt)‖2 + 2γ2β3ρ2 E ‖∇L(wt)‖2 + 2γ2β3ρ2σ2/b − γ(1− 2γβ)(1/2 + αρ)E ‖∇L(wt)‖2 + γ(1− 2γβ)σ2ρ2β2/2 + γ2σ2β/b ≤ L(wt)− γ
2 [1− 2ρβ(1− 2γβ(1− ρβ))]E ‖∇L(wt)‖2 + γσ2β/b[γ + ρ2L/2(1 + 2γβ)]
where we have used Lemma 13 and that ‖∇L(wt+1/2)−∇L(wt)‖2 ≤ β2‖wt+1/2 − wt‖2.
Using Lemma 14 we directly obtain the following convergence result.
Proposition 15. Assume (A1) and (A2’). For γ ≤ 1/(2β) and ρ ≤ 1/(2β), the iterates equation 5 satisfies:
1
T T−1∑ t=0 E ‖∇L(wt)‖2 ≤ 4 L(w0)− EL(wT ) Tγ + 4Tσ2β(γ + ρ2β)/b.
This proposition gives the first part of Proposition 12. The proof of the stronger result obtained when the function is in addition PL (Assumption (A3)) is similar to the proof of Theorem 3.2 of Gower et al. (2019), only the constants are changing.

heading: B.2.2 CONVERGENCE OF SUMMAX GRADIENT ALGORITHM
text: When the SAM algorithm is implemented to minimize the SumMax objective, the same batch is used in the ascent and descent steps. We obtain then iterates (5) for which we have stated the convergence result in Proposition 2. The proof follows the same lines as before with the minor different that we are assuming the individual gradient∇ft are Lipschitz (Assumption (A2)) to control the alignment of the expected SAM direction. Let us denote by∇Lt(w) = 1b ∑ i∈Jt ∇`i(w).
Lemma 16. Assume (A1-2). Then we have for all w ∈ Rd, ρ ≥ 0 and t ≥ 0
E〈∇Lt(w + ρ∇Lt(w)),∇L(w)〉 ≥ (1/2− ρβ)‖∇L(w)‖2 − β2ρ2σ2
2b .
The proof is very similar to the proof of Lemma 13. The only difference is that the Assumption (A2) is used instead of (A2’).
Proof. Let us denote by ŵ = w+ ρ∇L(w), the true gradient step. We first add and subtract∇Lt(ŵ)
〈∇Lt(w + ρ∇Lt(w)),∇L(w)〉 = 〈∇Lt(w + ρ∇Lt(w))−∇Lt(ŵ),∇L(w)〉 − 〈∇Lt(ŵ),∇L(w)〉.
We bound the two terms separately. We use the smoothness of Lt to bound the first term (Assumption (A2)):
−〈∇Lt(w + ρ∇Lt(w))−∇Lt(ŵ),∇L(w)〉 ≤ 1
2 ‖∇Lt(w + ρ∇Lt(w))−∇Lt(ŵ)‖2 +
1 2 ‖∇L(w)‖2
≤ β 2
2 E ‖w + ρ∇Lt(w)− ŵ‖2 +
1 2 ‖∇L(w)‖2
≤ β 2ρ2
2 ‖∇Lt(w)−∇L(w)‖2 +
1 2 ‖∇L(w)‖2.
And taking the expectation, we obtain:
−E〈∇Lt(w + ρ∇Lt(w))−∇Lt(ŵ),∇L(w)〉 ≤ β2ρ2σ2
2b +
1 2 E ‖∇L(w)‖2.
For the second term, we apply directly Lemma 7
E〈∇Lt(ŵ),∇L(wt)〉 = 〈∇L(ŵ),∇L(w)〉 ≥ (1− βρ)‖∇L(w)‖2.
Assembling the two inequalities yields to the result.
The next lemma shows that the decrease of function values of SAM algorithm can be controlled similarly as in the case of gradient descent. It is analogous to Lemma 8 where different batches are used in both the ascent and descent steps of SAM algorithm.
Lemma 17. Assume (A1-2). For all γ ≤ 1/β and ρ ≤ 1/(4β), the iterates (5) satisfies
EL(wt+1) ≤ EL(wt)− 3γ/8E ‖∇L(wt)‖2 + γβσ2/b(γ + 2ρ2β).
Proof. Let us define by wt+1/2 = wt + ρ∇Lt+1(wt). Using the smoothness of the function L which is implied by (A2), we obtain
L(wt+1) ≤ L(wt)− γ〈∇Lt+1(wt+1/2),∇L(wt)〉+ γ2β
2 ‖∇Lt+1(wt+1/2)‖2.
We still use the binomial squares
‖∇Lt+1(wt+1/2)‖2 = −‖∇L(wt)‖2 + ‖∇Lt+1(wt+1/2)−∇L(wt)‖2 + 2〈∇Lt+1(wt+1/2),∇L(wt)〉
and bound L(wt+1) by
L(wt+1) ≤ L(wt)− γ2β
2 ‖∇L(wt)‖2 +
γ2β
2 ‖∇Lt+1(wt+1/2)−∇L(wt)‖2 − γ(1− γβ)〈∇Lt+1(wt+1/2),∇L(wt)〉
≤ L(wt)− γ2β
2 ‖∇L(wt)‖2 + γ2β‖∇Lt+1(wt+1/2)−∇Lt+1(wt)‖2 + γ2β‖∇Lt+1(wt)−∇L(wt)‖2
− γ(1− γβ)〈∇Lt+1(wt+1/2),∇L(wt)〉
≤ L(wt)− γ2β
2 ‖∇L(wt)‖2 + γ2ββ2‖wt+1/2 − wt‖2 + γ2β‖∇Lt+1(wt)−∇L(wt)‖2
− γ(1− γβ)〈∇Lt+1(wt+1/2),∇L(wt)〉
= L(wt)− γ2β
2 ‖∇L(wt)‖2 + γ2β3ρ2‖∇Lt+1(wt)‖2 + γ2β‖∇Lt+1(wt)−∇L(wt)‖2
− γ(1− γβ)〈∇Lt+1(wt+1/2),∇L(wt)〉
= L(wt)− γ2β
2 (1− 4β2ρ2)‖∇L(wt)‖2 + γ2β(1 + 2β2ρ2)‖∇Lt+1(wt)−∇L(wt)‖2
− γ(1− γβ)〈∇Lt+1(wt+1/2),∇L(wt)〉
Taking the expectation and using Lemma 16 we obtain
EL(wt+1) ≤ EL(wt)− γ2β
2 (1− 4β2ρ2)E ‖∇L(wt)‖2 + γ2β(1 + 2β2ρ2)E ‖∇Lt+1(wt)−∇L(wt)‖2
− γ(1− γβ)E〈∇Lt+1(wt+1/2),∇L(wt)〉
≤ EL(wt)− γ2β
2 (1− 4β2ρ2)E ‖∇L(wt)‖2 + γ2β(1 + 2β2ρ2)σ2/b
− γ(1− γβ)(1/2− βρ)E ‖∇L(wt)‖2 + γ(1− γβ) ρ2σ2β2
2b
≤ EL(wt)− γ2β
2 (1− 4β2ρ2)E ‖∇L(wt)‖2 + γ2β(1 + 2β2ρ2)σ2/b
− γ 2 (1− 2βρ(1− γ(β − 2ρβ2)))E ‖∇L(wt)‖2 + γσ2/b[γβ + ρ2β2 2 (1 + 3γβ)].
Using Lemma 17 we directly obtain the following convergence result. Proposition 18. Assume (A1-2). For γ ≤ 1/β and ρ ≤ 1/4β, the iterates (5) satisfies:
1
T T−1∑ t=0 E ‖∇L(wt)‖2 ≤ 8/3 L(w0)− EL(wT ) Tγ + 8Tσ2β(γ + ρ2β)/(3b).
In addition, under (A3), with step-sizes γt = min{ 8t+43µ(t+1)2 , 1/(2β)} and ρt = √ γt/β:
E[L(wT )]− L∗ ≤ 3β2(L(w0)− L∗)
µ2T 2 +
22βσ2
µ2bT .
Proof. The first bound directly comes from Lemma 17. The second bound is similar to the proof of Theorem 3.2 of Gower et al. (2019), only the constants are changing.

heading: B.3 GRADIENT NORMALIZATION IN SAM
text: We consider here the implementation of SAM algorithm where the inner gradient step is normalized:
wt+1 = wt − γ∇L ( wt + ρ
∇L(wt) ‖∇L(wt)‖
) , (20)
We state here several claims showing that the normalized formulation in Eq.(20) can be less favorable than the unnormalized one of Eq.(18).
Fact I: Normalized SAM does not converge with constant step-size. The last iterate of SAM with constant step-size is never converging which can be easily seen on the function L(w) = 0.5w2 for w ∈ R. Indeed SAM iterates equation 20 become wt+1 = (1 − γ)wt − γρ sign(wt), and convergence implies that the limit w∞ would satisfy the fixed point w∞ = −ρ sign(w∞) which never holds for ρ > 0.
Fact II: Normalized SAM when averaged does not necessarily converge to a flatter minimum. Let us consider the following simple function
L(w) = { 1/2w2 if w ≥ 0 1/6w6 if w < 0.
When applied to this function with γ = ρ = 1, starting from w0 = −1, then it is experimentally observed that the averaged of the iterates (20) converge to a positive point w∞ = 0.0149 which is sharper than the origin, i.e, S1L(w∞) = 0.515 > S 1 L(0) = 0.5.
Fact III: Normalized SAM can stay stuck at non-stationary points of the function it optimizes. For any ρ, there exists a non-convex functionLρ and a pointwρ such that∇Lρ(wρ+ρ ∇Lρ(wρ)‖∇Lρ(wρ)‖ ) = 0 and ∇Lρ(wρ) 6= 0. To see that, just consider the function −0.5x2 and wρ = ρ.
Performance of the unnormalized SAM for deep networks. Given the potential shortcomings of the normalization step in SAM, it is interesting to compare the generalization properties of the normalized and unnormalized versions. For this, we repeat the experiments in Sec. 3.3 and Sec. 4 but for the unnormalized formulation in Fig. 10 with a grid search over ρ. We can observe that although the optimal ρ is not the same for both formulations (ρ = 0.3 vs ρ = 0.1 without label noise and ρ = 0.4 vs ρ = 0.2 with 60% label noise), the test error of the normalized vs unnormalized variants is very similar: 3.69% vs 3.70% without label noise, and 10.03% vs 10.29% with 60% label noise. Thus, given the same generalization performance, the unnormalized version may be preferred due to the lack of the shortcomings described above.

heading: C EXPERIMENTAL DETAILS
text: Experiments on deep networks. In all experiments, we train deep networks using SGD with step size 0.1 and momentum 0.9, and `2-regularization parameter λ = 0.0005. We use three datasets: CIFAR-10 Krizhevsky & Hinton (2009), CIFAR-100 Krizhevsky & Hinton (2009), and two-class CIFAR-10 where we select two random classes (horse and car). We train models with basic data augmentations: random image crops and mirroring. We use a pre-activation ResNet-18 (He et al., 2016) with a width factor 64 and piece-wise constant learning rates (with a 10-times decay at 50% and 75% epochs) for all experiments except Table 2 where we follow the setup of (Huang et al., 2020) and use ResNet-34 with cosine learning rates with one cycle. We train all models for 1000 epochs except the experiments in Tables 2, and Fig. 13 for which we use 200 epochs as they require more expensive training (due to adversarial training) or an expensive grid search over each dataset and label noise amount.
For all experiments involving SAM, we select the best ρ based on a grid search over ρ ∈ {0.025, 0.05, 0.1, 0.2, 0.3, 0.4}. In the fine-tuning ERM with SAM experiment, we used ρ = 0.2. We do not use the m-sharpness (Foret et al., 2021) since we were not able to reproduce exactly the results of the SAM paper with batch size of 256 and m = 32 (as suggested by Foret et al. (2021)) for the noisy label experiments. Instead, we observed that the same performance can be achieved using a smaller batch size 128 and m = 128, thus we opted for this setting which has an advantage of being significantly faster (thus, we could perform a detailed grid search for each setting) if only one GPU is available per run.
For model selection to determine the best epoch early stopping and the best ρ, we use 10% of examples from the noisy training set as the validation set following (Zhang & Sabuncu, 2018). We emphasize that we do not assume access to a clean validation set to perform early stopping.
For stochastic weight averaging (SWA), we use an exponential moving average with parameter τ = 0.999, and we update the moving average every iteration. Moreover, we start SWA from the very beginning of training (e.g., similarly to Rebuffi et al. (2021)) which is not the same setting as proposed in Izmailov et al. (2018) who start SWA from a specific epoch, use a different moving average coefficient, and a modified learning rate schedule. We deviate from the setting of Izmailov et al. (2018) since it has much more hyperparameters than the single moving average parameter τ , and those hyperparameters have to be chosen differently for the noisy label setting compared to the standard setting of Izmailov et al. (2018).
For experiments with the generalized cross-entropy loss (GCE) Zhang & Sabuncu (2018) we use q = 0.7 as recommended in Zhang & Sabuncu (2018) unless mentioned otherwise.
For experiments with `∞ adversarial training, we use = 8/255 as in Madry et al. (2018); Rice et al. (2020), 5 iterations of Projected Gradient Descent (PGD) for training and 10 iterations for testing (following Rice et al. (2020)), step size 2 /nsteps for training (as in Madry et al. (2018)) and step size /4 for evaluation (as in Croce & Hein (2020)). Note that we use PGD with 10 iterations to get an understanding of the qualitative behavior over epochs without using expensive computations. However, to strengthen the evaluation, we also use AutoAttack Croce & Hein (2020) to assess the robustness of the best model over training. The model selection is performed via the validation robust error obtained via 10-iteration PGD.
Experiments on a linear model. Here we provide details for the experiments shown in Fig. 4. We train a logistic regression model without regularization using full-batch gradient descent. We use 50 training examples which are sampled from two Gaussians in a 100-dimensional space. Each Gaussian has the standard deviation σ = 0.1, and the mean of each Gaussian µ is sampled randomly from the unit sphere. For optimization, we initialize the weight vector randomly on the sphere of radius 1/3 and use 50 iterations of gradient descent with a step size 2.0. We use ρ = 5.0 for SAM.
Computing infrastructure. We perform all our experiments with deep networks on NVIDIA V100 GPUs with 32GB of memory.

heading: D ADDITIONAL EXPERIMENTS
text: Differences between SAM and Stochastic Weight Averaging. Stochastic Weight Averaging (SWA) (Izmailov et al., 2018) is a weight averaging technique that improves generalization and, similarly to SAM, also motivated from the perspective of sharpness of the loss, thus it is natural to ask whether SAM shares some similarities with SWA. To study this, we plot the loss interpolations between the weights of SAM, ERM, and SWA models (wSAM , wERM , wSAM ) in Fig. 11 which suggest that SAM and ERM converge to different minima. This is further supported by the fact that ‖wSAM − wSWA‖2 is large, i.e., ≈ 1/3 of the distance of wSAM to its initialization. Interestingly, there is a peak between wSAM and wERM which is, however, not too high and the model at the peak still has non-trivial test error (around 40%). This is in contrast to wSWA which is located in the same valley as wERM with a very sharp increase close to werm (similarly to the interpolations from Izmailov et al. (2018)). We also note another difference between SAM and SWA in the context of the fine-tuning experiment from Fig. 2: fine-tuning an SWA model will lead to slight overfitting as it would converge to a standard ERM model (Izmailov et al., 2018) contrary to SAM which improves generalization.
Does the loss derivative change explain the behavior of a non-linear model? Here we provide more details on the experiment from Sec. 4 about the importance of changing the loss derivative `′(yi · fxi(w)) with SAM as opposed to the direction of the individual gradients∇fxi(w). We train a pre-activation ResNet-18 (He et al., 2016) with a width factor 16 (instead of the standard 64) on a two-class CIFAR-10 dataset (horse vs car classes) using batch size 256 and mSAM = 64 which refers to the m-sharpness defined in Foret et al. (2021)). We note that this experiment would be more difficult to perform for the multi-class setting due to the fact that `′(yi · fxi(wSAM )) will not be a scalar. In order to select ρ for each of the SAM-based methods, we performed a grid search over {0.001, 0.003, 0.01, 0.03, 0.1, 0.2}. We train each model three times and report the average test error (obtained with early stopping on the validation set) with the standard deviation.
In Fig. 12, we plot the test error of all the four methods over epochs. We can see that although the test error for SAM derivative is oscillating (unlike for other methods), the best test error over epochs (7.45% ± 0.28) is lower than for other methods and close to that of SAM (9.02% ± 0.34). We emphasize that all reported results are obtained via proper model selection over epochs using a validation set. To conclude, we can see that by discarding the direction of SAM the test errors start to oscillate but we still preserve the improved generalization if we use early stopping.
Beneficial effect of combining SAM with GCE. We check how predictive is our explanation about the advantage of using SAM with a robust loss such as GCE (Zhang & Sabuncu, 2018) instead of CE to better leverage the early-learning phenomenon. We validate this hypothesis experimentally using a ResNet-34 network trained on CIFAR-10 and CIFAR-100 datasets and present results in Table 2 for {40%, 60%} label noise (for the further details see App. C). We present results from the literature (Huang et al., 2020) where all of them use ResNet-34 except MentorNet which uses WRN-28-10. Early stopping based on a noisy validation set is done for each method except those reported from the literature. We can see that SAM indeed performs better with GCE than with CE which leads to competitive results. Moreover, the results can be further enhanced by using weight averaging (Izmailov et al., 2018) where it outperforms many recently proposed approaches including
Self-Adaptive Training (Huang et al., 2020). We also note that it can be further combined with other successful approaches against label noise such as MixUp (Zhang et al., 2017), bootstrapping (Reed et al., 2014), and semi-supervised learning approaches (Li et al., 2020a).
Mitigating robust overfitting. To illustrate how techniques from label noise can be transferred to adversarial training and improve upon robust overfitting, we benchmark here two approaches: (1) generalized cross-entropy (GCE) loss and (2) semi-supervised pairing terms (Laine & Aila, 2016; Luo et al., 2019). We train all methods for 200 epochs as in Rice et al. (2020). To the best of our knowledge, these approaches have not been used to prevent robust overfitting.
In Fig. 13a, we show the performance of adversarial training with the GCE loss (we use q = 0.9) instead of the standard cross-entropy loss. We observe that GCE is able to mitigate robust overfitting, however it does not improve the performance of the best model over epochs compared to the crossentropy loss, so its usefulness is limited.
In Fig. 13b, we show the results of adversarial training with an additional semi-supervised pairing term (Laine & Aila, 2016) that has been used to mitigate label noise in, e.g., Luo et al. (2019) (see also Li et al. (2020a) for a more advanced approach). Specifically, let f be the logits of the network, xadvi be the `∞ adversarial example generated for sample xi using the PGD attack and x̄i be a version of xi with a different augmentation. Then the additional term that we use in the objective is:
λ · 1 n n∑ i=1 ∥∥f(xadvi )− f(x̄i)∥∥22 , (21) which we use together with the cross-entropy on adversarial examples as in standard adversarial training (Madry et al., 2018). We use λ = 1.0 and a “warmup” scheme for the regularizer which consists in starting the regularizer after 50% of training epochs which is a standard practice (Laine & Aila, 2016; Luo et al., 2019). Both λ and the starting epoch were determined via a grid search. We show in Fig. 13b the robust error over epochs given by PGD with 10 iteration of adversarial training with and without the pairing term. We can observe that adding the pairing term mitigates the robust overfitting trend, and leads to improved robust error.
To make sure that the improvement is not due to gradient masking (Carlini et al., 2019), we evaluate the best model additionally with AutoAttack (Croce & Hein, 2020). The model trained with the pairing term achieves 51.15% robust error while the standard adversarial training leads to 52.56% robust error and adversarial training with SAM achieves 50.07% robust error. Thus, we conclude that the techniques from the label noise literature (beyond simple early stopping (Rice et al., 2020)) can be successfully transferred to improve upon robust overfitting. Moreover, we note that more advanced pairing terms are possible, e.g. involving temperature scaling of the logits (Berthelot et al., 2019), more advanced data augmentations (Sohn et al., 2020), and other improvements over the standard pairing (Li et al., 2020a). We leave applications of more advanced pairing techniques to improve upon robust overfitting as future work.
Finally, the improvement from the pairing term may also be connected to the fact that the TRADES objective (Zhang et al., 2019) often works better than standard adversarial training (see Gowal et al. (2020) for an exhaustive empirical study). The approach of TRADES also adds an additional pairing term that does not rely on the ground truth labels to the standard cross-entropy loss (although TRADES uses a different loss and does not use a different augmentation on x̄i). This suggests that even some of the existing successful techniques can be interpreted based on the connection between label noise and robust overfitting.

Review summaries:
The aim of the paper is to provide theoretical explanations for the recent successes of adversarial weight perturbation and sharpness-aware minimization (SAM) methods. In particular, these methods have been shown to significantly help robust generalization in deep learning and theoretical explanations for their success are missing so far.   The paper aims to fill this gap by getting a better understanding of SAM. It contributes a collection of miscellaneous results: * A proof that SAM provides better generalization for linear neural networks than SGD * The convergence of SAM is analyzed and generalization behavior is discussed * A new interpretation of SAM is presented in terms of gradient reweighing  * Finally, a connection of SAM to the noisy label literature is made

In this paper, the authors mainly analyzed a SAM-like algorithm with a square loss. For diagonal linear networks, this paper proves that the NNs trained by the algorithm has better generalizations than the ones trained by normal SGD. For deep NNs satisfies PL condition and other assumptions mentioned in section3, the authors can prove the algorithm will converge. The authors use a linear model to give the intuition why the SAM algorithm can generalize better on a dataset with label noise.

This is an ambitious paper that seeks to explain why SAM works the way it does. First, the authors try to justify the better generalization of SAM over SGD on deep networks. Next, the authors discuss why SAM is robust to noisy labels, and argue why SAM can improve generalization even for linear classifiers. Finally, the authors discuss how to further improve SAM by leveraging a gradient reweighting interpretation and combining with a robust loss.