The alignment property of SGD noise and how it helps select flat minima: A stability analysis
Authors: Lei Wu
Mingze Wang
Weijie J Su
Venue: NeurIPS 2022 Accept (2022)
Keywords: Flat minima
stochastic gradient descent
linear stability
implicit regularization
neural network

Abstract: The phenomenon that stochastic gradient descent (SGD) favors flat minima has played a critical role in   understanding the implicit regularization of SGD. In this paper, we provide an explanation of this striking phenomenon by relating the particular noise structure  of SGD to its \emph{linear stability}  (Wu et al., 2018). Specifically, we consider training over-parameterized models with square loss. We prove that if a global minimum $\theta^*$ is linearly stable for SGD, then it must satisfy $\|H(\theta^*)\|_F\leq O(\sqrt{B}/\eta)$, where $\|H(\theta^*)\|_F, B,\eta$ denote the Frobenius norm of Hessian at $\theta^*$, batch size, and learning rate, respectively. Otherwise, SGD will escape from that minimum \emph{exponentially} fast. Hence, for minima accessible to SGD, the sharpness---as measured by the Frobenius norm of the Hessian---is bounded \emph{independently} of the model size and sample size.   The key to obtaining these results is exploiting the particular structure of SGD noise: The noise concentrates in sharp directions of local landscape and  the magnitude is proportional to loss value. This alignment property of SGD noise provably holds for linear networks and random feature models (RFMs), and is empirically verified for nonlinear networks. Moreover, the validity and practical relevance of our theoretical findings are also justified by extensive experiments on CIFAR-10 dataset.

The alignment property of SGD noise and how it helps select flat minima: A stability analysis
Lei Wu
Mingze Wang School of Mathematical Sciences Peking University
School of Mathematical Sciences Peking University
Weijie J. Su
Department of Statistics and Data Science University of Pennsylvania
Abstract
The phenomenon that stochastic gradient descent (SGD) favors flat minima has played a critical role in understanding the implicit regularization of SGD. this paper, we provide explanation of this striking phenomenon by relating the particular noise structure SGD to its linear stability (Wu et al., 2018). Specifically, consider training over-parameterized models with square loss. We prove that if global minimum 8* linearly stable for SGD, then must satisfy |H( Il & O(VB/n). where (H r. B, denote the Frobenius norm of Hessian at 0*, batch size, and learning rate, respectively. Otherwise, SGD will escape from that minimum exponentially fast. Hence, for minima accessible SGD, the sharpness-as measured by the Frobenius norm of the Hessian-Ba is bounded independently of the model size and sample size. The key obtaining these results is exploiting the particular structure of SGD noise: The noise concentrates in sharp directions of local landscape and the magnitude is proportional loss value. This alignment property of SGD noise provably holds for linear networks and random feature models (RFMs), and is empirically verified for nonlinear networks. Moreover, the validity and practical relevance of our theoretical findings are also justified by extensive experiments on CIFAR- 10 dataset.
Introduction
Modern machine learning (ML) models often operated with far more unknown parameters than training examples, a regime referred to as over- pvrrparramtterization.. In this regime, there are many global minima, all of which have training loss but their test performance can be significantly different [47]. Fortunately, it is often observed that SGD converges to those generalizable ones, even without needing any explicit regularizations [49]. This suggests there must exist certain "implicit regularization" mechanism at work [32, 14, 47, 6).
More mysteriously, SGD solutions often generalize better than gradient descent (GD) solutions [21, 38). Therefore, the SGD noise must play a critical role in implicit regularization. The most popular explanation is that SGD favors flatter minima [21] and flatter minima generalize better [16]. This flat-minima principle has been extensively and successfully adopted in practice to tune the hyperparameters of SGD [44, 21] and to design new optimizers [17, 11, 43] for improving generalization. Therefore, understanding how SGD noise biases SGD towards flatter minima is of paramount importance, which is the main focus of this paper.
The works [44, 10, 48] show that SGD noise is highly anisotropic; [30, 41] find the magnitude of SGD noise to be loss dependent. Both structures are shown to be critical for SGD picking flat minima.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

However, these works all make unrealistic (even wrong) over-simplifications of SGD noise (see the related work section for more details) in their analysis. In addition, instead of studying SGD, they all consider the continuous-ting
stochastic differential equation (SDE), which is good modeling of SGD only in finite time and small learning rate (LR) regime [25]. It is generally unclear how this SDE modeling is relevant for understanding SGD with large L---a regime prefered in practice. Consequently, these works only provide intuitive and empirical analyses, lacking a quantitative characterization of when and how SGD favors flat minima.
Another line of works [46, 29] relate the selection bias of SGD to the dynamical stability. In over- parameterized case, all global minima are fixed points of SGD but their dynamical stabilities can be very different. At unstable minima, small perturbation will drive SGD to leave away, whereas, for stable minima, SGD can stay around and even converge back after initial perturbations. Thus SGD prefers stable minima over unstable ones. Specifically, [46, 29] analyze the linear stability [1] of SGD, showing that linearly stable minimum must be flat and uniform. Different from SDE-based analysis, this stability- -based analysis is relevant for large-LR SGD and even empirically accurate in predicting the properties of minima selected by SGD [46, 20, 8].
In this work, we follow the linear stability analysis in [46, 29] but take the particular structure of SGD noise into consideration. We establish a direct connection between linear stability and flatness, which allows us obtain a quantitative characterization of how the learning rate and batch size affect the flatness of minima accessible to SGD. In contrast, [46, 29] have to introduce another quantity: the non-uniformity together with flatness characterize linear stability because of neglecting the noise structure.
Setup Let ((1;, y:) Dim, with x; E Rd V/ E R be the training and f(;; 0) with @ E R be model. The model size is defined to be the number of parameters and in this paper. For simplicity, we will always assume K = and the extension to the case of K 1 is straightforward. Let L.(0)- = 1 (z;;; 0)- - yi be the fitting error at the i-th sample and L(0)- π. -i= (0) be the empirical risk. We shall focus the over-parameterized case in the sense that inf. L(0) To
minimize L(-), we consider the mini-batch SGD:
0.t1 0:
" B
Σ (O,) iEL
(1)
where 7 and are the learning rate and batch size, respectively. This SGD can be rewritten as 0.+1 0, (VI (O, +&,, where & is the noise, satisfying B&G, = 0 and B[c,{]] = E(0,) Here the noise covariance (0) = A, Σ (L) (L))) To characterize the local geometry of loss landscape, we consider the Fisher matrix: G(Ø) = Σ fz V) 8)T and the Hessian matrix: H
G(0) + 1 - (f -) v² V f(z; 0). When the loss value is small, (0) x G(O) and in particular, if 0+ is an global minimum, 11 (1)) G(O*).
Notations. For vector a, let ||a|| = v and [a]lw Va"Wa. For matrix A, denote by {A)} the eigenvalue of A in a decreasing order. For other notations, we refer to Appendix C.
Our main contributions are summarized as follows.
We first show that for many ML models, the SGD noise is geometry aware: 1) the noise magnitude is proportional to the loss value; 2) the noise covariance aligns well with the Fisher matrix. Specifically, to quantify the alignment strength, we define a loss-scaled alignment factor p))), which is proved to be bounded from below, i.e., there exists size- independent positive constant μio such that ((0)= 2 Poa for linear networks (Proposition 2.3) and RFMs (Proposition 2.5), and is also empirically justified for nonlinear networks. Moreover, we identify that it is the uniformity of model gradient norms [ () } that accounts for this alignment property of SGD noise.
We then provide thorough analysis of the linear stability of SGD by exploiting the alignment property of noise. We prove in Theorem 3.3 that if a global minimum 0* is linearly stable, then [[H(()))) ]| & n 1 /BB/ μo- Here the constant μ4 quantifies the alignment strength of SGD noise. Hence, for minima accessible to SGD, the Hessian's 's Frobenius norm the flatness perceived by SGD-is bounded independently of the model size and sample size. Moreover, if minimum is too sharp, violating the preceding stability condition, SGD will escape from it exponentially fast (Theorem 3.5). Together, we obtain a quantitative characterization of when and how much SGD dislikes sharp minima.
2

---

Our theoretical findings are also corroborated with well-designed experiments on a variety of models including linear networks. RFMs, convolutional networks, and fully-connected networks. In particular, the practical relevance is demonstrated in Section 4 by extensive experiments on classifying full CIFAR-10 dataset with VGG nets and ResNets.
1.1 Related work
Noise structures. [52, 19, 27] consider the Hessian- based approximation: Σ(0) = (0), where is a small constant. [53] proposes an improved version: Σ(0) R (0) (0). But these approxima- tions in general cannot be accurate since Hessian is not semi-positive definite (SPD) in non-convex region. More recently, [30] and [41] study SGD by assuming ()) = -2) ( where 0* is a
minimum of interest, and [() o² Ip)) respectively. These assumptions completely ignore the state-dependence of noise direction. In contrast, we assume (() = (L)))((0) with C(0) having nontrivial alignment with the Fisher matrix G(O), which does not impose any explicit structural assumption on ((O). As a result, our assumption is much weaker and can be rigorously justified for popular ML models both theoretically and empirically. More importantly, we show that this weak alignment property is sufficient for analyzing the linear stability of SGD. We anticipate that our alignment assumption can be also adopted to analyze other properties of SGD.
Escape from sharp minima. The escape behavior of SGD was first studied in [52, 46), as an indicator of how much SGD dislikes sharp minima. One of the most mysterious observation is that the escape happens in an unreasonably efficient way. However, the theoretical analysis there assumes the noise to be slate-independent, and consequently, the derived escape time depends polynomially on the loss barrier. Later [48, 30] attempt to study this issue using the classical diffusion- framework [12] (ItÔ-SDE), which cannot explain the unreasonable escape efficiency at all since the resulting escape rates depend on the loss barrier exponentially. See also [30] for an improved analysis. [37, 51] argues that the SGD noise is teavy-tailed and thus SGD should be modeled as Lévy-SDE instead of ItÓ-SDE. Moreover, it is shown that the heavy-tailedness can ensure the escape rate depends on the basin volume instead of the loss barrier. Unfortunately, the volume in high dimensions always scales with the dimension exponentially and consequently, this does not explain the escape efficiency in high dimensions. Moreover, whether SGD noise is really heavy-tailed and whether the heavy-tailedness is really important for generalization are still debatable for neural networks [40, 26]. In contrast, we show that the unreasonable escape efficiency comes from the particular geometry- awaare structure of SGD noise, regardless of whether the noise is heavy- or light-tailed.
Flatness metries In the literature, a variety of flatness metrics have been adopted, such as the largest eigenvalue of Hessian [21], the trace of Hessian [9, 6], the basin volume [5]], and the ones scaled by parameter norms [28, 39] in order to achieve the scaling invariance for ReLU nets. These metrics are proposed for either computation easiness or bounding generalization gaps. It is unclear if they are perceivable to SGD. let alone how the boundedness of them depends on the batch size, learning rate, as well as the model size and sample size. We show that for SGD solutions, the Frobenius norm of Hessian-a flatness perceived by SGD through the linear stability-is bounded by size-independent quantity. Note that a similar stability argument also applies GD but only yielding the boundedness of the largest eigenvalue of Hessian [46, 31].
Lastly, we particularly mention the work [33], which provides fine-grained analysis of the implicit bias of training two-layer diagonal linear networks. This work is related to ours since we both consider the magnitude and direction structure of SGD noise simultaneously. However, the analysis in [33] is limited to the specific toy model but ours is relevant for general models.
2 The alignment property of SGD noise
Since (L; (f(z;;0) v1 V. (z,; 0), we have the following intuitive approximation [30]:
2
( 2
E(0)
n d=1
(8) 2
L. (0))
V
= G()))
(2)
3

---

where (i) assumes that the full-batch gradient ( to be negligible compared with the sample gradients {0); (ti)
assumes that {Vf(z;;)) () and {0) are nearly decoupled. This approximation cannot be true in general but tells us that 1) The noise magnitude is proportional to the loss value; 2) the noise covariance aligns with the Fisher matrix.
Motivated by (2), we define al0) [ E 3(0) - [[ Here (x0) quantifies the similarity between Σ(A) and G(0). characterizing how much the noise concentrates in sharp directions of local landscape. ((0) characterizes the relative degenereccyy of noise (with respect to the loss value). Then we define the loss-scaled alignment factor:
F
Tr()(0)G(0)) m(0) ((0) = (L)))
which characterizes the direction and magnitude alignment simultaneously. Intuitively speaking, if p(0) is bounded below, SGD noise is non-degenerate in sharp directions of local landscape. In particular, a(+) B(0) = p(0) if the approximation (2) holds.
We say the noise satisfies the z-alignment if p(0)= > 0. Compared with the decoupling approximation (2), the y-alignment is much weaker condition. Note that this specific weak quantification of alignment is inspired for analyzing the linear stability of SGD, which is the focus of this paper. Specifically, Theorem 3.3 shows that u(0) along with the Frobenius norm of Hessian determines the linear stability of SGD. One may define other metrics to quantify alignment strength for studying other properties of SGD, but this beyond the scope of this paper.
A relaxed alignment. Let (0) E VI, E() = Σ1 (0) E²) It is often believed that the full-batch gradient V is relatively small compared to the sample gradients { [VI,);- As result, the influence of (²) should be negligible compared to (0). To disentangle the influences of them, we define
E2( ()(())V..()). Then
(O),(())(O)) #i (0) 2L(0) GOO
(2 Tr (O)(0)) 2000)
Then ((0)= - (0) #2 (0). Our linear stability analysis in Section show that μn (0) > Ln > 0, condition we refer to μ alignment, is sufficient to ensure that SGD only selects flat minima.
2.1 Why does the alignment property hold?
Definition 2.1 (Norm uniformity of model gradients). Let 9:(0) V fm;;0) xi(0) = x:(l)
(l.. (0) = of (O)), x(0) 1E-ixxi0 Define (()) == min.e
X(O)
The quantity (0) measures the uniformity of model gradient norms and this property can guarantee the μ aillligmment as stated below.
Lemma 2.2. μ1 (0) > (()
Proof. Noticing x(0) = 1 E. (()))(()))) Tr (G(O) have
(CC ((0)),))))))))))))) = [G)
(O)(()))()))) (,) (r(
2 L.(0)x 211()))
(.)))x >
Thus (1) = ( > (0).
The above proof suggests that the "decoupling" approximation in (2) holds in a weak sense }, are uniform. One can apply similar argument by assuming the uniformity of the fitting errors (L)) which, unfortunately, we find never hold in practice. In contrast, we will show that the norm uniformity of model gradients provably holds for linear networks and RFMs, and can be empirically justified for nonlinear networks.

---

Overmttrreeerrize linear models. Consider an over-parameterized linear model (OLM): f(;;0) F(0)²z, where F Ω + Rd denotes a general re- re-arammterization function. Note that f(: only represents linear functions but the corresponding landscape can be highly non- convex. Typical examples include the linear network: F(0)= = = W2...... ... Wr and the diagonal linear network: F(0) = (o] ff...... .. (82) Both have attracted lot of attention in analyzing the implicit bias of GD and SGD [3, 42, 33, 13, 4]. The following proposition provides a precise characterization of the noise covariance for OLM models, whose proof is deferred to Appendix D.
Proposition 2.3. Denote by S)) the Gaussian distribution with mean zero and covariance matrix S. Suppose f(-: 8) is a general OLM and x ~ NN(0, S). Consider the online SGD setting, i.e., n= 00. Then, Σ(() VL(O) + (())) and μm (0) 2 p(0) >
This proposition shows that the alignment property holds in the entire parameter space and moreover, the alignment strength is independent of model size. Here the alignment is only proved for the infinite-sample case. Similar results should hold for finite-sample cases by concentration inequalities as long as n is relatively large, but this straightforward extension does not bring any new insights. It is more interesting to consider the low-sample regime (i.e., d), where the alignments indeed hold (at least in typical regions explored by SGD) as demonstrated empirically in Figure 1. In addition, the above proposition provides closed-form expression of the noise covariance, which might be useful for analyzing other properties of SGD beyond the linear stability. A comprehensive analysis of these issues left to future work.
Feature-based models. Consider a feature-based model ((x;0) = - 019)(x) = (0, (())) In this case, the model gradients: 9: = g)) = (xx;) and the Hessian and Fisher matrix: G = H _ uniformity of model gradients also becomes constant: 7 = min, (=/(, Ex= Xi) with X = l1
n E 9:9 all are constant. But the noise covariance Σ(0) is still state-dep The norm
Lemma 2.4. (0) > 7: (0) < T(G) = = AR(G)/ E, X2 and p()) > 7 (G).
The above lemma suggests that the (2( term is negligible as long as the Fisher matrix is not nearly rank-1. By bounding the 7 and r(G), we can prove the μ11 and y-allinmmet for random ReLU feature models as stated in the following proposition. The proof is presented in Appendix F, where similar results for general RFMs are provided (see Proposition F.7).
Proposition 2.5. (x) = ReL.U(w] x) with wi iid Unif (vis(V/S"")) and x ~ Unif(Sa-1). ). Then, for any 5 E (0,1), if m 2 n Z d log(1/5) (5)), then w.p. at least - 8, (n( 2 1 and p(0) 2 d-1
Although feature-basse
models are linear, their analysis is still applicable to understand nonlinear models, as long as the nonlinear model locally behaves like the linearized one: fün (z; 0) := {(; 0*)+ (e 9", V fa 8*)) with fz 8*) learned from data. Hence, Proposition 2.5 can explain why the alignment property holds in a local region around global minima. Note that this is sufficient for characterizing the linear stability of SGD, which is a local property in nature.
2.2 Empirical validations
Figure Ia reports the values of a ((,,). (. 3(0),, (a((t) during the SGD training of four types of models, including the lincar networks and RFMs analyzed above, and fully-connected networks (FCN) and convolutional neural networks (CNN). First, one see that a(,,) I's are quite close to 1 during the entire training, suggesting the strong concentration of SGD noise in sharp directions of local landscape. Second, ((,)'ss 's keep bounded below, implying the noise magnitudes with respect to the training loss. As a result, ju(t))'s 's are significantly positive for all models examined. In particular, for the linear networks, the alignment holds in the law-sa regime, where < d.
sufficiently large
The size independence. Figure 1b further examines how the extent of over- parammettttriiinoonnn affects the alignment strength. One can see clearly that for linear networks and RFMs, ((f) 's are independent of the model size, which confirms our theoretical findings proved above. In addition, we also observe that for nonlinear networks, the alignment strength is also (nearly) independent of the model size. For instance, for CNN, the value of p(0) only decreases from around 1.05 to 1.0 as the model size increased by more than two orders of magnitude.
5

---

RFM
Linear
JCN
CNN
W W m
16-
1.2
144
200
12-
10.
700
T00
1903
2000
1000
10000
50660
10000
10000
of steps
stogge
steeps
RFM
PCN
1.03
1.12
1.3
1.9
3.55
1551
255
Logsy (model size)
Ingnddoode sioe)
Imode soo)
Imggee sirej
(b)
Figure 1: The alignment property of SGD noise. Four types of models, including the RFM, linear network, fully-connected network (PCN), and convolutional neural network (CNN), examined. We refer Appendix for the experimental setup. Note that the linear network trained in low-sample regime (n = 100, 50). (a) The alignment factors during training. (b) How the alignment strength of convergent solution changes with the over parameterization. The error bar corresponds to the standard deviation over 5 runs.
The norm uniformity of model gradients. Figure 2 shows the norm uniformity of model gradients, where we report the values of (0) the entire SGD trajectory. On can see that (0) indeed bounded below, implying that the SGD noise satisfies the /u m1-alignment according to Lemma 2.2.
RFM
PCN
CNN
0994
0:55
6.:22
1.54
wmmamm
300h
1.30
172
stogpe:
strep:
The norm uniformity of model gradients. The values of (0,)) in the entire SGD trajectory reported. is shown that the norms of model gradient norms are uniform during the entire training process.
Figure 2:
Note that in experiments, we only show that the alignment property is satisfied in typical regions explored by SGD, including the random initialization and the convergent region. In contrast, for OLMs and RFMs, we in fact prove the alignment property for the entire parameter space.
3 The linear stability analysis
Let 6* be a global minimum of L(.). When 0, is close to 8", the local dynamical behavior of SGD can be characterized by linearizing the dynamics around 0":
0,41 0,
! Σ (000000 0*),
B
VEL.
where (²I,+00)= (x;;; (2;; 0T This corresponds to the local quadratic approximation of the loss L(.) or the local linearization of the model around 6*:
fun (a; 6) (x; 8*) (0 0", vz(;0'))
(5)
Specifically, (4) is exactly the SGD of training the linearized model (5).
Definition 3.1 (Linear stability). A global minimum is said to be linearly stable if there exists a
C 0 such that it holds for the linearized dynamics (4) that [() & C')[ Vt > 0.
well-known dynamical system that the local behavior of the original nonlinear dynamics can be characterized by the linearized one if the local quadratic approximation is non-degenerate. However, over-parameterizer case, the local quadratic approximation is degenerate in flat directions. Consequently, one may be concerned about the relevance of local quadratic approximation and the

---

resulting linear stability analysis. Fortumately, the stability in Definition 1.1 is particularly measured with the change of loss value. Thus the instability mostly comes from noise perturbations in sharp directions and the alignment property guarantees that noise mostly concentrates in sharp directions. sharp directions, the local
Consequently, the flat directions contribute little to the instability. quadratic approximation is always valid, thereby explaining the relevance of linear stability analysis. The rigorous formulation of this intuition left to future work and we instead resort to numerical experiments to demonstrate the validity in this paper.
For simplicity. we will use 0g to denote de: let 0 = =0 and a- (x For the linearized model fin (-) we have L(0) .n Ei=1 18T 9.12 30" HO, G = 9:9." where we omit the dependence on 8 for simplicity. Note that the Fisher and Hessian matrix are constant but the noise covariance E(0) D lo f g. H000 H is still state-dependent.
Before considering the specific linearized SGD (4), we first have a general result.
Lemma 3.2. Consider a general SGD: 0.+1 = 0: - ((,)) + for the linearized model (5), 0, [3 S(O.). Then we have
where (5:)1a1 are any noises satisfying E[
[[L(0++1)] (lr(0)) +772v(0))]
(6)
where v(0) = (r((HS(0))22 and r(0) > 0. Moreover in 2/))(H,, then r(0) & 1.
Proof.
Using the fact R[Gr] = and [ = S(0,), we have
E[L(Or+1)] Ei (8: --- mET) H - -VL nnt)
=E(L(6,) - 7)V)())* H8; + 12
(L(0)) + r22
E(T((H((())) (7)
2
((rr (22((0)))
where r(0) = 1 - 2n 0" H2
+ +1²
26" 000 0 HT
since ((()) H0. By Lemma G.2, r(0) > 0 and if
ge
7 < (H), then r(0) 1.
The two terms ((,) and (0,) denote the contributions from the full-batch gradient (L(O.) and the noise εt, respectively. The stability is affected by both terms simultaneously. It well-known that if 6 is linearly stable for GD, then d1 (H"))) 2/n (see, e.g., [46, 31J). This also holds for SGD but SGD imposes a stricter condition because of the extra ((,)( term. Specifically,
[[ = (()()))((O))) + +),2((0) > x E[)((:) 0.5r/ Ir(H)(()).
(8)
Therefore, the more S(Ø) aligns with H. the more unstable that minimum is. Specifically, let S(Ø) 20 (()) Then, E[L(O++.)] > o² ((L(O,) Tr(H(((,)) Thus to ensure a stable convergence, we should roughly have Tr(HC(0)) < गे We next show that this can lead to a flatness control by utilizing the alignment between C(O) and H.
3.1 The linear stability imposes size-independent flatness constraints
For the mini-batch SGD, the following theorem characterizes how the batch size and learning rate affect the flatness as measured by the Frobenius norm of Hessian of minima accessible to SGD. Theorem 3.3. Let 8" be global minimum that is linearly stable. Denote by p(0) the alignment factors for the linearized SGD (4) and model (5), If ((0) Na, then (H(())))
)l & n ng
Proof. By (8) and the definition of μ(0), we have
B[L(0++1)] > 2)²
r | Blu(0,)(00)) > Han |
2B E[TT(HE(0M)
>
BLL(O))
B
B
To ensure the stability, we must have pon? H H BBB... < 1, leading to |||| & B/yo/n-
We have shown in Section 2 that μo is (nearly) size-independent, and thus the obtained upper bound of flatness also (nearly) size-independent. As comparison, for GD, the linear stability can only ensure Amaz (H (0)) 2/n. This gives bound of the Hessian's Frobenius norm: [0"))" & 2.//n.
7

---

depending on the model size explicitly. The comparison of two bounds partially explains why SGD tends select flatter minima than GD.
We show below that the j-alignment can be further relaxed to the p1-aligoment. The proof is similar to the one of Theorem 3.3 and deferred to Appendix G.
Proposition 3.4. Under the setting of Thhorrem 3.3, if the noise of linearized SGD satisfies jor (0) 2 ja,
then [0)) & min (B
2B
1
1)m1
M41
When B >> 1, the bound becomes B/ f/ of p-aliggmment Thus the influence of VI. is indeed negligible compared with [VLz).
/ (B 1)m1 & v BB/ui/n, n which is the same as the case
Note that the linear stability is local in nature and hence our analysis essentially only needs the Mi Kl-aligmment to hold locally around minima of interest. Experiments in Figure 2 shows that x(0)) is always bounded below, i.c., the norm uniformity of model gradients is satisfied. Combining with conclude that the alignment property assumed in Proposition 3.4 holds.
Proposition F.7,
Tightness our analysis. In the analysis above, we only inspect the instability caused by the noise, with the full-batch gradient completely ignored. Therefore, we anticipate that our bound is tighter in small-batch regime, where the noise term dominates the full-batch term. We will see that our numerical experiments indeed confirm this tightness in small-batch regime. However, for obtaining the tightest bound, one may need to consider both components simultancously; this is much more complicated and left to future work.
Numerical validations. Figure 3a numerically shows how the Frobenius norm of Hessian (not only the upper bound) changes with extent of porr--arammterzztionn where the trace of Hessian is also plotted for comparison. One can see that the Frobenius norm indeed keeps almost unchanged as increasing the model size but the Hessian trace increases significantly. Figure 3b further shows the ratio between the Frobenius norm and our upper bound the training process and two batch sizes examined. We have two observations. First, the correctness of our bound holds for the entire SGD trajectory, suggesting that the linear stability analysis relevant for the entire training process. Second, as expected, the theoretical bound is indeed tighter for the case with smaller batch size.
Linear De6
FCN
CNN
FCN
(:: |
[
|
D 40
02
1060
fogddd simee)
Iogudtmorkd sixe)
logaftoode sioe)
# of steps
(a)
(b)
Figure 3: (a) The Frobenius norm and trace Hessian vs. model size. The error bar corresponds to standard deviation estimated over 5 runs. (b) The ratio between the Frobenius norm and our upper bound in the training process. Here B = 5 and B = 40 are examined. For more experiment details, we refer to Appendix A.
3.2 SGD escapes from sharp minima exponentially fast
The following theorem shows that the pure noise-driven escape from a sharp minimum is exponentially fast, whose proof follows trivially from (9).
Theorem 3.5. Under the setting of Theorem 3.3, if (H) > V no (1 1. satisfies [[L().)] > n6 [[L( with 7o "
then the linearized SGD
v'Ho
Hence, linearized SGD takes roughly log (1/c) steps to escape from a O(e)-loss region to a O(1) loss region. The escape time depends on the loss barrier only logarithmically and is independent of the parameter space dimension. Due to the local closeness between linearized SGD and the original SGD, this partially explains the unreasonable escape efficiency of SGD for training big models. In contrast, the escape rates of existing works [48, 52, 51, 30] are either exponentially slow with respect to the loss barrier or suffer from the curse of dimensionality.

---

Figure 4 shows the trajectories of SGD escaping from sharp minima. It is demonstrated that the escape indeed exponentially fast and specifically, 10 steps are enough for SGD escaping to a high-loss region for all the models examined. In addition, we observe that the escape is still exponentially fast in the high-loss region, although our analysis only applies to a local region. How can we explain this nonlocal escape behavior? We leave the study of this interesting phenomenon to future work.
CNN
E
Ediki
ESAil
Look 2.08)
Loos: 2(00)
Locs 2.09)
10-3
Step:
Step:
Step:
Step:
Figure 4: The exponentially fast escape from sharp minima. The blue curves 200 trajectories of SGD; the curve corresponds to the average. The sharp minimum is found by GD. When GD nearly converge, we switch to SGD with the same learning rate. This choice ensures that the minimum is stable for GD, and thus the escape is purely driven by SGD noise. For more experimental details, we refer to Appendix A.
3.3 The importance of the noise structure
The magnitude structure. Theorem 3.5 together with its proof suggests that the loss dependence of noise magnitude is critical for obtaining the exponentially fast escape. The intuition is follows. When 0, is perturbed by noise to 8(+1 where ((f++1) > L(Oe). ). the noise magnitude becomes larger there and thus 0z+1 is easier be perturbed SGD to leave exponentially fast, On the contrary, the following lemma shows that the noise is uniformly bounded, the noise-driven escape is at most linear in time.
larger- region. This positive feedback drives
Lemma 3.6. Under the setting of Lemma 3.2, assume n & (H) and (H)) < 20 Then ([)(O.) L(0o)] & o/²
We set 17 (I) to avoid the exponential escape caused by the full-batch component.
Proof. By Lemma G.2, when 7) & (H), r(0) < 1. Thus Lemma 3.2 implies B[L(0++1)) ) & ([L( + 1700, which implies [E[((,)) < [[ + 1
The direction structure. We now turn to consider the impact of direction structure. Consider general SGDs: f.+1 = 8, ( ((.))) + +.) with E[Er S(0.) B for the linearized model (5).
We compare two type of (unrealistic) noises:
Geomelry-avare noise: (0) 2L(O)H.
Isotropic noise: S2(0) = 20 IO) with o² = Tr(H)/p. Here, the value of a² chosen ensure two types of noises have the same total variance for fair comparison [52].
Note that p denotes the model size. Analogous to Theorem 3.3, for the second isotropic SGD,
n² Tr(HS(O)) > B[L(O.))
o²n² B Tr(H) B[L(O.) n²² Tr(H)²
B[L(O-1)))
2B
pB
Hence, the instability decreases with the parameter-space dimension and the resulting flatness constraint is (0")) & vpB/nn depending on the model size explicitly. In contrast, for the fist noise, Theorem 3.3 implies [0r) < vB/n, independent of the model size. This difference can be intuitively explained follows. The isotropic noise wastes most energy in perturbing SGD along flat directions, which barely affects the instability. In contrast, the geometry- aware noise focuses most energy perturbing SGD along sharp directions, causing much more instability.
4 Larger-scale experiments
We have provided small-scale experiments to justify the validity of our theoretical findings for a variety of ML models. Here we turn to demonstrate the practical relevance by consider the classification of the CIFAR-10 dataset [22] with VGG nets [36] and ResNets [15]. In training, all explicit regularizations are removed to keep consistent with our theoretical analysis. More details of the experimental setup can be found in Appendix A.

---

The alignment property and escaping behavior. Figure 5a reports the alignment strength of SGD noise during training for VGG-19 and ResNet-110. One can see that the alignment factors are significantly positive and similar results are also observed for variety of VGG nets and ResNets of different depths and can be found in Figure of Appendix B. On can see that the alignment strength is nearly independent of the model size. Figure 5b shows the behavior of escaping from sharp minima for VGG- and ResNet-110. One can still observe that the escape is exponentially fast and similar observation for other ResNets and VGG nets can be found in Figure 9 in the appendix. These observations suggest that our theoretical findings also hold for this practical setting.
VGC-19 AtR
ResNet: 110
ResNet-110
12
18-1
#]10k]]
AIR?
Loos: L(A))
19-
Loos: L(Og))
Mr
Aww
200 of epochs
400
30
1000
of epochs
Step:
Step:
(h)
Figure 5: (a) The alignment factors and the escaping hehavior for VGG-19. Similar results are also observed for all examined ResNets and VGG nets, which can be found in Appendix B.2. (h) The actual flatness of SGD solution and the corresponding theoretical upper bound (Theorem 3.3). (c) The upper bound becomes tighter as decreasing the batch size.
The flatness and upper bound. Figure 6a reports the flatness of convergent solution and the corresponding upper bound predicted by our theory for ResNets and VGG nets. is again observed that the flatness is nearly independent of the model size. A surpisiin observation in Figure 6a is that our upper bounds are rather tight, see, e.g., VGG-16 and VGG-19. This tightness suggests that SGD runs (nearly) the edge of stability [46, 8]. Moreover, as expected, Figure 6b shows that the our bound becomes tighter as decreasing batch size. This is consistent with what we observe in small-scale experiment in Figure 3b.
ResNet
VGG
ResNet-38
VGG-16
300
300
[
123
W/A
LE
400
103
60-
53
48
Km-36
Bm-00
B1v113
VoG-13
VDG-27
56
22
06
model size
model size
batch size
batohh
sinc
(a)
(b)
Figure 6: (a) The flatness and upper bound for various ResNets and VGG nets. (b) How the tightness of upper bound changes with decreasing batch size.
Conclusion
We provide stability-basee based explanation of why SGD selects flat minima. Our current understanding is as follows. 1) For popular ML models, the SGD noise aligns very well with local landscape. 2) This alignment property ensures that the flatness of stable minima must be size-independent, This understanding is made rigorous and quantitative by introducing a loss-scaled alignment factor to characterize the alignment strength and analyze the linear stability. Obviously, many questions remains open. For example, can we understand what roles the stability plays in the whole dynamic process of SGD instead of only around global minima? Can we establish the connection between the Hessian's 's Frobenius norm and generalization? Can we provide fine-grainer characterization of the noise structure and how the structure is related to implicit regularization of SGD? We leave the discussion of these important questions to future work.
Acknowledgements
We thank Zhiqin Xu and the anonymous reviewers for helpful suggestions. The work of Lei Wu is supported by startup fund from Peking University. The work of Weijie J. Su is supported part by and an Alfred Sloan Research Fellowship.
NSF Grants CAREER
10

---

References
[1] Vladimir Igorevich Amold. Geometrical methods in the theory ordinary differential equations, volume 250. Springer Science Business Media, 2012. 2
[2] Nachman Aronszajn. Theory of reproducing kernels. Transactions of the American mathemati- cal society, 68(3):337-404, 1950. 22
[3] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix
factorization. Advances in Neural Information Processing Systems, 32, 2019. 5
[4] Shahar Azulay, Edward Moroshko, Mor Shpigel Nacson, Blake E Woodworth, Nathan Srebro, Amir Globerson, and Daniel Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference Machine Learning, pages 468-477,
PMLR, 2021. 5
[5] Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629--8 1012,668,,60 2017. 25
[6] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by Ornstein-Uhlenbeek like process. In Conference on learning theory, pages 483-513. PMLR, 2020. 1,
[7] Youngmin Cho and Lawrence K Saul. Kernel methods for deep learning. Proceedings of the 22nd International Conference on Neural Information Processing Systems, pages 342-350, 2009. 25
[8] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2020. 2. 10
[9] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34:27449- 27461, 2021.
[10] Yu Feng and Yuhai Tu. The inverse variance- flatness relation in stochastic gradient descent is critical for finding flat minima. Proceedings of the National Academy of Sciences, 118(9), 2021.
[11] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur: Sharpness-aware min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2020.
[12] Crispin Gardiner. Stochastic methods, volume 4. springer Berlin, 2009.
[13] Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory, pages 2315-2357. PMLR, 2021. 5
[14] Hangfeng He and Weijie Su. The local elasticity of neural networks. In International Conference on Learning Representations, 2020.
[15] Kaiming He, Xiangyu Zhang, Shaoging Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and paitern recognition, pages 770-778, 2016. 9, 15
[16] S. Hochreiter and J. Schmidhuber. Flat minima. Neural Computation, 9(1):1-42, 1997. 1
[17] Izmailov, AG Wilson, D Podoprikhin, Vetrov, and Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference Uncertainty in Artificial Intelligence 2018, UAI 2018, pages 976-885, 2018.
[18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pages 8571-8580, 2018. 21
[19] Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in SGD. arXiv preprint arXiv:1711.04623 2017.
[20] Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, 2019. 2
11

---

[21] N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In In International Conference on Learning Representations (ICI.R), 2017. 1,3
[22] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images, 2009.
[23] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applicd to document recognition. Proceedings of the IEEE, 86(11))2278-2)) 2324, 1998. 15
[24] Michel Ledoux and Michel Talagrand. Probability in Banach Spaces: Isoperimetry and processes. Springer Science & Business Media, 2013. 20
[25] Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pages 2101-2110. PMLR, Aug 2017. 2
[26] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling SGD with stochastic differential equations (SDEs). In Advances in Neural Information Processing Systems, volume 34, 2021.
[27] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022. 3
[28] Tengyuan Liang. Tomaso Poggio, Alexander Rukhlin, and James Stokes. Fisher-Rao metric, ge- ometry, and complexity of neural networks. In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89, pages 888-896. PMIR, 2019. 3
[29] Chao Ma and Lexing Ying. On linear stability of SGD and imput-smoothness of neural networks. Advances in Neural Information Processing Systems, 34, 2021. 2
[30] Takashi Mori, Liu Ziyin, Kangqiao Liu, and Masahito Ueda. Logarithmic landscape and power-law escape rate of SGD. arXiv preprint arXiv:2105. 095.09577 2021. 1,3,8
[31] Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: A view from function space. Advances in Neural Information Processing Systems, 34, 2021. 3, 7
[32]
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization in deep learning. arXiv preprint
2014. I
[33] Scott Pesme, Loucas Pillaud-Vives Vivien, and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: provable benefit of stochasticity. Advances in Neural Information Processing Systems, 34, 2021. 3,5
[34] Dino Sejdinovic and Arthur Gretton. What an RKHS? Lecture Notes, 2012. 25
[35] Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithmss Cambridge university press, 2014. 20
[36] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In 3rd International Conference Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. 9, 15
[37] Unut Simsekli, Levent Sagun, and Mert Gurbuzbalaban. A tail-index analysis of stochastic gradient noise in deep neural networks. In International Conference on Machine Learning, pages 5827-5837. PMLR, 2019. 3
[38] Weijie Su. Neurashed: A phenomenological model for imitating deep learning training. arXiv preprint arXiv:21
2021.
[39] Yusuke Tsuzuku, Issei Sato, and Masashi Sugiyama. Normalized flat minima: Exploring scale invariant definition of flat minima for neural networks using PAC-Bayesian analysis. In International Conference on Machine Learning, pages 9636-9647. PMLR, 2020. 3
[40] Xingyu Wang, Sewoong Oh, and Chang- Han Rhee. Eliminating sharp minima from SGD with truncated heavy-tailed noise. In International Conference on Learning Representations, 2022. 3
[41] Stephan Wojtowytsch. Stochastic gradient descent with noise of machine learning type. part II: Continuous time analysis. arXiv preprini arXiv:2 02502888 2021. 1,3
[42] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673, PMLR, 2020. 5
12

---

[43] Dongxian Wu, Sau-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in Neural Information Processing Systems, 33, 2020.
[44] Jingfeng Wu, Wenqing Hu, Haoyi Xiong. Jun Huan, Vladimir Braverman, and Zhanxing Zhu. On the noisy gradient descent that generalizes as SGD. In International Conference on Machine Learning, pages 10367-10376. PMIR, 2020.
[45] Lei Wu and Jihao Long. A spectral- analysis of the separation between two-layer neural networks and linear methods. Journal of Machine Learning Research, 23(119):1-34, 2022. 25
[46] Lei Wu, Chao Ma, and Weinan E. How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective. Advances Neural Information Processing Systems, 31:8279-8288 -2299,888,.80 2018. 2, 3, 7, 10
[47] Lei Wu, Zhanxing Zhu, and Weinan E. Towards understanding generalization of deep learning: Perspective of loss landscapes. arXiv preprint /70i: 10239, 2017.
[48] Zeke Xie, Issei Sato, and Masashi Sugiyama. diffusion theory for deep learning dynamies: Stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, 2020. 3,
[49] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
[50] Hongyi Zhang, Yann N. Dauphin, and Tengyu Ma. Residual learning without normalization via better initialization. In International Conference on Learning Representations, 2019. 15
[51] Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why SGD generalizes better than Adam in deep learning. Advances in Neural Information Processing Systems, 33, 2020. 8
[52] 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 International Conference on Machine Learning, pages 7654-7663. PMLR, 2019. 3,
8,9
[53] Liu Ziyin, Kangqiao Liu, Takashi Mori, and Masahito Ueda. Strength of minibatch noise in SGD. In International Conference on Learning Representations, 2022. 3
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] [Yes], [No], or [N/A]. You are strongly encouraged to include justification to your answer, either by referencing the appropriate section of your paper or providing a brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section ??.
Did you include the license the code and datasets? [No] The code and the data proprictary.
Did you include the license to the code and datasets? [N/A]
Please do not modify the questions and only the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questionssansseess below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Local analysis
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
13

---

2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or URL.)? [Yes] See Appendix
A.1
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they chosen)? [Yes] See Appendix A.
(c) Did you report error bars (e-g.. with respect the random seed after running experi- ments multiple times)? [Yes] Estimated over 5 runs
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [No]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "Tgkm-OYaoIV", "reviewer": "Reviewer_qyGZ", "summary": "The paper provides an explanation of the phenomenon of SGD selecting flat minima by relating the noise structure of SGD to its linear stability. In over-parameterized models with the square loss, it shows, exploiting the geometry awareness of SGD noise (that is provable for linear networks and RFMs), that the Hessian of an accessible minimum for SGD is bounded by a term depending only on the batch size an the learning rate.", "questions": "- figure 3: are the results robust w.r.t. other flatness notions?\n- figure 5, it could be useful to add VGG11 to the right panel\n- the paragraph \"notion of flatness\" at l110 may be made more clear. Also the role of the norm in the flatness notions is not clear. also, size-independence plays an important role\n\nMinors:\n- l96: initiated studied ", "limitations": "The role of overparameterization and the specific loss type (quadratic) are not discussed\n\npotential negative societal impact are not discussed", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strenghts:\n- the topic of study is relevant\n- the quantitative analysis is extensive, going from simple models s.t. RFMs and linear models to modern deep neural networks\n- the paper is well written and clear", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "IuekMrQS3Aw", "reviewer": "Reviewer_jEC4", "summary": "The authors of the paper present the following contributions:\n- They provide an in-depth study of the importance of the SGD noise both in term of geometry and in terms of scale.\n- They show that, if a certain alignment property is satisfied, a global minimum is linearly stable if and only if the Froebinius norm of the Hessian at optimum is upper bounded by constant independent of model and sample size  ", "questions": "In my opinion, the crux of the paper is to give a precise (or experimental) sense to the approximations of equation (2). Maybe the authors could try to provide estimation of how far each approximation is to the real covariance: \n-  First regarding the influence of the loss\n- Second regarding the independence between $\\nabla f(x_i, \\theta)$ and $L_i$\n\nThis corresponds to constants $\\mu_1$ and $\\mu_2$ defined in the paper but no experiments are shown concerning these variables.", "limitations": "I use this paragraph to conclude as I already discussed the limitations on the previous boxes. The paper tries to conduct both theoretical and experimental explanations of the wide minima selection of SGD. I'll weakly accept the paper for its clarity and detail about the noise. To raise my score, I would like the authors to justify more the fact that the escaping phenomenon is important and the alignment phenomenon of the covariance structures. \n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### **Strengths**\n\n- The first thing I have to say is that the paper is very well written. The exposition is clearly conducted and all assumptions and estimations are detailed, announced and commented. One may or may not validate the authors' noise model, but at least it is not hidden (as too often) and keys are given to appreciate the results\n- The main strength and difference of the article is the particular attention given on the noise model both in terms of noise and of geometry. From this, even if the results seem easy to prove, the exposition is cristal clear and more convincing with the previous literature.\n\n### **Weaknesses**\n\n- There are a lot of results and approximation stated in the article, but overall, the crux of the paper is two show that the approximations of the equation (2) at the beginning of the paper is valid: that is to say that either $\\alpha, \\beta$ or $\\mu$ are close to $1$. If Figure 1 is  pretty convincing from this point of view, I have to say that Figure 5 is not as the scale of these constant can be (way) smaller than 1. I would really appreciate if the authors comment more on this point because the sentence \"*This comparison suggests that the alignment strength significantly depends on the intrinsic complexity of the problem, (nearly) independent of the model size*\" is not very convincing.\n- This is a minor weakness but overall, even if the paper is not overselling and clear about their study, I am still not convinced by the stability argument of SGD. Indeed, on never see plots like Figure 4 in real training dynamics of neural networks and this suggests -at least to me- that the taking into account the noise is not a stability issue but a dynamical one. Maybe the authors could comment a bit on this fact. \n\n\n### **Minor typos**\n\n- lign 65: $\\lambda_i$ and not $\\lambda_1$\n- lign 227: In **the** current paper\n- ligns after 235: there is a confusion between $v$ and $\\nu$\n \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "45jEQ9uAYkA", "reviewer": "Reviewer_9D5P", "summary": "Deep learning folklore holds that the gradient noise in SGD causes it to prefer \"flat minima.\"  An important open question in deep learning theory is to make this folklore mathematically precise.  Taking a step in that direction, this paper describes a mechanism by which SGD escapes exponentially quickly from minima where the Frobenius norm of the Hessian is too large.  This mechanism is orthogonal to the curvature-driven process that causes full-batch gradient descent to escape from minima where the Hessian spectral norm exceeds 2 / step size. ", "questions": "-- I didn't understand why the submission keeps mentioning that the flatness is independent of model size.  Could the authors clarify why this is an important point?\n\n-- Are the authors sure that the linear stability analysis is only valid near a local minimum?  I ask because in the full-batch case, the curvature-driven escape when sharpness exceeds 2 / (step size) is valid generically, not just near a local minimum.\n\n-- I wonder if the bound would be tighter on the real networks (VGG, ResNet) if you considered a batch size smaller than 64.  Intuitively, I would expect the noise-driven escape to dominate when the batch size is small.", "limitations": "Yes, the authors adequately addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n  -- To the best of my knowledge, the idea of an exponentially fast escape that is driven purely by noise is novel.  It is interesting that a required condition for this phenomenon (noise magnitude proportional to loss value) is provably satisfied in linear nets and random feature models.\n\nWeaknesses:\n\n  --  The analysis only yields a sufficient condition for instability, not a necessary-and-sufficient condition.\n\n  -- While in general the stability will depend on both the full-batch component and the noise component, the analysis here considers only the noise component in isolation.  Accounting for both simultaneously is going to be hard, so I don't begrudge the authors for this simplification. \n\n  -- I think the paper would be clearer if the authors first presented the escape analysis (section 3) before the sufficient conditions (section 2).  When I read section 2, I spent a lot of time scratching my head wondering why alpha, beta, and mu are defined the way they are.  Later, when I got to section 3, I realized that \\mu is precisely what is needed to trigger exponentially fast escape.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept