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

Introduction

Deep neural networks (DNNs) owe their success in part to the broad availability of high-level, flexible, and efficient software libraries like Tensorflow (Abadi et al., 2015), Keras (Chollet et al., 2015), PyTorch.nn (Paszke et al., 2017), Chainer (Tokui et al., 2015; Akiba et al., 2017), JAX (Bradbury et al., 2018a), and others. These libraries enable researchers to rapidly build complex models by constructing them out of smaller primitives. The success of new machine learning approaches will similarly depend on developing sophisticated software tools to support them.

Recently, a new class of machine learning models has attracted significant attention, namely, deep infinitely wide neural networks. In the infinite-width limit, a large class of Bayesian neural networks become Gaussian Processes (GPs) with a specific, architecture-dependent, compositional kernel; these models are called Neural Network Gaussian Processes (NNGPs). This correspondence was first established for shallow fully-connected networks by Neal (1994) and was extended to multi-layer setting in (Lee et al., 2018; Matthews et al., 2018b). Since then, this correspondence has been expanded to a wide range of nonlinearities (Matthews et al., 2018a; Novak et al., 2019) and architectures including those with convolutional layers (Garriga-Alonso et al., 2019; Novak et al., 2019), residual connections (Garriga-Alonso et al., 2019), and pooling (Novak et al., 2019). The results for individual architectures have subsequently been generalized, and it was shown that a GP correspondence holds for a general class of networks that can be mapped to so-called tensor programs in (Yang, 2019). The recurrence relationship defining the NNGP kernel has additionally been extensively studied in the context of mean field theory and initialization (Cho & Saul, 2009; Daniely et al., 2016; Poole et al., 2016; Schoenholz et al., 2016; Yang & Schoenholz, 2017; Xiao et al., 2018; Li & Nguyen, 2019; Pretorius et al., 2018; Hayou et al., 2018; Karakida et al., 2018; Blumenfeld et al., 2019; Hayou et al., 2019).

2 Infinite-width neural networks trained by gradient descent

In addition to enabling a closed form description of Bayesian neural networks, the infinite-width limit has also very recently provided insights into neural networks trained by gradient descent. In the last year, several papers have shown that randomly initialized neural networks trained with gradient descent are characterized by a distribution that is related to the NNGP, and is described by the so-called Neural Tangent Kernel (NTK) (Jacot et al., 2018; Lee et al., 2019; Chizat et al., 2019), a kernel which was implicit in some earlier papers (Li & Liang, 2018; Allen-Zhu et al., 2018; Du et al., 2018a; b). In addition to this “function space” perspective, a dual, “weight space” view on the wide network limit was proposed in Lee et al. (2019) which showed that networks under gradient descent were well-described by the first-order Taylor series about their initial parameters.

3 Promise and practical barriers to working with infinite-width networks

Combined, these discoveries established infinite-width networks as useful theoretical tools to understand a wide range of phenomena in deep learning. Furthermore, the practical utility of these models has been proven by achieving state-of-the-art performance on image classification benchmarks among GPs without trainable kernels (Garriga-Alonso et al., 2019; Novak et al., 2019; Arora et al., 2019a), and by their ability to match or exceed the performance of finite width networks in some situations, especially for fully- and locally-connected model families (Lee et al., 2018; Novak et al., 2019; Arora et al., 2019b).

However, despite their utility, using NNGPs and NTK-GPs is arduous and can require weeks-to-months of work by seasoned practitioners. Kernels corresponding to neural networks must be derived by hand on a per-architecture basis. Overall, this process is laborious and error prone, and is reminiscent of the state of neural networks before high quality Automatic Differentiation (AD) packages proliferated.

4 Summary of contributions

In this paper, we introduce a new open-source software library called Neural Tangents targeting JAX (Bradbury et al., 2018a) to accelerate research on infinite limits of neural networks. The main features of Neural Tangents are:See §A for Neural Tangents comparison against specific relevant prior works.

A high-level neural network API for specifying complex, hierarchical, models. Networks specified using this API can have their infinite-width NNGP kernel and NTK evaluated analytically (§2.1, Listings 2.3, C, C, §B.2).

Functions to approximate infinite-width kernels by Monte Carlo sampling for networks whose kernels cannot be constructed analytically. These methods are agnostic to the neural network library used to build the network and are therefore quite versatile (§2.2, Figure 3, §B.5).

An API to analytically perform inference using infinite-width networks either by computing the Bayesian posterior or by computing the result of continuous gradient descent with an MSE loss. The API additionally includes tools to perform inference by numerically solving the ODEs corresponding to: continuous gradient descent, with-or-without momentum, on arbitrary loss functions, at finite or infinite time (§2.1, Figure 1, §B.4).

Functions to compute arbitrary-order Taylor series approximations to neural networks about a given setting of parameters to explore the weight space perspective on the infinite-width limit (§B.6, Figure 7).

Leveraging XLA, our library runs out-of-the-box on CPU, GPU, or TPU. Kernel computations can automatically be distributed over multiple accelerators with near-perfect scaling (§3.2, Figure 6, §B.3).

We begin with three short examples (§2) that demonstrate the ease, efficiency, and versatility of performing calculations with infinite networks using Neural Tangents. With a high level view of the library in hand, we then dive into a number of technical aspects of our library (§3).

5 Background

Examples

We begin by applying Neural Tangents to several example tasks. While these tasks are designed for pedagogy rather than research novelty, they are nonetheless emblematic of problems regularly faced in research. We emphasize that without Neural Tangents, it would be necessary to derive the kernels for each architecture by hand.

We begin by training an infinitely wide neural network with gradient descent and comparing the result to training an ensemble of wide-but-finite networks. This example is worked through in detail in the Colab notebook.www.colab.sandbox.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb

We train on a synthetic dataset with training data drawn from the process $y_{i}=\sin(x_{i})+\epsilon_{i}$ with $x_{i}\sim\text{Uniform}(-\pi,\pi)$ and $\epsilon_{i}\sim\mathcal{N}(0,\sigma^{2})$ independently and identically distributed. To train an infinite neural network with $\operatorname{Erf}$ activationsError function, a nonlinearity similar to $\tanh$ ; see §D for other implemented nonlinearities, including $\operatorname{Relu}$ . on this data using gradient descent and an MSE loss we write the following: