S HARPNESS -A WARE M INIMIZATION FOR E FFICIENTLY I MPROVING G ENERALIZATION
A BSTRACT
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 prior work connecting the geometry of the loss landscape and generalization, 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 minmax optimization problem on which gradient descent can be performed efﬁciently.
We present empirical results showing that SAM improves model generalization across a variety of benchmark datasets (e.g., CIFAR- { 10, 100 } , ImageNet, ﬁnetuning tasks) and models, yielding novel state-of-the-art performance for several.
Additionally, we ﬁnd that SAM natively provides robustness to label noise on par with that provided by state-of-the-art procedures that speciﬁcally target learning with noisy labels. We open source our code at https: //github.com/google-research/sam .
I NTRODUCTION
Modern machine learning’s success in achieving ever better performance on a wide range of tasks has relied in signiﬁcant part on ever heavier overparameterization, in conjunction with developing ever more effective training algorithms that are able to ﬁnd parameters that generalize well.
Indeed, many modern neural networks can easily memorize the training data and have the capacity to readily overﬁt (Zhang et al., 2016).
Such heavy overparameterization is currently required to achieve stateof-the-art results in a variety of domains (Tan & Le, 2019; Kolesnikov et al., 2020; Huang et al., 2018).
In turn, it is 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 training set is typically not sufﬁcient 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 an 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), [Figure]
The connection between the geometry of the loss landscape—in particular, the ﬂatness of minima— and generalization has been studied extensively from both theoretical and empirical perspectives (Shirish Keskar et al., 2016; Dziugaite & Roy, 2017; Jiang et al., 2019).
While this connection has held the promise of enabling new approaches to model training that yield better generalization, practical efﬁcient algorithms that speciﬁcally seek out ﬂatter minima and furthermore effectively improve generalization on a range of state-of-the-art models have thus far been elusive (e.g., see (Chaudhari et al., 2016; Izmailov et al., 2018); we include a more detailed discussion of prior work in Section 5).
We present here a new efﬁcient, scalable, and effective approach to improving model generalization ability that directly leverages the geometry of the loss landscape and its connection to generalization, and is powerfully complementary to existing techniques.
In particular, we make the following contributions:
• We introduce Sharpness-Aware Minimization (SAM), a 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 efﬁciently and easily.
• We show via a rigorous empirical study that using SAM improves model generalization ability across a range of widely studied computer vision tasks (e.g., CIFAR- { 10, 100 } , ImageNet, ﬁnetuning tasks) and models, as summarized in the lefthand plot of Figure 1.
For example, applying SAM yields novel state-of-the-art performance for a number of alreadyintensely-studied tasks, such as ImageNet, CIFAR- { 10, 100 } , SVHN, Fashion-MNIST, and the standard set of image classiﬁcation ﬁnetuning tasks (e.g., Flowers, Stanford Cars, Oxford Pets, etc).
• We show that SAM furthermore provides robustness to label noise on par with that provided by state-of-the-art procedures that speciﬁcally target learning with noisy labels.
• Through the lens provided by SAM, we further elucidate the connection between loss sharpness and generalization by surfacing a promising new notion of sharpness, which we term m-sharpness .
Section 2 below derives the SAM procedure and presents the resulting algorithm in full detail.
Section 3 evaluates SAM empirically, and Section 4 further analyzes the connection between loss sharpness and generalization through the lens of SAM.
Finally, we conclude with an overview of related work and a discussion of conclusions and future work in Sections 5 and 6, respectively.
S HARPNESS -A WARE M INIMIZATION (SAM)
Throughout the paper, we denote scalars as a , vectors as a , matrices as A , sets as A , and equality by deﬁnition as (cid:44) .
Given a training dataset S (cid:44) ∪ ni =1 { ( x i , y i ) } drawn i.i.d.
from distribution D , we seek to learn a model that generalizes well.
In particular, consider a family of models parameterized by w ∈ W ⊆ R d ; given a per-data-point loss function l : W ×X ×Y → R + , we deﬁne the training set loss L S ( w ) (cid:44) 1 n (cid:80) ni
=1 l ( w , x i , y i ) and the population loss L D ( w ) (cid:44) E ( x , y ) ∼ D
[ l ( w , x , y )] .
Having observed only S , the goal of model training is to select model parameters w having low population loss L D ( w ) .
Utilizing L S ( w ) as an estimate of L D ( w ) motivates the standard approach of selecting parameters w by solving min w L S ( w ) (possibly in conjunction with a regularizer on w ) using an optimization procedure such as SGD or Adam.
Unfortunately, however, for modern overparameterized models such as deep neural networks, typical optimization approaches can easily result in suboptimal performance at test time.
In particular, for modern models, L S ( w ) is typically non-convex in w , with multiple local and even global minima that may yield similar values of L S ( w ) while having signiﬁcantly different generalization performance (i.e., signiﬁcantly different values of L D ( 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 L S ( w ) , we seek out parameter values whose entire neighborhoods have uniformly low training loss value (equivalently, 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-wise training loss (full theorem statement and proof in Appendix A):
Theorem (stated informally) 1.
For any ρ > 0 , with high probability over training set S generated from distribution D ,
[Equation]
where h : R + → R + is a strictly increasing function (under some technical conditions on L D ( w ) ).
To make explicit our sharpness term, we can rewrite the right hand side of the inequality above as
[Equation]
The term in square brackets captures the sharpness of L S at w by measuring how quickly the training loss can be increased by moving from w to a 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 speciﬁc function h is heavily inﬂuenced by the details of the proof, we substitute the second term with λ || w || 22 for a hyperparameter λ , yielding a standard L2 regularization term.
Thus, inspired by the terms from the bound, we propose to select parameter values by solving the following SharpnessAware Minimization (SAM) problem:
[Equation]
where
[Equation]
where ρ ≥ 0 is a hyperparameter and p ∈ [1 , ∞ ] (we have generalized slightly from an L2-norm to a p -norm in the maximization over (cid:15) , though we show empirically in appendix C.5 that p = 2 is typically optimal).
Figure 1 shows 1 the loss landscape for a model that converged to minima found by minimizing either L S ( w ) or L SAM S ( w ) , illustrating that the sharpness-aware loss prevents the model from converging to a sharp minimum.
In order to minimize L SAM S ( w ) , we derive an efﬁcient and effective approximation to ∇ w L SAM S ( w ) by differentiating through the inner maximization, which in turn enables us to apply stochastic gradient descent directly to the SAM objective.
Proceeding down this path, we ﬁrst approximate the inner maximization problem via a ﬁrst-order Taylor expansion of L S ( w + (cid:15) )
w.r.t.
(cid:15) around 0 , obtaining
[Equation]
In turn, the value ˆ (cid:15) ( w ) that solves this approximation is given by the solution to a classical dual norm problem ( | · | q − 1 denotes elementwise absolute value and power) 2 :
[Equation]
where 1 /p + 1 /q = 1 .
Substituting back into equation (1) and differentiating, we then have
[Equation]
This approximation to ∇ w L SAM S ( w ) can be straightforwardly computed via automatic differentiation, as implemented in common libraries such as JAX, TensorFlow, and PyTorch.
Though this computation implicitly depends on the Hessian of L S ( w ) because ˆ (cid:15) ( w ) is itself a function of ∇ w L S ( w ) , the Hessian enters only via Hessian-vector products, which can be computed tractably without materializing the Hessian matrix.
Nonetheless, to further accelerate the computation, we drop the second-order terms.
obtaining our ﬁnal gradient approximation:
[Equation]
As shown by the results in Section 3, this approximation (without the second-order terms) yields an effective algorithm.
In Appendix C.4, we 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 a priority in future work.
We obtain the ﬁnal SAM algorithm by applying a standard numerical optimizer such as stochastic gradient descent (SGD) to the SAM objective L SAM S ( w ) , using equation 3 to compute the requisite objective function gradients.
Algorithm 1 gives pseudo-code for the full SAM algorithm, using SGD as the base optimizer, and Figure 2 schematically illustrates a single SAM parameter update.
Input: Training set S (cid:44) ∪ ni =1 { ( x i , y i ) } , Loss function l : W × X × Y → R + , Batch size b , Step size η > 0 , Neighborhood size ρ > 0 .
Output: Model trained with SAM
Initialize weights w 0 , t = 0 ; while not converged do
Sample batch B = { ( x 1 , y 1 ) , ... ( x b , y b ) } ; Compute gradient ∇ w L B ( w ) of the batch’s training loss; Compute ˆ (cid:15) ( w ) per equation 2;
Compute gradient approximation for the SAM objective (equation 3): g = ∇ w L B ( w ) | w +ˆ (cid:15) ( w ) ;
Update weights: w t +1 =
w t − η g ; t = t + 1 ; endreturn w t
Algorithm 1: SAM algorithm w t w t + w SAMt + w adv
L ( w t ) || L ( w t )||
2 L ( w t )
L ( w adv )
E MPIRICAL E VALUATION
In order to assess SAM’s efﬁcacy, we apply it to a range of different tasks, including image classiﬁcation from scratch (including on CIFAR-10, CIFAR-100, and ImageNet), ﬁnetuning pretrained models, and learning with noisy labels.
In all cases, we measure the beneﬁt of using SAM by simply replacing the optimization procedure used to 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.
3. I MAGE C LASSIFICATION F ROM S CRATCH
We ﬁrst evaluate SAM’s impact on generalization for today’s state-of-the-art models on CIFARand 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 overﬁtting; therefore, signiﬁcantly improving their generalization is quite non-trivial.
We have ensured that our implementations’ 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 ﬂip, padding by four pixels, and random crop).
We also evaluate in the setting of more advanced data augmentation methods such as cutout regularization (Devries & Taylor, 2017) and AutoAugment (Cubuk et al., 2018), which are utilized by prior work to achieve state-of-the-art results.
SAM has a single hyperparameter ρ (the neighborhood size), which we tune via a grid search over { 0 . 01 , 0 . 02 , 0 .
5 } using 10% of the training set as a validation set 3 .
Please see appendix C.1 for the values of all hyperparameters and additional training details.
As each SAM weight update requires two backpropagation operations (one to compute ˆ (cid:15) ( w ) and another to compute the ﬁnal 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 4 .
We run ﬁve independent replicas of each experimental condition for which we report results (each with independent weight initialization and data shufﬂing), reporting the resulting mean error (or accuracy) on the test set, and the associated 95% conﬁdence interval.
Our implementations utilize JAX (Bradbury et al., 2018), and we train all models on a single host having 8 Nvidia V100 GPUs 5 .
To compute the SAM update when parallelizing across multiple accelerators, we divide each data batch evenly among the accelerators, independently compute the SAM gradient on each accelerator, and average the resulting sub-batch
SAM gradients to obtain the ﬁnal SAM update.
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 to 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 easily-implemented, model-independent alternative.
Furthermore, SAM delivers improvements even when applied atop complex architectures that already use sophisticated regularization: for instance, applying SAM to a PyramidNet with ShakeDrop regularization yields 10.3% error on CIFAR-100, which is, to our knowledge, a new state-of-the-art on this dataset without the use of additional data.
Beyond CIFAR- { 10, 100 } , we have also evaluated SAM on the SVHN (Netzer et al., 2011) and
Fashion-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-MNIST.
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 of 0.1, and weight decay 0.0001.
When applying SAM, we use ρ = 0 .
(determined via a 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% conﬁdence interval across 5 independent runs). [Table]
As seen in Table
[Table]
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 overﬁtting.
In contrast [Table] the standard training procedure (without SAM) generally signiﬁcantly overﬁts as training extends from 200 to 400 epochs.
[Table]
3. F INETUNING
Transfer learning by pretraining a model on a large related dataset and then ﬁnetuning on a smaller target dataset of interest has emerged as a powerful and widely used technique for producing highquality models for a variety of different tasks.
We show here that SAM once again offers considerable beneﬁts in this setting, even when ﬁnetuning extremely large, state-of-the-art, already high-performing models.
In particular, we apply SAM to ﬁnetuning EfﬁcentNet-b7 (pretrained on ImageNet) and
EfﬁcientNet-L2 (pretrained on ImageNet plus unlabeled JFT; input resolution 475)
(Tan & Le, 2019;
Kornblith et al., 2018; Huang et al., 2018).
We initialize these models to publicly available checkpoints 6 trained with RandAugment (84.7% accuracy on ImageNet) and NoisyStudent (88.2% accuracy on ImageNet), respectively.
We ﬁnetune 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% conﬁdence interval of top-1 test error over 5 independent runs for each dataset.
As seen in Table [Table] SAM uniformly improves performance relative to ﬁnetuning without SAM. Furthermore
[Table] in many cases [Table] SAM yields novel state-of-the-art performance, including 0.30% error on CIFAR-10, 3.92% error on CIFAR-100, and 11.39% error on ImageNet
[Table]
3. R OBUSTNESS TO L ABEL N OISE
[Table]
The fact that SAM seeks out model parameters that are robust to perturbations suggests SAM’s potential to provide robustness to 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.
In particular, we measure the effect of applying SAM in the classical noisy-label setting for
CIFAR-10, in which a fraction of the training set’s labels are randomly ﬂipped; the test set remains unmodiﬁed (i.e., clean).
To ensure valid comparison to prior work, which often utilizes architectures specialized to the noisy-label setting, we train a simple model of similar size (ResNet-32) forepochs, following Jiang et al.
We evaluate ﬁve 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 ﬁrst trained as usual and then retrained from scratch on the labels predicted by the initially trained model).
When applying SAM, we use ρ = 0 .
1 for all noise levels except 80%, for which we use ρ = 0 .
for more stable convergence.
For the Mixup baselines, we tried all values of α ∈ { 1 , 8 , 16 , 32 } and conservatively report the best score for each noise level.
As seen in Table 4, SAM provides a high degree of robustness to label noise, on par with that provided by state-of-the art procedures that speciﬁcally target learning with noisy labels.
Indeed, simply training a model with SAM outperforms all prior methods speciﬁcally 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). [Figure]
S HARPNESS AND G ENERALIZATION T HROUGH THE L ENS OF SAM 4. m - SHARPNESS
Though our derivation of SAM deﬁnes 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 a subset of size m of a batch, as described in Section 3).
This latter setting is equivalent to modifying the SAM objective (equation 1) to sum over a set of independent (cid:15) maximizations, each performed on a sum of per-data-point losses on a disjoint subset of m data points, rather than performing the (cid:15) 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 m on SAM, we train a small ResNet on CIFAR-10 using SAM with a range of values of m .
As seen in Figure 3 (middle), smaller values of 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 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) 7 .
In particular, this implies that m -sharpness with m < n yields a better predictor of generalization than the full-training-set measure suggested by Theorem 1 in Section 2 above, suggesting an interesting new avenue of future work for understanding generalization.
4. H ESSIAN S PECTRA
Motivated by the connection between geometry of the loss landscape and generalization, we constructed SAM to seek out minima of the training loss landscape having both low loss value and low curvature (i.e., low sharpness).
To further conﬁrm that SAM does in fact ﬁnd minima having low curvature, we compute the spectrum of the Hessian for a WideResNet40-10 trained on CIFARfor 300 steps both with and without SAM (without batch norm, which tends to obscure interpretation of the Hessian), at different epochs during training.
Due to the parameter space’s dimensionality, we approximate the Hessian spectrum using the Lanczos algorithm of Ghorbani et al.
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 maximum eigenvalue ( λ max ) at convergence (approximately 24 without SAM, 1.0 with SAM), and the bulk of the spectrum (the ratio λ max /λ 5 , commonly used as a proxy for sharpness (Jastrzebski et al., 2020); up to 11.4 without SAM, and 2.6 with SAM).
R ELATED W ORK
The idea of searching for “ﬂat” minima can be traced back to Hochreiter & Schmidhuber (1995), and its connection to generalization has seen signiﬁcant study (Shirish Keskar et al., 2016; Dziugaite &
Roy, 2017; Neyshabur et al., 2017; Dinh et al., 2017).
In a 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 ﬁrst paper on penalizing the sharpness, regularizing a notion related to Minimum 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., 2016).
Another direction is to not penalize the sharpness explicitly, but rather average weights during training; Izmailov et al.
(2018) showed that doing so can yield ﬂatter minima that can also generalize better.
However, the measures of sharpness proposed previously are difﬁcult to compute and differentiate through.
In contrast, SAM is highly scalable as it only needs two gradient computations per iteration.
The concurrent work of Sun et al.
(2020) focuses on resilience to random and adversarial corruption to expose a model’s vulnerabilities; this work is perhaps closest to ours.
Our work has a different basis: we develop SAM motivated by a principled starting point in generalization, clearly demonstrate SAM’s
efﬁcacy via rigorous large-scale empirical evaluation, and surface important practical and theoretical facets of the procedure (e.g., m -sharpness).
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.
D ISCUSSION AND F UTURE W ORK
In this work, we have introduced SAM, a novel algorithm that improves generalization by simultaneously minimizing loss value and loss sharpness; we have demonstrated SAM’s efﬁcacy through a rigorous large-scale empirical evaluation.
We have surfaced a number of interesting avenues for future 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 in the past) suggests an interesting new lens through which to study generalization.
Methodologically, our results suggest that SAM could potentially be used in place of Mixup in robust or semi-supervised methods that currently rely on Mixup (giving, for instance, MentorSAM).
We leave to future work a more in-depth investigation of these possibilities.
A CKNOWLEDGMENTS
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 in replicating EfﬁcientNet using JAX, and Justin Gilmer for his implementation of the Lanczos algorithm 8 used to generate the Hessian spectra.
We thank Niru Maheswaranathan for his matplotlib mastery.
We also thank David Samuel for providing a PyTorch implementation of
SAM 9 .
A A PPENDIX A. PAC B AYESIAN G ENERALIZATION B OUND
Below, we state a generalization bound based on sharpness.
Theorem 2.
For any ρ > 0 and any distribution D with probability 1 − δ over the choice of the training set S ∼ D ,
[Equation]
where n = |S| , k is the number of parameters and we assumed L D ( w ) ≤ E (cid:15) i ∼N (0 ,ρ )
[ L D ( w + (cid:15) )] .
The condition L D ( w ) ≤ E (cid:15) i ∼N (0 ,ρ ) [ L D ( w + (cid:15) )] means that adding Gaussian perturbation should not decrease the test error.
This is expected to hold in practice for the ﬁnal solution but does not necessarily hold for any w .
Proof.
First, note that the right hand side of the bound in the theorem statement is lower bounded by (cid:112) k
log(1 + (cid:107) w (cid:107) 22 /ρ 2 ) / (4 n ) which is greater than 1 when (cid:107) w (cid:107) 22 >
ρ 2 (exp(4 n/k ) − 1) .
In that case, the right hand side becomes greater than 1 in which case the inequality holds trivially.
Therefore, in the rest of the proof, we only consider the case when (cid:107) w (cid:107) 22 ≤ ρ 2 (exp(4 n/k ) − 1) .
The proof technique we use here is inspired from Chatterji et al.
Using PAC-Bayesian generalization bound McAllester (1999) and following Dziugaite & Roy (2017), the following generalization bound holds for any prior P over parameters with probability 1 − δ over the choice of the training set S , for any posterior Q over parameters:
[Equation]
Moreover, if P = N ( µ P , σ 2 P I ) and Q = N ( µ Q , σ 2 Q I ) , then the KL divergence can be written as follows:
[Equation]
Given a posterior standard deviation σ Q , one could choose a prior standard deviation σ P to minimize the above KL divergence and hence the generalization bound by taking the derivative 10 of the above
KL with respect to σ P and setting it to zero.
We would then have σ ∗ P 2 =
σ 2 Q + (cid:107) µ P − µ Q (cid:107) 22 /k .
However, since σ P should be chosen before observing the training data S and µ Q , σ Q could depend on S , we are not allowed to optimize σ P in this way.
Instead, one can have a set of predeﬁned values for σ P and pick the best one in that set.
See Langford & Caruana (2002) for the discussion around this technique.
Given ﬁxed a, b > 0 , let T = { c exp((1 − j ) /k )
| j ∈ N } be that predeﬁned set of values for σ 2 P .
If for any j ∈ N , the above PAC-Bayesian bound holds for σ 2 P = c exp((1 − j ) /k ) with probability 1 − δ j with δ j = 6 δ π 2 j 2 , then by the union bound, all above bounds hold simultaneously with probability at least 1 − (cid:80) ∞
j =1 6 δ π 2 j 2 = 1 − δ .
Let σ Q = ρ , µ Q = w and µ P = 0 .
Therefore, we have:
[Equation]
We now consider the bound that corresponds to j = (cid:98) 1 − k log [Equation] + [Equation]
w [Equation] 22 /k ) /c ) (cid:99) .
We can ensure that j ∈ N using inequality equation 7 and by setting c = ρ 2 (1 + exp(4 n/k )) .
Furthermore, for σ 2 P = c exp((1 − j ) /k ) , we have:
[Equation]
Therefore, using the above value for σ P
[Equation]
KL divergence can be bounded as follows:
[Equation]
Given the bound that corresponds to j holds with probability 1 − δ j for δ j = 6 δ π 2 j 2 , the log term in the bound can be written as:
[Equation]
Therefore, the generalization bound can be written as follows:
[Equation]
In the above bound, we have (cid:15) i ∼ N (0 , σ ) .
Therefore, (cid:107) (cid:15) (cid:107) 22 has chi-square distribution and by
Lemma 1 in Laurent & Massart (2000), we have that for any positive t :
[Equation] [Equation]
Therefore, with probability 1 − 1 / √ n we have that:
[Equation]
Substituting the above value for σ back to the inequality and using theorem’s assumption gives us following inequality:
[Equation]
B A DDITIONAL E XPERIMENTAL R ESULTS B. SVHN AND F ASHION -MNIST
We report in table 5 results obtained on SVHN and Fashion-MNIST datasets.
On these datasets,
SAM allows a simple WideResnet to reach or push state-of-the-art accuracy (0.99% error rate for
SVHN, 3.59% for Fashion-MNIST).
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 on this dataset as described in (Cubuk et al., 2018) plus cutout (Devries & Taylor, 2017).
For Fashion-MNIST, the auto-augmentation line correspond to cutout only.
[Table]
C E XPERIMENT D ETAILS C. H YPERPARAMETERS FOR E XPERIMENTS
We report in table 6 the hyper-parameters selected by gridsearch for the CIFAR experiments, and the ones for SVHN and Fashion-MNIST in 7.
For CIFAR10, CIFAR100, SVHN and Fashion-MNIST, we use a 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 a batch size of 4096.
The initial learning rate is set to 1.0 and
[Table] [Table]
decayed using a cosine schedule.
Weight decay is set to 0.0001 with SGD optimizer and momentum = 0.9.
Finally
[Table]
for the noisy label experiments
[Table]
we also found ρ by gridsearch
[Table]
computing the accuracy on a (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
[Table]
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 is randomly initialized.
We train all models with weight decay 1 e − 5 as suggested in (Tan & Le, 2019), but we reduce the learning rate to 0.016 as the models tend to 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 as we would at test time 11 .
We train the models using SGD with momentum 0.9 and cosine learning rate decay.
For Efﬁcientnet-L2, we use this time a batch size 512 to save memory and adjusted the number of training steps accordingly. For CIFAR, we use the same autoaugment policy as in the previous experiments. We do not use data augmentation for the other datasets, applying the same preprocessing as for the Imagenet 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 . We do not randomly initialize the last layer as we did for the other datasets, but instead use the weights included in the checkpoint.
C. E XPERIMENTAL RESULTS WITH ρ = 0 .
A big sensitivity to the choice of hyper-parameters would make a method less easy to use.
To demonstrate that SAM performs even when ρ is not ﬁnely tuned, we compiled the table for the CIFAR and the ﬁnetuning experiments using ρ = 0 .
Please note that we already used ρ = 0 .for
all Imagenet experiments.
We report those scores in table 9 and
[Table] [Table]
[Table]
C. A BLATION OF THE S ECOND O RDER T ERMS
As described in section [Table] computing the gradient of the sharpness aware objective yield some second order terms that are more expensive to 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 ﬁgure [Table] along the training trajectory of both experiments.
We also report the training error rate (evaluated at w +ˆ (cid:15) ( w ) ) and the test error rate (evaluated at w ).
We observe that during the ﬁrst half of the training, discarding the second order terms does not impact the general direction of the training, as the cosine similarity between the ﬁrst 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 a lower test error, showing that the most efﬁcient method is also the one providing the best generalization on this example.
The reason for this is quite unclear and should be analyzed in follow up work. [Figure] [Figure]
C. C HOICE OF P - NORM
Our theorem is derived for p = 2 , although generalizations can be considered for p ∈ [ [Figure] + ∞ ] (the expression of the bound becoming way more involved).
Empirically [Figure] we validate that the choice p = 2 is optimal by training a wide resnet on cifar10 with SAM for p = ∞ (in which case we have ˆ (cid:15) ( w )
=
ρ sign ( ∇ w L S ( w )) ) and p = 2 (giving ˆ (cid:15) ( w ) =
ρ ||∇ w L S ( w ) || 22 ( ∇ w L S ( w )) ).
We do not consider the case p = 1 which would give us a perturbation on a single weight.
As an additional ablation study [Figure] we also use random weight perturbations of a ﬁxed Euclidean norm: ˆ (cid:15) ( w )
=
ρ ||
z
|| 22 z with z ∼ N ( 0 , I d ) .
We report the test accuracy of the model in ﬁgure
[Figure] We observe that adversarial perturbations outperform random perturbations, and that using p =yield superior accuracy on this example.
C. S EVERAL I TERATIONS IN THE I NNER M AXIMIZATION
To empirically verify that the linearization of the inner problem is sensible, we trained a WideResnet on the CIFAR datasets using a variant of SAM that performs several iterations of projected gradient ascent to estimate max (cid:15) L ( w + (cid:15) ) .
We report the evolution of max (cid:15) L ( w + (cid:15) )
− L ( 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 (cid:15) L ( w + (cid:15) )
− L ( w ) ) at the end of training in Table 11; we report means and standard deviations across 20 runs.
For most of the training, one projected gradient step (as used in standard SAM) is sufﬁcient to obtain a good approximation of the (cid:15) 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 a better (cid:15) (for example, on CIFAR-10, the maximum loss found
[Figure]
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 in Table 11, the test accuracy is not strongly affected by the number of inner maximization iterations, though on CIFAR-100 it does seem that several steps outperform a single step in a statistically signiﬁcant way. Number of projected gradient steps CIFARCIFAR
[Table]
[Table]