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 presents a method called ESAM for improving the efficiency of SAM. ESAM has two components: SWP and SDS. SWS accelerates the estimation of $\epsilon$ via random sampling a subset of parameters in backpropagation. SDS further improves the efficiency via sampling a subset of data points that is enough for calculating the upper bound of $L$. Combining these two techniques, they achieve a speed improvement for SAM while yielding comparable or even better performance.