Title: HOW NOISE AFFECTS THE HESSIAN SPECTRUM IN OVERPARAMETERIZED NEURAL NETWORKS

Abstract: Stochastic gradient descent (SGD) forms the core optimization method for deep neural networks. While some theoretical progress has been made, it still remains unclear why SGD leads the learning dynamics in overparameterized networks to solutions that generalize well. Here we show that for overparameterized networks with a degenerate valley in their loss landscape, SGD on average decreases the trace of the Hessian of the loss. We also generalize this result to other noise structures and show that isotropic noise in the non-degenerate subspace of the Hessian decreases its determinant. In addition to explaining SGDs role in sculpting the Hessian spectrum, this opens the door to new optimization approaches that may confer better generalization performance. We test our results with experiments on toy models and deep neural networks.

1. INTRODUCTION:

Deep neural networks have achieved remarkable success in the past decade on tasks that were out of reach prior to the era of deep learning. Yet fundamental questions remain regarding the strong performance of over-parameterized models and optimization schemes that typically involve only first-order information, such as stochastic gradient descent (SGD) and its variants.

Regarding generalization, it has been noted that flat minima with small curvature tend to generalize better than sharp minima (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017) . This has been argued by nonvacuous PAC-Bayes bounds (Dziugaite & Roy, 2017) and Bayesian evidence (Smith & Le, 2018) . But why does SGD bias learning towards flat minima? One possible explanation was proposed by Zhang et al. (2018) in terms of energy-entropy competition. Jastrzbski et al. (2018) also suggests that isotropic noise in SGD helps networks escape sharp minima with large Hessian determinant. However, previous theoretical analyses assume that minima are isolated and non-singular. In contrast, Sagun et al. (2017b) finds that most of the eigenvalues of the Hessian of the loss function at a minimum are close to zero, indicating highly degenerate minima. The degeneracy is further supported by results on mode connectedness (Garipov et al., 2018; Draxler et al., 2018) . Furthermore, it is shown that minima found from different initializations and the same optimization scheme are connected with essentially no barriers between them. Nguyen (2019) also shows theoretically that for a class of deep overparameterized neural nets with piecewise linear activation functions, all of the global minima are connected within a unique and potentially very large global valley.

In this paper, we prove that for models whose loss has a "minimal valley" structure, defined below, optimization via SGD decreases the trace of Hessian of the loss, by utilizing recent fluctuationdissipation relations (Yaida, 2019) . Furthermore, we derive the noise covariance matrix that would result in the reduction of other potentially desired quantities such as the Hessian determinant, leading towards the design of new optimization algorithms. We present experiments on toy models and deep neural networks to confirm our predictions.

2. MAIN THEOREM:

In this section, we present our main theorem -how noise during optimization affects the Hessian of the loss function when the loss landscape locally takes the shape of a degenerate valley. More specifically, let N and n be the dimension of the total parameter space and non-degenerate space, respectively, and consider a general loss function L(w) where w ∈ R N are the model parameters. 1 Since the curvature around the minima can vary, we approximate the loss function in such a valley around a degenerate minimum with a modified quadratic form L(w) = L * + 1 2 (w -P(w)) T H(P(w))(w -P(w)) where P is a function that projects a point w in parameter space to the nearest minimum, and H is the Hessian, a function of the location of the projected minimum. We have the following lemma: Lemma 1. For an arbitrary point w and its neighborhood in the valley, there exists an orthogonal transformation Q and a translation vector v such that the loss function in the new coordinate system θ = Qw -v has the following form,

EQUATION

where λ i s are the positive eigenvalues of the loss function Hessian for the non-degenerate space and depend on the position in the degenerate space. Also, note that the gradient descent equation is invariant under this transformation.

A detailed, constructive proof of this lemma can be found in Appendix A.1. Notice that this is a generalization of Morse's Lemma (Callahan, 2010) . In the rest of this section, we denote the nondegenerate and degenerate subspaces by θ = (θ 1 , ..., θ n ) T and θ = (θ n+1 , ..., θ N ) T , respectively. Similarly, the gradients of θ and θ are denoted by ∇ and ∇, respectively. Notice that at a minimum where θ = 0, the λ i s are the only nonzero Hessian eigenvalues.

Next we provide a quick review of the relevant fluctuation-dissipation relation formalism for stochastic gradient descent (Yaida, 2019) . We denote the loss for a random batch B of training sample as L B (θ).

Clearly we have 

EQUATION

EQUATION

We denote the two-point noise matrix as Ci,j (θ

) ≡ [[∂ θi L B (θ)∂ θj L B (θ)]] m.b. and the noise covari- ance matrix C i,j (θ) ≡ Ci,j (θ) -∂ θi L(θ)∂ θj L(θ).

Now we present our main theorem: Theorem 1. When the loss can be locally approximated as in Equation 1, assuming that the nondegenerate space θ is in a stationary state at time t and that the noise covariance matrix is aligned with the Hessian, we have

[[T t+1 -T t ]] m.b. ≤ 0 , where T = n i=1 λ i

is the trace of the Hessian and T t represents the trace at training step t. Equality holds when ∇L B (θ) = ∇L(θ) or ∇T ( θ) = 0.

Theorem 1 indicates that the change of the Hessian trace during SGD optimization is non-positive on average. A trivial condition for equality is optimization with gradient descent (GD) instead of its stochastic variant because the stationary state for noiseless optimization is a constant state with vanishing gradient. An alternate condition for equality is that the trace function of θ reaches a minimum or saddle point.

Before proving Theorem 1, we will first analyze the assumptions made. We emphasize at the outset that we expect our assumptions to be valid only after an early phase of training corresponding to ∼ 10 epochs in our experiments.

2.1. EVIDENCE FOR THE ASSUMPTIONS MADE:

In this subsection, we will analyze the assumptions in Theorem 1 and present some experiments that support their validity.

2.1.1. MINIMAL VALLEY:

Our theory relies on the existence of a degenerate valley in the loss function defined previously via Equation 1. In addition to degeneracy, we furthermore assume that there exists a region of connected degenerate minima in the landscape of the loss function, in particular for overparameterized models. SGD or its variants then selects one of the degenerate solutions. Such degeneracy has been observed empirically where most of the eigenvalues of the loss Hessian from various models and tasks are zero or close to zero (Sagun et al., 2017a; b) . Furthermore, the connectedness is further supported by results of mode connectedness (Garipov et al., 2018; Draxler et al., 2018) . Nguyen (2019) also shows theoretically that for a class of deep overparameterized neural nets with piecewise linear activation functions, all of the global minima are connected within a unique and potentially very large global valley.

Here we present additional experiments to support the existence of such minimal valleys, i.e. basins of attraction of connected, degenerate minima that can be locally described by Equation 1. These experiments will be also used to analyze and confirm our theory later. We consider classification of CIFAR10 with label-smoothing cross entropy (Szegedy et al., 2016) 2 , constant learning rate of 0.1, momentum 0.9, batch size 500 and total training epochs 150. The network architectures are Preact-ResNet18 (He et al., 2016) and VGG16 (Simonyan & Zisserman, 2015) . The training and validation loss are shown in Figure 1 (inset). After 150 training epochs, the networks fluctuates with small training loss, and the minimal training loss for both case is ∼ 0.005.

In order to confirm the existence of minimal valleys, we would like to project each point along the training path to the corresponding closest minimum3 and check whether the projected path crosses any barriers. To determine whether two minima can be connected by a line segment in a minimal valley implying no barriers in between, we use line interpolation to connect these two minima and measure the training loss along the interpolating line. In practice, we call it line connectedness if we sample evenly 10 points in the line segment between consecutive minima and the training loss is smaller than a threshold for all 10 points.

To project a state to a closest minimum, the model is trained from that state with the same hyperparameters except using GD instead of SGD to remove noise. The stopping criterion for these auxiliary training tasks is that the training loss is less than the minimal training loss, which is ∼ 0.005.

Ideally, we would project states after each iteration and check for line connectedness, but this involves large unnecessary calculations. In practice, if state B is obtained from n training iterations after state A, and A and B are not line connected, we insert C which is obtained from n/2 training iterations after A and check for line connectedness between A and C and between C and B. This process stops when there is no extra iteration between A and B or they are line connected. Starting from each pair of consecutive epochs as A and B, we obtain the projected path recursively. The training and validation losses on these two projected paths are shown in Figure 1 (main plot). The left path has 2700 projected points and the right has 4020. We see that after a few epochs, the training loss along the path remains close to zero (and hence minimal), which means that there exists a least one minimal valley connected to the model state found by regular SGD. Also SGD optimization happens within such a valley after the first few epochs. Furthermore, the validation loss varies along the valley.

More experiments can be found in Appendix A.2. Further discussion of violation of this minimal valley assumption can be found in Appendix A.6.

2.1.2. THE HESSIAN-NOISE COVARIANCE ALIGNMENT:

Empirical observations have found that the noise covariance matrix is aligned with the Hessian (Zhang et al., 2018; Zhu et al., 2019; Jastrzbski et al., 2018) . In fact, it was shown in Hoffer et al. (2017) that where S is the mini-batch size and M is the total number of training samples. Considering negative log-likelihood loss functions, which is also the loss function in our experiments for classification, the loss L can be written as

C i,j = 1 S 1 - S M 1 M M k=1 ∂ θi l(x k , θ)∂ θj l(x k , θ) ,

L = -1 M M i=k log f (x k , θ)

where f is the output of the network and

EQUATION

where the left-hand side is the positive Hessian matrix for i, j ≤ n and the first term on the righthand side is the positive matrix proportional to the noise covariance matrix. Empirically, negative eigenvalues are exponentially small in magnitude compared to major positive ones when a model is close to a minimum (Sagun et al., 2017b) and negative eigenvalues can only come from the contribution of the second term on the right-hand side of equation ( 4). Therefore unless the second term is extremely asymmetrical, we can assume that its contribution is small compared to the first term. We will return to this in Section 4.1 and show that eigenvalues with small magnitude arise from higher-order curvature in the degenerate space even when the Hessian with respect to the bottom of the valley is characterized by the positive λ i s for i = 1, ..., n or zero directions. Ignoring the second term on the right hand side, Equation 4becomes

EQUATION

We emphasize that this assumption will only be valid after the early phase of training, not from initialization. Further discussion and experimental support can be found in Appendix A.4.

2.1.3. TIMESCALE SEPARATION:

The final assumption is that the dynamics relaxes to a stationary state in the non-degenerate space θ but not in the degenerate space. This assumption is because there is a timescale separation between the dynamics of θ, which relax quickly and the dynamics of θ, which evolve much more slowly as the minimal valley is traversed.

Considering a simplified optimization problem where parameter directions undergo independent Levy-driven Ornstein-Uhlenbeck processes4 as for quadratic loss with noise, the decay to a stationary state is exponential with decay rate proportional the corresponding eigenvalues (Abdelrazeq et al., 2014) . Since the θ correspond to the leading order eigenvalues, their relaxation time is exponentially shorter than the evolution in the degenerate space. Therefore it is reasonable to assume a stationary state in the non-degenerate space θ when studying the degenerate θ dynamics. Further discussion and experimental support of this assumption using fluctuation-dissipation relations can be found in Appendix A.5.

2.2. PROOF OF THEOREM 1:

Proof. The stochastic gradient descent update for θ is θt+1 = θt -η ∇L B (θ t ). If we consider the average evolution of the eigenvalues λ i for i = 1, ..., n from step t to t + 1, we have

EQUATION

where we performed a Taylor expansion from the second to the third equality and the fourth equality holds from equation (2).

Notice that ∇L(θ) = 1 2 n i=1 θ 2 i ∇λ i ( θ)

. By keeping the leading order term, we have

EQUATION

Considering the observables O( θ) ≡ θ 2 i for i = 1, ..., n, the master equation ( 3) becomes

θ i ∂ θi L = η 2 Ci,i .

Notice that ∂ θi L = θ i λ i ( θ), we thus have

EQUATION

By the definition of the two-point noise matrix C and the noise covariance matrix C, we also have

EQUATION

Based on the assumption that the noise covariance matrix is aligned with the Hessian and scales as

1 S 1 -S M

where S is the mini-batch size and M is the total number of training samples (Hoffer et al., 2017) 5 , we have

EQUATION

Together with Eq. 8, 9 and 10, we have for i ≤ n,

EQUATION

Now Eq. 7 becomes

[[λ i,t+1 -λ i,t ]] m.b. = - η 2 4S 1 - S M n j=1 ∇λ i ( θt ) T ∇λ j ( θt ) + O(η 3 )

Next we consider the evolution of the trace of the Hessian T = n i=1 λ i . We have

[[T t+1 -T t ]] m.b. = - η 2 4S 1 - S M n i,j=1 ∇λ i ( θt ) T ∇λ j ( θt ) + O(η 3 ) ≈ - η 2 4S 1 - S M n i=1 ∇λ i ( θ) 2 ≤ 0 .

Thus we have shown that the trace of Hessian decreases on average during SGD optimization. Notice that equality holds when either one of the following conditions is satisfied: first, if S = M so that SGD becomes full-batch GD, and second that n i=1 ∇λ i ( θ) = 0, which implies ∇T ( θ) = 0.

3. OPTIMIZATION DYNAMICS WITH NOISE BEYOND SGD:

In Section 2, we concluded that, to leading order in the learning rate η, SGD dynamics reduces the Hessian trace. This result originates from the noise introduced by SGD with covariance given by (10).

Interestingly, other forms of noise can be introduced into gradient descent in order to decrease other desired quantities instead of the trace. Here we present a theorem that relates the expected decrease of certain functions of the Hessian spectrum to a corresponding noise covariance. And • represents an average over the stationary state and [[•]] averages the quantity over the corresponding noise.

Theorem 2. In a minimal valley with loss approximation in Equation 1, assuming that there exists a stationary state in the non-degenerate space θ, for any quantity f (λ) satisfying ∂f (λ) ∂λi ≥ 0 for i = 1, ..., n, where λ ≡ (λ 1 , ..., λ n ) T , if the loss is optimized with gradient descent along with external noise with diagonal covariance matrix

EQUATION

we have

[[f t+1 -f t ]] ≤ 0 ,

where f t represents f (λ) evaluated at training step t. Equality holds when ∇f (λ( θ)) = 0.

Proof. Similar to the proof of Theorem 1, we have

EQUATION

)

With the imposed noise with covariance matrix (12), the master equation ( 3) and the relation between noise covariance matrix and two-point noise matrix in equation ( 9), we have

EQUATION

Plugging in Equation 13we have

EQUATION

)

Thus we have shown that f decreases on average during optimization. Equality holds when ∇f (λ( θ)) = 0. Notice that Theorem 1 is a special case with f (λ) = Corollary 2.1 indicates that the non-degenerate determinant will decrease on average during optimization if we introduce isotropic noise in the non-degenerate space. It is known (Smith & Le, 2018) that the Bayesian evidence, minimization of which minimizes a PAC-Bayes generalization bound, has a contribution from the non-degenerate determinant. This contribution is also called the Occam factor. Our result thus leads to a new algorithm where one introduces isotropic noise into the non-degenerate space. Theorem 2 will allow future work to add designed noise that applies an entropic force on other properties of the Hessian yet to be determined.

4. EXTENDED ANALYSIS AND EXPERIMENTS:

4.1 NEGATIVE EIGENVALUES OF THE HESSIAN Empirically, small negative eigenvalues arise even with small training loss, which seems to indicate that the model is at a saddle point instead of a minimum (Sagun et al., 2017b; Ghorbani et al., 2019) . This however is consistent with the loss in Equation 1. The loss is minimal when θ i = 0 for i ≤ n. However, when we use a trained instance to estimate the eigenvalues of the Hessian, we only guarantee that θ i is close to zero (and thus nonzero training loss) for i ≤ n. Therefore, there could exist negative eigenvalues which originate from the second derivative of θ i for i > n. They are originally zero at the minimum with θ i = 0 for i ≤ n, and their magnitude is suppressed by θ 2 i for i ≤ n which is related to the small training loss.

In Appendix A.3, we predict that the negative eigenvalues on average are proportional to the training loss. To confirm this, we set up a simple experiment to classify MNIST data with a small training set of 500 randomly selected samples, with a single hidden layer fully-connected network with 6220 parameters and loss given by label-smoothed cross entropy, again to make a genuine minimum. We first obtain a minimum with gradient descent optimization at learning rate 0.1 and 10k training epochs and use this as initialization for later. The model is then further trained from this initialization with different batch sizes from 5 to 250 and learning rates from 0.001 to 0.1 in order to equilibrate to various training losses. Then the negative eigenvalues are calculated and the result is shown in Fig. 2(a) . The result shows a linear relationship between training loss and the sum of negative eigenvalues, as predicted. 

4.2. THE EVOLUTION OF TRACE AND DETERMINANT WITH A TOY MODEL:

In this section, we use two toy models to test our previous theoretical analysis. In Section 2, we showed that the trace evolution rate is proportional to η 2 . Notice that the two factors of η have different origins. One is from the expansion of λ and the gradient descent step as in Eq. 6. The other contribution is the equilibrium statistics in the non-degenerate directions as in Eq. 8. Furthermore, the coefficient of the proportionality relation also depends on ∇λ. Therefore to test this prediction for how learning rate affects the trace evolution, we train a model multiple times with different ηs. When doing this, the initialization for each run is in equilibrium in the non-degenerate space, and we choose a loss function for which ∇λ is constant. We design a simple four-dimensional toy model with the loss L(θ Next, we measure the evolution of the Hessian trace and determinant after introducing noise with different covariance structure. Our theory predicts that SGD noise decreases the trace while isotropic noise in the non-degenerate subspace decreases the determinant. To test this, we design a toy model where the behavior of these two quantities is anti-correlated. The loss can be written as

) = |θ 3 + θ 4 |θ 2 1 + |θ 3 + 2θ

L(θ) = (4 -2e -θ 2 3 -θ 2 4 )θ 2 1 + e -(θ3-θ4) 2 θ 2 2 .

We train the model with the same initialization and with the two different noise structures. The result is shown in Fig. 2(c)-(d) , consistent with our analysis.

4.3. TRACE EVOLUTION ALONG PROJECTED PATH:

In Section 2, we presented two experiments to support the existence of a minimal valley. Recall that the projected paths follow along the bottom of the valley associated with the SGD dynamics. Our theory predicts that the Hessian trace along the projected paths should decrease during optimization.

To confirm this, we use method in Bai et al. (1996) to estimate the Hessian trace and the results are shown in Figure 3 , in which the values and error bars are calculated by the mean and standard deviation of 10 different initializations. As seen in the figure, the Hessian trace indeed decreases on average in both cases, which is consistent with our theoretical prediction. Notice that test performance does not monotonically improve during optimization. Indeed the trace continues decreasing slightly even when validation loss goes up, indicating that the trace itself is not sufficient to describe generalization.

5. CONCLUSION:

In this paper, we showed that within a minimal valley, the trace of the loss function Hessian decreases on average. We also demonstrated how to design the covariance structure of added noise during optimization to affect other properties of the Hessian's evolution, opening the door to new algorithms. The analysis in this paper can be extended and generalized to models trained on much larger data sets. Sho Yaida. Fluctuation-dissipation relations for stochastic gradient descent. In ICLR, 2019.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In ICLR, 2017.

Yao Zhang, Andrew M. Saxe, Madhu S. Advani, and Alpha A. Lee. Energy-entropy competition and the effectiveness of stochastic gradient descent in machine learning. Molecular Physics, 116, 2018.

Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In ICML, 2019.

A APPENDIX:

A.1 PROOF OF LEMMA 1

To study the dynamics of SGD, we investigate the model behavior starting from θ t to θ t+1 by one-step training and assume that θ t+1 is close to θ t . In this subsection, we derive the property of the loss function in a small neighbourhood, B(θ t ), of θ t , which includes θ t+1 . We consider the model in a minimal valley with a set M of minima and define a projection P that maps any point θ in B(θ t ) to the minimum in M that is closest to θ, i.e.

P(θ) ≡ argmin:

EQUATION

)

Because M is connected by definition of minimal valley, we assume that P is also continuous in B(θ t ). We denote the dimension of whole parameter space and M by N and N -n, respectively. We can then make a quadratic approximation to the loss function in B(θ t ),

EQUATION

where we expand the loss function at the minimum that is closest to θ, and H(P(θ)) is the Hessian depending on the position P(θ) in the valley. L * is a constant by definition of a minimal valley and will be ignored in the following derivations.

We define a minimal path to be a smooth path in a minimal valley where every point on the path is a minimum. For any path passing through a minimum θ * , we have the tangent line of the minimal path at θ * being in the null space of the Hessian at θ * by definition. Precisely, for an arbitrary minimal path θ * (µ) : [0, 1] → M, we have

EQUATION

)

Next we have the following lemma, Lemma 2. Let J(θ) be the N × N Jacobian matrix of function P at θ defined as J ij = dP(θ)i dθj . We have

EQUATION

)

Proof. Notice that θ * (µ) ≡ P(θ + µ∆θ) forms a minimal path through θ where µ ∈ [0, 1] and ∆θ is an arbitrary direction in the parameter space. We have

EQUATION

By Eq. 18, we have

EQUATION

)

Since ∆θ is arbitrary, we have Eq. 19. In order to make calculate of the Hessian and covariance tractable, the networks were chosen to have one hidden layer with 70 nodes, and the dataset is transformed into grey scale and resized to 7 × 7 pixels and normalized. The network with label-smoothed cross-entropy loss is trained with SGD optimization with learning rate 0.001 and batch size 512 and the network with mean-squared loss is trained with SGD optimization with learning rate 0.1 and batch size 512. As shown in Figure 5 , in both cases the cosine similarity between the Hessian and noise covariance is larger than 0.9 after several epochs of initial training. Notice that two random matrices with the same size as H and C in these experiments are almost orthogonal with the cosine similarity of 0.

These experiments are consistent with our assumption of Hessian-noise covariance alignment-after the early phase of training-in the main text.

A.5 EXPERIMENT ON THE TIMESCALE SEPARATION ASSUMPTION:

In this section, we provide support for our assumption of timescale separation through experiments with a small convolutional neural network trained on CIFAR10. In order to observe the relaxation to steady-state, we adopt the method in Yaida (2019), which shows that

θ • ∇L = η(1 + µ) 2(1 -ν) v 2 (FDR)

is satisfied when the model is at stationarity, where µ and ν are the momentum and dampening in SGD optimization, respectively, and v is the velocity. It is easy to show that

P(θ) • P(∇L) = η(1 + µ) 2(1 -ν) P(v) 2 , (FDR')

when there exists a stationary state in a subspace S and P is the projection operator to S. We use O L and O R to denote the quantities of the left and right side respectively of FDR and FDR'. Therefore the quantity | ŌL / ŌR -1| can be used as a criterion to verify how close the model is to a stationary state in a specific subspace.

To show the existence of timescale separation, we train a small CNN on CIFAR10. This network has 3 convolutional layers with kernel size of 3 and output channel number of 15, 20, 10 respectively.

Throughout training, we estimate | ŌL / ŌR -1| on the raw parameter space and top-1 eigenspace of the Hessian. The result is shown in Figure 6 .

Figure 6 shows that the criterion in top-1 eigenvalue decays much faster than in the raw parameter space. After several epochs, | ŌL / ŌR -1| in top-1 subspace is already less then 10 -2 which is a threshold set in Yaida ( 2019) to indicate sufficient closeness to a stationary state, while the state in the whole parameter space is still far from equilibrium after 50 epoch of training. This experiment fully supports our assumption of timescale separation in the main text.

A.6 EXTENSION TO NON-DEGENERATE LANDSCAPE:

This section explores the situation where the minimal valley assumption is slightly violated, in which case the bottom of the landscape is not completely degenerate but has a small slope. The loss function 1 can then be written as

L(θ) = L * + 1 2 n i=1

θ 2 i λ i (θ n+1 , ..., θ N ) + g(θ n+1 , ..., θ N ) , (36) This reveals explicitly the effect of SGD noise as an entropic force during training. The training of gradient descent and its variant drives the model to states with θ = 0 and ∇g( θ) = 0 according to the energy function 36. This exert energetic force to the model moving forward to the lowest energy (smallest loss). However when training with noise, ∇g( θ) = 0 may not be satisfied when the locations with smallest trace and ∇g( θ) = 0 do not coincide because the right hand side of Equation 38is always greater or equal to zero. That being said the model is pushed away from ∇g( θ) = 0 as a entropic competitor to the energetic force.

We use a toy model to visualize this effect. Consider a loss function of L(θ 1 , θ 2 ) = 1 2 (θ 2 2 +1)θ 2 1 +α * θ 2 . We train this function by GD and SGD with α = 10 -4 , learning rate η = 0.01, noise strength 1 S -1 M = 1 and initialization (θ 1 , θ 2 ) = (0.1, 0.1). The result is shown in Figure 7 . The left column corresponds to training with gradient descent where θ 1 drops to zero very fast and θ 2 continues decreasing to negative infinity. As when noise is introduced, θ 2 stops decrease after equilibrium between entropic and energetic force. From 38 it is easy to show that θ 2 = -0.02, which is consistent with the right column of Figure 7 . 

:

Here the non-degenerate space corresponds to the subspace of the Hessian with non-zero eigenvalues. Note that we do not require any specific relationship between n and N , e.g. n N .

The label-smoothing technique is used here to make a true minimum of the cost function exist. Regular cross entropy gives similar results.

We use minimum here but specifically we mean a point with training loss less than a threshold.

Simsekli et al. (2019) proposes to analyze SGD as an SDE driven by a Levy motion.

Refer to Assumption 2.1.2 for details.

Figures and Tables:

where[[•]] m.b. represents the average over all mini-batch realizations. If there exists a stationary state for θ, p ss ( θ) and the stationary-average is defined as O( θ) ≡ d θp ss ( θ)O( θ) where O( θ) is any observable of θ, it is straightforward to derive the following master equation:

Figure 1: CIFAR10 trained on Preact-ResNet18 (left) and VGG16 (right). Main plots show training and validation loss of projected paths. Projected paths are obtained by projecting each model state during optimization to a corresponding state found by GD with small training loss. Each pair of consecutive projected states are then connected by line interpolation. The inset shows training and validation loss of training paths.

have the following corollary: Corollary 2.1 (Determinant-Decreasing Noise). With the same assumptions in Theorem 2, let f (λ) = log n i=1 λ i and C i,i = C where C is a arbitrary positive constant, we have [[Det t+1 -Det t ]] ≤ 0 , where Det ≡ n i=1 λ i is the non-degenerate determinant.

Figure 2: (a) The training loss v.s. the sum of negative eigenvalues. Red is the best fit line, y = wx+b with w = -40.07 and b = -1.22 * 10 -5 . The data is nearly a straight line and the y-intercept is almost zero, consistent with our prediction in Appendix A.3. (b) The trace evolution v.s. the squared learning rate. The result shows a linear relation. Red is best fit line. (c)-(d) Change of the Hessian trace and determinant during training using isotropic (blue) and SGD noise (red).

Figure 3: Hessian Trace estimation of projected paths found in Section 2 with network architecture of Preact-ResNet18 (left) and VGG16 (right). Red bars represent the errors in trace estimation. The decreasing trace during SGD optimization confirms our theory prediction.

4 |θ 2 2 to test this effect. We initialize the model to make sure that |θ 3 + θ 4 | and |θ 3 + 2θ 4 | don't change sign during a few iterations of training. For each η, we first train for 1000 iterations with SGD noise to equilibrate the model and calculate the trace. Then the model is trained for another 1000 iterations and we calculate the trace evolution. The result is shown in Fig 2(b), which shows a linear relation as predicted.

Figure 4: (a) Experiments performed as in the main text with DenseNet on CIFAR100. (b) Experiments performed as described in Section 1.2 using 4-layer ConvNet on MNIST and (c) Preact-ResNet18 on CIFAR10.

Figure 5: The cosine similarity between H and C and training curve of 2-layer networks trained on CIFAR10 with label-smoothing cross-entropy loss (left) and mean-squared loss (right).

Figure 6: | ŌL / ŌR -1| plots along training.The smaller this quantity is, the closer a state is to a equilibrium. This shows that a state in the top-1 eigenspace of the Hessian decays much faster to equilibrium then in the raw parameter space and supports the assumption of timescale separation.

Figure 7: Toy model L(θ 1 , θ 2 ) = 1 2 (θ 2 2 + 1)θ 2 1 + α * θ 2 trained with gradient descent (left) and noisy gradient descent (right).

Peer review (iclr2020):
This paper studies the behavior of SGD after it has reached a steady state and settled in a minimal valley. The authors show that even if SGD has reached a minimal valley, the noisy updates provided by a minibatch of samples result in the trace of the Hessians reducing with further SGD iterations. The authors also show that by changing the type of noise that is added to GD, one can get different functions of the Hessian to reduce with SGD iterations. The theoretical conclusions are then verified empirically with synthetic examples and real deep learning examples In my opinion this is an interesting paper and research direction that helps us understand how SGD biases solutions towards *flatter* minima as measured by the trace norm of the Hessian of the loss function. I have a few concerns though: 1. Is the steady state assumption valid for neural network training? Especially when training on exponential losses (like cross entropy) which drives the parameters towards large norm solutions? This seems to be an important yet unsupported assumption in the analysis. 2. Does the Hessian-noise covariance alignment exist for squared loss functions as well? Is this specific to log-likelihood models? 3. In the experiments to delineate how different types of noise can result in regularization of different quantities, is there a reason only a synthetic example is used? Can this be replicated for deep networks, or are the two quantities (trace vs determinant) closely related? It would be nice to see how this behavior extends to deep networks. Additional questions/comments: 1. While overparameterization seems to be important to the analysis (the parameters move in the degenerate subspace but not in the non-degenerate one), is there anyway one can amend this framework to analyze different levels of overparameterization? Is there any difference in the rates of decrease in the trace norms for highly overparameterized models vs lightly overparameterized ones? 2. This paper talks about how SGD has a bias towards flatter minima. Is there a reason to prefer flatter solutions over sharper ones to get better generalization? Can one prove this connection? 3. In the Densenet experiments in Figure 4a, it seems as though this relationship between flatness and generalization might not hold? There are points along the optimization trajectory where the validation loss seems to be better than at the last few points along the optimization path. However at the former set of points the trace norm of the hessian is larger than at the latter points. Is this from a stray experiment, or are the plots averaged over a number of different runs? Please add these details as well as other details such as learning rates, criteria to decide when steady state was reached, whether or not batch norm was used, etc. 4. The Fluctuation-Dissipation Relations need to be explained more clearly. For people unfamiliar with statistical physics (and Yaida 2019) the notation and formulation is not immediately clear and that makes the paper harder to read. Overall I believe this paper explains an important phenomenon. I am willing to update my score if my concerns are addressed/if there is any misunderstanding.