TorchNTK: A Library for Calculation of Neural Tangent Kernels of PyTorch Models
This provides a tool for researchers and practitioners in machine learning to analyze neural network behavior through NTK calculations, but it is incremental as it builds on existing NTK theory and PyTorch frameworks.
The authors introduced torchNTK, a Python library for efficiently calculating the empirical neural tangent kernel (NTK) of PyTorch models, focusing on multilayer perceptrons and extending to architectures like convolutional networks, with preliminary experiments demonstrating use cases and probing the NTK.
We introduce torchNTK, a python library to calculate the empirical neural tangent kernel (NTK) of neural network models in the PyTorch framework. We provide an efficient method to calculate the NTK of multilayer perceptrons. We compare the explicit differentiation implementation against autodifferentiation implementations, which have the benefit of extending the utility of the library to any architecture supported by PyTorch, such as convolutional networks. A feature of the library is that we expose the user to layerwise NTK components, and show that in some regimes a layerwise calculation is more memory efficient. We conduct preliminary experiments to demonstrate use cases for the software and probe the NTK.