Year: 2020

Abstract:
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 efficiency from local minima. However, these works make the over-simplified 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 efficiency 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 efficiently and tends to stop at flat minima that have lower generalization error. Finally, we conduct experiments to compare SGD and power-law dynamic, and the results verify our theoretical findings.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND
text: 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.

heading: 2.1 RELATED WORK
text: 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.

heading: 3 APPROXIMATING SGD BY POWER-LAW DYNAMIC
text: 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.

heading: 3.1 NOISE STRUCTURE OF STOCHASTIC GRADIENT DESCENT
text: 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

heading: 3.2 POWER-LAW DYNAMIC
text: According to CLT, the gradient noise follows Gaussian distribution if it has finite variance, i.e.,
g(w)− g̃(w)→d N (0, C(w)) 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:
wt+1 = wt − ηg(wt) + ηξt, ξt ∼ N (0, C(w)). (3)
Eq.3 can be treated as the discretization of the following SDE, which we call it power-law dynamic: dwt = −g(wt)dt+ √ ηC(w)dBt. (4)
Power-law dynamic characterizes how the distribution of w changes as time goes on. The distribution density of parameterw at time t (i.e., p(w, t)) is determined by the Fokker-Planck equation (Zwanzig’s type (Guo & Du, 2014)):
∂ ∂t p(w, t) = ∇p(w, t)g(w) + η 2 · ∇ (C(w) · ∇p(w, t)) . (5)
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
p(w) = 1
Z (C(w))
− H ησH exp H ( 4ρg,H ·ArcTan ( C′(w)/ √ 4σHσg − 4ρ2g,H )) ησH √ 4σHσg − 4ρ2g,H  , (6) whereC(w) = σg+2ρg,H(w−w∗)+σH(w−w∗)2, Z is the normalization constant andArcTan(·) 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∗, ) 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 infinite 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
1In the following context, we assume σg is positive number.
networks can be well approximated by quadratic curves, which supports Proposition 1. (2) The minimum of the quadratic curve is nearly located at the local minimum w∗. It indicates that the coefficient of the first-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 simplified form of C(w) that C(w) = σg + σH(w − w∗)2.
Corollary 3 If C(w) = σg + σH(w−w∗)2, the stationary distribution of 1-dimensional power-law dynamic (Eq.4) is
p(w) = 1
Z (1 + σHσ
−1 g (w − w∗)2)−κ, (7)
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 fluctuation 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.

heading: 3.3 MULTIVARIATE POWER-LAW DYNAMIC
text: In this section, we extend the power-law dynamic to d-dimensional case. We first 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Σ−1g (w−w∗)). Similarly as 1-dimensional case, we omit the first-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 = 1 κ ·H(w ∗). Then C(w) can be written as
C(w) = Σg(1 + 1 ηκ (w − w∗)TH(w∗)Σ−1g (w − w∗)). (8)
Theorem 4 Ifw ∈ Rd andC(w) has the form in Eq.(8) forw ∈ B(w∗, ). The stationary distribution density of power-law dynamic is
p(w) = 1
Z [1 +
1
ηκ (w − w∗)TH(w∗)Σ−1g (w − w∗)]−κ (9)
for w ∈ B(w∗, ), where Z is the normalization constant and κ satisfies ηΣ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.
4 ESCAPING EFFICIENCY OF POWER-LAW DYNAMIC
In this section, we analyze the escaping efficiency of power-law dynamic from local minima and its relation to generalization. Specifically, we analyze the mean escaping time for wt 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).
Definition 5 Suppose wt starts at the local minimum a, we denote the time for wt to first reach the saddle point b as inf{t > 0|w0 = a,wt = b}. The mean escaping time τ is defined as τ = Ewt [inf{t > 0|w0 = a,wt = b}].
We first 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∗)TH(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 ∆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) = σga +σHa(w− a)2, which can also be written as C(w) = σga + 2σHa Ha
(L(w)− L(a)). Furthermore, we suppose that C(w) = σga + 2σHa Ha
(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 6 Suppose that Assumption 1-3 are satisfied and C(w) = σga + 2σHa Ha (L(w)− L(a)) on the whole escaping path from a to b. The mean escaping time of 1-dimensional power-law dynamic is,
τ = 2π
(1− 1 2κ
) √ Ha|Hb|
( 1 + 2
κησga ∆L
)κ− 1 2
, (10)
where κ = HaησHa > 1 2 , Ha and Hb 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 flat minima, where the sharpness is measured by Ha and larger Ha 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 efficiency of SGD to escape from deep basin. (2) 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.,
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 infinite 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 definition 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. Specifically, 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 satisfied. For w ∈ Rd, we suppose C(w) = Σga + 2 ηκ
(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
τ = 2π √ −det(Hb)
(1− d 2κ
) √ det(Ha) 1 |Hbe|
( 1 + 1
ηκσe ∆L
)κ− 1 2
, (11)
where e indicates the most possible escaping direction, Hbe is the only negative eigenvalue of Hb, σe is the eigenvalue of Σga 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 flatness is measured by det(Ha). IfHa has zero eigenvalues, we can replace Ha by H+a in above theorem, where H + a is obtained by projecting Ha onto the subspace composed by the eigenvectors corresponding to the positive eigenvalues of Ha. This is because by Taylor expansion, the loss L(w) only depends on the positive eigenvalues and the corresponding eigenvectors of Ha, i.e., L(w) = L(a) + 12 (w − a) THa(w − a) = L(a) + 12 (P(w − a)) TΛ H+a P(w − a), where Λ H+a
is a diagonal matrix composed by non-zero eigenvalues of Ha and the operator P(·) operates the vector to the subspace corresponding to non-zero eigenvalues of Ha. 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 benefit 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 flatter minimum has smaller generalization error.
Theorem 8 Suppose that w ∈ Rd and κ > d2 . For δ > 0, with probability at least 1 − δ, the stationary distribution of power-law dynamic has the following generalization error bound,
Ew∼p(w),x∼P(x)`(w, x) ≤ Ew∼p(w)L(w) +
√ KL(p||p′) + log 1δ + log n+ 2
n− 1 ,
where KL(p||p′) ≤ 12 log det(H) det(Σg) + Tr(ηΣgH −1)−2d 4(1− 1κ ( d 2−1)) + d2 log 2 η , p(w) is the stationary distribution of d-dimensional power-law dynamic, p′(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+ 1 2κ
) , KL divergence is decreasing as H decreases. For d > 1 and fixed Tr(ΣgH−1)
and det(Σg), the generalization error (i.e., Ew∼p(w),x∼P(x)`(w, x)− Ew∼p(w)L(w)) is decreasing as det(H) decreases, which indicates that flatter minimum has smaller generalization error. Moreover, if 2d > Tr(ηΣgH−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 flatter model which generalizes better.

heading: 5 EXPERIMENTS
text: In this section, we conduct experiments to verify the theoretical results. We first study the fitness between parameter distribution trained by SGD and power-law κ distribution. Then we compare the escaping behavior for power-law dynamic, Langevin dynamic and SGD.

heading: 5.1 FITTING PARAMETER DISTRIBUTION USING POWER-LAW DISTRIBUTION
text: We investigate the distribution of parameters trained by SGD on deep neural networks and use power-law κ distribution to fit the parameter distribution. We first 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 fit 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 fit the histogram. 2 We have the following observations: (1) The distribution of the parameter trained by SGD can be well fitted 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.

heading: 5.2 COMPARISON ON ESCAPING EFFICIENCY
text: We use a 2-dimensional model to simulate the escaping efficiency from minima for power-law dynamic, Langevin dynamic and SGD. We design a non-convex 2-dimensional function written as L(w) = 1n ∑n i=1 `(w − xi), where `(w) = 15 ∑2 j=1 |wj − 1|2.5 · |wj + 1|3 and training data xi ∼ N (0, 0.01I2). We regard the following optimization iterates as the numerical discretization of the power-law dynamic, wt+1 = wt − ηg(wt) + ηλ2 √ 1 + λ1(wt − w∗)2 ξ, where ξ ∼ N (0, I2), λ1, λ2 are two hyper-parameters and 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 efficiency of SGD, while Langevin dynamic can not. We then scale the loss
2The training errors under the six settings are almost zero.
function by 0.9 to make the minima flatter 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 flatter minima slower.

heading: 6 CONCLUSION
text: 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 efficiency 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 findings.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.

heading: 7 APPENDIX

heading: 7.1 POWER-LAW DYNAMIC AND STATIONARY DISTRIBUTION
text: Theorem 9 (Theorem 2 in main paper) The stationary distribution density for 1-dimensional powerlaw dynamic (Eq.4) is
p(w) = 1
Z (C(w))
− H ησH exp H ( 4ρg,H ·ArcTan ( C′(w)/ √ 4σHσg − 4ρ2g,H )) ησH √ 4σHσg − 4ρ2g,H  , whereC(w) = σg+2ρg,H(w−w∗)+σH(w−w∗)2, Z is the normalization constant andArcTan(·) is the arctangent function.
Proof: We denote the function H(4ρg,H ·ArcTan(C′(w)/
√ 4σHσg−4ρg,H))
ησH √ 4σHσg−4ρ2g,H as h(w). According to the
Fokker-Planck equation, p(w) satisfies
0 = ∇p(w)g(w) + η 2 · ∇ · (C(w)∇p(w))
= ∇ · [ (p(w) · ∇L(w)) + η
2 C(w)∇p(w) ] = ∇ · [η 2 C(w) − HησH +1eh(w)∇(C(w) H ησH · e−h(w) · p(w))
] 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 ·eh(w) ·p(w). So we can get the conclusion in the theorem.
Theorem 10 (Corollary 3 in main paper) If C(w) = σg + σH(w−w∗)2, the stationary distribution density of power-law dynamic is
p(w) = 1
Z (1 + σHσ
−1 g (w − w∗)2)−κ, (12)
where Z = ∫ 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) satisfies
0 = ∇p(w)g(w) + η 2 · ∇ · (C(w)∇p(w))
= ∇(p(w) · ∇L(w)) + η 2 ∇ · (σg + 2σH H (L(w)− L(w∗)))∇p(w) = ∇ · η 2 C(w)(1 + 2σH Hσg (L(w)− L(w∗))) H −ησH ∇(1 + 2σH Hσg (L(w)− L(w∗))) H ησH p(w)
Because the left side equals zero, we have (1 + 2σHHσg (L(w)− L(w ∗))) H ησH p(w) equals to constant. So p(w) ∝ (1 + 2σHHσg (L(w)− L(w ∗))) H −ησH . So we can get the conclusion in the theorem.
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.1, 0 and use green, red,
purple and blue line to illustrate their corresponding density function, respectively. When κ = 0, it is Gaussian distribution. From the figure, 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 wt that satisfies power-law dynamic has analytic form, i.e., p(w, t) ∝ (1 + Hηκσ(t) (w −w(t))
2)−κ, where w(t) = w∗ + (w0 −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.

heading: 7.2 SGD AND MULTIVARIATE POWER-LAW DYNAMIC
text: 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 ∈ Rd, we use C(w) to denote the covariance matrix of stochastic gradient g̃(w) = g̃(w∗)+H̃(w−w∗) and Σ to denote the covariance matrix of g̃(w∗). IfCov(g̃i(w∗), H̃jk) = 0,∀i, j, k, we have
Cij(w) = Σij + (w − w∗)TA(ij)(w − w∗), (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 ∗) + ∑d a=1 H̃ia(wa − w∗a), g̃j(w) = g̃j(w∗) + ∑d b=1 H̃jb(wb − 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 6= 0, A (ij)
Σij are equal for all i ∈ [d], j ∈ [d]. The first assumption is
reasonable because both Σij andA(ij) reflect the dependence of the derivatives along the i-th direction and j-th direction. Let ΣH = A (ij)
Σij ,C(w) can be written asC(w) = Σg(1+(w−w∗)TΣH(w−w∗)).
The d-dimensional power-law dynamic is written as dwt = −H(w − w∗)dt+ √ ηC(w)dBt, (14)
where C(w) = Σg(1 + (w − w∗)TΣH(w − w∗)) which is a symmetric positive definite 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 = QTΛQ,ΣH = QTΓQ,H = QTΠQ. Then, the stationary distribution of power-law dynamic is
p(w) = 1
Z (1 + (w − w∗)TΣH(w − w∗))−κ, (15)
where Z is the normalization constant and κ = Tr(H)ηTr(ΣHΣg) .
Proof: Under the codiagonalization assumption on Σg,ΣH , H , Eq.15 can be rewritten as dvt = −Πvtdt+ √ ηΛ(1 + vTt Γvt)dBt if we let vt = Q(wt − w∗).
We use φ(v) = ηC(v)2 = η 2 Λ(1 + v TΓv), the stationary probability density p(v) satisfies the Smoluchowski equation:
0 = d∑ i=1 ∂ ∂vi (Πivi · p(v)) + d∑ i=1 ∂ ∂vi · ( φi(w) ∂ ∂vi p(v) ) (16)
= d∑ i=1 ∂ ∂vi (Πi·vi · p(v)) + d∑ i=1 ∂ ∂vi · ( ηΛi 2 (1 + vTΓv) ∂ ∂vi p(v) ) . (17)
According to the result for 1-dimensional case, we have the expression of p(v) is p(v) ∝ (1 + vTΓv)−κ. To determine the value of κ, we put p(v) in the Smoluchowski equation to obtain
d∑ i=1 Πip(v)− 2κ d∑ i=1 Πivi · Γivi · (1 + vTΓv)−κ−1
= d∑ i=1 ∂ ∂vi ( ηΛiκ(1 + v TΓv)−κ · Γivi )
= d∑ i=1 ( ηΛiκ(1 + v TΓv)−κ · Γi ) − 2 d∑ i=1 ( ηΛiκ 2(1 + vTΓv)−κ−1 · (Γivi)2 ) .
The we have ∑d i=1 Πi = ηκ ∑d i=1 ΛiΓi. So we have κ = Tr(H) ηTr(ΣHΣg) .
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 6= 0,A(ij) are equal for all i ∈ [d], j ∈ [d] and we denoteA(ij) = ΣH . We suppose η ·ΣH = κH . (3) Σg = σg · Id 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
p(w) = 1
Z [1 +
1
ηκ (w − w∗)THΣ−1g (w − w∗)]−κ (18)
where Z = ∫∞ −∞[1 + 1 ηκ (w − w ∗)THΣ−1g (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) satisfies 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 significantly larger than the covariance of two different dimensions of g̃(w). Consider a two layer fully-connected linear neural network fw,v(x) = wvx where w ∈ R1×m, v ∈ Rm×d, x ∈ Rd and h(·) is the ReLU activation. We consider the regression loss `(w, v) = 12 (y − fw,v(x))
2. The gradient of wi and vjk can be written as
∂`(w, v)
∂wi = (fw,v(x)− y) · vix (19)
∂`(w, v)
∂vjk = (fw,v(x)− y) · wjxk, (20)
where vi denotes the i-th row of matrix v. Suppose that the initialization of w and v is: wi i.i.d∼ N(0, δ1) and vij i.i.d∼ N(0, δ2) . We also assume that Exi = Exj = 0 and xi, xj are independent with each other for i 6= j where xi is the i-th dimension. We have
Ew,v ∂`(w, v)
∂wi
∂`(w, v)
∂wj = Ew,v(fw,v(x)− y)2 · vix · vjx (21)
= Ew,vy2 · vix · vjx+ Ew,v m∑ i=1 (wivix) 2 · vix · vjx− 2Ew,v( m∑ i=1 ywivix) · vix · vjx (22)
Because the independence of vi, vj and their expectations are zero, we can obtain Ew,v ∂`(w,v)∂wi ∂`(w,v) ∂wj = 0 for i 6= j. Similarly, we can get Ew,v ∂`(w,v)∂wi ∂`(w,v) ∂vjk = 0 and Ew,v ∂`(w,v)∂vj′k′ ∂`(w,v) ∂vjk = 0 for (j, k) 6= (j′, k′).
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. (2017). So, diagonalization dominant property of C(w) is reasonable.

heading: 7.3 SUPPLEMENTARY MATERIALS FOR RESULTS IN SECTION 4

heading: 7.3.1 PROOF FOR MEAN ESCAPING TIME
text: Lemma 14 (Lemma 6 in main paper) We suppose C(w) = σga + 2σHa Ha (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,
τ = 2π
(1− 1 2κ
) √ Ha|Hb|
( 1 + 2
κησga ∆L
)κ− 1 2
, (23)
where κ = HaησHa , Ha, Hb 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∈Va)∫ Ω JdΩ
, where Va is the volume of basin a, J is the probability current that satisfies
−∇J(w, t) = ∂ ∂w (g(w) · p(w, t)) + ∂ ∂w
( φ(w) ∂p(w, t)
∂w
)
= ∂
∂w φ(w) · (1 + µ σg ∆L(w) )−κ ∂ ((1 + µ σg ∆L(w) )κ p(w, t) ) ∂w  ,
where φ(w) = η2C(w) and µ = 2σHa Ha , σg = σga and ∆L(w) = L(w) − L(a). Integrating
both sides, we obtain J(w) = −φ(w) · (
1 + µ σg
∆L(w) )−κ ∂((1+ µσg ∆L(w))κp(w,t))
∂w . Because there
is no field source on the escape path, J(w) is fixed constant on the escape path. Multiplying φ(w)−1 · ( 1 + µσg ∆L(w) )κ on both sizes, we have
J · ∫ c a φ(w)−1 · ( 1 + µ σg ∆L(w) )κ dw = − ∫ c a
∂ (( 1 + µσg ∆L(w) )κ p(w, t) ) ∂w dw
= −0 + p(a).
Then we get J = p(a)∫ c a φ(w)−1· ( 1+ µσg ∆L(w) )κ dw
. As for the term ∫ c a φ(w)−1 · ( 1 + µσg ∆L(w) ) 1 κ dw,
we have ∫ c a φ(w)−1 · ( 1 + µ σg ∆L(w) )κ dw (24)
= 2
ησg ∫ c a ( 1 + µ σg ∆L(w) )−1+κ dw
= 2
ησg ∫ b c ( 1 + µ σg (∆L(b)− 1 2 |Hb|(w − b)2) )−1+κ dw
= 2
ησg ∫ b c ( 1 + µ σg (∆L(b)− 1 2 |Hb|(w − b)2) )−1+κ dw
= 2
ησg (1 +
µ
σg ∆L(b))−1+κ ∫ b c ( 1− µ σg · 1 2 |Hb|(w − b)2 1 + µ σg ∆L(b) )−1+κ dw
= 2
ησg (1 +
µ
σg ∆L(b))−1+κ ·
( 1 2 µ σg |Hb|
1 + µ σg ∆L(b) )−1/2 ∫ 1 0 y−1/2(1− y)−1+κdy
= 2
ησg (1 +
µ
σg ∆L(b))− 1 2 +κ √ 2σg µ|Hb| B( 1 2 , κ),
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 ∈ Va), we have P (w ∈ Va) = ∫ Va p(w)dV = ∫ w∈Va p(a)(1 + µ σg ∆L(w))−κ =
p(a) √
2σg µHa B( 1 2 , κ − 1 2 ), where we use Taylor expansion of L(w) near local minimum a. Then we
have τ = P (w∈Va)∫ Ω JdΩ
= P (w∈Va)J because J is a constant. Combining all the results, we can get the result in the lemma.
Theorem 15 (Theorem 7 in main paper) Suppose w ∈ Rd 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
τ = 2π √ −det(Hb)
(1− d 2κ
) √ det(Ha) 1 |Hbe|
( 1 + 1
ηκσe ∆L
)κ− 1 2
, (25)
where e indicates the most possible escape direction, Hbe is the only negative eigenvalue of Hb, σe is the eigenvalue of Σga 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∈Va)∫ Ω JdΩ , where Va is the volume of basin a, J is the probability current that satisfies −∇ · 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 Hbe, and the probability flux of other directions can be ignored. Then we have∫
Ω
JdΩ = Je · ∫
Ω
( 1 + 1
ηκ (w − b)T (HbΣ−1g )⊥e(w − b)
)−κ+ 12 dΩ, (26)
where Je = p(a) · η(1+µσe∆L(b))
−κ+ 1 2 √ µσe|Hbe|
2 √ 2B( 12 ,κ) which is obtained by the calculation of Je for
1-dimensional case in the proof of Lemma 13, and (·)⊥e denotes the directions perpendicular to the escape direction e.
Suppose HbΣ−1g are symmetric matrix. Then there exist orthogonal matrix Q and diagonal matrix Λ = diag(λ1, · · · , λd) that satisfy HbΣ−1g = QTΛQ. We also denote v = Q(w − b).
We define a sequence as Tk = 1 + 1ηκ · ∑d j=k λjv
2 j for k = 1, · · · , d. As for the term∫
Ω ( 1 + 1ηκ (w − b) T (HbΣ −1 g ) ⊥e(w − b) )−κ+ 12 dΩ, we have∫
Ω
( 1 + 1
ηκ (w − b)T (HbΣ−1g )⊥e(w − b)
)−κ+ 12 dΩ
= ∫ (1 + 1
ηκ · vTΛv)−κ+ 12 dw
= ∫ (1 + 1 ηκ · d∑ j 6=e λjv 2 j ) −κ+ 12 dv
=((ηκ)−1λ1) − 12 ∫ T −κ+ 12 2 B( 1
2 , κ)dv
= d−2∏ j=0 ((ηκ)−1λj) − 12B( 1 2 , κ− j 2 )
= d−2∏ j=0 ((ηκ)−1λj) − 12 · √ πdΓ(κ− d2 ) Γ(κ)
=
√ (ηκπ)d−1 · Γ(κ− d−22 )
Γ(κ+ 12 ) √ det((HbΣ −1 g )⊥e) .
As for the term P (w ∈ Va), we have P (w ∈ Va) = ∫ Va p(w)dV = p(a) ∫ w∈Va ( 1 + (w − w∗)THaΣ−1g (w − w∗) ) dw (27)
=p(a) · √ (ηκπ)d · Γ(κ− d2 )
Γ(κ) √ det((HaΣ −1 g ))
(28)
where we use Taylor expansion of L(w) near local minimum a.
Combined the results for P (w ∈ Va) and J , we can get the result.

heading: 7.3.2 FURTHER EXPLANATION ABOUT ASSUMPTION 1-3
text: 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.

heading: 7.3.3 EXTENSION TO MORE COMPLEX DYNAMIC ON THE ESCAPING PATH
text: In Lemma 6, we assume that C(w) = σga + 2σHa Ha (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) = σga + 2σHa Ha
(L(w) − L(a)) dominates the path a → z and the dynamic with C(w) = σgb + 2σHb Hb
(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. (2020). Since the proof is analogous, we omit the details here.

heading: 7.4 PAC-BAYES GENERALIZATION BOUND
text: We briefly introduce the basic settings for PAC-Bayes generalization error. The expected risk is defined as Ex∼P(x)`(w, x). Suppose the parameter follows a distribution with density p(w), the expected risk in terms of p(w) is defined as Ew∼p(w),x∼P(x)`(w, x). The empirical risk in terms of p(w) is defined as Ew∼p(w)L(w) = Ew∼p(w) 1n ∑n i=1 `(w, xi). Suppose the prior distribution over the parameter space is p′(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′(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 ∈ Rd, we select the prior distribution p′(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,
Ew∼p(w),x∼P(x)`(w, x) ≤ Ew∼p(w)L(w) +
√ KL(p||p′) + log 1δ + log n+ 2
n− 1 , (29)
whereKL(p||p′) ≤ 12 log det(H) det(Σg) + Tr(ηΣgH −1)−2d 4(1− 1κ ( d 2−1)) + d2 log 2 η andP(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′(w) = 1√
(2π)d det (I) exp{− 12 (w−w ∗)T I(w−w∗)}. The posterior distribution density is the
stationary distribution for power-law dynamic, i.e., p(w) = 1Z ·(1+ 1 ηκ ·(w−w ∗)THΣ−1g (w−w∗))−κ.
Suppose HΣ−1g are symmetric matrix. Then there exist orthogonal matrix Q and diagonal matrix Λ = diag(λ1, · · · , λd) that satisfy HΣ−1g = QTΛQ. We also denote v = Q(w − w∗). We have
log
( p(w)
p′(w) ) = −κ log(1 + 1
ηκ · (w − w∗)THΣ−1g (w − w∗))− logZ +
1 2 (w − w∗)T I(w − w∗) + d 2 log 2π
The KL-divergence is defined as KL(p(w)||p′(w)) = ∫ w p(w) log ( p(w) p′(w) ) dw. Putting v = Q(w − w∗) in the integral, we have
KL(p(w)||p′(w))
= d 2 log 2π − logZ + 1 2Z ∫ v vT v ( 1 + 1 ηκ · vTΛv )−κ dv − 1 Zη ∫ v vTΛv · (1 + 1 ηκ · vTΛv)−κdv,
(30)
where we use the approximation that log(1 + x) ≈ x. We define a sequence as Tk = 1 + 1ηκ ·∑d j=k λjv 2 j for k = 1, · · · , d. We first calculate the normalization constant Z.
Z = ∫ (1 + 1
ηκ · vTΛv)−κdw =
∫ (1 + 1 ηκ · d∑ j=1 λjv 2 j ) −κdv
=((ηκ)−1λ1) − 12 ∫ T −κ+ 12 2 B( 1
2 , κ− 1 2 )dv = d∏ j=1 ((ηκ)−1λj) − 12B( 1 2 , κ− j 2 )
= d∏ j=1 ((ηκ)−1λj) − 12 · √ πdΓ(κ− d2 ) Γ(κ)
We define Zj = ((ηκ)−1λj)− 1 2B ( 1 2 , κ− j 2 ) . For the third term in Eq.(30), we have
2Z · III
= ∫ v vT v(1 + 1 ηκ vTΛv)−κdv
= ∫ v2,···vd ∫ v1 v21 ( 1 + 1 ηκ · vTΛv )−κ dv1 + Z1 ( d∑ j=2 v2j )( 1 + 1 ηκ · d∑ j=2 λjv 2 j )−κ+ 1 2 dv2··· ,vd
= ∫ v2,···vd T−κ2 ∫ v1 v21 ( 1 + (ηκ)−1λ1v 2 1 T2 )−κ dv1 + Z1 ( d∑ j=2 v2j )( 1 + 1 ηκ · d∑ j=2 λjv 2 j )−κ+ 1 2 dv2··· ,vd
= ∫ v2,··· ,vd T−κ2 ∫ ( T2 (ηκ)−1λ1 ) 3 2 y 1 2 (1 + y)−κ dy + Z1 ( d∑ j=2 v2j )( 1 + 1 ηκ · d∑ j=2 λjv 2 j )−κ+ 1 2 dv2··· ,vd
= ∫ v2,··· ,vd ((ηκ)−1λ1) − 3 2 T −κ+ 3 2 2 B ( 3 2 , κ− 3 2 ) + Z1 ( d∑ j=2 v2j )( 1 + 1 ηκ · d∑ j=2 λjv 2 j )−κ+ 1 2 dv2··· ,vd
=( λ1 ηκ )− 3 2B
( 3
2 , κ− 3 2 )∫ v2,··· ,vd T −κ+ 3 2 2 dv2··· ,vd + ∫ v2,··· ,vd Z1 ( d∑ j=2 v2j )( 1 + 1 ηκ · d∑ j=2 λjv 2 j )−κ+ 1 2 dv2··· ,vd
For term ∫ v2,··· ,vd T − 1κ+ 3 2 2 dv2··· ,vd in above equation, we have
∫ v2,··· ,vd T −κ+ 32 2 dv2··· ,vd
= ∫ v3,··· ,vd T−κ+23 ((ηκ) −1λ2) − 12B ( 1 2 , κ− 2 ) dv3,··· ,vd
= ∫ v4,··· ,vd T −κ+ 52 4 ((ηκ) −1λ2) − 12 ((ηκ)−1λ3) − 12B ( 1 2 , κ− 5 2 ) B ( 1 2 , κ− 2 ) dv4,··· ,vd
= ∫ vd T −κ+ 12 + 1 2×d d d−1∏ j=2 ((ηκ)−1λj) − 12 d−1∏ j=2 B ( 1 2 , κ− ( j 2 + 1) ) dvd
= d∏ j=2 ((ηκ)−1λj) − 12 d∏ j=2 B ( 1 2 , κ− ( j 2 + 1) )
Let Aj = ((ηκ)−1λj)− 3 2B ( 3 2 , κ− ( j 2 + 1) ) . According to the above two equations, we can get the recursion
2Z ∫ vT vT−κ1 dv
=A1 · ∫ T −κ+ 32 2 + Z1 ∫ v2,··· ,vd  d∑ j=2 v2j T−κ+ 122 dv2··· ,vd =A1 · ∫ T −κ+ 3−12 2 dv2···vd + Z1 ·A2 ∫ T −κ+ 42 3 dv3··· ,vd + Z1Z2 ∫  d∑ j=3 v2j
T−κ+ 123 dv3··· ,vd =
d−1∑ j=1 Aj j−1∏ k=1 Zk ∫ T −κ+ j+1+12 j+1 dvj+1,··· ,vd + d−1∏ k=1 Zk ∫ v2dT −κ+ d−12 d dvd
= d−1∑ j=1 ( λj ηκ )− 3 2B ( 3 2 , κ− ( j 2 + 1) ) j−1∏ k=1 ( λk ηκ )− 1 2B ( 1 2 , κ− k 2 ) d∏ s=j+1 (( λs ηκ )− 1 2 d∏ s=j+1 B ( 1 2 , κ− (s 2 + 1) )
+ d−1∏ j=1 ( λj ηκ )− 1 2B( 1 2 , κ− j 2 − 1) · (λd ηκ )− 3 2B( 3 2 , κ− (d 2 + 1))
=
√ πdΓ(κ− d2 − 1)Tr(H −1Σg)
2Γ(κ) √ (ηκ)−(d+2) det(H−1Σg)
We have
III =
√ πdΓ(κ− d2 − 1)Tr(H −1Σg)
4Γ(κ) √ (ηκ)−(d+2) det(H−1Σg) · d∏ j=1 ((ηκ)−1λj) 1 2 · Γ(κ)√ πdΓ(κ− d2 )
= ηκTr(H−1Σg)
4(κ− d2 − 1)
Similarly, for the fourth term in Eq.(30), we have IV = κd 2(κ− d2−1) . Combining all the results together, we can get KL(p||p′) = 12 log det(H) (ηκ)d det(Σg) + log Γ(κ) Γ(κ− d2 ) + Tr(ηΣgH −1)−2d 4(1− 1κ ( d 2−1)) + d2 log 2. Using the fact that log Γ(κ) Γ(κ− d2 ) ≤ d2 log κ, we have KL(p||p ′) ≤ 12 log det(H) det(Σg) + Tr(ηΣgH −1)−2d 4(1− 1κ ( d 2−1)) + d2 log 2 η .

heading: 7.5 IMPLEMENTATION DETAILS OF THE EXPERIMENTS

heading: 7.5.1 OBSERVATIONS ON THE COVARIANCE MATRIX
text: In this section, 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 → Conv1→ maxpool → Conv2→ maxpool → fc1→ Relu→ fc2→ output. Both Conv1 and Conv2 use 5 × 5 kernels with 10 channels and no padding. Dimensions of full connected layer fc1 and fc2 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.1. After 3000 iterations, GD converges with almost 100% training accuracy and the training loss being 1e−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.001. After 10000 iterations, GD converges with 100% training accuracy and the training loss being 1e−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 fixed 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 the 20 points by changing the parameters at the Conv1 layer with Scale = 0.001 and Conv2 layer with Scale = 0.0001, respectively. For ResNet18, we choose the 20 points by changing the parameters for a convolutional layer at the first 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 figures), which is consistent with our theoretical results in Section 3.1.

heading: 7.5.2 SUPPLEMENTARY EXPERIMENTS ON PARAMETER DISTRIBUTIONS OF DEEP NEURAL NETWORKS
text: For Figure. 3(a), we train LeNet-5 on MNIST dataset using SGD with constant learning rate η = 0.03 for each batchsize till it converges. Parameters are conv2.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 = 4 and then a RandomHorizontalF lip. In training, momentum is set to be 0.9 and weight decay is set to be 5e− 4. Initial learning rate in SGD is set to be 0.1 and we using a learning rate decay of 0.1 on {150, 250}-th epoch respectively. We train it until converges after 250 epoch. Parameters are layer1.1.conv2.weight in ResNet-18.
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 fitted by powerlaw distribution. Parameters in this figure are all randomly selected to be features.10.weight, features.14.weight, features.5.expand3 × 3.weight, Mixed_6d.branch7 × 7_3.conv.weight, layer4.2.conv3.weight and features.denseblock2.denselayer1.conv2.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 fit" 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 fitted 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.

heading: 7.5.3 FURTHER EXPLANATION ON EXPERIMENTS IN SECTION 5.2
text: As for the experiments for 2-D model, we also calculate coefficient 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.

heading: 7.5.4 ESCAPING EFFICIENCY ON NEURAL NETWORK
text: 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, wt+1 = wt − η∇L(wt) + η · α∇L(wt) √ 1 + β(wt − w∗)2 ξ, where α, β are hyperparameters, ξ ∼ N (0, I), and stands for Hadamard product. Here we select α = 2.4, β = 2 after grid search. Expected sharpness is measured as Eν∼N(0,δ2I)[L(w+ ν)]−L(w) where δ = 0.01, and the expectation is computed by average on 1000 times sampling.
The numbers at the first 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 find flatter critical points than GLD and GD.

heading: 7.5.5 COMPARISON OF MEAN ESCAPING TIME WITH DIFFERENT BARRIER HEIGHTS
text: We design this 1-dimensional model to help to validate the theoretical results of escaping time in Table 1. Loss function L(w) = 1n ∑n i=1 `(w − xi), where `(w) = { w2 + bw , w < 0 −w2 + bw , w ≥ 0 and
xi ∼ N (0, 0.05). L(w) is plotted in Figure.10(a), and we can adjust barrier height through parameter b in `(w) without changing the Hessian on minima w∗ and saddle point.
For power-law dynamic (PLD),wt+1 = wt−η∇L(wt)+ηλ2 √
1 + λ1(wt − w∗)2 ξ, where λ1, λ2 are hyperparameters, ξ ∼ N (0, I), and 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.1 for both methods. We initialize w0 = w∗ and apply both methods on L(w) with different barrier heights. Then we record the number of iterations t when wt firstly 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.

Review summaries:
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.