Sharpness-Aware Minimization for Efficiently Improving Generalization — part 1/4
Cleaned-text code points [0, 15877). Bibliography omitted here; source.md includes it.
# Sharpness-Aware Minimization for Efficiently Improving Generalization
###### 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 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 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. We open source our code at https://github.com/google-research/sam .
## 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 state-of-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 ) , 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 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 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 a , vectors as 𝒂 \bm{a} , matrices as 𝑨 \bm{A} , sets as 𝒜 \mathcal{A} , and equality by definition as ≜ \triangleq . Given a training dataset 𝒮 ≜ ∪ i = 1 n { ( 𝒙 i , 𝒚 i ) } \mathcal{S}\triangleq\cup_{i=1}^{n}\{(\bm{x}_{i},\bm{y}_{i})\} drawn i.i.d. from distribution 𝒟 \mathscr{D} , we seek to learn a model that generalizes well. In particular, consider a family of models parameterized by 𝒘 ∈ 𝒲 ⊆ ℝ d \bm{w}\in\mathcal{W}\subseteq\mathbb{R}^{d} ; given a per-data-point loss function l : 𝒲 × 𝒳 × 𝒴 → ℝ + l:\mathcal{W}\times\mathcal{X}\times\mathcal{Y}\rightarrow\mathbb{R}_{+} , we define the training set loss L S ( 𝒘 ) ≜ 1 n ∑ i = 1 n l ( 𝒘 , 𝒙 i , 𝒚 i ) L_{S}(\bm{w})\triangleq\frac{1}{n}\sum_{i=1}^{n}l(\bm{w},\bm{x}_{i},\bm{y}_{i}) and the population loss L 𝒟 ( 𝒘 ) ≜ 𝔼 ( 𝒙 , 𝒚 ) ∼ D [ l ( 𝒘 , 𝒙 , 𝒚 ) ] L_{\mathscr{D}}(\bm{w})\triangleq\mathbb{E}_{(\bm{x},\bm{y})\sim D}[l(\bm{w},\bm{x},\bm{y})] . Having observed only 𝒮 \mathcal{S} , the goal of model training is to select model parameters 𝒘 \bm{w} having low population loss L 𝒟 ( 𝒘 ) L_{\mathscr{D}}(\bm{w}) .
Utilizing L 𝒮 ( 𝒘 ) L_{\mathcal{S}}(\bm{w}) as an estimate of L 𝒟 ( 𝒘 ) L_{\mathscr{D}}(\bm{w}) motivates the standard approach of selecting parameters 𝒘 \bm{w} by solving min 𝒘 L 𝒮 ( 𝒘 ) \min_{\bm{w}}L_{\mathcal{S}}(\bm{w}) (possibly in conjunction with a regularizer on 𝒘 \bm{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 𝒮 ( 𝒘 ) L_{\mathcal{S}}(\bm{w}) is typically non-convex in 𝒘 \bm{w} , with multiple local and even global minima that may yield similar values of L 𝒮 ( 𝒘 ) L_{\mathcal{S}}(\bm{w}) while having significantly different generalization performance (i.e., significantly different values of L 𝒟 ( 𝒘 ) L_{\mathscr{D}}(\bm{w}) ).
Motivated by the connection between sharpness of the loss landscape and generalization, we propose a different approach: rather than seeking out parameter values 𝒘 \bm{w} that simply have low training loss value L 𝒮 ( 𝒘 ) L_{\mathcal{S}}(\bm{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 \rho>0 , with high probability over training set 𝒮 \mathcal{S} generated from distribution 𝒟 \mathscr{D} , L 𝒟 ( 𝒘 ) ≤ max ‖ ϵ ‖ 2 ≤ ρ L 𝒮 ( 𝒘 + ϵ ) + h ( ‖ 𝒘 ‖ 2 2 / ρ 2 ) , L_{\mathscr{D}}(\bm{w})\leq\max_{\|\bm{{\epsilon}}\|_{2}\leq\rho}L_{\mathcal{S}}(\bm{w}+\bm{{\epsilon}})+h(\|\bm{w}\|_{2}^{2}/\rho^{2}), where h : ℝ + → ℝ + h:\mathbb{R}_{+}\rightarrow\mathbb{R}_{+} is a strictly increasing function (under some technical conditions on L 𝒟 ( 𝐰 ) L_{\mathscr{D}}(\bm{w}) ).
To make explicit our sharpness term, we can rewrite the right hand side of the inequality above as [ max ‖ ϵ ‖ 2 ≤ ρ L 𝒮 ( 𝒘 + ϵ ) − L 𝒮 ( 𝒘 ) ] + L 𝒮 ( 𝒘 ) + h ( ‖ 𝒘 ‖ 2 2 / ρ 2 ) . [\max_{\|\bm{{\epsilon}}\|_{2}\leq\rho}L_{\mathcal{S}}(\bm{w}+\bm{\epsilon})-L_{\mathcal{S}}(\bm{w})]+L_{\mathcal{S}}(\bm{w})+h(\|\bm{w}\|_{2}^{2}/\rho^{2}). The term in square brackets captures the sharpness of L 𝒮 L_{\mathcal{S}} at 𝒘 \bm{w} by measuring how quickly the training loss can be increased by moving from 𝒘 \bm{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 𝒘 \bm{w} . Given that the specific function h h is heavily influenced by the details of the proof, we substitute the second term with λ ‖ w ‖ 2 2 \lambda||w||_{2}^{2} for a hyperparameter λ \lambda , yielding a standard L2 regularization term. Thus, inspired by the terms from the bound, we propose to select parameter values by solving the following Sharpness-Aware Minimization (SAM) problem: min 𝒘 L 𝒮 S A M ( 𝒘 ) + λ ‖ 𝒘 ‖ 2 2 where L 𝒮 S A M ( 𝒘 ) ≜ max ‖ ϵ ‖ p ≤ ρ L S ( 𝒘 + ϵ ) , \min_{\bm{w}}L^{SAM}_{\mathcal{S}}(\bm{w})+\lambda||\bm{w}||_{2}^{2}\text{\penalty\ \penalty\ \penalty\ \penalty\ \penalty\ \penalty\ where\penalty\ \penalty\ \penalty\ \penalty\ \penalty\ \penalty\ }L^{SAM}_{\mathcal{S}}(\bm{w})\triangleq\max_{||\bm{\epsilon}||_{p}\leq\rho}L_{S}(\bm{w}+\bm{\epsilon}), (1) where ρ ≥ 0 \rho\geq 0 is a hyperparameter and p ∈ [ 1 , ∞ ] p\in[1,\infty] (we have generalized slightly from an L2-norm to a p p -norm in the maximization over ϵ \bm{\epsilon} , though we show empirically in appendix C.5 that p = 2 p=2 is typically optimal). Figure 1 shows 1 1 1 Figure 1 was generated following Li et al. (2017) with the provided ResNet56 (no residual connections) checkpoint, and training the same model with SAM. the loss landscape for a model that converged to minima found by minimizing either L 𝒮 ( 𝒘 ) L_{\mathcal{S}}(\bm{w}) or L 𝒮 S A M ( 𝒘 ) L^{SAM}_{\mathcal{S}}(\bm{w}) , illustrating that the sharpness-aware loss prevents the model from converging to a sharp minimum.
In order to minimize L 𝒮 S A M ( 𝒘 ) L^{SAM}_{\mathcal{S}}(\bm{w}) , we derive an efficient and effective approximation to ∇ 𝒘 L 𝒮 S A M ( 𝒘 ) \nabla_{\bm{w}}L^{SAM}_{\mathcal{S}}(\bm{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 L 𝒮 ( 𝒘 + ϵ ) L_{\mathcal{S}}(\bm{w}+\bm{\epsilon}) w.r.t. ϵ \bm{\epsilon} around 𝟎 \bm{0} , obtaining ϵ ∗ ( 𝒘 ) ≜ arg max ‖ ϵ ‖ p ≤ ρ L 𝒮 ( 𝐰 + ϵ ) ≈ arg max ‖ ϵ ‖ p ≤ ρ L 𝒮 ( 𝐰 ) + ϵ T ∇ 𝐰 L 𝒮 ( 𝐰 ) = arg max ‖ ϵ ‖ p ≤ ρ ϵ T ∇ 𝐰 L 𝒮 ( 𝐰 ) . {\bm{\epsilon}^{*}}(\bm{w})\triangleq\argmax_{\|\bm{\epsilon}\|_{p}\leq\rho}L_{\mathcal{S}}(\bm{w}+\bm{\epsilon})\approx\argmax_{\|\bm{\epsilon}\|_{p}\leq\rho}L_{\mathcal{S}}(\bm{w})+\bm{\epsilon}^{T}\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w})=\argmax_{\|\bm{\epsilon}\|_{p}\leq\rho}\bm{\epsilon}^{T}\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w}). In turn, the value ϵ ^ ( 𝒘 ) \hat{\bm{\epsilon}}(\bm{w}) that solves this approximation is given by the solution to a classical dual norm problem ( | ⋅ | q − 1 |\cdot|^{q-1} denotes elementwise absolute value and power) 2 2 2 In the case of interest p = 2 p=2 , this boils down to simply rescaling the gradient such that its norm is ρ \rho . :
ϵ ^ ( 𝒘 ) = ρ sign ( ∇ 𝒘 L 𝒮 ( 𝒘 ) ) | ∇ 𝒘 L 𝒮 ( 𝒘 ) | q − 1 / ( ‖ ∇ 𝒘 L 𝒮 ( 𝒘 ) ‖ q q ) 1 / p \hat{\bm{\epsilon}}(\bm{w})=\rho\mbox{ sign}\left(\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w})\right)\left|\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w})\right|^{q-1}/\bigg(\|\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w})\|_{q}^{q}\bigg)^{1/p} (2) where 1 / p + 1 / q = 1 1/p+1/q=1 . Substituting back into equation ( 1 ) and differentiating, we then have ∇ 𝒘 L 𝒮 S A M ( 𝒘 ) ≈ ∇ 𝒘 L 𝒮 ( 𝒘 + ϵ ^ ( 𝒘 ) ) = d ( 𝒘 + ϵ ^ ( 𝒘 ) ) d 𝒘 ∇ 𝒘 L 𝒮 ( 𝒘 ) | 𝒘 + ϵ ^ ( w ) = ∇ w L 𝒮 ( 𝒘 ) | 𝒘 + ϵ ^ ( 𝒘 ) + d ϵ ^ ( 𝒘 ) d 𝒘 ∇ 𝒘 L 𝒮 ( 𝒘 ) | 𝒘 + ϵ ^ ( 𝒘 ) . \begin{split}\nabla_{\bm{w}}L^{SAM}_{\mathcal{S}}(\bm{w})&\approx\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w}+\hat{\bm{\epsilon}}(\bm{w}))=\frac{d(\bm{w}+{\hat{\bm{\epsilon}}(\bm{w})})}{d\bm{w}}\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w})|_{\bm{w}+\hat{\bm{\epsilon}}(w)}\\ &={\nabla_{w}L_{\mathcal{S}}(\bm{w})|_{\bm{w}+\hat{\bm{\epsilon}}(\bm{w})}}+{\frac{d\hat{\bm{\epsilon}}(\bm{w})}{d\bm{w}}\nabla_{\bm{w}}L_{\mathcal{S}}(\bm{w})|_{\bm{w}+\hat{\bm{\epsilon}}(\bm{w})}}.\end{split} This approximation to ∇ 𝒘 L 𝒮 S A M ( 𝒘 ) \nabla_{\bm{w}}L^{SAM}_{\mathcal{S}}(\bm{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 𝒮 ( 𝒘 ) L_{\mathcal{S}}(\bm{w}) because ϵ ^ ( 𝒘 ) \hat{\bm{\epsilon}}(\bm{w}) is itself a function of ∇ 𝒘 L 𝒮 ( 𝒘 ) \nabla_{\bm{w}}L_{\mathcal{S}}(\bm{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: ∇ 𝒘 L 𝒮 S A M ( 𝒘 ) ≈ ∇ 𝒘 L 𝒮 ( w ) | 𝒘 + ϵ ^ ( 𝒘 ) . \nabla_{\bm{w}}L^{SAM}_{\mathcal{S}}(\bm{w})\approx{\nabla_{\bm{w}}L_{\mathcal{S}}(w)|_{\bm{w}+\hat{\bm{\epsilon}}(\bm{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 L 𝒮 S A M ( 𝒘 ) L^{SAM}_{\mathcal{S}}(\bm{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.
## 3 Empirical Evaluation