Implicitly regularized interaction between SGD and the loss landscape geometry
Authors: Sungyoon Lee
Cheongjae Jang
Venue: NeurIPS 2022 Submitted (2022)
Keywords: SGD
learning rate
batch size
optimization
generalization
implicit bias
implicit regularization
sharpness
scaling rule

Abstract: We study unstable dynamics of stochastic gradient descent (SGD) and its impact on generalization in neural networks. We find that SGD induces an implicit regularization on the interaction between the gradient distribution and the loss landscape geometry. Moreover, based on the analysis of a concentration measure of the batch gradient, we propose a more accurate scaling rule, Linear and Saturation Scaling Rule (LSSR), between batch size and learning rate.

SGD plays important role in the success of deep learning. However, still do not fully understand how SGD works from the perspectives of both optimization behavior and generalization performance. 10 To be specific, SGD is a stochastic approximation of full-batch gradient descent (GD), but SGD Implicitly regularized interaction between SGD and the loss landscape geometry
Anonymous Author(s)
Affiliation
Address email
Abstract
We study unstable dynamics of stochastic gradient descent (SGD) and its impact on generalization in neural networks. We find that SGD induces an implicit regularization the interaction between the gradient distribution and the loss landscape geometry. Moreover, based on the analysis of concentration measure of the batch gradient, we propose more accurate scaling rule, Linear and Saturation Scaling Rule (LSSR), between batch size and learning rate.
Introduction
11 generally yields better generalization with a small batch size [2] Moreover, GD is discretization
12 of gradient flow (GF) with a finite learning rate, i.c., GF is a GD the limit of vanishing learning
13 rate, but GD generally performs better with a large learning rate
28 There are some how tune the learning rate for varying batch sizes, but they fail when the batch size gets large 42 38 57
14 scaling rules [25
10 15 45
4 33 Especially for greater data- tataa-pralleeismm to accelerate the training process, we require a more accurate scaling rule for the large-batch regime.
17 There has been many studies to understand the SGD dynamics and its impacts on generalization in deep neural networks. While they provide some useful and intuitive explanations to help us understand these properties of SGD, unfortunately, some results often rely on impractical assumptions only apply certain range of learning rates and batch sizes. For example, some approximate SGD
18
19
20
stochastic differential equation (SDE) in the limit of vanishing learning rate [34
35
29
16]
30 31
18 44 4 Therefore, in practical finite learning rate regime, this may not properly describe the
23
SGD dynamics. Moreover, Yaida [52) raises some theoretical issues about the SDE approximation
24 and Li et al.
3 theoretically analyze a sufficient condition for the SDE approximation to fail.
25 In this paper, aim understand the dynamics and the implicit bias of SGD through the analysis of
26 the interaction between SGD and the loss landscape of neural network with minimal assumptions. To be specific, we investigate the unstable dynamics of SGD "at the edge of stability"
[6] (Section 4.22. This investigation leads to more refined characterization of the edge of stability by the inleraction-awatt sharpness which extends the previous findings for full-batch GD to a general SGD, 30 Then, we introduce a concentration measure of the the batch gradient distribution SGD. By doing
28
29
31 so, we find that SGD implicitly regularizes the interaction-awar sharpness and its regularization effect controlled by the ratio of the concentration measure to learning rate (Section Finally, accurate scaling rule between batch size and learning rate, based on novel
32
38
propose
Submitted to 36th Conference Neural Information Processing Systems (NeurIPS 2022). Do distribute.

---

34
38 In this section, we review some concepts required for further discussion. We also summarize the analysis of the implicit regularization and the concentration measure (Section 5.2 This can be applied to any batch size including the large-batch regime where the previous scaling rules fail We name Linear and Saturation Scaling Rule (LSSR).
35
36
18
38
57
Stochastic Gradient and Loss Landscape
37
39 notations in Appendix for a quick reference. We often omit the dependence on some variables and the subscript of the expectation operation when clear from the context.
40
41
For learning task, we use parameterized model (neural network) with model parameter 0
O C RT. Then train the model using training data D { J and loss function ((x; 0). We denote the (total) training loss by L(0) = Di= {(x;; f) for training data D. At time step
48
44 I. we update the parameter O, using GD: 0z+1 = 0, ((o)) with a learning rate 7), or using
45 SGD: 0,+1 = 0₂ mgh((p, with a mini-batch gradient (0.) Ezes, Vol(x; E R for
46
bath- B, C D size b (1 < < n).
ready introduce some important matrices, C., So. and H. First, define the (o) (0)) (as(0) Elgu(0))) R and the second moment (0) Elas (0)g((()) E mmxm of the mini-batch gradient (b) over batch sampling for a batch size 1 < b < n The covariance C and the second moment Sj satisfy not only
Now,
48
covariance (0) Varlgs(0) E
50
51 C. = S Sn but also the following equation 15 29 49]:
C = In. (S1 - (.) = 7mn C1, b
b
52
where Tn n-1
n-b for sampling without replacement and Tn.h for sampling with replacement
53 We provide self-contained proof of (T) in Appendix [.I] We note that, for sampling without replacement, many previous works approximate "nn. = assuming b < [8] [5
46). but
consider the whole range of 1 < b n (0 Tn.6 < 1 with Nn = and Yn, 0). Second, 56 we define the Hessian (0) (Q) = E (Vel(x 0)] E IRmx and the operator norm (the
57 top eigenvalue) ||| = [//l|/ Hul] of H. We also denote the i-th largest eigenvalue and its
58
corresponding normalized eigenvector by x E IR and qi E R", respectively, for i = =1, m.
59
Therefore, with these matrices, we can write one of our goals as follows:
80
We aim to understand how the gradient distribution (C,, and So) and the loss landscape geometry (H) interact with each other during SGD training.
We investigate this "interaction" in terms of matrix multiplication HS,, To be specific, we consider the trace tr( (H)))) or its normalized one (x(HSs) (will be denoted by || in Definition later). tr(Se)
63
Related Work
64
85
Some studies investigate the interaction between the gradient distribution and the loss landscape geometry represented by tr(H..)) in the context of escaping efficiency [58 Section 3. 3. stationarity Section 2.2], and convergence 48 Section 3.1.1]. However, they require some additional assumptions like SDE approximation of SGD (58). the existence of stationary-state distribution of the model parameter (52) Section 2.3.4], and strong convexity of the training loss function (48). respectively. In this paper, we provide a new insight into the interaction tr( (HS.) without these assumptions.
66
87
52
88
70
71
72
Convergence of full-bateh GD (b n) has been instead analyzed with an upper bound on the interaction (r) with further assumptions for the stable optimization, such 3-smoothness of
73
'These two matrices Cs and S,, are also called the second central and non-central moments, respectively. But avoid confusion, we the term "second moment" only for the non non-central S..

---

74 the objective and 0 < (e.g., 17 = )) 39 37 However, it may lose useful information
75 of the interaction between and S,- Moreover, when we train a standard neural network with GI)
76 in practice, |H]] ( (3) increases in the early phase of training and the iterate enters the regime called
77 the edge of stability 6 where [||| / I. i.e., '/ ~ [H 2 z. This contradicts with the assumption
78 for stable optimization and the iterate exhibits unstable behavior with a non- monoonootyy decreasing We further extend this discussion of unstable dynamics for GD the case of SGD.
79 loss
50 6
From the generalization perspective, many studies focus on the implicit bias of SGD toward better 80
81
1 46 There are mainly two factors known to correlate with the generalization performance: the batch gradient distribution during training [15] [8] 44 [88 and the sharpness of the loss landscape at the minimum
generalization [40 56, 47 20
21
82
83
[14] 23 22 26 We provide a link between the
84
batch gradient distribution and the sharpness that the model is implicitly regularized have low
sharpness when the second moment of the batch gradient is large (see Section 5.1
4
Optimization through Loss Landscape
86
We start by investigating the optimization behavior of SGD through the interaction between SGD and the loss landscape without the stochastic differential equation (SDE) approximation.
88
89 4.1
Unstable Optimization
90
Using the second-order Taylor expansion, the change in total training loss L2 = L(Or) as the SGD iterate moves from O, to Oc+1 at time step can be expressed as follows:
s1
2 2
L L111 L, = - 9b +
Hgb 0(16,,1)))
(2)
92
where 6, =0+41 0, -196.. Thus, we obtain the expected loss difference as follows:
E[4+t1l L1 = L [[ + Elg Hgl +
(3)
2
1/2
= ~11 + tr (E +
(4)
2
1?
= -nntr(Sm) + tr(HSp) +c
(5)
17² tr((..)
(H)())))
2
+,
2
tr()n)
n
where e (( and [[g.] VL is used. For the moment, we make minimal assumption that the training loss is locally quadratic, i.e., near 6., but will revisit this assumption later (see Section 4.2). Then, the expected loss increases when the following instability condition is met:
93
Definition (Instability Condition).
(r(H))) 2 tr(Sn) 1
96 We also define unstable regime U (0 E
sr(H)))) } and stable regime U For tr(S..)
97 standard non-quadratic loss function, will show in the following sections that the iterate tends not
to stay within the unstable regime U and operates near at the boundary 0S of the stable regime S called the edge of stability 6 Cohen et al. 6 mark the edge of stability with {0 E θ |||| )}
99
100
for GD, but we mark with as = (0 θ clear and generalized indication shown in Figure later. On the other hand, for a globally quadratic loss, when the GD iterate satisfies the instability condition, it diverges within the unstable regime [6]. We emphasize that many studies on the convergence of GD usually consider the optimization within
tun (
2) for both SGD and GD which provides more
102
103
2 (T(+.)))-((0) < VLL V (0e+1-0,)+ |0.+1 1,1 V 2n2 | -n(1 8n [| and thus the loss monotonically decreases when 1) 3.
3

---

Figure 1: [An empirical validation of (6) for SGD (top) and 99) for GD (bottom)) In the early phase, until the iterate enters the edge of stability, it validates (6) and ( with the blue line with the slope of and x-intercept n. For GD (bottom), they are plotted after || H|] exceeds 2 after which 7 |s,, starts to increase from 0 to in few steps. For cross-entropy loss, we mark the end point 1e 2 CE, b 12 0.02
1e 2CE, = 256. 17 0.02
1EL..1] Laltr(S,))
4
(ELL.. (.itt(s.)))
100
2-
slope-
80
2
slope
step
step
0-
80
= 100
-2
= 100
40
60
0.00
0.75
1.50
2.25
3.00
0.00
0.75
1.50
2.25
3,00
(rp15))
1e2
MinS
1e2
1ee2
CE GD, n 0.02
1e-2
MSE. GD, 0 0.02
(lee1- 4.itm(..))
120
(altriss))
opee
slope
270
115
0
260
-2
= 100
110
100
250
0.00
0.75
1.50
2.25
3.00
0.00
0.75
1.50 ||M|ls.
2.25
3.00
||H]|
1e2
1e2
with 'x' when the iterate enters the unstable regime. We train 6CNN with 1 0.02.
37 3, but GD mostly occurs the edge of stability after a few steps of training. We will argue that this behavior is crucial for generalization in neural networks.
the stable regime
39
106
For later use, we also define the interaction-aware sharpness as follows:
Definition (interaction-awar sharpness).
tr (H,)))
| |s,
(8)
tr(Ss)
107
Here, Sr. < || (r()),, i.e., || < || and the equality holds only when every go is aligned in the direction of the top cigenvector of H.
Figure
(top row) empirically validates showing the normalized loss difference
B(4,+1]-L
tx(S.)
110 against (H()) in the early phase of training before entering the unstable regime. This result implies
te(S.))
that the training loss L(0) is approximately locally quadratic, i.c., z O, in the early phase. Especially,
for full-batch GD (b = n.), the instability condition can be rewritten as ||, > 2 and we have the
112
following relationship between the loss difference Le+1- L, and ||,,, from (6
172 L - L, tr((s)) [1LE|]s.. 2 +
114 Figure (bottom row) shows [|s.. soars from in a few steps after || exceeds 2 6 satisfying approximately with 0, before the iterate enters the edge of stability. This result is consistent with the following Proposition for a quadratic training loss L. The proof is deferred to Appendix B.2
117
Proposition 4.1. For GD with a quadratic L, if [| > v
2 and 0 < A < 3 for all i * 1, then cos (L().)) l L(())) and [|s, increase to 1, oo and ||]]], respectively, as t + co.
118
119
4.2
Non-quadraticity, Asymmetric Valleys and the Edge of Stability
In the previous section, have shown that the training loss is approximately locally quadratic before the iterate enters the edge of stability. However, after the iterate enters the edge of stability, i.e.,
u (SH)) reaches and exceeds 7)* the step size is relatively large for the sharp loss landscape so that the iterate jumps across the valley [99. and the higher-order terms in (6) and ( become non-negli
122
tr.S.)
and cause different behavior of the iterate than in the stable regime.

---

1e
CE b 128, ==
1e XEE b= 256. 0.02
1EL..1] Laltr(S,))
4
2000
(EL... - (.itt(s.)))
2000
2-
slope
1500
slope
1500
1000
0-
= 100
500
=000
500
0
0.00
0.75
1.50
2.25
3,00
0.00
0.75
1.50
2.25
3,00
1e2
1e2
IrSul
1e-2
CE. GD, 17 0.02
1e-2 MSE, GD, n 0.02
1500
1500
(lee1- (.itm(s))))
0.e11 (altris,))
slope
1000
slope-
1000
scep
step
0
-2
100
-2
= 100
0.00
0.75
1.50
2.25
3.00
0.00
0.75
1.50
2.25
3,00
||H||s
1e2
||M||
1e2
Figure 2: NNon-quarraticity and overestimation] The normalized loss difference
B|L++1]-Le against
tr(S.))
tr(HS,)
during training. After the iterate enters the edge of stability, it often shows a more gentle tr(S.)
slope than n² 2 especially in the unstable regime.
125
Figure shows empirical evidences for the non-quadraticity. After the SGD/GD iterate enters the edge of stability, when the instability condition (r(HSs
2 is met, the normalized increase in the Tr(S.)] "
126
B| 1 often smaller than of (S,)
(.) from and (9) (blue line) when assuming te(s.))
127
loss
a locally quadratic function. This results gentle slope less than y
128
We hypothesise that due to this non-quadraticity of the training loss, the iterate is discouraged from staying within the unstable regime. Figure 3 demonstrates the asymmetric valley [12] that one side is
130
131
sharp and the other flat. In Figure (left), we evaluate the directional sharpness || |. along V
132
the gradient descent direction - Vn (() where Ha = H( - onj V for x [1, 21, 2, 4, 5]. and compare (.,)),,(0)) with s.. At the sharp side, has a high || s. > "j (blue) with
135
134
the gradient L and the top eigenvector qı (H) of the Hessian being highly aligned (cf. Prop. 4.1).
135
However, when the loss landscape gets far from being quadratic, the Hessian and its top eigenvector can change abruptly, q (Ha) would not always be aligned with qı (H) and VI.(0), and || Halls, tends decrease. This would be a possible explanation for the tendency of decreasing and then oscillating ||.... See Appendix C.3 for detailed empirical evidences of the above arguments.
196
137
Figure 3 (right) similarly shows that when the iterate is at sharp side of the valley, it tends to jump to the other side of a flatter area, and vice versa.
140
To summarize, we make the following observations for GD in order: (i) || increases in the beginning (the progressive sharpening (6),, (ii) || exceeds 2 (iii) the gradient V becomes more aligned with the top eigenvector qı (H) in few steps, (iv) ||H]|, reaches the threshold 3 and the iterate jumps across the valley, (v) ||, tends to decrease due to the non-quadraticity, and it repeats this process, while ||]]s, oscillating around 2 tr( (Ss)
141
142
'7'
144
")* We observe similar behavior with oscillating S.. around 2 for SGD. It requires further investigation into the exact underlying mechanisms 7
146
and we leave it as future work.
148
report the experimental results using vanilla SGD/GD without momentum and weight decay, constant learning rate, and no data augmentation. We train simple 6-layer CNN (6CNN, m 0.51M) CIFAR-10-8 where DATASET-n denotes subset of DATASET with (DI and k=2 1024. See Appendix C.1 C.3 for the results from other datasets,
Remark (Experiments in Section
149
150
152
learning rates and networks (ResNet-9 with m 2.3M [13 and WRN-28-2 with m = 36M (55)).
5

---

||l|s x
||M|ss >
||A||s. <
||H||s. >
10a
-min
IHdls TP
10
AL
-VL
0
10
0.25
0.50
0.75
1.00
1.25
min
zero
Figure [Asymmetric valleys] Left: The ratio ||Hals. where Ha = H V(()) for 1 (1,2,3,4,3 4,5] for each & during training. When |[]]s < (red), 11.s.. is usually larger than |, On the other hand, when || > (blue), ||Halls, is usually smaller than | Right: The training loss difference along the gradient descent direction, for each 6,. Each plot is normalized and translated to have the same minimum value and the same zero where AL 0. We also plot the quadratic baseline (cyan dashed curve). When || < 2 (red), it usually becomes sharper across the valley (right-shifted). On the other hand, when ||/ > (blue), it usually becomes flatter across the valley (left-shifted). We train 6CNN using GD with = 0.04.
|[
153
5
Generalization through Implicit Regularization
In the previous section, we have empirically demonstrated that the SGD iterate is implicitly discour- aged from staying within the unstable regime. Now, we are ready to further analyze this property from the regularization perspective.
155
156
157
5.1
Implicit Interaction Regularization (IIR)
158
First, to understand the effect of batch size the gradient distribution, we define the following ps: Definition 3 (a concentration measure of the batch gradient). We define Ps the ratio of the squared norm of the total gradient || to the expected squared norm of the batch gradients B[[lgh|[2], i.e.
|V ir(S.)
pb
E 11 (r()
(10)
161 Here, we can write || = || 1| and thus the ratio pb B.
S is similar to the square of the batch gradient Pb (36] especially when |I] is small compared to gl [I] (see Appendix C.5 for empirical evidences). Both Ps and R3 are concentration measures and have lower values when the batch gradients ÿb are more scattered. Therefore, is natural to expect that the ratio Ps is small for a small batch size b, and we will revisit this in more detail in the following section (cf. 122 We also note that Pa. = I2 1.
|| 2
of the mean resultant length R2 | I ]1I
162
163
st
184
165
186
Now, we can rewrite the instability condition S. (S.)
2h (multiplying both sides by Pb) as ||s >
167
168
2ph
In other words, the interaction- -wwar aware sharpness |Hss s. implicitly regularized to be less than
We name this Implicit Interaction Regularization (IIR).
Definition 4 (Implicit Interaction Regularization (IIR)).
H < 2pb
(1I)
We argue that the upper constraint 2n in IIR is crucial in determining the generalization performance. 171 With a low constraint, SGD strongly regularizes the interaction-awarr sharpness || ||]ss. s. We also 172 note that IIR affects not only the magnitde || but also the directional interaction. In other words, IIR discourages the batch gradients from aligning with the top eigensubspace of the Hessian that is spanned by few largest eigenvectors of the Hessian (cf. [T]..
y

---

ResNet: 9, CE, GD, 0.01
ResNet-9, CE, GD, n= 0.02
800
||PI|s.
[
IMIL
[H]
2.0
700
2.0
350
600
1.5-
500
1.5
250
400
200
1.0
1.0
300
0.5
200
0.5
100
0.0
3,00
1500
2000
0.0
1000
2500
2000
Step
Stop
(a) GD, ResNet-9, cross-entropy. 17 = 0.01
(b) GD, ResNet-9, cross-entropy. 1) = 0.02
6CNN, MSE, GD, 0.0.2
0.5
200
train loss
|Hll.
MM
0.4
115
25.0
0.3
20.0
100
0.2
775
15.0
50
10.0
0.1
25
500
1000
1500
2000
125
250
375
500
Step
Stop
(c) GD, 6CNN, MSE, 1) 0.02
(d) SGD, ResNet-9, cross-entropy, 77 0.08
Figure 4: IA clear indication of the edge of stability] (a)-(c): After few steps of full-batch training, || (blue) hovers above 2 [6) but [| ll.. (red, defined in (8) oscillates around _ (red dashed horizontal line). The edge of stability is more evident in the latter (red). Curves are plotted for every step. We train a model on CIFAR- 10-8k
7
21 using (a)/(b) cross-entropy loss with ? 0.01 /0.02, respectively, and (c) MSE with v 0.02. (d): We plot curves || |l when trained with various i's. After few steps (around 125), they reach the threshold which linearly increases as b becomes larger when b < 213 and saturates to 2ps = 7 when is large. Curves are smoothed
for visual clarity. We use SGD with b {23. .213 and 17 = 0.08.
175
Figures 4(a) 4(c) show that, for GD (pn 1), the interaction- aware sharpness || (red) oscillates around 2 and exhibits IIR. This result is consistent with Cohen et al. (6) that || ||| hovers above 2
176
177
|| S || and the equality holds only when the gradient V and the top eigenvector F1 of H are aligned, but generally they are not. For this reason, IIR provides tighter relation and more clearly identifies the edge of stability than Cohen
for GD. This because, as mentioned earlier,
178
et al. 6 These results are also consistent with Prop. 4.1 that [ ||s. suddenly increases from 0 to 2
180
181 in few steps after || exceeds (see Appendix C.3 C.4 for more). Moreover, IIR also applies to a general SGD training with 1 < 6 S n. Figure 4(a) shows IIR for SGD with different batch sizes b E (23, 2222. The upper bound (2ps/1 according to (11 of [| is higher when using
182
188
larger batch size, but limited to less than 2/n (ps < 1). We will further discuss this behavior with an investigation of Pb in the following section.
185
186
5.2
Linear and Saturation Scaling Rule (LSSR)
187
The ratio b/n of batch size b to learning rate rj has long been believed as an important factor in- fluencing the generalization performance, and the test accuracy has observed to be similar when trained with the same ratio b/1 b' /T/. i.e., b' kb and 1/' km for k >> O. This called the
188
189
linear scaling rule (LSR) [25 10
18
44 57). They argue that LSR holds because Oetk O,
B DI= EreB. (((x;0.4.)
DM EreR,,: Vl(z; 0,)
g' 3 ExeB,.iA Vl(x; 0,) as-
191
182
suming Vi(0044)) Ve(O.) for 0 <
where B = of Bre+ and |Bac+A| = kb R. However, the assumption is false and the gradient oscillates mostly with negative cosine value cos(gn (0t), go (0++1)) 0 between two consecutive gradients after entering the edge of stability
7

---

CRAA MMA..COO.. teali.
LSSR: 110)
....
LSR: R()) = bp
2²
SASR: (b) =
Inge/o
20
2
3 9303300
10g:101
STL-2044 Autberd
Try-rrugshee-02 Raskus
2-6
160
325
V10.00
2-8
20
22
2
26
29
22
2
batch size (b)
Figure 5: [Linear and Saturation Scaling Rule (LSSR) Left: LSSR (red) in (2) LSR (black dotted line) [T0] and SRSR (blue dotted line) [13] For LSSR,
observe both linear and saturation 2 2) Right: Heatmaps of test accuracy for models trained with a large number of pairs of (b, 17) on CIFAR-IO-8K CIFAR-100-8k STL-10-4k, and Any-ImmagNNe-332 (from left to right, from top to bottom). It does not follow either LSR or SRSR, but LSSR. We also plot f(b) Ps (yellow dashed curve) for some p on each heatmap. Note that they are all log-log plots and thus slope of means it is linear.
regions (n 8k,
195 (see Appendix C.S) Moreover, LSR fails when the batch size large 18 38 57
46 On the other hand, Krizhevsky (25), Hoffer al. [55] propose the square root scaling rule (SRSR) with another ratio vb/n to keep the covariance of the parameter update constant for < based on
43
196
197
Var [ng]] 7 2G Tn.an C1 q² C1. However, Shallue et al. 42 show that both LSR and SRSR
do not hold in general.
200
Based on the analysis of IIR with new ratio 2pb/n in the previous section, we explore why LSR fails the large-batch regime and provide more accurate rule to explain the generalization performance
201
202
of the models trained with various choices of batch size and learning rate pairs (b, 17).
To this end, we investigate the concentration measure Ps: = tr(S.). ) (r( (S))))) By combining two equations, C = Sb Sn (by definition) and C Ind (S1 S,) in we can obtain So = C + S, = - S + (1 >
Ind b )S,- Therefore, we have tr(S,) = Tn tr + In. tr(S.).
205
206
which leads to the following equation:
Pb = tr(S,)
tr(S,)
naib p = bp
if b is small if is large
ir(St)
Tnub tr((i)
Ins (rA)
Tri
p + (1
2m
(12)
207
from 10 where p pI = tr(Sa)/ (r())). Note that p (much) smaller than because V((x;)
has different direction for each v; and tr(S,) || 1 2 < L. I[V((x;))1
208
209
(r S.). In other words, 1/p is (much) larger than (see Appendix C.5
Figure 5) (left) demonstrates new scaling rule with the ratio Ps/n. called the Linear and Saturation
210
211 Scaling Rule (LSSR), with the two regimes that (i) Ps is almost linear when b < n (linear regime) and
212
(ii) Ps saturates when is large (saturation regime), which are also shown Figure 4(d) It depends on which part of the denominator (+) in 12) dominates the other. First, when < n, then 7n.b/ very small and the first term To.b dominates the second term 1 2nd since >>>>>>>>>> Second, as b becomes large, 7n.b/b zs 0 and the second term (* 1) dominates the first term. Thus, Pb saturates
213
214
215
216
to and is not linearly related and LSR is no longer valid. The above arguments also hold for
the batches sampled with replacement where the only modification is Tr 1, vb in 12 Figure (right) empirically supports LSSR with the test accuracies when trained with various combinations of pairs (b, 1)). To be specific, the optimal learning rate is almost linear when b is small, but it saturates

---

220
when b is large. We also plot f(b) = Ps (the yellow dashed curve) for some p. Note that Figure 8 of Shallue et al. 42 Section 4.7] shows similar "linear and saturation" behaviors supportive of LSSR on other datasets (see also Figure 7 of Zhang et al. [37, Section 4.3
221
222
223
Remark (Experiments in Section . We train models using vanilla SGD/GD without momentum and weight decay, constant learning rate, and no data augmentation. For Figure 5 we use subsets of the datasets CIFAR- 10 124), CIFAR- 124), STI.-10 [5], and Tiny-ImageNet (a subset of ImageNet KZy with 3 64 64 images and 200 object classes). We use large number of epochs (800) and batch normalization (77] achieve a zero training error even with a large b and a small 1/. However, in the lower right corner (red area) of each heatmap in Figure (right), when b is too large nj is too small so that [|0++1 0ell = n/lg61l is too small, it requires an exponentially large number of steps for the iterate to enter the edge of stability. Thus, in this case, the assumplion in Goyal et al. [10),
224
225
228
227
229
230
231
V [((e)) x VE(0++s) for 0 S i < k; approximately holds and the reasoning on ISR is valid. However,
232
this only holds for a non-practical (b, 11) which shows a suboptimal performance. See Appendix
C.5 for the results from other networks and hyperparameters. 233
Discussion
234
235
We provide new insight on the link between the batch gradient distribution and the sharpness of the loss landscape. In this section, reconcile our arguments with some previous studies.
236
237
Jastrzebski et al. [8] explain the optimization behavior of SGD with the SDE approximation
d0,- VL.O. )dt /TC/////00/ (W of the SGD where W is an m-dimensional Brownian motion. 238
Therefore, the same ratio E y leads to the same SDE, which implies LSR. Moreover, large b
240
implies a large diffusion in SDE, which has been linked with the escaping efficiency from a sharp local minimum in Zhu et al. (58) We instead argue that a large second moment tr(S,) (compared to (S,)) and a large 77 lead to a low constraint 2pb/n on the interaction aware sharpness. We
241
242
243
emphasize that we do not model SGD with SDE and thus our argument applicable to practical
244
learning rate regime.
Wu et al. 44 empirically show that what is important for the generalization performance of a neural network is not the class to which the gradient distribution belongs. but the second moment of the distribution. This is consistent with our arguments with the interaction tr (H.))) and the concentration measure Ph = tr(S.)/ (r(S)), because they depend on the second moment Sb, not on the class of the gradient distribution.
248
247
248
249
33 suggest necessary condition that the "noise-to-signs ratio" needs to be large for LSR (and the SDE assumption) to hold. This is consistent with our result on the linear regime (where 1 and pb are small) because the noise-to-signal ratio is approximately the inverse of the ratio Pb = =r((S,) tr( (r()),, but defined for an equilibrium distribution. We provide not only the necessary condition but also the sufficient condition for LSR with a novel scaling rule
Recently, Li et al.
251
252
253
"signal-to-noiss
254
255
LSSR applicable to every batch size including where LSR fails (the saturation regime).
Conclusion
296
257
From an analysis of unstable dynamics of SGD (Section 4.1) and the instability condition (Definition D.
we clearly mark the edge of stability (Figure 4 with the interaction-aware sharpness ||]] |.. (Definition 2 and show the presence of the implicit regularization effect on the interaction between the gradient distribution and the loss landscape geometry (IIR) (Section 5.1 Definition 4 Moreover,
260
281
introducing the concentration measure Ps of the batch gradient (Definition 3
12 we link the
262
second moment of the gradient distribution and the sharpness of the loss landscape, and propose
263
a new scaling rule called Linear and Saturation Scaling Rule (LSSSR) (Section 5.2) Figure Due
264
to the simplicity of the analysis, we hope that our insights will motivate the future work toward
285
understanding various learning tasks.

---

References
286
[1] David Barrett and Benoit Dherin. Implicit gradient regularization. In International Con- ference on Learning Representations, 2021. URL https://openreviiee wwt./frumm/fiiidaooooooo
268
269
3q5IqUrkcF)
270
[2] Nils Bjorck, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normalization. Advances in neural information processing systems, 31, 2018.
272
[3] Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-311, 2018.
273
274
[4] Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational infer- ence, converges limit cycles for deep networks. In International Conference on Learning Representations, 2018.
275
277
[5] Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsuper- vised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 215-223. JMLR Workshop and Conference Proceedings, 201
278
279
[6] 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 Learn- ing Representations, 2021. https::s::/
282
www./errummr Forum? hi=imiiiijjihheeeee
283 [7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248-255. lece, 2009.
284
285
286
[8]
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028. PMLR, 2017.
287
288
[9] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awart min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL httphtss///ppeerreiiee wwwp:eteeoooommddnniiiiioeoo
291
292
[10] Priya Goyal, Piotr Dollár, Ross Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour. arXiv preprint arXiv: (oo::7006007777 2017.
293
294
295
[11] Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint
296
2018.
297
[12] Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: Beyond sharp and flat local minima, Advances in neural information processing systems, 32, 2019.
298
299
[13] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
300
301
302
[14] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural compulation, 9(1):144,, 1997.
[15] Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generalization gap in large batch training of neural networks. Advances in neural information processing systems, 30, 2017.
304
305
308
[16] Wenqing Hu, Chris Junchi Li, Lei Li, and Jian-Guo Liu. On the diffusion approximation of nonconvex stochastic gradient descent. Annals of Mathematical Sciences and Applications, 4
(1), 2019.
10

---

309
[17] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448 456. PMLR, 2015.
310
311
312
[18] Stanistaw Jastrzebski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors influencing minima in sgd. arXiv preprint arX1v:1711.04633, 2017.
313
314
[19] Stanislaw Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, 2019. URL https
317
318
//openreview.o het.r/rweeeetooorrrrooooosogaaaa
319
[20]
Stanistaw Jastrzębski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho*, and Krzysztof Geras The break-even point on optimization trajecto- ries of deep neural networks. In International Conference on Learning Representations, 2020.
320
321
URL httpp:://ppenrrviee htt.:ppaeereeeeeeeeeeeeoooorrarr nett/foruumidd=1g57766
322
323
[21] Stanislaw Jastrzębski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras. Catastrophic fisher explosion: Farly phase fisher matrix impacts generalization. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pages 4772-4784, PMLR, 18-24 Jul 2021. URL mlr (tress//113//jasttzeebsst
325
326
327
328
html
329
[22]
Yiding Jiang, Bchnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them. In International Conference on Learning Representations, 2020. URL htttps://ppenrvvieee httppp/nweeeeeeeeeeefortrfttfeee mett/forummaddSSIIIPPU
330
331
332
[23]
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreviiee htt.:forumm.dddHHollRl....
334
335
336
337
[24] A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
338
339
[25] Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv
340
preprint arXiv: 14045544 2014.
[26]
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness- aware minimization for scale-invariant learning of deep neural networks, 2021.
[27]
Yann A LeCun, Léon Bottou, Genevieve B Orr, and Klaus-Robert Müller. Efficient backprop. In Neural networks: Tricks of the trade, pages 9-48. Springer, 2012.
344
345
[28]
Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint arXiv:2003.022 022118 2020.
346
347
348
[29] Chris Junchi Li, Lei Li, Junyang Qian, and Jian-Guo Liu. Batch size matters: diffusion
3/9
approximation framework on nonconvex stochastic gradient descent. stat, 1050:22, 2017.
350
[30] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and adaptive stochastic gradient algorithms. In International Conference Machine Learning, pages 2101-21
PMLR, 2017.
11

---

353
[31] Qianxiao Li, Cheng Tai, and E Weinan. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 20(1):1474-1500 2019.
354
355
[32] Yuanzhi LI, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in Neural Information Processing Systems, 32, 2019.
357
358
359
[33] Zhiyuan Li, Sadhika Malladi, and Sanjeev Arora. On the validity of modeling sgd with stochastic differential equations (sdes). Advances in Neural Information Processing Systems, 34, 2021.
361
[34] Stephan Mandt, Matthew Hoffman, and David Blei. variational analysis of stochastic gradient algorithms. International conference on machine learning. pages 354-363. PMLR, 2016.
362
383
[35] Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18:1-35, 2017.
364
[36] Kanti V Mardia, Peter E Jupp, and KV Mardia. Directional statistics, volume 2. Wiley Online Library, 2000.
367
[37] James Martens. New insights and perspectives on the natural gradient method. arXiv preprint 1121v:1412.1193, 2014.
368
389
[38] Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint arXiv: 2018.
370
[39] Yurii Nesterov. Introductory lectures convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
372
373
[40]
Behnam Neyshabur. Implicit regularization in deep learning. arXiv preprint arXiv:
2017.
375
[41]
Mark Schmidt. Convergence rate of stochastic gradient with constant step size. 2014.
[42] Christopher J Shallue, Jachoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint 1xXiv:8111.030000 1.03600, 2018.
377
378
379
[43] Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. In International Conference on Machine Learning, pages 955899067. PMLR, 2020.
382
[44] Samuel Smith and Quoc Le. A bayesian perspective on generalization and stochastic gradient descent. In International Conference on Learning Representations, 2018
383
384
[45] Samuel L Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V Le. Don't 't decay the learning rate, increase the batch size. arXiv preprint arXiv: 171.04899 2017.
385
[46] Samuel L Smith, Benoit Dherin, David Barrett, and Soham De. On the origin of implicit regular- ization in stochastic gradient descent. In International Conference on Learning Representations, 2021. BRLLhttts://ppeerreee http:p/////pppooooooooeoooooooaa
387
388
net http://ooormmmi-drrrr000001111
389
[47] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent separable data. The Journal of Machine Learning Research, 19(1):2822-22) 28788 2018
381
382
[48] Valentin Thomas, Fabian Pedregosa, Bart Merriënboer, Pierre-Antoine Manzagol, Yoshua Bengio, and Nicolas Le Roux. On the interplay between noise and curvature and its effect on optimization and generalization. In International Conference on Artificial Intelligence and Statistics, pages 3503-3513, PMLR, 2020.
383
12

---

396
[49]
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. PMLR, 2020.
387
398
[50] Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in ver-parametttized learning: A dynamical stability perspective. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 8289-8298, 2018.
400
401
[1 Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint
2018.
404
[52] Sho Yaida. Fluctuatioon-disiipaioon relations for stochastic gradient descent. In International Conference on Learning Representations, 2019. URL htpps://///peeereeeen net/forum?
405
408
id=SkNksoRctQ
407
[53] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE International Conference on Big Data (Big Data), pages 581-590. IEEE, 2020.
408
409
410
[54] Yang You, Igor Gitman, and Boris Ginsburg, Scaling sgd batch size to 32k for imagenet training. arXiv preprint arXiv: / 708.00888, 2017.
411
[55] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks.
arXiv preprint
arXiv:1605.07146, 2016.
414
[56] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):
415
107-115, 2021.
417
[57] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger Grosse. Which algorithmic choices matter which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019.
418
419
420
(58) 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.
421
422
13

---

Checklist
423
424
(a) For all authors...
425
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
426
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] We try to avoid theoretical analysis based on impractical assumptions. Therefore, some of our claims are supported by experiments and may require further theoretical investigation.
428
429
430
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
431
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
432
433
(b) If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Prop. 4.1
435
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix B
436
(c) If you ran experiments..
437
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes]
438
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
441
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [N/A]
442
443
(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)? [N/A] We do not propose any new algorithm which requires to report the computational cost.
444
445
446
(d) If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
447
(a) If your work uses existing assets, did you cite the creators? [Yes]
448
(b) Did you mention the license of the assets? [Yes]
449
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
450
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes]
451
452
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
453
454
(c) 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]
456
457
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
458
459
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
460
14


Reviews:
[{"review_id": "ReaYK30t2kK", "reviewer": "Reviewer_xxYn", "summary": "The paper studies the instability in SGD optimization dynamics (edge of stability) and how that impacts generalization. In particular, the authors propose the idea of interaction-aware sharpness, which I believe is original. Based on a definition of implicit interaction regularization, the authors propose the linear and saturation learning rate scaling rule (LSSR). Most of the findings are supported by empirical evidences but theoretical support is at an intuitive level.", "questions": "* My feeling is that the current theoretical setup may explain the gap between batch optimization and minibatch optimization but that is not the same as explaining generalization because generalization depends on the model class size. Can you agree or disagre?\n* More specifically, do I understand correctly that the theoretical setup does not assume anywhere that the model is a neural network? Then why are all the experiments on neural networks?\n* Can we motivate the linear and saturation scaling rule better? I am not sure why keeping the right-hand side of Eq. (11) constant leads to optimal scaling.", "limitations": "N/A", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n * The idea of interaction-aware sharpness seems to be original.\n * Empirical evidence to support the self-regulation of interaction-aware sharpness around the edge of stability.\n * Intuituive argument for learning rate scaling.\n\nWeaknesses:\n * Learning rate scaling only supported by a vague argument and empirical evidence.\n * Weak connection between the concepts proposed in the paper and generalization.\n * Nothing in the theory related to neural networks.", "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": "KqI0iGGtgC_", "reviewer": "Reviewer_fFDo", "summary": "This paper makes several contributions:\n 1.  The authors propose a new characterization of 'edge of stability' which is said to improve over the original characterization from Cohen et al. (2021) in that it (1) fits full-batch GD better, and (2) generalizes to SGD.  Specifically, whereas the original characterization is that the sharpness (maximum Hessian eigenvalue) hovers **just above** the value $2/\\eta$, the authors propose characterizing the EoS as a regime in which the expected loss difference that is predicted by a quadratic taylor approximation hovers **around** the value zero (here, the expectation is wrt randomness in minibatch sampling).  For SGD, this condition comes out to $\\frac{\\text{tr}(H S_b)}{\\|g\\|^2} \\approx \\frac{2}{\\eta}$, where g is the (full-batch) gradient, H is the Hessian, and $S_b$ is the second moment matrix of minibatch gradients.  For full-batch GD, the condition reduces to $\\frac{g^T H g}{\\|g\\|^2} \\approx \\frac{2}{\\eta}$.  In the case of full-batch GD, the authors argue that their condition $\\frac{g^T H g}{\\|g\\|^2} \\approx \\frac{2}{\\eta}$ improves over the original condition $\\|H\\|_2 \\gtrapprox 2/\\eta $ in that  $\\frac{g^T H g}{\\|g\\|^2}$ empirically hovers right around $2/\\eta$ (Figure 4), whereas there is usually a gap between $\\|H\\|_2$ and  $2/\\eta $.  More importantly, it is currently an open question how the 'edge of stability' phenomenon generalizes to SGD, and the authors argue that their condition $\\frac{\\text{tr}(H S_b)}{\\|g\\|^2} \\approx \\frac{2}{\\eta}$ holds during SGD. \n  \n 2.  The authors investigate the behavior of gradient descent at the EoS (Figures 1-3).   They find, essentially, that gradient descent is leaping back and forth across an asymmetric valley (https://arxiv.org/abs/1902.00744).  At every iteration where  $\\frac{g^T H g}{\\|g\\|^2}$ is above $2/\\eta$, it dips below $2/\\eta$ on the next step.  At every iteration where $\\frac{g^T H g}{\\|g\\|^2}$ is below $2/\\eta$, it rises above $2/\\eta$ on the next step.  Note that a similar point has been made (with no experiments) in the contemporaneous work of Chen & Bruna (https://arxiv.org/abs/2206.04172).\n\n 3.  The authors propose a new scaling rule that describes how the linear rate ought to be scaled with the batch size in order to preserve the same implicit regularization.  Namely, the authors suggest that two training runs have the same implicit regularization whenever $\\text{tr}(S_b) / [\\eta \\| g\\|^2]$ is the same, and they present evidence (Figure 5) in support of their hypothesis.  Their scaling rule reduces to the well-known linear scaling rule in the limit of small batches, and reduces to \"no scaling\" in the limit of large batches.", "questions": "- In Figure 5, where does the yellow dotted line come from? \n\n- The plots of $||H||$ in Figure 31 look weird to me -- typically, the sharpness along the optimization trajectory is smaller when the batch size is smaller.  Could there be an error here?\n", "limitations": "Yes", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "In what follows, I will reference the contributions I listed above by their numbers.\n\n- Weakness of contribution #1: I could be wrong, but I don't think the new EoS characterization will generalize to momentum gradient descent (either with Polyak-style or Nesterov-style momentum).  In the Cohen et al EoS paper, it is intuitively clear why the sharpness should hover just above $2/\\eta$: if the training objective is modeled by its quadratic Taylor approximation, we see that gradient descent cannot linger for long in any region where the sharpness exceeds $2/\\eta$, because it will oscillate with exponentially growing magnitude in a certain direction until it leaves the region.  The extensions in that paper to Nesterov and Polyak momentum are based on the same reasoning.  By contrast, the underlying logic behind the new proposed EoS characterization (which is based on the idea that the loss is always non-monotonic) is less clear to me, and one consequence is that I don't see how it can be generalized to momentum.\n\n- Strength of contribution #1: the proposed extension of EoS to SGD is interesting (note that it's basically a quadratic Taylor approximation version of the phenomenon reported in Cohen et al (2021) appendix H).  Nothing has been published in the literature addressing the question of how EoS might generalize to SGD, so any experimental result that seems to hold consistently across a wide range of networks is valuable.\n\n- Strength of contribution #2: I think that these experiments will be useful for people interested in how gradient descent manages to remain semi-stable at the EoS.\n\n- Weakness of contribution #3: both $\\text{tr}(S_b)$ and $\\|g\\|^2$ will vary over the course of training, and I don't see any reason why their ratio should remain constant.  Thus, the scaling rule is not even well-defined.\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": "CX6bRxAVzGG", "reviewer": "Reviewer_etxE", "summary": "The authors explore the notion of edge of stability for SGD dynamics, they introduce some quantities and explain / illustrate why they make sense for the understanding of the dynamics. They then propose a novel \"linear and saturation scaline rule\" and provide experimental support.", "questions": "- eq (3): over what are you taking the expectation ? if it is the full expectation then you should also have $\\mathbb{E}[L_t]$, I guess you are taking the conditional expectation but this is not clear. \n- line 96: a bit weird to name a set a \"regime\". It would also be clearer to keep the $\\theta$ index in $H$: $H_\\theta$, to make clear that the hessian depends on $\\theta$.\n- unclear experiments in fig 1: what is the precise setting ? what dataset ?\n- l110: \" This result implies that the training loss L(θ) is approximately locally quadratic, i.e., ε ≈ 0, in the early phase.\" Why and what does this mean ? isn't any twice differentiable function approximately locally quadratic ?\n- \"We hypothesise that due to this non-quadraticity of the training loss, the iterate is discouraged from staying within the unstable regime.\" I don't see why this is due to non-quadraticity. If you want to optimise $f(x, y) = |y| x^2$ for example, which is always \"locally quadratic\": then for a large step-size, the iterates are \"discouraged from staying within the unstable regime\" too.\n- Figure 3: what is $\\Delta L$ ?\n- Figure 4: very hard to read labels, especially in fig (d)\n- equation 12: \"if $b$ is small\", \"if $b$ is large\": how small ? how large ?", "limitations": "N/A", "rating": 3, "confidence": 3, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I believe there are some interesting ideas and observations. For example Figure 5 (right) (which illustrates the linear and saturation scaling rule) is rather convincing. Some quantities which are introduction seem to make sense. However the paper and the claims are not at all convincing enough. The abstract and introduction are very superficial and it is not even clear what the authors want to prove / show. The related work is neglected. \n\nI am also bothered by some very vague and imprecise claims and sentences. The title to start with: \"Implicitly regularized interaction between SGD and the loss landscape geometry\". I don't understand what this means: what is a \"regularised interaction\" ? \nThen in the text: \"we find that SGD induces an implicit regularization on the interaction between the gradient distribution and the loss landscape geometry.\", \"we find that SGD implicitly regularizes the interaction-aware sharpness\". I do not manage to understand these sentences, they are very vague and the meaning is not clear at all. \n\nThe fact that the paper is not well written with very unprecise statements and claims makes me heavily doubt on the sanity and significance of the presented results. In the following questions I give a few examples to support my opinion.", "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: Reject