S ECOND - ORDER REGRESSION MODELS EXHIBIT PROGRESSIVE SHARPENING TO THE EDGE OF STABILITY
Anonymous authors
Paper under double-blind review
A BSTRACT
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 a 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, we consider the next simplest class of predictive 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 a 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 speciﬁc structure of a neural network, suggesting that progressive sharpening and edge-of-stability behavior aren’t unique features of neural networks, and could be a more general property of discrete learning algorithms in high-dimensional non-linear models.
I NTRODUCTION
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 et 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 a function-space ODE which allows for explicit characterization of the network outputs (Jacot et al., 2018; Lee et 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 it describes a 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 large-width scaling regimes can preserve some non-linearity and allow for certain types of feature learning (Bordelon & Pehlevan, 2022; Yang et al., 2022), such approaches tend to 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 et al., 2019; Foret et al., 2022).
In 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 it stabilizes at a value equal to roughly two divided by the learning rate, corresponding to the largest eigenvalue for which gradient descent would converge in a quadratic potential (Wu et al., 2018;
Giladi et 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 understanding.
In particular, we propose a simple quadratic regression model and corresponding quartic loss function which fulﬁlls 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 limit.
Finally, we conduct a numerical analysis on the properties of a 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 as the theoretical models.
B ASIC QUARTIC LOSS FUNCTION 2. M ODEL DEFINITION
We consider the optimization of the quadratic loss function L ( θ ) =
z 2 / 2 , where z a quadratic function on the P × 1 -dimensional parameter vector θ and Q is a P × P symmetric matrix:
[Equation]
This can be interpreted either as a model in which the predictive function is 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 ﬂow (GF) dynamics with scaling factor η is given by
[Equation]
It is useful to re-write the dynamics in terms of z and the 1 × P -dimensional Jacobian J =
∂z/∂
θ :
[Equation]
The curvature is a scalar, described by the neural tangent kernel (NTK) JJ (cid:62) .
In these coordinates, we have E = JQ + J (cid:62) − 2 z , where Q + denotes the Moore-Penrose pseudoinverse.
The GF equations can be simpliﬁed by two transformations.
First, we transform to ˜ z
= ηz and ˜J =
η 1 / 2 J .
Next, we rotate θ so that Q is diagonal.
This is always possible since Q is symmetric.
Since the NTK is given by JJ (cid:62) , this rotation preserves the dynamics of the curvature.
Let ω 1 ≥ . .
≥ ω P be the eigenvalues of Q , and v i be the associated eigenvectors (in case of degeneracy, one can pick any basis).
We deﬁne ˜
J ( ω i ) =
˜
Jv i , the projection of ˜ J onto the i th eigenvector.
Then the gradient ﬂow equations can be written as:
[Equation]
The ﬁrst equation implies that ˜ z does not change sign under GF dynamics.
Modes with positive ω i ˜
z decrease the curvature, and those with negative ω i ˜
z increase the curvature.
In order to study edge-of-stability behavior, we need initializations which allow the curvature ( JJ (cid:62) in this case) to increase over time - a phenomenon known as progressive sharpening .
Progressive sharpening has been shown to be ubiquitous in machine learning models (Cohen et al., 2022a), so any useful phenomenological model should show it as well.
One such initialization for this quadratic regression model is ω 1 = − ω , ω 2 = ω ,
˜ J ( ω 1 ) =
˜ J ( ω 2 ) .
This initialization (and others) show progressive sharpening at all times.
2. G RADIENT 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 (cid:62) , remains close to the critical value 2 /η .
We deﬁne edge of stability with respect to the maximum NTK eigenvalue instead of the maximum loss Hessian eigenvalue from Cohen et al.
(2022a).
We will prove this form of EOS in our simpler models, and ﬁnd that it holds empirically in more complex models.
See Appendix
A.1 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 ﬂow analysis, this causes some trajectories to increase their curvature before convergence.
This causes the ﬁnal 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
6 θ       θ1010 -
Lo ss    4 θ6.
θη =1 . 90 η =1 .η
=1 .
30 η =1 .η =0 .1010 - 10101010 way to mitigate this issue is to consider only every other step (Figure 1, right).
We will use this observation to analyze the gradient descent (GD) dynamics directly to ﬁnd conﬁgurations where these trajectories show edge-of-stability behavior.
In the eigenbasis coordinates, the gradient descent equations are
[Equation]
[Equation]
We’ll ﬁnd it convenient in the following to write the dynamics in terms of weighted averages of ˜
J ( ω i ) 2 instead of the modes ˜
J ( ω i ) :
[Equation]
The dynamical equations become:
[Equation]
If Q is invertible, then we have E = T t ( − 1) − 2˜ z t .
Note that by deﬁnition T t (0) =
η J t J (cid:62) t is the (rescaled) NTK.
edge-of-stability behavior corresponds to dynamics which keep T t (0) near the value 2 as ˜ z t goes to 0 .
2.2. R EDUCTION TO CATAPULT DYNAMICS
If the eigenvalues of Q are {− ω, ω } , and E = 0 , the model becomes equivalent to a 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-stability behavior (Lewkowycz et al., 2020).
We will analyze this model in our ˜ z − T (0) variables as a warmup, with an eye towards analyzing a different parameter setting which does show sharpening and edge of stability.
We assume without loss of generality that the eigenvalues are {− 1 , 1 } - which can be accomplished by rescaling ˜ z .
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 ˜ z and the curvature T (0) only (Appendix B.1):
[Equation]
For E = 0 , we can see that sign(∆ T (0))
=
sign( T t (0) − 4) , as in Lewkowycz et al.
so convergence requires strictly decreasing curvature.
For E (cid:54) = 0 , there is a region where the curvature can increase (Appendix B.1).
However, there is still no edge-of-stability behavior - there is no set of initializations which starts with λ max far from 2 /η , which ends up near 2 /η .
In contrast, we will show that asymmetric eigenvalues can lead to EOS behavior.
2.2. E DGE 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 − (cid:15) , for 0 < (cid:15) ≤ 1 .
With this notation we can write the dynamical equations (Appendix B.1) as
[Equation] [Equation]
For small (cid:15) , there are trajectories where λ max is initially away from 2 /η 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 ( ηz 0 , ηT 0 (0)) to show the universality of the ˜ z - T (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 a large step size λ .
For every coordinate ˜ θ , the one-step and two-step dynamics are ˜
[Equation]
While the dynamics converge for λ < 2 , if λ > 1 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 (cid:15) , the dynamics shows the distinct phases described in (Li et al., 2022): an initial increase in T (0) , a slow increase in ˜ z , then a decrease in T (0) , and ﬁnally a slow decrease of ˜ z while T (0) remains near 2 (Figure 2, middle).
Unfortunately, the two-step version of the dynamics deﬁned by Equations 12 and 13 are more complicated – they are 3 rd order in T (0) and 9 th order in ˜ z ; see Appendix B.2 for a more detailed discussion.
However we can still analyze the dynamics as ˜ z goes to 0 .
In order to understand the mechanisms of the EOS behavior, it is useful to understand the nullclines of the two step dynamics.
The nullcline f ˜ z (˜ z ) of ˜ z and f T (˜ z ) of T (0) are deﬁned implicitly by
[Equation]
where ˜ z t +2 − ˜
z t and T t +2 (0) − T t (0) are the aforementioned high order polynomials in ˜ z and T (0) .
Since these polynomials are cubic in T (0) , there are three possible solutions as ˜ z goes to 0 .
We are particularly interested in the solution that goes through ˜ z
= 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 (cid:15) , so they become close as (cid:15) goes to 0 .
(Figure 2, middle).
In addition, the trajectories stay near f ˜ z - which gives rise to EOS behavior.
This suggests that the dynamics are slow near the nullclines, and trajectories appear to be approaching an attractor.
We can ﬁnd the structure of the attractor by changing variables to y t ≡ T t (0) − f ˜ z (˜ z t ) - the distance from the ˜ z nullcline.
To lowest order in ˜ z
and y , the two-step dynamical equations become (Appendix B.3):
[Equation] [Equation]
We immediately see that ˜ z changes slowly for small y - since we chose coordinates where ˜ z t +2 − ˜ z t = 0 when y = 0 .
We can also see that y t +2 − y t is O ( (cid:15) ) for y t
= 0 - so for small (cid:15) , the y dynamics is slow too.
Moreover, we see that the coefﬁcient of the (cid:15) ˜
z
2 t term is negative - the changes in ˜
z tend to drive y (and therefore T (0) ) to decrease.
The coefﬁcient of the y t term is negative as well; the dynamics of y tends to be contractive.
The key is that the contractive behavior takes y to an O ( (cid:15) ) ﬁxed point at a rate proportional to ˜ z
2 , while the dynamics of ˜ z are proportional to (cid:15) .
This suggests a separation of timescales if ˜ z 2 (cid:29) (cid:15) , where y ﬁrst equilibrates to a ﬁxed value, and then ˜ z converges to 0 (Figure 2, right).
This intuition for the lowest order terms can be formalized, and gives us a prediction of lim t →∞ y t
= − (cid:15)/ 2 , conﬁrmed numerically in the full model (Appendix B.5).
[Figure]
We can prove the following theorem about the long-time dynamics of ˜ z
and y when the higher order terms are included (Appendix B.4):
Theorem 2.1.
There exists an (cid:15) c > 0 such that for a quadratic regression model with E = 0 and eigenvalues {− (cid:15), 1 } , (cid:15) ≤ (cid:15) c , there exists a neighborhood U ⊂ R 2 and interval [ η 1 , η 2 ] such that for initial θ ∈ U and learning rate η ∈ [ η 1 , η 2 ] , the model displays edge-of-stability behavior:
[Equation]
for δ λ of O ( (cid:15) ) .
This neighborhood corresponds to the inverse image of the ˜ z − y space region [0 , ˜ z
c ) × [0 , y c ) , for (cid:15) -independent ˜
z c and y c . Therefore, unlike the catapult phase model, the small (cid:15) provably has EOS behavior - whose mechanism is well-understood by the ˜ z − y coordinate transformation.
Q UADRATIC REGRESSION MODEL 3. G ENERAL MODEL
While the model deﬁned in Equation 1 provable displays edge-of-stability behavior, it required tuning of the eigenvalues of Q to demonstrate it.
We can deﬁne a more general model which exhibits edge-of-stability behavior with less tuning.
We deﬁne the quadratic regression model as follows.
Given a P -dimensional parameter vector θ , the D -dimensional output vector z is given by
[Equation]
Here y is a D -dimensional vector, G is a D × P -dimensional matrix, and Q is a D × P × Pdimensional 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 ( θ , θ )
α = θ (cid:62) Q
α θ .
If Q = 0 , the model corresponds to linearized learning (as in the NTK regime).
When Q (cid:54) = 0 , we obtain the ﬁrst correction to NTK regime.
We note that:
[Equation]
for the D × P dimensional Jacobian J .
For D = 1 , we recover the model of Equation 1.
In the remainder of this section, we will study the limit as D and P increase with ﬁxed ratio D/P .
The quadratic regression model corresponds to a model with a 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 et al., 2022), and the perturbation theory for small Q is studied in Roberts et al.
Other related models are detailed in Appendix A. We will provide evidence that even random, unstructured quadratic regression models lead to EOS behavior.
3. G RADIENT FLOW DYNAMICS
We will focus on training with squared loss L ( z )
= 12 (cid:80)
α z 2 α .
We begin by considering the dynamics under gradient ﬂow (GF):
[Equation]
We can write the dynamics in the output space z and the Jacobian J as
[Equation]
When Q = 0 (linearized/NTK regime), J is constant, the dynamics are then linear in z , and are controlled by the eigenstructure of JJ (cid:62) , 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 dynamics of the maximum eigenvalue λ max of JJ (cid:62) at early times for random initializations.
In Appendix C.1, we prove the following theorem:
Theorem 3.1.
Let z , J , and Q be initialized with i.i.d.
elements with zero mean and variances σ 2 z , σ 2 J , and 1 respectively, with distributions invariant to rotation in data and parameter space, and have ﬁnite fourth moments.
Let λ max be the largest eigenvalue of JJ (cid:62) .
In the limit of large D and
P , with ﬁxed ratio D/P , at initialization we have
[Equation]
where E denotes the expectation over z , J , and Q at initialization.
Much like in the D = 1 case, Theorem 3.1 suggests that it is easy to ﬁnd initializations that show progressive sharpening - and increasing σ z makes sharpening more prominent.
3. G RADIENT DESCENT DYNAMICS
We now consider ﬁnite-step size gradient descent (GD) dynamics.
The dynamics for θ are given by:
[Equation]
In this setting, the dynamic equations can be written as
[Equation]
If Q = 0 , the dynamics reduce to discrete gradient descent in a quadratic potential - which converges iff λ max < 2 /η .
One immediate question is: when does the η 2 in Equation 25 affect the dynamics?
Given that it scales with higher powers of η and z than the ﬁrst term, we can conjecture that the ratio of the magnitudes of the terms, r NL , is proportional to ||
z
|| 2 and η .
A calculation in Appendix C.2 shows that, for random rotationally invariant initializations, we have:
[Equation] [Figure]
where as 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 σ z increases, so does the curvature λ max (as suggested by Theorem 3.1), and when σ z is O (1) , the dynamics is non-linear (as predicted by r NL ) and EOS behavior emerges.
This suggests that the second term in Equation 25 is crucial for the stabilization of λ max .
We can conﬁrm this more generally by initializing over various η , D , P , σ z , and σ J over multiple seeds, and plotting the resulting phase diagram of the ﬁnal λ max reached.
We can simplify the plotting with some rescaling of parameters and initializations.
For example, in the rescaled variables
[Equation]
the dynamics are equivalent to Equations 25 and 26 with η = 1 .
As in the ˜ z − T (0) model of
Equations 8–9, λ max in the rescaled coordinates is equivalent to ηλ max in the unscaled coordinates.
We can also deﬁne rescaled initializations for z and J .
If we set
[Equation]
then we have r NL = ˜
σ z which allows for easier comparison across ( D, P ) pairs.
Using this initialization scheme, we can plot the ﬁnal value of λ max reached as a function of ˜ σ z and ˜ σ J for 100 independent random initializations for each ˜ σ z , ˜ σ J pair (Figure 4).
We see that the key is for r NL = ˜
σ z to be O (1) - corresponding to both progressive sharpening and non-linear dynamics near initialization.
In particular, initializations with small ˜ σ J values which converge at the EOS correspond to trajectories which ﬁrst sharpen, and then settle near λ max = 2 /η .
Large ˜ σ z and large ˜ σ J dynamics diverge.
There is a small band of initial ˜ σ J over a wide range of ˜ σ z which have ﬁnal λ max ≈ 2 /η ; 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-linear models.
C ONNECTION 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 CIFAR10 with learning rate 10 − 2a setting which shows edge of stability behavior.
Close to the onset of EOS, we approximately computed λ 1 , the largest eigenvalue of JJ (cid:62) , and its corresponding eigenvector v 1 using a Lanczos method (Ghorbani et al., 2019; Novak et al., 2019).
We use v 1 to compute z 1 = v (cid:62) 1 z , where z is the vector of residuals f ( X , θ ) − Y for neural network function f , training inputs X , labels Y , and [Figure] parameters θ .
The EOS behavior in the NTK is similar to the EOS behavior deﬁned with respect to the full Hessian in Cohen et al.
(2022a)
(Figure 5, left and right).
Once again, plotting the trajectories 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 λ max = 2 /η line
[Figure]
η There is evidence that low-dimensional features of a quadratic regression model could be used to explain some aspects of EOS behavior.
We empirically compute the the second derivative of the output f ( x , θ ) by automatic differentiation.
We denote by Q ( · , · ) the resulting tensor.
We can use matrix-vector products to compute the spectrum of the matrix Q 1 ≡ v 1 · Q ( · , · ) , which is projection of the output of Q in the v 1 direction, without instantiating Q in memory (Figure 6, left).
This ﬁgure reveals that the spectrum does not shift much from step 3200 to 3900 (the range of our plots).
This suggests that Q doesn’t change much as these EOS dynamics are displayed.
We can also see that Q is much larger in the v 1 direction than a random direction.
Let y be deﬁned as y
= λ
1 η − 2 .
Plotting the two-step dynamics of z 1 versus 2 yz we see a remarkable agreement (Figure 6, middle).
This is the same form that the dynamics of ˜ z takes in our simpliﬁed model.
It can also be found by iterating Equation 25 twice with ﬁxed Jacobian for y
= λ 1 η − 2 and discarding terms higher order in η .
This suggests that during this particular EOS behavior, much like in our simpliﬁed model the dynamics of the eigenvalue is more important than any rotation in the eigenbasis.
The dynamics of y is more complicated; y t +2 − y t is anticorrelated with z 21 but there is no low-order functional form in terms of y and z 1 (Appendix D.1).
We can get some insight into the stabilization by plotting the ratio of η 2 Q 1 ( J z 1 v 1 , J z 1 v 1 ) (the non-linear contribution to the z 1 dynamics from the v 1 direction) and λ 1 z 1 (the linearized contribution), and compare it to the dynamics of y (Figure 6, right).
The ratio is small during the initial sharpening, but becomes O (1) shortly before the curvature decreases for the ﬁrst time.
It remains O (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.
[Figure]
D ISCUSSION 5. L ESSONS LEARNED FROM QUADRATIC REGRESSION MODELS
The main lesson to be learned from the quadratic regression models is that behavior like progressive sharpening (for both GF and GD) and edge-of-stability 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 our simpliﬁed 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 a low-dimensional description.
Quadratic approximations of real models quantitatively can capture the early features of EOS behavior (the initial return to λ max < 2 /η ), 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 the convergence of λ max to a limiting two-cycle that oscillates around 2 /η , with an average value below 2 /η .
In the simpliﬁed two-parameter model, it is possible to analytically predict the ﬁnal value at convergence, and indeed we ﬁnd that it deviates slightly from the value 2 /η .
A key feature of all the models studied in this work is that looking at every-other iterate (the twostep dynamics) greatly aids in understanding the models theoretically and empirically.
Near the edge of stability, this makes the changes in the top eigenmode small.
In the simpliﬁed model, the slow ˜ z dynamics (and related slow T (0) dynamics) allowed for the detailed theoretical analysis; in the CIFAR model, the two-step dynamics is slowly varying in both z 1 and λ max .
The quantitative comparisons of these small changes may help uncover any universal mechanisms/canonical forms that explain EOS behavior in other systems and scenarios.
5. F UTURE WORK
One avenue for future work is to quantitatively understand progressive sharpening and EOS behavior in the quadratic regression model for large D and P .
In particular, it may be possible to predict the ﬁnal deviation 2 − ηλ max in the edge-of-stability regime as a function of σ z , σ J , and D/P .
It would also be useful to understand how higher order terms affect the training dynamics.
One possibility is that a small number of statistics of the higher order derivatives of the loss function are sufﬁcient to obtain a better quantitative understanding of the oscillations around y = 2 .
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- lar the relationship between z and the eigenstructure of JJ (cid:62) .
Understanding how Q mediates the dynamics of these two quantities may provide a quantitative basis for understanding feature learning which is complementary to existing theoretical approaches (Roberts et al., 2022; Bordelon &
Pehlevan, 2022; Yang et al., 2022).
A C ONNECTION TO OTHER MODELS A. H ESSIAN 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 et al.
(2022a).
We note that a version of Theorem 2.1 is true for the maximum Hessian eigenvalue as well.
In general, the Hessian can be written as
[Equation] [Equation]
For squared loss in particular we have
[Equation]
As the loss gradient goes to 0 , the Hessian eigenvalues approach the eigenvalues of JJ T - whose non-zero eigenvalues are the same as those of the empirical NTK JJ T .
Since the theorem involves behavior as ˜ z 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 on CIFAR10), our experiments show that the maximum NTK eigenvalue shows edge of stability behavior.
The CIFAR model is the same as the one in Cohen et al.
(2022a) which was used to illustrate the edge of stability in terms of the maximum Hessian eigenvalues.
Therefore we focused on the NTK version of EOS in our paper, as we found it more amenable to theoretical analysis and explanation.
There are almost certainly cases where EOS behavior is displayed in the Hessian eigenvalues but not the NTK eigenvalues, particularly in cases where the loss is highly non-isotropic in the outputs (that is, ∂ 2 L ∂ z ∂ z (cid:48) 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 difﬁcult to analyze Cohen et al.
(2022a).
We leave understanding of EOS with more complicated loss functions for future work.
A. O NE - HIDDEN LAYER LINEAR NETWORK
Consider a one hidden layer network with a scalar output:
[Equation]
(32) where x is an input vector of length N , U is a K × N dimensional matrix, and v is a K dimensional vector.
We note that
[Equation]
(33) where δ ij is the Kroenecker delta.
For a ﬁxed training set, this second derivative is constant; therefore, the one-hidden layer linear network is a quadratic regression model of the type studied in
Section 3.
In the particular case of a single datapoint x , we can compute the eigenvectors of the Q matrix.
Let ( w , W ) be an eigenvector of Q , representing the v and U components respectively.
The eigenvector equations are
[Equation]
Simplifying, we have:
[Equation]
We have two scenarios.
The ﬁrst is that ω = 0 .
In this case, we have w = 0 , and W is a matrix with x in its nullspace.
The latter condition gives us M constraints on M × N equations - for a total of
M ( N − 1) of our M ( N + 1) total eigenmodes.
If ω (cid:54) = 0 , then combining the equations we have the conditions:
[Equation]
This gives us ω = ±√ x · x .
We know from Equation 37 that W is low rank.
Therefore, we can guess a solution of the form
[Equation]
where the e i are the M coordinate vectors.
This suggests that we have
[Equation]
This gives us our ﬁnal 2 M eigenmodes.
We can analyze the initial values of of the ˜ J ( ω i ) as well.
The components of the Jacobian can be written as:
[Equation] [Equation]
From this form, we can deduce that J is orthogonal to the 0 modes.
We can also compute the conserved quantity.
Let J 2+ be the total weight in the positive eigenmodes, and J 2 − be the total weight in the negative eigenmodes.
A direct calculation shows that
[Equation] [Equation]
which implies that E = 0 .
Therefore, the single-hidden layer linear model on one datapoint is equivalent to the quartic loss model with E = 0 and eigenvalues ±√ x · x .
A. C ONNECTION TO B ORDELON & P EHLEVAN (2022)
Since the one-hidden layer linear model has constant Q , the models in Section F.1 of Bordelon &
Pehlevan (2022) fall into the quadratic regression class.
In the case of Section F.1.1, Equation 67, we can make the mapping to a D = 1 model explicit.
The dynamics are equivalent to said model with a single eigenvalue ω
if we make the identiﬁcations
[Equation]
A. C ONNECTION TO NTH
The Neural Tangent Hierarchy (NTH) equations extend the NTK dynamics to account for changes in the tangent kernel by constructing an inﬁnite sequence of higher order tensors which control the non-linear dynamics of learning Huang & Yau (2020).
Truncation of the NTH equations at 3 rd order is related to, but not the same as the quadratic regression model, as we will show here.
The 3 rd order NTH equation describes the change in the tangent kernel JJ (cid:62) .
Consider the D × D ×
D -dimensional kernel K 3 whose elements are given by
[Equation]
46) where repeated indices are summed over.
In the NTH, for squared loss the change in the NTK JJ (cid:62) is given by
[Equation]
For ﬁxed Q = ∂ 2 z ∂ θ ∂ θ
cid:48) , this equation is identical to the GF equations for the NTK in the quadratic regression model.
We note that K 3 is not constant under the quadratic regression model.
Conversely, for ﬁxed K 3 , ∂ 2 z ∂ θ ∂ θ (cid:48) is not constant either.
Therefore, the two methods can be used to construct different low-order expansions of the dynamics.
B 2 PARAMETER MODEL B. D ERIVATION OF ˜ z - T (0) EQUATIONS
We can use the conserved quantity E to write the dynamics in terms of ˜ z and T (0) only.
Without loss of generality, let the eigenvalues are 1 and λ , with − 1 ≤ λ ≤ 1 .
(We can achieve this by rescaling ˜ z .)
Recall the dynamical equations
[Equation]
[Equation] [Equation]
We will ﬁnd substitutions for T (1) and T (2) in terms of ˜ z and T (0) .
Recall that we have
[Equation]
where E is conserved throughout the dynamics (and indeed is a property of the landscape).
We will use this deﬁnition to solve for T (1) and T (2) .
Since P = 2 , we can write T ( − 1) =
bT (0)
+ aT (1) , for coefﬁcients a and b which are valid for all combinations of ˜ J .
If ˜ J ( λ ) = 0 , we have b = 1 − a .
If ˜ J (1) = 0 , we have 1 = λ (1 − a )
+ λ 2 a .
Solving, we have:
[Equation]
The restrictions on λ translate to a / ∈ ( − 1 , 1) .
In terms of the conserved quantity E = T ( − 1) − 2˜ z , we have:
[Equation]
[Equation]
In order to convert the dynamics, we need to solve for T (1) and T (2) in terms of T (0) and ˜ z .
We have:
[Equation]
We also have
[Equation] [Equation]
This gives us
[Equation]
4 z
T
Final T , E =0 .0.0 0.
T f    4 z
T
Final T , E =0 .0.0 0.
T f     5 T f
N u m be r o f t r a j e c t
o r i e s
Distribution of T f , E =0 .
[Equation]
If λ = − (cid:15) (that is, a = (cid:15) − 1 ) we recover the equations from the main text.
The non-negativity of ˜ J 2 gives us constraints on the values of ˜ z and T .
For a > 1 (small negative second eigenvalue), the constraints are:
[Equation] [Equation]
This is an upward-facing cone with vertex at ˜ z
= − E/ 2 (Figure 8, left).
For a < − 1 , the constraints are
[Equation]
This is a sideways facing cone with vertex at ˜ z
= − 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 no convergence except at T (0) = 0 .
This why we focus on the case of one positive and one negative eigenvalue.
We can also solve for the nullclines - the curves where either ˜ z t +1 − ˜ z t =
(blue in Figure 8), or
T t +1 (0) − T t (0) =
(orange in Figure 8).
The nullcline (˜ z, f ˜ z (˜ z )) for ˜ z is given by
[Equation]
[Equation]
The nullcline (˜ z, f T (˜ z )) for T (0) is given by
[Equation] [Equation]
The line ˜
z
= 0 is also a nullcline.
For the symmetric model (cid:15) = 1 , the structure of the nullclines determines the presence or lack of progressive sharpening.
For E = 0 , there is no sharpening; the phase portrait (Figure 7, left) conﬁrms this as the nullcline in T t (0) divides the space into two halves, one which converges, and the other which doesn’t.
However, when E (cid:54) = 0 , the nullclines split, and there is a small region where progressive sharpening can occur (Figure 7, middle).
However, there is still no edge-ofstability behavior in this case - there is no region where the trajectories cluster near λ max = 2 /η
(Figure 7, right).
For the symmetric model (cid:15) = 1 , the structure of the nullclines determines the presence or lack of progressive sharpening.
For E = 0 , there is no sharpening; the phase portrait (Figure 7, left) conﬁrms this as the nullcline in T t (0) divides the space into two halves, one which converges, and the other which doesn’t.
However, when E (cid:54) = 0 , the nullclines split, and there is a small region where progressive sharpening can occur (Figure 7, middle).
However, there is still no edge-ofstability behavior in this case - there is no region where the trajectories cluster near λ max = 2 /η
(Figure 7, right).
[Figure] [Figure] B. T WO - STEP DYNAMICS
The two-step difference equations can be derived by iterating Equations 12 and 13.
We have
[Equation]
Here the p i and q i are polynomials in ˜ z , maximum 9 th order in ˜ z and 6 th order in (cid:15) .
They can be computed explicitly but we choose to omit the exact forms for now.
Numerical simulation of the dynamics for small (cid:15) reveals an edge of stability effect (Figure 9).
[-=
 dT =     4 z-]We see that the distribution of ﬁnal values of T for random initializations has a peak near T (0) =(right).
By plotting the two-step dynamics, we can see that the two-step nullclines which go in to
[-2 =0 .dz =
 dT-]T (0) = 2 almost coincide (left).
By studying these nullclines, we will be able to understand the edge of stability effect.
For ﬁxed (cid:15) , we can solve for the ˜ z two-step nullclines ( ˜ z t +2 − ˜ z t = 0 ) and the T nullclines ( T t +2 (0) −
T t (0) = 0 ) using Cardano’s formula to solve for T as a function of ˜ z .
In particular, each nullcline equation has a solution that goes through ˜ z
= 0 , T (0)
= 2 , independent of (cid:15) .
This is the family of solutions that we will focus on.
Let (˜ z, f ˜ z,(cid:15) (˜ z )) be the nullcline of ˜ z , and let (˜ z, f T,(cid:15) (˜ z )) be the nullcline of T (0) .
We will show that the T values of the nullclines, as a function of ˜ z and (cid:15) , is differentiable around ˜ z
= 0 , (cid:15) = 0 .
[Equation]
We omit the higher order terms for now in anticipation of differentiating at ˜ z
= 0 to use the implicit function theorem.
Dividing by ˜ z [Equation] we have the equations
[Equation]
We immediately see that ˜
z
= 0 , T = 2 solves both equations for all (cid:15) .
Let w ( (cid:15), ˜ z, T ) and v ( (cid:15), ˜ z, T ) be the right hand sides of Equations 65 and 66 respectively.
We have
[Equation]
In both cases the derivative is invertible.
Therefore, f ˜ z,(cid:15) (˜ z ) and f T,(cid:15) (˜ z ) are continuously differentiable in both ˜ z and (cid:15) in some neighborhood of 0 .
In fact, since w and v are analytic in all three arguments, f ˜ z,(cid:15)
(˜ z ) and f T,(cid:15) (˜ z ) 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 deﬁne the nullclines as formal power series:
[Equation]
We can then solve for the ﬁrst few terms of the series using Equations 65 and 66.
From this procedure, we have:
[Equation]
[Equation]
The difference f ∆ ,(cid:15) (˜ z ) between the two is:
[Equation]
As (cid:15) decreases, for the low order terms the distance between the nullclines also decreases.
We can show that the difference goes as (cid:15) .
The one-step dynamical equations for (cid:15) = 0 are
[Equation]
[Equation]
[Equation]
Therefore, ∆˜ z = 2∆ T .
This means that both the one step AND two-step nullclines are identical.
Since f ˜ z, 0 (˜ z )
=
f T,
(˜ z ) , and both are differentiable with respect to (cid:15) , we have: [Equation]
[Equation]
for some function f ∆ ,(cid:15) (˜ z ) which is analytic in (cid:15) and ˜ z in a neighborhood around (0 , 0) .
B. T WO - STEP DYNAMICS OF y
It is useful to deﬁne dynamical equations in coordinates (˜ z, y ) where y is the difference between
T (0) and the ˜ z nullcline:
[Equation]
To lowest order in ˜ z
and (cid:15) we have
[Equation]
We note that y = 0 , at ˜ z
= 0 corresponds to T (0)
= 2 .
y near but slightly less than 0 is equivalent to edge-of-stability behavior.
For positive ˜ z , y = 0 implies T (0) > 2 .
We can write the dynamics of ˜ z and y .
The dynamics for ˜ z are given by:
[Equation]
We know that the right hand side of this equation is analytic in ˜ z , (cid:15) , and (trivially)
y as well.
By evaluating the multiple continuous derivatives of f , we can write:
[Equation]
Here, f 1 ,(cid:15) and f 2 ,(cid:15) are analytic in ˜ z , (cid:15) , and y in some neighborhood around 0 .
This means that we have the bounds
[Equation] [Equation]
for (˜ z, (cid:15), y ) ∈ [ − ˜ z
d , ˜ z
d ] × [0 , (cid:15) d ] × [ − y d , y d ] for some non-negative constants F 1 and F 2 .
Note that this bound is independent of (cid:15) .
Now we consider the dynamics of y .
We have:
[Equation]
Since lim ˜ z → 0 ,y → 0 ˜ z t +2 = 0 , f ˜ z,(cid:15) (˜ z t +2 ) is analytic in some neighborhood of (0 , 0 , 0) .
Therefore y t +2 − y t is analytic as well.
Substituting, we have
[Equation]
If we write f ˜ {+z,(cid:15) (˜ z ) = f T,(cid:15) (˜ z ) + (cid:15)f ∆ ,(cid:15) (˜ z ) , then we can write: y t +2 − y t = q 
 (˜ z t , (cid:15) )+}
{++ q 1 (˜ z t , (cid:15) ) [ f T,(cid:15) (˜ z )] + q 2 (˜ z t , (cid:15) ) [ f T,(cid:15) (˜ z )] 2 + q 
+}
(˜ z t , (cid:15) ) [ f T,(cid:15) (˜ z {+)]2 q 
 (˜ z t , (cid:15)+} )
[ f T,(cid:15) (˜ z )( y + (cid:15)f ∆ ,(cid:15) (˜ z {+))] + 3 q 
 (˜ z t , (cid:15)+} ) [( f T,(cid:15) (˜ z ))
( y + (cid:15)f ∆ ,(cid:15) (˜ z )) 2 + ( f T,(cid:15)
{+(˜ z )) 
+}
( y + (cid:15)f ∆ ,(cid:15) (˜ z ))]
q 0 (˜ z t , (cid:15) ) + q 1 (˜ z t , (cid:15) ) [ y + (cid:15)f ∆ ,(cid:15) (˜ z )] + q 2 (˜ z t , (cid:15) )
[ y + (cid:15)f ∆ ,(cid:15) (˜ z )] 2 + q 3 (˜ z t , (cid:15) ) [ y + (cid:15)f ∆ ,(cid:15) (˜ z )]− f ˜ z,(cid:15) (˜ z t + 2 y t ˜ z t + y
[Equation]
f 1 ,(cid:15) (˜ z
[Equation]
y [Equation] f
[Equation] ,(cid:15) (˜ [Equation] f
[Equation]
(cid:15) (˜ z t ) (83) By the deﬁnition of the nullclines, the ﬁrst four terms vanish. Once again using the differentiability of the nullclines, as well as f 1 ,(cid:15) and f 2 ,(cid:15) , we can rewrite the dynamics in terms of the expansion:
[Equation]
Here g 1 ,(cid:15) and g 2 ,(cid:15) are analytic near zero in ˜ z , y , and (cid:15) .
We have the bounds
[Equation]
for (˜ z, (cid:15), y ) ∈ [ − ˜ z
d , ˜ z
d ] × [0 , (cid:15) d ] × [ − y d , y d ] for some non-negative constants G 1 and G 2 .
This bound is also independent of (cid:15) .
We can summarize these bounds in the following lemma:
Lemma B.1.
Deﬁne y
=
T − f ˜ z (˜ z ) .
The two step dynamics of ˜ z and y are given by
[Equation]
Where f 1 ,(cid:15) , f 2 ,(cid:15) , g 1 ,(cid:15) , g 2 ,(cid:15) are all analytic in ˜ z , y , and (cid:15) .
Additionally, there exist positive ˜
z
c , y c , and (cid:15) c such that
[Equation]
for all (˜ z, (cid:15), y ) ∈
[ − ˜ z d , ˜ z
d ] × [0 , (cid:15) d ] × [ − y d , y d ] , where F 1 , F 2 , G 1 , and G 2 are all non-negative constants.
We can use this Lemma to analyze the dynamics for small ﬁxed (cid:15) , for small initializations of ˜ z , y .
The control of the higher order terms will allow for an analysis which focuses on the effects of the lower order terms.
B. P ROOF OF T HEOREM 2.
Using Lemma B.1, the dynamics in ˜ z
and y can be written as:
[Equation]
Let (cid:15) <
(cid:15) d .
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 > [Equation] there exist ˜ z
c and y c such that for ˜ z ∈ [0 , 2˜ z c ] , y ∈ [ − y c
[Equation] y c ] [Equation] we have the bounds:
[Equation]
Proof.
We begin by the following decomposition:
[Equation]
From Lemma B.1, there exists a region [
− ˜ z d , ˜ z
d ] × [0 , (cid:15) d ] × [ − y d , y d ] where the magnitudes of f 1 ,(cid:15) , f 2 ,(cid:15) , g 1 ,(cid:15) , and g 2 ,(cid:15) are bounded by F 1 , F 2 , G 1 , and G 2 respectively.
[Equation]
Deﬁne ˜
z c and y c as
[Equation]
The desired bounds follow immediately.
We consider an initialization (˜ z 0 , y 0 ) such that ˜ z 0 ≤ ˜ z
c and y 0 ≤ y c , and y 0 ≤ ˜ z 20 .
Armed with
Lemma B.2, we can analyze the dynamics.
There are two phases; in the ﬁrst phase, ˜ z is increasing, and y is decreasing.
The ﬁrst phase ends when y becomes negative for the ﬁrst time - reaching a value of O ( (cid:15) ) .
In the second phase, ˜ z is decreasing, and y stays negative and O ( (cid:15) ) .
B.4. P HASE
[Equation]
Let t sm be the time such that for t ≤ t sm , ˜ z t ≤ 2˜ z 0 .
(We will later show that ˜ z t ≤ 2˜ z 0 over the whole dynamics.)
For t ≤ t sm , using Lemma B.2, the change in ˜ z can be bounded from below by
[Equation]
Therefore at initialization, ˜ z is increasing.
It remains increasing until y t becomes negative, or ˜
z t ≥ 2˜
z 0 .
We want to show that y t becomes negative before ˜ z t ≥ 2˜
z 0 .
For any t ≤ t sm , Lemma B.2 gives the following upper bound on y t +2 − y t :
[Equation]
Let t − be the ﬁrst time that y t becomes negative.
Since ˜ z t is increasing for t ≤ t − , we have
[Equation]
This gives us the following bound on y t :
[Equation]
valid for t ≤ t − and t ≤ t sm .
We will now show that t − < t sm .
Suppose that t sm ≤ t − .
Then at t sm + 2 , ˜
z t sm +2 > 2˜ z 0 for the ﬁrst time.
Summing the bound in Equation 99, we have:
[Equation]
where the second bound comes from the deﬁnition of t sm .
Using our bound on y t , we have:
[Equation]
Since y 0 ≤ ˜ z 20 , ˜ z t sm +2 ≤ 2˜ z 0 .
However, by assumption ˜
z t sm +2 > 2˜ z 0 .
We arrive at a contradiction; t sm is not less than or equal to t − .
There are three possibilities: the ﬁrst is that t − is well-deﬁned, and t − < t sm .
Another possibility is that t − is not well-deﬁned - that is, y t never becomes negative.
In this case the bounds we derived are valid for all t .
Therefore using Equation 102, there exists some time t (cid:15) where y t (cid:15) <
(4 − B ) (cid:15) ˜
z 20 .
Then, using Equation 101 we have y t (cid:15) +2 < 0 .
Therefore, we conclude that t − is ﬁnite and less than t sm .
Since the well deﬁned value t − < t sm ,, when y ﬁrst becomes negative, ˜ z t − ≤ 2˜ z 0 .
This means that we can continue to apply the bounds from Lemma B.2 at the start of the next phase.
At t = t − − 2 , applying Lemma B.2 and ˜ z t − ≤ 2˜ z 0 , we have
[Equation]
which gives us y t − ≥ − 4(4 + B ) (cid:15)
˜
z 20 .
This concludes the ﬁrst phase.
To summarize we have
[Equation]
Now consider the second phase of the dynamics.
We will show that y remains negative and O ( (cid:15) ) , and ˜ z decreases to 0 .
While y is negative, ˜ z decreases.
While y ≥ − y 0 , from Lemma B.2 we have
[Equation]
Therefore as long as − y 0 ≤ y < 0 , ˜ z t is decreasing.
If this is true for all subsequent t , ˜ z 0 will converge to 0 .
We will now show that y remains negative and O ( (cid:15) ) , concluding the proof.
Let y ∗ = −
[Equation]
We can re-write the dynamical equation for y as
[Equation]
Applying Lemma B.2 to the higher order terms, we have:
[Equation]
These inequalities are valid as long as | y t | < y c .
At t − , y ∗ < y t < 0 .
When y ∗ <
y t < 0 , then | y t | ≤ | y ∗ | .
Note that (cid:15) <
| y ∗ | .
From Equation 109, we have
[Equation]
From this inequality we can conclude that
[Equation]
If B < 1 , then both terms are negative.
We can conclude that if y ∗ <
y t < 0 , y t +2 < 0 .
In fact, from the last term we can conclude that y t +2 <
− 4 (cid:15) ˜
z 2 t .
Now we must show that when y ∗ <
y t < 0 , y t +2 does not become too negative (namely, smaller than − y c ).
Using Equation 110, we have:
[Equation]
This means that if y t starts larger than y ∗ , it will be at most 3 B ˜ z
20 y ∗ below y ∗ at the next step.
Since
B < 1 , y t +2 > − y c if y ∗ <
y t < 0 .
Finally, we will show that if y ∗ (1 + 3 B/ (8 − B )) <
y t < y ∗ , y ∗ (1 + 3 B/ (8 − B )) <
y t +2 < 0 .
Since y t − +2 ﬁts this condition, we can conclude that y t is negative for all t > t − , with magnitude bounded from below by y ∗ (1 + 3 B/ (8 − B )) , and complete the proof.
We will ﬁrst show that y ∗ (1 + 3 B/ (8 − B )) <
y t implies that y ∗ (1 + 3 B/ (8 − B )) <
y t +2 .
Let y t
= (1 + δ t ) y ∗ , for δ t < 3 B/ (8 − B ) .
We will show that δ t +2 < 3 B/ (8 − B ) .
Using Equation 110, we have:
[Equation]
Substituting y t +2 =
(1 + δ t +2 )
y ∗ , and dividing both sides by y ∗ we have
[Equation]
If δ t < 3 B/ (8 − B ) , then we have δ t +2 < 3 B/ (8 − B ) as desired.
Finally, we will show that 0 < δ t < 3 B/ (8 − B ) implies that δ t +2 > − 1 - that is, [1 + 3 B/ (8 −
B )] y ∗ < y t < y ∗ implies [1 + 3 B/ (8 − B )]
y ∗ <
y t +2 < 0 .
Equation 109 implies
[Equation] [Equation]
which gives us
[Equation]
[Equation]
If δ t > 0 implies
[Equation] [Equation]
If 3 B ˜ z 20 < 1 , then δ t +2 > − 1 .
This means that y t +2 <
if [1 + 3 B/ (8 − B )] y ∗ < y t < y ∗ .
Finally, we make some choices of B and ˜ z
0 to guarantee convergence.
Choose ˜
z 20 < 3 / 7 , and choose B < 12 .
Then in summary, what we have shown for phase two is:
• At the start of the phase (time t − ), y ∗ <
y t − < 0 .
• If y ∗ <
y t < 0 , t > t − , y ∗ (1 + 3 B ˜ z 20 ) <
y t +2 < − 4 (cid:15) ˜
z 2 t .
If [1 + 3 B/ (8 − B )] y ∗ < y t < y ∗ , t > t − , [1 + 3 B/ (8 − B )] y ∗ < y t +2 < 0 .
Through our choices of ˜ z 0 and B , we know that [1 + 3 B/ (8 − B )] y ∗ < y ∗ (1 + 3 B ˜ z 20 ) .
Therefore, the entire trajectory for t > t − is accounted for by these regions, and [1 + 3 B/ (8 − B )] y ∗ < y t <for all t > t − .
Additionally, we know that at least once every 2 steps, y t < − 4 (cid:15) ˜
z 2 t .
This means that the dynamics of ˜ z t can be bounded from above by
[Equation]
From this we can conclude that ˜ z t converges to 0 .
Therefore, for any positive initialization with ˜ z 0 ≤ ˜ z c , y 0 ≤ y c , and y 0 ≤ ˜ z 20 , we have:
[Equation]
where y f =
O
( (cid:15) ) .
Now we can prove the statement of Theorem 2.1.
Given a model with (cid:15) ≤ (cid:15) c , there is a continuous mapping between θ − η space and ˜ z − y space.
Since the region [0 , ˜ z c ) × [0 , y c ) in ˜ z − y space displays edge-of-stability behavior for (cid:15) -independent
˜
z c and y c (that is, T t (0) converges to within O ( (cid:15) ) of 2 in that region), the inverse image of that neighborhood is a neighborhood in θ − η space that displays edge-of-stability behavior.
This concludes the proof.
B. L OW ORDER DYNAMICS
In order to predict the ﬁnal value of y , and understand the convergence to the ﬁxed point, We can study the low order dynamics in ˜ z
and y .
The low order dynamical equations are:
[Equation]
[Equation]
For these reduced dynamics, we can show the following:
Theorem B.3.
For the dynamics deﬁned by Equations 121 and 122, for (cid:15) (cid:28) 1 , for positive inititializations ˜ z 0 (cid:28) 1 , y 0 (cid:28) 1 with the additional constraints − (cid:15) log( (cid:15) )
(cid:28) 16˜ z 20 and y 0 < 2˜ z 20 , we have
[Equation]
Proof.
The proof distinguishes two phases in the time evolution:
• Phase 1: ˜
z starts positive and increases, y starts positive and decreases.
At the end of the phase we want ˜ z t ≤ 2˜ z 0 and y to be negative but bounded by − 16˜ z 20 (cid:15) .
• Phase 2 : ˜
z decreases slowly, and y settles to the ﬁxed point (relatively) quickly, up to error O
( (cid:15) 2 ) .
Let (cid:15) (cid:28) 1 .
Consider an initialization (˜ z 0 , y 0 ) where both variables are positive, such that ˜ z 0 (cid:28) 1 , (cid:15) log( (cid:15) )
(cid:28) ˜
z 20 , and y 0 (cid:28) ˜
z 20 .
From Equations 121 and 122, we see that the dynamics of y will depend on the balance of the two terms.
Initially ˜ z increases and y decreases.
We analyze the dynamics of y assuming that ˜ z is ﬁxed, and then compute the corrections.
Phase 1.
At initialization, the ﬁrst term in the dynamics dominates, since by assumption (cid:15) ˜
z
2 t (cid:28) y t ˜ z
2 t (cid:28) .
Since ˜ z 20 (cid:28) 1 , y initially decreases exponentially with decay rate bounded from above by 8˜ z 20 .
Therefore within log( − (cid:15)/y 0 ) / 8˜ z 20 steps, y < (cid:15) .
At this point, the rate of change of y is at least − 4 (cid:15) ˜
z 20 .
Therefore, in no more than 1 / 4˜ z 20 additional steps, y becomes negative.
Let t − be the ﬁrst time that y becomes negative.
We note that y t − ≥ − 4 (cid:15) ˜
z 20 under this analysis - the ﬁrst term in Equation 122 is less than y t in magnitude, so the smallest value that y t +2 can take if y t is positive is − 4 (cid:15)
˜
z 20 .
We can now understand the corrections due to the change in ˜ z .
We note that e − 8˜ z
20 t is an upper bound for y - since ˜ z is increasing, and the − 4 (cid:15) ˜
z 2 t decreases y faster than exponential decay from the ﬁrst term.
Since ˜ z is increasing, y t ≥ e − 8˜ z 2 0 t as long as y remains positive ( t < t − ).
Let t sm be a time such that ˜ z t sm < 2˜
z 0 .
We can bound the change in ˜ z t for t < t sm .
We know that y t ≥ y
e − 8˜ z 20 t .
The change in ˜ z can be bounded by
[Equation]
If y 0 < 2˜ z 20 , then the bound holds independent of the value of t sm , as long as the bound on y is correct.
We know that the bound on y is correct until time t − ; therefore, t sm ≥ t − .
Phase 2.
This proves that there exists a time t − , such that ˜
z t − ≤ 2˜ z , and − 16˜ z 20 (cid:15) ≤ y t − ≤ 0 .
Now that y is negative, it will stay negative, and ˜ z will decrease until it reaches 0 .
In order to understand the dynamics, we will use a change of coordinates.
Consider solving Equation 122 for y t +2 − y t
=for ˜ z t (cid:54) = 0 .
We have
[Equation] [Equation]
Consider now the coordinate δ t deﬁned by the equation
[Equation]
[Equation]
The dynamics of δ t are given by
[Equation]
Since ˜ z t (cid:28) 1 , δ t is strictly decreasing in magnitude.
We can bound δ t from above by
[Equation]
[Equation] Since δ starts negative, and is decreasing in magnitude, we know that y t > −
[Equation]
This means that we can bound ˜ z t by
[Equation]
[Equation]
Substitution gives us the following bound on δ t :
[Equation]
Using the integral approximation for the sum, the bound becomes
[Equation]
[Equation]
From our previous analysis, we know that − 1 ≤ δ t − ≤ 0 .
In the limit of large t we have
[Equation] [Equation]
If we have the condition
[Equation]
[Equation]
then lim t →∞ | δ t | ≤ (cid:15) 2 .
If we want lim t →∞ y t = − (cid:15)/ 2 + O ( (cid:15) 2 ) , then we need the condition
[Equation]
(134) or equivalently − (cid:15) log( (cid:15) ) <
16˜ z 20 .
Under these conditions, lim t →∞ ˜ z t =
and lim t →∞ y t = − (cid:15)/ 2 + O ( (cid:15) 2 ) .
This result can be conﬁrmed numerically by running the dynamical equations from a variety of initializations, computing the median eigenvalue (restricted to the range [1 . 9 , 2 . 0] ), and plotting versus (cid:15) (Figure 10).We note that since the dynamics is slow, the ODE given by
[Equation]
[Equation]
[Equation]
also obtains the same limit (Figure 10).
The ODE suggests that the concentration relies on both the equal-orders in ˜ z of the y 0 and y 1 terms, as well as a separation of timescales
- ˜ z converges to 0 at a rate of (cid:15) , while y converges to the ﬁxed point at a rate ˜ z
2 t .
In both cases, the deviation from − (cid:15)/scales as O ( (cid:15) 2 ) (Figure 10, right).
[Figure]
B. P ARAMETER SPACE VS . ˜ z − T SPACE
Most of our analysis has been focused in the normalized ˜ z − T coordinate space.
In this section, we conﬁrm that the more usual setup in parameter space is consistent with the normalized coordinate space.
In particular, EOS behavior is often described by ﬁxing an initialization [Figure] and training with different learning rates - as in Figure 1.
We can plot the dynamics of T (0) for the trajectories from Figure 1.
We see that for small learning rates there is convergence to T (0) < 2 , large learning rates there’s divergence [Figure] and for intermediate learning rates there is convergence to 2 − (cid:15)/ 2 (Figure 11).
This conﬁrms that the theorem is useful to describe the more traditional method of discovering and exploring EOS behavior. [Figure]
behavior, where ﬁnal T is predicted by Theorem 2.1.
C Q UADRATIC REGRESSION MODEL DYNAMICS
We use Einstein summation notation in this section - repeated indices on the right-hand-side of equations are considered to be summed over, unless they show up on the left-hand-side.
C. P ROOF OF T HEOREM 3.
Let z , J , and Q be initialized with i.i.d.
random elements with 0 mean and variance σ 2 z , σ 2 J , and 1 respectively.
Furthermore, Let the distributions be invariant to rotations in both data space and parameter space, and have ﬁnite 4 th moment.
In order to understand the development of the curvature at early times [Equation] we consider coordinates which convert J into its singular value form.
In these coordinates, we can write:
[Equation]
[-=1 .η =
 .2 − †/ EOS-][Equation]
The singular values σ α are the square roots of the singular values of the NTK matrix.
We assume that they are ordered from largest ( σ 1 ) to smallest in magnitude.
By assumption, under this rotation the statistics of z and Q are left unchanged
[Equation]
The time derivatives at t = 0 can be computed directly in the singular value coordinates.
The ﬁrst derivative is given by
[Equation]
[Equation]
Using the diagonal coordinate system
[Equation]
we have
[Equation]
However, the average second derivative is positive.
Calculating, we have:
[Equation]
[Equation]
We can compute the average at initialization.
We have:
[Equation] [Equation]
[Equation]
To compute the second term
[Equation]
we compute ¨ J αi :
[Equation]
Expanding, we have:
[Equation]
In the diagonal coordinates J αα = σ α .
This gives us:
[Equation]
Averaging over the Q , we get:
[Equation]
Which evaluates to:
[Equation]
In the limit of large D and P , for ﬁxed ratio D/P , the statistics of the Marchenko-Pastur distribution allow us to compute the derivative of the largest eigenmode as
[Equation]
Taken together, this gives us
[Equation]
We conﬁrm the prediction numerically in Figure 12. [Figure]
That is, the second derivative of the maximum curvature is positive on average.
If we normalize with respect to the eigenvalue scale [Figure] in the limit of large D and P we have:
[Equation]
[Equation] Therefore [Figure] increasing σ z increases the relative curvature of the λ max trajectory.
This gives us the proof of Theorem 3.1.
This result suggests that as σ z increases, so does the degree of progressive sharpening.
This can be conﬁrmed by looking at GF trajectories (Figure 13).
The trajectories with small σ z don’t change their curvature much [Figure] and the loss decays exponentially at some rate.
However, when σ z is larger, the curvature increases initially [Figure] and then stabilizes to a higher value, allowing for faster convergence to the minimum of the loss [Figure]
C. T IMESCALES FOR GRADIENT DESCENT
Consider a random initialization of z , J [Equation] and Q , where the terms are i.i.d.
with zero mean variances σ 2 z , σ 2 J [Equation] and 1 respectively [Equation]
and ﬁnite fourth moments.
Furthermore, suppose that z , J [Equation] and Q are rotationally invariant in both input and output space.
Under these conditions
[Equation]
we hope to compute
[Equation]
at initialization, in the limit of large D and P [Equation]
The denominator is given by:
[Equation]
Evaluation gives us:
[Equation]
where C 4 is the 4 th moment of J αi .
To lowest order in D and P
[Equation]
Evaluating the numerator [Equation] we have:
[Equation]
where M 4 is the 4 th moment of Q αij .
This gives us:
[Equation]
Next, we perform the z averages.
We have
[Equation]
where C 4 is the 4 th moment of z .
Simpliﬁcation gives us:
[Equation]
[Equation]
For large D and P , the ﬁnal three terms are asymptotically smaller than the ﬁrst term.
Evaluating the ﬁrst term, to leading order we have:
[Equation]
This gives us:
[Equation]
to leading order, in the limit of large D and P .
C. D EPENDENCE ON D AND P
We can see empirically that the sharpening is more pronounced in the overparameterized regime where D > P .
Using the trajectories from Figure [Figure] we can make a scatter plot of the normalized maximum NTK eigenvalues ηλ max at both initialization and the ﬁnal point of the dynamics (Figure
C.3).
In all cases, a variety of initializations ( x -axis) lead to ﬁnal values which concentrate around 2 ( y -axis)
[Figure] 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.
[Figure]
D A NALYSIS OF REAL MODELS D. D YNAMICS OF y IN CIFAR10 MODEL
The dynamics of y in the CIFAR10 model analyzed in Section 4 are more complicated than the zdynamics.
We see from Figure 5 that there is a z 1 and y -independent component of the two-step change in y .
We can approximate this change b by computing the average value of y t +2 − y t for small z 1 (taking z 1 < 10 − 4 in this case).
We can then subtract off b from y t +2 − y t , and plot the remainder against z 2 t (Figure 15 left).
We see that y t +2 − y t − b is negatively correlated with z 2 t , particularly for large z t .
However, y t +2 − y t is clearly not simply function of z 1 .
The two-step model dynamics could be written as ( ay + c )˜
z 2 .
If we plot ( y t +2 − y t − b ) /z 21 versus y t , we again don’t have a single-valued function (Figure 15, right).
Therefore, the functional form of y t +2 − y t is not given by b + ayz 21 + cz 21 . [Figure]
D. Q UADRATIC EXPANSION OF 2 - CLASS CIFAR MODEL
We trained a CIFAR model using the ﬁrst two classes only with 5000 datapoints using the Neural
Tangents library (Novak et al., 2019) - which let us perform 2 nd and 3 rd 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 0 .
The targets were scalar valued - +1 for the ﬁrst class, − 1 for the second class.
A learning rate of 0 . 003204 was used in all experiments. All plots were made using ﬂoat-64 precision.
Taking a quadratic expansion at initialization, we see that the loss tracks the full model for the ﬁrst 1000 steps in this setting (Figure 16, left), but misses the edge-of-stability behavior.
We use Neural
Tangents to efﬁciently compute the NTK to get the top eigenvalue λ 1 (and consequently, y ).
We can also compute z 1 by computing the associated eigenvector v 1 and projecting residuals z .
If the quadratic expansion is taken closer to the edge of stability, the dynamics of z 1 well approximates the true z 1 dynamics, up to a shift associated with exponential growth of z 1 occurring at different times (Figure 16, middle).
We see that the shape of the ﬁrst peak in | z 1 | is the same for the full model and the quadratic model, but the subsequent oscillations are faster and more quickly damped in the full model.
This suggests that the initial EOS behavior may be captured by the quadratic model, but the detailed dynamics require an understanding of higher order terms.
For example, the 3 rd order
Taylor expansion improves the prediction of the magnitude and period of the oscillations, but still misses key quantitative features (Figure 16, right).
[Figure]