Title: THE ANISOTROPIC NOISE IN STOCHASTIC GRADIENT DESCENT: ITS BEHAVIOR OF ESCAPING FROM MIN-IMA AND REGULARIZATION EFFECTS Anonymous

Abstract: Understanding the behavior of stochastic gradient descent (SGD) in the context of deep neural networks has raised lots of concerns recently. Along this line, we theoretically study a general form of gradient based optimization dynamics with unbiased noise, which unifies SGD and standard Langevin dynamics. Through investigating this general optimization dynamics, we analyze the behavior of SGD on escaping from minima and its regularization effects. A novel indicator is derived to characterize the efficiency of escaping from minima through measuring the alignment of noise covariance and the curvature of loss function. Based on this indicator, two conditions are established to show which type of noise structure is superior to isotropic noise in term of escaping efficiency. We further show that the anisotropic noise in SGD satisfies the two conditions, and thus helps to escape from sharp and poor minima effectively, towards more stable and flat minima that typically generalize well. We verify our understanding through comparing this anisotropic diffusion with full gradient descent plus isotropic diffusion (i.e. Langevin dynamics) and other types of position-dependent noise.

1. INTRODUCTION:

As a successful learning algorithm, stochastic gradient descent (SGD) was originally adopted for dealing with the computational bottleneck of training neural networks with large-scale datasets (Bottou, 1991) . Its empirical efficiency and effectiveness have attracted lots of attention. And thus, SGD and its variants have become standard workhorse for learning deep models. Besides the aspect of empirical efficiency, recently, researchers started to analyze the optimization behaviors of SGD and its impacts on generalization.

The optimization properties of SGD have been studied from various perspectives. The convergence behaviors of SGD for simple one hidden layer neural networks were investigated in (Li & Yuan, 2017; Brutzkus et al., 2017) . In non-convex settings, the characterization of how SGD escapes from stationary points, including saddle points and local minima, was analyzed in (Daneshmand et al., 2018; Jin et al., 2017; Hu et al., 2017) .

On the other hand, in the context of deep learning, researchers realized that the noise introduced by SGD impacts the generalization, thanks to the research on the phenomenon that training with a large batch could cause a significant drop of test accuracy (Keskar et al., 2017) . Particularly, several works attempted to investigate how the magnitude of the noise influences the generalization during the process of SGD optimization, including the batch size and learning rate (Hoffer et al., 2017; Goyal et al., 2017; Chaudhari & Soatto, 2017; Jastrzębski et al., 2017) . Another line of research interpreted SGD from a Bayesian perspective. In (Mandt et al., 2017; Chaudhari & Soatto, 2017) , SGD was interpreted as performing variational inference, where certain entropic regularization involves to prevent overfitting. And the work (Smith & Le, 2018) tried to provide an understanding based on model evidence. These explanations are compatible with the flat/sharp minima argument (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017) , since Bayesian inference tends to targeting the region with large probability mass, corresponding to the flat minima.

However, when analyzing the optimization behavior and regularization effects of SGD, most of existing works only assume the noise covariance of SGD is constant or upper bounded by some constant, and what role the noise structure of stochastic gradient plays in optimization and generalization was rarely discussed in literature.

In this work, we theoretically study a general form of gradient-based optimization dynamics with unbiased noise, which unifies SGD and standard Langevin dynamics. By investigating this general dynamics, we analyze how the noise structure of SGD influences the escaping behavior from minima and its regularization effects. Several novel theoretical results and empirical justifications are made.

1. We derive a key indicator to characterize the efficiency of escaping from minima through measuring the alignment of noise covariance and the curvature of loss function. Based on this indicator, two conditions are established to show which type of noise structure is superior to isotropic noise in term of escaping efficiency;

2. We further justify that SGD in the context of deep neural networks satisfies these two conditions, and thus provide a plausible explanation why SGD can escape from sharp minima more efficiently, converging to flat minima with a higher probability. Moreover, these flat minima typically generalize well according to various works (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Neyshabur et al., 2017; Wu et al., 2017) . We also show that Langevin dynamics with well tuned isotropic noise cannot beat SGD, which further confirms the importance of noise structure of SGD;

3. A large number of experiments are designed systematically to justify our understanding on the behavior of the anisotropic diffusion of SGD. We compare SGD with full gradient descent with different types of diffusion noise, including isotropic and positiondependent/independent noise. All these comparisons demonstrate the effectiveness of anisotropic diffusion for good generalization in training deep networks.

The remaining of the paper is organized as follows. In Section 2, we introduce the background of SGD and a general form of optimization dynamics of interest. We then theoretically study the behaviors of escaping from minima in Ornstein-Uhlenbeck process in Section 3, and establish two conditions for characterizing the noise structure that affects the escaping efficiency. In Section 4, we show that the noise of SGD in the context of deep learning meets the two conditions, and thus explains its superior efficiency of escaping from sharp minima over other dynamics with isotropic noise. Various experiments are conducted for verifying our understanding in Section 5, and we conclude the paper in Section 6.

2. BACKGROUND:

In general, supervised learning usually involves an optimization process of minimizing an empirical loss over training data, L(θ) := 1/N N i=1 (f (x i ; θ), y i ), where {(x i , y i )} N i=1 denotes the training set with N i.i.d. samples, the prediction function f is often parameterized by θ ∈ R D , such as deep neural networks. And (•, •) is the loss function, such as mean squared error and cross entropy, typically corresponding to certain negative log likelihood. Due to the over parameterization and non-convexity of the loss function in deep networks, there exist multiple global minima, exhibiting diverse generalization performance. We call those solutions generalizing well good solutions or minima, and vice versa.

Gradient descent and its stochastic variants A typical approach to minimize the loss function is gradient descent (GD), the dynamics of which in each iteration t is, θ t+1 = θ t -η t g 0 (θ t ), where g 0 (θ t ) = ∇ θ L(θ t ) denotes the full gradient and η t denotes the learning rate. In non-convex optimization, a more useful kind of gradient based optimizers act like GD with an unbiased noise, including gradient Langevin dynamics (GLD), θ t+1 = θ t -η t g 0 (θ t ) + σ t t , t ∼ N (0, I), and stochastic gradient descent (SGD), during each iteration t of which, a minibatch of training samples with size m are randomly selected, with index set B t ⊂ {1, 2, . . . , N }, and a stochastic gradient is evaluated based on the chosen minibatch, g(θ t ) = i∈Bt ∇ θ (f (x i ; θ t ), y i )/m, which is an unbiased estimator of the full gradient g 0 (θ t ). Then, the parameters are updated with some learning rate η t as θ t+1 = θ t -η t g(θ t ). Denote g(θ) = ∇ θ ((f (x; θ), y), the gradient for loss with a single data point (x, y), and assume that the size of minibatch is large enough for the central limit theorem to hold, and thus g(θ t ) follows a Gaussian distribution (Mandt et al., 2017; Li et al., 2017) ,

g(θ t ) ∼ N g 0 (θ t ), 1 m Σ(θ t ) , where Σ(θ t ) ≈ 1 N N i=1 g(θ t ; x i ) -g 0 (θ t ) g(θ t ; x i ) -g 0 (θ t ) T .

(1) Note that the covariance matrix Σ depends on the model architecture, dataset and the current parameter θ t . Now we can rewrite the update of SGD as,

EQUATION

Inspired by GLD and SGD, we may consider a general kind of optimization dynamics, namely, gradient descent with unbiased noise,

EQUATION

For small enough constant learning rate η t = η, the above iteration in Eq. ( 3) can be treated as the numerical discretization of the following stochastic differential equation (Li et al., 2017; Jastrzębski et al., 2017; Chaudhari & Soatto, 2017) ,

EQUATION

Considering ησ 2 t Σ t as the coefficient of noise term, existing works (Hoffer et al., 2017; Jastrzębski et al., 2017) studied the influence of noise magnitude of SGD on generalization, i.e. ησ 2 t = η/m. In this work, we focus on studying the benefits of anisotropic structure of Σ t in SGD helping escape from minima by bridging the covariance matrix with the Hessian of the loss surface, and its implicit regularization effects on generalization, especially in deep learning context. For the purpose of eliminating the influence of the noise magnitude, we constrain it to be a constant when studying different structures of noise covariance. The noise magnitude could be evaluated as the expectation of the squared norm of the noise vector,

EQUATION

Thus, we introduce the following constraint,

EQUATION

From the statistical physics point of view, Tr(ησ 2 t Σ t ) characterizes the kinetic energy (Gardiner) , thus it is natural to force the energy to be unchanging, otherwise it is trivial that the higher the energy is, the less stable the system is.

For simplicity, we absorb ησ 2 t into Σ t , denoting ησ 2 t Σ t as Σ t . If not pointed out, the subscript t of matrix Σ t is omitted to emphasize that we are fixing t and discussing the varying structure of Σ.

3. THE BEHAVIORS OF ESCAPING FROM MINIMA IN ORNSTEIN-UHLENBECK PROCESS:

For a general loss function L(θ) = E X X (θ) (the expectation could be either population or empirical), where X denotes data example and θ denoted parameters to be optimized, under suitable smoothness assumptions, the SDE associated with the gradient variant optimizer as shown in Eq. ( 4) can be written as follows (Li et al., 2017; Jastrzębski et al., 2017; Chaudhari & Soatto, 2017; Hu et al., 2017) , with little abuse of notation,

EQUATION

Let L 0 = L(θ 0 ) be one of the minimal values of L(θ), then for a fixed t small enough (such that L t -L 0 ≥ 0), E θt [L t -L 0 ] characterizes the efficiency of θ escaping from the minimum θ 0 of L(θ).

It is natural to measure the escaping efficiency using E[L t -L 0 ] since it characterizes the increase of the potential, i.e., the increase of the loss L. And also note that L t -L 0 ≥ 0, for any δ > 0, the escaping probability P (L t -L 0 ≥ δ) can be controlled by the expectation E[L t -L 0 ] since by Markov's inequality, we have

P (L t -L 0 ≥ δ) ≤ E[Lt-L0] δ .

Proposition 1 (Escaping efficiency for general process). For the process (7), provided mild smoothness assumptions, the escaping efficiency from the minimum θ 0 is,

EQUATION

where H t denotes the Hessian of L(θ t ) at θ t .

We provide the proof in Appendix, and the same for the other propositions.

The escaping efficiency for general processes is hard to analyze due to the intractableness of the integral in Eq. ( 8). However, we may consider the second-order approximation locally near the minima θ 0 , where

L(θ) ≈ L 0 + 1 2 (θ -θ 0 ) T H(θ -θ 0 ).

Without losing generality, we suppose θ 0 = 0. Further, suppose that H is a positive definite matrix and the diffusion covariance Σ t = Σ is constant for t. Then the SDE ( 7) becomes an Ornstein-Uhlenbeck process,

EQUATION

)

Proposition 2 (Escaping efficiency of Ornstein-Uhlenbeck process). For Ornstein-Uhlenbeck process ( 9), with t small enough, the escaping efficiency from minimum θ 0 = 0 is,

EQUATION

)

Inspired by Proposition 1 and Proposition 2, we propose Tr (HΣ) as an empirical indicator measuring the efficiency for a stochastic process escaping from minima. Now we turn to analysis which kind of noise covariance structure Σ will benefit escaping sharp minima, under the constraint Eq. ( 6).

Firstly, for the isotropic loss surface, i.e., H = λI, the escaping efficiency is

E[L t -L 0 ] = λt 2

Tr Σ, which is invariant under the constraint that Tr Σ is constant (Eq. ( 6)). Thus it is only nontrivial to study the impact of noise structure when the Hessian of loss surface is anisotropic.

Secondly, H and Σ being semi-positive definite, to achieve the maximum of Tr(HΣ) under constraint (6), Σ should be Σ * = (Tr Σ) • λ 1 u 1 u T 1 , where λ 1 , u 1 are the maximal eigenvalue and corresponding unit eigenvector of H. Note that the rank-1 matrix Σ * is highly anisotropic. More generally, the following Proposition 3 characterizes one kind of anisotropic noise significantly outperforming isotropic noise in order of number of parameters D, given H is ill-conditioned. Proposition 3 (The benefits of anisotropic noise). With semi-positive definite H and Σ, assume

(1) H is ill-conditioned. Let λ 1 ≥ λ 2 ≥ . . . , ≥ λ D ≥ 0 be the eigenvalues of H in descent order, and for some constant k D and d > 1 2 ,

EQUATION

(2) Σ is "aligned" with H. Let u i be the corresponding unit eigenvector of eigenvalue λ i , for some projection coefficient a > 0,

EQUATION

)

then we have the benefit of the anisotropic noise over the isotropic one in term of escaping efficiency, which can be characterized by the follow ratio,

EQUATION

where Σ = Tr Σ D I denotes the covariance of isotropic noise, to meet the constraint Eq. ( 6).

To give some geometric intuitions on the left hand side of Eq. ( 12), let the maximal eigenvalue and its corresponding unit eigenvector of Σ be γ 1 , v 1 , then the right hand side has a lower bound as

u T 1 Σu 1 ≥ u T 1 v 1 γ 1 v T 1 u 1 = γ 1 u 1 , v 1 2 .

Thus if the maximal eigenvalues of H and Σ are aligned in proportion, γ 1 / Tr Σ ≥ a 1 λ 1 / Tr H, and the angle of their corresponding unit eigenvectors is close to zero, u 1 , v 1 ≥ a 2 , the second condition Eq. ( 12) in Proposition 3 holds for a = a 1 a 2 .

Typically, in the scenario of modern deep neural networks, due to the over-parameterization, Hessian and the gradient covariance are usually ill-conditioned and anistropic near minima, as shown by (Sagun et al., 2017) and (Chaudhari & Soatto, 2017) . Thus the first condition in Eq. ( 11) usually holds for deep neural networks, and we further justify it by experiments in Section 5.3. Therefore, in the following section, we turn to focus on how the gradient covariance, i.e. the covariance of SGD noise meets the second condition of Proposition 3 in the context of deep neural networks.

4. THE ANISOTROPIC NOISE OF SGD IN DEEP NETWORKS:

In this section, we mainly investigate the anisotropic structure of gradient covariance in SGD, and explore its connection with the Hessian of loss surface.

Around the true parameter According to the classic statistical theory (Pawitan, 2001, Chap. 8) , for population loss L(θ) = E X (θ), with being the negative log likelihood, when evaluating at the true parameter θ * , there is the exact equivalence between the Hessian H of the population loss and Fisher information matrix F ,

EQUATION

In practice, with the assumptions that the sample size N is large enough (i.e. indicating asymptotic behavior) and suitable smoothness conditions, when the current parameter θ t is not far from the ground truth, Fisher is close to Hessian. Thus we can obtain the following approximate equality between gradient covariance and Hessian,

Σ(θ t ) = F (θ t ) -∇ θ L T (θ t )∇ θ L(θ t ) ≈ F (θ t ) ≈ H(θ t ).

The first approximation is due to the dominance of noise over the mean of gradient in the later stage of SGD optimization, which has been shown in (Shwartz-Ziv & Tishby, 2017) . A similar experiment as (Shwartz-Ziv & Tishby, 2017) has been conducted to demonstrate this observation, which is left in Appendix due to the limit of space.

In the following, we theoretically characterize the closeness between Σ and H in the context of one hidden layer neural networks; and show that the gradient covariance introduced by SGD indeed has more benefits than isotropic one in term of escaping from minima, provided some assumptions.

One hidden layer neural network with fixed output layer parameters For binary classification neural network with one hidden layer in classic setups (with softmax and cross-entropy loss), we have following results to globally bound Fisher and Hessian with each other. Proposition 4 (The relationship between Fisher and Hessian in one hidden layer neural network).

Consider the binary classification problem with data {(x i , y i )} i∈I , y ∈ {0, 1}, and typical (either population or empirical) loss as

L(θ) = E[φ • f (x; θ)],

where f denotes the output of neural network, and φ denotes the cross-entropy loss with softmax,

φ(f (x), y) = -y log e f (x) 1 + e f (x) + (1 -y) log 1 1 + e f (x) , y ∈ {0, 1}.

If: (1) the neural network f is with one hidden layer and piece-wise linear activation. And the parameters of output layer are fixed during training; (2) the optimization happens on a set U such that, f (x; θ) ∈ (-C, C), ∀θ ∈ U, ∀x, i.e., the output of the classifier is bounded during optimization.

Then, we have the following relationship between (either population or empirical) Fisher F and Hessian H almost everywhere:

e -C F (θ) H(θ) e C F (θ). A B means that (B -A) is semi-positive definite.

There are a few remarks on Proposition 4. Firstly, as shown in (Brutzkus et al., 2017) , the considered neural networks in Proposition 4 are non-convex and have multiple minima, and thus it is still nontrivial to consider the escaping from minima. Secondly, the Proposition 4 holds in both population and empirical sense, since the proof does not distinguish the two circumstances. Thirdly, the bound between F and H holds "globally" in the set U where the output f is bounded, rather than merely around the true global minima as discussed previously.

By Proposition 4, the following relationship between gradient covariance and Hessian could be derived. Proposition 5 (The relationship between gradient covariance and Hessian in one hidden layer neural network). Assume the conditions in Proposition 4 hold, then for some small δ > 0 and for θ close enough to minima θ * (local or global),

u T Σu ≥ e -2(C+δ) λ Tr Σ Tr H (15)

holds for any positive eigenvalue λ and its corresponding unit eigenvector u of Hessian H.

As a direct corollary of Proposition 5, for such neural networks, the second condition Eq. ( 12) in Proposition 3 holds in a very loose sense.

Therefore, based on the discussion on population loss around the true parameters and one hidden layer neural network with fixed output layer parameters, given the ill-conditioning of H due to the over-parameterization of modern deep networks, according to Proposition 3, we can conclude the noise structure of SGD helps to escape from sharp minima much faster than the dynamics with isotropic noise, and converge to flatter solutions with a high probability. These flat minima typically generalize well (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017; Neyshabur et al., 2017; Wu et al., 2017) . Thus, we attribute such properties of SGD on its better generalization performance comparing to GD, GLD and other dynamics with isotropic noise (Hoffer et al., 2017; Goyal et al., 2017; Keskar et al., 2017) .

In the following, we conduct a series of experiments systematically to verify our understanding on the behavior of escaping from minima and its regularization effects for different optimization dynamics.

5. EXPERIMENTS:

To better understanding the behavior of anisotropic noise different from isotropic ones, we introduce dynamics with different kinds of noise structure to empirical study with, as shown on Table 1 . 

GLD 1st eigven(H):

t ∼ N 0, λ1u1u T 1 λ1, u1 are the maximal eigenvalue and its corresponding unit eigenvector of the Hessian matrix of loss L(θt).

5.1. TWO-DIMENSIONAL TOY EXAMPLE:

We design a 2-D toy example L(w 1 , w 2 ) with two basins, a small one and a large one, corresponding to a sharp and flat minima, (1, 1) and (-1, -1), respectively, both of which are global minima.

Please refer to Appendix for the detailed constructions. We initialize the dynamics of interest with the sharp minimum (w 1 , w 2 ) = (1, 1), and run them to study their behaviors escaping from this sharp minimum.

To explicitly control the noise magnitude, we only conduct experiments on GD, GLD const, GLD diag, GLD leading (with k = 2 = D in Table 1 , or in other words, the exactly covariance of SGD noise), GLD Hessian (k = 2) and GLD 1st eigven(H). And we adjust σ t in each dynamics to force their noise to share the same expected squared norm as defined in Eq. ( 6). Figure 1(a) shows the trajectories of the dynamics escaping from the sharp minimum (1, 1) towards the flat one (-1, -1), while Figure 1 (b) presents the success rate of escaping for each dynamic during 100 repeated experiments.

As shown in Figure 1 , GLD 1st eigvec(H) achieves the highest success rate, indicating the fastest escaping speed from the sharp minimum. The dynamics with anisotropic noise aligned with Hessian well, including GLD 1st eigvec(H), GLD Hessian and GLD leading, greatly outperform GD, GLD const with isotropic noise, and GLD diag with noise poorly aligned with Hessian. These experiments are consistent with our theoretical analysis on Ornstein-Uhlenbeck process shown Proposition 2 and 3, demonstrating the benefits of anisotropic noise for escaping from sharp minima. We empirically show that in one hidden layer neural network with fixed output layer parameters, the anisotropic noise induced by SGD indeed helps escape from sharp minima more efficiently than isotropic noise. Three networks are trained to binary classify 1, 000 linearly separable two-dimensional points. The number of hidden nodes for each network varies in {20, 200, 2000}. We plot the empirical indicator Tr (HΣ) in Figure 2 . We can easily observe that as the increase of the number of hidden nodes, the ratio Tr(HΣ) Tr(H Σ)

5.2. ONE HIDDEN LAYER NEURAL NETWORK WITH FIXED OUTPUT LAYER PARAMETERS:

is enlarged significantly, which is consistent with the Eq. ( 13) described in Proposition 3.

5.3. PRACTICAL DATASETS:

In this part, we conduct a series of experiments in real deep learning scenarios to demonstrate the behavior of SGD noise and its implicit regularization effects. We construct a noisy training set based on FashionMNIST dataset 1 . Concretely, the training set consist of 1000 images with correct labels, and another 200 images with random labels. All the test data are with clean labels. A small LeNet-like network is utilized such that the spectrum decomposition over gradient covariance matrix and Hessian matrix are computationally feasible. The network consists of two convolutional layers and two fully-connected layers, with 11, 330 parameters in total.

We firstly run the standard gradient decent for 3000 iterations to arrive at the parameters θ * GD near the global minima with near zero training loss and 100% training accuracy, which are typically sharp minima that generalize poorly (Neyshabur et al., 2017) . And then all other compared methods are initialized with θ * GD and run for optimization with the same learning rate η t = 0.07 and same batch size m = 20 (if needed) for fair comparison2 .

Verification of SGD noise satisfying the conditions in Proposition 3 To see whether the noise of SGD in real deep learning circumstance satisfies the two conditions in Proposition 3, we run SGD optimizer initialized from θ * GD , i.e. the sharp minima found by GD. Figure 3 along the trajectory of SGD. The plot indicates that the projection coefficient is in a descent scale comparing to D 2d-1 , thus satisfying the second condition in Proposition 3. Therefore, Proposition 3 ensures that SGD would escape from minima θ * GD faster than GLD in order of O(D 2d-1 ), as shown in Figure 3(c ). An interesting observation is that in the later stage of SGD optimization, Tr(HΣ) becomes significantly (10 7 times) smaller than in the beginning stage, implying that SGD has already converged to minima being almost impossible to escape from. This phenomenon demonstrates the reasonability to employ Tr(HΣ) as an empirical indicator for escaping efficiency. Behaviors of different dynamics escaping from minima and its generalization effects To compare the different dynamics on escaping behaviors and generalization performance, we run dynamics initialized from the sharp minima θ * GD found by GD. The settings for each compared method are as follows. The hyperparameter σ 2 for GLD const has already been tuned as optimal (σ = 0.001) by grid search. For GLD leading, we set k = 20 for comprising the computational cost and approximation accuracy. As for GLD Hessian, to reduce the expensive evaluation of such a huge Hessian in each iteration, we set k = 20 and update the Hessian every 10 iterations. We adjust σ t in GLD dynamic, GLD Hessian and GLD 1st eigvec(H) to guarantee that they share the same expected squred noise norm defined in Eq. ( 6) as that of SGD. And we measure the expected sharpness of different minima as E ν∼N (0,δ 2 I) L(θ + ν) -L(θ), as defined in ((Neyshabur et al., 2017) , Eq.( 7)). The results are shown in Figure 4 .

As shown in Figure 4 , SGD, GLD 1st eigvec(H), GLD leading and GLD Hessian successfully escape from the sharp minima found by GD, while GLD, GLD dynamic and GLD diag are trapped in the minima. This demonstrates that the methods with anisotropic noise "aligned" with loss curvature can help to find flatter minima that generalize well.

We also provide experiments on standard CIFAR-10 with VGG11 in Appendix. 

6. CONCLUSION:

We theoretically investigate a general optimization dynamics with unbiased noise, which unifies various existing optimization methods, including SGD. We provide some novel results on the behaviors of escaping from minima and its regularization effects. A novel indicator is derived for characterizing the escaping efficiency. Based on this indicator, two conditions are constructed for showing what type of noise structure is superior to isotropic noise in term of escaping. We then analyze the noise structure of SGD in deep learning and find that it indeed satisfies the two conditions, thus explaining the widely know observation that SGD can escape from sharp minima efficiently toward flat minina that generalize well. Various experimental evidence supports our arguments on the behavior of SGD and its effects on generalization. Our study also shows that isotropic noise helps little for escaping from sharp minima, due to the highly anisotropic nature of landscape. This indicates that it is not sufficient to analyze SGD by treating it as an isotropic diffusion over landscape (Zhang et al., 2017; Mou et al., 2017) . A better understanding of this out-of-equilibrium behavior (Chaudhari & Soatto, 2017 ) is on demand.

A PROOFS OF PROPOSITIONS IN MAIN PAPER:

A.1 PROOF OF PROPOSITION 1 Proof. The "mild smoothness assumptions" refers that L t = L(θ t ) ∈ C 2 . Then the Ito's lemma holds (Øksendal, 2003) .

And by Ito's lemma, the SDE of L t is

dL t = -∇L T ∇L + 1 2 Tr Σ 1 2 t H t Σ 1 2 t dt + ∇L T Σ 1 2 t dW t = -∇L T ∇L + 1 2 Tr (H t Σ t ) dt + ∇L T Σ 1 2 t dW t .

Taking expectation with respect to the distribution of θ t ,

EQUATION

for the expectation of Brownian motion is zero. Thus the solution of EY t is,

EL t = L 0 - t 0 E ∇L T ∇L + t 0 1 2 E Tr(H t Σ t ) dt.

A.2 PROOF OF PROPOSITION 2

Proof. Without losing generality, we assume that L 0 = 0.

For multivariate Ornstein-Uhlenbeck process, when θ 0 = 0 is an constant, θ t follows a multivariate Gaussian distribution (Øksendal, 2003) .

Consider change of variables θ → φ(θ, t) = e Ht θ t . Here, for symmetric matrix A, e A := U diag(e λ1 , . . . , e λn )U,

where λ 1 , . . . , λ n and U are the eigenvalues and eigenvector matrix of A. Note that with this notation,

de Ht dt = He Ht .

Applying Ito's lemma, we have

dφ(θ t , t) = e Ht Σ 1 2 dW t ,

which we can integrate form 0 to t to get

θ t = 0 + t 0 e H(s-t) Σ 1 2 dW s

The expectation of θ t is zero. And by Ito's isometry (Øksendal, 2003) , the covariance of θ t is, Tr e H(s-t) HΣe H(s-t) ds (for H is symmetric.) 

Eθ t θ T t = E   t 0 e H(s-t) Σ 1 2 dW s t 0 e H(r-t) Σ 1 2 dW r T   = E t 0 e H(s-t) Σ 1 2 Σ 1 2 e

EQUATION

)

The proof is finished.

A.4 PROOF OF PROPOSITION 4

Proof. Firstly compute the gradients and Hessian of φ,

∂ φ ∂f = e f 1 + e f -y = e f 1+e f > 0 y = 0, -1 1+e f < 0 y = 1. ∂ 2 φ ∂f 2 = e f (1 + e f ) 2 .

And note the Gauss-Newton decomposition for functions with the form of L = φ • f ,

H = E (x,y) ∂ 2 ((x,y);θ) ∂θ 2 = E (x,y) ∂ 2 φ ∂f 2 ∂ f ∂θ ∂ f T ∂θ + E (x,y) ∂ φ ∂f ∂ 2 f ∂θ 2 .

Since the output layer parameters for f is fixed and the activation functions are piece-wise linear, f (x; θ) is a piece-wise linear function on its parameters θ. Therefore ∂ 2 f ∂θ 2 = 0, a.e., and H = E (x,y)

∂ 2 φ ∂f 2 ∂ f ∂θ ∂ f T ∂θ . It is easy to check that e -C ∂ φ ∂f 2 ≤ ∂ 2 φ ∂f 2 ≤ e C ∂ φ ∂f 2

. Thus,

H = E (x,y) ∂ 2 φ ∂f 2 ∂ f ∂θ ∂ f T ∂θ E (x,y) e C ∂ φ ∂f 2 ∂ f ∂θ ∂ f T ∂θ = E (x,y) e C ∂ φ ∂f ∂ f ∂θ ∂ φ ∂f ∂ f ∂θ T = e C F. H = E (x,y) ∂ 2 φ ∂f 2 ∂ f ∂x ∂ f T ∂x E (x,y) e -C ∂ φ ∂f 2 ∂ f ∂x ∂ f T ∂x = E (x,y) e -C ∂ φ ∂f ∂ f ∂θ ∂ φ ∂f ∂ f ∂θ T = e -C F.

A.5 PROOF OF PROPOSITION 5

Proof. For simplicity, we define g := ∇ , g 0 := ∇L = E∇ .

The gradient covariance and Fisher has the following relationship,

F = Eg • g T = E(g 0 + )(g 0 + ) T = g 0 g T 0 + E T = g 0 g T 0 + Σ.

Applying Taylor's expansion to g 0 (θ),

g 0 (θ) = g 0 (θ * ) + H(θ * )(θ -θ * ) + o(θ -θ * ) = H(θ * )(θ -θ * ) + o(θ -θ * ).

Hence,

g 0 (θ) 2 2 ≤ H 2 2 θ -θ * 2 2 + o θ -θ * 2 2 = H 2 2 θ -θ * 2 2 + o θ -θ * 2 2 .

Therefore, with the condition θ

-θ * 2 ≤ √ δu T F u H 2 , we have g 0 (θ) 2 2 ≤ δu T F u + o |δ| . Thus, u T Σu Tr Σ = u T F u -u T g 0 g T 0 u Tr F -Tr(g 0 g T 0 ) ≥ u T F u -g 0 2 2 Tr F -g 0 2 2 ≥ u T F u -g 0 2 2 Tr F = u T F u Tr F 1 - g 0 2 2 u T F u ≥ u T F u Tr F 1 -δ -o |δ| ≥ u T F u Tr F e -2δ ,

for δ small enough.

On the other hand, Proposition 4 indicates that e -C F H e C F , which means, ∀u, u T (e C F -H)u ≥ 0

and Tr(He -C F ) ≥ 0.

Thus u T F u Tr F ≥ u T (e -C H)u

Tr(e C H) . Therefore, for λ, u being a positive eigenvalue and the corresponding unit eigenvector of H, we have Figure 6 shows the first 50 iterations of FashionMNIST experiments in main paper. We observe that SGD, GLD 1st eigvec(H), GLD Hessian and GLD leading successfully escape from the sharp minima found by GD, while GLD diag, GLD dynamic, GLD const and GD do not. These experiments are implemented by TensorFlow 1.5.0.

u T F u Tr F ≥ e -2C λ Tr H u T Σu Tr Σ ≥ u T F u Tr F e -2δ ≥ e -2(

B.3 ADDITIONAL EXPERIMENTS ON STANDARD CIFAR-10 AND VGG11:

Dataset Standard CIFAR-10 dataset without data augmentation.

Model Standard VGG11 network without any regularizations including dropout, batch normalization, weight decay, etc. The total number of parameters of this network is 9, 750, 922.

Training details Learning rates η t = 0.05 are fixed for all optimizers, which is tuned for the best generalization performance of GD. The batch size of SGD is m = 100. The noise std of GLD constant is σ = 10 -3 , which is tuned to best. Due to computational limitation, we only conduct experiments on GD, GLD const, GLD dynamic, GLD diag and SGD. Loss Surface The loss surface L(w 1 , w 2 ) is constructed by,

Estimation of Sharpness:

s 1 = w 1 -1 -x 1 , s 2 = w 2 -1 -x 2 ,

(w 1 , w 2 ; x 1 , x 2 ) = min{10(s 1 cos θs 2 sin θ) 2 + 100(s 1 cos θ + s 2 sin θ) 2 , (w 1 -x 1 + 1) 2 + (w 2 -x 2 + 1) 2 },

L(w 1 , w 2 ) = 1 N N k=1

(w 1 , w 2 ; x k 1 , x k 2 ), where θ = 1 4 π, N = 100,

x k ∼ N (0, Σ), Σ = cos θ sin θ sin θ cos θ .

Note that Σ is the inverse of the Hessian of the quadric form generalizeing the sharp minima. And the 3-dimensional plot of the loss surface is shown in Figure 8 .

Hyperparameters All learning rates are equal to 0.005. All dynamics concerned are tuned to share the same expected square norm, 0.01. The number of iteration during one run is 500. Model Network architecture: input ⇒ conv1 ⇒ max_pool ⇒ ReLU ⇒ conv2 ⇒ max_pool ⇒ ReLU ⇒ fc1 ⇒ ReLU ⇒ fc2 ⇒ output. Both two convolutional layers use 5 × 5 kernels with 10 channels and no padding. The number of hidden units between fully connected layers are 50. The total number of parameters of this network are 11, 330.

:

https://github.com/zalandoresearch/fashion-mnist

In fact, in our experiment, we test the equally spacing learning rates in the range [0.01, 0.1], and the final results are consistent with each other.

Figures and Tables:

2-D toy example. Compared dynamics are defined in Table 1, k = 2, σ 2 t is tuned to keep noise of all dynamics sharing same expected squared norm, 0.01. All dynamics are run by 500 iterations with learning rate 0.005. (a) The trajectory of each compared dynamics for escaping from the sharp minimum in one run. (b) Success rate of arriving the flat solution in 100 repeated runs. (c) Tr(HtΣt) of compared dynamics in one run.

Figure 2: One hidden layer neural networks with fixed output layer parameters. The solid and the dotted lines represent the value of Tr(HΣ) and Tr(H Σ), respectively. The number of hidden nodes varies in {20, 200, 2000}, the results of which are denoted in different colors.

(a)  shows the first 400 eigenvalues of Hessian at θ * GD , from which we see that the 140th eigenvalue has already decayed to about 1% of the first eigenvalue. Note that Hessian H ∈ R D×D , D = 11330, thus H around θ * GD approximately meets the ill-conditioning requirement in Proposition 3. Figure3(b) shows the projection coefficient estimated by â = u T 1 Σu1 Tr H λ1 Tr Σ

FashionMNIST experiments. (a) The first 400 eigenvalues of Hessian at θ * GD , the sharp minima found by GD after 3000 iterations. (b) The projection coefficient estimation â = u T 1 Σu 1 Tr H λ 1 Tr Σ , as shown in Proposition 3. (c)Tr(HtΣt) versus Tr(Ht Σt) during SGD optimization initialized from θ * GD , Σt = Tr Σ t D I denotes the isotropic noise with same expected squared norm as SGD noise.

FashionMNIST experiments. Compared dynamics are initialized at θ * GD found by GD, marked by the vertical dashed line in iteration 3000. The learning rate is same for all the compared methods, ηt = 0.07, and batch size m = 20. (a) Training accuracy versus iteration. (b) Test accuracy versus iteration. (c) Expected sharpness versus iteration. Expected sharpness is measured as E ν∼N (0,δ 2 I) L(θ + ν) -L(θ), and δ = 0.01, the expectation is computed by average on 1000 times sampling.

e -2Ht Σ . The last approximation is by Taylor's expansion. A.3 PROOF OF PROPOSITION 3 Proof. Firstly, Tr(HΣ) has the decomposition as Tr(HΣ) = D i=1 λ i u T i Σu i . Secondly, compute Tr(HΣ), Tr(H Σ) respectively, Tr(HΣ) ≥ u T 1 Σu 1 ≥ aλ 1

Figure 6: The fisrt 50 iterations of FashionMNIST experiments in main paper. Compared dynamics are initialized at θ * GD found by GD. The learning rate is same for all the compared methods, ηt = 0.07, and batch size m = 20. (a) Training accuracy versus iteration. (b) Test accuracy versus iteration.

+ ν j ) -L(θ), ν j ∼ N (0, δ 2 I), with M = 100 and δ = 0.01.Experiments Similar experiments are conducted as in main paper for CIFAR-10 and VGG11, as shown in Figure7. The observations and conclusions consist with main paper. These experiments are implemented by PyTorch 0.3.0.

Figure 7: CIFAR-10 experiments. Compared dynamics are initialized at θ * GD found by GD, marked by the vertical dashed line in iteration 3000. The learning rate is same for all the compared methods, ηt = 0.05, and batch size m = 100. (a) Training accuracy versus iteration. (b) Test accuracy versus iteration. (c) Expected sharpness versus iteration. Expected sharpness is measured as E ν∼N (0,δ 2 I) L(θ + ν) -L(θ), and δ = 0.01, the expectation is computed by average on 100 times sampling. All observations consist with main paper.

These experiments are implemented by PyTorch 0.3.0. C.2 FASHIONMNIST WITH CORRUPTED LABELS Dataset Our training set consists of 1200 examples randomly sampled from original FashionM-NIST training set, and we further specify 200 of them with randomly wrong labels. The test set is same as the original FashionMNIST test set.

Compared

Table content: <table><tr><td/><td>Noise t</td><td>Remarks</td></tr><tr><td>SGD</td><td>t ∼ N 0, Σ sgd t</td><td>Σ sgd t is defined as in Eq. (1), and σt = η t √ m</td></tr><tr><td>GLD constant</td><td>t ∼ N (0, I)</td><td>σt is a tunable constant</td></tr><tr><td>GLD dynamic</td><td>t ∼ N (0, I)</td><td>σt is adjusted to make σt t share the same expected</td></tr><tr><td/><td/><td>norm as that of SGD</td></tr><tr><td>GLD diagonal</td><td>t ∼ N 0, diag(Σt)</td><td>The covariance diag(Σt) is the diagonal of the covari-</td></tr><tr><td/><td/><td>ance of SGD noise.</td></tr><tr><td>GLD leading</td><td>t ∼ N 0, Σt</td><td>Σt = k i=1 γiviv T i . γi, vi are the first k leading eigen-</td></tr><tr><td/><td/><td>values and corresponding eigenvalues of the covariance</td></tr><tr><td/><td/><td>of SGD noise, respectively. (A low rank approximation</td></tr><tr><td/><td/><td>of Σ sgd t )</td></tr><tr><td>GLD Hessian</td><td>t ∼ N 0, Ht</td><td>Ht is a low rank approximation of the Hessian matrix</td></tr><tr><td/><td/><td>of loss L(θ) by its the first k leading eigenvalues and</td></tr><tr><td/><td/><td>corresponding eigenvalues.</td></tr></table>

Figure 5: L 2 norm of gradient mean, ∇L(θ t ) , and the expected norm of noise η t E[ T t t ]/m during the training using SGD. The dataset and model are same as the experiments of FashionMNIST in main paper, or as in Section C.2 B.2 THE FIRST 50 ITERATIONS OF FASHIONMNIST EXPERIMENTS IN MAIN PAPER

Table content: <table><tr><td>100</td><td/><td/><td/><td/></tr><tr><td>10 1</td><td/><td/><td/><td/></tr><tr><td>10 2</td><td/><td/><td/><td/></tr><tr><td>10 3</td><td colspan="2">gradient mean norm expected noise norm</td><td/><td/></tr><tr><td>0</td><td>2000</td><td>4000 iteration 6000</td><td>8000</td><td>10000</td></tr><tr><td/><td/><td/><td/><td>C+δ) λ Tr H</td><td>.</td></tr><tr><td>B ADDITIONAL EXPERIMENTS</td><td/><td/><td/><td/></tr><tr><td colspan="3">B.1 DOMINANCE OF NOISE OVER GRADIENT</td><td/><td/></tr></table>

Peer review (iclr2019):
 The paper studies the benefit of an anisotropic gradient covariance matrix in SGD optimization for training deep network in terms of escaping sharp minima (which has been discussed to correlate with poor generalization in recent literature). In order to do so, SGD is studied as a discrete approximation of stochastic differential equation (SDE). To analyze the benefits of anisotropic nature and remove the confounding effect from scale of noise, the scale of noise in the SDE is considered fixed during the analysis. The authors identify the expected loss around a minimum as the efficient of escaping the minimum and show its relation with the hessian and gradient covariance at the minimum. It is then shown that when all the positive eigenvalues of the covariance matrix concentrate along the top eigenvector and this eigenvector is aligned with the top eigenvector of the Hessian of the loss w.r.t. the parameters, SGD is most efficient at escaping sharp minima. These characteristics are analytically shown to hold true for a 1 hidden layer network and experiments are conducted on toy and real datasets to verify the theoretical predictions. Comments: I find the main claim of the paper intuitive-- at any particular minimum, if noise in SGD is more aligned with the direction along which loss surface has a large curvature (thus the minimum is sharp along this direction), SGD will escape this minimum more efficiently. On the other hand, isotropic noise will be wasteful because a sample from isotropic noise distribution may point along flat directions of the loss even though there may exist other directions along which the loss curvature is large. However, I have several concerns which I find difficult to point out because *many equations are not numbered*. 1. In proposition 2, it is assumed under the argument of no loss of generality that both the loss at the minimum L_0=0 and the corresponding theta_0 =0. Can the authors clarify how both can be simultaneously true without any loss of generality? 2. A number of steps in proposition 2 are missing which makes it difficult to verify. When applying Ito*s lemma and taking the integral from 0 to t, it is not mentioned that both sides are also multiplied with the inverse of exp(Ht). 3. In proposition 2, when computing E[L(theta_t)] on page 12, the equalities after line 3 are not clear how they are derived. Please clarify or update the proof with sufficient details. 4. It is mentioned below proposition 2 that the maximum of Tr(H. Sigma) under constraint (6) is achieved when Sigma* = Tr(Sigma). lambda_1 u1.u1^T, where lambda_1 is the top eigenvalue of H. How is lambda_1 a factor in Sigma*? I think Sigma* should be Tr(Sigma). u1.u1^T because this way the sum of eigenvalues of Sigma remains unchanged which is what constraint (6) states. 5. The proof of proposition 5 is highly unclear.Where did the inequality ||g_0(theta)||^2 <= delta.u^TFu + o(|delta|) come from? Also, the inequality right below it involves the assumption that u^Tg_0 g_0u <= ||g_0||^2 and no justification has been provided behind this assumption. Regarding experiments, the toy experiment in section 5.1 is interesting, but it is not mentioned what network architecture is used in this experiment. I found the experiments in section 5.3 and specifically Fig 4 and Fig 7 insightful. I do have a concern regarding this experiment though. In the experiment on FashionMNIST in Fig 4, it can be seen that both SGD and GLD 1st eigvec escapes sharp minimum, and this is coherrent with the theory. However, for the experiment on CIFAR-10 in Fig 7, experiment with GLD 1st eigvec is missing. Can the authors show the result for GLD 1st eigvec on CIFAR-10? I think it is an important verification of the theory and CIFAR-10 is a more realistic dataset compared with FashionMNIST. A few minor points: 1. In the last paragraph of page 3, it is mentioned that the probability of escaping can be controlled by the expected loss around minimum due to Markov*s inequality. This statement is inaccurate. A large expected loss upper bounds the escaping probability, it does not control it. 2. Section 4 is titled "The anisotropic noise of SGD in deep networks", but the sections analyses a 1 hidden layes network. This seems inappropriate. 3. In the conclusion section, it is mentioned that the theory in the paper unifies various existing optimization mentods. Please clarify. Overall, I found the argument of the paper somewhat interesting but I am not fully convinced because of the concerns mentioned above.