SALR: Sharpness-aware Learning Rates for Improved Generalization

1 INTRODUCTION Generalization in deep learning has recently been an active area of research . The efforts to improve generalization over the past two decades have brought upon many cornerstone advances and techniques ; be it dropout ( Gal & Ghahramani , 2016 ) , batch-normalization ( Ioffe & Szegedy , 2015 ) , data-augmentation ( Shorten & Khoshgoftaar , 2019 ) , weight decay ( Loshchilov & Hutter , 2019 ) , adaptive gradient-based optimization ( Kingma & Ba , 2015 ) , architecture design and search ( Radosavovic et al. , 2020 ) , ensembles and their Bayesian counterparts ( Garipov et al. , 2018 ; Iz- mailov et al. , 2018 ) , amongst many others . Yet , recently , researchers have discovered that the concept of sharpness/flatness plays a fundamental role in generalization . Though sharpness was first discussed in the context of neural networks in the early work of Hochreiter & Schmidhuber ( 1997 ) , it was only brought to the forefront of deep learning research after the seminal paper by Keskar et al . ( 2017 ) . While trying to investigate decreased generalization performance when large batch sizes are used ( LeCun et al. , 2012 ) in stochastic gradient descent ( SGD ) , Keskar et al . ( 2017 ) notice that this phenomena can be justified by the ability of smaller batches to reach flat minimizers . Such flat minimizers in turn , generalize well as they are robust to low precision arithmetic or noise in the parameter space ( Dinh et al. , 2017 ; Kleinberg et al. , 2018 ) , as shown in Figure 1 . Since then , the generalization ability of flat minimizers has been repeatedly observed in many recent works ( Neyshabur et al. , 2017a ; Goyal et al. , 2017 ; Li et al. , 2018 ; Izmailov et al. , 2018 ) . Indeed , flat minimizers can potentially tie together many of the aforementioned approaches aimed at generalization . For instance , ( 1 ) higher gradient variance , when batches are small , increases the probability to avoid sharp regions ( same can be said for SGD compared to GD ) ( Kleinberg et al. , 2018 ) ( 2 ) averaging over multiple hypotheses leads to wider optima in ensembles and Bayesian deep learning ( Izmailov et al. , 2018 ) ( 3 ) regularization techniques such as dropout or over-parameterization can adjust the loss landscape into one that allows first order methods to favor wide valleys ( Chaudhari et al. , 2019 ; Allen-Zhu et al. , 2019 ) . In this paper we study the direct problem of developing an algorithm that can converge to flat minimizers.Specifically , we introduce SALR : a sharpness aware learning rate designed to explore the loss-surface of an objective function and avoid undesired sharp local minima . SALR dynamically updates the learning rate based on the sharpness of the neighborhood of the current solution . The idea is simple : automatically increase the learning rates at relatively sharp valleys in an effort to escape them . One of the key features of SALR is that it can be fitted into any gradient based method such as Adagrad ( Duchi et al. , 2011 ) , ADAM ( Kingma & Ba , 2015 ) and also into recent approaches towards escaping sharp valleys such as Entropy-SGD ( Chaudhari et al. , 2019 ) . 1.1 RELATED WORK . From a theoretical perspective , generalization of deep learning solutions has been explained through multiple lenses . One of which is uniform stability ( Bottou & Le Cun , 2005 ; Bottou & Bousquet , 2008 ; Hardt et al. , 2016 ; Gonen & Shalev-Shwartz , 2017 ; Bottou et al. , 2018 ) . An algorithm is uniformly stable if for all data sets differing in only one element , nearly the same outputs will be produced ( Bousquet & Elisseeff , 2002 ) . Hardt et al . ( 2016 ) show that SGD satisfies this property and derive a generalization bound for models learned with SGD . From a different viewpoint , Choromanska et al . ( 2015 ) ; Kawaguchi ( 2016 ) ; Poggio et al . ( 2017 ) ; Mohri et al . ( 2018 ) attribute generalization to the complexity of the hypothesis-space . Using measures like Rademacher complexity ( Mohri & Rostamizadeh , 2009 ) and the Vapnik-Chervonenkis ( VC ) dimension ( Sontag , 1998 ) , the former works show that deep hypothesis spaces are typically more advantageous in representing complex functions . Besides that , the importance of flatness on generalization has been theoretically highlighted through PAC-Bayes bounds ( Dziugaite & Roy , 2017 ; Neyshabur et al. , 2017b ; Wang et al. , 2018 ) . These papers highlight the ability to derive non-vacuous generalization bounds based on the sharpness of a model class while arguing that relatively flat solutions yield tight bounds . From an algorithmic perspective , approaches to recover flat minima are still limited . Most notably , Chaudhari et al . ( 2019 ) developed the Entropy-SGD algorithm . Entropy-SGD defines a localentropy-based objective which smoothens the energy landscape based on its local geometry . This in turn allows SGD to attain flatter solutions . Indeed , this approach was motivated by earlier work in statistical physics ( Baldassi et al. , 2015 ; 2016 ) which proves the existence of non-isolated solutions that generalize well in networks with discrete weights . Such non-isolated solutions correspond to flat minima in continuous settings . The authors then propose a set of approaches based on ensembles and replicas of the loss to favor wide solutions . Not too far , recent methods in Bayesian deep learning ( BDL ) have also shown potential to recover flat minima . BDL basically averages over multiple hypotheses weighted by their posterior probabilities ( ensembles being a special case of BDL ( Izmailov et al. , 2018 ) ) . One example , is the stochastic weighted averaging ( SWA ) algorithm proposed by Izmailov et al . ( 2018 ) . SWA simply averages over multiple points along the trajectory of SGD to potentially find flatter solutions compared to SGD . Another example is the SWA-Gaussian ( SWAG ) . SWAG defines a Gaussian posterior approximation over neural network weights . Afterwards , samples are taken from the approximated distribution to perform Bayesian model averaging ( Maddox et al. , 2019 ) . Here we also note the recent work by Patel ( 2017 ) which partially motivates our method . Upon the aformentioned observations in Keskar et al . ( 2017 ) , Patel ( 2017 ) shows that the learning rate lowerbound threshold for the divergence of batch SGD , run on quadratic optimization problems , increases for larger batch-sizes . In general non-convex settings , given a problem with N local minimizers , one can compute N lower bound thresholds for local divergence of batch SGD . The number of minimizers for which batch SGD can converge is non-decreasing in the batch size . This is used to explain the tendency of low-batch SGD to converge to flatter minimizers compared to large-batch SGD . The former result links the choice of batch size and its effect on generalization to the choice of the learning rate . With the latter being a tunable parameter , to our knowledge , developing a dynamic choice of the learning rate that targets convergence to flat minimizers has not been studied before . 2 GENERAL FRAMEWORK . In this paper , we propose a framework that dynamically chooses a Sharpness-Aware Learning Rate to promote convergence to flat minimizers . More specifically , our proposed method locally approximates sharpness at the current iterate and dynamically adjusts the learning rate accordingly . In sharp regions , relatively large learning rates are attained to increase the chance of escaping that region . In contrast , when the current iterate belongs to a flat region , our method returns a relatively small learning rate to guarantee convergence . Our framework can be adopted by any local search descent method and is detailed in Algorithm 1 . Algorithm 1 : Sharpness-Aware Learning Rate ( SALR ) Framework Input : Starting point θ0 , initial learning rate η0 , Number of iterations K. for k = 0 , 1 , . . . , K do Estimate Ŝk , the local sharpness around the current iterate θk ; Set ηk = η0 Ŝk Median { Ŝi } k i=1 ; Compute θk+1 using some local search descent method ( Gradient Descent , Stochastic Gradient Descent , ADAM , . . . ) ; end Return θK . As detailed in Algorithm 1 , at every iterate k , we compute the learning rate as a function of the local sharpness parameters { Ŝk } k i=1 . The main intuition is to have the current learning rate to be an increasing function of the current estimated sharpness . Since the scale of the sharpness at different points can vary when using different networks or datasets ( Dinh et al. , 2017 ) , we normalize our estimated sharpness by dividing by the median of the sharpness of previous iterates . For instance in Figure 2 , despite having a similar sharpness measure , we consider the minimizer around θ = 1 to be sharp relative to the blue plot and flat relative to red plot . Normalization resolves this issue by helping our sharpness measure attain scale invariant properties . One can think of the median of previous sharpness values as a global sharpness parameter the algorithm is trying to learn . When k is sufficiently large , the variation in the global sharpness parameter among different iterates will be minimal . From an algorithmic perspective , SALR exploits a neighborhood around the current iterate to dynamically compute a desired learning rate while simultaneously exploring the sharpness of the landscape to refine this global sharpness parameter . 3 SHARPNESS MEASURE . Several sharpness/flatness measures have been defined in recent literature ( Rangamani et al. , 2019 ; Keskar et al. , 2017 ; Hochreiter & Schmidhuber , 1997 ) . For instance , Hochreiter & Schmidhuber ( 1997 ) computes flatness by measuring the size of the connected region in the parameter space where the objective remains approximately constant . In a more recent paper , Rangamani et al . ( 2019 ) proposed a scale invariant flatness measure based on the quotient manifold . Computing such notions for complex non-convex landscapes can be intractable in practice . In addition to the cited results , Keskar et al . ( 2017 ) quantifies flatness by finding the difference between the maximum value of the loss function within a small neighborhood around a given point and the current value . More specifically , they define sharpness as follows : φ ( ε , θ ) , S ( ε , θ ) 1 + f ( θ ) and S ( ε , θ ) = max θ′∈IBε ( θ ) f ( θ′ ) − f ( θ ) , ( 1 ) where IBε ( θ ) is a euclidean ball with radius ε centered at θ and 1+f ( θ ) is a normalizing coefficient . One drawback of equation 1 is that the sharpness value around a maximizer is nearly zero . To resolve this issue , one can simply modify the sharpness measure in equation 1 as follows : S ( ε , θ ) , max θ′∈IBε ( θ ) f ( θ′ ) − min θ′∈IBε ( θ ) f ( θ′ ) . ( 2 ) It can be easily shown that if θ is a local minimizer , equation 2 is equivalent to equation 1 . Both measures defined in equation 1 and equation 2 require solving a possibly non-convex function which is in general NP-Hard . For computational feasibility , we provide a sharpness approximation by running n1 gradient ascent and n2 gradient descent steps . The resulting solutions are used to approximate the maximization and minimization optimization problems . Here we note that our definition for sharpness does not include a normalizing coefficient , as median { Ŝi } k i=1 in Algorithm 1 plays this role . The details of the approximation are shown in the Definition 1 . Definition 1 . Given θ ∈ Rn , iteration numbers n1 and n2 , and step-size γ , we define the sharpness measure Ŝ ( θ ) , f ( θ ( n2 ) k , + ) − f ( θk ) + f ( θk ) − f ( θ ( n1 ) k , − ) + = f ( θ ( n2 ) k , + ) − f ( θ ( n1 ) k , − ) , where θ ( 0 ) k , + = θ ( 0 ) k , − = θk , θ ( n1 ) k , − = θk − n1−1∑ i=0 γ ∇f ( θ ( i ) k , − ) ∥∥∥∇f ( θ ( i ) k , − ) ∥∥∥ , and θ ( n2 ) k , + = θk + n2−1∑ i=0 γ ∇f ( θ ( i ) k , + ) ∥∥∥∇f ( θ ( i ) k , + ) ∥∥∥ . Remark 2 . In contrast to the measures defined in equation 2 and equation 1 , Definition 1 does not require a ball radius ε . However , our definition requires specifying the step-size γ and the number of ascent and descent iterations . Remark 3 . Running gradient descent/ascent with fixed step-size near a minimizer can return a very small sharpness value even if the minimizer is sharp . This is due to the small gradient norm around a minimizer . To resolve this issue , we normalize the gradient at every descent/ascent step . Moreover , normalizing by the norm of the gradient helps in understanding the radius of the ball containing the iterates { θ ( j ) k , − } n1 j=1 and { θ ( j ) k , + } n2 j=1 . Figure 3 shows the plots of the three different sharpness measures defined in this section when computed for a function f ( θ ) = 0.5θ sin ( 3θ ) + 1 . Notice that the blue plot corresponding to the sharpness measure φ ( · ) attains a zero value at local maximizers compared to a positive value for the other two sharpness plots . Moreover , notice that the sharpness value in these three plots attains a small value near the local minimizer . This can be explained by our choice of radius ε = 0.1 which limits the neighborhood being exploited . Increasing the radius for φ ( ε , · ) and S ( ε , · ) ( increasing n1 and n2 for Ŝ ) will provide higher values around the minimizer . We next show that using our sharpness measure in Definition 1 , gradient descent with SALR framework in Algorithm 1 , denoted as GD-SALR , escapes sharp local minima .

Summary: This work proposes an algorithm that aims at finding a flat minimizer. The high-level strategy in the design of the proposed algorithm is increasing the learning rate when the iterate is in the region of a sharp minimizer. The authors claim that by increasing the learning rate, the iterate can get out of the undesired region. To estimate the local landscape (local sharpness), the authors provide a heuristic, which requires running gradient descent and gradient ascent for some number of iterations. The proposed algorithm has a promising result empirically, compared to entropy SGD (Chaudhari et al.) which also aims at finding a solution that generalizes well.