A second order regression model shows edge of stability behavior
Authors: Atish Agarwala
Jeffrey Pennington
Fabian Pedregosa
Venue: Submitted to ICLR 2023 (2023)
Keywords: deep learning theory
non-linear dynamics
optimization

Abstract: Recent studies of learning algorithms have shown that there is a regime with an initial increase in the largest eigenvalue of the loss Hessian (progressive sharpening), followed by a stabilization of the eigenvalue near the maximum value which allows convergence (edge of stability). We consider a class of predictive models that are quadratic in the parameters, which we call second-order regression models. This is in contrast with the neural tangent kernel regime, where the predictive function is linear in the parameters. For quadratic objectives in two dimensions, we prove that this second order regression model exhibits both progressive sharpening and edge of stability behavior. We then show that in higher dimensions, the model shows this behavior generically without the structure of a neural network, due to a non-linearity induced in the learning dynamics. Finally, we show that edge of stability behavior in neural networks is correlated with the behavior in quadratic regression models.

Under review as a conference paper at ICLR 2023
SECOND-ORDER REGRESSION MODELS EXHIBIT PRO- GRESSIVE SHARPENING TO THE EDGE OF STABILITY
Anonymous authors Paper under double-blind review
ABSTRACT
Recent studies of gradient descent with large step sizes have shown that there is often a regime with an initial increase in the largest eigenvalue of the loss Hessian (progressive sharpening), followed by stabilization of the eigenvalue near the maximum value which allows convergence (edge of stability). These phenomena are intrinsically non-linear and do not happen for models in the constant Neural Tangent Kernel (NTK) regime, for which the predictive function is approximately linear in the parameters. As such, dictive models, namely those that are quadratic in the parameters, which we call second-order regression models. For quadratic objectives in two dimensions, we prove that this second-order regression model exhibits progressive sharpening of the NTK eigenvalue towards value that differs slightly from the edge of stability, which we explicitly compute. In higher dimensions, the model generically shows similar behavior, even without the specific structure of neural network, suggest- ing that progressive sharpening and edge-of-stabillic behavior aren't 't unique fea- tures of neural networks, and could be a more general property of discrete learning algorithms in high-dimensional non-linear models.
consider the next simplest class of pre-
INTRODUCTION
A recent trend in the theoretical understanding of deep learning has focused on the linearized regime, where the Neural Tangent Kernel (NTK) controls the learning dynamics (Jacot et al., 2018; Lee al., 2019). The NTK describes learning dynamics of all networks over short enough time horizons, and can describe the dynamics of wide networks over large time horizons. In the NTK regime, there is function-space ODE which allows for explicit characterization of the network outputs (Jacot et al., 2018; Lee al., 2019; Yang. 2021). This approach has been used across the board to gain insights into wide neural networks, but it suffers a major limitation: the model is linear in the parameters, so describes regime with relatively trivial dynamics that cannot capture feature learning and cannot accurately represent the types of complex training phenomena often observed in practice.
While other width scaling regimes can preserve some non- con-lincarity and allow for certain types of feature learning (Bordelon & Pehlevan, 2022; Yang et al., 2022), such approaches tend focus on the small learning-rate or continuous-time dynamics. In contrast, recent empirical work has highlighted a number of important phenomena arising from the non-linear discrete dynamics in training practical networks with large learning rates (Neyshabur et al., 2017; Gilmer et al., 2022; Ghorbani al., 2019; Foret al., 2022). particular, many experiments have shown the tendency for networks to display progressive sharpening of the curvature towards the edge of stability, in which the maximum eigenvalue of the loss Hessian increases over the course of training until stabilizes at value equal to roughly two divided by the learning rate, corresponding to the largest eigenvalue for which gradient descent would converge in quadratic potential (Wu al., 2018; Giladi al., 2020; Cohen et al., 2022b:a).
In order to build a better understanding of this behavior, we introduce a class of models which display all the relevant phenomenology, yet are simple enough to admit numerical and analytic understand- ing. In particular, we propose a simple quadratic regression model and corresponding quartic loss function which fulfills both these goals. We prove that under the right conditions, this simple model shows both progressive sharpening and edge-of-stability behavior. We then empirically analyze a more general model which shows these behaviors generically in the large datapoint, large model

---

Under review as a conference paper at ICLR 2023
limit. Finally, conduct numerical analysis the properties of real neural network and use tools from our theoretical analysis to show that edge-of-stability behavior "in the wild" shows some of the same patterns the theoretical models.
2 BASIC QUARTIC LOSS FUNCTION
2.1 MODEL DEFINITION
We consider the optimization of the quadratic loss function L(O) = =2²/2, where z quadratic function the P x 1-dimensional parameter vector 0 and is P symmetrri matrix:
2 00Q00-
Q0
(I)
This can be interpreted either as a model in which the predictive function quadratic in the input parameters, or as a second-order approximation to a more complicated non-linear function such as a deep network. In this objective, the gradient flow (GF) dynamics with scaling factor 17 is given by
Oz = 2
0 = Ve = 1)2 00
Te Q0 - E Q0.
(2)
It is useful to re-write the dynamics in terms of and the 1 x P-dimensional Jacobian J = 0z/00:
= [ )2 = =
(3)
The eurvature a scalar, described by the neural tangent kernel (NTK) JJ In these coordinates, we have E' JQ +Q+J* J 2z, where denotes the Moore-Penrose pseudoinverse.
The GF equations can be simplified by two transformations. First, we transform to ī = a/2 and J ry J. Next, we rotate 0 so that Q is diagonal. This is always possible since Q is symmetric. Since NTK given by JJJ this rotation preserves the dynamics of the curvature. Let w1 2 ... 2 wp be the eigenvalues of Q, and v, be the associated eigenvectors (in case of degeneracy, one can pick any basis). We define Jwx; = Jv:, the projection of J onto the ith eigenvector. Then the gradient
flow equations can be written as:
P dž = 3 --E))
dJ(w)) -w:,,; dt
(4)
The first equation implies that does not change sign under GF dynamics. Modes with positive w,ē decrease the curvature, and those with negative w, increase the curvature.
In order to study edge-of-stabilit behavior, need initializations which allow the curvature (JJ in this case) to increase over time a phenomenon known as progressive sharpening. Progressive sharpening has been shown be ubiquitous in machine learning models (Cohen et al., 2022a), so any useful phenomenoloiical model should show it as well. One such initialization for this quadratic regression model is W1 -w, W2 = w. J(w1) = M(wz). This initialization (and others) show progressive sharpening at all times.
2.2 GRADIENT DESCENT
We are interested in understanding the edge-of-stability (EOS) behavior in this model: gradient descent (GD) trajectories where the maximum eigenvalue of the NTK, JJ remains close the critical value 2/1). We define edge of stability with respect to the maximum NTK eigenvalue instead of the maximum loss Hessian eigenvalue from Cohen al. (2022a). We will prove this form of EOS in our simpler models, and find that holds empirically in more complex models. See Appendix A. I for further discussion.
When Q has both positive and negative eigenvalues, the loss landscape is the square of a hyperbolic parabaloid (Figure 1, left). As suggested by the gradient flow analysis, this causes some trajectories to increase their curvature before convergence. This causes the final curvature to depend on both the initialization and learning rate. One of the challenges in analyzing the gradient descent (GD) dynamics is that they rapidly and heavily oscillate around minima for large learning rates. One
2

---

Under review as a conference paper at ICLR 2023
-30
10a
-3.5
101
101
10
LOSS
100
9-100 1.30
8 0
4.5
4-
1,30
10
10
-5.0
1,50
4- 0.70
10
5.5
10
-8
4
10
-6.0
-2
10
01
0
Figure 1: Quartic loss landscape L(.) function of the parameters 0, where D 2, and Q has cigenvalues 1 and -0.1. The GD trajectories (initialized at (1.5, -4.32), marked with an x) converge minima with larger curvature than initialization and therefore show progressive sharpening (left). The two-step dynamics, in which we consider only even iteration numbers, exhibit fewer oscillations near the edge of stability (right).
way to mitigate this issue to consider only every other step (Figure 1, right). We will use this observation to analyze the gradient descent (GD) dynamics directly to find configurations where these trajectories show edge-of-stability behavior.
In the eigenbasis coordinates, the gradient descent equations are
Σ J(wi)? + 1
2411 2, -
(w,))
(5)
J(m))++ - J(w)){ -2 Tw:)) for all 1 sisp.
(6)
We'll find it convenient in the following write the dynamies in terms of weighted averages of [wx)2 instead of the modes [uz)):
T(a) = Σ (wr))
The dynamical equations become:
2+1 21 = (T) + (00(())))88)77) (1)))))
(8)
T (k:) = (2T, + +1 TT,
If Q is invertible, then we have E = T-(-1) 22,. Note that by definition Ti(0) = nJ.J. is the (rescaled) NTK. edge-of-stability behavior corresponds to dynamics which keep (1) near the value 2 as ,, goes to O.
2.2.1 REDUCTION TO CATAPULT DYNAMICS
If the eigenvalues of Q are {-w, W. and E = 0, the model becomes equivalent to single hidden layer linear network with one training datapoint (Appendix A.2) also known as the catapult phase dynamics. This model doesn't exhibit sharpening or edge-of-stabiliit behavior (Lewkowycz et al., 2020). We will analyze this model in Σ T(0) variables warmup, with eye towards analyzing different parameter setting which does show sharpening and edge of stability.
We assume without loss of generality that the eigenvalues by rescaling ē. The loss function is then the square of a hyperbolic parabaloid. Since there are only 2 variables, we can rewrite the dynamics in terms of ī and the curvature (()) only (Appendix B. 1):
(-1, which can be accomplished
EE+1 is = - + (2 E)
(10)
(0) T.(0) -2, F.) T.))
(11)
3

---

Under review as a conference paper at ICLR 2023
For E = O, can see that sign sign(A((0)) sig - 4), as in Lewkowyez et al. (2020) so convergence requires strictly decreasing curvature. For E * 0, there is a region where the curvature increase (Appendix B.1). However, there is still no edge-of-stabilin of initializations which starts with Amax far from 2/1,, which ends up near 2/1j. In contrast, we will show that asymmetric eigenvalues can lead to EOS behavior.
behavior there is no set
2.2.2 EDGE OF STABILITY REGIME
In this section, we consider the case in which Q has two eigenvalues one of which is large and positive, and the other one small and negative. Without loss of generality, we assume that the largest eigenvalue of Q is 1. We denote the second eigenvalue by -c, for 0 <esl. 1. With this notation we can write the dynamical equations (Appendix B. as
E(s 2, = -TT,( + (111c)1.(0) + c(2 E))
(12)
(0) (T(0) = 22 - (1-)) + # [T. - c (c- (T,(0) E 222-233 22g
(13)
For small c, there trajectories where Amax is initially away from 2/n but converges towards it (Figure 2, left) - in other words, EOS behavior. We used a variety of step sizes but initialized at pairs initialized at pairs (1750, (0)) to show the universality of the (0) coordinates.
In order to quantitatively understand the progressive sharpening and edge of stability, it is useful to look at the two-step dynamics. One additional motivation for studying the two-step dynamics follows from the analysis of gradient descent on linear least squares (i.e., linear model) with large step size A. For every coordinate 0, the one-step and two-step dynamics are
0.+1 - 02 -X0, and 0.42 0, (1 - 2)
(GD in quadratic potential).
(14)
While the dynamics converge for À < 2, if > the one-step dynamics oscillate when approaching minimum, whereas the the two-step dynamics maintain the sign of ê and the trajectories exhibit no oscillations.
Likewise, plotting every other iterate in the two parameter model more clearly demonstrates the phenomenology. For small c, the dynamics shows the distinct phases described in (Li et al., 2022): an initial increase in T(0), a slow increase in ž, then decrease in T()), and finally a slow decrease of z while T(0) remains near 2 (Figure 2, middle).
Unfortunately, the two-step version of the dynamics defined by Equations 12 and 13 more com- plicated they are 3rd order in (0) and 9th order in Σ; see Appendix B.2 for a more detailed discussion. However we still analyze the dynamics as 2 goes to O. In order to understand the mechanisms of the EOS behavior, it is useful to understand the nullclines of the two step dynamics. The nullcline (z) of and (5) of T(0) are defined implicitly by
(ž++2 f=()) = 0, (T++2(0) (z(0) fr())
(15)
where ž,42 3, and (0) - (O) are the aforementioned high order polynomials in ā and T(0). Since these polynomials are cubic T(0), there are three possible solutions as ā goes to O. We are particularly interested in the solution that goes through ă 0, T(0) 2 that is, the critical point corresponding to EOS.
Calculations detailed in Appendix B.2 show that the distance between the two nullclines is linear in so they become close as c goes to O. (Figure 2, middle). In addition, the trajectories stay near fa which gives rise to EOS behavior. This suggests that the dynamics and trajectories appear to be approaching an attractor. We can find the structure of the attractor by changing variables to y, = (0) f=(āt) the distance from the nulleline. To lowest order in and y, the two-step dynamical equations become (Appendix B.3):
slow near the nullclines,
2 2 = 2y, + +
(16)
Ur42 - = 2 - 3e + 4e )y 2e +=00)))0(()')
(17)
We immediately see that ē changes slowly for small y - since we chose coordinates where 2,42 -, = O when y = 0. We can also see that yr+2 y/ is ((e) for 21 for small the y dynamics
4

---

Under review as a conference paper at ICLR 2023
is slow too. Moreover, see that the coefficient of the cτî term is negative the changes in 2 tend to drive y (and therefore T()) to decrease. The coefficient of the yi term is negative as well; the dynamics of y tends to be contractive. The key is that the contractive behavior takes y to O(c) fixed point at a rate proportional to 8², while the dynamics of ă are proportional to e. This suggests separation of timescales if 8 >> 6, where y first equilibrates to 0 (Figure 2. right). This intuition for the lowest order terms can be formalized, and gives us prediction of lim 1/2 e/22, confirmed numerically in the full model (Appendix B.5).
fixed value, and then 2 converges
Trajactorios
dynamics,
Curvature (Ama))
03
10P
size
101
step
0.1
109
1000
100
200
300
400
500
5.0
03
0.5
-01
0.00
0.55
0.15
0.50
0.25
Number of lterations
Figure 2: For small two-eigenvalue model shows EOS behavior for various step sizes (c 5 10 left). Trajectories are the same up to scaling because corresponding rescaled coordinates and T(0) are the same at initialization. Plotting every other iterate, we see that for variety of initializations (black s), trajectories in ž TT() space stay near the nulleline (=))) the curve where ă+2 & (middle). Changing variables to y T(0) fe(2) shows quick concentration to a curve of coar-constant, small, negative y (right).
We can prove the following theorem about the long-time dynamics of 2 and y when the higher order terms are included (Appendix B.4);
Theorem 2.1. There exists an eo > O such that for a quadratic regression model with E 0 and eigenvalues -e, S to there exists neighborhood U C IR² and interval [11,))) such that for inilial 0 E U and learning rate r) E [h, M2], the model displays edge-of stabbiity behavior:
2/n 6x S lim Amax 22/n
(18)
t-22
for 8x of (e)). This neighborhood corresponds to the inverse image of the ă space region [0, 2c) [0. for e-independent žn and ye-
Therefore, unlike the catapult phase model, the small provably has EOS behavior whose mecha- nism well-understood by the Σ y coordinate transformation.
3 QUADRATIC REGRESSION MODEL
3.1 GENERAL MODEL
While the model defined in Equation provable displays edge-of-stability behavior, required tuning of the eigenvalues of Q to demonstrate it. We can define more general model which exhibits edge-of-stability behavior with less tuning. We define the quadratic regression model as follows. Given P-dimensional parameter vector 0, the D-dimensional output vector z is given by
1 Q(0,0)
z= y G
(19)
Here y is a D-dimensional vector, G is D P-dimensional matrix, and Q is D P P- dimensional tensor symmetric in the last two indices that is, Q(.,.) takes two P-dimensional vectors as input, and outputs a D-dimensional vector verifying Q(e, 0) 0 T Q. If Q O,
the model corresponds to linearized learning (as in the NTK regime). When Q * 0, we obtain the first correction NTKK regime. We note that:
aza
8za
G
Quij
+ G Q(0,)
(20)
00;
00,00,
o=0
5

---

Under review as a conference paper at ICLR 2023
for the D X P dimensional Jacobian J. For D 1, recover the model of Equation 1. In the remainder of this section, we will study the limit as and P increase with fixed ratio D/P.
The quadratic regression model corresponds to model with constant second derivative with respect to parameter changes or a second order expansion of a more complicated ML model. Quadratic expansions of shallow MLPs have been previously studied (Bai & Lee, 2020; Zhu al., 2022), and the perturbation theory for small Q is studied in Roberts et al. (2022). Other related mod- els are detailed in Appendix A. We will provide evidence that even random, unstructured quadratic regression models lead to EOS behavior.
3.2 GRADIENT FLOW DYNAMICS
We will focus on training with squared loss L(z) ies under gradient flow (GF):
Ea 2,g
We begin by considering the dynam-
aL(z)
00 =J
(21)
We can write the dynamics in the output space z and the Jacohian as
à = JØ JJ. z, j Q(JT,
(22)
When Q 0 (linearize//NTK regime), J is constant, the dynamics then linear in z, and controlled by the eigenstructure of JJ the empirical NTK. In this regime there is no EOS behavior.
We are interested in settings where progressive sharpening occurs under GF. We can study the dy- namics of the maximum eigenvalue Amax of JJ T at early times for random initializations. In Ap- pendix C.I, we prove the following theorem:
Theorem 3.1. Let z, J, and Q be initialized with ii.d. elements with zero mean and variances o². o3 and 1 respectively, with distributions invariant to rotation in data and parameter space, and have finite fourth moments. Let Amax be the largest eigenvalue of JJ In the limit of large D and P. with fixe ratio D/ P/P, initialization we have
EJAmas (0)] 0, (0)] Armax (0) =
(23)
where E denotes the expectation over J, and Q at initialization.
Much like in the D = 1 case, Theorem 3.1 suggests that it is easy to find initializations that show progressive sharpening and increasing σg makes sharpening more prominent.
3.3 GRADIENT DESCENT DYNAMICS
We now consider finite-step size gradient descent (GD) dynamics. The dynamics for 0 are given by:
0 = 0 - ryJ,
(24)
In this setting, the dynamic equations can be written as
Z Z, = -JJJJ, 2, + "" 2,, J.
(25)
- 2).
(26)
If Q = O, the dynamics reduce discrete gradient descent in quadratic potential which converges iff Xxmax < 2/7
One immediate question is: when does the n in Equation 25 affect the dynamics? Given that it scales with higher powers of r) and z than the first term, we can conjecture that the ratio of the magnitudes of the terms, P'L, proportional ||z||2 and n, A calculation in Appendix C.2 shows that, for random rotationally invariant initializations, we have:
1/2
E QJJ zo. 1T))
TNL =
11 D,
(27)
Elll J zo

---

Under review as a conference paper at ICLR 2023
Loss steps, D P-
Asass VS Sleps, D= 80. P=
2.5
00%.0 - 0.05
2.0
rj -0000
Loss
1.5
6, 00,, -0000
6, 1--00,43-0000
6,-007 63
5,00.45 43388
0.2
0.5
0.0
100
200
400 0.0
100
Steps
Steps
Figure 3: Gradient descent dynamics in the quadratic regression model. z initialization variance of increases, so does the curvature Amax upon convergence. As sharpening drives 1/Amax near 2, larger og allows for non-linear effects to induce edge-of-stabiliyy behavior (right). Resulting loss trajectories non-monotonic but still converge to 0 (left).
where before the expectation is taken over the initialization of z, J, and Q. This suggests that increasing the learning rate increases the deviation of the dynamics from GF (which is obvious), but increasing ||z| also increases the deviation from GF.
We can see this phenomenology in the dynamics of the GD equations (Figure 3). Here we plot different trajectories for random initializations of the type in Theorem 3.1 with D) 60, P 120, and = 1. As σ₂ increases, so does the curvature Amax (as suggested by Theorem 3.1), and when o, (1), the dynamics is non-linean (as predicted by I'N L.) and EOS behavior emerges. This suggests that the second term in Equation 25 is crucial for the stabilization of Xmax-
We can confirm this more generally by initializing over various 17, P. ag and σ, over multiple seeds, and plotting the resulting phase diagram of the final Amax reached. We can simplify the plotting with some rescaling of parameters and initializations. For example, in the rescaled variables
2 r/z, j= 1/2
(28)
the dynamics are equivalent to Equations 25 and 26 with = 1. As in the ă T(0) model of Equations 8-9, Amax in the rescaled coordinates is equivalent mamma in the unscaled coordinates. We can also define rescaled initializations for and J. If we set
0= /. a, /,/ (DP)
(29)
then we have TNL = Ō which allows for easier comparison across (D, pairs.
Using this initialization scheme, we can plot the final value of Amax reached as a function of o₂ and 5, for 100 independent random initializations for each o., 5. pair (Figure 4). We see that the key is for PNL = ش= to O(1) corresponding to both progressive sharpening and non-linear dynamics near initialization. In particular, initializations with small ã values which converge at the EOS correspond to trajectories which first sharpen, and then settle near Amax = 2/1,. Large ōg and large ō. dynamics diverge. There is a small band of initial ō over wide range of ō_ which have final Amax Rs 2/1; these correspond to models initialized near the EOS, which stay near it.
This suggests that progressive sharpening and edge of stability aren't uniquely features of neural network models, and could be a more general property of learning in high-dimensional, non-lincar models.
4 CONNECTION TO REAL WORLD MODELS
In this section we examine how representative is the proposed model and the developed theory to the behavior of "real world" models. Following Cohen et al. (2022a), we trained a 2-hidden layer tanh network using the squared loss on 5000 examples from CIFARIO with learning rate 10-2 a setting which shows edge of stability behavior. Close to the onset of EOS, we approximately computed X1, the largest eigenvalue of JJ and its corresponding eigenvector vı using a Lanczos method (Ghorbani et al., 2019; Novak et al., 2019). We use vi to compute 2 = v z, where z is the vector of residuals [X, 0) Y for neural network function f, training inputs X, labels Y, and
7

---

Under review as a conference paper at ICLR 2023
Phase plane, 12=
Phese plane. D
Phase plane. D 120
2.0
20
20
Up Modian Miososa
Amazz
Modian
0.24 0.13 8,2 0.39 0.460.57 0.57
0.04 1.130.22 9.51 6:48 0.57
0.0
0.04 0.13 6.22 6.31 6.8 0.57 0.65
0.0
83
83
83
Figure 4: áz/03 phase planes for quadratic regression models, for various D and P. Models initialized with 100 random seeds for each 0= ãJ pair and iterated until convergence. For each pair Ts oj we plot the median Amax of the NTK J J. For intermediate oz where both sharpening and non-linear z dynamics occur, trajectories tend to converge so Amax of the NTK is near 2/n1 (EOS).
parameters 0. The EOS behavior in the NTK is similar to the EOS behavior defined with respect to the full Hessian in Cohen at every other step gets rid of the high frequency oscillations (Figure 5, middle). Unlike the D = 1, P = 2 model, there are multiple crossings of the critical line Amax 2/11 line.
(2022a) (Figure 5, left and right). Once again, plotting the trajectories
VS. sleps, n= 0.01
NTK eig vo. slaps, 9= 0.01
10-
200
10.0
7.5
-10
2.5
170
0.0
2/1
3420
3800
3500
3200
9402
2002
5900
8200
8400
3000
steps
steps
steps
Figure 5: A FCN trained on CIFAR shows multiple cycles of sharpening and edge-of-stability be- havior. 21, the projection of the training set residuals f(X, 0) Y onto the top NTK eigenmode V1, increases magnitude and oscillates around 0 (left). Plotting dynamics every two steps removes high frequency oscillations (middle). The largest eigenvalue λ crosses the edge of stability multiple times, but the second largest eigenvalue A2 remains below the edge of stability.
There is evidence that low-dimensional features of a quadratic regression model could be used to explain aspects of EOS behavior. We empirically compute the the second derivative of the output fx, 0) by automatic differentiation. We denote by Q(,,)) the resulting tensor. We can use matrix-vector products compute the spectrum of the matrix Q vı Q(::), which projection of the output of Q in the direction, without instantiating Q in memory (Figure 6, left). This figure reveals that the spectrum does not shift much from step 3200 3900 (the range of our plots). This suggests that Q doesn't t change much as these EOS dynamics are displayed. We can also see that Q much larger in the Vı direction than a random direction.
Let y be defined as remarkable agreement (Figure 6, middle). This is the same form that the dynamics of takes in our simplified model. It can also be found by iterating Equation 25 twice with fixed Jacobian for y = \ι - and discarding terms higher order in r). This suggests that during this particular EOS behavior. much like in simplified model the dynamics of the eigenvalue is more important than any rotation in the cigenbasis.
A17 2. Plotting the two-step dynamics of 21 versus 2yz we see
The dynamics of is more complicated; y/+2 Ye anticorrelated with z but there is low-order functional form in terms of y and 21 (Appendix D. We can get some insight into the stabilization by plotting the ratio of 172 P (J21V1, J21 (the non-linear contribution to the 21 dynamics from the vi direction) and X1 (the linearized contribution), and compare to the dynamics of y (Figure 6, right). The ratio is small during the initial sharpening, but becomes O(1) shortly before the
8

---

Under review as a conference paper at ICLR 2023
curvature decreases for the first time. It remains ((1) through the rest of the dynamics. This suggests that the non-linear feedback from the dynamics of the top eigenmode onto itself is crucial to understanding the EOS dynamics.
0.100
Oriated sharpriing
End sharpninng Pansion prii
0.075-
0,000
16
Elgenvalues
0.085
0,000
-100
0.2
0.4
0.1
1/2
1200
3400
3620
3800
Relative rank
Steps
Figure 6: Q is approximately constant during edge-of-stability dynamics for FCN trained on CI- FARIO (left). Projection onto largest eigendirection V1 (blue and orange) larger than projec- tion onto random direction (green). Two step difference (21)++2 - (z1)² is well approximated by 2z1y/ (middle), leading order term of models with fixed eigenbasis. Non-linear dynamical contri- bution 17² Qq (J V1, J21vv is small during sharpening, but becomes large immediately preceding decrease top eigenvalue (right) as the case in the simple model.
DISCUSSION
5.1 LESSONS LEARNED FROM QUADRATIC REGRESSION MODELS
The main lesson be learned from the quadratic regression models is that behavior like progressive sharpening (for both GF and GD) and edge-of-stabilee behavior (for GD) may be common features of high-dimensional gradient-based training of non-linear models. Indeed, these phenomena can be revealed in simple settings without any connection to deep learning models: with mild tuning simplified model, which corresponds to 1 datapoint and 2 parameters can provably show EOS behavior. This combined with the analysis of the CIFAR model suggest that the general mechanism may have low-dimensional description.
Quadratic approximations of real models quantitatively can capture the early features of EOS be- havior (the initial return Xmaz 2/n). but do not necessarily capture the magnitude and period of subsequent oscillations these require higher order terms (Appendix D.2). Nevertheless, the quadratic approximation does correctly describe much of the qualitative behavior, including convergence of Amar to limiting two-cycle that oscillates around 2/n, with an average value be- low 2/n. In the simplified two-parameter model, it is possible to analytically predict the final value at convergence, and indeed we find that it deviates slightly from the value 2/17.
key feature all the models studied this work is that looking at every-other iterate (the two- step dynamics) greatly aids in understanding the models theoretically and empirically. Near the edge of stability, this makes the changes the top eigenmode small. the simplified model, the slow ī dynamics (and related slow T(0) dynamics) allowed for the detailed theoretical analysis; in the CIFAR model, the two-step dynamics slowly varying in both 2 and Amax- The quantitative comparisons of these small changes may help uncover any universal mechanismssccanniica forms that explain EOS behavior in other systems and scenarios.
5.2 FUTURE WORK
One avenue for future work is to quantitatively understand progressive sharpening and FOS behavior in the quadratic regression model for large D and P. In particular, it may be possible predict the final deviation also be useful to understand how higher order terms affect the training dynamics. One possibility is that a small number statistics of the higher order derivatives of the loss function are sufficient to obtain a better quantitative understanding of the oscillations around y = 2.
myymmx in the edge-of-stability regime as a function oz a.,, and D/P. would
Finally, our analysis has not touched on the feature learning aspects of the model. In the quadratic regression model, feature learning is encoded in the relationship between J and z, and in particu-

---

Under review as a conference paper at ICLR 2023
lar the relationship between and the eigenstructure of JJ Understanding how Q mediates the dynamics of these two quantities may provide a quantitative basis for understanding feature learn- ing which is complementary to existing theoretical approaches (Roberts et al., 2022: Bordelon & Pehlevan, 2022; Yang et al., 2022).
REFERENCES
Yu Bai and Jason D. Lee. Beyond Linearization: Quadratic and Higher-Order Approximation of Wide Neural Networks. In International Conference on Learning Representations, March 2020.
Blake Bordelon and Cengiz Pehlevan. Self-Consisten Dynamical Field Theory of Kernel Evolution in Wide Neural Networks, May 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 on Learning Representations, February 2022a.
Jeremy M. Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, 7achary Nado, George E. Dahl, and Justin Gilmer. Adaptive Gradient Methods the Edge of Stability, July 2022h.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware Minimiza- tion for Efficiently Improving Generalization. International Conference on Learning Repre- sentations, April 2022.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An Investigation into Neural Net Optimiza- tion via Hessian Eigenvalue Density, In Proceedings of the 36th International Conference on Machine Learning, pp. 2232-2241, PMLR, May 2019.
Niv Giladi, Mor Shpigel Nacson, Elad Hoffer, and Daniel Soudry. At Stability's Edge: How to Adjust Hyperparameters to Preserve Minima Selection in Asynchronous Training of Neural Net- works? In highth International Conference on Learning Representations, April 2020.
Justin Gilmer, Behrooz Ghorbani, Ankush Garg, Sneha Kudugunta, Behnam Neyshabur, David Car- doze, George Edward Dahl, Zachary Nado, and Orhan Firat. A Loss Curvature Perspective on Training Instabilities of Deep Learning Models. In International Conference on Learning Repre- sentations, March 2022.
Jiaoyang Huang and Horng-Tzer Yau. Dynamics of Deep Neural Networks and Neural Tangent Hierarchy. In Proceedings of the 37th International Conference on Machine Learning, pp. 4542- 4551. PMLR, November 2020.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. In Advances in Neural Information Processing Systems 31, pp.
Curran Associates, Inc., 2018.
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. In Advances in Neural Information Processing Systems 32, pp. 8570- 8581. Curran Associates, Inc., 2019.
Aitor Lewkowyez, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: The catapult mechanism. March 2020.
Zhouzi Li, Zixuan Wang, and Jian Li. Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge Stability, July 2022.
Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, and Nati Srebro. Exploring Gen- eralization in Deep Learning. In Advances in Neural Information Processing Systems 30, pp- 5947-5956. Curran Associates, Inc., 2017.
Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein and Samuel S. Schoenholz. Neural Tangents: Fast and Easy Infinite Neural Networks in Python. urXiv::9122 02803 [cs, stat], December 2019.
10

---

Under review as a conference paper at ICLR 2023
Daniel A. Roberts, Sho Yaida, and Boris Hanin. The Principles of Deep Learning Theory. May 2022. doi: 10.1017/978 100002330500
Lei Wu, Chao Ma, and Weinan How SGD Selects the Global Minima in Over-parameterinn Learning: A Dynamical Stability Perspective. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Greg Yang. Tensor Programs I: Wide Feedforward or Recurrent Neural Networks of Any Architec- ture are Gaussian Processes. arXiv: /2478 [cond-mat, physies: May 2021.
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, March 2022.
Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Quadratic models for understanding neural network dynamics, May 2022.
A CONNECTION TO OTHER MODELS
A.1 HESSIAN VERSUS NTK MAXIMUM EIGENVALUE
In this work we focus on EOS dynamics of the largest eigenvalue of the NTK, rather than the Hessian as in Cohen al. (2022a). We note that version of Theorem 2.1 is true for the maximum Hessian eigenvalue well. In general, the Hessian can be written
8
82
J 0
VL
(30)
d00g
0000'
driz'
For squared loss in particular, we have
8 C
82z
0000' VL.
JJJ
(31)
0000'
As the loss gradient goes to 0, the Hessian eigenvalues approach the eigenvalues of JJT whose non-zero eigenvalues are the same as those of the empirical NTK JJ Since the theorem involves behavior as ã goes to convergence, the maximum NTK and maximum Hessian eigenvalues are equal in the limit, and the same EOS behavior applied in both cases.
For the higher dimensional models (quadratic regression model and fully connected network CI- FAR10), our experiments show that the maximum NTK eigenvalue shows edge stability behavior. The CIFAR model is the same the one in Cohen et al. (2022a) which used to illustrate the edge of stability in terms of the maximum Hessian eigenvalues. Therefore we focused on the NTK version of EOS in paper,
found it more amenable to theoretical analysis and explanation.
There are almost certainly cases where EOS behavior is displayed in the Hessian eigenvalues but the NTK eigenvalues, particularly in cases where the loss is highly non-isotropic in the outputs 82c (that is, T2ozz is far from a multiple of the identity matrix). As pointed out in previous works in these cases even the Hessian-based EOS is more difficult to analyze Cohen et al. (2022a). We leave understanding of EOS with more complicated loss functions for future work.
A.2 ONE-HIDDEN LAYER LINEAR NETWORK
Consider hidden layer network with scalar output:
f(x) = v Txx
(32)
where x is an input vector of length N, U is a K x N dimensional matrix, and v is K dimensional vector. We note that
8² fxx
0 f(x) O. ov, 5k
8 f²x(x) 6.yX2
(33)
ōv,ōv,
JUUUU
where 8,j is the Kroenecker delta. For fixed training set, this second derivative is constant; there- fore, the one-hidden layer linear network a quadratic regression model of the type studied in Section 3.
11

---

Under review as a conference paper at ICLR 2023
In the particular case of single datapoint compute the eigenvectors of the Q matrix. Let (w, W) be an eigenvector of Q, representing the v and components respectively. The eigenvector equations are
ww; Xn W jm
(34)
W jm = di Wi
(35)
Simplifying, we have:
WW = Wx
(36)
W wx
(37)
We have two scenarios. The first is that w 0. this case, we have w = 0, and W a matrix with in its nullspace. The latter condition gives us M constraints on M x N equations for total of
M (N - 1) of our (N + 1) total eigenmodes.
If w A O, then combining the equations we have the conditions:
w²w w = (x .x)
(38)
w WW Wxx
(39)
This gives us w = x. We know from Equation 37 that W is low rank. Therefore, we can guess a solution of the form
W W+,+ =e,x
(40)
where the e, are the M coordinate vectors. This suggests that we have
W (Vx x)e;
(41)
This gives us our final 2M eigenmodes.
We can analyze the initial values of of the J(w. as well. The components of the Jacobian can be written as:
(J+); = 0f(x)
Unm Xm
(42)
Öv:
0f(x)
(Ju)sm aU,m
(43)
From this form, we can deduce that J is orthogonal to the O modes. We can also compute the conserved quantity. Let J² be the total weight in the positive eigenmodes, and J2 be the total weight in the negative eigenmodes. A direct calculation shows that
w (2) 2f(x)
(44)
which implies that E 0.
Therefore, the hidden layer linear model on one datapoint is equivalent to the quartic loss model with E and eigenvalues +/x-.
A.3 CONNECTION BORDELON & PEHLEVAN (2022)
Since the one-hidden layer linear model has constant Q the models in Section F.1 of Bordclon & Pehlevan (2022) fall into the quadratic regression class. In the case of Section F.1.1, Equation 67, we can make the mapping to D model explicit. The dynamics are equivalent to said model with single eigenvalue Wo if we make the identifications
A 2, Hy = Jh 70 /2w, y - F2/2
(45)
12

---

Under review as a conference paper at ICLR 2023
A.4 CONNECTION TO NTH
The Neural Tangent Hierarchy (NTH) equations extend the NTK dynamics to account for changes in the tangent kernel by constructing an infinite sequence of higher order tensors which control the non-linear dynamics of learning Huang & Yau (2020). Truncation of the NTH equations at 3rd order related to. but not the same as the quadratic regression model, as we will show here.
The 3rd order NTH equation describes the change in the tangent kernel JJ Consider the D x D D-dimensional kernel K3 whose elements are given by
8%a 00,00, JinJ 39 +
8zB
(K3)a87=
J ja
(46)
00,08,
where repeated indices are summed over. In the NTH, for squared loss the change in the NTK JJ T given by
JJ
aB (KKs)aar
(47)
dt
For fixed Q pe3er. this equation is identical to the GF equations for the NTK in the quadratic regression model. We note that K3 not constant under the quadratic regression model. Conversely, 8²2 for fixed K3, 30307 is not constant either. Therefore, the two methods can be used to construct different low-order expansions of the dynamics.
8'2
B 2 PARAMETER MODEL
B. DERIVATION OF T-T(0) EQUATIONS
We can use the conserved quantity E to write the dynamics in terms of ã and T(0) only. Without loss of generality, let the eigenvalues are and X, with -1 S s 1. (We can achieve this by rescaling ž.) Recall the dynamical equations
2o41 - 2 = -,T,( + 1 +
(48)
T2+1 (0) - Tr(0) = =2;((2(1)) =,TT(2))
(49)
We will find substitutions for T(1) and T(2) in terms of ă and (0).. Recall that we have
T(-1) = 2
(50)
where E is conserved throughout the dynamics (and indeed is property of the landscape). We will use this definition to solve for T(1) and T(2).
Since P = 2, can write T(-1) bT()) aT((1), for coefficients a and b which are valid for all combinations of J. If J(A) 0, we have Solving, we have:
a. If J(1) we have 1 A(1 a) X²a
T(-1) (1 -1 (1) + a1 for 1
(51)
The restrictions on À translate to a & (-1, In terms of the conserved quantity E T(-1) 22, we have:
T(-1) = E 2
(52)
In order to convert the dynamics, need solve for T(1) and T(2) in terms of T(0) and E. We have:
1 T(1) a
1 a + + (a 1)T(0)
1)T(0))
(E+
(53)
We also have
(1-2 (T(0) E -22
1 (1-a) a
T(2) = T(0) +
(54)
This gives us
1
2141 - 2 = - +
(2, (3)((a 1)Tx(0) + 22, + E)
(55)
2a
13

---

Under review as a conference paper at ICLR 2023
R-
T, 8-0.00
Distribution of 0.00
Number of trajectories
190-
35
30
80
102
25
25
20
20F7
10
1.0
35
0.5
3.0
0.0
Figure 7: Phase portraits for symmetric model. Arrows indicate signs of changes in and T, and grey area represents disallowed coordinates. Dynamics are run from an evenly spaced grid of initial- izations, and the final value of the curvature T(0) is recorded. Nullelines representing ăo④1 - (blue) and T (0) (o) = 0 (orange) depend on E. Trajectories show progressive sharpening but no edge-of-stability effect (right).
1 (1-a) (0) T0( - (25, + E (a 1)T+()) + 2 T.(0) +
2
(Tr(0) E 2-22)
a²
(56)
If A -e (that is, = we recover the equations from the main text.
The non-negativity of J gives us constraints on the values of 2 and T. For >> 1 (small negative second eigenvalue), the constraints are:
T 28 E, >> -(28 + E)/a
(57)
This is an upward-facing cone with vertex at -E/2 (Figure 8, left). For a < -1, the constraints are
(23 + E)/a < +
(58)
This is a sideways facing cone with vertex at 2 -E/2 (Figure 8, right). We see that in this case, there is a limited set of values of T to converge to. Indeed, for E = 0, there is convergence except at (0)) = 0. This why we focus on the case of one positive and one negative eigenvalue.
We also solve for the nullclines the curves where either 2441 2 0 (blue in Figure 8), or (0+ - (0) = 0 (orange in Figure 8). The nullcline (2, fz()) for 2 is given by
2(2 + E) 2a (a - 1)
fa(3)
(59)
The nullcline (r)) for T(0) is given by
(a - 1)2 2a
fr(3) (a² a + 1) 2a(a (23 (22+ E) E)
(60)
The line / = 0 is also a nullcline.
For the symmetric model e the structure of the nullclines determines the presence lack of progressive sharpening. For E' = 0, there is no sharpening; the phase portrait (Figure 7, left) confirms this the nullcline in (0) divides the space into halves, one which converges, and the other which doesn't. it. However, when E * 0, the nullclines split, and there is small region where progressive sharpening can occur (Figure middle). However, there is still no edge-of- stability behavior in this case there is no region where the trajectories cluster near Xmax 2/n (Figure 7, right).
For the symmetric model < 1, the structure of the nullclines determines the presence lack of progressive sharpening. For 0, there is no sharpening; the phase portrait (Figure 7, left) confirms this the nullcline in (0) divides the space into two halves, one which converges, and the other which doesn't. tt However, when E + 0, the nullclines split, and there a small region where progressive sharpening can occur (Figure 7, middle). However, there is still edge-of- stability behavior in this case there no region where the trajectories cluster near Xmax = 2/n (Figure 7, right).
14

---

Under review as a conference paper at ICLR 2023
B= 1.50, Ao 0.83
E= 1.50, As 0.31
Figure 8: Phase planes of D 1, 1, P 2 model. Grey region corresponds to parameters forbidden
by positivity constraints on J.,,)² For À > O, allowed region is smaller and intersects = 0 at small range only. Nullclines be solved for analytically.
Final T, c 0.17, E 0.00
Ty. 0.17 E 0.00
4.0
Porp day 0
//
Number of trajectories
3500
3.5
M0 =
3000
3.0
2.5
2000
2.0 E
1.5
1500
2
1.0
0.5
500
0
0.0
22
T5
Figure 9: Phase portraits for c 0.17, E' = 0. Arrows indicate signs of changes in step dynam- ics and T(0) (left), and grey area represents disallowed coordinates. Nullclines representing 242 2 (purple) and (0) (0) 0 (green) going in to T(0) almost overlap. Final T(0) for initializations taken uniformly over ž - 7((0) space show peak near T(0) = 2 (right).
B.2 TWO-STEP DYNAMICS
The two-step difference equations can be derived by iterating Equations and 13. We have
242 - 2, po( p1 e)T,(0)) P2 <)), ps( c)T
(61)
T(0)+22- - (. = 40
(7))
q (5,,
(62)
Here the P: and qa are polynomials in غ, maximum 9th order in ă and 6th order in They can be computed explicitly but we choose to omit the exact forms for now.
Numerical simulation of the dynamics for small reveals an edge of stability effect (Figure 9). We see that distribution of final values T for random initializations has peak near T(0) (right). By plotting the two-step dynamics, we can see that the two-step nullelines which go in to T(0) 2 almost coincide (left). By studying these nullclines, we will be able to understand the edge of stability effect.
For fixed e, we can solve for the 7 two-step nullelines (ă+42 - = 0) and the I nullclines (T++((0) T,(0) = using Cardano's 's formula to solve for T as function of ž. In particular, each nullcline equation has a solution that goes through solutions that will focus
0, T(0) = 2, independent c. This the family of
Let (ž, (2)) be the nullcline of غ, and let (ž, (ā)( be the nullcline of T(0). We will show that the 7' values of the nullclines, as function of ã and e, is differentiable around 0, 0.
15

---

Under review as a conference paper at ICLR 2023
The nullclines defined by the implicit equations
7
0 62 - -
2
(63)
1 T2:3 (9:² 10c + +9) / T
1
1 T 4e +33 +0(2))
2
0 = -882 12:²( (c- 1)c + 4Tä( - 1) 2T (3e² -c+ - (c- 1) (e + 4e +1++1)
- 22 1) - T x² (7é² - 87 + - T (e- (9e² - + + (é - -
+ (c -1 (3e² + + 0(**)
(64)
We omit the higher order terms for now in anticipation of diffecrentiating at ž 0 to use the implicit function theorem. Dividing by 2, we have the equations
0 6 27 (e 1) T3? 2) + + + 7 T -1)
2
(65)
+ - T22 (9e² - 10c + 9) - g 1)
1 _/ -36² - 4e + 3) + +((2))
2
88c 12: (c 1)c 47 - 1) 2T -c + 1) (22 + 4c
(e - 1) - 7 (7e² 8e + 7 (2 (9e² - 9) T (e² - (66)
+ (E - (3e² - + 0
+
We immediately see that 3 = 0, T 2 solves both equations for all Let w(c, 2, T) and v(c, 3, T)
be the right hand sides of Equations 65 and 66 respectively. We have
dw
Ov
2,
(67)
OT (0.0.2)
8T (0.0.2)
In both cases the derivative is invertible. Therefore, (2) and (2) are continuously differen-
tiable in both Σ and e in some neighborhood of 0. In fact, since w and v are analytic in all three arguments, (2) and fr. (2) are analytic as well.
We can use the analyticity to solve for the low-order structure of the nullclines. One way to compute the values of the derivatives is to define the nullclines as formal power series:
20 oc
f=())=
(68)
j=1k=1
fr( = 2
bj zh
(69)
j=1 k=1
We can then solve for the first few terms of the series using Equations 65 and 66. From this proce- dure, we have:
fz) 2 c)
c 2) =
(70)
fr,e(2) = (2 3c + 2c2) -
1
z + 2
(4 - 4 +(
(71)
The difference (3) between the two is:
: 2 + 0
3
fa($) = fz(3) - fr(
(72)
2
As decreases, for the low order terms the distance between the nullclines also decreases.
We can show that the difference goes as e. The one-step dynamical equations for e = are
2+11 2 = - + T (0)
(73)
2
(0) (0) - - + =TT,(0)
(74)
Therefore, AE 2AT. This means that both the one step AND two-step nullclines are identical. Since f2, = fr.)(2), and both are differentiable with respect to c, we have:
fa.( - fr,e(2) = (z)
(75)
for some function (3) which is analytic in and E in neighhorhood around (0, 0).
16

---

Under review as a conference paper at ICLR 2023
B.3 TWO-STEP DYNAMICS y
It is useful to define dynamical equations in coordinates (2, where is the difference between T(0) and the & nullcline:
T(0) - f2,c(2)
(76)
To lowest order in and we have
T()) - 2 - (1- - - 2 (1 - 11
0
(77)
We note that y 0, at 2 corresponds T(0) 2. y near but slightly less than is equivalent to edge-of-stabiligy behavior. For positive ž, = 0 implies T(0) > 2.
We can write the dynamics and y. The dynamics for are given by:
242 - ( <) + (E,, ( (=.)( + p2(2,, (we p
pa <) + (²,)
(78)
We know that the right hand side this equation is analytic in ž, e, and (trivially) y as well. By evaluating the multiple continuous derivatives of f. we can write:
_ ž 2y, fi, Me) 22,
(79)
Here, fi,e and f2.. are analytic in 2, e, and y some neighborhood around 0. This means that we have the bounds
(1,r(E,)) v)1 < Fi, (2,,(2,)) v)1 < F2
(80)
for (ž, E [-ēa, 2d] x [0, ea] I-ya, yal for some non-negative constants F1 and P2. Note that this bound is independent of c.
Now consider the dynamics of y. We have:
V++2 V1 (0) T,(0) f2.0(1122) + (2,)
(81)
Since lim;-00 2t2 (x+4() is analytic in some neighborhood of (0, 0). Therefore
Ve+2 is analytic as well. Substituting, we have
yt-2 - V1 qo(@,, e) + qı (3 e)ly (a)) + qa(ā, e)ly fe. qa(E,, e)ly (z)
- f2, + 2 (E. 22,352 S2,(())
(82)
we write f2,c(3) (3) + cfa..
then we can write:
Ut+2 - 14 http::/wwo.ooooooooo..../oh.ooh.
2q |4 (3)
- -
(2,)
(83)
By the definition of the nullelines, the first four terms vanish. Once again using the differentiability of the nullclines, well as fi,c and f2,r,
rewrite the dynamics in terms of the expansion:
Vi+2 - = 2( - 3e + 4e )u 2 + ur x 91, (Z,, ve) + E (Z))
(84)
Here I1.e and I2,c are analytic near zero in غ, and c. We have the bounds
(,, G1, (2, v)1 G2
(85)
for (2,,,)) E 1- 10, cal [- yal for some non-negative constants G₁ and Gg. This bound is also independent of e.
We summarize these bounds in the following lemma:
17

---

Under review as a conference paper at ICLR 2023
Lemma B.1. Define y T - fs()). The two step dynamics of and given by
2 - 24 = 2 + vi fi, vi) + y1 (2,((2))
(86)
11+2 - yt = - 2(14- - 3e + 4e 2 yı _ 2e + ui = gn (2t, ve) + 92, yh)
(87)
Where f1. f2.0- 91,0, I2, all analytic in 2, y, and Additionally, there exist positive žeo Ves and
<e such that
121, < Fi. (f2.c(2, v)1 < F2, (2,,((2 < G1, (2,10 < G2
(88)
for all (2,e,y) E [-ãa, Ra] x [0, eal x I-ya, Mal, where Fi, P2, CG, and G2 are all non-negative constants.
We can use this Lemma to analyze the dynamies for small fixed e, for small initializations of ž, y- The control of the higher order terms will allow for an analysis which focuses on the effects of the lower order terms.
B.4 PROOF OF THEOREM 2.1
Using Lemma B.1, the dynamics in ž and can be written as:
2 ãt 2 + v2 ž, ( ve) + 2 2 (3x)
(89)
1/+2 Ut 2(4
36 46 ) 2x + v 2f (E,, ve) (2 ye)
(90)
Let E < Ed. Then we can use the bounds from Lemma B.1 to control the contributions of the higher order terms to the dynamics:
Lemma B.2. Given constants A > 0 and B > 0, there exist āo and ye such that for ă E [0, 25c].
y E -- Vel. we have the bounds:
y²² f2,c 1 Al2y3
(91)
B
1 91,
- 8
2(4 - 3e + 4c²
(92)
B
|c0 92.0 (2, v) s 4422
(93)
Proof. We begin by the following decomposition:
1t² v) + y2 f2. < 1y (2,8)1 + fe.
(94)
From Lemma 1. there exists a region [- 2n] x [0, e.] x [- ya] where the magnitudes of fi,e
fa.e, 91.e. and I2.e bounded by F1, F2, G1, and G2 respectively.
(Z,) (2() < Fiy"2 + F2y=
(95)
(2, < Guy???
(96)
( < G5
(97)
Define ān and ye as
ye = (A/F,, wa), 2 = minn(A/2232, w-)
(98)
The desired bounds follow immediately.
We consider an initialization (žo, yo) such that &n < & and you I yc. and yo S 28 Armed with Lemma B.2, we can analyze the dynamics. There are two phases; in the first phase, 2 is increasing, and y is decreasing. The first phase ends when y becomes negative for the first time reaching a value of (c)) In the second phase, is decreasing, and y stays negative and O(e).
18

---

Under review as a conference paper at ICLR 2023
B.4.1
PHASE ONE
Let tem be the time such that for sm. & < 2%g. (We will later show that & I 250 over the whole dynamics.) For & S tam. using Lemma B.2, the change in ă can be bounded from below by
2+ 2 21 - A)
(99)
Therefore at initialization, & is increasing. It remains increasing until y/ becomes negative, or 2ạ 28g. We want to show that y/ becomes negative before & 2 2%0-
2
For any S tam. Lemma B.2 gives the following upper bound on yt+2 yb:
U - Yr S - B) - (4 e)
(100)
Let t_ be the first time that Yr becomes negative. Since & is increasing for t S t-, we have
Ut++ y/ v -8 B)y - (4 - B)
(101)
This gives us the following bound yu!
V1 S V (8 B)
(102)
valid for 1 S t- and f S 1ym-
We will now show that t- < tam- Suppose that tem S t... Then tsm + 2, 2tem > 22g for the first time. Summing the bound in Equation 99, we have:
tum
50
21 A) S 450 A)
yt
(103)
where the second bound comes from the definition of (sm- Using our bound on yr. we have:
Etmm +2 Eu S 48 + A)
Σ you B)
(1 + A)
yuu
(104)
2
Since yo 2, 2t St...t +2 22g. However, by assumption & +2 2.0- We arrive contradiction; tem is not less than or equal to t...
There are three possibilities: the first is that is well-defined, and i_ smm Another possibility is that t_ is not well-defined - that is, Yr never becomes negative, In this case the bounds we derived are valid for all Therefore using Equation 102, there exists some time t, where yt, < (4 B) Then, using Equation 101 we have t1,+2 < 0. Therefore, we conclude that t_ is finite and less than tsm
Since the well defined value t- < tsm. when y first becomes negative, 7 S 28p. This means that we can continue apply the bounds from Lemma B.2 the start of the next phase. At t. 2, applying Lemma B.2 and , I 2%0, we have
yt_ vt- 4(8 B)y 2 4(4 - B)
(105)
which gives us ye. 2 - ((4 + B) This concludes the first phase. To summarize we have
- (B4+BBeEž < yr_ S 0, _σ_ < 2%o
(106)
B.4.2 PHASE TWO
Now consider the second phase of the dynamics. We will show that y remains negative and O and ă decreases to O. While y is negative, decreases. While 2 -Vo from Lemma B.2 we have
Et42 < A)2yt
(107)
Therefore as long as -30 I y < 0, & is decreasing. If this is true for all subsequent t, Eu will converge to O.
We will now show that y remains negative and Occ), concluding the proof. Lety" = 2-
We can re-write the dynamical equation for y as
Vt y/ 224 3e 4 4c v²) 2 (5.. v1) =0
(108)
19

---

Under review as a conference paper at ICLR 2023
Applying Lemma B.2 to the higher order terms, we have:
1/++ 1/+ - - 244 3c + 4e ) - v²) + B(lyel + +)
(109)
y1+2 V/2 2(4 3e + 4E )) y²) B(yel + +c),2
(110)
These inequalities are valid as long as |3/e < Ve
At t.., y" < Vt < O. When y < y/ < 0, then yel - |//|. Note that e < 2/1"|. From Equation 109,
we have
V/+22 Yt < 2 3E + 4e )² (wr- y²) + B(- ve + +c
(111)
From this inequality we can conclude that
Mo42 S (1 2(4 3c + 4 B)y, + 2/4 - 3e 4c")y* + Bc] (112)
If B < 1, then both terms are negative. We can conclude that if y² < V < 0, Ve-2 < 0. In fact, from the last term we can conclude that 1/+2 < -
Now we must show that when y² < 1/: < 0, 1/++ does not become too negative (namely, smaller than We). Using Equation 110, we have:
ye+2 > 8 (1 + 3Bag) if > y²
(113)
This means that if yr starts larger than y", it will be at most 3BE below y" the next step. Since B < 1, 3/++2 > -e if y² < yt: < O.
Finally, we will show that if y"(1 (1 + 3B/(8- B)) < < < yn (1 + 3B/(8 B)) < M++2 < O. Since +22+2 fits this condition, we can conclude that yr is negative for all f > t.,, with magnitude bounded from below by y (1 + +///BB5 /S8/888-111 B)), and complete the proof.
We will first show that y² y1 3B/ - B)) < yt implies that y² (1 (/+BB/ - B)) < 2 Let
Yr (1 +1+00 6,)y", for de 3B/(8 B). We will show that de42 < 33B/(8- - B). Using Equation
1 110, we have:
V142 2 (1+00 5t)y² 82/6,0" Bē (1 + 8,)y")
(114)
Substituting y2+2 = (1 + 8, 2) 2)y" and dividing both sides by y" have
6,-2 6, S -(8 B) + 3B
(115)
If & < 38/(8- - B), then have 6,42 < 3B/(8 B) as desired.
Finally, we will show that 0 < 6, < 3B/ (8- B) implies that 6.42 > -1 - that is, [1 + 3B/(8
B)ly" V v" implies [1 + 3B/BB88 (8- - B)l < 3t+2 O. Equation 109 implies
yt+2 S (1 + 5,)y" 82,6,y" + B (c (1 + 5.)y")
(116)
which gives us
8e- - 8, 2 -(8 - B) - -3BE?
(117)
If 5, > 0 implies
8,42 > -3B
(118)
If 3B < 1, then 8e42 1. This means that 2142 0 if [1+ 3B/ B)yy* < yt < y²
Finally, we make some choices of B and &o to guarantee convergence. Choose z < 3/7, and choose B < 1. Then in summary, what we have shown for phase two is:
At the start of the phase (time t-), 1- Vo_ < 0.
If y" Ve 0, t > t-, (1 (1 + 3B < y142 < -4c27.
If + 33/(88 - B)]y* < < y², 1 L_, [1 + 3B/(8 (8 -Bily' < 9142 < 0.
Through our choices of 20 and B, we know that + the entire trajectory for > t_ is accounted for by these regions, and [1+ for all t > t_. Additionally, we know that at least once every 2 steps, y/ < - -422 This means that the dynamics of & can be bounded from above by
< 31 +3EE). Therefore,
< ye
E bn < 2z
(119)
20

---

Under review as a conference paper at ICLR 2023
From this conclude that & converges to 0. Therefore, for any positive initialization with ug S &c. y0 2 ye. and yo S 20. we have:
lim & + 0, lim y yf 10000
(120)
t1000
where yy O(e).
Now we can prove the statement of Theorem 2.1. Given a model with S Ec. there is a continuous
space and Σ - y space. Since the region [0, xc) x [0, ye) in ž space displays edge-of-stability behavior for c-independent & and Ye (that is, (0) converges to within
mapping between
(cc) of 2 in that region), the inverse image of that neighborhood is a neighborhood in 0 space that displays lige-of-stability behavior. This concludes the proof.
B.5 Low ORDER DYNAMICS
In order to predict the final value of y, and understand the convergence to the fixed point, We can study the low order dynamics in and The low order dynamical equations are:
E+2 3s = Zyr
(121)
Vo+2 = 2(4 3c -c 4e ) 4c
(122)
For these reduced dynamics, we can show the following:
Theorem B.3. For the dynamics defined by Equations 121 and 122, for < 1, for positive initi-
tializations Eg 1, y/ 1 with the additional constraints log(e) < 162g and yo < 2 we have
lim 2 =0 0, lim yr = - 11000
(123)
t-000
Proof. The proof distinguishes two phases in the time evolution:
Phase I: ă starts positive and increases, y starts positive and decreases. At the end of the phase we want ž, 22g and y to be negative but bounded by -665ee.
Phase 2: decreases slowly, and y settles the fixed point (relatively) quickly, up error O(e)).
Let e < Consider an initialization (žo, yo) where both variables are positive, such that Eo << 1, clog(E) < 23. and yo < z From Equations 121 and 122, we see that the dynamics of will depend on the balance of the two terms.
Initially increases and decreases. We analyze the dynamics of y assuming that ă is fixed, and then compute the corrections.
At initialization, the first term in the dynamics dominates, since by assumption eă? 2 <. Since z < 1, y initially decreases exponentially with decay rate bounded from above by 828. Therefore within log( (og(-</yo)/88) steps, y <
Phase 1.
At this point, the rate of change steps, y becomes negative. Let t_ be the first time that y becomes negative. We note that yt_ 2 -4 under this analysis the first term in Equation 122 is less than yı in magnitude, so the smallest value that 2 can take if yı is positive is -4c
is at least -4c28. Therefore, in no more than 1/42/ additional
We can now understand the corrections due to the change in ž. We note that e -8E3t is an upper bound for y - since z increasing, and the -4 decreases y faster than exponential decay from the first term. Since غ is increasing, 2 e -8 as long as remains positive (t < t-). Let tem be a time such that E,ym. 22g. We can bound the change in &, for t < tsm- We know that yo you SE The change in can be bounded by
1 y00
2t pm
20 <
2irgee < 450
2 y/ < 450y0
(124)
2
20
t=0
t=0
t=0
21

---

Under review as a conference paper at ICLR 2023
If Yo 220- then the bound holds independent of the value of tsm. long the bound on y is correct. We know that the bound on y is correct until time t-; therefore, tm 2 t..
Phase 2. This proves that there cxists a time t_, such that & 22, and -162ße yt_ 0. Now that y is negative, it will stay negative, and ă will decrease until it reaches In order to understand the dynamics, for 2 + O. We have
will use change of coordinates. Consider solving Equation 122 for y/+2 y
y 2 3/2c + 2c²
(125)
Consider now the coordinate o, defined by the equation
J/c = -1 61) - 3/22e + 2
(126)
The dynamics of 6, are given by
6642 = (1 - 2(4 - 3 422)
(127)
Since 2, < 1, 6, strictly decreasing in magnitude. We can bound 8, from above by
|6 exp
z2
6_1
(128)
Since starts negative, and is decreasing in magnitude, we know that ye > means that we can bound & by
This
2-3/2c+26"
5 2 2e 20
(129)
Substitution gives the following bound dz:
| S exp
8
- 26s x2 4e
20
(130)
s=t_
Using the integral approximation for the sum, the bound becomes
-3220 l e -22 )o: 6, = =exp (-162) /e(1 - E22()) 18:
J S exp
(131)
From previous analysis, we know that -1 1 d 0. In the limit of large have
lim 16e| S exp (-162)/e) |8, 61000
(132)
we have the condition
162 llg((e)
(133)
then limy-+00 16gl < e If we want lim Bt -e/2+ (c2). then we need the condition
16=] log(e)
(134)
or equivalently -c -oo(() < 162g. Under these conditions, lim, E. 0 and lim Vt
This result can be confirmed numerically by running the dynamical equations from variety of initializations, computing the median eigenvalue (restricted to the range [1.9, 2.0)), and plotting versus e (Figure 10). We note that since the dynamics is slow, the ODE given by
1 = 2y²
(135)
ÿ - - 3c + )e² 4e
(136)
also obtains the same limit (Figure 10). The ODE suggests that the concentration relies on both the equal-orders in ă of the 8d and y' terms. well separation of timescales - converges to 0 at a rate of e, while y converges to the fixed point at a rate 2? both cases, the deviation from -c/2 scales Oc (Figure 10, right).
22

---

Under review as a conference paper at ICLR 2023
10
ODE
CDE
Discrete dynamizs
Discrete dynamics 0 0.83e*
10
Theory
6/2
112
10a
2
10
10
10a
e
Figure 10: Final values of y. normalized deviation from critical value T(0) = 2, for discrete dy- namics and ODE approximation. Deviation well approximated by e/2 over large range (left). Deviations from c/2 are (E²2) (right).
B.6 PARAMETER SPACE VS. - - SPACE
Most of our analysis has been focused in the normalized 2 T coordinate space. In this section, we confirm that the more usual setup parameter space is consistent with the normalized coordinate space. In particular, EOS behavior is often described by fixing an initialization, and training with different learning rates as in Figure
We can plot the dynamics of T(0) for the trajectories from Figure 1. We see that for small learning rates there convergence T(0) 2, large learning rates there's divergence, and for intermediate learning rates there is convergence to 2 e/2 (Figure 11).
This confirms that the theorem is useful to describe the more traditional method of discovering and exploring EOS behavior.
7 1.90
5
1.60
y= 1.30
y 1.00
0.70.0
E
2-e/2
2
1
100
Steps
Figure 11: Dynamics of T for trajectories from Figure/ 1. For small learning rate trajectories converge to 2, and for large learning rates trajectories diverge. For intermediate trajectories, we have EOS behavior, where final T is predicted by Theorem 2.1.
C QUADRATIC REGRESSION MODEL DYNAMICS
We use Einstein summation notation in this section repeated indices on the hight-hand-side of equations considered be summed over, unless they show up the left-hand-side.
23

---

Under review as a conference paper at ICLR 2023
C.1 PROOF OF THEOREM 3.1
Let z, J, and Q be initialized with i.id. random elements with mean and variance o aj. and I respectively. Furthermore, Let the distributions be invariant to rotations in both data space and parameter space, and have finite 4th moment.
In order understand the development of the curvature early times, consider coordinates which convert J into its singular value form. In these coordinates, we can write:
a Ja a
Jai
(137)
The singular values da are the square roots of the singular values of the NTK matrix. We assume that they are ordered from largest (01) to smallest in magnitude. By assumption, under this rotation the statistics and are left unchanged.
The time derivatives at 0 can be computed directly in the singular value coordinates. The first derivative is given by
d
= 20.00
(138)
dt
Using the diagonal coordinate system, we have
d
E
E J, =
E
(139)
dt
However, the average second derivative is positive. Calculating, we have:
d²
dtz 0. 2(02 + toāa)
(140)
We can compute the average at initialization. We have:
E = E R.ajj 2 Qnsi. Js = JJJj
(141)
E E = Go و Q 5k
Σ E(J 2[13,2231 =
(142)
To compute the second term, we compute Ja::
Jan Qaas (Jaل + J3)2a)
(143)
Expanding, we have:
Jai Quij (J (Jaj Ja + Q Jo.
(144)
In the diagonal coordinates J Jan = aa- This gives us:
Bla.ā.] Eloa Qaay Q Q3nJJazzz2
(145)
Averaging over the Q we get:
Elo.öa]= = P = Elo Z..
(146)
Which evaluates to:
Ela,öa| = PE
(147)
In the limit of large D and P, for fixed ratio D) P. the statistics of the Marchenko-Paae distribution allow us to compute the derivative of the largest eigenmode as
F[ = 22 + 2
(148)
Taken together, this gives
L Amax
( +1)
+1)
(149)
di
We confirm the prediction numerically in Figure 12.
24

---

Under review as a conference paper at ICLR 2023
D= 1, P= 100, 0.03
D= 100, 100,000000 / 0.08
10r
Simulation
10a
Simulation
10
Theory
Theory
10
101
10
104
10%
10
10
10P
10a
10"
10P
2
2
D= 100, P= 200, J 0.03
D= 200, P= 100, J2 0.03
Simulation
Simulation
Theory
10a
Theory
10
10
10
10%
10a
10a
10
10P
10a
101
10P
2
2
Figure 12: Average Ämax (0) versus dz, various and P (100 seeds).
That is, the second derivative of the maximum curvature positive on average. we normalize with respect to the eigenvalue scale, in the limit of large D) and P we have:
d Amaz di2
/E|
/F[Amaz]
(150)
Therefore, increasing d= increases the relative curvature of the Amaz trajectory. This gives us the proof of Theorem 3.1.
This result suggests that as increases, so does the degree of progressive sharpening. This can be confirmed by looking at GF trajectories (Figure 13). The trajectories with small σN don't change their curvature much, and the loss decays exponentially at some rate. However, when ag is larger, the curvature increases initially, and then stabilizes to higher value, allowing for faster convergence to the minimum of the loss.
Loss, D = 100, P- 200, J 0.00
Max eigenvalue, 10-000 P- 200, -000
2.-
2- 00003
0.003
1.5
0.010
Loss
2- 0.032
10%
0/322
1.0
0.100
= 0.16
100
11000
0.0
10
15
30
15
20
25
Figure 13: Gradient flow trajectories of loss and max NTK eigenvalues for quadratic regression models for varying aI- As og increases, Amax changes more quickly, and is generally increasing. Models with higher σ₂ converge faster in GF dynamics.
25

---

1 Under review as a conference paper at ICLR 2023
C.2 TIMESCALES FOR GRADIENT DESCENT
Consider random initialization of z, J, and Q. where the terms are i.idd with zero mean variances og. 03. and respectively, and finite fourth moments. Furthermore, suppose that z, J, and Q rotationally invariant in both input and output space. Under these conditions, we hope to compute
INL 1 (J (o
4 0²
(151)
I (J(Joo)oo (0)))))))))
initialization, in the limit of large and P.
The denominator is given by:
E Bs (2))
b = [
(152)
Evaluation gives us:
E Bz (a
(153)
where C4 is the 4th moment of Joi. To lowest order in D and P
J as (r(zg)Jo, Joy [o)(2 o + O(DP)
(154)
Evaluating the numerator, have:
QamnJ Zo un ElJpizaJ 8y 2g pm Zry on Byo (dimOja in + (M4 1)6c5mm)
(155)
where M1 the 4th moment of Qas- This gives us:
D Blone Qomnd mm Z Jun
= 2,
(156)
(M1 1)
Next, we perform the averages. We have
D PaQaijJ Q Qamn ym Z un
= o
+ 53, + )30
+ (Ca -') و 1
(M4 1)o [ J. yo Jovil (5 83y 888
(M4 1) JE JJ:JJ Jnd 1000
(157)
where C4 is the 4th moment of z. Simplification gives us:
El ( Qomn J ym Z, Jun
[B]JJJJ J JndJ ss) + EIJ ni sj Ja:J (e)))
D
(C4 n أزد
(Ma 1)o J J + E[J 9f JJJJJJJJ JaJJ I + E[J sj 8; J
(M4 - 1) - -') sنJ ليج p:J
(158)
For large D and P, the final three terms are asymptotically smaller than the first term. the first term, to leading order we have:
Exaluating
D
(iQiijJaazaJJ. 2, Q.amJ z, Jun http://hhaaaaaoap-.----0------:3
(159)
E(Q.ijJa: en 2nn http://wwwahnnnaaooooo//-3-----0
(160)
This gives us:
2 D0 1
D2 4
(161)
DP
to leading order, in the limit of large D and P.
26

---

Under review as a conference paper at ICLR 2023
C.3 DEPENDENCE ON D AND P
We can empirically that the sharpening is more pronouneed in the overparameterrzzed regime where D > P. Using the trajectories from Figure 4, we can make a scatter plot of the normalized maximum NTK eigenvalues M) at both initialization and the final point of the dynamics (Figure C.3). In all cases, a variety of initializations (æ-axis) lead to final values which concentrate around (y-axis).
We can see that the concentration is tightest for the overparameterized regime where P > D (right plot). We hypothesize that for large D and P. the EOS behavior is stronger and more likely to happen when P > D. We leave futher exploration of this hypothesis for future work.
final, P-
and final, D-90 P-90
Final, P-
Am)
Anea (
00
0.5
1.0
2.0 Amsaz
15
25
300
35
00
8.5
1.0
15
2.0
25
300
3.5
0.0
0.5
1.0
15
25
25
30
35
(0)
Awmez (0)
Ausa (0)
Figure 14: Scatter plots of initial vs. final normalized maximum eigenvalues 7jmma for quadratic regression models. Trajectories are taken from the data used to generate Figure 4. For large D and P. as the model becomes overparameterized (P > D), a subset of trajectories show tighter EOS behavior where 77mmar concentrates close 2 for a variety of initializations.
D ANALYSIS OF REAL MODELS
D.1 DYNAMICS OF y IN CIFARIO MODEL
The dynamics of y in the CIFARI0 model analyzed in Section 4 are more complicated than the dynamics. see from Figure 5 that there 2 and y-independent component of two-step change in y. We approximate this change by computing the average value of yt+2 Be for small 21 (taking 21 < 10 in this case). We can then subtract off from V1+2 Vt, and plot the remainder against =? (Figure 15 left). We that 3++2 Y/- b is negatively correlated with =? particularly for large 24. However, V++2
is clearly not simply function of 21.
The two-step model dynamics could be written as (ay + c) If we plot (V1+2 y: b)/2/ versus yt. we again don't ton't have a single-valued function (Figure 15, right). Therefore, the functional form of yr+2 y/ is given by b azzt + cz7.
0.0004
0.000
gt M)
0.0002
-0,006
0.0000
-0,010
-0.015
-0.0004-
-0.05
0.o
0.05
0.10
Figure 15: Gradient flow trajectories of loss and max NTK eigenvalues for quadratic regression models for varying a., As σ; increases, Amax changes more quickly, and is generally increasing. Models with higher σg converge faster in GF dynamics.
27

---

Under review as a conference paper at ICLR 2023
D.2 QUADRATIC EXPANSION OF 2-CLASS CIFAR MODEL
We trained a CIFAR model using the first two classes only with 5000 datapoints using the Neural Tangents library (Novak al., 2019)- which let perform 2nd and 3rd order Taylor expansions of the model at arbitrary parameters. The models were 2-hidden layer fully-connected networks, with hidden width 256 and Erf non-linearities. Models were initialized with the NTK parameterization, with weight variance 1 and bias variance O. The targets were scalar valued +1 for the first class, for the second class. A learning rate of 0.003204 was used in all experiments. All plots were made using float-64 precision.
Taking a quadratic expansion initialization, we see that the loss tracks the full model for the first 1000 steps in this setting (Figure 16, left), but misses the edge-of-stability behavior. We use Neural Tangents to efficiently compute the NTK to get the top eigenvalue Àι (and consequently, y). We also compute 21 by computing the associated cigenvector V1 and projecting residuals If the quadratic expansion is taken closer to the edge of stability, the dynamics true dynamics, up to a shift associated with exponential growth of 21 occurring at different times (Figure 16, middle). We see that the shape the first peak in |21| is the same for the full model and the quadratic model, but the subsequent oscillations full model. This suggests that the initial EOS behavior may be captured by the quadratic model, but the detailed dynamics require understanding of higher order terms. For example, the 3rd order Taylor expansion improves the prediction of the magnitude and period of the oscillations, but still misses key quantitative features (Figure 16, right).
21 well approximates the
faster and more quickly damped in the
10
0.5
approc
101
0.4
Loss
107
0.2
10"
Rd model
Quad. slap 1400
-066
Quad,
1000
Steps
Steps
Steps
Figure 16: Quadratic expansion of FCN model trained two-class CIFAR. Expanding initial- ization gives good approximation to full model for 1000 steps, after which EOS behavior occurs in full model but approximate one (left). When z1 is small, quadratic model tracks full model; however, initial exponential increase may happen earlier in approximate model (middle). Magni- tude of 21 has larger oscillations in full model compared to approximate model. Third-order Taylor expansion better captures magnitude and period of oscillations, but still misses quantitative features (right).
28


Reviews:
[{"review_id": "oARyjNrwLk3", "reviewer": "Reviewer_T8Nt", "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 using a second order regression model. The authors start with a scalar-valued quadratic model and show that in the case of two-dimensional input it has the EoS behavior in the non-symmetric eigenvalue scenario. After that the authors consider a more general, vector-valued second order regression model. They prove that in the high-dimensional limit and under appropriate random initialization, this model shows progressive sharpening at initialization. Then, the paper discusses the scaling at which the EoS behavior occurs. Finally, the paper shows experiments with CIFAR10 and relates them to preceding theoretical findings.     ", "summary_of_the_review": "A good paper without any serious issues.", "strength_and_weaknesses": "I like this paper. It is clearly written, gives a concise demonstration of the EoS effect on a toy model, proposes a natural general class of models shown to exhibit EoS, and the theoretical results are connected to the real world CIFAR data.  \n\nOne very minor drawback that I see is that experiments with real world data are limited to a single network and data set. ", "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": "Not applicable"}, "extra_text": {}}, {"review_id": "SkDo3IYufAx", "reviewer": "Reviewer_GxcF", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "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 consider a class of predictive models that are quadratic in parameters (equation 1 and 19),  and prove progressive sharpening (section 2.1) and edge of stability (ESO, section 2.2) in two dimension.  The paper also claims that the two properties could be general property in high dimensional non-linear models. ", "summary_of_the_review": "The paper study some fundamental property of second order regression. ", "strength_and_weaknesses": "Strength\n[+] The paper studied a second order regression models and analyze the progressive sharpening and ESO in details.\n\nWeakness\n[-] The paper takes time to follow and some of the organization could be better.\n(a) \\title{z} is used before the definition, before and after equation (3)\n(b) curvature (JJ^\\top)  could be defined after the definition of JJ^\\top after (3)\n", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IkXZmvFOuS", "reviewer": "Reviewer_7fHN", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper made an effort towards understanding the edge of stability phenomenon. I think their major observation is that second order approximation is needed to analyze the gradient flow (the prediction is a quadratic function so the loss is a 4-th order function). Then they also propose a simplified quadratic regression model to capture the edge of stability behavior. \n\nI feel the paper is answering a quite important question but I am unable to get the big picture intuition --- the only two messages i got was (i) looking at 2nd order approximation is the right way to explain edge of stability, and (ii) it requires very heavy computation to even get some simplified case out. ", "summary_of_the_review": "it is a strong result but I have low confidence. ", "strength_and_weaknesses": "Strength: I believe the authors tackle a quite important open problem related to SGD and the key message of looking at quadratic regression is valuable. \nWeakness: one of those papers you need to sit down and parse the paper multiple times to get the intuition of the technical detail part. ", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "HB3FD6wg3uJ", "reviewer": "Reviewer_saA5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper aims to understand the recently discovered \"Edge of Stability\" phenomena in a second-order regression model. The authors consider second-order NTK models and show for simple settings, under proper conditions,  one can observe progressive sharpening close to initialization. Finally, the authors simulate the simple models to observe the EoS phenomena.", "summary_of_the_review": "Overall, my scores are on the borderline. The mathematical jargons in the paper make it really difficult to understand the main message (and the underlying idea) behind the important sections. Please find my questions above.", "strength_and_weaknesses": "The major strength of the paper is to show that second-order NTKs can show the Edge of stability behavior close to initialization, under proper conditions. \n\n\nHowever, I feel the main messages of sections 2.2.2, 3.2, and 3.3 are not clear at all. I have the following questions.\na) How should one be reading the plot in Figure 2 (middle)?\n\nWhen the authors mention \"For small, the dynamics show the distinct phases described in (Li et al., 2022): an initial increase in $T(0)$, a slow increase in $\\tilde{z}$, then a decrease in $T(0)$, and finally a slow decrease of $\\tilde{z}$ while $T(0)$ remains near 2 (Figure 2, middle)\", how should one read this phenomenon in Figure 2 (middle)?\n\nMoreover, what are the blue lines in figure 2 (middle)?\n\nb) How does the result of theorem 3.1 imply progressive sharpening? The theorem statement simply says that the expected time derivative of $\\lambda_{\\max}$ is $0$ at initialization, while the expected double derivative of $\\lambda_{\\max}$ depends on the initial scale of output. How does this imply that $\\lambda_{\\max}$ will increase to $2/\\eta$ after few steps? \n\nFurthermore, the jacobian $J = G + Q(\\theta, .)$, how do the authors initialize $J$ with variance $\\sigma_J^2$? Or are the authors initializing both $G$ and $Q$ such that the condition on $J$ holds true? Same question for $z$, how do the authors make sure $z$ has variance $\\sigma_z$?\n\n\nc) In theorem 2.1, the authors require the second eigenvalue to be much smaller than $1$, otherwise, the bound on $\\lim_{t \\to \\infty} \\lambda_{\\max}$ gets very loose. Can the authors comment on whether such a big gap between the eigenvalues is necessary for EoS to hold in practice? Or do the authors need such a gap to make sure the EoS occurs near initialization?\n\nFurthermore, Cohen et al. [2020] observe that the EoS phenomenon is robust to the scale of the learning rate. However, in theorem 2.1, we require the learning rate to stay in a range $[\\eta_1, \\eta_2]$. Can the authors comment on this necessity in the proof?\n\n\nd) Is the main message behind section 3.3 that we need the second-order term in NTK for the EoS phenomena? Can the authors comment on how the ratio D/P matters affect theorem 3.1? How will the plots in Figure 3 and Figure 4 change with different D/P values?\n\nFurthermore, in theorem 2.1, we needed a large gap between the top two eigenvalues of Q. Is that somehow related to $\\sigma_z$ and $\\sigma_J$?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-TefAlq0z-a", "reviewer": "Reviewer_n7HJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the dynamics of a quadratic model trained with quadratic loss. \nThe effective dynamics of NTK is derived.\nAuthors argue that edge of stability behavior in neural networks is correlated with the behavior in quadratic regression models.", "summary_of_the_review": "This paper focuses on quadratic model with quadratic loss and derive effective dynamics of loss and NTK for the model.\nThe topic is interesting.\nHowever, there are many unjustified claims. Moreover, the experiments and theory doesn't match well.", "strength_and_weaknesses": "Strength:\n\nThe topic should be interesting to optimization community.\n\nWeakness:\n\nThis article tampered with the definition of EOS by changing Hessian into NTK. More specific discussion is in following section.\n\nMany claims is not supported. More specific discussion is in following section.\n\nThere is mismatch between experiment and theory even in simple quadratic model.", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject