A Diffusion Theory For Deep Learning Dynamics: Stochastic Gradient Descent Exponentially Favors Flat Minima

1 INTRODUCTION . In recent years , deep learning ( LeCun et al. , 2015 ) has achieved great empirical success in various application areas . Due to the over-parametrization and the highly complex loss landscape of deep networks , optimizing deep networks is a difficult task . Stochastic Gradient Descent ( SGD ) and its variants are mainstream methods for training deep networks . Empirically , SGD can usually find flat minima among a large number of sharp minima and local minima ( Hochreiter & Schmidhuber , 1995 ; 1997 ) . More papers reported that learning flat minima closely relate to generalization ( Hardt et al. , 2016 ; Zhang et al. , 2017a ; Arpit et al. , 2017 ; Hoffer et al. , 2017 ; Dinh et al. , 2017 ; Neyshabur et al. , 2017 ; Wu et al. , 2017 ; Dziugaite & Roy , 2017 ; Kleinberg et al. , 2018 ) . Some researchers specifically study flatness itself . They try to measure flatness ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2017 ; Sagun et al. , 2017 ; Yao et al. , 2018 ) , rescale flatness ( Tsuzuku et al. , 2019 ; Xie et al. , 2020b ) , and find flatter minima ( Hoffer et al. , 2017 ; Chaudhari et al. , 2017 ; He et al. , 2019b ; Xie et al. , 2020a ) . However , we still lack a quantitative theory that answers why deep learning dynamics selects a flat minimum . The diffusion theory is an important theoretical tool to understand how deep learning dynamics works . It helps us model the diffusion process of probability densities of parameters instead of model parameters themselves . The density diffusion process of Stochastic Gradient Langevin Dynamics ( SGLD ) under injected isotropic noise has been discussed by ( Sato & Nakagawa , 2014 ; Raginsky et al. , 2017 ; Zhang et al. , 2017b ; Xu et al. , 2018 ) . Zhu et al . ( 2019 ) revealed that anisotropic diffusion of SGD often leads to flatter minima than isotropic diffusion . A few papers has quantitatively studied the diffusion process of SGD under the isotropic gradient noise assumption . Jastrzębski et al . ( 2017 ) first studied the minima selection probability of SGD . Smith & Le ( 2018 ) presented a Beyesian perspective on generalization of SGD . Wu et al . ( 2018 ) studied the escape problems of SGD from a dynamical perspective , and obtained the qualitative conclusion on the effects of batch size , learning rate , and sharpness . Hu et al . ( 2019 ) quantitatively showed that the mean escape time of SGD exponentially depends on the inverse learning rate . Achille & Soatto ( 2019 ) also obtained a related proposition that describes the mean escape time in terms of a free energy that depends on the Fisher Information . Li et al . ( 2017 ) analyzed Stochastic Differential Equation ( SDE ) of adaptive gradient methods . Nguyen et al . ( 2019 ) mainly contributed to closing the theoretical gap between continuous-time dynamics and discrete-time dynamics under isotropic heavy-tailed noise . However , the related papers mainly analyzed the diffusion process under parameter-independent and isotropic gradient noise , while stochastic gradient noise ( SGN ) is highly parameter-dependent and anisotropic in deep learning dynamics . Thus , they failed to quantitatively formulate how SGD selects flat minima , which closely depends on the Hessian-dependent structure of SGN . We try to bridge the gap between the qualitative knowledge and the quantitative theory for SGD in the presence of parameter-dependent and anisotropic SGN . Mainly based on Theorem 3.2 , we have four contributions : • The proposed theory formulates the fundamental roles of gradient noise , batch size , the learning rate , and the Hessian in minima selection . • The SGN covariance is approximately proportional to the Hessian and inverse to batch size . • Either a small learning rate or large-batch training requires exponentially many iterations to escape minima in terms of ratio of batch size and learning rate . • To the best of our knowledge , we are the first to theoretically and empirically reveal that SGD favors flat minima exponentially more than sharp minima . 2 STOCHASTIC GRADIENT NOISE AND SGD DYNAMICS . We mainly introduce the necessary foundation for the proposed diffusion theory in this section . We denote the data samples as { xj } mj=1 , the model parameters as θ and the loss function over data samples x as L ( θ , x ) . For simplicity , we denote the training loss as L ( θ ) . Following Mandt et al . ( 2017 ) , we may write SGD dynamics as θt+1 = θt − η ∂L̂ ( θt ) ∂θt = θt − η ∂L ( θt ) ∂θt + ηC ( θt ) 1 2 ζt , ( 1 ) where L̂ ( θ ) is the loss of one minibatch , ζt ∼ N ( 0 , I ) , and C ( θ ) represents the gradient noise covariance matrix . The classic approach is to model SGN by Gaussian noise , N ( 0 , C ( θ ) ) ( Mandt et al. , 2017 ; Smith & Le , 2018 ; Chaudhari & Soatto , 2018 ) . Stochastic Gradient Noise Analysis . We first note that the SGN we study is introduced by minibatch training , C ( θt ) 1 2 ζt = ∂L ( θt ) ∂θt − ∂L̂ ( θt ) ∂θt , which is the difference between gradient descent and stochastic gradient descent . According to Generalized Central Limit Theorem ( Gnedenko et al. , 1954 ) , the mean of many infinite-variance random variables converges to a stable distribution , while the mean of many finite-variance random variables converges to a Gaussian distribution . As SGN is finite in practice , we believe the Gaussian approximation of SGN is reasonable . Simsekli et al . ( 2019 ) argued that SGN is Lévy noise ( stable variables ) , rather than Gaussian noise . They presented empirical evidence showing that SGN seems heavy-tailed , and the heavy-tailed distribution looks closer to a stable distribution than a Gaussian distribution . However , this research line ( Simsekli et al. , 2019 ; Nguyen et al. , 2019 ) relies on a hidden strict assumption that SGN must be isotropic and obey the same distribution across dimensions . Simsekli et al . ( 2019 ) computed “ SGN ” across n model parameters and regarded “ SGN '' as n samples drawn from a single-variant distribution . This is why one tail-index for all parameters was studied in Simsekli et al . ( 2019 ) . The arguments in Simsekli et al . ( 2019 ) did not necessarily hold for parameter-dependent and anisotropic Gaussian noise . In our paper , SGN computed over different minibatches obeys a n-variant Gaussian distribution , which can be parameter-dependent and anisotropic . In Figure 1 , we empirically verify that SGN is highly similar to Gaussian noise instead of heavy-tailed Lévy noise . We recover the experiment of Simsekli et al . ( 2019 ) to show that gradient noise is approximately Lévy noise only if it is computed across parameters . Figure 1 actually suggests that the contradicted observations are from the different formulations of gradient noise . Simsekli et al . ( 2019 ) studied the distribution of SGN as a single-variant distribution , while we relax it as a n-variant distribution . Our empirical analysis in Figure 1 holds well at least when the batch size B is larger than 16 , which is common in practice . Similar empirical evidence can be observed for training ResNet18 ( He et al. , 2016 ) on CIFAR-10 ( Krizhevsky et al. , 2009 ) , seen in Appendix C. Panigrahi et al . ( 2019 ) also observed that for batch sizes 256 and above , the distribution of SGN is best described as Gaussian at-least in the early phases of training . Comparing our results with Panigrahi et al . ( 2019 ) , we noticed that the Gaussianity of SGN may depend on more unknown factors . First , SGN on random models is more Gaussian than well-trained models . Second , the layer/network matters . Because SGN on some layers/networks is more Gaussian than other layers/networks . The isotropic gradient noise assumption is too rough to capture the Hessian-dependent covariance structure of SGN , which we will study in Figure 2 later . Our theory that focuses on parameterdependent and anisotropic SGN brings a large improvement over existing parameter-independent and isotropic noise , although Simsekli et al . ( 2019 ) brought an improvement over more conventional parameter-independent and isotropic Gaussian noise . A more sophisticated theory is interesting under parameter-independent anisotropic heavy-tailed noise , when the batch size is too small ( B ∼ 1 ) to apply Central Limit Theorem . We will leave it as future work . SGD Dynamics . Let us replace η by dt as unit time . Then the continuous-time dynamics of SGD ( Coffey & Kalmykov , 2012 ) is written as dθ = −∂L ( θ ) ∂θ dt+ [ 2D ( θ ) ] 1 2 dWt , ( 2 ) where dWt ∼ N ( 0 , Idt ) and D ( θ ) = η2C ( θ ) . We note that the dynamical time t in the continuoustime dynamics is equal to the product of the number of iterations T and the learning rate η : t = ηT . The associated Fokker-Planck Equation is written as ∂P ( θ , t ) ∂t =∇ · [ P ( θ , t ) ∇L ( θ ) ] +∇ · ∇D ( θ ) P ( θ , t ) ( 3 ) = ∑ i ∂ ∂θi [ P ( θ , t ) ∂L ( θ ) ∂θi ] + ∑ i ∑ j ∂2 ∂θi∂θj Dij ( θ ) P ( θ , t ) , ( 4 ) where∇ is a nabla operator , and Dij is the element in the ith row and jth column of D. In standard SGLD , the injected gradient noise is fixed and isotropic Gaussian , D = I . The next question is how to formulate the SGN covariance C ( θ ) for SGD ? Based on Smith & Le ( 2018 ) , we can express the SGN covariance as C ( θ ) = 1 B 1 m m∑ j=1 ∇L ( θ , xj ) ∇L ( θ , xj ) > −∇L ( θ ) ∇L ( θ ) > ≈ 1 Bm m∑ j=1 ∇L ( θ , xj ) ∇L ( θ , xj ) > . ( 5 ) The approximation is true near critical points , due to the fact that the gradient noise variance dominates the gradient mean near critical points . We know the observed fisher information matrix satisfies FIM ( θ ) ≈ H ( θ ) near minima , referring to Chapter 8 of ( Pawitan , 2001 ) . Following Jastrzębski et al . ( 2017 ) ; Zhu et al . ( 2019 ) , we obtain C ( θ ) ≈ 1 Bm m∑ j=1 ∇L ( θ , xj ) ∇L ( θ , xj ) > = 1 B FIM ( θ ) ≈ 1 B H ( θ ) , ( 6 ) which approximately gives D ( θ ) = η 2 C ( θ ) = η 2B H ( θ ) ( 7 ) near minima . It indicates that the SGN covariance C ( θ ) is approximately proportional to the Hessian H ( θ ) and inverse to the batch size B . Obviously , we can generalize Equation 7 by D ( θ ) = ηC ( θ ) 2 = η 2B [ H ( θ ) ] + near critical points , when there exist negative eigenvalues in H along some directions . We use [ · ] + to denote the positive semidefinite transformation of a symmetric matrix : if we have the eigendecomposation H = U diag ( H1 , · · · , Hn−1 , Hn ) U > , then [ H ] + = U diag ( |H1| , · · · , |Hn−1| , |Hn| ) U > . We empirically verify this relation in Figure 2 for pretrained fully-connected networks , and a followup paper Xie et al . ( 2020c ) first verified this relation for randomly initialized fully-connected networks on real-world datasets . The Pearson Correlation is up to 0.999 for pretrained networks . We note that , the relation still approximately holds for even the randomly network , which is far from critical points . The correlation is especially high along the flat directions with small-magnitude eigenvalues of the Hessian ( Xie et al. , 2020c ) . We emphasize that previous papers with the isotropic Lévy or Gaussian noise approximation all failed to capture this core relation in deep learning dynamics .

Summary: The paper develops a density diffusion theory to reveal how minima selection quantitatively depends on the minima sharpness and the hyperparameters. It shows theoretically and empirically that SGD favors flat minima exponentially more than sharp minima. In particular, the paper analyzed the dependence of mean escape time from the valley with the Hessians on local minima and saddle points for both SGD and SGLD, and revealed the exponential dependence of the mean escape time with the sharpness. Experiments on real-world data have verified the theoretical results on the mean escape time.