The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Minima and Regularization Effects

1 INTRODUCTION . As a successful learning algorithm , stochastic gradient descent ( SGD ) was originally adopted for dealing with the computational bottleneck of training neural networks with large-scale datasets ( Bottou , 1991 ) . Its empirical efficiency and effectiveness have attracted lots of attention . And thus , SGD and its variants have become standard workhorse for learning deep models . Besides the aspect of empirical efficiency , recently , researchers started to analyze the optimization behaviors of SGD and its impacts on generalization . The optimization properties of SGD have been studied from various perspectives . The convergence behaviors of SGD for simple one hidden layer neural networks were investigated in ( Li & Yuan , 2017 ; Brutzkus et al. , 2017 ) . In non-convex settings , the characterization of how SGD escapes from stationary points , including saddle points and local minima , was analyzed in ( Daneshmand et al. , 2018 ; Jin et al. , 2017 ; Hu et al. , 2017 ) . On the other hand , in the context of deep learning , researchers realized that the noise introduced by SGD impacts the generalization , thanks to the research on the phenomenon that training with a large batch could cause a significant drop of test accuracy ( Keskar et al. , 2017 ) . Particularly , several works attempted to investigate how the magnitude of the noise influences the generalization during the process of SGD optimization , including the batch size and learning rate ( Hoffer et al. , 2017 ; Goyal et al. , 2017 ; Chaudhari & Soatto , 2017 ; Jastrzębski et al. , 2017 ) . Another line of research interpreted SGD from a Bayesian perspective . In ( Mandt et al. , 2017 ; Chaudhari & Soatto , 2017 ) , SGD was interpreted as performing variational inference , where certain entropic regularization involves to prevent overfitting . And the work ( Smith & Le , 2018 ) tried to provide an understanding based on model evidence . These explanations are compatible with the flat/sharp minima argument ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2017 ) , since Bayesian inference tends to targeting the region with large probability mass , corresponding to the flat minima . However , when analyzing the optimization behavior and regularization effects of SGD , most of existing works only assume the noise covariance of SGD is constant or upper bounded by some constant , and what role the noise structure of stochastic gradient plays in optimization and generalization was rarely discussed in literature . In this work , we theoretically study a general form of gradient-based optimization dynamics with unbiased noise , which unifies SGD and standard Langevin dynamics . By investigating this general dynamics , we analyze how the noise structure of SGD influences the escaping behavior from minima and its regularization effects . Several novel theoretical results and empirical justifications are made . 1 . We derive a key indicator to characterize the efficiency of escaping from minima through measuring the alignment of noise covariance and the curvature of loss function . Based on this indicator , two conditions are established to show which type of noise structure is superior to isotropic noise in term of escaping efficiency ; 2 . We further justify that SGD in the context of deep neural networks satisfies these two conditions , and thus provide a plausible explanation why SGD can escape from sharp minima more efficiently , converging to flat minima with a higher probability . Moreover , these flat minima typically generalize well according to various works ( Hochreiter & Schmidhuber , 1997 ; Keskar et al. , 2017 ; Neyshabur et al. , 2017 ; Wu et al. , 2017 ) . We also show that Langevin dynamics with well tuned isotropic noise can not beat SGD , which further confirms the importance of noise structure of SGD ; 3 . A large number of experiments are designed systematically to justify our understanding on the behavior of the anisotropic diffusion of SGD . We compare SGD with full gradient descent with different types of diffusion noise , including isotropic and positiondependent/independent noise . All these comparisons demonstrate the effectiveness of anisotropic diffusion for good generalization in training deep networks . The remaining of the paper is organized as follows . In Section 2 , we introduce the background of SGD and a general form of optimization dynamics of interest . We then theoretically study the behaviors of escaping from minima in Ornstein-Uhlenbeck process in Section 3 , and establish two conditions for characterizing the noise structure that affects the escaping efficiency . In Section 4 , we show that the noise of SGD in the context of deep learning meets the two conditions , and thus explains its superior efficiency of escaping from sharp minima over other dynamics with isotropic noise . Various experiments are conducted for verifying our understanding in Section 5 , and we conclude the paper in Section 6 . 2 BACKGROUND . In general , supervised learning usually involves an optimization process of minimizing an empirical loss over training data , L ( θ ) : = 1/N ∑N i=1 ` ( f ( xi ; θ ) , yi ) , where { ( xi , yi ) } Ni=1 denotes the training set with N i.i.d . samples , the prediction function f is often parameterized by θ ∈ RD , such as deep neural networks . And ` ( · , · ) is the loss function , such as mean squared error and cross entropy , typically corresponding to certain negative log likelihood . Due to the over parameterization and non-convexity of the loss function in deep networks , there exist multiple global minima , exhibiting diverse generalization performance . We call those solutions generalizing well good solutions or minima , and vice versa . Gradient descent and its stochastic variants A typical approach to minimize the loss function is gradient descent ( GD ) , the dynamics of which in each iteration t is , θt+1 = θt − ηtg0 ( θt ) , where g0 ( θt ) = ∇θL ( θt ) denotes the full gradient and ηt denotes the learning rate . In non-convex optimization , a more useful kind of gradient based optimizers act like GD with an unbiased noise , including gradient Langevin dynamics ( GLD ) , θt+1 = θt − ηtg0 ( θt ) + σt t , t ∼ N ( 0 , I ) , and stochastic gradient descent ( SGD ) , during each iteration t of which , a minibatch of training samples with size m are randomly selected , with index set Bt ⊂ { 1 , 2 , . . . , N } , and a stochastic gradient is evaluated based on the chosen minibatch , g̃ ( θt ) = ∑ i∈Bt ∇θ ` ( f ( xi ; θt ) , yi ) /m , which is an unbiased estimator of the full gradient g0 ( θt ) . Then , the parameters are updated with some learning rate ηt as θt+1 = θt − ηtg̃ ( θt ) . Denote g ( θ ) = ∇θ ` ( ( f ( x ; θ ) , y ) , the gradient for loss with a single data point ( x , y ) , and assume that the size of minibatch is large enough for the central limit theorem to hold , and thus g̃ ( θt ) follows a Gaussian distribution ( Mandt et al. , 2017 ; Li et al. , 2017 ) , g̃ ( θt ) ∼ N ( g0 ( θt ) , 1 m Σ ( θt ) ) , where Σ ( θt ) ≈ 1 N N∑ i=1 ( g ( θt ; xi ) − g0 ( θt ) ) ( g ( θt ; xi ) − g0 ( θt ) ) T . ( 1 ) Note that the covariance matrix Σ depends on the model architecture , dataset and the current parameter θt . Now we can rewrite the update of SGD as , θt+1 = θt − ηtg0 ( θt ) + ηt√ m t , t ∼ N ( 0 , Σ ( θt ) ) . ( 2 ) Inspired by GLD and SGD , we may consider a general kind of optimization dynamics , namely , gradient descent with unbiased noise , θt+1 = θt − ηtg0 ( θt ) + σt t , t ∼ N ( 0 , Σt ) . ( 3 ) For small enough constant learning rate ηt = η , the above iteration in Eq . ( 3 ) can be treated as the numerical discretization of the following stochastic differential equation ( Li et al. , 2017 ; Jastrzębski et al. , 2017 ; Chaudhari & Soatto , 2017 ) , dθt = −∇θL ( θt ) dt+ √ ησ2tΣt dWt . ( 4 ) Considering √ ησ2tΣt as the coefficient of noise term , existing works ( Hoffer et al. , 2017 ; Jastrzębski et al. , 2017 ) studied the influence of noise magnitude of SGD on generalization , i.e . ησ2t = η/m . In this work , we focus on studying the benefits of anisotropic structure of Σt in SGD helping escape from minima by bridging the covariance matrix with the Hessian of the loss surface , and its implicit regularization effects on generalization , especially in deep learning context . For the purpose of eliminating the influence of the noise magnitude , we constrain it to be a constant when studying different structures of noise covariance . The noise magnitude could be evaluated as the expectation of the squared norm of the noise vector , E [ ( √ ησt t ) T ( √ ησt t ) ] = ησ 2 tE [ T ] = ησ2t TrE [ T ] = ησ2t Tr Σt . ( 5 ) Thus , we introduce the following constraint , given time t , ησ2t Tr ( Σt ) is constant . ( 6 ) From the statistical physics point of view , Tr ( ησ2tΣt ) characterizes the kinetic energy ( Gardiner ) , thus it is natural to force the energy to be unchanging , otherwise it is trivial that the higher the energy is , the less stable the system is . For simplicity , we absorb ησ2t into Σt , denoting ησ 2 tΣt as Σt . If not pointed out , the subscript t of matrix Σt is omitted to emphasize that we are fixing t and discussing the varying structure of Σ . 3 THE BEHAVIORS OF ESCAPING FROM MINIMA IN ORNSTEIN-UHLENBECK PROCESS . For a general loss function L ( θ ) = EX ` X ( θ ) ( the expectation could be either population or empirical ) , where X denotes data example and θ denoted parameters to be optimized , under suitable smoothness assumptions , the SDE associated with the gradient variant optimizer as shown in Eq . ( 4 ) can be written as follows ( Li et al. , 2017 ; Jastrzębski et al. , 2017 ; Chaudhari & Soatto , 2017 ; Hu et al. , 2017 ) , with little abuse of notation , dθt = −∇θL ( θt ) dt+ Σ 1 2 t dWt . ( 7 ) Let L0 = L ( θ0 ) be one of the minimal values of L ( θ ) , then for a fixed t small enough ( such that Lt−L0 ≥ 0 ) , Eθt [ Lt−L0 ] characterizes the efficiency of θ escaping from the minimum θ0 of L ( θ ) . It is natural to measure the escaping efficiency using E [ Lt − L0 ] since it characterizes the increase of the potential , i.e. , the increase of the loss L. And also note that Lt − L0 ≥ 0 , for any δ > 0 , the escaping probability P ( Lt − L0 ≥ δ ) can be controlled by the expectation E [ Lt − L0 ] since by Markov ’ s inequality , we have P ( Lt − L0 ≥ δ ) ≤ E [ Lt−L0 ] δ . Proposition 1 ( Escaping efficiency for general process ) . For the process ( 7 ) , provided mild smoothness assumptions , the escaping efficiency from the minimum θ0 is , E [ Lt − L0 ] = − ∫ t 0 E [ ∇LT∇L ] + ∫ t 0 1 2 ETr ( HtΣt ) dt , ( 8 ) where Ht denotes the Hessian of L ( θt ) at θt . We provide the proof in Appendix , and the same for the other propositions . The escaping efficiency for general processes is hard to analyze due to the intractableness of the integral in Eq . ( 8 ) . However , we may consider the second-order approximation locally near the minima θ0 , where L ( θ ) ≈ L0 + 12 ( θ − θ0 ) TH ( θ − θ0 ) . Without losing generality , we suppose θ0 = 0 . Further , suppose that H is a positive definite matrix and the diffusion covariance Σt = Σ is constant for t. Then the SDE ( 7 ) becomes an Ornstein-Uhlenbeck process , dθt = −Hθt dt+ Σ 1 2 dWt , θ0 = 0 . ( 9 ) Proposition 2 ( Escaping efficiency of Ornstein-Uhlenbeck process ) . For Ornstein-Uhlenbeck process ( 9 ) , with t small enough , the escaping efficiency from minimum θ0 = 0 is , E [ Lt − L0 ] = 1 4 Tr ( ( I − e−2Ht ) Σ ) ≈ t 2 Tr ( HΣ ) . ( 10 ) Inspired by Proposition 1 and Proposition 2 , we propose Tr ( HΣ ) as an empirical indicator measuring the efficiency for a stochastic process escaping from minima . Now we turn to analysis which kind of noise covariance structure Σ will benefit escaping sharp minima , under the constraint Eq . ( 6 ) . Firstly , for the isotropic loss surface , i.e. , H = λI , the escaping efficiency is E [ Lt−L0 ] = λt2 Tr Σ , which is invariant under the constraint that Tr Σ is constant ( Eq . ( 6 ) ) . Thus it is only nontrivial to study the impact of noise structure when the Hessian of loss surface is anisotropic . Secondly , H and Σ being semi-positive definite , to achieve the maximum of Tr ( HΣ ) under constraint ( 6 ) , Σ should be Σ∗ = ( Tr Σ ) · λ1u1uT1 , where λ1 , u1 are the maximal eigenvalue and corresponding unit eigenvector of H . Note that the rank-1 matrix Σ∗ is highly anisotropic . More generally , the following Proposition 3 characterizes one kind of anisotropic noise significantly outperforming isotropic noise in order of number of parameters D , given H is ill-conditioned . Proposition 3 ( The benefits of anisotropic noise ) . With semi-positive definite H and Σ , assume ( 1 ) H is ill-conditioned . Let λ1 ≥ λ2 ≥ . . . , ≥ λD ≥ 0 be the eigenvalues of H in descent order , and for some constant k D and d > 12 , λ1 > 0 , λk+1 , λk+2 , . . . , λD < λ1D −d , ( 11 ) ( 2 ) Σ is “ aligned ” with H . Let ui be the corresponding unit eigenvector of eigenvalue λi , for some projection coefficient a > 0 , uT1 Σu1 ≥ aλ1 Tr Σ TrH , ( 12 ) then we have the benefit of the anisotropic noise over the isotropic one in term of escaping efficiency , which can be characterized by the follow ratio , Tr ( HΣ ) Tr ( HΣ̄ ) = O ( aD ( 2d−1 ) ) , ( 13 ) where Σ̄ = Tr ΣD I denotes the covariance of isotropic noise , to meet the constraint Eq . ( 6 ) . To give some geometric intuitions on the left hand side of Eq . ( 12 ) , let the maximal eigenvalue and its corresponding unit eigenvector of Σ be γ1 , v1 , then the right hand side has a lower bound as uT1 Σu1 ≥ uT1 v1γ1vT1 u1 = γ1 〈u1 , v1〉 2 . Thus if the maximal eigenvalues of H and Σ are aligned in proportion , γ1/Tr Σ ≥ a1λ1/TrH , and the angle of their corresponding unit eigenvectors is close to zero , 〈u1 , v1〉 ≥ a2 , the second condition Eq . ( 12 ) in Proposition 3 holds for a = a1a2 . Typically , in the scenario of modern deep neural networks , due to the over-parameterization , Hessian and the gradient covariance are usually ill-conditioned and anistropic near minima , as shown by ( Sagun et al. , 2017 ) and ( Chaudhari & Soatto , 2017 ) . Thus the first condition in Eq . ( 11 ) usually holds for deep neural networks , and we further justify it by experiments in Section 5.3 . Therefore , in the following section , we turn to focus on how the gradient covariance , i.e . the covariance of SGD noise meets the second condition of Proposition 3 in the context of deep neural networks .

Summary: The authors studied the effect of the anisotropic noise of SGD on the algorithm’s ability to escape from local optima. To this end, the authors depart from the established approximation of SGD in the vicinity of an optimum as a continuous-time Ornstein-Uhlenbeck process. Furthermore, the authors argue that in certain deep learning models, the anisotropic noise indeed leads to a good escaping from local optima.