TorchNTK: A Library for Calculation of Neural Tangent Kernels of PyTorch Models
Not on the map. Why?
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.
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
23 of 33 references.
References · 23
- 2015Deep Residual Learning for Image Recognitioncited by 52,096
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,621
- 2016Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learningcited by 1,512
- 2020Sharpness-Aware Minimization for Efficiently Improving Generalizationcited by 1,417
- 2021EfficientNetV2: Smaller Models and Faster Trainingcited by 1,135
- 2019Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descentcited by 738
- 2018On Lazy Training in Differentiable Programmingcited by 675
- 2019On Exact Computation with an Infinitely Wide Neural Netcited by 573
- 2021Neural Architecture Search on ImageNet in Four GPU Hours: A Theoretically Inspired Perspectivecited by 169
11 more
- 2020Finite Versus Infinite Neural Networks: an Empirical Studycited by 157
- 2019Neural Tangents: Fast and Easy Infinite Neural Networks in Pythoncited by 149
- 2020Predicting trends in the quality of state-of-the-art neural networks without access to training or testing datacited by 115
- 2019Harnessing the Power of Infinitely Wide Deep Nets on Small-data Taskscited by 87
- 2020Spectra of the Conjugate Kernel and Neural Tangent Kernel for linear-width neural networkscited by 77
- 2020Every Model Learned by Gradient Descent Is Approximately a Kernel Machinecited by 55
- 2020Dissecting Hessian: Understanding Common Structure of Hessian in Neural Networkscited by 45
- 2020A Deeper Look at the Hessian Eigenspectrum of Deep Neural Networks and its Applications to Regularizationcited by 44
- 2019Pathological spectra of the Fisher information metric and its variants in deep neural networkscited by 32
- 2019The Normalization Method for Alleviating Pathological Sharpness in Wide Neural Networkscited by 20
- 2021Understanding Layer-wise Contributions in Deep Neural Networks through Spectral Analysiscited by 3
Built on by
Citing papers per year: 2022: 1; 2023: 0; 2024: 1; 2025: 1.
3 papers.