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

arXiv: https://arxiv.org/abs/2205.12372 · HTML: https://mlsanity.com/paper/2205.12372/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 33 references that are in the map.
- First posted on arXiv 24 May 2022.

## Builds on (most cited first, 23 of 23)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning](https://mlsanity.com/paper/1602.07261/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [EfficientNetV2: Smaller Models and Faster Training](https://mlsanity.com/paper/2104.00298/) (2021)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [On Exact Computation with an Infinitely Wide Neural Net](https://mlsanity.com/paper/1904.11955/) (2019)
- [Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspective](https://mlsanity.com/paper/2102.11535/) (2021)
- [Finite Versus Infinite Neural Networks: an Empirical Study](https://mlsanity.com/paper/2007.15801/) (2020)
- [Neural Tangents: Fast and Easy Infinite Neural Networks in Python](https://mlsanity.com/paper/1912.02803/) (2019)
- [Predicting trends in the quality of state-of-the-art neural networks without access to training or testing data](https://mlsanity.com/paper/2002.06716/) (2020)
- [Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks](https://mlsanity.com/paper/1910.01663/) (2019)
- [Spectra of the Conjugate Kernel and Neural Tangent Kernel for linear-width neural networks](https://mlsanity.com/paper/2005.11879/) (2020)
- [Every Model Learned by Gradient Descent Is Approximately a Kernel Machine](https://mlsanity.com/paper/2012.00152/) (2020)
- [Dissecting Hessian: Understanding Common Structure of Hessian in Neural Networks](https://mlsanity.com/paper/2010.04261/) (2020)
- [A Deeper Look at the Hessian Eigenspectrum of Deep Neural Networks and its Applications to Regularization](https://mlsanity.com/paper/2012.03801/) (2020)
- [Pathological spectra of the Fisher information metric and its variants in deep neural networks](https://mlsanity.com/paper/1910.05992/) (2019)
- [The Normalization Method for Alleviating Pathological Sharpness in Wide Neural Networks](https://mlsanity.com/paper/1906.02926/) (2019)
- [Understanding Layer-wise Contributions in Deep Neural Networks through Spectral Analysis](https://mlsanity.com/paper/2111.03972/) (2021)

## Built on by (most cited first, 3 of 3)

- [Characterizing the Spectrum of the NTK via a Power Series Expansion](https://mlsanity.com/paper/2211.07844/) (2022)
- [Determinant Estimation under Memory Constraints and Neural Scaling Laws](https://mlsanity.com/paper/2503.04424/) (2025)
- [Train Faster, Perform Better: Modular Adaptive Training in Over-Parameterized Models](https://mlsanity.com/paper/2405.07527/) (2024)

## Abstract (arXiv)

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.