Year: 2020

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 www.github.com/google/neural-tangents We also provide an accompanying interactive Colab notebook1.

www.github.com/google/neural-tangents
We also provide an accompanying interactive Colab notebook1.
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
∗Equal contribution. †Work done during an internship at Google Brain. 1colab.research.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb
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).
2See §A for NEURAL TANGENTS comparison against specific relevant prior works.
We begin with three short examples (§2) that demonstrate the ease, efficiency, and versatility of performing calculations with infinite networks using NEURAL TANGENTS. With a high level view of the library in hand, we then dive into a number of technical aspects of our library (§3).
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 of nonlinearities. Let zli(x) describe the i th pre-activation following a linear transformation in lth layer of a neural network. At initialization, the parameters of 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, x′) = E[zli(x)zli(x′)]. This describes a NNGP with the kernel, K(x, x′). One can use the NNGP to make Bayesian posterior predictions at a test point, x, which are Gaussian distributed with with mean µ(x) = K(x,X )K(X ,X )−1Y and variance σ2(x) = K(x, x) − K(x,X )K(X ,X )−1K(X , x), where (X ,Y) is the training set of inputs and targets respectively. NTK. When neural networks are optimized using continuous gradient descent with learning rate η on mean squared error (MSE) loss, the function evaluated on training points evolves as ∂tft(X ) = −ηJt(X )Jt(X )T (ft(X )− Y) where Jt(X ) is the Jacobian of the output ft evaluated at X and Θt(X ,X ) = Jt(X )Jt(X )T is the NTK. In the infinite-width limit, the NTK remains constant (Θt = Θ) throughout training and the time-evolution of the outputs can be solved in closed form as a Gaussian with mean ft(x) = Θ(x,X )Θ(X ,X )−1 (I − exp [−ηΘ(X ,X )t])Y .
2 EXAMPLES
We begin by applying NEURAL TANGENTS to several example tasks. While these tasks are designed for pedagogy rather than research novelty, they are nonetheless emblematic of problems regularly faced in research. We emphasize that without NEURAL TANGENTS, it would be necessary to derive the kernels for each architecture by hand.
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-finite networks. This example is worked through in detail in the Colab notebook.3
We train on a synthetic dataset with training data drawn from the process yi = sin(xi) + i with xi ∼ Uniform(−π, π) and i ∼ N (0, σ2) independently and identically distributed. To train an infinite neural network with Erf activations4 on this data using gradient descent and an MSE loss we write the following:
from neural_tangents import predict, stax
init_fn, apply_fn, kernel_fn = stax.serial(
stax.Dense(2048, W_std=1.5, b_std=0.05), stax.Erf(),
stax.Dense(2048, W_std=1.5, b_std=0.05), stax.Erf(),
stax.Dense(1, W_std=1.5, b_std=0.05))
y_mean, y_var = predict.gp_inference(kernel_fn, x_train, y_train, x_test, 'ntk',
diag_reg=1e-4, compute_cov=True)
The above code analytically generates the predictions that would result from performing gradient descent for an 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:
3www.colab.research.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb 4Error function, a nonlinearity similar to tanh; see §D for other implemented nonlinearities, including Relu.
predict_fn = predict.gradient_descent_mse_gp(kernel_fn, x_train, y_train, x_test, 'ntk',
diag_reg=1e-4, compute_cov=True)
y_mean, y_var = predict_fn(t=100) # Predict the distribution at t = 100.
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 TANGENTSṄ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,
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.
def WideResNetBlock(channels, strides=(1, 1), channel_mismatch=False):
Main = stax.serial(stax.Relu(), stax.Conv(channels, (3, 3), strides, padding='SAME'),
stax.Relu(), stax.Conv(channels, (3, 3), padding='SAME'))
Shortcut = (stax.Identity() if not channel_mismatch else
stax.Conv(channels, (3, 3), strides, padding='SAME'))
return stax.serial(stax.FanOut(2), stax.parallel(Main, Shortcut), stax.FanInSum())
def WideResNetGroup(n, channels, strides=(1, 1)):
blocks = [WideResNetBlock(channels, strides, channel_mismatch=True)]
for _ in range(n - 1):
blocks += [WideResNetBlock(channels, (1, 1))]
return stax.serial(*blocks)
def WideResNet(block_size, k, num_classes):
return stax.serial(stax.Conv(16, (3, 3), padding='SAME'),
WideResNetGroup(block_size, int(16 * k)),
WideResNetGroup(block_size, int(32 * k), (2, 2)),
WideResNetGroup(block_size, int(64 * k), (2, 2)),
stax.GlobalAvgPool(), stax.Dense(num_classes))
init_fn, apply_fn, kernel_fn = WideResNet(block_size=4, k=1, num_classes=10)
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.
0 10
−4
5
−5
1
lo g1
0( D
ist an
ce )
NNGP
log2(#Samples)
lo g2
(w id
en in
g fa
ct or
)
0 10
−4
5
−5
1
lo g1
0( D
ist an
ce )
NTK
log2(#Samples)
lo g2
(w id
en in
g fa
ct or
)
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 − Kk,n‖2F / ‖K‖ 2 F , where k is the widening factor and n is the number of samples.
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 , θ) ∈ Rd×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
Kz = Eθ [ ziz T i ] , Θz = Eθ [ ∂zi ∂θ ( ∂zi ∂θ )T] (1)
where zi ∈ Rd is the ith 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 CNNGP 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 h = 1d1d. Then
Kpool1 ⌘ 2! d2
X
↵,↵0
⇥ KL+11 ⇤ ↵,↵0 + 2b . (17)
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 |2d2 memory to compute the sample-sample covariance of the GP (or O d2
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↵,
Ke↵1 ⌘ 2! ⇥ KL+11 ⇤ ↵,↵ + 2b . (18)
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
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.
7
Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes Roman Novak^, Lechao Xiao*^, Jeahoon Lee*†, Yasaman Bahri*†, Greg Yang°, Daniel A. Abolafia, Jeffrey Pennington, Jascha Sohl-Dickstein Google Brain. *Google AI residents. ^,†Equal contribution. °Microsoft Research AI. 1. Summary
3. Different NN-GP architectures
5. SGD-trained finite vs Bayesian infinite CNNs Under review as a conference paper at ICLR 2019 (a) (c) (b) No Pooling Global Average Pooling L C N C N N #Channels
Figure 3: (a): SGD-trained CNNs often perform better with increasing number of channels. Each line corresponds to a particular choice of architecture and initialization hyperparameters, with best learning rate and weight decay selected independently for each number of channels (x-axis). (b): SGD-trained CNNs often approach the performance of their corresponding CNN-GP with increasing number of channels. All models have the same architecture except for pooling and weight sharing, as well as training-related hyperparameters such as learning rate, weight decay and batch size, which are selected for each number of channels (x-axis) to maximize validation performance (y-axis) of a neural network. As the number of channels grows, best validation accuracy increases and approaches accuracy of the respective GP (solid horizontal line). (c): However, the best-performing SGD-trained CNNs can outperform their corresponding CNN-GPs. Each point corresponds to the test accuracy of: (y-axis) a specific CNN-GP; (x-axis) the best (on validation) CNN with the same architectural hyper-parameters selected among the 100%-accurate models on the full training CIFAR10 dataset with different learning rates, weight decay and number of channels. While CNN-GP appears competitive against 100%-accurate CNNs (above the diagonal), the
best CNNs overall outperform CNN-GPs by a significant margin (below the diagonal, right). For
further analysis of factors leading to similar or diverging behavior between SGD-trained finite CNNs
and infinite Bayesian CNNs see Tables 1 and 2. Experimental details: all networks have reached
100% training accuracy on CIFAR10. Values in (b) are reported on an 0.5K/4K train/validation sub-
set downsampled to 8 ⇥ 8 for computational reasons. See §A.7.5 and §A.7.1 for full experimental
details of (a, c) and (b) plots respectively.
10
Under review as a conference paper at ICLR 2019
Depth: 1 10 100 1000 Phase boundary
C N
N -G P FC N -G P
Table 3: Validation accuracy of CNN- and FCN- GPs as a function of weight ( 2! , horizontal axis) and bias ( 2b , vertical axis) variances. As predicted in §A.2, the regions of good performance concentrate around the critical line (phase boundary, right) as the depth increases (left to right). All plots share common axes ranges and employ the erf nonlinearity. See §A.7.2 for experimental details.
A.2 RELATIONSHIP TO DEEP SIGNAL PROPAGATION
The recurrence relation linking the GP kernel at layer l + 1 to that of layer l following from Equation 10 (i.e. Kl+11 = (C A) Kl1 ) is precisely the covariance map examined in a series of related papers on signal propagation (Xiao et al., 2018; Poole et al., 2016; Schoenholz et al., 2017; Lee et al., 2018) (modulo notational differences; denoted as F , C or e.g. A ? C in Xiao et al. (2018)). In those works, the action of this map on hidden-state covariance matrices was interpreted as defining a dynamical system whose large-depth behavior informs aspects of trainability. In particular, as l ! 1, Kl+11 = (C A) Kl1 ⇡ Kl1 ⌘ K⇤1, i.e. the covariance approaches a fixed point K⇤1. The convergence to a fixed point is problematic for learning because the hidden states no longer contain information that can distinguish different pairs of inputs. It is similarly problematic for GPs, as the kernel becomes pathological as it approaches a fixed point. Precisely, in the chaotic regime outputs of the GP become asymptotically decorrelated and therefore independent, while in the ordered regime they approach perfect correlation of 1. Either of these scenarios captures no information about the training data in the kernel and makes learning infeasible.
This problem can be ameliorated by judicious hyperparameter selection, which can reduce the rate of exponential convergence to the fixed point. For hyperpameters chosen on a critical line separating two untrainable phases, the convergence rates slow to polynomial, and very deep networks can be trained, and inference with deep NN-GP kernels can be performed – see Table 3.
A.3 STRIDED CONVOLUTIONS AND AVERAGE POOLING IN INTERMEDIATE LAYERS
Our analysis in the main text can easily be extended to cover average pooling and strided convolutions (applied before the pointwise nonlinearity). Recall that conditioned on Kl the pre-activation zlj (x) 2 Rd1 is a zero-mean multivariate Gaussian. Let B 2 Rd2⇥d1 denote a linear operator. Then Bzlj (x) 2 Rd2 is a zero-mean Gaussian, and the covariance is
E{!l,bl} h Bzlj (x) Bzlj (x 0) T Kl i = BE{!l,bl} h zlj (x) z l j (x 0) T Kl i BT . (20)
One can easily see that Bzlj Kl j are i.i.d. multivariate Gaussian as well. Strided convolution. Strided convolution is equivalent to a non-strided convolution composed with subsampling. Let s 2 N denote size of the stride. Then the strided convolution is equivalent to choosing B as follows: Bij = (is j) for i 2 {0, 1, . . . (d2 1)}. Average pooling. Average pooling with stride s and window size ws is equivalent to choosing Bij = 1/ws for i = 0, 1, . . . (d2 1) and j = is, . . . , (is + ws 1). ND convolutions. Note that our analysis in the main text (1D) easily extends to higher-dimensional convolutions by replacing integer pixel indices and sizes d,↵, with tuples (see also Figure 4).
18
4. Computing the NN-GP covariance
8. Large-depth behavior • CNNs suffer from exploding or vanishing
gradients at large depths unless they are initialized correctly. The phase boundary between order and chaos (rightmost plot for ERF nonlinearity) is the region where large-depth training is feasible.
• CNN-GP exhibits similar behavior. Initialized away from the phase boundary, the covariance between all different inputs becomes asymptotically constant and makes learning infeasible.
• 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 finite CNNs often perform better with increasing number of channels. (b) Performance of SGD-trained finite CNNs often approaches that of the respective CNN-GP with increa ing the number of channels. (c) However, the best performing SGD CNN outperforms the best CNN-GP by a 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 CNNGPs 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. Disentangling the CNN architecture
52
63
74
85
FCN FCN-GP LCN LCN w/pooling CNN-GP CNN CNN w/ pooling
Compositionality Local connectivity Equivariance InvarianceQuality
Model
Ac cu
ra cy
Ac cu
ra cy
53
64
75
86
FCN CNN w/ ERF CNN w/ small learning rate CNN w/ ReLU & large learning rate CNN w/ pooling
CNN-GP
FCN-GP
NN (underfitting allowed)
NN (no underfitting)
n1 → ∞
n1 → ∞
n2 → ∞
n2 → ∞
4x4 covariance per class y0 z0 z1 y2 = ϕ (z1) z̄2 4 3 2 1 10 classes per input 4 inputs, 3 features 4 3 2 1 4 3 2 1 ⊗ I10 K0 K0 K1 K2 K1 K2σ2ωK0 + σ2b σ2ωK1 + σ2b $ (K1)$ (K0) affine transform affine transform affine transform matrix-multiplymatrix-multiplymatrix-multiply nonlinearity nonlinearity % % 4x4 input covariance
4x4x10 input covariance
4x4 covariance per class
⊗ I10
y0
z0 z1 z̄2
10 classes per input
4 inputs, 3 features, 10 pixels
4 3 2 1
4 3
2 1
4 3 2 1
K0 K1 K2$ (K1)$ (K0) affine transform affine transform affine transform
convolution convolution matrix-multiply
nonlinearitynonlinearity
y2 = ϕ (z1)y1 = ϕ (z0)
%
n1 →
∞
n1 →
∞ n2 →
∞
n2 →
∞
C (K) = Ex⇠N (0,K) h (x) (x) T i <latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit><latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit><latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit><latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit> FCN (fully-connected): CNN (1D):
6. When do CNNs outperfor CNN-GPs?
• Outputs of a deep convolutional neural network (CNN) converge to a Gaussian Process ("CNN-GP") as the number of channels in hidden layers go to infinity. • CNN-GP achieves SOTA results on CIFAR10 among GPs without trainable kernels. • Absent pooling CNNs with and without weight sharing converge to the same GP. Translation equivariance is lost in infinite Bayesian CNNs (CNN-GPs), and finite SGD-trained CNNs can outperform them. 2. Convergence
(Monte Carlo estimate)
(C N N w /o w ei gh t s ha rin g)
(error function non-linearity)
(CNN w/o weight sharing)(Fully-connected network)
(Fully-connected)
(before the last readout layer)
Under review as a conference paper at ICLR 2019
M C
-C N
N -G
P
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 KL+1n,M when the number f ou er products contributing to K L+1 n,M approximately equals its rank. Values reported re for a 3-l yer 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.
random finite width networks and use the empirical uncentered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,
⇥
Kln,M
⇤
↵,↵0
(x, x0) ⌘ 1
Mn
MX
m=1
nX
c=1
ylc↵ (x; ✓m) y l c↵0 (x 0; ✓m) (19)
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, limn!1 Kln,M = K l 1 in probability. For finite width networks, the uncertainty in Kln,M is Var ⇥ Kln,M ⇤ = Var✓ ⇥ Kln (✓) ⇤ /M . From Daniely et al. (2016), we know that Var✓ ⇥ Kln (✓) ⇤ / 1n , which leads to Var✓[Kln,M ] / 1Mn . For finite n, Kln,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,
Kln,M (✓) KL1 2
F is nearly constant for constant Mn. We thus tr at Mn as the ffective 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 ti and bias.
In a non-distributed setting, the MC-GP reduces the memory requirements to compute GPpool from O ⇣ |X |2 d2 ⌘ to O ⇣ |X |2 + n2 + 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:⇥ KL1 ⇤LCN ↵,↵0
(x, x0) = 0, for ↵ 6= ↵0 and all x, x0 2 X and L > 0. However, LCNs preserve the covariances between input examples at every pixel: ⇥ KL1 ⇤LCN ↵,↵ (x, x0) = ⇥ KL1 ⇤CNN ↵,↵
(x, x0). As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical. Moreover, LCN-GPs with
8
Under review as a conf rence paper at ICLR 2019
M
C
-C
N
N
-G
P
Figure 2: Valid tion curacy (left) of an MC-CNN-GP increases with n ⇥ M (i.e. channel count times number of samples) and pproaches that of the exact CNN-GP (not shown), while the distance (right) t the exact ker el decreases. The d rk band in the left plot corre p nds t ill-c ditioning f KL+1n,M whe the number of outer roducts contributi g to K L+1 n,M approximately equals its rank. Values repo ted are for a 3-layer model applied to a 2K/4K tra n/validation subset of CIFAR10 downsampled to 8 ⇥ 8. See Figure 7 for similar results with other architectures and §A.7.2 for experiment l details.
random finite width networks and use the empiri al uncentered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,
⇥
Kln,M
⇤
↵,↵0
(x, x0) ⌘ 1
Mn
MX
m=1
nX
c=1
ylc↵ (x; ✓m) y l c↵0 (x 0; ✓m) (19)
where ✓ consists of M draws of the weights nd biases from their prior distribution, ✓m ⇠ p (✓), and n is the width or number of channels in hidd n layers. The MC-GP ker el converges to the analytic kernel with increas ng w dth, limn!1 Kln,M = K l 1 in probability. For finite width networks, th uncertainty in Kln,M is Var ⇥ Kln,M ⇤ = Var✓ ⇥ Kln (✓) ⇤ /M . From Daniely et al. (2016), we know that Var✓ ⇥ Kln (✓) ⇤ / 1n , which leads to Var✓[Kln,M ] / 1Mn . For finite n, Kln,M is also a biased estimate of K l 1, w re the bias epends s lely on net ork width. We do not currently have n analytic m for this ias, but we can see in Figures 2 and 7 that for th hyperparameters we probe it is sm ll r lative to the varia ce. In particular,
Kln,M (✓) KL1 2
F is nearly constant for consta t Mn. We thus treat Mn as the eff ct ve sample siz for the Monte Ca lo k rnel estimate. Increasing M and reducing n can reduce memory c st, though potentially at the expense of increased compute time and bi s.
I a non-distribu ed se ting, the MC-GP reduces the memory requirements to compute GPpool from O ⇣ |X |2 d2 ⌘ to O ⇣ |X |2 2 + nd ⌘ , making the evaluation of CNN-GPs with pooling practical.
5 DISCUSSION
5.1 BAYESIA CNNS WITH MANY CHANNELS ARE IDENTICA TO OCALLY CONNECTED NETWORKS, IN THE ABSENCE F POOLING
Locally Connected Networks (LCNs) (Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spati l locations. LCNs preserve the onnectivi y pattern, and thus topology, of a CNN. However, they do not poss ss the equivariance property of a CNN – if an nput i translated, the latent represe tatio in an LCN will be completely different, rather than also bei g translated.
The CNN-GP predictions without s atial pooling in §3.1 and §3.2.2 epend only on sample-sample covariances, and do not epend on -pixel covariances. LCNs destroy -pixel covariances:⇥ KL1 ⇤LCN ↵,↵0
(x, x0) = 0, for 6= ↵0 and all x, x0 2 X and L > 0. However, LCNs preserve the covariances between input exampl s at every pixel: ⇥ KL1 ⇤LCN ↵,↵ (x, x0) = ⇥ KL1 ⇤CNN ↵,↵
(x, x0). As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical. Moreover, LCN-GPs with
8
(#NN instantiations)
• Outputs of a CNN converge in distribution to a Gaussian process with a compositional kernel (under mild conditions on the weights and biases distribution and nonlinearity). (#channels in hidden layers) (convergence in distribution) (an affine transformation) (inputs covariance) (computes covariance of post-nonlinearity activations given normal pre-activations) outputs inputs D ! min{n1,...,nL}!1 N ⇣ 0, A (C A)L K0 ⌦ InL+1 ⌘ <latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit><latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit><latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit><latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit> (#channels in the outputs)
(F ul
lyco
nn ec
te d)
(all displayed models are 100% accurate on the training set)
(not always necessary to track the whole 4x4x10x10 covariance.)
(attaching a fully-connected layer only to the center pixel of the output)
n1 → ∞
n1 → ∞
n2 → ∞
n2 → ∞
CN N
(1 D)
ac
tiv at io ns
CN N-
GP
co va
ria nc
e
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 CNNGP 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 tr in/validation subsets of CIFAR10. See §A.7.3 for experimental d tails.
1. Global average pooling: take h = 1d1d. Then
Kpool1 ⌘ 2! d2
X
↵,↵0
⇥ KL+11 ⇤ ↵,↵0 + 2b . (17)
This approach corresponds to applying global average pooling right after the last convoluti nal layer.8 This appro ch takes all pixel-pixel covariances into consideration d makes the kernel translation invariant. However, it requires O |X |2d2 memory to compute the sample-sample covariance of the GP (or O d2
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↵,
Ke↵1 ⌘ 2! ⇥ KL+11 ⇤ ↵,↵ + 2b . (18)
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 Mont Ca lo estimation method for NN-GP kernels which are computationally i - practical to compute an lytically, or for which we d not know the analy ic form. Simi ar in spirit to traditional random feature methods (Rahimi & Recht, 2007), the core idea is to instantiate many
8 Spatially local av rag pooling in int rmediary lay rs 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.
7
Bayesian Deep Convolutional Networks with Many Cha els are Gaussian Processes Roman ovak^, Lechao Xiao*^, J ahoon Lee*†, Yasaman Bahri*†, Greg Yang°, Daniel A. Abolafia, Jeffrey Pennington, Jascha Sohl-Dickstein Google Brain. *Google AI residents. ^,†Equal contribution. °Microsoft Research AI. 1. Summary
3 Different NN-GP architectures
5. SGD-trained finite vs Bayesian infinite CNNs Under review as a conference paper at ICLR 2019 (a) (c) (b) No Pooling Global Average Pooling L C N C N N
#Chan els
Figure 3: (a): SGD-trained CNNs often perform better with increasing number of channels. Each line corresponds to a particular choice of architecture and initialization hyperparameters, with best learning rate and weight decay selected independently for each number of channels (x-axis). (b): SGD-trained CNNs often approach the performance of their corresponding CNN-GP with increasing number of channels. All models have the same architecture except for pooling and weight sharing, as well as training-related hyperparameters such as learning rate, weight decay and batch size, which are selected for each number of channels (x-axis) to m ximize validation performance (y-axis) of a neural network. As the number of channels grows, best validation accuracy increases and approaches accuracy of the respective GP (solid horizontal line). (c): However, the best-performing SGD-trained CNNs can outperform their corresponding CNN-GPs. Each point corresponds to the test accuracy of: (y-axis) a specific CNN-GP; (x-axis) the best (on validation) CNN with the same architectural hyper-parameters selected among the 100%-accurate models on the full training CIFAR10 dataset with different learning rates, weight decay and number of channels. While CNN-GP appears competitive against 100%-accurate CNNs (above the diagonal), the best CNNs overall outperform CNN-GPs by a significant margin (below the diagonal, right). For further analysis of factors leading to similar or diverging behavior between SGD-trained finite CNNs and infinite Bayesian CNNs see Tables 1 and 2. Experimental details: all networks have reached 100% training accuracy on CIFAR10. Values in (b) are reported on an 0.5K/4K train/validation subset downsampled to 8 ⇥ 8 for computational reasons. See §A.7.5 and §A.7.1 for full experimental details of (a, c) a d (b) plots respectively.
10
Under review as a conference paper at ICLR 2019
Depth: 1 10 100 1000 Phase boundary
C N
N -G P FC N -G P
able 3: Valid tion ccuracy of CNN- and FCN- GPs as a fun of weight ( 2! , horizontal axis) and bias ( 2b , vertical axis) variances. As predicted in §A.2, the regions of good perform nce conce te around the criti al line (phase boundary, right) as the depth incr ases (left to right). All plots share common axes ranges and employ the erf nonlin arity. See §A.7 2 for experimental details.
A.2 RELATIONSHIP TO DEEP SIGNAL PROPAGATION
The recurrence relation linking the GP kernel at layer l + 1 to that of layer l following from Equation 10 (i.e. Kl+11 = (C A) Kl1 ) is precisely the covariance map examined in a series of related papers on signal propagation (Xiao et al., 2018; Poole et al., 2016; Schoenholz et al., 2017; Lee et al., 2018) (modulo notational differences; denoted as F , C or e.g. A ? C in Xiao et al. (2018)). In those works, the action of this map on hidden-state covariance matrices was interpreted as defining a dynamical system whose large-depth behavior informs aspects of trainability. In particular, as l ! 1, Kl+11 = (C A) Kl1 ⇡ Kl1 ⌘ K⇤1, i.e. the covariance approaches a fixed point K⇤1. The convergence to a fixed point is problematic for learning because the hidden states no longer contain information that can distinguish different pairs of inputs. It is similarly problematic for GPs, as the kernel becomes pathological as it approaches a fixed point. Precisely, in the chaotic regime outputs of the GP become asymptotically decorrelated and therefore independent, while in the ordered regime they approach perfect correlation of 1. Either of these scenarios captures no information about the training data in the kernel and makes learning infeasible.
This problem can be ameliorated by judicious hyperparameter selection, which can reduce the rate of exponential convergence to the fixed point. For hyperpameters chosen on a critical line separating two untrainable phases, the convergence rates slow to polynomial, and very deep networks can be trained, and inference with deep NN-GP kernels can be performed – see Table 3.
A.3 STRIDED CONVOLUTIONS AND AVERAGE POOLING IN INTERMEDIATE LAYERS
Our analysis in the main text can easily be extended to cover average pooling and strided convolutions (applied before the pointwise nonlinearity). Recall that conditioned on Kl the pre-activation zlj (x) 2 Rd1 is a zero-mean multivariate Gaussian. Let B 2 Rd2⇥d1 denote a linear operator. Then Bzlj (x) 2 Rd2 is a zero-mean Gaussian, and the covariance is
E{!l,bl} h Bzlj (x) Bzlj (x 0) T Kl i = BE{!l,bl} h zlj (x) z l j (x 0) T Kl i BT . (20)
One can easily see that Bzlj Kl j are i.i.d. multivariate Gaussian as well. Strided convolution. Strided convolution is equivalent to a non-strided convolution composed with subsampling. Let s 2 N denote size of the stride. Then the strided convolution is equivalent to choosing B as follows: Bij = (is j) for i 2 {0, 1, . . . (d2 1)}. Average pooling. Average pooling with stride s and window size ws is equivalent to choosing Bij = 1/ws for i = 0, 1, . . . (d2 1) and j = is, . . . , (is + ws 1). ND convolutions. Note that our analysis in the main text (1D) easily extends to higher-dimensional convolutions by replacing integer pixel indices and sizes d,↵, with tuples (see also Figure 4).
18
4. Computing the NN-GP covariance
8. Large-depth behavior • CNNs suffer from exploding or vanishing
gradients at large depths unless they are initialized correctly. The phase boundary between order and chaos (rightmost plot for ERF nonlinearity) is the region where large-depth training is feasible.
• CNN-GP exhibits similar behavior. Initialized away from the phase boundary, the covariance between all different inputs becomes asymptotically constant and makes learning infeasible.
• 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 experim nts 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 ccuracy.
(a) SGD-trained finite CNNs often perform better with increasing number of channels. (b) Performance of SGD-trained finite CNNs often approache that of the respective CNN-GP with in rea ing the number of chann ls. (c) However, the best performing SGD CNN outperforms the best CNN-GP by a significant margin.
Various CNN architectural deci ions like pooling / vectorizing / subsampling, zero or no padding have a corresponding NN-GP.
• Pooling enforces translation-invariant predictions in both CNNs and CNNGPs 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. Disentangling the CNN architecture
52
63
74
85
FCN FCN-GP LCN LCN w/pooling CNN-GP CNN CNN w/ pooling
Compositionality Local connectivity Equivariance InvarianceQuality
Model
Ac cu
ra cy
Ac cu
ra cy
53
64
75
86
FCN CNN w/ ERF CNN w/ small learning rate CNN w/ ReLU & large learning rate CNN w/ pooling
CNN-GP
FCN-GP
NN (underfitting allowed)
NN (no underfitting)
n1 → ∞
n1 → ∞
n2 → ∞
n2 → ∞
4x4 covariance per class
y0 z0 z1 y2 = ϕ (z1)
z̄2 4 3 2 1 10 classes per input 4 inputs, 3 features
4 3
2 1
4 3 2 1
⊗ I10
K0
K0 K1 K2
K1 K2σ2ωK0 + σ2b σ2ωK1 + σ2b
$ (K1)$ (K0)
affine transform affine transform affine transform
matrix-multiplymatrix-multiplymatrix-multiply nonlinearity nonlinearity
%
%
4x4 input covariance
4x4x10 input covariance
4x4 covariance per class
⊗ I10
y0
z0 z1 z̄2
10 classes per input
4 inputs, 3 features, 10 pixels
4 3
2 1
4 3
2 1
4 3 2 1
K0 K1 K2$ (K1)$ (K0) affine transform affine transform affine transform
convolution convolution matrix-multiply
nonlinearitynonlinearity
y2 = ϕ (z1)y1 = ϕ (z0)
%
n1 →
∞
n1 →
∞ n2 →
∞
n2 →
∞
C (K) = Ex⇠N (0,K) h (x) (x) T i
<latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit><latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit><latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit><latexit sha1_base64="LKSS7iNm7WJo+g4I+BjNBlPbawo=">AAACeHicbVFdS8MwFE3r9/yaH2++BKeoINLuRV8EUQRBEAWnwlpHkqVbMGlLciuOMv+nz/orfDJbq+jmhcDhnHO5N+fSVAoDnvfmuBOTU9Mzs3OV+YXFpeXqyuqdSTLNeIMlMtEPlBguRcwbIEDyh1Rzoqjk9/TpbKDfP3NtRBLfQi/loSKdWESCEbBUq/q6FSgCXUZkftYPJI9g9zLQotOFPXyMhxql+Xm/lb8ERqgf81Vp9va/7QXRDNKuKKSXUhhnHm8LEG61qjXvwBsWHgd+CWqorOtW9T1oJyxTPAYmiTFN30shzIkGwSTvV4LM8JSwJ9LhTQtjorgJ82FOfbxtmTaOEm1fDHjI/u7IiTKmp6h1Dv5pRrUB+a9GqRoZDdFRmIs4zYDHrJgcZRJDggdXwG2hOQPZs4AwLezymHWJJgzsrSo2FX80g3FwVz/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit>
FCN (fully-connected):
CNN (1D):
6. When do CNNs outperfor CNN-GPs?
• Outputs of a deep convolutional neural network (CNN) converge to a Gaussian Process ("CNN-GP") as the number of channels in hidden layers go to infinity. • CNN-GP achieves SOTA results on CIFAR10 among GPs without trainable kernels. • Absent pooling CNNs with and without weight sharing converge to the same GP. Translation equivariance is lost in infinite Bayesian CNNs (CNN-GPs), and finite SGD-trained CNNs can outperform the . 2. Convergence
(Monte Carlo estimate)
(C N N w /o w ei gh t s ha rin g)
(error function non-linearity)
(CNN w/o weight sharing)(Fully-connected network)
(Fully-connected)
(before the last readout layer)
Under review as a conference paper at ICLR 2019
M C
-C
N
N -G
P
Figure 2: Validation accuracy (left) of an MC-CNN-GP increases with n ⇥ M (i.e. channel count times number of samples) a d appr aches that of the exact CNN-GP (no shown), while the distance (right) to the exact kernel decreases. The dark band in the left plot corresponds to ill-conditioning of KL+1n,M when the number f ou er products contributing to K L+1 n,M approximately equals its rank. Values reported re for a 3-l yer 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.
random finite width n tworks a d us the mpirical unc ntered c variances of activati ns to stimate the Monte Carlo-GP (MC-GP) kernel,
⇥ Kln,M ⇤ ↵,↵0 (x, x0) ⌘ 1 Mn MX
m=1
nX
c=1
ylc↵ (x; ✓m) y l c↵0 (x 0; ✓m) (19)
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, limn!1 Kln,M = K l 1 in probability. For finite width networks, the uncertainty in Kln,M is Var ⇥ Kln,M ⇤ = Var✓ ⇥ Kln (✓) ⇤ /M . From Daniely et al. (2016), we know that Var✓ ⇥ Kln (✓) ⇤ / 1 , which leads to Var✓[Kl ,M ] / 1Mn . For finite n, Kln,M is also a biased estimate of K l 1, where h bias d p nds solely on etwork wi th. 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,
Kln,M (✓) KL1 2
F is nearly constant for constant Mn. We thus tr at Mn as the ffective sample size for the Monte Carlo kernel estimate. Increasing M and reducing n can reduce memory cost, though potentially at the expense of incr ased com e ti and bias.
In a non-distributed setting, the MC-GP reduces the memory requirements to compute GPpool from O ⇣ |X |2 d2 ⌘ to O ⇣ |X |2 + n2 + 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:⇥ KL1 ⇤LCN ↵,↵0
(x, x0) = 0, for ↵ 6= ↵0 and all x, x0 2 X and L > 0. However, LCNs preserve the covariances between input examples at every pixel: ⇥ KL1 ⇤LCN ↵,↵ (x, x0) = ⇥ KL1 ⇤CNN ↵,↵
(x, x0). As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical. Moreover, LCN-GPs with
Under r view as a conf renc paper at ICLR 2019
M C
-C
N
N -G
P
Figure 2: Valid tion curacy (left) of an MC-CNN-GP increases with n ⇥ M (i.e. channel count times number of samples) a d pp oaches that of the exact CNN-GP (no shown), while the distance (right) t the exact ker el decreases. The d rk band in the left plot corre p nds t ill-c ditioning of KL+1n,M whe the number of outer roducts contributi g to K L+1 n,M approximately equals its rank. Values repo ted are for a 3-layer model applied to a 2K/4K tra n/validation subset of CIFAR10 downsampled to 8 ⇥ 8. See Figure 7 for similar results with other architectures and §A.7.2 for experiment l details.
ran om finite width n tworks and us the mpiri al unc ntered covariances of activations to estimate the Monte Carlo-GP (MC-GP) kernel,
⇥ Kln,M ⇤ ↵,↵0 (x, x0) ⌘ 1 Mn MX
m=1
nX
c=1
ylc↵ (x; ✓m) y l c↵0 (x 0; ✓m) (19)
where ✓ consists of M draws of the weights nd biases from their prior distribution, ✓m ⇠ p (✓), and n is the width or number of channels in hidd n layers. The MC-GP ker el converges to the analytic kernel with increas ng w dth, limn!1 Kln,M = K l 1 in probability. For finite width networks, th uncertainty in Kln,M is Var ⇥ Kln,M ⇤ = Var✓ ⇥ Kln (✓) ⇤ /M . From Daniely et al. (2016), we know that Var✓ ⇥ Kln (✓) ⇤ / 1n , which leads to Var✓[Kl ,M ] / 1Mn . For fi ite n, Kln,M is also a biased estimate of K l 1, w th bias ep nd s lely on et ork wi th. We do not currently have n analytic m for this ias, but we can see in Figures 2 and 7 that for th hyperparameters we probe it is sm ll r lative to the varia ce. In particular,
Kln,M (✓) KL1 2
F is nearly constant for consta t Mn. We thus treat Mn as the eff ct ve sample siz for the Monte Ca lo k rnel estimate. Increasing M and reducing n can reduce memory c st, though potentially at the xpense of incr as d compute tim and bi s.
I a non-distribu ed se ting, the MC-GP reduces the memory requirements to compute GPpool from O ⇣ |X |2 d2 ⌘ to O ⇣ |X |2 2 + nd ⌘ , making the evaluation of CNN-GPs with pooling practical.
5 DISCUSSION
5.1 BAYESIA CNNS WITH MANY CHANNELS ARE IDENTICA TO OCALLY CONNECTED NETWORKS, IN THE ABSENCE F POOLING
Locally Connected Networks (LCNs) (Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spati l locations. LCNs preserve the onnectivi y pattern, and thus topology, of a CNN. However, they do not poss ss the equivariance property of a CNN – if an nput i translated, the latent represe tatio in an LCN will be completely different, rather than also bei g translated.
The CNN-GP predictions without s atial pooling in §3.1 and §3.2.2 epend only on sample-sample covariances, and do not epend on -pixel covariances. LCNs destroy -pixel covariances:⇥ KL1 ⇤LCN ↵,↵0
(x, x0) = 0, for 6= ↵0 and all x, x0 2 X and L > 0. However, LCNs preserve the covariances between input exampl s at every pixel: ⇥ KL1 ⇤LCN ↵,↵ (x, x0) = ⇥ KL1 ⇤CNN ↵,↵
(x, x0). As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical. Moreover, LCN-GPs with
8
(#NN instantiations)
• Outputs of a CNN converge in distribution to a Gaussian process with a compositional kernel (under mild conditions on the weights and biases distribution and nonlinearity). (#channels in hidden layers) (convergence in distribution) (an affine transformation) (inputs covariance) (computes covariance of post-nonlinearity activations given normal pre-activations) outputs inputs D ! min{n1,...,nL}!1 N ⇣ 0, A (C A)L K0 ⌦ InL+1 ⌘ <latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit><latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit><latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit><latexit sha1_base64="rSefQ0IBgTGoBSFhXrbE0xA5Jbk=">AAAC4nicbVJNb9QwEHXCV1k+usCRi8UWqYhVlfQCx0I5gKiqIrFtpXV25XidXauOHdkT6CqYH8ANceWPcebCz8D52KpsGSnyy5s3euMZp4UUFqLoVxBeu37j5q2N2707d+/d3+w/eHhsdWkYHzEttTlNqeVSKD4CAZKfFobTPJX8JD3br/Mnn7ixQquPsCx4ktO5EplgFDw17f/ZIpJnsE2An4PJK11CUYLFJBXzL1io+scRI+YLeEbOm5Maoz+PSS5UU0oqNYmHZKbBDtXkoNUSR0AToTJYJhXJKSwYldUb5y7woWuNoyG+4F45woRhXUcrdr9jL6nafrxZo3w/iVYdahA5t/jd1DdVHTyP3Uq7Ne0Pop2oCXwVxB0YoC6Opv3f/k6szLkCJqm14zgqIKmoAcEkdz1SWl5QdkbnfOyhot44qZqNOPzUMzOcaeM/BbhhL1dUNLd2madeWV/Lrudq8r+5NM3XrCF7mVTNorhirXNWSgwa1/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit> (#channels in the outputs)
(F ul
lyco
nn ec
te d)
(all displayed models are 100% accurate on the training set)
(not always necessary to track the whole 4x4x10x10 covariance.)
(attaching a fully-connected layer only to the center pixel of the output)
n1 → ∞
n1 → ∞
n2 → ∞
n2 → ∞
<latexit sha1_base64="djKoC8a3j3Zgk4rwu2KpBm00p38=">AAAB+3icbVC7SgNBFL0bX3F9RS1tBoNgFWa10EYM2gg2Ec0DkjXMTmaTITO7y8ysEJZ8gq1+QDqx9Rv8BLH2R5wkFpp44MLhnHs5lxMkgmuD8aeTW1hcWl7Jr7pr6xubW4XtnZqOU0VZlcYiVo2AaCZ4xKqGG8EaiWJEBoLVg/7l2K8/MKV5HN2ZQcJ8SboRDzklxkq31/e4XSjiEp4AzRPvhxTP392zZPThVtqFr1YnpqlkkaGCaN30cGL8jCjDqWBDt5VqlhDaJ13WtDQikmk/m7w6RAdW6aAwVnYigybq74uMSK0HMrCbkpienvXG4r9eEMiZaBOe+hmPktSwiE6Tw1QgE6NxEajDFaNGDCwhVHH7PKI9ogg1ti7XtuLNdjBPakcl77iEb3CxfAFT5GEP9uEQPDiBMlxBBapAoQuP8ATPztAZOS/O63Q15/zc7MIfOG/fT3WXvg==</latexit>
A <latexit sha1_base64="O66ATy1aw1tYLYUpvpJ8FgNJCH8=">AAACA3icbVC7SgNBFL0bX3F9RS1tBoNgFXa10EaM2lhGMA9IljA7mSRDZmeWmVkhLCn9BFstLO3E1tpPEGt/xNkkhSYeuHA4517u4YQxZ9p43peTW1hcWl7Jr7pr6xubW4XtnZqWiSK0SiSXqhFiTTkTtGqY4bQRK4qjkNN6OLjK/PodVZpJcWuGMQ0i3BOsywg2Vmq2Imz6BPP0YtQuFL2SNwaaJ/6UFM8/3LP4+dOttAvfrY4kSUSFIRxr3fS92AQpVoYRTkduK9E0xmSAe7RpqcAR1UE6jjxCB1bpoK5UdoRBY/X3RYojrYdRaDeziHrWy8R/vTCMZl6b7mmQMhEnhgoy+dxNODISZYWgDlOUGD60BBPFbHhE+lhhYmxtrm3Fn+1gntSOSv5xybvxiuVLmCAPe7APh+DDCZThGipQBQISHuARnpx758V5dd4mqzlnerMLf+C8/wA5N5uk</latexit>
C <latexit sha1_base64="Bp3s0NXazI8dQfK7i5edA+N1hFs=">AAACA3icbVC7SgNBFL0bX3F9RS1tBoNgFXa10EYMprGMYB6QLGF2MkmGzMwuM7NCWFL6CbZaWNqJrbWfINb+iLNJCk08cOFwzr3cwwljzrTxvC8nt7S8srqWX3c3Nre2dwq7e3UdJYrQGol4pJoh1pQzSWuGGU6bsaJYhJw2wmEl8xt3VGkWyVszimkgcF+yHiPYWKnVFtgMCOZpZdwpFL2SNwFaJP6MFC8/3Iv4+dOtdgrf7W5EEkGlIRxr3fK92AQpVoYRTsduO9E0xmSI+7RlqcSC6iCdRB6jI6t0US9SdqRBE/X3RYqF1iMR2s0sop73MvFfLwzF3GvTOw9SJuPEUEmmn3sJRyZCWSGoyxQlho8swUQxGx6RAVaYGFuba1vx5ztYJPWTkn9a8m68YvkKpsjDARzCMfhwBmW4hirUgEAED/AIT8698+K8Om/T1Zwzu9mHP3DefwA8Y5um</latexit>
K11<latexit sha1_base64="XvvOkSjqlJ92PocDUdmzMoTY5m8=">AAACBHicbVC7SgNBFL0bXzG+opY2Q0RIFXa1iGXQRrCJYB6QXcPsZDYZMju7zMwKy5LWT7Cx0A+wE1t7P0G080ecPApNPHDhcM69nMvxY86Utu0PK7e0vLK6ll8vbGxube8Ud/eaKkokoQ0S8Ui2fawoZ4I2NNOctmNJcehz2vKH52O/dUulYpG41mlMvRD3BQsYwdpI7uWN081cJgKdjrrFQ7tiT4AWiTMjh7VS+euz+v5Q7xa/3V5EkpAKTThWquPYsfYyLDUjnI4KbqJojMkQ92nHUIFDqrxs8vMIHRmlh4JImhEaTdTfFxkOlUpD32yGWA/UvDcW//V8P5yL1sGplzERJ5oKMk0OEo50hMaNoB6TlGieGoKJZOZ5RAZYYqJNbwXTijPfwSJpHleck4p9Zeo5gynycAAlKIMDVajBBdShAQRiuIdHeLLurGfrxXqdruas2c0+/IH19gOqA5yR</latexit> 2 1<latexit sha1_base64="OUeUKe+lIobds6D6xEX6ikdejzA=">AAACBHicbVC7SgNBFJ2NrxhfUUubIUFIFXZjEcugjWATwTwgG8PsZDYZMju7zNwVliWtn2BjoR9gJ7b2foJo5484eRSaeODC4Zx7OZfjRYJrsO0PK7Oyura+kd3MbW3v7O7l9w+aOowVZQ0ailC1PaKZ4JI1gINg7UgxEniCtbzR+cRv3TKleSivIYlYNyADyX1OCRjJvbyp9FKXSx+ScS9ftMv2FHiZOHNSrBVKX5/V94d6L//t9kMaB0wCFUTrjmNH0E2JAk4FG+fcWLOI0BEZsI6hkgRMd9Ppz2N8bJQ+9kNlRgKeqr8vUhJonQSe2QwIDPWiNxH/9TwvWIgG/7SbchnFwCSdJfuxwBDiSSO4zxWjIBJDCFXcPI/pkChCwfSWM604ix0sk2al7JyU7StTzxmaIYuOUAGVkIOqqIYuUB01EEURukeP6Mm6s56tF+t1tpqx5jeH6A+stx+roZyS</latexit> Kvec1<latexit sha1_base64="6K3mCNDghmeg3opAxpsEWlsnFsk=">AAACHXicbVC7SgNBFJ2NrxhfUUubRRGswq4WsQzaCDYRTBTcGGYnd3VwZnaZuRsMy36A32DhJ9jqB1gIYiuinT/iJLHQxAMDh3PO5d45YSK4Qc97dwoTk1PTM8XZ0tz8wuJSeXmlaeJUM2iwWMT6NKQGBFfQQI4CThMNVIYCTsKr/b5/0gVteKyOsZdAS9ILxSPOKFqpXd4IJMVLRkV2mLezgKsIe/l5FiBco5ZZF1ie25RX8QZwx4n/QzZq61ufH9Xn23q7/BV0YpZKUMgENebM9xJsZVQjZwLyUpAaSCi7ohdwZqmiEkwrG3wmdzet0nGjWNun0B2ovycyKo3pydAm+6ebUa8v/uuFoRxZjdFuK+MqSREUG26OUuFi7ParcjtcA0PRs4Qyze3xLrukmjK0hZZsK/5oB+OkuV3xdyreka1njwxRJGtknWwRn1RJjRyQOmkQRm7IPXkgj86d8+S8OK/DaMH5mVklf+C8fQNy+KfT</latexit>A 0 <latexit sha1_base64="h4C+kvbRlWTUqodPHM9c9Ea/pzA=">AAACFXicbVC7TgJBFJ3FF+Jr1ZJmIjGBhuxqoSVqY2KDiTwSFsnsMAsTZh+ZuWuCGwq/wFK/QFv9ADtja23tjzgLFAqe5CYn59ybe+9xI8EVWNaXkVlYXFpeya7m1tY3NrfM7Z26CmNJWY2GIpRNlygmeMBqwEGwZiQZ8V3BGu7gLPUbN0wqHgZXMIxY2ye9gHucEtBSx8w7PoE+JSI5GTmCeVC8uLYcyXt9KHXMglW2xsDzxJ6SQqV0fxtlnx6rHfPb6YY09lkAVBClWrYVQTshEjgVbJRzYsUiQgekx1qaBsRnqp2Mnxjhfa10sRdKXQHgsfp7IiG+UkPf1Z3pyWrWS8V/Pdf1Z1aDd9xOeBDFwAI62ezFAkOI04hwl0tGQQw1IVRyfTymfSIJBR1kTqdiz2YwT+oHZfuwbF3qeE7RBFmUR3uoiGx0hCroHFVRDVF0h57RC3o1How34934mLRmjOnMLvoD4/MHS8CiBw==</latexit>
A <latexit sha1_base64="O66ATy1aw1tYLYUpvpJ8FgNJCH8=">AAACA3icbVC7SgNBFL0bX3F9RS1tBoNgFXa10EaM2lhGMA9IljA7mSRDZmeWmVkhLCn9BFstLO3E1tpPEGt/xNkkhSYeuHA4517u4YQxZ9p43peTW1hcWl7Jr7pr6xubW4XtnZqWiSK0SiSXqhFiTTkTtGqY4bQRK4qjkNN6OLjK/PodVZpJcWuGMQ0i3BOsywg2Vmq2Imz6BPP0YtQuFL2SNwaaJ/6UFM8/3LP4+dOttAvfrY4kSUSFIRxr3fS92AQpVoYRTkduK9E0xmSAe7RpqcAR1UE6jjxCB1bpoK5UdoRBY/X3RYojrYdRaDeziHrWy8R/vTCMZl6b7mmQMhEnhgoy+dxNODISZYWgDlOUGD60BBPFbHhE+lhhYmxtrm3Fn+1gntSOSv5xybvxiuVLmCAPe7APh+DDCZThGipQBQISHuARnpx758V5dd4mqzlnerMLf+C8/wA5N5uk</latexit>
C <latexit sha1_base64="Bp3s0NXazI8dQfK7i5edA+N1hFs=">AAACA3icbVC7SgNBFL0bX3F9RS1tBoNgFXa10EYMprGMYB6QLGF2MkmGzMwuM7NCWFL6CbZaWNqJrbWfINb+iLNJCk08cOFwzr3cwwljzrTxvC8nt7S8srqWX3c3Nre2dwq7e3UdJYrQGol4pJoh1pQzSWuGGU6bsaJYhJw2wmEl8xt3VGkWyVszimkgcF+yHiPYWKnVFtgMCOZpZdwpFL2SNwFaJP6MFC8/3Iv4+dOtdgrf7W5EEkGlIRxr3fK92AQpVoYRTsduO9E0xmSI+7RlqcSC6iCdRB6jI6t0US9SdqRBE/X3RYqF1iMR2s0sop73MvFfLwzF3GvTOw9SJuPEUEmmn3sJRyZCWSGoyxQlho8swUQxGx6RAVaYGFuba1vx5ztYJPWTkn9a8m68YvkKpsjDARzCMfhwBmW4hirUgEAED/AIT8698+K8Om/T1Zwzu9mHP3DefwA8Y5um</latexit>
<latexit sha1_base64="S97hAc1Xq6o+hnybyN3EFc8GOUA=">AAAB/HicbVC7SgNBFL3rM8ZX1FKRwSBYhV0ttAzaWCZgHpAsYXYymwyZmV1mZoWwxM7WVj/ATmwlv2Jt6U84m6TQxAMXDufcy7mcIOZMG9f9dJaWV1bX1nMb+c2t7Z3dwt5+XUeJIrRGIh6pZoA15UzSmmGG02asKBYBp41gcJP5jXuqNIvknRnG1Be4J1nICDaZ1I77rFMouiV3ArRIvBkplo/G1e/H43GlU/hqdyOSCCoN4VjrlufGxk+xMoxwOsq3E01jTAa4R1uWSiyo9tPJryN0apUuCiNlRxo0UX9fpFhoPRSB3RTY9PW8l4n/ekEg5qJNeOWnTMaJoZJMk8OEIxOhrAnUZYoSw4eWYKKYfR6RPlaYGNtX3rbizXewSOrnJe+i5FZtPdcwRQ4O4QTOwINLKMMtVKAGBPrwBM/w4jw4r86b8z5dXXJmNwfwB87HD0B7mPY=</latexit>
<latexit sha1_base64="S97hAc1Xq6o+hnybyN3EFc8GOUA=">AAAB/HicbVC7SgNBFL3rM8ZX1FKRwSBYhV0ttAzaWCZgHpAsYXYymwyZmV1mZoWwxM7WVj/ATmwlv2Jt6U84m6TQxAMXDufcy7mcIOZMG9f9dJaWV1bX1nMb+c2t7Z3dwt5+XUeJIrRGIh6pZoA15UzSmmGG02asKBYBp41gcJP5jXuqNIvknRnG1Be4J1nICDaZ1I77rFMouiV3ArRIvBkplo/G1e/H43GlU/hqdyOSCCoN4VjrlufGxk+xMoxwOsq3E01jTAa4R1uWSiyo9tPJryN0apUuCiNlRxo0UX9fpFhoPRSB3RTY9PW8l4n/ekEg5qJNeOWnTMaJoZJMk8OEIxOhrAnUZYoSw4eWYKKYfR6RPlaYGNtX3rbizXewSOrnJe+i5FZtPdcwRQ4O4QTOwINLKMMtVKAGBPrwBM/w4jw4r86b8z5dXXJmNwfwB87HD0B7mPY=</latexit>
X = y0 <latexit sha1_base64="w7fOKhzB63i59cdFS6XECob+pDw=">AAACCXicbVC7SgNBFL3rM8ZHVi0VGQyCVdjVQhshaGOZgHlAsobZyWwyZGZ3mZkVliWllZ9gqx9gF2z9CmtLf8LJo9DEAxcO59zLPRw/5kxpx/m0lpZXVtfWcxv5za3tnYK9u1dXUSIJrZGIR7LpY0U5C2lNM81pM5YUC5/Thj+4GfuNByoVi8I7ncbUE7gXsoARrI3UsQttgXWfYJ41h1fpvdOxi07JmQAtEndGiuXDUfX78WhU6dhf7W5EEkFDTThWquU6sfYyLDUjnA7z7UTRGJMB7tGWoSEWVHnZJPgQnRili4JImgk1mqi/LzIslEqFbzbHMdW8Nxb/9XxfzL3WwaWXsTBONA3J9HOQcKQjNK4FdZmkRPPUEEwkM+ER6WOJiTbl5U0r7nwHi6R+VnLPS07V1HMNU+TgAI7hFFy4gDLcQgVqQCCBZ3iBV+vJerNG1vt0dcma3ezDH1gfPxUHncI=</latexit>
z0 <latexit sha1_base64="U1l7Dtq33FgrRs54bTLLsjNbCyQ=">AAAB+3icbVC7SgNBFL0bX3F9RS1tBoNgFWa10EYM2lhGNA9I1jA7mU2GzOwuM7NCXPIJtvoB6cTWb/ATxNofcZJYaOKBC4dz7uVcTpAIrg3Gn05uYXFpeSW/6q6tb2xuFbZ3ajpOFWVVGotYNQKimeARqxpuBGskihEZCFYP+pdjv37PlOZxdGsGCfMl6UY85JQYK9083OF2oYhLeAI0T7wfUjx/d8+S0YdbaRe+Wp2YppJFhgqiddPDifEzogyngg3dVqpZQmifdFnT0ohIpv1s8uoQHVilg8JY2YkMmqi/LzIitR7IwG5KYnp61huL/3pBIGeiTXjqZzxKUsMiOk0OU4FMjMZFoA5XjBoxsIRQxe3ziPaIItTYulzbijfbwTypHZW84xK+xsXyBUyRhz3Yh0Pw4ATKcAUVqAKFLjzCEzw7Q2fkvDiv09Wc83OzC3/gvH0Dmi6X7Q==</latexit>
z1 <latexit sha1_base64="V7AStE6FLEtgSPN87apVgzmLNDU=">AAAB+3icbVC7SgNBFL0bX3F9RS1tBoNgFXa10EYM2lhGNA9I1jA7mU2GzMwuM7NCXPIJtvoB6cTWb/ATxNofcZJYaOKBC4dz7uVcTphwpo3nfTq5hcWl5ZX8qru2vrG5Vdjeqek4VYRWScxj1QixppxJWjXMcNpIFMUi5LQe9i/Hfv2eKs1ieWsGCQ0E7koWMYKNlW4e7vx2oeiVvAnQPPF/SPH83T1LRh9upV34anVikgoqDeFY66bvJSbIsDKMcDp0W6mmCSZ93KVNSyUWVAfZ5NUhOrBKB0WxsiMNmqi/LzIstB6I0G4KbHp61huL/3phKGaiTXQaZEwmqaGSTJOjlCMTo3ERqMMUJYYPLMFEMfs8Ij2sMDG2Lte24s92ME9qRyX/uORde8XyBUyRhz3Yh0Pw4QTKcAUVqAKBLjzCEzw7Q2fkvDiv09Wc83OzC3/gvH0Dm8OX7g==</latexit> y1 <latexit sha1_base64="1Y9jtIAcXc6WD6dG14lqrSZUMkY=">AAAB+3icbVC7SgNBFL0bXzG+opaKDAbBKuxqoWXQxjJB84BkDbOTSTJkZnaZmRWWJaWlrX6AndgK+RVrS3/CyaPQxAMXDufcy7mcIOJMG9f9dDJLyyura9n13Mbm1vZOfnevpsNYEVolIQ9VI8CaciZp1TDDaSNSFIuA03owuB779QeqNAvlnUki6gvck6zLCDZWuk3uvXa+4BbdCdAi8WakUDocVb4fj0bldv6r1QlJLKg0hGOtm54bGT/FyjDC6TDXijWNMBngHm1aKrGg2k8nrw7RiVU6qBsqO9Kgifr7IsVC60QEdlNg09fz3lj81wsCMRdtupd+ymQUGyrJNLkbc2RCNC4CdZiixPDEEkwUs88j0scKE2PrytlWvPkOFkntrOidF92KrecKpsjCARzDKXhwASW4gTJUgUAPnuAZXpyh8+q8Oe/T1Ywzu9mHP3A+fgAurphX</latexit>
y2 <latexit sha1_base64="M1NznfhXnMFosLH4HXe507vFGXw=">AAAB+3icbVC7SgNBFL3rM8ZX1FKRwSBYhd1YaBm0sUzQPCBZw+xkkgyZmV1mZoVlSWlpqx9gJ7ZCfsXa0p9w8ig08cCFwzn3ci4niDjTxnU/naXlldW19cxGdnNre2c3t7df02GsCK2SkIeqEWBNOZO0apjhtBEpikXAaT0YXI/9+gNVmoXyziQR9QXuSdZlBBsr3Sb3xXYu7xbcCdAi8WYkXzoaVb4fj0fldu6r1QlJLKg0hGOtm54bGT/FyjDC6TDbijWNMBngHm1aKrGg2k8nrw7RqVU6qBsqO9Kgifr7IsVC60QEdlNg09fz3lj81wsCMRdtupd+ymQUGyrJNLkbc2RCNC4CdZiixPDEEkwUs88j0scKE2PrytpWvPkOFkmtWPDOC27F1nMFU2TgEE7gDDy4gBLcQBmqQKAHT/AML87QeXXenPfp6pIzuzmAP3A+fgAwQ5hY</latexit>
z̄2 <latexit sha1_base64="4weGomNjpmJmWZL3O94Y1OsZs9o=">AAACAHicbVC7SgNBFL0bX3F9RS1tBoNgFXZjoY0YtLGMYB6QrGFmMpsMmZ1dZmaFGNL4Cbba2diJrT/gJ4i1P+LkUWjigQuHc+7lXA5JBNfG876czMLi0vJKdtVdW9/Y3Mpt71R1nCrKKjQWsaoTrJngklUMN4LVE8VwRASrkd7FyK/dMqV5LK9NP2FBhDuSh5xiY6V6k2CF7m6KrVzeK3hjoHniT0n+7MM9TZ4/3XIr991sxzSNmDRUYK0bvpeYYICV4VSwodtMNUsw7eEOa1gqccR0MBj/O0QHVmmjMFZ2pEFj9ffFAEda9yNiNyNsunrWG4n/eoREM9EmPAkGXCapYZJOksNUIBOjURuozRWjRvQtwVRx+zyiXawwNbYz17biz3YwT6rFgn9U8K68fOkcJsjCHuzDIfhwDCW4hDJUgIKAB3iEJ+feeXFenbfJasaZ3uzCHzjvPwxnmdI=</latexit>
A 1 1 <latexit sha1_base64="qKhejSfyk+54mP43KqY2NvyxdM8=">AAACHHicbVA9SwNBEN2LXzF+RS1tjgQhacKdFlpGbQSbCOYDcvHY2+wlS/b2jt05IR7p7WytLC1t9QfYia1g7R9xL0mhiQ8GHu/NMDPPizhTYFlfRmZhcWl5JbuaW1vf2NzKb+80VBhLQusk5KFseVhRzgStAwNOW5GkOPA4bXqDs9Rv3lCpWCiuYBjRToB7gvmMYNCSmy84AYY+wTw5GTmc+lC6uLZdhwkfho5kvT6U3XzRqlhjmPPEnpJitXx/G2WfHmpu/tvphiQOqADCsVJt24qgk2AJjHA6yjmxohEmA9yjbU0FDqjqJONfRua+VrqmH0pdAsyx+nsiwYFSw8DTnenlatZLxX89zwtmVoN/3EmYiGKggkw2+zE3ITTTpMwuk5QAH2qCiWT6eJP0scQEdJ45nYo9m8E8aRxU7MOKdanjOUUTZNEeKqASstERqqJzVEN1RNAdekYv6NV4NN6Md+Nj0poxpjO76A+Mzx834KUz</latexit> K̃0 <latexit sha1_base64="oBlV8IU3XcIIzHb4GLFJzufcj10=">AAACD3icbVDLSsNAFJ34rPEVFVdugkVwVZK60I1YdCO4qWAf0MQymUzaoTOTMDMRSshH+AniTj/AnbgTv0Bc+yNO2i609cCFwzn3cg8nSCiRynG+jLn5hcWl5dKKubq2vrFpbW03ZZwKhBsoprFoB1BiSjhuKKIobicCQxZQ3AoGF4XfusNCkpjfqGGCfQZ7nEQEQaWlrrXrKUJDnHkMqj6CNLvK81una5WdijOCPUvcCSmfvZunyeOnWe9a314Yo5RhrhCFUnZcJ1F+BoUiiOLc9FKJE4gGsIc7mnLIsPSzUfzcPtBKaEex0MOVPVJ/X2SQSTlkgd4sUspprxD/9YKATb1W0YmfEZ6kCnM0/hyl1FaxXZRjh0RgpOhQE4gE0eFt1IcCIqUrNHUr7nQHs6RZrbhHleq1U66dgzFKYA/sg0PggmNQA5egDhoAgQw8gCfwbNwbL8ar8TZenTMmNzvgD4yPH7hnoDk=</latexit>
<latexit sha1_base64="MR9rGQGBF6rkChHkVe95gk8IgI4=">AAACBXicbVC7SgNBFL3rM66vqKXNYhCswm4stBGDNoJNBPOAZA2zk9lkyMzsOjMrhCW1n2CrWNuJraWfINb+iLNJCk08cOFwzr3cwwliRpV23S9rbn5hcWk5t2Kvrq1vbOa3tmsqSiQmVRyxSDYCpAijglQ11Yw0YkkQDxipB/3zzK/fEaloJK71ICY+R11BQ4qRNpLf4kj3MGLp5fDGa+cLbtEdwZkl3oQUTj/sk/j5066089+tToQTToTGDCnV9NxY+ymSmmJGhnYrUSRGuI+6pGmoQJwoPx2FHjr7Ruk4YSTNCO2M1N8XKeJKDXhgNrOQatrLxH+9IOBTr3V47KdUxIkmAo8/hwlzdORklTgdKgnWbGAIwpKa8A7uIYmwNsXZphVvuoNZUisVvcNi6cotlM9gjBzswh4cgAdHUIYLqEAVMNzCAzzCk3VvvViv1tt4dc6a3OzAH1jvP39lnFM=</latexit>
K̃1 <latexit sha1_base64="Pi0I6+H4uEBVTV4rBnczV+vnOGM=">AAACD3icbVDLSsNAFJ34rPEVFVdugkVwVZK60I1YdCO4qWAf0MQymUzaoTOTMDMRSshH+AniTj/AnbgTv0Bc+yNO2i609cCFwzn3cg8nSCiRynG+jLn5hcWl5dKKubq2vrFpbW03ZZwKhBsoprFoB1BiSjhuKKIobicCQxZQ3AoGF4XfusNCkpjfqGGCfQZ7nEQEQaWlrrXrKUJDnHkMqj6CNLvK81u3a5WdijOCPUvcCSmfvZunyeOnWe9a314Yo5RhrhCFUnZcJ1F+BoUiiOLc9FKJE4gGsIc7mnLIsPSzUfzcPtBKaEex0MOVPVJ/X2SQSTlkgd4sUspprxD/9YKATb1W0YmfEZ6kCnM0/hyl1FaxXZRjh0RgpOhQE4gE0eFt1IcCIqUrNHUr7nQHs6RZrbhHleq1U66dgzFKYA/sg0PggmNQA5egDhoAgQw8gCfwbNwbL8ar8TZenTMmNzvgD4yPH7n8oDo=</latexit>
K̃2 <latexit sha1_base64="ngBeSUGob5edbBDod/0msER50A4=">AAACD3icbVDLSsNAFJ34rPEVFVdugkVwVZK60I1YdCO4qWAf0MQymUzaoTOTMDMRSshH+AniTj/AnbgTv0Bc+yNO2i609cCFwzn3cg8nSCiRynG+jLn5hcWl5dKKubq2vrFpbW03ZZwKhBsoprFoB1BiSjhuKKIobicCQxZQ3AoGF4XfusNCkpjfqGGCfQZ7nEQEQaWlrrXrKUJDnHkMqj6CNLvK89tq1yo7FWcEe5a4E1I+ezdPk8dPs961vr0wRinDXCEKpey4TqL8DApFEMW56aUSJxANYA93NOWQYelno/i5faCV0I5ioYcre6T+vsggk3LIAr1ZpJTTXiH+6wUBm3qtohM/IzxJFeZo/DlKqa1iuyjHDonASNGhJhAJosPbqA8FREpXaOpW3OkOZkmzWnGPKtVrp1w7B2OUwB7YB4fABcegBi5BHTQAAhl4AE/g2bg3XoxX4228OmdMbnbAHxgfP7uRoDs=</latexit>
<latexit sha1_base64="uk5eDtz7FwTCLBL89ubWC4LcRbU=">AAACBXicbVC7SgNBFL3rM66vqKXNYhCswm4stBGDNoJNBPOAZA2zk9lkyMzsOjMrhCW1n2CrWNuJraWfINb+iLNJCk08cOFwzr3cwwliRpV23S9rbn5hcWk5t2Kvrq1vbOa3tmsqSiQmVRyxSDYCpAijglQ11Yw0YkkQDxipB/3zzK/fEaloJK71ICY+R11BQ4qRNpLf4kj3MGLp5fCm1M4X3KI7gjNLvAkpnH7YJ/Hzp11p579bnQgnnAiNGVKq6bmx9lMkNcWMDO1WokiMcB91SdNQgThRfjoKPXT2jdJxwkiaEdoZqb8vUsSVGvDAbGYh1bSXif96QcCnXuvw2E+piBNNBB5/DhPm6MjJKnE6VBKs2cAQhCU14R3cQxJhbYqzTSvedAezpFYqeofF0pVbKJ/BGDnYhT04AA+OoAwXUIEqYLiFB3iEJ+veerFerbfx6pw1udmBP7DefwCA+pxU</latexit>
<latexit sha1_base64="rSjldedmulJSB31Muo592G+LDLg=">AAAB+3icbVC7SgNBFL0bX3F9RS1tBoNgFXZjoY0YtLGMaB6QrGF2MpsMmZldZmaFGPIJtvoB6cTWb/ATxNofcfIoNPHAhcM593IuJ0w408bzvpzM0vLK6lp23d3Y3Nreye3uVXWcKkIrJOaxqodYU84krRhmOK0nimIRcloLe1djv/ZAlWaxvDP9hAYCdySLGMHGSreP98VWLu8VvAnQIvFnJH/x4Z4no0+33Mp9N9sxSQWVhnCsdcP3EhMMsDKMcDp0m6mmCSY93KENSyUWVAeDyatDdGSVNopiZUcaNFF/Xwyw0LovQrspsOnqeW8s/uuFoZiLNtFZMGAySQ2VZJocpRyZGI2LQG2mKDG8bwkmitnnEelihYmxdbm2FX++g0VSLRb8k0LxxsuXLmGKLBzAIRyDD6dQgmsoQwUIdOAJnuHFGToj59V5m65mnNnNPvyB8/4DnfyX8Q==</latexit>
K0 = X T <latexit sha1_base64="80an4amqpQgZAcH3j/u4sWxX9PU=">AAACInicbVDLSgMxFM34rONr1KWbYBG6KjN1oZti0Y3gpkJf0BeZNNOGZjJDkhHKML/gT+gnuNWt4E5ciODOHzHTFqmtBwLnnnMv9+a4IaNS2fansbS8srq2ntkwN7e2d3atvf2aDCKBSRUHLBANF0nCKCdVRRUjjVAQ5LuM1N3hZerXb4mQNOAVNQpJ20d9Tj2KkdJS18q1fKQGGLH4OunYsAh/60YyQzuVrpW18/YYcJE4U5I9fzGL4f2HWe5a361egCOfcIUZkrLp2KFqx0goihlJzFYkSYjwEPVJU1OOfCLb8fhHCTzWSg96gdCPKzhWZydi5Es58l3dmR4p571U/NdzXX9utfLO2jHlYaQIx5PNXsSgCmCaF+xRQbBiI00QFlQfD/EACYSVTtXUqTjzGSySWiHvnOQLN3a2dAEmyIBDcARywAGnoASuQBlUAQZ34BE8gWfjwXg13oz3SeuSMZ05AH9gfP0AGq+nzA==</latexit>
y0 X <latexit sha1_base64="5RFGxLmik9P80wsdfTVluJuYoMw=">AAACC3icbVC7SgNBFJ2NrxhfGy0VGQyCVdiNhTZC0MYyAfOAJIbZyWwyZGZ2mZlVliWlpZ9gqx9gE8TWj7C29CecTVJo4oELh3Pu5R6OFzKqtON8Wpml5ZXVtex6bmNza3vHzu/WVRBJTGo4YIFsekgRRgWpaaoZaYaSIO4x0vCGV6nfuCNS0UDc6DgkHY76gvoUI22krp2Pbx14Adsc6QFGLGmOunbBKToTwEXizkihfDCufj8cjitd+6vdC3DEidCYIaVarhPqToKkppiRUa4dKRIiPER90jJUIE5UJ5lEH8Fjo/SgH0gzQsOJ+vsiQVypmHtmM42o5r1U/NfzPD73WvvnnYSKMNJE4OlnP2JQBzAtBvaoJFiz2BCEJTXhIR4gibA29eVMK+58B4ukXiq6p8VS1dRzCabIgn1wBE6AC85AGVyDCqgBDO7BE3gGL9aj9Wq9We/T1Yw1u9kDf2B9/ADGvJ4Y</latexit>
<latexit sha1_base64="aoqmHiiqO2ZBzCsL7LfI0ss8c6M=">AAACA3icbVC7SgNBFL3rM66vqKXNYBCswm4stBGDNpYR8oJkCbOTSTJkZmeZmRXCktJPsNXC0k5srf0EsfZHnE1SaOKBC4dz7uUeThhzpo3nfTlLyyura+u5DXdza3tnN7+3X9cyUYTWiORSNUOsKWcRrRlmOG3GimIRctoIh9eZ37ijSjMZVc0opoHA/Yj1GMHGSq22wGZAME+r406+4BW9CdAi8WekcPnhXsTPn26lk/9udyVJBI0M4Vjrlu/FJkixMoxwOnbbiaYxJkPcpy1LIyyoDtJJ5DE6tkoX9aSyExk0UX9fpFhoPRKh3cwi6nkvE//1wlDMvTa98yBlUZwYGpHp517CkZEoKwR1maLE8JElmChmwyMywAoTY2tzbSv+fAeLpF4q+qfF0q1XKF/BFDk4hCM4AR/OoAw3UIEaEJDwAI/w5Nw7L86r8zZdXXJmNwfwB877D1f9m7k=</latexit>
T <latexit sha1_base64="aoqmHiiqO2ZBzCsL7LfI0ss8c6M=">AAACA3icbVC7SgNBFL3rM66vqKXNYBCswm4stBGDNpYR8oJkCbOTSTJkZmeZmRXCktJPsNXC0k5srf0EsfZHnE1SaOKBC4dz7uUeThhzpo3nfTlLyyura+u5DXdza3tnN7+3X9cyUYTWiORSNUOsKWcRrRlmOG3GimIRctoIh9eZ37ijSjMZVc0opoHA/Yj1GMHGSq22wGZAME+r406+4BW9CdAi8WekcPnhXsTPn26lk/9udyVJBI0M4Vjrlu/FJkixMoxwOnbbiaYxJkPcpy1LIyyoDtJJ5DE6tkoX9aSyExk0UX9fpFhoPRKh3cwi6nkvE//1wlDMvTa98yBlUZwYGpHp517CkZEoKwR1maLE8JElmChmwyMywAoTY2tzbSv+fAeLpF4q+qfF0q1XKF/BFDk4hCM4AR/OoAw3UIEaEJDwAI/w5Nw7L86r8zZdXXJmNwfwB877D1f9m7k=</latexit>
Layer Tensor Op NNGP Op NTK Op
0 (input) y0 = X K0 = XX T Θ0 = 0 0 (pre-acti ations) z0 = Conv ( y0 ) K̃0 = A ( K0 ) Θ̃0 = K̃0 +A ( Θ0 ) 1 (activations) y1 = φ ( z0 ) K1 = T ( K̃0 ) Θ1 = Ṫ ( K̃0 ) Θ̃0
1 (pre-activations) z1 = Conv
(
y1
)
K̃1 = A
(
K1
)
Θ̃1 = K̃1 +A
(
Θ1
)
2 (activations) y2 = φ (
z1
)
K2 = T
(
K̃1
)
Θ2 = Ṫ
(
K̃1
)
Θ̃1
2 (readout) z2 = Dense ◦ Flatten
(
y2
)
K̃2 = Tr
(
K2
)
Θ̃2 = K̃2 + Tr
(
Θ2
)
Figure 4: A xample of the tra s atio of a convolutional neural ne work in o 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-hidd n-layer 1D CNN with nonline rity φ, performing regression on the 10-dimensional outputs z2 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 lay rs. Top: recursive output (z2) computation in the CNN (top) induces a respective recursive NNGP kernel (K̃2 ⊗ I10) computation (NTK computation being similar, not shown). Bottom: explicit listing of t sor and corres nding kernel ps in each l yer. See Table 1 for operation definitions. Illustratio and description ad ted from Figure 3 in Novak et al. (2019).
assume z is a mean zero multivaria e Gaussian. We wish to compute the kernel corresponding
to h = Dense (σω, σb) (φ(z)) by computing the kernel of = φ(z) and h = Dens (σω, σb) (y)
separately. Here,
h = Dens (σω, σb)(y) ≡ ( 1/ √ n ) σωWy + σbβ, (2)
and the variables Wij and βi are i.i.d. Gaussian N (0, 1). We will compute kernel operations - denoted φ∗ and Dense(σω, σb)∗ - induced by the tensor operations φ and Dense(σω, σb) 5. Finally, we will compute the kernel operation as ocia ed with the composition (D ns (σω, σb) ◦ φ)∗ = Dense(σω, σb)∗ ◦ φ∗. First we compute the NNGP and NT ker els f r y. T compu e Ky note h from its finit on,
Ky = Kφ(z) = Eθ [ φ(z)i φ(z) T i ] = Eθ[φ(zi)φ(zi)T ] = T (Kz). (3)
Since φ does not introduce any new variables Θy can be compute as,
Θy = Eθ
[ ∂φ(zi)
∂θ
( ∂φ(zi)
∂θ
)T] = Eθ [ diag(φ̇(zi)) ∂z
∂θ ∂zi ∂θ
)T diag(φ̇(zi)) ] = Ṫ (Kz) Θz.
Taken together these equations imply that,
(Ky, Θy) = φ∗ (Kz, Θz) ≡ ( T (Kz), Ṫ (Kz) Θz ) (4)
5T (Σ) ≡ E [ φ(u)φ(u)T ] , Ṫ (Σ) ≡ E [ φ′(u)φ′(u)T ] , u ∼ N (0,Σ), as in (Lee et al., 2019).
7
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 we consider the case of a dense operation. Using the independence between the weights, the biases, and h it follows that,
Kh = EW,β,θ[hihTi ] = σ2ωEθ[yiyTi ] + σ2b = σ2ωKy + σ2b . (5) Finally, the NTK of h can be computed as a sum of two terms:
Θh = EW,β,θ
[ ∂hi
∂(W,β)
( ∂hi
∂(W,β)
)T] + EW,β,θ [ ∂hi ∂θ ( ∂hi ∂θ )T] = σ2ωKy + σ2b + σ2ωΘy . (6)
This gives the translation rule for the dense layer in terms of Ky and Θy as, (Kh,Θh) = Dense(σω, σb)∗ (Ky,Θy) ≡ ( σ2ωKy + σ2b , σ2ωKy + σ2b + σ2ωΘy ) . (7)
3.2 PERFORMANCE
Our library performs a number of automatic performance optimizations without sacrificing flexibility.
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 Σ⊗ IC , where Σ has shape |X | × |X | and IC 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.
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 l, constructing the |X | × |X | output covariance matrix, Σ, involves computing l intermediate layer covariance matrices, Σl, of size |X | d× |X | d (see Listing 1 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); GarrigaAlonso et al. (2019) remarked, if no pooling is used in the network the output covariance Σ can be computed by only using the stack of d |X | × |X |-blocks of Σl, bringing the time and memory cost from O(|X |2 d2) down to O(|X |2 d) per layer (see Figure 4 and Listing 2 for models admitting this optimization). Finally, if the network has no convolutional layers, the cost further reduces to O(|X |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:
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.
batched_kernel_fn = nt.batch(kernel_fn, batch_size)
batched_kernel_fn(x, x) == kernel_fn(x, x) # True!
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!
ACKNOWLEDGMENTS
We thank Yasaman Bahri for frequent discussion and useful feedback on the manuscript. We additionally appreciate both Yasaman Bahri and Greg Yang for the ongoing contributions to improve the library. We thank Sergey Ioffe for feedback on the text, as well as Ravid Ziv, and Jeffrey Pennington for discussion and feedback on early versions of the library.
6These transpositions could not be automatically fused by the XLA compliler.
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, Yangqing 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, Fernanda 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://tensorflow.org/. Software available from tensorflow.org.
Takuya Akiba, Keisuke Fukuda, and Shuji Suzuki. ChainerMN: Scalable Distributed Deep Learning Framework. In Proceedings of Workshop on ML Systems in The Thirty-first Annual Conference on Neural Information Processing Systems (NIPS), 2017. URL http://learningsys.org/nips17/ assets/papers/paper_25.pdf.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In International Conference on Machine Learning, 2018.
Anonymous. Infinite attention: Nngp 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: The quantization-depth trade-off. arXiv preprint arXiv:1906.00771, 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne. JAX: composable transformations of Python+NumPy programs, 2018a. URL http://github.com/google/jax.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman-Milne. Stax, a flexible neural net specification library in jax, 2018b. URL https://github.com/google/jax/blob/master/jax/experimental/stax.py.
Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. arXiv preprint arXiv:1812.07956, 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://keras.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 S Du, Jason D Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. arXiv preprint arXiv:1811.03804, 2018a.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. arXiv preprint arXiv:1810.02054, 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 kernels. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp.
5724–5734. Curran Associates, Inc., 2019. URL http://papers.nips.cc/paper/ 8809-graph-neural-tangent-kernel-fusing-graph-neural-networks-with-graph-kernels. 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-Alonso, Carl Edward Rasmussen, and Laurence Aitchison. Deep convolutional networks as shallow gaussian processes. In International Conference on Learning Representations, 2019.
GPy. GPy: A gaussian process framework in python. http://github.com/SheffieldML/GPy, 2012.
Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the selection of initialization and activation function for deep neural networks. arXiv preprint arXiv:1805.08266, 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 in 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 on Learning Representations, 2018.
Jaehoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha SohlDickstein, 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://openreview.net/forum?id=HJx54i05tX.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. In Advances in Neural Information Processing Systems, pp. 8157–8166, 2018.
Alexander G. de G. Matthews, Mark van der Wilk, Tom Nickson, Keisuke. Fujii, Alexis Boukouvalas, Pablo Le‘on-Villagr‘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.org/papers/v18/16-537.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 arXiv:1804.11271, 2018a.
Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on 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 are gaussian processes. In International Conference on Learning Representations, 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, 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-Dickstein, 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 signal propagation in noisy rectifier neural networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 5717–5726. Curran Associates, Inc., 2018. URL http://papers.nips.cc/paper/ 7814-critical-initialisation-for-deep-signal-propagation-in-noisy-rectifier-neural-networks. pdf.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv preprint arXiv:1611.01232, 2016.
Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: a next-generation open source framework for deep learning. In Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Twenty-ninth Annual Conference on Neural Information Processing Systems (NIPS), 2015. URL http://learningsys.org/papers/LearningSys_2015_paper_33.pdf.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a 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 generalization in deep learning. arXiv preprint arXiv:1912.13053, 2019.
Ge Yang and Samuel Schoenholz. Mean field residual networks: On the edge of 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 arXiv:1902.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 overparameterized deep relu networks. Machine Learning, Oct 2019. ISSN 1573-0565. doi: 10.1007/s10994-019-05839-6. URL https://doi.org/10.1007/s10994-019-05839-6.
APPENDIX
A 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-Alonso et al., 2019; Arora et al., 2019a). While these efforts must have required implementing certain parts of our library, to our knowledge such prior efforts were either not open-sourced or not comprehensive / user-friendly / scalable enough to be used as a user-facing library. In addition, all of the works above used their own separate implementation, which further highlights a need for a more general approach.
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).nngp and kernel_fn(x1, x2).ntk respectively, which
can then be used to make posterior test set predictions as the mean of a conditional multivariate normal Ytest = K (Xtest,Xtrain)K (Xtrain,Xtrain)−1 Ytrain:
from jax.numpy.linalg import inv
y_test = kernel_fn(x_test, x_train).ntk @ inv(kernel_fn(x_train, 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 neural_tangents as nt
y_test = nt.predict.gp_inference(kernel_fn, x_train, y_train, x_test,
get='ntk', diag_reg=1e-4, compute_cov=False)
B.3 COMPUTING INFINITE NETWORK KERNELS IN BATCHES AND IN PARALLEL
Naively, the kernel_fn will compute the whole kernel in a 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 a batch decorator that takes a kernel_fn and returns a 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 as nt
kernel_fn = nt.batch(kernel_fn, batch_size=32)
Note that batching is often used to compute large covariance matrices that may not even fit on a GPU/TPU device, and require to be stored and used for inference using CPU RAM. This is easy to achieve by simply specifying nt.batch(..., store_on_device=False) . Once the matrix is
stored in RAM, inference will be performed with a CPU when nt.predict methods are called. As mentioned in §3.2, for many (notably, convolutional, and especially pooling) architectures, inference cost can be small relative to kernel construction, even when running on CPU (for example, it takes less than 3 minutes to execute jax.scipy.linalg.solve(..., sym_pos=True) on a 45, 000×45, 000 training covariance matrix and a 45, 000× 10 training target matrix).
B.4 TRAINING DYNAMICS OF INFINITE NETWORKS
In addition to closed form multivariate Gaussian posterior prediction, it is also interesting to consider network predictions following continuous gradient descent. To facilitate this we provide several functions to compute predictions following gradient descent with an MSE loss, for gradient descent with arbitrary loss, or for momentum with arbitrary loss. The first case is handled analytically, while the latter two 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_time .
import neural_tangents as nt
predictor = nt.predict.gradient_descent_mse(kernel_fn(x_train, x_train), y_train, kernel_fn(x_test, x_train))
fx_train, fx_test = predictor(training_time, fx_train, fx_test)
B.5 INFINITE NETWORKS OF ANY 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.
0 2 4 6 8 10 Training Epoch
10 2
10 1
Ba tc
h M
ea n
Sq ua
re d
Er ro
r
0th order (constant)
1st order (linear)
2nd order (quadratic)
Original function
While most of NEURAL TANGENTS is devoted to a function-space perspective—describing the distribution of function values on finite collections of training and testing points—we also provide tools to investigate a dual weight space perspective described in Lee et al. (2019). Convergence of dynamics to NTK dynamics coincide with networks being described by a 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
taylor_apply_fn = nt.taylor_expand(apply_fn, params, order)
fx_train_approx = taylor_apply_fn(new_params, x_train)
These act exactly like normal JAX functions and, in particular, can be plugged into gradient descent, which we demonstrate in Figure 6.
B.7 EXTENDING NEURAL TANGENTS
Many neural network layers admit a 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 on their inputs. To add such layer to NEURAL TANGENTS, one only has to implement it as a method in nt.stax with the following signature:
@_layer # an internal decorator taking care of certain boilerplate.
NewLayer(layer_params: Any) -> (init_fn: function, apply_fn: function, kernel_fn: function)
Here init_fn and apply_fn are 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 is no need to implement these methods and the user can simply return the respective methods from jax.experimental.stax (see nt.stax.Flatten for an example; in fact the majority of
nt.stax layers call the original jax.experimental.stax layers for finite width layer methods).
In this case what remains is to implement the kernel_fn method with signature
kernel_fn(input_kernel: nt.utils.Kernel) -> output_kernel: nt.utils.Kernel
Here both input_kernel and output_kernel are namedtuple s 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 user in the context of the particular operation that the finite width layer performs. This transformation could be as simple as an affine map on the kernel matrices, but could also be analytically intractable.
Once implemented, the correctness of the implementation can be very easily tested by extending the nt.tests.stax_test with the new layer, to test the agreement with large-widths empirical NNGP and NTK kernels.
C ARCHITECTURE SPECIFICATIONS
from neural_tangents import stax
def ConvolutionalNetwork(depth, W_std=1.0, b_std=0.0):
layers = []
for _ in range(depth):
layers += [stax.Conv(1, (3, 3), W_std, b_std, padding='SAME'), stax.Relu()]
layers += [stax.Flatten(), stax.Dense(1, W_std, b_std)]
return stax.serial(*layers)
Listing 2: All-convolutional model (ConvOnly) definition used in Figure 3.
from neural_tangents import stax
def FullyConnectedNetwork(depth, W_std=1.0, b_std=0.0):
layers = [stax.Flatten()]
for _ in range(depth):
layers += [stax.Dense(1, W_std, b_std), stax.Relu()]
layers += [stax.Dense(1, W_std, b_std)]
return stax.serial(*layers)
Listing 3: Fully-connected (FC) model definition used in Figure 3.
D IMPLEMENTED AND COMING SOON FUNCTIONALITY
The following layers7 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 padding8
• Relu
• LeakyRelu
• Abs
• ABRelu 9
• Erf
• Identity
• Flatten
• AvgPool
• GlobalAvgPool
• SumPool
• GlobalSumPool
• Dropout
• LayerNorm
• GlobalSelfAttention (Anonymous, 2020)
The following is in our near-term plans:
• Exp , Elu , Selu , Gelu
• Apache Beam support.
The following layers do not have known closed-form expressions for infinite network covariances, and respective infinite networks have to be estimated empirically (via nt.monte_carlo_kernel_fn ) or using other approximations (not currently implemented):
• Sigmoid , Tanh ,10 Swish ,11 Softmax , LogSoftMax , Softplus , MaxPool .
7 Abs , ABRelu , Erf , GlobalAvgPool , GlobalSumPool , LayerNorm , and GlobalSelfAttention
are only available in our library nt.stax and not in jax.experimental.stax . 8In addition to SAME and VALID , we support CIRCULAR padding, which is especially convenient for theoretical analysis and was used by Xiao et al. (2018) and Novak et al. (2019). 9 ABRelu is a more flexible extension of LeakyRelu , computed as amin (x, 0) + bmax (x, 0).
10 Sigmoid and Tanh are similar to Erf which does have a covariance expression and is implemented. 11 Swish is similar to Gelu .
h,h′ (x, x ′) = ∑ dh,dw [Σ] w+dw,w′+dw h+dh,h′+dh (x, x
′) /q2, 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 to other settings is trivial and supported by the library). [Σ] ∗ n = [Σ, . . . ,Σ] (n-fold replication). For LayerNorm, FanInConcat, and Attn (Anonymous, 2020) translation rules we refer the reader to our code at https://github. com/google/neural-tangents, as these ops are challenging to express concisely using current notation. See Figure 4 for an example of applying the translation rules to a specific model, and §3.1 for deriving a sample translation rule. See §D for the full list of currently implemented translations.

Review summaries:
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.