D YNAMIC OF S TOCHASTIC G RADIENT D ESCENT WITH S TATE -D EPENDENT N OISE
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic gradient descent (SGD) and its variants are mainstream methods to train deep neural networks.
Since neural networks are non-convex, more and more works study the dynamic behavior of SGD and its impact to generalization, especially the escaping efﬁciency from local minima.
However, these works make the over-simpliﬁed assumption that the distribution of gradient noise is stateindependent, although it is state-dependent.
In this work, we propose a novel power-law dynamic with state-dependent diffusion to approximate the dynamic of SGD.
Then, we prove that the stationary distribution of power-law dynamic is heavy-tailed, which matches the existing empirical observations.
Next, we study the escaping efﬁciency from local minimum of power-law dynamic and prove that the mean escaping time is in polynomial order of the barrier height of the basin, much faster than exponential order of previous dynamics.
It indicates that SGD can escape deep sharp minima efﬁciently and tends to stop at ﬂat minima that have lower generalization error.
Finally, we conduct experiments to compare SGD and power-law dynamic, and the results verify our theoretical ﬁndings.
I NTRODUCTION
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 ﬁnally 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 efﬁciency 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 coefﬁcient and proved its escaping efﬁciency from local minima.
These works make over-simpliﬁed assumption that the covariance matrix of gradient noise is constant, although it is state-dependent in general.
The simpliﬁed 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.
; ¸Sim¸sekli et al.
(2019) point that the variance of stochastic gradient may be inﬁnite, and they propose to approximate SGD by dynamic driven by α -stable process with the strong inﬁnite variance condition.
However, as shown in the work (Xie et al. , 2020; Mandt et al. , 2017), the gradient noise follows Gaussian distribution and the inﬁnite variance condition does not satisﬁed.
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 coefﬁcient 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 inﬁnite variance of stochastic gradient in (Simsekli et al. , 2019).
Second, we analyze the escaping efﬁciency 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 ﬂat minima.
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 coefﬁcient.
Furthermore, we provide a PAC-Bayes generalization bound and show power-law dynamic can generalize better than dynamic with constant diffusion coefﬁcient.
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 ﬁtted by power-law κ distribution.
Then, we compare the escaping efﬁciency 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 inﬁnite 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 ﬁndings.
B ACKGROUND
In empirical risk minimization problem, the objective is L ( w ) = 1 n (cid:80) ni =1 (cid:96) ( x i , w ) , where x i , i = 1 , · · · , n are n i.i.d.
training samples, w ∈ R d is the model parameter, and (cid:96) is the loss function.
Stochastic gradient descent (SGD) is a popular optimization algorithm to minimize L ( w ) .
The update rule is w t +1 =
w t − η · ˜ g ( w t ) , where ˜ g ( w t ) = 1 b (cid:80) x ∈ S b ∇ w (cid:96) ( x, w t ) is the minibatch gradient calculated by a randomly sampled minibatch S b of size b and η is the learning rate.
The minibatch gradient ˜ g ( w t ) is an unbiased estimator of the full gradient g ( w t ) =
∇ L ( w t ) , and the term ( g ( w t ) − ˜ g ( w t )) 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 ) −
[Equation]
, 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,
[Equation]
where B t is a standard Brownian motion in R d and √ ηC 1 /
dB t 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. R ELATED W ORK
There are many works that approximate SGD by Langevin dynamic and most of the theoretical results are obtained for Langevin dynamic with constant diffusion coefﬁcient.
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.
; Smith & Le (2017) show that SGD implicitly regularizes the negative entropy of the learned distribution.
Recently, the escaping efﬁciency 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 coefﬁcient 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.
; ¸Sim¸sekli et al.
(2019) assume the variance of stochastic gradient is inﬁnite and regard SGD as discretization of a stochastic differential equation (SDE) driven by an α -stable process.
The escaping efﬁciency 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 ﬂuctuations 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.
A PPROXIMATING SGD BY P OWER - LAW D YNAMIC
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 ﬁrst study 1-dimensional power-law dynamic in Section 3.2 and extend it to high dimensional case in Section 3.3.
3. N OISE S TRUCTURE OF S TOCHASTIC G RADIENT D ESCENT
For non-convex optimization, we investigate the noise structure of SGD around local minima so that we can analyze the escaping efﬁciency from it.
We ﬁrst 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 (cid:15) -ball B ( w ∗ , (cid:15) ) with center w ∗ and radius (cid:15) as a quadratic basin if the loss function for w ∈ B ( w ∗ , (cid:15) ) is equal to its second-order Taylor expansion as L ( w ) =
L ( w ∗ ) + 12 ( w − w ∗ ) T H ( w ∗ )( w − w ∗ ) .
Here, H ( w ∗ ) is the Hessian matrix of loss at w ∗ , which is (semi) positive deﬁnite.
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 (cid:101) H ( w ∗ ) , which reﬂects the ﬂuctuations of the ﬁrst-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 ∗ , (cid:15) )
⊂ 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 ﬁnite if ˜ g ( w ∗ ) and ˜ H ( w ∗ ) have ﬁnite variance because ρ (˜ g ( w ∗ ) , ˜ H ( w ∗ )) ≤ (cid:113) σ (˜ g ( w ∗ )) · σ ( ˜ H ( w ∗ )) according to Cauchy–Schwarz inequality.
For ﬁxed w ∗ , a sufﬁcient condition for that ˜ g ( w ∗ ) and ˜ H ( w ∗ ) have ﬁnite variance is that the training data x are sampled from bounded domain.
This condition is easy to be satisﬁed because the domain of training data are usually normalized to be bounded before training.
In this case, the inﬁnite variance assumption about the stochastic gradient in α -stable process is not satisﬁed.
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.
3. P OWER - LAW D YNAMIC
According to CLT, the gradient noise follows Gaussian distribution if it has ﬁnite variance, i.e., g
[Equation]
→
[Equation]
as b → ∞ , (2) where → d means “converge in distribution”.
Using Gaussian distribution to model the gradient noise in SGD, the update rule of SGD can be written as:
[Equation]
Eq.3 can be treated as the discretization of the following SDE, which we call it power-law dynamic: [Equation]
[Equation]
Power-law dynamic characterizes how the distribution of w changes as time goes on.
The distribution density of parameter w at time t (i.e., p ( w, t ) ) is determined by the Fokker-Planck equation (Zwanzig’s type (Guo & Du, 2014)):
[Equation]
The stationary distribution of power-law dynamic can be obtained if we let the left side of FokkerPlanck equation be zero.
The following theorem shows the analytic form of the stationary distribution of power-law dynamic, which is heavy-tailed and the tail of the distribution density decays at polynomial order of w − w ∗ .
This is the reason why we call the stochastic differential equation in
Eq.4 power-law dynamic.
Theorem 2 The stationary distribution density for 1-dimensional power-law dynamic (Eq.4) is
[Equation]
where C ( w ) = σ g +2 ρ g,H ( w − w ∗ )+ σ H ( w − w ∗ ) 2 , Z is the normalization constant and ArcTan ( · ) is the arctangent function.
We make discussions on property of p ( w ) .
The decreasing rate of p ( w ) as w goes away from the center w ∗ is mainly determined by the term C ( w ) − HησH (because the function ArcTan ( · ) is bounded) which is a polynomial function about w − w ∗ .
Compared with Gaussian distribution the probability density which follows exponential decreasing rate, power-law distribution is less concentrated in the quadratic basin B ( w ∗ , (cid:15) ) and heavy-tailed.
We call Hησ H the tail-index of p ( w ) and denote it as κ in the following context.
We can conclude that the state-dependent noise results in heavy-tailed distribution of parameters, which matches the observations in (Mahoney & Martin, 2019).
Langevin dynamic with constant diffusion can be regarded as special case of power-law dynamic when ρ H,g = 0 and σ H = 0 .
In this case, p ( w ) degenerates to Gaussian distribution.
Compared with α -stable process, we do not assume inﬁnite variance on gradient noise and demonstrate another mechanism that results in heavy-tailed distribution of parameters.
We empirically observe the covariance matrix around the local minimum of training loss on deep neural networks.
The results are shown in Figure.1.
Readers can refer more details in Appendix 7.1.
We have the following observations: (1) The traces of covariance matrices for the deep neural [Figure] networks can be well approximated by quadratic curves, which supports Proposition 1.
The minimum of the quadratic curve is nearly located at the local minimum w ∗ .
It indicates that the coefﬁcient of the ﬁrst-order term ρ g,H ≈ 0 .
Based on the fact that ρ g,H is not the determinant factor of the tail of the distribution in Eq.6 and the observations in Figure.1, we consider a simpliﬁed form of C ( w ) that C ( w ) =
σ g + σ H ( w − w ∗ ) 2 .
Corollary
If C ( w ) = σ g + σ H ( w − w ∗ ) 2 , the stationary distribution of 1-dimensional power-law dynamic (Eq.4) is
[Equation]
where Z is the normalization constant and κ = Hησ H is the tail-index.
The distribution density in Eq.7 is known as the power-law κ distribution (Zhou & Du, 2014)
(It is also named as q -Gaussian distribution in (Tsallis & Bukman, 1996)).
As κ → ∞ , the distribution density tends to be Gaussian, i.e., p ( w ) ∝ exp( − H ( w − w ∗ ) 2 ησ g ) .
Power-law κ distribution becomes more heavy-tailed as κ becomes smaller.
Meanwhile, it produces higher probability to appear values far away from the center w ∗ .
Intuitively, smaller κ helps the dynamic to escape from local minima faster.
In the approximation of dynamic of SGD, κ equals the signal (i.e., H ( w ∗ ) ) to noise (i.e., ησ H ) ratio of second-order derivative at w ∗ in SGD, and κ is linked with three factors: (1) the curvature
H ( w ∗ ) ; (2) the ﬂuctuation of the curvature over training data; (3) the hyper-parameters including η and minibatch size b .
Please note that σ H linearly decreases as the batch size b increases.
3. M ULTIVARIATE P OWER -L AW D YNAMIC
In this section, we extend the power-law dynamic to d -dimensional case.
We ﬁrst illustrate the covariance matrix C ( w ) of gradient noise in SGD.
We use the subscripts to denote the element in a vector or a matrix.
We use Σ g to denote the covariance matrix of ˜ g ( w ∗ ) and assume that Σ g is isotropic (i.e., Σ g = σ g · I ).
We also assume that Cov ( ˜ H i ( w ∗ ) , ˜ H j ( w ∗ )) are equal for all i, j .
It can be shown that C ( w ) =
Σ g (1 + ( w − w ∗ ) T Σ H Σ − 1 g ( w − w ∗ )) .
Similarly as 1-dimensional case, we omit the ﬁrst-order term ( w − w ∗ ) in C ( w ) .
Readers can refer Proposition 10 in Appendix 7.2 for the detailed derivation.
We suppose that the signal to noise ratio of ˜ H ( w ∗ ) can be characterized by a scalar κ , i.e., η Σ H =κ · H ( w ∗ ) .
Then C ( w ) can be written as
[Equation]
[Equation] Theorem 4 If w ∈ R d and C ( w ) has the form in Eq.(8) for w ∈ B ( w ∗ , (cid:15) ) .
The stationary distribution density of power-law dynamic is
[Equation]
[Equation] for w ∈ B ( w ∗ , (cid:15) ) , where Z is the normalization constant and κ satisﬁes η Σ H = 1 κ · H ( w ∗ ) .
Remark: The multivariate power-law κ distribution (Eq.9) is a natural extension of the 1-dimensional case.
Actually, the assumptions on Σ g and κ can be replaced by just assuming Σ g , H ( w ∗ ) , Σ H are codiagonalized.
Readers can refer Proposition 11 in Appendix 7.2 for the derivation.
E SCAPING E FFICIENCY OF P OWER - LAW D YNAMIC
In this section, we analyze the escaping efﬁciency of power-law dynamic from local minima and its relation to generalization.
Speciﬁcally, we analyze the mean escaping time for w t to escape from a basin.
As shown in Figure.2, we suppose that there are two basins whose bottoms are denoted as a and c respectively and the saddle point b is the barrier between two basins.
The barrier height is denoted as ∆ L = L ( b ) − L ( a ) .
Deﬁnition 5 Suppose w t starts at the local minimum a , we denote the time for w t to ﬁrst reach the saddle point b as inf { t > 0 | w
= a, w t = b } .
The mean escaping time τ is deﬁned as τ =
E w t [inf { t > 0 | w 0 = a, w t = b } ] .
We ﬁrst give the mean escaping time for 1-dimensional case in Lemma 6 and then we give the mean escaping time for high-dimensional power-law dynamic in Theorem 7. To analyze the mean escaping time, we take the following assumptions.
Assumption 1: The loss function around critical points can be written as L ( w ) =
L ( w ∗ ) + 12 ( w − w ∗ )
T H ( w ∗ )( w − w ∗ ) , where w ∗ is a critical point.
Assumption 2: The system is in equilibrium near minima, i.e., ∂p ( w,t ) ∂t = 0 . Assumption 3: (Low temperature assumption) The gradient noise is small, i.e., ησ g (cid:28) ∆ L . These three assumptions are commonly used in analyzing escaping time (Xie et al. , 2020;
Zhou & Du, 2014) for a dynamic.
Because both a and b are critical points, we can apply Assumption 1 to get the loss surface around them. We put more discussions about the assumptions in Appendix 7.3.2.
We suppose the basin a is quadratic and the variance of noise has the form that C ( w ) =
σ g a + σ H a ( w − a ) 2 , which can also be written as C ( w )
=
σ g a + 2 σ Ha H a ( L ( w ) − L ( a )) .
Furthermore, we suppose that
C ( w ) =
σ g a + 2 σ Ha H a ( L ( w ) − L ( a )) on the whole escaping path from a to b (not just near the local minimum a ).
It means that the variance of gradient noise becomes larger as the loss becomes larger.
The following lemma gives the mean escaping time of power-law dynamic for 1-dimensional case.
Lemma [-
 We suppose-]6 Suppose that Assumption 1-3 are satisﬁed and C ( w ) =
σ g a + 2 σ Ha H a ( L ( w ) − L ( a )) on the whole escaping path from a to b .
The mean escaping time of 1-dimensional power-law dynamic is,
[Equation]
where κ = H a ησ Ha > 12 , H a and H b are the second-order derivatives of training loss at local minimum a and at saddle point b , respectively.
The proof of Lemma 6 is based on the results in (Zhou & Du, 2014).
We provide a full proof in
Appendix 7.3.1.
For the dynamic near the saddle point, we just assume that its dynamic is the same as that near the local minimum for simplicity.
This assumption is not necessary and we put the extension to more complex dynamic in Appendix 7.3.3.
We summarize the mean escaping time of power-law dynamic and dynamics in previous works in Table 1.
Based on the results, we have the following discussions.
Comparison with other dynamics: (1) Both power-law dynamic and Langevin dynamic can escape sharp minima faster than ﬂat minima, where the sharpness is measured by H a and larger H a corresponds to sharper minimum.
Power-law dynamic improves the order of barrier height (i.e., ∆ L ) from exponential to polynomial compared with Langevin dynamic, which implies a faster escaping efﬁciency of SGD to escape from deep basin.
The mean escaping time for α -stable process is independent with the barrier height, but it is in polynomial order of the width of the basin (i.e.,
[Table] width= | b − a | ). Compared with α -stable process, the result for power-law dynamic is superior in the sense that it is also in polynomial order of the width (if ∆ L ≈ O ( | b − a | 2 ) ) and power-law dynamic does not rely on the inﬁnite variance assumption.
Based on Lemma 6, we analyze the mean escaping time for d -dimensional case.
Under the low temperature condition, the probability density concentrates only along the most possible escaping paths in the high-dimensional landscape.
For rigorous deﬁnition of most possible escaping paths, readers can refer section 3 in (Xie et al. , 2020).
For simplicity, we consider the case that there is only one most possible escaping path between basin a and basin c. Speciﬁcally, the Hessian at saddle point b has only one negative eigenvalue and the most possible escaping direction is the direction corresponding to the negative eigenvalue of the Hessian at b . Theorem 7 Suppose that Assumption 1-3 are satisﬁed.
For w ∈ R d , we suppose C ( w ) =
Σ g a +ηκ ( L ( w ) − L ( a )) on the whole escaping path from a to b and there is only one most possible path path between basin a and basin c .
The mean escaping time for power-law dynamic escaping from basin a to basin c is
[Equation]
where e indicates the most possible escaping direction, H be is the only negative eigenvalue of H b , σ e is the eigenvalue of Σ g a that corresponds to the escaping direction, ∆ L = L ( b ) − L ( a ) , and det( · ) is the determinant of a matrix.
Remark:
In d -dimensional case, the ﬂatness is measured by det( H a ) .
If H a has zero eigenvalues, we can replace H a by H + a in above theorem, where H + a is obtained by projecting H a onto the subspace composed by the eigenvectors corresponding to the positive eigenvalues of H a .
This is because by
Taylor expansion, the loss L ( w ) only depends on the positive eigenvalues and the corresponding eigenvectors of H a , i.e., L ( w ) = L ( a ) [Equation]
[Equation]
) , where Λ H + a is a diagonal matrix composed by non-zero eigenvalues of H a and the operator P ( · ) operates the vector to the subspace corresponding to non-zero eigenvalues of H a .
Therefore, the dimension d in Theorem 7 can be regarded as the dimension of subspace that is composed by directions with large eigenvalues.
It has been observed that most of the eigenvalues in H is very small (Sagun et al. , 2016).
Therefore, d will not be a large number and power-law dynamic in multi-dimensional case will inherit the beneﬁt of that in 1-dimensional case compared with Langevin dynamic and α -stable process.
The next theorem give an upper bound of the generalization error of the stationary distribution of power-law dynamic, which shows that ﬂatter minimum has smaller generalization error.
Theorem 8 Suppose that w ∈ R d and κ > d 2 .
For δ > 0 , with probability at least 1 − δ , the stationary distribution of power-law dynamic has the following generalization error bound,
[Equation]
where KL ( p || p (cid:48) ) ≤ 12 log det( H ) det(Σ g )
+ Tr ( η Σ g H − 1 ) − 2 d 4 ( 1 − 1 κ ( d 2 − 1 ))
+ d 2 log 2 η , p ( w ) is the stationary distribution of d -dimensional power-law dynamic, p (cid:48)
( w ) is a prior distribution which is selected to be standard
Gaussian distribution, and P ( x ) is the underlying distribution of data x , det( · ) and Tr ( · ) are the determinant and trace of a matrix, respectively.
We make the following discussions on results in Theorem 8. For 1-dimensional case, we have if H > η 2(1+ 12 κ ) , KL divergence is decreasing as H decreases. For d > 1 and ﬁxed Tr (Σ g H − 1 ) and det(Σ g ) , the generalization error (i.e., E w ∼ p ( w ) ,x ∼P ( x ) (cid:96) ( w, x ) − E w ∼ p ( w ) L ( w ) ) is decreasing as det( H ) decreases, which indicates that ﬂatter minimum has smaller generalization error.
Moreover, if 2 d > Tr ( η Σ g H − 1 ) , the generalization error is decreasing as κ increases.
When κ → ∞ , the generalization error tends to that for Langevin dynamic.
Combining the mean escaping time and the generalization error bound, we can conclude that state-dependent noise makes SGD escape from sharp minima faster and implicitly tend to learn a ﬂatter model which generalizes better.
E XPERIMENTS
In this section, we conduct experiments to verify the theoretical results.
We ﬁrst study the ﬁtness between parameter distribution trained by SGD and power-law κ distribution.
Then we compare the escaping behavior for power-law dynamic, Langevin dynamic and SGD.
5. F ITTING P ARAMETER D ISTRIBUTION USING P OWER -L AW D ISTRIBUTION
We investigate the distribution of parameters trained by SGD on deep neural networks and use power-law κ distribution to ﬁt the parameter distribution.
We ﬁrst use SGD to train various types of deep neural networks till it converge.
For each network, we run SGD with different minibatch sizes over the range { 64 , 256 , 1024 } .
For the settings of other hyper-parameters, readers can refer
Appendix 7.5.2.
We plot the distribution of model parameters at the same layer using histogram.
Next, we use power-law κ distribution to ﬁt the distribution of the parameters and estimate the value of κ via the embedded function " TsallisQGaussianDistribution [] " in Mathematica software.
We show results for LeNet-5 with MNIST dataset and ResNet-18 with CIFAR10 dataset (LeCun et al. , 2015;
He et al. , 2016b) in this section, and put results for other network architectures in Appendix 7.5.2.
In Figure 3, we report the generalization error (i.e., Test error - Training error) and the values of κ that best ﬁt the histogram.
We have the following observations: (1) The distribution of the parameter trained by SGD can be well ﬁtted by power-law κ distribution (blue curve).
(2) As the minibatch size becomes larger, κ becomes larger.
It is because the noise σ H linearly decreases as minibatch size becomes larger and κ = Hησ H .
(3) As κ becomes smaller, the generalization error becomes lower.
It indicates that κ also plays a role as indicator of generalization.
These results are consistent with the theory in Section 4.
5. C OMPARISON ON E SCAPING E FFICIENCY
We use a 2-dimensional model to simulate the escaping efﬁciency from minima for power-law dynamic, Langevin dynamic and SGD.
We design a non-convex 2-dimensional function written as L ( w ) = 1 n (cid:80) ni
=1 (cid:96) ( w − x i ) , where (cid:96) ( w ) =
15 (cid:80)
j =1 | w j − 1 | 2 . 5 · | w j + 1 | 3 and training data x i ∼ N (0 , 0 .
01 I 2 ) .
We regard the following optimization iterates as the numerical discretization of the power-law dynamic, w t +1 =
w t − ηg ( w t ) + ηλ 2 (cid:112) 1 + λ 1 ( w t − w ∗ ) 2 (cid:12) ξ , where ξ ∼ N (0 , I 2 ) , λ 1 , λ 2 are two hyper-parameters and (cid:12) stands for Hadamard product.
Note that if we set λ 1 = 0 , it can be regarded as discretization of Langevin dynamic.
We set learning rate η = 0 . 025 , and we take 500 iterations in each training.
In order to match the trace of covariance matrix of stochastic gradient at minimum point w ∗ with the methods above, λ 2 is chosen to satisfy Tr ( Cov ( λ 2 ξ )) =
Tr ( Cov ( g ( w ∗ ))) .
We compare the success rate of escaping for power-law dynamic, Langevin dynamic and SGD by repeating the experiments 100 times.
To analyze the noise term λ 1 , we choose different λ 1 and evaluate corresponding success rate of escaping, as shown in Figure.4(c).
The results show that: (1) there is a positive correlation between λ 1 and the success rate of escaping; (2) power-law dynamic can mimic the escaping efﬁciency of SGD, while Langevin dynamic can not.
We then scale the loss [Figure] [Figure] function by 0 .
9 to make the minima ﬂatter and repeat all the algorithms under the same setting.
The success rate for the scaled loss function is shown in Figure.4(d).
We can observe that all dynamics escape ﬂatter minima slower.
[Figure]
C ONCLUSION
In this work, we study the dynamic of SGD via investigating state-dependent variance of the stochastic gradient.
We propose power-law dynamic with state-dependent diffusion to approximate the dynamic of SGD.
We analyze the escaping efﬁciency from local minima and the PAC-Bayes generalization error bound for power-law dynamic.
Results indicate that state-dependent noise helps SGD escape from poor local minima faster and generalize better.
We present direct empirical evidence to support our theoretical ﬁndings.
This work may motivate many interesting research topics, for example, nonGaussian state-dependent noise, new types of state-dependent regularization tricks in deep learning algorithms and more accurate characterization about the loss surface of deep neural networks.
We will investigate these topics in future work.
A PPENDIX 7. P OWER - LAW D YNAMIC AND S TATIONARY D ISTRIBUTION
Theorem 9 (Theorem 2 in main paper)
The stationary distribution density for 1-dimensional power- law dynamic (Eq.4) is
[Equation]

 [Equation] where C ( w ) = σ g +2 ρ g,H ( w − w ∗ )+ σ H ( w − w ∗ ) 2 , Z is the normalization constant and ArcTan ( · ) is the arctangent function.
Proof: We denote the function H ( 4 ρ g,H · ArcTan ( C (cid:48)
[Equation]
[Equation]
H ))
[Equation]
as h ( w ) .
According to the
Fokker-Planck equation, p ( w ) satisﬁes
[Equation]
Readers can check the third equality by calculating ∇ ( C ( w ) HησH · e − h ( w ) · p ( w )) with C ( w )
=
σ g + 2 ρ g,H ( w − w ∗ )+ σ H ( w − w ∗ ) 2 .
Because the left side equals zero, we have C ( w ) HησH · e − h ( w ) · p ( w ) equals to constant.
So p ( w ) ∝ C ( w ) − HησH · e h ( w ) · p ( w ) .
So we can get the conclusion in the theorem.
(cid:3)
Theorem 10 (Corollary 3 in main paper)
If C ( w ) = σ g + σ H ( w − w ∗ ) 2 , the stationary distribution density of power-law dynamic is
[Equation]
where Z = (cid:82) w (1 + σ H σ − 1 g ( w − w ∗ ) 2 ) − κ dw is the normalization constant and κ =
Hησ H is the tail-index.
Proof: According to the Fokker-Planck equation, p ( w ) satisﬁes
[Equation]
Because the left side equals zero, we have (1 + 2 σ H [Equation]
p ( w ) equals to constant.
So p ( w ) ∝ (1 + [Equation] σ H
[Equation] g ( L ( w ) − L ( w ∗ )))
H − ησH .
So we can get the conclusion in the theorem.
(cid:3)
We plot the un-normalized distribution density for 1-dimensional power-law dynamics with different κ in Figure 5.
For the four curves, we set β = 10 .
We set κ = 1 , 0 . 5 , 0 .
and use green, red, purple and blue line to illustrate their corresponding density function, respectively.
When κ = 0 , it is
Gaussian distribution.
From the ﬁgure, we can see that the tail for power-law κ -distribution is heavier than Gaussian distribution.
Actually, for any given time t , the distribution p ( w, t ) for w t that satisﬁes power-law dynamic has analytic form, i.e., p ( w, t ) ∝ (1 +
H ηκσ ( t ) ( w − w ( t )) 2 ) − κ , where w ( t ) =
w ∗ + ( w 0 − w ∗ ) e − Ht
and σ ( t ) is a function of σ g and t .
Readers can refer Eq.18 - Eq.23 in (Tsallis & Bukman, 1995) for the detailed expression.
7. SGD AND M ULTIVARIATE P OWER - LAW D YNAMIC
The following proposition shows the covariance of stochastic gradient in SGD in d -dimensional case.
We use the subscripts to denote the elements in a vector or a matrix.
Proposition 11 For w ∈ R d , we use C ( w ) to denote the covariance matrix of stochastic gradient ˜ g ( w ) =
˜ g ( w ∗ )
+ (cid:101) H ( w − w ∗ ) and Σ to denote the covariance matrix of ˜ g
( w ∗ ) .
If Cov (˜ g i ( w ∗ ) , (cid:101) H jk ) = 0 , ∀ i, j, k , we have
[Equation]
(13) where Σ ij =
Cov (˜ g i ( w ∗ ) , ˜ g j ( w ∗ )) , A ( ij ) is a d × d matrix with elements A ( ij ) ab
=
Cov ( ˜ H ia , ˜ H jb ) with a ∈ [ d ] , b ∈
[ d ] .
Eq.13 can be obtained by directly calculating the covariance of ˜
g
i ( w ) and ˜ g j ( w ) where ˜ g
i ( w ) = ˜
g i ( w ∗ ) + (cid:80) da =1 ˜ H ia ( w a − w ∗ a ) , ˜ g j ( w ) = ˜ g j ( w ∗ ) + (cid:80) db =1 ˜ H jb ( w b − w ∗ b ) .
In order to get a analytic tractable form of C ( w ) , we make the following assumptions: (1) If Σ ij = 0 ,
A ( ij ) is a zero matrix; (2) For Σ ij (cid:54)
= 0 , A ( ij ) Σ ij are equal for all i ∈
[ d ] , j ∈ [ d ] .
The ﬁrst assumption is reasonable because both Σ ij and A ( ij ) reﬂect the dependence of the derivatives along the i -th direction and j -th direction.
Let Σ H =
A ( ij ) Σ ij , C ( w ) can be written as C ( w ) =
Σ g (1+( w − w ∗ ) T Σ H ( w − w ∗ )) .
The d -dimensional power-law dynamic is written as
[Equation]
(14) where C ( w ) =
Σ g (1 + ( w − w ∗ ) T Σ H ( w − w ∗ )) which is a symmetric positive deﬁnite matrix that
C ( w ) 1 / 2 exists.
The following proposition shows the stationary distribution of the d -dimensional power-law dynamic.
Proposition 12 Suppose Σ g , Σ H , H are codiagonalizable, i.e., there exist orthogonal matrix Q and diagonal matrices Λ , Γ , Π to satisfy Σ g = Q T Λ Q, Σ H
= Q T Γ Q, H = Q T Π Q .
Then, the stationary distribution of power-law dynamic is
[Equation]
(15) where Z is the normalization constant and κ =
[Equation]
Proof: Under the codiagonalization assumption on Σ g , Σ H , H , Eq.15 can be rewritten as dv t = − Π v t dt + (cid:112)
η Λ(1 + v Tt Γ v t )
dB t if we let v t = Q ( w t − w ∗ ) .
We use φ ( v ) =
ηC ( v )=
η 2 Λ(1 + v T Γ v ) , the stationary probability density p ( v ) satisﬁes
the
Smoluchowski equation:
[Equation]
According to the result for 1-dimensional case, we have the expression of p ( v ) is p ( v )
∝ (1 + v T Γ v ) − κ .
To determine the value of κ , we put p ( v ) in the Smoluchowski equation to obtain
[Equation]
= (cid:16) η Λ
i κ 2 (1 + v T Γ v ) − κ − 1 · (Γ i v i ) 2 (cid:17) .
The we have (cid:80) di =1 Π i = ηκ (cid:80) di
=1 Λ i Γ i .
So we have κ =
Tr ( H ) ηTr (Σ H Σ g ) .
(cid:3)
According to Proposition 11, we can also consider another assumption on Σ g , Σ H , H without assuming their codiagonalization.
Instead, we assume (1) If Σ ij = 0 , A ( ij ) is a zero matrix; (2) For Σ ij (cid:54) = 0 , A ( ij ) are equal for all i ∈
[ d ] , j ∈ [ d ] and we denote A ( ij ) = Σ H .
We suppose η · Σ H = κH .
(3) Σ g
= σ g · I d which is isotropic.
Under these assumptions, we can get the following theorem.
Theorem 13 (Theorem 4 in main paper)
If w is d -dimensional and C ( w ) has the form in Eq.(8).
The stationary distribution density of multivariate power-law dynamic is
[Equation]
(18) where Z = (cid:82) ∞
−∞ [1 + 1 ηκ ( w − w ∗ ) T H Σ − 1 g ( w − w ∗ )] − κ dw is the normalization constant.
The proof for Theorem 12 is similar to that for Proposition 11.
Readers can check that p ( w ) satisﬁes the Smoluchowski equation.
An example to illustrate why C ( w ) is diagonally dominant. In Theorem 13, C ( w ) is assumed to be diagonally dominant. Diagonally dominant indicates that the variance of each dimension of ˜ g ( w ) is signiﬁcantly larger than the covariance of two different dimensions of ˜ g ( w ) .
Consider a two layer fully-connected linear neural network f w,v ( x ) = wvx where w ∈ R 1 × m , v ∈ R m × d , x ∈ R d and h ( · ) is the ReLU activation.
We consider the regression loss (cid:96) ( w, v ) =
12 ( y − f w,v ( x )) 2 .
The gradient of w i and v jk can be written as
[Equation]
) where v i denotes the i-th row of matrix v .
Suppose that the initialization of w and v is: w i i.i.d ∼
N (0 , δ 1 ) and v ij i.i.d ∼ N (0 , δ 2 ) .
We also assume that E x i
=
E x j = 0 and x i , x j are independent with each other for i (cid:54) =
j where x i is the i -th dimension.
We have
E w,v ∂(cid:96)
[Equation]
[Equation]
[Equation]
[Equation]
Because the independence of v i , v j and their expectations are zero, we can obtain
E w,v ∂(cid:96) ( w,v ) ∂w i ∂(cid:96)
( w,v )
∂w
j = 0 for i (cid:54)
=
j .
Similarly, we can get E w,v ∂(cid:96) ( w,v ) ∂w i ∂(cid:96) ( w,v ) ∂v jk = 0 and
E w,v ∂(cid:96) ( w,v ) ∂v j (cid:48) k (cid:48)
∂(cid:96) ( w,v ) ∂v jk
= 0 for ( j, k ) (cid:54) =
( j (cid:48) , k (cid:48) ) .
The above analyses show that the gradients for different dimensions are independent at initialization.
It has been observed that many weights are kept random during training because of the over-parameterization Balduzzi et al.
So, diagonalization dominant property of C ( w ) is reasonable.
7. S UPPLEMENTARY MATERIALS FOR RESULTS IN S ECTION7.3. P ROOF FOR M EAN E SCAPING T IME
Lemma 14 (Lemma 6 in main paper)
We suppose C ( w ) =
σ g a + 2 σ Ha H a ( L ( w ) − L ( a )) on the whole escaping path from a to b .
The mean escaping time of the 1-dimensional power-law dynamic is,
[Equation]
where κ = H a ησ Ha , H a , H b are the second-order derivatives of training loss at local minimum a and saddle point b .
Proof: According to (Van Kampen, 1992), the mean escaping time τ is expressed as τ = P ( w ∈
[Equation]
Ω , where V a is the volume of basin a , J is the probability current that satisﬁes −∇
[Equation]
where φ ( w ) =
η 2 C ( w ) and µ = 2 σ Ha H a , σ g =
σ g a and ∆ L ( w ) =
L ( w ) − L ( a ) .
Integrating both sides, we obtain J ( w ) =
− φ ( w ) · (cid:16) 1 + µσ g ∆ L
w ) (cid:17) − κ ∂ [Equation] µσg
[Equation] L [Equation] κ p
[Equation] ∂w . Because there is no ﬁeld source on the escape path, J ( w ) is ﬁxed constant on the escape path.
Multiplying φ ( w ) − 1 · (cid:16) 1 + µσ g ∆ L ( w ) (cid:17) κ on both sizes, we have
[Equation]
Then we get J =
[Equation]
As for the term
[Equation]
we have
[Equation]
, where the third formula is based on the second order Taylor expansion.
Under the low temperature assumption, we can use the second-order Taylor expansion around the saddle point b .
As for the term P ( w ∈ V a ) , we have P ( w ∈ V a ) =
(cid:82) V a p ( w ) dV
= (cid:82) w ∈ V a p ( a )(1 + µσ g ∆ L ( w ))
− κ = p ( a ) (cid:113) 2 σ g µH a B ( 12 , κ − 12 ) , where we use Taylor expansion of L ( w ) near local minimum a .
Then we have τ = P ( w ∈ V a ) (cid:82)
[Equation]
) J because J is a constant.
Combining all the results, we can get the result in the lemma.
[Equation]
Theorem 15 (Theorem 7 in main paper) Suppose w ∈ R d
and there is only one most possible path path between basin a and the outside of basin a .
The mean escaping time for power-law dynamic escaping from basin a to the outside of basin a is
[Equation]
where e indicates the most possible escape direction, H be is the only negative eigenvalue of H b , σ e is the eigenvalue of Σ g a corresponding to the escape direction and ∆ L = L ( b ) − L ( a ) .
Proof: According to (Van Kampen, 1992), the mean escaping time τ is expressed as τ = P ( w ∈ V a ) (cid:82) Ω Jd Ω , where V a is the volume of basin a , J is the probability current that satisﬁes −∇ · J ( w, t ) =
∂p ( w,t ) ∂t .
Under the low temperature assumption, the probability current J concentrates along the direction corresponding the negative eigenvalue of H be , and the probability ﬂux of other directions can be ignored.
Then we have
[Equation] [Equation]
where J e = p ( a )
[Equation]
) which is obtained by the calculation of J e for 1-dimensional case in the proof of Lemma 13, and ( · ) ⊥ e denotes the directions perpendicular to the escape direction e.
Suppose H b Σ − 1 g are symmetric matrix.
Then there exist orthogonal matrix Q and diagonal matrix Λ = diag ( λ 1 , · · · , λ d ) that satisfy H b Σ − 1 g = Q T Λ Q .
We also denote v = Q ( w − b ) .
We deﬁne a sequence as T k
= 1 + 1 ηκ · (cid:80) dj =
k λ j v 2 j for k = 1 , · · · , d .
As for the term (cid:82) Ω (cid:16) 1 + 1 ηκ ( w − b ) T ( H b Σ − 1 g ) ⊥ e ( w − b ) (cid:17) − κ + 12 d Ω , we have (cid:90)
[Equation]
(cid:18)
[Equation]
As for the term P ( w ∈ V a ) , we have
[Equation]
[Equation]
where we use Taylor expansion of L ( w ) near local minimum a .
Combined the results for P ( w ∈ V a ) and J , we can get the result.
(cid:3)
7.3. F URTHER E XPLANATION ABOUT A SSUMPTION 1-
We adopt the commonly used assumptions to analyze mean escaping time for dynamic system (Xie et al. , 2020; Smith & Le, 2017; Zhou & Du, 2014).
Assumption 2 can be replaced by weaker assumption that the system is quasi-equilibrium which is adopted in (Xie et al. , 2020).
For the differences between quasi-equilibrium and equilibrium, readers can refer to (Xie et al. , 2020) for detailed discussions.
Assumption 3 is commonly used (Xie et al. , 2020; Zhou & Du, 2014).
Under
Assumption 3, the probability densities will concentrate around minima and the most possible paths.
Assumption 3 will make the second order Taylor approximation more reasonable.
7.3. E XTENSION TO MORE COMPLEX DYNAMIC ON THE ESCAPING PATH
In Lemma 6, we assume that C ( w ) =
σ g a + 2 σ Ha H a ( L ( w ) − L ( a )) on the whole escaping path from a to b for ease of comparison and presentation.
This assumption is not necessary and we can assume a different dynamic near saddle point b .
Specially, we can assume the point z is the midpoint on the most possible path beween a and b , where L ( z ) =
(1 − z ) L ( a ) + zL ( b ) .
The dynamic with C ( w ) =
σ g a + 2 σ Ha H a ( L ( w ) − L ( a )) dominates the path a → z and the dynamic with
C ( w ) =
σ g b + 2 σ Hb H b ( L ( b ) − L ( w )) dominates the path z
→ b .
Then only two things will be changed in proof of Lemma 6.
First, we need to change the stationary distribution near saddle points according to its own dynamic in Eq.20.
Second, we need to change the integral about probability density on the whole path to sum of integrals on these two sub-paths.
Similar proof techniques are adopted for analyzing escaping time of Langevin dynamic in proof of Theorem 4.1 in the work Xie et al.
Since the proof is analogous, we omit the details here.
7. PAC-B AYES G ENERALIZATION B OUND
We brieﬂy introduce the basic settings for PAC-Bayes generalization error.
The expected risk is deﬁned as E x ∼P ( x ) (cid:96) ( w, x ) .
Suppose the parameter follows a distribution with density p ( w ) , the expected risk in terms of p ( w ) is deﬁned as E w ∼ p ( w ) ,x ∼P ( x ) (cid:96) ( w, x ) .
The empirical risk in terms of p ( w ) is deﬁned as E w ∼ p ( w ) L ( w ) =
E w ∼ p ( w ) 1 n (cid:80) ni =1 (cid:96) ( w, x i ) .
Suppose the prior distribution over the parameter space is p (cid:48)
( w ) and p ( w ) is the distribution on the parameter space expressing the learned hypothesis function.
For power-law dynamic, p ( w ) is its stationary distribution and we choose p (cid:48)
( w ) to be Gaussian distribution with center w ∗ and covariance matrix I .
Then we can get the following theorem.
Theorem 16 (Theorem 8 in main paper) For w ∈ R d , we select the prior distribution p (cid:48) ( w ) to be standard Gaussian distribution.
For δ > 0 , with probability at least 1 − δ , the stationary distribution of power-law dynamic has the following generalization error bound,
[Equation]
where KL ( p || p (cid:48) [Equation] ≤ 12 log det
[Equation]
and P ( x ) is the underlying distribution of data x .
Proof: Eq.(29) directly follows the results in (McAllester, 1999).
Here we calculate the Kullback–Leibler (KL) divergence between prior distribution and the stationary distribution of power-law dynamic.
The prior distribution is selected to be standard Gaussion distribution with distribution density p (cid:48)
( w ) =
[Equation]
det( I ) exp {
[Equation]
12 ( w
[Equation]
w
[Equation]
) T I ( w − w ∗ ) } .
The posterior distribution density is the stationary distribution for power-law dynamic, i.e.,
[Equation]
Suppose H Σ − 1 g are symmetric matrix.
Then there exist orthogonal matrix Q and diagonal matrix Λ = diag ( λ 1 , · · · , λ d ) that satisfy H Σ − 1 g = Q T Λ Q .
We also denote v = Q ( w − w ∗ ) .
We have log
[Equation]
The KL-divergence is deﬁned as KL ( p ( w ) || p (cid:48) ( w )) =
(cid:82) w p ( w ) log (cid:16) p ( w ) p (cid:48)
( w ) (cid:17) dw .
Putting v = Q ( w − w ∗ ) in the integral, we have
KL
[Equation]
p
[Equation]
where we use the approximation that log(1 + x ) ≈ x .
We deﬁne a sequence as T k
= 1 + 1 ηκ · (cid:80) dj =
k λ j v 2 j for k = 1 , · · · , d .
We ﬁrst calculate the normalization constant Z .
[Equation]
We deﬁne Z j =
(( ηκ ) − 1 λ j )
− 12 B (cid:0) 12 , κ − j 2 (cid:1) .
For the third term in Eq.(30), we have
[Equation]
[-j = v 2 j (cid:33) (cid:32) 
-]
[-+ 1 ηκ · d (cid:88) j = λ j v 2 j (cid:33) − κ +d v 2 ··· ,v d = (cid:90) v 2 , ··· ,v d T − κ 2 (cid:90) (cid:18) T( ηκ ) − 1 λ(cid:19) 32 y 12 (1 + y ) − κ dy + Z 1 (cid:32) d (cid:88) j = v 2 j (cid:33) (cid:32) 
-][Table]
[-j = λ j v 2 j (cid:33) − κ +d v 2 ··· ,v d = (cid:90) v 2 , ··· ,v d (( ηκ ) − 1 λ 1 ) − 32 T − κ + 
-]
[-+ Z 1 (cid:32) d (cid:88) j = v 2 j (cid:33) (cid:32) 
 + 1 ηκ · d (cid:88) j = λ j v 2 j (cid:33) − κ-] +d v 2 ··· ,v d
For term (cid:82) v 2 [-j (cid:33) (cid:32) 
 + 1 ηκ ·-], ··· ,v d T − 1 κ + 32d v 2 ··· ,v d in above equation, we have (cid:90)
[Equation]
=
[Equation]
Let A j = (( ηκ ) − 1 λ j ) − 32 B (cid:0) 32 , κ − ( j 2 + 1) (cid:1) .
According to the above two equations, we can get the recursion
[Equation]
We have
[Equation]
Similarly, for the fourth term in Eq.(30) [Equation] we have IV = κd 2( κ − d 2 − 1) .
Combining all the results together, we can get KL ( p || p (cid:48) )
= 12 log
[Equation]
2 log 2 .
Using the fact that log Γ( κ ) Γ( κ − d 2 ) ≤ d 2 log κ
[Equation]
we have KL ( p || p
[Equation]
≤
[Equation]
log
[Equation]
7. I MPLEMENTATION D ETAILS OF THE E XPERIMENTS 7.5. O BSERVATIONS ON THE C OVARIANCE M ATRIX
In this section [Equation] we introduce the settings on experiments of the quadratic approximation of covariance of the stochastic gradient on plain convolutional neural network (CNN) and ResNet.
For each model, we use gradient descent with small constant learning rate to train the network till it converges.
The converged point can be regarded as a local minimum, denoted as w ∗ .
As for the detailed settings of the CNN model, the structure for plain CNN model is input →
Conv 1 → maxpool → Conv 2 → maxpool → fc 1 → Relu → fc 2 → output .
Both Conv 1 and
Conv 2 use 5 × 5 kernels with 10 channels and no padding.
Dimensions of full connected layer fc 1 and fc 2 are 1600 × 50 and 50 × 10 respectively.
We randomly sample 1000 images from
FashionMNIST (Xiao et al. , 2017) dataset as training set.
The initialization method is the Kaiming initialization (He et al. , 2015) in PyTorch.
The learning rate of gradient descent is set to be 0 .
After 3000 iterations, GD converges with almost 100% training accuracy and the training loss being 1 e − 3 .
As for ResNet, we use the ResNet-18 model (He et al. , 2016b) and randomly sample 1000 images from
Kaggle’s dogs-vs-cats dataset as training set.
The initialization method is the Kaiming initialization (He et al. , 2015) in PyTorch.
The learning rate of gradient descent is set to be 0 .
Afteriterations, GD converges with 100% training accuracy and the training loss being 1 e − 3 .
We then calculate the covariance matrix of the stochastic gradient at some points belonging to the local region around w ∗ .
The points are selected according to the formula: w ∗ layerL ± ( i × Scale ) , where w ∗ layerL denotes the parameters at layer L , and i × Scale,
i ∈
[ N ] determines the distance away from w ∗ layerL .
When we select points according to this formula by changing the parameters at layer L , we ﬁxed the parameters at other layers.
For both CNN model and ResNet18 model, we select 20 points by setting i = 1 , · · · , 10 .
For example, for CNN model, we choose thepoints by changing the parameters at the Conv 1 layer with Scale = 0 .
001 and Conv 2 layer with
Scale = 0 . 0001 , respectively.
For ResNet18, we choose the 20 points by changing the parameters for a convolutional layer at the ﬁrst residual block with Scale = 0 .
0001 and second residual block with Scale = 0 . 0001 , respectively.
The results are shown in Figure.1.
The x-axis denotes the distance of the point away from the local minimum and the y-axis shows the value of the trace of covariance matrix at each point.
The results show that the covariance of noise in SGD is indeed not constant and it can be well approximated by quadratic function of state (the blue line in the ﬁgures), which is consistent with our theoretical results in Section 3.1.
. S UPPLEMENTARY E XPERIMENTS ON P ARAMETER D ISTRIBUTIONS OF D EEP N EURAL N ETWORKS
[Figure]
Figure. 3(a), which could help to observe the heavy-tailed properties among different batchsize.
[Figure] For Figure. 3(a), we train LeNet-5 on MNIST dataset using SGD with constant learning rate η = 0 .for
each batchsize till it converges.
Parameters are conv 2 .weight in LeNet-5.
For Figure 3(b), we train ResNet-18 on CIFAR10 using SGD with momentum.
We do a RandomCrop on training set scaling to 32 × 32 with padding =
and then a RandomHorizontalFlip .
In training, momentum is set to be 0 .
9 and weight decay is set to be 5 e − 4 .
Initial learning rate in SGD is set to be 0 .
and we using a learning rate decay of 0 .
1 on { 150 , 250 } -th epoch respectively train it converges
We also observe the parameter distribution on many pretrained models.
Details for pre-trained models can be found on https://pytorch.org/docs/stable/torchvision/models.html .
Figure.7 shows the distribution of parameters trained by SGD can be well ﬁtted by powerlaw distribution.
Parameters in this ﬁgure are all randomly selected to be features.10.weight , features.14.weight , features.
5 .expand 3 × 3 .weight , Mixed _ 6 d.branch 7 × 7 _ 3 .conv.weight , layer 4 .
2 .conv 3 .weight and features.denseblock 2 .denselayer 1 .conv 2 .weight for VGG-16,
AlexNet, SqueezeNet 1.0, Inception v3, Wide ResNet-50-2 and DenseNet-121 respectively.
A Q-Q plot is created by plotting quantiles of two probability distributions against one another, which can provide an assessment of "goodness of ﬁt" by how much the solid line close to the dashed line.
From Figure.8, it is clear that the solid lines in bottom pictures are closer to dashed lines on most cases, which indicates network parameters can be better ﬁtted by power-law distribution.
Moreover, solid lines in the upper plots severely deviate from dashed lines on the tail of distribution but those in the bottom plot do not, which means the distribution of parameters is indeed heavy-tailed.
F URTHER E XPLANATION ON E XPERIMENTS IN S ECTION 5.
As for the experiments for 2-D model, we also calculate coefﬁcient of the second-order term for the quadratic curve shown in Figure.4(b), and its value is roughly 30 , which matches the result in
Figure.4(c) in the sense that the result for SGD is similar with the result for power-law dynamic with λ 1 ≈ 32 . E SCAPING E FFICIENCY ON N EURAL [Figure]
We follow the settings in (Zhu et al. , 2019).
For convenience of the readers, here we give the details of this setting again.
We use corrupted FashionMNIST dataset which contains 1000 images with correct labels and another 200 images with random labels to be training data.
A small LeNet-like network with 11,330 parameters is used.
Firstly we run the full gradient decent to reach the parameters w ∗ near the global minima.
Then we continue training using both Langevin dynamic(GLD)
and power-law dynamic(PLD).
Following Zhu’s setting, the learning rates for GD, GLD and PLD are η GD = 0 .
1 , η GLD = 0 .
07 and η PLD = 0 . 07 , respectively.
For GLD, noise std σ = 10 − 4 as Zhu already tuned.
For our PLD, w t +1 =
w t − η ∇ L ( w t ) + η · α ∇ L ( w t ) (cid:12) (cid:112) 1 + β ( w t − w ∗ ) 2 (cid:12) ξ , where α, β are hyperparameters, ξ ∼ N (0 , I ) , and (cid:12) stands for Hadamard product.
Here we select α = 2 .
4 , β
= 2 after grid search.
Expected sharpness is measured as E ν ∼ N (0 ,δ 2 I )
[ L ( w + ν )] − L ( w ) where δ = 0 .
01 , and the expectation is computed by average on 1000 times sampling.
The numbers at the ﬁrst column of the legend show the test accuracy and the numbers in the bracket show the sharpness of the model trained by the three algorithms.
From Figure 9, we can conclude that PLD generalizes better than GLD and GD.
Moreover, PLD can ﬁnd ﬂatter critical points than
GLD and GD.
7.5. C OMPARISON OF M EAN ESCAPING TIME WITH DIFFERENT BARRIER HEIGHTS
[Figure]
We design this 1-dimensional model to help to validate the theoretical results of escaping time in
Table 1.
Loss function L ( w ) = 1 n (cid:80) ni
=1 (cid:96) ( w − x i ) , where (cid:96) ( w ) = (cid:26) w 2 + bw , w <
− w 2 + bw , w ≥ 0 and x i ∼ N (0 , 0 . 05) .
L ( w ) is plotted in Figure.10(a), and we can adjust barrier height through parameter b in (cid:96) ( w ) without changing the Hessian on minima w ∗ and saddle point.
For power-law dynamic (PLD), w t +1 =
w t − η ∇ L ( w t )+ ηλ 2 (cid:112) 1 + λ 1 ( w t − w ∗ ) 2 (cid:12) ξ , where λ 1 , λare hyperparameters, ξ ∼ N (0 , I ) , and (cid:12) stands for Hadamard product.
Here we let λ 1 = 1 , λ 2 = 4 .
For Langevin dynamic (GLD), we set noise std σ = 4 in consistence with PLD.
Learning rate η = 0 .for
both methods.
We initialize w 0 = w ∗ and apply both methods on L ( w ) with different barrier heights.
Then we record the number of iterations t when w t ﬁrstly escaping from the barrier.
We repeat this procedure 100 rounds for each method and each barrier height and utilize the average to estimate the mean escaping time, of which the results are shown in Figure.10(b).
From Figure.10(b), the mean escaping time of GLD grows much faster than PLD along barrier height, which validates that power-law dynamic improves the order of barrier height compared with Langevin dynamic.