PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Tangents: Fast and Easy Infinite Neural Networks in Python

Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, Samuel S. Schoenholz

arXiv:1912.02803v1stat.MLcs.LG

TL;DR

Researchers lacked an efficient, general way to work with infinite-width neural networks because deriving kernels for each architecture was laborious and error prone. Neural Tangents provides a high-level JAX library for specifying architectures, computing analytical or Monte Carlo kernels, performing inference, and studying training dynamics. It aims to make infinite-width networks easier to explore while supporting CPU, GPU, TPU, and distributed execution, subject to affine-before-nonlinearity requirements for many infinite-width results.

  • Problem

    Using NNGPs and NTK-GPs requires laborious, error-prone, architecture-specific kernel derivations that can take weeks to months even for experienced practitioners.

  • Method

    Neural Tangents is an open-source JAX library with high-level network composition, analytical and Monte Carlo kernel computation, inference APIs, ODE solvers, and Taylor approximations.

  • Results

    The library supports exact infinite-width inference that agrees closely with ensembles of finite-width networks, while providing portable and distributed kernel computation.

  • Takeaways & Limitations

    Neural Tangents is intended to democratize infinite-width networks for research and support their use as practical and theoretical tools, especially in data-limited settings.

  • Takeaways & Limitations

    Many infinite-width results require nonlinear transformations to be preceded by dense or convolutional affine transformations, a condition enforced by the library.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Infinite-width neural networks offer useful theoretical and practical capabilities, but using them has been laborious because kernels must be derived manually for each architecture. Neural Tangents addresses this barrier with a high-level JAX library for specifying networks, computing kernels and inference, and studying finite- and infinite-width dynamics.

  • Motivation: Infinite-width networks connect Bayesian neural networks to architecture-dependent Gaussian processes and characterize gradient-descent training through the Neural Tangent Kernel.These developments established infinite-width networks as tools for understanding deep-learning phenomena and achieving strong benchmark performance in some settings.
  • Motivation: Using NNGPs and NTK-GPs can require weeks to months of work because architecture-specific kernels must be derived by hand, making the process laborious and error prone.The paper compares this situation to neural-network development before high-quality automatic-differentiation packages became widespread.
  • Contributions: Neural Tangents introduces an open-source JAX library that accelerates research on infinite limits of neural networks.Its high-level API supports complex, hierarchical model specifications and analytical evaluation of infinite-width NNGP kernels and NTKs.
  • Contributions: The library approximates otherwise intractable infinite-width kernels with Monte Carlo sampling using methods agnostic to the neural-network library.This broadens the architectures and software contexts for which infinite-width calculations can be explored.
  • Contributions: Neural Tangents provides analytical Bayesian and continuous-gradient-descent inference, numerical ODE solvers for broader losses and dynamics, and arbitrary-order Taylor approximations for weight-space analysis.The ODE tools cover continuous gradient descent with or without momentum at finite or infinite time.
  • Contributions: Leveraging XLA, the library runs on CPUs, GPUs, and TPUs and distributes kernel computations across multiple accelerators with near-perfect scaling.The paper presents this portability and distributed execution as infrastructure for efficient infinite-width research.

2 EXAMPLES

The examples show how Neural Tangents supports exact infinite-width inference, finite-network comparisons, kernel approximation, and architecture experiments. Across these tasks, infinite-width predictions agree closely with finite-width ensembles, empirical kernels converge toward analytic kernels, and WideResNets outperform simpler architectures as datasets grow.

  • Inference with an infinitely wide neural network: Exact infinite-width inference shows excellent agreement with an ensemble of one-hundred finite-width networks trained by gradient descent.The comparison uses training curves and output predictions; hidden width affects the finite models but not infinite-width inference.
  • An infinitely wide WideResNet: Monte Carlo estimates of finite-channel WideResNet kernels approach the corresponding analytic infinite-channel kernels as widening factor or sample count increases.The experiment studies both NNGP and NTK estimates for WRN-28-k relative to WRN-28-∞.
  • An infinitely wide WideResNet: For finite widening factors, Monte Carlo kernel estimates remain biased, but the bias is small relative to variance and empirical distance decreases with more samples.The convergence experiment uses averaged finite-network kernels to study this sampling behavior.
  • Inference and architecture experiments: Neural Tangents supports constructing complex architectures and performing exact infinite-time inference with analytic infinite-width NNGP or NTK kernels.The examples compare fully connected, convolutional, and wide-residual architectures while selecting depth by training-set marginal log likelihood.
  • Comparison of neural network architectures and training set sizes: CIFAR-10 performance increases approximately logarithmically with dataset size, with architecture ranking FC < CONV < WRESNET, especially for larger datasets.The comparison uses infinite-time NTK inference and full Bayesian NNGP inference across fully connected, convolutional, and WideResNet families.

3 IMPLEMENTATION: TRANSFORMING TENSOR OPS TO KERNEL OPS

NEURAL TANGENTS translates compositional tensor operations into corresponding NNGP and NTK kernel operations, while optimizing covariance computation and distributing workloads across accelerators.

  • Tensor-to-kernel translation: Translation rules map finite-width tensor operations, including affine transforms, nonlinearities, pooling, and normalization, to infinite-width NNGP and NTK kernel transformations.The library applies these rules compositionally to supported architectures.
  • Tensor-to-kernel translation: Affine transformations before nonlinearities support the approximately Gaussian pre-activations assumed by many infinite-width results.Random weights and biases in infinite affine layers provide this Gaussianizing mechanism.
  • Performance: Block-diagonal output covariances reduce computational complexity and storage by an order of magnitude in many common classification cases, making exact inference feasible.The structure is Σ ⊗ I_C when fully connected readout logits are conditionally independent.
  • Performance: Covariance tracking reduces convolutional intermediate-layer costs from O(|X|^2 d^2) to O(|X|^2 d) per layer when no pooling is used, and to O(|X|^2) without convolutional layers.The optimization tracks only the covariance blocks needed to construct the output covariance.
  • Performance: Covariance propagation for convolutional layers can be expressed as 2D convolutions, enabling use of hardware accelerators optimized for those operations.The representation applies to both full covariance matrices and diagonal blocks.
  • Performance: Batching and automatic parallelization support analytic or empirical kernels across CPUs, GPUs, and TPUs, with near-perfect scaling across accelerators.Figure 5 reports near-perfect scaling with the number of accelerators for a 21-layer convolutional network.

4 CONCLUSION

NEURAL TANGENTS aims to make infinite-width neural networks easier to explore by providing a broader, more general library than earlier released implementations. The library remains limited by its lack of approximate inference tools and planned support for additional layers and performance improvements.

  • Conclusion: The authors expect NEURAL TANGENTS to enable researchers to explore infinite-width networks more quickly and use them as theoretical tools in deep learning.They especially highlight data-limited settings and theoretical investigations.
  • Future work: Future work includes adding more layers and implementing further performance improvements for larger models and datasets.The authors invite community contributions and feedback.
  • Prior work: Earlier infinite-width studies were often not open-sourced, comprehensive, user-friendly, or scalable enough to serve as general-purpose libraries.Separate implementations across prior works motivated a more general approach.
  • Prior work: The code released by Lee et al. was described as a strict and minor subset of NEURAL TANGENTS, with about +9,500/−2,500 lines separating the implementations.The comparison specifically notes that most notably nt.stax was new in this submission.
  • Limitations: NEURAL TANGENTS does not provide tools for approximate inference with its kernels.This is an explicit scope limitation of the library.

B LIBRARY DESCRIPTION

The library extends a hierarchical neural-network API with kernel functions for infinite-width models. These functions expose NNGP and NTK covariances for analytic inference and can be used with efficient prediction routines.

  • Library interface: NEURAL TANGENTS provides a high-level interface for specifying analytic infinite-width Bayesian and gradient-descent-trained neural networks as Gaussian Processes.Its interface follows the JAX stax API.
  • Library interface: stax components expose init_fn and apply_fn functions that can be composed serially or in parallel to specify complicated architectures hierarchically.init_fn generates parameters, while apply_fn evaluates a finite network on inputs.
  • Kernel functions: NEURAL TANGENTS adds kernel_fn as a third layer function representing the covariance functions of the architecture's infinite NNGP and NTK networks.Because infinite networks are Gaussian Processes, their covariance functions fully define them under the stated zero-mean convention.
  • Kernel functions: kernel_fn accepts two input batches and returns NNGP and NTK covariance matrices that can support posterior test predictions.The library also provides an efficient GP inference method in the predict submodule.

B.3 COMPUTING INFINITE NETWORK KERNELS IN BATCHES AND IN PARALLEL

The library batches kernel computation for large datasets and automatically parallelizes it across available devices, while supporting CPU-based inference for matrices that exceed accelerator memory.

  • The batch decorator returns a kernel_fn with the same signature that computes kernels in batches and parallelizes across available devices.This preserves the surrounding API while changing computation and distribution behavior.
  • Near-perfect speedup scaling with the number of devices is reported for automatically parallelized kernel calculations.
  • Batching supports covariance matrices that may not fit on a GPU or TPU by storing them in CPU RAM for inference.
  • A 45,000×45,000 training covariance matrix and a 45,000 × 10 target matrix can be solved in under 3 minutes on a CPU.For many convolutional and especially pooling architectures, inference cost can be small relative to kernel construction.

B.4 TRAINING DYNAMICS OF INFINITE NETWORKS

The library provides prediction functions for continuous gradient-descent training dynamics, covering MSE, arbitrary losses, and momentum.

  • Neural Tangents provides functions for gradient descent with MSE loss, arbitrary loss functions, and momentum with arbitrary loss.
  • MSE gradient-descent predictions are handled analytically, while arbitrary-loss and momentum cases are computed by numerically integrating the corresponding differential equations.
  • The functions compute train and test function evaluations after continuous gradient descent for a specified training time.

B.5 INFINITE NETWORKS OF ANY ARCHITECTURE THROUGH SAMPLING

For architectures without analytically computable kernels, Neural Tangents provides architecture-agnostic Monte Carlo estimators for NNGP covariance and NTK, with convergence demonstrated on WideResNet.

  • Monte Carlo utilities estimate NNGP covariance and NTK when the analytic kernel cannot be computed.They work with neural networks constructed using any neural network library.
  • The Monte Carlo kernel function can be constructed from initialization and apply functions with a specified number of random samples.
  • Convergence of Monte Carlo kernel estimates to closed-form analytic kernels is demonstrated for a WideResNet.
  • In a 5-layer width-512 Erf network trained on MNIST, lower-order Taylor expansions deviate from the original function faster than higher-order expansions.

B.6 WEIGHTS OF WIDE BUT FINITE NETWORKS

Neural Tangents supports weight-space analysis of wide finite networks through linearization and arbitrary-order Taylor expansions, implemented as ordinary JAX-compatible functions.

  • The library provides linearize and taylor_expand decorators to study weight-space behavior and approximations around initial parameters.
  • These decorators take an apply_fn and return a new apply_fn that computes the corresponding series approximation.
  • Taylor-expanded apply functions behave like normal JAX functions and can be plugged into gradient descent.
  • Adding a layer requires initialization, forward-pass, and kernel-propagation methods, with the kernel transformation potentially analytically intractable.
  • New layer implementations can be tested by comparing their kernels with empirical NNGP and NTK kernels from large-width networks.

C ARCHITECTURE SPECIFICATIONS

The examples define convolutional and fully connected architectures in Neural Tangents using composable layer lists and serial construction.

  • The all-convolutional model uses Flatten and repeated Dense layers before a final one-unit Dense output layer.
  • The fully connected model repeatedly applies one-unit Dense layers followed by ReLU, then adds a final one-unit Dense layer.

D IMPLEMENTED AND COMING SOON FUNCTIONALITY

Neural Tangents translates tensor operations into NNGP and NTK kernel operations while supporting selected layers, empirical approximations, and additional planned functionality.

  • Implemented functionality: Translation rules convert tensor operations into corresponding NNGP and NTK kernel operations for compositional network architectures.The rules cover kernels with dataset, spatial, and channel dimensions and preserve kernel structure through operations such as nonlinearities and convolutions.
  • Implemented functionality: Convolutional layers support arbitrary filter shapes, strides, and padding, while GlobalSelfAttention is also implemented.
  • Coming soon: Near-term plans include additional nonlinearities, Apache Beam support, and broader data-format support beyond NHWC.The listed planned nonlinearities are Exp, Elu, Selu, and Gelu; NHWC is currently the supported data format.
  • Limitations: Networks containing Sigmoid, Tanh, Swish, Softmax, LogSoftMax, Softplus, or MaxPool lack known closed-form infinite-network covariances and require empirical or approximate treatment.The library provides nt.monte_carlo_kernel_fn for empirical estimation, while other approximations are not currently implemented.
  • Evaluation: Figure 7 compares predictive negative log-likelihoods in its top panel and covariance condition numbers in its bottom panel for CIFAR-10 models.The models are fully connected and convolutional without pooling; the test set contains 2000 points.
Loading 1912.02803v1…