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: In this paper,the authors proposed a new method for sharpness aware training. The idea is to note a surrogate gap as a measure of sharpness, which motivates to minimize the perturbed loss and the surrogate gap simultaneously. The authors proposed to use a gradient descent to minimize the perturbed loss and gradient ascent in an orthogonal direction to minimize the surrogate gap. The use of the orthogonal direction is to avoid the influence of the perturbed loss in the ascent step. The paper gives theoretical results to show the connection between the proposed surrogate gap and the sharpness. Convergence rates and generalization bounds are also presented. Experimental results are reported which seem to be convincing.