[Paper Review] Neural Tangents: Fast and Easy Infinite Neural Networks in Python

Roman Novak, Lechao Xiao+5 more|arXiv (Cornell University)|Dec 5, 2019

Model Reduction and Neural Networks28 references57 citations

TL;DR

Neural Tangents provides a JAX-based API to define, train, and analyze infinite-width neural networks via analytic NNGP/NTK kernels or finite-width approximations, with scalable CPU/GPU/TPU support.

ABSTRACT

Neural Tangents is a library designed to enable research into 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. Infinite-width networks can be trained analytically using exact Bayesian inference or using gradient descent via the Neural Tangent Kernel. Additionally, Neural Tangents provides tools to study gradient descent training dynamics of wide but finite networks in either function space or weight space. 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. Neural Tangents is available at www.github.com/google/neural-tangents. We also provide an accompanying interactive Colab notebook.

Motivation & Objective

• Introduce Neural Tangents, an open-source library for researching infinite-width neural networks.
• Provide a high-level API to specify complex architectures and obtain analytic infinite-width kernels (NNGP/NTK).
• Enable both exact Bayesian inference and gradient-descent-inspired training in the infinite-width limit.
• Support Monte Carlo approximations for kernels that lack analytic forms.
• Demonstrate scalability and practical usefulness across architectures and datasets.

Proposed method

• Define a translation from standard tensor operations (dense, conv, nonlinearities) to corresponding kernel operations for NNGP and NTK.
• Provide analytic kernel computations for infinite-width networks and MC-based kernel approximations when analytic forms are unavailable.
• Offer an API to perform Bayesian inference or gradient-descent-like training in the infinite-width limit via exact kernels.
• Incorporate continuous-time training dynamics and finite-width weight-space explorations through Taylor series approximations.
• Leverage XLA/JAX for CPU/GPU/TPU execution with automatic device distribution and batching.

Experimental results

Research questions

• RQ1 How can one efficiently compute and use NNGP and NTK kernels for a wide range of neural architectures in the infinite-width limit?
• RQ2 Can finite-width networks be accurately analyzed via their infinite-width counterparts using analytic or Monte Carlo kernel estimates?
• RQ3 What is the performance and scalability of analytic kernel computations across devices and architectures?
• RQ4 How do different architectures (e.g., fully-connected, CNNs, WideResNet) compare under infinite-width inference and kernel-based training?
• RQ5 What practical tools and workflows can accelerate research into infinite-width neural networks?

Key findings

• Neural Tangents enables exact inference for infinite-width networks via analytic NNGP/NTK kernels for many architectures.
• Monte Carlo kernel estimation provides flexible approximations when analytic kernels are unavailable, with convergence toward analytic values as width or sample count grows.
• In experiments, infinite-width inference can closely match finite-width ensemble training on synthetic data, illustrating practical fidelity.
• The library demonstrates convergence and hierarchy of performance across architectures when evaluated on CIFAR-10, with WideResNet-like structures performing best among tested variants.
• The implementation achieves near-linear scaling with additional accelerators and supports batch processing, op fusion, and 2D-convolution-based covariance computations for efficiency.