Efficient Sharpness-aware Minimization for Improved Training of Neural Networks

1 INTRODUCTION . Deep learning has achieved astounding performances in many fields by relying on larger numbers of parameters and increasingly sophisticated optimization algorithms . However , DNNs with far more parameters than training samples are more prone to poor generalization . Generalization is arguably the most fundamental and yet mysterious aspect of deep learning . Several studies have been conducted to better understand the generalization of DNNs and to train DNNs that generalize well across the natural distribution ( Keskar et al. , 2016 ; Neyshabur et al. , 2017 ; Chaudhari et al. , 2019 ; Zhang et al. , 2019 ; Wu et al. , 2020 ; Foret et al. , 2020 ; Zhang et al. , 2021 ) . For example , Keskar et al . ( 2016 ) investigate the effect of batch size on neural networks ’ generalization ability . Zhang et al . ( 2019 ) propose an optimizer for training DNNs with improved generalization ability . Specifically , Hochreiter & Schmidhuber ( 1995 ) , Li et al . ( 2017 ) and Dinh et al . ( 2017 ) argue that the geometry of the loss landscape affects generalization and DNNs with a flat minimum can generalize better . The recent work by Foret et al . ( 2020 ) proposes an effective training algorithm Sharpness Aware Minimizer ( SAM ) for obtaining a flat minimum . SAM employs a base optimizer such as Stochastic Gradient Descent ( Nesterov , 1983 ) or Adam ( Kingma & Ba , 2014 ) to minimize both the vanilla training loss and the sharpness . The sharpness , which describes the flatness of a minimum , is characterized using eigenvalues of the Hessian matrix by Keskar et al . ( 2016 ) . SAM quantifies the sharpness as the maximized change of training loss when a constraint perturbation is added to current weights . As a result , SAM leads to a flat minimum and significantly improves the generalization ability of the trained DNNs . SAM and its variants have been shown to outperform the state-of-the-art across a variety of deep learning benchmarks ( Kwon et al. , 2021 ; Chen et al. , 2021 ; Galatolo et al. , 2021 ; Zheng et al. , 2021 ) . Regrettably though , SAM and its variants achieve such remarkable performance at the expense of doubling the computational overhead of the given base optimizers , which minimize the training loss with a single forward and backward propagation step . SAM requires an additional propagation step compared to the base optimizers to resolve the weight perturbation for quantifying the sharpness . The extra propagation step requires the same computational overhead as the single propagation step used by base optimizers , resulting in SAM ’ s computational overhead being doubled ( 2× ) . As demonstrated in Figure 1 , SAM achieves higher test accuracy ( i.e. , 84.46 % vs. 81.89 % ) at the expense of sacrificing half of the training speed of the base optimizer ( i.e. , 276 imgs/s vs. 557 imgs/s ) . SGD SAM ESAM 80 81 82 83 84 85 86 A cc ur ac y ( % ) 100 200 300 400 500 600 700 Tr ai ni ng S pe ed ( i m ag es /s ) Accuracy Training Speed sharpness of DNNs . As a result , the sharpness calculated over the subsets can serve as an upper bound of the SAM ’ s sharpness , ensuring that SDS ’ s performance is comparable to that of SAM ’ s . We verify the effectiveness of ESAM on the CIFAR10 , CIFAR100 ( Krizhevsky et al. , 2009 ) and ImageNet ( Deng et al. , 2009 ) datasets with five different DNN architectures . The experimental results demonstrate that ESAM obtains flat minima at a cost of only 40 % ( vs. SAM ’ s 100 % ) extra computational overhead over base optimizers . More importantly , ESAM achieves better performance in terms of the test accuracy compared to SAM . In a nutshell , our contributions are as follows : • We propose two novel and effective training strategies Stochastic Weight Perturbation ( SWP ) and Sharpness-sensitive Data Selection ( SDS ) . Both strategies are designed to improve efficiency without sacrificing performance . The empirical results demonstrate that both of the proposed strategies can improve both the efficiency and effectiveness of SAM . • We introduce the ESAM , which integrates SWP and SDS . ESAM improves the generalization ability of DNNs with marginally additional computational cost compared to standard training . The rest of this paper is structured in this way . Section 2.1 introduces SAM and its computational issues . Section 2.2 and Section 2.3 discuss how the two proposed training strategies SWP and SDS are designed respectively . Section 3 verifies the effectiveness of ESAM across a variety of datasets and DNN architectures . Section 4 presents the related work and Section 5 concludes this paper . 2 METHODOLOGY . We start with recapitulating how SAM achieves a flat minimum with small sharpness , which is quantified by resolving a maximization problem . To compute the sharpness , SAM requires additional forward and backward propagation and results in the doubling of the computational overhead compared to base optimizers . Following that , we demonstrate how we derive and propose ESAM , which integrates SWP and SDS , to maximize efficiency while maintaining the performance . We introduce SWP and SDS in Sections 2.2 and 2.3 respectively . Algorithm 1 shows the overall proposed ESAM algorithm . Throughout this paper , we denote a neural network f with weight parameters θ as fθ . The weights are contained in the vector θ = ( θ1 , θ2 , . . . , θN ) , where N is the number of weight units in the neural network . Given a training dataset S that contains samples i.i.d . drawn from a distribution D , Algorithm 1 Efficient SAM ( ESAM ) Input : Network fθ , θ = ( θ1 , θ2 , . . . , θN ) , Training set S , Batch size b , Learning rate η > 0 , Neighborhood size ρ > 0 , Iterations A , SWP hyperparameter β , SDS hyperparameter γ . Output : A flat minimum solution θ̂ . 1 : for a = 1 to A do 2 : Sample a mini-batch B ⊂ S with size b . 3 : for n = 1 to N do 4 : if θn is chosen by probability β then 5 : n ← ρ1−β∇θnLB ( fθ ) . SWP in B1 6 : else 7 : n ← 0 8 : ̂← ( 1 , ... , N ) . Assign Weight Perturbation 9 : Compute ` ( fθ+̂ , xi , yi ) and construct B+ with selection ratio γ ( Equation 6 ) 10 : Compute gradients g = ∇θLB+ ( fθ+̂ ) . SDS in B2 11 : Update weights θ ← θ − ηg the network is trained to obtain optimal weights θ̂ via empirical risk minimization ( ERM ) , i.e. , θ̂ = arg min θ { LS ( fθ ) = 1 |S| ∑ ( xi , yi ) ∈S ` ( fθ , xi , yi ) } ( 1 ) where ` can be an arbitrary loss function . We take ` to be the cross entropy loss in this paper . The population loss is defined as LD ( fθ ) , E ( xi , yi ) ∼D ( ` ( fθ , xi , yi ) ) . In each training iteration , optimizers sample a mini-batch B ⊂ S with size b to update parameters . 2.1 SHARPNESS-AWARE MINIMIZATION AND ITS COMPUTATIONAL DRAWBACK . To improve the generalization capability of DNNs , Foret et al . ( 2020 ) proposed the SAM training strategy for searching flat minima . SAM trains DNNs by solving the following min-max optimization problem , min θ max : ‖ ‖2≤ρ LS ( fθ+ ) . ( 2 ) Given θ , the inner optimization attempts to find a weight perturbation in Euclidean ball with radius ρ that maximizes the empirical loss . The maximized loss at weights θ is the sum of the empirical loss and the sharpness , which is defined to be RS ( fθ ) = max : ‖ ‖2 < ρ [ LS ( fθ+ ) − LS ( fθ ) ] . This sharpness is quantified by the maximal change of empirical loss when a constraint perturbation is added to θ . The min-max problem encourages SAM to find flat minima . For a certain set of weights θ , Foret et al . ( 2020 ) theoretically justifies that the population loss of DNNs can be upper-bounded by the sum of sharpness , empirical loss , and a regularization term on the norm of weights ( refer to Equation 3 ) . Thus , by minimizing the sharpness together with the empirical loss , SAM produces optimized solutions for DNNs with flat minima , and the resultant models can thus generalize better ( Foret et al. , 2020 ; Chen et al. , 2021 ; Kwon et al. , 2021 ) . LD ( fθ ) ≤ RS ( fθ ) + LS ( fθ ) + λ‖θ‖22 = max : ‖ ‖2≤ρ LS ( fθ+ ) + λ‖θ‖22 ( 3 ) In practice , SAM first approximately solves the inner optimization by means of a single-step gradient descent method , i.e. , ̂ = arg max : ‖ ‖2 < ρ LS ( fθ+ ) ≈ ρ∇θLS ( fθ ) . ( 4 ) The sharpness at weights θ is approximated by RS ( fθ ) = LS ( fθ+̂ ) − LS ( fθ ) . Then , a base optimizer , such as SGD ( Nesterov , 1983 ) or Adam ( Kingma & Ba , 2014 ) , updates the DNNs ’ weights to minimize LS ( fθ+̂ ) . We refer to LS ( fθ+̂ ) as the SAM loss . Overall , SAM requires two forward and two backward operations to update weights once . We refer to the forward and backward propagation for approximating ̂ as F1 and B1 and those for updating weights by base optimizers as F2 and B2 respectively . Although SAM can effectively improve the generalization of DNNs , it additionally requires one forward and one backward operation ( F1 and B1 ) in each training iteration . Thus , SAM results in a doubling of the computational overhead compared to the use of base optimizers . To improve the efficiency of SAM , we propose ESAM , which consists of two strategies—SWP and SDS , to accelerate the sharpness approximation phase and the weight updating phase . Specifically , on the one hand , when estimating ̂ around weight vector θ , SWP efficiently approximates ̂ by randomly selecting each parameter with a given probability to form a subset of weights to be perturbed . The reduction of the number of perturbed parameters results in lower computational overhead during the backward propagation . SWP rescales the resultant weight perturbation so as to assure that the expected weight perturbation equals to ̂ , and the generalization capability thus will not be significantly degraded . On the other hand , when updating weights via base optimizers , instead of computing the upper bound LB ( fθ+̂ ) over a whole batch of samples , SDS selects a subset of samples , B+ , whose loss values increase the most with respect to the perturbation ̂ . Optimizing the weights based on a fewer number of samples decreases the computational overhead ( in a linear fashion ) . We further justify that LB ( fθ+̂ ) can be upper bounded by LB+ ( fθ+̂ ) and consequently the generalization capability can be preserved . In general , ESAM works much more efficiently and performs as well as SAM in terms of its generalization capability .

Summary: This paper investigates the crucial efficiency issue of the sharpness-aware minimizer (SAM). SAM improves the generalization of DNNs but results in a double training time compared to vanilla training. By analyzing the min-max procedure of SAM, the authors observe the computational redundancy and then propose a method ESAM to improve the efficiency from the data and parameters perspectives. The authors argue that SAM can be approximated properly with fewer computations. Empirical results show that ESAM can reduce the extra training time in CIFAR10/100 and ImageNet datasets with improved accuracy compared to SAM.