Dynamic of Stochastic Gradient Descent with State-dependent Noise
Authors: Qi Meng
Shiqi Gong
Wei Chen
Zhi-Ming Ma
Tie-Yan Liu
Venue:  (2021)
Keywords: state-dependent noise
power-law dynamic
stochastic gradient descent
generalization
deep neural network
heavy-tailed
escape time

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 state-independent, 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.

Under review as a conference paper at ICLR 2021
DYNAMIC OF STOCHASTIC GRADIENT DESCENT WITH STATE- DEPENENTT NOISE
Anonymous authors Paper under double-bling review
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 state- independent, although 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.
INTRODUCTION
Deep learning has achieved great success various AI applications, such computer vision, natural language processing, and speech recognition (He er al.) [2016b] Vaswani et al. 2017; He er al. (0166a). Stochastic gradient descent (SGD) and its variants are the mainstream methods 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, one hand, researchers are investigating the loss surface of the neural networks with variant architectures (Choromanska el al.) 2015 Li el al. 2018b He el al., 2019b Draxler et al. 2018: Li al. (2018)); on the other hand, researchers illustrate that the noise stochastic algorithm may make it escape from local minima (eeskar& et al. 2016: He et al.) 2019a Zhu et al., 2019) Wu et al. 2019a HaoChen al. 20200. It is clear that whether stochastic algorithms from poor local minima and finally stop at minimum with low generalization error 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.
escape
To study the dynamic behavior of SGD, most of the works consider SGD 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 al.) 2017; Zhou al. 2019: Liu al.
2018 Chaudhari & Soatto 2018 He et al.) 2019a] Zhu et al. 2019 Hu er 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 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, Simsselli et al. (2019); (Simşekli et al. 2019) point that the variance of stochastic gradient may be infinite, and they propose to approximate SGD by dynamic driven by o-stable process with the strong infinite variance condition. However, as shown in the work (Xie

---

Exes, v,, w/) is the minibatch Under review as a conference paper at ICLR 2021
et al. 2020 Mandt et al.) 2017). the gradient noise follows Gaussian distribution and the infinite variance condition does not satisfied. Therefore it 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 a quadratic function 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 K distribution, where the signal to noise ratio of the second order derivatives at local minimum. Compared with Gaussian distribution, power-law i. distribution is heavy-tailed with tail-index k. matches the empirical observation that the distribution of parameters becomes heavy-tailed after SGD training without assuming infinite variance of stochastic gradient (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 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 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 parameters trained by SGD without assuming infinite variance of gradient noise. (2) We analyze the mean escaping time and PAC-B BAC-Bayes gencralization bound for power-law dynamic and results show that power-law dynamic escape sharp local minima faster and generalize better compared with the dynamics with constant diffusion. Our experimental results can support the theoretical findings.
BACKGROUND
In empirical risk minimization problem, the objective is - = [a,,, w), where a,, 1.... , are i..dd training samples, E Rd is the model parameter, and is the loss function. Stochastic gradient descent (SGD) is popular optimization algorithm to minimize L(w). The update rule is W++1 = w; 11 ((w,), where g(w,) gradient calculated by randomly sampled minibatch S,, of size b and n the learning rate. The minibatch gradient ğ( ((w)) is an unbiased estimator of the full gradient g(w:) = (L(w)) and the term (g(w)) ((w))) called gradient noise SGD.
Langevin Dynamic In (He er 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) ~ (0, C*), where covariance matrix C' is constant matrix for all w. Then SGD can be regarded as the numerical discretization of the following Langevin dynamic,
dw, = -g(we)dt + V/O
(1)
where B, is standard Brownian motion in Rd and 1/2 dB, is called the diffusion term.
a-stable Process Simsekli et al. (2019) assume the variance of gradient noise is unbounded. By generalized CLT, the distribution of gradient noise is atsbaall distribution S(a, σ),, where a is the a-th moment of gradient noise for given a with a E (0, 2]. Under this assumption, SGD is approximated by the stochastic differential equation (SDE) driven by an o-stable process.
2

---

Under review as a conference paper at ICLR 2021
2.1
RELATED WORK
There are many works that approximate SGD by Langevin dynamic and most of the theoretical results are obtained for Langevin dynamic with constant diffusion coefficient. From the aspect of optimization, the convergence rate of SGD and its optimal hyper-parameters have been studied in eL al. 2017] He el al. 2018|[Luu et al.) 2018 He et al. (2018) via optimal control theory. From the aspect of generalization, Chaudhari & Soatto (2018) 8): 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 er al. 2019
Hu et al.) 2019 Xie et al. (020). He et al. (2019a) analyze the PAC-Baaes generalization error of Langevin dynamic 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 heavy-tailed (Mahoney & Martin, [019]] Hodgkinson & Mahoney] 2020] Gurbuzbalaban et al. 2220)) Simsekli er (2019); Şimş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 a-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 SGD state-dependent. There are some related works analyze state-dependent noise structure in SGD, such as label noise in (HaoChen et al. (220)) and multiplicative noise in (Wu er al,, 2019b). These works propose new algorithms motivated by the noise structure, but they do not analyze the escaping behavior of dynamic of SGD and the impact to generalization. [Wu et al. (2018) analyze the escaping behavior of SGD with considering the fluctuations of the second order derivatives and propose the concept linearly stability. In our work, we propose power-law dynamic to approximate SGD and analyze the stationary distribution and the mean escaping time for it.
3 APPROXIMATING SGD BY POWER-LAW DYNAMIC
In this section, study the (state-dependent) noise structure of SGD (in Section 3.1) and propose power-law dynamic to approximate the dynamic of SGD. We first study 1-dimensional power-law dynamic in Section 3.2 and extend it to high dimensional case in Section 3.3.
3.1 NOISE STRUCTURE OF STOCHASTIC GRADIENT DESCENT
For non-convex optimization, we investigate the noise structure of SGD around local minima so that we can analyze the escaping efficiency from it. We first describe the quadratic basin where the local minimum located. Suppose w" is a local minimum of the training loss L(w) and g(w") 0. We name the e-ball B(w", e) with center w² and radius the loss function for w E B(w",,) is equal to its second-order Taylor expansion as L(w) L(w*)- + {w w²) w*)TH((" (w w²). Here, H (w*) ) is the Hessian matrix of loss at w", which is (semi) positive definite.
a quadratic basin if
Then we start to analyze the gradient noise of SGD. The full gradient of training loss is g(w)
(ww")ww w²). ). The stochastic gradient is g(w) = 5(w") + A (w=))w -ww--w") w²) by Taylor expansion where g(-) and (-) are stochastic version of gradient and Hessian calculated by the minibatch. The randomness of gradient noise comes from two parts: j(w") and (ww)), which reflects the fluctua- tions 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 E B(w",e) C IR, the variance of gradient noise is o(g(w) - ō(w) = (ī((")) + 2p(ğ(w"), H(w") () - w²)++ (H(()))) - w where o()) and p(...) 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 ((w") and (ww)) have finite variance because p(ğ(w"), ) (ww")) s V (ū(w*)) (((())))) according to Cauchy-Schwar inequality. For fixed w", sufficient condition for that ğ(w") and A (ww") have finite variance is that
3

---

Under review as a conference paper at ICLR 2021
the training data x 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 a-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, use C(w)), Gg. Ht Pg. denote a(g(w) g(w")) a(ā("") )), (( }), p(ğ(w"), (w"")) in the following context, respectively,
3.2 POWER-LAW DYNAMIC
According to CLT, the gradient noise follows Gaussian distribution it has finite variance, i.e.,
( g(w) -a N((, as b xc,
where -d means "converge in distribution". Using Gaussian distribution to model the gradient noise SGD, the update rule of SGD can be written as:
W++1 We (wg) + rpEe,
& ~~ C(w)).
Eq 3 can be treated as the discretization of the following SDE, which we call it power-law dynamie:
dw - (g(ww)d )dt + VnC(ww)BB.
Power-law dynamic characterizes how the distribution of w changes as time goes on. The distribution density of parameter w at time (i.e., p(w, t)) is determined by the Fokker-Planck equation (Zwanzig's type (Guo & Du 2014)):
① δe P(w,t) t) = Vp(w, + " (C((w) Vp(w,t)).
(5)
The stationary distribution of power-law dynamic can be obtained if we let the left side of Fokker- Planck 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 power-law dynamic.
Theorem 2 The stationary distribution density for I-dimensional power-law dynamic (Eq 4 is
(4pp. Arel'nn (C'(w)/ /4400o - -4p
p(w)= = 7 (C(w)) WHB exp
(6)
D / A - 4p
where C(w) og + 2p, (w-w") +oH[w=="}+ Z is the normalization constant and Arcl'an())
the arctangent function.
We make discussions on property of p(w). The decreasing rate of p(w) as goes away from the f
center w" is mainly determined by the term C(w) w*H (because the function ArcT (nn)) is bounded) which is 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 ((w", ,) and heavy-tailed. We call H the tail-index of p(w) and denote it as k. 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 pawer-law dynamic when PH.s 0 and d 0. In this case, p(w) degenerates to Gaussian distribution. Compared with a-stable process, we do not assume infinite variance on gradient noise and demonstrate another mechanism that results in heavy-tailed distribution parameters.
We empirically observe the covariance matrix around the local minimum of training loss deep neural networks. The results are shown in Figure [ Readers can refer more details in Appendix 7.1. We have the following observations: (1) The traces of covariance matrices for the deep neural
In the following context, assume 0, is positive number.
4

---

Under review as a conference paper at ICLR 2021
0.5
080
2
0.0
639
to vis?t
vat
(a) CNN layer-1
(b) CNN layer-2
ResNet layer-1
(d) ResNet layer-2
Figure 1: Trace of covariance matrix of gradient noise in a region around local minimum w". w selected by running gradient descent with small learning rate till converges. The number at horizontal axis shows the distance of the point away from w". (a),(b): Results for plain CNN. (c),(d):Reuull for ResNet18.
networks can be well approximated by quadratic curves, which supports Proposition ! (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 Py.H O.
Based on the fact that H is not the determinant factor of the tail of the distribution in 6 and the observations Figure we consider a simplified form of ((w) that C(w) 0g + oh (u w w²)?
Corollary 3 IfC(w) = ag + ww w²) the stationary distribution - I-dimensional power-law
dynamic (Eq is
p(w) = 7 o (w (wc- )2)
where Z is the normalization constant and
H is the tail-index.
no
The distribution density in Eq is known as the power-law x distribution (Zhou & Du 2014) (It is also named q-Gaussian distribution in (Tsaliis & Bukman, 1996)). As - oc, the distribution density tends be Gaussian, i.e., p(w) x exp(- H(w-=" ndw
Power-law к distribution becomes more heavy-tailed as x becomes smaller. Meanwhile, it produces higher probability to appear values far away from the center w". Intuitively, smaller x helps the dynamic to escape from local minima faster.
In the approximation of dynamic of SGD, K equals the signal (i.e., H(w") )) to noise (i.e., qo ratio of second-order derivative w" in SGD, and is linked with three factors: (1) the curvature H (ww); (2) the fluctuation of the curvature over training data; (3) the hyper- arper-parameters including and minibatch size 6. Please note that он linearly decreases as the batch size b increases.
3.3 MULTIVARIATE POWER-LAW DYNAMIC
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 matrix. We use Mg to denote the covariance matrix of g(w*) ) and assume that Mg is H, (w")) ") are equal for all i, It can w w))). Similarly as I-dimensional case, we omit the first-order term (w - w²) in (ww)). Readers can refer Proposition 10 in Appendix 7.2 for the detailed derivation.
isotropic (i.e., Σg 0g I). We also assume that Cou(H.(w) be shown that C(w) Eg(1 (w w " En'' (w
We suppose that the signal to noise ratio of Hw") can be characterized by scalar K, i.c., 7/C Hww"). Then C(w) can be written as
1 (w w W")T " (w - w)))
C(w)= E1 +
(8)
Theorem W E & and C(w) has the form in Eq. density of power-law dynamic is
8 for E B(w", ,e). The stationary distribution
p(w)= = 7
1 (w w w
-wwwwn
J -
(9)
Tj
for B(w", {), where Z is the normalization constant and satisfies ETE 1.2 H(w").
5

---

Under review as a conference paper at ICLR 2021
The multivariate power-law k distribution (Eq 9 is natural extension of the 1-dimensional case. Actually, the assumptions on Eg and λ can be replaced by just assuming Eg. H(w"), Σh are codiagonalized. Readers can refer Proposition 11 in Appendix 7.2 for the derivation.
Remark:
4 ESCAPING EFFICIENCY OF POWER-LAW DYNAMIC
In this section, analyze the escaping efficiency of power-law dynamic from local minima and its relation to generalization. Specifically, we analyze the mean escaping time for w, escape from basin. As shown in Figure 2 we suppose that there are two basins whose bottoms are denoted and respectively and the saddle point b is the barrier between two basins. The barrier height is denoted as AL L(b) L(a).
Figure 2
Definition 5 Suppose we starts at the local minimum a, we denote the time for we to forst reach the saddle point b in > Olurp T == Fw, [inf(t > 0/ = a, We =b]
a, w, b}. The mean escaping time T defined as
We first give the mean escaping time for 1-dimensional case in Lemma 6 and then we give the mean escaping time for nigh-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 L(w) L(w") ) { H (w w"), where w" critical point.
Assumption 2: The system is in equilibrium near minima, i.e., Op(un.)) 0.
dt
Assumption 3: (Low temperature assumption) The gradient noise is small, i.e., n1, < AL.
These three assumptions are commonly used in analyzing escaping time (Xie et al.) 2020: Zhou & Du 2014) for dynamic. Because both and b are critical points, we can apply Assumption to get the loss surface around them. We put more discussions about the assumptions in Appendix 7.3.2.
We suppose the basin is quadratic and the variance of noise has the form that C(w) = asa to (w
a)², which can also be written C(w) Ts. + 20
suppose that Ha (I(w)- - [(a)) on the whole escaping path from a to b (not just near the local
H. (L(w) L(a)). Furthermore,
((w)) Ag. + 20H,, 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 I-dimensional case.
20 (L(w) - L(a)) on the Lemma 6 Suppose that Assumption 1-3 are satisfied and C(w) = 0g. Ha whole escaping path from a to b. The mean escaping time of 1-dimensional power-law dynamic is,
2m
1
_AL AL
(10)
v Ha Hul
KY/O
Ha. no.. 2. Ha and Hr are the second-order derivatives of training loss local minimum a and at saddle point b, respectively.
where &
The proof of Lemma 6 is based on the results in (Zhou & Du 2014). We provide full proof in Appendix 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 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 I 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 H. and larger H. corresponds to sharper minimum. Power-law dynamic improves the order of barrier height (i.e., AL) from exponential to polynomial compared with Langevin dynamic, which implies faster escaping efficiency of SGD to escape from deep basin. (2) The mean escaping time for a-stable process is independent with the barrier height, but it is in polynomial order of the width of the basin (i.e.,

---

Under review as a conference paper at ICLR 2021
Table Summary of related works and ours. Here, we only show I-dimensional result for escaping time in the table for all the three dynamics for ease of the presentation.
Noise distribution
Dynamic
Stationary solution
Escaping time
o0
Langevin
Gaussian
WHalHal
exp
2
S(a.0)
a-stable
Hcavy-tailed
O (7a
[lbaa
N0,, 0g + on w ))
Power-law (ours)
Heavy-tailed
(1
+
AL )
v Ha 1.lllllll
width=|b al). Compared with a-stable process, the result for power-law dynamic is superior in the sense that it is also in polynomial order of the width (if AL xe O(lb a )) and power-law dynamic does rely on the infinite variance assumption.
Based on Lemma 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-ding eigh-dimensional landscape. For rigorous definition of most possible escaping paths, readers can refer section (Xie al. 2020). For simplicity, we consider the case that there 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 Suppose that Assumption 1-3 are satisfied. For Rd. we suppose C(w) Lga 2. (L(w) L(a)) on the whole escaping path from a to b and there is only one most possible path nn
path between basin and basin The mean escaping time for power-law dynamic escaping from basin a to basin is
2T V Hp) 2N ) del(H.) Hoel
(1+ (1))) TIKO
R-
(11)
where e indicates the most possible escaping direction, Hhe is the only negative eigenvalue of H de is the eigenvalue of Ega that corresponds to the escaping direction, AL = L(b) L(a), and det(-) the determinant of matrix.
Remark: In d-dimensional case, the flatness measured by det (Ha). If Ha has zero eigenvalues, we can replace H, by H in above theorem, where H is obtained by projecting H, onto the subspace composed by the eigenvectors corresponding to the positive eigenvalues Taylor expansion, the loss L(w) only depends on the positive eigenvalues and the corresponding eigenvectors Ha, i.c., L(w) = L(a) + {w - a)² Ha where A md 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 the dimension of subspace that composed by directions with large eigenvalues. It has been observed that most of the eigenvalues small (Sagun al.) 2016). Therefore, d will not be large number and power-law dynamic in multi-diment dilei-ddmensionall case will inherit the benefit of that in I-dimensional case compared with Langevin dynamic and a-stable process.
Ha. This is because by
a) = L(a) + 1/ (P(w w - a))² A P(w a),
H is very
The next theorem give an upper bound of the generalization error of the stationary distribution of paw dynamic, which shows that flatter minimum has smaller generalization error.
Theorem 8 Suppose that w E Ied and > 4. For 8 > O, with probability at least 1 - 8, the stationary distribution of power-law dynamic has the following generalization error bound,
KT(pllo') + log + logn +
Bew. S Erewp(w) L(w)+ + +V
log 2 p(w) is the stationary distribution of d-dimensional power-law dynamic, p' (w) is prior distribution which is selected to be standard
where KL(pllp') < log det(H)) de, + Tr(n H 1)-2d (
7

---

Under review as a conference paper at ICLR 2021
Gaussian distribution, and ((x) is the underlying distribution of data z, 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 > 21+2555 KI. divergence is decreasing as H decreases. For d > 1 and fixed Tr H)
and det(Eg) , the generalization enror (i.e., P(x= (())) Ew--p()) ((w)) is decreasing as det(H) decreases. which indicates that flatter minimum has smaller generalization error. Moreover, if 2d > Tr HI the generalization error is decreasing as increases. When K - 0o, the generalization error tends to that for Langevin dynamic. Combining the mean escaping time and the generalization error bound, can conclude that state-dependent noise makes SGD escape from sharp minima faster and implicitly tend to learn a flatter model which generalizes better.
5 EXPERIMENTS
In this section, we conduct experiments to verify the theoretical results. We first study the fitmess between parameter distribution trained by SGD and pawer-law distribution. Then we compare the escaping behavior for power-law dynamic, Langevin dynamic and SGD.
5.1 FITTING PARAMETER DISTRIBUTION USING POWER-LAW DISTRIBUTION
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, 10243. For the settings of other hyper-parameterr 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 "T mallis@GaussionDDD. wiw.-iiiittriiiionn [[" in Mathematica software.
We show results for LeNet-5 with MNIST dataset and ResNet-18 with CIFAR 10 dataset (LeCun et al. 2015 He al. 2016b) in this section, and put results for other network architectures in Appendix 7.3.2. In Figure S we report the generalization error (i.e., Test error Training error) and the values of x that best fit the histogram. [ 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, x becomes larger. is because the noise σh linearly decreases as minibatch size becomes larger and K H (3) As k becomes smaller, the generalization error yaa
becomes lower. It indicates that also plays role as indicator of generalization. These results are consistent with the theory in Section
5.2 COMPARISON ON ESCAPING EFFICIENCY
We use a 2-dimensional model simulate the escaping efficiency from minima for power-law dynamic, Langevin dynamic and SGD, We design non-convex 2-dimensional function written as L(w) Ei= f(w x.). where ((w) 15 Σ;= I0 1 w, + 1/3 and training data n, NO 0. We regard the following optimization iterates as the numerical discretization of the plwer-law dynamic, W:+1 w; ng(w)) + 1/ /1 \₂ (w, )² @ &, where d X2 are two hyper-parameter, and @ stands for Hadamard product. Note that if we set X = O. can be regarded as discretization of Langevin dynamic. We set learning rate 1) = 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 wt with the methods above, X2 is chosen to satisfy Tr(Cou (A2E)) = Tr (T(CCo(ggw"))
I2).
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 X1, we choose different λ₃ and evaluate corresponding success rate of escaping, as shown in Figure The results show that: (1) there is positive correlation between λ and the success rate of escaping: (2) power-law dynamic can mimic the escaping efficiency of SGD, while I Langevin dynamic can not. We then scale the loss
The training errors under the six settings are almost zero.
8

---

Under review as a conference paper at ICLR 2021
*- 3.52
A+ 4.50
5,5
A- 1.05
x- 1.59
X=217
(a) LeNet-5 conv layer
(b) ResNet-18 layer
Figure Approximating distribution of parameters (trained by SGD) by power-law dynamic. Training batchsize, generalization error (i.e., Test emrr- Traiiing error) and approximated tail-index are shown in the title of each plot. (a): Results for LeNet-5. (b):Results for ResNet-18.
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 [dd We can observe that all dynamics escape flatter minima slower.
3.45
0.34
-10-10 100-10 &d
-151 105
250
InPoweninn
(a) 2-D
(b) Trace for covariance
(c) Success Rate (sharp)
(d) Success Rate (flat)
Figure 4: Lo):Loss surface of L(w) for 2-D model. (b):Trace covariance matrix around minimum
(1,1). (c)/(d): Success rate of escaping from the basin of L(w) (0.9L(w) in repeated 100 runs.
CONCLUSION
In this work, 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 eиror 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, non- Gaussian dependentndnt noise, new types of state dependet 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.
REFERENCES
Balduzzi, David, Frean, Marcus, Leary, Lennox, Lewis, JP, Ma, Kurt Wan-Duo, & McWilliams, Brian. 2017. The shattered gradients problem: If resnets are the answer, then what is the question? arXiv preprini arXiv: 1002.0822.
Bottou, Léon, & Bousquet, Olivier. 2008. The tradeoffs of large scale learning. Pages 161-168 of: Advances in neural information processing systems.
Chaudhari, Pratik, & Soatto, Stefano. 2018. Stochastic gradient descent performs variational inference. converges to limit cycles for deep networks. Pages 1-10 of: 2018 Information Theory and Applications Workshop (ITA). IEEE.
Choromanska, Anna, Henaff, Mikael, Mathieu, Michael, Arous, Gérard Ben, & LeCun, Yann. 2015. The loss surfaces of multilayer networks. Pages /92-204 of: Artificial intelligence and statistics.
Draxler. Felix. Veschgini, Kambis, Salmhofer, Manfred, & Hamprecht, Fred. 2018. Essentially No Barriers in Neural Network Energy Landscape. Pages 1309-1318 of: International Conference on Machine Learning.

---

Under review as a conference paper at ICLR 2021
Guo, Ran, & Du, Jiulin. 2014. Are power-law distributions equilibrium distribution stationary nonequilibrium distribution? Physica A: Statistical Mechanics and its Applications, 406, 281 281-286.
Gurbuzbalaban, Mert, Simsekli, Umut, & Zhu, Lingjiong. 2020. The Heavy-Tail Phenomenon in SGD. arXiv preprint arXiv:2006.04744
HaoChen, Jeff Z, Wei, Colin, Lee, Jason D, & Ma, Tengyu. 2020. Shape Matters: Understanding the Implicit Bias of the Noise Covariance, arXiv preprint arXiv:2006.0866 08680.
He, Di, Xia, Yingce, Qin, Tao, Wang, Liwei, Yu, Nenghai, Liu, Tie-Yan, Ma, Wei-Ying. 2016a.
Dual learning for machine translation. Pages 820-828 of: Advances in neural information processing systems.
He, Fengxiang, Liu, Tongliang, & Tao, Dacheng. 2019a. Control Batch Size and Learning Rate to Generalize Well: Theoretical and Empirical Evidence. Pages 1141-1150 of: Advances Neural Information Processing Systems.
He, Haowei, Huang, Gao, Yuan, Yang. 2019b. Asymmetric Valleys: Beyond Sharp and Flat Local Minima. Pages 2549-25295551 2-60 of: Advances in Neural Information Processing Systems.
He, Kaiming, Zhang, Xiangyu, Ren, Shaoqing, & Sun, Jian. 2015. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. Pages 1026-1034 of: Proceedings of the IEEE international conference on computer vision.
He, Kaiming, Zhang, Xiangyu, Ren, Shaoqing, & Sun, Jian. 2016b. Deep residual learning for image recognition. Pages 770-778 of: Proceedings of the IEEE conference computer vision and pattern recognition.
He, Li, Meng, Qi, Chen, Wei, Ma, Zhi-Ming, & Liu, Tie-Yan, 2018. Differential equations for modeling asynchronous algorithms. Pages 2220-2226 of: Proceedings of the 27th International Joint Conference Artificial Intelligence.
Hodgkinson, Liam, & Mahoney, Michael W. 2020. Multiplicative noise and heavy tails in stochastic optimization. arXiv preprint arXiv:2006.062: 062933.
Hu, Wenqing, Li, Chris Junchi, Li, Lei. & Liu, Jian-Guo. 2019. On the diffusion approximation of nonconvex stochastic gradient descent. Annals of Mathematical Sciences and Applications, 4(1),
3-32.
Keskar, Nitish Shirish, Mudigere, Dheevatsa, Nocedal, Jorge, Smelyanskiy, Mikhail, Tang, Ping Tak Peter. 2016. On largh-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv: 1609.04836.
LeCun, Yann, et al. 2015. LeNet-5, convolutional neural networks. URL: http://yann. lecun. com/exdb/tenet, 20(5), 14.
Li, Dawei, Ding, Tian, & Sun, Ruoyu. 2018a. Over-parameterizee deep neural networks have no striet local minima for any continuous activations. arXiv preprint arXiv: 18/2. /0039.
Li, Hao, Xu, Zheng, Taylor, Gavin, Studer, Christoph, & Goldstein, Tom. 2018b. Visualizing the loss landscape of neural nets. Pages Systems.
of: Advances in Neural Information Processing
Li Qianxiao, Tai, Cheng, et al. 2017. Stochastic modified equations and adaptive stochastic gradient algorithms. Pages 2101-2110 of: Proceedings of the 34th International Conference on Machine Learning- Bearning-Vollme Volume 70. JMLR. org.
Liu, Tianyi, Chen, Zhehui, Zhou, Enlu, & Zhao, Tuo. 2018 Toward deeper understanding of nonconvex stochastic optimization with momentum using diffusion approximations. arXiv preprint arXiv:1802.00515..
Mahoney, Michael, & Martin, Charles. 2019. Traditional and heavy tailed self regularization neural network models. Pages 4284 4293 of: International Conference Machine Learning.
10

---

Under review as a conference paper at ICLR 2021
Mandt, Stephan, Hoffman, Matthew D, & Blei, David M. 2017. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, 18(1), 4873 49077.
McAllester, David A. 1999. PAC-Bayesian model averaging. Pages 164-170 of: Proceedings of the !welfth annual conference on Computational learning theory.
Rakhlin, Alexander, Shamir, Ohad, & Sridharan, Karthik. 2012. Making gradient descent optimal for strongly convex stochastic optimization. Pages 1571-1578 of: Proceedings of the 29th International Coference on International Conference on Machine Learning.
Sagun, Levent, Bottou, Leon, & LeCun, Yann. 2016. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07477
Simşekli, Umut, Gürbüzbalaban, Mert, Nguyen, Thanh Huy, Richard, Gaël, & Sagun, Levent. 2019. On the Heavy-Tailed Theory of Stochastic Gradient Descent for Deep Neural Networks. arXiv preprint
Simsekli, Umut, Sagun, Levent, Gurbuzbalaban, Mert. 2019, A Tail-Index Analysis of Stochastic Gradient Noise in Deep Neural Networks. Pages 5827-5837 of: International Conference on Machine Learning.
Smith, Samuel L,, & Le, Quoe V. 2017. bayesian perspective on generalization and stochastic gradient descent. arXiv preprint arXiv:1710.06445
Tsallis, Constantino, Bukman, Dirk Jan. 1995. Anomalous diffusion in the presence of external forces: exact time-dependent solutions and entropy. arXiv preprint cond-mat/9511007
Tsallis, Constantino, Bukman, Dirk Jan. 1996. Anomalous diffusion in the presence of external forces: Exact time-dependent solutions and their thermostatistical basis. Physical Review E, 54(3), R2197.
Van Kampen, Nicolaas Godfried. 1992. Stochastic processes in physics and chemistry. Vol. 1. Elsevier.
Vaswani, Ashish, Shazeer, Noam, Parmar, Niki, Uszkoreit, Jakob, Jones, Llion, Gomez, Aidan N, Kaiser, Lukasz, & Polosukhin, Illia. 2017. Attention is all you need. Pages 5998 6008 of: Advances in neural information processing systems.
Wu, Jingfeng, Hu, Wenqing, Xiong, Haoyi, Huan, Jun, & Zhu, Zhanxing. 2019a. The Multiplicative Noise in Stochastic Gradient Descent: Dena-Deeendent Regularization, Continuous and Discrete Approximation. CoRR.
Wu, Jingfeng, Hu, Wenqing, Xiong, Haoyi, Huan, Jun, Braverman, Vladimir, & Zhu, Zhanxing. 2019b. On the Noisy Gradient Descent that Generalizes as SGD. arXiv preprint 1rXiv:1906.0744
Wu, Lei, Ma. Chao, et al. 2018. How sgd selects the global minima in wer-parametterized learning: A dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 8279-8288 8288.
Xiao, Han, Rasul, Kashif, Vollgraf, Roland. 2017. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07744
Xie. Zeke, Sato, Issei, & Sugiyama, Masashi. 2020. A Diffusion Theory for Deep Learning Dynamies: Stochastic Gradient Descent Escapes From Sharp Minima Exponentially Fast. arXiv preprint 45X1v:2002..03495.
Zhang, Yao, Saxe, Andrew M, Advani, Madhu S, & Lee, Alpha A. 2018. Energy-entropy competition and the effectiveness of stochastic gradient descent in machine learning. Molecular Physics, 116(21-22), 3214-3223.
Zhou, Mo, Liu, Tianyi, Li, Yan, Lin, Dachao, Zhou, Enlu, & Zhao. Tuo. 2019. Toward Understanding the Importance of Noise in Training Neural Networks. In: International Conference on Machine Learning.
11

---

Under review as a conference paper at ICLR 2021
Zhou, Yanjun, & Du, Jiulin. 2014. Kramers escape rate in overdamped systems with the power-law distribution. Physica A: Statistical Mechanics and its Applications, 402, 299-305.
Zhu, Zhanxing, Wu, Jingfeng, Yu, Bing, Wu, Lei, & Ma, Jinwen. 2019. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. Pages 7654-7663 of: Proceedings of International Conference on Machine Learning.
7 APPENDIX
7.1 POWER-LAW DYNAMIC AND STATIONARY DISTRIBUTION
Theorem (Theorem 2 main paper) The stationary distribution density for 1-dimensional power- law dynamic (Eq() is
H (Apg.1 Arel'an (C'(w)/ Aonoo - -00-44/44 402,1)
p(w) = 7 (C(w)) HHH
H
no n 40 40g.u
where C(w) = -gg 12-+2pp0 2p (w- w") + +o 2 is the normalization constant and AreTan(())
the arctangent function.
Proof: We denote the function (40... AreTan (w))/ 1apg,,44ppll as h(w). According to the you //rnn 40g.n
Lokker-Planche equation, p(w) satisties
0 vp(w)g(w) + ? V. (C(w)V((())
(p(w) VL(W)) -/2 (())v((w)
I 2 (oww) रमक h(w)) (CC(w) रमम e M(r) p(w))
Readers can check the third equality by calculating v(C(w) म. e h(w) p(w)) with C(w) 0g + 2pq. Hr (w w²) + H (w- w² Because the left side equals zero, we have C(w) गhw e - .p(w) equals constant. So p(w) C(w) M(w) -p(w). So we can get the conclusion in the theorem.
H_
Theorem 10 (Corollary in main paper) If C(w) = 0g + + w²)?, the stationary distribution density of power-law dynamic is
p(w)= 2 01+ (w w )²))
(12)
where Z = [1 + Too (w - w+)² - w is the normalization constant and k no tail-index.
H
is the
Proof:
According to the Fokker-Planck equation, p(w) satisfies
Vp(w)g(w) " (C(w)Vp(w))
(p(w)) VL(w)) + " (ag H
20 (L(w) [(w")))p((v))
V F. + 20H (L(w)- L(w"))) non ((1 + 20H (L(w) - L(w")())) "F" p(w)
Ho,
Ho,
# Because the left side equals zero, we have (1 + 20 [Lww) - L(w" ))) 7* p(w) equals to constant.
Ho,
So p(w) x (1 +
20H (I(w)) I(w")))
ma - So we can get the conclusion in the theorem.
Hoy
We plot the un-orrmalized distribution density for I-dimensional power-law dynamics with different
in Figure 5 For the four curves, set 10. We set 1, 0.5, 0.1, O and green, red,
12

---

Under review as a conference paper at ICLR 2021
edd = 1o00
Figure Probability density for power-law dynamic.
purple and blue line to illustrate their corresponding density function, respectively. When x 0, is Gaussian distribution. From the figure, we see that the tail for power-law k-distribution is heavier than Gaussian distribution.
Actually, for any given time t, the distribution p(w, t) for w, that satisfies power-law dynamic has H (w ) where w(t) w" + (wo w²) )e IIt and
analytic form, i... p(w, t) x (1 +1+ ())((B)
(1) is function of og and L. Readers can refer Eq. 18 Eqq23 in (Talllss & Bukman, 1995 for the detailed expression.
7.2 SGD AND MULTIVARIATE POWER-LAW DYNAMIC
The following proposition shows the covariance stochastic gradient in SGD in d-dimensional case. We the subscripts to denote the elements in vector a matrix.
Proposition 11 For E Rdd, use ((w) denote the covariance matrix of stochastic gradient g(w) = (W and denote the covariance matrix of g(w"). IfCov(@ H₂x)) 0, Vi, k, we have
Ciy(w) = Eij + (w - -0-===)'- ( w²).
(13)
where Σij- = Cov(āi (w))) (,);("))), Ali is d d matrix with elements A (()) Cov (Hia: H3²) with a E [d].be [d].
Eq 13 can be obtained by directly calculating the covariance of ār (w)) and ß. (w) where ge(w)
Ge( (e(w") Ed= (ua wg). Oj (w) タン (2w"w) Σi Hys(ws wg))
In order to get a analytic tractable form of Cww), we make the following assumptions: (1) If Eig = 0.
A61) is zero matrix; (2) For Σij o 0, Ali3) reasonable because both Σij and (i5) reflect the dependence of the derivatives along the i-th direction
Eag are equal for all E [d]. [d]. The first assumption is
Aa5) and j-th direction. Let Σµ = F5 C(w) can be written as C(w) Eg (w+ En (1-wwu) w²))
The d-dimensional power-law dynamic is written
dw - w" dt + CCwddB...
(14)
where C(w) Σ, (w- In w)) which is symmetric positive definite matrix that C(w))/2 exists. The following proposition shows the stationary distribution the d-dimensional power-law dynamic.
Proposition 12 Suppose Eg, EH, H are codiagonalizable, i.e., there exist orthogonal matrix Q and diagonal matrices A,T, IT to satisfy E.g = OF NQ. = QTT QTQ Then, the stationary distribution of power-law dynamic is
p(w) 7 (1+ (1++wu (w w0 "T E - -*
(15)
where Z is the normalization constant and k
Tr(H)
13

---

Under review as a conference paper at ICLR 2021
Proof:
Under the codiagonaliiation assumption 2g. En H, Eq.15 can be rewritten as de, -Teddt + Vn + o/ Te;) we let 8; = Q - w").
We use ((e) _C(n) = _ + e" r'),,, the stationary probability density p(u) satisfies the
Smoluchowski equation
8
(I,;; p()))(())) +
(i(ww) p(o)
(16)
Ou
Dor
(II. p(v))
a
(1
+ Te) Sv:
a p(u)
(17)
Dus
According the result for I-dimensional case, we have the expression of p(v) is p(u) x (1 y Te) -* To determine the value of put p(t) in the Smoluchowski equation to obtain
T.p(1)- - 2x Σ T.u. Puys (1 + TTT)
i (MiM(1 - To) -..m
Ooi
Σ (1/nn( + +TTT (n* r.) - mnAnn(1 + Tu) (Tmm)))). ((T.)))))))
The we have Σ IT, m/ Ei= A, So have [Tr(E
Tr(mf)
II we can also consider another assumption on Eg. En, H without assuming their codiagonalization. Instead, we assume (1) Eng = 0, A(0 is zero matrix; (2) For Σi + 0, (()) are equal for all i E [d] E [d] and we denote A = Σn- We suppose 7) 7 = kH. get the following theorem.
According to Proposition
(3) Eg 0g Ia which is isotropic. Under these assumptions,
Theorem 13 (Theorem 4 in main paper) If is d-dimensional and C(w) has the form Eq. The stationary distribution density of multivariate power-law dynamic is
1 p(w) z
1
[1 + (w - (w===)" "2 (w w w²]
(18)
7/K.
where Z 1_
L (w w )² 2 " (w w 1 is the normalization constant.
The proof for Theorem 12 is similar to that for Proposition 11. Readers can check that p(w) satisfies 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 (w. = wez where w E glxm E mg x Rd and ((()) is the ReLU activation. We consider the regression loss E(w, v) = 3 (v fw.e The gradient of w, and Uyg can be written
dl(w, v)
(f -v) ex
(19)
Ow;
dl(w,v)
(fm,o(x)) y)) wiykk
(20)
Öuzk
where v, denotes the i-th row of matrix v. Suppose that the initialization of w and v is: wi Lid
(O, 81) and Vij i.ddd N(0, 62) We also assume that Ex, Er, and ny; xy are independent with each other for i * where x; is the i-th dimension. We have
de(w, v) de(w, v)
E E
(fr, y)² -0 2yæ
(21)
Ow,
Ow,
EE,wy? 020 U5 z + E.g
(wwyar) U3a dyz 2E ywuuaa) Uaa eyæ (22)
14

---

Under review as a conference paper at ICLR 2021
Because the independence of va by and their expectations zero, for * j. Similarly, we can get Ew, dd(w.o) 0c(w.w) and
obtain
Eu, @t(w.o) d[(w,t)
Bhe Buy
Juyk
Eu,0 Ot(w.o) d6(w,r) = 0 for (i, + (J', :)).
00y'2' Bugk
The above analyses show that the gradients for different dimensions are independent at initial- ization. It has been observed that many weights are kept random during training because of the prrrtearamttriization Balduzzi er al. (2017). So, diagonalization dominant property of C(w) is reasonable.
7.3
SUPPLEMENTARY MATERIALS FOR RESULTS SECTION 4
7.3.1 PROOF FOR MEAN ESCAPING TIME
Lemma 14 (Lemma 6 in main paper) We suppose C(w) gg. + 20 (L(w) L(a)) the whole escaping path from to b. The mean escaping time of the I-dimensional power-law dynamic is,
Ha +
2m
2
6-
14
AL
(23)
H. Hi
KYJO go
where x = Ha Hay Hy are the second-order derivatives of training loss at local minimum a and saddle point b.
nos."
Proof: According (Van Kampen 1992>, the mean escaping time expressed P(wEV,) Ja JAQ
where Va is the volume of basin is the probability current that satisfies
(b(w) Op(w.1) Ouc
Jun, ow
(g(w) p(u, t))
Thu
o(w) (1+ sl(w)) -x =00(1+ "LL(w) p((e,)) t)
a
ow
0g
where d(w) (C(w) and p
2ane Ha 0g dg. and AL(W) L(w) L(a). Integrating both sides, we obtain ((w) = -ww)) (1+ a AL(()) of (1+ AL(w)) p(w.t)
Because there
field source the escape path, (w) is fixed constant the escape path. Multiplying ((w) + 0g AL(w) both sizes, have
J. ( ((w)-1 (1))=.(1+ * dw - 1 (11 AL(W)) p(w, dw
0,
og
Ow
-
+ p(a)-
15

---

Under review as a conference paper at ICLR 2021
As for the term La ([w)w= (1+ of (L(ww)) "
Then we get
p(a)
dw,
I d(w)
1- (1+ 11 ALw))
du
we have
[( d(w)
(1 (1 1 sL(ww)
(24)
0g
2
P SL(w)
-Ite
1
dw
10
0g
2
1 (AL(6) 2 [HHal(w - (1Hal(w-b)()
-1th
dw
g0
2
(AL(b)
Hbl(w - -0)3)
dw
70g
Fg
2
(1+
P
AL())
1+x L"
1 1
2 | (w b)²
-14x
du
1
0g
Tg
1
+ "
AL(b)
2 E HHul
-1/2
2
(1
+ AT(())
1+#
1/
1// 8) - dy
n0g
0g
1+ AT(b)
2
(1
M AT(()) 1+
B 1 M
20g
3 5- k),
110g
0g
where the third formula is based 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 E Va), have P(w E Vo) Sv. p(w)dV = Swev, p(a)(1 + _ (w))
p(a) V nHH B K - }), where we use Taylor expansion of L(w) near local minimum a. Then we have P(wEV,) P(weV=) because J is constant. Combining all the results, we can get the
20
Jg Jd
result in the lemma.
Theorem 15 (Theorem 7 in main paper) Suppose w E IRd and there is only one most possible path puth between basin and the outside of basin The mean escaping time for power-law dynamic escuping from basin a to the outside of basin
2T -et/-dttHH
(1+1AL AL
(25)
(1 ) det(Ha) Hiel
where indicates the most possible escape direction, Hhe is the only negative eigenvalue of Hy, de is the eigenvalue of Eg,. corresponding the escape direction and AL L(b) L(a).
Proof: According Van Kampen, 1992s. the mean escaping time expressed
P(WeV,) Jg Jal2"
where Va is the volume of basin J is the probability current that satisfies - V Jww, t) = 3p(w.()
0t
Under the low temperature assumption, the probability current concentrates along the direction corresponding the negative eigenvalue of Hbe, and the probability flux of other directions can be ignored. Then we have
-h+
I JAN = Je. 1/ (1+
1 (w (H)
dl,
(26)
where Je = p(a) n(1+Ho. AL()) *1 |Hbe which obtained by the calculation of Je for 2v2B(3.* 3
1 1-dimensional case in the proof of Lemma 13, and (+)Le L denotes the directions perpendicular to the escape direction e.
Suppose H₂Σ are symmetric matrix. Then there exist orthogonal matrix Q and diagonal ma- trix A "ii.. >aa) that satisfy H QT AQ. We also denote Q(w b).
16

---

Under review as a conference paper at ICLR 2021
We define sequence
1 me
Σj= AzUJ for k 1,... d.
As for the term
Sa (1
7p2 (w
(H -b)
-xt
dS2 we have
-xt
1
1
(w -on-bb" (H=Σg') "
"w (()w-b) -b)
dl
12
7/pc
+
x dw -²" Av)
/ +
x dv
j&c
((1))) {11) - / T2 x B( k)dv
d-2
II (1Tx) X1) Bl K 2
d-2
II ((T)) -1,) dT 4)
I(Kx)
f=0
V(NET)d
[' d-2
T(K 2) (H
As for the term P(w E Va), we have
P(we Va)= / p(w)dV p(a)) / =p(a)-
(w-)"" (w w²)) dw
(27)
vv.
wev.
(InET) I( [) [(K) det))
(28)
where we use Taylor expansion of I(w)) near local minimum a. Combined the results for P(w E Va) and J, we can get the result.
7.3.2 FURTHER EXPLANATION ABOUT ASSUMPTION 1-3
We adopt the commonly used assumptions to analyze mean escaping time for dynamic system Xie al.) 2020; Smith & Le 2017 Zhou & Du 2014). Assumption can be replaced by weaker assumption that the system is quasi-equee equilibrium which is adopted in (Xie et al.) 2020). For the differences between cuusi--quulibrium and equilibrium, readers can refer to (Xie er 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.3 EXTENSION TO MORE COMPLEX DYNAMIC ON THE ESCAPING PATH
In Lemma 6, assume that C(w) Jga + Ha (L(w) I(a)) the whole escaping path from to b for case 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 2 is the midpoint on the most possible path beween and b, where L(z) (1 =)T((a) + =L(b). The dynamic with C(w) dg. 20
H.
H. (L(w)- L(a)) dominates the path - and the dynamic with
((w) dap + 25Ha (L(b) L(w)) dominates the path = + b. Then only two things will be changed in proof of Lemma First, we need to change the stationary distribution near saddle points according to its dynamic in 20. Second, need change the integral about probability density
H.
17

---

Under review as a conference paper at ICLR 2021
the whole path sum of integrals these two sub-paths. Similar proof techniques adopted for analyzing escaping time Langevin dynamic in proof Theorem 4. in the work work[Xie et al. (2020). Since the proof is analogous, omit the details here.
7.4 PAC-BAYES GENERALIZATION BOUND
We briefly introduce the basic settings for PAC-Bayes generalization enror. The expected risk is defined (=) r). Suppose the parameter follows distribution with density p(w), the expected risk in terms of p(w) is defined as E P (,) æ). The empirical risk in terms of p(w) is defined as Eu--(()) (ww) Eu-pp(w) 1. - ((w, Suppose the prior distribution 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 /. Then can get the following theorem.
over the parameter space is p' (w) and p(w) is the distribution
Theorem 16 (Theorem 8 in main paper) For E Rd, we select the prior distribution pl (w) to be standard Gaussian distribution. For > O. with probability at least 6, the stationary distribution of power-law dynamic has the following generalization error bound,
KL(pllp') log log n + 2
E Sun= S Ewwpt(n) L(w) +
(29)
where K KLpllpp)) S log det(H)
(E,)).,)) + Tr(n 1)-2d
4(1 (1- + log and P(x)) is the underlying distribution
of data
Proof: Eq. (9) directly follows the results in McAllester 1999). Here we calculate the Kull- back-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 den- sity (w) v (2m) exp( 2 (w )² I(w w* The posterior distribution density the
stationary distribution for power-law dynamic, i.e., p(w) 1/2 .1+ 1 (w-ww" - T
- -w (w- ))
Suppose HΣg1 are symmetric matrix. Then there exist orthogonal matrix Q and diagonal matrix A = diag(A1, Aa) that satisfy HΣ, Q" AQ. We also denote = Q(w w²).
We have
p(w)
log
(w(ww)
==xlog(1 + 1 (w - (w-w">" H w²)) - -log Z + 1
(w w I - w²) + 2
log 2m
npc
2
The KL-divergence is defined as K (Lpp(w)lp((w)) = S.n p(w) log w²) in the integral, we have
plw) p(() dw. Putting v = Q(w
(T(p(w)llz'(v))
d
-x
1
log 2n log Z + 2z Nę
TT 1+
oT TTAu) AU)
du
v Au (1
TT Au)
2
Zn
mp
(30)
18

---

Under review as a conference paper at ICLR 2021
where we the approximation that log(1 + x) = We define sequence as T3 1 Σjzx 230 for ls 1,.., .d. We first calculate the normalization constant Z.
1 88
2 = 2= / (1 + q/
n² Av) =dw = / (1+
1
Σ
d
X50
*
ddv
19 j=1
(11))= / T Bi 1
1
II ((n))
1ン=)= B(2,'*
1 3 2
2 k
2: 15
2
1=1
II ((ppc) -2,) V - 9)
I(x)
3=1
We define 2₁ = ((Tx)) 2 (2,xx K- - (4).
For the third term in Eq. (0 we have
2Z.111
v'v(1+ + 1 o² Ao "dv
9/
-x+ 1
[ of 1+
1
o" (0)
doi + 211
Σ
d
1+ n/
d Σ
dog.n.ea
nini
1=2
y=2
(1 1 +
-1117
d
1
-x+
T Ju
del + 21
1
25
de .d
T2
3=2
-x 7
12 /
((mi))
T2
(1xc)
4 (1+
- dy - 1/
Σ
+
2
dog Rd
1/
1=2
((TN)) -1 *11)= TT,
x
B
2 K
3 2
+ 21
Σ
d
+
1
d
Ao
1/2
des 26
Tpp
1=
)=2
-nt
A1
t BB
3
2
n8
2
3
|
12 Nid
12
doy - nd
Zi
1+
GK
1
Xuu 2
dug Pd
j=2
For term
Jez....ou v0a
T7= #/2 +
dog....ou ma in above equation, we have
Tg k.+ doz... :01
12
2 11 -2) d dis,"r,us 1ts
T3 x+2 ((1)))=2 22 2: k
J vid
TT
mx+ (1T)) -1 x22) - - ((7x) -1 as)) $B
22:--2)
1
5
B
2 1. 2 dea...
1
2
Day va
4 xd d-1
d-1
Ta
II ((mm))-111))
B 2.
12 +1) do,
2
j=2
j=2
d
B
1
k
( 2
+ +1)
2
1=
1=2
19

---

Under review as a conference paper at ICLR 2021
Let A, ((nk)) recursion
(10))=11))) 12y) ( (2 + 1)). According to the above two equations, we can get the
27 / UT TTTT""d
d
=A1
/
Tz2 + Z1
T2
Σ
Vjj
2
T2
x
V2r vdd
f=2
=A1-
/
T2
3-1
de + 11 A2 / T5
deg va + ZZZZ2 /
Σ Uz
T33 x+
dug-
j-1
d-1
Aj
1I
Zx
T
TJ+1
desti 1n
II
18111 x+ d=1
dud
k=1
Σ A/
(3
2:
3
[ +1
3 1/2 B
11
2 18
k
)
II
B
2
2
1 +
17.
77
k=1
a=j+1
s=1+1
4-1
X
- B( 2
1
- 1)
Ad
2 B
3
d 12
+ 1))
176
2
7716
2
1=1
V
d
1)TT(H
2
2T(K) (11x) del(H- (H-1
We have
VT - d 11Tr11
[(x)
111
II (17x) 1z)
(T()) / (d++) del (H-1
V () )
1=1
NKTr (H- 4(K 4 1)
Similarly, for the fourth term in Eq. 30 we have IV
Combining all the results together,
2(x 4 13
we can get K KLppllp') 1 log (wee)" det + log P(c- 27
det(if)
[(x)
TrinE, 1-22 4 log Using the
111-
fact that log [(x- 9 log K, we have KL(pllp)) S log dct().,)
T(x)
det(H)
+ 4(1 (1-1())11)
Tr(nE,1-1) + d log "
7.5
IMPLEMENTATION DETAILS OF THE EXPERIMENTS
7.5.
OBSERVATIONS ON THE COVARIANCE MATRIX
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 imput + Convl - maxpool + Conv2 - mazpool - fcl - Relu - fe2 output. Both Convl and Conw2 use 5 x 5 kernels with 10 channels and no padding. Dimensions of full connected layer fel and fc2 are 1600 x 50 and 50 x 10 respectively. We randomly sample 1000 images from FaShionMNIST (Xiao et al. 2017) dataset 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 O. After 3000 iterations, GD converges with almost 100% training accuracy and the training loss being le 3
As for ResNet, we use the ResNet- 18 model (He et al. 2016b) and randomly sample 1000 images from Kaggle's 's dogs-vs-on vs-cats dataset training set. The mitializztion method is the Kaiming initialization (He et al. 2013) in Py Porch. The learning rate of gradient descent set to be 0.001, After 10000 iterations, GD converges with 100% training accuracy and the training loss being 1
20

---

Under review as a conference paper at ICLR 2021
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 WiayerL = (i x Scale), where Wlagers denotes the parameters at layer I,, and x Scale, i E [M] determines the distance away from W W.ayert.- 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 R8 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 Conel 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 [ig 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 SGI 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.
7.5.2
SUPPLEMENTARY EXPERIMENTS ON PARAMETER DISTRIBUTIONS OF DEEP NEURAL NETWORKS
Batchsize 64 Generalization Enror 4.59% K 1.08
Batchsize 236 Generalization Error 4.64% k= 1.69
Batchsize 1024 Generalization Error 5.74% K= 2.17
0.03
0.04
0.05
0.06
0.05
0.04
0.05
0.06
0.03
0.04
0.05
0.06
Figure 6: A close-up of right tail distribution of the result for ResNet18 in Figure. 3(a) which could help to observe the heavy-tailed properties among different batchsize.
VGG16 Top-5 Error 9.62%
AlexNet
SqueezaNet 1.0
Top-5 Error 20.91% K = 3.937
Top- Error 19.58%
K 3.547
k= 2.912
20
20
10
15
6-
6
-0
-0
0.05 0.00 0.05 0.10
-0.05 0.00 0.05 0.10
-0.1
0.0
0.1
0.2
Inception v3 Top-5 Error 6.44%
Wide ReeNet-50-22
MobileNet v2
Top-5 Error 5.91%
Top-S Error 9.71%
== 3.352
K == 2.4
k= 14.92
20
10:
15
10
40
5.
20
-0
-0.05 0.05
0.000.000000 000.000000..0.0000033
0.2
Figure 7: Approximating distribution of parameters (trained by SGD) by power-law dynamic. These networks use pre-trained models offered by PyTorch, and all of them are pre-trained on ImageNet dataset using SGD. The second line in each title shows Top-5 test error and the third line shows approximated tail-index k.
we train LeNet-5 on MNIST dataset using SGD with constant learning rate 1) 0.03 for each batchsize till it converges. Parameters are weightt in LeNet-5. For Figure (bb) we train ResNet-18 18 on CIFARI0 using SGD with momentum. We do a RandomCrop on training set scaling to 32 x 32 with podding = and then a Random H orizonnta Flip. 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 epoch respectively. We train it until converges after 250 epoch. Parameters are layer 1 1.come2. com.2.eiight in ResNet-18.
For Figure. 3(a)
21

---

Under review as a conference paper at ICLR 2021
VGG16
AlexNet
SqueezeNet
Inception
Wide ResNet-50-2
DenseNet-121 v2
0.0
6:1
62
030
-321
+0.0
-6.0
-330
0.0
6:1
301
-22
50
01 32
Figure 8: Comparison between Q-Q plots of network parameters versus normal distribution and power-law distribution. (upper): Q-Q plots of parameters versus normal distribution. (bottom): Q-Q plots of parameters versus power-law distribution.
We also observe the parameter distribution on many pretrained models. Details for pre-trained models can be found on wtttts::/ppytorrtii wttp:chh.rg//oo///sseeeess
org/ Figure shows the distribution of parameters trained by SGD can be well fitted by power- law distribution. Parameters in this figure are all randomly selected to be features. 10.es.....eet weiggh, 3. 3.weight, Mixed_6dll 6d..brnchTT features. /4. weight, feeatrres..ecranndee denseblock2 oenneellleckdddneelayeyr www.sgerrr..oo......m conar2.wiigt 2ager4...ccmma' weigght and AlexNet, SqueezeNet 1.0, Inception v3, Wide ResNet-50-2 and DenseNet-122 respectively.
html
7_3.comv. weigt,, for VGG-16,
A Q-Q plot is created by plotting quantiles of two probability distributions against one another, which can provide an assessment of "gooddnes of fit" by how much the solid line close the dashed line. From Figure g it is clear that the solid lines in bottom pictures are closer to dashed lines 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 not, which means the distribution of parameters is indeed heavy-tailed.
7.5.3 FURTHER EXPLANATION ON EXPERIMENTS IN SECTION 5.2
As for the experiments for 2-D model, 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 A1 Rs 32.
7.5.4 ESCAPING EFFICIENCY ON NEURAL NETWORK
GD : 67.62 (0.0152)
GLD 67.66 (0.0150)
PLD 67.75 (0.0145)
Text Accurary
0.68
2500
5000
7500
10000
12500 15600 17500
iteratioon
Figure 9: Escaping experiment corrupted FashionMNIST. Test accuracy versus iteration after pretraining by GD. Model is pretrained by GD before the vertical dashed line and continued by GD, GLD and PLD (ours). Numbers in brackets are expected sharpness after model converging.
22

---

Under review as a conference paper at ICLR 2021
We follow the settings in (Zhu al. 2019). For convenience of the readers, here give the details of this setting again. We use corrupted PashionMNIST 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 1 1,,300 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(GLDD and power-law dynamic(PLD), Following Zhu's setting, the learning rates for GD, GLD and PLD are MGD HGLD 0.07 and 1/LDD 0.07, respectively. For GLD, noise std σ 10 as Zhu already tuned. For our PLD, W1+11 Wi [()) + 1 aV c / + B(w; w²)² &. where 5 are hyperparameters, Ç (0, f). and @ stands for Hadamard product. Here we select 2.4, 3 after grid search. Expected sharpness is measured as (0. [L(w + w)] ((w) where 5 = 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.
7.3.5
COMPARISON OF MEAN ESCAPING TIME WITH DIFFERENT BARRIER HEIGHTS
700
2.0
GLD
1.5
Mean Escaping Iteration
PLD
saddle
500
1.0
400
0.5
300
0.0
200
-0.
100
-1.0
minima
1
2
3
5
-2
Barrier Height
(a) I-D Model
(b) Mean Escaping Time Barrier Height
Figure 10: (a): Loss curve L(w) for 1-D model. (b): Mean escaping time versus different barrier heights. Mean escaping time is computed by average on 100 rounds, in which we record the number of iterations when firstly escaping from the saddle point.
We design this 1-dimensional model to help to validate the theoretical results escaping time in w2 + bw w 0
Table Loss function L(w) = Σm= f(w v.), where ((w) = 2 + bw
and
w 2 0 I N (7, 00,0. L(w) is plotted in Figure 10(a) and we can adjust barrier height through parameter b in ((w) without changing the Hessian minima w" and saddle point.
(WL))=++)aa For power-law dynamic (PLD), We+1 w/ - V/ +1/22 1 + A (wr, w² 2 where λ A2
are hyperparameters, (O, 7), and @ stands for Hadamard product. Here we let X 1, X2 4. For Langevin dynamic (GLD), we set noise std a = in consistence with PLD. Learning rate = 0.1 for both methods. We initialize wo w" and apply both methods L(w) with different barrier heights. Then we record the number of iterations when w, 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 [O(b)
From Figure Toxb) 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.
23


Reviews:
[{"review_id": "qevyo7b-1p", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes to use power-law dynamics to approximate the state-dependent gradient noise in SGD, and analyses its escaping efficiency compared with previous dynamics. \n\nStrength:\n1.\tTo the best of my knowledge, it is novel to use power-law dynamics to analyze the state-dependent noise in SGD. \n2.\tStill with strong assumptions on covariance structure, the analytical results based on power-dynamics are interesting. For example, it indicates that so-called kappa distribution highly depends on the fluctuations to the curvature over the training data. This is consistent with following work. So I suggest authors provide some discussion with the following work.\nWu et.al 2018. How sgd selects the global minima in over-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems (pp. 8279-8288).\n\nWeakness & Issues \n1.\tThe analytical results seem that they strongly depend on the covariance structure assumption, i.e. C(w) is diagonally dominant according to empirical observation. Does it have any theoretical justifications, or even in simplified cases? \n2.\tThe delivered PAC generalization bound and the followed analysis are a little ambiguous.  Firstly, in current deep learning theory community, the relationship between flatness (even how to define a proper flatness) and generalization is still mysterious and controversial, which depends many factors. This work uses one type of flatness measure, the determinant of H, and shows that flatter minima generalize better by only considering the KL term. However, the first term also includes the Hessian and might also affect generalization bound. Thus, the conclusion appears a little problematic. \nThe authors said that generalization error will decrease w.r.t. kappa’s increase and infinite kappa results in Langevin dynamics. Then the question is what are the difference between the power-law dynamics and Langevin dynamics in term of generalization? \nMy view on the ambiguous analysis is that the authors attempt to answer extremely challenging questions but left with many questionable concerns. \n3.\tThe experiments might not be sufficient. \nI don’t think fitting the parameter distribution according to limited empirical observations is an appropriate way to make justifications. At least, from visual observation, there are many other alternatives besides power-law distribution to fit, as Fig 3 shows. \nAbout comparing the escaping efficiency, the result only shows the success rate, and the evidence about the polynomial and exponential difference should be provided. Also, practical networks and datasets should also be considered to provide more strong evidence.\n\nIf the authors can resolve these issues carefully, I would raise the score. \n\nTypos\n“Eq. 4” should be “Eq.3” below equation 3\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "juO-JoFhad", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: the paper studies the effect of SGD noise near a local minimum of the loss by using a novel Taylor expansion to estimate the distribution of gradient noise in the neighborhood of that minimum. They use this to derive closed-form equations describing the distribution of the iterates, which they use to characterize properties such escaping times and generalization. \n\nPros: \n- To my knowledge, the mathematical analysis appears to be quite novel and insightful. It is particularly interesting that the authors show that the second order effect of the SGD noise in the Hessian induces a power law distribution over the iterates. \n- Some empirical support is provided for the theory.\n\nCons:\n- In general, a clearer statement (and justification) of the assumptions is required. For example, it appears to be implicit throughout the paper that we only consider the neighborhood of a local minimum, so the analysis is essentially for a quadratic in this neighborhood. This should be stated more explicitly. \n- I also have some concerns about mathematical precision in the theorem statements. It is sometimes unclear which computations are rigorous equalities and which are not - for example, in Lemma 6 about escaping times, exact equality is used. However, the proof relies on Taylor expansion and uses approximate equalities in the steps. This is potentially misleading. \n\nIn general, the results seem interesting, and it is understandable that certain assumptions/heuristics must be used because this area of research is technically challenging. However, I would like to see the clarity of the presentation improved before recommending acceptance. \n\nI have more specific questions regarding the details in the paper below:\n- Could the authors elaborate on the relationship between kappa and generalization? From the paper my understanding was that smaller kappa meant flatter curvature and better generalization, but this doesn't seem to be supported by Figure 3. \n- How is the value of kappa contained in Figure 3? Is it computed via computing the Hessian and its covariance or chosen to best fit the histograms in the figure?\n- Eqn 6: It seems like this closed form computation is specifically for the case when the function is quadratic (e.g. we take 2nd order Taylor approximation around a local min. Can the authors confirm?) If this is the case, what happens to the dependency on w - w^* and why is there no such explicit term in eqn 6? It would appear that g(w) should depend on (w - w^*). \n- In the overparameterized regime, it would appear that \\sigma_g could go to 0 if each training example is overfit by the model. It appears that plugging in \\sigma_g = 0 would introduce some degeneracy in equation 6 and 7, however. Can the authors comment on this?\n- Intuition on the term \\sigma_H: what do we expect this to look like in practice and do the authors have a sense on whether this term only matters around local minimum?\n- The definition of \\Sigma_H in the multivariate case: in the first paragraph of section 3, the definition on the LHS has no mention of i, j but the RHS does.\n- Assuming the signal to noise ratio of \\tilde{H} can be characterized by a scalar - why is this assumption reasonable?\n\n*********\nEDIT: Changed my score from 5 to 6 after the author response/revision.", "extra_scores": {}, "extra_text": {}}, {"review_id": "XaiXgknkrs5", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes power-law dynamic of SGD which considers state-dependent noise. The power-law distributed derived from this  dynamic explains the heavy-tailed distribution of parameters trained by SGD. Besides, this dynamic also shows efficiency of escaping local minima.\n\nConcerns:\n1. The proof of theorem 2 is not provided in the appendix. But I doubt if C(w) is well-defined. It is not clear how w* is selected considering there are multiple local minima. It does not make sense to me if w* is fixed when taking x-->\\infty, as the quadratic approximation should be used in the neighborhood of w*. \n\n2. The escaping efficiency of the power-law dynamic is only analyzed in low-dimension case. I wonder how if performs in high-dimensional space. Does it provide more benefits than Langevin/alpha-stable dynamic in the expense of calculating sigma_g and sigma_H.\n\nMinor comments:\n1. I think  [Li et al., 2017] also proposed state-dependent noise in Theorem 1.", "extra_scores": {}, "extra_text": {}}, {"review_id": "0KhY94zUnkY", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThis 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.\n\n### Strong/Weak points\n- The paper is built up on simple principles\n- It first gives a one-dimensional analysis then generalize, helping the reader to understand\n- Except a few points, in general the paper is well-written.\n- The assumptions made are somewhat strong and may not hold in some cases, see below.\n\nIn general I have a tendency to accept this paper. Even though there are crucial assumptions that are made, it can be considered as a first step towards a more rigorous and general argument.  \n\nHere are a few points that I have problems with in the paper:\n- On page 3, paragraph 2, it is written that the solution of Langevin equation is Gaussian distribution. What does it mean? The solution of a SDE is a Markov process, and considering the distribution of the process at time $t$, it is not necessarily Gaussian; the Fokker-Planck equation governs the change of distribution, having the Gibbs distribution as its stationary distribution, which is not Gaussian in general.\n- The whole argument is made through assuming that near the basin, everything is quadratic (not approximately, equal!). This is completely reflected in Proposition 1 and the further analysis.\n- In Theorem 4, it is not stated that $H = H(w^*)$, and I don't see why should one be interested when $w\\to\\infty$? This is because we are talking about an $\\epsilon$ ball around $w^*$ and tending $w \\to \\infty$ has no meaning.... Maybe I am missing something here? Also, it seems that the distribution is defined only for positive $w$.\n- In the argument for Section 4 (escaping), it is assumed that the basin is quadratic and __stays__ quadratic, even when it reaches the saddle point. I find this assumption flawed, or I am missing something.\n- At the bottom of page 3, it is said \"in this case, .... is satisfied\", while I think it should be \"not satisfied\".\n- On page 4, the notion $\\rightarrow_p$ is used for convergence in distribution, which is not usual and is reserved for convergence in probability.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The average review rating is 5.5 which means it’s somewhat borderline. One of the reviewers planned to increase the score but apparently didn’t do so formally. A subset of the main pros and cons the reviewers pointed out are: 

Pros: 
“Some empirical support is provided for the theory.”
“ It is particularly interesting that the authors show that the second order effect of the SGD noise in the Hessian induces a power law distribution over the iterates.”

Cons: 
“The escaping efficiency of the power-law dynamic is only analyzed in low-dimension case. ...” The author responded that Theorem 7 proves the multi-dimensional case. But the AC noted that it’s very likely that escaping time is exponential in dimension (because kappa needs to be larger than d as the author noted and the det() might also be exponential in d. The author did say in the revision that the dimension should be considered as the effective dimension of the hessian, but the AC couldn’t find a formal argument about it.)
“The assumptions made are somewhat strong and may not hold in some cases...”

The reviewers also had a few clarity questions which the author addressed in revisions with re-organized writing. The AC weighed the pros and cons and found that the unclarity and potential exponential escaping time in the multi-dimensional case outweigh the pros.   


Decision: Reject