Eliminating Sharp Minima from SGD with Truncated Heavy-tailed Noise

1 INTRODUCTION . Stochastic gradient descent ( SGD ) and its variants have seen unprecedented empirical successes in training deep neural networks . The training of deep neural networks is typically posed as a nonconvex optimization problem , and even without explicit regularization the solutions obtained by SGD often perform surprisingly well on test data . Such an unexpected generalization performance of SGD in deep neural networks are often attributed to SGD ’ s ability to avoid sharp local minima1 in the loss landscape , which tends to lead to poor generalization ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2016 ; Li et al. , 2018b ; Jiang et al. , 2019 ) ; see Appendix D for more details . Despite significant efforts to explain such phenomena theoretically , understanding how SGD manages to avoid sharp local minima and end up with flat local minima within a realistic training time still remains as a central mystery of deep learning . Recently , the heavy-tailed dynamics of SGD received significant attention , and it was suggested that the heavy tails in the stochastic gradients may be a key ingredient that facilitates SGD ’ s escape from sharp local minima : for example , Şimşekli et al . ( 2019a ; b ) report the empirical evidence of heavy-tails in stochastic gradient noise in popular deep learning architectures ( see also ( Hodgkinson & Mahoney , 2020 ; Srinivasan et al. , 2021 ; Garg et al. , 2021 ) ) and show that SGD can escape sharp local minima in polynomial time under the presence of the heavy-tailed gradient noise . More specifically , they view heavy-tailed SGDs as discrete approximations of Lévy driven Langevin equations and argue that the amount of time SGD trajectory spends in each local minimum is proportional to the width of the associated minimum according to the metastability theory ( Pavlyukevich , 2007 ; Imkeller et al. , 2010a ; b ) for such heavy-tailed processes . In this paper , we study the global dynamics and long-run behavior of heavy-tailed SGD and its practical variant in depth . In particular , we consider an adaptive version of SGD , where the stochastic gradient is truncated above a fixed threshold . Such truncation scheme is often called gradient clipping and employed as default in various contexts ( Engstrom et al. , 2020 ; Merity et al. , 2018 ; Graves , 2013 ; 1We use the terminology sharpness in a broad sense ; we refer to Appendix C for a detailed discussion . Pascanu et al. , 2013 ; Zhang et al. , 2020 ; Gorbunov et al. , 2020 ) . We uncover a rich mathematical structure in the global dynamics of SGD under this scheme and prove that the long-run behavior of such SGD is fundamentally different from that of the pure form of SGD : in particular , under a suitable structural condition on the geometry of the loss landscape , gradient clipping completely eliminates sharp minima from the trajectory of SGDs . This provides a critical insight into how heavy-tailed dynamics of SGD can be utilized to find a local minimum that generalizes better . Figure 1 ( Left , Middle ) clearly illustrates these points with the histograms of the sample trajectories of SGDs . Note first that SGDs with light-tailed gradient noise— ( c ) and ( d ) of Figure 1 ( Left , Middle ) — never manages to escape a ( sharp ) minimum regardless of gradient clipping . In contrast , SGDs with heavy-tailed gradient noise— ( a ) and ( b ) of Figure 1 ( Left , Middle ) —easily escapes from local minima . Moreover , there is a clear difference between SGDs with gradient clipping and without gradient clipping . In ( a ) of Figure 1 ( Left ) , SGD without gradient clipping spends a significant amount of time at each of all four local minima ( { m1 , m2 , m3 , m4 } ) , although it spends more time around the wide ones ( { m2 , m4 } ) than the sharp ones ( { m1 , m3 } ) . On the other hand , in ( b ) of Figure 1 ( Left ) , SGD with gradient clipping not only escapes from local minima but also avoids sharp minima ( { m1 , m3 } ) almost completely . This means that if we stop training SGD at an arbitrary time point , it is almost guaranteed that it won ’ t be at a sharp minimum , effectively eliminating sharp minima from its training trajectories . We also propose a novel computational strategy that takes advantage of our newly discovered global dynamics of the heavy-tailed SGD . While the evidence of heavy tails were reported in many deep learning tasks ( Şimşekli et al. , 2019b ; a ; Garg et al. , 2021 ; Gurbuzbalaban et al. , 2020 ; Hodgkinson & Mahoney , 2020 ; Nguyen et al. , 2019 ; Mahoney & Martin , 2019 ; Srinivasan et al. , 2021 ; Zhang et al. , 2020 ) , there seem to be plenty of deep learning contexts where the stochastic gradient noises are light-tailed ( Panigrahi et al. , 2019 ) as well . Guided by our new theory , we propose an algorithm that injects heavy-tails to SGD by inflating the tail distribution of the gradient noise and facilitating the discovery of a local minimum that generalizes better . Our experiments with image classification tasks , reported in Tables 1 and 2 , illustrate that the tail-inflation strategy we propose here can indeed improve the generalization performance of the SGD as predicted by our theory . The rest of the paper is organized as follows . Section 2 formulates the problem setting and characterizes the global dynamics of the SGD driven by heavy-tailed noises . Section 3 presents numerical experiments that confirm our theory . Section 4 proposes a new algorithm that artificially injects heavy tailed gradient noise in actual deep learning tasks and demonstrate the improved performance . Technical Contributions : 1 ) We rigorously characterize the global behavior of the heavy-tailed SGD with gradient clipping . We first focus on the case where the loss function is in R1 with some simplifying assumptions on its geometry . Even with such assumptions , our theorem involves substantial technical challenges since the traditional tools for analyzing SGD fail in our context due to the adaptive nature of its dynamics and non-Gaussian distributional assumptions . For example , while the unclipped pure SGD can be analyzed by partitioning its trajectory at arrival times of large noises ( as in Pavlyukevich ( 2005 ) and Imkeller et al . ( 2010a ) ) , such an approach falls short in our context . Instead , we developed a set of delicate arguments for dealing with SGD ’ s ( near ) regeneration structure and the return times to the local minima , as well as controlling the probability of atypical scenarios that would not arise in the unclipped case . Moreover , as evidenced by our Rd results in Appendix G , the approach developed here is critical in extending the analysis to general loss landscapes . 2 ) We propose a novel computational strategy for improving the generalization performance of SGD by carefully injecting heavy-tailed noise . We test the proposed algorithm with deep learning tasks and demonstrate its superiority with an ablation study . This also suggests that the key phenomenon we characterize in our theory— elimination of sharp local minima—manifests in real-world tasks . 2 THEORETICAL RESULTS . This section characterizes the global dynamics of SGD with gradient clipping when applied to a non-convex objective function f . In Section 2.1 and 2.2 , we make the following assumptions for the sake of the simplicity of analysis . However , as our multidimensional result in Section 2.3 and the experiments in Section 3 and 4 suggest , the gist of the phenomena we analyze—elimination of sharp local minima—persists in general contexts where the domain of f is multi-dimensional , and the stationary points are not necessarily strict local optima separated from one another . Assumption 1 . Let f : R → R be a C2 function . There exist a positive real L > 0 , a positive integer nmin and an ordered sequence of real numbers m1 , s1 , m2 , s2 , · · · , snmin−1 , mnmin such that ( 1 ) −L < m1 < s1 < m2 < s2 < · · · < snmin−1 < mnmin < L ; ( 2 ) f ′ ( x ) = 0 iff x ∈ { m1 , s1 , · · · , snmin−1 , mnmin } ; ( 3 ) For any x ∈ { m1 , m2 , · · · , mnmin } , f ′′ ( x ) > 0 ; ( 4 ) For any x ∈ { s1 , s2 , · · · , snmin−1 } , f ′′ ( x ) < 0 . As illustrated in Figure 2 ( Left ) , the assumption above requires that f has finitely many local minima ( to be specific , the count is nmin ) , all of which contained in some compact domain [ −L , L ] . Moreover , the points s1 , · · · , snmin−1 naturally partition the entire real line into different regions Ωi = ( si−1 , si ) ( here we adopt the convention that s0 = −∞ , snmin = +∞ ) . We call each region Ωi the attraction field of the local minimum mi , as the gradient flow in Ωi always points to mi . Throughout the optimization procedure , given any location x ∈ R we assume that we have access to the noisy estimator f ′ ( x ) − Zn of the true gradient f ′ ( x ) , and f ′ ( x ) itself is difficult to evaluate . Specifically , in this work we are interested in the case where the iid sequence of noises ( Zn ) n≥1 are heavy-tailed . Typically , the heavy-tailed phenomena are captured by the concept of regular variation : for a measurable function φ : R+ 7→ R+ , we say that φ is regularly varying at +∞ with index β ( denoted as φ ∈ RVβ ) if limx→∞ φ ( tx ) /φ ( x ) = tβ for all t > 0 . For details on the definition and properties of regularly varying functions , see , for example , chapter 2 of Resnick ( 2007 ) . In this paper , we work with the following distributional assumption on the gradient noise . Let H+ ( x ) , P ( Z1 > x ) , H− ( x ) , P ( Z1 < −x ) , H ( x ) , H+ ( x ) +H− ( x ) = P ( |Z1| > x ) . Assumption 2 . EZ1 = 0 . Furthermore , there exists some α ∈ ( 1 , ∞ ) such that function H ( x ) is regularly varying ( at +∞ ) with index −α . Besides , regarding the positive and negative tail for distribution of the noises , we have lim x→∞ H+ ( x ) H ( x ) = p+ , lim x→∞ H− ( x ) H ( x ) = p− = 1− p+ where p+ and p− are constants in interval ( 0 , 1 ) . Roughly speaking , Assumption 2 means that the shape of the tail for the distribution of noises Zn resembles a polynomial function x−α , which is much heavier than the exponential tail of Gaussian distributions . Therefore , large values of Zn are much more likely to be observed under Assumption 2 compared to the typical Gaussian assumption . The index α of regular variation encodes the heaviness of the tail—the smaller the heavier—and we are assuming that the left and right tails share the same index α . The purpose of this simplifying assumption is clarity of presentation , but our Rd results in Appendix G relax such a condition and allow different regular variation indices in different directions . Our work concerns a popular variant of SGD where the stochastic gradient is truncated . Specifically , when updating the SGD iterates with a learning rate η > 0 , rather than using the original noisy gradient descent step η ( f ′ ( Xn ) − Zn ) , we will truncate it at a threshold b > 0 and use ϕb ( η ( f ′ ( Xn ) − Zn ) ) instead . Here the truncation operator ϕ· ( · ) is defined as ϕc ( w ) , w ·min { 1 , c/|w| } ∀w ∈ R , c > 0 . ( 1 ) Besides truncating the stochastic gradient , we also project the SGD into [ −L , L ] at each iteration ; recall that L is the constant in Assumption 1 . That is , the main object of our study is the stochastic process { Xηj } j≥0 driven by the following recursion Xηj = ∆ ϕL ( Xηj−1 − ϕb ( η ( f ′ ( Xηj−1 ) − Zj ) ) ) . ( 2 ) The projection ϕL and truncation ϕb here are common practices in many learning tasks for the purpose of ensuring that the SGD does not explode and drift to infinity . Besides , the projection also allows us to drop the sophisticated assumptions on the tail behaviors of f that are commonly seen in previous works ( see , for instance , the dissipativity conditions in Nguyen et al . ( 2019 ) ) . For technical reasons , we make the following assumption about the truncation threshold b > 0 . Note that this assumption is a very mild one , as it is obviously satisfied by ( Lebesgue ) almost every b > 0 . Assumption 3 . For each i = 1 , 2 , · · · , nmin , min { |si −mi| , |si−1 −mi| } /b is not an integer .

Summary: The paper studies gradient descent with injected power-law tail noise. The work shows that in the infinitesimal learning rate regime, the heavy tail noise can cause GD to not to converge to sharp minima.   Based on this theory, the paper proposes a technique to inject noise to GD to help training