Limitations of the NTK for Understanding Generalization in Deep Learning

TL;DR

This work critically evaluates the Neural Tangent Kernel (NTK) as a proxy for understanding generalization in deep learning by comparing data-scaling laws of real networks, empirical NTKs, and infinite NTKs. It shows that finite-width networks can attain significantly better data-scaling exponents than both NTK variants at initialization, and that the empirical NTK continues to evolve throughout training, while the after-kernel improves with more data but cannot fully explain neural-network scaling. The findings reveal concrete limitations of NTK-based explanations for real-world datasets and architectures, particularly in the width/train-size regime, and motivate developing theories that account for representation learning and dynamic kernel evolution beyond the NTK limit.

Abstract

The ``Neural Tangent Kernel'' (NTK) (Jacot et al 2018), and its empirical variants have been proposed as a proxy to capture certain behaviors of real neural networks. In this work, we study NTKs through the lens of scaling laws, and demonstrate that they fall short of explaining important aspects of neural network generalization. In particular, we demonstrate realistic settings where finite-width neural networks have significantly better data scaling exponents as compared to their corresponding empirical and infinite NTKs at initialization. This reveals a more fundamental difference between the real networks and NTKs, beyond just a few percentage points of test accuracy. Further, we show that even if the empirical NTK is allowed to be pre-trained on a constant number of samples, the kernel scaling does not catch up to the neural network scaling. Finally, we show that the empirical NTK continues to evolve throughout most of the training, in contrast with prior work which suggests that it stabilizes after a few epochs of training. Altogether, our work establishes concrete limitations of the NTK approach in understanding generalization of real networks on natural datasets.

Figures (9)

• Figure 1: Summary of results: (A) Neural network scales better than NTK at initialization: We compare the scaling exponent of a neural network, its corresponding infinite and empirical NTK at initialization. Details in Section \ref{['sec:data-scaling']}. (B) After-kernel continues to improve with more training samples: We train a neural network with $m=\{1K, 2K...1024K\}$ samples, extract the empirical NTK at completion, and use this kernel to fit 500 samples. Details in Section \ref{['sec:ker_data']}. (C) Empirical NTK improves with constant rate with respect to training time: We extract the empirical NTK at various times in training and use it to fit the full train dataset. Details in Section \ref{['sec:timedyn']}.
• Figure 2: On the Effect of width. In Figure \ref{['fig:width-a']} we plot data scaling laws of the Myrtle-CNN at small (16) and large (1024) widths and its the infinite NTK. We observe that both finite widths have similar scaling constant which is better than that of the infinite NTK. In Figure \ref{['fig:width-b']} we plot the performance of Myrtle-CNN and its Empirical NTK for a fixed training size while varying width. In Figure \ref{['fig:width-svhn']} we do the same for a 5-layer CNN and the SVHN-parity task. Both Figure \ref{['fig:width-b']} and \ref{['fig:width-svhn']} we observe that (a) the Empirical NTK performance continues to improve with width, moving towards the infinite NTK performance while (b) neural network performance improves initially and then starts to deteriorate towards the infinite NTK performance. Error bars represent estimated standard deviation. See Appendix \ref{['app:exprm-details']} for more details.
• Figure 3: Neural networks continue to have a better scaling constant under various hyperparameter choices. We compare the data-scaling for a Myrtle-CNN, its Empirical NTK and its infinite NTK on CIFAR-5m-bin task under various hyperparameter changes: (a) Higher and lower learning rate compared to Figure \ref{['fig:main']} (b) GD instead of SGD (c) SGD without momentum (d) Training until convergence (no early stopping)
• Figure 4: After-Kernel continues to improve with dataset size. In Figure \ref{['fig:knn_vs_net']} we plot data scaling curves of $K_n(n), K^F_n(n)$ and the neural network and observe that they behave very similarly. In Figure \ref{['fig:cnn_ker_learning']} we plot $K^F_m(500)$ versus $m$ and observe that the performance improved with increasing $m$. In Figure \ref{['fig:k_m_vs_net']} we plot data scaling curves of Empirical NTK at initialization, $K_{16k}$, $K_{32k}$ and the neural network. We observe that the neural network has the best scaling law amongst these.
• Figure 5: Empirical NTK keeps improving uniformly throughout most of the training. In Figure \ref{['fig:ker_time_overp_simple']} we plot the test error of Myrtle-CNN, its empirical NTK at initialization and $K^t_{fit}$ at time $t$. We observe that the slope $K^t_{fit}$ does not decrease with time suggesting that the change in kernel does not slow down after an initial part of training. Using this same setup, we plot the data scaling curves of $K^t$ for various $t$ and the data scaling of Myrtle-CNN in Figure \ref{['fig:k_m_t']}. We observe that the Myrtle-CNN has the best scaling law.

Theorems & Definitions (1)

• Claim 3.1