Neural Tangents: Fast and Easy Infinite Neural Networks in Python
Authors: Roman Novak
Lechao Xiao
Jiri Hron
Jaehoon Lee
Alexander A. Alemi
Jascha Sohl-Dickstein
Samuel S. Schoenholz
Venue:  (2020)
Keywords: Infinite Neural Networks
Gaussian Processes
Neural Tangent Kernel
NNGP
NTK
Software Library
Python
JAX

Abstract: 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. 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. 

In addition to the repository below, we provide an accompanying interactive Colab notebook at
https://colab.research.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb


Published as a conference paper at ICLR 2020
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'
Google Brain,
Unniversity of Cambridge
(romann, xlc)@google.com jh20840cam.ac ac.uk, Cjaehlee, alemi, jaschasd, schsam)@ggoglee.co
ABSTRACT
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 infimite- wiith limit. lnfinite-width networks can be trained analytically using exact Bayesian inference using gradient descent via the Neural Tangent Kernel. Additionally, NEURAL TANGENTS provides tools 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 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
www gigghhthhhh com/ hwwp///oglleeeeeeraaaannngnnnn
We also provide an accompanying interactive Colab notebook'.
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.n (Paszke et al., 2017), Chainer (Tokui et al., 2015; Akiba et al., 2017), JAX (Bradbury et al., 2018a), and others. These libraries enable researchers rapidly build complex models by constructing them out of smaller primitives. The success of new machine learning approaches will similarly depend developing sophisticated software tools support them.
1.1 INFINITE-WIDTH BAYESIAN NEURAL NETWORKS
Recently, 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 specific, architeture-dependdet. 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 multi- layer setting in (Lee et al., 2018; Matthews et al., 2018b). Since then, this correspondence has been expanded a wide range of nonlinearities (Matthews 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), pooling (Novak et al., 2019), as well as graph neural networks (Du et al., 2019). The results for individual architectures have subsequently been generalized, and it was shown that 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
Equul contribution. Work done during an internship at Google Brain. colah google
tangents cookbook.ipynh

---

Published as a conference paper at ICLR 2020
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).
1.2 INFINITEWIDTHH NEURAL NETWORKS TRAINED BY GRADIENT DESCENT
In addition to enabling 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 related to the NNGP, and is described by the so-called Neural Tangent Kernel (NTK) (Jacot al., 2018 8; Lee al., 2019; Chizat et al., 2019), kernel which was implicit in some earlier papers (Li & Liang, 2018; Allen-Zhu et al., 2018; Du et al., 2018;;b; Zou et al., 2019). In addition 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.
1.3 PROMISE AND PRACTICAL BARRIERS TO WORKING WITH NFINITE-WIDTH NETW ORKS
Combined, these discoveries established infinite-width networks as useful theoretical tools to under- stand 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 Alonso et al., 2019; Novak et al., 2019; Arora et al., 2019a),
GPs without trainable kernels and by their ability to match or exceed the performance of finite width networks in some situations, especially for fully- and locally- pooneey- model families (Lee et al., 2018; Novak 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.
1.4 SUMMARY OF CONTRIBUTIONS
In this paper, we introduce a new open-source software library called NEURAL TANGENTS targeting JAX (Bradbury et al., 201 8a) to accelerate research on infinite limits of neural networks. The main features of NEURAL TANGENTS are:
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 1, Listings 1, 2, 3, $B.2).
Functions to approximate infinite-width kernels by Monte Carlo sampling for networks whose kernels cannot be constructed analytically. These methods are agnostic the neural network library used to build the network and are therefore quite versatile ($2.2, Figure 2, $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 perform inference by numerically solving the ODES corresponding to: continuous gradient descent, with-or-withour momentum, on arbitrary loss functions, finite or infinite time 1$2.1, Figure $B.4).
Functions to compute arbitrary-order Taylor series approximations to neural networks about a given setting of parameters explore the weight space perspective on the wifth limit ($B.6, Figure 6).
Leveraging XLA, our library runs out-of-the-box on CPU, GPU, or TPU. Kernel computa- tions can automatically be distributed over multiple accelerators with near-perfect scaling ($3.2, Figure 5, $B.3).
"ee $A for NEURAL TANGENTS comparison against specific relevant prior works.
2

---

Published as a conference paper at ICLR 2020
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, then dive into number of technical aspects of our library ($3).
1.5 BACKGROUND
We briefly describe the NNGP ($ 1.1 and NTK ($1.2). NNGP. Neural networks are often structured as affine transformations followed by pointwise applications nonlinearities. Let (ar) describe the 7th pre-activation following linear transformation in n layer of neural network. At initialization, the parameters the network are randomly distributed and so central-limit theorem style arguments can be used to show that the pre-activations become Gaussian distributed with mean zero and are therefore described entirely by their covariance matrix K(x, = E[a] (a: (x))) This describes a NNGP with the kernel, (x), One can the NNGP make Bayesian posterior predictions point, I, which are Gaussian distributed with with mean uux)) - K(x, [(rr,& ())((,)))= and variance o² (xx) K(æ,, x) K(xx,, XX) X) - where (X))) the training set of inputs and targets respectively. NTK. When neural networks are optimized using continuous gradient descent with learning rate 77 on mean squared eπror (MSE) loss. the function evaluated on training points evolves as (X)) - (X) (fa(X) ))) where (+())) is the Jacobian of the output fr evaluated at X' and remains constant (o2 = +) throughout training and the time-evolutiog of the outputs can be solved in closed form as a Gaussian with mean fe(x) = Oz ()) (1 (x), y
test
J+ ( is the NTK. the infinite-wide
limit, the NTK
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 research. We emphasize that without NEURAL TANGENTS, would be necessary to derive the kernels for each architecture by hand.
2.1 INFERENCE WITH AN INFINITELY WIDE NEURAL NETWORK
We begin by training an infinitely wide neural network with gradient descent and comparing the result to training an ensemble of wide-but-fnite networks. This example is worked through in detail in the Colab notebook."
We train on synthetic dataset with training data drawn from the process y4 sin();; + G with r, ~ Uniform(-s) т) and (O0 o²) independently and identically distributed. To train an infinite neural network with Erf activations" on this data using gradient descent and an MSE loss we write the following:
from neural_ tuural_tnngent import predict, stax
init i_i,,fn, apply_fn, kernel_fn stax. serial(
stax.Densee (2048, W_std=1.5, 0.5t stax.Erf(0 stax. (2050 W_std=1.5, b_std=a.05) Erax. (ense(1, M_std=1.5, b_std=0.05))
y_mean, y_var = predict. (po_innerrencckernel.. x_train, y_train, x_test, 'ntk', digg_reg-1e-4, compute_cov=True)
The above code analytically generates the predictions that would result from performing gradient descent for infinite amount of time. However, it is often desirable to investigate finite-time learning dynamics of deep networks. This is also supported in NEURAL TANGENTS as illustrated in the following snippet:
hawpeeurraattggnnsshhhhoee www.eaahaaisseeeeeeeessicceeeeee Langents conkbook.iyynh EEror function, nonlincarity similar tanh; $D for other implemented nonlinearities, including Relu.
www.colah research. eesearch.google..o
3

---

Published as a conference paper at ICLR 2020
10
intrite Test
Bayesan inference
10
Finite Ensembbre
1.0
Gradient Descent
Fintle Ensemble
10
0.5
Loss
10
0.0
10"1
-0.5
10
-1.0
10
-1.5
10
10
10
10
-3
-2
-1
1
2
3
Step
Figure 1: Training dynamics for an ensemble of finite-width networks compared with an infi- nite network. Left: Mean and variance of the train and test MSE loss evolution throughout training. Right: Comparison between the predictions of the trained infinite network and the respective ensem- ble of finite-wide width networks. The shaded region and the dashed lines denote two standard deviations of uncertainty in the predictions for the infinite network and the ensemble respectively.
x_train, y_train, x_test, 'ntk', diag.reg=le-4, compute_cov=True) y_mean, y_var predict_fn(t=110 Predict the distribution 100.
predict_fn predict.gree seacceen_ddecent_mmeeeggg
The above specification set the hidden layer widths 2048, which has no effect on the infinite width network inference, but the init_fn and apply_fn fn here correspond to ordinary finite width networks. In Figure we compare the result of this exact inference with training an ensemble of one-hundred of these finite-width networks by looking at the training curves and output predictions of both models. We see excellent agreement between exact inference using the infinite-width model and the result of training ensemble using gradient descent.
2.2
INFINITELY WIDERESNET
The above example considers relatively simple network synthetic task. In practice may want to consider real-world architectures, and see how close they are to their infinite-width limit. For this task study variant of infinite-channel Wide Residual Network (Zagoruyko Komodakis, 2016) (WRN-28-oo). We first define both finite and infinite models within Listing 1.
We now study how quickly the kernel of the finite-channel WideResNet approaches its infinite channel limit. We explore two different axes along which convergence takes place: first, as a function of the number of channels measured by the widening factor, k) and second as a function of the number of lnite-network Monte Carlo samples we average over. NEURAL TANGENTS makes it easy to compute MC averages of finite kernels using the following snippet:
kernel_fn
cente_ccrlo__eenee feermel_fninnt_ffal apply_fn, rg__ey, n_samples sampled kerneeleeeeeeeeeeeeeeeeeeee kernel_fn(x,, x)
The convergence is shown in Figure 2. We that both the number of samples is increased or the network is made wider, the empirical kernel approaches the kernel of the infinite network. As noted in Novak et al. (2019), for any finite widening factor the MC estimate is biased. Here, however, the bias small relative to the variance and the distance to the empirical kernel decreases with the number of samples.
2.3 COMPARISON NEURAL NETWORK ARCHITECTURES AND TRAINING SET SIZES
The above examples demonstrate how might construct complicated architecture and perform inference using NEURAL TANGENTsNext we train range of architectures on CIFAR-10 and compare their performance function of dataset size. In particular, compare fully-connected network,
4

---

Published as conference paper at ICLR 2020
convolutional network whose penultimate layer vectorizes the image, and the wide-residual network described above. In each case, we perform exact infinite-time inference using the analytic infinite- width NNGP NTK. For each architecture perform lyperparameter search the depth of the network, selecting the depth that maximizes the marginal log likelihood on the training set.
def dideeeestettllooccccheeeee strides=(1, 1), channel_mismme (mismmtch==Falsse
Main stax. serial Relu(), stax. Conv(channels, (3, 3), strides, padding= SAME'),
stax. Relu(), stax. Conv(channels, (3, 3), padding=
Shortcut (stax. Identity) if not channel_mismen
else
stax, Cony (3, 3), strides, padding badding='SMME') "SMM"))
return stax. seriaa stax.serrallssea..aanottt)))) stax parallel (Main, Shortcut), stax. FaxxIaanIS))))
def MidefessNetGrouu(n, channels, strides= 1)):
blocks [hideResNetBlooccchameeLl, strides, channel smatch-True) for in range(n
blocks += (WideResNetBLooccCcaamiii (1, 1))J
return stax seriial
def (idoeesstet k, nun_classes):
return stax. (stxx Conv(16, (3, 3), padding= SAHE'),
KiceReskkrtrrouuoblockk.sszz int(16 k)),
WidbResstertGroup(bllokk.kk int k), (2, 2)),
WieeReswweerrruupbboookkssizz int(64 k), (2, 2)),
stax, GlobalAvgPool stax Dense(num_classes)
init_fn, ap,ly_fn, kernel_fn (ideRResNettblook.ssieee k=1, num_claasees classeess)))))
Listing 1: Definition of an infinitely WideResNet. This snippet simultaneously defines a finite (init_fn, apply_fn) and an infinite kernel_fn] model. This model used in Figures 2 and 3.
NNGP
NTK
20 widening fector) 51
1og21 widening factor)
log1 (DDistanne))
log Distance))
log2( #Saplees
10
(og2(##Smmples))
Figure 2: Convergence of the Monte Carlo (MC) estimates of the WideResNer WRN-28-k (where A is the widening factor) NNGP and NTK kernels (computed with monte_carlo waw.rooeeeeeeeennnnceeaam
to their analytic values (WRN-28-c0, computed with kernel_fn as the network gets wider by increasing the widening factor (vertical axis) and more random networks are averaged over (horizontal axis). Experimental detail. The kernel is computed in 32-bit precision batch of 8 x 8-downsampled CIFAR (Krizhevsky, 2009) images. For sampling efficiency, for NNGP the output of the penultimate layer was used, and for NTK the output layer was assumed to be of dimension (all logits are i.i.d. conditioned on a given input). The displayed distance relative Frobenius norm squared, i.c. ||A Kx,ml1 ||A]] where the widening factor and is the number of samples.
100 x 50

---

Published as a conference paper at ICLR 2020
0.80
0.095
3.5
0.75
0.090
3.0
WRESETT
0.085
0.65
0.60
0.080
Classificatio
0.55
0.50
TORTC
0.070
COMNNTC
0.40
COWON9C7
0.065
1.0
0.35
2' 2" 2 2/2
0.060
2' 2'
2'h 20221 22
0.5
2" 2 2n 2222
Training Set Size
Training Set Size
Training Set Size
Figure 3: CIFAR-10 classification with varying neural network architectures. NEURAL TAN- GENTS simplify experimentation with architectures. Here we use infinite time NTK inference and full Bayesian NNGP inference for CIFAR- 10 for Fully Connected (FC, Listing 3), Convolutional network without pooling (CONV, Listing 2), and Wide Residual Network w/ pooling (WRES- NET, Listing 1). As common prior work (leee et al., 2018; Novak et al., 2019), the classification task is treated as MSE regression on zero-mean targets like (-0.1,.... -0.1, 0.9. -..1,
0.1)
For each training set size, the best model in the family is selected by minimizing the mean negative marginal log- lellikklihood (NLL, right) on the training set.
The results shown in Figure 3. We that in each case the performance of the model increases approximately logarithmically in the size of the dataset. Moreover, we observe a clear hierarchy of performance, especially at large dataset size, in terms of architecture (FC CONV < WRESNET w/ pooling).
3 IMPLEMENTATION: TRANSFORMING TENSOR OPS TO KERNEL OPS
Neural networks are compositions of basic tensor operations such as: dense or convolutional affine transformations, application of pointwise nonlinearities, pooling, or normalization. For most networks without weight tying between layers the kernel computation can also be written compositionally and there is a direct correspondence between tensor operations and kernel operations (see $3.1 for an example). The core logic of NEURAL TANGENTS a set of translation rules, that sends each tensor operation acting finite-width layer to corresponding transformation of the kernel for an infinite-width network. This is illustrated in Figure for a simple convolutional architecture. In the associated table, we compare tensor operations (second column) with corresponding transformations of the NT and NNGP kernel tensors (third and fourth column respectively). See $D for a list of all tensor operations for which translation rules are currently implemented.
One subtlety to consider when designing networks is that most infinite-width results require nonlinear transformations be preceded by affine transformations (either dense or convolutional). This is be- cause infinite-width results often assume that the activatass of nonlinear layers are approximately Gaussian. Randomness in weights and biases causes the output of infinite affine layers to satisfy this Gaussian requirement. Fortunately, prefacing nonlinear operations with affine transformations is common practice when designing neural networks and NEURAL TANGENTS will raise an error if this requirement is not satisfied.
3.1 A TASTE OF TENSOR-TO TO- KERKELL OPS TRANSLATION
To get some intuition behind the translation rules, we consider the case of nonlinearity followed by dense layer. Let z (X, 0) E [ be the preactivations resulting from d distinct inputs at node in some hidden layer neural network. Suppose 2 has NNGP kernel and NTK given by
0z:
O22
T
k, = Eg [2
O = Eg
(1)
00
a0
where 2; E Rd is the th neuron and 9 are the parameters the network up until the cardinality of the network inputs X' and is the number of neurons in the node. We
z. Here

---

Published as a conference paper at ICLR 2020
y'
CNN (1D) activation
10 classes per input
convekrtin
inputs, features,
X
R
R
ONN- covariance
per class
464x10 Input covariance
(15 ming recooaay track covenanoi)
Layer
Tensor Op
NNGP Op
NTK Op
(input)
y X'
K XX
θ =
O (pre-activations)
2 Conv (v")
K = A (1C)
@ K + A (09)
1 (activations)
y (=")
K T (K9)
O f (CO 0
(pre-activations)
Conv (v")
Kt A (KC)
Ô R + (0!)
(activations)
2² & (x")
K² (R))
@ F (C)) 0
2 (readout)
2 Dense o Flatten (v²)
R2 = Tr (K3)
@ R2 + Tr (0?)
Figure 4: An example of the translation of a convolutional neural network into a sequence of kernel operations. We demonstrate how the compositional nature of a typical NN computation on its inputs induces a corresponding compositional computation on the NNGP and NT kernels. Presented 2-hidden-layer ID CNN with nonlinearity é. performing regression on the 10-dimensional outputs 2 for each of the 4 (1, 2, 3, 4) inputs r from the dataset X' To declutter notation, unit weight and zero bias variances are assumed in all layers. Top: recursive output (z²) computation in the CNN (top) induces respective recursive NNGP kernel (K:2 & J10) computation (NTK computation being similar, not shown). Bottom: explicit listing of tensor and corresponding kernel ops in each layer. See Table for operation definitions. Illustaaioon and description adapted from Figure 3 in Novak et al. (2019).
assume z is a mean zero multivariate Gaussian, We wish to compute the kernel corresponding h Dense (aw, σb) (0())) by computing the kernels of y Φ(z) and h Dense (aw, σb) (y) separately. Here,
= Dense(oo,, = (1/vR) aw W +
(2)
and the variables Wiz and B, are i.i.d Gaussian (1, 1). We will compute kernel operations denoted 6² and Dense(aw, (x))"""""" induced by the tensor operations e and Dense(a op) S. FFnally, we will compute the kernel operation associated with the composition (Dense σb) o o)) Dense(aw. 0b)" d".
First we compute the NNGP and NT kernels for y. To compute Kr note that from its definition,
Ky = Ko(()) = Ea [+(2), (<))7] = F. (d(z;) ((2)) T(K.).
(3)
Since & does not introduce any new variables O, can be computed as,
(b((2))
7
0z,
00z: T
θ, = E.
E, (())()(=)))
diag(ó(z;))
= +(KC.) @ Oz.
00
00
00
00
Taken together these equations imply that,
(Ky, θ,)) = 0- (Kz, O2) (T(K.). =T)=(T(()).. T(K,) o0.)
(4)
$T (T.) = F [((u))(())}] =())) [dd) N(O,
(0. E), in (Lee al., 2019).
7

---

Published as a conference paper at ICLR 2020
will be the translation rule for a pointwise nonlinearity. Note that Equation Equation 4 only has an analytic expression for a small set of activation functions ø.
Next consider the case of dense operation. Using the independence between the weights, the biases, and h it follows that,
K Ew.a,o [hch]]] - Bo
or Ky + oz
(5)
Finally, the NTK of h can be computed as a sum of two terms:
ah, ah; @(W,3) 3)
)
Oh
I Ew,,,
Ew,3,
ouky + oī + 00 (6)
a(W, B)
00
a0
This gives the translation rule for the dense layer in terms of K.y and O as,
(Kns, Oa)
Dense(aw, ob)" (Ky, Oyy (o + og Ka + +o +
(7)
3.2 PERFORMANCE
Our library performs number of automatic performance optimizations without sacrificing flexibility.
Leveraging block-diagonal covariance structure. A common computational challenge with GPs is inverting the training covariance matrix. Naively, for classification task with C classes and training set X, NNGP and NTK covariances have the shape of | C' x |X| C. For CIFAR-10 10, this would be 500, 000 x 500, 000. However, if a fully-connected readout layer is used (which is an extremely common design in classification architectures), the C' logits are i.i.d. conditioned on the input x. This results in outputs that are normally distributed with block-diagonal covariance matrix of the form Σ @ Ic. where Σ has shape | x [ and Ic is the C' x C identity matrix. This reduces the computational complexity and storage many common cases by an order of magnitude, which makes closed-fone exact inference feasible in these cases.
Automatically tracking only the smallest necessary subset of intermediary covariance entries. For most architectures, especially convolutional, the main computational burden lies in constructing the covariance matrix (as opposed to inverting it). Specifically for a convolutional network of depth I, constructing the | x | output covariance matrix, Σ, involves computing intermediate layer covariance matrices, Σ', of size | d x X d (see Listing for a model requiring this computation) where d is the total number of pixels in the intermediate layer outputs (e.g. d 1024 in the case of CIFAR-10 with SAME padding). However, as Xiao et al. (2018); Novak et al. (2019); Garriga- Alonso et al. (2019) remarked, if no pooling is used in the network the output covariance Σ can be computed by only using the stack of dl [ |blocks of Σ', bringing the time and memory cost from ( d2) down to ( d) per layer (see Figure and Listing for models admitting this optimization). Finally, the network has no convolutional layers, the cost further reduces to ( (see Listing 3 for example). These choices are performed automatically by NEURAL TANGENTS to achieve efficient computation and minimal memory footprint.
Expressing covariance computations as 2D convolutions with optimal layout. A key insight to high performance in convolutional models is that the covariance propagation operator for convo- lutional layers can be expressed in terms of 2D convolutions when it operates on both the full |X| d x | d covariance matrix Σ, and on the d diagonal |X| x |lll-blocks. This allows utilization of modern hardware accelerators, many of which target 2D convolutions as their primary machine leaming application.
Simultaneous NNGP and NT kernel computations. As NTK computation requires the NNGP covariance as an intermediary computation, the NNGP covariance is computed together with the NTK at no extra cost. This especially convenient for researchers looking to investigate similarities and differences between these two infinite-width NN limits.
Automatic batching and parallelism across multiple devices. In most cases as the dataset or model becomes large, it is impossible to perform the entire kernel computation at once. Additionally, in many cases is desirable to parallelize the kernel computation across devices (CPUs, GPUs, or TPUs). NEURAL TANGENTS provides easy way to perform both of these common tasks using single batch decorator shown below:
8

---

Published as a conference paper at ICLR 2020
2
2
GPU Count
Time / Sample (s)
23
2
GPU Count
2.
GPU Count
Time Sample
2
27
2
2.
2
2
2
2
2
2
2
2a
2
2
2a
2b
Batch Size
GPU Count
Figure 5: Performance scaling with batch size (left) and number of GPUs (right). Shows time per entry needed to compute the analytic NNGP and NTK covariance matrices (using kernel_fn in 21 21-layer ReLU network with global average pooling. Left: Increasing the batch size when computing the covariance matrix in blocks allows for a significant performance increase until certain threshold when all cores in single GPU are saturated. Simpler models expected have better scaling with batch size. Right: Time-per-sample scales linearly with the number of GPUs, demonstrating near-perfect hardware utilization.
Batched_kerrne__f nt. nt..bacchkkeree_.f batch_size) kntchee_.kerne__fntk kernel_fn(x, x) True!
This code works with either analytic kernels empirical kernels. By default, it automatically shares the computation over all available devices. We plot the performance as function of batch size and number of accelerators when computing the theoretical NTK of 21-layer convolutional network in Figure 5, observing near-perfect scaling with the number of accelerators.
Op fusion. JAX and XLA allow end-to-end compilation of the whole kernel computation and/or inference. This enables the XLA compiler to fuse low-level ops into custom model-specific accelerator kernels, as well as eliminating overhead from op-by-op dispatch an accelerator. In similar vein, we allow the covariance tensor to change its order of dimensions from layer to layer, with the order tracked and parsed additional metadata under the hood. This eliminates redundant transpositions' by adjusting the computation performed by each layer based on the input metadata.
4 CONCLUSION
We believe NEURAL TANGENTS will enable researchers to quickly and easily explore infinite-width networks. By democratizing this previously challenging model family, hope that researchers will begin to use infinite neural networks, in addition to their finite counterparts, when faced with a new problem domain (especially in cases that are data-limited,. In addition, we are excited to see novel uses of infinite networks as theoretical tools to gain insight and clarity into many of the hard theoretical problems in deep learning. Going forward, there are significant additions to NEURA TANGENTS that we are exploring. There are more layers we would like to add in the future (SD) that will enable an even larger range of infinite network topologies. Additionally, there are further performance improvements we would like to implement, to allow experimenting with larger models and datasets. We invite community to join our efforts by contributing new layers to the library ($B.7), by using for research and providing feedback!
ACKNOWLEDGMENT
We thank Yasaman Bahri for frequent discussion and usefuul feedback on the manuscript. We additionally appreciate both Yasaman Bahri and Greg Yang for the ongoing contributions improve the library. We thank Sergey loffe for feedback on the text, as well as Ravid Ziv, and Jeffrey Pennington for discussion and feedback on early versions of the library.
'These transpositions could not be automatically fused by the XI.A compliler.

---

Published as a conference paper at ICLR 2020
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangging Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke. Vijay Vasudevan, Fermanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg. Martin Wicke. Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL http://tennorffoon org/. Software available from tensorflow.org.
Takuya Akiba, Keisuke Fukuda, and Shuji Suzuki. ChainerMN: Scalable Distributed Deep Learning Framework. In Proceedings of Workshop on MI. Systems in The Thirty-forst Annual Conference on Neural Information Processing Systems (NIPS), 2017. URL assets/papers blsstts/papess/appee _____.5pdd. 2d5.f
nips///
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. International Conference on Machine Learning. 2018.
Anonymous. Infinite attention: Nugp and ntk for deep attention networks. In International Conference on Machine Learning (ICML), 2020. submission under review.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances In Neural Information Processing Systems, 2019a.
Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets on small-data tasks, 2019b.
Yaniv Blumenfeld, Dar Gilboa, and Daniel Soudry. A mean field theory of quantized deep networks: 2019.
The quantization-depth trade-off. arXiv preprint
James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, and Skye Wanderman-Milnn JAX: composable transformations of Python+NumPP programs, 2018a. URL http://gi thub. com cow.:oogl///oeoo////ooooocooocce google
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, and Skye Wanderman-Milne. Stax, a flexible neural net specification library in jax, 2018b. URL https: gi thub. com www.///ww////o//oooooooooooooooo
http:///ttteepeeea/aa/aaaaaaaa/a imental/stax....mm py.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. arXiv preprint arXiv: 18/207995 2019.
Youngmin Cho and Lawrence K Saul. Kernel methods for deep learning. In Advances In Neural Information Processing Systems, 2009.
François Chollet et al. Keras. https: //kerss// io, 2015.
Amit Daniely, Roy Frostig, and Yoram Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. In Advances In Neural Information Processing Systems, pp. 2253-2261, 2016.
Simon Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv: /8/ 10881..80 1.08044 2018a.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv: 1810.0054, 2018b.
Simon S Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kemels.. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp.
10

---

Published as a conference paper at ICLR 2020
Curran Associates, Inc., 2019. URL http://pppers.nnpsccccppppp -erra---urrl-ttangee---- htw.://wwww............ooootoona pdf.
Jacob R Gardner, Geoff Pleiss, David Bindel, Kilian Q Weinberger, and Andrew Gordon Wilson Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration. In Advances in Neural Information Processing Systems, 2018.
Adrià Garriga-Alnso,, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks as shallow gaussian processes. In International Conference on Learning Representations, 2019.
GPy. GPy: 2012.
gaussian process framework in python. http:///ggttbb..cc////fffeeeell
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the selection of initialization and activation function for deep neural networks. arXiv preprint 1805 080666, 2018.
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. Mean-field behaviour of neural tangent kernel for deep neural networks, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances neural information processing systems, 2018.
Ryo Karakida, Shotaro Akaho, and Shun-ichi Amari. Universal statistics of fisher information in deep neural networks: mean field approach. 2018.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Sam Schoenholz, Jeffrey Pennington, and Jascha Sohl-dickstein. Deep neural networks as gaussian processes. In International Conference Learning Representations, 2018.
Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in neural information processing systems, 2019.
Ping Li and Phan-Minh Nguyen. On random deep weight-tied autoencoders: Exact asymptotic analysis, phase transitions, and implications to training. In International Conference on Learning Representations, 2019. URL https: /ttt:::///peneevieee oet.://ooruum....H)5441111..
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in Neural Information Processing Systems, pp. 8157-8166. 2018.
Alexander G. de G. Matthews, Mark der Wilk, Tom Nickson, Keisuke. Fujii, Alexis Boukouvalas, Pablo Le" on- Son-Villagr's a, Zoubin Ghahramani, and James Hensman. GPflow: A Gaussian process library using TensorFlow. Journal of Machine Learning Research, 18(40): 1-6, apr 2017. URL http://jmlr.. ///p.//////ppep///////////// 81181/37.37771 html.
Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint 17X.v:1804.1222 2018a.
Alexander G. G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference Learning Representations, 2018b.
Radford M. Neal. Priors for infinite networks (tech. rep. no. crg-tr-94-1). University of Toronto, 1994.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels gaussian processes. In International Conference Learning Representations, 2019.

---

Published as a conference paper at ICLR 2020
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. In NIPS- W, 2017.
Ben Poole. Subhaneil Lahiri, Maithra Raghu, Jascha Sohl- Dickksics,,,n, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. In Advances In Neural Information Processing Systems, 2016.
Arnu Pretorius, Elan van Biljon, Steve Kroon, and Herman Kamper. Critical initialisation for deep sig- nal propagation in noisy rectifier neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grau- man, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31. pp- 5717-5726. Curran Associates, Inc., 2018. 8. URL 781 pdf.
nips. cc/paper/ for- ooapllpppppaaattooooooooooooeeee
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint 1611v:1611.0/220 2016.
Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: next-generation open source framework for deep learning. In Proceedings of Workshop on Machine Learning Systems (LearningSys) The Twenty-ninth Annual Conference on Neural Information Processing Systems (NIPS), 2015. URL http: /wte//eaaannnnns.eeaera org/ hwtpssssssrggppps//eeeeee/ waapers/Lerrinngyyrr SeaarninnSys_215__ppp. 2aapsss__21__pppprrrr 33. 33.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Penning- ton. Dynamical isometry and mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks. In International Conference on Machine Learning, 2018.
Lechao Xiao, Jeffrey Pennington, and Samuel S Schoenholz. Disentangling trainability and general- ization deep learning. arXiv preprint arXiv:1912.1305: 2019.
Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge chaos. In Advances In Neural Information Processing Systems, 2017.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint 1rXY:11902.04760. 2019.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference (BMVC), 2016.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine Learning, Oct 2019. ISSN 1573-0565. doi: 10.11077/400994401995533.. URL https: 0ttp:///..00///1110400110-----00
APPENDIX
NEURAL TANGENTS AND PRIOR WORK
Here we briefly discuss the differences between NEURAL TANGENTS and the relevant prior work.
1. Prior benchmarks in the domain of infinitely wide neural networks. Various prior works have evaluated convolutional and fully-connected models on certain datasets (Lee et al., 2018; Matthews et al., 2018b;a; Novak et al., 2019; Garriga-Alons, et al., 2019; Arora et al., 2019a). While these efforts must have required implementing certain parts our library, to our knowledge such prior efforts were either not open-sourced or not comprehensive user-friendly scalable enough to be used addition, all of the works above used their own separate implementation, which further highlights& need for more general approach.
user-facing library. In
12

---

Published as a conference paper at ICLR 2020
2. Code released by Lee et al. (2019). Lee et al. (2019) have released code along with their paper submission, which strict and minor subset our library. More specif- ically, the time of the submission, Lee et al. (2019) have released code equivalent nt linearize, et.empprical_ttk_f
nt. predict. gr arriient_descent_meee
gradient_descent and predict. momentum Every other part of the nt.preddct.graadienn_.deee
library (most notably, stax is new in this submission and was not used by Lee et al. (2019) or any other prior work. At the time of writing, NEURAL TANGENTS differs from the code released by Lee et al. (2019) by about / 2, 500 lines of code.
3. GPy (2012), GPFlow (Matthews et al., 2017), GPy Toreh (Gardner al., 2018), and other GP packages. While various packages allowing for kernel construction, optimization, and inference with Gaussian Processes exist, none of them allow easy construction of the very specific kernels corresponding infinite neural networks (NNGP/NTK: stax nor do they provide the tools and convenience for studying wide but finite networks and their training dynamics ( nt. taylor_expand nt. predict,, mottme careeeccallloo caarlokkeneelff ) On the other hand, NEURAL TANGENTS does not provide any tools for approximate inference with these kernels.
B LIBRARY DESCRIPTION
NEURAL TANGENTS provides high-level interface for specifying analytic, infinite-width, Bayesian and gradient descent trained neural networks as Gaussian Processes. This interface closely follows the stax API (Bradbury al., 2018b) in JAX.
B. NEURAL NETWORKS WITH JAX
stax represents each component of network as two functions: init_fn and apply_fn These components can be composed in serial or in parallel to produce new network components with their own init_fn and apply_fn In this way, complicated neural network architectures can be specified hierarchically.
Calling init_fn parameters for neural network. Calling apply_fn on these parameters and batch of inputs returns the outputs of the given finite neural network.
random seed and input shape generates random draw of trainable
from jax. experrieimm import stax
init_fn, apply_fn = stax. ccax.serial (ss..rrss(sst...eees(((111 stax. Relu, stax. Dense(18))
params = init_fn(key, (-1, 32 32 * 3))
fx_train, fx_test apply_fn(parans, x_train), apply_fn(params, x_test)
B.2 INFINITE NEURAL NETWORKS WITH NEURAL TANGENTS
We extend stax layers to return third function kernel_fn which represents the covariance functions the infinite NNGP and NTK networks of the given architecture (recall that since intinite networks GPs, they are fully defined by their covariance functions, assuming O mean is common in the literature).
from neural_ targgentttngnssst import stax
init_fn, apply_fn, kernel_fn stax. (stax ( stax. Relu(), stax Dense
We demonstrate specification more complicated architecture (WideResNet) in Listing
kernel_fn fn accepts two batches of inputs x1 and x2 and returns their NNGP covariance and NTK matrices as kernel_fnxx1, .nngnp and kernel_fn(x1, x2) ntk respectively, which
13

---

Published as a conference paper at ICLR 2020
then be used to make posterior test set predictions the mean of conditional multivariate normal Viess IC (Kistt Xirain) K. (Airain: Virain) M Virain-
from jax. nox.nnumy. 1inalg import inv y_test kernel_fn(x_test, x_train) ntk invckernel_fn(xxtrrinn x_train) .ntk) y_train
Note that the above code does not do Cholesky decomposition and is presented merely to show the mathematical expression. We provide efficient GP inference method in the predict submodule:
import gural__angents nt y_test nt. wwwpeptt.gp.infeeeneee....eeeeee pe_iineerencekerne_L... x_train, y_train, x_test,
net='nt',, diag_reg: compute_cov==alse)
B.3 COMPUTING INFINITE NETWORK KERNELS IN BATCHES AND IN PARALLEL
Naively, the kernel_fn fn will compute the whole kernel in single call on one device. However, for large datasets or complicated architectures, it is often necessary to distribute the calculation in some way. To do this, we introduce batch decorator that takes a kernel_fn and returns new kernel_fn with the exact same signature. The new function computes the kernel in batches and automatically parallelizes the calculation over however many devices are available, with near-perfect speedup scaling with the number of devices (Figure 5, right).
import neural_tangents kernel_fn nt. bstchhhe batch_size=32)
Note that batching is often used to compute large covariance matrices that may not even fit on GPU/TPU device, and require to be stored and used for inference using CPU RAM. This is easy Once the matrix is stored in RAM, inference will be performed with CPU when nt. itteicc methods are called. As mentioned in $3.2, for many (notably, convolutional, and especially pooling) architectures, inference cost can be small relative kernel construction, even when running on CPU (for example, it takes less Linalg.solve(.................._
to achieve by simply specifying nt. at..batch....
storr_on_ddeice=Faal....
than minutes to execute jax. scipy linalg. solve(.... jom_pos=True) on 45, 000 000
training covariance matrix and a 45, 000 x 10 training target matrix).
B.4 TRAINING DYNAMICS OF INFINITE NETWORKS
In addition closed form multivariate Gaussian posterior prediction, it also interesting consider network predictions following continuous gradient descent. To facilitate this we provide several functions compute predictions following gradient descent with MSE loss, for gradient descent with arbitrary loss, or for momentum with arbitrary loss. The first case is handled analytically, while the latter are computed by numerically integrating the differential equation. For example, the following code will compute the function evaluation on train and test points following gradient descent for some time training_tinm time
import neural_tangents nt predictor www.pedict.grddeettdeeeee
x_train), y_train,
fx_train, fx_test predictor (eafiitto(traaiingg.t... time, fx_train, fx_test)
B.5 INFINITE NETWORKS OF ARCHITECTURE THROUGH SAMPLING
Of course, there are cases where the analytic kernel cannot be computed. To support these situations, we provide utility functions to efficiently compute Monte Carlo estimates of the NNGP covariance and NTK. These functions work with neural networks constructed using any neural network library.
14

---

Published as a conference paper at ICLR 2020
Batch Mean Squared Error
Oth order (constant)
10-1
1st order (linear)
2nd order (quadratic)
10-2
Original function
0
4 Training Epoch
8
10
Figure 6: Training a neural network and its various approximations using nt. tayloayyyllpppoeeee Presented is a 5-layer Erf-neural network of width 512 trained on MNIST using SGD with mo- mentum, along with its constant (0th order), linear (1st order), and quadratic (2nd order) Taylor expansions about the initial parameters. As training progresses (left to right), lower-order expansions deviate from original function faster than higher-order ones.
from jax import random from eax.experimental import stax import neural_tangents as nt
init_fn, apply_fn stax, ( stax. BatchNorm(), stax Signoid, stax Dense(1)) apply_fn, key=random. PRNGKey(1), n_samples=128)
kernel_fn nt. kernel kernel_fn(x_train, x_train)
We demonstrate convergence the Monte Carlo kernel estimates to the closed-fone form analytic kernels in the case of WideResNet in Figure 2.
B.6 WEIGHTS OF WIDE BUT FINITE NETWORKS
While most of NEURAL TANGENTS is devoted to function-space perspective-deserbbiing the distribution of function values on finite collections training and testing points points-we ww also provide tools to investigate dual weight space perspective described in Lee et al. (2019). Convergence of dynamics NTK dynamics coincide with networks being described by linear approximation about their initial set of parameters. We provide decorators linearize and taylor_expand to approximate functions to linear order and to arbitrary order respectively. Both functions take an apply_fn and returns a new apply_fn that computes the series approximation.
import neural_tangents as nt petbo__expand(pply_ffi params, order)
taylor_apply_fnn www.ttylll___ppe
fx_train_approx
x_train)
These act exactly like normal JAX functions and, in particular, can be plugged into gradient descent, which we demonstrate Figure 6.
B.7 EXTENDING NEURAL TANGENTS
Many neural network layers admit sensible infinite-width limit behavior in the Bayesian and continuous gradient descent regimes as long as the multivariate central limit theorem applies to their outputs conditioned their inputs. To add such layer NEURAL TANGENTS, only has to implement it as method in nt. stax with the following signature:
15

---

Published as a conference paper at ICLR 2020
8_layer an internal decorator taking care certain bollerplate. New.ayer(layer_parae Any) -> (init_fn: function, apply_fn: function, kernel_fn: function)
Here init_fn and apply_fn initialization and the forward pass methods of the finite width layer implementation (see $B.1). If the layer of interest already exists in JAX, there no need to implement these methods and the user can simply return the respective methods from eex.eeiremmettll. stax (see nt stax. Flatten for an example; in fact the majority of nt. stax layers call the original jax. exper stax layers for finite width layer methods).
In this case what remains is to implement the kernel_fn method with signature
nt. ut.uttis. kernel_fnnchnuut_kkrnll
kernel output_kernel: nt. ur..util... Kernel
Here both input_kernel and output_kerne are namedtuple containing the NNGP and NTK covariance matrices, as well as additional metadata useful for computing the kernel propagation operation. The specific operation to be performed should be derived by the in the context of the particular operation that the finite width layer performs. This transformation could be as simple as an affine map the kernel matrices, but could also be analytically intractable.
Once implemented, the correctness the implementation can be very easily tested by extending the nt. tests. stax_ test with new layer, to test the agreement with large-widths empirical NNGP and NTK kernels.
C ARCHITECTURE SPECIFICATIONS
from neural_tangents import stax
def
W_std=1.0, b_std=0.0) 0):
layers []
for in range(depth):
layers [stax. Conxx,Covv1,,
W_std, b_std, padding='SAME Relu()
layers [stax [stax.Flatted Flatten(), stax. Dense(1, W_std, b_std)]
return stax. (*layers))
Listing 2: All-convolutional model (ConvOnly) definition used in Figure 3.
from neural_tangents import stax
def (cllyCConnnttodWeeeorkddddhhhh W_std=1.0, b_std=0.0) .):
layers [stax.Flatten, ()]
for in range(:)
layers += [stax. Dense(1, W_std, b_std), stax. stax. Relu()]
layers += Dense(1, W_std, b_std)]
return stax. (slayer))
Listing 3: Fully-connected (FC) model definition used in Figure 3.
16

---

Published as a conference paper at ICLR 2020
D IMPLEMENTED AND COMING SOON FUNCTIONALITY
The following layers' are currently implemented, with translation rules given in Table 1:
serial
parallel
FanOut
FanInSum
FanInConcat
Dense
Conv with arbitrary filter shapes, strides, dimension numbers, and padding
Relu
LeakyRelu
Abs
ABRelu
Erf
Identity
Flatten
AvgPool
GlobalAvgPool
SumPool
GumPol
Dropout
LayerNorm
PettallsllAtttettio (Anonymous, 2020)
The following is in our near-term plans:
Exp..
Elu
Selu.
Gelu
Apache Beam support.
The following layers not have known dlosed-form expressions infinite network covariances, and respective infinite networks have to be estimated empirically (via nt. coont__call_kkenella. or using other approximations (not currently implemented):
Sigmoid
Tanh
Swish
Sof tmft.ax
LogSof
Softplus
MaxPool
7 Abs ABRelu Erf GlobalAvgPool
GlobalSumPool
LayerNorm and GlobalseeFAttentiion
only available in library stax and not in jax. experimental. stax
'In addition to SAME and VALID we support CIRCULAR padding, which especially convenient for theoretical analysis and was used by Xiao et al. (2018) and Novak et al. (2019).
ABRelu is more flexible extension of LeakyRelu computed as a min (a, 0) b max (a, 0). 10 Sigmoid and Tanh are similar to Erf which does have covariance expression and is implemented. Swish similar to Gelu
17

---

Published as a conference paper at ICLR 2020
Tensor Op
NNGP Op
NTK Op
X
K
θ
Donse(ow, σb)
o + oo
(o + +ot) +o
T(K)
T(K) 0
K+ (1- -1) Diag(K))
+ + (1 (1-1) -01)) Diag(H)
Dropout(p)
Conv(ow, σb)
A (IC) + +f
2 (K) + o + A
Flatten
Tr(K)
Tr(K + +)
AvgPool(s, q.P)
AvgPool(s, (s,9,))
AvgPool(s.
GlobalAvgPool
GlobalAvgPool(k)
GlobalAvgPool +
SumPool q.
SumPool(s.g.P)(3
SumPool(s, q. +
GlobalSumPool
GlobalSumPool(K)
GlobalSumPPol((3 + +
Att(OQQ, gov) (Anonymous, 2020)
Atn((QQ,
20tn(ook,
Aun(oqk, 0ov (()
FanInSum ,,
An) Σz= Kr
Σ θ
FanOut(n)
KK] * n
[ n
Table 1: Translation rules ($3) converting tensor operations into operations on NNGP and NTK kernels. Here the input tensor X' is assumed to have shape | x H x W x C' (dataset size, height, width, number of channels), and the full NNGP and NT kernels K. and T are considered to be of shape (Ix) H W)*² (in practice shapes of | x2 H W and | X2 are also possible, depending on which optimizations in $3.2 are applicable). Notation details. The Tr, GlobalAvgPoo and GlobalSumPool ops are assumed to act on all spatial axes (with sizes H and W in this example), producing a |x x2 -ernell Similarly, the AvgPool and SumPool ops is assumed act on spatial axes well, applying the specified strides pooling window sizes and padding strategy p to respective axes pairs in and T (acting as 4D pooling with replicated parameters of 2D version). T and i are defined identically to Lee (2019) as T (Σ) = E [o(u) . (Σ) = (('(w))((w))
~ A (0, E). These expressions can be evaluated nonlinearities, and preserve the shape of the kernel. The A op is defined similarly to Novak al. wrr. hihh
closed form for many
hh (x,x')/q² where the summation is performed over the convolutional filter receptive field with q pixels (we assume unit strides and circular padding in this expression, but generalization other settings is trivial and supported by the library). [S] = Σ Σ) (n-fold replication). For LayerNorm, FanInConcat, and Attn (Anonymous, 2020) translation rules refer the reader to code https: / //gi thub. com/ google//nerralll welle//nurllttngeets as these ops are challenging to express concisely using current notation. See Figure for an example of applying the translation rules to specific model, and $3.1 for deriving sample translation rule. See $D for the full list of currently implemented translations.
(2019); Xiao et al. (2018) as [A (")) 2x, = La. dun +
tow +dw
18

---

Published as a conference paper at ICLR 2020
Independent
Full Test Coyariance
3.0
3.0
FCNNTK FCNNGP
CONV-NTI
2.5
2.5
Test NLL (nats)
Test NLL (nats)
2.0
2.0
1.5-
1.5-
1.0
1.0
2
2
2
2
11
22
13
2
14
2
26
27
2
2
20
2 12
2
e1
2
2
Training Set Size
Training Set Size
NTK
NNGP
109
10
108
108
Condition Number
107
Condition Number
10
106
10
10
10
10
10
FC
FC-Test
Conv
10
10
Cont-Test-Co Cov
WRasNet-TesttCCo
102
102
2
2
2
2
2
10
11
2 12
2
2
2
2
2
2
2
21z 213
Training Set Size
Training Set Size
Figure Predictive negative log-likelihoods and condition numbers. Top. Test negative log- likelihoods for NNGP posterior and Gaussian predictive distribution for NTK at infinite training time for CIFAR-10 (test set of 2000 points). Fully Connected (FC, Listing 3) and Convolutional network without pooling (CONV, Listing 2) models selected based train marginal negative log-likelihoods in Figure 3. Bottom. Condition numbers for covariance matrices corresponding to NTK/NNGP well respective predictive covaraince the test set. Ill-conditioning of Wide Residual Network kernels due to pooling layers (Xiao et al., 2019) could be the cause of numerical issues when evaluating predictive NLL for this kernels.
19


Reviews:
[{"review_id": "ryxYhr0aKr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: A Jax based neural tangents kernel library is introduced, with native GPU, TPU, and XLA support. Due to the correspondences with infinite neural network kernels (NNGPs), these kernels are also able to be computed for (essentially) free. Layers which do not emit an analytical form (e.g. Tanh or Softplus) can be implemented using Monte Carlo estimates. Several engineering-based experiments are performed demonstrating the potential scalability of their library.\n\nWhile I really enjoyed reading the paper and believe that this library could be extremely practically useful, I vote to reject this paper because I do not feel that it has sufficient novelty to be a paper on its own in light of Lee et al, 2019. \n\nEdit: post rebuttal, I'm bumping my score to a weak reject but would have minimal qualms if this paper were to be accepted. I find the further experiments performed by the authors of very good quality overall, but I'm still not particularly satisfied by their `argument that the codebase itself is distinct enough from separate related work. It's unfortunately a bit hard, from a machine learning researcher side, to review the quality of a codebase in and of itself. \n\nSignificance: Having played around with the code a bit, I find that the library itself is of very high quality and is pretty straightforward to use. I could definitely see myself using this library in the future for research work.\n\nHowever, my primary concern with this paper is that it’s not sufficiently distinct from the previous work of Lee et al, 2019. After all, most of the experiments in that paper would have required the type of implementation that is described in greater detail in this paper. \n\nTo be able to vote to accept this paper, I will have to see an experiment that is practically performed with the current library in order to distinguish it from previous work (specifically Lee et al, 2019). In recent work, Arora et al, 2019 (Note: I do not consider this reference in my review other than to be mentioned as an example of an experiment that could be run with your library) run neural tangent kernels on tabular data using kernel SVMs. One other potential example would be a kernel SVM in this manner on CIFAR-10. An alternative example would be to exploit the Gaussian process representation and test out both NTKs and NNGPs in comparison to standard kernels for GPs and NNs on UCI regression tasks.\n\nOriginality: Again, a very efficient and easy to use implementation of neural tangent kernels would be a great boost to the community. This is doubly so as Jax is easy and pretty straightforward to use and is quite numpy like. \n\nAgain, I am very concerned with originality in comparison to Lee et al, 2019. Even checking out the link to their codebase provides a github repo that is quite similar to this one. Given that ICLR is a venue of similar domain to NeurIPS, it’s not clear to me why this paper ought to be anything other than a separate supporting tech report. If this paper had been submitted to something like SysML (edit: or JMLR MLOSS), I would see the distinctness instead.\n\nClarity: I find the paper to be extremely well-written and easy to follow. The addition of code snippets throughout is very well done, even if it’s a bit overkill. I don’t know what adding a half page long description of an infinitely wide WideResNet adds to the paper when that space could be better used by another experiment. \n\nQuality: I find the experiments performed to be very well constructed. Below are a few mostly minor comments on the experiments:\n\nIn Figure 1 on the right, I would have liked to have seen the posterior predictive for a NNGP with the same kernel as well. \n\nIn Figure 2, why is the NNGP slower to converge to the analytic values here? Obviously, the rates of convergence are the same, but the constants seem different.\n\nIn Figure 3 (and throughout the experiments), does “full Bayesian inference for CIFAR10” mean that you treated the classification labels as regression targets? If so, how was classification error measured.\n\nIn Section 3.1, you mention that the library “leverages block-diagonal structure” to invert the CIFAR10 covariance matrix for each class (still 50k x 50k). Possibly this is because I haven’t had the chance to use TPUs, but I’m currently struggling to see how one could form and invert (via Choleskys) matrices of this size (50k x 50k) on a standard GPU (or CPU). Could the authors please clarify how they did this (whether through iterative methods, another structure exploiting trick, lots of memory, etc.)?\n\nsecond edit: I was also unable to respond to the final comment about the UCI experiments in its own comment, but thank you for providing the estimated depths. These results definitely show the potential software promise of the codebase and open some interesting new research questions as a result.\n\nReferences:\n\nArora, S., et al., Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks, https://arxiv.org/abs/1910.01663\n\nLee, J., et al., Wide Neural Networks of any Depth Evolve as Linear Models Under Gradient Descent, NeurIPS, 2019, https://arxiv.org/abs/1902.06720\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "H1emYsn6Yr", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "POST-REBUTTAL COMMENTS\n\nI appreciate the response from the authors. \n\nI particularly like the comparison table in the response to the other reviewer and ought to be highlighted in the paper.\n\nIf I were to start this line of research, I would be inclined to expand on the codebase. The contribution is significant. Hence, I am bumping up my score to accept.\n\n\nPRIOR FEEDBACK\n\nThe contribution of this work lies in providing a library for working with the existing variants of infinite-width neural networks and avoiding the need to derive the NNGP and NT kernels for each architecture by hand. The authors have firstly shown performance comparisons between inferences between finite vs. infinitely wide neural networks. The authors then go into some implementation details with their library. The authors have provided the code and cookbook in the links provided in the abstract. On the overall, I like this effort which is timely.\n\nSome additional suggestions below:\n\nI would like to see an additional metric for performance comparison of probabilistic models, which is often used in the GP literature: mean negative log probability.\n\nIt would also be interesting to see how the posterior variance (e.g., Fig. 1 right) evolves over the entire space during training. \n\nI would have preferred a more detailed discussion about the implementation on transforming tensor ops to kernel ops in Section 3.\n\nFor the summary of contributions, can you give the corresponding section number to refer to when you demonstrate each feature? For example, is the 4th feature (i.e., exploring weight space perspective) demonstrated in the paper?\n\nCan the authors elaborate on the ease of expanding their library for the new developments in this field?\n\n\nMinor issues:\n\nPage 1: Gaussian Procesesses?\nPage 4: it’s infinite?\nFig. 4: I would have preferred the indices to be placed as subscripts instead of superscripts.\nPage 8: it’s order of dimensions?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "BJgNczUaFB", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work develops a library for working with a class of infinitely wide neural networks, in particular those corresponding to neural tangent kernels (NTKs) and neural network Gaussian processes (NNGPs). The theory for these two kernels was well developed in a series of recent papers, and this library provides an automatic way to transform any appropriate neural net architecture into its corresponding NTK and NNGP.\n\nInfinitely wide neural networks have been a popular subject of theoretical research and have been observed to have highly nontrivial performance on a variety of tasks (e.g. CIFAR-10 classification). It's really nice to see the development of such a library, which I believe could benefit the deep learning community a lot, especially for theoretical research on NTK.\n\nI appreciate this work a lot. Currently I can only give weak accept instead of accept for a couple of reasons:\n1. The theory and formulae of NTKs and NNGPs were well developed. This work mostly consists of implementing and modularizing them. The research contribution is relatively low.\n2. As commented in the paper and if I understand correctly, the current library cannot scale to large datasets for CNNs with pooling. This would make the computation much more expensive (and probably infeasible without additional techniques and huge computing power) as mentioned in [Novak et al. 2019] and [Arora et al. 2019]. However pooling seems extremely useful for NTKs and NNGPs on image datasets. I think this makes this work somewhat less exciting than it may sound.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a software library for dealing with neural networks either in the (usual) finite limit or in the infinite limit. The latter is obtained by using the Neural Tangent Kernel theory. 

There is variance in the reviewers' scores, however there has also been quite a lot of discussion, which has been facilitated by the authors' elaborate rebuttal. The main points in favor and against are clear: on the positive side, the library is demonstrated well (especially after rebuttal) and is equipped with desirable properties such as usage of GPU/TPU, scalability etc. On the other hand, a lot of the key insights build heavily on prior work of Lee et al, 2019. However, judging novelty when it comes to a software paper is more tricky to do, especially given that not many such papers appear in ICLR and therefore calibration is difficult. This has been discussed among reviewers. 

It would help if some further theoretical insights were included in this paper; these insights could come by working backwards from the implementation (i.e. what more can we learn about infinite width networks now that we can experiment easily with them?).

Overall, this paper should still be of interest to the ICLR community.

Decision: Accept (Spotlight)