L OGARITHMIC LANDSCAPE AND POWER - LAW ESCAPE RATE OF SGD
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ∆ L = L ( θ s ) −
L ( θ ∗ ) between a minimum θ ∗ and a saddle θ s but by the logarithmized loss barrier height ∆ log L = log
[ L ( θ s ) /L
( θ ∗ )] .
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 ﬂat minima with low effective dimensions, which gives an insight into implicit biases of SGD.
I NTRODUCTION
Deep learning has achieved breakthroughs in various applications in artiﬁcial intelligence such as image classiﬁcation (Krizhevsky et al., 2012; LeCun et al., 2015), speech recognition (Hinton et al., 2012), natural language processing (Collobert & Weston, 2008), and natural sciences (Iten et al., 2020; Bapst et al., 2020; Seif et al., 2021).
Such unparalleled success of deep learning hinges crucially on stochastic gradient descent (SGD) or its variants as an efﬁcient training algorithm.
Although the loss landscape is highly nonconvex, the SGD often succeeds in ﬁnding a global minimum.
It has been argued that the SGD noise plays a key role in escaping from local minima (Jastrze¸bski et al., 2017; Wu et al., 2018; 2020; Zhu et al., 2019;
Meng et al., 2020; Xie et al., 2021;
Liu et al., 2021).
It has also been suggested that SGD has an implicit bias that is beneﬁcial for generalization.
That is, SGD may help the network to ﬁnd ﬂat minima , which are considered to imply good generalization (Keskar et al., 2017; Hoffer et al., 2017; Wu et al., 2018).
How and why does SGD help the network escape from bad local minima and ﬁnd ﬂat minima?
These questions have been addressed in several works, and it is now recognized that the SGD noise strength and structure importantly affect the efﬁciency 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 at temperature T has been thoroughly discussed (Kramers, 1940; Langer, 1969).
When the (free) energy barrier is given by ∆ E , the escape rate is proportional to e − ∆ E/T , 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 ∆ L = L ( θ s ) − L ( θ ∗ ) , where L ( θ ) denotes the loss function at the network parameters θ , θ ∗ stands for a local minimum of L ( θ ) , and θ s denotes a saddle point that separates θ ∗ from other minima.
If we assume that the SGD noise is uniform and isotropic, which is often assumed in machine-learning literature (Jastrze¸bski et al., 2017), the escape rate is proportional to e − ∆ L/D , where D denotes the SGD noise strength.
In this paper, we show that inhomogeneity of the SGD noise strength brings about drastic modiﬁcation for the mean-square loss.
It turns out that the escape rate is determined by the logarithmized loss barrier height ∆ log L = log L ( θ s ) − log L ( θ ∗ ) =
log[ L ( θ s ) /L
( θ ∗ )] .
In other words, the escape rate is determined by not the difference but the ratio of L ( θ s ) and L ( θ ∗ ) .
This result means that even if the loss barrier height ∆ L
Moreover, given the fact that the eigenvalue spectrum of the Hessian at a 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 ﬂat minima with a low effective dimension , where the effective dimension is deﬁned as the number of outliers (MacKay, 1992) and ﬂatness is measured by a typical magnitude of outlier eigenvalues (Keskar et al., 2017).
The previous theories (Jastrze¸bski et al., 2017; Wu et al., 2018; Zhu et al., 2019; Meng et al., 2020; Xie et al., 2021; Liu et al., 2021) have also successfully explained that SGD prefers ﬂat 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.
Main contributions : We obtain the following main results:
• We derive an equation for approximating the SGD noise in Eq.
Remarkably, the SGD noise strength in the mean-square loss is shown to be proportional to the loss function, which is experimentally veriﬁed in Sec.
A key ingredient in deriving Eq.
(6) is the decoupling approximation given in Eq.
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.
Although the original SDE (4) has a multiplicative noise, the transformed SDE (14) 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.
Remarkably, the escape rate depends on the ratio between L ( θ ∗ ) and L ( θ s ) .
In Sec.
5.3, we experimentally test the validity of this result for linear regressions.
• Our escape rate crucially depends on the ﬂatness and the effective dimension, which shows that SGD has implicit biases towards ﬂat minima with low effective dimension.
We also show in Eq.
(18) that a local minimum with an effective dimension n greater than a certain critical value n c 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., 2021; Liu et al., 2021; Meng et al., 2020; Wojtowytsch, 2021).
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 beneﬁcial for escape from sharp minima (Zhu et al., 2019; Xie et al., 2021; Liu et al., 2021).
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, 2021), the SGD dynamics under a statedependent noise is discussed.
However, in these previous works, the connection between the noise strength and the loss function was not theoretically established, and the logarithmized loss landscape was not discussed.
The instability due to large effective dimensions was also not shown.
Another recent work (Pesme et al., 2021) observed that the noise is proportional to the loss for speciﬁc simple models.
In our paper, such a result is derived for more generic models.
G¨urb¨uzbalaban et al.
(2021) showed that SGD will converge to a heavy-tailed stationary distribution due to a multiplicative nature of the SGD noise in a simple linear regression problem.
Our paper strengthens this result: we argue that such a heavy-tailed distribution generically appears for the mean-square loss.
B ACKGROUND 2. S ETUP
We consider supervised learning.
Let D =
( x ( µ ) , y ( µ ) ) : µ = 1 , 2 , . . .
, N } be the training dataset, where x ( µ ) ∈ R
d denotes a data vector and y ( µ ) ∈ R be its label.
The network output for a given input x is denoted by f ( θ, x ) ∈ R , where θ ∈ R P stands for a set of trainable parameters with P being the number of trainable parameters (Extension to the multi-dimensional label and output is straightforward).
In this work, we focus on the mean-square loss
[Equation]
The training proceeds through optimization of L ( θ ) .
In most machine-learning applications, the optimization is done via SGD or its variants.
In SGD, the parameter θ k
[Equation]
1 at the time step k
+ 1 is determined by
[Equation]
where η > 0 is the learning rate, B k ⊂ { 1 , 2 , . .
. , N } with | B k | =
B is a mini-batch used at the k th time step, and L B k denotes the mini-batch loss.
Since the training dataset D is randomly divided into mini-batches, the dynamics deﬁned by Eq.
(2) is stochastic.
When B = N , the full training data samples are 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 ξ k = − [ ∇ L B k ( θ k ) − ∇ L ( θ k )] , Eq.
(2) is rewritten as
[Equation]
Obviously, (cid:104) ξ k (cid:105) = 0 , where the brackets denote the average over possible choices of mini-batches.
The noise covariance matrix is deﬁned as Σ( θ k ) :=
(cid:104) ξ k ξ T k (cid:105) .
The covariance structure of the SGD noise is important in analyzing the SGD dynamics, which will be discussed in Sec. 3.1.
2. S TOCHASTIC DIFFERENTIAL EQUATION FOR SGD
When the parameter update for each iteration is small, which is typically the case when the learning rate η is small enough, we can consider the continuous-time approximation (Li et al., 2017; Smith & Le, 2018).
By introducing a continuous time variable t ∈ R and regarding η as an inﬁnitesimal time step dt , we have a SDE
[Equation]
where dW t ∼ N (0 , I P dt ) with I n being the n -by- n identity matrix, and the multiplicative noise (cid:112) η Σ( θ t ) · dW t is interpreted as Itˆo since the noise ξ k in Eq.
(3) depends on θ k
but not on θ k +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, in which the isotropic and uniform Gaussian noise is injected into the GD as
[Equation]
where D > 0 corresponds to the noise strength (it is also called the diffusion coefﬁcient)
(Sato & Nakagawa, 2014; Zhang et al., 2017b; Zhu et al., 2019).
The stationary probability distribution
P GLD ( θ ) of θ for GLD is given by the Gibbs distribution P GLD ( θ ) ∝ e − L ( θ ) /D .
We will see in
Sec. 4 that the SGD noise structure, which is characterized by Σ( θ ) , drastically alters the stationary distribution and the escape rate from a local minimum.
T HEORETICAL FORMULATION 3. S TRUCTURE OF THE SGD NOISE COVARIANCE
The SGD noise covariance matrix Σ( θ ) signiﬁcantly affects the dynamics (Jastrze¸bski et al., 2017;
Smith & Le, 2018; Zhu et al., 2019; Ziyin et al., 2021).
In this section, under some approximations, we derive the following expression of Σ( θ ) for the mean-square loss within the the basin of attractions (or the “valley”) of a local minimum θ ∗ :
[Equation]
where H ( θ )
= ∇ 2 L ( θ ) is the Hessian.
It should be noted that ∇ L ( θ ∗ ) =
and H ( θ ∗ ) is positive semideﬁnite at any local minima θ ∗ .
We give a derivation below and the list of the approximations and their justiﬁcations in Appendix A.
In particular, the decoupling approximation is a new tool and plays a key role in the derivation.
It states that the quantities (cid:96) µ and C ( µ ) f ( θ ) :=
[Equation] f [Equation]
θ, x [Equation] µ ) ) ∇ f ( θ, x ( µ ) )
T are uncorrelated, which implies
[Equation]
This approximation is promising for large networks in which ∇ f looks a 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 (Jastrze¸bski et al., 2017; Zhu et al., 2019; Xie et al., 2021; Liu et al., 2021).
If the loss landscape has ﬂat directions, which correspond to the directions of the Hessian eigenvectors belonging to vanishingly small eigenvalues, the SGD noise does not work along those directions.
Consequently, SGD dynamics is frozen along those ﬂat directions, which effectively reduces the dimension of the parameter space explored by SGD dynamics.
This plays an important role in the escape efﬁciency.
Indeed, we will see that the escape rate crucially depends on the effective dimension of a given local minimum.
Second, the noise is proportional to the loss function, which is indeed experimentally conﬁrmed in Sec.
This property has not been pointed out and not been taken into account in previous studies (Jastrze¸bski et al., 2017; Zhu et al., 2019; Xie et al., 2021; Liu et 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 in Sec. 3.2.
This new formalism yields the power-law escape rate, i.e. Eqs.
(15) and (16), and the importance of the effective dimension of local minima for their stability.
Derivation of Eq.
We start from an analytic expression of Σ( θ ) , which reads
[Equation]
where B (cid:28) N is assumed in the second equality.
The derivation of Eq.
(8) is found in Jastrze¸bski et al.
(2017); Smith & Le (2018).
Usually, the gradient noise variance dominates the square of the gradient noise mean, and hence the term ∇ L ∇ L T in Eq.
(8) is negligible.
For the mean-square loss, we have ∇ (cid:96) µ =
[ f ( θ, x ( µ ) )
− y ( µ ) ] ∇ f ( θ, x ( µ ) ) , and hence Σ( θ ) ≈
[Equation]
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 is written as
[Equation]
It is shown by Papyan (2018) that the last term of Eq.
(10) does not contribute to outliers (i.e. large eigenvalues) of the Hessian.
Dynamics near a local minimum is governed by outliers, and hence we can ignore this term.
At θ = θ ∗ , we therefore obtain
[Equation]
[Equation]
Let us assume C ( µ )
f ( θ ) ≈ C ( µ )
f
( θ ∗ ) for θ within the valley of a local minimum θ ∗ .
We then obtain the desired expression (6) by substituting it into Eq.
3. L OGARITHMIZED LOSS LANDSCAPE
Let us consider the Itˆo SDE (4) with the SGD noise covariance (6) near a local minimum θ ∗ , which is written as
[Equation]
Let us consider a stochastic time t ( τ ) for τ ≥ 0 as
[Equation]
and perform a random time change from t to τ (Øksendal, 1998).
Correspondingly, we introduce the Wiener process d ˜ W τ ∼ N (0 , I P dτ ) .
Since dτ = L ( θ t ) dt , we have d ˜ W
τ
=
(cid:112) L ( θ t ) · dW t .
In terms of the notation ˜ θ τ = θ t , Eq.
(12) is expressed as
[Equation]
We should note that at a global minimum with L ( θ ) = 0 , which is realized in an overparameterized regime (Zhang et al., 2017a), the random time change through Eq.
(13) is ill-deﬁned since τ is frozen at a ﬁnite value once the model reaches a global minimum.
We can overcome this difﬁculty by adding an inﬁnitesimal constant (cid:15) > 0 to the loss, which makes the loss function positive without changing the ﬁnite-time dynamics like the escape from a local minimum θ ∗ with L ( θ ∗ ) > 0 .
In this way, the Langevin equation on the loss landscape L ( θ ) with multiplicative noise is transformed to that on the logarithmic loss landscape U ( θ ) =
log L ( θ ) with simpler additive noise.
This formulation indicates the importance of considering the logarithmized loss landscape U ( θ )
=
log L ( θ ) .
In the following, we use Eq.
(14) to discuss the escape efﬁciency from local minima.
E SCAPE RATE FROM LOCAL MINIMA
Let us consider a local minimum θ ∗ and its basin of attraction
A θ ∗ , which is the set of all the starting points θ 0 such that θ t tends to θ ∗ as t → ∞ if there is no noise.
In order to escape from the basin of θ ∗ , θ t must go through a saddle θ s with the help of noise.
The escape time is deﬁned as τ = inf { t > 0 : θ t / ∈
A θ ∗ , θ
= θ ∗ } .
The escape rate κ is then deﬁned as the inverse of the mean escape time: κ = (cid:104) τ (cid:105) − 1 .
Our idea for evaluating the escape rate is applying the Kramers formula (Kramers, 1940) or its high-dimensional generalization (Langer, 1969; Bovier et al., 2004;
Berglund, 2013) to Eq.
It should be noted that the Kramers formula is applicable only for additive noise, and hence it cannot be used for the original SDE (4) with multiplicative noise.
First, we present main results [Eqs.
(15), (16), and (17)], and give their derivations later.
For P =(i.e. the single-variable case θ ∈ R ), the escape rate through a saddle θ s is given by
[Equation]
(15) where h ∗ =
H ( θ ∗ ) .
This is a variant of the Kramers formula (Kramers, 1940), which is accurate when κ is small.
We can derive Eq.
(15) from Eq.
(14) without any further approximations.
For θ ∈ R P with P > 1 , the analysis is more involved and needs some assumptions and approximations, which are listed in Appendix A.
We assume that the eigenvalue spectrum of the Hessian
H ( θ ∗ ) at the local minimum θ ∗ 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 ﬂatness) and the number of outlier eigenvalues (i.e. the effective dimension), respectively.
As we discussed in Sec. 3.1, the SGD dynamics is frozen along ﬂat directions.
Therefore SGD dynamics around the local minimum θ ∗ is restricted to the n -dimensional manifold spanned by the outlier eigenvectors v 1 , v 2 , . .
, v n ∈ R P of H ( θ ∗ ) .
Now we parameterize θ by using n parameters z 1 , z 2 , . . .
, z n ∈ R as θ = θ ∗ + (cid:80) ni
=
z i v i .
The Hessian restricted to this outlier subspace is then written as ˆ H ( θ ) :
= ∇ 2 z L ( θ ) ∈ R n × n .
We then obtain the following escape rate formula for P > 1 :
[Equation]
where h se is the negative eigenvalue of H ( θ s ) corresponding to the escape direction.
Again, Eq.
(16) is valid when κ is small enough.
It should be noted that Eq.
(16) is reduced to Eq.
(15) when n = 1 .
For any P , the quasi-stationary distribution P ss ( θ ) within the valley of a local minimum θ ∗ , which can be identiﬁed as the stationary distribution restricted to the valley (Bianchi & Gaudilli`ere, 2016), is written as
[Equation]
Remarkably, it depends on L ( θ ) polynomially rather than exponentially as in the standard
GLD (Jastrze¸bski et al., 2017; Sato & Nakagawa, 2014; Zhang et al., 2017b).
This polynomial dependence on L ( θ ) is a key feature leading to the escape rate formula.
From Eq.
(16) we can obtain some implications.
The factor [ L ( θ s ) /L
( θ ∗ )] − ( Bηh ∗ +1 − n 2 ) increases with h ∗ and n , which indicates that sharp minima (i.e. minima with large h ∗ ) or minima with large n are unstable.
This fact explains why SGD ﬁnds ﬂat minima with a low effective dimension n .
Equation (16) also implies that the effective dimension of any stable minima must satisfy
[Equation]
The instability due to a 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 logarithmizedloss Hessian: det ∇ 2 z U ( θ ∗ ) / | det ∇ 2 z U ( θ s ) | = [det ˆ H ( θ ∗ ) / | det ˆ H ( θ s ) | ] · [ L ( θ s ) /L
( θ ∗ )] n/ 2 .
Derivation of Eq.
(16) and Eq.
Now we give a derivation of Eqs.
(16) and (17).
Equation (15) is straightforwardly obtained by putting P = 1 in the derivation below (some approximations are introduced below, but all of them are not necessary for P = 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: (cid:112)
[Equation]
, where h ∗ ∈ R + characterizes the magnitude of the Hessian outliers.
This assumption is justiﬁed 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 justiﬁed 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 a Hessian eigenvector v e , where e ∈ { 1 , 2 , . . . , n } identiﬁes the escape direction, and h ∗ corresponds to the Hessian eigenvalue at the minimum along the escape direction, i.e. h ∗ =
h ∗ e .
See Appendix C for the details.
Under this approximation, Eq.
(14) becomes dz τ =
−∇ z Udτ + √ 2 Td ˜ W τ , where T = ηh ∗ /B .
Its quasi-stationary distribution ˜ P ss ( θ ) within the valley including the local minimum θ = θ ∗ (i.e. z = 0 ) is then given by a Gibbs distribution with respect to U ( θ ) :
˜ P ss ( θ )
∝ e − U ( θ ) /T
∝ L ( θ ) − B/ ( ηh ∗ ) .
This is the distribution function of ˜ θ τ for a ﬁxed τ .
However, what we want is the quasi-stationary distribution of θ t for a ﬁxed t .
In Appendix B, by using Eq.
(13), we show that the two distributions are related with each other as P ss ( θ ) ∝ L ( θ ) − 1 ˜ P ss ( θ ) .
We thus obtain Eq.
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 κ τ per unit τ is given by
[Equation]
where u se is the negative eigenvalue of ∇ 2 z U ( θ s ) corresponding to the escape direction, ∆ U =
U ( θ s ) − U ( θ ∗ ) is called the potential barrier, and ∆ U/T is assumed to be large enough.
What we really want is the escape rate per unit time t .
It is a reasonable assumption that θ t stays close to θ ∗ for most times before escape, and hence τ is approximately given by τ (cid:39) L ( θ ∗ ) t .
The escape rate κ per unit t is then given by κ =
L ( θ ∗ ) κ τ .
By combining it with Eq.
(19), and substituting
U ( θ ) = log L ( θ ) , we ﬁnally obtain Eq.
E XPERIMENTS
Our key theoretical observation is that the SGD noise strength is proportional to the loss function, which is obtained as a result of the decoupling approximation.
This property leads us to the Langevin
[Figure] [Figure]
equation (14) with the logarithmized loss gradient and an additive noise through a 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 is valid during entire training dynamics.
In Sec.
5.2, we measure the SGD noise strength and conﬁrm that it is indeed proportional to the loss function near a minimum.
In Sec.
5.3, we experimentally test the validity of Eq.
(17) for the stationary distribution and Eq.
(16) for the escape rate.
We will see that numerical results for a linear regression and for a non-linear neural network agree with our theoretical results.
5. E XPERIMENTAL VERIFICATION OF THE DECOUPLING APPROXIMATION
Let us compare the eigenvalue distribution of the exact matrix (1 /N )
(cid:80) Nµ =1 (cid:96) ( µ ) C ( µ ) f with that of the decoupled one L ( θ ) · (1 /N )
(cid:80) Nµ =1 C ( µ ) f with C ( µ ) f
= ∇ f ( θ, x ( µ ) ) ∇ f
( θ, x ( µ ) ) T .
We consider a binary classiﬁcation problem using the ﬁrst 10 4 samples of the MNIST dataset such that we classify each image into even (its label is y = +1 ) or odd number (its label is y = − 1 ).
The network has two hidden layers, each of which has 100 units and the ReLU 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 we ﬁx η = 0 .
01 and B = 100 .
Figure 1 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. M EASUREMENTS OF THE SGD NOISE STRENGTH
As a measure of the SGD noise strength, let us consider the norm of the noise vector ξ given by (cid:104) ξ T ξ (cid:105) =
Tr Σ ≈ N /B , where N := (1 /N )
(cid:80) Nµ =1 ∇ (cid:96) T µ ∇ (cid:96) µ − ∇ L T ∇ L .
Here we present experimental results for two architectures and datasets.
First, we consider training of the Fashion-MNIST dataset by using a fully connected neural network with three hidden layers, each of which has 2 × 10units and the ReLU activation, followed by the output layer of 10 units with no activation (classiﬁca-
[Figure]
e xpon e n t φ
B η =0.01 η =0.1 η =        4 5 6m ea n f i r s t p a ss a g e ti m e c d
[Figure]
1 d
[Figure]
5 d
[Figure] [Figure]
[Figure]
i
[Figure]
tion labels are given in the one-hot representation).
Second, we consider training of the CIFARdataset by using a convolutional neural network.
Following Keskar et al.
(2017), let us denote a stack of n convolutional layers of a ﬁlters and a kernel size of b × c with the stride length of d by n × [ a, b, c, d ] .
We use the conﬁguration: 3 × [64 , 3 , 3 , 1] , 3 × [128 , 3 , 3 , 1] , 3 × [256 , 3 , 3 , 1] , where a 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-batch size B =and η = 0 .
1 for the mean-square loss.
During the training, we measure the training loss and the noise strength N for every epoch.
Numerical results are given in Fig.
We see that roughly N ∝ L at a later stage of the training, which agrees with our theoretical prediction.
Although N is not proportional to L at an early stage of training, it does not 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 ≈ 2 L ( θ ) × (1 /N )
(cid:80) Nµ =1 ∇ f ( θ, x ( µ ) )
T ∇ f
( θ, x ( µ ) ) , but the quantity (1 /N ) (cid:80) Nµ =1 ∇ f ( θ, x ( µ ) )
T ∇ f
( θ, x ( µ ) ) also changes during training.
Although Eq.
(9) is derived for the mean-square loss, the relation N ∝ L ( θ ) holds in more general loss functions; see Appendix D for general argument and experiments on the cross entropy loss.
5. E XPERIMENTAL TEST OF STATIONARY DISTRIBUTION AND ESCAPE RATE FORMULA
We experimentally verify our theoretical predictions.
Below, we ﬁrst present numerical results for a simple linear regression problem, and then for a nonlinear model, i.e., a neural network.
Let us start from the following linear regression problem: each entry of x ( µ ) ∈ R d and its label y ( µ ) ∈ R are i.i.d.
Gaussian random variables of zero mean and unit variance.
We focus on the case of d (cid:28) N .
The output for an input x is given by f ( θ, x ) =
θ T x , where θ ∈ R
d is the trainable network parameter.
We optimize θ via SGD.
The mean-square loss
L ( θ ) =
(1 / 2 N ) (cid:80) Nµ =1 (cid:0) θx ( µ ) − y ( µ ) (cid:1) 2 is quadratic and has a unique minimum at θ ≈ 0 .
First, we test Eq.
(17), i.e. the stationary distribution, for d = 1 and N = 10 5 .
We sampled the value of θ k at every 100 iterations ( k = j × 100 , j = 1 , 2 , . . .
, 10 4 ) and made a histogram.
We then ﬁt the histogram to the form P ss ( θ ) ∝ L ( θ ) − φ and determine the exponent φ .
Our theory predicts φ = 1 + B/ ( ηh ∗ ) .
Numerical results for the exponent φ are presented in Fig.
3 (a) against B for three ﬁxed learning rates η .
In the same ﬁgure, theoretical values of φ are plotted in dashed lines.
The agreement between theory and experiment is fairly well.
For a large learning rate η = 1 , the exponent slightly deviates from its theoretical value.
This is due to the effect of a ﬁnite learning rate (recall that η is assumed to be small in deriving the continuous-time stochastic differential equation).
Next, we test our formula on the escape rate, Eq.
Although the mean-square loss is quadratic and no barrier crossing occurs, we can measure the ﬁrst passage time, which imitates the escape time for a non-convex loss landscape.
Let us ﬁx a threshold value of the loss function.
The ﬁrst passage time t p is deﬁned as the shortest time at which the loss exceeds the threshold value.
Here, time t is identiﬁed as ηk , where k denotes the number of iterations in discrete SGD (2).
We identify the threshold value as L ( θ s ) , i.e., the loss at the saddle in the escape problem.
It is expected that t p is similar to the escape time and proportional to κ − 1 .
The Hessian H = (1 /N )
(cid:80) Nµ =1 x ( µ ) x ( µ )T has d nonzero eigenvalues, all of which are close to unity.
We can therefore identify h ∗ = 1 and n = d .
The mean ﬁrst passage time over 100 independent runs is measured for varying threshold values which are speciﬁed by c =
L ( θ s ) /L
( θ ∗ ) > 1 .
Experimental results for N = 10 4 are presented in Fig.
3 (b).
Dashed straight lines have slope B/ ( ηh ∗ ) + 1 − n/ 2 .
Experiments show that the ﬁrst passage time behaves as t p ∝
[ L ( θ s ) /L
( θ ∗ )]
B/ ( ηh ∗ )+1 − n/ 2 , which agrees with our theoretical evaluation of κ − 1 [see Eq.
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 et al., 2021; Meng et al., 2020).
Our escape-rate formula (16) is also veriﬁed in a non-linear model.
As in Sec. 5.1, we consider the binary classiﬁcation problem using the ﬁrst 10 4 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 a single unit with no activation.
We always use the mean-square loss.
Starting from thee Glorot initialization, the network is pre-trained via SGD with η = 0 .and
B = 100 for 10 5 iterations.
We ﬁnd that after pre-training, the loss becomes almost stationary around at L ( θ ) ≈ 0 .
We regard that the pre-trained network is near a local minimum.
We then further train the pre-trained network via SGD with a new choice of η and B (here we ﬁx B = 10 ), and measure the ﬁrst passage time t p .
The mean ﬁrst-passage time over 100 independent runs is measured for varying threshold values which are speciﬁed by c =
L ( θ s ) /L
( θ ∗ ) > 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 (16), we also compare the power-law exponent for the mean ﬁrst-passage time with our theoretical prediction B/ ( ηh ∗ )
+
1 − n/ 2 .
Here, h ∗ and n are estimated by the Hessian eigenvalues.
In Appendix E, we present a numerical result for eigenvalues of the approximate Hessian given by the right-hand side of Eq.
By identifying the largest eigenvalue of the Hessian as h ∗ , we have h ∗ ≈ 95 .
On the other hand, it is difﬁcult 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.
3 (c) correspond to our theoretical prediction κ − 1 ∝ c B/ ( ηh ∗ )
+1 − n/ 2 with h ∗ = 95 . 6 and n = 9 .
This excellent agreement shows that our theoretical formula (16) is also valid in non-linear models.
C ONCLUSION
In this work, we have investigated SGD dynamics via a Langevin approach.
With several approximations 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 differential equation (14) with additive noise near a minimum via a 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 ( θ )
=
log L ( θ ) .
This stochastic differential equation has a quasi-stationary distribution that decays polynomially with L ( θ ) near a minimum (17), not exponentially as in the usual Gibbs distribution.
This new formalism yields the power-law escape rate formula (16) whose exponent depends on η , B , h ∗ , and n .
Our escape-rate formula explains an empirical fact that SGD favors ﬂat minima with low effective dimensions.
The effective dimension of a minimum must satisfy Eq.
(18) 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 non-Gaussianity can also play an important role.
For example, S¸ims¸ekli et al.
(2019) approximated SGD as a L´evy-driven SDE, which explains why SGD ﬁnds wide minima.
It would be an interesting future problem to take the non-Gaussian effect into account.
A L IST OF APPROXIMATIONS AND THEIR JUSTIFICATIONS
In Sec. 3.1, we made several approximations to derive the result (6).
For clarity, we list the approximations made and their justiﬁcations below.
• In Eq.
(8), we make the approximation of B << N , which simpliﬁes the expression but is not essential.
The main conclusion is not affected by this approximation.
• We ignore ∇ L ∇ L T in Eq.
(8), which is justiﬁed near a local minimum.
• In Eq.
(9), we make the decoupling approximation (7), which is one of the key heuristic approximations in our work.
This approximation is veriﬁed experimentally in Sec. 5.1.
• We ignore the last term of Eq.
(10), which is a common approximation (Sagun et al., 2017).
This approximation is justiﬁed if we are only interested in the outliers of the Hessian eigenvalues.
Indeed, outliers of the Hessian eigenvalues, which play dominant roles in escape from local minima, are known to be attributable to the ﬁrst term of the right-hand side of Eq.
(10) (Papyan, 2018).
• In Eq.
(11), we assume that the matrix (1 /N ) (cid:80)
Nµ =1 ∇ f
( θ, x ( µ ) ) ∇ f ( θ, x ( µ ) )
T does not change so much in a valley with a given local minimum.
This approximation is indirectly veriﬁed in Fig.
2 (the proportionality between the loss and the noise strength implies this matrix is actually constant).
In deriving the escape rate for multi-variable cases in Sec. 4, we further make the following assumptions 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 justiﬁed because the SGD dynamics is almost frozen along ﬂat directions of the loss landscape.
• We approximate the SGD noise as an isotropic one within the n -dimensional outlier subspace.
This approximation is justiﬁed 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 ( θ ∗ ) t , which is justiﬁed because the model stays near a local minimum until it escapes from that minimum.
B S TATIONARY DISTRIBUTION
Since ˜ θ
τ obeys a simple Langevin equation
[Equation]
[Equation]
the stationary distribution of ˜ θ τ is given by the Gibbs distribution
˜ P ss ( θ )
∝ e − U ( θ ) /T .
On the other hand, what we want is the stationary distribution
P ss ( θ ) of θ t , where θ t = ˜
θ τ with τ = (cid:82) t 0 dt (cid:48)
L ( θ t (cid:48) ) .
In this section, we show the relation between the two distributions: P ss ( θ ) ∝ L ( θ ) − 1 ˜ P ss ( θ ) .
We express the stationary distributions in terms of the long-time average of the delta function:
[Equation]
[Equation]
By using the relation τ = (cid:82) t 0 dt (cid:48)
L ( θ t (cid:48) )
, we have dτ =
L ( θ t ) dt .
For a sufﬁciently large t , we also obtain τ ∼ t ¯ L , where ¯ L :=
lim s →∞ (1 /s ) (cid:82) s 0 dt (cid:48)
L ( θ t (cid:48) ) denotes the long-time average of L ( θ t ) .
By using them, P ss ( θ ) is rewritten as
[Equation]
We thus obtain the desired relation, P ss ( θ ) ∝ L ( θ ) − 1 ˜ P ss ( θ ) .
C J USTIFICATION OF THE ISOTROPIC - NOISE APPROXIMATION WITHIN THE n - DIMENSIONAL SUBSPACE
We now show that the isotropic-noise approximation is valid when the directions of the eigenvectors of the Hessian do not change within the valley of a given local minimum.
In this case, ∂ 2 U/∂z i ∂z j = 0 for any i (cid:54) =
j , where the Hessian at the minimum is given by H ( θ ∗ )
≈ (cid:80) ni
=1 h ∗ i v i v T
i ( h i is an eigenvalue and v i is the corresponding eigenvector of H ( θ ∗ ) ) and the displacement vector z ∈ R n is deﬁned by θ =
θ ∗ + (cid:80) ni
=
z i v i .
The stochastic differential equation
[Equation]
˜ W τ is then equivalent to the following Fokker-Planck equation for the distribution function P ( z, τ ) of z at τ :
[Equation]
Let us assume that the direction of e th eigenvector v e corresponds to the escape direction.
We denote by z ⊥ ∈ R n − 1 the displacement perpendicular to the escape direction, and write z =
( z e , z ⊥ ) .
At the saddle z s , ∇ z U ( z s ) = 0 , h se < 0 and h si > 0 for all i (cid:54) =
e , where { h si } is the set of eigenvectors of the Hessian at z s .
Under the above setting, we derive the escape rate formula following Kramers (1940).
The steady current J ∈ R n is aligned to the escape direction, and hence J e (cid:54) = 0 and J ⊥ = 0 .
Let us denote by
P ∗ the total probability within the valley of a given minimum θ ∗ and by J the total current ﬂowing to outside of the valley through the saddle θ s = θ ∗ + z s .
It is assumed that J is small and P ∗ is almost stationary.
The escape rate κ τ is then given by
[Equation] [Equation]
The escape rate κ per unit time is given by
[Equation]
We now evaluate P ∗ and J .
To evaluate P ∗ , it is necessary to know about the stationary distribution near the minimum θ ∗ .
Near the minimum θ ∗ (i.e. small z ),
[Equation] [Equation]
and
[Equation]
(27) where H ∗ =
H ( θ ∗ ) .
By substituting it into the Fokker-Planck equation (23), we obtain
[Equation]
i
[Equation]
∂ ∂z
i
[Equation]

[Equation]
(cid:88)
[Equation]
η
B
[Equation] [Equation]
[Figure]
If for all j
[Equation] [Equation]
P ss ( z ) is a stationary distribution near the minimum θ ∗ .
Indeed,
[Equation]
satisﬁes this condition.
It should be noted that the stationary distribution is independent of the
Hessian eigenvalues { h ∗ i } .
By using Eq.
(30), P ∗ is evaluated as
[Equation]
Next, let us evaluate J .
The probability current J along the escape direction e is given by
[Equation]
At the saddle, the current vector J ⊥ perpendicular to the escape direction is zero, and hence the probability distribution near the saddle with z e = z se and z ⊥ (cid:54) = 0 is also evaluated in a similar way as in Eq.
[Equation]
[Equation]
By substituting it into Eq.
(32), we obtain
[Equation]
(34) where the current along the escape path ( z ⊥ = 0 ) is denoted by J path :
=
J ( z se , z ⊥ = 0) .
The total current through the saddle is then evaluated as
[Equation]
When the distribution function is almost stationary, Eq.
(23) yields ∂J
e ( z e , z ⊥ = 0) /∂z
e ≈ 0 , and hence the current along the escape path is approximately constant J e ( z e , z ⊥ = 0)
≈ J path .
Since
J e ( z e , z ⊥ = 0) is given by Eq.
(32) by putting z ⊥ = 0 , we have
[Equation]
By multiplying e Bηh ∗ e U in both sides and integrating over z e from 0 to z ce , where z ce deﬁned as
U ( θ
[Equation]
+ z ce v e ) =
U ( θ ∗ ) (see Fig. 4), we obtain
[Equation]
where it is assumed that the probability at z ce is small and negligible, P ( z e =
z ce , z ⊥ = 0)
≈ 0 .
By using the saddle-point method, the integral in the left-hand side of Eq.
(37) is evaluated as
[Equation]
By substituting this result in Eq.
(37), we obtain
[Equation]
where ∆ U =
U ( θ s ) − U ( θ ∗ ) .
The total current J in Eq.
(35) is then expressed as
[Equation]
By using Eqs.
(31) and (40)
[Equation]
the escape rate κ in Eq.
(25) is evaluated as
[Equation]
Since ∆ U = log[ L ( θ s ) /L
( θ ∗ )] , we ﬁnally obtain κ
[Equation]
which is exactly identical to the escape rate formula derived in the main text using the isotropic-noise approximation with h ∗
= h ∗ e .
In this way, the isotropic-noise approximation is justiﬁed even when the loss landscape is not isotropic near the minimum.
We have assumed that the directions of the eigenvectors of the Hessian do not change within the valley of a given local minimum.
Under this assumption, the Fokker-Planck equation along the escape path is decoupled from the perpendicular directions z ⊥ .
Moreover, we have seen that the stationary distribution near a minimum is independent of the Hessian eigenvalues.
These properties explain the reason why the isotropic-noise approximation is justiﬁed in this case.
D O THER LOSS FUNCTIONS
In our paper, we mainly focus on the mean-square loss, for which we can analytically derive the relation between the loss L ( θ ) and the SGD noise covariance Σ( θ ) .
An important observation is that the SGD noise strength N is proportional to the loss, i.e., N ∝ L ( θ ) (see Sec.
5.2 of the main text for the deﬁnition of N ).
Here, we argue that the relation N ∝ L ( θ ) also holds in more general situations.
During the training, the value of (cid:96) µ will ﬂuctuate from sample to sample.
At a certain time step of SGD, let us suppose that N − M samples in the training dataset are already classiﬁed correctly and hence (cid:96) µ ∼ 0 , whereas the other M samples are not and hence (cid:96) µ ∼ 1 .
The loss function is then given by L ( θ ) =
(1 /N )
(cid:80) Nµ =1 (cid:96) µ ∼ M/N .
When (cid:96) µ is small, ∇ (cid:96) µ will also be small.
Therefore, for N − M samples with (cid:96) µ ∼ 0 , ∇ (cid:96) µ ∼ 0 also holds.
The other M samples will have nonsmall gradients: (cid:107)∇ (cid:96) µ (cid:107) 2 ∼ g , where g > 0 is a certain positive constant.
We thus estimate N as
N ≈ (1 /N )
(cid:80) N µ =1 ∇ (cid:96) T µ ∇ (cid:96) µ ∼ gM/N ∼ gL ( θ ) .
In this way, N ∝ L ( θ ) will hold, irrespective of the loss function.
However, we emphasize that this is a crude argument.
In particular, the above argument will not hold near a global minimum because all the samples are correctly classiﬁed there, which implies that (cid:96) µ is small for all µ , in contrast to the above argument relying on the existence of M wrongly classiﬁed samples with (cid:96) µ ∼ 1 and (cid:107)∇ (cid:96) µ (cid:107) 2 ∼ g .
[Figure]
We now experimentally test the relation N ∝ L ( θ ) 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 a convolutional neural network trained by CIFAR-10 (see Sec.
for the detail).
We ﬁx B = 100 in both cases, and η = 0 .
1 for the fully connected network and η = 0 .
for the convolutional network.
Experimental results are presented in Fig.
We ﬁnd that the relation N ∝ L ( θ ) seems to hold true at an intermediate stage of the training dynamics, although the proportionality is less clear compared with Fig.
2 in the main text for the mean-square loss.
We also ﬁnd that for sufﬁciently small values of the loss, N ∝ L ( θ ) 2 [see Fig.
5 (c) and (d)], whose implications should merit further investigation in future studies.
E H ESSIAN EIGENVALUES FOR A NEURAL NETWORK IN S EC . 5.
We present numerical results on Hessian eigenvalues in a pre-trained neural network studied in
Sec.
Instead of the exact Hessian, we consider an approximate Hessian given on the right-hand side of Eq.
(11), i.e.,
[Equation] [Equation]
Eigenvalues { h i } are arranged in descending order as h 1 ≥ h 2 ≥ · · · ≥ h P (in our model P = 7861 ).
A histogram of the Hessian eigenvalues is presented in Fig.
6 (a).
We see that most eigenvalues are close to zero, which corresponds to the bulk, but there are some large eigenvalues, which correspond to the outliers.
The largest eigenvalue is λ 1 = 95 . 6 , which is identiﬁed as h ∗ in our theoretical formula (16).
[Figure]
Another important quantity is the effective dimension n corresnding to the number of outliers.
Since the outliers and the bulk are not sharply separated, it is difﬁcult to precisely determine n .
In Fig.(b), we plot h i up to i = 30 .
From this ﬁgure, it seems reasonable to estimate n ∼ 10 .
As a heuristic method of determining n , we can consider the following identiﬁcation: ﬁrst we deﬁne the weight p
i for i th eigenmode as p i
= h i / (cid:80)
Pj =1 h j .
We then determine n as
[Equation]
which gives n = 12 . 7 in our case.
In Sec.
5.3, our formula (16) with n = 9 explains numerical results on the ﬁrst-passage time.