Source-linked AI summary
Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M. Roy, Surya Ganguli
TL;DR
The paper examines why practical deep-network training diverges from fixed-initialization NTK linearization and how basin geometry, kernel evolution, and training dynamics relate. Through large-scale measurements across architectures and datasets, it finds a rapid chaotic-to-stable transition: basin fate is set within 2–3 epochs, the NTK improves over initialization by a factor of 3 within 3–4 epochs, and reaches full-training performance in 15%–45% of training.
Problem
The paper addresses the limited understanding of how nonlinear training, loss-landscape geometry, and time-varying data-dependent NTKs interact in practical deep networks.
Method
The authors perform large-scale phenomenological measurements of diverse geometry, training-dynamics, and NTK metrics using full, linearized, and data-dependent-NTK training with parent–child spawning.
Results
Across architectures and datasets, a chaotic-to-stable transition occurs within 2–3 epochs; the learned NTK outperforms the initial NTK by a factor of 3 within 3–4 epochs and matches full-training performance in 15%–45% of training.
Takeaways & Limitations
The results provide an integrative empirical picture and motivate richer theories of deep learning and potentially more rational learning-rate schedules.
Takeaways & Limitations
Fixed-initialization NTK theory is a poor description of practical finite-width training, including nonlinear training at extremely low learning rates during the early chaotic phase.
Abstract
from arXiv · showhide
In suitably initialized wide networks, small learning rates transform deep neural networks (DNNs) into neural tangent kernel (NTK) machines, whose training dynamics is well-approximated by a linear weight expansion of the network at initialization. Standard training, however, diverges from its linearization in ways that are poorly understood. We study the relationship between the training dynamics of nonlinear deep networks, the geometry of the loss landscape, and the time evolution of a data-dependent NTK. We do so through a large-scale phenomenological analysis of training, synthesizing diverse measures characterizing loss landscape geometry and NTK dynamics. In multiple neural architectures and datasets, we find these diverse measures evolve in a highly correlated manner, revealing a universal picture of the deep learning process. In this picture, deep network training exhibits a highly chaotic rapid initial transient that within 2 to 3 epochs determines the final linearly connected basin of low loss containing the end point of training. During this chaotic transient, the NTK changes rapidly, learning useful features from the training data that enables it to outperform the standard initial NTK by a factor of 3 in less than 3 to 4 epochs. After this rapid chaotic transient, the NTK changes at constant velocity, and its performance matches that of full network training in 15% to 45% of training time. Overall, our analysis reveals a striking correlation between a diverse set of metrics over training time, governed by a rapid chaotic to stable transition in the first few epochs, that together poses challenges and opportunities for the development of more accurate theories of deep learning.
1 Diverse aspects of deep learning phenomenology
Prior work examines global and local loss-landscape geometry, linearized NTK dynamics, and early-training effects, but these strands remain insufficiently integrated. The paper asks how basin selection, kernel learning, and curvature interact during practical deep-network training.
- Low-loss regions can be high-dimensional, allowing low-dimensional affine parameter subspaces to intersect them.
- After sufficiently early parent training, independently stochastic child networks can converge to the same low-loss basin and remain linearly connected.
- In the infinite-width, low-learning-rate limit, training is approximated by linear motion in the tangent space at the initial function.
- Hessian studies report a few large eigenvalues and many small ones, supporting flat directions in low-loss regions.
- The paper investigates when basin fate becomes fixed, when the NTK learns data-dependent structure, when kernel learning stabilizes, and how curvature changes alongside these events.
2 Definition of measurement metrics for geometry and training
The paper formalizes training, function-space, kernel, and geometric measurements, then compares full SGD with linearized and data-dependent-NTK dynamics. Parent–child spawning provides a hierarchical probe of how training-time perturbations affect final networks.
- Definitions: The training set contains m examples with K-class labels, and fw(x) denotes a K-dimensional logit output for weights w.
- Training dynamics: SGD updates weights by Δt := wt+1 − wt = −η∇wR_S̄(wt), while a second-order Taylor expansion tracks logit changes.
- Linearized training: For infinitesimal learning rates, higher-order Taylor terms vanish and parameter-space steepest descent corresponds to function-space descent using an NTK.
- NTK dynamics: The NTK is κt(x,x′) = Jt(x)Jt(x′)^T; a time-constant kernel yields dynamics equivalent to training the network linearized at that time.
- Parent–child spawning: Parent–child spawning copies a parent at time ts, trains children with independent minibatch stochasticity, and measures their distances as a function of spawn time.
- Distance metrics: Function-space distance is approximated by the normalized fraction of test examples on which two networks disagree, with unit distance indicating uncorrelated errors.
3 An integrative view of learning dynamics
Across the studied networks and datasets, error barriers, kernel evolution, and child distances exhibit a common early-to-late transition. An initially chaotic phase gives way to slower, more stable dynamics, including under a constant learning rate.
- Standard training: Within two to five epochs, the error barrier on a linear path between children falls close to zero.
- Standard training: The NTK changes rapidly early in training and more slowly later, while final function, kernel, and ReLU distances between children decrease with spawn time.
- Constant learning rate: Constant-learning-rate experiments preserve the rapid early error-barrier drop, fast early NTK evolution, slow later evolution, and stabilized child distances.
- Chaotic-to-stable transition: The early phase is highly sensitive to SGD noise and determines basin selection, whereas the later phase is more stable and retains smaller child distances.
- Robustness: Equivalent results appear for other networks in Appendix C, alongside additional properties tracked across epochs.
4 The local and global geometry of the loss landscape surrounding children
Parent–child experiments connect spawn timing with basin geometry and function-space diversity. Early-spawned children can occupy distinct basins where tangent-plane approximations fail, while later-spawned children share a basin with smaller diversity and better local approximations.
- Early-spawned children reach well-separated loss basins, whereas later-spawned children reach the same basin while retaining smaller function diversity.
- Tangent-plane test-error approximations fail across different basins but approximately describe the landscape when children occupy the same basin.
- Within two to three epochs, the error barrier drops to zero; later spawning produces children that remain progressively closer in function space.
- Figure 5 compares test error and child function distance on a projected 2D section with tangent-plane error views around the parent and a child.
- Figure 6 relates final child function distance to spawn epoch and plots function distance against the error barrier for CIFAR10 and CIFAR100.
5 NTK velocity slows down and stabilizes after basin fate is determined
The error barrier and kernel velocity decrease together during early chaotic training, after which kernel velocity stabilizes at a low non-zero value once basin fate is determined.
- The error barrier and kernel velocity fall together and stabilize early, exhibiting strongly correlated dynamics.
- After the error barrier disappears, the NTK continues evolving at a low non-zero velocity.
6 The data-dependent NTK rapidly learns features useful for performance
The paper evaluates data-dependent NTKs by switching from nonlinear to linearized training at different times. These kernels rapidly learn useful features and later approach full-network performance.
- The protocol trains the nonlinear network until onset time ˜t, then Taylor expands it with respect to weights for subsequent linearized training.
- Within 3 to 4 epochs, the data-dependent NTK reduces error by at least a factor of 3 relative to the classic NTK at initialization.
- By 30 to 90 epochs, the data-dependent NTK essentially matches a network trained with full nonlinear training for 200 epochs.
7 NTK and nonlinear training remain different even at low learning rates
Finite-width nonlinear training remains distinct from learned-kernel linearized training even at very low learning rates, with the difference concentrated in the early chaotic phase.
- A significant nonlinear performance advantage over linearized training persists at finite width and low learning rate, but only during the first few epochs.
- The low-learning-rate nonlinear advantage disappears in tight correlation with the disappearance of the error barrier.
- The data-dependent NTK describes low-learning-rate dynamics well after the first few epochs, but not the early chaotic nonlinear phase.
8 Summary of contributions and discussion
Across datasets and architectures, the paper identifies a universal transition from chaotic to stable training within two to three epochs, while emphasizing limits of current NTK theory and opportunities for learning-rate theory.
- Within two to three epochs, diverse measurements reveal a universal chaotic-to-stable training transition across datasets and architectures.
- During the chaotic transient, basin fate is determined, the NTK changes rapidly, and learned-kernel training soon outperforms the initial NTK.
- After the transient, kernel velocity stabilizes while data-dependent kernel performance matches full-network training by 30 to 90 epochs.
- At finite widths and large learning rates, and even during early low-learning-rate training, NTK theory poorly describes practical nonlinear dynamics.
- The observed learning sequence could inform theories for rational learning-rate schedule design.
Broader Impact
The work aims to improve understanding of deep neural networks, with the potential long-term benefit of making machine-learning applications more reliable and transparent.
- The work seeks to better understand deep neural networks.
- Improved understanding could potentially make machine-learning applications more reliable and transparent in the long run.
Funding Sources
This section combines funding acknowledgments with methodological descriptions of function-space comparisons and visualization during training.
- Funding Sources: The research received support from NSERC, the Ontario Early Researcher Award, the Charles Simonyi Endowment, the Simons Foundation, the James S. McDonnell Foundation, NTT Research, and an NSF Career award.The work was also conducted partly during an Institute of Advanced Studies special year.
- Definitions of additional metrics: Function distance between child runs was measured for ResNet20 on CIFAR-10 and CIFAR-100, with raw data smoothed using a ±5-epoch window.
- Definitions of additional metrics: A t-SNE embedding visualized parent and child trajectories by flattening test-set predicted-probability vectors across images and classes.
B Definitions of additional metrics
The section defines metrics linking logit-gradient geometry, Hessian structure, function-space motion, and escape from local quadratic basins. It also reports early training changes, basin stabilization, and correlated metric dynamics.
- B.1 Logit gradient centroid alignment: The logit gradient centroid for class k is the average parameter gradient of its logit across training examples.
- B.2 Logit gradient centroids and the top Hessian eigenvectors: The span of the K logit gradient centroids approximately tracks the top K Hessian-curvature directions and describes basin-wall orientation.
- B.2 Logit gradient centroids and the top Hessian eigenvectors: For mutually almost orthogonal gradient centroids, the MSE Hessian is approximately represented by centroid directions with squared-norm singular values.
- B.1 Logit gradient centroid alignment: Logit gradient centroid alignment estimates overlap between the low-dimensional subspaces containing the sharpest Hessian directions of two networks.
- B.3 Escape threshold: The escape threshold determines whether a second-order loss trajectory descends within a quadratic basin or increases loss and escapes it.
- B.3 Escape threshold: For gradient descent, the escape-threshold condition simplifies to 2−ηλt, where λt is the Hessian spectral norm.
- C.1 Diverse metrics for loss landscape and training are highly correlated: Around five epochs, test and training error, the top Hessian eigenvalue, and several distance measures change rapidly before gradually freezing, with significant kernel learning.
- C.1 Diverse metrics for loss landscape and training are highly correlated: At spawn time ts = 1, children’s basin fate is sealed; beyond ts = 1, they end in the same basin with no loss barrier along a linear path.
C.2 Further discussion
The analysis presents a consistent training-time pattern: early dynamics determine the final low-loss basin, while the data-dependent NTK rapidly improves and later approaches full-network performance.
- C.2 Further discussion: Within 1–10% of training time, SGD-noise-dependent basin selection ends and independently spawned networks converge to the same linearly connected low-loss basin.The error barrier between spawned children falls below 0 at around 10 epochs across networks.
- C.2 Further discussion: Before one epoch, the data-dependent NTK rapidly learns useful information across networks and datasets.This is evidenced by the decline of the linearized-training error curves.
- C.2 Further discussion: Less than halfway through training, the data-dependent NTK performs nearly as well as the full network.This contrasts with the NTK at random initialization, which performs worse than the early data-dependent kernel.
- C.2 Further discussion: Early-built NTK features are sufficient for low errors competitive with full nonlinear networks across the studied settings.The authors therefore distinguish the usefulness of an early-trained data-dependent NTK from training the NTK at random initialization.
- C.2 Further discussion: Because the early phase determines the network’s final fate, understanding it may support earlier diagnosis and reduce training cost.The paper frames this as a potential application rather than a demonstrated intervention.