Understanding the Evolution of the Neural Tangent Kernel at the Edge of Stability
Authors: Kaiqi Jiang
Jeremy Cohen
Yuanzhi Li
Venue: NeurIPS 2025 poster ()
Keywords: Neural Tangent Kernel
training dynamics
Edge of Stability
Primary area: deep_learning
TL;DR: We reveal that Edge of Stability dynamics enhance the alignment between the Neural Tangent Kernel and the target.

Abstract: The study of Neural Tangent Kernels (NTKs) in deep learning has drawn increasing attention in recent years. NTKs typically actively change during training and are related to feature learning. In parallel, recent work on Gradient Descent (GD) has found a phenomenon called Edge of Stability (EoS), in which the largest eigenvalue of the NTK oscillates around a value inversely proportional to the step size. However, although follow-up works have explored the underlying mechanism of such eigenvalue behavior in depth, the understanding of the behavior of the NTK *eigenvectors* during EoS is still missing. This paper examines the dynamics of NTK eigenvectors during EoS in detail. Across different architectures, we observe that larger learning rates cause the leading eigenvectors of the final NTK, as well as the full NTK matrix, to have greater alignment with the training target.  We then study the underlying mechanism of this phenomenon and provide a theoretical analysis for a two-layer linear network. Our study enhances the understanding of GD training dynamics in deep learning.

Understanding the Evolution of the Neural Tangent Kernel at the Edge of Stability
Kaiqi Jiang
Jeremy Cohen
Department of Electrical and Computer Engineering Princeton University
Machine Learning Department Carnegie Mellon University Pittsburgh, PA 15213
Princeton, NJ 08540
kaiqi raiqij@princetonn @erpeerrtiiinnneemmmmeeme edu
jeremycohen&cmu. edu
Yuanzhi
Machine Learning Department Carnegie Mellon University
Pittsburgh, PA 15213 yuanzhi1@andrez cmu. edu
Abstract
The study of Neural Tangent Kernels (NTKs) deep learning has drawn increasing attention in recent years. NTKs typically actively change during training and are related feature learning. parallel, recent work on Gradient Descent (GD) has found phenomenon called Edge of Stability (EoS), in which the largest eigenvalue of the NTK oscillates around value inversely proportional to the step size. However, although follow-up works have explored the underlying mechanism of such eigenvalue behavior in depth, the understanding of the behavior of the NTK eigenvectors during EoS is still missing. This paper examines the dynamics of NTK cigenvectors during EoS detail. Across different architectures, we observe that larger learning rates cause the leading eigenvectors of the final NTK, as well as the full NTK matrix, to have greater alignment with the training target. We then study the underlying mechanism of this phenomenon and provide theoretical analysis for two-layer linear network. Our study enhances the understanding GD training dynamics in deep learning.
Introduction
Gradient Descent (GD) is a canonical minimization algorithm widely used in optimization and machine learning problems. While its behavior and guarantee in convex settings has been fully studied, there still enormous challenges when understanding its dynamics in non-convex settings, especially for deep learning, where we usually encounter high-dimensional and highly non-convex losses. In recent years, the study Neural Tangent Kernel (NTK) proposed in [IGH88 has been more and more popular when understanding the learning mechanism of GD, for example, (DZPS18) first proved the convergence of GD in the so-called lazy regime, where NTK remains unchanged during training. More and more follow-up studies point out the limitation of this lazy regime and advocate for the analysis in the rich regime where the NTK actively evolves during training. They argue that neural networks can learn more efficiently than kernel methods (GMMM20) WGL +20) DLS222). These results motivate to track the learning dynamics and the evolution of NTK more carefully.
There has also been line research on the behavior of NTK eigenvalues when training neural networks using GD. [CKL+21] observed phenomenon called Edge of Stability (EoS) where the largest eigenvalue of loss Hessian first increases and then hovers around a value inversely proportional
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

to the step size during training. This can also translate to similar behavior of the largest eigenvalue of NTK if we use MSE loss. There is lot of follow-up work studying the underlying mechanism of this eigenvalue behavior (e.g. [WLL22 DNL223) Despite the success of those works in studying the behavior of eigenvalues of NTK, the understanding of the rotation of its eigenvectors, especially those leading ones, is still lacking. Inspired by the evolution equation under Gradient Flow (GF),
d _ -) - K, -y)
dt
where is the target vector, ft and K2 represent the model output and NTK at time i, we are interested in studying the alignment between the NTK and the target, and how this alignment evolves over time.
The study of the alignment of NTK and the target not a new story [K120] SB21 WES $ 23 and many different alignment metries have been developed, e.g. Kernel Target Alignment (KTA) (CMR22).
Definition 1.1 (Full Kernel Target Alignment (KTA)). Let y I." be the target vector and K k²
Rmxm be the NTK. Then we say that the kernel target alignment |
However, there has been lack of research on how the EoS interacts with the evolution of the NTK eigenvectors and the kernel target alignment. This motivates us to study the following question:
How do eigenvectors of NTK evolve over time and align with the target, especially during EoS when eigenvalues of NTK have nontrivial oscillation?
1.1 Overview of Our Main Contributions
This paper is dedicated to addressing the above question. Apart from the KTA in Definition 1.1 we are also interested in the individual eigenvector target alignment (see Definition 1.2 We aim to track the evolution of these two metrics during GD training and study their connection the EoS.
Definition 1.2 (Individual Eigenvector Target Alignment). For the target vector y Rm. if E R
is the k-th eigenvector of the NTK, then we say that its alignment with the target is (v" wx) -Note that the alignment values sum up to 1, so they can be regarded as a kind of "distribution".
l|aili
We make the following contributions:
1. Across different architectures, observe that larger GD learning rates cause the final NTK matrix to have higher KTA alignment (Definition 1.2) with the target, as shown in Figure Ia] Looking into the mechanism, we observe that larger learning rates cause the leading eigenvectors of the final NTK to have greater individual alignment with the target (Definition 1.2) Namely, larger learning rates cause the "distribution" of the Individual Eigenvector Target Alignment shift towards leading eigenvectors, as shown in Figure lb
We refer to this phenomenon as alignment shift.
2. We study the underlying mechanism of the alignment shift phenomenon by tracking the detailed training dynamics. We find that during the phases of EoS when the sharpness reduces, we usually observe sudden fast increases in alignment, as shown in Figure Tc.
3. To better understand these empirical findings, we theoretically analyze the alignment shift phenomenon in a 2-layer linear network. This analysis reveals that the phase of EoS when the sharpness reduces is not only correlated with but essentially a contributing factor to the alignment shift. We also partially generalize our theoretical analysis to nonlinear networks in Appendix A.6
4. To further link sharpness reduction with increased kernel-target alignment, we use the central flow framework 24), which models gradient descent at EoS as implicitly following sharpness- penalized gradient flow. We show empirically that this sharpness penalty leads increased KTA.
2 Related Work
Study on eigenvalues of NTK [BD**20 showed that when wide neural networks are trained with a learning rate large enough to trigger initial instability, the largest eigenvalue of the NTK drops until stability is restored, an effect referred to as the "catapult mechanism." [CKI++21 showed that when standard neural networks are trained using GD with a fixed learning rate, the largest Hessian eigenvalue (which is nearly equivalent to the largest NTK eigenvalue up to a constant under MSE

---

Tukiin Log
132
131
122
114
112
Gumulatte Sarmerso Epennes
Evcls bnckeretutt EIA
(a)
(b)
(c)
Figure 1: (a) KTA across learning rates, (b) the alignment shift, and (c) connection between the alignment dynamics and Edge of Stability in a fully-connected network (see Section B.1p. Note that (b) only plots the alignment top 20 eigenvectors as later ones have much smaller alignment.
loss WLL22) rises until reaching the value 2/1 (the threshold beyond which instability is triggered), and then oscillates or fluctuates around that value. This phenomenon was referred to as "Edge of Stability" (EoS). Later, [WLL22] analyzed the underlying mechanism of EoS in two-layer linear networks, and [DNL23] proposed more general framework understand its mechanism in more general settings. [GKW 25 studied the ranges of 17 under which EoS oscillates with different ranks.
Alignment between NTK and the target Prior works have argued that the evolution of NTK and its aligmment with the target is related to feature learning [ES+ 23). [OJMDF2 [B21] observed that KTA is trending to increase during training. [K[20] found that top NTK eigenvectors tend to align with the learned target function, which improves the optimization. However, existing works either lack theoretical analysis of underlying mechanisms or focus on the GF setting GD case before EoS, lacking study on the NTK behavior at EoS during which its eigenvalues have nontrivial oscillation.
Relation between optimization dynamics and feature learning [ZLRB23) argued that "catapult dynamics," " which they defined as spikes in the training loss and decreases the NTK maximum eigenvalue, promote better feature learning, as quantified by the alignment of the AGOP matrix [BPP222] with the ground truth. The AGOP is an outer product matrix of the network gradients with respect to the input, whereas the NTK (studied here) is an outer product matrix of the network gradients with respect to the weights. On the two-layer linear network which we use in Section [.]] and for theoretical analysis, we find that large learning rates and sharpness reduction are much more associated with enhanced NTK alignment than with enhanced AGOP alignment (see Appendix (A.4). Note also that [ZLRB23 was limited to empirical observations and did not attempt to theoretically explain the mechanism by which catapult dynamics enhance AGOP alignment.
3 Empirical Observations about Alignment
3.1 Alignment Shift Phenomenon in Different Settings
Linear networks We start with a simple task on two-layer linear network. We use Gaussian data X as input where X = Zdiag(3, (3. 1.5, 1.2, 0.8, ..., 0.8) where each entry of Z E Rdxx is sampled ii.d from the standard normal (0, 1). Here,
200 is the number of training examples, and = 400 is the input dimension. The target a linear function of the input: Y = - all-one vector B. We use MSE loss and GD to train the model. Figure 2a the evolution of the KTA under different learning rates starting from the same initialization, showing that larger learning rates cause more significant increases of KTA. Figure 2b plots the alignment between Y and top NTK cigenvectors (U1, Ug) for the final models, showing that large learning rates cause alignment shift towards leading eigenvectors, while do not substantially change the normalized eigenvalue distribution (where the normalization is by the lo-morm of the eigenvalue list). Figure 2c plots the dynamics of the training loss, sharpness, and individual eigenvector target alignment under learning rate 0.014, showing that the alignment shift occurs during the period when the sharpness decreases.
3 X with an
Fully-connected networks We trained a 4-layer GeLU network using GD under MSE loss on subset of CIFAR-10 (K) where we randomly selected about 2000 examples in two classes (roughly 1000 for each class) and relabeled the target values +. We trained the model from the
3

---

comal Thrgpr
Neevaliee eqeesl
Eaditiin artte Pares
010
4igener seneerteeweennggg e geneecoos ondY
(a)
(b)
(c)
Figure 2: (a) KTA across learning rates, (b) the alignment shift, and (c) connection between the
alignment dynamics and Edge of Stability in 2-layer linear network (see Section 3.1
same initialization under different learning rates to roughly the same training loss. We plotted the evolution of KTA (every 50 steps) in Figure and and the individual eigenvector target alignment for the final models in Figure Tb As is mentioned in Section [ we see the increase of KTA and the alignment shift. Figure Tc zooms in the EoS period and plots the evolution of the sharpness and KTA (every 3 steps). We see that during the periods when the sharpness decreases, the KTA usually has sudden or fast increase. During the periods when the sharpness increases, the KTA increases slower or can decrease, i.e. the "speed" of increase is lower. Table [I] provides a quantitative analysis of the above connection, where we compute the total and average change (per step) of KTAA during each sharpness-decreasing period and sharpness-increasing period. We see that the average change in KTA per step, which can be regarded as the average changing "speed" of KTA, is positive and much larger in the caarness--derreasee above connection from another perspective using a tool called central flow developed by [CDT 24] and demonstrate that adding sharpness penalty to gradient flow can promote the increase of KTA.
period than in the Sharpness-increasimg period. In Section [5] we study the
Vision Transformer On the same dataset as above, we also trained simple Vir" using GD under MSE loss from the same initialization under different learning rates. Figure 3a plots the evolution of KTA (every 20 iterations) when training to roughly the same loss. Figure 3b shows the individual cigenvector target alignment well its cumulative summation for the final models. Again, we see the increase of KTA and the alignment shift phenomenon. Figure 3c zooms in the EoS period and plots the evolution of sharpness and the evolution of KTA every steps. Here, we see a stair-wise behavior of KTA where, after each period when the sharpness decreases, the KTA value goes up to the next "stair". We also add numerical evidence in Table to make a quantitative analysis, where we again see the connection between the decrease of sharpness and the fast increase of KTA per step.
Gemell Algrmant
Sormed eleense
Tiiing Lass
10
Ielibrad Poewede
Ewikter
0.1
0.7
120
landoo
(b)
(c)
Figure 3: (a) KTA across learning rates, (b) the alignment shift, and (c) connection between the aligmment dynamics and Edge of Stability in simple VIT (see Section 3.1 Similar to Figure Ib we plot the alignment of the first 20 eigenvectors in (b).
3.2 Summary and Discussion
We demonstrate that larger learning rates cause the NTK matrix as well as its leading eigenvectors to have higher alignment with the targets, and reveal the connection between this alignment behavior and EoS. In Appendix A.22, we provide supplementary results to demonstrate that our observations also
MIT license.

---

Table 1: Comparison of the changing speed of KTA between the Saarppess-ddereasingg periods and
sharpness innresiis-necraaiingg periods on (a) fully-connected network and (b) VIT. See Section 3.1 for details.
(a)
(b)
treaticn intrvvll
13-25
75:46
96-101
107-11
117-131
Heratoo iniee
195148
145-233
233-351
755-345
545-955
355-75
Stappteo
Starpaon
Total change of KTI
0.0086
0.0055
0,007
00021
0L01985
Tosal change KTA
0,0100
.00034
0.0065
00002
0.0066
400024
Aug thange dep
0.0029
0.0003
0,0488
-00014
000KS
Aug change KTA per step
8.80c4
-3.36-5
3.2644
48865
2.3024
hold in many other settings, such as different architectures, multi-clas, tasks, language transformers. Although we focus on GD for simplicity, in Appendix AA.] we show that our observations also generalize to the SGD setting. Now we add more discussion on the implication of our findings.
Connection to feature learning. Our work provides an interesting perspective on the beneficial effects of large learning rates on deep learning. While common perspective is that large learning rates find "flatter minima" which generalize better (HS97), our findings support the (potentially complementary) view that large learning rates improve feature learning [AVPVF23 ZLRB23 in our case by enhancing the alignment between the NTK and the target vector.
Connection to the generalization ability. ADH proposed a generalization bound in terms of the alignment between inverse NTK K- and the target y. which is proportional to V y" -""""y where is the number of training examples. Hence, if y gets more alignment with earlier eigenvectors of
K, i.e. later eigenvectors of K"1, we will expect to have smaller value of V y"K-'p and better generalization. This fact, together with our alignment shift observation, can help us understand the generalization ability of large learning rates. Appendix A.5 provides empirical results to demonstrate that the reducing sharpness during EoS and the alignment shift towards earlier eigenvectors are correlated with sudden decrease of the normalized y" K''K y, and that larger learning rates lead to more significant decreases. We also directly examine the change of the test loss over time and observe the improvement of the generalization ability after the zigzag period when alignment shift occurs.
"
Results on overlarge learning rates. Finally, we want to point out that the correlation between better generalization ability and better kernel target alignment discussed above also holds for overlarge learning rates. [able?]show> the change of the kernel target alignment under several large learning rates. We observe that 1) an overlarge learning rate may lead to bad test performance, 2) in the meantime, the alignment gets worse, too. That means our focus on the kernel target alignment is a reasonable choice to understand the effect of the learning rate on the generalization ability.
Table Generalization ability and KTA on large learning rates. Experiments were conducted in the fully-connected setting using SGD with batch size 256. All models were trained to the same training loss of 0.01.
Learning rates
0.02
0.025
0.03
0.035
Test losses
0.4290
0.4244
0.4241
0.4664
Kernel Target Alignment (KTA)
0.3607
0.3669
0.3878
0.3432
4 Theoretical Analysis on Two-layer Linear Networks
To better understand the mechanism of the empirical results presented in Section 3 this section we provide a theoretical analysis of the alignment behavior during EoS on a two-layer linear network.
4.1 Setup
Notation We use ||- ||2 to denote the norm a vector, and || || to denote the Frobenius norm of matrix. Let (.-)) be the inner product between vectors or matrices. Let X[i] (resp. X X(i,j) denote the i-th (resp. (ih)) component of a vector (resp. matrix) X. For a quantity which evolves throughout training, we use A, to denote its value at iteration t.
5

---

Model setup Our training dataset consists of n datapoints in d dimensions, represented by the data matrix X Rdx= and the label matrix Y E R We will use {g; and {A.] (Λ1 2 X 2 ... 2 An)) to denote the eigenvectors and eigenvalues, respectively, of the data kernel Kr == x X. We also use {p:} to denote the eigenvectors of the covariance matrix X XX".
We consider two-layer linear network with one-dimensional output which maps R- + I (()) with parameters W := (W(2), W()),, where W(()) E Rd x and W E Rixd and dp is the hidden layer size. We train the network using GD with the square loss function L(W)= = 1. | X YI By simple calculation, the NTK at iteration / is given by:
|w!e
2 X + x2 x w! X
K2
Data distribution Many prior works (e.g. [AGCHI9|]] ABP22 analyzing GD in linear networks use whitened data. On the other hand, most EoS analyses assume that the spectrum of the NTK/Hessian has just one outlier eigenvalue, as this setting results in only one direction of oscillation (which is more tractable to analyze). This typically requires the data kernel X X to have one or several outlier cigenvalues, e.g. assumed in [WLL22). Therefore in our case, we use the following data with several unwhitened leading eigenvalues, but a roughly whitened tail.
Assume that X (n).. Suppose there exists k such that 1) for the first k eigenvalues, Vi S k. X S 1/-1/7 with some y > 1; 2) the later n - 1 eigenvalues have the order ((n") with a E (0. 1). and that 2 n Σ _i= A; = 5xπ² with 6x o(1). We use linear function of our input the target, i.c. Y = f with ||| = e (n).. Further, recalling that {pi} denotes the eigenvectors of XX"", we assume Vi S n: < | ..)p e
vn This to assume that 3 is evenly distributed across the principal components of the data, reflecting the intuition that / is independent of X.
R 1 structure of w( Existing papers have shown that weight matrices approach low ranks during training 1GWB+17 LMZ18 CGMR20) For example, [ABP22 [ML244] focus on whitened data (i.e. 1
n XX is identity) and argued that if the initialization is small, then the weights will have
the following approximate updates since W is small during early training iterations:
w( ++1 & w(O") + n//n- ( Agaz
wx(2) *+1 R wr(2) + n//n Ayz W ((
where 17 is the learning rate and Aye YXT. Based on these approximate updates, [ML24] proved that under small initialization, the weights have the approximate rank-1 structure w(() x =
w!!?) x ct. u// with some cto > after a short time lo- They also proved that after time to. this approximate structure is preserved and that the u, remains unchanged and ce, are quantities which actively change. This allows us to write W!"), w!??) for > to as W:!") uzt. w!!) q.uT.
Our paper focuses on the general case with unwhitened XXT. Since eq. involves W: X, we now focus on the update of w( x and wx(2) and apply similar argument as above. Then we have 1+1 X = w(")x X + m//n. w(2)
W+1 w! + +n/n Y(W("))"
We can use a similar argument as in ABP22 JML244 to show that W!!") x and w! quickly converge to the approximate structure w(()) X = uot w = a,,uT after a short time tu with v. = = zt. In the following analysis, we will assume this approximate structure is exact.
TT
Assumption 4.1. Assume that some early iteration to, we have w, x uo? to Coo uT. with Ctp E E. u E Rdh,, Utg R.. Without loss of generality, assume that Cto and ||ll2 1. Also assume small network output at to such that Vi S n, Fr, [t] = o(Y
In Appendix B.2 we prove that if this rank-1 structure holds at time to, then it will be preserved for I >> 4o This allows us to focus on the dynamics of e and V,, In particular, the network output F; and the NTK K, can be written F; w(()) X crV,, K₁ = cxxTx x + v.oe. We also define the error vector E, == F₂ - Y. Here and throughout, we use standard notations ((),, @ and o(-) to represent order of quantities when - xc.
Thhi assumes that to in early training. which has been proven verified before, e.g. in ABP22 JML24

---

4.2 Four Phases of the Training Dynamics
WL.122 DNL revealed that when there is one eigenvalue at EoS, the training dynamics can be divided into four phases based on the position and the sign of change in the sharpness (see Ap- pendix [B.T] for a cartoon illustration). In the two-layer linear network setting, [WLL22] demonstrated that the largest eigenvalue of the Hessian is nearly equivalent to the largest eigenvalue of the NTK, and that this is in turn nearly equivalent to the norm squared of the second-layer weight matrix times the largest eigenvalue of the data kernel: Amax ( (W:)(W)) RJ 1o A s A Since X is constant, the four phases can therefore be specified by the position and the sign of change in of.
While [WLL22] essentially ignored the quadratic term of the step size, i.c. 17² our results reveal that 17² plays a crucial role on the alignment shift. More details can be found later in Phase III analysis.
Phase I: e7 < and is increasing During this phase, the training dynamies are stable and approximately follow the GF trajectory. Under GF updates, we have that
dt G = (E, - 12
d
n 21
n2 Σ (E., (F...))
1
11,11=--2
2n " n
Σ
(eE 9:) (2)
where 9: is the i-th eigenvector of the data kernel X mentioned before. In Appendix B.3
prove that under GF, vi S n: (E,, q.) (F.. < 0 for t > with some t', making (E,, Fi) < 0 and of increase. If the learning rate is very small, the loss will converge before the sharpness increases to the stability threshold 2/1,. Otherwise, the sharpness will increase to 2/1 during training, causing the dynamics to go unstable and move to Phase II discussed below. To ensure that this occurs, need to carefully pick our learning rate. Appendix discusse the intuition on the choice of 1j in detail, where we pick increases as well during Phase A more careful analysis can be found in Appendix B.3
On A2²) ). Noticing the ||,, 2 update in eq. 2 we also have that ||ve||2
Here we see the behavior of cî and ||v,|| is consistent, in that they are both increasing. This will not always be true in other phases, especially in Phase III. It turns out that the different behavior of of and [|e/ 2 in Phase III is the key contributing factor to the alignment shift phenomenon. For later phases, we make the following assumptions.
Assumption 4.2. Assume that during the whole training procedure, A 2 S n
C for some constant
C> O and that do c7 1 for
The first part of Assumption 4.2 stipulates that the largest eigenvalue of the approximate NTK is
smaller than C. an assumption previously made in [WL22 with C 4. The second part stipulates that other eigenvalues are small, which was also assumed in WLL22].
7)
Assumption 4.3. Assume that during the training procedure, || 12 < O (n/1)). Also assume that at the end of Phase I, the loss has undergone a nontrivial decrease: | = 6ol|Y|| with do o(1).
Further denote 51 := coos(Y, (Y,g)) q1)1 = 1YY, //>l//|YY|I 52 := max V81}.
The first part of Assumption 4.3 assumes that the loss does not diverge too much during training.
Because we pick 77 = O poly(n) mentioned before and the initial || |Eol is at the order ([|
On)), the assumption that E 2 < O(/nn) is saying that the loss no larger than poly(n) times the initial loss, which is not strong assumption.
17 and is still increasing. During this phase, the approximate sharpness has risen above 2. which will lead to the oscillation of the network output along some direction. More precisely, we prove in Lemma 4.4 that the error vector E, will oscillate with increasing magnitude along the qı direction, whereas (E,, for i 2 will remain small in magnitude.
Phase II: Ai ct
n
Lemma 4.4. Under our data distribution (Section 4. and Assumption 4.1 4.3 as discussed before. Write A Aa
Pick n O(n -
2 with A
0. Then for during Phase II, have that
a
R
(Fe, 2 (1 -
(E, )² < 0 (B3IT))))) -
7

---

Recall from Phase I that under the gradient flow, the time derivative of o was determined by the sign of Ft). We prove in Appendix B.5 that in Phase II and later phases, even though GD no longer follows the gradient flow, the sign of c of is still determined by the sign of {Es, Fr)- That means in Phase II, ct will keep increasing as long as {Ez, Fi) < O. However, since Lemma 4.4 shows the divergence of E, along the qI direction, once [(E,, q1)l grows sufficiently large, we have [|E;]] > ||Y||, which yields (E,, Ft) = | + (Er, Y) 2 |E l12 || || > 0, causing
E+1 < of and we move to Phase III. The detailed analysis of Phase II is presented in Appendix B.6
Phase III: q and starts to decrease. As discussed in Phase II, the divergence of E, along the ai direction eventually causes of to decrease. In Phase III, of will keep decreasing until the sharpness (which is approximately A in our case) goes below 2 and (1), starts to shrink. " Then we will move to Phase IV. As for || we can prove that ||w,||² is increasing in Phase III. In particular, the one-step evolution of |u,|| 2 is given by the following lemma. Lemma 4.5. We have that
|| 112 - -110011 2n + Aoj
Aot (1 n
20
2n E, K,Y n
(E, 9i)²
A Ae
n
n
where A is defined in Lemma
4.4
By the divergence of kE,, q1) we can prove that the RHS is dominated by the second term Ande Ai 91) (see Appendix B.7 which does not exist under GF update and was usually ignored in prior works. Note that in Phase III, d o7 2. and hence A, > 0. This leads to the increase of ||vell?. In the following section, we will see that this increasing behavior of ||w,/1 which is different from the decrease of ct: is the key contributing factor the alignment shift, suggesting that the typically ignored 1j² term actually plays a crucial role on the alignment behavior during EoS.
Phase IV: 2 = and is still decreasing. In Phase IV, E, starts to shrink along the qi direction. [WLL22] showed that we eventually have (E,, F2) < O and go back to Phase I. Our paper will not analyze Phase IV in detail but will study the overall behavior after going through Phase IV.
n
4.3
Aligmment Shift During the Sharpness Reduction Phases (III and IV)
Now we are ready to analyze the alignment shift phenomenon. We begin by defining the following ratio between of and ||v,1l 2 which turns out to be crucial quantity when analyzing alignment shift.
Definition 4.6. Define ag == ||
We note that the sharpness reduction phases (III and IV) are typically correlated with the spikes of the loss (as is shown empirically in Section 3.1) 11 WLL22 have observed a non-monotonically decreasing behavior of the loss where it has decreasing trend in spite of oscillation. In other words, the loss after the spike is typically around or smaller than the value before the spike. This motivates us to introduce the following assumption on the change of loss after going through phases III and IV.
Assumption 4.7. Let t1 be the start of Phase III and t2 the end of Phase IV. Assume |E,,) 1| s ||E,, 1|2. Theorem 4.8. Denote A. 21 c
2. Under our data distribution (see Section 4.1 1 and Assump-
n
tion 4.1 4.3 and pick 17 0 12))) as discussed before.
(A). For in Phase III such that [Fa,, q 2 (xb2ll with &n defined in Lemma 4.4 and D > (1max[1, m }), have that ct+1 < og and | > |0:11
(B). Suppose we further have Assumption 4.7 For t,, t2 defined in Assumption 4.7 if A., 2 (4
then will get that 042, > ai,, and that (U, Y) 2 1 co82)).os(vv) Y) 2 O62))
Theorem 4.8 (A) implies a one-step increasing behavior of o; in Phase III. In Phase IV, although Or may not have this one-step increase, Theorem 4.8 (B) shows that it will have an overall increase after going through the whole Phase IIII and IV, i.e. atz > Ctt, It also shows that at times t, and
l2, v. approximately aligns with the training target Y, which allows us to write K2 at l and l2 as
K, x xxTx x + ||vell jy"" where ÿ == Y /||Y|| E 12d is the unit vector.
8

---

We now argue that in the above approximate structure at t1 and ł2, a larger de causes ÿ to align with earlier NTK cigenvectors, as opposed to the later ones. Intuitively, if || is large relative to c² (i.e. if = |0| 1e/l// / is large), then early eigenvectors of K will align with ÿ. rather than with the top eigenvectors of X X. The following lemma gives a sufficient condition (involving a and the eigenvalue spectrum of X X) for the j-th NTK eigenvector to be the one with the highest alignment with ÿ. The lemma is just a property of eigenvalue spectra, and does not involve training dynamics. Hence for ease of notation, we drop the iteration subscript in the lemma statement. We provide an illustration of it in a warm-up setting in Appendix C.TT proof' in the general case in Appendix C.2
Lemma 4.9. Let (1...... In be the eigenvectors of the approximate NTK matrix e² xx X + with corresponding non- non-increasing eigenvalues, and let a == c Consider the k and 5₂ defined
in our data distribution in Section 4.1 For any j S k - 1, if satisfies Az-
1-016Atn w/2) < aaz
A--1-de 14066xtn-0/7) we have that arg maxigign Kq: 9) = =.
Theorem 4.8 (B) and Lemma 4.9 suggest that Y tends to align more with earlier NTK eigenvectors at l2 than at f1. This is the alignment shift trend after sharpness reduction periods, i.e. Phase III and IV.
The above discussion reveals the alignment shift trend in single sharpness reduction period. We highlight that for the long-term behavior after many oscillation phases, this trend still holds. To see this, note that during EoS, of oscillates around a fixed value 2n, मズ. As for ||well", the analysis in Appendix B.3 proves its increasing behavior in Phase I. By proof details in Appendix (B.7 we know that the behavior of | in Theorem 4.8 (A) also applics to Phase II, which gives the increasing trend of ||v during Phase II and III. During Phase IV it may decrease for some steps but will have an overall increase after going through Phase III and IV, according to Theorem 4.8 (B). Then we know that the long-term trend of || 1 is to increase, which gives long-term increasing trend of ar and thus a long-term tendency to increase the alignment between Y and earlier NTK eigenvectors.
In Appendix A.6 we generalize the above crucial quantity az in Definition 4.6 to nonlinear networks.
5 Leveraging Central Flows
In this section, we present further evidence that sharpness reduction at EoS improves the kernel target alignment. We leverage the recent central flows technique [CDT+22), which models the time-averaged (i.e. locally smoothed) GD trajectory using an ODE called a central flow of the form:
dW dt = VL(W) + VEE(), VLL(W))
(3)
sharpness penalty
Here, (t) is a particular matrix defined as the solution convex optimization problem, and ( denotes the Frobenius inner product between two matrices (i.e. the dot product of the two matrices when they are flattened into vectors), so the quantity (E(I), L(W()) is a metric of sharpness where each entry of the Hessian V²L(W) is weighted by the corresponding entry of [t)). The central flow eq. 33 penalizes this quantity and is therefore a sharpness-per penaaiieee gradient flow. [CDT* 24 shows that central flow averages out the oscillations while retaining their effect on the time-averaged gradient descent trajectory, which takes the form of this sharpness penalty. They demonstrate that central flow is a good approximation to the time-averaged GD. See more evidence in Appendix A.7
In Figure 4a we run both GD and the central flow different learning rates in the fully-connected setting described in Section 3.1 and show that higher learning rates lead to higher KTA even under the central flow. Since the only difference between the central flow trajectories at different learning rates is the differing sharpness penalty/ this experiment supports our contention that sharpness reduction leads to higher KTA. In Figure (4b] starting at various points along the central flow trajectory for " 0.1, we branch off and run gradient flow dW de V L(W) for 100 units of time (red). We find that gradient flow takes a trajectory with lower KTA than the central flow. Since gradient flow differs from the central flow only by omitting the sharpness penalty term, this experiment further supports
'The prefactor in eq. 3 is just time rescaling and does not affect the overall trajectory, only the speed at which the flow traverses that trajectory. It is only via the sharpness penalty that 1) influences the trajectory.
9

---

our contention that sharpness reduction leads to higher KTA. Results under more settings can be found in Appendix A.7
KTA different learning rates
KTA on central flow vs. gradient flow
kemel taagee aligument
Nemel allanment
0.3
0.20
0.2
0.15
GD. 0.05
GD o_ 0.1
0.10
central fow
central flows
0.05
gradient flow
100
200
step time
step time
Figure 4: (a) The KTA under the central flows (dashed black) match those of gradient descent (colors). In particular, the KTA larger when 7j larger. the central flow at times {100, 300, 500), takes a lower-KTA trajectory than the central flow (dashed black). The experiments are conducted the fully-connected setting described in Section 3.1
(b) Gradient flow (red), when branched off from
Conclusion and Future Work
This paper analyzes the evolution of NTK on EoS by examining its alignment effect with the training target. We show that the EoS dynamics enhance the alignment of NTK with the target and provide a theoretical analysis the underlying mechanism this phenomenon on 2-layer linear network. There several potential future directions: 1. It is known that the evolution of NTK is related to feature learning. Hence it would interesting to connect our findings to more feature learning properties. 2. Although we provide detailed theoretical analysis 2-layer linear networks and partially generalize it to 2-layer ReLU networks Appendix AA6 a more general theoretical understanding of multi-layer nonlinear networks is still needed in the future.
References
[ABP22] Alexander B. Atanasov, Blake Bordelon, and Cengiz. Pehlevan. Neural networks as kernel learners: The silent alignment effect. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022, 2022.
[ADH 19] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332 222-332. PMLR, 2019.
[AGCH19] Sanjeev Arora, Noah Golowich, Naday Cohen, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In 7th International Conference on Learning Representations, ICLR 2019, 2019.
[AVPVF23] Maksym Andriushchenko, Aditya Vardhan Varre, Loucas Pillaud- Vivien, and Nicolas Flammarion. SGD with large step sizes learns sparse features. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learn- ing. volume 202 of Proceedings of Machine Learning Research, pages 903-925. PMLR, 23-29 Jul 2023.
[BDR21] Prajjwal Bhargava, Aleksandr Drozd, and Anna Rogers. Generalization in nli: Ways (not) to go beyond simple heuristies, 2021.
[BNS78] James Bunch, Christopher P Nielsen, and Danny C Sorensen. Rank-one modification of the symmetric eigenproblem. Numerische Mathematik, 31(1):31-44 1978.
10

---

[CD Jeremy M. Cohen, Alex Damian, Ameet Talwalkar, Zico Kolter, and Jason D. Lee. Understanding optimization in deep learning with central flows. In The Thirteenth International Conference on Learning Representations, 2024.
[CGMR20] Hung-Hsu Chou, Carsten Gieshoff, Johannes Maly, and Holger Rauhut. Gradient descent for deep matrix factorization: Dynamics and implicit bias towards low rank. arXiv preprint (rXiv:2011.137/71 2020.
[CKL*21] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, Zico Kolter, and Ameet Talwalkar. Gradient descent neural networks typically occurs at the edge of stability. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
[CMR12] Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research, 13:795-828, 2012.
[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
[DNL23] Alex Damian, Eshaan Nichani, and Jason D. Lee. Self-stabilization:: The implicit bias of gradient descent the edge of stability, In The Eleventh International Conference on Learning Representations, 2023.
[DZPS18] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterrized neural networks. In International Conference on Learning Representations, 2018.
[5KW+25] Avrajit Ghosh, Soo Min Kwon, Rongrong Wang, Saiprasad Ravishankar, and Qing Qu. Learning dynamics of deep matrix factorization beyond the edge of stability. In The Thirteenth International Conference on Learning Representattins, 2025.
[GMMM20] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? Advances in Neural Information Processing Systems, 111148820844,,, 2020.
[GWB+17] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
[HS97] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1442, 1997.
[HZRS16] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Conver- gence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[JMI.24] Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural network geometry? Advances in Neural Information Processing Systems, 36, 2024.
[KH+09] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[K120] Dmitry Kopitkov and Vadim Indelman. Neural spectrum alignment: Empirical study, In Artificial Neural Neiworks and Machine Learning-ICANN 2020: 29th International Conference on Artificial Neural Networks, Bratislava, Slovakia, September 15-18, 2020, Proceedings, Part Il 29, pages 168-179. Springer, 2020.

---

[LBD*20] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur- Ari. The large learning rate phase of deep learning: the catapult mechanism. arXiv preprint UrXiv:2003.02218. 2020.
[LMZ18] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over- parameterized matrix sensing and neural networks with quadratic activations. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet. editors, Proceedings the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 2-47. PMLR, 06-09 Jul 2018.
[MDP+ 11] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguisties: Human Lan- guage Technologies, pages 142-150, Portland, Oregon, USA, June 2011. Association for Computational Linguistics.
[OJMDF21] Guillermo Ortiz-Jiménez, Seyed-Mohsen Mohsen Moosavi-Dezfooli, and Pascal Frossard. What can linearized neural networks actually say about generalization? Advances in Neural Information Processing Systems, 34:8998-9010, 2021.
[RBPB22] Adityanarayanan Radhakrishnan, Daniel Beaglehole, Parthe Pandit, and Mikhail Belkin. Mechanism of feature learning in deep fully connected networks and kernel machines that recursively learn features. arXiv preprint arXiv:2212.1388/ 2022.
[SB21] Haozhe Shan and Blake Bordelon. A theory of neural tangent kernel alignment and its influence on training. arXiv preprint arXiv:2105. (4t:::1055.1301,, 2021.
[Sim14] Karen Simonyan. Very deep convolutional networks for large-scale image recognition. arXiv preprint 1409. 2014.
[TCLT119 Iulia Ture, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read stu- dents learn better: On the importance of pre-training compact models. arXiv preprint 2 2019.
[WES+23] Zhichao Wang, Andrew Engel, Anand D Sarwate. Ioana Dumitriu, and Tony Chiang. Spectral evolution and invariance lincar-width neural networks. Advances in neural information processing systems, 010220695-07722 2023.
[WGL+20] Blake Woodworth, Suriya Gunasekar, Jason Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in over- parametrized models. In Conference on Learning Theory, pages 3635-3673. PMLR. 2020.
[WLL22] Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along gd trajectory: Pro- gressive sharpening and edge of stability. Advances Neural Information Processing Systems, 35:9983-9994, 2022.
[ZLRB23] Libin Zhu, Chaoyue Liu, Adityanarayanan Radhakrishnan, and Mikhail Belkin. Cat- apults in sgd: spikes in the training loss and their impact on generalization through feature learning. arXiv preprint 2306v.2306.0488 2023.
12

---

A Supplementary Empirical Results
A.1 Experimental details
Here we would like to provide more details of our experimental setup.
Training resources The experiment the linear network (see Section 3.1 for details) trained on a CPU since the model is small. All the other experiments, including those described in Section and additional experiments presented later in Appendix A.2 GB of memory.
were conducted on V100 GPUs with 40
Dataset The main dataset use is CIFAR-10 [KH+09] As is described in Section 3.1 due to GPU memory constraints, we randomly selected about 2000 examples in classes and (roughly 1000 for each class) and relabeled the target values as +1. In Appendix A.2.3 we will use subset of the IMDB dataset [MDP+ 11] where we randomly selected 1024 examples.
Network architectures Here provide more details on the network architectures use.
Linear network. As described in Section 3.1 we use a 2-layer linear network with input dimension 400 on Gaussian data. The size of the hidden layer is also 400.
2. Fully connected network. We trained 4-layer fully connected network with GeLU activation. The input dimension 322 * 32 * 3, which the size after flattening a CIFAR-10 image. The sizes of the three hidden layers 128,128, 128,64, respectively.
Vision Transformer. We also trained simple Vision Transformer (ViT) using an online imple- mentation it
We use their "SimpleVii" model with detailed parameters patch_ size = 4, num classes 1, dim 64, depth = 4, heads 8, mlp dim = 256).
4. VGG-11. We also conducted experiments on VGG[Siml4], We use an online implementatior with the
__nn" architecture in their implementation.
5. ResNet-20. In Appendix A.2.2 we conducted experiments on a ResNe((HZRS16) with 20 layers and GeLU activation. We use Group Normalization instead of Batch Normalization in this ResNet.
conducted experiments on language tasks where we fine-tuned BERT-small (TCLT19, BDR21 from an online implementation
BERT-small. Appendix A.2.3
A.2
Experiments in more settings
A.2.1
VGG
On the same dataset as in Section 3.1 we also trained a VGG-11 network using GD under MSE loss from the same initialization under different learning rates to roughly the same training loss. We again plotted the evolution of KTA (every 20 iterations) in Figure 5a and the individual eigenvector target alignment for the final models in 5b where we see the increase of KTA and the alignment shift phenomenon. Similar to Section 3.1 we also observe consistent correlation between a fast increase in KTA and the sharpness -eepeessnde period of EoS (see Figure and provide quantitative analysis in Table 3 to show this connection. Similar to Section 5
we also study the above connection this VGG setting using the central flow flow|CDT+22 2CDT+224 in Appendix A.7 and demonstrate the benefit of the sharpness penalty to the increase of KTA.
A.2.2 ResNet
On the same dataset as in Section 3.1 we also trained ResNet with 20 layers using GeLU activation under GD and MSE loss from the same initialization. Again, we trained the model under different learning rates roughly the same training loss. We plotted the evolution of KTA (every 20 iterations) in Figure 6a and the individual eigenvector target alignment (along with its cumulative summation) for the final models in Figure 6b where we see the increase of KTA and the alignment shift. In
MIT license. http:///icchbbicmm//eeeraaattttp 1iffrr MIT license.
16. 20.22/e 2/ eeaniiriiinnnn///eeaennnnnnninit Apache-2.0 license.
13

---

021
Camgateo Suemetional Pr (paveter Barge Aaynnrr
606
180
(a)
(c)
Figure 5: (a) KTA across different learning rates, (b) the alignment shift, and (c) connection between the alignment dynamics and Edge of Stability in VGG-11 (see Section A.2. I). Similar to Figure Tb we plot the alignment of the first 20 eigenvectors in (b).
Table Comparison of the changing speed of KTA between the sharpness-decreasinng periods and sharpness-increasingg periods on VGG-11. See Appendix 1.2.1 for the detailed experimental setup.
Iteration interval
12-15
15-42
42-48
48-78
sharpness behavior
decrease
increase
decrease
gradually stablized
total change of KTT
0.1008
0.0163
0.0260
0.0083
avg change of KTA per step
0,0336
0.0006
0.0043
0,0003
Section A.7 we study the connection between the increase of KTA and EoS using the central flow CD 24
and demonstrate the benefit of the sharpness penalty to the increase of KTA.
Norma lzed oigenwal ubs
1w:023
Kernel Target Aligment
0.16
0,005
3.00
Ir 0.01
Indiiduaal Eigeweddor Torget Migementt
0.02
0.14
1-0022
b-00.2
0.12
Cumulot ve Sunvmatton of indiviial Elyenvector Agnnent
330
0.10
322
1000
2000
3000
4000
5000
120
step
Ogenvocttr inde
(a)
(b)
Figure 6: (a) KTA across different learning rates, (b) the alignment shift in ResNet-20. More details can be found in Appendix A.2.2
A.2.3 Language transformer
In this section, provide experiments demonstrate that results the alignment behavior also generalize to language tasks. We fine-tuned BERT-small [TCLI99] [ on subset (1024 randomly selected examples) of the IMDB dataset [MDP 11). The task is to classify whether movie reviews are positive or negative"] and the labels are +1. Again, we trained the model from the same initialization under different learning rates to roughly the same training loss. We plotted the evolution of KTA (every 20 steps) in Figure 7a and the cumulative summation of the individual eigenvector target alignment for the final models in Figure 7b Here, we see the increase of KTA and the alignment shift phenomenon, similar to other tasks.
14

---

Samell brget Aigrmenc
0930
HEAKES
675
15.422
6008
23
189
Cuimalatr Bwwiiaihhuuggrreeccoooo Eggrieecir Target Aigrment
6A28
02-
601
0.05
60
602
392
201
Exrracts hdm
(b)
Figure 7: (a) KTA across different learning rates, (b) the alignment shift in the sentence classification task using BERT-small, Detailed setup is described Appendix A.2.3
A.2.4 Multi-class tasks
Section 3 presents empirical results binary classes. Here provide more experiments multi- class tasks. More precisely, we randomly selected 2000 examples in all ten classes in CIFAR-10 10 and trained a VGG-11 I network using one-hot labels and MSE loss. We computed the cumulative summation of the individual eigenvector target alignment amd made a comparison after the same effective steps: 400 (or 200) steps for a learning rate of 0.01
200 (or 100) steps for a learning rate of 0.02, as is shown in Figure X Here we see that a larger learning rate has a larger cumulative alignment value, which indicates the aligmment shift towards earlier eigenvectors and is consistent with the alignment shift phenomenon in the binary-class case.
Curulative Bigenvector Target Alignment After the Same Effective Steps
Cumulative Eigenvector Target Aligrment After the Same Effective Steps
3.
3.
35.
3A-
0.2
0.1
01
3.0
250
220
230
200
52
100
Figure 8: Cumulative summation of the Individual Eigenvector Target Alignment values after the same effective steps VGG-11 with CIFAR-10.
A.3 Generalization to the SGD setting
This paper mainly focuses full-batch GD, since the dynamics of GD simpler than SGD, and the FoS of GD is well defined. In particular, for GD, the EoS period often consists of distinct phases (e.g., the four phases discussed in our theoretical analysis in Section 4 where the sharpness rises above and then falls below 2/(step size). However, for SGD, [CKL+2 21 pointed out that "the sharpness does not always settle at any fixed value, let alone one that can be numerically predicted from the hyperparameters" That means for SGD, it's hard to link the increase in KTT to the period where the sharpness falls, as there no clear division of the cycles and the so-called caarpness-decreasee phase. Nevertheless, we find that our observation that larger learning rates result in stronger KTA alignment still holds for SGD, We added the following experiments on the fully connected setting described in Section 3.1 with batch sizes 1024 and 256. Again, we trained the models under different learning rates to roughly the same training loss. Figure and and Figure T0a show the increasing trend of KTA where larger learning rates lead to larger KTA values. Figure 9b and Figure 10b plot the
15

---

individual eigenvector target alignment for the final models as well as their cumulative summations, where we see the alignment shift phenomenon.
hommiizzed gemaaaiis
Kamal lageet Algrment
625
0.5
025
Indeidual
033
636-
Cummlativv Summation
rell dall Exeeeeetor Theget Aliremet
0.5
Swch
(a)
Figure 9: SGD experiments in the fully connected network described in Section 3.1 with batch size 1024: (a) KTA across different learning rates, (b) the alignment shift. Similar to Figure 1b. we plot the alignment of the first 20 eigenvectors in (b).
epenvolues
Kamnel Target Aigrment
625-
0.2
626-
Indluitiol Eggereerter Burgrr Alcommet
625
2 620-
0.55
6.5-
6.5
20
309
320
140
Docch
(a)
Figure 10: SGD experiments in the fully connected network described in Section 3.1 with batch size 256: (a) KTA across diffcrent learning rates, (b) the alignment shift. Similar to Figure lb we plot the alignment of the first 20 cigenvectors in (b).
A.4 Comparison with the AGOP alignment in ZLRR223
Here we provide evidence on the claim in Section 2 that on the two-layer linear network setting, the enhanced AGOP alignment is less correlated with large learning rates and sharpness reduction in EoS than the enhanced NTK alignment. We trained the 2-layer linear network described in Section 3. using GD with learning rates 0.005 and 0,01 and plotted the evolution of the AGOP alignment and KTA conditioned on the same effective iterations (defined as the learning rate times the number of iterations). Figure IIa shows the evolution of these two alignment values during the whole training procedure and Figure lb zooms in the EoS period where we see the sudden increase of KTA but no significant increase of AGOP.
A.5 Connection to the generalization ability
proposed generalization bound proportional y K n where K is the NTK and n is the number of the training examples. This can be viewed as an inverse version of the KTA. Here we focus on the 2-layer linear network setting described in Section 3.1 and plot the evolution of the normalized y K -y. We demonstrate that the sharpness reduction period of EoS and the alignment shift phenomenon towards earlier eigenvectors are correlated with sudden decrease of the normalized " K-1 y, see Figure 12a Figure 12b demonstrates that larger learning
As is discussed in Section 3
ADH+19
rates lead to more significant decrease of the normalized y K y than smaller learning rates.
16

---

376
0.0
(b)
Figure 11: Comparison between the AGOP alignment and KTA in a 2-layer linear network described in Section 3.1
We also conducted experiments the fully-connected network described in Section 3.1 Here, notice that the tail NTK cigenvalues are much smaller than earlier ones, as is shown in Figure T3a which make K- "noisy" and cause numerical issues. Hence add c/ to the original K and then calculate the inverse. More formally, we compute the following "inverse KTA":
y (KC+))
[k+]]]((K+e] "IF"
We set e = m..0005mmax The top and middle subfigures of Figure 13b plot training loss and the evolution of the sharpness. The bottom subfigure of Figure T3b shows the change of the inverse KTA every 3 steps, where we see negative values, and that means the inverse KTA decreasing over time. By comparing the middle and bottom subfigures, we see that the carppess-ddecreaiin periods correspond to larger decreasing values of the inverse KTA than the sharpness-incree increassin periods.
Evolition ofthe Inese KTA
Expston elta depn
140,000
0.314
3014
3000
0.02
2002
Busdion
Taraton
(a)
(b)
Figure 12: (a) Evolution of the sharpness, alignment between the target and the top 5 NTK eigenvec- tors, and alignment between the target and inverse NTK. (b) The evolution of the alignment between the inverse NTK and the target under different learning rates.
Now provide empirical results direct connection between the alignment shift phenomenon and the generalization ability. Figure 14a and [4b] show the alignment shift phenomenon under different learning rates where oscillation and a larger alignment shift towards earlier eigenvectors. To compare the generalization ability, we measure the decrease of the test loss in two ways.
that larger learning rate leads to significant sharpness
1. Measures the test loss Liest after the same effective iterations, i.e. for different learning rates "I,..., Tu. train the network for ti,..., fn iterations such that 7 1222
2. Measures the change of test loss A during the oscillation period when sharpness decreases and the alignment shift phenomenon occurs. For example in Figure 14a we compute the change between the 55th and the 75th step, and Figure 14b compute the change between the 10th and the 20th step.
17

---

Trunng
104
Cartirtie ofe
16.
Banserr
113
100
100
16
300
1233
1533
7033
132
Coondraaa Indee
Ianster
(a)
(b)
Figure 13: (a) NTK spectrum in the Culnecco network described in Section 3.1 after 60 steps. The spectra for other steps have roughly the same shape. (b) Evolution of the traningg loss, the sharpness, and change of the inverse KTTA every 3 steps.
Traning
Triinin
Peration
Coppess
Scheosn up digeneect'' a
Algrneetbbeuuenn to 5 oigenveceess
0.2
10
20
30
50
60
30
15
23
50
Benabor
(a) Small learning rate
(b) Large learning rate
Figure 14: Alignment shift under different learning rates
As is shown in Table 4 larger learning rate leads to smaller test loss after the same effective steps, as well as a larger decrease in the test loss during the period when sharpness reduces and the alignment shift phenomenon occurs.
Table 4: Generalization ability under different learning rates
Learning rate
0.8
0.9
1.2
Trest after the same effective steps
1.211
1.206
1.199
A Lesa during the oscillation period when alignment shift occurs
-0.008
-0.020
-0.062
A.6 Potential Generalizations to Two-layer Nonlinear Networks
In this section, try to generalize crucial quantity in Definition 4.6 the two-layer network with ReLU activation. We again use the MSE loss L(W)- I the elementwise ReLU operator. By simple calculation, the NTK iteration in this case given by
2n ( x) - y| where a--)
K; = K(") + k!?,, where K.d) = (x x), K!² D. Dr.
(1)) (w p.p
with D
w! ( [diag p {
18

---

We first demonstrate that K( o xy² (W; X) still exhibits low-rank structure where it close to rank-2 matrix. Here we train 2-layer ReLU network and plot the eigenvalue distribution of K/") after and 40 steps in Figure 15 matrix after only steps.
We can see that it quickly becomes an approximate rank-2
Expensiic ESTRLDOO
Coondrate
Figure 15: The approximate rank-2 structure of o (W!"x)y o(w!")) x) in two-layer ReLU network after (a) steps and (b) 40 steps
This approximately rank-2 structure allows us define a new version of oî as the leading eigenvalue of K,²) and ||æ, 112 as the summation of the top 2 leading eigenvalues of K<(1) and calculate the corresponding Cty as
04 | XI(K) )
(1(K(
Figure 16 compares the evolution of the sharpness, our new version of on, and the alignment of Y with top NTK eigenvectors (denoted as W1,... us) between small and large learning rate settings. The setup is the same as in the linear network setting described in Section 3.1 except that we use ReLU activation now. Here we see that in the large learning rate casc, the aligmment shift phenomenon happens towards carlier eigenvectors during the sharpness reduction period, and in the meantime, there is sudden increase of at. This is consistent with the behavior a and the alignment trend discussed in the theoretical analysis in Section
Evoludon
10
Abgrrwet betoeer tpp aquectes
Algrmmntt belnewen ebpaetton
0.5
10
12
20
20
30
40
35
25
33
35
Heration
(a) Small learning rate case
(b) Large learning rate case
Figure 16: The alignment shift and evolution of a in two-layer ReLU network
A.7
Supplementary Results on Central Flow
First, we provide evidence demonstrate that the central flow is good approximation the time- averaged gradient descent. On the fully-connected network, VGG-11 and ResNet-20 described in Section 3.] and Appendix A.2 we run both central flow and gradient descent under different learning rates. Figure 17 plots the evolution of the training losses, sharpness and the KTA, where we see that central flow approximately matches the time average of GD trajectories. (For the dashed lines the training loss subplots, report the "central flow prediction" for the training loss (CDT+24) 24), not the
19

---

training loss along the central flow.) For the VGG-11 in these experiments, we use average pooling (rather than maxpooling) and GeLU activation (rather than ReLU), so that the training objective smooth, as is needed in order for the central flows to work well.
loss
eigenvalue
targett
0C-
1.0
0.75
GD. CA5
6B: 6.1
0.2
0.25
0.1-
flwx
200
-00
000
200
400
622
200
400
602
des time
step /ime
5060 time
(a) Fully connected network
top eigenvalue
/ tanget alignment
500
60.4-0.30
03.0-0021
0.4-
200
CD. g-0.22
0.2
100
central fewm
0.0
(b) VGG-11
trah loss
Hess lan eigerivale
Karnel target alignment
1.25
432
1.00
60.0-0
0.72
60.0-0033
0.50
02.000022
central
0.13
Umn
atop U0m
Unn
(c) ResNet-20
Figure 17: Demonstration that the KTA along the central flow reasonably matches that along the GD trajectory under different network architectures. In the subplots the training losses, the thin/light solid lines are the actual training loss curves under GD, and the thick/dark solid lines are the smoothed (tim--averaged) version.
We now provide more results on the evolution of the KTA under central flow compared to gradient flow. Same in Section[5] in Figure [8] starting at various points along the central flow trajectory in different architectures, branch off and run gradient flow dW V for some time (red). Again, we find that gradient flow takes a trajectory with lower KTA than the central flow.
B Detailed Theoretical Analysis and Proofs for the Four Phases
B.1 Cartoon Illustration
As is discussed in Section 4.2 the training dynamics at EoS can be divided into four phases based on the position and the sign of change in the sharpness. We plot a cartoon figure to illustrate this division in Figure 19 A similar illustration figure also provided in [WLL.2]
B.2 Rank-1 structure
In this section, we verify that the structure in Assumption 4.1 will be preserved for to where la defined in Assumption 4.1
20

---

m 0.05
11
0.3
0.2
0.30
0.23
0.25
0.225
0.20
0.315
0.10
0.15
0.250
centra al How
gradient Ton
central How
gradient Fox
0.10
0.05
gracient Ron
central Row
800
/imm
step time
step time
(a) Fully connected network
0.005
0.01
0.02
5.34
6.12
0.125
0.125
0.100
0.100
0.075
central gracient
0.80
central
0.050
5.02
420
0.225
gradient
step
(h) VGG-11
0.005
Stange algrnoo
5:16
0.15
0.13
0.11
0.11
1200
2000
2000
4003
5000
500
1000
1500
2000
2300
3000
500
1000
15:00
stop
stop
(c) ResNet-20
Figure 18: Gradient flow (red), when branched off from the central flow at different times, takes lower-KTA trajectory than the central flow (dashed black).
Phase
sharpoos
Figure 19: A Cartoon illustration of the four phases during EoS
We prove it by induction. Suppose at time after to, we have that w/ x uo?" w(9) cut. Then at time t + 1, by the update rule of wr(")x X and wx(2) we have that:
with) x = (1) x
E,X = uo? G = =(e 2 (2)T
T -EXX"
n
( w (2) " (,,.)) x) = c/u " E =(a (Fe, vi)
Hence see that the rank-l structure is preserved with unchanged and
C-1 (Eg v.) 1/1
T _ e FxXX"
This completes the proof by induction.
21

---

B.3 Analysis of Phase I
Note that under the rank-l structure in Assumption 4.1 the NTK at iteration & can be written as
K2 I 1 X W.
x x x + iara
We want to prove that during Phase I, both of and | increase over time. We first introduce another assumption on Phase I where we assume that the dynamics follow the GF trajectory. This assumption was also made and empirically verified in [WLL22]
Assumption B.1. During Phase I, the dynamics follow GF trajectory.
W Recall that we denote F; = w ( X= C, under the rank-1 structure, and E, = Fi and use {g: Jiz represent the eigenvectors of K2 X Then plugging in the update rule of c and W; in Section B.2 we immediately get the following results.
d "1
(Fr, - ve)
Ct
dt
n ? =,,E, X0GxxxXx=- _nd a E,,g)) n
d
(v
dt
n
By Assumption 4.1 we have that at time to. Ct > 0 and E - for all S n, which gives us that (Fto q:) = (Y,, gi)) for all i S n. That means d (v,, Tt) has the same sign as (Y,4) at L = ln. which promotes (we, q:) to move towards (Y,4), direction deviating from (Fa, 9:) since (Fr, q.) has the opposite sign as (Y, (.) at to and for short time period later.
(Ei, Vt,) < O, then c > 0 at = la. which helps keep the positive sign o for a short time period later. On the other hand, the behavior of (Ee, vz). Repeating this procedure, we will get consistent increase of ot and consistent movement of (V,, q:) towards (Y, qi) for a short time period after to.
a (v,, qs) will help keep the negative sign of
If (Et, vta)) > O, then de ct < 0 at lg. which makes e, decrease for a short time period later. Note that in this case, the behavior of di (ur, q:) will promote (E,, v,) to decrease and become negative during the following training time. Repeating this procedure, we have that:
1) If (Ee, vz) becomes negative before o changes the sign, then e will start to increase and keep positive and we will go back to the case with (E,, v2) < 0 and C >> 0.
2) If e becomes negative before (Ee, vx) changes the sign, then (wr, will start to move towards -(Y, the same sign as (E,, which will prevent (F4, vx) from decreasing and changing the sign. Then (Fe, v2) will keep positive, Ct < and will keep decreasing, which will in turn promote (v,, %) to move towards - (Y). If we rewrite -c,- D;; et and U, respectively, then we essentially get the same dynamics as in the case with (E,, v.) and e > O.
In summary, the dynamics will eventually move to a stage where we have an increase of e and movement of (v,, 9:) towards (Y, g:), which has the opposite sign as (E,, q)) Now we assume that at some time point E > to in Phase I, for all < (Dr, q:) will have the opposite sign as (Ev, qt).
Assumption B.2. Assume that at some time t' in Phase I, cv > 0 and that (ve, qi) has the opposite sign as (Ev, q.) for all < n.
The following lemma ensures that for later training time /. >>>>>>>>> L' in phase I, before [ [E,,, g:)l becomes very small, it always holds that e > O and that (w): q:) has the opposite sign as (Fe, q:) for S n. When (Ee, q:)l becomes very small, it already implics convergence of E, along qi direction. The detailed analysis is beyond the scope of this paper and has been studied in [WLL22] They show that we will either enter EoS or have small | which implies convergence.
Lemma B.3. Under Assumption B.I B.2
for 1 l' with l' defined in Assumption B.2 For i S when K (Fr, 4)l is not very small, we will have that ot >> O and that (we, qi) has the opposite sign as EE, q.). Then we immediately get that (Fi, q.) (c.Ve, 9i) has the opposite sign as (Ei, q.)
22

---

Proof. We have that
d
" ni ((), " (E,,, (,) 9:)
[E 91)
dt
n
d
Ct (Ei, vz)
di
d (ve, 9i) - FF Xq:
77,3 ct Fe, 9s)
dt
Suppose for some T > 1' in Phase I such that for all t E 11, > 0 and that (E,, qi) doesn't I change the sign for all i S n. By the definition of t', we know that Vi S n: (ve, qi) has the opposite sign as EEg, (t). Then the dynamic of (vi, qi) gives us that (ve, 9i) has the opposite sign as (E,, 9:) Then for E [T, T + dt] we have that
fort E [t', T). Hence (E,, v,)
cc+dt C (Fe, v,) > 0
and that (Eg, ue) 9:) = dr,(EE, 4) for some d > O. When (E,, q:)l is not very small, we will have a bounded di which yields
(b) = 77 (c) + (Ee
dt
For E (T, + di) we have that
11: (c2 di, 9i) = exp
(- yyi (c? + dr.) (Fe, gi)
(Fe+de, (6) = (Fe, 4))
n
n
which has the same sign as (Ee, q:). Processing the above steps completes the proof.
Now notice that we have the following dynamics for oj and ||0.112.
d -
21
2n7 Σ (E, q.) q.)
(c))
2cr dt
=
EE, Fi) =
dt
|| 20 T d
12 = 27 Σ Ae (FF
di
dl
Then Lemma B.3 implics that of and 1/ 2| are increasing for > t' in Phase I before KEE, q.)l becomes very small. As mentioned above, more detailed analysis in [WLL22] gives us that we will either enter EoS due to the increase of of or get convergence before EoS.
B.4 Intuition on the choice of the learning rate
Let's consider the time period after e 2211 //inn reaches (1/n)) (denoted as To). From the proof of Lemma B.3 we have that with di, 0,
d
EE, qi) = nd (e + de.)) 9i) < 17 2 (Et, 9i) = - d (Ee, qi)
dt
X1
Then after T = O log n) steps, (Ee, qb) will shrink by O (oly(n)) Let's consider some index k such that A&- = (n for some small > Note that for S W,, X, On---a-s). hence after T = O(n''an'--- 1-a-t log n) steps, (E(+1, (+),,.9), < ke' will shrink by O poll(n))
That means after T' compared to those i > k.. Moreover, by our data distribution (see Section 4.1). we know that defined in our data distribution), which gives us k - K = Olog n) and hence the number of terms (-E,, (F) q) for k' < < k will be negligible compared to the number of terms when k. Combining the above two facts yields that, Σn X (x(-F,, u.) (F,,g) will be dominated by Eisk dr(- Et, q:) (F., Therefore,
steps, (-E,, q:) (Fi, qi) for i < k2' will be negligible
k' = - n) and i (logn) (where
d
2n
Xr- E, (. = 21
हi 11o.112
0
n
n
i=1
S
2n n
E,, q.) (Ft, gi)) 0 na
(c)1)
n
dt
23

---

Denote the end of Phase I as Ti. Then we have
| 11 - | @ n2 )cc o - = O T2"h
)
"/
Here we check that To + T will not exceed the length of Phase I, Ti. Note that
d. (c:) (c))= - 21 (E,, ( < n )(111))) = O())
dt
Starting from To Ti, of has a total increase of order ( which means that Ti To + S (
By the choice of learning rate discussed below, we have that s7² @ ( and hence Ti To + M(n'=), which is at higher order than To + = To + O( logn).
The above calculation implies that the increasing speed of | is O x times the increasing speed of os for later period (after T + To of Phase I). If we further assume that for early training period, we have ||UT++T|? = O ワ ||vell /n))
Choice of learning rate
then we can get that at the end of Phase I,
Note that at the end of Phase I, i.e. the beginning of EoS, we have ci = ((1nn). then for the output norm, we get that || = of 21,11112 (nn /r/)). When [e|] goes to order of ||Y we want it to go pass the above order (nn /n") so that we will have EoS. (Otherwise F₃ will approach Y before reach the order O(n" /m²) and we will have convergence before EoS) That means we require
n² n²/12² S ((II1 On)), +
1/² = 22 n
So we can pick n²² = 0 (,,+))
B.5 Dynamics of of
In Phase under the gradient flow, we prove that the time derivative of ct is determined by the sign of {E,, Fi). Now we prove that in Phase II and later phases, even though gradient descent no longer follows the gradient flow, the sign of c - of is still determined by the sign of (E,, Ft). Note that the update of &f under gradient descent is given by
2m (Fu, Fi) + n²
2007 (Fr, vz) + n² (Pe, v2)²
2++ - ot
(Pe, wz)2
n
n
" (E ve)
(") (Ee, vz) -20z) =
1 (Fe, vi)
11 (cGe - Y, v.v)) - 20x
n
n
n
" (E,, vi)
(ce
1
|0 -2)
(Y, v2,2
n
By Assumption 4.3 we have that c2|1 2| = ||Fs| 1 < E + 2| < O(+). Combining with
the fact that of ( during EoS yields that |/ < O(n). Then we have "
n |1,2, - (1)
and that (Y, u.)1 < ||||||| || < On)) Hence c (3 1/,11 - 2) (Y, =1), which
implies that the update of of has the same sign EF,, v.)c. -(E, Ft).
B.6 Analysis of Phase II and Proof of Lemma 4.4
First note that during Phase II, we have that (E,, Ft) which means that [ < |||| because otherwise, we would have that (E,, F2) = E + (E, Y) 2 12
E
| > 0. Similarly,
we have that Fill S || Hence we have (E,, F.Y. & O(l = (n))
We prove Lemma 4.4 by induction using the following two lemmas.
Lemma B.4. Under Assumption 4./ 4.3 during Phase II and pick n = 9(n- Write A, ==
A c
2 with A, > O. Suppose for S t - 1, we have (Eg, F.)| = (52 (Yl?), we can get that
17
KE,, 2 (1 (E -
[F, 9:)l 50 6212 d1 (.
(Fm. + 2 25ign
X,
where T1 is the end of Phase I.
24

---

Proof. During EoS we have that &t = ² Combining with \ = (n) assumed in our data
distribution (see Section 4.1 gives us that of 2 For the update of (Er, qz), we have that
(E,, 9b) (Ei-1, 93) + (Er - Et-1, 9:) ( qt) + (are a-100-1, 9s)
(E 9t) + (ct - -c 4) + c-- (o) 11-1.%) + (c+ ct 10 02-1
" (E U2 0/- (0,,)))
1/ 2 1 (E)-1,4))
(B)-1, -
n
n
^/
+ n²
(E 1)-1){ (1)
(Ei-1, - π
(E Fi ((-))
n7 c (E
n
(Br-1, 1, Pe- ()
n2
77
7
1E-11 Fi-1) +
d2 (Ei-1, P(-)) (EI-1. 4i)
n
(-1
ne
2
n²
1E--1,
Fr-)),,
no--1 1
(4)
For i = L, we have that
7221 + +nde-a) (Ba)) (E-1.91))
(Ec, q1)
1 1--2-122 - n²(2 + -1
(Er-1; P.- F.-1)(X,9))
ne
(i)
-1 A
11 + n
(B1-)) (E-_1. .n)
n² + 11)
112
Fy- (E
where (i) uses o 2 + A:-1 and (it)
uses d n c = =0()) and a = ((n) and the fact that
during phase II, (E,-1, F1-1) < 0.
Then by the assumption
(Fa--1, Ft-1)1 1)1 (622 (ō2n), it immediately follows that 91)1 - 02
(( 2 (1 +
Now we bound [EE,, q:)1 for any S with i First notice that and d² ((E,,1,F--1)) = o (8 3) = (62n). then we can merge these two terms together and write
gd o2 0(22)
12
|| = O e( 126222 by the assumption
(E.-1, Fo-1) S
n²
X1
nd
a n²
(E, F. -
(5)
A
n
Substituting into eq. and replacing by s give us
(180)
A A1
112
(E,, g.)l
1- θ
+
(E.-1,,F-1)) (E
n
172
+
KEE,1, F. F-ii)lK
7
A
+ 0(776))
(E ((76)) q))l
A1
A
1- 0
[E + +07
A1
25

---

where (ii) substitutes the order ci-1 = (1,,, (dv) uses (E,-1, & ((gnn again and (v) is
because under our data distribution (see Section 4.1 and the choice of learning rate 1f O ((= ) we have 1/ O -- 52) - 62) of (+ Telescoping this from Ti to t completes the proof. (
Lemma B.5. Under the conditions of Lemma R.4 we have that at time t, KE, Fi) 0
Proof. The upper bound of (Ee, qe) for i 2 2 in Lemma
B.4 can be further upper bounded by
o 82nl
A1 A
X
M + ((E,,,)))))
0 02 Vx, Vn!+ + +1E,, (Bt..ni))
(0)
n
||
Σx x Ak
which gives us
Σ (Fs 11)
- Σ
|| + (Fr, 9;)²)
Xinl+a
1=2
0
(8 n2² 2a Y 1 + || Err, 1| - 0 0
where the last equality uses 17 = O(n 1 and that Er; || 53 | -
By the analysis of Phase I, we know that (Ec.9:)) for 2 < i S k shrink faster than those for i > k in Phase Hence Ti, we have that
A
((ET, 10:=00 0 t1 log n
Ar ai)² k i>k
ET, ..)aa))2² = Ollggn) i>k
Σ EET,
1=2
i>k
Therefore,
Σ Xrn, S Eollogn + (xr,, a1)² 22==a=a)=00 O(nn 112) = 0663 -0(6
Then for time in Phase II, we can get that
82 n2 n-+a
Σ (:) <
| + de A<(ET,,,))
(6)
i=2
O (52.1 22-aa + hnn || 0635n "|
We have that 81 0
(V
A1
O(n a/2) and that sin(Y, q1)
Σx Au
In the subspace spanned by q and Y, choose y = Y/[[ and و as orthonormal basis. We have that
|| - 2 = 2 cos(g) ()) = 2sin Y) = O(6))
Hence we know that
(F,, 2
(1, 2(E, (1) - D) (F,, 91 -))²
2 (Ee, 91) (E,E,, g | E /1 -
اl1 - Fr /| || - 211
2 (Ei, )1
2 (E,, ) 12
(E,, 812 < | (),, < ||E.| (1
Σ 4)2 + ((
1=2
< 0 111))
<0
In the subspace spanned by F, and Y, denote Y and én as orthonormal basis. Note that F, E. + Y also lies in this subspace. Since (E,, Fi) < 0, by geometry relationships between E, Y, F, (note
26

---

that they form a triangle with cos Fi) < 0), we know that A. & which yields that
( (2 P) S ()I, 00211
Now we are ready to prove Lemma 4.4
Combining Lemma
B.4 and B.5 together, and noting that
the beginning of phase II, we have (E), S E
Fel
[ Y S ( by
Assumption 4.3 we can prove by induction that Lemma 4.4 holds for t in Phase II.
By the proof details of Lemma B.5 that in Phase II,
we know that (E,, Y) S ((22l|Y|))
II). This allows us to prove
CYs(Y, vi) = cYs((YFF) (FI, p
(Y + (E,, Y)
(E Y) |
2
21 - 0(62)
||
|I
B.7 Analysis of Phase III and Proof of Theorem 4.8 Part (A)
Since in Phase III, we have (E,, Ft) < O, then by the dynamics of of discussed in Appendix B.5 we know that 2:+1 < <...
Now we analyze the dynamics of ||,,122 We start by introducing the following lemma on the dynamics of (E,.q)).
Lemma B.6. Under Assumption 4.1 4.3 during Phase III and pick = O(n 0nn="")). Write A :=
d & 2 with A, > 0. Then we have that
KE 91)1 2 (1 (L -
) ( \11
11
(EE,,9)) 1- (()) -
25isn
where &; = 3 (F,, F.).
Proof. The proof of this lemma is similar to that of Lemma B.4 In particular, eq.
4 still holds.
For i = 1, we have that
(E.,g) = (1-2 21-2-nA--1 A2 7)--1)
(2 + +7,11) (EI-1, Fi-l) (Ec-1.91)
" (F,-1, (Y.1)
net
(-1 111-1A A + (m)) (E,-1. 91) - -9 (1") (E, ((-),,
where (1) is because under Assumption 4.3 we have that |(E,,1, Fr S ( and thus
Aa + +),,1)
1/111
n 22 1)) [(F,-1, F F1-1)
O())
By choice of A,-1, have 774-1 N(poly(n)). Plugging in this and the definition of A gives us that
(E,,g)) 2 (1 7 - -1 -14
2 (1 + -
| F-1) (0<2 (82n), by ولا Assumption 4.3 we have that (Er,,, 0 22 O
For 2 2, although we do not have the assumption
which is at lower order
27

---

than 2 O ). That means we can still merge these two terms together and eq. 5
still
holds. Then it follows that
A
(zx)
(E (Bin) ((E-1, +
(zi) (E-11, F(-)
[1²
(E,, (1,l s
1 - θ
\1
A
θ (d/1)) ((E-1, + (X,
< θ
\1
Note that in the last term of the first inequality, we use (E Pi-1) > 0 in Phase III.
The first inequality in the above lemma tells us that (Fe, 91)1 keeps the increasing trend as in Phase II. When it increases to [[62[|Y]]), we have the following lemma.
Lemma B.7. During time intervals in Phase II and III when ||E,|| = O(V)) and D 2 (1))
denote T == inf (EE 91)1 > 52][Y[] and 6 == 1(E,, l))l/||Y|I ||. then we have that for T, (E q.1l < (8em² Aa (Y, 9 + KET, 9/1) where T1 is the end of Phase I.
Proof. We first prove that for t > T. (E,, g1)l increases monotonically, which gives that 8, 2 62. To see this, note that both Lemma B 4. and Lemma B.6 us (E+,1, qi)l 2 (1+C2nA,) 91)1
OIIAII ())l with A; ي²
n (E,, Fi). Suppose at time t 2 T, T, 6, 2 82, then we have KE,, q1) 2
|2 2 KY, by 1(Y, [1)) = 6, || and 82 2 v 51. Combining with the condition A, 2
2 ((v51) gives us that C27/21/ (Ea, q.)1 2 (1 (Y, 2 (V by AJ s o 5) S
(r)) Hence ((E++,,q)l I 2 (Ee, g1)) and 6,+1 2 82.
Now we prove this lemma by induction.
Let's first prove the base case t = T. For t < if is in Phase II, then we know that Lemma B.5 holds. If is in Phase III, we can modify the proof of Lemma BB5 get that Lemma B.5 still holds for t < T. To see this, note that the only part we need to modify is the part using the geometry relationship between the vectors E,, F2 and Y. The geometry relationship
(.), S
(E, (>) doesn't 't hold in Phase III because now we have (E) > 0. However, the inequality (E q.)² < (66 || 12) still holds. That means when t & T, i.e. when (E,, 91)
62 we have that || = Σx Lin2 (E,, q:)² 2 < 0 [1a and hence || E& < (62)) (Y]) Then we get that ||FF|| < | + |Y|| < O([[I) and that (E,, < Fll.F < (bg))lY)))
That means for I T, specially / = T, the proof of Lemma 4.4 still holds, and we have that KE,, a3l S Od2n?² X. (X, ((ET, l)) Note that we have proved 2 62. then we prove the base case 1 = T.
Now let's deal with the case when > T. Suppose for T < t, we have that |EE, g. S O8 A1 [Y, q.لl + (),,)))).. Then applying the calculation in the proof of Lemma B.5 we
get that En (E., g:) < (6 which means (E.,qi>. > 2( (E..g))))))))) 9))). Hence for
T < < t. |. = ((EE,,g)) < (O((E.,9))) = ( which implies ||E,||
(IIIIII) Note that
|F_]|| = O(v)) =0([|]]), we have that [FFI1| S || +F.1++1Y11
Therefore, KE,, F.)l
< || S (b.)..)))) Then at time t + 1, imitating the calcu-
lation in Lemma B.4 we get that (E)+1, I $6inn X (Y, 4:9l + ((E7,.)))) Note that
8.17 A (Y, u))l
ge)l 2 (Y, S o KErng n 22 IY |
066.+117 KY, ())l.
Then we have (E 9
< ( X g +KEE, 1) & 2 10Y, ( + (Er, 9))) Hence we can complete the proof by induction.
Now we ready to prove the first part of Theorem The goal is give sufficient condition of A to ensure the increase of || 1|2. Substituting the results in Lemma B.7 to the calculation in Lemma B.5
we know that during time intervals in Phase II and III when E, = ((v)) and A 2 (2( 81). for t 2 T, we have (E, < (1)))) EL X, E, 9.)² I 0 Y and that E (.) q) 0(rr7 (1) ((
28

---

Now we start analyzing the update of 11211 Recall that we denote K₁ x"X X. We have for |
27 K, + T²
11++1112 - 1|00012
nh GEKKET
2n
21 E, + 2n Aun
E.
E ETT
n
A1
2n A.12
n
X2 (Ei, qi)²+
Aup A1 (E,, q1)²
2n
Σ Ar qi)²
n
2n E,K,YT
n
n
n
n
We want to have that
Ard \1 (E,, (E,,91)²
27 E,K,, > 0, which means that Li Nr (E)
" R Li=z 12
27
(B) 2 (E,,
(1), - 2 Σ (X(,9)) >0
1=2
1=2
When < & (n) and 12 T. we can use Lemma B.7 and get that \ A = 06
Jo. Er (Y, where we use 1(Y, = 81 |Y] and o2 2 Vo1. We also get that Σi=2 Ar qi)² - O6fFn")) || and Ei= Ar,, 9:) ()) 9)) - 0(63n))11 ||
Y1 (E,, q 2 1
Then we need lower bound of A. A. max {S2( voi ) (mi max (S pm074 9,
m/ mm--a) )}
When |Ed| = ((2)), Lemma B.6 telle us that for 2 2.
A Au θ (ALi)) 1(E,-1, 0(4
[( S 1 1 θ
$ (1-
(EE-1, 9111 + O
Denote the start of Phase III. which is also the end of Phase II, as T2. We have proved in Lemma 4.4 that [ET,, 93)l S O(52n² X, KY, 9))l 1 + ())) = o(KYY, q:)l) by our data distribution (see
Section 4.1 and our choice of n. For > T2, i.c. in Phase III. |E q:)l may increase but will
not exceed the order q))). This is because once KE,, 9.)l 2 2 ( 4))l such that RHS is smaller than (E,-1. then we will have (E,,%)
1)) < (E Hence we will have a stable ((((Y,.)))) Hence E 1i=2 (E2,g)² ((1Y)) = O()). In other words, when ||E = =n()), (Ee. )2 dominates E (E, 9.)2. More precisely, denote T3 as the first
upper bound (E,,
time in Phase III when 12 > > ( 9.) Consider time L such that / 2 Ta. we get that (E,, 91)² > Σm (E,, )1 Then we have (E,, F2) E O((E,, 92).. Now for s in
Phase III such that s S f, by Lemma B.6 we know that for 2,
((Fs, 9:)l S
1 1
θ
A
(A'-1))
1(F.-1,
A1
A
1 1-
1(E.-1, a1 + +044
A
Telescoping this inequality from time defined in Lemma B.7 / yields that
maXTs A 11
[E,, 931l so
1YY + (Kx,
A
maxT<ssi (Es, 2172 X2 1(Y + +(ET, 9i)l
nd;;
(1)(FE))))>>) nd; 00 ( (1P(PF n 1(Y,a))+ KET, ( = 172 (Fe, ((Y, + + ((Er,
d
Analyses before already give us that Ei, AEEr qi)2 Dmn {x (Y, S (1 Then we have that
S 0(8 and that
N (E, 9.) (Y, q)) 063 + (O(PEE..)))X
s 00 -
(E.,
29

---

where the last equality uses n7² = O(n (0(n((1-a) (1=a)) and (E,, q1)² = =S((n). We also have
"
(Ee, 91)*
Σ aE
Σ
(Y.
X
1
= =((3x")|X1 (E, 91))
Σ
(Y, 9:)
1=2
||X111
- 0003))))1
- (1)
Σ
ΣE An
122
+ n /1 (Be
<
nlta
(10
0633n 01000=)|111++000000000000 2- ( q1)221))2 S (Ee)
where (i) uses Assumption 4.3 that (E,, q1)² & | - (n//n) and (ii) plugs in 12² = O(n (((=((-a and (Ee, q1)²² = Q(n).
Combining these two cases together yields that
9:2 + Σ Ar ()) (Y,g) gi) s O(mxx(n 1))n)) = (F1) (00)) 91)
On the other hand, (1, (1) .n) = 61A1 2n [I]] S (E,, q1)². That means we require
that A 2 Ω maxx {1, 61 = ma mm=/s/ )}
Combining all the bounds of A, together and noticing the choice 1) O(n - ), we can get that A, 2 max M ().
B.8 Analysis of Phase IV and Proof of Theorem 4.8 Part (B)
By definition of t and tg, we know that they are the turning points for (E,, Fi) to change the
sign, which means these two time points, we have (E,, Fi,) RJ 0 and (Ee, Fx,) so 0. More
precisely, by the analysis in Appendix B.6 and Lemma B.5 we know that E. S o{ and
(E Fu < (62 ( (62)) Fr, 2 The final part of Appendix B.6 also implies that
(Y, v. 2 (())())) (2). Under Assumption 4.7 we know that at time tz. II | || || o(|]]])
and hence (E,,, Fi,)l < O ( (2) = O
(o 12,12 = 0(62) F 11 By a similar argument in
the final part of Appendix B.6, we also get that cos v,,) 2 1 O(62).
Note that | F|;
1= = || + 1 F, Ft). we have
Fr 1|2 = ||81| - || + 2(,,,, Fi,) 2 | E, 11? + 2(E,, Ft,))
|/,,, 1|2 - 2(E,, 1T2) 2 B(,,, Fr,)
FF, 11 (18)) [F, 112 (1 (12))) 2
Hence A,, 2 0 for some constant C 0, we can get that
A1
X1
X1
41 = A., 2 2((2)) X1 c7
>
1/2
n
and hence
= 1
+ 42
2 S (1 (6
which gives us
| 2
> - 062) F 1 1 166(6)) ((02)
||
1|
vz 11?
1- 2()2)
If the constant C' in the inequality D., 2 0 is sufficiently large, we can ensure that ||
||,,,1? 12 and therefore 0sg ass
30

---

Aligoment siggwectr
0.9
0.8
Abelule valte of co
0.7
0.
0.3
0.2
0.1
50
100
150
200
250
Figure 20: Illustration of Lemma 4.9 in a warm-up setting described in Appendix C.1
C Proof of Lemma 4.9
Let x Qdiag(A1, An An) be the eigenvalue decomposition of the data kernel X. The approximate NTK matrix in Lemma 4.9 can be written exx ağğ²" @Q(diag(A),..., An) + a(Q" () ()(("))))) """ Then it suffices to consider the matrix A := diag()),.... An) + abb" with > 0 and b == Q" y. Denote λι 2 )2 2 ...2 Ån as the eigenvalues of A with corresponding eigenvectors q. Let bli] be the i-th component of b.
C.1 Warm-up: A Simple Case
Let's first start with warm-up case where À1 X2 As A An ie there only 3 distinct eigenvalues for the data kernel. We first conduct an experiment with n = 100, A1 300, X2 150 and λ X
An 100 and b is random vector with unit norm. Figure 20 plots the absolute values of the cosine alignment between b and the first 3 leading eigenvectors of A, denoted Ø1, 92 Ø3- We observe that as increases, b gradually gets larger alignment values with earlier eigenvectors. We also notice sharp transitions around values 50 and 200, which equal A2 A3 and Å₁ λs, respectively. Now start the theoretical analysis in this warm-up setting.
Assumption C.I. Let A diag(A1, A2 ......, A3) abb Raxn with >> 0, 1bbl2 1. We assume that λ - A2 = On)), X2 As = (n)), A3 (n) and [b]i @(1/ I// E [n]]
Lemma C.2. Under Assumption C.I we have that
1. When a < 140(n-1/33 we have that arg maxi Kqr b)
2. When d2-d2 <
A1-da we have that arg maxigisn K0; b) 2. 1+0
1-0(n-1/2)
3. When a > I-Oin 1725.
A1- we have that arg maxigiça Ki by
1.
Proof. By the Bunch Mumch-Nielsen-Sorennse formula BNS78 the eigenvalues of have the following structures.
The first 3 largest eigenvalues A,, do, Ag are the roots the following secular equation
6[112
6/2]
R22
1
where Rº = b[12 b[2)
d
t
A
They also showed that Aa îg d2 R2 -()).
dn
D do Note that under Assumption C.1
2. The other 3 eigenvalues are all equal to Ag.
31

---

By Wikipedia 8
the k-th component of q: (i = 1, 3) is given by
b[k]
i [k]
(Ax N,)
where N; a number that makes the vector qi normalized. Hence it follows that
n
b|4j2
b[2]2 X,)2
R22
N
(Dx
)₂)
V
(A1 X.)2
(A2
(As 2,)2
where R² = 1- - b[1}2 b|2]² is defined in eq.
7 Then for the alignment (āt, b) with i = 1,
We have that
Σx=1 Ae-A.
-1/a
(i: 6)
(8)
6
R2
4/222
R?2
(A1 A,)
(A2- -x]²
(Ag A,)a
(A1 (AI-A,)
(A2-A,)²
(Ag
A,}²
Case 1: small a. Since when + O, we have d + Ao, then we can choose a sufficiently small such that Aa(1+ +2)
< X2.
Then by Assumption
C.I we know that
b b[1|
b/2]
+ -
1 nvn
=0
=
+
O
A
A
A Au3
n²
nyn
Hence eq.
7 is approximated by
R²
1
1
A A
a
nyn
Note that
R²
R
( )
then we can further write the approximated equation as
As
Ag
A2
R²
A A3 On =
which gives us Ao = A3 + al + O( 1/2)) = A +a(1 + (nn Then the condition
As(1 (1 + 1/2) < Xa can be satisfied for certain choice of a such that < 1+0
x2-22 Then by
eq. 8 have
R²
(/3, b)
1/a
Ag Ag
R 1 -0
1 +
6/2
(Ar (A1-Ag]2 Aa)2
(Aa (Aa-
(As -s)3
That means arg max ssisn K@; b)1 = 3.
Case 2: medium a. For the choice of a such that d2 < d2(1 ((n11/2)) and Ae On (nn <
Å1. Then we know that 21 approximated by
b/2
( m while
O
اوا O(+). Hence eq. 7 is
R2
A1
R
(1 O( 1/2)) (1121)
1
A3 - X22
a
which gives us de = As + RR + 11/2)) = As + a(1 = O(n1)/))) Then the condition
Ag < 1 - O(n (122) and d₂ On 11/2)) < Xı can be satisfied for certain choice of a such - that I-Oin-172 17 < a I+0
- 1/2 Again by eq. (8). we have (@ b) ss R==1 1- OOn -1). Hence
arg maxi (q. b) = 2.
Case 3: large For the choice of a such that X₁ -1110(m 11/2)). By similar argument as in the previous two cases, we know that arg maxicisn (@1. = 1. In this case, the condition d > A + O(n (1/2)) can be satisfied for certain choice of a such that a > 1-
stttp:://e.wwikipedia.a
32

---

C.2 Analysis of General Cases and Proof of Lemma 4.9
Recall that we denote p3 to be the left singular vectors of X. Further denote o, as the i-th singular value diag(A1, ..... An) abb" with := Q" y. By our data distribution (see Section 4.1 we have b[i] x (Y,4) = o; (8, pi). then we know that [bfill G Vi,j j E [M]. [Пाा
of X. Note that we write
i.. 0
(2 Combining with the order of {>: in our data distribution yields that (22) for 2 where the index lt defined in our data distribution. We know
b]oj2 O [im
that k- = O (log, n²).
Now are ready to prove Lemma 4.9 Recall aht we enote À 2 d2 2 ...> A. the eigenvalues of A with corresponding eigenvectors 9:- Again, by the Samch-NielsenSSorence formula [BNS78]
A (1 S i 2 n) satisfies the secular equation
6
A Xj
BNS78 also proved that Xy X/ < A--1 for j and X₁ \1.
By our data distribution (see Section 4.1 for the eigenvalue Y₁ with S -, we have Nú A (Aş),, > k. Note that we already proved b[i]² = ( for 2 k. Using these two facts, we can replace Map 2 t: by
M/j2 A with approximation error
1 \;
d A =0
A di
b|i}2
b|0)2
As À
(7,
(x) -x,)
X3
Hence write the secular equation
Σ - ~1
Mij2
as
b
R²a
Σ
A
d - X3
where R²
Σn Ri
b[i}2
and that
1
SAnz"
A <
O
A n h ==++1
X = 0
X
13
v=++1
where the last equality uses the assumption that 1 2i=k+1 X; du dx with 82 = o(1)
By the order of {A 1 in our data distribution, we have that
Σ k
A 0
nlogn n na+1
+
R
Li=k d
12
and hence
Σ 6012 = Ō 1
-6 Ō -(-)) a1 )
1
The order of P in our data distribution also tells us that Vjj < 1 11 - Q(n").
Then for certain choice of a such that \3- < 2y n -0/2) and X + +n a A--1 for certain
j S k 1, then we have AΛ: X = 2(2nn a/2) for i I k 1. Using these facts, we get that
6 X
n/2
1 Xynar2
O
X,
Hence the approximated secular equation can further be written
R²
1
O
+
VVJ Visk
\
X2
33

---

Recall that by discussion above, we have X, -
d = O(2y) and that X, 92(n)). That means
A A 0 5,) and
ōxn"
This allows us write the above approximated equation X2 X;
R
(1 O 0(6x + +n )
vjsk -
A
Xj
We get that 2; dE + a(1 ( + n 2)))) Then the condition X, < 11 - n and 2
25 + n na/2) < A--1 can be satisfied for certain choice of a such that (00(xta-0/2) < < - 140 6/2)
Moreover,
R²
Kāz 2)
1/a
Ax
X
R.
E
607
R22
(A:-A,)²
(Ax A2)²
That means arg maxigign K b)1
34

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:
[Yes]
Justification: The paper's contributions are clearly stated in the introduction, along with brief overview of the main results/contrrbutions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations and potential future directions in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
35

---

The full set of assumptions is provided in Section Complete proofs are provided in Appendix B and
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We described the experimental details in Section 3.1 and Appendix A. to help others reproduce the main results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
36

---

reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We provide code of our experiments in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https for more details.
/nips
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
All details are provided in Section 3.1 and Appendix A.I
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: This paper focuses on full-batch gradient descent, and hence, there is no randomness during training. Different initializations cause randomness, but the resulting training trajectories can be quite different, making it unreasonable to calculate the average and standard deviation. Therefore, instead of calculating the error bars, we report lots of plots under different experimental settings to demonstrate the universality of our results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
37

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources are reported Appendix A.I
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification: All parts conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
This paper is foundational research and not tied to particular applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
38

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses safety risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: In Section 3.1 and Appendix A. credited creators owners of assets, and also mentioned and respected the licenses.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
39

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
40

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
41


Reviews:
[{"review_id": "zRJnk618Jz", "reviewer": "Reviewer_6a2Z", "summary": "This article takes a novel attempt to describe NTK eigenvector behavior during the Edge of Stability phase. They observe that larger learning rates cause the top eigenvectors to align more closely with the target during EoS, coinciding with a decrease in model sharpness, as supported by empirical observations. The authors support these findings with a theoretical analysis of a linear network trained with gradient flow on MSE and further analyze the relationship between sharpness reduction and kernel alignment through the central flow framework.", "questions": "- Figures 1(c) (bottom) and 3(c) (bottom): The curves do not appear to overlap with their respective counterparts (lr = 0.1 & 0.01) from 1(a) and 3(a). Is this alignment between the top NTK eigenvector and Y instead (as in 2(c))?\n- Why is the Individual Eigenvector Target Alignment vs. Eigenvector index peak in Figs. 1 - 3 (b) close to but $\\neq0$?", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n- The article provides a clear, cohesive, and detailed approach to studying the behavior of the NTK eigenvector during the Edge of Stability phase in DNNs. This serves as a first-of-its-kind attempt to investigate this property, contrary to existing literature, which has focused on deciphering the reasons behind and the nature of the eigenvalue spectra instead\n- The authors conduct experiments over a) 4-Layer GeLU networks trained on (a binary class subset of) CIFAR-10, b) 2-layer linear models, c) vision transformer setups to provide empirical evidence for the argument that increasing learning rate causes\n  - An increase in overall NT-kernel alignment with the target $\\to$ indicating better generalization\n  - A shift of the individual NTK-eigenvector -- target alignment towards the top eigenvectors of the NTK $\\to$ supporting the point that there is an increased bias of the model on the leading eigenvector directions of the NTK (which capture the signal(s))\n  - Sudden drops in the sharpness over training coincide with a sharp rise in the KTA \n- Leveraging established connections between the Hessians' spectra and the (norm-squared of the) 2nd layer weight matrix of the 2-layer linear model under consideration, the authors perform a careful study of the evolution of the latter in terms of arriving at the EoS phase. Through this analysis, they demonstrate that a decrease in this proxy object (around the stability threshold) coincides with a shift in the alignment weights towards the top eigenvectors of the NTK", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their comments and would like to answer their questions below.\n\n> Figures 1(c) (bottom) and 3(c) (bottom): The curves do not appear to overlap with their respective counterparts (lr = 0.1 & 0.01) from 1(a) and 3(a). Is this alignment between the top NTK eigenvector and Y instead (as in 2(c))?\n\nFigures 1(c) bottom and 3(c) bottom are both about kernel target alignment, same as in 1(a) and 3(a). The curves do not overlap mainly because in 1(a) and 3(a), we plot the curves every 50 and 20 iterations respectively in a long time range to get an overall behavior of the evolution of KTA, while in 1(c) bottom and 3(c) bottom, we zoom in the EoS period and plot the curves every 3 and 5 steps respectively to study the detailed connection between the evolution of KTA and that of the sharpness.\n\n> Why is the Individual Eigenvector Target Alignment vs. Eigenvector index peak in Figs. 1 - 3 (b) close to but $\\neq0$?\n\nThis is somewhat common behavior, and we don't have a complete understanding why.   Note that in these situations, the NTK's eigenvector with index 0 usually corresponds to the all 1's vector, i.e. the \"DC component.\"   This is true both before and after feature learning.  Since it is always the all 1's vector, it cannot align with the target."}, {"review_id": "nFSFgNH3pO", "reviewer": "Reviewer_chUD", "summary": "This paper investigates the evolution of the Neural Tangent Kernel (NTK) during gradient descent (GD) training, particularly focusing on the alignment between NTK eigenvectors and the target output. \nIt highlights a phenomenon termed \"alignment shift,\" that higher learning rates induce stronger alignment of leading NTK eigenvectors with the target labels. \nThe authors validate this empirically across multiple architectures (FCN, ViT,) and support it with theoretical analysis on two-layer linear networks.", "questions": "Why the loss spikes always co-occur with the sharpness reduction? any intuition here?", "limitations": "1. The experiments are limited to the binary classification.\n2. The theoretical analysis relies on many assumptions with most assumptions being un-validated.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "[Strength]\n1. The authors study a quite under-explored problem: how the eigenvectors of NTK rotate during the GD training process, and how does these connect to the EoS phenomenon.\n2. The authors analyzes the phenomenon from a theoretical analysis on a two-layer linear network. Their theory strengthen their empirical findings.\n\n[Weakness]\n1. The paper’s main empirical observation, that higher learning rates lead to stronger alignment of leading NTK eigenvectors with the target, is somewhat intuitive. \n   - Larger learning rate (within a reasonable range) usually leads to better generalization, thus indicating better alignment between NTK and the target. The NTK at the local iterate somewhat represents the neural network in function space. \n   - Better alignment between NTK and the target indicates better alignment between leading eigenvector of NTK with the target. \n2. The theoretical analysis largely follows WLL22. Both study the NTK evolution of a two-layer linear network and approximate the NTK evolution with the norm of the second-layer weights. Novelty of the theory should be hightlighted.\n3. It is unclear that if the increase of KTA always associated with the sharpness reduction. For example, in fig 3 (c), during the sharpness reduction phase, KTA first decrease and then increases. Also, the experiments of Central Flow cannot fully support this claim. Central flow is a good approximation of GD. GD generalizes better than GF. Thus, KTA of Central flow should be larger than GF.\n4. Clearly, more experiments on more general settings are needed. The datasets are limited to binary classification, and more LRs should be tested. What if we consider the multi-class classifiction? and what if consider larger/smaller LRs?", "was_revised": true, "final_justification": "Bordline paper. Given the insight that sharpness reduction induces kernel target alignment, I am lean to acceptance.", "rebuttal": "We thank the reviewer for their comments and would like to address their concerns below.\n\n> The paper’s main empirical observation, that higher learning rates lead to stronger alignment of leading NTK eigenvectors with the target, is somewhat intuitive.\n\nWe agree that our main observation is somewhat intuitive, but we would argue that it is far from obvious. Although it is known that a larger learning rate (within a reasonable range) usually leads to better generalization, the underlying mechanism is not fully understood yet, and the connection between the generalization ability and the kernel target alignment is not obvious either. The traditional viewpoint focuses on the curvature perspective and argues that large learning rates find “flatter minima”, which generalize better. Our work supports an alternative perspective on the beneficial effects of large learning rates, in that they enhance feature learning by improving the alignment between the NTK and the target vector. We also study the connection between our findings and the generalization ability in lines 147-156 and Appendix A.5 in detail. We think such studies are valuable since the connections are not that obvious.\n\n> The theoretical analysis largely follows WLL22. Both study the NTK evolution of a two-layer linear network and approximate the NTK evolution with the norm of the second-layer weights. Novelty of the theory should be highlighted.\n\nA key innovation of this work, relative to prior works, is to demonstrate the role of the $O(\\eta^2)$ term in contributing to the alignment shift (as noted by reviewer GAQw).  This term was neglected in WLL22 and other prior works, but is necessary for understanding the alignment shift phenomenon that we study.\n\n> It is unclear that if the increase of KTA always associated with the sharpness reduction. For example, in fig 3 (c), during the sharpness reduction phase, KTA first decrease and then increases.\n\nOur claim is not that each step during the sharpness reduction phase will cause the KTA to rise, but rather that the overall phase will cause the KTA to rise (relative to its behavior under the gradient flow).  Note that in this figure, we computed the KTA every 5 steps.  If we had instead computed the KTA at every step, then you would see period-2 oscillations in the KTA.  These are intuitively caused by the sign-flipping oscillations in $\\langle E_t, q_1 \\rangle$, where $E_t$ is the error vector and $q_1$ is the top eigenvector of the data kernel. (See the theoretical analysis for more detailed intuition.) The important story in Figure 3 is that over the whole course of a sharpness reduction phase, the KTA increases in a 'stairwise' fashion after each phase.\n\nSuppose we really want to study the connection between the evolution of KTA and sharpness at each step. In that case, we will need to focus on the time-averaged trajectory to get rid of the influence of the oscillation. That's why we introduce the experiments under Central Flow.  In Figure 4(b), we 'branch off' gradient flow from the central flow at three different points in training, and observe that the KTA of the gradient flow is lower than the KTA of the central flow at each step after branching off.\n\n\n> Also, the experiments of Central Flow cannot fully support this claim. Central flow is a good approximation of GD. GD generalizes better than GF. Thus, KTA of Central flow should be larger than GF.\n\nAs mentioned above, our central flow experiments show that KTA under central flow is larger than that under GF. To see why this can support our claim, we want to highlight that Central Flow is a **sharpness penalized** gradient flow.  Therefore, our central flow experiments demonstrate that adding a sharpness penalty to the gradient flow promotes the increase of KTA, thus highlighting the connection between sharpness and KTA evolution from another perspective.\n\n> Clearly, more experiments on more general settings are needed. The datasets are limited to binary classification, and more LRs should be tested. What if we consider the multi-class classifiction? and what if consider larger/smaller LRs?\n\nWe would like to clarify that we do include more experiments in more settings in the appendices. The datasets include image tasks (in the main part of the paper and Appendix A.2.1 and A.2.2) and language tasks (Appendix A.2.3), and multi-class classification in Appendix A.2.4. For more experiments with more LRs, please see the first point of the responses to Reviewer GAQw.\n\n> Why the loss spikes always co-occur with the sharpness reduction? any intuition here?\n\nThe main intuition and underlying mechanism are discussed in detail in Section 4.2, the four phases of the training dynamics. We would like to make a brief summary here. First, we show that the evolution of the approximate sharpness, i.e. the norm square of the second-layer weights, is determined by $-\\langle E_t, F_t\\rangle$, where $F_t$ is the network output and $E_t$ is the error vector. At first, $\\langle E_t, F_t\\rangle<0$, causing $-\\langle E_t, F_t\\rangle>0$ and the sharpness to increase (phase I). When the sharpness is large enough, it will cause $E_t$ to oscillate with increasing magnitude along the $q_1$ direction, where $q_1$ is the top eigenvector of the data kernel. (see Lemma 4.4) That will lead to the growth of $\\Vert E_t\\Vert$. When $\\Vert E_t\\Vert$ grows large enough, we will see the loss spike, and in the meantime, $\\langle E_t, F_t\\rangle$ will become positive, i.e. $-\\langle E_t, F_t\\rangle<0$, making the sharpness decrease. That's why the loss spike is usually co-occur with the sharpness reduction.\n\n> The theoretical analysis relies on many assumptions with most assumptions being un-validated.\n\nFor more explanations on the intuition and validation of our assumptions, please see the first point of the response to Reviewer fnJU."}, {"review_id": "AKBdrjKj7s", "reviewer": "Reviewer_GAQw", "summary": "This paper studies the evolution of the Neural Tangent Kernel (NTK) at the Edge of Stability (EoS) during gradient descent training. It focuses on how the NTK eigenvectors align with the target and how this alignment evolves over time, especially during periods of sharpness reduction. The authors present empirical results across various architectures, showing that larger learning rates lead to greater NTK–target alignment. The paper also provides a theoretical analysis in the case of two-layer linear networks, describing the 4 phases of the sharpness dynamics and highlighting the key quantities that determine the increase in the alignment.", "questions": "See weaknesses above, in particular, regarding the assumption, and what happens to the alignment close to optimal learning rates.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- Relevance: the paper tackles an important problem: understanding the evolution of the NTK at the Edge of Stability (EoS). This is related to the evolution of the NTK, and thus has obvious connections  to the mechanisms of feature learning through alignment with the target.  \n- The empirical study spans multiple architectures, and provides convincing evidence that larger learning rates correlate with increased kernel-target alignment (KTA).  \n- In my view, the main theoretical contribution is the description of an interesting theoretical mechanism: that the $\\eta^2$ term drives the growth of $\\|v_t\\|^2$ during the sharpness reduction phases, contributing to alignment shift.  \n\n## Weaknesses\n\n- Experiment on the alignment between eigenvectors and the target (Figure 1, 2, 3) should be repeated with a few more learning rates, to see if the alignemnt purely happens in co-occurrences with the model achieving better performances. Hence, the statement that larger learning rate bring more alignment would have to refined/circumstantiated a bit more. What happens for the optimal learning rate? \n- The theoretical section (Section 4) is dense and hard to follow. The narrative around the four phases could be streamlined to focus more directly on the connection between sharpness reduction and alignment increase.  I would suggest re-structure it focusing on the connection with the first empirical part: prove that during the sharpness reduction period, the alignment increases, and discuss the role of the learning rate. I think the full description of the 4 phases of the early training dynamics is not necessary to convey the main points of the paper. \n- There is a confusing inconsistency in the learning rate scaling: the main text introduces $\\eta = \\Theta(n^{-(1-a)/2})$, but some lemmas (e.g., Lemma B.4, B.6) refer to $\\eta = \\Theta(n^{-(1-a)/n})$. This appears to be a typo but adds to the confusion.  \n- At times it is hard to comment on how realistic the assumptions are, and how necessary they are for the argument to hold. For instance, the preservation of a rank-1 structure of $W^{(1)}X$ for unwhitened data. Unfortunately, there is no experimental or theoretical evidence that the rank 1 assumption holds. It would also be interesting to know whether the assumption of not using whitened data affects the rank.  \nAlso, in the theoretical part the authors transition from gradient flow to discrete GD, introducing $\\eta^2$ terms that appear later without clear transition.\n- Minor: figures suffer from small font sizes and thin lines, which may hinder readability.  \n- Minor: the individual eigenvector target alignment plots are on a linear scale, making it hard to assess contributions from low-alignment modes. A log-scale plot would provide a clearer picture of the alignment shift across all modes.  \n- Minor: the sentence *Note that the alignment values sum up to 1, so they can be regarded as a kind of “distribution”* is a bit informal a not necessary in the definition.", "was_revised": true, "final_justification": "The issue on the learning rate is completely resolved. The one about the rank of the weight matrix only partially. Hence, I raised my score from 3 to 4.", "rebuttal": "We thank the reviewer for their comments and would like to address their concerns below.\n\n> Experiment on the alignment between eigenvectors and the target (Figure 1, 2, 3) should be repeated with a few more learning rates, to see if the alignment purely happens in co-occurrences with the model achieving better performances. Hence, the statement that larger learning rate bring more alignment would have to refined/circumstantiated a bit more. What happens for the optimal learning rate?\n\nWe would like to add experiments with more learning rates here. In the fully-connected setting, we tried SGD with batch size 256 under more learning rates. Here we focus on the test performances of the models. We trained the models with different learning rates to the same training loss of 0.01 and calculated the corresponding test losses and KTA.\n\n| Learning rates | 0.02   | 0.025  | 0.03   | 0.035  |\n| -------------- | ------ | ------ | ------ | ------ |\n| test loss      | 0.4290 | 0.4244 | 0.4241 | 0.4664 |\n| KTA            | 0.3607 | 0.3669 | 0.3878 | 0.3432 |\n\nHere we see that lr=0.03 is roughly the optimal learning rate in terms of the test performance, and we do see a correlation between better test performance and better KTA. We would like to argue that this doesn't conflict with our main message, because the goal of this work is to understand the benefit of large learning rates on feature learning and generalization ability. Here we observe that 1) an overlarge learning rate may lead to bad test performance, 2) in the meantime, the alignment gets worse, too. That means our focus on the alignment is a reasonable choice to understand the effect of the learning rate on the generalization ability.\n\n> The theoretical section (Section 4) is dense and hard to follow. \n\nThank you for this suggestion.  We agree that the organization you suggest may make the paper more streamlined. We would also like to clarify that the full description of the 4 phases is not just in the early training dynamics, but a division of the whole Edge of Stability cycle. We describe the 4 phases in detail, mainly to help readers who are not familiar with Edge of Stability to better understand its dynamics. After introducing the 4 phases, we will be able to clearly define the sharpness reduction period and prove the increase of alignment during this period. We agree that we should highlight the main message of the theoretical analysis in a better way and will consider re-structuring it.\n\n> There is a confusing inconsistency in the learning rate scaling.\n\nThis was indeed a typo. It should be $\\eta=\\Theta(n^{-\\frac{1-a}{2}})$ We apologize --- thank you for catching it.\n\n>At times it is hard to comment on how realistic the assumptions are, and how necessary they are for the argument to hold.\n\nFor more explanations on the intuition of our assumptions, please see the first point of the response to Reviewer fnJU.\n\n> For instance, the preservation of a rank-1 structure of $W^{(1)}X$ for unwhitened data. Unfortunately, there is no experimental or theoretical evidence that the rank 1 assumption holds. It would also be interesting to know whether the assumption of not using whitened data affects the rank.  \n\nIn lines 185-198 of the submission, we give a theoretical argument for why the rank-1 structure should hold, where lines 194-198 discuss the case with unwhitened data. The key reasons here are that 1) under small initialization, the network output is small during the early training period, making the error vector close to $-Y$, 2) we consider the one-dimensional output case where $W^{(2)}$ is a row vector (rank 1). Whether we use whitened or unwhitened data doesn't matter. We also want to emphasize that we only make this rank 1 assumption at a short initial time point, far before entering EoS. If this holds at any time $t_0$, it provably holds for all $t>t_0$ (Appendix B.2).\n\nNow we add experimental evidence on the approximate rank 1 structure. In the 2-layer linear network considered in the theoretical analysis, we train the model under lr=0.01 and calculate the eigenvalues of $W^{(1)T}W^{(1)}$. The top 6 eigenvalues at initialization are [3.9297, 3.8144, 3.8057, 3.7141, 3.6817, 3.6167], where there is no low-rank structure. The top 6 eigenvalues after 20 steps (early training) and 60 steps (later training) are [13.0442, 3.9242, 3.8099, 3.7926, 3.7138, 3.6643] and [20.9646, 3.9208, 3.8095, 3.7950, 3.7115, 3.6542] respectively, where we see the trend to become approximately rank 1. Note that if $W^{(1)}$ is rank 1, $W^{(1)}X$ will also be rank 1.\n\n> Also, in the theoretical part the authors transition from gradient flow to discrete GD, introducing $\\eta^2$ terms that appear later without clear transition.\n\nIn Phase I, i.e. the progressive sharpening phase, the dynamic approximately follows the gradient flow trajectory. This has been empirically verified by [CKL+ 21] and assumed by [WLL22]. Hence, we analyze gradient flow in Phase I for simplicity. However, for later phases, the gradient descent trajectory differs from that of gradient flow, especially during the sharpness reduction phase where the $\\eta^2$ term plays a key role. Then it is not valid to focus on gradient flow and therefore we transition to GD and analyze its dynamics in detail.\n\n[CKL+ 21] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In 9th International Conference on Learning Representations, ICLR 2021,\n\n[WLL22] Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems, 35:9983–9994, 2022.\n\n>Minor: figures suffer from small font sizes and thin lines, which may hinder readability.\n\n> Minor: the individual eigenvector target alignment plots are on a linear scale, making it hard to assess contributions from low-alignment modes. A log-scale plot would provide a clearer picture of the alignment shift across all modes.\n\n> Minor: the sentence _Note that the alignment values sum up to 1, so they can be regarded as a kind of “distribution”_ is a bit informal a not necessary in the definition.\n\nWe thank the reviewer for these helpful suggestions and will make edits in our paper."}, {"review_id": "HdkH03UiFF", "reviewer": "Reviewer_fnJU", "summary": "The paper investigates the behavior of the NTK during the training of neural networks, particularly focusing on the phenomenon known as the EoS. The authors observe that during EoS, the leading eigenvectors of the NTK align more with the training target when larger learning rates are used. They provide empirical evidence across various architectures and theoretically analyze this alignment shift phenomenon for a two-layer linear network.", "questions": "See weakness.", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n\n1. The paper presents comprehensive empirical results across different architectures (e.g., linear networks, fully connected networks, Vision Transformers) and datasets, demonstrating the alignment shift phenomenon. \n\n2. The authors provide a detailed theoretical analysis for a two-layer linear network, which is a significant contribution. This analysis offers insights into why the alignment shift occurs during the sharpness reduction phases of EoS, enhancing the understanding of the underlying mechanisms.\n\nWeakness: \n\n1. Some of the theoretical assumptions may be unrealistic, and the authors also do not verify them in the real-world data and network settings, making their application doubtful.\n\n2. The paper primarily focuses on full-batch gradient descent. While the authors mention that the observations generalize to SGD, a more detailed theoretical analysis in the SGD setting would be beneficial, as SGD is more commonly used in practice.\n\n3. Some presentation should be improved. For e.g., the legend and description axes in the Figures are too small.", "was_revised": true, "final_justification": "I have read the response and found some points are considerable, so I raise my rating.", "rebuttal": "We thank the reviewer for their comments and would like to address their concerns below.\n\n> Some of the theoretical assumptions may be unrealistic, and the authors also do not verify them in the real-world data and network settings, making their application doubtful.\n\nWe would like to elaborate more on the intuitions behind our assumptions and argue that although many, they are either mild assumptions or have been made or verified by prior works.\n\n- Data distribution assumptions: In lines 173-178, we explained the intuition of our data distribution. Here, we add more discussions. For the input $X$, we assume that the first several eigenvalues of the data kernel have large gaps while the remaining ones are close to each other. This captures many datasets in practice where its data kernel eigenvalues decrease quickly for the first several and then slowly for later ones. [WLL22] has a similar assumption that \"the gap between two adjacent eigenvalues is not very large, and there is a gap between the largest and the second largest eigenvalue.\" For the target $Y$, we assume that our target is a linear function of the input, which is common in prior theoretical works.\n- Assumption 4.1 assumes that after a short initial period, $W^{(1)}$ converges to a rank-1 matrix whose left singular vector aligns with the right singular vector of $W^{(2)}$.  We explain the intuition behind this assumption in lines 185-198.  We also want to emphasize that we only make this assumption at a short initial time point, which could be far before entering EoS. If this holds at any time $t_0$, it provably holds for all $t>t_0$ (Appendix B.2).\n- Assumption 4.2 is also a common assumption made in prior works. Please see lines 233-235 for more discussions.\n- Assumption 4.3: The first part of it assumes that the loss during training will not grow larger than poly(n) of the initial loss. Here, n is the number of training examples, which is typically large in practice. So, this is a weak assumption. This can also be seen in our loss curves (e.g. Figure 2(c), 5(c)) where the loss during training is not too large compared to the initial loss. The second part of this assumption stipulates that the loss has undergone a nontrivial decrease when it enters EoS. This has been observed and verified by prior EoS papers and captures the fact that the loss usually decreases quickly during the initial training period.\n- Assumption 4.7 is an additional assumption only required in the second point of our theorem. It captures the non-monotonically decreasing behavior of the loss, where the loss after the spike is typically around or smaller than the value before the spike. This phenomenon was also observed and studied by [CKL+ 21] and [WLL22].\n\n[CKL+ 21] Jeremy M. Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In 9th International Conference on Learning Representations, ICLR 2021,\n\n[WLL22] Zixuan Wang, Zhouzi Li, and Jian Li. Analyzing sharpness along GD trajectory: Progressive sharpening and edge of stability. Advances in Neural Information Processing Systems, 35:9983–9994, 2022.\n\n> While the authors mention that the observations generalize to SGD, a more detailed theoretical analysis in the SGD setting would be beneficial, as SGD is more commonly used in practice.\n\nWe agree that in principle, it would be helpful to analyze SGD.  However, the edge of stability dynamics of SGD remains poorly understood -- see [1] for some recent findings.  For example, for SGD, there is no clean four-stage delineation of the EOS dynamics, as there is for GD. For SGD, even theoretically analyzing the eigenvalue evolution is challenging, let alone the evolution of NTK eigenvectors and the alignment with the target. We study GD because we believe that understanding GD is a necessary stepping stone to understanding SGD.\n\n[1] Arseniy Andreyev and Pierfrancesco Beneventano. \"Edge of Stochastic Stability: Revisiting the Edge of Stability for SGD.\"\n\n> Some presentation should be improved. For e.g., the legend and description axes in the Figures are too small.\n\nWe thank the reviewer for this helpful suggestion and will increase the corresponding font sizes in the later version."}]

Author rebuttal:

Meta-review: This paper theoretically and empirically investigate an important problem how the NTK evolves during optimization at the edge of stability. It is shown that the top eigenvector of the NTK aligns more to the feature direction with a larger step size, and in such a setting, the model's sharpness decreases. The theoretical analysis was conducted on a two layer linear model.   
  
The paper is well written. The NTK dynamics was carefully investigated through a detailed and comprehensive numerical experiments with several models. The theoretical analysis is performed only on a linear network but it gives a precise description of the dynamics. Then, the paper successfully resolved the problem that has not been well addressed in previous work.   
  
I think this paper is beneficial to the community, hence recommend acceptance.
Decision: Accept (poster)