Understanding Neural Tangent Kernel Dynamics Through Its Trace Evolution

Zhiquan Tan, Tengyuan Liu, Weiran Huang

Keywords: Neural Tangent Kernel, Representation Learning, Training Dynamics

Abstract:

The Neural Tangent Kernel (NTK) has emerged as a valuable tool for analyzing the training and generalization properties of neural networks. While the behavior of the NTK in the infinite-width limit is well understood, a comprehensive investigation is still required to comprehend its dynamics during training in the finite-width regime. In this paper, we present a detailed exploration of the NTK's behavior through the examination of its trace during training.

By conducting experiments on standard supervised classification tasks, we observe that the NTK trace typically exhibits an increasing trend and stabilizes when the network achieves its highest accuracy on the training data. Additionally, we investigate the phenomenon of "grokking'', which has recently garnered attention, as it involves an intriguing scenario where the test accuracy suddenly improves long after the training accuracy plateaus. To shed light on this phenomenon, we employ the NTK trace to monitor the training dynamics during grokking. Furthermore, we utilize the NTK trace to gain insights into the training dynamics of semi-supervised learning approaches, including the employment of exponential moving average mechanisms. Through these investigations, we demonstrate that the NTK, particularly its trace, remains a powerful and valuable tool for comprehending the training dynamics of modern finite-width neural networks.

Primary Area: interpretability and explainable AI

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 9718