Year: 2021

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 ∆L = L(θ) − L(θ∗) between a minimum θ∗ and a saddle θ but by the logarithmized loss barrier height ∆ logL = 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 flat minima with low effective dimensions, which gives an insight into implicit biases of SGD.

text: 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 ∆ logL = 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 flat minima with low effective dimensions, which gives an insight into implicit biases of SGD.

heading: 1 INTRODUCTION
text: Deep learning has achieved breakthroughs in various applications in artificial intelligence such as image classification (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 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 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 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., 2017; 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 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 modification for the mean-square loss. It turns out that the escape rate is determined by the logarithmized loss barrier height ∆ logL = logL(θs)−logL(θ∗) = 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 is the same, minima with smaller values of L(θ∗) are more stable.
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 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 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 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.
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 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. (16). 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 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 a local minimum with an effective dimension n greater than a certain critical value nc 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 beneficial 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 specific simple models. In our paper, such a result is derived for more generic models. Gürbü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 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.

heading: 2 BACKGROUND

heading: 2.1 SETUP
text: We consider supervised learning. Let D = {(x(µ), y(µ)) : µ = 1, 2, . . . , N} be the training dataset, where x(µ) ∈ Rd 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 θ ∈ RP 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
L(θ) = 1
2N N∑ µ=1 [ f(θ, x(µ))− y(µ) ]2 =: 1 N N∑ µ=1 `µ(θ). (1)
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+1 at the time step k + 1 is determined by
θk+1 = θk − η∇LBk(θk), LBk(θ) = 1
2B ∑ µ∈Bk `µ(θ), (2)
where η > 0 is the learning rate, Bk ⊂ {1, 2, . . . , N} with |Bk| = B is a mini-batch used at the kth time step, and LBk denotes the mini-batch loss. Since the training dataset D is randomly divided into mini-batches, the dynamics defined 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 = −[∇LBk(θk)−∇L(θk)], Eq. (2) is rewritten as
θk+1 = θk − η∇L(θk) + ηξk. (3) Obviously, 〈ξk〉 = 0, where the brackets denote the average over possible choices of mini-batches. The noise covariance matrix is defined as Σ(θk) := 〈ξkξTk 〉. The covariance structure of the SGD noise is important in analyzing the SGD dynamics, which will be discussed in Sec. 3.1.

heading: 2.2 STOCHASTIC DIFFERENTIAL EQUATION FOR SGD
text: 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 infinitesimal time step dt, we have a SDE
dθt = −∇L(θt)dt+ √ ηΣ(θt) · dWt, (4)
where dWt ∼ N (0, IP dt) with In being the n-by-n identity matrix, and the multiplicative noise√ ηΣ(θt) · dWt is interpreted as Itô 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
dθt = −∇L(θt)dt+ √
2DdWt, (5) where D > 0 corresponds to 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 PGLD(θ) of θ for GLD is given by the Gibbs distribution PGLD(θ) ∝ 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.

heading: 3 THEORETICAL FORMULATION

heading: 3.1 STRUCTURE OF THE SGD NOISE COVARIANCE
text: The SGD noise covariance matrix Σ(θ) significantly 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 θ∗:
Σ(θ) ≈ 2L(θ) NB H(θ∗), (6)
where H(θ) = ∇2L(θ) is the Hessian. It should be noted that ∇L(θ∗) = 0 and H(θ∗) is positive semidefinite at any local minima θ∗. We give a derivation below and the list of the approximations and their justifications 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 `µ and C
(µ) f (θ) := ∇f(θ, x(µ))∇f(θ, x(µ))T are uncorrelated, which implies
1
N N∑ µ=1 `µC (µ) f (θ) ≈
( 1
N N∑ µ=1 `µ
) · ( 1
N N∑ µ=1 C (µ) f
) = L(θ) 1
N N∑ µ=1 C (µ) f (θ). (7)
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 flat 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 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 a given local minimum.
Second, the noise is proportional to the loss function, which is indeed experimentally confirmed in Sec. 5.2. 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. (6): We start from an analytic expression of Σ(θ), which reads
Σ(θ) = 1
B N −B N − 1
( 1
N N∑ µ=1 ∇`µ∇`Tµ −∇L∇LT ) ' 1 B ( 1 N N∑ µ=1 ∇`µ∇`Tµ −∇L∇LT ) , (8)
where B 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∇LT in Eq. (8) is negligible.
For the mean-square loss, we have∇`µ = [f(θ, x(µ))− y(µ)]∇f(θ, x(µ)), and hence
Σ(θ) ≈ 2 BN N∑ µ=1 `µ∇f(θ, x(µ))∇f(θ, x(µ))T = 2 BN N∑ µ=1 `µC (µ) f (θ) ≈ 2L(θ) NB N∑ µ=1 C (µ) f (θ), (9)
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
H(θ) = ∇2L(θ) = 1 N N∑ µ=1 C (µ) f (θ) + 1 N N∑ µ=1 [ f(θ, x(µ))− y(µ) ] ∇2f(θ, x(µ)). (10)
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
H(θ∗) ≈ 1 N N∑ µ=1 C (µ) f (θ ∗). (11)
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. (9).

heading: 3.2 LOGARITHMIZED LOSS LANDSCAPE
text: Let us consider the Itô SDE (4) with the SGD noise covariance (6) near a local minimum θ∗, which is written as dθt = −∇L(θt)dt+ √ 2ηL(θt)
B H(θ∗) · dWt. (12)
Let us consider a stochastic time t(τ) for τ ≥ 0 as
τ = ∫ t(τ) 0 dt′ L(θt′), (13)
and perform a random time change from t to τ (Øksendal, 1998). Correspondingly, we introduce the Wiener process dW̃τ ∼ N (0, IP dτ). Since dτ = L(θt)dt, we have dW̃τ = √ L(θt) · dWt. In
terms of the notation θ̃τ = θt, Eq. (12) is expressed as
dθ̃τ = − 1
L(θ̃τ ) ∇L(θ̃τ )dτ +
√ 2ηH(θ∗)
B dW̃τ = −
[ ∇ logL(θ̃τ ) ] dτ + √ 2ηH(θ∗)
B dW̃τ . (14)
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-defined since τ is frozen at a finite value once the model reaches a 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 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(θ) = logL(θ) with simpler additive noise. This formulation indicates the importance of considering the logarithmized loss landscape U(θ) = logL(θ). In the following, we use Eq. (14) to discuss the escape efficiency from local minima.

heading: 4 ESCAPE RATE FROM LOCAL MINIMA
text: 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 defined as τ = inf{t > 0 : θt /∈ Aθ∗ , θ0 = θ∗}. The escape rate κ is then defined as the inverse of the mean escape time: κ = 〈τ〉−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. (14). 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 = 1 (i.e. the single-variable case θ ∈ R), the escape rate through a saddle θs is given by
κ = 1
2π
√ h∗|hs| [ L(θs)
L(θ∗)
]−( 12 + Bηh∗ ) , (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 θ ∈ RP 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 flatness) 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 flat directions. Therefore SGD dynamics around the local minimum θ∗ is restricted to the n-dimensional manifold spanned by the outlier eigenvectors v1, v2, . . . , vn ∈ RP of H(θ∗). Now we parameterize θ by using n parameters z1, z2, . . . , zn ∈ R as θ = θ∗ + ∑n i=1 zivi. The Hessian restricted to this outlier subspace is then written as Ĥ(θ) := ∇2zL(θ) ∈ Rn×n. We then obtain the following escape rate formula for P > 1:
κ = |hse| 2π
√ det Ĥ(θ∗)
|det Ĥ(θs)|
[ L(θs)
L(θ∗)
]−( Bηh∗+1−n2 ) , (16)
where hse is the negative eigenvalue ofH(θ s) corresponding to the escape direction.1 Again, Eq. (16) is valid when κ is small enough. It should be noted that Eq. (16) is reduced to Eq. (15) when n = 1.
1If eigenvalues of Ĥ(θ∗) and those of Ĥ(θs) coincide with each other except for the escape direction e, det Ĥ(θ∗)/ det Ĥ(θs) = h∗e/hse holds, and Eq. (16) is simplified as κ =
(1/2π) √ h∗e |hse|[L(θs)/L(θ∗)] − ( B ηh∗ +1− n 2 ) .
For any P , the quasi-stationary distribution Pss(θ) within the valley of a local minimum θ∗, which can be identified as the stationary distribution restricted to the valley (Bianchi & Gaudillière, 2016), is written as
Pss(θ) ∝ L(θ)−φ, φ = 1 + B
ηh∗ . (17)
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 finds flat minima with a low effective dimension n. Equation (16) also implies that the effective dimension of any stable minima must satisfy
n < nc := 2
( B
ηh∗ + 1
) . (18)
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∇2zU(θ∗)/|det∇2zU(θs)| = [det Ĥ(θ∗)/|det Ĥ(θs)|] · [L(θs)/L(θ∗)]n/2. Derivation of Eq. (16) and Eq. (17): 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: √ 2ηH(θ∗)/BdW̃τ ≈ √ 2ηh∗/BdW̃τ , where h∗ ∈ 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 a Hessian eigenvector ve, where e ∈ {1, 2, . . . , n} identifies 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τ = −∇zUdτ + √
2TdW̃τ , 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 fixed τ . However, what we want is the quasi-stationary distribution of θt for a fixed t. In Appendix B, by using Eq. (13), we show that the two distributions are related with each other as Pss(θ) ∝ L(θ)−1P̃ss(θ). 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 κτ per unit τ is given by
κτ = |use| 2π √ det∇2zU(θ∗) |det∇2zU(θs)| e−∆U/T , (19)
where use is the negative eigenvalue of ∇2zU(θ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 τ ' L(θ∗)t. The escape rate κ per unit t is then given by κ = L(θ∗)κτ . By combining it with Eq. (19), and substituting U(θ) = logL(θ), we finally obtain Eq. (16).

heading: 5 EXPERIMENTS
text: 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
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 confirm 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.

heading: 5.1 EXPERIMENTAL VERIFICATION OF THE DECOUPLING APPROXIMATION
text: Let us compare the eigenvalue distribution of the exact matrix (1/N) ∑N µ=1 ` (µ)C (µ) f with that of
the decoupled one L(θ) · (1/N) ∑N µ=1 C (µ) f with C (µ) f = ∇f(θ, x(µ))∇f(θ, x(µ))T. We consider a binary classification problem using the first 104 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 fix η = 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.

heading: 5.2 MEASUREMENTS OF THE SGD NOISE STRENGTH
text: As a measure of the SGD noise strength, let us consider the norm of the noise vector ξ given by 〈ξTξ〉 = Tr Σ ≈ N/B, where N := (1/N) ∑N µ=1∇`Tµ∇`µ −∇LT∇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×103 units and the ReLU activation, followed by the output layer of 10 units with no activation (classifica-
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 a kernel size of b × c with the stride length of d by n× [a, b, c, d]. We use the configuration: 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 sizeB = 100 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. 2. 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 thatN ≈ 2L(θ)× (1/N) ∑N µ=1∇f(θ, x(µ))T∇f(θ, x(µ)), but the quantity
(1/N) ∑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.

heading: 5.3 EXPERIMENTAL TEST OF STATIONARY DISTRIBUTION AND ESCAPE RATE FORMULA
text: We experimentally verify our theoretical predictions. Below, we first 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(µ) ∈ Rd 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 N . The output for an input x is given by f(θ, x) = θTx, where θ ∈ Rd is the trainable network parameter. We optimize θ via SGD. The mean-square loss L(θ) = (1/2N) ∑N µ=1 ( θx(µ) − y(µ) )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 = 105. We sampled the value of θk at every 100 iterations (k = j × 100, j = 1, 2, . . . , 104) and made a histogram. We then fit the histogram to the form Pss(θ) ∝ 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 fixed learning rates η. In the same figure, 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 finite 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. (16). 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 a non-convex loss landscape. Let us fix a threshold value of the loss function. The first
passage time tp is defined as the shortest time at which the loss exceeds the threshold value. Here, time t is identified 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 tp is similar to the escape time and proportional to κ−1.
The HessianH = (1/N) ∑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 first passage time over 100 independent runs is measured for varying threshold values which are specified by c = L(θs)/L(θ∗) > 1. Experimental results for N = 104 are presented in Fig. 3 (b). Dashed straight lines have slope B/(ηh∗) + 1− n/2. Experiments show that the first passage time behaves as tp ∝ [L(θs)/L(θ∗)]B/(ηh
∗)+1−n/2, which agrees with our theoretical evaluation of κ−1 [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 et al., 2021; Meng et al., 2020).
Our escape-rate formula (16) is also verified in a non-linear model. As in Sec. 5.1, we consider the binary classification problem using the first 104 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.01 and B = 100 for 105 iterations. We find that after pre-training, the loss becomes almost stationary around at L(θ) ≈ 0.035. 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 fix B = 10), and measure the first passage time tp. The mean first-passage time over 100 independent runs is measured for varying threshold values which are specified 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 first-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. (11). By identifying the largest eigenvalue of the Hessian as h∗, 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. 3 (c) correspond to our theoretical prediction κ−1 ∝ cB/(η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.

heading: 6 CONCLUSION
text: 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(θ) = logL(θ). 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 flat 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 Lévy-driven SDE, which explains why SGD finds wide minima. It would be an interesting future problem to take the non-Gaussian effect into account.

heading: A LIST OF APPROXIMATIONS AND THEIR JUSTIFICATIONS
text: In Sec. 3.1, we made several approximations to derive the result (6). For clarity, we list the approximations made and their justifications below.
• In Eq. (8), 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 ∇L∇LT in Eq. (8), which is justified 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 verified 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 justified 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 first term of the right-hand side of Eq. (10) (Papyan, 2018).
• In Eq. (11), we assume that the matrix (1/N) ∑N µ=1∇f(θ, x(µ))∇f(θ, x(µ))T does not
change so much in a valley with a 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 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 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 subspace. 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(θ∗)t, which is justified because the model stays near a local minimum until it escapes from that minimum.

heading: B STATIONARY DISTRIBUTION
text: Since θ̃τ obeys a simple Langevin equation
dθ̃τ = −U ′(θ̃τ ) + √ 2TdW̃τ , (20)
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 Pss(θ) of θt, where θt = θ̃τ with τ = ∫ t 0 dt′ L(θt′). In this section, we show the relation between the two distributions: Pss(θ) ∝ L(θ)−1P̃ss(θ). We express the stationary distributions in terms of the long-time average of the delta function:
Pss(θ) = lim s→∞
1
s ∫ s 0 dt δ(θt − θ), P̃ss(θ) = lim s→∞ 1 s ∫ s 0 dτ δ(θ̃τ − θ). (21)
By using the relation τ = ∫ t
0 dt′ L(θt′), we have dτ = L(θt)dt. For a sufficiently large t, we also obtain τ ∼ tL̄, where L̄ := lims→∞(1/s) ∫ s 0 dt′ L(θt′) denotes the long-time average of L(θt). By
using them, Pss(θ) is rewritten as
Pss(θ) ≈ lim s→∞
1
s ∫ sL̄ 0 dτ δ(θ̃τ − θ) L(θ̃τ )
= 1
L(θ) lim s→∞
L̄
sL̄ ∫ sL̄ 0 dτ δ(θ̃τ − θ)
= L̄
L(θ) P̃ss(θ). (22)
We thus obtain the desired relation, Pss(θ) ∝ L(θ)−1P̃ss(θ).
C JUSTIFICATION 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, ∂2U/∂zi∂zj = 0 for any i 6= j, where the Hessian at the minimum is given by H(θ∗) ≈∑n i=1 h ∗ i viv T i (hi is an eigenvalue and vi is the corresponding eigenvector of H(θ
∗)) and the displacement vector z ∈ Rn is defined by θ = θ∗ + ∑n i=1 zivi. The stochastic differential equation
dz = −∇zUdτ + √
2ηH(θ∗)/BdW̃τ is then equivalent to the following Fokker-Planck equation for the distribution function P (z, τ) of z at τ :
∂P (z, τ)
∂τ = n∑ i=1 [ ∂ ∂zi ( ∂U ∂zi P ) + ηh∗i B ∂2 ∂z2i P ] . (23)
Let us assume that the direction of eth eigenvector ve corresponds to the escape direction. We denote by z⊥ ∈ Rn−1 the displacement perpendicular to the escape direction, and write z = (ze, z⊥). At the saddle zs,∇zU(zs) = 0, hse < 0 and hsi > 0 for all i 6= e, where {hsi} is the set of eigenvectors of the Hessian at zs.
Under the above setting, we derive the escape rate formula following Kramers (1940). The steady current J ∈ Rn is aligned to the escape direction, and hence Je 6= 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 flowing to outside of the valley through the saddle θs = θ∗ + zs. It is assumed that J is small and P ∗ is almost stationary. The escape rate κτ is then given by
κτ = J P ∗ . (24)
The escape rate κ per unit time is given by
κ = L(θ∗)κτ = L(θ ∗) J P ∗ . (25)
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),
U(θ) = logL(θ) ≈ U(θ∗) + (θ − θ ∗)TH(θ∗)(θ − θ∗)
2L(θ∗) = U(θ∗) +
zTH(θ∗)z
2L(θ∗) (26)
and ∂U
∂zi ≈ 1 L(θ∗) n∑ j=1 H∗ijzj , (27)
where H∗ = H(θ∗). By substituting it into the Fokker-Planck equation (23), we obtain
∂P (z, τ)
∂τ = n∑ i=1 ∂ ∂zi  n∑ j=1 H∗ij ( zj L(θ∗) + η B ∂ ∂zj ) P  . (28)
If for all j ( zj
L(θ∗) + η B ∂
∂zj
) Pss(z) = 0, (29)
Pss(z) is a stationary distribution near the minimum θ∗. Indeed,
Pss(z) = P (θ ∗)e− B 2ηL(θ∗)
∑n i=1 z
2 i (30)
satisfies 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
P ∗ ≈ ∫ dz1dz2 . . . dznPss(z) = P (θ ∗) [ 2ηL(θ∗)
B
]n/2 . (31)
Next, let us evaluate J . The probability current J along the escape direction e is given by
Je = − ∂U ∂ze P − ηh
∗ e
B
∂
∂ze P. (32)
At the saddle, the current vector J⊥ perpendicular to the escape direction is zero, and hence the probability distribution near the saddle with ze = zse and z⊥ 6= 0 is also evaluated in a similar way as in Eq. (30):
P (ze = z s e , z⊥) ≈ P (ze = zse , z⊥ = 0)e
− B 2ηL(θs) z2⊥ . (33) By substituting it into Eq. (32), we obtain
Je(z s e , z⊥) = Jpathe
− B 2ηL(θs) z2⊥ , (34)
where the current along the escape path (z⊥ = 0) is denoted by Jpath := J(zse , z⊥ = 0). The total current through the saddle is then evaluated as
J = ∫ dz⊥ Je(z s e , z⊥) = Jpath [ 2πηL(θs)
B
](n−1)/2 . (35)
When the distribution function is almost stationary, Eq. (23) yields ∂Je(ze, z⊥ = 0)/∂ze ≈ 0, and hence the current along the escape path is approximately constant Je(ze, z⊥ = 0) ≈ Jpath. Since Je(ze, z⊥ = 0) is given by Eq. (32) by putting z⊥ = 0, we have
Jpath = − ∂U
∂ze ∣∣∣∣ z⊥=0 P − ηh ∗ e B ∂P ∂ze = −ηh ∗ e B e − B ηh∗e U ∂ ∂ze ( e B ηh∗e U P ) . (36)
By multiplying e B ηh∗e U in both sides and integrating over ze from 0 to zce, where z c e defined as U(θ∗ + zceve) = U(θ ∗) (see Fig. 4), we obtain
Jpath ∫ zce 0 dze e B ηh∗e U = ηh∗e B e B ηh∗e U(θ∗) P (θ∗), (37)
where it is assumed that the probability at zce is small and negligible, P (ze = z c e, z⊥ = 0) ≈ 0. By using the saddle-point method, the integral in the left-hand side of Eq. (37) is evaluated as∫ zce 0 dze e B ηh∗e U ≈ ∫ ∞ −∞ dze exp [ B ηh∗e ( U(θs) + hse 2L(θs) (ze − zse)2 )]
=
( 2πηh∗e
B|hse|L(θs)
)1/2 e B ηh∗e U(θs) . (38)
By substituting this result in Eq. (37), we obtain
Jpath =
( ηh∗e|hse|
2πBL(θs)
)1/2 e − B ηh∗e ∆U P (θ∗), (39)
where ∆U = U(θs)− U(θ∗). The total current J in Eq. (35) is then expressed as J = √ h∗e|hse|
2πL(θs)
( 2πηL(θs)
B
)n/2 e − B ηh∗e ∆U P (θ∗). (40)
By using Eqs. (31) and (40), the escape rate κ in Eq. (25) is evaluated as
κ = L(θ∗) J P ∗ = √ h∗e|hse| 2π [ L(θs) L(θ∗) ]n 2−1 e− B ηhe∗∆U . (41)
Since ∆U = log[L(θs)/L(θ∗)], we finally obtain
κ = √ h∗e|hse| 2π [ L(θs) L(θ∗) ]−( B ηh∗e +1−n2 ) , (42)
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 justified 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 justified in this case.

heading: D OTHER LOSS FUNCTIONS
text: 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 definition of N ). Here, we argue that the relation N ∝ L(θ) also holds in more general situations. During the training, the value of `µ will fluctuate from sample to sample. At a certain time step of SGD, let us suppose that N −M samples in the training dataset are already classified correctly and hence `µ ∼ 0, whereas the other M samples are not and hence `µ ∼ 1. The loss function is then given by L(θ) = (1/N) ∑N µ=1 `µ ∼ M/N . When `µ is small, ∇`µ will also be small. Therefore, for N − M samples with `µ ∼ 0, ∇`µ ∼ 0 also holds. The other M samples will have nonsmall gradients: ‖∇`µ‖2 ∼ g, where g > 0 is a certain positive constant. We thus estimate N as N ≈ (1/N) ∑N µ=1∇`Tµ∇`µ ∼ 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 classified there, which implies that `µ is small for all µ, in contrast to the above argument relying on the existence of M wrongly classified samples with `µ ∼ 1 and ‖∇`µ‖2 ∼ g.
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. 5.2 for the detail). We fix B = 100 in both cases, and η = 0.1 for the fully connected network and η = 0.05 for the convolutional network. Experimental results are presented in Fig. 5. We find that the relationN ∝ 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 find that for sufficiently small values of the loss,N ∝ L(θ)2 [see Fig. 5 (c) and (d)], whose implications should merit further investigation in future studies.

heading: E HESSIAN EIGENVALUES FOR A NEURAL NETWORK IN SEC. 5.3
text: We present numerical results on Hessian eigenvalues in a pre-trained neural network studied in Sec. 5.3. Instead of the exact Hessian, we consider an approximate Hessian given on the right-hand side of Eq. (11), i.e.,
H(θ∗) ≈ 1 N N∑ µ=1 ∇f(θ, x(µ))∇f(θ, x(µ))T. (43)
Eigenvalues {hi} are arranged in descending order as h1 ≥ h2 ≥ · · · ≥ hP (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 identified as h∗ in our theoretical formula (16).
(a) (b)
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 difficult to precisely determine n. In Fig. 6 (b), we plot hi up to i = 30. From this figure, it seems reasonable to estimate n ∼ 10. As a heuristic method of determining n, we can consider the following identification: first we define the weight pi for ith eigenmode as pi = hi/ ∑P j=1 hj . We then determine n as
n = ( P∑ i=1 p2i )−1 , (44)
which gives n = 12.7 in our case.
In Sec. 5.3, our formula (16) with n = 9 explains numerical results on the first-passage time.

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

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.