Understanding Sharpness-Aware Minimization

1 INTRODUCTION . Understanding generalization of overparametrized deep neural networks is a central topic of the current machine learning research . Their training objective has many global optima where the training data are perfectly fitted ( Zhang et al. , 2016 ) , but different global optima lead to dramatically different generalization performance ( Liu et al. , 2019 ) . However , it has been observed that stochastic gradient descent ( SGD ) tends to converge to well-generalizing solutions , even without any explicit regularization methods ( Zhang et al. , 2016 ) . This suggests that the leading role is played by the implicit bias of the optimization algorithms used ( Neyshabur et al. , 2014 ) : when the training objective is minimized using a particular algorithm and initialization method , it converges to a specific solution with favorable generalization properties . However , even though SGD has a very beneficial implicit bias , significant overfitting can still occur , particularly in the presence of label noise ( Nakkiran et al. , 2019 ) and adversarial perturbations ( Rice et al. , 2020 ) . Recently it has been observed that the sharpness of the training loss , i.e. , how quickly it changes in some neighborhood around the parameters of the model , correlates well with the generalization error ( Keskar et al. , 2016 ; Jiang et al. , 2019 ) , and generalization bounds related to the sharpness have been derived ( Dziugaite & Roy , 2018 ) . The idea of minimizing the sharpness to improve generalization has motivated recent works of Foret et al . ( 2021 ) and Wu et al . ( 2020 ) which propose to use worst-case perturbations of the weights on every iteration of training in order to improve generalization . We refer to this method as Sharpness-Aware Minimization ( SAM ) and focus mainly on the version proposed in Foret et al . ( 2021 ) that performs only one step of gradient ascent to approximately solve the weight perturbation problem before updating the weights . Despite the fact that SAM significantly improves generalization in various settings , the existing justifications based on the generalization bounds provided by Foret et al . ( 2021 ) and Wu et al . ( 2020 ) do not seem conclusive . The main reason is that their generalization bounds do not distinguish the worst-case robustness to worst-case weight perturbation from average-case robustness to Gaussian noise . However the latter does not sufficiently improve generalization as both Foret et al . ( 2021 ) and Wu et al . ( 2020 ) report . Moreover , their analysis does not distinguish whether the weight perturbation is applied based on some or on all training examples which , as we will discuss , has an important impact on generalization . We aim at further investigating the reasons for SAM ’ s success and make the following contributions : • We first discuss that the SAM formulation can lead to two different objectives . We show that the gradient flows on these objectives over diagonal linear networks are implicitly biased towards solutions which enjoy better generalization properties compared to the standard gradient flow . • We provide convergence results for the SAM algorithm used with stochastic gradients for non-convex objectives and discuss the convergence and generalization behavior of SAM for deep networks . • We discuss why SAM can prevent overfitting in the noisy label setting by interpreting SAM as a gradient reweighting scheme and showing the further benefits of combining it with a robust loss . • Finally , we draw parallels between overfitting in learning with noisy labels and in adversarial training where SAM also improves generalization . This connection suggests that techniques from the noisy label literature can be more generally useful to improve robust generalization . 2 RELATED WORK . Here we discuss relevant works on robustness in the weight space and cover the main references on overfitting in the noisy label setting and in adversarial training , two settings where weight-space robustness also improves generalization . Weight-space robustness . Works on weight-space robustness of neural networks date back at least to 1990s ( Murray & Edwards , 1993 ; Hochreiter & Schmidhuber , 1995 ) . Random perturbations of the weights are used extensively in deep learning ( Jim et al. , 1996 ; Graves et al. , 2013 ) , and most prominently in approaches such as dropout ( Srivastava et al. , 2014 ) . Many practitioners have observed that using SGD with larger batches for training leads to worse generalization ( LeCun et al. , 2012 ) , and Keskar et al . ( 2016 ) have shown that this degradation of performance is correlated with the sharpness of the found parameters . This observation has motivated many further works which focus on closing the generalization gap between small-batch and large-batch SGD ( Wen et al. , 2018 ; Haruki et al. , 2019 ; Lin et al. , 2020 ) . More recently , Jiang et al . ( 2019 ) have shown a strong correlation between the sharpness and the generalization error on a large set of models under a variety of different settings hyperparameters , beyond the batch size . This has motivated the idea of minimizing the sharpness during training to improve standard generalization , leading to Sharpness-Aware Minimization ( SAM ) ( Foret et al. , 2021 ) . SAM modifies SGD such that on every iteration of training , the gradient is taken not at the current iterate but rather at a worst-case point in its vicinity . Zheng et al . ( 2021 ) concurrently propose a very similar weight perturbation method which also successfully improves standard generalization on multiple deep learning benchmarks . Wu et al . ( 2020 ) have also proposed a similar algorithm with the same motivation , although Wu et al . ( 2020 ) focuses on improving robust generalization of adversarial training . Overfitting in learning with noisy labels . Arpit et al . ( 2017 ) have observed that deep networks noticeably overfit in the presence of mislabeled samples but early stopping can mitigate the problem . There are multiple other approaches that can mitigate this overfitting behavior : removing ( Song et al. , 2020 ) or downweighting noisy points ( Jiang et al. , 2018 ; Huang et al. , 2020 ) , using robust losses ( Ghosh et al. , 2017 ; Zhang & Sabuncu , 2018 ; Menon et al. , 2019 ) . Most of these approaches explicitly or implicitly leverage the early-learning phenomenon ( Liu et al. , 2020 ) where the network tends to fit first correctly labeled points and noisy points later in training . In the context of adversarial robustness , Sanyal et al . ( 2020 ) have discussed that fitting label noise can aggravate adversarial vulnerability , although removing label noise is not sufficient to achieve adversarial robustness . Overfitting in adversarial training . Adversarial training in deep learning has been formulated as a robust optimization problem by Madry et al . ( 2018 ) where the worst-case perturbations are typically found via projected gradient descent . Rice et al . ( 2020 ) have described the robust overfitting phenomenon in adversarial training suggesting that early stopping is highly beneficial and provides a competitive baseline . However , they do not provide explanations about the reasons behind the overfitting phenomenon . The state-of-the-art approaches propose different ways to improve robust generalization and mitigate robust overfitting such as using additional training data ( Carmon et al. , 2019 ) , weight averaging ( Gowal et al. , 2020 ; Chen et al. , 2021 ) , or advanced data augmentations ( Rebuffi et al. , 2021 ) . Additionally , Wu et al . ( 2020 ) have shown that combining adversarial training with a method similar to sharpness-aware minimization also improves robust generalization . 3 IMPLICIT BIAS AND CONVERGENCE OF SAM . In this section , we discuss different objectives related to Sharpness-Aware Minimization ( SAM ) and their generalization properties , then derive convergence results for the stochastic SAM algorithm , and further discuss the convergence and generalization of SAM for deep networks . 3.1 IMPLICIT BIAS OF SAM . SAM objectives . Let { xi , yi } ni=1 be the training data points and ` i ( w ) be the loss of a classifier parametrized by weights w ∈ R|w| and evaluated at point ( xi , yi ) . Foret et al . ( 2021 ) theoretically base their approach on the following objective which we denote as MaxSum SAM : MaxSum SAM : min w∈R|w| max ‖δ‖2≤ρ 1 n n∑ i=1 ` i ( w + δ ) . ( 1 ) They justify this objective via a PAC-Bayesian generalization bound , although they show empirically ( see Fig . 3 therein ) that the following objective , denoted as SumMax SAM , leads to better generalization : SumMax SAM : min w∈R|w| 1 n n∑ i=1 max ‖δ‖2≤ρ ` i ( w + δ ) . ( 2 ) The update rule of SAM for these objectives amounts then to a variation of gradient descent with step size γt where the gradients are taken at intermediate points wit+1/2 , i.e. , wt+1 = wt − γtn ∑n i=1∇ ` i ( wit+1/2 ) . The two objectives , however , differ in how the points w i t+1/2 are computed since they approximately maximize different losses with inner step sizes ρt : MaxSum : wit+1/2 = wt + ρt n n∑ j=1 ∇ ` j ( wt ) vs. SumMax : wit+1/2 = wt + ρt∇ ` i ( wt ) . ( 3 ) We next show formally that SumMax SAM has a better implicit bias than ERM and MaxSum SAM . Implicit bias of SumMax and MaxSum SAM . The implicit bias of gradient methods is well understood for linear models where all gradient-based algorithms enjoy the same implicit bias . For diagonal linear neural networks , first-order algorithms have a richer implicit bias . We consider here a sparse regression problem and a predictor parametrized as fu , v = u u−v v. Before understanding how SAM induces a preferable bias , we first recall the seminal result of Woodworth et al . ( 2020 ) : assuming global convergence , the solution selected by the gradient flow initialised at α ∈ Rd and denoted βα∞ solves the constrained optimisation problem : βα∞ = arg min β∈Rd s.t . Xβ=y φα ( β ) . ( 4 ) The potential φα whose precise expression is given Eq . ( 12 ) in App . A interpolates nicely between the ` 1 and the ` 2 norms according to the initialization scale α : Small initialisations lead to low ` 1-type solutions which are known to induce good generalisation properties . Large initialisations lead to low ` 2-type solutions . Our main result is that both MaxSum and SumMax dynamics bias the flow towards solutions which still minimise the potential φα . However it does so with effective parameter αMaxSum and αSumMax which are strictly smaller than α for suitable inner step size ρ . Thus the chosen solution has better sparsity-inducing properties than the solution of the gradient flow . Proposition 1 ( Informal ) . Assuming global convergence , the solutions selected by the MaxSum and the SumMax algorithms defined Eq . ( 3 ) , taken in the infinitesimally-small-γt limit and initialised at α , solve the optimisation problem ( 4 ) with effective parameters αSumMax and αSumMax which satisfy : αSumMax = αe −ρ∆SumMax+O ( ρ2 ) and αMaxSum = αe−ρ∆MaxSum+O ( ρ 2 ) , where ∆SumMax and ∆MaxSum are two entrywise positive vectors for which typically : ‖∆SumMax‖1 ≈ d ∫ ∞ 0 L ( w ( s ) ) ds and ‖∆MaxSum‖1 ≈ d n ∫ ∞ 0 L ( w ( s ) ) ds . The results are formally stated in Proposition 4 and 5 in App . A . The SumMax implementations has better bias properties since its effective scale of α is considerably smaller than the one of MaxSum . It is worth noting that the vectors ∆SumMax and ∆MaxSum are linked with the integral of the loss function along the flow . Thereby , the speed of convergence of the training loss impacts the magnitude of the biasing effect : the slower the convergence , the better the bias , similarly to what observed for SGD by Pesme et al . ( 2021 ) . Empirical evidence for the implicit bias . We compare the performance of different methods ( ERM , MaxSum and SumMax SAM ) on the training and test losses in Fig . 1 . As predicted , the methods enjoy various generalization abilities : ERM and MaxSum SAM enjoy the same performance whereas SumMax benefits from a better implicit bias . We also note that the training loss of all the variants is converging to zero but as alluded before the convergence of SumMax SAM is slower . We show a similar experiments in App . A with stochastic variants of the algorithms . As expected , their performances are better than their deterministic counterparts ( Keskar et al. , 2016 ; Pesme et al. , 2021 ) . The results on the implicit bias presented above require that the algorithm converges to zero training error . Therefore we analyze next the convergence of the SAM algorithm for general nonconvex functions in the practically relevant stochastic case . Note that we can not expect rates as fast as the one observed for the previous simple model . They were due to its special structure for which fast convergence rates have been proven ( Yun et al. , 2021 ) .

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