Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability
Authors: Alex Damian
Eshaan Nichani
Jason D. Lee
Venue: ICLR 2023 poster (2023)
Keywords: gradient descent
optimization
edge of stability
implicit regularization
implicit bias

Abstract: Traditional analyses of gradient descent show that when the largest eigenvalue of the Hessian, also known as the sharpness $S(\theta)$, is bounded by $2/\eta$, training is "stable" and the training loss decreases monotonically. Recent works, however, have observed that this assumption does not hold when training modern neural networks with full batch or large batch gradient descent. Most recently, Cohen at al. (2021) detailed two important phenomena. The first, dubbed \emph{progressive sharpening}, is that the sharpness steadily increases throughout training until it reaches the instability cutoff $2/\eta$. The second, dubbed \emph{edge of stability}, is that the sharpness hovers at $2/\eta$ for the remainder of training while the loss continues decreasing, albeit non-monotonically. We demonstrate that, far from being chaotic, the dynamics of gradient descent at the edge of stability can be captured by a cubic Taylor expansion: as the iterates diverge in direction of the top eigenvector of the Hessian due to instability, the cubic term in the local Taylor expansion of the loss function causes the curvature to decrease until stability is restored. This property, which we call \emph{self-stabilization}, is a general property of gradient descent and explains its behavior at the edge of stability. A key consequence of self-stabilization is that gradient descent at the edge of stability implicitly follows \emph{projected} gradient descent (PGD) under the constraint $S(\theta) \le 2/\eta$. Our analysis provides precise predictions for the loss, sharpness, and deviation from the PGD trajectory throughout training, which we verify both empirically in a number of standard settings and theoretically under mild conditions. Our analysis uncovers the mechanism for gradient descent's implicit bias towards stability.

Published as a conference paper at ICLR 2023
SELF-STABILIZATION THE IMPLICIT BIAS OF GRADI- ENT DESCENT AT THE EDGE OF STABILITY
Alex Damian*, Eshaan Nichani* & Jason D. Lee Princeton University (ad27, eshni ch, Jason lee ( edu
ABSTRACT
Traditional analyses of gradient descent show that when the largest eigenvalue of the Hessian, also known as the sharpness S(Ø), is bounded by 2/n, training is "stable" and the training loss decreases monotonically. Recent works, however. have observed that this assumption does not hold when training modern neural networks with full batch or large batch gradient descent. Most recently, Cohen et al. (2021) detailed two important phenomena. The first, dubbed progressive sharpening, is that the sharpness steadily increases throughout training until it reaches the instability cutoff 2/1,. The second, dubbed edge of stability, is that the sharpness hovers at 2/1 for the remainder of training while the loss continues de- creasing, albeit non- monotonically. We demonstrate that, far from being chaotic, the dynamics of gradient descent at the edge of stability can be captured by a cu- the iterates diverge in direction of the top eigenvector
bic Taylor expansion: of the Hcssian due to instability, the cubic term in the local Taylor expansion of the loss function causes the curvature to decrease until stability restored. This property, which we call self-stabilization, is general property of gradient de- scent and explains its behavior the edge of stability. key consequence of self- saffssaabilization is that gradient descent at the edge of stability implicitly follows projected gradient descent (PGD) under constraint S()) S 2/11. Our analysis provides precise predictions for the loss, sharpness, and deviation from the PGD trajectory throughout training, which we verify both empirically in number standard settings and theoretically under mild conditions. Our analysis uncovers the mechanism for gradient descent's 's implicit bias towards stability.
INTRODUCTION
1.1 GRADIENT DESCENT THE EDGE OF STABILITY
Almost all neural networks are trained using a variant of gradient descent, most commonly stochas- tic gradient descent (SGD) or ADAM (Kingma & Ba 2015). When deciding on an initial learning rate, many practitioners rely on intuition drawn from classical optimization. In particular, the fol- lowing classical lemma, known as the "descent lemma," provides a common heuristic for choosing learning rate in terms of the sharpness of the loss function:
Definition 1. Given loss function L(O), the sharpness at O is defined to be S(0) Amaz L(0). When this eigenvalue is unique, the associated eigenvector is denoted by u(O).
Lemma (Descent Lemma), Assume that S(Ø) S for all If = 0r - L()))
L(0(4)) S L(O.) 17(2 mf() |[V((0)))
Here, the loss decrease is proportional the squared gradient, and is controlled by the quadratic (22 nt) in 1). This function is maximized at = 1/E, popular learning rate criterion. For any 1 < 2/2, the descent lemma guarantees that the loss will decrease. As result, learning rates below 2/6 are considered "stable" while those above 2/E are considered "unstable." For quadratic
Equal contribution

---

Published as a conference paper at ICLR 2023
loss functions, e.g. from linear regression, this is tight. Any learning rate above 2/0 provably leads to exponentially increasing loss.
However, it has recently been observed that in neural networks, the descent lemma is not predictive of the optimization dynamics. Recently Cohen et al. (202T) observed two important phenomena for gradient descent, which made more precise similar observations in Jastrzebski et al. (2019); Jastrzebski et al. (2000) for SGD:
Progressive Sharpening Throughout most of the optimization trajectory, the gradient of the loss negatively aligned with the gradient of sharpness, í.e. VI((0) (S(() < 0. As result, for any reasonable learning rate 7,, the sharpness increases throughout training until it reaches S(0) = 2/n.
Edge of Stability Once the sharpness reaches 2/11 (the "break-even" point in Jastrzebski et al. 2020 ), ceases increase and remains around 2/1 for the rest of training. The descent lemma no longer guarantees the loss decreases but the loss still continues decreasing, albeit non nonmoootonically.
Training Lowe
Progressive Sharpening
Edge of Stability
200
Progressive Sharpening
Edge of Stability
Sharpness
0.4
100
0.0
1000
2000
5000
1000
2000
3000
4000
5000
Training Steps
Training Steps
1.2 SELF-STABILIZATION: THE IMPLICIT BIAS OF INSTABILITY
In this work we explain the second stage, "edge of stability." " We identify a new implicit bias of gradient descent which we call self-stabillzatiin.. Self-stabilization is the mechanism by which the sharpness remains bounded around 2/1,, despite the continued force of progressive sharpening, and by which the gradient descent dynamics do not diverge, despite instability. Unlike progressive sharpening, which is only true for specific loss functions (e.g. those resulting from neural network optimization (Cohen et al. 202T)), self stabilization is general property of gradient descent.
Traditional non-convex optimization analyses involve Taylor expanding the loss function to second order around to prove loss decrease when 17 S 2/ 2/0( When this is violated, the iterates diverge exponentially in the top eigenvector direction, u, thus leaving the region in which the loss function is locally quadratic. Understanding the dynamics thus necessitates a cubic Taylor expansion.
Our key insight is that the missing term in the Taylor expansion of the gradient after diverging in the u direction is 3 (()(u,, u), which is conveniently equal to the gradient of the sharpness at 6:
Lemma 2 Self-Stabillizttion Property). If the top eigenvalue of 2 ((0) is unique, then the sharp- ness S(0) is differentiable at O and V (() = V (()((u()) u(0)).
As the iterates move in the negative gradient direction, this term has the effect of decreasing the sharpness. The story self-stabilization is thus that as the iterates diverge in the u direction, the strength of this movement in the (() direction grows until it forces the sharpness below 2/n, at which point the iterates in the direction shrink and the dynamics re-enter the quadratic regime.
This negative feedback loop prevents both the sharpness S(O) and the movement in the top eigen- vector direction, u. from growing out of control. As a consequence, we show that gradient descent implicitly solves the constrained minimization problem:
min L()) such that S(0) < 2/n-
Specifically, if the stable set MM is defined by : (0 S()) < 2/n and V (0) 0 the gradient descent trajectory {0.} tracks the following projected gradient descent trajectory {0.) which solves the constrained problem (Barber & Ha 2017):
0 .+ = proj. 6! --)V(()))
where
proj (0) = arg min 1|0 0'
(2)
0'EM
The condition that (Z()) m(()) is necessary ensure the stability of the constrained trajectory.
2

---

Published as a conference paper at ICLR 2023
Our main contributions are as follows. First, we explain self-stabilization generic property of gradient descent for a large class of loss functions, and provide precise predictions for the loss, sharpness, and deviation from the constrained trajectory {0] throughout training (Section 4). Next, we prove that under mild conditions on the loss function (which we verify empirically for standard architectures and datasets), our predictions track the true gradient descent dynamics up to higher order error terms (Section [5). Finally, we verify our predictions by replicating the experiments in Cohen et al. (202T) and show that they model the true gradient descent dynamics (Section
2 RELATED WORK
Xing et al. (2018) observed that for some neural networks trained by full-batch gradient descent, the loss is not monotonically decreasing. [Wu et al. (2018) remarked that gradient descent cannot converge to minima where the sharpness exceeds 2/n but did not give a mechanism for avoiding such minima. [ewkowycz et al. (2020) observed that when the initial sharpness is larger than 2/ gradient descent "catapults" into a stable region and eventually converges. Dasttzebski et al. (2019) studied the sharpness along stochastic gradient descent trajectories and observed an initial increase (i.e. progressive sharpening) followed by a peak and eventual decrease. They also observed inter- esting relationships between the dynamics in the top eigenvector direction and the sharpness. Das- trzebski et al. (2020) conjectured a general characterization of stochastic gradient descent dynamics asserting that the sharpness tends to grow but cannot exceed a stability criterion given by their eq (1), which reduces to S(0) S 2/1 in the case of full batch training. Cohen et al. (2021) demonstrated that for the special case of (full batch) gradient descent training, the optimization dynamics exhibit a simple characterization. First, the sharpness rises until it reaches S(O) = 2/1 at which point the dynamiis transition into an "edge of stability" (EOS) regime where the sharpness oscillates around 2/1 and the loss continues to decrease, albeit non-monotonically.
Recent works have sought to provide theoretical analyses for the EOS phenomenon. Ma et al. (002)) analyzes EOS when the loss satisfies "subquadratic growth" assumption. [Ahn et al. (2022) argues that unstable convergence is possible when there exists a "forward invariant subset" near the set of minimizers. [Arora et al. (2022) analyzes progressive sharpening and the EOS phenomenon for normalized gradient descent close to the manifold of global minimizers. Lyu et al. (2022) uses the EOS phenomenon to analyze the effect of normalization layers sharpness for scale-invariant loss functions. Chen & Bruna (2022) show global convergence despite instability for certain 2D toy problems and in --neuron ctudent-teacher setting. The concurrent work [Li et al.) (2022b) proves progressive sharpening for a two-layer network and analyzes the EOS dynamics through four stages similar ours using the norm of the output layer as proxy for sharpness.
Beyond the EOS phenomenon itself, prior work has also shown that SGD with large step size or small batch size will lead to a decrease in sharpness (Keskar et al. 2017 Jastrzebski et al. 2017 Jastrzębski et al. 2019;] Jastrzebski et al.) 2020). Gilmer et al. (2021) also describes connections between EOS, learning rate warm-up, and gradient clipping.
At a high level, our proof relies on the idea that oscillations in an unstable direction prescribed by the quadratic approximation of the loss cause a longer term effect arising from the third-order Taylor expansion of the dynamics. This overall idea has also been used to analyze the implicit regularization of SGD (Blane et al. 20200 Damian et al. 202T; [i et al., 2022a). In those settings, oscillations come from the stochasticity, while in our setting the oscillations stem from instability.
3 SETUP
We denote the loss function by L E (R Let 0 E Rdd follow gradient descent with learning rate
7, i.e. 0.+1 := 0, 1/V (O..). Recall that M == (0 5(0) S 2/n and VI(()) mn()= = 0} is the set of stable points and proj.v := arg mingrem ||0 0'll is the orthogonal projection onto M. For notational simplicity, we will shift time so that 0g is the first point such that (()) = 2/n-
The constrained trajectory 01 is initialized with of == (00) after which it follows eq. (2)
Our key assumption is the existence of progressive sharpening along the constrained trajectory, which is captured by the progressive sharpening coefficient (0) == LVL(0)- VV)(0):
Assumption (Progressive Sharpening). Let (()) == (T) (0)) Then a0) > 0.
3

---

Published as a conference paper at ICLR 2023
Stage Progressive Sharpening
Stage Blowup
Stage. 3: Self-Stailiiation
Stage Return to Stability
Change in storpioss
VS
X- 04 8 8"1
1n (0 8%)
0 81
Figure 1: The four stages of edge of stability, demonstrated on a toy loss function (Appendix B
We focus on the regime in which there is single unstable eigenvalue, and we leave understanding multiple unstable eigenvalues future work. We thus make the following assumption V V²I(0)):
Assumption (Eigengap). For some absolute constant we have X2 (()) < c/n.
4 THE SELF-STABILIZATION PROPERTY OF GRADIENT DESCENT
In this section, we derive a set of equations that predict the displacement between the gradient de- scent trajectory }0.} and the constrained trajectory {0. Viewed as dynamical system, these equations give rise to a negative feedback loop, which prevents both the sharpness and the displace- ment in the unstable direction from diverging. These equations also allow us to predict the values of the sharpness and the loss throughout the gradient descent trajectory.
4.1 THE FOUR STAGES OF EDGE OF STABILITY: HEURISTIC DERIVATION
The analysis in this section proceeds by cubic Taylor expansion around fixed reference point 0- of 2 For notational simplicity, will define the following quantities 0*:
VL = VL(0"), V2L (() u == u(0)), VS == VS(0"), a == a(0),, := |V
where VL-VS is the progressive sharpening coefficient at 6*. For simplicity, in Section we assume that VS 1 u and VI. VS E ker (V).. and ignore higher order error terms Our main argument in Section 5 does not require these assumptions and tracks all error terms explicitly.
We want to track the movement the unstable direction u and the direction of changing sharpness VS, and thus define I2 (0, - 0*) and y/ == VS (0, 0*). Note that ye is approximately equal to the change in sharpness from 0* to 0,, since Taylor expanding the sharpness yields
S(6.) = S(0") + VS (0: 8') 2/n ye-
At a high level, the mechanism for edge of stability can be described in 4 stages (see Figure
Stage Progressive Sharpening While x, y are small, VL(0,) VL. In addition, because VL VS < O, gradient descent naturally increases the sharpness at every step. In particular,
1t41 l/- = VS (0t+1 0,) s - V VS = njar
The sharpness therefore increases linearly with rate njck
Stage 2: Blowup As Ig measures the deviation from 0* the u direction, the dynamies of x can be modeled by gradient descent on a quadratic with sharpness S(O,) Rs 2/1 + ye- In particular, the rule for gradient descent on a quadratic gives
It+1 I, nuu V ((,,) 2, n) Re R - + w]x, -(1 my:)
When the sharpness exceeds 2/n, i.e. when > O, |.z|| begins to grow exponentially.
Beginning in Section 5 the reference points for our Taylor expansions change at every step to minimize errors. However, fixing the reference point in this section simplifies the analysis, better illustrates the negative feedback loop, and motivates the definition of the constrained trajectory.
A give an explicit example a loss function satisfying these assumptions Appendix A rigorous derivation of this update in terms of S(..) instad of S(0*) requires third-order Taylor expan- sion around 8*; see Appendix [ for more details.
4

---

Published as a conference paper at ICLR 2023
Slare
XX/
Figure 2: The effect of (() (left): We plot the evolution of the ODE in eq. 44 with a= == for varying X(0), Observe that smaller X (0)'s correspond to larger curves. The four stages of edge of stability (right): We show how the four stages of edge of stability described in Section 4. and Figure correspond to different parts of the curve generated by the ODE in eq.
Stage 3: Ceff-Stabilization Once the movement the u direction sufficiently large, the loss no longer locally quadratic. Understanding the dynamiis necessitates third order Taylor expansion.
The missing cubic term in the Taylor expansion of V ((O)) is L u) V by Lemma 2 This biases the optimization trajectory in the - VS direction, which decreases sharpness. Recalling
//] the new update for y becomes:
Di+1 1ja VS. (-nV²
2
Therefore once 2t > 2a/B, the sharpness begins to decrease and continues to do so until the sharpness goes below 2/n and the dynamics return to stability.
Stage 4: Return to Stability At this point |ay is still large from stages and 2. However, the self-
stabilization of stage 3 eventually drives the sharpness below 2/n so that Yt < 0. Because the rule for gradient descent on a quadratic with sharpness S(O.) = 2/1 9t < 2/1 is F++1 -(1 (1+4 ) [as begins to shrink exponentially and the process returns to stage 1.
Combining the update for Zt, Us all four stages, we obtain the following simplified dynamics:
zt 3 Je+1 R - -(1 + 19t) )x, and 3o+1 RJ yt + 17
(3)
2
where we recall a == VL- VS is the progressive sharpening coefficient and B = ||
4.2 ANALYZING THE SIMPLIFIED DYNAMICS
We now analyze the dynamics in eq. (3. First, note that changes sign every iteration. and that, Tt-1 ss -x, due to the instability in the u direction. While eq. (3) cannot be directly modeled by an ODE due these rapid oscillations, As consequence, we can couple the dynamics of [x,], yt to the following ODE ((1), (t)::
instead model |æxel, Dt. whose update is controlled by n-
3 X(1)²
X'(I) X())(()) and Y'())
This system has the unique fixed point (X, Y) (8.0) 0) where d /2a/B. We also that this ODE can be written as Volka-Votteraa predator-prey model after a change of variables, which is classical example of a negative feedback loop. In particular, the following quantity conserved:
Lemma 3. Let ((()) == z log z - 1. Then ((X(t), Y()) = A (x(())
Y(O)
conserved.
2a
Proof. d (o(((t), ( 2Y -
3 yt) a xx
= 0. J
As result we can use the conservation of to explicitly bound the size of the trajectory:
Corollary 1. For all i, x(0) s x()) $ & X and (Y()) & Va log (0).
5

---

Published as a conference paper at ICLR 2023
The fluctuations in sharpness are O(va), while the fluctuations in the unstable direction are ()).
Moreover, the normalized displacement in the VS direction, i.e. T (0 0*) is also bounded by (()), so the entire process remains bounded by (6)). Note that the fluctuations increase as the progressive sharpening constant grows, and decrease as the self ceff-stabilization force / grows.
VS
4.3 RELATIONSHIP WITH THE CONSTRAINED TRAJECTORY 0f
60 Equation completely determines the displacement 0. - in the u, VS directions and Section 4.2 shows that these dynamics remain bounded by ( where = /2a/3. However, progress still made in all other directions. Indeed, 6, evolves in these orthogonal directions by -, VI. at every step where P is the projection onto this orthogonal subspace. This can be interpreted as first taking gradient step of -7 and then projecting out the VS direction to ensure the sharpness does not change. Lemma 13 given in the Appendix, shows that this is precisely the update for of
(eq. ( up to higher order terms. The preceding derivation thus implies that 00 0,11 s Ö(6) and
that this (6) error term is controlled by the self-stabilizing dynamics in eq.
5
THE PREDICTED DYNAMICS AND THEORETICAL RESULTS
We now present the equations governing edge of stability for general loss functions.
5.1
NOTATION
Our general approach Taylor expands the gradient of each iterate θg around the corresponding iterate of of the constrained trajectory. We define the following Taylor expansion quantities at 07:
Definition (Taylor Expansion Quantities at 0,).
VL = (L())) L == VZ)(
L*L)),>L(()) VS. = VS(O)). us = u(f))
Furthermore, for any vector-valued function v(0), we define vi Pul v(t)) where P.,, is the projection onto the orthogonal complement of u.
We also define the following quantities which govern the dynamics near 07 20x Definition 3. Let == L VS. 3r = |s and & = V22 Furthermore, we define
8 V I (1-
VS! and 8 sup, 8,-
Recall that a is the progressive sharpening force, 3r is the strength of the stabilization force, and ò controls the size of the deviations from 6; and was the fixed point in the æ direction in Section 4.2 The scalars 8,-1 capture the effect of the interactions between V.S and the Hessian.
5.2 THE EQUATIONS GOVERNING EDGE OF STABILITY
We now introduce the equations governing edge of stability. We track the following quantities:
Definition 4. Define U2 := 0 0/, x; := Ut Ut y4 == VS, Ut-
Our predicted dynamics directly predict the displacement and the full definition is deferred to Appendix C However, they have relatively simple form the ut, VSS.t directions:
Lemma 4 (Predicted Dynamics for x. y)- Let ét denote our predicted dynamics (defined in Ap- pendix Cj Letting à, = u4 1t and Jt S it, we have
82
444 -1 nge)) and Mtt1
Σ
B st
(5)
2
Note that when 3r-1 constant, update reduces the simple discussed in Section 4 which we analyze fully. When F, is large, eq. (5) demonstrates that there a self-stabilization force which acts to decrease y4 however, unlike in Section 4 the strength of this force changes with t.

---

Published as a conference paper at ICLR 2023
5.3 COUPLING THEOREM
We now show that, under a mild set of assumptions which we verify to hold empirically in Ap- pendix P the true dynamics are accurately governed by the predieted dynamics. This lets us the predicted dynamics to prediet the loss, sharpness, and the distance to the constrained trajectory of.
Our emrors depend on the unitless quantity which we verify is small in Appendix E Definition 5. Let ft == n//00 and c == sup, Ct.
To control Taylor expansion errors, we require upper bounds on V and its Lipschitz constant Assumption 3. Let p3. PA to be the minimum constants such that for all 0, ||v S P3 and T is Px-lipschitz with respect to I1-Ilop Then we assume that P1 (
Next, we require the following generalization of Assumption D
Assumption 4. For all t, | |
VS. (1) and V = (
Finally, we require of "non- worst-caset assumptions, which are that the quantities L V and Amin LV are nicely behaved the directions orthogonal to ugs which generalizes the eigen- gap assumption. We verify the assumptions on V and V empirically in Appendix E
(n..n))
Amiin (V S O()).
Assumption 5. For all t and u, w uu
| L.ll
IV
With these assumptions in place, we can state main theorem which guarantees i. ê predict the loss, sharpness, and deviation from the constrained trajectory up to higher order terms:
Theorem 1. Let (c)²) and assume that min; & | 2 c1ó. Then for any t S T, we have L(0,) L(e)) / /n)
(Loss)
S(O,) = 2/1 + + (S. u) + 0(2/n)
(Sharpness)
0g 0: é. ((8))
(Deviation from fi)
The sharpness is controlled by the slowly evolving quantity ju and the period-2 oscillations of (VS- U) This combination of gradual and rapid periodic behavior observed by Cohen et al. (2021) and appears in our experiments. Theorem also shows that the loss at 0g spikes whenever de is large. On the other hand, when å; is small, (0:)) approaches the loss of the constrained trajectory.
EXPERIMENTS
We verify that the predicted dynamics defined in eq. (5) accurately capture the dynamics of gradient
descent at the edge stability by replicating the experiments in (Cohen et al.) 202T] and tracking the
theory 3-layer MLP and a 3-layer CNN trained with mean squared error (MSE) on a 5k subset of CIFAR10
deviation of gradient descent from the constrained trajectory. In Figure 3l evaluate
and 2-yeyer Transformer (Vaswani et al., 2017) trained with MSE SST2 provide additional experiments varying the learning rate and loss function in Appendix G which use
al. (2013). We
the generalized predicted dynamics described Section 7.2 For additional details, see Appendix D
Figure 3 confirms that the predicted dynamics eq. (5) accurately predict the loss, sharpness, and distance from the constrained trajectory. In addition, while the gradient flow trajectory diverges from the gradient descent trajectory at a linear rate, the gradient descent trajectory and the constrained trajectories remain close throughout training. In particular, the dynamics converge to the fixed point
([ær): yt) = (5,, 0) described Section 4.2 and ||0, - 0!1l -> 5,. This confirms our claim that
gradient descent implicitly follows the constrained trajectory eq. (2).
In Section 5 various assumptions on the model were made to obtain the EOS behavior. In Ap- pendix numerically verify these assumptions to ensure the validity of our theory.
For simplicity of exposition, we make these bounds L globally, however our proof only requires
them in small neighhorhood of the constrained trajectory 01.

---

Published as a conference paper at ICLR 2023
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
MLP+MSE on CIFAR1O, 77 = 0.002
0.150
1020
1010
Distance gt
Sharpness
Sharpness
0.004
Loss
0.145
1000
0.002
1000
995
0.140
0.000
200
400
200
400
0.000
0.002
0.004
200
400
0.15
Distance 0
0.02
Sharpness
1050
Sharpness
1005
Loss
0.14
1025
1000
0.13
1000
995
0.00
0
1000
2000
1000
2000
0.000
0,002
0.004
1000
2000
Step
Step
[X]
Step
CNN+MSE on CIFAR1O, 17 = 0.005
420
Distance to 0
0.0075
0.6250
Sharpness
Sharpness
Loss
410
0.0050
0.6225
400
0.0025
395
0.0000
100
200
100
200
0.000
0.002
0.004
100
200
0.62
Sharpness
550
Sharpness
4005
Distance 6
0.06
Loss
0.04
0.61
0.60
395
400
0.00
1000
2000
1000
2000
0.000
0.002
0.004
1000
2000
Step
Step
[X]
Step
Transformer+MSE on SST2, n = 0.001
2010
0.004
0.670
Sharpness
Sharpness
Distunce to
Loss
2000
0.002
0.665
0.000
0
200
400
200
400
0.000
0.001
0.002
200
400
0.67
2010
Sharpness
Sharpness
Distunce to
Loss
2050
0.66
0.005
0.65
2000
0.000
0
500
1000
0
500
1000
0.000
0.001
0.002
500
1000
Step
Step
[X]
Step
Figure 3: We empirically demonstrate that the predicted dynamics given by eq. 55) track the true EOS dynamics. For each learning rate, the top row is zoomed in copy of the bottom row which isolates one cycle and represented by the dashed rectangle. Reported sharpnesses are two-step averages for visual clarity. For additional experimental details, see Section 0 and Appendix D
8

---

Published as a conference paper at ICLR 2023
7 DISCUSSION
7.1 TAKEAWAYS FROM THE PREDICTED DYNAMICS
A key consequence of the predicted dynamics is that the loss and sharpness only depend on (de. it which are governed by the 2D dynamical system eq. (5). Therefore, understanding the EOS dy- namies only requires analyzing this dynamical system, which is generally well behaved (Figure Furthermore, we expect žt, de to converge to (=0,, 0). the fixed point of eq. (5). In fact, Figure shows that after a few cycles, Be) indeed converges to this fixed point. We can accurately predict
its location as well as the loss increase from the constrained trajectory due to é. o O.
7.2 GENERALIZED PREDICTED DYNAMICS
In order for our cubic Taylor expansions to track the true gradients, we require a bound on the fourth derivative of the loss (Assumption B) This is usually sufficient to capture the EOS dynamics as demonstrated by Figure 3 and Appendix E However, this condition was violated in some of our
experiments, especially when using logistic loss. To overcome this challenge, we developed a gen- eralized form of the predicted dynamics whose definition we defer to Appendix [F] These generalized predictions are qualitatively similar to those given by the predicted dynamics in Section however, they precisely track the dynamics of gradient descent in a wider range of settings (see Appendix G
7.3 IMPLICATIONS FOR NEURAL NETWORK TRAINING
Non-Monotonic Loss Decrease An important property of EOS is that the loss decreases over long time scales, albeit non-monotonically. Our theory provides a clear explanation for this phenomenon. We show that the gradient descent trajectory remains close to the constrained trajectory (Sections 4 and its
and 5). Since the constrained trajectory is stable, it satisfies a descent lemma (Lemma loss monotonically decreases. Over short time periods, the loss is dominated by the rapid fluctua- tions of z; described in Section 4 Over longer time periods, the loss decrease of the constrained trajectory overpowers the bounded fluctuations of xr, leading to an overall loss decrease.
Generalization & the Role of Large Learning Rates Prior work has shown that in neural net- works, both decreasing sharpness of the learned solution (Keskar et al., 2017 Dziugaite & Roy 2017 Neyshabur et al. 2017 Jiang et al. 2020) and increasing the learning rate (Smith et al. ] 2018 Li et al. 2019 Lewkowycz et al., 2020) are correlated with better generalization. Our analysis shows that gradient descent implicitly constrains the sharpness to stay below 2/1,, which suggests larger learning may improve generalization by reducing the sharpness. In Figure 4 we confirm that standard setting, full-batch gradient descent generalizes better with large learning rates.
Training Speed Additional experiments in (Cohen et al. 202 1, Appendix F) show that, despite the
instability in the training process, larger learning rates lead to faster convergence, This phenomenon is explained by our analysis. Gradient descent is coupled to the constrained trajectory which min- imizes the loss while constraining movement in the u, V.S,- directions. Since only two directions "off limits," " the constrained trajectory can still move quickly in the orthogonal directions, using the large learning rate to accelerate convergence. We demonstrate this empirically in Figure 4
We defer additional discussion of our work, including the effect of multiple unstable eigenvalues and connections to Sharpness Aware Minimization (Foret et al., (02I), warm-up (Gilmer et al. 2021).
and scale-invariant loss functions (Lyu et al. 2022 Appendix
7.4 FUTURE WORK
An important direction for future work is understanding the EOS dynamics when there are multiple unstable eigenvalues, which we briefly discuss in Appendix H Another interesting direction is un- derstanding the global convergence properties at EOS, including convergence to a KKT point of the constrained update eq. (2). Next, our analysis focused on the EOS dynamics but left open the ques- tion of why neural networks exhibit progressive sharpening. Finally, we would like to understand the role of self-staabliization in stochastic- ttocasstic-radieen gradient descent and how interacts with the implicit biases of SGD (Blane et al. 2020 Damian et al. 2021
Li et al. 2022a

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
AD acknowledger support from a NSF Graduate Research Fellowship. EN acknowledges support from National Defense Science & Engineering Graduate Fellowship, and NSF grants 19F-1907661 and DMS-2014279, JDL, AD, EN acknowledge support of the Sloan Research Fellowship. NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, and NSF-CAREER under award #21 14994.
The authors would like to thank Jeremy Cohen, Kaifeng Lyu, and Lei Chen for helpful discussions throughout the course of this project. We would especially like thank Jeremy Cohen for suggesting the term "self- sooff-stabilization" to describe the negative feedback loop derived in this paper.
REFERENCES
Kwangjun Ahn, J. Zhang, and Suvrit Sra. Understanding the unstable convergence of gradient descent. ArXiv, abs/2204.01050, 2022.
Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi. Understanding gradient descent on the edge of stability in deep learning. In Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 948-1024. PMLR, 2022.
Rina Foygel Barber and Wooseok Ha. Gradient descent with nonconvex constraints: local concavity determines convergence. arXiv: Optimization and Control, 2017.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference on Learning Theory, pp. 483-513, 2020.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnn and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. cww wwwpo/////oo////ooooooooooo/e//e
Lei Chen and Joan Bruna. On gradient descent convergence beyond the edge of stability. ArXiv, abs/2206.04172. 2022.
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference Learning Representations, 2021.
Alex Damian, Tengyu Ma, and Jason D. Lee. Label noise SGD provably prefers flat global mini- mizers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence, 2017.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-awar minimiza- tion for efficiently improving generalization. International Conference on Learning Represen- tations, 2021.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Car- doze, George Dahl, Zachary Nado, and Orhan Firat. A loss curvature perspective on training instability in deep learning, 2021. URL httpst::: arxiv http/raxxvvvrrrraa////
04369
Stanislaw Jastrzebski. Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Ben- gio, and Amos Storkey. Three factors influencing minima in sgd. ArXiv, abs/1711.04623, 2017.
Stanislaw Jastrzebski, Maciej Szymezak, 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 Learning Representations, 2020.
10

---

Published as a conference paper at ICLR 2023
Stanistaw Jastrzebski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amost Storkey. On the relation between the sharpest directions DNN loss and the SGD step length. In International Conference on Learning Representations, 2019.
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.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017.
Diederik P. Kingma and Jimmy Ba. 0bss/412.6990, 2015.
Adam: A method for stochastic optimization. CoRR,
Aitor Lewkowyez, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. ArXiv, 0tts22003.02188, 2020.
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate training neural networks. In Advances Neural Information Processing Systems, 2019.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022a.
Zhouzi Li, Zixuan Wang, and Jian Li. Analyzing sharpness along gd trajectory: Progressive sharp- ening and edge of stability. arXiv, abs/2207.12675. 2022b.
Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normal- ization layers: Sharpness reduction. ArXiv, abs/2206.07085, 2022.
Chao Ma, Lei Wu, and Lexing Ying. The multiscale structure of neural network loss functions: The effect on optimization and origin. ArXiv, abs/2204.1 1326, 2022.
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, and Nati Srebro. Exploring general-
ization in deep learning. In Advances in Neural Information Processing Systems, 2017.
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Swish: a self-gated activation function. arXiv: Neural and Evolutionary Computing, 2017.
Samuel L. Smith, Pieter-Jan Kindermans, and Quoc V. Le. Don't decay the learning rate, increase the batch size. International Conference on Learning Representations, 2018
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. tional Linguistics. URL htt.s: http://www..lclw...........c acweeb. org/ htt.:/wwww.....//.ooooooo///l---
Seattle, Washington, USA, October 2013. Association for Computa-
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NIPS, 2017.
Lei Wu, Chao Ma, and E Weinan. How sgd selects the global minima over-parameterizee learning: A dynamical stability perspective. In NewrlPS, 2018.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. ArXiv, abs/1802.08877, 2018.
11

---

Published as a conference paper at ICLR 2023
TABLE OF CONTENTS
Introduction
Gradient Descent at the Edge Stability
1.2
Self-stabilization: The Implicit Bias of Instability
2
Related Work
3
Setup
4
The Self-stabilization Property of Gradient Descent
The Four Stages of Edge of Stability: A Heuristic Derivation
4.2
Analyzing the simplified dynamics
4.3
Relationship with the constrained trajectory 0
5
The Predicted Dynamics and Theoretical Results
5.1
Notation
5.2
The equations governing edge of stability
5.3
Coupling Theorem
6
Experiments
7
Discussion
7.1
Takeaways from the Predicted Dynamics
7.2
Generalized Predicted Dynamics
7.3
Implications for Neural Network Training
7.4
Future Work
Notation
13
B A Toy Model for Self-Stabilization
13
Definition of the Predicted Dynamics
13
D Experimental Details
14
D. Architectures
14
D.2 Datal
14
D.3 Experimental Setup
14
Empirical Verification of the Assumptions
15
The Generalized Predicted Dynamics
21
F1 Deriving the Generalized Predicted Dynamies
21
12

---

Published as a conference paper at ICLR 2023
F22
Properties of the Generalized Predicted Dynamics
21
Additional Experiments
22
The Benefit of Large Learning Rates: Training Time and Generalization
22
G.2
Experiments with the Generalized Predicted Dynamics
23
Additional Discussion
28
H.T
Scale Invariant Lemmas
29
Proofs
32
Properties of the Constrained Trajectory
32
1.2 Proof of Theorem
34
L.3 Proof of Auxiliary Lemmas
35
A
NOTATION
We denote by V L(0) the k-th order derivative of the loss L at 0. Note that V L(O) is symmetric k-tensor in (IRd) when IRd.
For symmetric k-tensor T, and vectors - Uj E IRd we will use T(u Uj) to denote the tensor contraction of T" with uj...., Ug i.c.
[T(i..... U/ (k)],,m.d
T (U1)ie ... (uy)ie-
We use Pai....uy to denote the orthogonal projection onto span(u1,.... ux) and P....... u is the projection onto the corresponding orthogonal complement.
For matrices A A Au. we define
II k=1
Ax = A....A A, and II k=t
A A1.
B A Toy MODEL FOR SELF-STABILIZATION
For a, 3 > 0, consider the function
L(x,y,2) (2
x² 2 vB
By
initialized the point ( 0, 0). Note that the constrained trajectory will follow of 0. v o. of = -rjt as it cannot decrease y without increasing the sharpness past 2/n. We therefore have:
VL, = 0. α
V3" 1. = [1, S, 2/n + BBy V = Suuzt VS, 0, / 3,
Note that this satisfies all of the assumptions in Section 4 and it satisfies = - VL. VS, and 8 ||? This process will then follow eq. in the x, y directions while tracks the constrained trajectory 0f moving linearly in the - -p V L [0,0, direction.
C DEFINITION OF THE PREDICTED DYNAMICS
Below, we present the full definition of the predicted dynamics:
13

---

Published as a conference paper at ICLR 2023
Definition 6 (Predicted Dynamics, full). Define do uo, and ler i;
82 Ue D1
VS
-i Then
07 x7 2
vit1 Purci (1) np 8e-1 VS
(1 m 1141
(6)
For convenience, we will define the map step, R + Rd as follows: Definition 7. Given a vector and timestep define (t)p,(v) by
x²
P u0+1
(t) -Pu (1- VSS
M++1 (1) -(1 my)ar.
(8)
where x = Ut and y S n.
It is casy to see that i/41 step, (ét).
Proof of Lemma 4 Defining A, = (I L we can unfold the recursion in eq. the following formula for dt.
6 to obtain
s+1
8 x
PL Mal1
Ag
VS+
- + (1++1 myt) )=t u
(9)
2
Lk=t
It is then immediate to see that i; = Ut ue d1-
V.S.t
.y have the following simple update:
83 x 2
L -1 )z) and vi+1
1
8 -t
where we recall that we have defined
Base = V
II A VS.
(10)
D EXPERIMENTAL DETAILS
D.1 ARCHITECTURES
We evaluated our theory on four different architectures. The 3-layer MLP and CNN are exact copies of the MLP and CNN used in (Cohen et al., The MLP has width 200, the CNN has width 32, and both are using the swish activation |Ramachandran et al. 2017). We also evaluate on a ResNet18 with progressive widths 16, 32, 64, 128 and on 2-layer Transformer with hidden dimension 64 and two attention heads.
D.2 DATA
We evaluated our theory on three primary tasks: CIFARIO multi-class classification with both cate- gorical MSE loss and cross-entropy loss, CIFARIO binary classification (cats vs dogs) with binary MSE loss and logistic loss, and SST2 (Socher al. 2013) with binary MSE loss and logistic loss.
D.3 EXPERIMENTAL SETUP
For every experiment, we tracked the gradient descent dynamics until they reached instability and then began tracking the constrained trajectory, gradient descent, gradient flow, and both our pre- dicted dynamics (Section(5) and our generalized predicted dynamics (Appendix F). In addition, tracked the various quantities on which we made assumptions for Section order to validate these assumptions. We also tracked the second eigenvalue of the Hessian at the constrained trajectory
14

---

Published as a conference paper at ICLR 2023
throughout training and stopped training once it reached 1.9/1, ensure the existence of a single unstable eigenvalue. Finally, as the edge of stability dynamics are very sensitive to small perturba- tion when [z] is small (see Figure we switched to computing gradients with 64-bit precision after first reaching instability to avoid propagating floating point errors.
Eigenvalues were computed using the LOBPCG sparse eigenvalue solver in JAX (Bradbury al., 2018). To compute the constrained trajectory, we computed a linearized approximation for proj M inspired by Lemma 13 along with Newton step in the u, direction ensure that VL 0. Each linearized approximation step required recomputing the sharpness and top eigenvector and each projection step then consisted of three linearized projection steps, for total of three eigenvalue computations per projection step.
Our experiments were conducted in JAX BBradbury al. 2018), using gttps: com/ locus lab/ http://ggeoeftaabaaeeeee staaiiaaiiiii lyy as a reference for replicating the experimental setup used in (Cohen et al. (021)) All experiments were conducted on two servers, each with NVIDIA GPUs. Our code can be found at httttts://gtttiiii
com/ adamian98/EOS
E EMPIRICAL VERIFICATION OF THE ASSUMPTIONS
For each of the experimental settings considered (MLP+MSE, CNN+MSE, CNN+Logistic, M+MSE, Transforme++MSEE, Transformer-LLgisstig,, we plot a number of quantities along the constrained trajectory to verify that the assumptions made in the main text hold. For each learn- ing rate we have plots tracking various quantities, which verify the assumptions as follows: D is verified by the 1st plot, being small is verified by the 2nd plot, Assumption verified by the 3rd and 4th plots, Assumption 3 is verified by the 5th plot, and Assumption 5 is verified by the last 3 plots. As described in the experimental setup, training is stopped once the second eigenvalue is 1.9/n. n, so Assumption always holds with c 1.9 as well.
Assumption
15

---

Published as a conference paper at ICLR 2023
MLP+MSE on CIFARIO 17 0.02
NS,
V
0
V4. VS,
& = ny V4
FL FS
1.0
0.5
0.00
0.0
0
100
100
100
100
Ps
( op
L'L,VLL VL)
n
| op li
V VI,
2
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
0
100
0
100
100
0
100
17
0.01
VS
||
04 = V4 VS,
& = n 04
VL P55
0.05
1.0
1.0
20
0.5
0.00
0.0
100
200
100
200
100
200
100
200
P
(1
V'L, VL)
V lop [
F44 VL.I
1.0
1.0
1.0
2-
0.5
0.5
0.5
0
0.0
0.0
0.0
0
100
200
0
100
200
100
200
0
100
200
77
0.005
VS,
0g
V4. VS,
& = 17 //4 04
VL PS-
1.0
25
0.02
0.5
0.0
0
500
500
500
500
02
v (1 o
IV,, V4.)
V7 4 bp
1|
F4 VL.I
V
||
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
0
500
0
500
500
0
500
77
0.002
VL, VS,
VVS, ||
0 = VL, VS,
& : 17 aV
C
1.0
1.0
Pa
50
0.01
0.5
0.0
0
1000
2000
1000
2000
1000
2000
1000
2000
P4
V1, V1,)
7P
L Lop
L lop V
VV |
1.0
1.0
0.5
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
0
1000
2000
16

---

Published as a conference paper at ICLR 2023
CNN+MSE on CIFARIO
17 0.02
VI,
04
V,, VS,
& nVa
1.0
0.05
0.00
0.0
0.0
2001
400
200
400
200
400
0
200
400
P4
VVL,
V.L a
1
V'L 10llopp1111
V4L op VLJI
V 1l
1.0
1.0
1.0
0.5
0.5
0.5
0.5-
0.0
0.0
0.0
400
200
400
17
0.005
VL.VV,
VS
V4,- VS,
& 17
0
1.0
VL VS
1.0
P2
20
0.02
0.5
0.5
0-
0.00
0.0
0.0
0
1000
1000
1000
0
1000
P4
V (v
VVL. VL)
n
V- Leloplyy
V VL
VV
1.0
0.5
0.0
0.0
0.0
1000
1000
1000
0
1000
77
0.002
VL. VS,
VS, I|
0g =
V44 VS,
= 17 C4
1.0
VI, VS
1.0
P
0.01
0.5
0.00
0.0
0.0
1000
2000
1000
2000
1000
2000
1000
2000
P4
V V'i
V'L VL)
LV +
np?
L lop VL
V
V.
1.0
1.0
1.0
MN
0.5
0.5
0.5
0.0
0.0
0.0
1000
2000
1000
2000
1000
2000
1000
2000
17

---

Published as a conference paper at ICLR 2023
CNN+Logistic on CIFARIO (cats vs dogs)
n
0.02
VS
V
0g
V4
VS,
8. = 7y 04
ISS
0.05
1.0
1.0
5-
0.5
0.00
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
1000
2000
P4
VLL.VL VL)
(v
| op
by
| VL,
V 1|
1.0
1.0
1.0
5
0.5
0.5
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
0
1000
2000
1000
2000
17
0.01
VS
V
0
V4 VS,
& = ny 04
1.0
1.0
P7
10
0.02
0.5
0.00
0.0
0
1000
2000
1000
2000
1000
2000
1000
2000
P4
V Liy
(VL. VL.)
(v, wf
n
V.4 lon [v
V.L. lop VL 1
V 1
1.0
1.0
1.0
10
0.5
0.5
0.5
0
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
0
1000
2000
77
0.005
VS,
V ||
0
VL VS,
& = 17 1
1.0
V
1.0
P3
0.02
0.5
0.00
0.0
0
1000
2000
1000
2000
1000
2000
1000
2000
04
VVL,,VL
(v, wf)
V.L lon [v
V L lop VL
V 1
1.0
1.0
1.0
20
0.5
0.5
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
0
1000
2000
77
0.002
VV.VV
VS.
4- == VL VS,
E, : 17 nVa
1.0
|V4|l VS-11
1.0
P3
0.02
100
0.00
0.0
1000
2000
1000
2000
1000
2000
0
1000
2000
P4
?P
V
L. lep
V4
V 1||
20
1.0
0.5
0.5
0.0
0.0
0.0
0
1000
2000
1000
2000
1000
2000
0
1000
2000
18

---

Published as a conference paper at ICLR 2023
ResNet18+MSE on CIFARIO (cats vs dogs)
77
0.002
VV
V.S II
44
VL VS,
& = my 04
1.0
FLIIIVS
P3
0.1
0.5
0.0
0.0
0.0
500
1000
500
1000
0
500
1000
0
500
1000
P4
LVL.V VL)
nP
V oo by
V-4 oop VL
| |
1.0
1.0
2.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
500
1000
500
1000
0
500
1000
0
500
1000
17
0.001
VS
V
0g
44. VS,
E 17
04
|V4 VS.
Ps
5000
0.05
0.00
0.0
500
1000
500
1000
500
1000
500
1000
P4
VV,,,VI, VV,)
nPi
V 1y
V Tap
V4
V+
1.0
1.0
1.0
5-
0.5
0.5
0.5-
0.0
0.0
0.0
500
1000
500
1000
500
1000
500
1000
Transformer+MSE on SST2
0.002 17
VV,
VS;
0₂
44 VS,
& 7/
1.0
VL VS
1.0
Pa
200
0.02
0.5
0.5
0.00
0.0
0.0
1000
1000
1000
1000
Ps
(1
LVL. VL)
V2
V oo by
V op VL
Z
1.0
1.0
1.0
20
0.5
0.5
0.5
0.0
0.0
0.0
1000
1000
1000
1000
77 = 0.001
VS.
VS-
04
VI, VS,
E 17 04
1.0
VL.
VS
1.0
Pa
0.01
0.5
0.5
0.00
0.0
0.0
0
1000
1000
0
1000
0
1000
P4
1',, V
v L. w
x 100
np5
1.0
on
1.0
V lop VL
1.0
V
2.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
1000
1000
1000
1000
19

---

Published as a conference paper at ICLR 2023
Transtrrmer-Logisic on SST2
17 0.02
VS,
04
VL, VS,
& = nya
VL I
1.0
1.0
0.1
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
1000
2000
P1
4 (1y
LLL. VL)
V lop 1
V3r V7,
V |=
1.0
1.0
1.0
5
0.5
0.5
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
1000
2000
17
0.01
V.,
VSS ||
0₂ = VL. VS,
& := n
1.0
V4. VS- ||
1.0
P1
0.5
0.0
0
1000
2000
1000
2000
1000
2000
1000
2000
P4
v 1.
VLL VL)
L. m
nP
L.L lop VL
V 1
1.0
1.0
1.0
10
0.5
0.5
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
0
1000
2000
77
0.005
VI,
VS
||
0
VL. VS,
E, := 17 //
1.0
VL. V5- 11
1.0
P1
50
0.5
0.0
1000
2000
1000
2000
1000
2000
1000
2000
P1
v 11
VV,, V1+)
nPs
V Lilloo 1u
V lop V
| |1
20
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
0
1000
2000
0
1000
2000
1000
2000
0
1000
2000
77
0.002
VS,
V ||
0 = VL, VS,
0.02
E, := TVOk
1.0
| ||VS|| II
1.0
P3
0.00
0.0
0
1000
2000
1000
2000
1000
2000
0
1000
2000
P4
(V..v.
7P
V 7
V 1p
V4
V
1.0
10
0.5
0.5
0.0
0.0
0.0
1000
2000
1000
2000
0
1000
2000
0
1000
2000
20

---

Published as a conference paper at ICLR 2023
F THE GENERALIZED PREDICTED DYNAMICS
Our analysis relies on a cubic Taylor expansion of the gradient. However, in order for this Taylor expansion to accurately track the gradients we need a bound on the fourth derivative of the loss (Assumption 3 Section 6 and Appendix L show that this approximation is sufficient to capture the dynamics of gradient descent at the edge of stability for many standard models when the loss criterion is the mean squared error. However, for certain architectures and loss functions, including 8 and models trained with the logistic loss, this condition is often violated.
In these situations, the loss function in the top eigenvector direction is either sub-quadratic, mean- ing that the quadratie Taylor expansion overestimates the loss and sharpness or super- super-quadratic, meaning that the quadratic Taylor expansion underestimates the loss and sharpness. To capture this phenomenon, we derive a more general form of the predicted dynamics which reduces to the stan- dard predicted dynamics in Section 5 when the loss in the top eigenvector direction is approximately quadratic. In addition, Appendix G shows that the generalized predicted dynamics capture the dy- namics of gradient descent at the edge of stability for both mean squared enror and cross-entropy in all settings we tested.
F.1 DERIVING THE GENERALIZED PREDICTED DYNAMICS
To derive the generalized predicted dynamies, we will abstract away the dynamics the top eigen- vector direction. Specifically, for every / we define
x²
(x) L(A) + zuu) L(0))
7
We say that L is sub-quadratic at t if (1) < 0 and super-quadratic if F((x) > 0.
Note that knowing F; is not sufficient to capture the dynamics in the us direction. Specifically,
L++1 = 1742 V + v2) + x, 1744 L + ru.).
It is still critically important to track the effect that the movement in the VS,- direction has the dynamics of x. As in Section 4.1 the effect of the movement in the VS, direction on the dynamics of is changing the sharpness by yt- This gives us the generalized predicted dynamics update:
8
xx²
2 = P ue (I v1) + TP
VS.
2 u4+1
2
where 2741 -11 )xt (x))
Note that when Fr(xr) 0 is exactly quadratic, this reduces the standard predicted dynamics update in eq. (6). Note that the update for is completely unchanged:
Lemma 5. Restricted to the ue, VS, directions, the generalized predicted dynamics oz imply:
62 - x
# = -1 ) (xt)) and Vi+1 =
8
(11)
2
8=0
The proof is identical to the proof of Lemma
F.2 PROPERTIES OF THE GENERALIZED PREDICTED DYNAMICS
Note that due the sign flipping argument in Appendix can assume that F is an even function as the odd part will only influence the dynamics through additional oscillations of period 2, so throughout the remainder of this section we will assume that F((x) F((-z). Otherwise, we can simply redefine F by its even part.
Next, note that the fixed point of eq. (1]) is still when 2, ò,, regardless of the shape of F,, due to the need to stabilize the S.S.- direction. This contradicts previous 1-dimensional analyses of edge
TThis sub-quadratic phenomenon was also observed (a et al. 2022
21

---

Published as a conference paper at ICLR 2023
of stability in which the fixed point in the top eigenvector direction strongly depends the shape of Ft, the loss the u, direction.
The limiting value of y/ therefore be read from the update for zt. If (ōe, y) is orbit of period 2 of eq. (TT) then
F())))
6, = (1 + 19)6, ().) y
6t
In addition, note that the sharpness longer be approximated as S(6,) 2/1 y as the sharpness now changes along the u, direction. In particular, it changes by (x) so that
S(0,) x 2/n + Ur + Fx)).
Therefore, the limiting sharpness of eq. is
(5))
S(O.) - 2/n
(6.)-
5
When Fg
and the loss is exactly quadratic in the u direction, this update reduces to fixed point
predictions in Section 4.1
One interesting phenomenon observed by Cohen et al. (0021) is that with cross-entropy loss, the sharpness
exactly 2/n, but usually hovered above it. This contradicts the predictions of the standard predicted dynamics which predict that the fixed point has sharpness O. However, using the generalized predicted dynamics eq. (II) we can give a clear explanation.
When the loss sub-quadratic, e.g.
when Fe(x)) -1 we have
S(,,) + 2/n 87
8
P4 2/n
2/n
so the sharpness will converge to value below 2/ 2/1z. On the other hand if the loss is super-quadratic, the sharpness converges to a value above 2/1j. More generally, whether the loss converges to a value above or below 2/1) depends the sign of F." (6.) (((.)).
In our experiments in Appendix G we observed both sub-quadratic and super-quadratic loss func- tions. In particular, the loss was usually sub-quadratic when first reached instability but gradually became super-quadratic as training progressed at the edge of stability.
G ADDITIONAL EXPERIMENTS
G.1 THE BENEFIT OF LARGE LEARNING RATES: TRAINING TIME AND GENERALIZATION
We trained ResNet18 with full batch gradient descent on the full 50k training set of CIFARIO with various learning rates, in addition to the commonly proposed learning rate schedule := 1/S(0,). We show that despite entering the edge of stability, large learning rates converge much faster. In addition, due to the self-stabilization effect of gradient descent, the final sharpness is bounded by 2/n which is smaller for larger learning rates and leads to better generalization (see Figure 4p.
Training Accuracy
100
Training Loss
Test Accuracy
1 0.2
60
2-
n = 0.1
50
50
0.05
25
40
n = 0.02
0.005
30
5000
10000
5000
10000
5000
10000
n 1/S(0,)
Step
Step
Step
Figure 4: Large learning rates converge faster and generalize better (ResNet18 and CIFAR10).
22

---

Published as a conference paper at ICLR 2023
G.2
EXPERIMENTS WITH THE GENERALIZED PREDICTED DYNAMICS
MLP+MSE on CIFAR10
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
n = 0.02
110
0.04
0.64
120
Sharpness
Sharpness
105
Distance to
0.03
Loss
0.62
110
100
0.02
100
95
0.01
90
0.58
0.00
0
100
100
0.00
0.02
100
Step
Step
Step
17
0.01
0.53
230
0.03
0.52
Loss
0.51
Sharpness
220
Sharpness
Distance to
0.02
200
0.50
195
0.49
0.48
190
190
0.00
100
200
100
200
0.00
0.01
100
200
Step
Step
[x]
Step
n7
0.005
460
0.38
Sharpness
440
Sharpness
Loss
Distance
0.02
400
420
0.36
0.01
400
390
0.34
250
sûo
250
500
0.000
0.005
0.010
250
500
Step
Step
[x]
Step
17
0.002
1010
0.150
0.020
1060
Distance g²
Sharpness
Sharpness
1005
0.015
Loss
1040
0.1440
1000
0.010
1020
0.005
0.130
1000
995
1000
2000
1000
2000
0.000
0.002
0.004
1000 Step
2000
Step
Step
X
23

---

Published as a conference paper at ICLR 2023
CNN+MSE on CIFAR10
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
0.02
110
0.08
0.74
0.73
140
0.06
Sharpness
Sharpness
Loss
0.72
Distance
0.04
120
95
0.70
100
0.00
90
200
400
200
400
0.02
200
400
Step
Step
Step
17
0.005
410
0.05
Distance to 01
0.62
Sharpness
500
Sharpness
Loss
0.04
0.61
450
0.02
0.60
AAAAmm
0.00
0
1000
1000
0.0000
0.0025
0.0050
1000
Step
Step
Step
17 0.002
0.5375
1080
1005.0
0.015
0.5350
1060
Distance to 0
Sharpness
Sharpness
1002.5
Loss
0.5325
1040
0.010
1000.0
0.5300
1020
997.5
0.005
0.5275
0.5250
1000
995.0
0.000
2000
2000
0.000
0.002
2000
Step
Step
[x]
Step
24

---

Published as a conference paper at ICLR 2023
LogiLoi on CIFARIO (cats vs dogs)
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
17
0.02
104
0.3
0.58
Distance O
Sharpness
Sharpness
102
0.2
Loss
100
0.54
0.1
0.52
0.0
1000
2000
1000
2000
0.01
1000
2000
Step
Step
[x]
Step
17
0.01
204
0.10
260
202
Distance 0
Sharpness
Sharpness
Loss
0.53
240
0.06
200
0.04
0.52
220
198
0.51
200
196
0.00
1000
2000
1000
2000
0.000
0.005
1000
2000
Step
Step
Step
0.005
460
0.475
405
Loss
Sharpness
440
Sharpness
Distance to
400
0.465
420
0.02
0.450
395
400
0.455
0.00
1000
2000
1000
2000
0.000
0.005
1000
2000
Step
Step
[x]
Step
77
0.002
1060
1004
0.025
Distunce 0
0.020
0.3775
Sharpness
1040
Sharpness
Loss
1000
0.015
0.3750
1020
0.010
0.3725
996
0.005
994
0.000
1000
2000
0
1000
2000
0.000
0.002
1000
2000
Step
Step
[x]
Step
25

---

Published as a conference paper at ICLR 2023
ResNet 8+MSE on CIFAR10 (cats vs dogs)
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
n
0.002
1,25
5000
1.00
Distance to 6
Sharpness
Sharpness
1080
Loss
0.75
0.4
3000
1000
0.50
2000
950
0.2
0.25
1000
900
0
500
1000
500
1000
0.00
0.02
500
1000
Step
Step
[x]
Step
17
0.001
0.3
4500
2100
Sharpness
4000
2050
Distance to 01
Sharpmess
0.2
Loss
3500
2000
0.4
3000
1950
2500
0.2
2000
1900
500
1000
500
1000
0.00
0.01
0.02
500
1000
Step
Step
[x]
Step
Transformer+ on SST2
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
n
0.002
0.88
1300
1020
Distance to 6
0.06
Sharpness
Sharpness
Loss
0.86
1200
0.04
1000
0.84
1100
0.02
1000
980
0.00
1000
1000
0.000
0.005
1000
Step
Step
[x]
Step
17
0.001
2100
0.0100
0.68
2075
2020
Distance to 0
Sharpness
Sharpness
0.0075
Loss
0.67
2050
0.66
2025
2000
0.0050
0.0025
0.65
1975
1000
1000
0.002
0.004
1000
Step
Step
[x]
Step
26

---

Published as a conference paper at ICLR 2023
Transformer-Logiitic on SST2
Gradient Descent
Predicted Dynamics
Constrained Trajectory
Gradient Flow
Predicted Fixed Point
17 0.02
0.02
0.7
1250
110-
Sharpness
1000
Sharpness
Distance to eT
1.0
Loss
100
0.5
500
0.4
2000
1000
2000
0.00
1000
2000
Step
Step
[x]
Step
17
0.01
0.65
210-
0.4
Distance to 6
Sharpness
600
Sharpness
0.3
Loss
0.60
200
400
190-
0.50
200
0.0
1000
2000
1000
2000
0.00
0.01
1000
2000
Step
Step
[x]
Step
0.005
0.62
700
420
0.15
0.60
Sharpness
600
Sharpness
Distance to 6
Loss
0.58
400
0.10
390
0.05
0.54
400
380
0.52
0.00
0
1000
2000
1000
2000
0.000
0.005
0.010
1000
2000
Step
Step
X
Step
77
0.002
0.44
1150
0.04
1010
0.43
0.42
Sharpness
1100
Sharpness
1005
Distance to
Loss
0.41
1000
0.02
1050
0.40
0.01
0.39
1000
990
0.00
0
1000
2000
1000
2000
0.000
0.002
0.004
1000
2000
Step
Step
[x]
Step
27

---

Published as a conference paper at ICLR 2023
H ADDITIONAL DISCUSSION
A Precise Criterion for Self-Stabilization Our theoretical results in Sections and give suf- ficient conditions for self-stabilization as evidenced by Theorem [ However, these assumptions may not be strictly necessary and we believe that self-stabilization may hold under significantly weaker assumptions. An important open question is understanding the precise conditions on the loss function and learning rate that enable self-stabilization
Furthermore, in Theorem in Section 3 we give quantitative bound on the time for which the EOS trajectory and the constrained trajectory be coupled. It is an interesting future direction to understand whether coupling can be strengthened to hold for longer time periods to allow for convergence KKT point of the constrained trajectory eq. (2)
Multiple Unstable Eigenvalues Our work focuses on explaining edge of stability in the pres- ence of a single unstable eigenvalue (Assumption 2). However, Cohen et al. (2021) observed that progressive sharpening appears to apply to all eigenvalues, even after the largest eigenvalue has be- come unstable. As result, all of the top eigenvalues will successively enter edge of stability (see Figure 55. In particular, Figure [ shows that the dynamics are fairly well behaved in the period when only single eigenvalue is unstable, yet appear to be significantly more chaotic when multiple eigenvalues are unstable.
0.26
100
Training Loss
Sharpness
80
0.22
0.20
60
40
500
1000
1500
2000
500
1000
1500
2000
Training Steps
Training Steps
Figure 5: Edge of stability with multiple unstable eigenvalues. Each vertical line is the time which the corresponding eigenvalue of the same color becomes unstable.
One technical challenge with dealing with multiple cigenvalues is that, when the top eigenvalue is not unique, the sharpness is longer differentiable and it is unclear how to generalize our analysis. However, one might expect that gradient descent can still be coupled to projected gra- dient descent under the non-differentiable constraint S(O))) S 2/n- When there are unstable cigenvalues, with corresponding eigenvectors u.....,,4 u the constrained update is roughly equiv- alent to projecting out the subspace spantV S Le ut) i.j E [k]] from the gradient update L,, Demonstrating self-stabilizztiio thus requíres analyzing the dynamics in the subspace
span {ut [x]} U Lu,,, w/) i, [K]}). We leave investigating the dynamics of mul-
tiple unstable cigenvalues for future work.
Connection to Sharpness Aware Minimization (SAM) Foret et al. 2021 introduced the sharpness-aware minimization (SAM) algorithm, which aims to control sharpness by solving the optimization problem ming maxis <e L).+ This is roughly equivalent to minimizing S(0) over all global minimizers, and thus SAM tries explicitly minimize the sharpness. Our analysis shows that gradient descent implicitly minimizes the sharpness, and for a fixed v) looks to minimize L(0) subject to S(O) = 2/1)
Connections to Warmup. Gilmer et al. 2021 demonstrated that learning rate warmup, which consists of gradually increasing the learning rate, empirically leads to being able to train with larger learning rate. The self-stabblizztion property of gradient descent provides a plausible explanation for this phenomenon. If too large of an initial learning rate V)o is chosen (so that S((b) much greater
28

---

Published as a conference paper at ICLR 2023
than 2/n)), then the iterates may diverge before self stabilization decrease the sharpness 2/nu- On the other hand, if the learning rate is chosen that S(16) is only slightly greater than 2/1o.. self- stabilization will decrease the sharpness to 2/To- Repeatedly increasing the learning rate slightly could then lead to small decreases in sharpness without the iterates diverging, thus allowing training proceed with large learning rate.
Connection to Weight Decay and Sharpness Reduction. Lyu et al. (0022) proved that when the
loss function is scale-invariant, gradient descent with weight decay and sufficiently small learning rate converges leads to reduction of the normalized sharpness S(0 ((0/1101). In fact, the mechanism behind the sharpness reduction exactly the self-stabilization force described in this paper restricted to the setting in (yu et al. 2022). We present here heuristic derivation of this equivalence.
First, we show that any scale invariant satisfies our assumptions when trained with weight decay. Lemma 6. Let f be scale invariant loss function, i.e. f(0 f(()). Let L(0) f(0) + 2 [101 Then for any local minimizer of [(0) such that S(O) = 2/1
VI(0) 1 u(0)
P4 O(np))
a<0) > 0
a(0) [VL(O)|lVV(el] (1)
|VS())]] = (ps))
Our primary result is that gradient descent solves the constrained problem ming L(O) such that (() 2/n. Let (y) denote the largest eigenvalue of f2) To prove equivalence to the sharpness reduction, we will need the following lemma from (Lyu et al. 2022) which follows from scale invariance of the f:
() = |0i| (0(011))
Let = ||0]
9 Then we have the following equality between minimization problems:
min L(0)
such that
S(0) $ 2/m
min f(0) + A 1101 such that Sy 2/n
min f(0) |1011
1 Sy(0) -
2 1T
such that |1lal
"
min (00 +
2 n (
where the last line follows from the scale-invariance of the loss function. In particular if yyλ is sufficiently small and the dynamics are initialized near global minimizer of the loss, this will converge to the solution of the constrained problem:
min SQ) such that f(0) = 0.
|0
H.1 SCALE INVARIANT LEMMAS
Let 0 denote arbitrary parameter and let Ø 0/||0||- Throughout this section, let f be scale invariant function with non-vanishing Hessian.
Lemma 7.
VS(0) V((O) |@0|
(2 P V SOPP
P
V -
V/0)
10012
29

---

Published as a conference paper at ICLR 2023
Proof. We start with the equality:
-(0)=(00)
P
Differentiating with respect to and using that V.0
|
gives,
(()
|
Differentiating this again gives:
F(0) Pe P(
- P/ V/ - -(P))
|001
A few corollaries immediately follow:
Corollary 2. For any critical point 0 of J.
Pd 42 P(0)))
v2 f(0)
1011
Corollary 3. For any critical point 0 of f.
u()) 1 VL(0)
Proof. Note that from Corollary 2 the top eigenvector of V is perpendicular to In addition.
V2 = f() + AI
so this is also the top eigenvector of 4 T²)) i.e. u0). Finally,
VL(0) ( X0 X0
which is parallel to and concludes the proof.
Lemma 8.
VS(0)= = P VS(0) - (S(0) X)0
||@|$
Proof. Let (0) denote the largest eigenvalue of V ((). Then by scale invariance, (2f(0) (²)(0)/ This implies that
Sy( Sy(
|1011
Differentiating this gives:
SSy P (0) Sy)0
||0|
Finally, we have from (BL( 4 f(0) + A7 that S(0) Sy( A so
P V - (S( - X)Ø
VS(0)
|
30

---

Published as a conference paper at ICLR 2023
Lemma 9. Let
P3 = sup 5(0)54/1
[ op
and
P4 = sup
8(0)54/1
Then
P3 = 8/7 and = O
In particular, P4 O(ng))
Proof. Note that S(Ø) < 4/7 implies that
4/n Sy( |101 2
+ x = ||0|| 1n/21/2 1/2
Therefore,
V2L(O) 0(0-3/2)
L()
1011
and
VI = V2L(0) = 0
1101
by compactness of Sd-1, In addition, note that L 2)]l |V5(0)) so P3 -3/2)) Lemma 10. At any second order stationary point 0 of
8(8)
O 3/2)
a(0) ((S(0) X)
> 0.
|1011
Proof.
(S(Ø) x)Ø
X(S(0) A) |1011
a()) VL() VS( -
|101|
Lemma 11. Al any second order stationary point 0 of J.
a0))
x(S(0) A)
|1012
Proof.
a
A(S(0) A)
1
|1011
Iella
p1 P (S(0) X)2
P V
+
-(Ø)-A
(1))
where the last step follows from compactness of 5d-1 and the fact that 2 f is non-vanishing.
Proof of Lemma The lemma is simply restatement of Corollary
Lemma 10 Lemma and
Lemma 11
31

---

Published as a conference paper at ICLR 2023
1 PROOFS
I.1 PROPERTIES OF THE CONSTRAINED TRAJECTORY
We next prove several nice properties of the constrained trajectory. First, require the following auxiliary lemma, which shows that several quantities are Lipschitz in a neighborhood around the constrained trajectory:
Definition 8 (Lipschitz Sets). St == B(0 imps) where c is the constant in Assumption and (xx r) denotes the ball of radius centered at z.
Lemma 12 (Lipschitz Properties).
1. 0 + VL(0) is -s in each set St.
2. 0 + (0) is p3-Lipschitz with respect to I1-112-
3. 0 + (( is ps-Lipschitz.
4. 0 + (0) is (npss in each set St.
5. 0 + VS()) is Lippsciips in each set S
Proof. The Lipschitzness of L() follows immediately from the bound ||) S P3-
Weil's inequality then immediately implies the desired bound on the Lipschitz constant of the eigen- values of 2 (0). Therefore for any 4, we have for all 0 E St:
2 2 2n
2-c
X1 V L(O)) >2(V² Z A (V >2 \2(
2p3 4mps
Next, from the derivative of eigenvector formula:
[|vn(0)]l2 = -
P3
< (V())((0)) (V
2Mps $
c
= O(pss)
which implies the bound on the Lipschitz constant of restricted to S.. Finally, because VS((0) L(O))) (()). - { +1] - Op1 + mp²) S 0(1p5)
where the second to last inequality follows from the bound on || Vu(O)]l, restricted to S and the last inequality follows from Assumption
Lemma 13 (First-order approximation the constrained trajectory update (0])). For all t S I,
0 = 0; - p) u.. vS. VL, + O n//LLel) and S 2/n
Proof. We will prove by induction that S; = 2/1 for all t. The base case follows from the definitions of co. f%. Next, assume S(e!) = 0 for some t 2 0. Let 0' = 8! V, Then because 0. E M we have |0].| - 0²|| < of 0/| = ni L. Then because 0f (O') the KKT conditions for this minimization problem imply that there exist x, y with y 2 0 such that
0/41 0: V- x vV. VL v())
V
0=
0f V. - x S1S ue+1 + Vu [
1/ V
0: V 21S U111 0 |V ||] SV
00 VL 2(S,u, + 0 - VS, O(7 (((
00 V. =S yVS. + +0 w/1 |
L..)).
32

---

Published as a conference paper at ICLR 2023
Next, note that we can decompose VS. uvS. u. + VS!
6 = 0 Vv
- [x:; + (VS uz)]u4 yvS. 0([[2|mps+ + 2 VL,I). VL.l).))...
Let 8t =
We can now perform the change of variables
| ||
x 1/' V.
Vss
(xS. + vVS. wd) (VH,)) (a.y)
St
V
to get
0/11 = 07 - V, d'u, - y/ + +0'2p / V1 [/1()=/) -1-1
Note that
O p3l [L]| le)) V y
(12)
2
for sufficiently small because
[0 r+1 00
we have
2 +3² $ |0l 1+ 0¹ / VLel
so z, y VL.II. Therefore,
o +1 0: 17 L1 I've 8t + 0 p3l Psl|v
= 0g - 1/ T, I'ue - s1 + 0 !n V
Then Taylor expanding VL around 0: gives
L Ue41 = VL Ue + (VL++1 - VL) Ur + V . (u++1 we)
-/ v (. - z'ue- 3 + + O(=² n| /l + +(E² -IV
-I'S, ( |VLll)
So x = 0 // V II). We can also Taylor expand S.+1 around 0f and use that S, = 2/n to get
S.+1 = 2/1 + V. -1 - ue yf + 0 pal pallVVL]]*] + +0(( -onllV
= 2/n + na, - y/ | +| + 002 p371| V1,
Now note that for sufficiently small we have
0 panll VT+ll) < 0 .na) S n0
Therefore if y' = 0, we would have S.+1 > 2/n which contradicts 0._1 M. Therefore y' > 0
and therefore y > 0, which by complementary slackness implies S 2/n. This then implies that
VL VS! v V,,, || +0 pall VL. =0 - --
+(e n VL.I).
VS!1
Putting all together gives
0, = 0: nP VL, + + n/ V
= 0] nP P.... vS, VL, + 0(2 n/ v
where the last line follows from We VL, 0.
Lemma 14 (Descent Lemma for 01). For all 1 < T,
1(D < L(0)) - 22 p VS.
VS. Vi.ll²).
33

---

Published as a conference paper at ICLR 2023
Proof. Taylor expanding L(f)+1) around L(6,) and using Lemma 13 gives
- ~ (0 01) + 00000l00 - o2ll²)
1(0 = =((,)) + Vl+ (0(+1 0f) + 2 10/ -
22²2 LX))
Puu, vs. V
= =(0)) n v.. VL
2
-(O) (1(2 c)
I I. VS. V1 1 VL.
VL.I"). 1*)
2
Next, note that because = (1) we have /L = O U...v. vs. VL, [. Therefore for c suffi- ciently small,
0 ps''ssllVLl11 = O ///1/ - 1/22 - c) | VLII?
Therefore,
1/2 L(0 S L(O))
-c)
|| VS.
VL L(O)) M6n
Vs. Pul ue. Vs, VLell")
4
which completes the proof.
Corollary 4. Let L = ming I(0). Then there exists 1 S g such that
P v. VL, 00 [(O)) I.*
11
Proof. Inductively applying Lemma 14 we have that there exists an absolute constant e such that
L S L(0 < L(O))
Σ
P vs. vL,b
en
tKg
which implies that
Ener
PPl.
Vs.
VL 1
L(6b) L'
min Vs. VLol/²
<0
g
n.TT
1.2 PROOF OF THEBRRE 1
We first require the following three lemmas, whose proofs are deferred to Appendix I.3 Lemma 15 (2-Step Lemma). Let
T=: U142 (ste,+
Assume that ||ve|| S 118 Then
[rlll s o
26 max || 3
Lemma 16. Assume that there exists constants C1, og such that for all t S T, 1 ozd. 12 c. Then, for all t S T, we have
||,
J < O(e))
Lemma 17. Fort < 9, J. 5
With these lemmas in hand, prove Theorem
1
34

---

Published as a conference paper at ICLR 2023
Proof of Theorem First, by Lemma 17 we have || S 000).
Next, by Lemma 16 we have
02 - 0 = = 02 + Ocd))
Next, we Taylor expand to calculate S(
S(O.) S(0)) VS, 12 0
2/n + VS 1 VS: 4444 Dt + O
-2/n+ + Vs i22 + VS: 4444 0q + (p3cc + n0302)
2/n + y: + (VS, 4n) + 1
Finally, we Taylor expand the loss:
L(O,) L(O)) + VL, 1 +
~² 2
Le, + Opsslvu,11))
L(O)) + _ +
1
T
L 10 pollo,)))))
2º²
2
L(O)) + 1 +
1
T
L4 + +nn 18%
2 7
L(()) + 3 + - E8).
where the last line follows from Assumption 5
1.3
PROOF OF AUXILIARY LEMMAS
Proof of Lemma
15
Taylor expanding the update for 6,41 about o/, we get
0,1 -.-nnVL(9))
= 0: V - 1 I1 vz) + 0 (npaallall))
Additionally, recall that the update for 0:41 is
0+ = 03 - VT, VI, + (00.00/2//V))))) v/[l/.
4s
35

---

Published as a conference paper at ICLR 2023
Subtracting the previous 2 equations and expanding we obtain
V L(vr, 0i) via the non-worst-case bounds,
24+1
(I - Post VL,)
2 ni n VS, - no V (o 2 n Le(vi,et
+ +0(pplllee + + 1 n/1V1.1))
VL. V
1
(1 LL)
2 N
maa V (o(
|12
0 (ppaelluell mpaalllu 10: 22 ./.n/VVA VLiI)
C -2
(1 1- + 5
/n Uedt mF V Le(u,
+0 2. |1d01 +e²
||//| + 2.
82
& - xt
(I Li)ur + S
n 7 VS. ui nr V Le(ue, v
+0 28 max (1, |lceel| 3)
)
We would first like to compute the magnitude of ut+1-
|| = O |Ioell + n + vl[V + 28 max (1, |rtlll) )
Observe that by definition of and 8, and since ||e/|| S 26
Ompps|Iv:l 12) (llvell - 1133) < ( . -1 s O(1|=
0 max (1, ||
8 ) - 0 11uull 2 - 0 + ||val))-
Hence
|| 011v:11 / +0 0(11=21 || + c6).
Note that we can bound
|| |4+11 - uell ||041a|| = 0(7 Polpall VL (qwee) + e5))
= =0(2² (1 + +o))
< 0 max([|ur]] 8)).
Therefore, the one-step update in the u. direction is:
It-1 Ut+1 . Ue+1
utt1 ut + O maxx(()=:11.6))
= - -4 2 non VS. 2 - - Vs of + +0(2 maax((lleell.lb + c28 max (1, ||vee] (1xal
))
= =-x;(1 + myc))
2
1
V un + +0
28 max
1, ||
8
1
||ve|| 3
= =-x,(1 + vigr)
SS. in + +
(28 max
1,
マ
1
=1, + 19)))
nat
SS + O(E)),
36

---

Published as a conference paper at ICLR 2023
where we have defined the error term Eg as
E, = max (1,
(1,1%-1)
1,
||1x
3
8
The update in the vt direction
4 x²²
21411
Pal.
[[-1nV22 L4) + t/S
mn22 V.S
us P.l, u4s1 ut non Paep I'In
#1-1
28. mmx (1, |Ivell
(1,l1ral) 3
+0
c x²
1
pl
(1-1
- 1/ Lc) u. + 7
1z P.
ui+1 ue 2 n
us ue ue mare wen v'Lo(ue, VS, g
u1-1
(1, [1el) 3
+ 28 max
First, observe that
pul uuis Ue|
= |lue 1 well
S llue
< <0([ 2
ue+ill)
Therefore we can control the first of the error terms as
r rrP me+1 us + 22 na VS. ue Pul u
$ $0 - -4+1 (11vel + mpalleel
S 50( |t - u+iil -1xel))
$ (e2ll)))
As for the second error term, can decompose
I deen Le of') b|| s n/|tz|l |l (1 v Lo(u, ot || | - pr...
||||||| Lo
we have || Vi Lo of I| s O(epallurl).. Additionally, Iledl Poten 5
By Assumption
0lI Me+1 Therefore
I Puen, LOLe(ue utt < O nl|uell + n/lbx I|lue11 well pallvell)
o enpslleel| nps|lea 12:2)
2 ||02||
< 0
+ e
8
|/well 3) )²)
0 28 max
I,
where we used /pss|vu/|| = O(1). Altogether, we have
2 x
288 max (I, (1.llall) 3
utt = PL
(I (1- (2) S
+ +0
2
{ x2
=P -
+ V.
+ O(Fs))
2
We next compute the two-step update for xe:
Je-2 = -I m1)++) - 1 VS u141 O(F44))
=xx(1 + (1 + 78++1) + (mmna VS, + +
2441 SS41 u.+1) + O((1 + 7/3/|/u,1)EE + E.t1).
37

---

Published as a conference paper at ICLR 2023
We previously obtained nps||e O(1). Furthermore,
Ex+1 = 20 max (1, |1u
11, |10e|1 |l
5 + ))
3
0 26 max
|112||
)
0 28 max
1,
8
O(Fa)-
Hence
Tt-2 - xr + (1 + 1)(+1) + ? 2 [ dr + x S. ur - 2bit SS41 ue+1) + OE,))
The first of these error terms can be bounded as
ut I (1 so s0 2 l1x1l)
2 " n
ya Vs:
82
As for the second term, we can bound
SS++ u++1
VS. uel
1 U1411 (VS++1 - VSA (u++1 ue)
I ||V5++1 - VS + 0(ps) ||1411 uell
I (1' (+)
1 (c²ppa)
Additionally, we have
Fe41 = x (mpplllllll Ez).
Therefore
7/ VS M441 - S ue| S q VS 1+ 441 VS. . uel + n -.,)/VS.. u411
$ 0 np3 1 112 Pe2 + npslluel| l| (mpalleel? + Ee))
|1//| 3
$ 0
2 2/1lell + t²².
+ Eg
82
O(E))
Altogether, the two-step update for Ie is
Te+2 (1 + (1 +
Additionally, the two-step update for vz is
Pul,, ett =PLL (I 2 Le+1) Puc. te41 + V
Ei+ x7+1 O(E++)
{ x?
pat. 2 (I - 1/ Le+1) p.t. uell (1 - v/ Le P.l or + mpt (1 - L(41) uii VS.
+ +mP 1.
VS
2t41
+0(E)).
uets
2
Define U241 = (te), 0,42 = step, (Vt), and F; = U. u4 D, = VS! . for (t+1,t+2). By the definition of step, one sees that
| r O(E.)
38

---

Published as a conference paper at ICLR 2023
and
1 J++1 M141 S 2
St vel + O(B)) S Orpssllvel1 + Ex)
2
The update for a after applying step is
Tx+2 -I++1 + 1/1r--1) =z,(1 + my:: + 111t+11)
Therefore
2-42 - T442 - M + +OEi)
s 0npsssleell) - - +(EE))
(
Additionally, the update for vー is
-
2242 =P.L.
Puuta (I 1/ L++1) P.l. (I (1, + n (I (I- 11/
VS,
m4t1
I?
+ tyPPl 4442 VS
9+1
2
Therefore
||0t2
v/42||
s (nl |1 xz+1) Et)
O(133 10 J F2+1l + F.))
17 13 [//l|l Es)
0 2. |10 1| Ee
O(E))
Altogether, get that
0 28 max
|eell 3)
[rd/ S O(E.)
1,
8
as desired.
Proof of Lemma 16
Define
0 rt-1
is even is odd
Wt
and define the auxiliary trajectory T by 0. and Ge+1 step():) I first claim that 6 0 for all even < T, which we will prove by induction on The base case given by assumption so assume the result for some even f. 2 0. Then,
Ut+2 = s (step; +
step:+ (stop, + rt
stept+1 (i++1) W(+1
0
which completes the induction.
Next, we will prove by induction that for < T,
|
3 ||| |r: - del S ((E)) S oz8.
39

---

Published as a conference paper at ICLR 2023
By definition, 0 Uo uu. the claim is clearly true for t 0. Next, assume the claim holds for t. If f is even then |wyll = 0: otherwise ||uall I 2026, and thus
|wee| - 0 28 max (1, c2) ()))).
First observe that
|6/41 - -6/1111 $ |II - my - ||| 7P: ب7 - 2??
+ ||wlll
2
S (1 + 1// Amein 2 Lx) || - + + +O(e)
1 - +
- (1 + 1/ Amonn (2 || - +
+ +((20)
3/
Next, note that
Eo+1
(1 + (1+mi)) + 0((2))
Eo+1
m/)) + 0(e28)
(1 + +)t)); + 0(e2)) + O(e) |0 ot
+ (1+10)) + (c28)
:1 + 0 2 [1 -6,10 ll:
Therefore
Fet1
it+1
Ft - 2t
+ + + 8 1110t -6-1 |).
xc+1
dt
Let d, max (107 - 2 ||. 5)
8, 8.-
1) Then
2,
|ott1
S (1 + 7/ Amin (
+ +
8++1 ic+1
8
$ (1+00 +1+0
Le+1
Therefore
d++1 $ (1 +1 Ammin [L + +(())
+ so for t s I we have de+1 5 ((E)). Therefore
$ (1 + (()))((()(((+))) 0(2)),
| - الاي| |5441 ir41l S O(ed) c2d
so the induction is proven. Altogether, we get l|0₂ bell S (cd)) for all such t, as desired.
Proof of Lemma 17 Recall that
03
2
ait1 (1 + )z] and vit1 1) Σ
3 s-
8
2
0=
Since /
1 maxe Amin (V*L,)] we have that 3r-ss Opp)), and thus
S (pp)) (
Therefore
|41l = + n < (1 + (2.)
Since t < <O(E=)), | grows by at most a constant factor, and thus |
< 06)). Finally, recall that
x+1
82
x x2
A
s 0
k=
By the triangle inequality,
| || S {O(ntpso)) 50())
Therefore || < (88).
40


Reviews:
[{"review_id": "pGCBtErD1S", "reviewer": "Reviewer_UEL4", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The paper studies the Edge-of-Stability effect from a rather general perspective. It identifies progressive sharpening and the oscillatory behavior of optimization trajectories at EoS as resulting simply from the alignment of the loss and sharpness gradients. This alignment is determined by a particular cubic term of the Taylor expansion of the loss. The resulting dynamics is related to the well-known Lotka-Volterra predator-prey model. The paper first discusses a simplified scenario where this effect can be easily demonstrated and understood. After that a theorem is stated addressing a general setting. Finally, a series of experiments with real world data sets and neural networks is presented, showing agreement with theory. ", "summary_of_the_review": "A good paper, accept.", "strength_and_weaknesses": "I quite like this paper. There are several other current studies of EoS (e.g., https://openreview.net/forum?id=p7EagBsMAEO, https://openreview.net/forum?id=R2M14I9LEwW), but the present paper paints an especially general and appealing picture of this effect. The exposition is very clear. The theory is confirmed by multiple experiments with real data.   \n\nOne natural question that does not seem to have been answered in the paper is why can we generally expect Assumption 1, the key alignment assumption, to hold? It seems that the authors leave this question for future work?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "f1q60hYqJY", "reviewer": "Reviewer_dCeH", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "The authors study a phenomenon they call \"self-stabilization\", which is proposed to explain common empirical observations such as sharpening and convergence at the edge of stability. \n\nFormally, the authors consider the iterates of (discrete time) gradient descent and assume that the product between the gradient of the objective and the gradient of sharpness is negative, where sharpness is defined as the largest eigenvalue of the Hessian. For simplicity, the authors also assume that the Hessian has at most one negative eigenvalue. Under these assumptions, the authors present a simplified local analysis of the self-stabilization phenomenon. Then, they extend this analysis to general loss functions, but under much more restrictive assumptions, such as Assumption 4, on the gradients of sharpness and loss. The main result is given in Theorem 1, which shows how, up to high-order terms, the loss and sharpness coevolve.\n\nFinally, the authors present numerical evaluations serving as a justification for their assumptions and validating some of their claims.", "summary_of_the_review": "I enjoyed reading the paper and I believe it will be of interest to the community. While the assumptions behind the theory are far from perfect, the derivation is still of interest. The empirical evaluations are meaningful and support the theory.", "strength_and_weaknesses": "## Strengths  \n1. The work sheds some light on the recently observed phenomena in deep learning, which may help us understand how deep networks learn the data.\n2. The proposed explanation is quite intuitive and is supported by empirical evaluations.\n3. The paper has a good flow and starts with simple ideas that finally lead to the main result.\n\n## Weaknesses  \n1. From the mathematical point of view, this paper is not strong. Assumption 4 intuitively implies self-stabilization because it guarantees a significant negative correlation between the gradients of loss and sharpness. While the conclusions still require quite a bit of technical work to do, which the authors did well, the assumptions are very stringent. I find this to be the main weakness of this paper.\n2. The studied problem rarely comes with deterministic gradients, and the impact of noise in stochastic training is not considered.\n3. The implications of the results do not seem to change our understanding of how neural networks need to be trained. The notion of self-stabilization does not seem to suggest much more than the edge of stability already tells us.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "A4ZBvwsX7A", "reviewer": "Reviewer_qpmy", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 studies the edge of stability phenomenon. They propose to incorporate the cubic term in the convergence analysis and show that there is a self-stabilization property caused by this cubic term for general nonconvex optimization problems. Moreover, they show that GD is inherently regularizing the sharpness, i.e., it is close to a constraint optimization problem that explicitly adds the sharpness as the constraint. Last, they provide very promising simulations showing that the predicted trajectory mimics the actual GD trajectory on many realistic tasks.", "summary_of_the_review": "Please see the above comments.", "strength_and_weaknesses": "## Strength\n- This work provides a new direction to understand the edge of stability problem. Though the author argues that their method can only explain the EoS phase. By the argument in Stage 1, I feel like it can also be used to explain the progressive sharpening phase. It would be very great if the author can help me to understand what's the gap to preclude them to claim that this method can be used to study the first phase.\n- I really like the motivation section (Section 4), which is very clear and convincing to me.\n\n## Weaknesses\n- At the end of the first paragraph in Section 1.2, they say that \"Unlike progressive sharpening ..., self stabilization is a general property of gradient descent.\" However, in Section 5, the main theorem relies on many assumptions, which are only verified on neural network training. Hence, I am a little bit concerned about its generality. Apart from the numerical verification, I can only find the toy example provided in Appendix B, which is just a cubic function. It would be great if the author can provide either more general examples, or verify that those assumptions hold for a general function class. \n- As argued in Section 4.3, they show that the distance between $\\theta_t$ and $\\theta_t^{\\dagger}$ is bounded by $\\delta$. However, it is unclear how large this $\\delta$ can be. Note that $\\delta=\\sqrt{2\\alpha/\\beta}$. In the appendix, it can be seen that $\\alpha$ is not small. Hence, it is not clear to me why $\\theta_t$ and $\\theta_t^{\\dagger}$ are close.", "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