Sharpness-aware Minimization for Efficiently Improving Generalization
Authors: Pierre Foret
Ariel Kleiner
Hossein Mobahi
Behnam Neyshabur
Venue: ICLR 2021 Spotlight (2021)
Keywords: Sharpness Minimization
Generalization
Regularization
Training Method
Deep Learning

Abstract: In today's heavily overparameterized models, the value of the training loss provides few guarantees on model generalization ability. Indeed, optimizing only the training loss value, as is commonly done, can easily lead to suboptimal model quality. Motivated by the connection between geometry of the loss landscape and generalization---including a generalization bound that we prove here---we introduce a novel, effective procedure for instead simultaneously minimizing loss value and loss sharpness.  In particular, our procedure, Sharpness-Aware Minimization (SAM), seeks parameters that lie in neighborhoods having uniformly low loss; this formulation results in a min-max optimization problem on which gradient descent can be performed efficiently. We present empirical results showing that SAM improves model generalization across a variety of benchmark datasets (e.g., CIFAR-{10, 100}, ImageNet, finetuning tasks) and models, yielding novel state-of-the-art performance for several.  Additionally, we find that SAM natively provides robustness to label noise on par with that provided by state-of-the-art procedures that specifically target learning with noisy labels.

Published as a conference paper at ICLR 2021
SHARPNESS-AWARY MINIMIZATION FOR EFFICIENTLY IMPROVING GENERALIZATION
Pierre Foret
Ariel Kleiner
Hossein Mobahi Google Research hmobah @googgoooe..ooooooo com
Google Research
Google Research ak e ner e com
eierreforet@@gogg com
Behnam Neyshabur Blueshift, Alphabet
neyshabur@gooole.com
ABSTRACT
In today's heavily overparameterized models, the value of the training loss pro- vides few guarantees model generalization ability. Indeed, optimizing only the training loss value, as is commonly done, can easily lead to suboptimal model quality. Motivated by prior work connecting the geometry of the loss landscape and generalization, we introduce a novel, effective procedure for in- stead simultaneously minimizing loss value and loss sharpness. In particular, our procedure, Sharpness-Awar Minimization (SAM), seeks parameters that lie in neighborhoods having uniformly low loss; this formulation results in a min- max optimization problem which gradient descent can be performed effi- ciently. We present empirical results showing that SAM improves model gen- eralization across variety of benchmark datasets (e.g.. CIFAR- -10, 100), Ima- geNet, finetuning tasks) and models, yielding novel state-of-the-arr performance for several. Additionally, bel noise on par with that provided by state-of-the-art procedures that specifi- cally target learning with noisy labels. We open source our code at ht https: / / //github. com/ google research/ sam.
find that SAM natively provides robustness to la-
INTRODUCTION
Modern machine learning's success in achieving ever better performance on wide range of tasks has relied in significant part on ever heavier overparameterization, in conjunction with developing ever more effective training algorithms that are able to find parameters that generalize well. Indeed, many modern neural networks can easily memorize the training data and have the capacity to readily overfit (Zhang et al., 2016). Such heavy overparameterizztion currently required to achieve state- of-the-ar art results in variety of domains (Tan & Le, 2019; Kolesnikov et al., 2020; Huang et al., 2018). turn, it essential that such models be trained using procedures that ensure that the parameters actually selected do in fact generalize beyond the training set.
Unfortunately, simply minimizing commonly used loss functions (e.g., cross-entropy) on the train- ing set is typically not sufficient to achieve satisfactory generalization. The training loss landscapes of today's models are commonly complex and non-convex, with a multiplicity of local and global minima, and with different global minima yielding models with different generalization abilities (Shirish Keskar et al., 2016). As a result, the choice of optimizer (and associated optimizer settings) from among the many available (e.g., stochastic gradient descent (Nesterov, 1983), Adam (Kingma
& Ba, 2014), RMSProp (Hinton et al.), and others (Duchi et al., 2011: Dozat, 2016; Martens & Grosse, 2015)) has become important design choice, though understanding of its relationship to model generalization remains nascent (Shirish Keskar et al., 2016: Wilson et al., 2017; Shirish Keskar & Socher, 2017: Agarwal et al., 2020; Jacot et al., 2018). Relatedly, a panoply of methods for modifying the training process have been proposed, including dropout (Srivastava et al., 2014),
WWrk done part of the Google AI Residency program.

---

Published as a conference paper at ICLR 2021
Ci
0000
Cifar100
o0 amee
000
Imagenet
oamo o
Finetuning
SVHN
F-MNIST
00:00
Noisy Cilar
Error reduction (%)
Figure 1: (left) Error rate reduction obtained by switching to SAM. Each point a different dataset / model / data augmentation. (middle) A sharp minimum to which ResNet trained with SGD converged. (right) wide minimum to which the same ResNet trained with SAM converged.
batch normalization (loffe & Szegedy, 2015), stochastic depth (Huang et al., 2016), data augmenta- tion (Cubuk et al., 2018), and mixed sample augmentations (Zhang et al., 2017; Harris et al., 2020).
The connection between the geometry of the loss landscape-in particular, the flatness of minima- and generalization has been studied extensively from both theoretical and empirical perspectives (Shirish Keskar et al., 2016; Dziugaite & Roy, 2017; Jiang al., 2019). While this connection has held the promise enabling new approaches to model training that yield better generalization, practical efficient algorithms that specifically seek out flatter minima and furthermore effectively improve generalization on range of state-of-the-aat models have thus far been elusive (e.g., see (Chaudhari et al., 2016; Izmailov al., 2018); we include more detailed discussion of prior work in Section 5).
We present here new efficient, scalable, and effective approach improving model generalization ability that directly leverages the geometry of the loss landscape and its connection to generaliza- tion, and powerfully complementary existing techniques. In particular, we make the following contributions:
We introduce Sharpness-Awar Minimization (SAM), novel procedure that improves model generalization by simultaneously minimizing loss value and loss sharpness. SAM functions by seeking parameters that lie in neighborhoods having uniformly low loss value (rather than parameters that only themselves have low loss value, as illustrated in the middle and righthand images of Figure 1), and can be implemented efficiently and easily.
We show via a rigorous empirical study that using SAM improves model generalization ability across range of widely studied computer vision tasks (e.g., CIFAR-(10, 100). ImageNet, finetuning tasks) and models, as summarized in the lefthand plot of Figure 1. For example, applying SAM yields novel state-of-the-un performance for number of already- intensely-studied tasks, such as ImageNet, CIFAR- 10, 100), SVHN, Fashion-MNIST,, and the standard set of image classification finetuning tasks (e.g.. Flowers, Stanford Cars, Oxford Pets, etc).
We show that SAM furthermore provides robustness label noise on par with that provided by state-of-the-art procedures that specifically target learning with noisy labels.
Through the lens provided by SAM, further elucidate the connection between loss sharpness and generalization by surfacing a promising new notion of sharpness, which we term m-sharpness.
Section below derives the SAM procedure and presents the resulting algorithm in full detail. Sec- tion 3 evaluates SAM empirically, and Section 4 further analyzes the connection between loss sharp- ness and generalization through the lens SAM. Finally, we conclude with an overview related work and a discussion of conclusions and future work in Sections 5 and 6, respectively.
2

---

Published as a conference paper at ICLR 2021
2 SHARPNESS-AWAER MINIMIZATION (SAM)
Throughout the paper, we denote scalars as a, vectors as a, matrices as A, sets as A, and equality by definition as _ Given a training dataset S _ U. { y/)} drawn i..d. from distribution L, we seek to learn a model that generalizes well. In partieular, consider a family of models parameterized by w E W C Ie:; given a per-data-point loss function / W x X' x ) + IR,,, we define the training set loss (s(w) _ 1 Li=l i= l(w, 24 yi) and the population loss L (w) - E(a.y)>
æ, y)]- Having observed only S, the goal of model training to select model parameters w having low
population loss Lg(w).
Utilizing (g(w) as an estimate of (ww motivates the standard approach of selecting parameters w by solving min Ls(w) (possibly in conjunction with regularizer on w) using an optimization procedure such as SGD or Adam. Unfortunately, however, for modern overparameterized mod- els such as deep neural networks, typical optimization approaches can easily result in suboptimal performance at test time. In particular, for modern models, (us)w) is typically non-convex w, with multiple local and even global minima that may yield similar values of Ls(w) while having significantly different generalization performance (i.e., significantly different values (w)).
Motivated by the connection between sharpness of the loss landscape and generalization. we propose a different approach: rather than seeking out parameter values w that simply have low training loss value Ls(w), we seek out parameter values whose entire neighborhoods have uniformly low training loss value (cquivalently, neighborhoods having both low loss and low curvature). The following theorem illustrates the motivation for this approach by bounding generalization ability in terms of neighborhood-ww wise training loss (full theorem statement and proof in Appendix A):
Theorem (stated informally) 1. For any p > O, with high probability over training set S generated from distribution 9,
(g(w) S max Ls(w + + + h([lw/l// |ell25p
where h: R + R., is a strictly increasing function (under some technical conditions on Ly (w)).
To make explicit our sharpness term, we can rewrite the right hand side of the inequality above
max Ls(w + <) Ls(w)) + (wr) + ++[[ (13pp)
|
The term in square brackets captures the sharpness of Lg at ur by measuring how quickly the training loss can be increased by moving from w to nearby parameter value; this sharpness term is then summed with the training loss value itself and a regularizer on the magnitude of w. Given that the specific function h is heavily influenced by the details of the proof, substitute the second term with >/ |w|] 2 for a hyperparameter x, yielding a standard L2 regularization term. Thus, inspired by the terms from the bound, we propose select parameter values by solving the following Sharpness- Aware Minimization (SAM) problem:
SAM (w) + A|| w| where min L
LSAM (wr)
max Ls(w + E), ||
(1)
where P 2 0 is a hyperparameter and p E oo] (we have generalized slightly from an L.2-norm to p-norm in the maximization over c, though show empirically in appendix C.5 that p 2 is typically optimal). Figure shows! the loss landscape for a model that converged to minima found by minimizing either Is(w)) (w) I. (wr). illustrating that the sharpness-aware loss prevents the model from converging to sharp minimum.
In order minimize L (w), we derive an efficient and effective approximation to V L.. (w) by differentiating through the inner maximization, which in turn enables to apply stochastic gradient descent directly to the SAM objective. Proceeding down this path, we first ap- proximate the inner maximization problem via a first-order Taylor expansion of Lg (w + e) w.r.t.
around O, obtaining
E" (w) A arg max Lsw ||
(w + 6) RJ arg max Ls(w) + E T VwLssw) = arg max E T V Vw Lsw). ||
||
FFiuure 1 was generated following Li et al. (2017) with the provided ResNet56 (no residual connections) checkpoint, and training the same model with SAM.
3

---

Published as a conference paper at ICLR 2021
In turn, the value É(w) that solves this approximation is given by the solution to classical dual norm problem () 1 denotes elementwise absolute value and power)":
((w) sign (Is()() (V [is (1V 1/P
where 1/p + 1/g- 1. Substituting back into equation (1) and differentiating, we then have
V V SAM (w) VwLs(w ((w)) d(w + E(w)) V (ow)ss()leee w+t(w))
dw
dē((w) V (.(lss()))))))) w+E(w)-
(Volss()) (+((w)
dw
This approximation V. L. (10) (w) can be straightforwardly computed via automatic differentia- tion, as implemented in common libraries such as JAX, TensorFlow, and PyTorrc. Though this com- putation implicitly depends the Hessian Ls(w) because E(w) is itself function of wsw(ww) the Hessian enters only via Hessian-vector vector products, which can be computed tractably without ma- terializing the Hessian matrix. Nonetheless, to further accelerate the computation, we drop the second-order terms. obtaining our final gradient approximation:
L (w) x Ls w+t(w))
As shown by the results in Section 3, this approximation (without the second-order terms) yields an effective algorithm. In Appendix C.4, additionally investigate the effect of instead including the second-order terms; in that initial experiment, including them surprisingly degrades performance, and further investigating these terms' effect should be priority in future work.
We obtain the final SAM algorithm by applying a standard numerical optimizer such as stochastic gradient descent (SGD) to the SAM objective LSAM (w), using equation 3 to compute the requisite objective function gradients. Algorithm gives pseudo-code for the full SAM algorithm, using SGD the base optimizer, and Figure schematically illustrates single SAM parameter update.
Input: Training set S -ー (ar; y. Loss function I: W x X' x + R., Batch size b, Step size 0, Neighborhood size p Output: Model trained with SAM
Wr
Initialize weights wu. L = 0:
(VL((W))
while not converged do
(L((W))
W:
SAM
Sample batch 8 = ( V1 ...),,..(..,, yb)};
Compute gradient V (te) the batch's training loss;
Compute (ww) per equation 2:
Wadv
(WaaWo
Compute gradient approximation for the SAM objective
(equation 3): V.. [
Update weights: W
1!::
1
Figure Schematic of the SAM param- eter update.
end
return wi
Algorithm SAM algorithm
EMPIRICAL EVALUATION
In order to assess SAM's efficacy, we apply it to a range of different tasks, including image clas- sification from scratch (including on CIFAR-10, CIFAR-100, and ImageNet), finetuning pretrained models, and learning with noisy labels. In all cases, we measure the benefit of using SAM by simply replacing the optimization procedure used train existing models with SAM, and computing the resulting effect on model generalization. As seen below. SAM materially improves generalization performance in the vast majority of these cases.
In the of interest p this hoils down to simply rescaling the gradient such that its p.
4

---

Published as a conference paper at ICLR 2021
3.1 IMAGE CLASSIFICATION FROM SCRATCH
We first evaluate SAM's impact on generalization for today's 's state-of-the-ar models on CIFAR-10 and CIFAR-100 (without pretraining): WideResNets with ShakeShake regularization (Zagoruyko & Komodakis, 2016; Gastaldi, 2017) and PyramidNet with ShakeDrop regularization (Han et al., 2016; Yamada et al., 2018). Note that some of these models have already been heavily tuned in prior work and include carefully chosen regularization schemes to prevent overfitting: therefore, significantly improving their generalization is quite non- non-trivial. We have ensured that our imple- mentations' generalization performance in the absence of SAM matches or exceeds that reported in prior work (Cubuk et al., 2018; Lim et al., 2019)
All results use basic data augmentations (horizontal flip, padding by four pixels, and random crop). We also evaluate in the setting of more advanced data augmentation methods such as cutout regu- larization (Devries & Taylor, 2017) and AutoAugment (Cubuk al., 2018), which are utilized by prior work to achieve state-of-the-art results.
SAM has single hyperparameter p (the neighborhood size), which we tune via a grid search over {0.01, 0.02, 0.05, 0.2, 0.5} using 10% of the training set as a validation set'. Please see ap- pendix C.I for the values of all hyperparameters and additional training details. As each SAM weight update requires two backpropagation operations (one to compute é( ((w) and another to com- pute the final gradient), we allow each non-SAM training run to execute twice as many epochs as each SAM training run, and we report the best score achieved by each non-SAM training run across either the standard epoch count or the doubled epoch count'. We run five independent replicas of each experimental condition for which we report results (each with independent weight initialization and data shuffling), reporting the resulting mean error (or accuracy) on the test set, and the associ- ated 95% confidence interval. Our implementations utilize JAX (Bradbury et al., 2018), and we train all models on a single host having 8 Nvidia V100 GPUs To compute the SAM update when parallelizing across multiple accelerators, independently compute the SAM gradient on each accelerator, and average the resulting sub-batch SAM gradients to obtain the final SAM update.
divide each data batch evenly among the accelerators,
As seen in Table 1, SAM improves generalization across all settings evaluated for CIFAR-10 and CIFAR-100. For example, SAM enables a simple WideResNet attain 1.6% test error, versus 2.2% error without SAM. Such gains have previously been attainable only by using more complex model architectures (e.g.. PyramidNet) and regularization schemes (e.g., Shake-Shake, ShakeDrop); SAM provides an -mplllemmpeeeeteed, model-indepene improvements even when applied atop complex architectures that already use sophisticated regular- ization: for instance, applying SAM to PyramidNet with ShakeDrop regularization yields 10.3% error on 100, which is, to our knowledge, new Matee-ofthe-aa on this dataset without the use of additional data.
alternative. Furthermore, SAM delivers
Beyond CIFAR-(10, 100), we have also evaluated SAM on the SVHN (Netzer et al., 2011) and Mashion-MNIST datasets (Xiao et al., 2017). Once again, SAM enables a simple WideResNet to achieve accuracy at or above the state-of-the-art for these datasets: 0.99% error for SVHN and 3.59% for Fashion- Details are available in appendix B.1.
To assess SAM's performance at larger scale, we apply it to ResNets (He et al., 2015) of different depths (50. 101, 152) trained on ImageNet (Deng et al., 2009). In this setting, following prior work (He et al., 2015; Szegedy et al., 2015). we resize and crop images to 224-pixel resolution, normalize them, and use batch size 4096, initial learning rate 1.0, cosine learning rate schedule, SGD optimizer with momentum 0.9, label smoothing 0.1, and weight decay 0.0001. When applying SAM, we use p 0.05 (determined via grid search on ResNet-50 trained for 100 epochs). We train all models on ImageNet for up to 400 epochs using a Google Cloud TPUv3 and report top-1 and top-5 test error rates for each experimental condition (mean and 95% confidence interval across 5 independent runs).
We found p = 0.05 to solid default value, and we report in appendix C.3 the scores for all our experiments, obtained with = 0.05 without further tuning.
'Training for longer generally did not improve accuracy significantly, except for the models previously trained for only 200 epochs and for the largest, most regularized model (PyramidNet + ShakeDrop).
'Because SAM's performance is amplified by not syncing the perturbations, data parallelism is highly recommended to leverage SAM's full potential (see Section 4 for more details).
5

---

Published as a conference paper at ICLR 2021
CIFAR-IO
CIFAR
100
Model
Augmentation
SAM
SGD
SAM
SGD
WRN-28-10 (200 epochs)
Basic
2.740.1
3.5-01
16.5=ao2
18.820.2
WRN-28-10 (200 epochs)
Cutout
2.310.1
2.6=0.1
14,9
16.910.1
WRN-28-11 (200 epochs)
AA
2.1 2.1
2.3=0.1
13.6
15.8
WRN-28-10 (1800 epochs)
Basic
2.4
3.5=0.1
16.3
19.1
WRN-28-10 (1800 epochs)
Cutout
2,140.1
2.7=01
14.0=0.1
17.4.0.1
WRN-28-10 (1800 epochs)
AA
1.640.1
2.24c0.1
12.8
16.1
Shake-Shake (26 2x96d)
Basic
2.3
2.7
15.1
17.0=o.1
Shake-Shake (26 2x96d)
Cutout
2.0
2.3=0.1
14.220.2
15.720.2
Shake-Shake (26 2x96d)
AA
1.64
1.9-01
12.8
14.1+0.2
PyramidNet
Basic
2.700.1
4.0_a.1
14.6=0.4
19.720.3
PyramidNet
Cutout
1940.1
2.5
12.6-0.2
16.420.1
PyramidNet
AA
1640.11
1.9-u11
11.6-0.1
14.6+u.1
PyramidNet-ShakeDDop
Basic
2.110.1
2.5
13.3-0.2
14.520.1
Shaakid
Cutout
1.6:0..1
1.9_0.1
11.3.0.1
11.820.2
PyamidNNe+-ShakeDDop AA
1.4
1.64c0.1
10.3
10.6.0.1
Table 1: Results for SAM state-of-the-art models CIFAR-(1O, 100) (WRN WideResNet; AA = AutoAugment; SGD is the standard non-SAM procedure used to train these models).
As seen in Table 2, SAM again consistently improves performance, for example improving the ImageNet top-1 error rate of ResNet-152 from 20.3% to 18.4%. Furthermore, note that SAM enables increasing the number of training epochs while continuing to improve accuracy without overfitting. In contrast, the standard training procedure (without SAM) generally significantly overfits as training extends from 200 to 400 epochs.
SAM
Standard Training (No SAM)
Model
Epech
Top-1
Top-5
Top-1
Top-5
ResNer-50
100
22.540.1
6.2840.08
22.920.1
6.62+0.11
200
21.440.1
5.82+0.03
22.3
6.37
400
20.9
5.5140.03
22.3
6.40
ResNet- 101
100
20.240.1
5.1240.03
21.220.1
3.66
200
19.4
4.7640.03
20.920.1
5.6620.04
400
19.0+<0.01
4.6540.05
22.3
6.41
ResNet-152
100)
19.2
4.69
20.44
5.39
200
18.540.1
4.37+0.03
20.3
5.39
400
18.4
4.35
20.9+-c0.0
5.84
Table 2: Test error rates for ResNets trained on ImageNet, with and without SAM.
3.2 FINETUNING
Transfer learning by pretraining model on large related dataset and then finetuning on smaller target dataset of interest has emerged as powerful and widely used technique for producing high- quality models for a variety of different tasks. We show here that SAM once again offers con- siderable benefits in this setting, even when finetuning extremely large, state-of-the-arr, already gigh-performing models.
In particular, we apply SAM to finetuning EfficentNer-b7 (pretrained on ImageNet) and EfficiemNett-22 (pretrained on ImageNet plus unlabeled JFT; input resolution 475) (Tan & Lc, 2019; Kornblith et al., 2018; Huang et al., 2018). We initialize these models to publicly available check- points trained with RandAugment (84.7% accuracy on ImageNet) and NoiseStudent (88.2% ac- curacy ImageNet), respectively. We finetune these models on each of several target datasets by training each model starting from the aforementioned checkpoint; please see the appendix for details of the hyperparameters used. We report the mean and 95% confidence interval of top-1 test error over 5 independent runs for each dataset.
httpp: http://www....aaaooooooooooaa//a efficientnet.

---

Published as a conference paper at ICLR 2021
As seen in Table 3, SAM uniformly improves performance relative to finetuning without SAM. Furthermore, in many cases, SAM yields novel state-of-the-ar the-art performance, including 0.30% error on CIFAR-10, 3.92% error CIFAR-100, and 11.39% error ImageNet.
Et
Prev. SOTA
EffNet-L2
Dutasct
+ SAM
EfiNet-b77
(ImageNet only)
+ SAM
EftNet-L2
Prev. SOTA
FiveaAireratt
6.8020.06
8.1040.08
5.3 (TBMSL--Ne)
4.82
5.80
(TBMSLLNNNl
Flowers
0.63
1.1640.05
0.7 (BiT-M)
0.35
0.40
0.37 (EffNet)
Oxford IIIT Pets
10.04
1..09
4.1 (Gpipe)
2.90
3.0810.04
(Gpipe)
Stanford. Cars
5.1820.02
5.94
5.0 (TBMSL-Net)
4.0420.03
4.93+0.04
(DAT)
CIFAR-10 10
0.00
0.95
(Gpipe)
0.30
0.34
0.63 (BiT-L)
CIFAR-100 100
7.44
7.68++.06
7.83 (BiT-M)
3.92
4.07
6.49 (BiT-L)
Birdsnap
13.64
10.10
15.7 (ERNeet)
10.15
10.31 10.15
14.5 (DAT)
Food101
7.02=0002
7.17400.0
7.0 (Gpipe)
3.82.g0.01
3.97
4.7 (DAT)
ImageNet
15.1440:03
15.3
14.2 (KDIorAA)
11.39:00.02
11.8
11.45 (VIT)
Table 3: Top-1 error rates for finetuning EfficientNet-b7 (left; ImageNet pretraining only) and EfficientNet-L22 (right; pretraining on ImageNet plus additional data, such as JFT) on various down- stream tasks. Previous state-of-the-art (SOTA) includes EfficientNet (EffNet) (Tan & Le, 2019), Gpipe (Huang al., 2018), DAT (Ngiam et al., 2018), BiT-M/L (Kolesnikov al., 2020), KD- forAA (Wei et al., 2020). TBMSL-Ne (Zhang et al., 2020), and ViT (Dosovitskiy et al., 2020).
3.3 ROBUSTNESS TO LABEL NOISE
The fact that SAM seeks out model parameters that are robust to perturbations suggests SAM's poten- tial provide robustness noise in the training set (which would perturb the training loss landscape). Thus, we assess here the degree of robustness that SAM provides to label noise.
Mcthod
Noise rate
20
40
60
80
Sanchez al. (2019)
94.0
92.8
90.3
74.1
Zhang & Sabuncu (2018) 8)
89.7
87.6
82.7
67.9
Lee et al. (2019)
87.1
81.8
75.4
Chen et (2019)
89.7
52.3
Huang
(2019)
90.3
43.4
In particular, we measure the effect of apply- ing SAM in the classical noisy-label setting for CIFAR-10, in which fraction of the training set's labels are randomly flipped; the test set remains unmodified (i.e., clean). To ensure valid compar- ison to prior work, which often utilizes architec- tures specialized to the nobsy-label setting, we train simple model of similar size (ResNet-32) for 200 epochs, following Jiang et al. (2019). We evalu- five variants of model training: standard SGD, SGD with Mixup (Zhang et al., 2017), SAM, and "bootstrapped" variants of SGD with Mixup and SAM (wherein the model is first trained as usual and then retrained from scratch on the labels pre- dicted by the initially trained model). When apply- ing SAM, we use p = 0.1 for all noise levels except 80%, for which we use p = 0.05 for more stable convergence. For the Mixup baselines, tried all values of a E {1,8, 16, 32) and conservatively report the best score for each noise level.
MentorNet (2017)
92.0
91.2
74.2
60.0
Mixup (2017)
94.0
91.5
86.8
76.9
MentorMix (2019)
95.6
94.2
91.3
81.0
SGD
84.8
68.8
48.2
26.2
Mixup
93.0
90.0
83.8
70.2
Boutstrap Mixup
93.3
92.0
87.6
72.0
SAM
93.4
90.5
77.9
Boolstrap SAM
95.4
94.2
91.8
79.9
Table 4: Test accuracy on the clean test set for models trained CIFAR-10 with noisy la- bels. Lower block is our implementation. up- per block gives scores from the literature, per Jiang et al. (2019).
As seen in Table 4, SAM provides high degree of robustness to label noise, par with that provided by ofate-of-the art procedures that specifically target learning with noisy labels. Indeed, simply training model with SAM outperforms all prior methods specifically targeting label noise robustness, with the exception of MentorMix (Jiang et al., 2019). However, simply bootstrapping SAM yields performance comparable to that of MentorMix (which is substantially more complex).

---

Published as a conference paper at ICLR 2021
Epoch:
Araa 9
Amray As 2.5
Araa 18.6
Amaa Aog 3.6
0.08
0.050
20
60
60
Enror (%)
0.07
Mutual information
0.045
Task 2
0.17
information
Epoch: 50
Amaz 12.5
Amaz Po 1.7
Aras 8.9
Amad/Ao 1.9
0.06
64
16
0.040
0.16
258
10
10
0.035
0.15
Epech: 300
Amas 24.2
Amas 1.0
0.030
Amau//A - 11.4
Amass/A - 2.6
0.00
0.05
0.10
0.15
256
m
Figure 3: (left) Evolution of the spectrum of the Hessian during training of a model with standard SGD (lefthand column) or SAM (righthand column). (middle) Test error as a function p for dif- ferent values of m. (right) Predictive power of m-sharpness for the generalization gap, for different values of m (higher means the sharpness measure more correlated with actual generalization gap).
SHARPNESS AND GENERALIZATION THROUGH THE LENS OF SAM
4.1
M-SSHARPNESS
Though our derivation of SAM defines the SAM objective over the entire training set, when utilizing SAM in practice, we compute the SAM update per-batch (as described in Algorithm 1) or even by averaging SAM updates computed independently per-accelerator (where each accelerator receives subset of size m of batch, as described in Section 3). This latter setting is equivalent to modifying the SAM objective (equation 1) sum over a set of independent e maximizations, each performed disjoint subset of m data points, rather than performing the
on sum of per-data-pon -ardttt-pointtt losses e maximization over a global sum over the training set (which would be equivalent to setting m to the total training set size). We term the associated measure of sharpness of the loss landscape m-sharpness.
To better understand the effect of SAM, we train small ResNet CIFAR-10 10 using SAM with a range of values of m. As seen in Figure (middle), smaller values m tend to yield models having better generalization ability. This relationship fortuitously aligns with the need to parallelize across multiple accelerators in order to scale training for many of today's 's models.
Intriguingly, the m-sharpness measure described above furthermore exhibits better correlation with models' actual generalization gaps as m decreases, as demonstrated by Figure 3 (right)². In partic- ular, this implies that m-sharpness with < yields better predictor of generalization than the full-training-ss set measure suggested by Theorem in Section 2 above, suggesting an interesting new avenue of future work for understanding generalization.
4.2 HESSIAN SPECTRA
Motivated by the connection between geometry of the loss landscape and generalization, we con- structed SAM to seek out minima of the training loss landscape having both low loss value and low curvature (i.e., low sharpness). To further confirm that SAM does in fact find minima having low curvature, compute the spectrum of the Hessian for WideResNet40-111 trained CIFAR-10 10 for 300 steps both with and without SAM (without batch norm, which tends to obscure interpretation of the Hessian), different epochs during training. Due to the parameter space's dimensionality, approximate the Hessian spectrum using the Lanczos algorithm of Ghorbani et al. (2019).
Figure 3 (left) reports the resulting Hessian spectra. As expected. the models trained with SAM converge to minima having lower curvature, as seen in the overall distribution of eigenvalues, the
We follow the rigorous framework Jiang et al. (2019), reporting the mutual information between the m-sharpness measure and generalization on the two publicly available tasks from the Predicting gen- eralization in deep learning NeurIPS2020 competition. https: / compet i i ona. coda ab. /rg/ compet eicen//2 cns/25301
8

---

Published as a conference paper at ICLR 2021
maximum eigenvalue (Amax) at convergence (approximately 24 without SAM, 1.0 with SAM), and the bulk of the spectrum (the ratio \max/ \s. commonly used as a proxy for sharpness (Jastrzebski et al., 2020); up to 11 11.4 without SAM, and 2.6 with SAM).
RELATED WORK
The idea of searching for "flat" minima can be traced back Hochreiter & Schmidhuber (1995), and its connection to generalization has seen significant study (Shirish Keskar et al., 2016; Dziugaite & Roy, 2017; Neyshabur al., 2017; Dinh et al., 2017). In recent large scale empirical study, Jiang et al. (2019) studied 40 complexity measures and showed that a sharpness-based measure has highest correlation with generalization, which motivates penalizing sharpness. Hochreiter & Schmidhuber (1997) was perhaps the first paper on penalizing the sharpness, regularizing notion related to Min- imum Description Length (MDL). Other ideas which also penalize sharp minima include operating on diffused loss landscape (Mobahi, 2016) and regularizing local entropy (Chaudhari et al., 201 6). Another direction is to not penalize the sharpness explicitly, but rather average weights during train- yield flatter minima that can also generalize
ing; Izmailov et al. (2018) showed that doing better. However, the measures sharpness proposed previously are difficult to compute and differ- entiate through. In contrast, SAM is highly scalable it only needs gradient computations per iteration. The concurrent work Sun et al. (2020) focuses on resilience to random and adversarial corrpption to expose model's vulnerabilitiie; this work is perhaps closest to ours. Our work has different basis: we develop SAM motivated by a principled starting point in generalization, clearly demonstrate SAM's efficacy via rigorous large-scale empirical evaluation, and surface important practical and theoretical facets of the procedure (e.g., m-sharppess). The notion of all-layer margin introduced by Wei & Ma (2020) is closely related to this work; one is adversarial perturbation over the activations of a network and the other over its weights, and there is some coupling between these two quantities.
DISCUSSION AND FUTURE WORK
In this work, we have introduced SAM, novel algorithm that improves generalization by simulta- neously minimizing loss value and loss sharpness; we have demonstrated SAM's efficacy through rigorous large-scale empirical evaluation. We have surfaced a number interesting avenues for fu- ture work. On the theoretical side, the notion of per-data point sharpness yielded by m-sharpness (in contrast to global sharpness computed over the entire training set, as has typically been studied the past) suggests interesting new lens through which to study generalization. Methodologicallly results suggest that SAM could potentially be used in place of Mixup in robust or semi eemissupervised methods that currently rely on Mixup (giving, for instance, MentorSAM). We leave to future work a more in-depth investigation of these possibilities.
7 ACKNOWLEDGMENTS
We thank our colleagues at Google Atish Agarwala, Xavier Garcia, Dustin Tran, Yiding Jiang, Basil Mustafa, Samy Bengio- for their feedback and insightful discussions. We also thank the JAX and FLAX teams for going above and beyond to support our implementation. We are grateful to Sven Gowal for his help replicating EfficientNet using JAX, and Justin Gilmer for his implementation of the Lanczos algorithm" used to generate the Hessian spectra. We thank Niru Maheswaranathan for his matplotlib mastery. We also thank David Samuel for providing PyTorc implementation of SAM".
REFERENCES
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Revisiting the gener- alization of adaptive gradient methods, 2020. URL /ttp://///pe/eeew/e/// net/forum 1d=BJ16t 64tvr.
Phttps://gt thub. wtw.thh.....oo...oooooooooott goog/ specttall-dessitty Phttps: //gi thub. com haabb..oo/dvvaaa/ 54 sam
9

---

Published as a conference paper at ICLR 2021
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milnn JAX: composable transformations of Python+NumPy programs, 2018. URL http:
com/ google/ jax.
Niladri S Chatterji, Behnam Neyshabur, and Hanie Sedghi. The intriguing role of module criticality in the generalization of deep networks. In International Conference on Learning Representations, 2020.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD: Biasing Gradi- November 2016.
ent Descent Into Wide Valleys. arXiv e-prints, art.
Pengfei Chen, Benben Liao, Guangyong Chen, and Shengyu Zhang. Understanding and utilizing deep neural networks trained with noisy labels. CoRR, abs/1905.05040, 2019. URL http: / //arxiv.org/abll www.ass///aaaaaeooooooo.....eeea hatp//bb//190 1905. 05040.
Ekin Dogus Cubuk, Barret Zoph, Dandelion Mané, Vijay Vasudevan, and Quoc V. Le. Au- toaugment: Learning augmentation policies from data. CoRR, abs/1805.09901 2018. URL
http:// http//aaxxv....bba.a./
.95501.
J. Deng, W. Dong, R. Socher, L.-. Li, K. Li, and L. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR09, 2009.
Terrance Devries and Graham W. Taylor. Improved regularization of convolutional neural networks with cutout. CoRR, Abs/1708.04552 2017. URL http::ttt///. //apii///xiiv// .rgg woorg/ab//iie /aa///b////// 04552.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. arXiv preprint arXiv:1703.04490 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv e-prints, art. arXiv:2010.1 1929, October 2020.
Timothy Dozat. Incorporating nesterov momentum into adam. 2016.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data, arXiv preprint / 1008, 2017.
Xavier Gastaldi.
Shake-shake regularization. CoRR, abs/1705.07485. 2017.
URL http:
/
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An Investigation into Neural Net Optimiza- tion via Hessian Eigenvalue Density. arXiv e-prints, art. 1rX1v:1901.10155: January 2019.
Dongyoon Han, Jiwhan Kim, and Junmo Kim. Deep pyramidal residual networks. CoRR, abs/1610.02915, 2016. URL http: http:://wxxxvvooooeeooo ottp:/.ooo////////////. habbs///....91..
Ethan Harris, Antonia Marcu, Matthew Painter, Mahesan Niranjan, Adam Prügel-Bennett, and Jonathon Hare. FMix: Enhancing Mixed Sample Data Augmentation. arXiv e-prints, art. 12447, February 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual leamnig for image recog- nition. CoRR, 0bs/15122.03385 2015. URL http: /ttpppp///axiv.../e org/ aaa/abbs/1151 03385.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In Advances in neural information processing systems, pp. 529-536, 1995.
10

---

Published as a conference paper at ICLR 2021
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Weinberger. Deep Networks with Stochastic Depth. arXiv e-prints, art.
March 2016.
J. Huang, L. Qu, R. Jia, and B. Zhao. O2u-net: A simple noisy label detection approach for deep neural networks. In 2019 IEEE/CVF International Conference Computer Vision (ICCV), pp- 3325-3333 2019.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen. GPipe: Ef- ficient Training of Giant Neural Networks using Pipeline Parallelism. arXiv e-prints, art. 1.X,v.1811.06965, November 2018.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167, 2015.
Pavel Izmailov, Dmitrii Podoprikhin. Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wil- son. Averaging Weights Leads Wider Optima and Better Generalization. arXiv e-prints, art. 1eXiv:1803.05400 March 2018.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. CoRR, abs/1886.00572, 2018. URL 1806. 07572.
/
Stanislaw Jastrzebski, Maciej Szymezak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The Break-Even Point on Optimization Trajectories of Deep Neural Networks. arXiv e-prints, art. arXivv2002.09572, February 2020.
Lu Jiang, Zhengyuan Zhou, Thomas Leung, Li-Jia Li, and Li Fei-Fei. Mentomet: Regularizing very deep neural networks on corrupted labels. CoRR, abs/1712.05055, 2017. URL http: / //arxiv.or www. abwaabaa// 111.//1111.0 05055.
Lu Jiang, Di Huang, Mason Liu, and Weilong Yang. Beyond Synthetic Noise: Deep Learning on Controlled Noisy Labels. arXiv e-prints, art. arXiv:1911 1.9981,, November 2019.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. arXiv preprint 1rXiv:111..02178, 2019.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv e-prints, art. 4tXiv:1412.69800 December 2014.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning, 2020.
Simon Kornblith, Jonathon Shlens, and Quoe Le. Do Better ImageNet Models Transfer Better? arXiv e-prints, art. arXiv: www::1805009995 May 2018.
John Langford and Rich Caruana. (not) bounding the true error. In Advances in Neural Information Processing Systems, pp. 809-816, 2002.
Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selec- tion. Annals of Statistics, pp. 1302-1338, 2000.
Kimin Lee, Sukmin Yun, Kibok Lee, Honglak Lee, Bo Li, and Jinwoo Shin. Robust inference via generative classifiers for handling noisy labels, 2019.
Hao Li, Zheng Xu, Gavin Taylor, and Tom Goldstein. Visualizing the loss landscape neural nets. CoRR, abs/1712.00911. 2017. URL
org/
wogg/ab/////ooooooooo/oooo//. abs 1721.1111 09913.
Sungbin Lim, Ildoo Kim, Taesup Kim, Chiheon Kim, and Sungwoong Kim. Fast autoaugment. CoRR, abs/1 0111999559097. 2019. URL http:://aaaee org/ abs/aabss 00397.
James Martens and Roger Grosse. Optimizing Neural Networks with Kronecker-factoonn Approxi- mate Curvature. arXiv e-prints, art. 1rXiv:1503.0555 March 2015.

---

Published as a conference paper at ICLR 2021
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual confer- ence on Computational learning theory, pp. 164-170, 1999.
Hossein Mobahi. Training recurrent neural networks by diffusion. CoRR, Abs/1601.041) 14, 2016.
URL pett:://axxvy
abs/1601, 16011.4114..
Y. E. Nesterov. A method for solving the convex programming problem with convergence rate o11/k²). Dokl. Akad. Nauk SSSR, 269:543-5444 1983. URL https: att./.iiiiii.................... neid/1002996622/em
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. 2011.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring general- ization deep learning. In Advances in neural information processing systems, pp. 5947-5956,
2017.
Jiquan Ngiam, Daiyi Peng, Vijay Vasudevan, Simon Kornblith, Quoc Le, and Ruoming Pang. Domain adaptive transfer learming with specialist models. CoRR, abs/1811.07056, 2018. URL http://arxiv org/ aabs/1811 07056.
Eric Arazo Sanchez, Diego Ortego, Paul Albert, Noel E. O'Connor, and Kevin McGuinness. Un- supervised label noise modeling and loss correction. CoRR, abs/1904. 1123/, 2019. URL
httt:://arxivorgg/
19044.9..... 11238.
Nitish Shirish Keskar and Richard Socher. Improving Generalization Performance by Switching from Adam to SGD. arXiv e-prints, art. arXiv:1712.07622. December 2017.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On Barge-Batch Training for Deep Learning: Generalization Gap and Sharp Minima. arXiv e-prints, art. arXiv:1609) 04866, September 2016.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Iya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15()):1929-1988 2014.
Xu Sun, Zhiyuan Zhang, Xuancheng Ren, Ruixuan Luo, and Liangyou Li. Exploring the Vul- nerability of Deep Neural Networks: A Study of Parameter Corruption. arXiv e-prints, art. arXiv:2006.05620, June 2020.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jonathon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision, 2015.
Mingxing Tan and Quoc V. Le. EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks. arXiv e-prints, art. 190. May 2019.
Colin Wei and Tengyu Ma. Improved sample complexities for deep neural networks and robust classification via an all-layer margin. In International Conference on Learning Representations, 2020.
Longhui Wei, An Xiao, Lingxi Xie, Xin Chen, Xiaopeng Zhang, and Qi Tian. Circumventing Outliers AutoAuument with Knowledge Distillation. arXiv e-prints, art. arXiv:2003.11342. March 2020.
Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. In Advances in neural information pro- cessing systems, pp-
2017.
Han Xiao, Kashif Rasul, and Roland Vollgraf.
Fashion-mnist a novel image dataset for benchmark- ing machine learning algorithms. CoRR, abs/1708.077.0 2017. URL httt::///axxii org/ abs/1708. 07747.
Yoshihiro Yamada, Masakazu Iwamura, and Koichi Kise. Shakedrop regularization. CoRR,
abs/1 00200100.,2335 2018. URL http: /tatp://rxxiii www.v/vw.v...o........o....oo... orgg abs/ /10.s////...... 02.35.
12

---

Published as a conference paper at ICLR 2021
Sergey Zagonuyko and Nikos Komodakis. Wide residual networks. CoRR, abs/ 2016. URL http:://axii/ org/ abs/1603 07146.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. CoRR, abs/1611.03530, 2016. URL http: // //arxiv. orq. abs 11161111111 03530.
Fan Zhang, Meng Li, Guisheng Zhai, and Yizhao Liu. Multi-branch and multi-scale attention learn- ing for fine-gre geeigraieed visual categorization, 2020.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412 2017.
Zhilu Zhang and Mert R. Sabuncu. Generalized cross entropy loss for training deep neural networks with noisy labels. CoRR, abs/1805.07836, 2018, URL htttp://axxi.... org/ abs 1805. 07836.
A APPENDIX
A. PAC BAYESIAN GENERALIZATION BOUND
Below, state generalization bound based sharpness.
Theorem 2. For any p > and any distribution 2, with probability 1 8 over the choice of the training set S 2.
k klog
1
||al1
1
log(n) b
+ +4log 3/ + Ö(1
(w) S max Ls(w + E) |ellaSp
(4)
where n = (S), the number of parameters and we assumed Lg(w) S (ee)s(() (o) TTg (w + é)}.
The condition Lg (w) S (e,=N((,e)) [L not decrease the test error. This is expected hold in practice for the final solution but does not necessarily hold for any w.
(Lg + E)] means that adding Gaussian perturbation should
Proof. First, note that the right hand side of the bound in the theorem statement is lower bounded by v log(1 + ||wll/p²)/ that case, the right hand side becomes greater than in which case the inequality holds trivially. 1).
which is greater than when || > p² ('(cx(AAn/H) - 1). In
Therefore, in the rest of the proof, we only consider the case when ||w|| S (exp((4n
here is inspired from Chatterji al. (2020). Using PAC-Bayesian generalization bound McAllester (1999) and following Dziugaite & Roy (2017), the following gen- cralization bound holds for any prior S over parameters with probability the training set S, for any posterior 2 over parameters:
The proof technique
over the choice of
KL(2) 9) log
E Ew ( < E [
(5)
2(n - 1)
Moreover, if 3 -NWNupp of1) and 2 NNuo ( I2 then the KL divergence can be written as follows:
1|pp Hell kog
KL
k + k kog
(6)
of
Given a posterior standard deviation σQ, one could choose a prior standard deviation ap to minimize the above KI. divergence and hence the generalization bound by taking the derivative 10 of the above
Despite the nonconvexity of the function here in of it has unique stationary point which happens to be its minimizer.
13

---

Published as a conference paper at ICLR 2021
KL with respect to op and setting it to zero. We would then have pp = o However, since op should be chosen before observing the training data S and NQ: could depend on S. we not allowed to optimize op in this way. Instead, can have set of predefined values for op and pick the best one in that set. See Langford & Caruana (2002) for the discussion around this technique. Given fixed b >> let T {cexp((1 j) N) be that predefined
11pe يQ /
set of values for oz. If for any j E N, the above PAC-Bayesian bound holds for of = c cexp((1
j)/k) with probability 1 - 5, with 6, T4jz. then by the union bound, all above bounds hold simultaneously with probability at least 1 ブド=1 デッ
68
6.
Let ae = Ps p = un and μP 0. Therefore, we have:
a + 1pp / 1| 12/k s p² exp(4n/k))
We now consider the bound that corresponds to j = [1 - k log + || 3/ /k)/ We can ensure that j E N using inequality equation and by setting
p² exp( 4n/k)). Furthermore, for
oi. = cexp((1 - j)/k), we have:
p² + o. exp(1/k) (0²
| kl/kk
Therefore, using the above value for σp, KL divergence can be bounded follows:
1
[ko + up - 142l
of
(1(9112)
k loog
(9)
2
oz
0%
1
Tkpp²= + || l1//k]
exp(1/k)
(p² +
|00
(3/k)
k + klog
(10)
2
p²
[
12/k
Ø²
1
exp(11k) (02 + | //0k//k) )
k loo
(11)
2
D²
1
2 1+
12
1 ls klog
(12)
2
k
Given the bound that corresponds to j holds with probability 1 - 6g for 6, the bound can be written as:
zzzz. the log term in
n2j²
log n
n + log
log
n 22 log
2
|12+||001///////
n
< log
+ log
5
6
22 log 2 (c/p²)
n σ
I log
+ log
6 22 2og2(1 (1 + exp(4n/k))
n
< log
+ log
8
6
22 (2 + 4n/k)²
< log
+ log
6
6
n
< log
log (6n + 3k)
Therefore, the generalization bound can be written as follows:
(Lgwwte) & (Ts kloo (1
|| bog
+ 1 + log n + +2log (6n
3k)
(13)
In the above bound, we have e; NN, a). Therefore, || has chi-square distribution and by Lemma Laurent & Massart (2000), we have that for any positive t:
P([lel1: ko² > 20² / + 210²) exp(-()
(14)
14

---

Published as a conference paper at ICLR 2021
Therefore, with probability 1 - 1/v/ we have that:
2 In(n) (
|| - (2ln(() + + 2vknn(V)))) S s
S
p
k
Substituting the above value for a back the inequality and using theorem's assumption gives us following inequality:
(w) S (1 - 1/v1)
max Lssw + E) + 1/vn ||l|
} k log
1+ ||
1 +
(og(n))
+ log 7/ + 2 +l2og (6n + 3k)
k
71
max c SP
Ls(w + c)+
I log
|| 1
(
1
log(n) k
2
+ 4 log 3/ + 8ogg (6n + 3k)
D
1
B ADDITIONAL EXPERIMENTAL RESULTS
B.I SVHN AND FASHION-MNIST
We report in table 5 results obtained SVHN and Fashion-MNIST datasets. On these datasets, SAM allows simple WideResnet to reach or push state-of-the-art accuracy (0.99% error rate for SVHN, 3.59% for Fashion-MNISTT..
For SVHN, we used all the available data (73257 digits for training set + 531131 additional samples). For auto-augment, we use the best policy found this dataset described in (Cubuk et al., 2018) plus cutout (Devries & Taylor, 2017). For Fashion-Mi MNNIS, the auto-augmentation line correspond cutout only.
Table 5: Results on SVHN and Fashion-MNIST.
SVHN
Pashion-MNIS"
Model
Augmentation
SAM
Baseline
SAM
Baseline
Reale-RRNNe--2-110
Basic
1.42 1.0220.02
1.58 1..8
3..8
4.57
Wide-ResNNet-28-11
Auto augment
0.99...01 00.01
1.14
3..0
3.8610.14
Shake-Shake (26 2x96d)
Basic
1.44-0.u2
1.58
3.97 20.09
4.37
Shake-Shake (26 2x96d)
Auto augment
1.07
1.03
3.59
3.59+0.01
#0.01
3.7640.07
C EXPERIMENT DETAILS
C.1 HYPERPARAMETERS FOR EXPERIMENTS
We report in table 6 the hyper-parameters selected by gridsearch for the CIFAR experiments, and the ones for SVHN and Fashion-MNISS in 7. For CIFARIO, CIFAR100, SVHN and Fashion-MNIST. use batch size of 256 and determine the learning rate and weight decay used to train each model via a joint grid search prior to applying SAM; all other model hyperparameter values are identical to those used in prior work.
For the Imagenet results (Resnet models), the models are trained for 100, 200 or 400 epochs on Google Cloud TPUv3 32 cores with batch size of 4096. The initial learning rate is set 1.0 and
15

---

Published as a conference paper at ICLR 2021
Table 6: Hyper-paramette used produce the CIFAR-(10,100) results
CIFAR Dataset
LR
WD
(CIFAR-10)
(CIFAR-100)
WRN 28- 10 (200 epochs)
0.1
0.0005
0.05
0.1
WRN 28- 10 (1800 epochs)
0.05
0.001
0.05
0.1
WRN 26-2x6 ShakeShake
0.02
0.0010
0.02
0.05
Pyramid vanilla
0.05
0.0005
0.05
0.2
Pyramid ShakeDrop (CIFAR-10) 10)
0.02
0.0005
0.05
Pyramid Drop (CIFAR- 100)
0.05
0.0005
0.05
Table 7: Hyper-parameter used to produce the SVHN and MasSion-MNIST results
LR
WD
p
WRN
0.01
0.0005
0.01
SVHN
ShakeShake
0.01
0.0005
0.01
WRN
0.1
0.0005
0.05
Fashion
ShakeShake
0.1
0.0005
0.02
decayed using a cosine schedule. Weight decay is to 0.0001 with SGD optimizer and momentum 0....
Finally, for the noisy label experiments, we also found p by gridsearch, computing the accuracy on (non-noisy) validation set composed of a random subset of 10% of the usual CIFAR training samples. We report the validation accuracy of the bootstrapped version of SAM for different levels of noise and different in table 8
20%
40%
60%
80%
0
15.0%
31.2%
52.3%
73.5%
0.01
13.7%
28.7%
50.1% %
72.9%
0.02
12.8%
27.8%
48.9%
73.1%
0.05
11.6%
25.6%
47.1%
21.0%
0.1
4.6%
6.0%
8.7%
56.1%
0.2
5.3%
7.4%
23.3%
77.1%
0.5
17.6%
40.9%
80.1%
89.9%
Table 8: Validation accuracy of the bootstrapped-SAM for different levels of noise and different, p
C.2 FINETUNING DETAILS
Weights are initialized to the values provided by the publicly available checkpoints, except the last dense layer, which change size to accomodate the new number of classes, that randomly initial- ized. We train all models with weight decay 1e as suggested in (Tan & Le, 2019), but we reduce the learning rate to 0.016 as the models tend diverge for higher values. We use a batch size of 1024 on Google Cloud TPUv3 64 cores and cosine learning rate decay. Because other works train with batch size of 256, we train for 5k steps instead of 20k. We freeze the batch norm statistics and use them for normalization, effectively using the batch norm we would at test time !! We train the models using SGD with momentum 0.9 and cosine learning rate decay. For Efficientnet-L2,, we use this time batch size 512 to save memory and adjusted the number of training steps accord- ingly. For CIFAR, we use the same autoaugment policy as the previous experiments. We do not data augmentation for the other datasets, applying the same preprocessing experiments. We also scale down the learning rate to 0.008 as the batch size is now twice as small. We used Google Cloud TPUv3 128 cores. All other parameters stay the same. For Imagenet, we trained both models from checkpoint for 10 epochs using a learning rate of 0.1 and = 0.05. We do not randomly initialize the last layer as did for the other datasets, but instead use the weights included in the checkpoint.
for the Imagenet
We found anecdotal evidence that this makes the finetuning more robust overtraining.
16

---

Published as a conference paper at ICLR 2021
C.3 EXPERIMENTAL RESULTS WITH p 0.05
A big sensitivity to the choice of hyper-parameters would make method less easy to use. To demonstrate that SAM performs even when is not finely tuned, we compiled the table for the CIFAR and the finetuning experiments using p 0.05. Please note that already used p 0.05
for all Imagenet experiments. We report those scores table and 10.
Cifar 10
CifarI00
Model
Augmentation
p
0.05
SGD
rho=0.05
SGD
WRN-28-10 10 (200 epochs)
Basic
2.7
3.5
16.5
18.8
WRN-28-10 10 (200 epochs)
Cutout
2.3
2.6
14.9
16.9
WRN-28-10 10 (200 epochs)
AA
2.1
2.3
13.6
15.8
WRN-28- (1800 epochs)
Basic
2.4
3.5
16.3
19.1
WRN-28- 10 (1800 epochs)
Cutout
2.1
2.7
14.0
17.4
WRN-28- (1800 epochs)
AA
1.6
2.2
12.8
16.1
WRN 26-2x6
Basic
2.4
2.7
15.1
17.0
WRN 26-2x6 ss
Cutout
2.0
2.3
14.2
15.7
WRN 26-2x6
AA
1.7
1.9
12.8
14.1
PyramidNet
Basic
2.1
4.0
15,4
19.7
PyramidNet
Cutout
1.6
2.5
13.1
16.4
PyramidNet
AA
1.4
1.9
12.1
14.6
PyramidNer+-SakkeDrop
Basic
2.1
2.5
13.3
14.5
PyramidNet+SSakeeDDop
Cutout
1.6
1.9
11.3
11.8
PyramidNet+ShakeDDropp
AA
1.4
1.6
10.3
10.6
Table 9:
Results for the 1ifar10/Cifar100 experiments,
using
0.05 for all mod-
dlsdaataeets/aggmeenaaiioss
Efficientnet-b7
Efficientnet-b7
Dataset
Efficientnet-b7
SAM (optimal)
SAM (p = 0.05)
FirCaAireraft
6.80
7.06
8.15
Flowers
0.63
0.81
1.16
Oxford Pets
3.97
4.15
4.24
Stanford Cars
5.18
5.57
5.94
cifar10
0.88
0.88
0.95
c0far100
7.44
7.56
7.68
Birdsnap
13.64
13.64
14.30
Food101
7.02
7.06
7.17
Table 10: Results for the the finetuning experiments, using 0.05 for all datasets.
C.4 ABLATION OF THE SECOND ORDER TERMS
As described in section 2, computing the gradient of the sharpness aware objective yield some second order terms that more expensive compute. To analyze this ablation more in depth, we trained a Wideresnet-40x2 on CIFAR-10 using SAM with and without discarding the second order terms during training. We report the cosine similarity of the two updates in figure along the training trajectory of both experiments. We also report the training error rate (evaluated at ur + +(ww)) and the test error rate (evaluated at w)).
We observe that during the first half of the training, discarding the second order terms does not im- pact the general direction of the training, as the cosine similarity between the first and second order updates are very close to 1. However, when the model nears convergence, the similarity between both types of updates becomes weaker. Fortunately, the model trained without the second order terms reaches lower test error, showing that the most efficient method is also the providing the best generalization on this example. The reason for this is quite unclear and should be analyzed in follow up work.
17

---

Published as a conference paper at ICLR 2021
0.10
order
1.0
second first
Train error rate
0.08
cos( W2nd
0.8
metric
0.06
train, eerorr rate
0.6
error_rate
0.4
0.04
0.2
0.02
second
0.0
first
0.00
10000
20000
30000
40000
0
10000
20000 step
30000
40000
step
Figure 4: Training and test error for the first and second order version of the algorithm.
Figure 5: Cosine similarity between the first and second order updates.
C.5 CHOICE OF P-NORM
Our theorem is derived for p 2, although generalizations can be considered for p E [oo] (the expression of the bound becoming way more involved). Empirically, we validate that the choice p = is optimal by training a wide resnet on cifarl0 with SAM for p = oo (in which case we have WwLs(w)) We do not consider the case = which would give us a perturbation on a single weight. As an additional ablation study, we also use random weight perturbations of fixed Euclidean norm: E(w) = T
É(w) p sign (V and p 2 (giving É(w) Ls
with z ~ NO, Ia). We report the test accuracy of the model in figure
0.048
0.046
0.044
rat
0.042
0.040
0.088
Constrait telepp
0.086
0.084
tw/p=p
10:4
10.9
10-1
10:3
10:4
10-1
10P
10'
10P
10a
Figure 6: Test accuracy for a wide resnet trained on CIFARIO with SAM, for different perturbation norms.
We observe that adversarial perturbations outperform random perturbations, and that using p 2 yield superior accuracy on this example.
C.6 SEVERAL ITERATIONS IN THE INNER MAXIMIZATION
To empirically verify that the linearization the inner problem is sensible, we trained WideResnet on the CIFAR datasets using variant of SAM that performs several iterations of projected gradient ascent estimate max, I(w++).. c). We report the evolution of max, [[w + [(w)) during training (where L stands for the training error rate computed on the current batch) in Figure 7, along with the test accuracy and the estimated sharpness (max, L(w+es - L(w)) at the end training Table 11; report means and standard deviations across 20 runs.
For most of the training, one projected gradient step (as used standard SAM) sufficient to obtain good approximation of the found with multiple inner maximization steps. We however observe that this approximation becomes weaker near convergence, where doing several iterations of projected gradient ascent yields better (for example, CIFAR-10, the maximum loss found
18

---

Published as a conference paper at ICLR 2021
Cifar10
Citar100
0.05
stepe
max maxL(w++) [[ -L(w)
[() L(w)
2 stope
3 stops steps
0.03
steps
0.04
steps
0.01
0.02
stops
10000
20000
30000
10000
20000
30000
Updales
Updales
Figure 7: Evolution of max, L(w + e) - L(w) vs. training step, for different numbers of inner projected gradient steps.
on each batch is about 3% more when doing 5 steps of inner maximization, compared to when doing a single step). That said, as seen Table 1I, the test accuracy is not strongly affected by the number of inner maximization iterations, though CIFAR-100 it does seem that several steps outperform a single step in a statistically significant way.
Number of projected gradient steps
CIFAR-10
CIFAR- 100
Test error
Estimated sharpness
Test error
Estimated sharpness
2.77 +0.03
0.17
16. 16.7200.00
0.82=0.05
2.7640.03
00.2
16.5940.00
1.83_0.05
3
2.73 10.04
1.49 10.05
16.62 10.09
2.36 _0.03
20.03
2.26
16.60
2.82
Table
Test error rate and estimated sharpness (max, L (w+ + - L(w)) at the end of the training.
19


Reviews:
[{"review_id": "y8CyTyenKg6", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary ##\nThe paper proposes a modified loss function for supervised learning, in which the original loss at w is replaced with a maximum of the loss in a small p-norm ball around w. An approximate way to compute gradients for this loss is presented, and evaluated in high detail on a variety of supervised learning problems where it is shown to consistently improve the overall generalization error. Furthermore, based on PAC-Bayes theory, a generalization bound for learning under that loss is presented. \n\n## Explanation of Rating ## \nThe main strengths of the paper are the simplicity and convincing evaluation of the method. Another strength is the sound theoretical generalization bound. The paper is also very well written, and I therefore clearly recommend acceptance. \n\nPerhaps a weakness of the paper is the lack of a \"broader/high-level perspective\" on the method (see detailed comment #1) and its comparison to variational inference methods which also optimize PAC-Bayes bounds (see detailed comment #2). \n\n## Detailed Comments ##\n\n1. In the past, Gaussian convolution smoothings of loss functions have been considered a lot in the context of homotopy continuation methods or variational inference. To me, the proposed loss function \\max_\\eps L(w + \\eps) - \\delta_{|\\eps| < \\rho} can be seen as a convolution of the loss, but on the tropical semiring (max, +) with the convolution kernel being the indicator function \\delta_{|\\eps| < \\rho}. Such types of convolution have been heavily studied in the field of convex analysis, where they are known under the name infimal convolution / epi-addition. This opens up the question on why this specific convolution kernel has been chosen. The practical results in this paper are quite strong and have been elusive so far for variational inference and Gaussian smoothing methods, even though they have been around for a long time. Do these good results mainly stem from the efficiently implementable algorithm, or are they more due to favorable geometrical properties of the loss due to performing the convolution in the tropical geometry? \n\n2. What is the main advantage of the proposed approach over a variational inference (VI) method (e.g. with mean-field Gaussian approximation) to minimize the PAC-Bayes bound? In variational inference, the free parameter \\rho is also be optimized by minimizing the right-hand side of the PAC-Bayes bound. Have you tried minimizing over \\rho, or perhaps consider a diagonal approximation of \\rho? \n\nMinor comments / typos:\n- p.12 \"Adding a Gaussian perturbation should increase the test error\" -> \"Adding a Gaussian perturbation should not decrease the test error.\"  This is expected to hold in practice at a minimum. But it is not clear if it holds for any w. \n- p.12 Typo: \"Then we KL divergence\" -> \"Then the KL divergence\". \n- The steps (12) and (13) in the proof are not easy to follow. In particular, it was unclear to me what meant with \"each bound\" before Eq. 12, and why the bounds should hold with this specific choice of probability. Furthermore, I couldn't follow what happens from (12) to (13). \n- I did not notice where the assumption ||w|| >= 1 is used, perhaps it can be mentioned in the proof. \n- \\lambda is another hyperparameter, but it is not mentioned later on how it is chosen. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "rAkZKbef7T4", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "- Overview\n\nThis paper proposed a learning algorithm using the idea of flatness.\nBasically, the algorithm was constructed by a first-order Taylor approximation of flatness and using only the main term.\nThe developed algorithm uses the gradient of the loss with steepest direction when updating the parameters.\nThe proposed algorithm is simple and fast to work.\nIt is also easy to incorporate into existing networks and algorithms and has updated accuracy in many models.\n\n- Comments.\n\nThis is an interesting study, as it is the first application of flatness to a practical algorithm.\nIt is also excellent in that it is easy to implement and immediately applicable. The performance is also nice.\n\nThe validity of this approximation is a concern.\nWhen the authors perform a first-order Taylor approximation, they are looking only at the gradient of the parameter to determine the adversarial direction.\nSo, if the loss function is very spiky, there is a concern that the approximation will not work because the local gradient alone will not find the appropriate direction.\nThe nice networks used in their experiments is supposed to use smoothing techniques for the loss surface, such as the batch normalization and the skip connections, but the concern is whether the proposed method works properly in neural networks with a more less smooth loss surface.\n\nTo address this concern, my suggestion is to use the algorithm in without the Taylor approximation make comparisons on a more primitive network with a less smooth loss surface.\nThat way, users would feel more comfortable using the method.", "extra_scores": {}, "extra_text": {}}, {"review_id": "SMQLOQzD-T", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Motivated by the connection between the flatness of minima and its generalization ability, the authors propose Sharpness-aware Minimization (SAM), which explicitly minimizes both loss value and loss sharpness during training deep neural networks. They find SAM improves generalization for a range of image classification tasks and provide robustness to label noise as well. They also introduce a new notion of sharpness named m-sharpness.\n\nStrength: \n* The paper is overall well written with clear motivation.\n* The experiments are comprehensive and the results show clear improvement over non-SAM approaches or previous SOTA.\n\nWeakness:\n* There is no clear definition of the “sharpness“ that the algorithm tries to optimize. Given many existing definitions of the sharpness (e.g., [1]), it is not clear how the proposed measurement connects or differs with previous works. \n\n* My major concern is about the usage of hyperparameter $\\rho$:\n\na) The introduction of the dataset and model dependent hyperparameter $\\rho$ and the need of grid-search before training makes the algorithm more tricky to work and sensitive to other hyperparameters and scale of $w$, e.g., when weight decay is applied, the norm of $w$ usually shrinks during training, and the same radius $\\rho$ could be too large for a small scaled $w$ at the end of training in comparison with the $w$ at the beginning. This discrepancy would become larger when the number of epochs training gets larger. \n\nb) The details for how to obtain the optimal $rho$ is not quite clear, e.g., smaller $\\rho$ in sec 3.3. An ablation study on the sensitivity of $\\rho$ regarding different dataset, model and noisy level would be useful. \n\nc) The wall-clock training time of the SAM method is not discussed. A comprehensive of the cost (including hyperparameter search for $rho$.) would be helpful to have for evaluating the complexity of the method. \n\n\n* The message conveyed in section 4.1 is not quite clear. Does each accelerator perform independent $\\epsilon$ estimation? Is $epsilon$ obtained on each accelerator synchronized after their estimation? Does it indicate the SAM training is done better in model-parallel in small batches rather than data-parallel with large batches? \n\n\nSuggestions:\n1) To avoid the scaling issue of $\\rho$, one suggestion would be considering optimizing the sharpness metric on the normalized loss landscape as described in [2]. In Figure 1, the authors adopt [2] for comparing the landscape of minimas obtained by non-SAM and SAM, so it might be intuitive to optimize this normalized sharpness directly, in which $\\rho$ can be fixed and random direction is sufficient?\n\n\n2) The benefit of flatness to the robustness to label noise is not well discussed. What is the performance when the label noise is over 90% or even 100%. Eventually all models should not generalize given 100% corruption but it would be interesting to know where the limit of SAM is.\n\nMinor:\n* Some figures are not well described, e.g., the meaning of Figure 1 left is not quite clear.  Figure 2 is not intuitive as the loss contour value is not clear. It is not straightforward to know why w_{t+1}^{SAM} is a better or “flatter” move. The notion $w_{adv}$ is also not defined anywhere.\n\n[1] Keskar et al, On large-batch training for deep learning: Generalization gap and sharp minima, ICLR 2017\n[2] Li et al, Visualizing the Loss Landscape of Neural Nets, NIPS 2018\n\n======\nAfter Rebuttal\n\nThanks for the detailed reply and additional experiments. I increased my score accordingly and I hope the authors could further address following issues:\n\n- While the results in C.3 shows default $\\rho$ improves over SGD on most experiments (may also add SVHN and Fashion), I can still see its sensitivity to datasets, architecture, noise level and number of accelerators as shown in Table 6, 7, 8 and Fig. 3. For example, 0.05 is not close to optimal with labe noise 20%~60% in Table 8. It is unclear whether $\\rho$ is robust to other hyperparameter changes (e.g., weight decay that controls weight scales).  So an ablation study on the sensitivity of $\\rho$ and further explanation would be necessary and much valuable for practitioners. \n\n- It would be also helpful if the authors can provide more details about how to get the flat minima of Fig .1 (right) when optimizing deep non-residual networks, such as $\\rho$ and other hyperparameters.\n\n- Minor: Table 8 should be validation errors rather than accuracy.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "CcKY1_krQM3", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "## Summary\n\nThis paper proposes and empirically evaluates SAM, an optimization method that is designed to seek out regions of uniformly low training loss. The method is derived from a bound on the generalization performance of parameters $w$ in terms of the maximal training loss in a region around $w$. After various approximations, minimizing this upper bound gives rise to a simple method, which first performs a (normalized) gradient ascent step; computes the gradient at that perturbed location; and uses that gradient to update the weights. The empirical evaluation shows that SAM improves generalization performance across a wide range of settings.\n\n## Rating\n\nThe paper is well-structured and easy to read. The proposed method is motivated from a PAC-Bayesian generalization bound, but involves approximations which are only vaguely justified. The empirical evaluation is extensive and rigorous. Given its simplicity, scalability, and the convincing results, I believe that this method could have significant impact. With a small exception (see below) all derivations are, to the best of my knowledge, mathematically sound. Overall, this is a good paper and I recommend acceptance, but there are various aspects that could be improved. If these are addressed during the rebuttal, I will consider increasing my score.\n\n## Major Comments\n\n1) The proposed method is motivated by a generalization bound, but this bound is never evaluated numerically. Prior work has shown that many generalization bounds are numerically vacuous in deep learning settings. The bound in Theorem 1 should be really straight-forward to evaluate after the end of training and I think this would be a vital addition to this paper. The paper shows convincing practical benefits of SAM, but it is important to know whether this benefit actually stems from the (approximate) minimization of a non-vacuous bound. This is my single biggest concern with this work.\n\n2) The derivation of the SAM update relies on various approximations to turn the intractable bi-level optimization problem of Eq. (1) into a practical method. My concern here is that the paper does not even attempt to justify these approximations and/or to discuss their limitations and possible pitfalls. I’m not saying that these approximations are not justified; I just think that a bit of discussion would significantly strengthen the paper. A few thoughts:\n    a) The linear approximation that facilitates the closed-form solution of the inner optimization problem (top of page 4) is really drastic. It would be accurate if and only if the curvature in a $\\rho$-sized region around w would be negligibly small. That, however, is exactly the type of region the method wants to find, so this is somewhat circular. Contrasting the first-order with a second-order approximation may be a useful exercise. An interesting ablation would be to solve the inner loop optimization in Eq. (1) more accurately, say, by performing multiple iterations of (projected) gradient descent. (When dropping the second-order term, this wouldn’t be exceedingly costly and could provide further insight.)\n    b) Later on, the “second order” term, which derives through the solution $\\hat{\\epsilon}$ of the inner-loop optimization, is dropped. The authors write that this is done to “further accelerate the computation” and “does not adversely affect model training and in fact improves numerical stability”. However, as shown in Appendix C.3, dropping the second-order term yields a quite considerable *improvement*. It always makes me a bit nervous, when a (relatively crude) approximation works better than the quantity you actually want to compute. I don’t think it is adequate to sweep this under the rug with a vague reference to numerical stability. At least, I thank that (i) the fact should be stated clearly in the main text, and (ii) the reference to numerical stability should be marked as speculation, unless there’s evidence to back it up.\n\n3) Even if the approximations involved are only vaguely justified, the SAM update itself makes a lot of intuitive sense, and I think the reader might benefit from discussing that in a bit more detail. For example, one simple observation would be that, for a sharp minimum whose basin of attraction is of radius smaller than rho, the ascent step of SAM (by virtue of the normalisation to length rho) would take you outside of that basin. If the basin of attraction is larger than rho, the SAM update stays inside. Maybe this intuition could be formalized?\n\n4) Connection to Entropy-SGD:\n    a) Is it really fair to characterize EntropySGD as “only suitable to small models and datasets”? Of course, it requires a number L of additional gradient steps to obtain an MC-estimate of the gradient of the local entropy. But, if my understanding is correct, these could be parallelized. The original paper uses L=20 but, to my knowledge, nobody has tested the (lower) limits of that number. If L<=5 does something useful, it wouldn’t be too outlandish in terms of computational cost.\n    b) On that note: It would have been great to see an empirical comparison to EntropySGD for at least one of the experiments. \n    c) Looking more closely at the proof of the generalization bound, the term involving maximization over a local region actually originates from a bound on the *integral* over that region, which makes the connection to EntropySGD even closer. Is it fair to say that the maximization step in SAM is actually just a clever and cheap surrogate for integrating the gradient the over the local region as done in EntropySGD?\n    d) An interesting piece of related work that you might want to add is [1], which shows that EntropySGD (in a specific sense) also optimizes a generalization bound.\n\n5) You clearly show the benefits of SAM in pushing the state-of-the-art, which is fantastic! Personally, I would have also valued experiments that are more illustrative and maybe test the limits of the method:\n    a) What does SAM do in a convex setting, say, a logistic regression on some (maybe synthetic) dataset that is easily overfit? Does it converge to the unique minimizer of the training loss? Does it plateau before reaching that?\n    b) What happens if you apply SAM in the setting of [2] with random labels assigned to images? Will it converge to a point that closely fits each (nonsense) training sample?\n    c) With regards to the finding on m-sharpness, it would be great to see a small-scale experiments with full-batch optimization.\n\n6) The finding on m-sharpness is intriguing, though somewhat counter to the motivation for SAM laid out in the paper. The results in Figure 3 make me wonder, whether SAM would even gives practical improvements with large m? This is an important question, but I totally agree with the authors decision to leave it for future work.\n\n## Minor Comments\n\n7) SAM is motivated from the generalization bound in Theorem 1. The bound includes a term $h(||w||^2 / \\rho^2)$, which is later replaced with a simple L2 regularization term with coefficient lambda. From the Theorem, that coefficient $\\lambda$ should be coupled with the neighborhood size $\\rho$ in an inversely proportional manner. Of course, for a fixed choice of $\\rho$, this can be subsumed into the regularization parameter lambda. However, it would make a difference when comparing various choices for rho, as is done in the experiments. Did you keep the regularization parameter lambda constant in those experiments or did it vary with rho, as suggested by the Theorem? Furthermore, it would be interesting to know to what extent SAM depends on that L2 regularization. Did all experiments use L2 regularization and did you, by any chance, test SAM’s performance without the regularization parameter?\n\n8) The bib file could really need some love:\n    a) You cite arXiv versions of of multiple papers that have been published at peer-reviewed venues.\n    b) Capitalize properly: “nesterov momentum”, “pac-bayesian”, “journal of machine learning research”, etc…\n    c) Cite identical venues consistently.\n\n9) An interesting connection (or non-connection) to discuss in the paper would be adversarial training. Some of these methods perform an adversarial perturbation step to the *data*, followed by a weight update computed on that perturbed data. This is an interesting parallel to SAM, where the first step is sort of an adversarial perturbation of the weights.\n\n10) If I am not mistaken, there is a minor flaw in the proof of Theorem 1. The very last equation gives a high-probability bound on the norm of $||\\epsilon||^2$, which is then plugged into the high-probability generalization bound of Eq. (). Both bounds hold with probability 1-delta individually, but that does not imply that the combined bound holds with probability 1-delta. However, that should easily be fixed with a slight adaptation of the constants.\n\n## Typos / Style\n\n- At least according to some style guides, I think that you should write “Sharpness-Aware Minimization” instead of “Sharpness-aware Minimization”.\n- Capitalize references to sections, equations, figures: “in Appendix C.3”, “using Equation 3”, et cetera…\n- Typo “opefrations” near the bottom of page 4\n- Footnote markers that refer to an entire sentence go after the punctuation mark (e.g., for footnotes 3 and 5).\n- In Table 3, the boldface is assigned to the wrong method in the Stanford Cars row.\n\n## References\n\n[1] Dziugaite, G. K., & Roy, D. (2018). Entropy-SGD optimizes the prior of a PAC-Bayes bound: Generalization properties of Entropy-SGD and data-dependent priors. In International Conference on Machine Learning (pp. 1377-1386). PMLR.\n\n[2] Zhang, C., Bengio, S., Hardt, M., Recht, B., & Vinyals, O. (2016). Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530.\n\n## Update after Rebuttal\n\nThanks for the detailed replies to my questions and comments. I think the paper has been improved substantially and I have increased my rating. Congratulations on the good work!", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes to minimize the loss while regularizing its sharpness: so that the minimum will lie in a region with uniformly low loss.
The reviewers uniformly appreciated the paper. They have made a number of suggestion for improving the paper, which the authors should consider incorporating in their final version.


Decision: Accept (Spotlight)