Dynamic of Stochastic Gradient Descent with State-dependent Noise

1 INTRODUCTION . Deep learning has achieved great success in various AI applications , such as computer vision , natural language processing , and speech recognition ( He et al. , 2016b ; Vaswani et al. , 2017 ; He et al. , 2016a ) . Stochastic gradient descent ( SGD ) and its variants are the mainstream methods to train deep neural networks , since they can deal with the computational bottleneck of the training over large-scale datasets ( Bottou & Bousquet , 2008 ) . Although SGD can converge to the minimum in convex optimization ( Rakhlin et al. , 2012 ) , neural networks are highly non-convex . To understand the behavior of SGD on non-convex optimization landscape , on one hand , researchers are investigating the loss surface of the neural networks with variant architectures ( Choromanska et al. , 2015 ; Li et al. , 2018b ; He et al. , 2019b ; Draxler et al. , 2018 ; Li et al. , 2018a ) ; on the other hand , researchers illustrate that the noise in stochastic algorithm may make it escape from local minima ( Keskar et al. , 2016 ; He et al. , 2019a ; Zhu et al. , 2019 ; Wu et al. , 2019a ; HaoChen et al. , 2020 ) . It is clear that whether stochastic algorithms can escape from poor local minima and finally stop at a minimum with low generalization error is crucial to its test performance . In this work , we focus on the dynamic of SGD and its impact to generalization , especially the escaping efficiency from local minima . To study the dynamic behavior of SGD , most of the works consider SGD as the discretization of a continuous-time dynamic system and investigate its dynamic properties . There are two typical types of models to approximate dynamic of SGD . ( Li et al. , 2017 ; Zhou et al. , 2019 ; Liu et al. , 2018 ; Chaudhari & Soatto , 2018 ; He et al. , 2019a ; Zhu et al. , 2019 ; Hu et al. , 2019 ; Xie et al. , 2020 ) approximate the dynamic of SGD by Langevin dynamic with constant diffusion coefficient and proved its escaping efficiency from local minima.These works make over-simplified assumption that the covariance matrix of gradient noise is constant , although it is state-dependent in general . The simplified assumption makes the proposed dynamic unable to explain the empirical observation that the distribution of parameters trained by SGD is heavy-tailed ( Mahoney & Martin , 2019 ) . To model the heavy-tailed phenomenon , Simsekli et al . ( 2019 ) ; Şimşekli et al . ( 2019 ) point that the variance of stochastic gradient may be infinite , and they propose to approximate SGD by dynamic driven by α-stable process with the strong infinite variance condition . However , as shown in the work ( Xie et al. , 2020 ; Mandt et al. , 2017 ) , the gradient noise follows Gaussian distribution and the infinite variance condition does not satisfied . Therefore it is still lack of suitable theoretical explanation on the implicit regularization of dynamic of SGD . In this work , we conduct a formal study on the ( state-dependent ) noise structure of SGD and its dynamic behavior . First , we show that the covariance of the noise of SGD in the quadratic basin surrounding the local minima is a quadratic function of the state ( i.e. , the model parameter ) . Thus , we propose approximating the dynamic of SGD near the local minimum using a stochastic differential equation whose diffusion coefficient is a quadratic function of state . We call the new dynamic power-law dynamic . We prove that its stationary distribution is power-law κ distribution , where κ is the signal to noise ratio of the second order derivatives at local minimum . Compared with Gaussian distribution , power-law κ distribution is heavy-tailed with tail-index κ . It matches the empirical observation that the distribution of parameters becomes heavy-tailed after SGD training without assuming infinite variance of stochastic gradient in ( Simsekli et al. , 2019 ) . Second , we analyze the escaping efficiency of power-law dynamic from local minima and its relation to generalization . By using the random perturbation theory for diffused dynamic systems , we analyze the mean escaping time for power-law dynamic . Our results show that : ( 1 ) Power-law dynamic can escape from sharp minima faster than flat minima . ( 2 ) The mean escaping time for power-law dynamic is only in the polynomial order of the barrier height , much faster than the exponential order for dynamic with constant diffusion coefficient . Furthermore , we provide a PAC-Bayes generalization bound and show power-law dynamic can generalize better than dynamic with constant diffusion coefficient . Therefore , our results indicate that the state-dependent noise helps SGD to escape from sharp minima quickly and implicitly learn well-generalized model . Finally , we corroborate our theory by experiments . We investigate the distributions of parameters trained by SGD on various types of deep neural networks and show that they are well fitted by power-law κ distribution . Then , we compare the escaping efficiency of dynamics with constant diffusion or state-dependent diffusion to that of SGD . Results show that the behavior of power-law dynamic is more consistent with SGD . Our contributions are summarized as follows : ( 1 ) We propose a novel power-law dynamic with state-dependent diffusion to approximate dynamic of SGD based on both theoretical derivation and empirical evidence . The power-law dynamic can explain the heavy-tailed phenomenon of parameters trained by SGD without assuming infinite variance of gradient noise . ( 2 ) We analyze the mean escaping time and PAC-Bayes generalization bound for power-law dynamic and results show that power-law dynamic can escape sharp local minima faster and generalize better compared with the dynamics with constant diffusion . Our experimental results can support the theoretical findings . 2 BACKGROUND . In empirical risk minimization problem , the objective is L ( w ) = 1n ∑n i=1 ` ( xi , w ) , where xi , i = 1 , · · · , n are n i.i.d . training samples , w ∈ Rd is the model parameter , and ` is the loss function . Stochastic gradient descent ( SGD ) is a popular optimization algorithm to minimize L ( w ) . The update rule is wt+1 = wt − η · g̃ ( wt ) , where g̃ ( wt ) = 1b ∑ x∈Sb ∇w ` ( x , wt ) is the minibatch gradient calculated by a randomly sampled minibatch Sb of size b and η is the learning rate . The minibatch gradient g̃ ( wt ) is an unbiased estimator of the full gradient g ( wt ) = ∇L ( wt ) , and the term ( g ( wt ) − g̃ ( wt ) ) is called gradient noise in SGD . Langevin Dynamic In ( He et al. , 2019a ; Zhu et al. , 2019 ) , the gradient noise is assumed to be drawn from Gaussian distribution according to central limit theorem ( CLT ) , i.e. , g ( w ) − g̃ ( w ) ∼ N ( 0 , C ) , where covariance matrix C is a constant matrix for all w. Then SGD can be regarded as the numerical discretization of the following Langevin dynamic , dwt = −g ( wt ) dt+ √ ηC1/2dBt , ( 1 ) where Bt is a standard Brownian motion in Rd and √ ηC1/2dBt is called the diffusion term . α-stable Process Simsekli et al . ( 2019 ) assume the variance of gradient noise is unbounded . By generalized CLT , the distribution of gradient noise is α-stable distribution S ( α , σ ) , where σ is the α-th moment of gradient noise for given α with α ∈ ( 0 , 2 ] . Under this assumption , SGD is approximated by the stochastic differential equation ( SDE ) driven by an α-stable process . 2.1 RELATED WORK . There are many works that approximate SGD by Langevin dynamic and most of the theoretical results are obtained for Langevin dynamic with constant diffusion coefficient . From the aspect of optimization , the convergence rate of SGD and its optimal hyper-parameters have been studied in ( Li et al. , 2017 ; He et al. , 2018 ; Liu et al. , 2018 ; He et al. , 2018 ) via optimal control theory . From the aspect of generalization , Chaudhari & Soatto ( 2018 ) ; Zhang et al . ( 2018 ) ; Smith & Le ( 2017 ) show that SGD implicitly regularizes the negative entropy of the learned distribution . Recently , the escaping efficiency from local minima of Langevin dynamic has been studied ( Zhu et al. , 2019 ; Hu et al. , 2019 ; Xie et al. , 2020 ) . He et al . ( 2019a ) analyze the PAC-Bayes generalization error of Langevin dynamic to explain the generalization of SGD . The solution of Langevin dynamic with constant diffusion coefficient is Gaussian process , which does not match the empirical observations that the distribution of parameters trained by SGD is a heavy-tailed ( Mahoney & Martin , 2019 ; Hodgkinson & Mahoney , 2020 ; Gurbuzbalaban et al. , 2020 ) . Simsekli et al . ( 2019 ) ; Şimşekli et al . ( 2019 ) assume the variance of stochastic gradient is infinite and regard SGD as discretization of a stochastic differential equation ( SDE ) driven by an α-stable process . The escaping efficiency for the SDE is also shown in ( Simsekli et al. , 2019 ) . However , these theoretical results are derived for dynamics with constant diffusion term , although the gradient noise in SGD is state-dependent . There are some related works analyze state-dependent noise structure in SGD , such as label noise in ( HaoChen et al. , 2020 ) and multiplicative noise in ( Wu et al. , 2019b ) . These works propose new algorithms motivated by the noise structure , but they do not analyze the escaping behavior of dynamic of SGD and the impact to generalization . Wu et al . ( 2018 ) analyze the escaping behavior of SGD with considering the fluctuations of the second order derivatives and propose the concept linearly stability . In our work , we propose power-law dynamic to approximate SGD and analyze the stationary distribution and the mean escaping time for it . 3 APPROXIMATING SGD BY POWER-LAW DYNAMIC . In this section , we study the ( state-dependent ) noise structure of SGD ( in Section 3.1 ) and propose power-law dynamic to approximate the dynamic of SGD . We first study 1-dimensional power-law dynamic in Section 3.2 and extend it to high dimensional case in Section 3.3 . 3.1 NOISE STRUCTURE OF STOCHASTIC GRADIENT DESCENT . For non-convex optimization , we investigate the noise structure of SGD around local minima so that we can analyze the escaping efficiency from it . We first describe the quadratic basin where the local minimum is located . Suppose w∗ is a local minimum of the training loss L ( w ) and g ( w∗ ) = 0 . We name the -ball B ( w∗ , ) with center w∗ and radius as a quadratic basin if the loss function for w ∈ B ( w∗ , ) is equal to its second-order Taylor expansion as L ( w ) = L ( w∗ ) + 12 ( w − w ∗ ) TH ( w∗ ) ( w − w∗ ) . Here , H ( w∗ ) is the Hessian matrix of loss at w∗ , which is ( semi ) positive definite . Then we start to analyze the gradient noise of SGD . The full gradient of training loss is g ( w ) = H ( w∗ ) ( w − w∗ ) . The stochastic gradient is g̃ ( w ) = g̃ ( w∗ ) + H̃ ( w∗ ) ( w − w∗ ) by Taylor expansion where g̃ ( · ) and H̃ ( · ) are stochastic version of gradient and Hessian calculated by the minibatch . The randomness of gradient noise comes from two parts : g̃ ( w∗ ) and H̃ ( w∗ ) , which reflects the fluctuations of the first-order and second-order derivatives of the model at w∗ over different minibatches , respectively . The following proposition gives the variance of the gradient noise . Proposition 1 For w ∈ B ( w∗ , ) ⊂ R , the variance of gradient noise is σ ( g ( w ) − g̃ ( w ) ) = σ ( g̃ ( w∗ ) ) + 2ρ ( g̃ ( w∗ ) , H̃ ( w∗ ) ) ( w − w∗ ) + σ ( H̃ ( w∗ ) ) ( w − w∗ ) 2 , where σ ( · ) and ρ ( · , · ) are the variance and covariance in terms of the minibatch . From Proposition 1 , we can conclude that : ( 1 ) The variance of noise is finite if g̃ ( w∗ ) and H̃ ( w∗ ) have finite variance because ρ ( g̃ ( w∗ ) , H̃ ( w∗ ) ) ≤ √ σ ( g̃ ( w∗ ) ) · σ ( H̃ ( w∗ ) ) according to Cauchy–Schwarz inequality . For fixed w∗ , a sufficient condition for that g̃ ( w∗ ) and H̃ ( w∗ ) have finite variance is that the training data x are sampled from bounded domain . This condition is easy to be satisfied because the domain of training data are usually normalized to be bounded before training . In this case , the infinite variance assumption about the stochastic gradient in α-stable process is not satisfied . ( 2 ) The variance of noise is state-dependent , which contradicts the assumption in Langevin dynamic . Notations : For ease of the presentation , we use C ( w ) , σg , σH , ρg , H to denote σ ( g ( w ) − g̃ ( w∗ ) ) , σ ( g̃ ( w∗ ) ) , σ ( H̃ ( w∗ ) ) , ρ ( g̃ ( w∗ ) , H̃ ( w∗ ) ) in the following context , respectively . 1

Summary: This paper proposes an analysis for an approximate dynamic of SGD which captures the heavy-tailed noise distributions seen practically at local minima. The authors derive this new dynamic (which they call Power-law dynamic) using basic principles and the assumption that the noise variance depends on the state. The dynamics becomes a modified Langevin equation. They prove also that the expected time to escape a barrier is polynomial in the parameters, as well as a generalization error bound.