Neural Tangents: Fast and Easy Infinite Neural Networks in Python

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.

References in corpus (12)

• Wide Residual Networks
• Neural Tangent Kernel: Convergence and Generalization in Neural Networks
• A Convergence Theory for Deep Learning via Over-Parameterization
• Gradient Descent Provably Optimizes Over-parameterized Neural Networks
• Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data
• Toward Deeper Understanding of Neural Networks: The Power of Initialization and a Dual View on Expressivity
• Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation
• Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes
• ChainerMN: Scalable Distributed Deep Learning Framework
• Universal Statistics of Fisher Information in Deep Neural Networks: Mean Field Approach
• On the Selection of Initialization and Activation Function for Deep Neural Networks
• A Mean Field Theory of Quantized Deep Networks: The Quantization-Depth Trade-Off

Cited by in corpus (16)

• A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges
• Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains
• Machine Learning and Deep Learning -- A review for Ecologists
• Spectral Bias and Task-Model Alignment Explain Generalization in Kernel Regression and Infinitely Wide Neural Networks
• Improved machine learning algorithm for predicting ground state properties
• Asymptotics of representation learning in finite Bayesian neural networks
• Simple, Fast, and Flexible Framework for Matrix Completion with Infinite Width Neural Networks
• Wide and Deep Neural Networks Achieve Optimality for Classification
• Towards NNGP-guided Neural Architecture Search
• On the Neural Tangent Kernel of Deep Networks with Orthogonal Initialization
• Temperature check: theory and practice for training models with softmax-cross-entropy losses
• The Surprising Simplicity of the Early-Time Learning Dynamics of Neural Networks
• Rethinking Influence Functions of Neural Networks in the Over-parameterized Regime
• Metric Flows with Neural Networks
• Exploring the Uncertainty Properties of Neural Networks' Implicit Priors in the Infinite-Width Limit
• Constrained Policy Gradient Method for Safe and Fast Reinforcement Learning: a Neural Tangent Kernel Based Approach