Title: Understanding the Evolution of the Neural Tangent Kernel at the Edge of Stability
Venue: NeurIPS 2025
Accepted: yes
Keywords: Neural Tangent Kernel

training dynamics

Edge of Stability
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.


Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths

1. 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. 

2. 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.

Weakness: 

1. 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.

2. 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.

3. Some presentation should be improved. For e.g., the legend and description axes in the Figures are too small.

QUESTIONS:
See weakness.

LIMITATIONS:
N/A
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nSee weaknesses above, in particular, regarding the assumption, and what happens to the alignment close to optimal learning rates.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n[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?\n\nQUESTIONS:\nWhy the loss spikes always co-occur with the sharpness reduction? any intuition here?\n\nLIMITATIONS:\n1. The experiments are limited to the binary classification.\n2. The theoretical analysis relies on many assumptions with most assumptions being un-validated.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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\n\nQUESTIONS:\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- Why is the Individual Eigenvector Target Alignment vs. Eigenvector index peak in Figs. 1 - 3 (b) close to but $\\neq0$?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-09"}