Research map / Optimization / Implicit bias and generalization
Understanding the Evolution of the Neural Tangent Kernel at the Edge of Stability
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that larger learning rates at the edge of stability shift NTK eigenvector alignment with the training target toward leading eigenvectors, with theory on two-layer linear networks and central-flow evidence.
Problem
How NTK eigenvectors evolve and align with the training target during gradient descent at the Edge of Stability, where NTK eigenvalues oscillate, is not understood.
Approach
Larger learning rates shift the kernel-target alignment (both full KTA and individual eigenvector alignment) toward leading NTK eigenvectors; the increase occurs during sharpness-decreasing EoS phases, explained by a growing ||v_t||^2 relative to c_t^2 in a two-layer linear network and by a sharpness penalty from central flow.
Key results
- Across architectures (linear, FC, ViT, VGG-11, ResNet-20, BERT-small), larger learning rates yield higher final KTA and shift individual eigenvector alignment toward leading eigenvectors.
- Per-step KTA gain is much larger in sharpness-decreasing EoS periods (e.g., 0.0029 vs 0.0003 per step in FC network intervals 72-75 vs 75-96).
- Central flow reproduces higher KTA for larger learning rates, and gradient flow branched from central flow yields lower KTA, supporting a sharpness-penalty mechanism.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase the GD learning rate | 2-layer linear network on Gaussian data · Gaussian linear-regression task | KTA | Helped |
| increase the GD learning rate | 4-layer GeLU network · CIFAR-10 subset, binary classification | KTA | Helped |
| increase the GD learning rate | SimpleViT · CIFAR-10 subset, binary classification | individual eigenvector target alignment | Helped |
| increase the GD learning rate | VGG-11 · CIFAR-10 subset, binary classification | KTA | Helped |
| increase the GD learning rate | ResNet-20 · CIFAR-10 subset, binary classification | KTA | Helped |
| increase the GD learning rate | BERT-small · IMDB subset, sentiment classification | KTA | Helped |
| increase the GD learning rate | VGG-11 · CIFAR-10, 10-class task | cumulative individual eigenvector target alignment | Helped |
| increase the learning rate from 0.8 to 1.2 | 4-layer GeLU network · CIFAR-10 subset, binary classification | test loss after same effective steps: 1.211 -> 1.199 | Helped |
| increase the learning rate from 0.8 to 1.2 | 4-layer GeLU network · CIFAR-10 subset, binary classification | test loss change during oscillation period: -0.008 -> -0.062 | Mixed |
| sharpness decreases rather than increases | 4-layer GeLU network · CIFAR-10 subset, iterations 72-120 | average KTA change per step: 0.0029; 0.0048; 0.0065 vs 0.0003; -0.0014 | Helped |
| sharpness decreases rather than increases | SimpleViT · CIFAR-10 subset, iterations 130-425 | average KTA change per step: 6.68e-4; 3.26e-4; 2.30e-4 vs -3.83e-5; -4.67e-5; -3.99e-5 | Helped |
| sharpness decreases rather than increases | VGG-11 · CIFAR-10 subset, iterations 12-48 | average KTA change per step: 0.0336; 0.0043 vs 0.0006 | Helped |
Surprising findings
- The increase in target alignment is driven by the typically-ignored eta^2 term, causing ||v_t||^2 to increase while c_t^2 decreases during sharpness reduction.
Limitations
- Theory is limited to two-layer linear networks; nonlinear extension (ReLU) is only partial and heuristic.
- Experiments use small subsets (about 2000 CIFAR-10 examples, 1024 IMDB reviews) and GD mostly, with SGD only partially checked.
Experimental setup
Datasets: CIFAR-10 IMDB synthetic Gaussian data
Benchmarks: CIFAR-10 IMDB
Metrics: kernel target alignment individual eigenvector target alignment sharpness test loss
Baselines: gradient flow AGOP alignment
Builds on: Edge of Stability central flow catapult mechanism AGOP alignment
Closest papers by meaning
- Understanding Edge-of-Stability Training Dynamics with a Minimalist Example (ICLR 2023 · similarity 0.91)
- Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability (NeurIPS 2022 · similarity 0.90)
- Trajectory Alignment: Understanding the Edge of Stability Phenomenon via Bifurcation Theory (NeurIPS 2023 · similarity 0.89)
- Neural Tangent Kernels Motivate Cross-Covariance Graphs in Neural Networks (ICML 2024 · similarity 0.87)
- Shortcut Features as Top Eigenfunctions of NTK: A Linear Neural Network Case and More (NeurIPS 2025 · similarity 0.87)
- The Challenges of the Nonlinear Regime for Physics-Informed Neural Networks (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.