Logarithmic landscape and power-law escape rate of SGD
Authors: Takashi Mori
Liu Ziyin
Kangqiao Liu
Masahito Ueda
Venue: ICLR 2022 Submitted (2022)
Keywords: stochastic gradient descent
noise structure
escape rate
flat minima
statistical physics

Abstract: Stochastic gradient descent (SGD) undergoes complicated multiplicative noise for the mean-square loss. We use this property of the SGD noise to derive a stochastic differential equation (SDE) with simpler additive noise by performing a random time change. In the SDE, the loss gradient is replaced by the logarithmized loss gradient. By using this formalism, we obtain the escape rate formula from a local minimum, which is determined not by the loss barrier height $\Delta L=L(\theta^s)-L(\theta^*)$ between a minimum $\theta^*$ and a saddle $\theta^s$ but by the logarithmized loss barrier height $\Delta\log L=\log[L(\theta^s)/L(\theta^*)]$. Our escape-rate formula strongly depends on the typical magnitude $h^*$ and the number $n$ of the outlier eigenvalues of the Hessian. This result explains an empirical fact that SGD prefers flat minima with low effective dimensions, which gives an insight into implicit biases of SGD.

Under review as a conference paper at ICLR 2022
LOGARITHMIC LANDSCAPE AND POWER-LAW ESCAPE RATE OF SGD
Anonymous authors Paper under double-bling review
ABSTRACT
Stochastic gradient descent (SGD) undergoes complicated multiplicative noise for the mean-square loss. We use this property of the SGD noise to derive a stochastic differential equation (SDE) with simpler additive noise by performing random time change. In the SDE, the loss gradient is replaced by the logarithmized loss gradient. By using this formalism, we obtain the escape rate formula from a local minimum, which is determined not by the loss barrier height AL = L(0*) L(0") between minimum 6* and a saddle 0 but by the logarithmized loss barrier height A log L = ((0))/ the typical magnitude h" and the number n of the outlier eigenvalues of the Hessian. This result explains an empirical fact that SGD prefers flat minima with low effective dimensions, which gives an insight into implicit biases of SGD.
Our escape-rate formula strongly depends
INTRODUCTION
Deep learning has achieved breakthroughs in various applications in artificial intelligence such as image classification (Krizhevsky al.
et al. 2013), speech recognition (Hinton al. 2012) natural language processing (Collobert & Weston 2008), and natural sciences (Iten et al,,
2020 Bapst et al.. 2020 Seif et al.] (021). Such unparalleled success of deep learning hinges
crucially on stochastic gradient descent (SGD) or its variants as an efficient training algorithm.
Although the loss landscape is highly nonconvex, the SGD often succeeds in finding a global minimum. It has been argued that the SGD noise plays a key role in escaping from local min- ima (Jastrzebski al. 20177; Wu et al., 2018 2020 Zhu et al./ 2019 Meng al. 2020: Xie al., 2021 Liu et a 20221. It has also been suggested that SGD has an implicit bias that is beneficial for generalization. That is, SGD may help the network to find flat minima, which are considered to imply good generalization (Keskar et al. 2017 Hoffer et al. [017]] Wu et al., 2018). How and why does SGD help the network escape from bad local minima and find flat minima? These questions have been addressed in several works, and it is now recognized that the SGD noise strength and structure importantly affect the efficiency of escape from local minima. Our work follows this line of research, and add new theoretical perspectives.
In physics and chemistry, escape from a local minimum of the (free) energy landscape due to thermal noise temperature T has been thoroughly discussed (Kramers) 1940 Langer 1969). When the (free) energy barrier is given by AE, the escape rate is proportional to AE77 which is known as the Arrhenius law. By analogy, in machine learning, escape from a local minimum of the loss function is considered to be determined by the loss barrier height AL = L(0*)- - L(0*), where L(0) denotes the loss function at the network parameters 8, 8" stands for a local minimum L(#). and 8" denotes a saddle point that separates 8* from other minima. If we assume that the SGD noise is uniform and isotropic, which is often assumed in machine-learning literature (Jastrzębski et al., 2017). the escape rate is proportional to e -AL/D, where D denotes the SGD noise strength.
In this paper, show that inhomogeneity of the SGD noise strength brings about drastic modifi- cation for the mean mean-square loss. It turns out that the escape rate determined by the logarithmized loss barrier height A log L log L((*)- -lo") L(#*) log ( (8)))((0)) In other words, the escape rate is determined by not the difference but the ratio L(#") and L((*). This result means that even if the loss barrier height A T, is the same, minima with smaller values of T((0)) are more stable.

---

Under review as a conference paper at ICLR 2022
Moreover, given the fact that the eigenvalue spectrum of the Hessian at minimum consists of a bulk of almost zero eigenvalues and outliers (Sagun et al. 2017: Papyan 2019), our escape-rate formula implies that SGD prefers flat minima with a low effective dimension, where the effective dimension is defined as the number of outliers (MacKay 1992) and flatness is measured by a typical magnitude
2017). The previous theories (Jastrzebski et al. 2017 Wu et al.) 2018 Zhu et al., 2019; Meng et al. 2020; Xie et al. [021] Liu et al.) (221)) have also successfully explained that SGD prefers flat minima, but not shown the preference of small effective dimensions. The logarithmized loss naturally explains the latter, and sheds light on implicit biases of SGD.
of outlier eigenvalues (Keskar al.,
Main contributions: We obtain the following main results:
We derive an equation for approximating the SGD noise in Eq. (6) Remarkably, the SGD noise strength in the mean-square loss is shown to be proportional to the loss function, which is experimentally verified in Sec. [5.2] A key ingredient in deriving Eq. (6) is the
decoupling approximation given in Eq. 7 This is a novel approximate method introduced
in our analysis, and hence we experimentally verify it in Sec. 5.1
We derive a novel stochastic differential equation (SDE) in Eq. 14 via a random time change introduced in Eq. (13). Although the original SDE (4) has multiplicative noise, the transformed SDE (4) has a simple additive noise with the gradient of the logarithmic loss. This shows the convenience of the logarithmic loss landscape for understanding SGD,
We derive a novel form of SGD escape rate from a local minimum in Eq. (16). Remarkably, the escape rate depends on the ratio between L(0*) and L(0*). In Sec. 5.3 we experimen- tally test the validity of this result for linear regressions.
Our escape rate crucially depends the flatness and the effective dimension, which shows that SGD has implicit biases towards flat minima with low effective dimension. We also show in Eq. (18) that local minimum with an effective dimension greater than a certain critical value ne becomes unstable.
Related works: The role of the SGD noise structure has been discussed in some previous works (Zhu et al.) 2019 Xie et al. [021 Liu et al. 2021 Meng et al. 2020; Wojtowytsch| (021). It was pointed out that the anisotropic nature of the SGD noise is important: the SGD noise covariance matrix is aligned with the Hessian of the loss function, which is beneficial for escape from sharp minima (Zhu et al. 2019, Xie et al. 202T: Liu et al. 202T). These previous works, however, do not take the inhomogeneity of the SGD noise strength into account, and consequently, escape rates derived there depend exponentially on the loss barrier height, which differs from our formula.
Compared with the anisotropy of the SGD noise, the inhomogeneity of the SGD noise strength has been less explored. In (Meng et al.] 2020, Wojtowytsch] [021]. the SGD dynamics under a state- dependent noise is discussed. However, in these previous works, the connection between the noise strength and the loss function was not thcoretically established, and the logarithmized loss landscape was not discussed. The instability due to large effective dimensions
also not shown. Another recent work (Pesme et al. (021) observed that the noise is proportional to the loss for specific simple models. In our paper, such result is derived for more generic models. Gürbüzbalaban et al. (2021) showed that SGD will converge to a heavy-tailed stationary distribution due to a multiplicative nature of the SGD noise in simple linear regression problem. Our paper strengthens this result: we argue that such heaved distribution generically appears for the mean-squan loss.
2 BACKGROUND
2.1 SETUP
We consider supervised learning. Let D ((x(())))) = 1,2,,,... be the training dataset, where r((1) E 1 denotes data vector and y(u) E IR be its label. The network output for given input x is denoted by f(,,x:) E TR where A E R. stands for a set of trainable parameters with P being the number of trainable parameters (Extension to the multi-dimennional label and output is straightforward). this work, we focus on the mean-square loss
1 L( 2N
Σ
[(0,z(1)) - )/00)² =;
(0)
(1)
2

---

Under review as a conference paper at ICLR 2022
The training proceeds through optimization of L(0). In most machine-leaming applications, the optimization is done via SGD or its variants. In SGD, the parameter 0x+1 at the time step l: + is determined by
Bk+1 02 vVV (0x), LBB (0) 2B
Σ (0), ME
where 7j > 0 the learning rate, Bg C (1.2.... N) with [B&] = B is mini-batch used the kth time step, and L.B, denotes the mini-batch loss. Since the training dataset 'D is randomly divided into mini-batchest the dynamics defined by Eq. (2) is stochastic. When B = N, the full training data samples used for every iteration. In this case, the dynamics is deterministic and called gradient descent (GD). SGD is interpreted as GD with stochastic noise. By introducing the SGD noise & LVL, (8%) VL()))), Eq. (2) is rewritten as
0k+1 02 V ((p)) T
Obviously, (Ex) = O, where the brackets denote the average over possible choices of mini-batches. The noise covariance matrix defined as Σ(f²)) == (ExEZ) The covariance structure of the SGD noise is important in analyzing the SGD dynamics, which will be discussed in Sec. 3.1
2.2 STOCHASTIC DIFFERENTIAL EQUATION FOR SGD
When the parameter update for each iteration is small, which typically the case when the learning rate 7 is small enough, we can consider the continuous-time approximation (Li al. 2017; |mithh & Le 2018). By introducing a continuous time variable t E IR and regarding 1) as an infinitesimal time step df, we have SDE
d0; = - (L( + VnI((,)) dW,,
(4)
where dW, (OO with I, being the n-by-n identity matrix, and the multiplicative noise v (0.) dW, is interpreted as Itō since the noise & in Eq. (3) depends on 0g but not on 0x+1- Throughout this work, we consider the continuous-time approximation 4 with Gaussian noise.
In machine learning, the gradient Langevin dynamics (GLD) is also considered, which the isotropic and uniform Gaussian noise is injected into the GD
d0, = ((. + v2DdW,,
(5)
where D >> corresponds the noise strength (it is also called the diffusion coefficient) (Sato & Nakagawa 2014 Zhang et al., 2017b Zhu et al., 2019). The stationary probability distribution PaLD (0) of 0 for GLD is given by the Gibbs distribution PCLD(() ε L00)DD We will see in Sec. that the SGD noise structure, which is characterized by E()), drastically alters the stationary distribution and the escape rate from a local minimum.
3 THEORETICAL FORMULATION
3.1 STRUCTURE OF THE SGD NOISE COVARIANCE
The SGD noise covariance matrix ( significantly affects the dynamics (Jastrzębski et al. 2017 Smith & Le 2018 Zhu et al., 2019 Ziyin et al.] 2021). In this section, under some approxima- tions, derive the following expression of E() for the mean-square loss within the the basin of attractions (or the "valley") of a local minimum 0":
[()) ~ 20(0)
NB H(O)),
(6)
where H(0) = L() is the Hessian. It should be noted that VL(0*) = 0 and H(0") is positive semidefinite at any local minima 0" We give a derivation below and the list of the approximations and their justifications in Appendix [A] In particular, the decoupling approxima- tion is a new tool and plays key role in the derivation. It states that the quantities ey and c(uu) (0) vf (l)))/(( f (w)) y are uncorrelated, which implies
1 N Σ
Σen 4Mx
(
N Σ
cf cf
L()) 1 Σ clP
3

---

Under review as a conference paper at ICLR 2022
This approximation is promising for large networks in which ff looks random vector. In Sec. 5.1 we experimentally verify the decoupling approximation for the entire training dynamics.
Our formula (6) possesses two important properties. First, the noise is aligned with the Hessian, which has been well known and pointed out in the literature (Jastrzebski et al. 2017 Zhu et al, 2019 Xie al., 202T Liu et al. 202T). If the loss landscape has flat directions, which correspond to the directions of the Hessian eigenvectors belonging to vanishingly small eigenvalues, the SGD noise does work along those directions. Consequently, SGD dynamics is frozen along those flat directions, which effectively reduces the dimension of the parameter space explored by SGD dynamics. This plays an important role in the escape efficiency. Indeed, we will see that the escape rate crucially depends on the effective dimension of given local minimum.
Second, the noise is proportional in Sec. 5.2 This property has not been pointed out and not been taken into account in previous studies (Jsttzzbski et al. 2017 Zhu et al.) 2019 Xie et al.] [202T] [Lu al. 2021) and therefore gives new insights into the SGD dynamics. Indeed, this property allows us to formulate the Langevin equation on the logarithmized loss landscape with simple additive noise as discussed Sec. 3.2 This new formalism yields the power-law escape rate, i.e. Eqs. [5 and (16) and the importance of the effective dimension of local minima for their stability.
the loss function, which is indeed experimentally confirmed
Derivation of Eq.
(6): We start from analytic expression of (0)), which reads
N Σ
N B Σ(0) B N -
Ve
-
VLVLT
1 N
Σ
VL.ve
-LVLL
VLVLT
(8)
B
where B < N is assumed in the second equality. The derivation of Eq. ( is found in Jastrzebski
ct al. (2017); Smith & Le (2018). Usually, the gradient noise variance dominates the square of the gradient noise mean, and hence the term VLV L in Eq. (8) is negligible.
For the mean-square loss, we have Vl 1f(0,xx00)
y rl()), and hence
(0) ~
N
(0.00 (
2 BN
20(0)
Σ
N
(9)
BN
NB
where the decoupling approximation 7 is used in the last equality. Equation (9) is directly related to the Hessian of the loss function near a (local or global) minimum. The Hessian written as
H(0) = v
Σ
cle O
Σ
[1 200 gnm) 42 ( f(0,x(1)).
(10)
is shown by Papyan 2018 that the last term of Eq. (T)) does not contribute outliers (i.e. large cigenvalues) of the Hessian. Dynamics near a local minimum is governed by outliers, and hence we can ignore this term. At 8", we therefore obtain
H(O0)
Σ cl (0))
(11)
Let us assume c( (0) ((") for 8 within the valley of local minimum 6*, We then obtain
the desired expression (6) by substituting it into Eq.
3.2
LOGARITHMIZED LOSS LANDSCAPE
Let us consider the Itô SDE with the SGD noise covariance 6 near a local minimum 0*, which written as
d0,- VL (m..(
H((0)) dW..
(12)
B
Let us consider a stochastic time {(7) for 2 0 as
((T)
dt' I(Ox),
(13)
4

---

Under review as a conference paper at ICLR 2022
and perform random time change from to (Øksendal 1998). Conrespondingly, introduce the Wiener process dW, (0, Ipdr). Since dr (() dt, have dW, ((0,) dW. In
terms of the notation 8 f:, Eq. 12) is expressed as
db, - 1 (((0,) + I(0+))
2 (0)) dW, = - V log L(o)) dr + (0+)
B dW,. (14)
B
We should note that at global minimum with L( O. which realized in an overparameterizee regime (Zhang et al.) (017a), the random time change through Eq. (13) is ill-defined since т is frozen finite value once the model reaches global minimum. We can overcome this difficulty by adding an infinitesimal constant > 0 to the loss, which makes the loss function positive without changing the finite-time dynamics like the escape from local minimum 8 with L(#*)= > 0.
In this way, the Langevin equation on the loss landscape L( with multiplicative noise is trans- formed to that on the logarithmic loss landscape U()) log I.(0) with simpler additive noise. This formulation indicates the importance of considering the logarithmized loss landscape U(0)
log L()). In the following, we use Eq. (14 to discuss the escape efficiency from local minima.
ESCAPE RATE FROM LOCAL MINIMA
Let us consider a local minimum 0 and its basin of attraction Age, which is the set of all the starting points Oo such that O, tends to 0* as t + oc if there is no noise. In order to escape from the basin of 8", , Oe must go through a saddle 4a with the help of noise. The escape time is defined as T inf > 0 0g &f A.a..B = }. The escape rate K is then defined as the inverse of the (T)
-1 Our idea for evaluating the escape rate is applying the Kramers formula (Kamers 1940) or its high-dimensional generalization (Langer 1969 Bovier et al. [004] Berglund 2013) to Eq. (4). It should be noted that the Kramers formula is applicable only for
mean escape time:
additive noise, and hence it cannot be used for the original SDE (4) with multiplicative noise.
First, we present main results [Eqs. (i.e. the single-variable case A E RR, the escape rate through a saddle 0 is given by
(16). and (I7)]. and give their derivations later. For P=
T(00) L(O+)]
B
k
hir ha
(15)
2T
where 4 = -O0).. This is a variant of the Kramers formula (Kramers) 1940). which is accurate when i is small. We can derive Eq. (15) from Eq. (4) without any further approximations.
For O E I²' with P > 1, the analysis is more involved and needs some assumptions and approx- imations, which are listed in Appendix [A] We assume that the eigenvalue spectrum of the Hessian (0*) at the local minimum 0* consists of bulk of almost zero eigenvalues and outliers, which is indeed empirically correct (Sagun et al., 2017 Papyan 2019). We denote by h* and n a typical magnitude of outlier eigenvalues (i.e. the flatness) and the number of outlier eigenvalues (i.e. the ef- fective dimension), respectively. As we discussed in Sec. Sec.B.1 the SGD dynamics frozen along flat directions. Therefore SGD dynamics around the local minimum 0* is restricted to the n-dimensional manifold spanned by the outlier eigenvectors U1, U2...., Un E R²" of H (8"). Now we parameterize .....Uyy.
0 by using n parameters 21, 22, ... 2n R as 0 0" + + Z,U. The Hessian restricted to this outlier subspace is then written as (()) V L()) Rax..
We then obtain the following escape rate formula for P > 1:
|ng| det (0*)
L(0*)
(L 11+1--11
K
(16)
2x det ((*))
L(0*)
where lgg is the negative eigenvalue of H(0*) corresponding to the escape direction Again, Eq. (16) valid when e is small enough. It should be noted that Eq. (6 reduced to Eq. (5) when n =T.
IIf eigenvalues of (00) and those of (()0)) coincide with each other except for the es- cape direction e. det (0r) det (0*) = hg/h2 holds, and Eq. (16) is simplified as K
hh2
+1-3).
5

---

Under review as a conference paper at ICLR 2022
For any P. the quasi-stationary distribution (0) within the valley of local minimum 0-, which can be identified as the stationary distribution restricted to the valley (Bianchi & Gaudillière 2016).
is written
B
(0) x L(0) Φ -=1+
(17)
1//-"
Remarkably, it depends L(O) polynomially rather than exponentially as in the standard GLD (Jsstzzbbkki et al.) 2017; Sato & Nakagawa 2014; Zhang et al. 2017b). This polynomial dependence on L(B) key feature leading to the escape rate formula.
From Eq. (16) we can obtain some implications. The factor [[())/ with h- and n, which indicates that sharp minima (i.e. minima with large h²) or minima with large are unstable. This fact explains why SGD finds flat minima with low effective dimension n. Equation (16) also implies that the effective dimension of any stable minima must satisfy
increases
B
ng ==
17h- +
(18)
The instability due to large effective dimension is a new insight naturally explained by the picture of the logarithmized loss landscape. It arises from the ratio of the determinants of the logarithmized-
loss Hessian: det (2 //1 det ((*) = |det H (
/ det (( [L(+') [[)()))))))))))))))))))))))))[
Derivation of Eq. (16) and Eq. (17): Now we give derivation of Eqs. (16) and (17). Equation (15)
straightforwardly obtained by putting P in the derivation below (some approximations are introduced below, but all of them are not necessary for P =1 1).
As we already noted, SGD dynamics is restricted to the n-dimensional outlier subspace. First, we assume that the anisotropy of the SGD noise within this n-dimensional space is not relevant, and approximate the SGD noise in Eq. (14) as an isotropic one:
//dW,, R /2nh"/Bdl T. where h' E R- characterizes the magnitude of the Hessian outliers. This assumption is justified when the loss landscape is isotropic within the n-dimensional subspace near the minimum. Even if the Hessian at the minimum is not isotropic, this approximation is justified when the directions of the Hessian eigenvectors do not change within the valley. In the latter case, the escape path is a straight line along the direction of Hessian eigenvector Ue, where E (1, n) identifies the escape direction, and h* corresponds to the Hessian eigenvalue at the minimum along the escape direction, i.e. h" = ha- See Appendix for the details.
Under this approximation, Eq. (14) becomes dz, -UU V2TdW,, where T= =n/h² 3ph"/B. Its
quasi-stationary distribution Pe. (0) within the valley including the local minimum 0 = 0 (i.e.
0) is then given by a Gibbs distribution with respect U(8): P... (0) x UU0)TT x L())-B/(nh'),
This is the distribution function of 0- for a fixed т. However, what we want is the quasi-stationary distribution of Og for a fixed l. In Appendix
B by using Fq. (13). we show that the two distributions are related with each other as Pss (0) -
We thus obtain Eq. 17
The escape rate under the Langevin equation with isotropic additive noise is evaluated by using celebrated Kramers formula (Kramers 1940) or its high-dimensional generalization (Langer) 1969.
Bovier et al. 2004; Berglund 2013). According to it, the escape rate Ni per unit T is given by
|
det VZU(0~) AU/T
Kr
2x
det VZU(0=)
(19)
where U., is the negative eigenvalue of (2T( corresponding to the escape direction, AU U(e*)= - U(8*) ) is called the potential barrier, and AU/T is assumed to be large enough. What we really want is the escape rate per unit time L. It is a reasonable assumption that O, stays close to O" for most times before escape, and hence T is approximately given by ~ L(O").. The escape rate K per unit t is then given by K = I(A*),, By combining it with Eq. (19 and substituting (0) log L(0), we finally obtain Eq. (66).
5 EXPERIMENTS
Our key theoretical observation is that the SGD noise strength is proportional to the loss function, which is obtained result of the decoupling approximation. This property leads to the Langevin

---

Under review as a conference paper at ICLR 2022
(a) initialization
(b) epochs
(C) 500 epochs
10a
10'
102
10a
10a
10a
107
109
10a
-20
-10
-20
-20
-10
10
logleigenv alues)
logt eggeeigenvalues
lnglegenvalues)
Figure 1: Comparison of the eigenvalue distributions left-hand side (exact expression) the right- hand side (decoupled onc) of Fq. (12) main text. They agree with each other except for very small cigenvalues during the entire training dynamics.
10
3CD abronghht223 189
wiangron25 100
10a
an Contend
500 votse Freenc
10-4
10a
10-2
10
10a
10-1
opoch
opoch
Figure 2: Training dynamics the loss and the SGD noise strength N. In the figure, we multiplied N by a numerical factor to emphasize that N is actually proportional to the loss in a later stage of the training. (Left) Fully-connected network trained by the Fashion-MNIST MNIST dataset. (Middle) Convolutional network trained by the CIFAR-10 dataset. (Right) the loss A in the training of the convolutional network. The dashed line straight line of slope I, which implies N x L(O).
equation 14p with the logarithmized loss gradient and an additive noise through random time change (13) Equation (14) implies the stationary distribution (17) and the escape rate (16).
In Sec. 5.1 we show that the decoupling approximation valid during entire training dynamics. In Sec. 5.2 we measure the SGD noise strength and confirm that it is indeed proportional to the loss function near minimum. In Sec. 5.3] we experimentally test the validity of Eq. (17) for the stationary distribution and Eq. (T6) for the escape rate. We will see that numerical results for linear regression and for a non-linear neural network agree with our theoretical results.
5.1 EXPERIMENTAL VERIFICATION OF THE DECOUPLING APPROXIMATION
Let us compare the eigenvalue distribution of the exact matrix (1/N) EN (lr) with that of the decoupled one (( (I/N) Σ cfs) with c(u) = V0 x(1)) (O (T.
consider a binary classification problem using the first 10 samples of the MNIST dataset such that we classify each image into even (its label is = +1) or odd number (its label is = - 1). The network has two hidden layers, each of which has 100 units and the ReI.U activation, followed by the output layer of a single unit with no activation. Starting from the Glorot initialization, the training is performed via SGD with the mean-square loss, where fix 7 0.01 and B = 100.
Figure D shows histograms of their eigenvalues at different stages of the training: (a) at initialization, (b) after 50 epochs, and (c) after 500 epochs. We see that the exact matrix and the approximate one have statistically similar eigenvalue distributions except for very small eigenvalues during the training dynamics. This means that the decoupling approximation always holds during training.
5.2 MEASUREMENTE OF THE SGD NOISE STRENGTH
As measure of the SGD noise strength, let consider the norm of the noise vector & given by (5Tς) Tr Σ ~ A/B, where / (1/N) Σ Ve, V Here present experi- mental results for two architectures and datasets. First, consider training of the FaShion-MNIST dataset by using a fully connected neural network with three hidden layers, each of which has 2 x 10 units and the ReI.U activation, followed by the output layer of 10 units with no activation (classifica-

---

Under review as a conference paper at ICLR 2022
Exponent @ (0)
(b) by for the linear regression
(c) tp for the neural network
300000
1000
100
10000
TyHELE
30
10
1
1-00000
0.1
1-00012
1000
Figure 3: (a) Exponent for the stationary distribution (0) L( for d = in the linear regression. Dashed lines show the theoretical prediction e = 1 + B (i/h ). (b) Log-log plot of the mean first-passage time tp vs c = L((*) ((8") for B = and 1 = 0.1 in the linear regression. Dashed lines show the theorctical prediction, tp x cdd//2
d/2, (c) Log-log plot of tp vs e for B 10 and various 7 in the neural network. Dashed lines show the theoretical prediction, tp x K=- BB((xh")) (0/(("44+11000 with h- 95.6 and 9.
tion labels are given in the one-hot representation). Second, we consider training of the CIFAR-10 dataset by using a convolutional neural network. Following Keskar et al. (2017). , let us denote a stack of n convolutional layers of a filters and kernel size of b x with the stride length of d by n x [a. b, c, d]. We use the configuration: 3 x [64, 3, 3, 1). 3 x [128, 3, 3, 1), 3 x [256, 3. 3, 1], where MaxPool(2) is applied after each stack. To all layers, the ReLU activation is applied. Finally, an output layer consists of 10 units with no activation.
Starting from the Glorot initialization, the network is trained by SGD of the mini- size B 100 and = 0.1 for the mean-square loss. During the training, we measure the training loss and the noise strength A for every epoch. Numerical results are given in Fig. 2 We see that roughly A L at later stage of the training, which agrees with our theoretical prediction.
Although A is not proportional to L at early stage of training, it does mean that Eq. 9 is invalid there. Since the decoupling approximation is valid for the entire training dynamics, Eq. (9) always holds. The reason why the SGD noise strength does not decrease with the loss function in the early-stage dynamics is that N x 2L(0)= x (1/N) Ex- (0,, x(H)) (0,, (x))), but the quantity
(1/N) Σ N
f(,, a((")) f0, f x(()) also changes during training.
Although Eq. p derived for the mean-square loss, the relation W xx T.(0) holds in more general loss functions; see Appendix D for general argument and experiments on the cross entropy loss.
5.3 EXPERIMENTAL TEST OF STATIONARY DISTRIBUTION AND ESCAPE RATE FORMULA
We experimentally verify our theoretical predictions. Below, we first present numerical results for a simple linear regression problem, and then for nonlinear model, i.c., a neural network.
Let start from the following linear regression problem: each entry of 2(1)) E Rd and its la- bel y(u) E I are i.i.d. Gaussian random variables of zero mean and unit variance. We fo- cus on the case of d & N. The output for an input x is given by f(0,x)) =QTr. 0Tx, where 0 E Rd is the trainable network parameter. We optimize 0 via SGD. The mean-square loss L(0) = (1/2N) Em= (0z(w) - y(m))² 2 is quadratic and has unique minimum at 1 x 0.
First, we test Eq. (77. i.e. the stationary distribution, for d = 1 and N = 10a We sampled the value of 0% at every 1O0 iterations (A: j 100, j 1. 104) and made histogram. We then fit the histogram to the form (0) x T.(0)-* and determine the exponent db Our theory predicts 1+ (.(n")) Numerical results for the exponent & are presented Fig- 3 (a) against for three fixed learning rates 7). In the same figure, theoretical values of 6 are plotted in dashed lines. The agreement between theory and experiment is fairly well. For a large learning rate 77 = 1, the exponent slightly deviates from its theoretical value. This is due the effect of finite learning rate (recall that 7) is assumed to be small in deriving the continuous-tin time stochastic differential equation).
Next, test our formula the escape rate, Eq. (6). Although the mean-square loss is quadratic and no barrier crossing occurs, we can measure the first passage time, which imitates the escape time for non-convex loss landscape. Let fix threshold value of the loss function. The first
8

---

Under review as a conference paper at ICLR 2022
passage time ty is defined as the shortest time at which the loss exceeds the threshold value. Here, time is identified as 1pk, where & denotes the number of iterations in discrete SGD 2 We identify the threshold value as L(0*). i... the loss at the saddle in the escape problem. It is expected that tp is similar to the escape time and proportional to x=
The Hessian H = ( 2z(1))((17) has d nonzero eigenvalues, all of which are close to unity. We can therefore identify --= and d. The mean first passage time over 100 independent runs is measured for varying threshold values which are specified by tal results for N 10t are presented in Fig. (b). Dashed straight lines have slope B/(nh") +
L[0*)/L((e
> 1. Experimen-
n/2. Experiments show that the first passage time behaves as 1p x [T(0*)/ ( [] /// (//(h/++1///2 ++1-n/2 which agrees with our theoretical evaluation of x= [see Eq. (16)]. We conclude that the escape rate crucially depends on the effective dimension n, which is not explained by the previous results (Zhu
et al. 2019 Xie et al., 2021 Liu al. 2021 Meng et al. 2020).
Our escape-rate formula (6p is also verified in a non-linear model. As in Sec. 5.1 we consider the binary classification problem using the first 10 samples of MNINST dataset such that we classify each image into even or odd. The network has one hidden layer with 10 units activated by ReLU,, followed by the output layer of single unit with activation. We always use the mean-square loss. Starting from thee Glorot initialization, the network is pre-trained via SGD with = 0.01 and B 100 for 10a iterations. We find that after pre-training, the loss becomes almost stationary around at L()) x 0.035. We regard that the pre-trained network is near local minimum. We then further train the pre-trained network via SGD with new choice of r) and B (here we fix = 10), and measure the first passage time tp- The mean first fass--assaae time over 100 independent runs is measured for varying threshold values which are specified by e = L(0*)/L(0') > 1. Experimental results are presented in Fig. 3 (c). We see the power-law behavior, which is consistent with our theory.
To further verify our theoretical formula (66.. we also compare the power-law exponent for the mean first-passage time with our theoretical prediction B/ + 1 n/2. Here, ht and n are estimated by the Hessian eigenvalues. In Appendix E we present numerical result for eigenvalues of the approximate Hessian given by the right-hand side of Eq. (T). By identifying the largest eigenvalue of the Hessian ht, we have h" 95.6. On the other hand, it is difficult to precisely determine the effective dimension n, but it seems reasonable to estimate n ~ 10. It turns out that theory and experiment agree with each other by choosing n 9. Dashed lines in Fig. (c) correspond to our theoretical prediction n=1 x (B/(nh")+1-m) with h" 95.6 and 9. This excellent agreement shows that our theoretical formula (16) is also valid in non- non-linear models.
6 CONCLUSION
In this work, we have investigated SGD dynamics via a Langevin approach. With several approx- imations listed in Appendix A we have derived Eq. (6), which shows that the SGD noise strength is proportional to the loss function. This SGD noise covariance structure yields the stochastic dif- ferential equation (14) with additive noise near minimum via random time change (13). The original multiplicative noise is reduced to simpler additive noise, but instead the gradient of the loss function is replaced by that of the logarithmized loss function U(O) log ((). This stochastic differential equation has a quasi-stationary distribution that decays polynomially with L(0) near minimum (17), not exponentially as in the usual Gibbs distribution. This new formalism yields the power-law escape rate formula (16) whose exponent depends v), B, h²" and n.
Our escape-rate formula explains an empirical fact that SGD favors flat minima with low effective dimensions. The effective dimension of minimum must satisfy Eq. (8) for its stability. This result as well as the formulation of SGD dynamics using the logarithmized loss landscape should help understand more deeply the SGD dynamics and its implicit biases in machine learning problems.
Although the present work focuses on the Gaussian noise, the -on-Gaussianity can also play an important role. For example, (Simşekli et al. (2019) approximated SGD as Lévy-driven SDE, which explains why SGD finds wide minima. would be an interesting future problem to take the non-Gaussian effect into account.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Victor Bapst, Thomas Keck, A Grabska-Barwinske Craig Donner, Ekin Dogus Cubuk, Samuel Schoenholz, Annette Obika, Alexander WR Nelson, Trevor Back, Demis Hassabis, and Pushmeet Kohli. Unveiling the predictive power of static structure in glassy systems. Nature Physics, 16 (4):448-454, 2020.
Nils Berglund. Kramers' Law: Validity, Derivations and Generalisations. Markov Processes and Related Fields, 19:459 490, 2013.
Alessandra Bianchi and Alexandre Gaudillière. Metastable states, quasi-stationary distributions and soft measures. Støchastic Processes and their Applications, 126: 166116221680, 2016.
Anton Bovier, Michael Eckhoff, Véronique Gayrard, and Markus Klein. Metastability in reversible diffusion processes I. Sharp asymptotics for capacities and exit times. Journal of the European Mathematical Society, 6:399-424, 2004.
Ronan Collobert and Jason Weston. A unified architecture for natural language processing: Deep neural networks with multitask learning. In International Conference on Machine Learning. 2008
Mert Gürbüübalaban, Umut Şimşekli, and Lingjiong Zhu. The Heavy-Tail Phenomenon in SGD. In International Conference on Machine Learning, 2021.
Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, and Others. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. JEEE Signal processing magazine, 22(6)822-97, 2012.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. In Advances in Neural Information Processing Systems, 2017.
Raban Iten, Tony Metger, Henrik Wilming, Lídia Del Rio, and Renato Renner. Discovering Physical Concepts with Neural Networks. Physical Review Letters, 124(1): 10508, 2020.
Stanislaw Jastrzębski. Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Ben- gio, and Amos Storkey. Three Factors Influencing Minima in SGD. arXiv:1711.04623, 2017.
Nitish Shirish Keskar, Jorge Nocedal, Ping Tak Peter Tang. Dheevatsa Mudigere, and Mikhail Smelyanskiy. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
Hendrik Anthony Kramers. Brownian motion in a field of force and the diffusion model of chemical reactions. Physica, 7(4):284-304, 1940.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. In Advances in Neural Information Processing Systems, pp. 1097-1105, 2012.
James S. Langer. Statistical theory of the decay of metastable states. Annals of Physics, 54(2): 258-275, 1969.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 521(7553))446-44, 2015.
Qianxiao Li, Cheng Tai, and E. Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference on Machine Learning, 2017.
Kangqiao Liu, Liu Ziyin, and Masahito Ueda. Noise and Fluctuation of Finite Learning Rate Stochastic Gradient Descent. In International Conference on Machine Learning, 2021.
Dje MacKay. Bayesian model comparison and backprop nets. In Advances in Neural Information Processing Systems, 1992.
10

---

Under review as a conference paper at ICLR 2022
Qi Meng, Shiqi Gong, Wei Chen, Zhi Ming Ma, and Tie Yan Liu. Dynamic of Stochastic Gradient Descent with Satee-eeenddett Noise. arXiv:2006. 2020.
Bernt Øksendal. Stochastic differentiall equations: an introduction with applications. Springer, Berlin, 1998.
Vardan Papyan. The Full Spectrum of Deepnet Hessians Scale: Dynamics with SGD Training and Sample Size. 1rXiv:/8/1.07066 2018.
Vardan Papyan. Mcasurements of three-level hierarchical structure in the outliers in the spectrum of deepnet hessians. In International Conference on Machine Learning, 2019.
Scott Pesme, Loucas Pillaud-Vivies and Nicolas Flammarion. Implicit Bias of SGD for Diagonal Linear Networks: Provable Benefit of Stochasticity. arXiv:2106.09522 2021.
Levent Sagun, Utku Evci, Ugur Güney, Yann Dauphin, and Léon Bottou. Empirical analysis of the hessian of over wee-ppaamettrized neural networks. arXiv:/706.044555 2017.
Issei Sato and Hiroshi Nakagawa. Approximation analysis of stochastic gradient langevin dynam- ics by using fokker-planck equation and ito process. In International Conference on Machine Learning. 2014.
Alireza Seif, Mohammad Hafezi, and Christopher Jarzynski. Machine learning the thermodynamic arrow of time. Nature Physics, 17:105-113, 2021.
Umut Şimşekli, Levent Sagun, and Mert Gürbiizbalaban. A tail-index analysis of stochastic gradient noise deep neural networks. International Conference on Machine Learning, 2019.
Samuel L. Smith and Quoe Le. A Bayesian perspective on generalization and stochastic gradient descent. In International Conference Learning Representations, 2018.
Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. Part II: Continuous time analysis. arXiv:2106.0253: 2021.
Jingfeng Wu, Wenging Hu, Haoyi Xiong, Jun Huan, Vladimir Braverman, and Zhanxing Zhu On the Noisy Gradient Descent that Generalizes as SGD. In International Conference on Machine Learning. 2020.
Lei Wu, Chao Ma, and E. Weinan. How SGD selects the global minima in wer-parameterrize learn- ing: A dynamical stability perspective. In Advances in Neural Information Processing Systems, 2018.
Zeke Xie, Issei Sato, and Masashi Sugiyama. A Diffusion Theory For Deep Learning Dynamies: Stochastic Gradient Descent Exponentially Favors Flat Minima. In International Conference on Learning Representations, 2021.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding Deep Learning Requires Rethinking of Generalization. In International Conference on Learning Representations, 2017a
Yuchen Zhang, Percy Liang, and Moses Charikar. A hitting time analysis of stochastic gradient Langevin dynamics. In Proceedings of Machine Learning Research, 2017b.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochas- tic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning, 2019.
Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. On Minibatch Noise: Discrete-Time SGD, Overparametrizatton, and Bayes. arXiv:2102.0537. 2021.
11

---

Under review as a conference paper at ICLR 2022
A LIST OF APPROXIMATIONS AND THEIR JUSTIFICATIONS
In Sec. 3.1 we made several approximations to derive the result 6 For clarity, we list the approxi- mations made and their justifications below.
In Eq. (), we make the approximation of B << N, which simplifies the expression but is not essential. The main conclusion is not affected by this approximation.
We ignore V in Eq. (8), which is justified near local minimum.
In Eq. (9), make the decoupling approximation 7 which is one of the key heuristic approximations in our work. This approximation is verified experimentally in Sec. 5.T
We ignore the last term of Eq. TO). which is a common approximation (Sagun et al. 2017 This approximation is justified if are only interested in the outliers of the Hessian eigen- values. Indeed, outliers of the Hessian cigenvalues, which play dominant roles in escape from local minima, are known to be attributable to the first term of the right-hand side of Eq. 10
(Papyan 2018).
In Eq. 11 we assume that the matrix (1/N) Ex
Vf(0, x(()) V x()) does not change so much in a valley with given local minimum. This approximation is indirectly verified in Fig. 2 (the proportionality between the loss and the noise strength implies this matrix is actually constant).
In deriving the escape rate for nulti-variable cases in Sec. we further make the following assump- tions and approximations:
First of all, we discard a bulk of near-zero eigenvalues of the Hessian and restricted learning dynamics within the subspace spanned by the outlier eigenvectors. This approximation is justified because the SGD dynamics is almost frozen along flat directions of the loss landscape.
We approximate the SGD noise as an isotropic one within the n-dimensional outlier sub- space. This approximation is justified either when the loss landscape in the outlier subspace is isotropic near the minimum or when the directions of the Hessian eigenvectors do not change within the valley.
At the last part of the derivation, we approximate L(O"),, which is justified because the model stays near a local minimum until escapes from that minimum.
B STATIONARY DISTRIBUTION
Since 0, obeys simple Langevin equation
dū_ = -U'((,)) + 22TW,,
(20)
the stationary distribution of 6, is given by the Gibbs distribution (f)( x e /T On the other
hand, what we want is the stationary distribution (o) of 0r, where 0g = 0, with r = = di' T.(Oe).
In this section, we show the relation between the two distributions: (x) x L(8))-1Pss(!)
We express the stationary distributions in terms of the long-time average of the delta function:
(0) lim 1 /" dt 8(0, 0), Paa(O) lim [" dr 6(0, 0).
J
(21)
By using the relation = dt' I(Or), we have dr L(O,) dt. For sufficiently large t, we also obtain
LT, where T. = lim (1//) (1/8) Jo dt' T(Ox)) denotes the long-time average of T(..).. By
12

---

Under review as a conference paper at ICLR 2022
using them, (0) is rewritten as
ss dr
8(0, - 8)
(0) Rc lim
80000
L(O+)
L at.
lim
dr ö(0, -0)
I(0))
5-000 1h
T, (o)
(22)
L(0)
We thus obtain the desired relation, (0) x L(o)
C JUSTIFICATION OF THE ISOTROPIC-NOIST APPROXIMATION WITHIN THE 1-DIMENSIONAL SUBSPACE
We now show that the isotropic-noise approximation is valid when the directions of the eigen- vectors of the Hessian do not change within the valley of a given local minimum. In this case, 0
for any * where the Hessian at the minimum is given by H(0)) z E h2 Fa'n (h, is an eigenvalue and v, is the corresponding eigenvector of (())) and the dis- placement vector R" is defined by 0 = 0" + Z,0% The stochastic differential equation d= UUUdT (0*) is then equivalent to the following Fokker-Planck equation for the distribution function P)=,T) of at
0P(z,T) ÖT
11
OU
THE 8² B Oz?
(23)
Ozh
02i
Let us assume that the direction of eth eigenvector Ue corresponds to the escape direction. We denote by zy E IR-11 the displacement perpendicular to the escape direction, and write 2 = (Ze, =1). At the saddle 2². (="( of the Hessian at 2..
} is the of eigenvectors he < 0 and hद > for all o e, where {hā}
Under the above setting, we derive the escape rate formula following Kramers 1940). The steady current J E IR' is aligned to the escape direction, and hence Je * O and J1 Let denote by Pt the total probability within the valley of a given minimum 0" and by J the total current flowing to outside of the valley through the saddle 0" = 0* = It is assumed that J is small and P* is almost stationary. The escape rate K, is then given by
J
N₂ p
(24)
The escape rate per unit time is given by
K = L(8")(())) = L(0") J
(25)
p.
We now evaluate Pa and J. To evaluate Pr, it necessary to know about the stationary distribution near the minimum 0*. Near the minimum 0* (i.e. small 2).
U(0) = log ((0) R U(0") + (0 (0- 0T) H(O)) - 0*) ((0)(000(000000
=0H(0')z
=U(0-) 2T((0*)
(26)
2L(0-)
and
aU 1
n
Hiyay.
(27)
0z.
L(0-)
j=1
where H = (80)). By substituting it into the Fokker-Planne equation
23
we obtain
0P(2.T)
Σ
Σ
H
2j
+
"7
P
(28)
Iza
L(0-)
B
oz,
1=1
Ø=1
13

---

Under review as a conference paper at ICLR 2022
U(0))U(0)):0 =
P())
J path
0 (0 = P*)
(0 0*)
Figure 4: Schematic illustration of the escape from potential barrier.
for j
2j +
17 θ
Pss -0,
(29)
T.(0*) B 0zz)
Psa(2) is stationary distribution near the minimum 0" Indeed,
(s)(2) P(0")e (,, Σ
(30)
satisfies this condition. It should be noted that the stationary distribution is independent of the Hessian eigenvalues {hē ). By using Eq. (30). p is evaluated as
pt ~ dzndzz ... de, (z) P(O)) (n n/2
(31)
B
Next, let evaluate J. The probability current J along the escape direction is given by
OU
1 θ P. B O2.
Je az.
(32)
At the saddle, the current vector Jn perpendicular the escape direction zero, and hence the probability distribution near the saddle with = and + 0 is also evaluated in similar way as in Eq.
P( 21, R. P(ze 2 21 0)e Sallory
(33)
By substituting it into Fq.
32
we obtain
16(2,,21 Jpathe .nc(or =
(34)
where the current along the escape path (= 0) denoted by Jpath == J(z.. = =). The total current through the saddle is then evaluated as
(00)) 1n/1)/2
J=
J
f=fdez tzp Je(=e, Jpach
(35)
B
When the distribution function is almost stationary, Eq. 23 yields 21 0)/Oz. s 0. and hence the current along the escape path is approximately constant (Ze = 0) Jpath. Since
Je [z 21 = 0) is given by Eq. (2) by putting æ1 0. we have
Jpath = 0U
Tha OP =
Th n U
(ene (em=Up)
(36)
Oze
B Oze
B
Oze
=1=0
By multiplying e in both sides and integrating over ie from to 2g, where 2g defined as
U + 25ve) = (((r)) (see Fig. 4h, we obtain
Jpath J
1
= dz wh.
BU
Th
U(8')
P(0"),
(37)
B
14

---

Under review as a conference paper at ICLR 2022
where it is assumed that the probability 2 is small and negligible, P(ze 2 21 0) = 0. By using the saddle point method, the integral in the left- beft-hand side of Eq. (37 is evaluated as
dže e
U Fd
dže exp
B mH
U(00)
+
he (ze =2)²
2T.
2xnh
1/2
B
U
th
(38)
B L(0*)
By substituting this result in Eq. (77.. we obtain
1/2
AU P(O*).
Jpath
2xz BL(00)
e
(39)
where AU = UU())) (.²). The total current J Eq.
(35
is then expressed as
v h
(0m
n/2
J
mh
AU
P(O*).
(40)
2n ((0))
By using Eqs.
31 and (40), the escape rate in Eq.
25 is evaluated as
v hal hel
L(0*) I(0*)
K L(00) p.
AU
(41)
2x
Since AU log)L(())/L((0)} we finally obtain
Vha hel
[(00)
+1-
(42)
2x
L(0*)
which is exactly identical to the escape rate formula derived in the main text using the isotropic- noise approximation with h = hi:
In this way, the isotropic-noise approximation is justified even when the loss landscape is not isotropic near the minimum. We have assumed that the directions the eigenvectors the Hessian do not change within the valley of given local minimum. Under this assumption, the Fokker-Planck equation along the escape path is decoupled from the perpendicular directions 2 Moreover, we have seen that the stationary distribution near minimum is independent of the Hessian eigenvalues. These properties explain the reason why the isotropic- noise approximation is justified in this case.
D OTHER LOSS FUNCTIONS
In our paper, we mainly focus on the mean-square loss, for which we can analytically derive the relation between the loss (()) and the SGD noise covariance E(O). An important observation is that the SGD noise strength N proportional to the loss, i.c., A' x (()) (see Sec. 5.2 of the main text for the definition of N).
Here, we argue that the relation A x L(0) also holds in more general situations. During the training. the value of l, will fluctuate from sample to sample. At certain time step of SGD, let us suppose that N - . samples in the training dataset are already classified correctly and hence l ~ 0. whereas the other M samples are not and hence ly 1. The loss function is then given by L(A) = (1/N) EN ly ~ M/N. When ly is small, Vl,, will also be small. Therefore, for N M samples with la ~ O, Vlp ~ also holds. The other M samples will have non- small gradients: [[Ve,, g. where g > is a certain positive constant. We thus estimate A as N = (1/N) E VCTVE, ~ gM/N ~ gL()). In this way, N x L(0) will hold, irrespective of the loss function.
However, we emphasize that this is crude argument. In particular, the above argument will hold near global minimum because all the samples are correctly classified there, which implies that lu is small for all μ. in contrast to the above argument relying on the existence of M wrongly classified samples with l, 1 and Ve,, g.
15

---

Under review as a conference paper at ICLR 2022
(a)
(b)
10a
10
10-2
10
10 4.
10-3
10-6
10
loss
10-4.
loss
10
noise
10")
60
80
100
60
80
100
epoch
epoch
(c)
(d)
10a
5GD strength
strength
10a
102
10
10'
10
109
10-3
10-1
10-2
10
10 10-2
10e
loss
loss
Figure 5: Training dynamics the loss function and the SGD noise strength N for (a) a fully connected network trained by the Fashion-MNIST dataset and (b) convolutional network trained by the CIFAR-10 dataset. In the figure, we multiplied A by numerical factor to emphasize that N actually proportional to the loss in an intermediate stage of the training. Loss N for (c) fully connected network trained by the CIFAR-10. 10. Dashed lines in (c) and (d) are straight lines of slope 1, which imply A L(A).
and (d) a convolutional network trained by
We experimentally test the relation N L(0) for the cross-entropy loss. We consider the same architectures and datasets in Sec. 5.2 of the main text: a fully connected neural network trained by Fashion-MNIST and convolutional neural network trained by CIFAR-10 (see Sec. 5.2 for the detail). We fix B = 100 in both cases, and 17 0.1 for the fully connected network and 7 0.05 for the convolutional network. Experimental results are presented Fig. We find that the relation /V x [() seems to hold true at an intermediate stage of the training dynamics, although the proportionality is less clear compared with Fig. 2 the main text for the mean-square loss.
We also find that for sufficiently small values of the loss, x L()) [see Fig. 5 (c) and (d)), whose
implications should merit further investigation in future studies.
E HESSIAN EIGENVALUES FOR NEURAL NETWORK IN SEC. 5.3
We present numerical results on Hessian eigenvalues in a pre-trained neural network studied in Sec. (.3 Instead of the exact Hessian, we consider an approximate Hessian given on the right-hand side of Eq. (1] i.e.
H(0))
VS(0, xz (00,
(43)
N
Eigenvalues } are arranged in descending order as h 2 h2 2 hp (in our model 7861).
histogram of the Hessian eigenvalues presented in Fig. (a). We see that most eigenvalues are close to zero, which corresponds to the bulk, but there some large eigenvalues. which correspond to the outliers. The largest eigenvalue is X₁ 95.6, which is identified as h² in our theoretical formula (16).
16

---

Under review as a conference paper at ICLR 2022
(b)
10
80
10a
60
10
10
20
10
0
20
40
60
80
100
1
4
7
10
13
16
19
22
25
28
einenvalues
Figure 6: Hessian eigenvalues for pre-trained neural network studied in Sec. 5.3 of the Hessian cigenvalues. Most cigenvalues are close to zero, but there are some large cigenvalues, which correspond to outliers.
(a) The histogram
Another important quantity is the effective dimension corresnding to the number of outliers. Since the outliers and the bulk are not sharply separated, it is difficult to precisely determine In Fig, 6 (b), we plot h; up to 30. From this figure, it seems reasonable to estimate 10.
As heuristic method of determining we can consider the following identification: first we define the weight p; for ith eigenmode p; h4/ Σド=1 hj- We then determine
11
1
Σ
(44)
which gives n = 12.7 our case.
In Sec. 5.3 our formula (16) with 9 explains numerical results on the first-passage time.
17


Reviews:
[{"review_id": "IXfDllMZD2J", "reviewer": "Reviewer_zZQR", "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": "This paper studies the linear dependency of SGD noise with the loss. Theoretically, this linear dependency comes from a simplification of the noise covariance matrix. Numerically, it is verified for neural network models. With the linear dependency, SGD is formulated as a process with isotropic noise on the landscape of log-loss. Thus, the escape of SGD from minimum no longer depends exponentially on the barrier height. Rather, it depends polynomially on the height. Further, the escape rate depends on the effective non-vanishing dimension of the landscape. The connection of the stability of SGD around minima with the effective dimension is novel. \n\nThis paper is well written and well organized. Though the idea of studying the location-dependent, anisotropic noise of SGD is not new, this paper proposes a new approach to the problem by smartly simplify the noise covariance. The major concern of the reviewer is the lack of numerical support to the power law escape rate and stationary distribution for nonlinear models. Currently, models used in figure 3 is a linear regression. Is it possible to conduct similar experiments for nonlinear models such as neural networks? Will the results be different? I am willing to increase the score if strong numerical results are provided. \n\nSome minor comments are:\n1. In page 4, below equation (11), the authors mentioned \"the last term in equation (11) does not contribute to outliers of Hessians\". Could the authors provide some explanations? Also, it is good to explain what is an outlier of the Hessian. (though later we will know it means nonzero eigenvalues of the Hessian).\n\n2. In page 6, above equation (20), it says \"we assume anisotropy of the SGD noise within this n-dimensional space is not relevant\". Why it is not relevant? The nonzero eigenvalues of Hessian can vary a lot, giving quite different noise strength along different directions. Could the authors provide justifications (either theoretically or numerically), that the behavior is not changed by considering an isotropic noise here. ", "summary_of_the_paper": "This paper studies the behavior of SGD around the minimum. Unlike many other works that simply treat SGD noise as a fixed noise, the authors characterizes the location-dependence of SGD noise, which gives drastically different escaping behavior. By some simplification of the noise covariance matrix, the authors are able to formulate an SGD with isotropic noise with a time change. The new SDE is a dynamics on the log-loss landscape with a simple additive noise. With this result, the escape rate and stationary distributions are derived, depending polynomially on the loss, instead of exponentially. Numerical experiments are conducted to justify the assumptions made in the analysis, and verify the theoretical conclusions in the case of linear regression.", "summary_of_the_review": "This paper gives more realistic mathematical characterization of SGD's behavior around minima, by studying the location-dependence of SGD noise. The noise covariance is simplified in a smart way (with numerical justification) so that SGD dynamics can be formulated as a simple diffusion process on a log-loss landscape. Then, escape rate and stationary distribution are derived depending on a power law, instead of an exponential law. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7_4xA_mto1D", "reviewer": "Reviewer_hfnj", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "The paper gives an interesting characterization of the covariance matrix of SGD (note the noise due to data subsampling is structured) and uses a SDE persective, with the corresponding covariance to dervie power-law escape rates of SGD. The paper provides an intuitive introduction to the theoretical understanding of the papers and good sectional clarity in the way main concepts are presented.\n\nHowever, the paper does not clearly define a lot of the variables used (such as the covariance matrix Sigma, which is a central quantity and deserves its own line) and much approximations are justified via empirical (and not mathematical) reasons. It's hard to understand what exactly is the novelty and main contribution of the paper, since many steps in the derivations are cited. It seems like the application of existing techniques on a restrictive class of functions (with mean-squared loss) may be somewhat interesting, but I'm unclear why (it's not clearly stated in the introduction).  Furthermore, the escape rate of SGD is with respect to a local minima, theta^*, so it is unclear how you define theta^* when it is not a global minima.\n\n", "summary_of_the_paper": "The paper characterizes the escape rate of SGD for the mean-square loss via the perspective of SDE. ", "summary_of_the_review": "Due to the lack of novelty and the use of empirical justifications/lack of rigorous definitions, I would recommend a significant rewrite.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "7-VpmJpo9Jt", "reviewer": "Reviewer_bWLZ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "Strengths:\n1. Overall, this paper is well written and easy to follow. The contribution of this paper is clearly presented.\n2. To the best of my knowledge, the escape rate analysis of the SGD under the logarithmized loss landscape is novel as it partially explains why SGD prefers ﬂat minima with low effective dimensions.\n\nWeakness:\n1. There is a lack of clear definition of \"escape rate\". While I can somewhat guess the meaning of this quantity, I would prefer this critical concept to the explicitly and rigorously defined. \n2. I also did not see a clear statement/reference of the escape rate (16). Maybe it is trivial to the authors, but I think for the broader audience, it would be very helpful if more reference/discuss could be included. I could not follow the discussion for the multidimensional case as well for the same reason.\n3. In section 3.2, when the change of time scale is performed, are we still analyzing the (continuous time limit of) SGD iterates? I could not understand the meaning of the random time change from $t$ to $\\tau$ as the loss $L$ is involved in (14). If we understand the variable $t$ as time, what is the variable $\\tau$?\n\nI think this paper is potentially important to the understanding of the bias of SGD, but I cannot fully comprehend the discussion in section 4.", "summary_of_the_paper": "This paper considers the rate at which the SGD iterations will escape the valley around a local minimum. Under some approximation assumptions, this paper shows that the SGD noise covariance is highly structured as it aligns with the Hessian at the local minimum in the immediate vicinity. By considering the Ito SDE with the approximate SGD noise covariance and through a random change of time scale, the the Langevin equation on the loss landscape is transformed to that on the logarithmic loss landscape, but with a simpler additive noise. Such a logarithmized loss landscape is then exploited to derive the escape rate of SGD from a local minimum.", "summary_of_the_review": "A potentially novel perspective for understanding the bias of SGD, but more discussion on the preliminary results are needed for non-expert readers to understand the derivation.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies rate at which SGD escapes local minima and provides a potential justification for the "flat minima" observation. 
Reviewers agree that the paper studies an interesting problem and provides a nice result. But it seems like that paper in it's current shape is not ready for publication at ICLR. Issue is that the paper's writing is not up to bar, and requires a fair bit of work. In particular, the paper doesn't define the key quantities formally, doesn't provide all the assumptions in one place and justify why they might be reasonable. Finally, it would be great if the final result about escape rate is provided clearly with a self contained theorem/lemma that define/describe most of the key quantities in the rate.
Decision: Reject