Sharpness-aware Minimization for Efficiently Improving Generalization

1 INTRODUCTION . Modern machine learning ’ s success in achieving ever better performance on a 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 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 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 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 ) , ∗Work done as part of the Google AI Residency program . batch normalization ( Ioffe & Szegedy , 2015 ) , stochastic depth ( Huang et al. , 2016 ) , data augmentation ( 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 et al. , 2019 ) . While this connection has held the promise of 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 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 efficient , 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 efficiently 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 , finetuning 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 classification finetuning 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 specifically 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 . 2 SHARPNESS-AWARE MINIMIZATION ( SAM ) . Throughout the paper , we denote scalars as a , vectors as a , matrices asA , sets asA , and equality by definition as , . Given a training dataset S , ∪ni=1 { ( xi , yi ) } drawn i.i.d . from distribution D , we seek to learn a model that generalizes well . In particular , consider a family of models parameterized byw ∈ W ⊆ Rd ; given a per-data-point loss function l : W×X ×Y → R+ , we define the training set loss LS ( w ) , 1n ∑n i=1 l ( w , xi , yi ) and the population loss LD ( w ) , 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 LD ( w ) . Utilizing LS ( w ) as an estimate of LD ( w ) motivates the standard approach of selecting parameters w by solving minw LS ( w ) ( possibly in conjunction with a regularizer onw ) 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 , LS ( w ) is typically non-convex in 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 of LD ( 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 valueLS ( 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 , LD ( w ) ≤ max ‖ ‖2≤ρ LS ( w + ) + h ( ‖w‖22/ρ2 ) , where h : R+ → R+ is a strictly increasing function ( under some technical conditions on LD ( w ) ) . To make explicit our sharpness term , we can rewrite the right hand side of the inequality above as [ max ‖ ‖2≤ρ LS ( w + ) − LS ( w ) ] + LS ( w ) + h ( ‖w‖22/ρ2 ) . The term in square brackets captures the sharpness ofLS atw 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 specific function h is heavily influenced 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 : min w LSAMS ( w ) + λ||w||22 where LSAMS ( w ) , max|| ||p≤ρ LS ( w + ) , ( 1 ) where ρ ≥ 0 is a hyperparameter and p ∈ [ 1 , ∞ ] ( we have generalized slightly from an L2-norm to a p-norm in the maximization over , though we show empirically in appendix C.5 that p = 2 is typically optimal ) . Figure 1 shows1 the loss landscape for a model that converged to minima found by minimizing either LS ( w ) or LSAMS ( w ) , illustrating that the sharpness-aware loss prevents the model from converging to a sharp minimum . In order to minimize LSAMS ( w ) , we derive an efficient and effective approximation to ∇wLSAMS ( 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 first approximate the inner maximization problem via a first-order Taylor expansion of LS ( w + ) w.r.t . around 0 , obtaining ∗ ( w ) , argmax ‖ ‖p≤ρ LS ( w + ) ≈ argmax ‖ ‖p≤ρ LS ( w ) + T∇wLS ( w ) = argmax ‖ ‖p≤ρ T∇wLS ( w ) . 1Figure 1 was generated following Li et al . ( 2017 ) with the provided ResNet56 ( no residual connections ) checkpoint , and training the same model with SAM . In turn , the value ̂ ( 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 : ̂ ( w ) = ρ sign ( ∇wLS ( w ) ) |∇wLS ( w ) |q−1 / ( ‖∇wLS ( w ) ‖qq ) 1/p ( 2 ) where 1/p+ 1/q = 1 . Substituting back into equation ( 1 ) and differentiating , we then have ∇wLSAMS ( w ) ≈ ∇wLS ( w + ̂ ( w ) ) = d ( w + ̂ ( w ) ) dw ∇wLS ( w ) |w+̂ ( w ) = ∇wLS ( w ) |w+̂ ( w ) + d̂ ( w ) dw ∇wLS ( w ) |w+̂ ( w ) . This approximation to ∇wLSAMS ( 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 ofLS ( w ) because ̂ ( w ) is itself a function of∇wLS ( 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 final gradient approximation : ∇wLSAMS ( w ) ≈ ∇wLS ( w ) |w+̂ ( w ) . ( 3 ) 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 final SAM algorithm by applying a standard numerical optimizer such as stochastic gradient descent ( SGD ) to the SAM objective LSAMS ( 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 . end returnwt Algorithm 1 : SAM algorithm

Summary: This 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.