N EURAL T ANGENTS : F AST AND E ASY I NFINITE N EURAL N ETWORKS IN P YTHON
A BSTRACT
N EURAL T ANGENTS is a library designed to enable research into 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 ﬁnite-width as usual or in their inﬁnite-width limit.
Inﬁnite-width networks can be trained analytically using exact Bayesian inference or using gradient descent via the Neural Tangent Kernel.
Additionally, N EURAL T ANGENTS provides tools to study gradient descent training dynamics of wide but ﬁnite 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.
N EURAL T ANGENTS is available at w ww.github.com/google/neural-tangents
We also provide an accompanying interactive Colab notebook 1 .
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), 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 deﬁning the NNGP kernel has additionally been extensively studied in the context of mean ﬁeld 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. 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; 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 ﬁ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; 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. 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 accelerate 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: continuous 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 li ( x ) describe the i th pre-activation 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 li ( x ) z li ( x (cid:48) )] .
This describes a NNGP with the kernel, K ( x, x (cid:48) ) .
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 ) − 1 Y and variance σ 2 ( x ) =
K ( x, x ) − K ( x, X ) K ( X , X ) − 1 K
( 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 ∂ t f t ( X ) = − ηJ t ( X ) J t ( X ) T ( f t ( X ) − Y ) where J t ( X ) is the Jacobian of the output f t 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 , σ 2 ) independently and identically distributed.
To train an inﬁnite neural network with Erf activations 4 on this data using gradient descent and an MSE loss we write the following:
[Equation]
[Equation]
ntk
[Equation]
diag
[Equation]
, compute
[Equation]
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:
[Figure]
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]
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 noted in Novak et al.
(2019), for any ﬁnite 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. C OMPARISON OF N EURAL N ETWORK ARCHITECTURES AND TRAINING SET SIZES
The above examples demonstrate how one might construct a complicated architecture and perform inference using N EURAL T ANGENTS ˙Next 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 inﬁnite-time inference using the analytic inﬁnitewidth 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 =
Listing 1: Deﬁnition of an inﬁnitely WideResNet .
This snippet simultaneously deﬁnes a ﬁnite ( init_fn, apply_fn ) and an inﬁnite ( kernel_fn ) model. This model is used in Figures 2 and 3. [Figure] [Figure]
l [Figure] ( D i s t a n ce )
NNGP log [Figure] ) l og ( w i d e n i n g f a c t o r ) [Figure]
[Figure]
n
[Figure]
NTK [Figure] [Figure]
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 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 N EURAL T ANGENTS is a set of translation rules , that sends each tensor operation acting on a ﬁnite-width layer to a corresponding transformation of the kernel for an inﬁnite-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 inﬁnite-width results require nonlinear transformations to be preceded by 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
To get some intuition behind the translation rules, we consider the case of a nonlinearity followed by a dense layer.
Let z
= z ( X , θ ) ∈ R
d × n be the preactivations resulting from d distinct inputs at a node in some hidden layer of a neural network. Suppose z has NNGP kernel and NTK given by
[Equation]
where z i ∈ R
d is the i th neuron and θ are the parameters in the network up until z .
Here d is the cardinality of the network inputs X and n is the number of neurons in the z node.
We assume z is a mean zero multivariate Gaussian.
We wish to compute the kernel corresponding
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) SGD-trained CNNs often perform better with increasing number of channels Each line corresponds to a particular choice of architecture and initialization hyperparameters, with learning rate and decay selected independently for each number of channels x axis) (b)
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 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]
[Figure]
G
P
F C N
G P
Validation share common
[Figure] [Table]
for experimental details
[Table] [Table]
S
[Table]
P
The
[Table]
GP kernel at to l following from e
+
[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
[Table]
[Table]
Invariance Quality Model
[Table]
u r
[Table]
[Table]
[Table]
[Table]
CNN
[Table]
GP
FCN-GP
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]
4x4x10 input covariance covariance
[Table]
[Table]
[Table]
a ﬃ ne transform transform
[Table] [Table]
2. Convergence
[Figure]
) ( CNN w / o w e i g h t s h a r i n g ) (
[Figure]
non-linearity)
(CNN w/ [Figure] [Figure] -connected) (before the last readout layer)
Under review as a conference paper at ICLR
[Figure]
random ﬁ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,
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]
random ﬁ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 ⇥
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: 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
[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
We 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]
[Table]
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
behavior
[Table]
converge
[Table]
a tractable analytic form
[Table]
of a CNN-GP with no pooling can
[Table]
computed
[Table]
resources
[Table]
other
[Table]
we use a Monte Carlo approach
[Table]
of a given
[Table]
computing the
[Table] allows to [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 between
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 [-
-]210 classes per input
⊗ I
[Figure]
[Figure]
ﬃ ﬃ matrix-multiply matrix-multiply matrix-multiply nonlinearity nonlinearity
[Figure]
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[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 ) (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
[Figure]
random ﬁ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]
[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]
random ﬁnite widthnetworksandusetheempiricaluncenteredcovariancesofactivationstoestimate the Monte Carlo-GP (MC-GP) kernel,
[Equation]
n [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-pixel
[Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
+
[Table]
+EP/
[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 a fully-
[Table]
only to
[Table]
of
[Table]
n 1 →∞ n 1 →∞ n 2 →∞ n 2 →∞
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
z
[Equation]
[Table]
[Equation]
to h = Dense ( σ ω , σ b ) ( φ ( z )) by computing the kernels of y =
φ ( z ) and h =
Dense ( σ ω , σ b ) ( y ) separately. Here,
[Equation]
and the variables W ij and β i are i.i.d.
Gaussian N (0 , 1) .
We will compute kernel operationsdenoted φ ∗ and Dense ( σ ω , σ b ) ∗ - induced by the tensor operations φ and Dense ( σ ω , σ b ) 5 . Finally, we will compute the kernel operation associated with the composition ( Dense ( σ ω , σ b ) ◦ φ ) ∗ =
Dense ( σ ω , σ b ) ∗ ◦ φ ∗ . First we compute the NNGP and NT kernels for y . To compute K y note that from its deﬁnition,
[Equation]
Since φ does not introduce any new variables Θ y can be computed as,
[Equation]
Taken together these equations imply that,
[Equation]
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 φ .
[Equation]
Next we consider the case of a dense operation.
Using the independence between the weights, the biases, and h it follows that,
[Equation]
Finally, the NTK of h can be computed as a sum of two terms:
[Equation]
This gives the translation rule for the dense layer in terms of K y and Θ y as,
[Equation]
3. P ERFORMANCE
Our library performs 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, 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 . This results in outputs that are normally distributed with a block-diagonal covariance matrix of the form Σ ⊗ I C , where Σ has shape |X| × |X| and I C is the C × C identity matrix. This reduces the computational complexity and storage in many common cases by an order of magnitude, which makes closed-form exact inference feasible in these cases.
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). Speciﬁcally for a convolutional network of depth l , constructing the |X| × |X| output covariance matrix, Σ , involves computing l intermediate layer covariance matrices, Σ l , of size | [Equation] ×
[Equation]
| 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.
; Novak et al.
; 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 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:
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 ) [Figure]
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.
[-GPU Count = 4 GPU Count =
-]
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 (§D) 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 CKNOWLEDGMENTS
We thank Yasaman Bahri for frequent discussion and useful feedback on the manuscript, Sergey Ioffe for feedback on the text, as well as Greg Yang, Ravid Ziv, and Jeffrey Pennington for discussion and feedback on early versions of the library.
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.
(2019) by about +9 , 500 / − 2 , 500 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 easy construction of 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. from
[Table]
init_fn, apply_fn =
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
_, params =
[Table]
[Equation]
key
[Equation]
[Table]
[Table]
[Table]
[Table]
[Table]
=
[Table]
[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). from
[Table]
init
[Table]
[Table] , kernel [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]
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]
[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 .
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation]
[Table]
[Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
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]
t c h M e a n
Sq u a r e d
E rr o r 0th order (constant) 1st order (linear) 2nd order (quadratic)
[Figure]
[Figure]
init_fn, apply_fn = stax.serial(stax. Dense(64), stax. BatchNorm(), stax. Sigmoid, stax.Dense(1))
[Figure]
_fn = nt [Figure] _carlo_kernel_fn(init_fn, apply_fn, key=random.
PRNGKey(1), n_samples=128)
[Figure]
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. [Figure] tangents
[Figure]
nt taylor_apply_fn
[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:
@_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 ﬁ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 kernel_fn(
[Table] [Table] [Table] [Table]
) ->
[Table]
[Table]
[Table] [Table] .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 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 A RCHITECTURE SPECIFICATIONS from
[Table]
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
) deﬁnition used in Figure 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 in Figure 3 .
D
The following layers 7 are currently implemented, with translation rules given in Table 1:
• 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:
• SumPool • Dropout • FanInConcat • 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 .
[Table] Op
NNGP Op
NTK Op
[Table]
[Equation]
[Table]
( σ QK , σ OV )
[Table]
[Table]
[Table]
(Hron et al., 2019)
[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 (cid:48) )
= (cid:80) 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]
Training [Figure] LL
[Figure]
Full Test Covariance [Figure] [Figure] N LL ( n a
[Figure]
Independent [Figure] C o n d i t i o n N u m b e r
[Figure]