N EURAL T ANGENTS : F AST AND E ASY I NFINITE N EURAL N ETWORKS IN P YTHON
Anonymous authors
Paper under double-blind review
A BSTRACT
N EURAL T ANGENTS is a library for working with inﬁnite-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 ﬁnitewidth as usual, or in their inﬁnite-width limit.
For the inﬁnite-width networks,
N EURAL T ANGENTS performs exact inference either via Bayes’ rule or gradient descent, and generates the corresponding Neural Network Gaussian Process and
Neural Tangent kernels.
Additionally, N EURAL T ANGENTS provides tools to study gradient descent training dynamics of wide but ﬁnite networks.
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.
N EURAL T ANGENTS is available at w ww.github.com/neural-tangents/neural-tangents
We also provide an accompanying interactive Colab notebook .
I NTRODUCTION
Deep neural networks (DNNs) owe their success in part to the broad availability of high-level, ﬂexible, and efﬁcient software libraries like Tensorﬂow (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. I NFINITE - WIDTH B AYESIAN NEURAL NETWORKS
Recently, a new class of machine learning models has attracted signiﬁcant attention, namely, deep inﬁnitely wide neural networks.
In the inﬁnite-width limit, a large class of Bayesian neural networks become Gaussian Processes (GPs) with a speciﬁc, architecture-dependent, compositional kernel; these models are called Neural Network Gaussian Processes (NNGPs).
This correspondence was ﬁrst 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), and pooling (Novak et al., 2019).
The results for individual architectures have subsequently been generalized, and it was shown that a GP correspondence holds for a general class of networks that can be mapped to so-called tensor programs in (Yang, 2019).
The recurrence relationship deﬁning the NNGP kernel has additionally been studied separately from the GP correspondence in (Cho & Saul, 2009; Daniely et al., 2016; Poole et al., 2016;
Schoenholz et al., 2016; Yang & Schoenholz, 2017; Xiao et al., 2018).
1. I NFINITE - WIDTH NEURAL NETWORKS TRAINED BY GRADIENT DESCENT
In addition to enabling a closed form description of Bayesian neural networks, the inﬁnite-width limit has also very recently provided insights into neural networks trained by gradient descent.
In the last year, several papers have shown that randomly initialized neural networks trained with gradient descent are characterized by a distribution that is related to the NNGP, and is described by the so-called Neural Tangent Kernel (NTK)
(Jacot et al., 2018; Lee et al., 2019; Chizat et al., 2019), a kernel which was implicit in some earlier papers (Li & Liang, 2018; Allen-Zhu et al., 2018;
Du et al., 2018a;b).
In addition to this “function space” perspective, a dual, “weight space” view on the wide network limit was proposed in Lee et al.
(2019) which showed that networks under gradient descent were well-described by the ﬁrst-order Taylor series about their initial parameters.
1. P ROMISE AND PRACTICAL BARRIERS TO WORKING WITH INFINITE - WIDTH NETWORKS
Combined, these discoveries established inﬁnite-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 classiﬁcation 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 ﬁnite width networks in some situations, especially for fully- and locally-connected model families (Lee et al., 2018; Novak et al., 2019).
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.
Moreover, to make the code performant, it may be necessary to introduce custom, specialized, CUDA kernels (Arora et al., 2019a).
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. S UMMARY OF CONTRIBUTIONS
In this paper, we introduce a new open-source software library called N EURAL T ANGENTS targeting
JAX (Bradbury et al., 2018a) to facilitate research on inﬁnite limits of neural networks.
The main features of N EURAL T ANGENTS are:
• A high-level neural network API for specifying complex, hierarchical, models.
Networks speciﬁed using this API can have their inﬁnite-width NNGP kernel and NTK evaluated analytically (§2.1, Listings 1, 2, 3, §B.2).
• Functions to approximate inﬁnite-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 inﬁnite-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: gradient descent, with-or-without momentum, on arbitrary loss functions, at ﬁnite or inﬁnite 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 inﬁnite-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).
We begin with three short examples (§2) that demonstrate the ease, efﬁciency, and versatility of performing calculations with inﬁnite networks using N EURAL T ANGENTS .
With a high level view of the library in hand, we then dive into a number of technical aspects of our library (§3).
1. B ACKGROUND
We brieﬂy describe the NNGP (§1.1) and NTK (§1.2).
NNGP.
Neural networks are often structured as afﬁne transformations followed by pointwise applications of nonlinearities.
Let z l ( x ) describe the pre-activations following a linear transformation in l th 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 (cid:48) )
=
E [ z l ( x ) z l ( x (cid:48) )] .
One can therefore use the NNGP to make Bayesian posterior predictions which are Gaussian distributed with mean
K ( x, X ) K ( X , X ) − 1 Y and covariance K ( x, x ) − K
( x, X ) K ( X , X ) − 1 K
( X , x ) , where ( X , Y ) is the training set of inputs and targets.
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 ∂ t f t ( X ) = − ηJ t ( X ) J t ( X ) T ( f t ( X ) − Y ) where J t ( X ) is the Jacobian of f evaluated at
X and Θ t ( X , X ) =
J t ( X ) J t
( X ) T is the NTK.
In the inﬁnite-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 f t ( x ) =
Θ( x, X )Θ( X , X ) − 1 ( I − exp [ − η Θ
( X , X ) t ]) Y .
E XAMPLES
We begin by applying N EURAL T ANGENTS 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 N EURAL T ANGENTS , it would be necessary to derive the kernels for each architecture by hand.
2. I NFERENCE WITH AN I NFINITELY W IDE N EURAL N ETWORK
We begin by training an inﬁnitely wide neural network with gradient descent and comparing the result to training an ensemble of wide-but-ﬁnite networks.
This example is worked through in detail in the Colab notebook .
We train on a synthetic dataset with training data drawn from the process y i
= sin( x i ) + (cid:15)
i with x i ∼ Uniform ( − π, π ) and (cid:15) i ∼ N (0 , σ ) independently and identically distributed.
To train an inﬁnite neural network with Erf 4 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 inﬁnite amount of time.
However, it is often desirable to investigate ﬁnite-time learning dynamics of deep networks.
This is also supported in N EURAL T ANGENTS as illustrated in the following snippet:
[Equation]
[Table]
The above speciﬁcation set the hidden layer widths to 2048, which has no effect on the inﬁnite width network inference, but the init_fn and apply_fn here correspond to ordinary ﬁnite width networks.
In Figure 1 we compare the result of this exact inference with training an ensemble of one hundred of these ﬁnite-width networks by looking at the training curves and output predictions of both models.
We see excellent agreement between exact inference using the inﬁnite-width model and the result of training an ensemble using gradient descent.
2. A N I NFINITELY W IDE R ES N ET
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 inﬁnite-width limit.
For this task we study a variant of an inﬁnite-channel Wide Residual Network (Zagoruyko & Komodakis, 2016) (WRN-28- ∞ ).
We ﬁrst deﬁne both ﬁnite and inﬁnite models within Listing 1.
We now study how quickly the kernel of the ﬁnite-channel WideResNet approaches its inﬁnite channel limit.
We explore two different axes along which convergence takes place: ﬁrst, as a function of the number of channels (as measured by the widening factor, k ) and second as a function of the number of ﬁnite-network Monte Carlo samples we average over.
N EURAL T ANGENTS makes it easy to compute MC averages of ﬁnite kernels using the following snippet:
[Equation]
[Table]
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 inﬁnite network.
As remarked in Novak et al.
(2019), for any ﬁnite widening factor the MC estimate is biased, but here the bias is small relative to the variance, and the distance to the empirical kernel decreases with the number of samples.
2. C OMPARISON OF N EURAL N ETWORK ARCHITECTURES AND TRAINING SET SIZES
The above examples demonstrated the ﬂexibility N EURAL T ANGENTS provides in comparing various forms of inference for the same architecture.
Now, we further leverage this ﬂexibility to demonstrate training 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 inﬁnite-time inference using the analytic inﬁnite-width NNGP or NTK.
For a given architecture class, there are multiple depth.
For a given training set size we choose speciﬁc depth based on marginal log likelihood on the training set.
from neural_tangents import stax 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) model is used in Figures 2 and 3.
− − l og ( D i s t a n ce )
NNGP log2(#Samples) l og ( w i d e n i n g f a c t o r )  − − l og ( D i s t a n ce )
NTK log2(#Samples) l og ( w i d
e n i n g f a c t o r )
[Figure]
FCCONVWRESNET
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).
I MPLEMENTATION : T RANSFORMING T ENSOR O PS TO K ERNEL O PS
Neural networks are compositions of basic tensor operations such as: dense or convolutional afﬁne transformations, application of pointwise nonlinearities, pooling, or normalization.
For most networks without weight tying between layers, the kernel computation can be written compositionally, with a direct correspondence between each tensor operation and a kernel operation (see §3.1 for an example).
The core logic of N EURAL T ANGENTS is thus a set of translation rules , that sends each function acting on a ﬁnite-width layer to a function acting on the kernel for an inﬁnite-width network.
This is illustrated in Figure 4 for a simple convolutional architecture.
The function applied to the data tensor is given in the second column, and the corresponding transformations of the NTK and NNGP kernel tensors are given in the third and fourth column.
See §E for a list of all tensor operations for which translation rules are currently implemented.
One subtlety to consider when designing networks is that most inﬁnite-width results require nonlinear transformations to be prefaced with afﬁne transformations (either dense or convolutional).
This is because inﬁnite-width results often assume that the pre-activations of nonlinear layers are approximately
Gaussian.
Randomness in weights and biases causes the output of inﬁnite afﬁne layers to satisfy this Gaussian requirement.
Fortunately, prefacing nonlinear operations with afﬁne transformations is common practice when designing neural networks, and N EURAL T ANGENTS will raise an error if this requirement is not satisﬁed.
3. A TASTE OF T ENSOR - TO -K ERNEL O PS T RANSLATION
We demonstrate the intuition behind the translation rules that map tensor ops into kernel ops, and the factorization of kernel ops on a simple example.
We compute the kernel op corresponding to the composition of a dense layer and a nonlinearity layer.
Let z
= z ( X , θ ) ∈ R
d × n be a node in some hidden layer with NNGP kernel and NTK given by
[Equation]
where z i ∈ R
d is the i th neuron and θ is the collection of weights that z depends on.
Here d is the cardinality of the network inputs X and n is the number of neurons in the z node.
We assume z is a mean zero Gaussian, and we compute the kernels of y =
φ ( z ) and h =
Dense ( σ ω , σ b ) ( y ) , where
[Equation]
Global average pooling: take h = 1 d 1 d .
Then
[Equation]
This approach corresponds to applying global average pooling right after the last convolu- tional layer.
This approach takes all pixel-pixel covariances into consideration and makes the kernel translation invariant.
However, it requires O  |X| 2 d 2  memory to compute the sample-sample covariance of the GP (or O  d per covariance entry in an iterative or dis- tributed setting).
It is impractical to use this method to analytically evaluate the GP, and we propose to use a Monte Carlo approach (see § 4).
Subsampling one particular pixel: take h = e ↵ ,
[Equation]
This approach makes use of only one pixel-pixel covariance, and requires the same amount of memory as vectorization ( § 3.1) to compute.
We compare the performance of presented strategies in Figure 1.
Note that all described strategies admit stacking additional FC layers on top while retaining the GP equivalence, using a derivation analogous to § 2 (Lee et al., 2018; Matthews et al., 2018b).
4 M ONTE C ARLO EVALUATION OF INTRACTABLE GP KERNELS
We introduce a Monte Carlo estimation method for NN-GP kernels which are computationally im- practical 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 layers can be constructed in fashion ( §
We focusonglobalaveragepoolinginthisworktomoreeffectivelyisolatetheeffectsofpoolingfromotheraspectsofthemodellikelocalconnectivityorequivariance.
Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes Roman Novak ^ , Lechao Xiao* ^ , Jeahoon Lee* † , Yasaman Bahri* † , Greg Yang°, Daniel A. Abolaﬁa, Je ﬀ rey Pennington, Jascha Sohl-Dickstein Google Brain.
Google AI residents.
^ , † Equal contribution.
°Microsoft Research AI.
Summary
3. Di ﬀ erent NN-GP architectures
SGD-trained ﬁnite vs Bayesian inﬁnite CNNs Under review as a conference paper at ICLR 2019 (a) (c) (b)
No Pooling Global Average Pooling L C N C NN #Channels Figure
(a) trained CNNs often perform better with increasing number of channels particular choice of architecture initialization hyperparameters decay selected independently for each number of axis
(b often of corresponding have the same architecture except for pooling 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 accu- racy 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 speciﬁc 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 chan
[Figure]
P
[Figure]
G
P
[Figure] [Figure]
[Figure] [Figure]
experimental
[Figure]
[Figure] [Figure]
[Table] [Table]
l following from
[Figure]
[Table]
-meanGaussian,andthecovarianceis
E { !
l ,b l } h Bz lj ( x
)  Bz lj ( x 0 )  T
K l i =
B E { !
l ,b l } h z lj ( x ) z lj ( x 0 ) T
K l i B T .
Onecaneasilyseethat
Bz lj
K l j arei.i.d.
multivariateGaussianaswell.
Stridedconvolution .
Stridedconvolutionisequivalenttoanon-stridedconvolutioncomposedwith subsampling.
Let s 2 N denote size of the stride.
Then the strided convolution is equivalent to choosing B asfollows: B ij =
( is  j ) for i 2 { 0 , 1 ,...
( d 2  1) } .
Average pooling .
Average pooling with stride s and window size ws is equivalent to choosing B ij = 1 /ws for i = 0 , 1 ,...
( d 2  1) and j = is,..., ( is + ws  1) .
NDconvolutions.
Notethatouranalysisinthemaintext(1D)easilyextendstohigher-dimensional convolutions by replacing integer pixel indices and sizes d,↵, with tuples (see also Figure 4).
4. Computing the NN-GP covariance
[Table]
(a) SGD-trained ﬁnite CNNs often perform better with increasing number of channels.
(b) Performance of SGD-trained ﬁnite CNNs often approaches that of the respective CNN-GP with increasing the number of channels.
(c) However, the best performing SGD CNN outperforms the best CNN-GP by a signiﬁcant margin.
[Table]
Disentangling the CNN architecture
FCN FCN
[Table]
GP
[Table]
LCN w/pooling CNN
[Table]
GP CNN CNN
[Table]
/ pooling
Compositionality Local connectivity Equivariance Invariance Quality
Model
A cc u r ac
[Table]
A cc u r
FCN
CNN w/ ERF CNN w/ small learning rate CNN w/ ReLU & large learning rate CNN w/ pooling
CNN-GP
FCN-GP
NN (
NN (no underﬁtting) n 1 →∞ n 1 →∞ n 2 →∞ n 2 →∞ 4x4 covariance per class yzzy 2 = ϕ
( z 1 ) ¯ z 210 classes per input 4 inputs, 3 features 4321⊗
I 10 K 0 K
K
K 2 K
K 2 σ 2 ω K 0 + σ 2 b σ 2 ω K 1 + σ 2 b 𝒜 ( K 1 )
𝒜 ( K 0 ) a ﬃ ne transform a ﬃ ne transform a ﬃ ne
[Table]
nonlinearity nonlinearity
[Table]
4x4x covariance covariance
[Table]
[Table] [Table]
[Table]
ne
[Table]
𝒦
[Table] [Table]
2. Convergence
[Figure]
( CNN w / o w e
i g h t s h a r i n g )
[Figure]
(Fully-connected) (before the last readout layer)
Under review as a conference paper at ICLR
[Figure]
random ﬁnite widthnetworksandusetheempiricaluncenteredcovariancesofactivationstoestimate the Monte Carlo
[Equation]
(19) where ✓ consistsof M drawsof theweights andbiasesfrom theirprior distribution, ✓ m ⇠ p ( ✓ ) , and n is the width or number of channels in hidden layers.
The MC-GP kernel converges to the analytic kernel with increasing width, lim n !
1 K ln,M = K l 1 in probability.
For ﬁnite width networks, the uncertainty in K ln,M is Var ⇥
K ln,M ⇤ =
Var ✓ ⇥
K ln ( ✓ ) ⇤ /M .
From Daniely et al.
, we know that Var ✓ ⇥
K ln ( ✓ ) ⇤ / 1 n , which leads to Var ✓ [ K ln,M ] / 1 Mn .
For ﬁnite n , K ln,M is also a biased estimate of K l 1 , where the bias depends solely on network width.
We do not currently have an analytic form for this bias, but we can see in Figures 2 and 7 that for the hyperparameters we probe it is small relative to the variance.
In particular,
K ln,M ( ✓ )  K L 2 F is nearly constant for constant Mn .
We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate.
Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.
In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from
O ⇣ |X| 2 d 2 ⌘ to O ⇣ |X| 2 + n 2 + nd ⌘ , making the evaluation of CNN-GPs with pooling practical.
5 D ISCUSSION 5.1 B AYESIAN CNN S WITH MANY CHANNELS ARE IDENTICAL TO LOCALLY CONNECTED NETWORKS , IN THE ABSENCE OF POOLING
Locally Connected Networks (LCNs) (Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spatial locations.
LCNs preserve the connectivity pattern, and thus topology, of a
CNN.
However, they do not possess the equivariance property of a CNN – if an input is translated, the latent representation in an LCN will be completely different, rather than also being translated.
The CNN-GP predictions without spatial pooling in § 3.1 and § 3.2.2 depend only on sample-sample covariances, and do not depend on pixel-pixel covariances.
LCNs destroy pixel-pixel covariances: ⇥
K L⇤
LCN ↵,↵ 0 ( x,x 0 ) = 0 , for ↵ 6 = ↵ 0 and all x,x 0 2 X and L > 0 .
However, LCNs preserve the covariances between input examples at every pixel: ⇥
K L⇤
LCN ↵,↵ ( x,x 0 ) = ⇥
K L⇤
CNN ↵,↵ ( x,x 0 ) .
As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical.
Moreover, LCN-GPs with
Under review as a conference paper at ICLR
[Figure]
reported to
See Figure 7 for similar with random ﬁnite widthnetworksandusetheempiricaluncenteredcovariancesofactivationstoestimate the Monte Carlo-GP (MC-GP) kernel,
[Equation]
where ✓ consistsof M drawsof theweights andbiasesfrom theirprior distribution, ✓ m ⇠ p
[Equation]
✓ ) , and n is the width or number of channels in hidden layers.
The MC-GP kernel converges to the analytic kernel with increasing width, lim n !
1 K ln,M = K l 1 in probability.
For ﬁnite width networks, the uncertainty in K ln,M is Var ⇥
K ln,M ⇤ =
Var ✓ ⇥
K ln ( ✓ ) ⇤ /M .
From Daniely et al.
, we know that Var ✓ ⇥
K ln ( ✓ ) ⇤ / 1 n , which leads to Var ✓ [ K ln,M ] / 1 Mn .
For ﬁnite n , K ln,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,
[Equation]
is nearly constant for constant Mn .
We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate.
Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.
In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from
O ⇣ |X| 2 d 2 ⌘ to O ⇣ |X| 2 + n 2 + nd ⌘ , making the evaluation of CNN-GPs with pooling practical.
5 D ISCUSSION 5.1 B AYESIAN CNN S WITH MANY CHANNELS ARE IDENTICAL TO LOCALLY CONNECTED NETWORKS , IN THE ABSENCE OF POOLING
Locally Connected Networks (LCNs) (Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spatial locations.
LCNs preserve the connectivity pattern, and thus topology, of a
CNN.
However, they do not possess the equivariance property of a CNN – if an input is translated, the latent representation in an LCN will be completely different, rather than also being translated.
The CNN-GP predictions without spatial pooling in § 3.1 and § 3.2.2 depend only on sample-sample covariances, and do not depend on pixel-pixel covariances.
LCNs destroy pixel-pixel covariances: ⇥
K L⇤
LCN ↵,↵ 0 ( x,x 0 ) = 0 , for ↵ 6 = ↵ 0 and all x,x 0 2 X and L > 0 .
However, LCNs preserve the covariances between input examples at every pixel: ⇥
K L⇤
LCN ↵,↵ ( x,x 0 ) = ⇥
K L⇤
CNN ↵,↵ ( x,x 0 ) .
As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical.
Moreover, LCN-GPs with
[Table]
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 a ﬃ ne transformation) (inputs covariance) (computes covariance of post-nonlinearity activations given normal pre-activations)  outputs  inputs
D !
min { n 1 ,...,n L
[Table]
/vGM2E4A7n0gDIjfPOYLaihDPyr6PmpxOszuAqOd3dijz/sDvZed/PZQI/RE7SNYvQC7aG36AiNEAsOAwhc8DWchd/C7+GPVhoGXc0j9E+EP/8C5Jrtng==</latexit> (#channels in the outputs)
( F u ll y - c
o nn ec t e d ) (all displayed models are 100% accurate on the training set)
[Figure]
n 1 →∞ n 1 →∞ n 2 →∞ n 2 →∞
CNN ( 1 D ) ac t
i v a t
i o n s
CNN - G
P c o v a r i a n ce
Under review as a conference paper at ICLR
Global average pooling:
h
= 1 d 1 d .
Then
[Equation]
This approach corresponds to applying global average pooling right after the last convolu- tional layer.
This approach takes all pixel-pixel covariances into consideration and makes the kernel translation invariant.
However, it requires O  |X| 2 d 2  memory to compute the sample-sample covariance of the GP (or O  d per covariance entry in an iterative or dis- tributed setting).
It is impractical to use this method to analytically evaluate the GP, and we propose to use a Monte Carlo approach (see § 4).
Subsampling one particular
pixel
: take h
[Equation]
↵
[Equation]
This approach makes use of only one pixel-pixel covariance, and requires the same amount of memory as vectorization ( § 3.1) to compute.
We compare the performance of presented strategies in Figure 1.
Note that all described strategies admit stacking additional FC layers on top while retaining the GP equivalence, using a derivation analogous to § 2 (Lee et al., 2018; Matthews et al., 2018b).
4 M ONTE C ARLO EVALUATION OF INTRACTABLE GP KERNELS
We introduce a Monte Carlo estimation method for NN-GP kernels which are computationally im- practical 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
Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes Roman Novak ^ , Lechao Xiao* ^ , Jeahoon Lee* † , Yasaman Bahri* † , Greg Yang°, Daniel A. Abolaﬁa, Je ﬀ rey Pennington, Jascha Sohl-Dickstein Google Brain.
Google AI residents.
^ , † Equal contribution.
°Microsoft Research AI.
Summary
3. Di ﬀ erent NN-GP architectures
SGD-trained ﬁnite vs Bayesian inﬁnite CNNs Under review as a conference paper at ICLR 2019 (a) (c) (b)
No Pooling Global Average Pooling L C N C NN #Channels
[Figure] [Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Table]
2  1) and j = is
[Table]
..., ( is + ws  1) .
NDconvolutions.
Notethatouranalysisinthemaintext(1D)easilyextendstohigher-dimensional convolutions by replacing integer pixel indices and sizes d,↵, with tuples (see also Figure 4).
4. Computing the NN-GP covariance
[Table]
• Various CNN architectural decisions like pooling / vectorizing / subsampling , zero or no padding have a corresponding NN-GP.
• Pooling enforces translation-invariant predictions in both CNNs and CNN- GPs and allows for the best performance.
• Shallow CNN models can perform worse than fully-connected alternatives due to failing to capture non-linear interactions
Disentangling the CNN architecture
FCN
FCN-GP LCN LCN w/pooling CNN-GP CNN CNN w/ pooling
Compositionality Local connectivity Equivariance Invariance Quality
Model
A cc u r ac y A cc u r ac y
FCN
CNN w/ ERF CNN w/ small learning rate CNN w/ ReLU & large learning rate CNN w/ pooling
CNN-GP
FCN-GP
NN (underﬁtting allowed)
NN (no underﬁtting) n 1 →∞ n 1 →∞ n 2 →∞ n 2 →∞
[Figure]
yzzy 2 = ϕ
( z 1 ) ¯ z
⊗ I
[Figure]
[Figure]
matrix-multiply matrix-multiply matrix-multiply nonlinearity nonlinearity
[Figure]
[Figure]
[Table]
/wLb6p105Oy3xm0QbaRLvIR4foBF2ga9RADH04FWfNWXc+XezuuHuF1XXKnjX0p9z6F6yiwSY=</latexit>
FCN (fully-connected):
CNN (1D): 6.
When do
• 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 inﬁnity.
• 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 inﬁnite Bayesian CNNs (CNN-GPs), and ﬁnite SGD-trained CNNs can outperform them.
Convergence
[Figure]
( CNN w / o w e
i g h t s h a r i n g ) (
[Figure]
non
[Figure]
linearity) (CNN w/o weight sharing
[Figure] [Figure]
(Fully
[Figure]
connected) (before the last readout layer)
Under review as a conference paper at ICLR
[Figure] [Figure]
ﬁnite widthnetworksandusetheempiricaluncenteredcovariancesofactivationstoestimate the Monte Carlo-GP (MC-GP) kernel,
[Equation]
where ✓ consistsof M drawsof theweights andbiasesfrom theirprior distribution, ✓ m ⇠ p ( ✓ ) , and n is the width or number of channels in hidden layers.
The MC-GP kernel converges to the analytic kernel with increasing width, lim n !
1 K ln,M = K l 1 in probability.
For ﬁnite width networks, the uncertainty in K ln,M is Var
[Equation]
M .
From Daniely et al.
, we know that Var ✓
[Equation]
/ 1 n , which leads to Var ✓ [ K ln,M ] / 1 Mn .
For ﬁnite n , K ln,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,
[Equation]
is nearly constant for constant Mn .
We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate.
Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.
In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from
O ⇣ |X| 2 d 2 ⌘ to O ⇣ |X| 2 + n 2 + nd ⌘ , making the evaluation of CNN-GPs with pooling practical.
5 D ISCUSSION 5.1 B AYESIAN CNN S WITH MANY CHANNELS ARE IDENTICAL TO LOCALLY CONNECTED NETWORKS , IN THE ABSENCE OF POOLING Locally Connected Networks (LCNs)
(Fukushima, 1975; Lecun, 1989) are CNNs without weight sharing between spatial locations.
LCNs preserve the connectivity pattern, and thus topology, of a CNN.
However, they do not possess the equivariance property of a CNN – if an input is translated, the latent representation in an LCN will be completely different, rather than also being translated.
The CNN-GP predictions without spatial pooling in § 3.1 and § 3.2.2 depend only on sample-sample covariances, and do not depend on pixel-pixel covariances.
LCNs destroy pixel-pixel covariances: ⇥ K L 1 ⇤ LCN ↵,↵ 0 ( x,x 0 ) =
0 , for ↵ 6 = ↵ 0 and all x,x 0 2 X and L > 0 .
However, LCNs preserve the covariances between input examples at every pixel: ⇥
K L 1 ⇤ LCN ↵,↵ ( x,x 0 ) = ⇥
K L 1 ⇤ CNN ↵,↵ ( x,x 0 ) .
As a result, in the absence of pooling, LCN-GPs and CNN-GPs are identical.
Moreover, LCN-GPs with
Under review as a conference paper at ICLR
[Figure] [Figure]
ﬁnite widthnetworksandusetheempiricaluncenteredcovariancesofactivationstoestimate the Monte Carlo-GP (MC-GP) kernel,
[Equation]
(19) where ✓ consistsof M drawsof theweights andbiasesfrom theirprior distribution, ✓ m ⇠ p ( ✓ ) , and n is the width or number of channels in hidden layers.
The MC-GP kernel converges to the analytic kernel with increasing width, lim n !
1 K ln,M = K l 1 in probability.
For ﬁnite width networks, the uncertainty in K ln,M is Var ⇥
K ln,M ⇤ =
Var ✓ ⇥
K ln ( ✓ ) ⇤ /M .
From Daniely et al.
, we know that Var ✓ ⇥
K ln ( ✓ ) ⇤ / 1 n , which leads to Var ✓ [ K ln,M ] / 1 Mn .
For ﬁnite n , K ln,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,
[Equation]
[Equation]
is nearly constant for constant Mn .
We thus treat Mn as the effective sample size for the Monte Carlo kernel estimate.
Increasing M and reducing n can reduce memory cost, though potentially at the expense of increased compute time and bias.
In a non-distributed setting, the MC-GP reduces the memory requirements to compute GP pool from
O ⇣ |X| 2 d 2 ⌘ to O ⇣ |X| 2 + n 2 + nd ⌘ , making the evaluation of CNN-GPs with pooling practical.
5 D ISCUSSION 5.1 B AYESIAN CNN S 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-
[Table]
+
[Table]
+
[Table] [Table]
==</latexit> (#channels in the outputs)
( F u ll y - c
o nn ec t e d ) (all displayed models are 100% accurate on the training set)
(attaching
[Table]
n 1 →∞ n 1 →∞ n 2 →∞ n 2 →∞
[Table]
[Equation]
[Table]
[Equation]
(cid:0)
[Equation]
(cid:0)
[Equation]
(cid:1)
[Table]
[Equation]
(cid:0)
[Equation]
(cid:1)
[Equation]
(cid:16)
[Equation]
(cid:17)
[Equation]
(cid:16)
[Equation]
(cid:17) (cid:12)
[Equation]
[Table] [Equation] (cid:0) [Equation]
(cid:1)
[Equation]
(cid:0)
[Equation]
(cid:1)
[Equation]
(cid:0)
[Equation]
(cid:1)
[Table]
[Equation]
(cid:0)
[Equation]
(cid:1)
[Equation]
(cid:16) ˜
[Equation]
(cid:17)
[Equation]
(cid:16)
[Equation]
(cid:17) (cid:12)
[Equation]
[Table]
z
[Equation]
= Dense ◦
[Table]
(cid:0) y
[Equation]
(cid:1) ˜
[Equation]
(cid:0)
K
[Equation]
(cid:1) ˜Θ
[Equation]
and the variables W ij and β i are i.i.d.
Gaussian N (0 , 1) .
We will show that kernel op φ ∗ and
Dense ( σ ω , σ b ) ∗ induced by the tensor op φ and Dense ( σ ω , σ b ) are
[Equation]
which implies
[Equation]
First notice that
[Equation]
To compute the NTK of h , note that the φ op does not introduce any new variables.
Thus
[Equation]
This gives Equation 3.
Now consider the Dense op.
Using independence between the variables
[Equation]
Finally, to derive Equation 4, we simply remark that the NTK of h is a sum of two terms:
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
3. P ERFORMANCE
Our library does a number of automatic performance optimizations without sacriﬁcing ﬂexibility.
Leveraging block-diagonal covariance structure.
A common computational challenge with GPs is inverting the training set covariance matrix.
Naively, for a classiﬁcation 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 ,× 500 , 000 .
However, we highlight that if a fully-connected readout layer is used (which is an extremely common design in classiﬁcation architectures), the C logits are i.i.d.
conditioned on the input x , resulting in outputs being normally distributed with a block-diagonal covariance matrix of form Σ ⊗ I C , where Σ has the shape |X| × |X| , which makes closed-form exact inference feasible in these tasks.
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 (as opposed to inverting it), as construction of the |X| × |X| output covariance Σ involves computing intermediary layer l covariances Σ 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 intermediary layer outputs, which is 1024 in the case of CIFAR-10 and SAME padding.
However, as Xiao et al.
Novak et al.
Garriga-Alonso et al.
(2019) remarked, if no pooling is used in the network the output covariance Σ can be computed by only using the stack of d |X| × |X| -blocks of Σ l , bringing the time and memory cost from O ( |X| 2 d 2 ) 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 N EURAL T ANGENTS to achieve efﬁcient 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 inﬁnite-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).
N EURAL T ANGENTS provides an easy way to perform both of these common tasks using a single batch decorator shown below:
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-speciﬁc 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 transpositionsby adjusting the computation performed by each layer based on the input metadata.
Batch Size 22 - 22 - 22 - 2T
i m e /
S a m p l e ( s )
GPU Count = 1 GPU Count =
GPU Count = 4 GPU Count =
GPU Count 22 - 22 - 22 - 2T i m e /
S a m p l e ( s )
C ONCLUSION
We believe N EURAL T ANGENTS will enable researchers to quickly and easily explore inﬁnite-width networks.
By democratizing this previously challenging model family, we hope that researchers will begin to use inﬁnite neural networks, in addition to their ﬁnite 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 inﬁnite networks as theoretical tools to gain insight and clarity into many of the hard theoretical problems in deep learning.
Going forward, there are signiﬁcant additions to N EURAL
T ANGENTS that we are exploring.
There are more layers we would like to add in the future (§E) that will enable an even larger range of inﬁnite 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!
A PPENDIX A N EURAL T ANGENTS AND PRIOR WORK
Here we brieﬂy discuss the differences between N EURAL T ANGENTS and the relevant prior work.
Prior benchmarks in the domain of inﬁnitely 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.
Code released by Lee et al.
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, N EURAL T ANGENTS differs from the code released by Lee et al.
by about +9500/–2500 lines of code.
GPy (2012), GPFlow (Matthews et al., 2017), and other GP packages.
While various packages allowing for kernel construction, optimization, and inference with Gaussian Processes exist, none of them allow to easily construct the very speciﬁc kernels corresponding to inﬁnite neural networks (NNGP/NTK; nt.stax ), nor do they provide the tools and convenience for studying wide but ﬁnite networks and their training dynamics ( nt.taylor_expand , nt.predict , nt.monte_carlo_kernel_fn ).
On the other hand, N EURAL T ANGENTS does not provide any tools for approximate inference with these kernels.
B L IBRARY DESCRIPTION
N EURAL T ANGENTS provides a high-level interface for specifying analytic, inﬁnite-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. N EURAL 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 speciﬁed 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 ﬁnite neural network.
[Table]
B. I NFINITE NEURAL NETWORKS WITH N EURAL T ANGENTS
We extend stax layers to return a third function kernel_fn , which represents the covariance functions of the inﬁnite NNGP and NTK networks of the given architecture (recall that since inﬁnite networks are GPs, they are fully deﬁned by their covariance functions, assuming 0 mean as is common in the literature).
[Table]
We demonstrate a speciﬁcation 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:
[Equation]
Note that the above code does not do Cholesky decomposition and is presented merely to show the mathematical expression.
We provide efﬁcient GP inference method in the predict submodule:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
B. C OMPUTING 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).
[Figure]
[Equation]
[Table]
[Figure]
[Equation]
[Table]
[Equation]
Note that batching is often used to compute large covariance matrices that may not even ﬁt 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 ,training covariance matrix and a 45 , 000 × 10 training target matrix).
B. T RAINING 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 ﬁrst 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)
[Figure]
r
[Figure]
d
E rr o r 0th order (constant) 1st
[Figure]
(linear) 2nd
[Figure]
(quadratic)
[Figure]
B. I NFINITE 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 efﬁciently compute Monte Carlo estimates of the NNGP covariance and NTK.
These functions work with neural networks constructed using any neural network library.
[Figure]
from jax.experimental import stax import neural_tangents as
We demonstrate convergence of the Monte Carlo kernel estimates to the closed-form analytic kernels in the case of a WideResNet in Figure 2.
B. W EIGHTS OF WIDE BUT FINITE NETWORKS
While most of N EURAL T ANGENTS is devoted to a function-space perspective—describing the distribution of function values on ﬁnite collections of training and testing points—we also provide tools to investigate a dual weight space perspective described in Lee et al.
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.
[Table]
These act exactly like normal JAX functions and, in particular, can be plugged into gradient descent, which we demonstrate in Figure 6.
B. E XTENDING N EURAL T ANGENTS
Many neural network layers admit a sensible inﬁnite-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 N EURAL T ANGENTS , one only has to implement it as a method in nt.stax with the following signature:
[Equation]
Here init_fn and apply_fn are initialization and the forward pass methods of the ﬁnite 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 ﬁnite width layer methods).
In this case what remains is to implement the kernel_fn method with signature
[Equation]
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 speciﬁc operation to be performed should be derived by the user in the context of the particular operation that the ﬁnite width layer performs.
This transformation could be as simple as an afﬁne 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 C OMPARISON OF N EURAL N ETWORK A RCHITECTURES ON S MALL D ATASETS
Recent work by Arora et al.
(2019b) showed that the NTK of rectiﬁed linear networks combined with
Support Vector Machines can outperform other techniques (both classical kernels as well as neural networks) on small tabular datasets.
To do this the authors considered a collection of 90 UCI datasets along with a protocol established by Fernández-Delgado et al.
On each dataset, the authors conducted a search over SVM hyperparameters as well as the number of frozen layers of the network that were left untrained.
Final performance was reported following 4-fold cross-validation.
We adapted code provided by Arora et al.
(2019b)
and used N EURAL T ANGENTS to additionally search over three new architectures: fully-connected Erf networks, ReLU residual networks, and ReLU residual networks with Fixup (Zhang et al., 2019)
inspired initialization (here the weights before each activation function are scaled by a factor of 1 / √ L but the weights following the activation function have unit variance).
In each case, for each dataset, we optimized over which layers of the network were frozen and the variance of the initial weights.
Our results are summarized in Table 1.
We observe that NTKs derived from Erf networks seem to perform the worst while those from residual networks with Fixup exhibit a negligible improvement in performance.
[Table]
D A RCHITECTURE SPECIFICATIONS
[Table] [Table]
Listing 2: All-convolutional model (ConvOnly) deﬁnition used in Figure 3 .
[Table] [Table]
[Table] [Table]
std)
[Table] [Table]
layers
[Table] [Table]
Listing 3: Fully-connected (FC) model deﬁnition used in Figure 3 .
E
I
The following layers 7 are currently implemented, with translation rules given in Table 2:
• serial • parallel • FanOut • FanInSum • Dense • Conv 8 with arbitrary ﬁlter shapes, strides, and padding• Relu • LeakyRelu • Abs • ABRelu 10 , • Erf • Identity • Flatten • AvgPool • GlobalAvgPool • GlobalSelfAttention (Hron et al., 2019) • LayerNorm
The following is in our near-term plans:
• Dropout • SumPool • Exp , Elu , Selu , Gelu • Apache Beam support.
The following layers do not have a known closed-form solution for inﬁnite network covariances, and networks with them have to be estimated empirically (provided with out implementation via nt.monte_carlo_kernel_fn ) or using other approximations (not currently implemented):
• Sigmoid , Tanh 11 , Swish 12 , Softmax , LogSoftMax , Softplus , MaxPool .
Tensor Op
NNGP Op
NTK Op
X
K Θ
[Table]
[Equation]
[Table]
[Equation]
Dropout
[Table]
[Equation]
[Table]
( σ
[Equation]
, σ b )
[Equation]
[Table]
Attn ( σ QK , σ OV )
Attn
[Equation]
K )
[Equation]
[Table]
[Equation]
K
[Equation]
(Hron et al., 2019)
[Table]
FanInSum
[Table]
FanOut
[Table]
in §3.2 are applicable).
Notation details.
The Tr and GlobalAvgPool ops are assumed to act on all spatial axes (with sizes H and W in this example), producing a |X| × 2 -kernel.
Similarly, the AvgPool op is assumed to act on all spatial axes as well, applying the speciﬁed strides s , pooling window sizes p and padding strategy p to the respective axes pairs in K and T (acting as 4D pooling with replicated parameters of the 2D version).
T and ˙ T are deﬁned identically to Lee et al.
as T
[Equation]
Σ) =
E
[Equation]
These expressions can be evaluated in closed form for many nonlinearities, and preserve the shape of the kernel.
The A op is deﬁned similarly to Novak et al.
Xiao et al.
(2018) as [ A (Σ)] w,w (cid:48) h,h
[Equation]
( x, x
[Equation]
) =
[Equation]
dh
[Equation]
( x, x (cid:48) ) /q 2 , where the summation is performed over the convolutional ﬁlter receptive ﬁeld 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).
See Figure 4 for an example of applying the translation rules to a speciﬁc model, and §3.1 for deriving a sample translation rule. [Figure] i [Figure]
i
[Figure]
n N
[Figure]