Year: 2019

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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 MAIN THEOREM
text: 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) wherew ∈ RN 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∗ + 12 (w − P(w))
TH(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,
L(θ) = L∗ + 1
2 n∑ i=1 θ2i λi(θn+1, ..., θN ) , (1)
where λis 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 λis 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 batchB of training sample as LB(θ). Clearly we have [[∇LB(θ)]]m.b. = ∇L(θ) , (2) where [[•]]m.b. represents the average over all mini-batch realizations. If there exists a stationary state for θ̄, pss(θ̄) and the stationary-average is defined as 〈O(θ̄)〉 ≡ ∫ dθ̄pss(θ̄)O(θ̄) where O(θ̄) is any observable of θ̄, it is straightforward to derive the following master equation:
〈O(θ̄)〉 = 〈[[O[θ̄ − η∇̄LB(θ)]]]m.b.〉 . (3)
We denote the two-point noise matrix as C̃i,j(θ) ≡ [[∂θiLB(θ)∂θjLB(θ)]]m.b. and the noise covariance matrix Ci,j(θ) ≡ C̃i,j(θ)− ∂θiL(θ)∂θjL(θ). 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 〈[[Tt+1 − Tt]]m.b.〉 ≤ 0 , where T = ∑n i=1 λi is the trace of the Hessian and Tt represents the trace at training step t. Equality holds when∇LB(θ) = ∇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.

heading: 2.1 EVIDENCE FOR THE ASSUMPTIONS MADE
text: In this subsection, we will analyze the assumptions in Theorem 1 and present some experiments that support their validity.
1Here 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 .

heading: 2.1.1 MINIMAL VALLEY
text: 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 minimum 3 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.

heading: 2.1.2 THE HESSIAN-NOISE COVARIANCE ALIGNMENT
text: 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
Ci,j = 1
S
( 1− S
M
) 1
M M∑ k=1 ∂θi l(xk,θ)∂θj l(xk,θ) ,
2The label-smoothing technique is used here to make a true minimum of the cost function exist. Regular cross entropy gives similar results.
3We use minimum here but specifically we mean a point with training loss less than a threshold.
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 L = − 1M ∑M i=k log f(xk,θ) where f is the output of the network and l(xk,θ) = − log f(xk,θ). Notice that
∂θi∂θjL(θ) = 1
M M∑ k=1 ∂θi l(xk,θ)∂θj l(xk,θ)− 1 M M∑ k=1
1
f(xk,θ) ∂θi∂θjf(xk,θ) , (4)
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 λis for i = 1, ..., n or zero directions. Ignoring the second term on the right hand side, Equation 4 becomes
Ci,j = 1
S
( 1− S
M
) Hi,j . (5)
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.

heading: 2.1.3 TIMESCALE SEPARATION
text: 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.
4Simsekli et al. (2019) proposes to analyze SGD as an SDE driven by a Levy motion.

heading: 2.2 PROOF OF THEOREM 1
text: Proof. The stochastic gradient descent update for θ̂ is θ̂t+1 = θ̂t − η∇̂LB(θt). If we consider the average evolution of the eigenvalues λi for i = 1, ..., n from step t to t + 1, we have
〈[[λi,t+1 − λi,t]]m.b.〉 = 〈[[λi(θ̂t+1)− λi(θ̂t)]]m.b.〉 = 〈[[λi(θ̂t − η∇̂LB(θt))− λi(θ̂t)]]m.b.〉 = −η〈[[∇̂λi(θ̂t)T∇̂LB(θt))]]m.b.〉+O(η2) = −η〈∇̂λi(θ̂t)T∇̂L(θt))〉+O(η2) ,
(6)
where we performed a Taylor expansion from the second to the third equality and the fourth equality holds from equation (2).
Notice that ∇̂L(θ) = 12 ∑n i=1 θ 2 i ∇̂λi(θ̂). By keeping the leading order term, we have
〈[[λi,t+1 − λi,t]]m.b.〉 ' − η
2 n∑ j=1 〈θ2j 〉∇̂λi(θ̂t)T∇̂λj(θ̂t) . (7)
Considering the observables O(θ̄) ≡ θ2i for i = 1, ..., n, the master equation (3) becomes
〈θi∂θiL〉 = η
2 〈C̃i,i〉 .
Notice that ∂θiL = θiλi(θ̂), we thus have
〈θ2i 〉 = η〈C̃i,i〉 2λi(θ̂) . (8)
By the definition of the two-point noise matrix C̃ and the noise covariance matrix C, we also have
C̃i,i = Ci,i + ∂θiL(θ)∂θiL(θ)
= Ci,i + θ 2 i λ 2 i (θ̂)
(9)
Based on the assumption that the noise covariance matrix is aligned with the Hessian and scales as 1 S ( 1− SM ) where S is the mini-batch size and M is the total number of training samples (Hoffer et al., 2017)5, we have
Ci,i = 1
S
( 1− S
M
) λi(θ̂) . (10)
Together with Eq. 8, 9 and 10, we have for i ≤ n,
〈θ2i 〉 = η
S(2− ηλi)
( 1− S
M ) = η
2S
( 1− S
M
) +O(η2)
(11)
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
〈[[Tt+1 − Tt]]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 . 5Refer to Assumption 2.1.2 for details.
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.

heading: 3 OPTIMIZATION DYNAMICS WITH NOISE BEYOND SGD
text: 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
Ci,i = λi ∂f(λ)
∂λi , (12)
we have 〈[[ft+1 − ft]]〉 ≤ 0 ,
where ft represents f(λ) evaluated at training step t. Equality holds when ∇̂f(λ(θ̂)) = 0.
Proof. Similar to the proof of Theorem 1, we have
〈[[ft+1 − ft]]〉 = − η
2 n∑ i=1 ∂f ∂λi n∑ j=1 〈θ2j 〉∇̂λTi ∇̂λj +O(η2) . (13)
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
〈θ2i 〉 = η (2− ηλi) ∂f ∂λi
= η
2
∂f ∂λi +O(η2) .
(14)
Plugging in Equation 13 we have
〈[[ft+1 − ft]]〉 = − η2
4 ∥∥∥ n∑ i=1 ∂f ∂λi ∇̂λi(θ̂) ∥∥∥2 +O(η3) ≈ −η 2
4 ∥∥∥ n∑ i=1 ∇̂f(λ(θ̂)) ∥∥∥2
≤ 0 .
(15)
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(λ) = ∑n i=1 λi.
Furthermore, we 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 Ci,i = C where C is a arbitrary positive constant, we have
〈[[Dett+1 −Dett]]〉 ≤ 0 , where Det ≡ ∏n i=1 λi is the non-degenerate determinant.
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.

heading: 4 EXTENDED ANALYSIS AND EXPERIMENTS

heading: 4.1 NEGATIVE EIGENVALUES OF THE HESSIAN
text: 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 θ2i 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.

heading: 4.2 THE EVOLUTION OF TRACE AND DETERMINANT WITH A TOY MODEL
text: 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(θ) = |θ3 + θ4|θ21 + |θ3 + 2θ4|θ22 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.
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 L(θ) = (4− 2e−θ23−θ24 )θ21 + e−(θ3−θ4) 2
θ22 . 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.

heading: 4.3 TRACE EVOLUTION ALONG PROJECTED PATH
text: 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.

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

heading: A APPENDIX

heading: A.1 PROOF OF LEMMA 1
text: 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 setM of minima and define a projection P that maps any point θ in B(θt) to the minimum inM that is closest to θ, i.e.
P(θ) ≡ argmin θ∗∈M
‖θ − θ∗‖2 . (16)
BecauseM 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 andM by N and N − n, respectively. We can then make a quadratic approximation to the loss function in B(θt),
L(θ) = L∗ + 1
2 (θ − P(θ))TH(P(θ))(θ − P(θ)) , (17)
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
H(θ∗(µ)) dθ∗(µ)
du = 0 . (18)
Next we have the following lemma,
Lemma 2. Let J(θ) be the N ×N Jacobian matrix of function P at θ defined as Jij = dP(θ)idθj . We have
H(P(θ))J(θ) = 0 . (19)
Proof. Notice that θ∗(µ) ≡ P(θ + µ∆θ) forms a minimal path through θ where µ ∈ [0, 1] and ∆θ is an arbitrary direction in the parameter space. We have
dθ∗(µ)
du ∣∣∣ µ=0 = J(θ)∆θ . (20)
By Eq. 18, we have H(P(θ))J(θ)∆θ = 0 . (21)
Since ∆θ is arbitrary, we have Eq. 19.
Next we consider the diagonalization of the Hessian as H(P(θ)) = V (P(θ))TS(P(θ))V (P(θ)) , (22)
where S(P(θ)) = diag(λ1(P(θ)), ..., λn(P(θ)), 0, ..., 0) and λi > 0 for 1 ≤ i ≤ n. As suggested by Gur-Ari et al. (2018), we assume that the eigenspace change is negligible in the small neighbourhood B(θt) and denote constant orthogonal matrix V ≡ V (P(θt)) ≈ V (P(θ)). We then perform the coordinate transformation
φ = V θ , (23) and the loss function becomes
L(φ) = 1
2 n∑ i=1 (φi − P̃(φ)i)2λi(P̃(φ)) , (24)
where P̃(φ) is the same projection map as P(θ) but in the coordinate of φ. Notice here that P̃(φ) is the projection to the minimum with the least L-2 distance to φ, because the coordinate transformation is orthogonal and distance preserved. We have
P̃(φ) = V P(θ) . (25) We have the following theorem, Proposition 3. P̃(φ)i is constant for i ≤ n.
Proof. In φ-coordinate, by Eq. 23 and 25 the Jacobian of P̃(φ), J̃(φ)ij ≡ dP̃(φ)idφj , is related to J(θ) as J̃(φ) = V J(θ)V T . (26) Therefore Lemma 2 becomes S(P(θ))J̃(φ) = 0 . (27) Notice that [S(P(θ))J̃(φ)]ij = ∑n k=1 λkδikJ̃(φ)kj where δik is the Kronecker delta function, we have [S(P(θ))J̃(φ)]ij = λiJ̃(φ)ij , i ≤ n . (28) From Eq. 27 and λi > 0, we have J̃(φ)ij = 0 , i ≤ n . (29)
Notice that J̃(φ)ij ≡ dP̃(φ)idφj = 0 for i ≤ n and all j, we have that P̃(φ)i is constant for i ≤ n.
We denote P̃(φ)i as φ∗i for i ≤ n. Therefore we can define a global translation vector φ∗ = (φ∗1, ..., φ ∗ n, 0, ..., 0) T and consider a new set of coordinate ψ = φ− φ∗ and have
L(ψ) = 1
2 n∑ i=1 ψ2i λi(P̄(ψ)) , (30)
where P̄(ψ) is also the projection to the minimum with the least L-2 distance to ψ. Clearly we have the set of minima of the loss function 30 as {ψ|ψi = 0, i ≤ n} and the closest minimum to ψ in the L-2 distance is ψ∗ = (0, ..., 0, ψn+1, ..., ψN )T. Therefore we have
P̄(ψ) = (0, ..., 0, ψn+1, ..., ψN )T , (31) and
L(ψ) = 1
2 n∑ i=1 ψ2i λi(ψn+1, ..., ψN ) , (32)
Notice that the coordinate transformations from θ to ψ only involve orthogonal rotation and translation, the SGD rule is conservative. To see this, we first have
∇θ = V T∇φ = V T∇ψ . (33) And the SGD on θ becomes
θt+1 = θt − ηV T∇ψL(ψ) ⇒ φt+1 = φt − η∇ψL(ψ) ⇒ ψt+1 = ψt − η∇ψL(ψ) .
(34)
In the main text, we replace the notation ψ with θ for symbol simplicity.

heading: A.2 ADDITIONAL EXPERIMENTS
text: To further support our assumption of connected minima, we conducted extra experiments including DenseNet (Huang et al., 2017) on CIFAR100, and explored the landscape by introducing randomized labels on a subset of the training data to initialize subsequent training on purely clean data (details in Section A.2.2).

heading: A.2.1 DENSENET ON CIFAR100
text: We trained a DenseNet with batch normalization (number of blocks = [6,12,24,16] and growth rate=12) on CIFAR100 and followed the same pipeline as described in the main text experiments. The learning rate is 0.1 and the batch size is 512. The result is shown in Figure 4(a) which supports the connectedness assumption as well as the result of the Hessian trace decreasing under SGD.
A.2.2 4-LAYER CONVNET ON MNIST AND RESNET18 ON CIFAR10 WITH RANDOMIZED LABELS
The setup of these experiments is as follows: we randomly split the training set into two sets, A and B. For MNIST and CIFAR10 we let the number of samples in A be 40000. Then we shuffle the labels in set B. For the rest of this section, we call A the training set and B the shuffle set. The loss minimum, Hessian trace, etc. are all with respective to the training set A.
First we will obtain minima with worse generalization than the minima found by regular SGD. To do so, we train the model with the union of the training set A and shuffle set B until convergence. Then we further train the model with only A using gradient descent so the model reaches a minimum. This minimum corresponds to the left end point of the path in Figure 4 (b) and (c).
We then train the model from this minimum with training set A alone using SGD. We find that the model starts to diffuse and validation loss becomes smaller even though it is at minimum already. Using the same method described in main text, we project this path to minimal and the results are shown in Figure 4 (b) and (c).
In both cases, we found that minima found by SGD are connected, and we further estimated the trace along both paths and found that it is decreasing on average. These results support the connectedness and degeneracy assumption and our result of SGD decreasing Hessian on average. It also illustrates how SGD helps model explore the loss landscape and escape bad minima.

heading: A.3 THE SUM OF NEGATIVE EIGENVALUES
text: In fact, if we consider an equilibrium of θi for i ≤ n, we have 〈L〉 = L∗ + C ∑n i=1 λi and
〈∂j∂kL〉 = C ∑n i=1 ∂j∂kλi for j, k > n from Eq. 11, where C = η 4S (1− S M ). Thus as we change learning rate or batch size, and hence C, 〈∂j∂kL〉 ∝ 〈L〉 −L∗ for states nearby in the minimal valley because both are proportional to C and ∑n i=1 λi and ∂j∂kλi are approximately constant for states nearby in the minimal valley. Usually in deep neural networks training loss can go to nearly zero with a large number of training iterations and small learning rate (Zhang et al., 2017; Du et al., 2019), so we ignore L∗ and have that 〈∂j∂kL〉 ∝ 〈L〉. This result is tested and verified in Figure 2(a).

heading: A.4 EXPERIMENTS ON THE HESSIAN-NOISE COVARIANCE ALIGNMENT ASSUMPTION
text: To further support our assumption of the Hessian-noise covariance alignment, we trained 2-layer neural networks on the CIFAR10 dataset with label-smoothed cross-entropy loss and mean-squared loss (MSE) and calculated the cosine similarity between the Hessian of the loss (H) and the noise covariance (C).
The cosine similarity between two matrices A and B is defined as
CS(A,B) = Tr(ABT )√
Tr(AAT )Tr(BBT ) , (35)
where Tr is the trace operator. It is easy to show that A and B are aligned as A = kB, k > 0 if and only if CS(A,B) = 1.
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.

heading: A.5 EXPERIMENT ON THE TIMESCALE SEPARATION ASSUMPTION
text: 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− ν) 〈v2〉 (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 OL and OR to denote the quantities of the left and right side respectively of FDR and FDR’. Therefore the quantity |ŌL/Ō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 |ŌL/Ō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, |ŌL/Ō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.

heading: A.6 EXTENSION TO NON-DEGENERATE LANDSCAPE
text: 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 θ2i λi(θn+1, ..., θN ) + g(θn+1, ..., θN ) , (36)
where the function g describes how the bottom of the valley varies in the previously degenerate directions. Following the same proof as in Section 2.2, it is easy to obtain
〈[[Tt+1 − Tt]]m.b.〉 ≈ − η2
4S
( 1− S
M )∥∥∥ n∑ i=1 ∇̂λi(θ̂) ∥∥∥2 − η n∑ i=1 ∇̂λi(θ̂)∇̂g(θ̂) . (37)
At the end of training, the model reaches equilibrium, where
〈[[Tt+1 − Tt]]m.b.〉 = 0 .
Thus, from Equation 37 we have
∇̂g(θ̂) n∑ i=1 ∇̂λi(θ̂) = − η 4S ( 1− S M )∥∥∥ n∑ i=1 ∇̂λi(θ̂) ∥∥∥2 (38)
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 38 is 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 ofL(θ1, θ2) = 12 (θ 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.

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