Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability
Authors: Zixuan Wang
Zhouzi Li
Jian Li
Venue: NeurIPS 2022 Accept (2022)
Keywords: gradient descent trajectory
progressive sharpening
edge of stability
sharpness
deep learning theory

Abstract: Recent findings demonstrate that modern neural networks trained by full-batch gradient descent typically enter a regime called Edge of Stability (EOS). In this regime, the sharpness, i.e., the maximum Hessian eigenvalue, first increases to the value 2/(step size) (the progressive sharpening phase) and then oscillates around this value (the EOS phase). 
This paper aims to analyze the GD dynamics and the sharpness along the optimization trajectory.
Our analysis naturally divides the GD trajectory into four phases depending on the change in the sharpness value. We empirically identify the norm of output layer weight as an interesting indicator of the sharpness dynamics. Based on this empirical observation, we attempt to theoretically and empirically explain the dynamics of various key quantities that lead to the change of the sharpness in each phase of EOS. Moreover, based on certain assumptions, we provide a theoretical proof of the sharpness behavior in the EOS regime in two-layer fully-connected linear neural networks. We also discuss some other empirical findings and the limitation of our theoretical results.

Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability
Zhouzi Li III,, Tsinghua University 2houz211886630gaa11 com
Zixuan Wang" IIIS, Tsinguua University
vanggxx2010012322gmmm.....m
Jian Li
III, Tsinghua University
1apordge@gmail.com
Abstract
Recent findings demonstrate that modern neural networks trained by full-batch gradient descent typically enter regime called Edge of Stability (EOS). In this regime, the sharpness, i.c., the maximum Hessian eigenvalue, first increases to the value 2/(step size) (the progressive sharpening phase) and then oscillates around this value (the EOS phase). This paper aims to analyze the GD dynamics and the sharpness along the optimization trajectory. Our analysis naturally divides the GD trajectory into four phases depending on the change in the sharpness value. We empirically identify the norm of output layer weight as an interesting indicator of the sharpness dynamics. Based on this empirical observation, we attempt to theoretically and empirically explain the dynamics of various key quantities that lead to the change of the sharpness in each phase of EOS. Moreover, based on certain assumptions, we provide a theoretical proof of the sharpness behavior in the EOS regime in two-layer cullecconnected linear neural networks. We also discuss some other empirical findings and the limitation of our theoretical results.
Introduction
Deep learning has achieved great success in a variety of machine learning applications, and gradient- based algorithms are the prevailing optimization methods for training deep neural networks. However, mathematically understanding the behavior of the optimization methods for deep learning is highly challenging, due to convexix,, over parameterrization and complicated architectures. In particular, some recent empirical findings in deep networks contradict the traditional understandings of gradient methods. For example, Wu al. (30 observed that the solution found by gradient descent has sharpness approximately equal to 2/1 instead of just being smaller than 2/1j. Also, Jastrzebski et al. 14 observed that there a break-even point in the SGD trajectory, and after this point, there is regularization effect on the loss curvature.
One recent well-known example the phenomenon called "Edge of Stability" (EOS) (Cohen et al. 6)). Based on the classical optimization theory, the learning rate 1 of gradient- based method should be smaller than 2/A so that the loss can decrease, where À the largest eigenvalue of the Hessian of the objective, also called "sharpness" in the literature. Otherwise, the loss diverges (even for simple quadratic functions). However, the empirical findings in Cohen et al. (6] show that under various
"Contributed equally, listed in alphabetical order.
'The authors are supported in part by the National Natural Science Foundation of China Grant 62161 1166144004, Turing Al Institute of Nanjing and Xi'an an Institute for Interdisciplinary Information Core Technology.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

network settings, the EOS phenomena typically occurs along the gradient descent trajectory: (1) the sharpness first increases until reaches 2/m (called "progressive sharpening") (2) the sharpness starts hovering around 2/n (the EOS regime) and (3) the loss non-monotonically decreases without diverging.
Although (1) seems to be consistent with the traditional belief's about optimization, rigorous mathematical explanation for it still open. Moreover, phenomena (2) and (3) are more mysterious because they violate the v) < 2/X "rule" in traditional optimization theory, yet the training loss does not completely diverge. Instead, the loss may oscillate but still decrease in the long run, while the sharpness seems to be restrained from further increasing.
In this paper, we aim to provide a theoretical and empirical explanation for the mystery of EOS. Towards the goal, we focus on the dynamics of these key quantities when EOS happens and attempt to find out the main driving force explain these phenomena along the gradient descent trajectory from both theoretical and empirical perspectives.
1.1 Our Contributions
Our contributions can be summarized as follows.
We analyze the typical sharpness behavior along the gradient descent trajectory when EOS happens, and propose a four-phass division of GD trajectory. based on the dynamics of some key quantities such as the loss and the sharpness, for further understanding this phenomenon.
(Section
3.1)
(Section 3.2 We empirically identify the weight norm of the output layer of the sharpness dynamics. We show that analyzing the dynamics of this surrogate can qualitatively explain the dynamics of sharpness. By assuming this relation, together with some additional simplify- ing assumptions and approximations, we can explain the dynamics of the sharpness, the loss, and the output layer norm in each phase of EOS (Section 3.3 In this context, we also offers an interesting explanation for the non-monotonic loss decrement (also observed in Cohen et al. (6), Xing et al. (Section 3.4.
effective indicator
(Section Following similar ideas, we provide a more rigorous proof for the progressive sharpening and EOS phenomena in a two-layer fully-connected linear neural network setting based on certain assumptions. The assumptions made here are either weaker or arguably less restrictive.
1.2 Related work
The structure of Hessian The Hessian matrix carries the second order information of the loss landscape. Several prior works have empirically found that the spectrum of Hessian has several "outliers" and a continuous "bulk" (Sagun al. 28 29). Papyan 25 26 Typically, each outlier corresponds to one class in multi-class classification. As consider the binary classification setting, there is typically one outlier (i.c., the largest eigenvalue) that is much larger than other eigenvalues. is consistent with our Assumption 4.1 The Gauss Newton decomposition of the Hessian was used in several prior works (Martens [23,, Bottou et al. [4], Papyan (25) 261). Papyan [25] empirically showed that the outliers of Hessian can be attributed to a "G component", which is also known as
Fisher Information Matrix (FIM) in Karakida et al. 16) Also, Wu et al. [31] analyzed the leading Hessian eigenspace by approximating the Hessian with Kronecker factorization and theoretically proved the outliers structure under some random setting assumption.
Neural Tangent Kernel A recent line of work studied the learning of over-parameterized neural networks in the so-called. "neural tangent kernel (NTK) regime or the lazy training regime (Jacot et al. [13, Lee et al. [18). Du et al. (8, 77). Arora et al. (2). Chizat et al. (5)). A main result in this regime is that if the neural network wide enough, gradient flow can find the global optimal empirical minimizer very close to the initialization. Moreover, the Hessian does not change much in the NTK regime. Our findings go beyond NTK setting to analyze the change of sharpness.
Edge of Stability regime The Edge of Stability phenomena was first formalized by Cohen et al. [6] Similar phenomena were also identified in Jastrzebski et [14 as the existence of the "break-even" point on SGD trajectory after which loss curvature gets regularized. Xing et al. (2) observed that gradient descent eventually enters regime where the iterates oscillate on the leading curvature direction and the loss drops non mon-ooootoni,llyy Recently Ahn et al. [ studied the non-monotonic

---

decreasing behavior of GD which they called unstable convergence, and discussed the possible causes of this phenomenon. Ma al. [22] proposed a special subquadratic landscape property and proved that EOS occurs based on this assumption. Arora et al. [3] studied the implicit bias on the sharpness of deterministic gradient descent in the EOS regime. They proved in some specific settings with a varying learning rate (called normalized GD) or with a modified loss vI, gradient descent enters EOS and further reduces sharpness. They mainly focus on the analysis near the manifold of minimum loss, but our analysis also applies to the early stage of the training when the loss is not close to the minimum. In particular, our analysis provides an explanation of non-monotonic loss decrease that cannot be explained by their theory. Another difference is that they consider VL (for constant learning rate) where L is a fairly general MSE loss independent of any neural network structure, while our analysis is strongly tied with the MSE loss of neural network. Very recently, Lyu et al. [21] explained how GD enters EOS for normalized loss (e.g.. neural networks with normalization layers), and analyzed the sharpness reduction effect along the training trajectory. The notion of sharpness in their work is somewhat different due to normalization. In particular, they consider the so- so-called spherical sharpness, that is the sharpness of the normalized weight vector. They also mainly studied the regime where the parameter is close to the manifold of minimum loss as in 3 and proved that GD approximately tracks continuous tharpness-reducttee a similar regime called "catapult phase" where loss does not diverge even if the largest Hessian eigenvalue larger than 2/17. Our work mainly considers training this regime and assumes that the training is not in the "divergent phase" in Lewkowycz et al. [19]. Compared with Lewkowycz et al.
flow. Lewkowyez et al. [19] proposed
19 we provide more detailed analysis in more general settings along gradient descent trajectory.
2
Preliminaries
Notations: We denote the training dataset as (x, 3i C R x {1, and the neural network as f I x R" + R. The network f(0, x) maps the input x E Rd and parameter E IR" to an output in R. In this paper, we mainly consider the case of binary classification with mean square error (MSE) loss f(2, y) = (= (5=-y) y)²
o.. Denote the input matrix as X = X2., X.) E Rddx and the label vector as Y (11, 1/2 ..... Mn) E R". We let F(t) = (f(o(t))x)) x1), (1), f(0(t),
x,)) E IR" and ((t) Ft)) Y be the (output) prediction vector, and the residual vector respectively at time t. The training objective is: L(f(O)) = Σm= x)) = (f x.). v.)
Hessian, Fisher information matrix and NTK: In this part, we apply previous works to show that the largest eigenvalue of Hessian is almost the same as the largest eigenvalue of NTK. We use the latter as the definition of the sharpness in this paper. Further details can be found in Appendix [F]
As shown in Papyan 266, Martens 223,, Bottou et al. 4 the Hessian can be decomposed into two components, where the term known as Ctauss-Newtor matrix", G-term Fisher information matrix (FIM), dominates the second term in terms of the largest eigenvalue. Meanwhile, Karakida et al. [16]
pointed out the duality between the FIM and a Gram matrix M, defined as M = (0) aF(e) is also known the neural tangent kernel NTK (Karakida et al. [16 15),, which has been studied extensively in recent years (see e.g., 133 8 2. Note that in this paper, we do not assume the training is in NTK regime, in which the Hessian does not change much during training. It is not hard to see that M and FIM share the same non-zero eigenvalues: if Gu = Au for some eigenvector E IR, M ar(0) u aF(0) Gu = X 0 u, i.c., X is also an eigenvalue of M.
80 0e
In this paper, we use O(t) to denote the parameter iteration (or time t) and the sharpness at time t as A(t) = A(O(t)). We similarly define (t), (t), ((t,, Ltt).
Here show the gradient flow dynamics of the residual vector D(t):
(t) at( dd(() dt
(t aC(t)
2 (t) OF T
D(t) (M()))(() (1)
00 dt
00
00
n
00 00
3 A Four-phase Analysis of GD Dynamics
In this section, we study the dynamics of gradient descent and the change of sharpness along the optimization trajectory. We divide the whole training process into four phases, occurring repeatedly
3

---

in the EOS regime. In Section 3.1, we introduce the four phases. In Section 3.2 that the change of the norm of the oulput layer weight vector almost coincides with the change of the sharpness. In Section 3.3 using this observation, we attempt to explain the dynamics of each phase and provide a mathematical explanation for the changes in the sharpness. Section [3.4] we explain why the loss decreases but non-monotonically We admit that a completely rigorous theoretical explanation is still beyond our reach and much our argument is based on various simplifying assumptions and is somewhat heuristic at some points. Due to space limits, we defer all the proofs in this section to Appendix E.T
we show empirically
3.1 A Four-phase Division
To further understand the properties along the trajectory when EOS happens, we study the behaviors of the loss and the sharpness during the training process. As illustrated in Figure 1, we train a shallow neural network by gradient descent on subset of 1,000 samples from CIFAR-10 (Krizhevsky et al. [17), using the MSE loss as the objective. Notice that the sharpness keeps increasing while the loss decreases until the sharpness reaches 2/1.. Then the sharpness begins to oscillate around 2/n while the loss decreases non monotonically.. This is a typical sharpness behavior in the EOS regime, and consistent with the experiments in
We divide the training process into four phases according to the evolution of the loss, the sharpness, and their correlation, shown in Figure 1. The four phases happen cyclically along the training trajectory. We first briefly describe the properties of each phase and explain the dynamics in more detail in Section 3.3
Phase I: Sharpness A < 2/4. In this stage, all the eigenvalues of Gram matrix M are below the threshold 2/n. In particular, using standard initialization, the training typically starts from this phase, and during this phase the loss keeps decreasing and the sharpness keeps growing along the trajectory. This initial phase is called progressive sharpening (PS) in prior work Cohen et al. [6]. Empirically, the behavior of GD trajectory (as well as the loss and the sharpness) is very similar to that of gradient
flow, until the sharpness reaches 2/n (this phenomena is also observed in Cohen al. (6) See Figure 5 or Appendix J.1 in their paper). We note that GD may come back to this phase from Phase IV later.
Phase II: Sharpness A > 2/1z. In this phase, the sharpness exceeds 2/n and may keep increasing. We will show shortly that the fact that A > 2/11 causes |D '1 (where U1 the first eigenvector of M) to increase exponentially (Lemna 3.2 This would quickly lead D to exceed || |Y|| in a few iterations, which leads the sharpness to start decreasing by Proposition 3.1 hence the training process enters Phase III.
Phase II: Sharpness A > 2/1 yet begins to gradually drop. Before A drops below 2/17. Lemma 3.2 still holds, so D v1 keeps increasing. Proposition 3.1 still holds and thus the sharpness keeps decreasing until it is below 2/1, at which point we enter Phase IV. A distinctive feature of this phase is that the loss may increase due to the exponential increase of D Vvl.
Phase IV: Sharpness A < 2/7. When the sharpness is below 2/1,. Tu begins to decrease quickly, leading the loss to decrease quickly. At the same time, the sharpness keeps oscillating and gradually decreasing for some iterations. This lasts until the loss decrease to a level that is around its value right before Phase III. The sharpness is still below 2/n and our training process gets back to Phase I.
3.2
The Norm of the Output Layer Weight
It is difficult to rigorously analyze the dynamics of the sharpness A(t). In this subsection, we make an interesting observation, that the change of the norm of the output layer of the network (usually a fully-connected linear layer) is consistent with the change of the sharpness most of the time.
In particular, for a general neural network f(x) A h(W, x), where A E Rmm is the output layer weight and the feature extractor h Rº x _d + TM outputs mn-dimensional feature vector (h corresponds to all but the last layers). W E IR is the parameter vector of the extractor h.
Note that M-
BF 00
T(
OF 06
be decomposed follows:
0F
(FF))
T
OF
OF
T
OF
OF T OW
M =
00
00
=
0A
0A
OW
= MA MW

---

400
sharpness
Phase
Phase
Phono
Phase III
200
300
Phase IV
200
100
100
sharpness
100
200
100
400
500
00
50
60
70
60
5-
1055
5-
loss
3-
2-
(a) Sharpness and Loss
(b) Phase Division Illustration
Figure 1: In Figure (a), we show the sharpness and the loss when training on a llye linear activated network by gradient descent. Shortly after the sharpness crosses 2/n, it drops quickly back to some value below 2/ n/. In Figure (b) we illustrate how the four phases are divided according to the evolution of the sharpness.
where the (i.j)-entry of Mw is (Mw) ky
(f(x) oW
8f(x,) SW
= A
Oh X.) B., OW
A.
OW
In this expression, intuitively |A| should be positively related M| ||- We empirically observe that the part MA (ax)
0A has much smaller spectral norm compared to the whole Gram
är
matrix M (see Figure ((a) and Appendix D). which means |Mw/l dominates |M Therefore,
||A|| should be positively correlated with [M]].
The benefit of analyzing || is that the gradient flow of || ||A] enjoys the following clean formula:
/|A 2
aL 0A
(FF
dt = 2
A DT
A D.F.
(2)
n
0A
n2
In this work, we do experiments on two-layer linear networks, fully connected deep neural networks, and Resner18, and all of them have such output layer structures. From Figure 3(a) we can observe that the output layer norm || and the sharpness A change of the time along the gradient descent trajectory, i.e., they both increase or decrease at the same time. We note that they may change in different directions very occasionally around the time when ||(t 1) ||A(t)] changes its sign (see the experiments in Figure
same direction most
3.3 Detailed Analysis of Each Phase
In this section, explain the dynamics of each phase in more detail. For clarity, we first list the assumptions we need in this section. For different phases, we may need some different assumptions to simplify the arguments. Most of the assumptions are consistent with the experiments or the findings in the literature. Some of them are somewhat stronger, and we also discuss how to relax them.
3.3.1 Assumptions Used in Section
Assumption 3.1. (A-norm and sharpness) Along the gradient descent training trajectory, for all time i, the norm A()]] of the output layer and the sharpness M1) moves the same directions, i.e., (Att(A(t.on) + 1) (t)) sign(])
1)ll [tt)))
is the key observation that we have discussed in Section 3.2 The following are two assumptions about the gradient descent trajectory. The first one assumes that D(t) and ||||| are updated according their first order approximations. Empirical justification of this approximation can be found in Appendix|D D.1.3]
Assumption 3.2. (First Order Approximation of GD) Along the gradient descent trajectory. the update rule is assumed as the first order approximation
41 D(t)
D(t + 1) D(t -n || 1)1 - =1
F(t))
(3)
n
Assumption 3.3. (Gradient flow for the PS phase) When A(t) < 2/1, follows the gradient D(t)
flow trajectory: dD(t)
5

---

121
230
Suteow
8300 loartam)
630
642
180
200
25
125
150
115
24
15
100
15
(a) Two-layer Linear Network
(b) Deep 5-layer tanh network
Resnet18
Figure 2: Sharpness and A norm are highly correlated; Anomaly points (orange points) in the figure correspond to the iterations when sign(A(t +) A()) * sign(]]A(t + 1)ll [t)ll) More
discussions can be found in Section and Appendix
Assumption 3.3 holds empirically, especially in the progressive sharpening phase (see Figure 5 or Appendix J. in Cohen et (6]) when the networks are continuously differentiable. We include these experimental details in Appendix D See also (Theorems 4.3 and 4.5) in Arora et al. [ for further theoretical justification. We need this assumption for the proof in the progressive sharpening phase.
Then we state an assumption on the upper bound of the sharpness to restrict the regime we discuss: Assumption 3.4. (Sharpness upper bound) If the training does not diverge, there exists some constant Bx, such that 0 < A(t) < Ba. for all t.
This assumption states that there is an upper bound of the sharpness throughout the optimization process. Actually, in Lewkowycz et al. [19]. they proved that 4/n) is an upper bound of the sharpness two-layer linear network with one datapoint, otherwise the training process (loss) would diverge. They empirically found that similar upper bounds exist also for nonlinear activations, albeit with somewhat larger constant BĄ, In the work, We focus on the case when the loss does not diverge and hence we make Assumption 3.4
The main set of assumptions we need is about the change of M's eigendirections. Assumption 3.5. Denote {v:} to be the set of eigenvectors of Mt). We have three levels of assumptions M's 's eigenspace.
(i) (fixed eigendirections) the set {u:) is fixed throughout the phase under consideration;
(ii (eigendirections move slowly) at all time and for any i, Ftt) dvi(t) < (() (t);;;
(iii) (principal directions moves slowly) at all time I, there is a small constant e2 2 0 such that (1,( 01 1)) 2 2-
Clearly, these three assumptions are increasingly weaker from (i) (iii). Assumption 3.5 on the eigenvectors is somewhat strong, and the eigenvectors corresponding to small eigenvalues may change notably in our experiments. We use it illustrate a basic proof idea of the progressive sharpening phase, but later we relax this assumption to Assumption 3.5 (ii).
Moreover, for the proof in Phase II and III, Assumption 3.5 (iii), which only assume that the main direction changes slightly, is sufficient for our proof. Actually, we note that 01 (t) (the eigenvector corresponding to the largest eigenvalue) changes slowly and the inner product of its initial direction and its direction at the end of the phase is also large (see Appendix D for the empirical verification).
For the proof in Phase II, need another small technical assumption:
Assumption 3.6. Assume D(I) v2 (1) 2 062 D for some e for some = to at the beginning of this phase. Here e2 is defined in Assumption 3.5 (tii).
Assumption 3.6 says that D(t) has non-negligible component in the direction of i₁ Since 62 >> 0 is a small constant, this is not a strong assumption as some small perturbation (due to discrete updates) would make the assumption hold for some 1.
3.3.2 Detailed Analysis
In each phase, we attempt to explain the main driving force of the change the sharpness and the loss.

---

Phase I: In this phase, we show that D(t)
F(t) < 0 under certain assumptions (detailed shortly)
on the spectral properties of (t) (see Lemma 3.1 below). By Assumption 3.2 we have A
1)1l3 - A
O. implying that the sharpness A() also increases based on Assumption 3.1 This
phase stops if Att) grows larger than 2/1/-
We assume the output vector F(t) is initialized to be small (this is true if we use very small initial weights). For simplicity, we assume F(0) in the following argument.
Lemma 3.1. For all in Phase I, under Assumption 3.5 (i) and 3.3 it holds that D(t) F(t) < 0.
From this lemma, || keeps increasing by Assumption 3.2, hence the sharpness keeps increasing by Assumption 3.1 until it reaches 2/1 or the loss converges to 0. In the former case, the training process enters Phase II, while the latter case is also possible when 17 is very small (e.g.. even the largest possible sharpness value is less than 2/n). We admit that Assumption 3.5 (i) is somewhat strong. In fact, the assumption can be relaxed significantly to Assumption [35 (.5 We show in Appendix (.2 that under Assumption 3.5 (ii) and Assumption 3.3 we can still guarantee D( F(!) < 0. Moreover,
we provide a dynamical system view of the dynamics of D(t) Ft)) in that Appendix.
When the training process just enters Phase II, the sharpness keeps increasing. We show shortly that ((t) v1 (t) starts to increase geometrically, and this causes the sharpness to stop increasing at some point, thus entering Phase III.
Phase II:
In this phase, we adopt a weaker assumption on the sharpness direction in Assumption 3.5 (iii).
This assumption holds in our experiments (See Figure 17]. Also, Assumption 3.6 is necessary.
Lemma 3.2. Suppose Assumption 3.5 (iii) and 3.6 hold during this phase (with constants c2 > 0 and
<> 1). A(t = (2 + 1) and T > 1-c2-1/c 1/c 1, then ((t) V1(() increases geometrically with
factor (1 + ++ E2 1/c) > f t 2 to this phase.
will exceed ||| eventually, Next, the following proposition states that when this happens, ((t) T F(t) > 0. Consequently, || decreases by Assumption 3.2 leading to the decrement of the sharpness based on our Assumption 3.1
Since D(t) T 21 (t) increases geometrically, [|] 2 D(t)
v1(t) D(t)Tv1(t)
Proposition 3.1. If D (t) ||||] then F(t) O.
The sharpness is still larger than 2/n, but it starts decreasing. Meanwhile, the loss continues to increase rapidly due to Lemma 3.2 Eventually, the sharpness will fall below 2/n and then the training process enters phase IV.
Phase III:
By Lemma 3.2 if the sharpness stays above 2/1,, then we can have an arbitrarily large loss. According to Proposition 3.1 if the loss is large enough, the sharpness keeps decreasing.
Now we show that if the sharpness stays above 2/n, [t)(1) will decrease by a significant amount. This partially explains that the sharpness should also decrease significantly until it drops below 2/n (instead of decreasingly converging to a value above 2/n without ever entering the next phase).
Proposition 3.2. Under Assumption 3.2 if [() >
||||, then || + 1)112 (1))1) <
An ((t)l - 2
From the above argument, we can see that if D(t) vn (t) is larger than ||Y|],
I]. then D(t) vn (1) does
not decrease in Phase III, and according Proposition 3.2 At)] decreases significantly, implying the sharpness drops below 2/1 eventually.
Remark:
The fact that the sharpness can provably drop below 2/m in this phase can be proved more rigorously in Section 3 for the two-layer linear setting. See Theorem 2
Phase IV: First, since the training process has just left phase III, F(())) is still positive and large, hence || [t)]] keeps decreasing and the sharpness decreases as well. Since the sharpness stays
below 2/1,, the loss decreases due to the following descent lemma (with u replaced by D(t)..
Lemma 3.3. If A(t) < 2/T,, then for any vector u E R", [ (t) < (1 na) n
where a (2/1 A(t), Xmin (M))))))))) In particular, replacing u with D(t), we can see [ + 1)Il < (1 - =1())1(((())
Next we argue that D( F()) will become negative eventually, which indicates that [|A() and hence the sharpness will grow again. Since the sharpness is below 2/n, D(t) vz (t) decreases geometrically due to Lemma 3.3 (replacing with (())s)())
7

---

390
1.8
argest digenvoucc MLA
narmi"2h
300
the second torgest eigenvalue 2ieta
230
100
0.8
50-
3.6
0.4
100
125
1.50
200
400
600
800
1000
(a) Sharpness, Amax MM)), and (M).
(h) The loss (+)) /n) and [R(t)]] /n
Figure 3: In Figure (a), we verify the assumptions in Section 3.2 and 3.3 empirically. The steel blue line represents the largest eigenvalue of M, i.e. the sharpness Amx M The orange line: the second largest eigenvalue is far below 1/11, illustrating Assumption 3.7 The blue line: the largest eigenvalue of MA is negligible compared to the sharpness. In Figure (b), we show that even though the total loss [(t) = [t) II? /n oscillates considerably, || [|((t)] /n decreases more steadily.
In fact, D can be decomposed into the et componntttt vn v. and the remaining part defined as R(t) (I 1n (t)u, (t).. Then have
D(1) F()) (v) (1) (t) D(1)) (1) (1) D(1) + Y) + R(I) (R(()) + +Y)
As shown in the next subsection, R(t) almost follows similar gradient descent trajectory (t)
(Lemma 3.5 More precisely, R' (t) is defined as (t +1++1 -I (t) 2y (t)xr
R) While D's dynamics is D 1) (I () R' follows similar dynamics R(t+1) (I n (t)) R'(),, where (t)
(Lemma
(t) - (t) Note that M' (t) has eigenvalues smaller than 1/1 for any time t (by Assumption 3.7 hence with an assumption similar to Assumption 3.5 (i) (or similar version of our relaxed assumption in Appendix E.2 for M'), can prove that F' (R) Y) < 0 for any time (See Appendix E.1 for the rigorous proof).
Since R(t) R (t),, the second term in the decomposition 44) is always negative and the first term (#) direction term) is decreasing geometrically. Therefore, there are only two possible cases. The first possibilityy is that the first term decreases to a small value near 0 and the second term remains largely negative. Then their sum will be negative, which is (t) F()) < O, thus implying the training enters Phase I. The second possibility is that when the first term decreases to small value near O, the second term also small negative value. In this case both R(t) and D(t) T "1(() are small, implying the loss is almost 0, which is indeed the end of the training.
3.4 Explaining Non-monot Loss Decrement
In this subsection, attempt explain the non-monotonic decrement of the loss during the entire GD trajectory. See Figure (bb As defined in the last section, we decompose D into the i cy-momnnn vivd D and the remaining part R. Below, prove that R(() is not affected much by the exponential growth of the loss (Proposition 3.5) in Phase II and III, and almost follows a converging trajectory (which is defined as R'(t) later in this section).
The arguments in this subsection need Assumption 3.4 and Assumption 3.3 (iii), both very consistent with the experiments. We need additional assumption the spectrum of M.
Assumption 3.7. All M (1) 's eigenvalues except A(I) = Amax M (1)) are smaller than f/rrfor t.
Recall that the largest eigenvalue is at most BA by Assumption 3.4 Empirically, the largest eigenvalue an outlier in the spectrum, i.c., it is much larger than the other cigenvalues. Hence, we make Assumption 3.7 which states that all other eigenvalues experiments. See Figure Bbb). Similar fact is also mentioned in 28 29].
most 1/1,. which is consistent with
First, we let Bp be an upper bound of Dtt), i.e., for all t, | < Bp. In the two-layer linear network case, we can have an explicit form BD. (see Lemma C.9 Assumption 3.4. BA is the upper bound of 1/A
Appendix C) Recall that in
Lemma 3.4. Suppose Assumption 3.5 (iii) holds. R(t) satisfies the following:
R(t + 1) (I 10---1a0 (t) R(t) ee(()) where [[er(t)] < 6:eel D(t))( -
8

---

Lemma 3.5. Define an auxiliary sequence R (t) by R' R(O), and R'{t+1)= - (I-
vi (t))x(()) R) (t). If Assumption 3.4 Assumption 3.5 (iii, Assumption 3.7 hold, and for any time t there exists a quantity X, > such that the smallest eigenvalue of M (t) i.c. Amin(M Xp,
then there exists a constant C+ > O such that || [[(() R'II S Cr BD(BB-1)yee
1/A
Now, in light of Lemma 3.2 and Lemma 3.5 we arrive at an interesting explanation of the phenomena
of non-monotonic decrease of the loss. Basically, D can be decomposed into the 2i -1-ommonnent
V D and the remaining part R = (I - V D.D. The V1 -componnttt may increase geometrically during the EOS (Lemma 3.2 but the behavior of the remaining part R(t) is close to Rt),), which follows the simple updating rule (t = (I () (t),,, so Lemma 3.3 implies that the R part almost keeps decreasing during the entire trajectory (here Lemma 3.3 applies with u replaced by R' (t), noticing that the eigenvalues except the first are well below 2/n). Hence, the non-monotonicity of the loss mainly due to the ei -component of D, and the rest part R is optimized in the classical regime (step size well below 2/(the operator norm)) and hence steadily decreases. See Figure ((b)
A Theoretical Analysis for 2-Layer Linear NN
In this section, we aim to provide more rigorous explanation of the EOS phenomenon in two-layer linear networks. The proof ideas follow similar high-level intuition as the proofs in Section [.3] In particular, we remove replace the assumptions in Section 3.3 with arguably weaker assumptions. Due to space limit, we state our main theoretical results and elaborate their relation with the proofs in The detailed settings and proof are more tedious and can be found in Appendix C
Section 3.3
4.1 Setting and basic notations
Model: In this section, we study a two-layer neural network with linear activation, i.e. f(x) Vm AT Wa where W [wi, Wm. E R A [al, (01,-...mm Rm
D - Vm V. W,
Dataset: For simplicity, we assume y; = +1 for all i E [n], and X X||2 = O(n). We assume X has rank r, and we decompose X X and Y according to the orthonormal basis (v:). the eigenvectors of X X: x Ex Ae,;e," Y Y= = (Y v.)u, = E 2;0; where vs is the eigenvector corresponding to the i-th largest eigenvalue A; of X F. = Y v; is the projection of Y onto the direction v, Here we suppose >> and the global minimum (A*, W*) exists.
Update rule: We write explicitly the GD dynamics of D(t): D(t+1) (I M ((),, where (t) my (t]) 2 X STWT(C (t))) mPm F(t)) T is the Gram matrix combined with second order terms.
47
4.2 Main Theorem and The Proof Sketch
Phase I and Progressive Sharpening:
Assumption 4.1. There exists some constant X > 1, s.1. for all E Ir - 11, x(XTx)= S
x (X X). Moreover, A (X X) Z 2X X XX.
Assumption 4.2. There exists K = Q(r ') such that min / /
2 k.
The first assumption is about the eigenvalue spectrum of X xTx x. 3 The second assumes that all component 2; = Y "V; are not too small.
Theorem (Informal). Suppose Assumption 4.1 Assumption 4.2 hold, the smallest nonzero eigen-
value X, = =x x)= > 0 and A1 >maxx X x) = cin. Then for any 0, if m = 22 s ), and
n === Q
x2 x+2z), ), we have the progressive sharpening property: A(t+1) > 0 fort t = 1, 2, to
where to is the time when
|| D(t))l? < O(2)) or Amax (M)(t)) > 1/1 the first time.
In the proof of this theorem, show that the Gram matrix M(t) = which serves as a justification of Assumption 3.3 we made in Section 3.3
2
()))(()) + T X That shows all (t)
'It guarantees the gap between two adjacent eigenvalues is not very large, and there is a gap between the largest and the second largest eigenvalue. Note the second part of the assumption is a relaxed version of Assumption 3.7 In our CIFAR- 10 Ik-subset with samples' mean subtracted, /1/De x 3 (See Figure 19.

---

approximately share the same set of eigenvectors X. In our proof, we also prove more rigorously that [t)( is an indicator of the sharpness in this simpler setting.
Edge of Stability (Phase II IV):
Assumption 4.3. There exists some constant 62 > such that [[ S m
This assumption is based on Theorem [ In Theorem we state that in the progressive sharpening phase, [[]()]] has upper bound of O(l/m). Now in the EOS phase, we assume that [[(()] grows
larger by at most a constant factor. Further discussions refer to Appendix[D22.
Assumption 4.4. There exists some constant / > 0, such that A S n - B).
This assumption is consistent with Assumption 3.4 which assumes an upper bound of the sharpness.
Assumption 4.5. There exist some constant C3 such that ((t) vi > csyn//m m for some to at the beginning of phase II.
This assumption is in the same spirit of Assumption 3.6 with the only change of the bound in terms of m and n. Now, we are ready to state our theorem in this stage.
Theorem 2. Denote the smallest nonzero eigenvalue as X, eigenvalue as X1 - XXT X). Under Assumption 4.3 4.4 4.5 and X1 (X X) > 222 X) in Assumption 4.1 there exist constants ca, Co2 CG such that in c6 m max
A ( X) 0 and the largest
Cid'n"
A can then
There exists p O()) which depends on C3 such that if A(to) > (1 + p) for some to, there must exist some 1 > to such that A(t,) < (1+
A(), 1) 2 (1 + p). then there is a constant e9 > 0 (depending on ca) such that +(++ T) 01 [t) 01 11111++)) (7).
Define R(t) = (I fitia D(t). and == -I (I- (M
1).
holds that [[([)) (1)l 0
Vnnd
We can conclude the following from Theorem 2 (1) The first statement of the theorem states that the progressive sharpening phase causes the sharpness to grow over 2/1,. then the sharpness eventually goes below 2/n). This illustrates the regularization effect of gradient descent on the sharpness (this consistent with the analysis of Phase III in Section 3.3). (2) The second states that [t) (ttTT, vil
geometrically increases in Phase II and III. Note that we proved a similar Lemma 3.2 for Phase II the more general setting in Section 3.3 (3) The third conclusion gives an upper bound for the
distance between R(t)'s trajectory and R (t)s.. This bound helps illustrate why R(t)'s trajectory is
similar with R' (t) in Phase IV of Section 3.3
Discussions and Open Problems
In this section, we discuss the limitation of our theory and some related findings. First, argument crucially relies on the assumption that ||||| changes in the same direction as A does most of the time. Here, elaborate more this point. Seeing from longer time scale, || and the sharpness may have very different overall trends (See Figure (c) in 2). i.e., the sharpness oscillates around 2/ /n but || ||2 increases. Moreover, the sharpness may oscillate more frequently than |l while the low-frequency trends seem to match well (See the late training phases in Figure (b) in 2 Currently, our theory cannot explain the high-frequency oscillation of the sharpness in Figure (b).
While we still believe the change of || is a major driving force of the change of the sharpness, other factors (such as other layers) must be taken into consideration for complete understanding and explanation of the sharpness dynamics. We also carry out some experiments that reveal some interesting relation between the inner layers and the sharpness, which is not yet reflected in our theory. Due to space limit, we defer to Appendix D.3
We conclude with some open problems. It would be very interesting to remove some of our assump- tions or replace them (especially those related the spectrum of M) by weaker or more natural assumptions the data or architectures, or make some of the heuristic argument more rigorous (e.g.. first order approximation of the dynamics (3). Extending our results in Section D to deeper neural networks with nonlinear activation function an intriguing and challenging open problem.
10

---

References
[1] Kwangjun Ahn, Jingzhao Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. arXiv preprint D.Xiv:2204.010520 2022.
[2] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332 332. PMLR, 2019.
[3] Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on edge of stability in deep learning. arXiv preprint arXiv:2205.09745. 2022.
[4] Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-313 2018.
[5] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances in Neural Information Processing Systems, 32, 2019.
[6] Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gra- dient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2/03.00065 2021.
[7] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pages 1675-1685, PMLR, 2019.
[8] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes wwer-parameterized neural networks. arXiv preprint arXiv:1810.02054 2018.
[9] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware mini- mization for efficiently improving generalization. arXiv preprint arXiv:2010.01412, 2020.
[10] Noah Golmant, Zhewei Yao, Amir Gholami, Michael Mahoney, and Joseph Gonzalez. sytorchhessian--iiennngg efficient pytorch hessian eigendecomposition, october 2018. URL https://github. com/hoodggolman//yytrrhh----o seathh-hhsii-eeeeeee -iggeinnnignnnninggs...
[11] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[12] Wei Hu, Lechao Xiao, Ben Adlam, and Jeffrey Pennington. The surprising simplicity of the early-time learning dynamics of neural networks. Advances in Neural Information Processing Systems, 33:17116-17123 2020.
[13] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[14] Stanislaw Jastrzebski, Maciej Szymczak, Stanislav Fort, Devansh Arpit, Jacek Tabor, Kyunghyun Cho, and Krzysztof Geras. The break-even point on optimization trajectories of deep neural networks. arXiv preprint arXiv:2002.09572 2020.
[15] Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. The normalization method for alleviating pathological sharpness in wide neural networks. Advances in neural information processing systems, 32, 2019.
[16] Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Pathological spectra of the fisher in- formation metric and its variants in deep neural networks. arXiv preprint 2019.
[17] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[18] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
[19] 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 a0X3v:2003.02218, 2020.
11

---

[20] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. arXiv preprint arXiv:2110.06914, 2021.
[21] Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. arXiv preprint arXiv:2206.07085, 2022.
[22] Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions:
The effect optimization and origin. arXiv preprint arXfv:220..11326 2022.
[23] James Martens. Second-order optimization for neural networks. University of Toronto (Canada), 2016.
[24] Rotem Mulayoff, Tomer Michaeli, and Daniel Soudry. The implicit bias of minima stability: A view from function space. Advances in Neural Information Processing Systems, 34, 2021.
[25] Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with sgd training and sample size. arXiv preprint vrXiv:1811.07006 2018.
[26] Vardan Papyan. Measurements of three-level hierarchical structure in the outliers in the spectrum 2019.
of deepnet hessians. arXiv preprini
[27] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high pegfpperformanc deep learning library. Advances in neural information processing systems, 32, 2019.
[28] Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint 1rX1::661. 2016.
[29] Levent Sagun, Utku Evei, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks. arXiv preprint arXiv: 044544 2017.
[30] Lei Wu, Chao Ma, et al. How sgd selects the global minima in over pereparraettrizee learning: dynamical stability perspective. Advances in Neural Information Processing Systems, 31, 2018.
[31] Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Under- standing common structure of hessian in neural networks. arXiv preprint arXiv:2010.04200 2020.
[32] Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. walk with sgd. arXiv preprint arXiv: arXiv: 1802.08770, 2018.
12


Reviews:
[{"review_id": "wKY3wKPS_ON", "reviewer": "Reviewer_z1Jj", "summary": "The work aims to understand the phenomena of edge of stability along with progressive sharpening (rising of leading eigenvalues of the Hessian). As a foundation, it provides an interesting observation wrt the correlation between sharpness and the output layer's norm. Then it divides the analysis into four stages where the behavior of the leading eigenvalue is studied. Moreover, it proves the theorems in the setting of two-layer linear neural networks.  Some assumptions are verified with numerical experiments.\n\n", "questions": "Please see the above concerns.\n\nTypo:\n\n1. Shall the fraction in Eq(3) be $\\frac{4}{n}$?", "limitations": "No. But it is a pure theoretical work in a standard setting of optimization, so there is no significant need to do so.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n\n1. The analysis is clear and convicing generally, with a detailed investigation into the four stages.\n2. The observation of the correlation between the sharpness and output layer's norm is quite interesting, which makes it a good start point for the following analysis.\n\nConcerns:\n\n1. How to handle several eigenvalues close to $2/\\eta$? This phenomena is observed by Cohen et al., but the assumption 3.6 and 4.1 rule it out.\n\nGenerally speaking, I would like to recommend an acceptance score for the complete analysis to resolve the problem of EoS.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "VRPrY_nOpJP", "reviewer": "Reviewer_Nthg", "summary": "This paper takes some steps towards theoretically understanding the \"Edge of Stability\" (EoS) phenomenon.  Two big open questions are: (1) why does the sharpness tend to rise whenever gradient descent is stable, and (2) once gradient descent is destabilized, why does the sharpness go back down / stop rising?\n\nThe paper has two parts: (1) a rigorous mathematical analysis of the two-layer linear network case (section 4), and (2) a nonrigorous mathematical analysis of the general neural net case combined with thorough experiments (section 3).\n\n**Rigorous analysis of two-layer linear network**:   For a two-layer linear network, the NTK is the sum of two terms: (a) a term originating from the gradients w.r.t the first layer, and (b) a term originating from the gradients w.r.t the second layer.  Term (a) has a simple form: it is the product of the second-layer weight norm ( a scalar) and the input data Gram matrix.  Sharpness is defined as the maximum eigenvalue of the NTK (rather than the Hessian).\n\nThe authors first prove (in Lemma C3) that term (b) of the NTK will barely change during the first phase of training (the initial sharpening phase) -- rigorously, they prove that the delta of this term will be O(1/m) where m is the width of the network.  This implies that the NTK will approximately stay diagonalizable by the eigenvectors of the input data Gram matrix, and also that the sharpness is essentially totally determined by the second-layer weight norm.  Thus, to prove that progressive sharpening occurs, it suffices to prove that the second-layer weight norm increases during gradient descent.  In the update equation of the second-layer weight norm, the first-order term is the negative inner product between the predictions and the residual.  The authors are able to rigorously prove that the higher-order terms don't matter, and thus to prove progressive sharpening it suffices to prove that the inner product between the predictions and residual is always negative.  This is proved in Lemma C5.\n\nThe authors next prove that after becoming destabilized, gradient descent will move to a flatter region where the sharpness is below 2 / eta.  To use the authors' own words, the proof is \"long and tedious,\" so I didn't attempt to parse the proof; but I assume that the intuition is similar to that of Proposition 3.2, which is that if the residual ever grows too big, then the last-layer weight norm will automatically decrease.  \n\n**Beyond two-layer linear networks** The authors next attempt to show experimentally that the insights from the two-layer linear NN analysis  carry over in a very literal way to the general neural network case.  (Here, I write \"next,\" but in the current draft this section occurs first chronologically.)  In particular, they demonstrate experimentally that in many networks, the dynamics of the sharpness correlate well with those of the last-layer weight norm: when the sharpness rises, so does the last-layer weight norm, and when the sharpness drops, so does the last-layer weight norm.   This seems to be more true in the initial phase of training, and becomes less true towards the end.  To their credit, the authors admit that there are counterexamples to this trend.\n", "questions": "For the two-layer NN:\n - does the degree of sharpening depend on the input dataset?  If we trained on random Gaussian data, would there still be sharpening?\n - does the degree of sharpening depend on the last-layer initialization (which is much larger in scale than the standard initialization)?\n\n", "limitations": "The paper is honest about its limitations.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- Strength: this paper is the first to prove that some architecture (in this case, a two-layer linear network) undergoes progressive sharpening.  Indeed, this is the first paper to even give any kind of explanation for why progressive sharpening might occur.  (Actually, it was not even previously known that 2-layer linear nets can exhibit progressive sharpening.)\n\n- Strength: I think that this paper is the first to prove that for some architecture (in this case, a two-layer linear network), instability causes the sharpness to decrease.  (The catapult paper includes a handwavy explanation for this phenomenon, but not a literal proof.)\n\n- Weakness: the paper is at its weakest when it attempts to argue that its EoS analysis for two-layer linear networks carries over in a very literal manner to general neural networks.  First, I would point out that the correlation between sharpness and last-layer weight norm is not very robust: in Appendix B.1.2, we see that after the first few cycles of instability, there are a huge number of 'anomaly points' (steps where the change in the sharpness is not positively correlated with the change in last-layer weight norm).  Second, I would point out that \"all layers seem to work together to influence the sharpness,\" as the authors write.\n\nOverall, I would recommend framing this paper very differently.  I suggest centering the two-layer linear network analysis rather than the debatable claims about general neural networks.  If the authors are concerned that this analysis is very tedious, I would recommend just providing the intuition in the main paper (e.g. the fact that the leading term in the change in the sharpness is the inner product between the residual and the predictions) while deferring the complete proofs to the appendix.  Then, _after_ discussing the two-layer linear network analysis, you could mention that some of the patterns might carry over to general networks.  The authors are of course free to take or leave this advice, but I think that this restructuring would make the paper more compelling.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "UvNgamhpZ1-", "reviewer": "Reviewer_VhEp", "summary": "This paper proposes a theoretical explanation for the progressive sharpening and edge of stability phenomena observed in Cohen et al (2021). Specifically, this paper claims that the gradient descent dynamics can be divided into four stages based on the value of the sharpness. For a general neural network, the paper gives a heuristic derivation for the four stages, where the proxy for sharpness is the last layer weight. It also gives a more rigorous derivation for a two-layer linear neural network.", "questions": "As mentioned in the weaknesses section, I find a number of the assumptions to be unrealistic, and don’t find the empirical justification sufficient. Can you please elaborate on why these assumptions are reasonable? In particular, why the last layer norm is a good proxy for sharpness, and why one can assume $W(t)$ changes very little for the two-layer linear network?\n\nFurthermore, can you please explain why the argument in section 3 guarantees a return to phase I?\n\nFrom a clarity perspective, I find the assumptions and heuristic derivations being stated within the main theorems hard to follow. One suggestion is to have an explicit assumption section before the main results in section 3.", "limitations": "The paper does admit that the analysis is highly heuristic, which is the main limitation of this work.", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 1, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n\nThe edge of stability phenomenon contradicts much of classical optimization, and as of yet does not have a satisfactory theoretical explanation. Therefore I find this work to be an important research direction. Furthermore, the four stage division of dynamics proposed here appears to be novel.\n\nWeaknesses:\n\nMy central issue with this paper is that the derivations are far too heuristic. As a result, I do not find the claims to be theoretically sound or a convincing explanation of the edge of stability phenomenon. Some specific instances are the following (all page references are for the version in the supplementary material):\n\nSection 3:\n- Rather than tracking the sharpness, this paper tracks $\\|A(t)\\|^2$, where $A(t)$ is the last layer weight. I don’t believe that in general the last layer weight is an accurate proxy for sharpness. The justification given here is that the difference in norms is usually the same sign as the change in sharpness; however, the empirical justification in Figure 2 is unconvincing as this correlation is only strong for the two layer linear network. Furthermore, while the analysis in Section 3 shows that $\\|A(t)\\|^2$ increases or decreases, this does not necessarily imply that the sharpness will decrease below $2/\\eta$, which is a central component of the empirical analysis in Cohen et al (2021).\n- This analysis does not track the full GD dynamics, but rather a first order Taylor expansion of the gradient. Appropriately dealing with the full dynamics seems far more challenging. This should also be stated explicitly as an assumption.\n- Assumption 3.2 is very strong, and changing eigenvectors could have a nontrivial effect on the dynamics.\n- I don’t find the argument for why the dynamics eventually return to stage I convincing. This argument is given at the bottom of page 8 + top of page 9 and is quite handwavy / unrigorous.\n- Lemma 3.5 claims to explain the non-monotonic loss decrease, by saying $R(t)$, the dynamics with the $v_1$ direction projected out, approximately follows $R’(t)$. However, it seems that the error term in this Lemma can be very large, since $\\eta$ is small and $\\lambda_r$ can be very small as well.\n\nSection 4:\n- I don’t find Assumption 4.3 to be realistic. First, the fact that $\\Gamma(t)$ is bounded during the progressive sharpening phase does not mean it will stay bounded during the EOS phase. In fact, Figure 18 clearly shows that $\\|\\Gamma(t)\\|$ spikes during the edge of stability. Another issue is that assumption 4.3 amounts to the weights $W(t)$ changing very little during training, which is essentially equivalent to saying the network is in the lazy training/NTK regime. It is known that neural networks in practice don’t follow lazy training and weights do more far from initialization, and since edge of stability is inherently a non-quadratic phenomenon, I thus find it unreasonable to assume that weights do not move far from initialization. This assumption appears to be key to the claims in section 4, and therefore I do not find the main claims of this section to be convincing.\n\nFurthermore, the supplementary material contains an alternate version of the paper where the main text is much longer (13 pages). I am not sure if this is allowed under the submission guidelines.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept