Neural Tangents: Fast and Easy Infinite Neural Networks in Python

By Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Narain Sohl-Dickstein, Sam Schoenholz. 2019

261 citations recorded by the source corpus.

Neural Tangents is a library for working with infinite-width neural networks. It provides a high-level API for specifying complex and hierarchical neural network architectures. These networks can then be trained and evaluated either at finite-width as usual, or in their infinite-width limit. For the infinite-width networks, Neural Tangents performs exact inference either via Bayes' rule or gradient descent, and generates the corresponding Neural Network Gaussian Process and Neural Tangent kernels. Additionally, Neural Tangents provides tools to study gradient descent training dynamics of wide but finite networks. The entire library runs out-of-the-box on CPU, GPU, or TPU. All computations can be automatically distributed over multiple accelerators with near-linear scaling in the number of devices. In addition to the repository below, we provide an accompanying interactive Colab notebook at https://colab.sandbox.google.com/github/neural-tangents/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Statistical and Nonlinear Physics, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Statistical and Nonlinear Physics

In Model Architectures

Citations in the corpus

Wide neural networks of any depth evolve as linear models under gradient descent (2019, methodology): “…limit, the output of the network at initialization is a draw from a Gaussian process (GP); moreover, the network output remains governed by a GP after exact Bayesian training using squared loss (Neal, 1994; Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019; Garriga-Alonso et al., 2018).”
The large learning rate phase of deep learning: the catapult mechanism (2020, methodology): “We are using JAX (Bradbury et al., 2018) and the Neural Tangents Library for our experiments (Novak et al., 2020).”
Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains (2020, methodology): “All experiments are implemented in JAX [8] and the NTK functions are calculated automatically using the Neural Tangents library [30].”
Coresets via Bilevel Optimization for Continual Learning and Streaming (2020, methodology): “We compare the performance of our method with RBF proxy and with NTK proxy corresponding to the CNN architecture chosen in Section 6.1 without dropout and max-pooling calculated using the library of [45], on SplitMNIST in continual learning and streaming settings.”
Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes (2018, methodology): “, 2018) CNNs benefit from overparameterization Due to computational limitations, the performance of CNN-GP with pooling on the full dataset was only verified after the publication using the Neural Tangents library (Novak et al., 2020).”
Disentangling trainability and generalization in deep learning (2019, methodology): “We compute the exact NTK with activation function Erf using the Neural Tangents library (Novak et al., 2019a).”
On the infinite width limit of neural networks with a standard parameterization (2020, methodology): “All of the experiments in this section were done using Neural Tangents library [24].”
Explaining neural scaling laws (2021, methodology): “We have presented a framework for categorizing neural network scaling laws, along with derivations that help to explain their origins.”
Machine learning and deep learning—A review for ecologists (2022, result): “…deeper and wider decreases generalization error, suggesting that model size, contrary to our general expectation, can actually be beneficial for reducing total prediction error (Arora, Du, et al., 2019; Huh et al., 2021 ; Nakkiran et al., 2019 ; Novak et al., 2019 ; Shwartz-Ziv & Alemi, 2020 ).”
Mean-field Behaviour of Neural Tangent Kernel for Deep Neural Networks (2019, methodology): “These approximations are implemented in Python library (Novak et al., 2020).”
Spectrum Dependent Learning Curves in Kernel Regression and Wide Neural Networks (2020, methodology): “For these experiments, we use the Neural-Tangents Library (Novak et al., 2020) which supports training and inference for both finite and infinite width neural networks.”
Provably efficient machine learning for quantum many-body problems (2021, methodology): “We examine metrics k x; xl ð Þ∈R based on Gaussian kernel, Dirichlet kernel, and neural tangent kernel (50).”
Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks (2020, unknown): “Our results are directly applicable to in ﬁ nite-width limit of neural networks that admit a kernel description (including feedforward, convolutional and recurrent neural networks) 13,55 – 59 , and explain their inductive bias towards simple functions 35,51,60 – 64 .”
Dataset Meta-Learning from Kernel Ridge-Regression (2020, methodology): “All our kernel-based experiments use the Neural Tangents library (Novak et al., 2020), built on top of JAX (Bradbury et al., 2018).”
A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges (2020, background): “[373] proposed neural tangents; a library intended to facilitate research into infinite-width neural networks.”
Dataset Distillation with Infinitely Wide Convolutional Networks (2021, methodology): “As an illustration, when using the neural_tangents [Novak et al., 2020] library, one is able to use roughly the following kernel batch sizes9 on a Tesla V100 GPU (16GB RAM): FC1 kernel batch size ∼ 104 Conv-Vec1 kernel batch size ∼ 103 ConvNet kernel batch size ∼ 101.”