Title: NEURAL TANGENTS: FAST AND EASY INFINITE NEURAL NETWORKS IN PYTHON

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 finitewidth 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
We also provide an accompanying interactive Colab notebook 1 .

1. 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.

1.1. INFINITE-WIDTH BAYESIAN NEURAL NETWORKS:

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 multilayer 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) , 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 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 1 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. 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; Zou et al., 2019) . 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.

1.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-tomonths 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., 2018a) to accelerate research on infinite limits of neural networks. The main features of NEURAL TANGENTS are:2 

• 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 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 to 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 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 6 ). • 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 5 , §B.3). The above specification set the hidden layer widths to 2048, which has no effect on the infinite width network inference, but the init_fn and apply_fn here correspond to ordinary finite width networks. In Figure 1 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 an ensemble using gradient descent.

2.2. AN INFINITELY WIDERESNET:

The above example considers a relatively simple network on a synthetic task. In practice we may want to consider real-world architectures, and see how close they are to their infinite-width limit. For this task we study a variant of an infinite-channel Wide Residual Network (Zagoruyko & Komodakis, 2016) (WRN-28-∞). 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 (as measured by the widening factor, k) and second as a function of the number of finite-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 = nt.monte_carlo_kernel_fn(init_fn, apply_fn, rng_key, n_samples) sampled_kernel = kernel_fn(x, x)

The convergence is shown in Figure 2 . We see that as 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 is small relative to the variance and the distance to the empirical kernel decreases with the number of samples.

2.3. COMPARISON OF NEURAL NETWORK ARCHITECTURES AND TRAINING SET SIZES:

The above examples demonstrate how one might construct a complicated architecture and perform inference using NEURAL TANGENTS Ṅext we train a range of architectures on CIFAR-10 and compare their performance as a function of dataset size. In particular, we compare a fully-connected network, Published as a conference paper at ICLR 2020 a 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 infinitewidth NNGP or NTK. For each architecture we perform a hyperparameter search over the depth of the network, selecting the depth that maximizes the marginal log likelihood on the training set.

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 is used in Figures 2 and 3 to their analytic values (WRN-28-∞, computed with kernel_fn ), as the network gets wider by increasing the widening factor (vertical axis) and as more random networks are averaged over (horizontal axis). Experimental detail. The kernel is computed in 32-bit precision on a 100 × 50 batch of 8 × 8-downsampled CIFAR10 (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 1 (all logits are i.i.d. conditioned on a given input). The displayed distance is the relative Frobenius norm squared, i.e.

K -K k,n 2 F / K 2 F ,

where k is the widening factor and n is the number of samples.

Published as a conference paper at ICLR 2020 2 6 2 7 2 8 2 9 2 10 2 11 2 12 2 13 2 14 2 15 Training Set Size As is common in prior work (Lee 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, -0.1, . . . , -0.1) .

For each training set size, the best model in the family is selected by minimizing the mean negative marginal log-likelihood (NLL, right) on the training set.

The results are shown in Figure 3 . We see 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 is a set of translation rules, that sends each tensor operation acting on a finite-width layer to a corresponding transformation of the kernel for an infinite-width network. This is illustrated in Figure 4 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 to be preceded by affine transformations (either dense or convolutional). This is because infinite-width results often assume that the pre-activations 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-KERNEL OPS TRANSLATION:

To get some intuition behind the translation rules, we consider the case of a nonlinearity followed by a dense layer. Let z = z (X , θ) ∈ R d×n be the preactivations resulting from d distinct inputs at a node in some hidden layer of a neural network. Suppose z has NNGP kernel and NTK given by

K z = E θ z i z T i , Θ z = E θ ∂z i ∂θ ∂z i ∂θ T (1)

where z i ∈ R d is the i th neuron and θ are the parameters in the network up until z. Here d is the cardinality of the network inputs X and n is the number of neurons in the z node. We Published as a conference paper at ICLR 2020

Under review as a conference paper at ICLR 2019 Figure 1 : Different dimensionality collapsing strategies described in §3. Validation accuracy of an MC-CNN-GP with pooling ( §3.2.1) is consistently better than other models due to translation invariance of the kernel. CNN-GP with zero padding ( §3.1) outperforms an analogous CNN-GP without padding as depth increases. At depth 15 the spatial dimension of the output without padding is reduced to 1 ⇥ 1, making the CNN-GP without padding equivalent to the center pixel selection strategy ( §3.2.2) -which also performs worse than the CNN-GP (we conjecture, due to overfitting to centrally-located features) but approaches the latter (right) in the limit of large depth, as information becomes more uniformly spatially distributed (Xiao et al., 2018) . CNN-GPs generally outperform FCN-GP, presumably due to the local connectivity prior, but can fail to capture nonlinear interactions between spatially-distant pixels at shallow depths (left). Values are reported on a 2K/4K train/validation subsets of CIFAR10. See §A.7.3 for experimental details.

1.. Global average pooling: take:

EQUATION

This approach corresponds to applying global average pooling right after the last convolutional layer. 8 This approach takes all pixel-pixel covariances into consideration and makes the kernel translation invariant. However, it requires O |X | 2 d 2 memory to compute the sample-sample covariance of the GP (or O d 2 per covariance entry in an iterative or distributed setting). It is impractical to use this method to analytically evaluate the GP, and we propose to use a Monte Carlo approach (see §4).

2. Subsampling one particular pixel: take h = e ↵ ,

EQUATION

This approach makes use of only one pixel-pixel covariance, and requires the same amount of memory as vectorization ( §3.1) to compute.

We compare the performance of presented strategies in Figure 1 . Note that all described strategies admit stacking additional FC layers on top while retaining the GP equivalence, using a derivation analogous to §2 (Lee et al., 2018; Matthews et al., 2018b) .

4. MONTE CARLO EVALUATION OF INTRACTABLE GP KERNELS:

We introduce a Monte Carlo estimation method for NN-GP kernels which are computationally impractical to compute analytically, or for which we do not know the analytic form. Similar in spirit to traditional random feature methods (Rahimi & Recht, 2007) , the core idea is to instantiate many • CNN-GP exhibits similar behavio from the phase boundary, the cov all different inputs becomes asym constant and makes learning infe

• Many different NN architectures converge to a GP.

• Some admit a tractable analytic form. E.g. the covariance of a CNN-GP with no pooling can be computed using O(#samples^2 x #pixels) resources.

• For other architectures we use a Monte Carlo approach.

Sampling finite random networks of a given architecture and empirically computing the output covariance allows to approximate the CNN-GP covariance. In our experiments this (biased) estimate converges to the true CNN-GP covariance in #(instantiated networks) and #channels, both in terms of covariance Frobenius distance and the GP accuracy.

• Various CNN architectural decisions like pooling / vectorizing / subsampling, zero or no padding have a corresponding NN-GP.

• Pooling enforces translation-invariant predictions in both CNNs and CNN-GPs and allows for the best performance.

• Shallow CNN models can perform worse than fully-connected alternatives due to failing to capture non-linear interactions between distant pixels. 

y 2 = ϕ ( z 1 ) y 1 = ϕ ( z 0 ) n 1 → ∞ n 1 → ∞ n 2 → ∞ n 2 → ∞

6. When do random finite width networks and use the empirical uncentered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,

EQUATION

where ✓ consists of M draws of the weights and biases from their prior distribution, ✓ m ⇠ p (✓), and n is the width or number of channels in hidden layers. The MC-GP kernel converges to the analytic kernel with increasing width, lim n!1 K l n,M = K l 1 in probability. For finite width networks, the uncertainty in K l n,M is Var

⇥ K l n,M ⇤ = Var ✓ ⇥ K l n (✓) ⇤ /M . From Daniely et al. (2016), we know that Var ✓ ⇥ K l n (✓) ⇤ / 1 n , which leads to Var ✓ [K l n,M ] / 1

Mn . For finite n, K l n,M is also a biased estimate of K l 1 , where the bias depends solely on network width. We do not currently have an analytic form for this bias, but we can see in Figures 2 and 7 that for the hyperparameters we probe it is small relative to the variance. In particular,

K l n,M (✓) K L 1 2 F

is nearly constant for constant Mn. We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate. Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.

In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from

O ⇣ |X | 2 d 2 ⌘ to O ⇣ |X | 2 + n 2 + nd ⌘

, making the evaluation of CNN-GPs with pooling practical.

5. DISCUSSION:



5.1. BAYESIAN CNNS WITH MANY CHANNELS ARE IDENTICAL TO LOCALLY CONNECTED NETWORKS, IN THE ABSENCE OF POOLING:

Locally Connected Networks (LCNs) (Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spatial locations. LCNs preserve the connectivity pattern, and thus topology, of a CNN. However, they do not possess the equivariance property of a CNN -if an input is translated, the latent representation in an LCN will be completely different, rather than also being translated.

The CNN-GP predictions without spatial pooling in §3.1 and §3.2.2 depend only on sample-sample covariances, and do not depend on pixel-pixel covariances. LCNs destroy pixel-pixel covariances:

⇥ K L 1 ⇤ LCN ↵,↵ 0 (x, x 0 ) = 0, for ↵ 6 = ↵ 0 and all x,

x 0 2 X and L > 0. However, LCNs preserve the covariances between input examples at every pixel: n,M when the number of outer products contributing to K L+1 n,M approximately equals its rank. Values reported are for a 3-layer model applied to a 2K/4K train/validation subset of CIFAR10 downsampled to 8 ⇥ 8. See Figure 7 for similar results with other architectures and §A.7.2 for experimental details.

⇥ K L 1 ⇤ LCN ↵,↵ (x, x 0 ) = ⇥ K L 1 ⇤ CNN ↵,↵ (x, x 0 ). As a result,

random finite width networks and use the empirical uncentered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,

EQUATION

where ✓ consists of M draws of the weights and biases from their prior distribution, ✓ m ⇠ p (✓), and n is the width or number of channels in hidden layers. The MC-GP kernel converges to the analytic kernel with increasing width,

lim n!1 K l n,M = K l 1 in probability. For finite width networks, the uncertainty in K l n,M is Var ⇥ K l n,M ⇤ = Var ✓ ⇥ K l n (✓) ⇤ /M . From Daniely et al. (2016), we know that Var ✓ ⇥ K l n (✓) ⇤ / 1 n , which leads to Var ✓ [K l n,M ] / 1 Mn . For finite n, K l n,M is also a biased estimate of K l 1 ,

where the bias depends solely on network width. We do not currently have an analytic form for this bias, but we can see in Figures 2 and 7 that for the hyperparameters we probe it is small relative to the variance. In particular,

K l n,M (✓) K L 1 2 F

is nearly constant for constant Mn. We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate. Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.

In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from

O ⇣ |X | 2 d 2 ⌘ to O ⇣ |X | 2 + n 2 + nd ⌘

, making the evaluation of CNN-GPs with pooling practical.



Locally Connected Networks (LCNs) (Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spatial locations. LCNs preserve the connectivity pattern, and thus topology, of a CNN. However, they do not possess the equivariance property of a CNN -if an input is translated, the latent representation in an LCN will be completely different, rather than also being translated.

The CNN-GP predictions without spatial pooling in §3.1 and §3.2.2 depend only on sample-sample covariances, and do not depend on pixel-pixel covariances. LCNs destroy pixel-pixel covariances:

⇥ K L 1 ⇤ LCN ↵,↵ 0 (x, x 0 ) = 0, for ↵ 6 = ↵ 0 and all x,

x 0 2 X and L > 0. However, LCNs preserve the covariances between input examples at every pixel:

⇥ K L 1 ⇤ LCN ↵,↵ (x, x 0 ) = ⇥ K L 1 ⇤ CNN ↵,↵ (x, x 0

). As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical. Moreover, LCN-GPs with 8 (#NN instantiations)

(not always necessary to track the whole 4x4x10x10 covariance.) (attaching a fully-connected layer only to the center pixel of the output)

CNN (1D) activations:



CNN-GP covariance:

Under review as a conference paper at ICLR 2019 Figure 1 : Different dimensionality collapsing strategies described in §3. Validation accuracy of an MC-CNN-GP with pooling ( §3.2.1) is consistently better than other models due to translation invariance of the kernel. CNN-GP with zero padding ( §3.1) outperforms an analogous CNN-GP without padding as depth increases. At depth 15 the spatial dimension of the output without padding is reduced to 1 ⇥ 1, making the CNN-GP without padding equivalent to the center pixel selection strategy ( §3.2.2) -which also performs worse than the CNN-GP (we conjecture, due to overfitting to centrally-located features) but approaches the latter (right) in the limit of large depth, as information becomes more uniformly spatially distributed (Xiao et al., 2018) . CNN-GPs generally outperform FCN-GP, presumably due to the local connectivity prior, but can fail to capture nonlinear interactions between spatially-distant pixels at shallow depths (left). Values are reported on a 2K/4K train/validation subsets of CIFAR10. See §A.7.3 for experimental details.

EQUATION

This approach corresponds to applying global average pooling right after the last convolutional layer. 8 This approach takes all pixel-pixel covariances into consideration and makes the kernel translation invariant. However, it requires O |X | 2 d 2 memory to compute the sample-sample covariance of the GP (or O d 2 per covariance entry in an iterative or distributed setting). It is impractical to use this method to analytically evaluate the GP, and we propose to use a Monte Carlo approach (see §4).

2.. Subsampling one particular pixel: take:

EQUATION

This approach makes use of only one pixel-pixel covariance, and requires the same amount of memory as vectorization ( §3.1) to compute.

We compare the performance of presented strategies in Figure 1 . Note that all described strategies admit stacking additional FC layers on top while retaining the GP equivalence, using a derivation analogous to §2 (Lee et al., 2018; Matthews et al., 2018b) .

We introduce a Monte Carlo estimation method for NN-GP kernels which are computationally impractical to compute analytically, or for which we do not know the analytic form. Similar in spirit to traditional random feature methods (Rahimi & Recht, 2007), the core idea is to instantiate many • CNN-GP exhibits similar behavio from the phase boundary, the cov all different inputs becomes asym constant and makes learning infe

3.. Different NN-GP architectures:

• Many different NN architectures converge to a GP.

• Some admit a tractable analytic form. E.g. the covariance of a CNN-GP with no pooling can be computed using O(#samples^2 x #pixels) resources.

• For other architectures we use a Monte Carlo approach.

Sampling finite random networks of a given architecture and empirically computing the output covariance allows to approximate the CNN-GP covariance. In our experiments this (biased) estimate converges to the true CNN-GP covariance in #(instantiated networks) and #channels, both in terms of covariance Frobenius distance and the GP accuracy.

(a) SGD-trained finit (b) Performance of S CNN-GP with increa (c) However, the bes significant margin.

• Various CNN architectural decisions like pooling / vectorizing / subsampling, zero or no padding have a corresponding NN-GP.

• Pooling enforces translation-invariant predictions in both CNNs and CNN-GPs and allows for the best performance.

• Shallow CNN models can perform worse than fully-connected alternatives due to failing to capture non-linear interactions between distant pixels. 

7.. Disentang:

⊗ I10 K 0 K 0 K 1 K 2 K 1 K 2 σ 2 ω K 0 + σ 2 b σ 2 ω K 1 + σ 2 b ( K 1 ) ( K 0 )

K 0 K 1 K 2 ( K 1 ) ( K 0 )

affine transform affine transform affine transform convolution convolution matrix-multiply nonlinearity nonlinearity 

y 2 = ϕ ( z 1 ) y 1 = ϕ ( z 0 ) n 1 → ∞ n 1 → ∞ n 2 → ∞ n 2 → ∞ C (K) = E x⇠N (0,K) h (x) (x) T i < l a t e x i t s h a 1 _ b a s e 6 4 = " L K S S 7 i N m 7 W J o + g 4 I + B j N B l P b a w o = " > A A A C e H i c b V F d S 8 M w F E 3 r 9 / y a H 2 + + B K e o I N L u R V 8 E U Q R B E A W n w l p H k q V b M G l L c i u O M v + n z / o r f D J b q + j m h c D h n H O 5 N + f S V A o D n v f m u B O T U 9 M z s 3 O V + Y X F p e X q y u q d S T L N e I M l M t E P l B g u R c w b I E D y h 1 R z o q j k 9 / T p b K D f P 3 N t R B L f Q i / l o S K d W E S C E b B U q / q 6 F S g C X U Z k f t Y P J I 9 g 9 z L Q o t O F P X y M h x q l + X m / l b 8 E R q g f 8 1 V p 9 v a / 7 Q X R D N K u K K S X U h h n H m 8 L E G 6 1 q j X v w B s W H g d + C W q o r O t W 9 T 1 o J y x T P A Y m i T F N 3 0 s h z I k G w S T v V 4 L M 8 J S w J 9 L h T Q t j o r g J 8 2 F O f b x t m T a O E m 1 f D H j I / u 7 I i T K m p 6 h 1 D v 5 p R r U B + a 9 G q R o Z D d F R m I s 4 z Y D H r J g c Z R J D g g d X w G 2 h O Q P Z s 4 A w L e z y m H W J J g z s r S o 2 F X 8 0 g 3 F w V z / w L b 6 p 1 0 5 O y 3 x m 0 Q b a R L v I R 4 f o B F 2 g a 9 R A D H 0 4 F W f N W X c + X e z u u H u F 1 X X K n j X 0 p 9 z 6 F 6 y i w S Y = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " L K S S 7 i N m 7 W J o + g 4 I + B j N B l P b a w o = " > A A A C e H i c b V F d S 8 M w F E 3 r 9 / y a H 2 + + B K e o I N L u R V 8 E U Q R B E A W n w l p H k q V b M G l L c i u O M v + n z / o r f D J b q + j m h c D h n H O 5 N + f S V A o D n v f m u B O T U 9 M z s 3 O V + Y X F p e X q y u q d S T L N e I M l M t E P l B g u R c w b I E D y h 1 R z o q j k 9 / T p b K D f P 3 N t R B L f Q i / l o S K d W E S C E b B U q / q 6 F S g C X U Z k f t Y P J I 9 g 9 z L Q o t O F P X y M h x q l + X m / l b 8 E R q g f 8 1 V p 9 v a / 7 Q X R D N K u K K S X U h h n H m 8 L E G 6 1 q j X v w B s W H g d + C W q o r O t W 9 T 1 o J y x T P A Y m i T F N 3 0 s h z I k G w S T v V 4 L M 8 J S w J 9 L h T Q t j o r g J 8 2 F O f b x t m T a O E m 1 f D H j I / u 7 I i T K m p 6 h 1 D v 5 p R r U B + a 9 G q R o Z D d F R m I s 4 z Y D H r J g c Z R J D g g d X w G 2 h O Q P Z s 4 A w L e z y m H W J J g z s r S o 2 F X 8 0 g 3 F w V z / w L b 6 p 1 0 5 O y 3 x m 0 Q b a R L v I R 4 f o B F 2 g a 9 R A D H 0 4 F W f N W X c + X e z u u H u F 1 X X K n j X 0 p 9 z 6 F 6 y i w S Y = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " L K S S 7 i N m 7 W J o + g 4 I + B j N B l P b a w o = " > A A A C e H i c b V F d S 8 M w F E 3 r 9 / y a H 2 + + B K e o I N L u R V 8 E U Q R B E A W n w l p H k q V b M G l L c i u O M v + n z / o r f D J b q + j m h c D h n H O 5 N + f S V A o D n v f m u B O T U 9 M z s 3 O V + Y X F p e X q y u q d S T L N e I M l M t E P l B g u R c w b I E D y h 1 R z o q j k 9 / T p b K D f P 3 N t R B L f Q i / l o S K d W E S C E b B U q / q 6 F S g C X U Z k f t Y P J I 9 g 9 z L Q o t O F P X y M h x q l + X m / l b 8 E R q g f 8 1 V p 9 v a / 7 Q X R D N K u K K S X U h h n H m 8 L E G 6 1 q j X v w B s W H g d + C W q o r O t W 9 T 1 o J y x T P A Y m i T F N 3 0 s h z I k G w S T v V 4 L M 8 J S w J 9 L h T Q t j o r g J 8 2 F O f b x t m T a O E m 1 f D H j I / u 7 I i T K m p 6 h 1 D v 5 p R r U B + a 9 G q R o Z D d F R m I s 4 z Y D H r J g c Z R J D g g d X w G 2 h O Q P Z s 4 A w L e z y m H W J J g z s r S o 2 F X 8 0 g 3 F w V z / w L b 6 p 1 0 5 O y 3 x m 0 Q b a R L v I R 4 f o B F 2 g a 9 R A D H 0 4 F W f N W X c + X e z u u H u F 1 X X K n j X 0 p 9 z 6 F 6 y i w S Y = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " L K S S 7 i N m 7 W J o + g 4 I + B j N B l P b a w o = " > A A A C e H i c b V F d S 8 M w F E 3 r 9 / y a H 2 + + B K e o I N L u R V 8 E U Q R B E A W n w l p H k q V b M G l L c i u O M v + n z / o r f D J b q + j m h c D h n H O 5 N + f S V A o D n v f m u B O T U 9 M z s 3 O V + Y X F p e X q y u q d S T L N e I M l M t E P l B g u R c w b I E D y h 1 R z o q j k 9 / T p b K D f P 3 N t R B L f Q i / l o S K d W E S C E b B U q / q 6 F S g C X U Z k f t Y P J I 9 g 9 z L Q o t O F P X y M h x q l + X m / l b 8 E R q g f 8 1 V p 9 v a / 7 Q X R D N K u K K S X U h h n H m 8 L E G 6 1 q j X v w B s W H g d + C W q o r O t W 9 T 1 o J y x T P A Y m i T F N 3 0 s h z I k G w S T v V 4 L M 8 J S w J 9 L h T Q t j o r g J 8 2 F O f b x t m T a O E m 1 f D H j I / u 7 I i T K m p 6 h 1 D v 5 p R r U B + a 9 G q R o Z D d F R m I s 4 z Y D H r J g c Z R J D g g d X w G 2 h O Q P Z s 4 A w L e z y m H W J J g z s r S o 2 F X 8 0 g 3 F w V z / w L b 6 p 1 0 5 O y 3 x m 0 Q b a R L v I R 4 f o B F 2 g a 9 R A D H 0 4 F W f N W X c + X e z u u H u F 1 X X K n j X

EQUATION

where ✓ consists of M draws of the weights and biases from their prior distribution, ✓ m ⇠ p (✓), and n is the width or number of channels in hidden layers. The MC-GP kernel converges to the analytic kernel with increasing width, lim n!1 K l n,M = K l 1 in probability. For finite width networks, the uncertainty in K l n,M is Var

⇥ K l n,M ⇤ = Var ✓ ⇥ K l n (✓) ⇤ /M . From Daniely et al. (2016), we know that Var ✓ ⇥ K l n (✓) ⇤ / 1 n , which leads to Var ✓ [K l n,M ] / 1

Mn . For finite n, K l n,M is also a biased estimate of K l 1 , where the bias depends solely on network width. We do not currently have an analytic form for this bias, but we can see in Figures 2 and 7 that for the hyperparameters we probe it is small relative to the variance. In particular,

K l n,M (✓) K L 1 2 F

is nearly constant for constant Mn. We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate. Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.

In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from 

O ⇣ |X | 2 d 2 ⌘ to O ⇣ |X | 2 + n 2 + nd ⌘ ,

EQUATION

where ✓ consists of M draws of the weights and biases from their prior distribution, ✓ m ⇠ p (✓), and n is the width or number of channels in hidden layers. The MC-GP kernel converges to the analytic kernel with increasing width, lim n!1 K l n,M = K l 1 in probability. For finite width networks, the uncertainty in K l n,M is Var

⇥ K l n,M ⇤ = Var ✓ ⇥ K l n (✓) ⇤ /M . From Daniely et al. (2016), we know that Var ✓ ⇥ K l n (✓) ⇤ / 1 n , which leads to Var ✓ [K l n,M ] / 1

Mn . For finite n, K l n,M is also a biased estimate of K l 1 , where the bias depends solely on network width. We do not currently have an analytic form for this bias, but we can see in Figures 2 and 7 that for the hyperparameters we probe it is small relative to the variance. In particular,

K l n,M (✓) K L 1 2 F

is nearly constant for constant Mn. We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate. Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.

In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from

O ⇣ |X | 2 d 2 ⌘ to O ⇣ |X | 2 + n 2 + nd ⌘

, making the evaluation of CNN-GPs with pooling practical.

(#NN instantiations)

(not always necessary to track the whole 4x4x10x10 covariance.)

(attaching a fully-connected layer only to the center pixel of the output)

K 0 < l a t e x i t s h a _ b a s e = " d j K o C a j Z g k r w u K p B m p = " > A A A B + i c b V C S g N B F L b X F R S t B o N g F W a E Y M g g E c D k j X M T m a T I T O y y s E J Z g q + Q D q x R v B L H R w k F p p M L h n H s l x M k g m u D a e T W h c W l J r p r x u b W X t n Z q O U V Z l c Y i V o A a C Z x K q G G E a i W J E B o L V g / l K / M K V H N Z Q c J S b o R D z k l x k q / e X S j i E p A z R P v h x T P z Z P T h V t q F r Y n p q l k k a G C a N c G L j C j D q W B D t V q l h D a J W t D Q i k m k / m w R A d W a A w V n Y i g y b q u M S K H M r C b k p i e n v X G r e E M i Z a B O e + h m P k t S w i E T w Q g E N x E a j D F a N G D C w h V H H P K I o g g t i X t u L N d j B P a k c l i E b C x f A F T G E P u E Q P D i B M l x B B a p A o Q u P A T P z t A Z O S / O Q / z c M I f O G / f T W X v g = = < / l a t e x i t > A < l a t e x i t s h a 1 _ b a s e 6 4 = " O 6 6 A T y 1 a w 1 t Y L Y U p v p J 8 F g N J C H 8 = " > A A A C A 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X a 1 0 E a M 2 l h G M A 9 I l j A 7 m S R D Z m e W m V k h L C n 9 B F s t L O 3 E 1 t p P E G t / x N k k h S Y e u H A 4 5 1 7 u 4 Y Q x Z 9 p 4 3 p e T W 1 h c W l 7 J r 7 p r 6 x u b W 4 X t n Z q W i S K 0 S i S X q h F i T T k T t G q Y 4 b Q R K 4 q j k N N 6 O L j K / P o d V Z p J c W u G M Q 0 i 3 B O s y w g 2 V m q 2 I m z 6 B P P 0 Y t Q u F L 2 S N w a a J / 6 U F M 8 / 3 L P 4 + d O t t A v f r Y 4 k S U S F I R x r 3 f S 9 2 A Q p V o Y R T k d u K 9 E 0 x m S A e 7 R p q c A R 1 U E 6 j j x C B 1 b p o K 5 U d o R B Y / X 3 R Y o j r Y d R a D e z i H r W y 8 R / v T C M Z l 6 b 7 m m Q M h E n h g o y + d x N O D I S Z Y W g D l O U G D 6 0 B B P F b H h E + l h h Y m x t r m 3 F n + 1 g n t S O S v 5 x y b v x i u V L m C A P e 7 A P h + D D C Z T h G i p Q B Q I S H u A R

n p x 7 5 8 V 5 d d 4 m q z l n e r M L f + C 8 / w A 5 N 5 u k < / l a t e x i t > C < l a t e x i t s h a 1 _ b a s e 6 4 = " B p 3 s 0 N X a z I 8

d Q f K 7 i 5 e d A + N 1 h F s = " > A A A C A 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X a 1 0 E Y M p r G M Y B 6 Q L G F 2 M k m G z M w u M 7 N C W F L 6 C b Z a W N q J r b W f I N b + i L N J C k 0 8 c O F w z r 3 c w w l j z r T x v C 8 n t 7 S 8 s r q W X 3 c 3 N r e 2 d w q 7 e 3 U d J Y r Q G o l 4 p J o h 1 p Q z S W u G G U 6 b s a J Y h J w 2 w m E l 8 x t 3 V G k W y V s z i m k g c F + y H i P Y W K n V F t g M C O Z p Z d w p F L 2 S N w F a J P 6 M F C 8 / 3 I v 4 + d O t d g r f 7 W 5 E E k G l I R x r 3 f K 9 2 A Q p V o Y R T s d u O 9 E 0 x m S I + 7 R l q c S C 6 i C d R B 6 j I 6 t 0 U S 9 S d q R B E / X 3 R Y q F 1 i M R 2 s 0 s o p 7 3 M v F f L w z F 3 G v T O w 9 S J u P E U E m m n 3 s J R y Z C W S G o y x Q l h o 8 s w U Q x G x 6 R A V a Y G F u b a 1 v x 5 z t Y J P W T k n 9 a 8 m 6 8 Y v k K p s j D A R z C M f h w B m W 4 h i r U g E A E D / A I T 8 6 9 8 + K 8 O m / T 1 Z w z u 9 m H P 3 D e f w A 8 Y 5 u m < / l a t e x i t > K 1 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " X v v O k S j q l J 9 2 P o c D U d m z M o T Y 5 m 8 = " > A A A C B H i c b V C 7 S g N B F L 0 b X z G + o p Y 2 Q 0 R I F X a 1 i G X Q R r C J Y B 6 Q X c P s Z D Y Z M j u 7 z M w K y 5 L W T 7 C x 0 A + w E 1 t 7 P 0 G 0 8 0 e c P A p N P H D h c M 6 9 n M v x Y 8 6 U t u 0 P K 7 e 0 v L K 6 l l 8 v b G x u b e 8 U d / e a K k o k o Q 0 S 8 U i 2 f a w o Z 4 I 2 N N O c t m N J c e h z 2 v K H 5 2 O / d U u l Y p G 4 1 m l M v R D 3 B Q s Y w d p I 7 u W N 0 8 1 c J g K d j r r F Q 7 t i T 4 A W i T M j h 7 V S + e u z + v 5 Q 7 x a / 3 V 5 E k p A K T T h W q u P Y s f Y y L D U j n I 4 K b q J o j M k Q 9 2 n H U I F D q r x s 8 v M I H R m l h 4 J I m h E a T d T f F x k O l U p D 3 2 y G W A / U v D c W / / V 8 P 5 y L 1 s G p l z E R J 5 o K M k 0 O E o 5 0 h M a N o B 6 T l G i e G o K J Z O Z 5 R A Z Y Y q J N

b w X T i j P f w S J p H l e c k 4 p 9 Z e o 5 g y n y c A A l K I M D V a j B B d S h A Q R i u I d H e L L u r G f r x X q d r u a s 2 c 0 + / I H 1 9 g O q A 5 y R < / l a t e x i t > K 2 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " O U e U K e + l I o b d s 6 D 6 x E X 6 i k d e j z A = "

> A A A C B H i c b V C 7 S g N B F J 2 N r x h f U U u b I U F I F X Z j E c u g j W A T w T w g G 8 P s Z D Y Z M j u 7 z N w V l i W t n 2 B j o R 9 g J 7 b 2 f o J o 5 4 8 4 e R S a e O D C 4 Z x 7 O Z f j R Y J r s O 0 P K 7 O y u r a + k d 3 M b W 3 v 7 O 7 l 9 w + a O o w V Z Q 0 a i l C 1 P a K Z 4 J I 1 g I N g 7 U g x E n i C t b z R + c R v 3 T K l e S i v I Y l Y N y A D y X 1 O C R j J v b y p 9 F K X S x + S c S 9 f t M v 2 F H i Z O H N S r B V K X 5 / V 9 4 d 6 L / / t 9 k M a B 0 w C F U T r j m N H 0 E 2 J A k 4 F G + f c W L O I 0 B E Z s I 6 h k g R M d 9 P p z 2 N 8 b J Q + 9 k N l R g K e q r 8 v U h J o n Q S e 2 Q w I D P W i N x H / 9 T w v W I g G / 7 S b c h n F w C S d J f u x w B D i S S O 4 z x W j I B J D C F X c P I / p k C h C w f S W M 6 0 4 i x 0 s k 2 a l 7 J y U 7 S t T z x m a I Y u O U A G V k I O q q I Y u U B 0 1 E E U R

u k e P 6 M m 6 s 5 6 t F + t 1 t p q x 5 j e H 6 A + s t x + r o Z y S < / l a t e x i t > K vec 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " 6 K 3 m C N D g h m e g 3 o p A x p s E W l s n F s

k = " > A A A C H X i c b V C 7 S g N B F J 2 N r x h f U U u b R R G s w q 4 W s Q z a C D Y R T B T c G G Y n d 3 V w Z n a Z u R s M y 3 6 A 3 2 D h J 9 j q B 1 g I Y i u i n T / i J L H Q x A M D h 3 P O 5 d 4 5 Y S K 4 Q c 9 7 d w o T k 1 P T M 8 X Z 0 t z 8 w u J S e X m l a e J U M 2 i w W M T 6 N K Q G B F f Q Q I 4 C T h M N V I Y C T s K r / b 5 / 0 g V t e K y O s Z d A S 9 I L x S P O K F q p X d 4 I J M V L R k V 2 m L e z g K s I e / l 5 F i B c o 5 Z Z F 1 i e 2 5 R X 8 Q Z w x 4 n / Q z Z q 6 1 u f H 9 X n 2 3 q 7 / B V 0 Y p Z K U M g E N e b M 9 x J s Z V Q j Z w L y U p A a S C i 7 o h d w Z q m i E k w r G 3 w m d z e t 0 n G j W N u n 0 B 2 o v y c y K o 3 p y d A m + 6 e b U a 8 v / u u F o R x Z j d F u K + M q S R E U G 2 6 O U u F i 7 P a r c j t c A 0 P R s 4 Q y z e 3 x L r u k m j K 0 h Z Z s K / 5 o B + O k u V 3 x d y r e k a 1 n j w x R J G t k n W w R n 1 R J j R y Q O m k Q R m 7 I P X k g j 8 6 d 8 + S 8 O K / D a M H 5 m V k l f + C 8 f Q N y + K f T < / l a t e x i t > A K 0 < l a t e x i t s h a 1 _ b a s e 6 4 = " h 4 C + k v b R l W T U q o d P H M 9 c 9 E a / p z A = " > A A A C F X i c b V C 7 T g J B F J 3 F F + J r 1 Z J m I j G B h u x q o S V q Y 2 K D i T w S F s n s M A s T Z h + Z u W u C G w q / w F K / Q F v 9 A D t j a 2 3 t j z g L F A q e 5 C Y n 5 9 y b e + 9 x I 8 E V W N a X k V l Y X F p e y a 7 m 1 t Y 3 N r f M 7 Z 2 6 C m N J W Y 2 G I p R N l y g m e M B q w E G w Z i Q Z 8 V 3 B G u 7 g L P U b N 0 w q H g Z X M I x Y 2 y e 9 g H u c E t B S x 8 w 7 P o E + J S I 5 G T m C e V C 8 u L Y c y X t 9 K H X M g l W 2 x s D z x J 6 S Q q V 0 f x t l n x 6 r H f P b 6 Y Y 0 9 l k A V B C l W r Y V Q T s h E j g V b J R z Y s U i Q g e k x 1 q a B s R n q p 2 M n x j h f a 1 0 s R d K X Q H g s f p 7 I i G + U k P f 1 Z 3 p y W r W S 8 V / P d f 1 Z 1 a D d 9 x O e B D F w A I 6 2 e z F A k O I 0 4 h w l 0 t G Q Q w 1 I V R y f T y m f S I J B R 1 k T q d i z 2 Y w T + o H Z f u w b F 3 q e E 7 R B F m U R 3 u o i G x 0 h C r o H F V R D V F 0 h 5 7 R C 3 o 1 H o w 3 4 9 3 4 m L R m j O n M L v o D 4 / M H S 8 C i B w = = < / l a t e x i t > A < l a t e x i t s h a 1 _ b a s e 6 4 = " O 6 6 A T y 1 a w 1 t Y L Y U p v p J 8 F g N J C H 8 = " > A A A C A 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X a 1 0 E a M 2 l h G M A 9 I l j A 7 m S R D Z m e W m V k h L C n 9 B F s t L O 3 E 1 t p P E G t / x N k k h S Y e u H A 4 5 1 7 u 4 Y Q x Z 9 p 4 3 p e T W 1 h c W l 7 J r 7 p r 6 x u b W 4 X t n Z q W i S K 0 S i S X q h F i T T k T t G q Y 4 b Q R K 4 q j k N N 6 O L j K / P o d V Z p J c W u G M Q 0 i 3 B O s y w g 2 V m q 2 I m z 6 B P P 0 Y t Q u F L 2 S N w a a J / 6 U F M 8 / 3 L P 4 + d O t t A v f r Y 4 k S U S F I R x r 3 f S 9 2 A Q p V o Y R T k d u K 9 E 0 x m S A e 7 R p q c A R 1 U E 6 j j x C B 1 b p o K 5 U d o R B Y / X 3 R Y o j r Y d R a D e z i H r W y 8 R / v T C M Z l 6 b 7 m m Q M h E n h g o y + d x N O D I S Z Y W g D l O U G D 6 0 B B P F b H h E + l h h Y m x t r m 3 F n + 1 g n t S O S v 5 x y b v x i u V L m C A P e 7 A P h + D D C Z T h G i p Q B Q I S H u A R n p x 7 5 8 V 5 d d 4 m q z l n e r M L f + C 8 / w A 5 N 5 u k < / l a t e x i t > C < l a t e x i t s h a 1 _ b a s e 6 4 = " B p 3 s 0 N X a z I 8 d Q f K 7 i 5 e d A + N 1 h F s = " > A A A C A 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X a 1 0 E Y M p r G M Y B 6 Q L G F 2 M k m G z M w u M 7 N C W F L 6 C b Z a W N q J r b W f I N b + i L N J C k 0 8 c O F w z r 3 c w w l j z r T x v C 8 n t 7 S 8 s r q W X 3 c 3 N r e 2 d w q 7 e 3 U d J Y r Q G o l 4 p J o h 1 p Q z S W u G G U 6 b s a J Y h J w 2 w m E l 8 x t 3 V G k W y V s z i m k g c F + y H i P Y W K n V F t g M C O Z p Z d w p F L 2 S N w F a J P 6 M F C 8 / 3 I v 4 + d O t d g r f 7 W 5 E E k G l I R x r 3 f K 9 2 A Q p V o Y R T s d u O 9 E 0 x m S I + 7 R l q c S C 6 i C d R B 6 j I 6 t 0 U S 9 S d q R B E / X 3 R Y q F 1 i M R 2 s 0 s o p 7 3 M v F f L w z F 3 G v T O w 9 S J u P E U E m m n 3 s J R y Z C W S G o y x Q l h o 8 s w U Q x G x 6 R A V a Y G F u b a 1 v x 5 z t Y J P W T k n 9 a 8 m 6 8 Y v k K p s j D A R z C M f h w B m W 4 h i r U g E A E D / A I

T 8 6 9 8 + K 8 O m / T 1 Z w z u 9 m H P 3 D e f w A 8 Y 5 u m < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " S 9 7 h A c 1 X q 6 o + h n y b y

N 3 E F c 8 G O U A = " > A A A B / H i c b V C 7 S g N B F L 3 r M 8 Z X 1 F K R w S B Y h V 0 t t A z a W C Z g H p A s Y X Y y m w y Z m V 1 m Z o W w x M 7 W V j / A T m w l v 2 J t 6 U 8 4 m 6 T Q x A M X D u f c y 7 m c I O Z M G 9 f 9 d J a W V 1 b X 1 n M b + c 2 t 7 Z 3 d w t 5 + X U e J I r R G I h 6 p Z o A 1 5 U z S m m G G 0 2 a s K B Y B p 4 1 g c J P 5 j X u q N I v k n R n G 1 B e 4 J 1 n I C D a Z 1 I 7 7 r F M o u i V 3 A r R I v B k p l o / G 1 e / H 4 3 G l U / h q d y O S C C o N 4 V j r l u f G x k + x M o x w O s q 3 E 0 1 j T A a 4 R 1 u W S i y o 9 t P J r y N 0 a p U u C i N l R x o 0 U X 9 f p F h o P R S B 3 R T Y 9 P W 8 l 4 n / e k E g 5 q J N e O W n T M a J o Z J M k 8 O E I x O h r A n U Z Y o S w 4 e W Y K K Y f R 6 R P l a Y G N t X 3 r b i z X e w S O r n J e + i 5 F Z t P d c w R Q 4 O 4 Q T O w I N L K M M t V K A G B P

r w B M / w 4 j w 4 r 8 6 b 8 z 5 d X X J m N w f w B 8 7 H D 0 B 7 m P Y = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " S 9 7 h A c 1 X q 6 o + h n y b y 

N 3 E F c 8 G O U A = " > A A A B / H i c b V C 7 S g N B F L 3 r M 8 Z X 1 F K R w S B Y h V 0 t t A z a W C Z g H p A s Y X Y y m w y Z m V 1 m Z o W w x M 7 W V j / A T m w l v 2 J t 6 U 8 4 m 6 T Q x A M X D u f c y 7 m c I O Z M G 9 f 9 d J a W V 1 b X 1 n M b + c 2 t 7 Z 3 d w t 5 + X U e J I r R G I h 6 p Z o A 1 5 U z S m m G G 0 2 a s K B Y B p 4 1 g c J P 5 j X u q N I v k n R n G 1 B e 4 J 1 n I C D a Z 1 I 7 7 r F M o u i V 3 A r R I v B k p l o / G 1 e / H 4 3 G l U / h q d y O S C C o N 4 V j r l u f G x k + x M o x w O s q 3 E 0 1 j T A a 4 R 1 u W S i y o 9 t P J r y N 0 a p U u C i N l R x o 0 U X 9 f p F h o P R S B 3 R T Y 9 P W 8 l 4 n / e k E g 5 q J N e O W n T M a J o Z J M k 8 O E I x O h r A n U Z Y o S w 4 e W Y K K Y f R 6 R P l a Y G N t X 3 r b i z X e w S O r n J e + i 5 F Z t P d c w R Q 4 O 4 Q T O w I N L K M M t V K A G B P

w = " > A A A C C X i c b V C 7 S g N B F L 3 r M 8 Z H V i 0 V G Q y C V d j V Q h s h a G O Z g H l A s o b Z y W w y Z G Z 3 m Z k V l i W l l Z 9 g q x 9 g F 2 z 9 C m t L f 8 L J o 9 D E A x c O 5 9 z L P R w / 5 k x p x / m 0 l p Z X V t f W c x v 5 z a 3 t n Y K 9 u 1 d X U S I J r Z G I R 7 L p Y 0 U 5 C 2 l N M 8 1 p M 5 Y U C 5 / T h j + 4 G f u N B y o V i 8 I 7 n c b U E 7 g X s o A R r I 3 U s Q t t g X W f Y J 4 1 h 1 f p v d O x i 0 7 J m Q A t E n d G i u X D U f X 7 8 W h U 6 d h f 7 W 5 E E k F D T T h W q u U 6 s f Y y L D U j n A 7 z 7 U T R G J M B 7 t G W o S E W V H n Z J P g Q n R i l i 4 J I m g k 1 m q i / L z I s l E q F b z b H M d W 8 N x b / 9 X x f z L 3 W w a W X s T B O N A 3 J 9 H O Q c K Q j N K 4 F d Z m k R P P U E E w k M + E R 6 W O J i T b l 5 U 0 r 7 n w H i 6 R + V n L P S 0 7 V 1 H M N U + T g A I 7 h F F y 4 g D L c Q g V q Q C C B Z 3 i B V + v J e r N G 1 v t 0 d c m a 3 e z D H 1 g f P x U H n c I = < / l a t e x i t > z 0

< l a t e x i t s h a 1 _ b a s e 6 4 = " U 1 l 7 D t q 3 3 F g r R s 5

4 b T L L s j N b C y Q = " > A A A B + 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F W a 1 0 E Y M 2 l h G N A 9 I 1 j A 7 m U 2 G z O w u M 7 N C X P I J t v o B 6 c T W b / A T x N o f c Z J Y a O K B C 4 d z 7 u V c T p A I r g 3 G n 0 5 u Y X F p e S W / 6 q 6 t b 2 x u F b Z 3 a j p O F W V V G o t Y N Q K i m e A R q x p u B G s k i h E Z C F Y P + p d j v 3 7 P l O Z x d G s G C f M l 6 U Y 8 5 J Q Y K 9 0 8 3 O F 2 o Y h L e A I 0 T 7 w f U j x / d 8 + S 0 Y d b a R e + W p 2 Y p p J F h g q i d d P D i f E z o g y n g g 3 d V q p Z Q m i f d F n T 0 o h I p v 1 s 8 u o Q H V i l g 8 J Y 2 Y k M m q i / L z I i t R 7 I w G 5 K Y n p 6 1 h u L / 3 p B I G e i T X j q Z z x K U s M i O k 0 O U 4 F M j M Z F o A 5 X j B o x s I R Q x e 3 z i P a I I t T Y u l z b i j f b w T y p H Z W 8 4 x K + x s X y B U y R h z 3 Y h 0 P w 4 A T K c A U V q A K F L j z C E z w 7 Q 2 f k v D i v 0 9 W c 8 3 O z C 3 / g v H 0 D m i 6 X 7 Q = = < / l a t e x i t > z 1

< l a t e x i t s h a 1 _ b a s e 6 4 = " V 7 A S t E 6 F L E t g S P N 8 7 a p

V g z m L N D U = " > A A A B + 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X a 1 0 E Y M 2 l h G N A 9 I 1 j A 7 m U 2 G z M w u M 7 N C X P I J t v o B 6 c T W b / A T x N o f c Z J Y a O K B C 4 d z 7 u V c T p h w p o 3 n f T q 5 h c W l 5 Z X 8 q r u 2 v r G 5 V d j e q e k 4 V Y R W S c x j 1 Q i x p p x J W j X M c N p I F M U i 5 L Q e 9 i / H f v 2 e K s 1 i e W s G C Q 0 E 7 k o W M Y K N l W 4 e 7 v x 2 o e i V v A n Q P P F / S P H 8 3 T 1 L R h 9 u p V 3 4 a n V i k g o q D e F Y 6 6 b v J S b I s D K M c D p 0 W 6 m m C S Z 9 3 K V N S y U W V A f Z 5 N U h O r B K B 0 W x s i M N m q i / L z I s t B 6 I 0 G 4 K b H p 6 1 h u L / 3 p h K G a i T X Q a Z E w m q a G S T J O j l C M T o 3 E R q M M U J Y Y P L M F E M f s 8 I j 2 s M D G 2 L t e 2 4 s 9 2 M E 9 q R y X / u O R d e 8 X y B U y R h z 3 Y h 0 P w 4 Q T K c A U V q A K B L j z C E z w 7 Q 2 f k v D i v 0 9 W c 8 3 O z C 3 / g v H 0 D m 8 O X 7 g = = < / l a t e x i t > y 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " 1 Y 9 j t I A c X c 6 W D 6 d G 1 4 l q r S Z U M k Y = " > A A A B + 3 i c b V C 7 S g N B F L 0 b X z G + o p a K D A b B K u x q o W X Q x j J B 8 4 B k D b O T S T J k Z n a Z m R W W J a W l r X 6 A n d g K + R V r S 3 / C y a P Q x A M X D u f c y 7 m c I O J M G 9 f 9 d D J L y y u r a 9 n 1 3 M b m 1 v Z O f n e v p s N Y E V o l I Q 9 V I 8 C a c i Z p 1 T D D a S N S F I u A 0 3 o w u B 7 7 9 Q e q N A v l n U k i 6 g v c k 6 z L C D Z W u k 3 u v X a + 4 B b d C d A i 8 W a k U D o c V b 4 f j 0 b l d v 6 r 1 Q l J L K g 0 h G O t m 5 4 b G T / F y j D C 6 T D X i j W N M B n g H m 1 a K r G g 2 k 8 n r w 7 R i V U 6 q B s q O 9 K g i f r 7 I s V C 6 0 Q E d l N g 0 9 f z 3 l j 8 1 w s C M R d t u p d + y m Q U G y r J N L k b c 2 R C N C 4 C d Z i i x P D E E k w U s 8 8 j 0 s c K E 2 P r y t l W v P k O F k n t r O i d F 9 2 K r e c K p s j C A R z D K X h w A S W 4 g T J U g U A P n u A Z X p y h 8 + q 8 O e / T 1 Y w z u 9 m H P 3 A + f g A u r p h X < / l a t e x i t > y 2 < l a t e x i t s h a 1 _ b a s e 6 4 = " M 1 N z n f h X n M F o s L H 4 H X e 5 0 7 v F G X w = " > A A A B + 3 i c b V C 7 S g N B F L 3 r M 8 Z X 1 F K R w S B Y h d 1 Y a B m 0 s U z Q P C B Z w + x k k g y Z m V 1 m Z o V l S W l p q x 9 g J 7 Z C f s X a 0 p 9 w 8 i g 0 8 c C F w z n 3 c i 4 n i D j T x n U / n a X l l d W 1 9 c x G d n N r e 2 c 3 t 7 d f 0 2 G s C K 2 S k I e q E W B N O Z O 0 a p j h t B E p i k X A a T 0 Y X I / 9 + g N V m o X y z i Q R 9 Q X u S d Z l B B s r 3 S b 3 x X Y u 7 x b c C d A i 8 W Y k X z o a V b 4 f j 0 f l d u 6 r 1 Q l J L K g 0 h G O t m 5 4 b G T / F y j D C 6 T D b i j W N M B n g H m 1 a K r G g 2 k 8 n r w 7 R q V U 6 q B s q O 9 K g i f r 7 I s V C 6 0 Q E d l N g 0 9 f z 3 l j 8 1 w s C M R d t u p d + y m Q U G y r J N L k b c 2 R C N C 4 C d Z i i x P D E E k w U s 8 8 j 0 s c K E 2 P r y t p W v P k O F k m t W P D O C 2 7 F 1 n M F U 2 T g E E 7 g D D y 4 g B L c Q B m q Q K A H T / A M L 8 7 Q e X X e n P f p 6 p I z u z m A P 3 A + f g A w Q 5 h Y < / l a t e x i t > z2 < l a t e x i t s h a 1 _ b a s e 6 4 = " 4 w e G o m N j p m J m W Z L 3 O 9 4 Y 1 O s Z s 9 o = " > A A A C A H i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X Z j o Y 0 Y t L G M Y B 6 Q r G F m M p s M m Z 1 d Z m a F G N L 4 C b b a 2 d i J r T / g J 4 i 1 P + L k U W j i g Q u H c + 7 l X A 5 J B N f G 8 7 6 c z M L i 0 v J K d t V d W 9 / Y 3 M p t 7 1 R 1 n C r K K j Q W s a o T r J n g k l U M N 4 L V E 8 V w R A S r k d 7 F y K / d M q V 5 L K 9 N P 2 F B h D u S h 5 x i Y 6 V 6 k 2 C F 7 m 6 K r V z e K 3 h j o H n i T 0 n + 7 M M 9 T Z 4 / 3 X I r 9 9 1 s x z S N m D R U Y K 0 b v p e Y Y I C V 4 V S w o d t M N U s w 7 e E O a 1 g q c c R 0 M B j / O 0 Q H V m m j M F Z 2 p E F j 9 f f F A E d a 9 y N i N y N s u n r W G 4 n / e o R E M 9 E m P A k G X C a p Y Z J O k s N U I B O j U R u o z R W j R v Q t w V R x + z y i X a w w N b Y z 1 7 b i z 3 Y w T 6 r F g n 9 U 8 K 6 8 f O k c J s j C H u z D I f h w D C W 4 h D J U g I K A B 3 i E J + f e e X F e n b f J a s a Z 3 u z C H z j v P w x n m d I = < / l a t e x i t > A K 1 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " q K h e j S f y k + 5 4 m P 4 3 K q Y 2 N v y x d M 8 = " > A A A C H H i c b V A 9 S w N B E N 2 L X z F + R S 1 t j g Q h a c K d F l p G b Q S b C O Y D c v H Y 2 + w l S / b 2 j t 0 5 I R 7 p 7 W y t L C 1 t 9 Q f Y i a 1 g 7 R 9 x L 0 m h i Q 8 G H u / N M D P P i z h T Y F l f R m Z h c W l 5 J b u a W 1 v f 2 N z K b + 8 0 V B h L Q u s k 5 K F s e V h R z g S t A w N O W 5 G k O P A 4 b X q D s 9 R v 3 l C p W C i u Y B j R T o B 7 g v m M Y N C S m y 8 4 A Y Y + w T w 5 G T m c + l C 6 u L Z d h w k f h o 5 k v T 6 U 3 X z R q l h j m P P E n p J i t X x / G 2 W f H m p u / t v p h i Q O q A D C s V J t 2 4 q g k 2 A J j H A 6 y j m x o h E m A 9 y j b U 0 F D q j q J O N f R u a + V r q m H 0 p d A s y x + n s i w Y F S w 8 D T n e n l a t Z L x X 8 9 z w t m V o N / 3 E m Y i G K g g k w 2 + z E 3 I T T T p M w u k 5 Q A H 2 q C i W T 6 e J P 0 s c Q E d J 4 5 n Y o 9 m 8 E 8 a R x U 7 M O K d a n j O U U T Z N E e K q A S s t E R q q J z V E N 1 R N A d e k Y v 6 N V 4 N N 6 M d + N j 0 p o x p j O 7 6 A + M z x 8 3 4 K U z < / l a t e x i t > K0 < l a t e x i t s h a 1 _ b a s e 6 4 = " o B l V 8 I U 3 X c I I z H b 4 G L F J z u f c j 1 0 = " > A A A C D 3 i c b V D L S s N A F J 3 4 r P E V F V d u g k V w V Z K 6 0 I 1 Y d C O 4 q W A f 0 M Q y m U z a o T O T M D M R S s h H + A n i T j / A n b g T v 0 B c + y N O 2 i 6 0 9 c C F w z n 3 c g 8 n S C i R y n G + j L n 5 h c W l 5 d K K u b q 2 v r F p b W 0 3 Z Z w K h B s o p r F o B 1 B i S j h u K K I o b i c C Q x Z Q 3 A o G F 4 X f u s N C k p j f q G G C f Q Z 7 n E Q E Q a W l r r X r K U J D n H k M q j 6 C N L v K 8 1 u n a 5 W d i j O C P U v c C S m f v Z u n y e O n W e 9 a 3 1 4 Y o 5 R h r h C F U n Z c J 1 F + B o U i i O L c 9 F K J E 4 g G s I c 7 m n L I s P S z U f z c P t B K a E e x 0 M O V P V J / X 2 S Q S T l k g d 4 s U s p p r x D / 9 Y K A T b 1 W 0 Y m f E Z 6 k C n M 0 / h y l 1 F a x X Z R j h 0 R g p O h Q E 4 g E 0 e F t 1 I c C I q U r N H U r 7 n Q H s 6 R Z r b h H l e q 1 U 6 6 d g z F K Y A / s g 0 P g g m N Q A 5 e g D h o A g Q w 8 g C f w b N w b L 8 a r 8 T Z e n T M m N z v g D 4 y P H 7 h n o D k = < / l a t e x i t > K 1 < l a t e x i t s h a 1 _ b a s e 6 4 = " M R 9 r G Q G B F 6 r k C h H k V e 9 5 g k 8 I g I 4 = " > A A A C B X i c b V C 7 S g N B F L 3 r M 6 6 v q K X N Y h C s w m 4 s t B G D N o J N B P O A Z A 2 z k 9 l k y M z s O j M r h C W 1 n 2 C r W N u J r a W f I N b + i L N J C k 0 8 c O F w z r 3 c w w l i R p V 2 3 S 9 r b n 5 h c W k 5 t 2 K v r q 1 v b O a 3 t m s q S i Q m V R y x S D Y C p A i j g l Q 1 1 Y w 0 Y k k Q D x i p B / 3 z z K / f E a l o J K 7 1 I C Y + R 1 1 B Q 4 q R N p L f 4 k j 3 M G L p 5 f D G a + c L b t E d w Z k l 3 o Q U T j / s k / j 5 0 6 6 0 8 9 + t T o Q T T o T G D C n V 9 N x Y + y m S m m J G h n Y r U S R G u I + 6 p G m o Q J w o P x 2 F H j r 7 R u k 4 Y S T N C O 2 M 1 N 8 X K e J K D X h g N r O Q a t r L x H + 9 I O B T r 3 V 4 7 K d U x I k m A o 8 / h w l z d O R k l T g d K g n W b G A I w p K a 8 A 7 u I Y m w N s X Z p h V v u o N Z U i s V v c N i 6 c o t l M 9 g j B z s w h 4 c g A d H U I Y L q E A V M N z C A z z C k 3 V v v V i v 1 t t 4 d c 6 a 3 O z A H 1 j v P 3 9 l n F M = < / l a t e x i t >

K1:

< l a t e x i t s h a 1 _ b a s e 6 4 = " P i 0 

I 6 + H 4 u E B V T V 4 r B n c z V + v n O G M = " > A A A C D 3 i c b V D L S s N A F J 3 4 r P E V F V d u g k V w V Z K 6 0 I 1 Y d C O 4 q W A f 0 M Q y m U z a o T O T M D M R S s h H + A n i T j / A n b g T v 0 B c + y N O 2 i 6 0 9 c C F w z n 3 c g 8 n S C i R y n G + j L n 5 h c W l 5 d K K u b q 2 v r F p b W 0 3 Z Z w K h B s o p r F o B 1 B i S j h u K K I o b i c C Q x Z Q 3 A o G F 4 X f u s N C k p j f q G G C f Q Z 7 n E Q E Q a W l r r X r K U J D n H k M q j 6 C N L v K 8 1 u 3 a 5 W d i j O C P U v c C S m f v Z u n y e O n W e 9 a 3 1 4 Y o 5 R h r h C F U n Z c J 1 F + B o U i i O L c 9 F K J E 4 g G s I c 7 m n L I s P S z U f z c P t B K a E e x 0 M O V P V J / X 2 S Q S T l k g d 4 s U s p p r x D / 9 Y K A T b 1 W 0 Y m f E Z 6 k C n M 0 / h y l 1 F a x X Z R j h 0 R g p O h Q E 4 g E 0 e F t 1 I c C I q U r N H U r 7 n Q H s 6 R Z r b h H l e q 1 U 6 6 d g z F K Y A / s g 0 P g g m N Q A 5 e g D h o A g Q w 8 g C f

A 4 = " > A A A C D 3 i c b V D L S s N A F J 3 4 r P E V F V d u g k V w V Z K 6 0 I 1 Y d C O 4 q W A f 0 M Q y m U z a o T O T M D M R S s h H + A n i T j / A n b g T v 0 B c + y N O 2 i 6 0 9 c C F w z n 3 c g 8 n S C i R y n G + j L n 5 h c W l 5 d K K u b q 2 v r F p b W 0 3 Z Z w K h B s o p r F o B 1 B i S j h u K K I o b i c C Q x Z Q 3 A o G F 4 X f u s N C k p j f q G G C f Q Z 7 n E Q E Q a W l r r X r K U J D n H k M q j 6 C N L v K 8 9 t q 1 y o 7 F W c E

c W k 5 t 2 K v r q 1 v b O a 3 t m s q S i Q m V R y x S D Y C p A i j g l Q 1 1 Y w 0 Y k k Q D x i p B / 3 z z K / f E a l o J K 7 1 I C Y + R 1 1 B Q 4 q R N p L f 4 k j 3 M G L p 5 f C m 1 M 4 X 3 K I 7 g j N L v A k p n H 7 Y J / H z p 1 1 p 5 7 9 b n Q g n n A i N G V K q 6 b m x 9 l M k N c W M D O 1 W o k i M c B 9 1 S d N Q g T h R f j o K P X T 2 j d J x w k i a E d o Z q b 8 v U s S V G v D A b G Y h 1 b S X i f 9 6 Q c C n X u v w 2 E + p i B N N B B 5 / D h P m 6 M j J K n E 6 V B K s 2 c A Q h C U 1 4 R 3 c Q x J h b Y q z T S v e d A e z p F Y q e o f F 0 p V b K J / B G D n Y h T 0 4 A A + O o A w X U I E q Y L i F B 3 i E J + v e e

L o v Q r s p s O n q e W 8 s / u u F o Z i L N t F Z M G A y S Q 2 V Z J o c p R y Z G I 2 L Q G 2 m K D G 8 b w k m i t n n E e l i h Y m x d b m 2 F X + + g 0 V S L R b 8 k 0 L x x s u X L m G K L B z A I R y D D 6 d Q g m s o Q w U I d O A J n u H F G T o j 5 9 V 5 m 6 5 m n N n N P v y B 8 / 4 D n f y X 8 Q = = < / l a t e x i t > K 0 = X X T

< l a t e x i t s h a 1 _ b a s e 6 4 = " 8 0 a n 4 a m q p Q g Z A c H 3 j / u 4 s W x X 9 P U = " > A

A A C I n i c b V D L S g M x F M 3 4 r O N r 1 K W b Y B G 6 K j N 1 o Z t i 0 Y 3 g p k J f 0 B e Z N N O G Z j J D k h H K M L / g T + g n u N W t 4 E 5 c i O D O H z H T F q m t B w L n n n M v 9 + a 4 I a N S 2 f a n s b S 8 s r q 2 n t k w N 7 e 2 d 3 a t v f 2 a D C K B S R U H L B A N F 0 n C K C d V R R U j j V A Q 5 L u M 1 N 3 h Z e r X b 4 m Q N O A V N Q p J 2 0 d 9 T j 2 K k d J S 1 8 q 1 f K Q G G L H 4 O u n Y s A h / 6 0 Y y Q z u V r p W 1 8 / Y Y c J E 4 U 5 I 9 f z G L 4 f 2 H W e 5 a 3 6 1 e g C O f c I U Z k r L p 2 K F q x 0 g o i h l J z F Y k S Y j w E P V J U 1 O O f C L b 8 f h H C T z W S g 9 6 g d C P K z h W Z y d i 5 E s 5 8 l 3 d m R 4 p 5 7 1 U / N d z X X 9 u t f L O 2 j H l Y a Q I x 5 P N X s S g C m C a F + x R Q b B i I 0 0 Q F l Q f D / E A C Y S V T t X U q T j z G S y S W i H v n O Q L N 3 a 2 d A E m y I B D c A R y w A G n o A S u Q B l U A Q Z 3

4 B E 8 g W f j w X g 1 3 o z 3 S e u S M Z 0 5 A H 9 g f P 0 A G q + n z A = = < / l a t e x i t > See Table 1 for operation definitions. Illustration and description adapted from Figure 3 in Novak et al. (2019) .

y 0 = X < l a t e x i t s h a _ b a s e = " R F G x L m i k P w s d f T V l u J u Y o M w = " > A A A C C i c b V C S g N B F J N r x h f G y V G Q y C V d i N h T Z C M Y y A f O A J I b Z y W w y Z G Z m Z l V l i W l p Z g q x g E T W j C C e c T V J o o E L h P u R O F z K q t O N W p m l Z X V t e x b m N z a v H z u / W V R B J T G o Y I F s e k g R R g W p a a o Z a Y a S I O x v C G V n f u C N S U D c D g k H Y g v o U I k r p P b x A d s c Q F G L G m O u n b B K T o T w E X i z k i h f D C u f j c j i t d + v d C D E i d C Y I a V a r h P q T o K k p p i R U a d K R I i P E R j J U I E U J l E H F j o / S g H g z Q s O J + v s i Q V y p m H t m M o r U / N f z P D W v v n n Y S K M N J E O l n P J Q B z A t B v a o J F i z B C E J T X h I R g i b A e V M K + B u k X i q p V S d R z C a b I g n w B E A C A G V y D C q g B D O B E g G L a j W q W e / T Y w u k D f B / A D G v J Y < / l

0 (input) y 0 = X K 0 = X X T Θ 0 = 0 0 (pre-activations) z 0 = Conv y 0 K0 = A K 0 Θ0 = K0 + A Θ 0 1 (activations) y 1 = φ z 0 K 1 = T K0 Θ 1 = Ṫ K0 Θ0 1 (pre-activations) z 1 = Conv y 1 K1 = A K 1 Θ1 = K1 + A Θ 1 2 (activations) y 2 = φ z 1 K 2 = T K1 Θ 2 = Ṫ K1 Θ1 2 (readout) z 2 = Dense • Flatten y 2 K2 = Tr K 2 Θ2 = K2 + Tr Θ 2

assume z is a mean zero multivariate Gaussian. We wish to compute the kernel corresponding to h = Dense (σ ω , σ b ) (φ(z)) by computing the kernels of y = φ(z) and h = Dense (σ ω , σ b ) (y) separately. Here,

EQUATION

and the variables W ij and β i are i.i.d. Gaussian N (0, 1). We will compute kernel operationsdenoted φ * and Dense(σ ω , σ b ) * -induced by the tensor operations φ and Dense(σ ω , σ b )5 . Finally, we will compute the kernel operation associated with the composition

(Dense(σ ω , σ b ) • φ) * = Dense(σ ω , σ b ) * • φ * .

First we compute the NNGP and NT kernels for y. To compute K y note that from its definition,

EQUATION

Since φ does not introduce any new variables Θ y can be computed as,

Θ y = E θ ∂φ(z i ) ∂θ ∂φ(z i ) ∂θ T = E θ diag( φ(z i )) ∂z i ∂θ ∂z i ∂θ T diag( φ(z i )) = Ṫ (K z ) Θ z .

Taken together these equations imply that,

EQUATION

Published as a conference paper at ICLR 2020 will be the translation rule for a pointwise nonlinearity. Note that Equation Equation 4only has an analytic expression for a small set of activation functions φ.

Next we consider the case of a dense operation. Using the independence between the weights, the biases, and h it follows that,

EQUATION

Finally, the NTK of h can be computed as a sum of two terms:

Θ h = E W,β,θ ∂h i ∂(W, β) ∂h i ∂(W, β) T + E W,β,θ ∂h i ∂θ ∂h i ∂θ T = σ 2 ω K y + σ 2 b + σ 2 ω Θ y . (6)

This gives the translation rule for the dense layer in terms of K y and Θ y as, 4 and Listing 2 for models admitting this optimization). Finally, if the network has no convolutional layers, the cost further reduces to O(|X |

EQUATION

2 ) (see Listing 3 for an 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 convolutional layers A 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 |-blocks. This allows utilization of modern hardware accelerators, many of which target 2D convolutions as their primary machine learning 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 is 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 it is desirable to parallelize the kernel computation across devices (CPUs, GPUs, or TPUs). NEURAL TANGENTS provides an easy way to perform both of these common tasks using a single batch decorator shown below:

2 0 2 1 2 2 2 3 2 4 2 5 Batch Size 2 -10 2 -9 2 -8 2 -7 2 -6 2 -5 2 -4 Time / Sample (s) GPU Count = 1 GPU Count = 2 GPU Count = 4 GPU Count = 8 2 -1 2 0 2 1 2 2 2 3 GPU Count 2 -11 2 -10 2 -9 2 -8 2 -7 2 -6 2 -5

Time / Sample (s) This code works with either analytic kernels or empirical kernels. By default, it automatically shares the computation over all available devices. We plot the performance as a function of batch size and number of accelerators when computing the theoretical NTK of a 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 to 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 as additional metadata under the hood. This eliminates redundant transpositions6 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, we 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 NEURAL TANGENTS that we are exploring. There are more layers we would like to add in the future ( §D) 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 the community to join our efforts by contributing new layers to the library ( §B.7), or by using it for research and providing feedback! 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 is a strict and minor subset of our library. More specifically, at the time of the submission, Lee et al. (2019) have released code equivalent to nt.linearize , nt.empirical_ntk_fn , nt.predict.gradient_descent_mse , nt.predict.gradient_descent , and nt.predict.momentum . Every other part of the library (most notably, nt.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 +9, 500/ -2, 500 lines of code.

3. GPy (2012), GPFlow (Matthews et al., 2017 ), GPyTorch (Gardner et 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 to infinite neural networks (NNGP/NTK; nt.stax ), nor do they provide the tools and convenience for studying wide but finite networks and their training dynamics ( nt.taylor_expand , nt.predict , nt.monte_carlo_kernel_fn ).

On the other hand, NEURAL TANGENTS does not provide any tools for approximate inference with these kernels.

B LIBRARY DESCRIPTION:

NEURAL TANGENTS provides a 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 et al., 2018b) in JAX.

B.1 NEURAL NETWORKS WITH JAX:

stax represents each component of a 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 on a random seed and an input shape generates a random draw of trainable parameters for a neural network. Calling apply_fn on these parameters and a batch of inputs returns the outputs of the given finite neural network.

from jax.experimental import stax init_fn, apply_fn = stax.serial(stax.Dense(512), stax.Relu, stax.Dense(10)) _, params = init_fn(key, (-1, 32 * 32 * 3)) fx_train, fx_test = apply_fn(params, x_train), apply_fn(params, x_test)

B.2 INFINITE NEURAL NETWORKS WITH NEURAL TANGENTS:

We extend stax layers to return a third function kernel_fn , which represents the covariance functions of the infinite NNGP and NTK networks of the given architecture (recall that since infinite networks are GPs, they are fully defined by their covariance functions, assuming 0 mean as is common in the literature).

from neural_tangents import stax init_fn, apply_fn, kernel_fn = stax.serial(stax.Dense(512), stax.Relu(), stax.Dense(10))

We demonstrate a specification of a more complicated architecture (WideResNet) in Listing 1.

kernel_fn accepts two batches of inputs x1 and x2 and returns their NNGP covariance and NTK matrices as kernel_fn(x1, x2 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.

:

See §A for NEURAL TANGENTS comparison against specific relevant prior works.

Spatially local average pooling in intermediary layers can be constructed in a similar fashion ( §A.3). We focus on global average pooling in this work to more effectively isolate the effects of pooling from other aspects of the model like local connectivity or equivariance.

T (Σ) ≡ E φ(u)φ(u) T , Ṫ (Σ) ≡ E φ (u)φ (u) T , u ∼ N (0, Σ), as in(Lee et al., 2019).

These transpositions could not be automatically fused by the XLA compliler.

Figures and Tables:

Figure 1: Training dynamics for an ensemble of finite-width networks compared with an infinite 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 ensemble of finite-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.

.

Figure 2: Convergence of the Monte Carlo (MC) estimates of the WideResNet WRN-28-k (where k is the widening factor) NNGP and NTK kernels (computed with monte_carlo_kernel_fn )to their analytic values (WRN-28-∞, computed with kernel_fn ), as the network gets wider by increasing the widening factor (vertical axis) and as more random networks are averaged over (horizontal axis). Experimental detail. The kernel is computed in 32-bit precision on a 100 × 50 batch of 8 × 8-downsampled CIFAR10(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 1 (all logits are i.i.d. conditioned on a given input). The displayed distance is the relative Frobenius norm squared, i.e. K -K k,n

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 is common in prior work(Lee 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, -0.1, . . . , -0.1) . For each training set size, the best model in the family is selected by minimizing the mean negative marginal log-likelihood (NLL, right) on the training set.

Figure 2: Validation accuracy (left) of an MC-CNN-GP increases with n ⇥ M (i.e. channel count times number of samples) and approaches that of the exact CNN-GP (not shown), while the distance (right) to the exact kernel decreases. The dark band in the left plot corresponds to ill-conditioning of K L+1n,M when the number of outer products contributing to K L+1 n,M approximately equals its rank. Values reported are for a 3-layer model applied to a 2K/4K train/validation subset of CIFAR10 downsampled to 8 ⇥ 8. See Figure7for similar results with other architectures and §A.7.2 for experimental details.

Figure 2: Validation accuracy (left) of an MC-CNN-GP increases with n ⇥ M (i.e. channel count times number of samples) and approaches that of the exact CNN-GP (not shown), while the distance (right) to the exact kernel decreases. The dark band in the left plot corresponds to ill-conditioning of K L+1n,M when the number of outer products contributing to K L+1 n,M approximately equals its rank. Values reported are for a 3-layer model applied to a 2K/4K train/validation subset of CIFAR10 downsampled to 8 ⇥ 8. See Figure7for similar results with other architectures and §A.7.2 for experimental details. random finite width networks and use the empirical uncentered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,

Figure 2: Validation accuracy (left) of an MC-CNN-GP increases with n ⇥ M (i.e. channel count times number of samples) and approaches that of the exact CNN-GP (not shown), while the distance (right) to the exact kernel decreases. The dark band in the left plot corresponds to ill-conditioning of K L+1n,M when the number of outer products contributing to K L+1 n,M approximately equals its rank. Values reported are for a 3-layer model applied to a 2K/4K train/validation subset of CIFAR10 downsampled to 8 ⇥ 8. See Figure7for similar results with other architectures and §A.7.2 for experimental details. random finite width networks and use the empirical uncentered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,

r w B M / w 4 j w 4 r 8 6 b 8 z 5 d X X J m N w f w B 8 7 H D 0 B 7 m P Y = < / l a t e x i t > X = y 0 < l a t e x i t s h a 1 _ b a s e 6 4 = " w 7 f O K h z B 6 3 i 5 9 c d F S 6 X E C o b + p D

w b N w b L 8 a r 8 T Z e n T M m N z v g D 4 y P H 7 n 8 o D o = < / l a t e x i t > K2 < l a t e x i t s h a 1 _ b a s e 6 4 = " n g B e S U G o b 5 e d b B D o d / 0 m s E R 5 0

e 5 a 4 E 1 I + e z d P k 8 d P s 9 6 1 v r 0 w R i n D X C E K p e y 4 T q L 8 DA p F E M W 5 6 a U S J x A N Y A 9 3 N O W Q Y e l n o / i 5 f a C V 0 I 5 i o Y c r e 6 T + v s g g k 3 L I A r 1 Z p J T T X i H + 6 w U B m 3 q t o h M / I z x J F e Z o / D l K q a 1 i u y j H D o n A S N G h J h A J o s P b q A 8 F R E p X a O p W 3 O k O Z k m z W n G P K t V r p 1 w 7 B 2 O U w B 7 Y B 4 f A B c e g B i 5 B H T Q A A h l 4 A E / g 2 b g 3 X o x X 4 2 2 8 O m d M b n b A H x g f P 7 u R o D s = < / l a t e x i t > K 2 < l a t e x i t s h a 1 _ b a s e 6 4 = " u k 5 e D t z 7 F w T C L B L 8 9 u b W C 4 L c R b U = " > A A A C B X i c b V C 7 S g N B F L 3 r M 6 6 v q K X N Y h C s w m 4 s t B G D N o J N B P O A Z A 2 z k 9 l k y M z s O j M r h C W 1 n 2 C r W N u Jr a W f I N b + i L N J C k 0 8 c O F w z r 3 c w w l i R p V 2 3 S 9 r b n 5 h

r F e r b f x 6 p w 1 u d m B P 7 D e f w C A + p x U < / l a t e x i t > z 2 < l a t e x i t s h a 1 _ b a s e 6 4 = " r S j l d e d m u l J S B 3 1M u o 5 9 2 G + L D L g = " > A A A B + 3 i c b V C 7 S g N B F L 0 b X 3 F 9 R S 1 t B o N g F X Z j o Y 0 Y t L G M a B 6 Q r G F 2 M p s M m Z l d Z m a F G P I J t v o B 6 c T W b / A T x N o f c f I o N P H A h c M 5 9 3 I u J 0 w 4 0 8 b z v p z M 0 v L K 6 l p 2 3 d 3 Y 3 N r e y e 3 u V X W c K k I r J O a x q o d Y U 8 4 k r R h m O K 0 n i m I R c l o L e 1 d j v / Z A l W a x v D P 9 h A Y C d y S L G M H G S r e P 9 8 V W L u 8 V v A n Q I v F n J H / x 4 Z 4 n o 0 + 3 3 M p 9 N 9 s x S Q W V h n C s d c P 3 E h M M s D K M cD p 0 m 6 m m C S Y 9 3 K E N S y U W V A e D y a t D d G S V N o p i Z U c a N F F / X w y w 0

a t e x i t > T < l a t e x i t s h a 1 _ b a s e 6 4 = " a o q m H i i q O 2 Z B z C s L 7 L f I 0 s s 8 c 6 M = " > A AA C A 3 i c b V C 7 S g N B F L 3 r M 6 6 v q K X N Y B C s w m 4 s t B G D N p Y R 8 o J k C b O T S T J k Z m e Z m R X C k t J P s N X C 0 k 5 s r f 0 E s f Z H n E 1 S a O K B C 4 d z 7 u U e T h h z p o 3 n f T l L y y u r a + u 5 D X d z a 3 t n N 7 + 3 X 9 c y U Y T W i O R S N U O s K W c R r R l m O G 3 G i m I R c t o I h 9 e Z 3 7 i j S j M Z V c 0 o p o H A / Y j 1 G M H G S q 2 2 w G Z A M E + r 4 0 6 + 4 B W 9 C d A i 8 W e k c P n h X s T P n 2 6 l k / 9 u d y V J B I 0 M 4 V j r l u / F J k i x M o x w O n b b i a Y x J k P c p y 1 L I y y o D t J J 5 D E 6 t k o X 9 a S y E x k 0 U X 9 f p F h o P R K h 3 c w i 6 n k v E / / 1 w l D M v T a 9 8 y B l U Z w Y G p H p 5 1 7 C k Z E o K w R 1 m a L E 8 J E l m C h m w y M y w A o T Y 2 t z b S v + f A e L p F 4 q + q f F 0 q 1 X K F / B F D k 4 h C M 4 A R / O o A w 3 U I E a E J D w A I / w 5 N w 7 L 8 6 r 8 z Z d X X J m N w f w B 8 7 7 D 1 f 9 m 7 k = < / l a t e x i t > T < l a t e x i t s h a 1 _ b a s e 6 4 = " a o q m H i i q O 2 Z B z C s L 7 L f I 0 s s 8 c 6 M = " > A A A C A 3 i c b V C 7 S g N B F L 3 r M 6 6 v q K X N Y B C s w m 4 s t B G D N p Y R 8 o J k C b O T S T J k Z m e Z m R X C k t J P s N X C 0 k 5 s r f 0 E s f Z H n E 1 S a O K B C 4 d z 7 u U e T h h z p o 3 n f T l L y y u r a + u 5 D X d z a 3 t n N 7 + 3 X 9 c y U Y T W i O R S N U O s K W c R r R l m O G 3 G i m I R c t o I h 9 e Z 3 7 i j S j M Z V c 0 o p o H A / Y j 1 G M H G S q 2 2 w G Z A M E + r 4 0 6 + 4 B W 9 C d A i 8 W e k c P n h X s T P n 2 6 l k / 9 u d y V J B I 0 M 4 V j r l u / F J k i x M o x w O n b b i a Y x J k P c p y 1 L I y y o D t J J 5 D E 6 t k o X 9 a S y E x k 0 U X 9 f p F h o P R K h 3 c w i 6 n k v E / / 1 w l D M v T a 9 8 y B l U Z w Y G p H p 5 1 7 C k Z E o K w R 1 m a L E 8 J E l m C h m w y M y w A o T Y 2 t z b S v + f Ae L p F 4 q + q f F 0 q 1 X K F / B F D k 4 h C M 4 A R / O o A w 3 U I E a E J D w A I / w 5 N w 7 L 8 6 r 8 z Z d X X J m N w f w B 8 7 7 D 1 f 9 m 7 k = < / l a t e x i t >

Figure4: 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 is a 2-hidden-layer 1D CNN with nonlinearity φ, performing regression on the 10-dimensional outputs z 2 for each of the 4 (1, 2, 3, 4) inputs x from the dataset X . To declutter notation, unit weight and zero bias variances are assumed in all layers. Top: recursive output (z 2 ) computation in the CNN (top) induces a respective recursive NNGP kernel ( K2 ⊗ I 10 ) computation (NTK computation being similar, not shown). Bottom: explicit listing of tensor and corresponding kernel ops in each layer. See Table1for operation definitions. Illustration and description adapted from Figure3inNovak et al. (2019).

Figure5: 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 a 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 a certain threshold when all cores in a single GPU are saturated. Simpler models are expected to have better scaling with batch size. Right: Time-per-sample scales linearly with the number of GPUs, demonstrating near-perfect hardware utilization.

Figure 7: Predictive negative log-likelihoods and condition numbers. Top. Test negative loglikelihoods 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 are selected based on train marginal negative log-likelihoods in Figure3. Bottom. Condition numbers for covariance matrices corresponding to NTK/NNGP as well as respective predictive covaraince on 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.



Table content: <table><tr><td>3the best-performing SG point corresponds to the tion) CNN with the same on the full training CIFAR nels. While CNN-GP ap best CNNs overall outpe further analysis of factors and infinite Bayesian CN 100% training accuracy o set downsampled to 8 ⇥ details of (a, c) and (b) pl</td></tr><tr><td>8 Spatially local average pooling in intermediary layers can be constructed in a similar fashion ( §A.3). We focus on global average pooling in this work to more effectively isolate the effects of pooling from other aspects of the model like local connectivity or equivariance.</td></tr><tr><td>7</td></tr></table>



Table content: <table><tr><td>• CNNs suffer from exploding or va</td></tr><tr><td>gradients at large depths unless t</td></tr><tr><td>initialized correctly. The phase bo</td></tr><tr><td>order and chaos (rightmost plot f</td></tr><tr><td>nonlinearity) is the region where l</td></tr><tr><td>training is feasible.</td></tr></table>

Leveraging block-diagonal covariance structure. A common computational challenge with GPs is inverting the training set covariance matrix. Naively, for a classification task with C classes and training set X , NNGP and NTK covariances have the shape of |X | C × |X | C. For CIFAR-10, this would be 500, 000 × 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 a block-diagonal covariance matrix of the form Σ ⊗ I C , where Σ has shape |X | × |X | and I C is the C × C identity matrix. This reduces the computational complexity and storage in many common cases by an order of magnitude, which makes closed-form exact inference feasible in these cases.

Table content: <table><tr><td>)</td></tr><tr><td>3.2 PERFORMANCE</td></tr><tr><td>Our library performs a number of automatic performance optimizations without sacrificing flexibility.</td></tr></table>

).nngp and kernel_fn(x1, x2).ntk respectively, which

Table content: <table><tr><td/><td colspan="2">2.5 3.0</td><td>Independent</td><td>FC-NTK FC-NNGP CONV-NTK CONV-NNGP</td><td/><td>2.5 3.0</td><td>Full Test Covariance</td></tr><tr><td>Test NLL (nats)</td><td colspan="2">1.5 2.0</td><td/><td/><td colspan="2">Test NLL (nats)</td><td>1.5 2.0</td></tr><tr><td/><td colspan="4">Training Set Size 2 6 2 7 2 8 2 9 2 10 2 11 2 12 2 13 2 14 2 15 1.0</td><td/><td>2 6 2 7 2 8 2 9 2 10 2 11 2 12 2 13 2 14 2 15 Training Set Size 1.0</td></tr><tr><td/><td>9</td><td/><td>NTK</td><td/><td/><td>10 9</td><td>NNGP</td></tr><tr><td/><td>8</td><td/><td/><td/><td/><td>10 8</td></tr><tr><td>Condition Number</td><td>2 3 4 5 6 7</td><td colspan="3">2 6 2 7 2 8 2 9 2 10 2 11 2 12 2 13 Training Set Size</td><td>Condition Number</td><td>10 2 10 3 10 4 10 5 10 6 10 7</td><td>Training Set Size 2 6 2 7 2 8 2 9 2 10 2 11 2 12 2 13 FC FC-Test-Cov Conv Conv-Test-Cov WResNet WResNet-Test-Cov</td></tr></table>

Peer review (iclr2020):
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. While 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. Edit: 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. Significance: 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. However, 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. To 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. Originality: 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. Again, 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. Clarity: 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. Quality: I find the experiments performed to be very well constructed. Below are a few mostly minor comments on the experiments: In Figure 1 on the right, I would have liked to have seen the posterior predictive for a NNGP with the same kernel as well. In 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. In 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. In 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.)? second 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. References: Arora, S., et al., Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks, https://arxiv.org/abs/1910.01663 Lee, J., et al., Wide Neural Networks of any Depth Evolve as Linear Models Under Gradient Descent, NeurIPS, 2019, https://arxiv.org/abs/1902.06720