Surrogate Gap Minimization Improves Sharpness-Aware Training

1 INTRODUCTION . Modern neural networks are typically highly over-parameterized and easy to overfit to training data , yet the generalization performances on unseen data ( test set ) often suffer from a gap from the training performance ( Zhang et al. , 2017a ) . Many studies try to understand the generalization of machine learning models , including the Bayesian perspective ( McAllester , 1999 ; Neyshabur et al. , 2017 ) , the information perspective ( Liang et al. , 2019 ) , the loss surface geometry perspective ( Hochreiter & Schmidhuber , 1995 ; Jiang et al. , 2019 ) and the kernel perspective ( Jacot et al. , 2018 ; Wei et al. , 2019 ) . Besides analyzing the properties of a model after training , some works study the influence of training and the optimization process , such as implicit regularization of stochastic gradient descent ( SGD ) ( Bottou , 2010 ; Zhou et al. , 2020 ) , the learning rate ’ s regularization effect ( Li et al. , 2019 ) , and the influence of the batch size ( Keskar et al. , 2016 ) . These studies have led to various modifications to the training process to improve generalization . Keskar & Socher ( 2017 ) proposed to use Adam in early training phases for fast convergence , then switch to SGD in late phases for better generalization . Izmailov et al . ( 2018 ) proposed to average weights to achieve a wider local minimum , which is expected to generalize better than sharp minima . A similar idea was later used in Lookahead ( Zhang et al. , 2019 ) . Entropy-SGD ( Chaudhari et al. , ∗Work was done during an internship at Google 2019 ) derived the gradient of local entropy to avoid solutions in sharp valleys . Entropy-SGD had a nested Langevin iteration , inducing much higher computation costs than vanilla training . The recently proposed Sharpness-Aware Minimization ( SAM ) ( Foret et al. , 2020 ) is a generic training scheme that improves generalization and has been shown especially effective for Vision Transformers ( Dosovitskiy et al. , 2020 ) when large-scale pre-training is unavailable ( Chen et al. , 2021 ) . Suppose vanilla training minimizes loss f ( w ) ( e.g. , the cross-entropy loss for classification ) , where w is the parameter . SAM minimizes perturbed loss defined as fp ( w ) , max||δ||≤ρ f ( w+ δ ) , which is the maximum loss within radius ρ centered at model parameter w. Intuitively , vanilla training seeks a single point with a low loss , while SAM searches for a neighborhood within which the maximum loss is low . However , we show that a low perturbed loss fp could appear in both flat and sharp minima , implying that only minimizing fp is not always sharpness-aware . Although the perturbed loss fp ( w ) might disagree with sharpness , we find surrogate gap defined as h ( w ) , fp ( w ) − f ( w ) agrees with sharpness — Lemma 3.3 shows that the surrogate gap h is an equivalent measure of the dominant eigenvalue of Hessian at a local minimum . Inspired by this observation , we propose the Surrogate Gap Guided Sharpness Aware Minimization ( GSAM ) which jointly minimizes the perturbed loss fp and the surrogate gap h : a low perturbed loss fp indicates a low training loss within the neighborhood , and a small surrogate gap h avoids solutions in sharp valleys and hence narrows the generalization gap between training and test performances ( Thm . 5.3 ) . When both criteria are satisfied , we find a generalizable model with good performance . GSAM consists of two steps for each update : 1 ) descend gradient∇fp ( w ) to minimize the perturbed loss fp ( this step is exactly the same as SAM ) , and 2 ) decompose gradient ∇f ( w ) of the original loss f ( w ) into components that are parallel and orthogonal to ∇fp ( w ) , i.e. , ∇f ( w ) = ∇‖f ( w ) + ∇⊥f ( w ) , and perform an ascent step in ∇⊥f ( w ) to minimize the surrogate gap h ( w ) . Note that this ascent step does not change the perturbed loss fp because∇f⊥ ( w ) ⊥ ∇fp ( w ) by construction . We summarize our contribution as follows : • We define surrogate gap , which measures the sharpness at local minima and is easy to compute . • We propose the GSAM method to improve generalization of neural networks . GSAM is widely applicable and incurs negligible computation overhead compared to SAM . • We demonstrate the convergence of GSAM and its provably better generalization than SAM . • We empirically validate GSAM over image classification tasks with various neural architec- tures , including ResNets ( He et al. , 2016 ) , Vision Transformers ( Dosovitskiy et al. , 2020 ) , and MLP-Mixers ( Tolstikhin et al. , 2021 ) . 2 PRELIMINARIES . 2.1 NOTATIONS . • f ( w ) : A loss function f with parameter w ∈ Rk , where k is the dimension of parameters . • ρt ∈ R : A scalar value controlling the amplitude of perturbation at step t. • ∈ R : A small positive constant ( to avoid division by 0 , = 10−12 by default ) . • wadvt , wt + ρt ∇f ( wt ) ||∇f ( wt ) ||+ : The solution to max||w′−wt||≤ρt f ( w ′ ) when ρt is small . • fp ( wt ) , max||δ||≤ρt f ( wt + δ ) ≈ f ( wadvt ) : The perturbed loss induced by f ( wt ) . For each wt , fp ( wt ) returns the worst possible loss f within a ball of radius ρt centered at wt . When ρt is small , by Taylor expansion , the solution to the maximization problem is equivalent to a gradient ascent from wt to wadvt . • h ( w ) , fp ( w ) − f ( w ) : The surrogate gap defined as the difference between fp ( w ) and f ( w ) . • ηt ∈ R : Learning rate at step t. • α ∈ R : A constant value that controls the scaled learning rate of the ascent step in GSAM . • g ( t ) , g ( t ) p ∈ Rk : At the t-th step , the noisy observation of the gradients ∇f ( wt ) , ∇fp ( wt ) of the original loss and perturbed loss , respectively . • ∇f ( wt ) = ∇f‖ ( wt ) +∇f⊥ ( wt ) : Decompose ∇f ( wt ) into parallel component ∇f‖ ( wt ) and vertical component∇f⊥ ( wt ) by projection∇f ( wt ) onto∇fp ( wt ) . 2.2 SHARPNESS-AWARE MINIMIZATION . Conventional optimization of neural networks typically minimizes the training loss f ( w ) by gradient descent w.r.t . ∇f ( w ) and searches for a single pointw with a low loss . However , this vanilla training often falls into a sharp valley of the loss surface , resulting in inferior generalization performance ( Chaudhari et al. , 2019 ) . Instead of searching for a single point solution , SAM seeks a region with low losses so that small perturbation to the model weights does not cause significant performance degradation . SAM formulates the problem as : minw fp ( w ) where fp ( w ) , max||δ||≤ρ f ( w + δ ) ( 1 ) where ρ is a predefined constant controlling the radius of a neighborhood . This perturbed loss fp induced by f ( w ) is the maximum loss within the neighborhood . When the perturbed loss is minimized , the neighborhood corresponds to low losses ( below the perturbed loss ) . For a small ρ , using Taylor expansion around w , the inner maximization in Eq . 1 turns into a linear constrained optimization with solution arg max||δ||≤ρ f ( w + δ ) = arg max||δ||≤ρ f ( w ) + δ > ∇f ( w ) +O ( ρ2 ) = ρ ∇f ( w ) ||∇f ( w ) || ( 2 ) As a result , the optimization problem of SAM reduces to minw fp ( w ) ≈ minw f ( wadv ) where wadv , w + ρ ∇f ( w ) ||∇f ( w ) ||+ ( 3 ) where is a scalar ( default : 1e-12 ) to avoid division by 0 , and wadv is the “ perturbed weight ” with the highest loss within the neighborhood . Equivalently , SAM seeks a solution on the surface of the perturbed loss fp ( w ) rather than the original loss f ( w ) ( Foret et al. , 2020 ) . 3 THE SURROGATE GAP MEASURES THE SHARPNESS AT A LOCAL MINIMUM . 3.1 THE PERTURBED LOSS IS NOT ALWAYS SHARPNESS-AWARE . Despite that SAM searches for a region of low losses , we show that a solution by SAM is not guaranteed to be flat . Throughout this paper we measure the sharpness at a local minimum of loss f ( w ) by the dominant eigenvalue σmax ( eigenvalue with the largest absolute value ) of Hessian . For simplicity , we do not consider the influence of reparameterization on the geometry of loss surfaces , which is thoroughly discussed in ( Laurent & Massart , 2000 ; Kwon et al. , 2021 ) . Lemma 3.1 . For some fixed ρ , consider two local minima w1 and w2 , fp ( w1 ) ≤ fp ( w2 ) 6=⇒ σmax ( w1 ) ≤ σmax ( w2 ) , where σmax is the dominant eigenvalue of the Hessian . We leave the proof to Appendix . Fig . 1 illustrates Lemma 3.1 with an example . Consider three local minima denoted as w1 to w3 , and suppose the corresponding loss surfaces are flatter from w1 to w3 . For some fixed ρ , we plot the perturbed loss fp and surrogate gap h around each solution . Comparing w2 with w3 : Suppose their vanilla losses are equal , f ( w2 ) = f ( w3 ) , then fp ( w2 ) > fp ( w3 ) because the loss surface is flatter around w3 , implying that SAM will prefer w3 to w2 . Comparing w1 and w2 : fp ( w1 ) < fp ( w2 ) , and SAM will favor w1 over w2 because it only cares about the perturbed loss fp , even though the loss surface is sharper around w1 than w2 . 3.2 THE SURROGATE GAP AGREES WITH SHARPNESS . We introduce the surrogate gap that agrees with sharpness , defined as : h ( w ) , f ( wadv ) − f ( w ) ( 4 ) Intuitively , the surrogate gap represents the difference between the maximum loss within the neighborhood and the loss at the center point . The surrogate gap has the following properties . Lemma 3.2 . Suppose the perturbation amplitude ρ is sufficiently small , then the approximation to the surrogate gap in Eq . 4 is always non-negative , h ( w ) = f ( wadv ) − f ( w ) ≥ 0 , ∀w . Lemma 3.3 . For a local minimum w∗ , consider the dominate eigenvalue σmax of the Hessian of loss f as a measure of sharpness . Considering the neighborhood centered at w∗ with a small radius ρ , the surrogate gap h ( w∗ ) is an equivalent measure of the sharpness : σmax ≈ 2h ( w∗ ) /ρ2 . The proof is in Appendix . Lemma 3.2 tells that the surrogate gap is non-negative , and Lemma 3.3 shows that the loss surface is flatter as h gets closer to 0 . The two lemmas together indicate that we can find a region with a flat loss surface by minimizing the surrogate gap h ( w ) .

Summary: This paper builds on the main idea of SAM. Namely, the authors noticed that minimizing the perturbed loss $f_p$ is insufficient for guaranteeing flat minima. The authors then defined the notion of surrogate gap $h$, which they propose to minimize together with the perturbed loss. By using GSAM which consists of two gradient descent/ascent steps that include the minimizing $f_p$ and $h$, the authors show theoretically and experimentally that this results in a model with better generalization than SAM.