A new characterization of the edge of stability based on a sharpness measure aware of batch gradient distribution
Authors: Sungyoon Lee
Cheongjae Jang
Venue: ICLR 2023 poster (2023)
Keywords: edge of stability
SGD
learning rate
batch size
optimization
generalization
implicit bias
implicit regularization
sharpness
scaling rule

Abstract: For full-batch gradient descent (GD), it has been empirically shown that the sharpness, the top eigenvalue of the Hessian, increases and then hovers above $2/\text{(learning rate)}$, and this is called ``the edge of stability'' phenomenon. However, it is unclear why the sharpness is somewhat larger than $2/\text{(learning rate)}$ and how this can be extended to general mini-batch stochastic gradient descent (SGD). We propose a new sharpness measure (interaction-aware-sharpness) aware of the \emph{interaction} between the batch gradient distribution and the loss landscape geometry. This leads to a more refined and general characterization of the edge of stability for SGD. 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.

Published as a conference paper at ICLR 2023
A NEW CHARACTERIZATION OF THE EDGE OF STA- BILITY BASED ON A SHARPNESS MEASURE AWARE OF BATCH GRADIENT DISTRIBUTION
Sungyoon Lee & Cheongjae Jang Hanyang University, Seoul
.ac. (sungyocnlee, cjjang) ) @hanyang ac kr
ABSTRACT
For full-batch gradient descent (GD), it has been empirically shown that the sharpness, the top eigenvalue of the Hessian, increases and then hovers above 2/ //(eerning rate), and this is called "the edge of stability" phenomenon. However, it is unclear why the sharpness is somewhat larger than 2/(1carning rate) and how this can be extended to general mini-batch stochastic gradient descent (SGD). We propose a new sharpness measure enteraationnwwar-shaaaee interaction between the batch gradient distribution and the loss landscape geometry. This leads to more refined and general characterization of the edge of stability for SGD. Moreover, based on the analysis of 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.
aware of the
INTRODUCTION
For full-batch GD, it has been empirically observed that the sharpness, the top eigenvalue of the Hessian, increases and then hovers above 2/(lernning rate) (Cohen al. 2021 the training proceeds. This observation can provide link between two empiniall results regarding generalization. (i) using larger learning rates for GD can generalize better (Bjorek al., 2018 Li et al.) [2019b] Lewkowycz. et al., 2020 |mmith et al. 2020) and (ii) minima with low sharpness tend to generalize better (Hochreiter & Schmidhuber 1997 Keskar et al., 2017). This observation has a significant implication in existing neural network optimization convergence analyses since it is contrary to the frequent assumption that 'the learning rate is less than 2/3 (here 3 is an upper bound the Hessian top eigenvalue)", which ensures the decrease in the training loss (Nesterov) 2003 Schmidt 2014 Martens] 2014, Bottou et al., 2018). Even though the training loss evolves non- mon-monotoniaally over short timescales due to the violation of the assumption, interestingly, the loss is observed to decrease over long timescales consistently. This regime in which GD typically occurs has been referred to as 'the edge of stability (EoS)' (Cohen et al. 2021).
There remain many aspects that are not clearly explained about the EoS regime. For example, it is not clear why and to what extent the sharpness hovers above 2/ /leearning rate). Moreover, the inherent mechanism is not yet elucidated for the unstable optimization to occur at the EoS consistently while prevented from entirely diverging. How this phenomenon can be generalized beyond GD, especially to mini-batch SGD, is still an open question.
In this paper we provide new characterization of the EoS for SGD, which can serve as an answer to the above questions. As tool to analyze the optimization process of SGD, we first propose a sharpness measure of neural network loss landscape aware of SGD batch gradient distribution (hence capturing the interaction between SGD and the loss landscape), which we refer to as the interaction-aue aaarrttiin-aaar--saappeeeee (IAS) (Section 2 Based on this measure, we define the stable and unstable regions in the neural network parameter space. We then scrutinize both theoretically and empirically the transition process of the iterate from the stable to the unstable region (Section I and the mechanism to escape from the unstable region, i.c., how the optimization can occur at the EoS. We interpret the latter mechanism based on the non-quadraticity of the loss and the presence of asymmetric valleys in the loss landscape (He et al. 2019) (Section 4.2)

---

Published as a conference paper at ICLR 2023
Based on these analyses, we propose the notion of implicit interaction regularization (IIR), i... the IAS is implicitly bounded during SGD, as an implicit bias of SGD (Section 4.3). The value that IAS is bounded by is the ratio of concentration measure of the batch gradient distribution of SGD the learning rate. This is a more refined characterization of the EoS, as it shows that IAS does not hover above certain value, but rather hovers around. More importantly, it can be naturally applied to SGI) since we do not make any impractical assumptions on the batch size or learning rate.
Our new characterization of the EoS leads to a novel scaling rule between batch size and learning rate, from the idea of preserving a similar level of IR (Section [5) This scaling rule, referred to as the Linear and Saturation Scaling Rule (LSSR), recovers the well-known linear scaling rule (LSR) (Jastrzebski et al., [2117; Masters & Luschi 2018: Zhang et al.) 2019 (Shallue et al.] [2018] Smith et al. 2020 2021) for small batch sizes and reduces to no scaling (due to saturation) for large batch sizes.
2 GRADIENT DISTRIBUTION AND LOSS LANDSCAPE
In this section, we review some concepts required for further discussion. See Appendix A for quick reference for the notations. To simplify the notations, we often omit the dependence on some variables and the subscript of the expectation operation when clear from the context.
For a learning task, we use parameterized model with model parameter 0 E O C R Then we train the model using training data D {>x) and a loss function [(x; 0). We denote the (total)
training loss by L()) = Σ =-=1 (x;;; 0) for training data D. At time step t, we update the parameter
9, using GD: 0t+1 = 8t (o,,)) with a learning rate r) > O, or using SGD: 8(41 = 02 - (6s) with a mini-batch gradient (0)) = EeEB Vol(2; 0₂) E R for a mini-batch B C D of size b (1 b < n). Here, we the subseript & to denote the random batch sampling procedure.
Now, we are ready to introduce some important matrices, Co, Sb. and H. First, we define the covariance (0) Varg (0) Eg (ge (0) Eglg (0)(0) Eg ls (0))) E <m and the second moment (0) Eg [g (0) E IR m of the mini-batch gradient (0) over batch sampling for a batch size S S n The covariance G, and the second moment S, satisfy not only C S. Sn but also the following equation (Hoffer et al. 2017 Li et al. 2017 Wu et al. 2020):
C Vn.b (S1 Sa) Yn. C1, where Tn,6
for sampling without replacement (1)
R
for sampling with replacement
We provide a self-contained proof of (T) in Appendix B.I We note that, for sampling without replacement, many previous works approximate Tooh & I assuming b << (Jastrzebski et al.
Hoffer et al., 2017; Smith et al. 2021), but we consider the whole range of b and Tn.n 0). Second, we define the Hessian (0) V2L(0) (0 Vn. 1 with Yx.1 mm and denote the i-th largest cigenvalue and its corresponding normalized 1 Vzl(z,;)) 0)
2017,
eigenvector by (H) ) E IR and (H) E TR,, respectively, for i = 1, ... m. The operator norm |||| = suP Hul] of H is equivalent to the top eigenvalue X1. We emphasize that C, and S, represent the stochasticity of the batch gradients, and H represents the loss landscape geometry.
Therefore, we can write one of our goals as follows: we aim to understand how the loss landscape geometry (H) and the gradient distribution (Sy) interact with each other during SGD training. We investigate this "interaction" in terms of matrix multiplication S,, To be specific, we consider the trace St) and its normalized value H.( and we call the latter interaction-awee aware sharpness:
tr(S,)
Definition (Interaction- Sharpness (IAS)).
tr (.))
| ||,,
(2)
(r(Ss)
Here, tr(H...) < ||]]]] tr(S)), i.e., ||, < || [],, and the equality holds only when every de is aligned in the direction of the top eigenvector 9i of H.
'These two matrices CG and Sy are often called the second central and non-central moments, respectively. But avoid confusion, we the term "second moment" only for the non-central Ss.
2

---

Published as a conference paper at ICLR 2023
3 RELATED WORK
Some studies investigate the interaction between the gradient distribution and the loss landscape
geometry represented by tr( St) in the context of escaping efficiency (Zhu et al. 2019 Section
3.1), stationarity (Yaida 2019) Section 2.2), and convergence (Thomas et al. 2020 Section 3.1.1). However, they require some additional assumptions like stochastic differential equation (SDE) approximation of SGD (Zhu et al. 2019). the existence of a stationary-state distribution of the model
parameter Yaida 2019] Section 2.3.4), and strong convexity of the training loss function (Thomas et al. 2020), respectively. In this paper. we provide a new insight into the interaction (H.)) without these assumptions.
Convergence of full-batch GD (b n) has been instead analyzed with an upper bound on the interaction tr( (H,) with further assumptions for the stable optimization, such as 3-smoothness
of the objective and 0 < <
3 (e.g.. = 3) (Nesterov) 2003; Schmidt 2014 Martens 2014
Bottou et al. 2018 However, it may lose useful information of the interaction between H and Sn. Moreover, when we train standard neural network with GD in practice, || B) increases in
the early phase of training and the iterate enters the EoS where | ~ 2 " i.e., 77 & |i]] 2 2 This
2
3-
contradicts with the assumption for stable optimization and the iterate exhibits unstable behavior
with non-monotonicaly decreasing loss (Xing et al.) 2018 Wu et al., 2018 Cohen al. (021). We
further extend this discussion of unstable dynamics for GD at the EoS to the case of SGD.
From the generalization perspective, many studies focus on the implicit bias of SGD toward better generalization (Neyshabur 2017 Zhang et al. 22021] Gunasekar et al.) (017]: (oudry et al. 2018 Jastrzębski et al. 2020 [021]] Barrett& & Dherin [021]] | /mmihh et al., 2021). There are mainly two
factors known to correlate with the generalization performance: the batch gradient distribution during training (Hoffer et al., 2017 Jastrzzbski et al.) 2017] /mith Le| 2018 Zhuu et al. 2019) and the sharpness of the loss landscape at the minimum (Hochreiter & Schmidhuber 1997 Keskar et al.)
2017; Dinh et al. 2017: Diang et al. 2020 Foret et al. [021] Kwon et al.) (202I) I) We provide a link
between the batch gradient distribution and the sharpness that the model implicitly regularized to
have a low sharpness when the second moment of the batch gradient is large (see Section
We provide further discussion reconcile our arguments with some previous studies in Appendix D
4 A NEW CHARACTERIZATION OF THE EDGE OF STABILITY
4.1 UNSTABLE OPTIMIZATION
Using the second-orde Taylor expansion of the total training loss L(0) at 9G, the change in the loss L, L(0)) the SGD iterate moves from 0, to 0.41 at time step can be expressed as follows:
1? DE H + 0(1
L411 L, = L 96 +
(3)
2
where 8, = 0 - 0, -nn.. Given 6,, the expected loss difference over batch sampling is
Ee|Le+1] - Lt
(4)
- Eelgel + 1
Lg H(::)ge) +e
(Taking Eg for the both sides of 3
(5)
2
n/2
+
ir (Eg (H..))))) te
(Eglge) VL and u (r(vuT)
(6)
2
n (r((.)
+
(by definition of S, and Sn),
2
tr((s))
where E ( I1)) and EE [X] is the conditional expectation of X given O.. We consider bounded region with finite maximum loss near and including 8₁ in which the loss is approximately quadratic. Then, when the following instability condition is met within the region, the loss (close to the expected loss) tends to increase and the iterate tends to escape from the region:
2 LL((++1) - L(O,) S VL (lt+1 (0)+1-0)++ + |0++1 1 - 1,l ~
2n |LL = - -n(1 fn [|
and thus the loss monotonically decreases when 17 < 3.
3

---

Published as a conference paper at ICLR 2023
1e-2 2 CE. b 128.n 0.02
1e-2CE, 256. r 0.02
1 11- LD/tris))
[E(4e.1] -.ittrss))))
100
2-
slope-?
-80
2
slope-
step
step
0-
0
-80
-2-
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
1e2
MiH5
1e2
TRAAT
1e-2
CE, GD, n 0.02
1e-2
MSE, GD, 0 0.02
120
slope
slope
270
115
0
260
-2
100
110
-2
100
250
0.00
0.75
1.50
2.25
0.75
1.50
2.25
3.00
||M||s.
1e2
||M|ls.
1e2
[An empirical validation of (7) for SGD (top) and 12 for GD (bottom)) In the early phase, until the iterate enters the unstable region, it validates 7 and (12) with the blue line with the slope 5/2 and x-intercept 2 For GD (bottom), they are plotted after | exceeds 2 after which ||lS. starts to increase from to point with 'x" when the iterate enters the unstable region. right), we plot the graph for a few more steps after the iterate enters the EoS, We train 6CNN on CIFAR-10- 10-8k with = 0.02 (see Remark at the end of Section 4.2).
Figure 1:
7 few steps. For cross-entropy (CE) loss, we mark the end
mean squared error (MSE) loss (bottom
Theorem 1. For SGD on quadratic L, the expected loss increases, i.e., Eg|Le+1] - L, > 0, if and only if 0g satisfies the instability condition sr(HSt) > normally distributed, then the following inequalities hold for any positive 0:
tt(S.)) " Furthermore, if the batch gradient gę
P [1t1 TT+il 2 /28r 7mbb |
01)
<exp(-xr)
(8)
EglLe+1] --V288 0.) S Sexp(-xi) (L++1
where B = 1/'in (v (T
n' tr( tr(HC, and v = HHVI.
2b
The proof is deferred to Appendix region:
B.2
From the above theorem, define unstable and stable
tr(H.S)
2 and S U Ir(S,) "/
== (0 @
(10)
respectively. It has been empirically shown that, for full-batch GD, ||]|| increases and then hovers above 2/1 and Cohen et al. (2021) mark the EoS with {0 E θ [0)] = 2 ). but we mark with
tr(H 2
OS (0 θ
(11)
tr(Sw)
We emphasize the superiority of our Os over the previous EoS that (i) 8S provides a clearer "edge" since it considers the interaction between the gradient direction and the Hessian and that (ii) is more general since it applies to SGD with any batch size 1 < b < n.
Figure (top row) empirically validates 7 showing the normalized loss difference
Ee bs (r((S.)
against (r(H)s) in the early phase training before entering the unstable region. This result implies
(r(S.))
that the training loss T.()) is approximately locally quadratic, i.e., R 0, in the early phase.
Especially, for full-batch GD (b n), the instability condition can be rewritten as || ||s,, and we have the following relationship between the loss difference L - Tu and H s. from 7
T - I = y
tr(Sw)) [IFA] 17
2 +
(12)
4

---

Published as a conference paper at ICLR 2023
1e- CE, 128. 0.02
1e-2CE, b 256. 0.02
1 11 LD/tris)))
2000
[E(c.1]]- Liitils))))
2000
2-
slope
2
slope
1500
1000
1000
0-
0
-2-
100
500
-2
- 100
500
0.00
0.75
1.50
2.25
3.00
0.00
0.75
1.50
2.25
3.00
NCKS)
1e2
le2
1e-2 CE, GD, 1) 0.02
1e-2 MSE, GD, n 0.02
1500
1500
(L--1 1 La)tr((,))
Ø slope
slope
step
step
0
500
-2
100
-2
100
0.00
0.75
1.50 |ll
2.25
3.00
0.75
1.50 ||M|s
2.25
3.00
1e2
1e2
Eg Ee b,
Figure 2: (Non-quarraticity and overestimation] The normalized loss difference
tr(S..))
against (r(HSs) during training. After the iterate enters the EoS, it often shows more gentle slope tr(S.)
than 2 especially in the unstable region. See the caption of Figure
Figure (bottom row) shows | soars from 0 in few steps after ||]] exceeds 2
n satisfying
with R3 0, before the iterate enters the unstable region. This result is consistent with the following theorem for a quadratic loss L and generalized momentum GD with (2);
8, -:8..-11 L(. B28.-1);
(13)
0c+1 = 0, 0, 8.,
(14)
where $1, B2 E 1). Here, we have vanilla GD when 31 = 32 = 0, Polyak momentum when 31 (0, 1) and B2 = (Polyak 1963), and Nesterov momentum when 81 = 33 (0, 1) (Nesterov
1983). We will focus on the vanilla GD as it can be easily extended to the generalized momentum variants (see Appendix (f.oodixC.6fforr details). The proof deferred Appendix [.3
Theorem 2. For generalized momentum GD with (31, on a quadratic L, i 1+81
(2) < Au for all i * 1 where 7(1,, .2) = I+23,' then qi 5z oscillates and diverges with the exponential growth of |u] &el = O(eat) for some e > Moreover cos((n, 8,)1 and [| increase to I and λy. as t co, respectively, with 1 1- cos(g) 8,)1, |H]| = =0 2nt).
Note that 7(31, B2) for vanilla GD. To summarize, if exceeds 3, then || increases
towards H|| with the exponential convergence and also exceeds v) in few steps, i.e., the iterate enters the unstable region U. Together with Theorem [ if consider bounded subregion V C U with a finite maximum loss near and including 0, in which the loss is approximately quadratic, then
the iterate tends to escape from the region (Nar & Saatry 2018 Wu al., 2018 Cohen al.
2021).
4.2 NON-QUADRATICITY, ASYMMETRIC VALLEYS AND THE EDGE OF STABILITY
In the previous section, we have shown that the training loss is approximately locally quadratic before the iterate enters the unstable region. However, after the iterate enters the unstable region, i.e., tr(HS,)
((S..) reaches and exceeds 2, the step size is relatively large for the sharp loss landscape that the iterate jumps across the valley (Jastrzşbski et al. 2019), and the higher-order terms in (7) and (12)
become non-negligible and cause a different behavior of the iterate than in the stable region.
Figure shows empirical evidences for the non non-quadraticity. After the SGD/GD iterate enters the
unstable region, when the instability condition (xHSss) > is met, the normalized increase in the
tr(S.)
[g[[<+1]-44 is often smaller than 7² cr(H.ss tr(S.,)
from 7 and (2) (blue line), which the
loss
tr(S..)
case for a locally quadratic function. This results in gentle slope less than 7
5

---

Published as a conference paper at ICLR 2023
||M|ss 7
||M||s, i
||Hl|s, =
||H||s, =
10
Hldds IMIls
10
AL
10:
0.25
0.50
0.75
1.00
1.25
min
6.
zero
Figure 3: [Asymmetric valleys] Left: The ratio ||H.Il.. where Ha = H(0 V()) for ||
1 x 31,2,3,4, 4,5] for each t during training. When |s... < 2 (red), |HH is usually 17 (blue), ||H is usually smaller than || [ls. Right: The training loss difference along the gradient descent direction, for each Q.. Each plot is normalized and translated to have the same minimum value and the same zero where A AT(ee))) T.(zero) T.(O,) = O. We also plot the quadratic baseline (cyan dashed curve). When |s]... < 2 (red), it usually becomes sharper across the valley (right-shifted), while the opposite is observed (left-shifted) when ||H]s... > (blue). We train 6CNN using GD with 17 0.04.
larger than || On the other hand, when || > 2
We hypothesize that due to this non-quadraticity of the training loss, the iterate is discouraged from staying within the unstable region. Note that, for globally quadratic loss, when the iterate in the unstable region, it diverges within the unstable region. Figure F demonstrates the asymmetric valley He et al. 2019) that one side is sharp and the other is flat. In Figure (left) we evaluate the directional sharpness [o s. along the gradient descent direction L()) where H. = (0- an V L()) for a E ] x [1, 2, 4, 5], and compare (lHllls.. (0) with [|s. At the sharp side, it has a high ||..... > 2 (blue) with the gradient V and the top eigenvector (H) of the Hessian being highly aligned (cf. Theoremp] However, when the loss landscape gets far from being quadratic, the Hessian and its top eigenvector can change abruptly, 91 (Ha) would not always be aligned with qı (H) and (L))), and l.a tends to 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. Figure 3 (right) similarly shows that when the iterate is at sharp side of the valley, tends to jump to the other side of a flatter area, and vice versa.
To summarize, make the following observations for GD in order: (i) || increases in the beginning (the progressive sharpening (Jastrzebski et al., 2019 2020 Cohen et al. 2021)), (ii) [[
exceeds 2, (iii) the gradient L becomes more aligned with the top eigenvector 91 (/I) in a few
2 steps, (iv) ||s. reaches the threshold 7) and the iterate jumps across the valley, (v) || Hs. tends to decrease due to the non-quadraticity, and it repeats this process, while ||, oscillating around 2. We observe a similar behavior with oscillating tr(HSss around 2 for SGD. It requires further investigation into the exact underlying mechanisms (e.g., the progressive sharpening) (Arora et al.)
'1'
tr(S..)
7
2022
Li ct al. 2022 Damian al. 2022
Zhu et al. 2022) and we leave it as future work.
Remark (Experiments in Section 4.1 and 4.2 We 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) on CIFAR- 10-8k where DATASET-n denotes a subset of DATASET with |D] = n and k = 210 = 1024. See Appendix C.1 C.1 for the results from other datasets, learning rates and networks (ResNet-9 with m 2.3. (He al. 2016 and WRN-28-2 with m 36M (Zagoruyko & Komodakis 2016) where m = dim()))
4.3 IMPLICIT INTERACTION REGULARIZATIO (IRR)
In the previous sections, we have shown the SGD iterate is implicitly discouraged from staying in the unstable region. Now, ready investigate this property from the regularization perspective. First, to understand the effect of batch size b on the batch gradient distribution, we define the following concentration measure ps:

---

Published as a conference paper at ICLR 2023
ResNet-9,
n.0.00
ResNet-9, n= 0.02
800
400
train loss
|M
[[|]]
700
train
[P1]
350
500
250
300
150
0.5
200
0.5
300
50
1000
1500
2000
0.0
1000
1500
2000
(a) GD, ResNet-9, cross-entropy, 0.01
(h) GD, ResNet-9, cross-entropy, 0.02
b-222
D-266
D-
b-24
2-
b-
D-
b=
b-
30.0
60.0
25.0
20.0
|HH|s,/
10.0
20.0
5.0
0.0
125
250
500
0.0
375
500
Step
Step
(c) SGD, ResNet-9, cross-entropy, 0.08
(d) SGD, ResNet-9,, cross-entropy, 17 0.08
Figure 4: [Clearer indication of the EoS] (a)-(b): After few steps of full-batch training, ||]]] (blue) hovers above 2 (Cohen et al. 2021), but ||s.. (red, defined in (2b) oscillates around (red dashed horizontal line). The EoS is more evident in the latter (red). We also observe sharp increase in ||s, right after || exceeds 2.
" Curves are plotted for every step. We train model on CIFAR-104 10-8k (n 2 using cross-entropy loss with = 0.01/0.02, respectively. (c)-(d); We plot curves (c) | |ss, and (d)
|ل when trained with various b. After a few steps, the curves in (c) reach the threshold (see (d) together) which increases as b becomes larger when < 213
and saturates to 2pm x when b is large. Curves are smoothed for visual clarity.
Definition 2 (a concentration measure of the batch gradient). We define Ps the ratio of the squared norm of the total gradient || 2|2 to the expected squared norm of the batch gradients Eellle || i.c.,
|
tr(Sn)
Pb Eg 11 tr (r()))
(15)
is similar the square of the mean resultant length 12, = || Ioe ]1| < of the batch gradient d (Mardia et al. 2000). especially when stdę |||g:]] |] < E 9e (see Appendix C.5 for empirical evidences). Both ps and RC are concentration measures and have lower values when the batch gradients 9ę are more scattered. Therefore, it is natural to expect that the ratio is small for a small batch size b, and we will revisit this in more detail in the following section (cf. 17 We also note that pn. R Yin al. (2018) call 1/ (((p1)) the gradient diversity.
Here, we write 2 =1 [2 and thus the ratio ps
| Iec 1I
Eg
11 PP
95
Now, we rewrite the instability condition
tr( tr(S.)
4 (multiplying both sides by (b) follows:
2pb
H|s,,
(16)
From Theorem ! and the instability condition (16) implies that IAS H S, implicitly regularized
and bounded to be less than 20s, We name this Implicit Interaction Regularization (IIR). We argue that the upper constraint
2ps in IIR is crucial in determining the SGD dynamics. With low constraint,
"
7

---

Published as a conference paper at ICLR 2023
SGD strongly bounds IAS ||,, We also note that IIR affects not only the magnitude ||]]| 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.
Gur-Ari et al. (2018).
show that, for GD (p,, = 1), IAS ||,,, (red) oscillates around 2 and exhibits IRR. This result is consistent with Cohen et al. (2021) that []] hovers above for GD. This is because,
Figures
4(a) 4(b)
7)
as mentioned earlier. 2 = || S | and the equality holds only when the gradient VL and the 1) top eigenvector 91 of H are aligned, which is in general not the case. For this reason, IIR provides a tighter relation and more clearly identifies the EoS than [Cohen et al. (2021). These results are also consistent with Theorem that ||//|s. suddenly increases from to in a few steps after [|/||
17 (see Appendix C.3 for more). Moreover, IIR also applies to a general SGD training with 1 < b < Figure 4(c) 4(d) show IIR for SGD with different batch sizes. The upper bound Bea of || [ss, according to IIR is higher when using larger batch size, but limited to less than 2 7
exceeds
(pb 1). We will further discuss this behavior with an investigation of Pb in the following section. "
LINEAR AND SATURATION SCALING RULE (LSSR)
In this section, we first introduce two previous scaling rules on how to tune the learning rate for varying batch sizes. Then we explain why they fail and propose new scaling rule based on IIR.
The ratio b/1 17 of batch size to learning rate 7) has long been believed as an important factor influencing the generalization performance, and the test accuracy has observed to be similar when trained with the same ratio b/1 = b' /1//. i.e., b' = kb and r/' ki for k > 0. This is called the linear scaling
rule (LSR) Krizhevsky 2014; Goyal et al., 2017 Jastrzebski et al. 2017, Smith & Le 2018 Zhang
ct al. 2019). They argued that LSR holds because lz+k 0₂ = - Σ 1 EnEBIi V Vel(z; 014s) Rd
2 D EeeB Vol(x; 0.) - E ExEBRecA Vol(a; 0c) assuming Vol(Or4s) Vol(O.) for 0 Isie < k where Blith = B and 1B = kb = b. However, the assumption is not accurate since the gradient oscillates mostly with a negative cosine value cos( ge (Oz), 9s (0++1)) < 0
between two consecutive gradients after entering the EoS (see Figure 24 25 in Appendix C.3
Moreover, LSR is known to fail when the batch size is large (Jastrzębski et al. 2017 Masters & Luschi 2018 Zhang et al. 2019: Smith et al. 2020 202T). On the other hand, Krizhevsky 40 et al..
(017)) proposed the square root scaling rule (SRSR) with another ratio vb/1 to keep the covariance
b of the parameter update constant for b < based on Varg [vgc] r Tn C1 C
However, Shallue et al. (2018) showed that both LSR and SRSR do not hold in general.
Based on the analysis of IIR with a new ratio 2pb/1 in the previous section, we explore why LSR fails in the large-batch regime and provide a more accurate rule to achieve similar generalization performance of the models trained with various choices of batch size and learning rate pairs (b. n).
To this end, we investigate the concentration measure Pb = tr(S,). tr( (S)). By combining two equations, C₁ Sb Sn (by definition) and C In (S1 Sn) in (I. we can obtain S.
b C. + Sn = "m. S1 + (1 6. )S,,, Therefore, we have tr(S,) which leads to the following equation:
T tr(St) + (1 "/.6 tr(n),
b
tr(Sm)
(r(S.)
b p ~ hp
if b is small
Pb tr(55)
Tn.b
tr((1) + (1 b. tr(()) b
Iosh h
+ (1
Tn
if is large
b
(+)
(17) from (15) where p = P1 =0:=trr(e tr(Ss)/ tr(S)). Note that p is (much) smaller than because Vol(z;) has different direction for each x; and tr(S.) | I| Σ. Vol(x:) < Σ, (o.)))))
tr(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/m. called the Linear and Saturation Scaling Rule (LSSR), with the two regimes that (i) Ps is almost linear when < n (linear regime) and (ii) Pb saturates when b is large (saturation regime). It depends which part of the denominator
8

---

Published as a conference paper at ICLR 2023
CEAA-CCO Sadier
24
LSSR: (1))
LSR: ((D) bp
2²
SRSR: (15) ppyb
log:/0)
20
Rb)
log.(0)
logu101
ST1-304L. Ramoodd
2
2-6
11gg 1o0)0o
-15
308
2-8
2
2
2
2
2
22
batch size (b)
kg:
log,
Figure 5: [Linear and Saturation Scaling Rule (LSSR)) Left: LSSR (red) in (7)) LSR (black dotted line) (Goyal al. 2017) and SRSR (blue dotted line) (Hoffer al. 2017). For LSSR, we can observe both linear and saturation regions (n 8k, p 2-7). Right: Heatmaps of test accuracy for models trained with large number of pairs of (b, 77) on CIFAR-I0-Sk CIFAR- 100-8- 101,-10-4k, and Tiny- Imane--aageN (from left to right, from top to bottom). It does not follow either LSR or SRSR, but tends to follow LSSR. We also plot 0o/11 1) = (yellow dashed curve) for some p and C' on each heatmap. Note that they are all log-log plots and thus a slope of indicates linear relationship.
(+) in dominates the other. First, when < n, then "/,b/b is not very small and the first term 26.h dominates the second term 1 -
26.6 since >>>> 1. Second, as b becomes large, on,bbb 0 and the second term (=r 1) dominates the first term. Thus, Pb saturates to and is not linearly related b, and LSR is no longer valid. The above arguments also hold for the batches sampled with replacement where the only modification is "h,b 1, vb in [7). Figure (right) empirically supports LSSR with the test accuracies when trained with various combinations of pairs (b, 17). To be specific, the optimal learning rate almost linear when is small, but saturates when b is large. We also plot pb/n= C' (the yellow dashed curve) for some and C which shows a theoretical prediction of pairs (n,b) b) that yield the optimal performance. Note that Figure 8 of Shalluu et al. (2018) Section 4.7) shows similar "linear and saturation" behaviors supportive of LSSR on other datasets (see also Figure Zhang et al. (2019 Section 4.3)).
Remark (Experiments in Section 4.3 and 5 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-I0 (Krizhevsky & Hinton 2009). CIFAR-I00 (Krizhevsky & Hinton 2009), STL-10 (Coates et al. 2011). and Tiny-ImageNet (a subset of ImageNet (Deng et al. 2009) with 3 64 x 64 images and 200 object classes). We use large number of epochs (800) and batch normalization (loffe & Szegedy 2015) to achieve a zero training error even with large b and a small r7. In the lower right corner (red area) of each heatmap in Figure [ (right), when b is too large 7) is too small so that |0++1 0.1 = vllggl || is too small, it requires an exponentially large number of steps for the iterate to enter the BoS. Thus, in this case, the assumption in Goyal al. (2117). (o,( x vg((0++)) for 0 < k. approximately holds and the reasoning LSR is valid. However, this only holds for a non-practical (b, 17) which shows suboptimal performance. See Appendix C.4 C.5 for the results from other networks and hyperparameters.
6 CONCLUSION
From an analysis of unstable dynamics of SGD and the instability condition, we clearly mark the edge of stability with the interaction-aware sharpness ||/||s and show the presence the implicit regularization effect on the interaction between the gradient distribution and the loss landscape geometry (II)). Moreover, introducing the concentration measure Ps the batch gradient, we link the second moment of the gradient distribution and the sharpness of the loss landscape. and propose new scaling rule called Linear and Saturation Scaling Rule (LSSR). Due to the simplicity of the analysis, hope that our insights will motivate the future work toward understanding various learning tasks.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
S. Lee is also the lead author and the corresponding author. This work was done while S. Lee was at Korea Institute for Advanced Study. S. Lee was supported by a KIAS Individual Grant (AP083601) via the Center for AI and Natural Sciences at Korea Institute for Advanced Study. C. Jang was supported in part by ITTP Artificial Intelligence Graduate School Program for Hanyang University funded by MSIT (Grant No. 2020-0-01373) and NRE/MSIT (Grant No. 20 This work was supported by the Center for Advanced Computation at Korea Institute for Advanced Study. This work was supported by the research fund of Hanyang University (HY-202330000000555))
REPRODUCIBILITY STATEMENT
We refer the reader to the following pointers for reproducibility:
Codes to reproduce some Figures: Supplementary Material.
Proofs of the claims: Appendix B
Experimental settings: Remarks at the end of Sections 4.2 and 5 and Appendix C
Additional Figures for other hyperparameter settings: Appendix C.I C.6
REFERENCES
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of
stability in deep learning. arXiv preprint arXiv:2205.09745, 2022. 6
David Barrett and Benoit Dherin. Implicit gradient regularization. In International Confer- ence on Learning Representations, 2021. URL https:://ppenreeweeee http://wwwwwwweeeeeeeeeeeeeeeeee net net/forummiciccm
3q5IqUrkcF
Lucien Birgé and Pascal Massart. Minimum contrast estimators sieves: exponential bounds and
rates of convergence. Bernoulli, pp. 329-375, 1998. 19
Nils Bjorek, Carla P Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normaliza- tion. Advances in neural information processing systems, 31, 2018. []
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-311 2018.
Pratik Chaudhari and Stefano Soatto. Stochastic gradient descent performs variational inference, con- verges to limit cycles for deep networks. In International Conference on Learning Representations, 2018. 42
Adam Coates, Andrew Ng. and Honglak Lee. An analysis of single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp. 215-223. JMLR Workshop and Conference Proceedings, 2011.
24
Jeremy Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2021. URL
20
Alex Damian, Fshaan Nichani, and Jason D Lee. Self-stabilization: The implicit bias of gradient descent at the edge of stability. arXiv preprint arXiv:2209. 15594, 2022. 6
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, pp. 248-255. Icee, 2009,
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
10

---

Published as a conference paper at ICLR 2023
Saber Elaydi. An Introduction Difference Equation. Springer, 01 2005. 20
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URLLtttss /ttp://////peeeeeeeeee net/ www.m/oormmmm/ http://orummmmmmmmmmmmeeeemmmeme
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet
in hour. arXiv preprint 1rXiv:1706.022 02677, 2017. 8
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapallii, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
Guy Gur-Ari, Daniel A Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv preprint 18X2::18/2.04744 2018. g
Haowei He, Gao Huang, and Yang Yuan. Asymmetric valleys: Beyond sharp and flat local minima. Advances in neural information processing systems, 32, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
24
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. Advances in neural information processing systems, 30, 2017.
T6
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. 42
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448 456. PMLR, 2015. 9
Leon Isserlis. On formula for the product-moment coefficient of any order of a normal frequency distribution in any number of variables. Biometrika,
1918. 17
Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. Three factors inftuencing minima in sgd. arXiv preprint arXiv: 1.1V:1771..0622
2017. 2 42
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 Ihttps://opett www.opeeevviiieee ew
net./ meet/forumiion d= j05t 7 7
Stanistaw Jastrzębski, 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, 2020. URL https://
openreview net/
Stanistaw Jastrzębski, Devansh Arpit, Oliver Astrand, Giancarlo B Kerg, Huan Wang, Caiming Xiong, Richard Socher, Kyunghyun Cho, and Krzysztof J Geras. Catastrophic fisher explosion: Early phase fisher matrix impacts generalization. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 4772-4784 PMLR, 18-24 Jul 2021. URL https: htttps::/pprceeedng: mlr. wwwpppessssvv.... /ttpe///wl////arteeetet.e
html
B
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where to find them. In International Conference on Learning Representations, 2020. http://ee//ppeeeeeeeeeeee//rrrrr

---

Published as a conference paper at ICLR 2023
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 Learning Representations. ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL htttts::///ppeeeeertrre
A. Krizhevsky and G. Hinton. Leamiing multiple layers of features from tiny images. Master's thesis,
Department of Computer Science, University of Toronto, 2009. 24
Alex Krizhevsky, One weird trick for parallelizing convolutional neural networks. arXiv preprint arXiv:1404.5997 2014.
Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi. Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks, 2021.
Beatrice Laurent and Pascal Massart. Adaptive estimation of quadratic functional by model selection. Annals of Statistics, pp. 1302-1338, 2000. [8
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 2020.
Chris Junchi Li, Lei Li, Junyang Qian, and Jian-Guo Liu. Batch size matters: diffusion ap- proximation framework on nonconvex stochastic gradient descent. stat, 1050:22, 2017. 2 16
42
Qianxiao Li, Cheng Tai, and Weinan. Stochastic modified equations and dynamics of stochastic
gradient algorithms i: Mathematical foundations. The Journal of Machine Learning Research, 20 (1):1474-1520, 2019a. 42
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, 2019b.
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. 42
43
Zhouzi Li, Zixuan Wang. and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharpening and edge of stability. arXiv preprint arXiv:2207. 12678 2022. 6
Stephan Mandt, Matthew Hoffman, and David Blei. A variational analysis of stochastic gradient algorithms. In International conference on machine learning, pp. 354 363. PMLR, 2016. 42
Stephan Mandt, Matthew Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18:1-35, 2017. 42
Kanti V Mardia, Peter E Jupp, and KV Mardia. Directional statistics, volume Wiley Online
Library, 2000. 7
James Martens. New insights and perspectives on the natural gradient method. arXiv preprint arXiv: 1412. 2014.
Dominic Masters and Carlo Luschi. Revisiting small batch training for deep neural networks. arXiv preprint 1XX4::8804.07111 2018. 2|
Kamil Nar and Shankar Sastry. Step size matters in deep learning. Advances in Neural Information Processing Systems, 31, 2018. 5
Yurii Nesterov. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate o
(1/k* 2). In Dokl. akad. nauk Sssr, volume 269, pp. 543-547, 1983.
42
12

---

Published as a conference paper at ICLR 2023
Behnam Neyshabur. Implicit regularization in deep learning. arXiv preprint arXiv: 2017.
Boris Teodorovich Polyak. Gradient methods for minimizing functionals. Zhurnal Vychislitel noi Matematiki Matematicheskoi Fiziki, ((4):643-653, 1963.
42
Mark Schmidt. Convergence rate of stochastic gradient with constant step size, 2014.
Christopher Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George Dahl. Measuring the effects of data parallelism on neural network training. arXiv preprint arXiv:1811.03600 2018. 2
Samuel Smith, Erich Elsen, and Soham De. On the generalization benefit of noise in stochastic gradient descent. International Conference on Machine Learning, pp. 9058-9067. PMLR, 2020.
Samuel L Smith and Quoc V Le. A bayesian perspective on generalization and stochastic gradient descent. International Conference on Learning Representations, 2018.
42
Samuel Smith, Benoit Dherin, David Barrett, and Soham De. On the origin of implicit regularization in stochastic gradient descent. International Conference on Learning Representations, 2021.
URL
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-2878, 2018.
Valentin Thomas, Fabian Pedregosa, Bart Merriënboer, Pitrre-Antoine Manzagol, Yoshua Bengio, and Nicolas Le Roux. On the interplay between noise and curvature and its effeet optimization and generalization. In International Conference on Artificial Intelligence and Statistics, pp. 3503-3513.
PMLR, 2020. 3
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. pp. 10367-10376. PMLR, 2020.
42
Lei Wu, Chao Ma, and Weinan E. How sgd selects the global minima in over- over-parameterized learning: A dynamical stability perspective. In Proceedings of the 32nd International Conference Neural Information Processing Systems, pp. 2289-8298, 2018.
Chen Xing. Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. walk with sgd. arXiv preprint arXiv:1802.08770, 2018.
Sho Yaida. Fluctuation dissssptionnn relations for stochastic gradient descent. In International Confer- Learning Representations, 2019. URL https://openreeiew... net forum?id
SkNksoRct
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),
pp. 581-590. IEEE, 2020. 24
Dong Yin, Ashwin Pananjady, Max Lam, Dimitris Papailiopoulos, Kannan Ramchandran, and Peter Bartlett. Gradient diversity: a key ingredient for scalable distributed learning. In International Conference on Artificial Intelligence and Statistics, pp. 1998-2007. PMLR, 2018.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146. 2016.
24
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.5
3
13

---

Published as a conference paper at ICLR 2023
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger В Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. Advances in neural information processing systems, 32, 2019.
Xingyu Zhu, Zixuan Wang, Xiang Wang, Mo Zhou, and Rong Ge. Understanding edgaoof-tt---s-- training dynamics with a minimalist example. arXiv preprint arXiv:2210.03299. 2022. 6
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise stochastic gradient descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning, pp. 7654-7663, PMLR, 2019,
14

---

Published as a conference paper at ICLR 2023
A NOTATIONS
We summarize the notations for a quick reference.
I EER
time step
0 E θ C I (or indexed 0e); d0m = model parameter
x E X' (or indexed n;)
training sample
D {a:)==ii [D]
training data
(x;0)
loss function (total) training loss
L(O) Σ ((2;00)
1 D
{(x; L. L(O,)
7 0
learning rate
D (or indexed B B²:)
batch
6 [B] 1 n
batch size
1 Σ
1 vl(x;0)
(0) b xEB:
v((x;0)~ =
batch gradient
[
.EBE
0+11 0. vV
GD
O+11 O. 7IE (0+)
SGD
= 0,+1 O,
displacement/velocity vector the covariance of the batch gradient the second moment of the batch gradient
(0) = Var; Lag (0)]
R mmxm
S,(0) = Eg [gg (0)g (0)
E Rmxm
for sampling without replacement
7n,6 =
sampling coefficient
for sampling with replacement
(0)) V (()) = EIND (gl(;; 0)) E mm
Hessian
1/2
||l||
Σ 4f
the Euclidean e -nnor of vector
Au
|||||| = sup
the spectral (operator) norm of a matrix A
|||||
A = q: (H) E Rm
X (H)
E IR
the i-th largest eigenvalue of the Hessian the corresponding i-th eigenvector of the Hessian the trace square matrix
tr( Σ Al
tr( Hb))
1s.
interaction-aware sharpness
(r()s)
tr(sw)
P6 (r6)
concentration measure of the batch gradient
tr(Sn)
P (T))
concentration measure of the per-example gradient
= (0 E 0 |||s. 2pb }
unstable regime
2pb
S = - (0 E e:: |H1|ss 5 17
stable regime
05 () E : | 2ph
the edge of stability
8, 32 E [0
hyperparameters for generalized momentum GD
15

---

Published as a conference paper at ICLR 2023
B PROOFS AND REMARKS
B.1 PROOF OF
We provide proof of (T) to make the paper self-contained. Similar proofs given in Li al.
2017): Hoffer et al. (2017); Wu et al. (2020).
Proof. We start with
Co = Belgegel L
(18)
VCEe/wwr T IVL - VC
-1-1
VC
(19)
n
VC [F/lww 1 1 n² 11
1
VL
(20)
(21)
6a VCV
(22)
where VL = /e..... Ven. E Rmmx & = E(x:;). ). the random vector w = [wi,... w/ [T each
element of which represents 8 "how many times the index is sampled in B and bw.
In case of sampling with replacement, we have v (1)) +... v(() where v(() represents sampling of single sample. Thus, Varg[u] bVar[u()). We have EB[(1)] 11 and
{
Pli BBy]
if
Elu (L)T Ik
(23)
Pli E B()), E B0y =
else
where [B(1)] 1. Thus,
Varşlu] bVarlu (1) = b (1 1 n 11
(24)
In case of sampling without replacement, we have IEg [u] = 1 and
Pli E Bg] 1.n-1
if
Eg[vv li.
Pli E B E B&]
C(n
2(n-2.b2)
b(b-1) wn=-I
else
(25)
C(.
where C(ni, r1) the number of rı mn-oomiinations from a set n2 elements. This leads to
Eg[wv ] = b(b - 1) 11 +
b
b6-1))
b(b
(26)
n(n 1)
n(n-1)
and
b²
b(b 1)
b2
b(n b)
Varg/u] = [ve]]
11
11 +
(27)
n²
n(n 1)
n²
n(n
b(b n)
b(n 6)
11
+
(28)
n²(n 1)
n(n
b(n - b)
1
11T
(29)
n
n
Putting the two cases together, from
22p.
(24g
and 29
we have
C 12
1
11
VCT
nn,b
V (VCI)(EC)T)
1
(30)
n
nn.b 1
VLVE VIALT)
(31)
71
Tn.b
(Si Sn),
(32)
b
16

---

Published as a conference paper at ICLR 2023
where
for the sampling with replacement
mn.b
(33)
for the sampling without replacement
B.2 PROOF THEOREM
First, expect the variance of the loss at the next iteration +++ 1) to be small for sufficiently large
Lemma 3. If the batch gradient ge is normally distributed and the loss L is quadratic, then
Tm. (TCCu=+
T/inb
Varg L [0,]
tr(HC,HC,1)
(34)
b
26
where = V
Proof. Since L is quadratic, we have the following:
n T Hgg
T1+11 La "9 DE
9
(g VI.
(35)
17
constant "1'" Tgg( (g =)
(g +))+ H(g+
(E~ Co~,0 C2))
(36)
constant - -ng - n/'g" HE + "¹
(37)
2
constant - (n((1 - () T +
HH,
(38)
2
Then, the variance of L4+1 is Varg| 10,]
Egg Ec
Eg|Qg1 where Q== -nu Te E
vi TTTH He.
First, obtain the square of the expected value of Qe follows:
Ec[Qal = -ne Ee +
n
Belsel
= Σ
n
1/2 Hig = tr(HE),
(39)
2
2
Ex 11ah tr(H))² 4
(40)
where the last equation holds since Σ is symmetric and Σi,/ Hi En, Hng [. = tr(H)). Second, we have the expected value of the square of Qę as follows:
T Σ H 4
Q = 17 Σ 1.
+ UOUEEEE - T" Σ &.H
(41)
and, since Ee Re " by Isserlis' theorem (Isserlis 1918) for zero-mean normal random vector
and ELGG
Bleye]]Eleyey
Rleyey [[exen]]
+ 1
Eg [0 = 17² Σ
U. Ex,
Σ
H3 (E + Egg Ex, + Exx
(42)
4
= Ee (tr(H2) (tH(H))))
(43)
Finally, we have the variance
Var. [L+1 [0,] = Eg [ Eg 22
(44)
= 172
r T Ev + tr(HH)))
(45)
2
722mb (uTCu++
tr (HHCHHC))).
(46)
26
17

---

Published as a conference paper at ICLR 2023
From Lemma we can easily derive the following theorem from Chebyshev's inequality:
Theorem 4. For SGD on a quadratic loss L, the expected loss at the next step (t + 1) is
" 2 Ee[Le+1]= = Le
(i(H.)) 2
tr(Sn)
(47)
ir(S.)
7n
Further, the batch gradient 95 normally distributed, then given 8 > 0 and 2 v3/8, we have, with probability of at least 1 - 6,
1l+11 - [be]l++]]] Sa,
(48)
where 3 = Lnnn (o Ciu
T' (((HCHHCi) and v = (1-NH)VL. 2b
However, we can replace this inequality with a better exponential inequality. To do so, we need generalized Lemma of the following Lemma [ (Laurent& & Massart, (000):
Lemma (Laurent & Massart (2000).. For i.i.d. Gaussian variables , with mean 0 and
variance 1, the Jollowing inequality holds for any positive x::
PiZ- - [E[2| 2 2||al12v x + 2] S exp(-xx)
(49)
P(Z Z - 22 S exp(-a))
(50)
where Z- = Y.Y?
Lemma 6. For i..dd Gaussian variables (Y,,.., Yn) with mean O and variance 1, the the following inequality holds for any positive x:
||c|2 V + 2/10//b00
P
7
-
B[
2
1a11 + 2/
S eex((-I),
(51)
2
P Z - 2-E s 2 |a1 + | V
I exp(-Ir)
(52)
where 7 = aY? Y.
Proof. For a standard normal random variable Y, ~ NO, 1), let vb be the Cramér transform of Y? c;,Y 1:
(A) log Elexp(A(X? +aY 1))1
(53)
= log
p(y:N(o. 1)) exp(A(2² + c.y 1))dy
(54)
y²
-log
exp
exp(\(x² + cay 1))dy
(55)
JR
2
2
y²
+ A(y² + +ay -1) dy
-log
exp
(56)
2m
2
log
exp
1
(1 2X)y² + cdyy -2) dy 2)
(57)
JR
V
2T
2
log
(op(-A)exp
(-1(1 22)9² exp (-c;Ay) dy
(58)
2
exp(-X)o 1
(-
1
-log
V2T0² exp
202 U²2
exp (-c;dy) dy (0 (1 - (1-2X) -1)
(59)
- + log(o) + -By~N((,00)) [exp (-c:Ay)
(60)
A 1
log(1 2\) + (cioo)²
(61)
X2 )
(62)
1-2M
18

---

Published as a conference paper at ICLR 2023
where the last inequality holds for 0 < X < \ since -1 - 2 log(1 - 21) = - 1 2 k21
(2A)*
2
k22 = 2X² k20 k-2
12A)
(2X)*
S 2X² Σ k20 (21)* = 12 Ex)0(22)6 = =
-T-
Therefore,
log [oxpp(AA2]] Σ log E
exp aid
y, + cs y
11) - Σ
(1+ {a
(63)
1 20-1
(Ilall + ||[| )a?
2
(64)
21|a||1a0
We can obtain (69 from the following (Birgé & Massart 1998): log E (exp(AZ)) - b, then,
1-bbX
for any positive 0, P( 2 vdjx + bya) S exp(-r). Also, for < X < 0, (A) - X² A 2
and thus log E E[(xp(AZ)) (Ilal12 + lell ) which leads to (70).
2
Theorem 1. For SGD on quadratic L, the expected loss increases, i... Eg|Le+1] - L, > 0, if and only if 0t satisfies the instability condition tr(HSs)
tr(S..) > 2. Furthermore, the batch gradient ge is normally distributed, then the following inequalities hold any positive 0:
7mb | b
P L1+1 Lei1 2 /28r +
01
I exp(-xz)
(8)
P (1+11-
Eglletal S 2B2x 8.) S exp(-22
(9)
where B = = (o Cl
e (((HHHHHi) and v= = (1 - M 2h
Proof. The first part of the statement is trivial from with e = 0 for a quadratic L. Now, we
focus the concentration of L1+1 from its expected value: L++1 E[L++1] Z E Z for
Z = Q: = -mu E + n T HE where v = (I nH)VL and & ~ N/ Cb) from 38 For
C. Inch BBT, we can represent E
V 2 Ju BE with ē ~ N'00 I) and
n
Z = mp ++ He
(65)
2
= =-
Innh b
(B v)
TE
1 &T B
HBE
(66)
2h
Y T diag(a)Y
(67)
Eay2+ +cY.
(68)
where A 77'm. B HB Q diag(a) c -nV / v and Y (Y,,... Yym = T g.
Q
2b
Therefore, from the rotation invariance, (Y,) are i.i.d. Gaussian variables with mean 0 and variance 1 like č. From Lemma we have the following inequalities:
|cl
PP L++1 E[Le+1] 2 2/
|| +
Vx 2||A||
I <xp(-xx).
(69)
2
|
IP L - s 22/
|A +
Vx
I exp(-xx).
(70)
2
Further, we have
"P tr(B HBB HB) = I' p (HBBBB HBB
|| = tr(A A) =
(71)
2h
2b
12 n'mm ir (HHHH
(72)
2h
19

---

Published as a conference paper at ICLR 2023
and
T |B u11² Im 0 Clu. |c|2 p Iin 110 B u12
(73)
b
These lead to ||-11|- + leel2 (rCHH + " 7m. u CTCCU= 9 and then 99
follows
from 70p. Moreover, we have ||
rm 26
|||[[ which leads (8) from 69
B.3 PROOF OF THEOREM
2
We consider the following linear homogeneous second-order equation for the sequence {ye}:
V/42 + P19/+1 + P2y/ = 0
(74)
and its the characteristic equation
(<x) x² + pr + P2 =0 0.
(75)
Lemma 7 (Theorem 2.35 in Elaydi 2005)). All solutions of (74) (i) oscillate (about zero) if and only
if (75) has no positive real roots, and (ii) converges to if and only if the solutions of (75) x and x2
satisfy | |22| < 1.
Lemma (Theorem 2.37 in Elaydi 2005
) The conditions
1 1 P1 P2
(CI)
1 +1+
(C1)
1 P >0
(CI)
are necessary and sufficient for the equilibrium point (solution) of (74) to be asymptotically stable (i.g., all solutions converges to 0).
Theorem 2. For generalized momentum GD with (2),
quadratic L, 0 < X; 1+.9 then qī 5e oscillates and diverges with the exponential growth of |o] del = O(e=t) for some c > 0. Moreover, cos((1,ōe) and || increase and \₁ as co, respectively, with 1- cos(g1, A1 [| =(e 2ct).
2 (31, < A1 for all i 1 where 7(1,, 82)
T+282'
We provide a proof of Theorem 2 which modified from the proofs in Appendix A in Cohen et al. (022T)
Proof. Put a quadratic training loss L(0) = 3 H0 + bTo + c. The update rules for generalized
momentum GD with (81, 32) this quadratic function L(0) are:
8, = (3,1 - ( - 17b nHO, 0+41 0 6c
(76)
(77)
and thus, for pair of eigenvalue/cigeenvccoor (q. A) the update rules for the quantities 0, q 0: + 1q Tb and 8, = q² 8, are:
5. = (3) 32a)ō. a0,
(78)
8.t1 0g + 5c
(79)
where a = 1/i > and this leads to:
de+1 0. + de
(80)
0. + (B1 - 32a)ō, - a0;
(81)
0, + (B) 32a)((: 0,-1) a0,
(82)
(1 + 31 (1 + B2)a) 0. (31 32a)
(83)
0 = 0.+1 (1+ + 31 - (1 (1+ 3))) 0, + (31 - :2a
(84)
0- 0.41 - 1,00011100 (1+ - (1 Br (0: 0.-1) + (B1 (2 - 0c22),
(85)
3, (1 + 31 (1 32)a) _ (Br 32a) 8,
(86)
P1
20

---

Published as a conference paper at ICLR 2023
Note that O, and 8, have the same characteristic equation. Thus, by Lemma stable when the following three conditions (C1-3) hold:
6 is asymptotically
-1 82)a) + B1 -32a > 0,
(CI)
(87)
+ -0 (12)a) + B1 32a 0,
(C2)
(88)
p- 1 320 >0. 0.
(C3)
(89)
C3 and CI always hold because
1 - B1 > 0 and
1 (1+ 31 (1 82)a) + B1 B2a a> 0.
(90)
Lastly, since
1+ (1 - -1 (2)) + 31 B2a = 2 231 - +1++ (1+223)a > 0,
(91)
The asymptotic convergence condition (C2) is equivalent to
2+ 231 2 1 31 (1+232)n
2 1 17
<
(31.88))
(92)
Therefore, along the direction of qı, the sequence [4]55.] diverges, and along the direction of q: (i > 1), the sequence {g. de) converges to O. The discriminant D(ao) for the characteristic equation 5(2;0 ag) positive as shown below where ao = 2y(31, 82), i.c., o(x; an) has two distinct real solutions.
D(a) pi 4p2
(93)
(1+ + B1 (1+ 32)a)² - 1(3 B2a),
(94)
1 + 31
$1 283
D(ao)=
(ao = 27(B1, (2))
(95)
1 2/2
1 1.22
(1 + 31)²
4(81 2112 + 2,2)) 21+232)²
(96)
(1 31)²
8((31 1)32 2,33) O.
(97)
(1+ 23
94 implies that (a) is convex quadratic with respect to a. Thus, to show that D(a) > 0 (6(2; a)
has two distinet real solutions) for all ag. it is sufficient show that (a(a) for ag- The following inequality holds
D'(a) = 22) 2(1 81 - B2 + +13)) 2 0
(98)
if a 2
1+8 - (1+/2)²
And, D'(a) 2 f a > since > ag >
1+.9 -4, from (1+82)²
1 31
B2 + B182
2(1 + B1) 1+ B1 - B2 + 31,32
ag
(99)
(1+ B2)
1+
(1+32)²
2(1+ B1) + 32)² - (1 + 212) + 31 - B2 + BB
(100)
(1 211 + 32)
1 + 31 + 38g + B,33 432
> 0.
(101)
(1+ 21 + 32)
Now, for a > ag we want to show that 21 for the dominant solution (x11] > [x2] ) is negative so that the general solution o, = qrI + c22g z
c1 + C2
22
oscillates where C; 24 E and
m1
i 2). The of the solutions of (x;;a) is + 22 Pt and this is negative for >> do
21

---

Published as a conference paper at ICLR 2023
from the following:
Pl - (1 (1+ +1+
(102)
> - [1+ 61) (1+
(103)
(1 + -1+++
(104)
2(1 + 31)
+ B1 - (1+ 32)
(105)
1+ 1+2
(1 + B1)(1 + 232) (1+ 82)2 + B1)
(106)
1 2/2
+ 31 1 + 2
> 0.
(107)
Thus, we have that v1 for the dominant solution is negative which leads to an oscillatory behavior of since
qi dr. And it has the exponential growth of |q 6,|
=0([=11)) = O(e"t) for e = In |x1
10 dcl = |-11"
C1 + C2
J2
(108)
c ||11|| S 1 5.1 S 2/ca
||11| for 2 to for some to
(109)
(A) violates 992 with 32) which implies |x]1| >> 1).
This makes be asymptotically aligned with the top eigenvector @i of H, i.e.
lom cos((1) q1, 8t) lim
(g. lim //m 8012 1 where g q &
6.
Moreover, we can obtain the exponential convergence (1 cos q1, 6.)l =
as follows:
|3:00
[cos(92, 5:)1
(s
5(6) + 0),
(110)
i>
V 86112 + +5 188
=|cos((1, 0,)l
(111)
151122 11)2+00 2s
((1)2
+ 83 - 18(2)
5(2)22 2
(112)
8(112 $7
8(12
2(1)
1+1
of
87
(113)
8
1()
+ 27
1
0 5 1 - coss(1,,0000000 6,)1 S
for 2 & for some t1.
(114)
3(112
The last inequality holds because az = /(1)/ss) Moreover, we have
/syl diverges to co and lim x² xV
1++ 2
VeL(O,) = HQ, +
(115)
a aD
Σ Xaaiq 0x + Σ q:q qgq.b= = Σ Eaai 0r
b)
0,
X
4:.
(116)
[V2T100))1 Σ X²0102
(117)
HVOL(0) Σ Pbq.rk Ex 9: = Σ 220 g::
(118)
(.T (() Σ 1 - Σ 200001
9b Σ 230
(119)
22

---

Published as a conference paper at ICLR 2023
where 0/50 = =q, 0/ + b.b. X
Therefore, we have d1 |A|ss = =O 2nt)
from the following:
E, E:1001122
01 +
Na
x 0
ss.. Vol. HVOL
(120)
E, 130
0
AB (
X
(x)
0
x (A1 )2)
A
A1 |
(121)
02
X2 0(2)2
0= +
x 0(212
0 0 A1 ss
for t 2 to for some t
(122)
since lim 81:0 = 0 for i > 1.
Remark. Due to the exponential convergence of ||s,, to X1, it only takes a fee steps (5-20) for ||,, to exceed i) (3,, (22) (see Appendix C.3
Remark. In practice, ||]|| A1 keeps increasing after it exceeds ,,, Therefore, may relax the assumption as the eigenvalues A (0:) > " (31,,, + + and 62 < (e())) < 2 (2)
for & 1 may change within the bounds over f.ror some e1, +3, 0 (g: 's are fixed). We can draw the same conclusion except that the limit s...s may not exist because of varying || according to t, but we can ensure that || [ss. eventually exceeds (,,,,, Bz).
23

---

Published as a conference paper at ICLR 2023
C EXPERIMENTAL SETTINGS AND ADDITIONAL FIGURES
We report the experimental results using vanilla SGD/GD without momentum and weight decay, constant learning rate, and no data augmentation. We use simple 6-laver CNN (6CNN, m 0.51M),
ResNet-9 (He et al. 2016 (m = 2.3.M), WRN-28-2 (Zagorukko & Komodakis 2016) (m 36M)
We use subsets of the datasets CIFAR- 1IFAR-10/100 (Krizhevsky Hinton 2011 and LinyImmaeNee 0
2009 STL-10 (Coates et al. (a subset of ImageNNet (Deng et al. 2009 with 3 x 64 x 64 images and 200 object classes) where DATASET-n denotes subset of DATASET with [D] = n and k=2 1024.
Moreover, we use PyHessian (Yao et al. 2020 to compute the Hessian-vector product (e.g., HV L.,, the top eigenvalue Å and its corresponding eigenvector 9i of the Hessian. For these computations, we use the power iterations with batch size of 2k, the tolerance of 0.001, and the maximum iteration of 100.
C.I FIGURE
In Figure and Figure we plot Belhe-1l-Le against tr(HSs) tr(S.)
tr(S.)) which is equivalent
L++1 Le tr(S.)
against
||.. for GD. Therefore, we expect the following linear relationship with the slope 2 and the x x-intercept 2 when the training loss L is locally quadratic, i.e., = 0:
Ee|Li+1] L1
172
tr(HS))
2
(123)
tr(Sn)
2
tr(Sp)
17
L1+1 - L1
7?2
[H|ss
(124)
tr(Sn)
2
F the behavior in the early phase, until the iterate enters the edge of stability. For GD, they plotted after ||]]| exceeds after which || [|s, starts increase from 0 to 2 in few steps. 7 For cross-entropy loss, mark the end point with 'x' when the iterate enters the unstable regime.
1e 1e-2CE, b 128, 0.01
1e-2CE, b 256, = 0.01
[E[L:+1] L.))rr(.)))
[EIL..1] L.Mtt((,))
175
200
slope=?
slope= o'
150
0
150
125
-1-
-00
-1
200
100
100
0.0
1.5
4.5
6.0
0.0
1.5
3.0
4.5
MIAS
1e2
1e2
TR30
2
1e
CE, GD, n= 0.01
0.01
300
1e-2 MSE, GD, = 0.01
(L--1 (.trr(s))
(b Liltr((..)
1
slope-
295
slope
980
step
290
-60
-1
200
285
200
0.0
1.5
3.0 ||M||s.
6.0
0.0
1.5
4.5
6.0
1e2
||H||s.
1e2
Figure 6: 6CNN with 7 = 0.01. See caption of Figure 1 for more details.
httrasss/ww.eescccoooooooooaaorr Apache-2.0 license
MIT license
24

---

Published as a conference paper at ICLR 2023
le 2
CE, GD, 0 = 0.04
le
MSE, GD, 1) 0.04
50
(L++1 - La)/tr(S.))
(b+11 Li/ttt
100
slope
48
slope= e²
step
-80
50
44
50
60
0.000
0.375
0,750 [|H]
1.125
1.500
0.000
0.375
0.750 ||M|ls.
1.125
1.500
1e2
1e2
1e-1
CE, GD, n= 0.08
1e-1
MSE, GD, 0.08
50
(1--1 Laltr(s.))
[ Lalttnss.)
slope
slope =
40
21
30
20
25
25
19
20
0.000
1.875
3.750 ||H||
5.625
7.500
0.000
1.875
3.750
5.625
7.500 lel
lel
||H||,
Figure 7: 6CNN with CE/MSE (lefVright) and r) = 0.04/0.08 (top/bottom). See caption of Figure for more details.
Res-9, CE, GD, 0.01
Res-9, CE, GD, n 0.02
186
(be -..trs..)))
420
L.tttt...)
slope
dape-
184
-882
410
-1
200
-2
100
178
400
0.75
[H|].
1e2
[|ib.
1e2
Res-9, CE, GD, 0.04
Res-9, CE, GD, n= 0.08
(4:11 (a)t)(5
86
(4e11 La)tr(5
slope-
slope
40
84
82
-38
500
0.000
0.375
0.750 [|M]s,
1.125
1.500
80
0.000
1.875
3.750 [|M]s
5.625
7.500
1e2
le1
Figure ResNet-9 with CE and 7 = 0.01/0.02/0..04/0.00 See caption of Figure for more details.
WRN-28-2. CE, GD, n 0.005
1e-2 WRN-28-2, CE, GD, 0.0.1
1.0
445
slope- ?
442
Sope-
440
437
-0.5
400
200
0.3
1.5
|M|
|W]]
1e2
Figure WRN-28-2 with n 0.005/0.011. See caption of Figure 1 for more details.
25

---

Published as a conference paper at ICLR 2023
1e-2
Res-9, MSE, GD, n= 0.02
1e-2 Res-9, MSE, GD, Q= 9= 0.04
160
slope
440
5
slope "
8
150
420
1.50 -2-
100
50
410
140
0.75
0.750 ||i]a
1.125
[|
Figure 10: ResNet-9 with MSE and 17 0.02/0.04. See caption of Figure for more details.
1e-2
STL-10-4k, CE, GD, n00.0
STL-10-4K, CE, GD, q 0.02
1.5
1,0050
1.0-
slope
Sope
17
5000
x
11
16
-0.5
-1.0
203
100
3.0
1.5
3.0
4.5
6.0
0.00
0.75
1.50
225
3.00
[|]]
162
[MK
102
CIFAR-100-44 CE, GD. n 0.02
CIFAR-100-4k, CE, GD, p- 0.04
105
223
slopo
104
222
-1-
221
-2
-000
220
102
0.20
0.75
1.30
5.00
0.000
0.750
1.125
1.500
[|M]s
142
[M]ss
192
MNIST-BK, CE, GO, 0.0.1
1..-2
MNIST-8K, CE. GD. / 0.02
1.5
3
176
362
LibrS)))
Mapan
162
0.0
173
172
200
154
-00
354
-1.
170
6.0
0.00
0.75
1.50
5,00
[M
1x2
162
Figure 11:
(DATASET, 7) (STL-10-4k, 0.01/0.02). (CIFAR-100-4K, 0.02/0.04), (MNIST-8K.
0.02/0.04). See caption of Figure D for more details.
1e-2 STL-10-4k,
20
400
(EL1-11) Latens.))
2-
15
1.
s
10
$ -100
=00
100
0.00
0.75
1.50
2.25
3.00
0.0
1.5
3.0
4.5
6.0
mitr
1e2
1e2
Figure 12: (DATASET, b. 7) (STL-10-44 128, 0.02), (CIFAR-100-4k, 512, 0.01), (MNIST-8K. 0.02/0.04). See caption of Figure for more details.
26

---

Published as a conference paper at ICLR 2023
C.2 FIGURE
Figure 2 shows the non-quadratictty (represented by deviation from the blue line in Figure 2 and the following figures) of the training loss function L after the iterate enters the edge of stability.
1c 1e-CCE 1
1e- 2 0.01
[E[L.+:] Laitrt..))
3000
[EIL...] Liitris,))
3000
1-
slope
2000
1-
slope
2000
step
1000
-1
200
-1
200
0.0
1.5
3.0
4.5
6.0
0.0
1.5
3.0
4.5
1e2
IIMS)) 07551
1e2
8954
1e-2
CE, GD, -= 0.01
1e-2 MSE, GD, 1= 0.01
2
(L.+1 - -LMttris.
(L. - Lilitris,
1
slope
1500
slope 4
1500
-1000
step
1000
step
500
500
1
200
200
0
0.0
3.0 ||H|s,
4.5
6.0
1.5
3.0
4.5
6,0
1e2
|Hi|ss
1e2
Figure 13: 6CNN with 0.01. See caption of Figure for more details.
1e
CE, GD, ==
800
1e-2 MSE, GD, 0 0.04
(L++1 Li
(L. Li/tr(S..)
slope
600
slope
1500
step
step
400
1000
200
500
50
0.000
0.375
0.750
1.125
1.500
0.000
0.375
0.750 ||||s.
1.125
1.500
||H||s.
1e2
1e2
1e-1
CE, GD, n= 0.08
1e-1 MSE, GD, n= 0.08
(Le+1 - L.)er((s))
(Le+1 Laltt((s.))
slope=0
400
slope
1500
gtep
1000
step
200
500
25
0
0.000
1.875 ||H|ss.
3.750
5.625
7.500
0.000
1.875
3.750
5.625
7.500
1e1
||H|s..
1e1
Figure 14: 6CNN with CE/MSE (left/right) and r) = 0.04/0.08 (top/bottom). See caption of Figure
2 for more details.
27

---

Published as a conference paper at ICLR 2023
1e-2
Res-9, CE, 7 0.01
1e-2
Res-9, CE, GD, 7 0.02
dle - -1ltrti
L.ttr(.))))
stope-
1500
slope
-1000
1000
-1000
Step
500
500
200
- 100
0.0
1.5
3.0
4.5
6.0
0,00
0.75
1.50
2.25
3,00
||i|
||]|
1e2
1e-2 2
n=
1500
1e
Res-9, CE, r== 0.08
1000
11 L.ttriS
Ree Laitris.)
slope
1000
slope
step
600
step
400
50
-5
-5
0.000
0.375
0.750
1.125
1.500
0,000
1.875
3.750 ||||s
5.625
7.500
|H||s
Figure 15: ResNet-9 with CE and 7 0.01/00.22/0.04/0.008 See caption of Figure 2 for more details.
WRN-28-2, CE, GD, 0.005
1e-2 WRN-28-2, CE, GD, 0.001
01-1-((105 (.1065
0.5
slope- df
1500
(oau - (,0055
slope
1500
1000
1000
-0.5
400
200
1's
[M]ss,
Figure 16: WRN-28-2 with = 0.005/0.01. See caption of Figure for more details.
1e
Res-9, MSE, GD, 0.01
1c
Res-9, MSE, GD, 0.02
-1-1- Littris)))))
1500
Ree L)ltriS
1500
slope=
slope=
5top
1000
1000
500
500
200
- 100
0.0
1.5
3.0
4.5
6.0
0.00
0.75
1.50
2.25
3.00
[|i],
1c2
1c2
1e-2 Res-9, MSE, GD, 0.04
(L.+) - L./R1(.)))
1500
slope=
1000
500
50
-5
0.000
0.375
0.750
1.125
1.500
|||]s
1e2
Figure 17: ResNet-9 with MSE and r = 0.01/0.02/0.04 See caption of Figure for more details.
28

---

Published as a conference paper at ICLR 2023
STL-10-4k, CE, n0
STL-10-4K, CE, GD, ( 0.02
2.0
700
400
230
1.0
500
LaMS5
300
230
0.5
400
200
-0.5-
200
-1.0
100
-2
3.0
3.0
4.5
6.0
0.00
0.75
3.50
2:25
3.00
[M]s.
[M]
CIFAR-100-4K, CE. GD. o= 0.02
CIFAR-100-4k CE GD, n= 0.04
800
730
400
2
Slope-
620
Hope- e
500
400
3.0
200
230
100
-2-
0.20
0.75
1.30
2.25
3.00
0.000
0:375
0.730
1.255
1,500
[HH|
[HI
MNIST-BK 8K. CE, GD, 0.00
MNIST-BK CE. GD, - 0.02
500
B... La005d
1.0
400
2-
Mopen
250
100
150
200
100
-05
-10
0.0
1.5
3.0
4.5
6.0
0.00
0.75
1.50
2.25
3.00
[Mi]
162
[M]]
162
Figure 18:
(DATASET, 7) (STL-10-4k, 0.01/0.02). (CIFAR-100-4K, 0.02/0.04), (MNIST-8K.
0.01/0.02). See caption of Figure for more details.
STL-10-4k, CE, d= 126, 0.02
-IFARR-10--K CE, ==
500
600
stope
slope
0
200
200
100
0.00
0.75
1.50
2.25
3.00
0.0
1.5
3.0
5.0
31O0
Figure 19: (DATASET, b, =STLL-1000 128, 0.02), (CIFAR-100-4K, 512, 0.01). See caption of Figure 2 for more details.
29

---

Published as a conference paper at ICLR 2023
C.3 FIGURE AND THEOREM 2
Figures 21 provide additional information of Figure 3 After these figures, we focus on providing
some empirical evidences of Theorem 2 that | s., cos(11, VL) and |q] V increase in a few steps after
H|| exceeds 2.
||ss ;
||M|s,, >
|| =
||H||s. >
10*
-min
10"
7L
AL
10-1
0.25
0.50
0.75
1.00
1.25
8:
zero
Figure 20:
6CNN with 1) = 0.02. See caption of Figure 3 for more details.
0.10
TPil 2.00
0.05
[Maila, THIL
0.00
quadratic
-0.55
0.0
0.2
0.4
0.6
0.8
1.0
1.2
Figure 21: To gain some intuitions of Figure and 20 (right), we plot three graphs - quadratic function fx) f/(x 1) (black dashed line). and cubic functions fi (x) a1 1) fz)
in red and (2()) 02 (i: 1)(1 2z)) in blue. They chosen satisfy 8 "tx" / 8
ou v=-0
az Iz=0 0.5. We also choose 0r and a2 for fi and f2 to have the same minimum with f in x [0, 1]. i.e., min.e(o.11 f(x) min [011] fi(x) for 2.
and 8 02 1 18
ccosta. Rt)
lR VL
=(0)= + 16a03
(n= 2 a- 1.1)
1.0
150
30
0.8
125
100
20
-15
0-
0.4
50
10
25
0.2
-3
-2 -1
2
10
15
20
25
step
Figure 22: Left: A toy two-dimensional loss function L(0) = 0 + 16a02 where a = 1.1. We 32. We also plot the GD trajectory in yellow starting from 02) = 0.1). Right: We show the exponential increase in lq] V (purple) and the S-shape increase || (red) and cos(y), VL) (green) to ||/||| and respectively. which empirically demonstrates Theorem 2 We also note that they start to increase in the order of cos(g1, VL), || lls. and lq VL)
optimize the loss by GD with " 3/2 so that ||]]| = 32a
30

---

Published as a conference paper at ICLR 2023
ResNet- 0.04
ResNet CE, n= 0.02
||M|
[[|]]]
Kosiaz A.
140
||]|]
(costay 901
80
0.8
120
0.8
80
20
0.2
0.2
70
110
970
1.0
Step
Step
ResNet-9, CE. GD. n 0.01
le/v1]
ResNet-9, MSE, GD, o 0.01
14/91)
300
[[H]]
Kosta
ccosta;
250
0.8
250
0.8
200
200
0.6
0.6
150
100
100
$30
400
0.0
1220
1260
0.0
10
Step
16/91]
CE. GD, / 0.08
6CNN. CE. GD. n 0.04
100
1.0
2.00
1.0
|ML
kosia, VII
[|I|]
[costar OL]
175
150
125
100
40
0.4
0.4
2
50
20
40
60
50 Step
70
Step
[0/
10]
6CNN, n
6CNN, n
200
1.0
1.0
[|H|e
|M]
[coslas VLJI
[|i].
[[]]]
[coo|a [1
175
350
0.8
44
0.8
125
250
100
200
75
0.4
2
0.4
22
50
100
0.2
1
0.2
-1
25
50
900
120
0
250
260
2.70
200
290
300
310
Step
Step
Figure 23: After H|| (bluc) exceeds (red dashed line) at 2 80/176/ 40// 12250/19/43/111
in few steps (< 5/6/1//5/566//////////////////// 5/5). || (red) starts increase. As expected in Theorem 2
||s, increases together with cos((e VL) (green) and i VLI (purple). They are observed to
start to increase in the order of cos q VL). H|s.. and qi VL shown in Figure 22
31

---

Published as a conference paper at ICLR 2023
|||lll
[cos(qı, VL)
cos(VL(0)), VL(0
train loss
(H++ 1), 1H)
||Hlls,
1.0
250
2.2
200
0.5
2.0
150
0.0
100
1.8
0.5
50
0
1.6
1.0
50
100 Step
150
||H:lls//|HHll
Icos(q1. VL)I
cos( vL(@r+1))
train loss
Icos(q) (H) (2, (H+))
||Hells,
1.0
250
2.2
200
0.5
2.0
150
0.0
100
1.8
0.5
50
0
1.6
1..
0
200
400
600
800
Step
23 the cosine cos(gn (H,,, VL(),)) (green) between the sharpest direction and the gradient gets large, where H, (0.). Simultaneously, || ||Hel,,, increases and exceeds 2 (red solid > red dashed). the iterate entering the unstable regime and oscillating with cos L(,,). LL(0.+1)) R -1 (orange). However, due to the non-quadraticity, the sharpest direction changes with cos(g) (H++1) 1), q1 (H+)) (cyan) close to O. We train ResNet-9 CIFAR-10-1 with 1) 0.04 (top: steps 50-150. 150. bottom: steps 0-800).
Figure 24: After ||||| exceeds 2 (not shown in this Figure, see Figure
32

---

Published as a conference paper at ICLR 2023
/|1Hll
[cos(qı, VL)I
cos(VL(0)), VL(0
train loss
Icos(q) (H++ 1), 91(H)))
||Hills,
1.0
150
2.2
0.5
100
2.0
0.0
50
1.8
0.5
1
1.6
1.0
25
50
75
100
Step
||H:|ls//||Hlll
Icos(q1, VL)I cos(VL(0)), vL(()+1))
train loss
Icos(qı (H 1), (1()))
||Hrls
1.0
150
2.2
0.5
100
2.0
0.0
50
1.8
0.5
0
1.6
1.0
0
125
250
375
500
Step
We train ResNet-9 on SIFAR-10-88 with r7 = 0.08 (top: steps 0-100, bottom: steps 0-500). See caption of Figure? 24 for more details.
Figure 25:
33

---

Published as a conference paper at ICLR 2023
C.4
FIGURE 4
C.4.1
GD
6CNN, MSE, /1 0.04
BCNN, MSE, 7 0.08
0.5
200
train loss
[IMI
0.4
0.4
40
20
0.2
0.1
10
0.0
500
1000
1500
2000
0.0
500
1000
1500
2000
Step
Stap
Figure 26: 6CNN, MSE, and 1) = 0.04/0.08. See caption of Figure 4 for more details.
n
800
CE. n
400
train loss
[|Н]]
train loss
IlM
[H]]
600
300
1.5
1.5
200
1.0
1.0
300
150
100
100
50
500
1000
1500
2008
0.0
500
1000 Step
1500
2008
Stap
ResNet CE, 0.04
ResNet n=
200
train loss
|M]ls
MMI
1at lcss
|Mile
[MI
375
80
250
60
100
1.0-
1.0
40
0.5
50
0.3
0.0
0.0
Figure 27: ResNet-9, CE, and 0.01/0..2/0..4//0.08 See caption of Figure for more details.
34

---

Published as a conference paper at ICLR 2023
WRN-28-2, CE, n= 0.005
WRN-28-2 CE, / 0.01
1600
800
train loss
[|M
||||
train loss
[MI
2.0
2.0
700
1200
600
1.5
1000
1.5
500
400
1.0
600
300
200
100
500
1000
1500
2000
560
1000
1500
2000
Step
Step
WRN-28-22 A
400
Diin toos
MMb.
[M]
300
1.5
200
1.0
0.5
50
500
2000
Step
Figure 28: WRN-28-2, CE, and rj 0.005/0.01/0.02 See caption of Figure 4 for more details.
WRN-28-2, MSE, n= 0.01
WRN-28-2, MSE, 0.0.0
400
200
[14
[[]]
1|Mb
[|]]
175
0.50
150
250
125
0.45
0.45
100
0.40
150
0.35
50
0.35
25
1000
2000
3000
4000
500
1000
1500
2000
Step
Step
WRN-28-2, n
100
0.55
[|M|s
[|]]
0.50
0.45
40
0.30
2008
Step
Figure 29: WRN-28-2, MSE, and 0.01/0.02/0.04. See caption of Figure for more details.
35

---

Published as a conference paper at ICLR 2023
C.4.2 SGD
We demonstrate types of IIR, (i) || 2 30h and (ii)
tr(
in Figure 30 and 31
7
U.(...) "
respectively. They are equivalent, but the latter shows the effects of IR more clearly as has fixed threshold regardless of b.
2.5
2.5
2.0
10
1.-
Shep
Shep
ResNet-9. b 1024. n- 0.08
ResNett b 512. 0.08
2.5
130
2.5
130
IM
111
2.0
2.0
0.5
0.5-
400
600
600
ResNet-9, 256 0.08
ResNet-9 128, 00088
0.08
tran 15ss
[M
[M]
tmm 1ocs
[Mi
[4]]
-80
40
15
1.5
400
6000
803
220
460
6:00
803
32.0
2.5
2.5
1.0
1.0
Stcp
Step
Res 16, 0.08
Restit-9 D n 0.08
2.5
100
2.5
100
2.0
2.5
1.0
5,5
3.5
3.5
5.5
200
450
800
0.0
200
430
800
Figure 30: IIR of [H]s, S for SGD. We plot 2 (red dashed line) which is not the threshold n
for ||sss. ResNet-9, CE, 0.08, and E {2 211 23.. See caption of Figure for more details.
36

---

Published as a conference paper at ICLR 2023
ResNet-9. 4096. 0.08
ResNet-9. 2048.
2.5
Ians
1.5
1.5
1.0
1.0
Step
Step
RecNet-9, n
ResNert-O, $12. -
2.5
130
2.5
100
[H]
2.0
2.0-
1.0
1.0
0.5
0.5
Rocli at-O, b 256, n- 0.08
at-9, D 128, n 0.08
102
100
[P]
Tram
[M4
2.04
30
2.5-
-30
3.5
3.5
0.5
3.0
203
450 Stap
400
9.0
203
400
cin
Stop
Figure 31: IIR of (T()) 7
< 2 for SGD. With the upper bound 2 (red dashed line), this shows the effects of IIR more clearly than Figure 30 ResNet-9, CE, n7 0.08, and b {2 2 ,.... ,2?)
(r(S.)
WRN: 28 2, 00001
WRN WRN-28-2. 2, 128, n= 0.01
800
800
train loss
IiM
3.0
700
3.0
700
2.5
2.5
600
500
500
2.0
2.0
300
300
100
100
0.0
1200
2400
3600
4600
0.0
1200
2400
3600
4808
Step
Step
WRN-28 2, n=
WAN-28-22,66 8.
800
800
train loss
[M|]s
[M
train loss
||||
MM
2.5
600
2.5
600
2.0
2.0
400
400
1.5
1.5
1.0
1.0
200
100
100
1200
2400
3600
4500
1200
2000
3600
4800
Step
Step
Figure 32: IRR of | < 3ph caption of Figure 4 for more details.
for SGD. WRN-28-2, CE, 0.01, and b (29, 2,,27,25 25.33. See
"
37

---

Published as a conference paper at ICLR 2023
2000
350.
303.0
159.0
259.0
200.
0.0
1500
Step
Step
Figure 33: Left/Right: See caption of Figure 4(c) 4(d) for more details. WRN-28-2, CIFAR-10-8K ?) 0.01.
38

---

Published as a conference paper at ICLR 2023
C.5
FIGURE
Figure 34 35
provide some additional information of Figure 5 Figure 36 shows that 1/p 2 100
much larger than and that std[llge [] is 2-3x smaller than Eg [llsel]] even in the case of b 1. II]. and thus the square of the mean resultant length is similar the concentration measure pb as shown in the following approximation:
Therefore, we use the approximation l ~ l8g selllgelll
Eg [Tc] E. 9
R
E
||9c
Eg
Eg gg
Pb-
(125)
ResNet-9
Tiny imagenet-32k,
innylmagelNet-222 WRN-28-2
17.5
15.0
150
(0g2lm)
40
12.5
125
10.0
10.0
75
50
19 -10
40
101112131415
logz(b)
logz(b)
log:(b)
Figure 34: (DATASET, MODEL, EPOCHS) CCIFAR-100322K, ResNet-9, 800 epochs), (Tiny- ImageNet-32k, ResNet-9, 400 epochs), model is trained for 400 epochs, which is short compared to Figure ( (bottom right heatmap). See
WRN-28-2, 800 epochs). Middle: The
caption of Figure 5 (right) for more details.
24
24
log:o-
Ps
2
bpp
2
logoo=
SRSR: (tb) pv/D
log.o
29
29
Rb)
Rb)
2-2
2-2
2-4
2-4
2-6
2-6
2-8
2-8
20
24
26
2
210
2
2
2
24
2
2
2.0
212
batch size (b)
batch size
Figure 35: Left: We Yn,s for sampling with replacement ('wr', dotted purple curve), which is almost equivalent to the "without replacement" counterpart (red). Right: LSSR for different p values with the corresponding LSR (dotted lines). See caption of Figure (left) for more details.
39

---

Published as a conference paper at ICLR 2023
0=00.2
0.04
van
92in loss
E(jon]
100
(
2.0
2.0
sestilve)
19
1.9
1.8
100
10-1
50
1.00
150
50
100
150
200
Siop
2.3
n=
2.3
BIJ
36*
22
[[X]]]
2.2
1.0
2.1
302
2.1
2.0
104
101
109
10>
1.2
1.0
1.6
Med
10-4
1.6
150
200
1.5
10-0
Step
Shop
Figure 36: To further understand the concentration measure p of the per-example gradient, plot [[][ / |[[[V]]]|, vtr(S,) [[[ and std [[]]] We use 100 samples compute the expectation values and the standard deviation. Here, V is about 100. IE[] [e]] is 2-3x
larger than We train 6CNN with 7 0.02/0.04/0.08// /0.2 (GD).
40

---

Published as a conference paper at ICLR 2023
C.6 GENERALIZED MOMENTUM VARIANTS
For generalized momentum variants of GD with (B1, 32), we have
6: = - - 1 VoL( + 82)
(126)
When ||]] very small the beginning. we may approximate that L is linear where VgL(0) = g is constant with respect 0 in region. Then & converges to å satisfying the following equations:
ô = 31 19.
(127)
8
31 9.
(128)
and thus
L41 L4 2 -VoI(0.)6
(129)
1)
(130)
1-
Therefore, together with Theorem
we generalize
12
with the following approximation:
Li+1 - Le
tr(S..)
(131)
23,,,, 32) B1)
where (31, 82)
1+8, 142,92
Figures C.6 and
C.6
we empirically validate 131
similar Figures
and respectively.
16-2
0.01 (0.1.0)
384
10-2
1
0.01,
(B., (b) (0.1,
387
1
380
Literis..
slopes
367
365
378
362 0
376
24Q 183
-1
357
372
355
M|]s
||
1o-2
0.01 18, B2) (0.2. 0)
1e-2
0.01, (8,, 8,) (0.2, 0.2)
348
325
(le.1 L.br5)))
slope 399 R-T
346
1 L.triS.)
slope Alll
322
344
320
342
317
2:00,40 240
2: 171
340
315
338
312
1e2
MMME
1e2
0 = 0.01. (6, B2) (0.4,0)
1e-2
n= 0.01, (8, 02) (0.4. 0..4.
242
280
1-
slope slopemyag.at goT
278
(ree Latrii5l
slope 770.
240
238
274
236
2nbu/t 200
272
232
IM
1e2
INNI
1c2
Figure 37: We use 31 = 00.100.2/0.4. 0). and 3g = (Right, Polyak momentum Polyak 1963)) or
32 = Bi (Left, Nesterov momentum (Nesterov) 1983)) with DATASET CIFAR- 10-8k, = 0.01 on ResNet-9 without BN. They follows [31] before entering the edge stability. See caption of Figure for more details.
41

---

Published as a conference paper at ICLR 2023
0.01, (8,, (0.1,,,
(0.1,
0.01, 1.. 8.) (0.1.
600
||H|s
||
2.0
500
1.5
1.5
0.5
0.5
0.0
250
500
750
1000
250
500
750
1000
Step
Step
n
(8,, 02)
600
n
(PI. P2) (0.2. 0.2)
600
|Hi
||
||A|.
|||
2.0
500
2.0
500
400
400
300
300
1.0
1.0
200
200
0.5
100
100
0.0
0.0
500
1000
250
1000
= 0.01, (B., Pz = 10.4, 0)
0.01, (8.. Px) (0.4, 0.4)
[H]
[M]
Ma
500
1.5
1.5
0.5
0.5
250
500
750
1000
250
500
750 1000
Step
Step
Figure 38: We use 31 = 0.1/0.2/0... and Bg = 0 (Right, Polyak momentum (Polyak) 1963) ) or
32 = 31 (Left, Nesterov momentum (Nesterov) 1983)) with DATASET CIFAR-IO-8k. 17 0.01 on
ResNet-9 without BN. The red horizontal lines indicate 77 (,1,,, where y(1,, 32) T+2222 caption of Figure for more details.
1+&
See
D DISCUSSION
We provide insight the link between the batch gradient distribution and the sharpness of the loss landscape. In this section, we reconcile our arguments with some previous studies.
Jastrzebski et al. (017) explain the optimization behavior of SGD with the SDE approximation /² of SGD where W an m-dimensiooall Brownian motion.
d0, = L(O,)
Therefore, the same ratio z y leads the same SDE, which implies LSR. Moreover, large y implies large diffusion in SDE, which has been linked with the escaping efficiency from sharp local minimum in Zhu al. (2019). Our arguments are free from any other problems raised for the SDE-based analyses of SGD which assume vanishing learning rates (Mandt et al., 2016 2017 Hu al. 2019 Li al. 2017 7 2019a| Jastrzębski et al.) 2017, [mmthh & [e| 2018: Chaudhari Soatto 2018 e.g., the mismatch to practical finite learning rate regime the inherent theoretical issues in the SDE approximations (Yaida) 2019 Li et al. 2021). We instead argue that a large second moment tr(S,) (compared to tr(S..) and large lead to low constraint 2pt/n n on the interaction-awarr sharpness.
Wu et al. (2020) empirically show that what is important for the generalization performance of neural network is not the class to which the gradient distribution belongs, but the second moment of the distribution. This is consistent with arguments with the interaction ir (r(H..)) and the
42

---

Published as a conference paper at ICLR 2023
concentration measure ps tr/( tr(S)), because they depend on the second moment So, not the class of the gradient distribution.
Recently, Li et al. 2021 suggest necessary condition that the "noise-to-signa ratio" needs to be large for LSR (and the SDE assumption) to hold. This is consistent with our result on the linear regime (where b and ps small) because the noise-to-signal ratio is approximately the inverse of the "signal-to-noise ratio Ps = tr(S,)/ (S,). but defined for an equilibrium distribution. We provide not only the necessary condition but also the sufficient condition for LSR with novel scaling rule LSSR applicable to every batch size including where LSR fails (the saturation regime).
43

---

Published as a conference paper at ICLR 2023
E
OTHERS
= 2 12
b 2 10
b = 2
b 26 b = 25
b 2
b 2
b = 2
0.03
0.02
P1
0.01
0.00
300
600
900
1200
Step
Figure 39: P1 is not constant during training, but it does not change much at the EoS and shows similar evolution for different batch sizes, especially when they are large.
12
b 10
b 2
b = 26
= 2
b 29
b = 2
b 25
1.0
a 0.5
0.0
300
600
900
1200
Step
Figure 40: Pb is not constant during training, but it does not change much at the EoS. pb is small for a small batch size.
44


Reviews:
[{"review_id": "qLTfzpNZ33", "reviewer": "Reviewer_9x87", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper provides a fine-grained analysis of the edge of stability (EoS) phenomenon and in particular, their characterization also applies to SGD. In my understanding, the new characterization not only uses the largest eigenvalue of Hessian, which corresponds to the sharpest direction but also considers the curvature along other directions. To this end, an interaction-aware sharpness is defined. The numerical experiments show that the characterization is quite \"accurate\" and the authors provide an interesting application in explaining the linear and saturating scaling rule in large-batch training. ", "summary_of_the_review": "The more precise characterization of EoS and the empirical findings are interesting, but this work is still somewhat preliminary and more work is needed. \n\n\n=====\nThe author's response partially addressed my concern. I generally think that the analysis and numerical experiments are useful for future study on relevant topics. Hence, I increased my score to 6. ", "strength_and_weaknesses": "### Strong point\n\nThe derivation of new characterization is simple but powerful. The new sharpness also seems to capture the EoS phenomenon more precisely. The application of explaining the learning rate rescaling in large-batch training is consistent with empirical observations, which might be practically relevant. In general, I think this paper makes some interesting observations and I believe they are useful for further study on relevant topics. \n\n\n----\n### Weak point\n\nHowever, I feel that this paper is written in a hurry and the work is still somewhat preliminary at this stage. Listed below are a few concerns in my mind. \n\n1. **The discussions of prior works are unfair**\n\n   The phenomenon of  EoS was first observed by (Wu et al., 2018) at the end of training; (Jastrze ̨bski et al., 2020) further observed the progressive sharpening phenomenon in training neural nets. (Cohen et al., 2021) followed these works by providing a systematical investigation of these phenomena and also additionally found the non-monotonic convergence at EoS.  All three works are crucially important in this line of research. Only attributing the discovery of EoS to (Cohen et al., 2021) is unfair (at least in my opinion). \t\t\n\n2.  **The explanation of implicit interaction regularization is quite hand-waving and not convincing**\n   - First, the definition of $\\rho_b$ (a concentration measure of batch gradient)  is the same as the gradient diversity defined in (Yin et al., 2018)  and is never mentioned. \n   - I find equation  (16) misleading since both sides depend on where the iterate locates. In iterations, both sides keep changing. The statement that *$\\|H\\|_{S_b}$ is implicitly regularized and bounded to be less than $2\\rho_b/\\eta$* is technically correct. But it implies nothing since both sides keep changing. For example, we cannot claim that SGD tends to find solutions with smaller $\\|H\\|_{S_b}$ since $\\rho_b$ is not a constant. \n\n \t3. In Section 5, I do not understand why keeping $2\\rho_b/\\eta$ to be constant is a good choice for generalization. This is never explained. \n\n\n\n### Other comments\n\n1.  In the second paragraph, the authors ask three questions. But it seems that this paper only answers the third one. Therefore, claiming \"we provide a new characterization of the EoS for SGD, which can serve as an answer to the above questions\" is not appropriate. \n\n2.  In Section 2, I suggest focusing on the case of sampling with replacement for simplicity. Otherwise, the statement is much more complicated and no new insights are brought. \n\n3.  In Section 3, the authors claim that existing works using tr(HS_b) make some additional assumptions, but they do not. In my opinion, this claim is quite misleading, since this paper also needs to make some assumptions, e.g., the validity of local quadratic approximation. In my opinion, the real difference is that this paper uses this quantity to analyze EoS, which is not carried out previously.  \n\n4.  In Figure 1, I do not understand why the authors only mark the endpoint with 'x' for cross-entropy loss. Also, it should be clarified in the caption the meaning of MSE and CE. \n\n5. In Theorem 1, the batch gradient never follows a normal distribution, thus the Gaussian assumption is too strong. Also, I do not understand why  (8) and (9) are interesting.\n\n6. In the definition of the new edge (11), I think the author should explain the intuitive difference between $||H||_{S_n}$ and $||H||$. In my understanding, the former considers the curvatures in typical directions but the latter only considers the curvature in the sharpness direction:\n\n$$\n\\qquad ||H||_{S_n} = \\frac{\\theta^TH^2\\theta}{\\theta^T H \\theta}, \\qquad ||H|| = \\sup \\frac{\\theta^TH^2\\theta}{\\theta^T H \\theta},\n$$\nwhere the supremum is taken wrt the direction theta.\n\n8. In Figure 2, what do you mean by **overestimation**? In addition, are the \"iteratre enter the EoS\" and \"in the unstable region\" the same thing?\n\n9. In Equations (13) and (14) and Theorem 2, suddenly momentum GD is studied. I do not understand why? Do the authors want to say the EoS is the same as vanilla (S)GD? Theorem 2 is also repeatedly referred to argue the convergence of directions, but it is stated for momentum GD.  In addition, in Theorem 2, what is $\\lambda_i$? What do you mean by quadratic $L$? Is each sample loss also quadratic?\n\n\n\n\n---\n### Reference\n\n[1] Yin et al., Gradient diversity: a key ingredient for scalable distributed learning\n\n[2] Wu et al., How SGD selects the global minima in over-parameterized learning: A dynamical stability perspective\n\n[3] Jastrzebski et al.,The break-even point on optimization trajectories of deep neural networks", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Xj6ch-1frZp", "reviewer": "Reviewer_mqQV", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper aims at gaining more understanding of the Edge of Stability phenomenon. Some theoretical understanding was obtained via Theorem 1, which assumes a quadratic loss function. Although EoS does not manifest for quadratic loss, this theorem was then used to motivate empirical claims. Such claims are backed by experiments.", "summary_of_the_review": "Overall, this is an interesting paper and a helpful contribution toward an understanding of EoS. I have some concerns at this moment, but my score can be changed.", "strength_and_weaknesses": "Strengths:\n\n(1) Problems studied in this paper are very interesting.\n\n(2) Empirical investigations are detailed.\n\nWeaknesses:\n\nThe demonstration is still largely empirical, and I wish the theoretical component could be stronger. Some reasons are: \n\n(1) some of the claims appear to be rather general, but it is actually known that EoS does not occur for all loss functions. I’d appreciate it very much if the authors could be more specific about their conjectures (which I nevertheless found interesting), such as the hypothesized conditions under which they could be true.\n\n(2) The abstract suggests the question of “why the sharpness is somewhat larger than 2/(learning rate)” will be answered. Could the answer be more explicitly and concisely provided? Is there any quantification of how much larger it is than 2/(learning rate)? Is it loss function dependent? In fact, I wonder how the authors relate this claim (and other empirical observations) to the recent line of analytical work based on matrix factorization, e.g., [Wang et al. Large Learning Rate Tames Homogeneity. ICLR 22], [Zhu et al. Understanding Edge-of-Stability Training Dynamics with a Minimalist Example. arXiv]. The latter is concurrent so it is not a missing reference, but I’m curious if there is any contradiction: for example, these results seem to give sharpness not larger than 2/(learning rate). Although the authors seem unaware, theories in these papers do seem to be aligned with many of the empirical observations made here. However, do they align with all observations? It could be that neither is wrong, and instead it is just a matter of different loss functions. I wish these aspects could be explored more.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Af7o107mjld", "reviewer": "Reviewer_bH22", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "In this paper, the authors proposed a novel characterization of the \"edge of stability (EoS)\" phenomenon, referred to as the \"interaction-aware-sharpness\",  for both gradient descent (GD) and its stochastic variant SGD.\nThe proposed theory is of particular interest in its ability to provide novel and more precise insight into the EoS of SGD. Based on the theory, the authors further proposed a novel scaling rule (Linear and Saturation Scaling rule, LSSR) on the learning rate and SGD batch size, which extends the popular linear scaling rule (LSR).\nNumerical experiments on a few (not so deep) neural networks and across a few widely used datasets were provided to empirically support the theoretical results.", "summary_of_the_review": "The paper is in general well written and proposes an interesting and novel viewpoint of the EoS for SGD.\nI do not see any major flaws in the paper. \n\nDetailed comments:\n* Definition 2: the term \"concentration measure\" is used without providing more context. This can be misleading to some readers with probability or statistics background who may be more familiar with the term \"concentration of measure\".\n* It would be helpful to discuss the limitation of the proposed analysis (including the LSSR), from both theoretical and empirical perspectives possibly.\n* can one go beyond SGD to more involved settings? It would be of interest to have a short discussion on this point.", "strength_and_weaknesses": "**Strength**: the paper is in general well written and the main ideas are easy to follow. The problem under study is important and the proposed theory provides novel insights and is of interest. The numerical experiments look compelling.\n\n**Weaknesses**: the theoretical/technical depth of the paper is somewhat limited, the empirical contribution is rather significant, though. The limitations of the proposed theory/analysis are not sufficiently discussed.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster