TorchNTK: A Library for Calculation of Neural Tangent Kernels of PyTorch Models

TL;DR

TorchNTK is a Python library that efficiently computes neural tangent kernels for PyTorch models, supporting various architectures and providing insights into layerwise contributions and memory efficiency.

Contribution

The paper introduces torchNTK, a library that enables efficient calculation of NTKs for diverse PyTorch models, including layerwise analysis and comparison of differentiation methods.

Findings

Explicit differentiation and autodifferentiation methods are compared.

Layerwise NTK components can be more memory efficient.

Preliminary experiments demonstrate utility and insights into NTK properties.

Abstract

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.