P REDICTING THE O UTPUTS OF F INITE N ETWORKS T RAINED WITH N OISY G RADIENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
A recent line of works studied wide deep neural networks (DNNs) by approximating them as Gaussian Processes (GPs).
A DNN trained with gradient ﬂow was shown to map to a GP governed by the Neural Tangent Kernel (NTK), whereas earlier works showed that a DNN with an i.i.d.
prior over its weights maps to the socalled Neural Network Gaussian Process (NNGP).
Here we consider a DNN training protocol, involving noise, weight decay and ﬁnite width, whose outcome corresponds to a certain non-Gaussian stochastic process.
An analytical framework is then introduced to analyze this non-Gaussian process, whose deviation from a GP is controlled by the ﬁnite width.
Our contribution is three-fold: (i)
In the inﬁnite width limit, we establish a correspondence between DNNs trained with noisy gradients and the NNGP, not the NTK.
(ii)
We provide a general analytical form for the ﬁnite width correction (FWC) for DNNs with arbitrary activation functions and depth and use it to predict the outputs of empirical ﬁnite networks with high accuracy.
Analyzing the FWC behavior as a function of n , the training set size, we ﬁnd that it is negligible for both the very small n regime, and, surprisingly, for the large n regime (where the GP error scales as O (1 /n ) ).
(iii)
We ﬂesh-out algebraically how these FWCs can improve the performance of ﬁnite convolutional neural networks (CNNs) relative to their GP counterparts on image classiﬁcation tasks.
I NTRODUCTION
Deep neural networks (DNNs) have been rapidly advancing the state-of-the-art in machine learning, yet a complete analytic theory remains elusive.
Recently, several exact results were obtained in the highly over-parameterized regime ( N → ∞ where N denotes the width or number of channels for fully connected networks (FCNs) and convolutional neural networks (CNNs), respectively)
(Daniely et al., 2016).
This facilitated the derivation of an exact correspondence with Gaussian Processes (GPs) known as the Neural Tangent Kernel (NTK)
(Jacot et al., 2018).
The latter holds when highly over-parameterized DNNs are trained by gradient ﬂow, namely with vanishing learning rate and involving no stochasticity.
The NTK result has provided the ﬁrst example of a DNN to GP correspondence valid after end-to-end
DNN training.
This theoretical breakthrough allows one to think of DNNs as inference problems with underlying GPs (Rasmussen & Williams, 2005).
For instance, it provides a quantitative description of the generalization properties (Cohen et al., 2019;
Rahaman et al., 2018) and training dynamics (Jacot et al., 2018; Basri et al., 2019) of DNNs.
Roughly speaking, highly over-parameterized DNNs generalize well because they have a strong implicit bias to simple functions, and train well because low-error solutions in weight space can be reached by making a small change to the random values of the weights at initialization.
Despite its novelty and importance, the NTK correspondence suffers from a few shortcomings: (a) Its deterministic training is qualitatively different from the stochastic one used in practice, which may lead to poorer performance when combined with a small learning rate (Keskar et al., 2016).
(b) It under-performs, often by a large margin, convolutional neural networks (CNNs) trained with SGD (Arora et al., 2019).
(c) Deriving explicit ﬁnite width corrections (FWCs) is challenging, as it requires solving a set of coupled ODEs (Dyer & Gur-Ari, 2020; Huang & Yau, 2019).
Thus, there is a need for an extended theory for end-to-end trained deep networks which is valid for ﬁnite width DNNs.
Our contribution is three-fold.
First, we prove a correspondence between a DNN trained with noisy gradients and a Stochastic Process (SP) which at N → ∞ tends to the Neural Network Gaussian
Process (NNGP)
(Lee et al., 2018; Matthews et al., 2018).
In these works, the NNGP kernel is determined by the distribution of the DNN weights at initialization which are i.i.d.
random variables, whereas in our correspondence the weights are sampled across the stochastic training dynamics, drifting far away from their initial values.
We call ours the NNSP correspondence , and show that it holds when the training dynamics in output space exhibit ergodicity.
Second, we predict the outputs of trained ﬁnite-width DNNs, signiﬁcantly improving upon the corresponding GP predictions.
This is done by deriving leading FWCs which are found to scale with width as 1 /N .
The accuracy at which we can predict the empirical DNNs’ outputs serves as a strong veriﬁcation for our aforementioned ergodicity assumption.
In the regime where the GP RMSE error scales as 1 /n , we ﬁnd that the leading FWC are a decaying function of n , and thus overall negligible.
In the small n regime we ﬁnd that the FWC is small and grows with n .
We thus conclude that ﬁnite-width corrections are important for intermediate values of n (Fig. 1).
Third, we propose an explanation for why ﬁnite CNNs trained on image classiﬁcation tasks can outperform their inﬁnite-width counterparts, as observed by Novak et al.
The key difference is that in ﬁnite CNNs weight sharing is beneﬁcial.
Our theory, which accounts for the ﬁnite width, quantiﬁes this difference (§4.2).
Overall, the NNSP correspondence provides a rich analytical and numerical framework for exploring the theory of deep learning, unique in its ability to incorporate ﬁnite over-parameterization, stochasticity, and depth.
We note that there are several factors that make ﬁnite SGD-trained DNNs used in practice different from their GP counterparts, e.g. large learning rates, early stopping etc.
(Lee et al., 2020).
Importantly, our framework quantiﬁes the contribution of ﬁnite-width effects to this difference, distilling it from the contribution of these other factors.
1. R ELATED WORK
The idea of leveraging the dynamics of the gradient descent algorithm for approximating Bayesian inference has been considered in various works (Welling & Teh, 2011;
Mandt et al., 2017; Teh et al., 2016; Maddox et al., 2019; Ye et al., 2017).
However, to the best of our knowledge, a correspondence with a concrete SP or a non-parametric model was not established nor was a comparison made of the
DNN’s outputs with analytical predictions.
Finite width corrections were studied recently in the context of the NTK correspondence by several authors.
Hanin & Nica (2019) study the NTK of ﬁnite DNNs, but where the depth scales together with width, whereas we keep the depth ﬁxed.
Dyer & Gur-Ari (2020) obtained a ﬁnite N correction to the linear integral equation governing the evolution of the predictions on the training set.
Our work differs in several aspects: (a) We describe a different correspondence under different a training protocol with qualitatively different behavior.
(b) We derive relatively simple formulae for the outputs which become entirely explicit at large n .
(c) We account for all sources of ﬁnite N corrections whereas ﬁnite N NTK randomness remained an empirical source of corrections not accounted for by Dyer & Gur-Ari (2020).
(d) Our formalism differs considerably: its statistical mechanical nature enables one to import various standard tools for treating randomness, ergodicity breaking, and taking into account non-perturbative effects.
(e) We have no smoothness limitation on our activation functions and provide FWCs on a generic data point and not just on the training set.
Another recent paper (Yaida, 2020) studied Bayesian inference with weakly non-Gaussian priors induced by ﬁnite-
N DNNs.
Unlike here, there was no attempt to establish a correspondence with trained DNNs.
The formulation presented here has the conceptual advantage of representing a distribution over function space for arbitrary training and test data, rather than over speciﬁc draws of data sets.
This is useful for studying the large n behavior of learning curves, where analytical insights into generalization can be gained (Cohen et al., 2019).
A somewhat related line of work studied the mean ﬁeld regime of shallow NNs (Mei et al., 2018;
Chen et al., 2020; Tzen & Raginsky, 2020).
We point out the main differences from our work: (a) The NN output is scaled differently with width. (b) In the mean ﬁeld regime one is interested in the dynamics (ﬁnite t ) of the distribution over the NN parameters in the form of a PDE of the Fokker-Planck type. In contrast, in our framework we are interested in the distribution over function
space at equilibrium, i.e. for t → ∞ .
(c) It seems that the mean ﬁeld analysis is tailored for two-layer fully-connected NNs and is hard to generalize to deeper nets or to CNNs.
In contrast, our formalism generalizes to deeper fully-connected NNs and to CNNs as well, as we showed in section 4.2.
In this section we show that ﬁnite-width DNNs, trained in a speciﬁc manner, correspond to Bayesian inference using a non-parametric model which tends to the NNGP as N → ∞ .
We ﬁrst give a short review of Langevin dynamics in weight space as described by Neal et al.
Welling &
Teh (2011), which we use to generate samples from the posterior over weights.
We then shift our perspective and consider the corresponding distribution over functions induced by the DNN, which characterizes the non-parametric model.
Recap of Langevin-type dynamics - Consider a DNN trained with full-batch gradient descent while injecting white Gaussian noise and including a weight decay term, so that the discrete time dynamics of the weights read
[Equation]
where w t is the vector of all network weights at time step t , γ is the strength of the weight decay,
L ( z w ) is the loss as a function of the output z w , T is the temperature (the magnitude of noise), dt is the learning rate and ξ t ∼ N (0 , I ) .
As dt → 0 these discrete-time dynamics converge to the continuous-time Langevin equation given by ˙ w ( t ) =
−∇ w (cid:0) γ 2 || w ( t ) ||
2 + L ( z w ) (cid:1) + √ 2 Tξ
( t ) with (cid:104) ξ i ( t )
ξ j ( t (cid:48) )
(cid:105) =
δ ij δ ( t − t (cid:48) ) , so that as t → ∞ the weights will be sampled from the equilibrium distribution in weight space, given by (Risken & Frank, 1996)
[Equation]
The above equality holds since the equilibrium distribution of the Langevin dynamics is also the posterior distribution of a Bayesian neural network (BNN) with an i.i.d.
Gaussian prior on the weights w ∼ N (0 , σ 2 w I ) .
Thus we can map the hyper-parameters of the training to those of the
BNN: σ 2 w = T/γ and σ 2 = T/ 2 .
Notice that a sensible scaling for the weight variance at layer (cid:96) is σ 2 w,(cid:96) ∼ O
(1 /N (cid:96) − 1 ) , thus the weight decay needs to scale as γ (cid:96) ∼ O ( N (cid:96) − 1 ) .
A transition from weight space to function space - We aim to move from a distribution over weight space Eq.
2 to a one over function space.
Namely, we consider the distribution of z w ( x ) implied by the above P ( w ) where for concreteness we consider a DNN with a single scalar output z w ( x ) ∈ R on a regression task with data { ( x α , y α ) } nα =1 ⊂ R d × R .
Denoting by P [ f ] the induced measure on function space we formally write
[Equation]
(3) where (cid:82) dw denotes an integral over all weights and we denote by δ [ f − z w ] a delta-function in function-space.
As common in path-integrals or ﬁeld-theory formalism (Schulman, 2012), such a delta function is understood as a limit procedure where one chooses a suitable basis for function space, trims it to a ﬁnite subset, treats δ [ f − z w ] as a product of regular delta-functions, and at the end of the computation takes the size of the subset to inﬁnity.
To proceed we decompose the posterior over functions Eq.
3 as P [ f ] ∝ e − 12 σ 2 L [ f ] P
[ f ] where the prior over functions is P
[ f ] ∝ (cid:82) dwe − 12 σ 2 w || w ||
δ [ f − z w ] .
The integration over weights now obtains a clear meaning: it yields the distribution over functions induced by a DNN with i.i.d.
random weights chosen according to the prior P 0 ( w )
∝ e − 12 σ 2 w || w || 2 .
Thus, we can relate any correlation function in function space and weight space, for instance ( D f is the integration measure over function space)
[Equation]
As noted by Cho & Saul (2009), for highly over-parameterized DNNs the r.h.s.
of 4 equals the kernel of the NNGP associated with this DNN, K ( x, x (cid:48) ) .
Moreover P
[ f ] tends to a Gaussian and can be written as
[Equation]
where µ ( x ) is the measure of the input space, and the O (1 /N ) scaling of the ﬁnite- N correction will be explained in §3.
If we now plug 5 in 3, take the loss to be the total square error {+1 L [ f ] = (cid:80) nα =
+}
( y α − f ( x α )) 2 , and take N → ∞ we have that the posterior P [ f ] is that of a GP. Assuming ergodicity, one ﬁnds that training-time averaged output of the DNN is given by the posterior mean of a GP, with measurement noise 2 equal to σ 2 = T/ 2 and a kernel given by the NNGP of that DNN.
We refer to the above expressions for P
[ f ] and P [ f ] describing the distribution of outputs of a DNN trained according to our protocol – the NNSP correspondence .
Unlike the NTK correspondence, the kernel which appears here is different and no additional initialization dependent terms appear (as should be the case since we assumed ergodicity).
Furthermore, given knowledge of P 0 [ f ] at ﬁnite N , one can predict the DNN’s outputs at ﬁnite N .
Henceforth, we refer to P
[ f ] as the prior distribution, as it is the prior distribution of a DNN with random weights drawn from P 0 ( w ) .
Evidence supporting ergodicity - Our derivation relies on the ergodicity of the dynamics.
Ergodicity is in general hard to prove rigorously in non-convex settings, and thus we must revert to heuristics. The most robust evidence of ergodicity in function space is the high level of accuracy of our analytical expressions w.r.t. to our numerical results.
This is a self-consistency argument: we assume ergodicity in order to derive our analytical results and then indeed ﬁnd that they agree very well with the experiment, thus validating our original assumption.
Another indicator of ergodicity is a small auto-correlation time (ACT) of the dynamics.
Although short ACT does not logically imply ergodicity (in fact, the converse is true: exponentially long
ACT implies non-ergodic dynamics).
However, the empirical ACT gives a lower bound on the true correlation time of the dynamics.
In our framework, it is sufﬁcient that the dynamics of the outputs z w be ergodic, even if the dynamics of the weights converge much slower to an equilibrium distribution.
Indeed, we have found that the ACTs of the outputs are considerably smaller than those of the weights (see Fig. 2b).
Full ergodicity may be too strong of a condition and we don’t really need it for our purposes, since we are mainly interested in collecting statistics that will allow us to accurately compute the posterior mean of the distribution in function space.
Thus, a weaker condition which is sufﬁcient here is ergodicity in the mean (see App. F), and we believe our self-consistent argument above demonstrates that it holds.
In a related manner, optimizing the train loss can be seen as an attempt to ﬁnd a solution to n constraints using far more variables (roughly M · N 2 where M is the number of layers).
From a different angle, in a statistical mechanical description of satisﬁability problems, one typically expects ergodic behavior when the ratio of the number of variables to the number of constraints becomes much larger than one (Gardner & Derrida, 1988).
I NFERENCE ON THE RESULTING NNSP
Having mapped the time-averaged outputs of a DNN to inference on the above NNSP, we turn to analyze the predictions of this NNSP in the case where N is large but ﬁnite, such that the NNSP is only weakly non-Gaussian (i.e. its deviation from a GP is O (1 /N ) ).
The main result of this section is a derivation of leading FWCs to the standard GP regression results for the posterior mean ¯ f GP ( x ∗ ) and variance Σ GP ( x ∗ ) on an unseen test point x ∗ , given a training set { ( x α , y α ) } nα =1 ⊂ R d × R , namely (Rasmussen & Williams, 2005)
[Equation]
[Equation] where ˜ K αβ := K ( x α , x β ) + σ 2 δ αβ ; [Equation]
[Equation]
3. E DGEWORTH EXPANSION AND PERTURBATION THEORY
Our ﬁrst task is to ﬁnd how P [ f ] changes compared to the Gaussian ( N → ∞ ) scenario.
As the data-dependent part e −L [ f ] / 2 σ 2 is independent of the DNNs, this amounts to obtaining ﬁnite width corrections to the prior P
[ f ] .
One way to characterize this is to perform an Edgeworth expansion {+of P 
 [ f ]+} (Mccullagh, 2017; Sellentin et al., [-2017) of P 
-]2017).
We give a short recap of the Edgeworth expansion to elucidate our derivation, beginning with a scalar valued RV.
Consider continuous iid
RVs { Z i } and assume WLOG (cid:104) Z i (cid:105) = 0 , (cid:10) Z 2 i (cid:11) = 1 , with higher cumulants κ Zr for r ≥ 3 .
Now consider their normalized sum Y N = 1 √ N (cid:80) Ni
=1 Z i .
From additivity and homogeneity of cumulants we have κ
[Equation]
=
[Equation]
Zr N
[Equation]
Now, let ϕ ( y ) :
= (2 π ) − 1 / 2 e − y 2 / 2 .
The characteristic function of Y is ˆ f ( t ) : = F [ f ( y )] = exp (cid:16)(cid:80) ∞ r =1 κ r ( it ) r r ! (cid:17)
=
exp (cid:16)(cid:80) ∞
r =3 κ r ( it ) r r ! (cid:17) ˆ ϕ ( t ) .
Taking the inverse Fourier transform F − 1 has the effect of mapping it (cid:55)→ − ∂ y thus we get f ( y ) =
exp (cid:16)(cid:80) ∞
r =3 κ r ( − ∂ y ) r r ! (cid:17) [Equation] y ) (cid:1) where H r ( y ) is the r th Hermite polynomial , e.g. H 4 ( y ) = y 4 − 6 y 2 + 3 .
If we were to consider vector-valued RVs, then the r ’th order cumulant becomes a tensor with r indices, and the Hermite polynomials become multi-variate polynomials.
In our case, we are considering random functions deﬁned by our stochastic process (the
NNSP), thus the cumulants are functional tensors, i.e. are continuously indexed by the inputs x α .
This is especially convenient here since for all DNNs with the last layer being fully-connected, all odd cumulants vanish and the 2 r th cumulant scales as 1 /N r − 1 .
Consequently, at large N we can characterize P
[ f ] up to O ( N − 2 ) by its second and fourth cumulants, K ( x 1 , x 2 ) and U ( x 1 , x 2 , x 3 , x 4 ) , respectively.
Thus the leading order correction to P 0 [ f ] reads
[Equation]
(7) where the GP action S GP and the ﬁrst FWC action S U are given by
[Equation]
Here, H is the 4th functional Hermite polynomial (see App.
A), U is the 4th order functional cumulant of the NN output 3 , which depends on the choice of the activation function φ U
[Equation]
φ γ φ δ (cid:105) ) + 2 other perms . of ( α, β, γ, δ ) ∈ { 1 , . . . , 4 } (9) where φ α :=
φ ( z (cid:96) − 1 i ( x α )) and the pre-activations are z (cid:96)i ( x ) =
b (cid:96)i + (cid:80) N (cid:96) j =1 W ij φ ( z (cid:96) − 1 j ( x )) .
Here we distinguished between the scaled and non-scaled weight variances: σ 2 a = ς 2 a /N , where a are the weights of the last layer.
Our shorthand notation for the integration measure over inputs means e.g. dµ 1:4 :
= dµ ( x 1 ) · · · dµ ( x 4 ) .
Using perturbation theory, in App.
B we compute the leading FWC to the posterior mean ¯ f ( x ∗ ) and variance (cid:10) ( δf ( x ∗ )) 2 (cid:11) on a test point x ∗ ¯ f ( [Equation] = [Equation]
GP ( x
[Equation]
+ N [Equation] ¯ f U ( x [Equation]
+ O ( N [Equation] 2 ) (cid:10) ( δf ( x ∗ )) 2 (cid:11) = Σ GP ( x ∗ )
+ N − 1 Σ U ( x ∗ )
+ O ( N − 2 ) (10) with Σ U ( x ∗ ) =
(cid:10) ( f ( x ∗ )) 2 (cid:11) U − 2 ¯ f GP ( x ∗ ) ¯ f U ( x ∗ ) and
[Equation]
(11) where all repeating indices are summed over the training set (i.e. range over { 1 , . . . , n } ), denoting: ˜
y
α :
=
˜ K − 1 αβ y β , and deﬁning
[Equation]
Equations 11, 12 are one of our key analytical results, which are qualitatively different from the corresponding GP expressions Eq. 6.
The correction to the predictive mean ¯ f U ( x ∗ ) has a linear term in y , which can be viewed as a correction to the GP kernel, but also a cubic term in y , unlike ¯ f GP ( x ∗ ) which is purely linear.
The correction to the predictive variance Σ U ( x ∗ ) has quartic and quadratic terms in y , unlike Σ GP ( x ∗ ) which is y -independent.
˜
U ∗ α 1 α
α 3 has a clear interpretation in terms of GP regression: if we consider the indices
α 1 , α 2 , α 3 as ﬁxed, then U ∗ α 1 α 2 α 3 can be thought of as the ground truth value of a target function (analogous to y ∗ ), and the second term on the r.h.s.
U α
α 2 α 3 α 4 ˜ K − 1 α 4 β K ∗ β is then the GP prediction of U ∗ α
α 2 α 3 with the kernel K , where α 4 runs on the training set (compare to ¯ f GP ( x ∗ ) in Eq. 6).
Thus ˜ U ∗ α 1 α
α 3 is the discrepancy in predicting U
α
α 2 α 3 αusing a GP with kernel K .
In §3.2 we study the behavior of ¯ f U ( x ∗ ) as a function of n .
The posterior variance Σ( x ) = (cid:68) ( δf ( x )) 2 (cid:69) has a clear interpretation in our correspondence: it is a measure of how much we can decrease the test loss by averaging.
Our procedure for generating empirical network outputs involves time-averaging over the training dynamics after reaching equilibrium and also over different realizations of noise and initial conditions (see App. F).
This allows for a reliable comparison with our FWC theory for the mean.
In principle, one could use the network outputs at the end of training without this averaging, in which case there will be ﬂuctuations that will scale with Σ( x α ) .
Following this, one ﬁnds that the expected MSE test loss after training saturates is n −
[Equation] [Equation] where n ∗ is the size of the test set.
3. F INITE WIDTH CORRECTIONS FOR SMALL AND LARGE DATA SETS
The expressions in Eqs.
6, 11 for the GP prediction and the leading FWC are explicit but only up to a potentially large matrix inversion, ˜ K − 1 .
These matrices also have a random component related to the largely arbitrary choice of the particular n training points used to characterize the target function.
An insightful tool, used in the context of GPs, which solves both these issues is the Equivalent Kernel (EK) (Rasmussen & Williams, 2005; Sollich & Williams, 2004).
The EK approximates the GP predictions at large n , after averaging on all draws of (roughly) n training points representing the target function.
Even if one is interested in a particular dataset, the EK result captures the behavior of speciﬁc dataset up to small corrections.
Essentially, the discrete sums over the training set appearing in Eq.
6 are replaced by integrals over all input space, which together with a spectral decomposition of the kernel function K ( x, x (cid:48) )
= (cid:80) i λ i
ψ i ( x ) ψ
i ( x (cid:48) ) yields the well known result
[Equation]
Here we develop an extension of Eq. 13 for the NNSPs
we ﬁnd at large but ﬁnite N .
In particular, we ﬁnd the leading non-linear correction to the EK result, i.e. the "EK analogue" of Eq.
To this end, we consider the average predictions of an NNSP trained on an ensemble of data sets of size n (cid:48) , corresponding to n (cid:48) independent draws from a distribution µ ( x ) over all possible inputs x .
Following the steps in App.
J we ﬁnd
[Equation]
where an integral (cid:82) dµ ( x ) is implicit for every pair of repeated x coordinates.
We introduced the discrepancy operator ˜
δ xx (cid:48) which acts on some function ϕ as (cid:82) dµ ( x (cid:48) )
˜ δ xx (cid:48) ϕ
( x (cid:48) ) :
= ˜
δ xx (cid:48) ϕ
( x (cid:48) )
=
ϕ ( x ) − ¯ f EKGP ( x ) .
Essentially, Eq.
14 is derived from Eq.
11 by replacing each ˜ K − 1 by ( n/σ 2 )
˜ δ and noticing that in this regime ˜
U ∗ x 2 ,x 3 ,x 4 in Eq.
12 becomes ˜ δ x ∗ x 1 U x 1 ,x 2 ,x 3 ,x 4 .
Interestingly, ¯ f EK U ( x ∗ ) is written explicitly in terms of meaningful quantities: ˜ δ xx (cid:48) g
( x (cid:48) ) and ˜ δ x ∗ x 1 U x 1 ,x 2 ,x 3 ,x 4 .
Equations 13, 14 are valid for any weakly non-Gaussian process, including ones related to CNNs (where N corresponds to the number of channels).
It can also be systematically extended to smaller values of n by taking into account higher terms in 1 /n , as in Cohen et al.
At N → ∞ , we obtain the standard EK result, Eq. 13.
It is basically a high-pass linear ﬁlter which ﬁlters out features of g that have support on eigenfunctions ψ i associated with eigenvalues
λ i that are small relative to σ 2 /n .
We stress that the ψ i , λ i ’s are independent of any particular size n dataset but rather are a property of the average dataset.
In particular, no computationally costly data dependent matrix inversion is needed to evaluate Eq. 13.
Turning to our FWC result, Eq.
14, it depends on g ( x ) only via the discrepancy operator ˜
δ xx (cid:48) .
Thus these FWCs would be proportional to the error of the DNN, at N → ∞ .
In particular, perfect performance at N → ∞ , implies no FWC.
Second, the DNN’s average predictions act as a linear transformation on the target function combined with a cubic non-linearity.
Third, for g ( x ) having support only on some ﬁnite set of eigenfunctions ψ i of K , ˜ δ xx (cid:48) g ( x (cid:48) ) would scale as σ 2 /n at very large n .
Thus the above cubic term would lose its explicit dependence on n .
The scaling with n of this second term is less obvious, but numerical results suggest that ˜ δ x 2 x 3 also scales as σ 2 /n , so that the whole expression in the {· · · } has no scaling with n .
In addition, some decreasing behavior with n is expected due to the ˜ δ x ∗ x 1 U x 1 ,x 2 ,x 3 ,x 4 factor which can be viewed as the discrepancy in predicting U x,x 2 ,x 3 ,x 4 , at ﬁxed x 2 , x 3 , x 4 , based on n random samples ( x α ’s) of U x α ,x 2 ,x 3 ,x 4 .
In Fig.
1 we illustrate this behavior at large n and also ﬁnd that for small n the FWC is small but increasing with n , implying that at large N FWCs are only important at intermediate values of n .
[Figure]
T Ax = O (1) with x ∈ S d − 1 ( √ d ) so the number of parameters to be learnt is d ( d + 1) / 2 (vertical grey dashed line).
The GP discrepancy is monotonically decreasing with n whereas | ¯ f U ( x ∗ ) | increases linearly for small n (dashed-dotted lines in (a)) before it decays (best illustrated for larger d and σ 2 ).
For sufﬁciently large n , both the GP discrepancy and | ¯ f U ( x ∗ ) | scale as 1 /n (diagonal dashed black lines in (b),(c)).
This veriﬁes our prediction for the scaling of FWCs with n , Eq. 14 in the large n regime . Notably, it implies that at large N FWCs are only important at intermediate values of n .  N UMERICAL EXPERIMENTS
In this section we numerically test our analytical results.
We ﬁrst demonstrate that in the limit
N → ∞ the outputs of a FCN trained in the regime of the NNSP correspondence converge to a GP with a known kernel, and that the MSE between them scales as ∼ 1 /N 2 which is the scaling of the leading FWC squared.
Second, we show that introducing the leading FWC term N − 1 ¯ f U ( x ∗ ) , Eq. 11, further reduces this MSE by more than an order of magnitude.
Third, we study the performance gap between ﬁnite CNNs and their corresponding NNGPs on CIFAR-10.
4. T OY EXAMPLE : FULLY CONNECTED NETWORKS ON SYNTHETIC DATA
We trained a 2 -layer FCN f ( x ) = (cid:80) N i =1 a i φ ( w ( i ) · x ) on a quadratic target y ( x ) = x T Ax where the x ’s are sampled with a uniform measure from the hyper-sphere S d − 1 ( √ d ) , see App.
G.1 for more details.
Our settings are such that there are not enough training points to fully learn the target: Fig.
2a shows that the time averaged outputs (after reaching equilibrium)
¯ f DNN ( x ∗ ) is much closer to the GP prediction ¯ f GP ( x ∗ ) than to the ground truth y ∗ .
Otherwise, the convergence of the network output to the corresponding NNGP as N grows (shown in Fig.
2c) would be trivial, since all reasonable estimators would be close to the target and hence close to each other.
In Fig.
2c we plot in log-log scale (with base 10 ) the MSE (normalized by ( ¯ f DNN ) 2 ) between the predictions of the network ¯ f DNN and the corresponding GP and FWC predictions for quadratic and ReLU activations.
We ﬁnd that indeed for sufﬁciently large widths ( N (cid:38) 500 ) the slope of the GP-DNN MSE approaches − 2 .
(for both ReLU and quadratic), which is expected from our theory, since the leading FWC scales as 1 /N .
For smaller widths, higher order terms (in 1 /N ) in the
Edgeworth series Eq.
7 come into play.
For quadratic activation, we ﬁnd that our FWC result further reduces the MSE by more than an order of magnitude relative to the GP theory.
We recognize a regime where the GP and FWC MSEs intersect at N (cid:46) 100 , below which our FWC actually increases the MSE, which suggests a scale of how large N needs to be for our leading FWC theory to hold.
[Figure]
DNN MSE is further improved by more than an order of magnitude.
4. P ERFORMANCE GAP BETWEEN FINITE CNN S AND THEIR NNGP
Several papers have shown that the performance on image classiﬁcation tasks of SGD-trained ﬁnite
CNNs can surpass that of the corresponding GPs, be it NTK (Arora et al., 2019) or NNGP (Novak et al., 2018).
More recently, Lee et al.
(2020) emphasized that this performance gap depends on the procedure used to collapse the spatial dimensions of image-shaped data before the ﬁnal readout layer: ﬂattening the image into a one-dimensional vector (CNN-VEC) or applying global average pooling to the spatial dimensions (CNN-GAP).
It was observed that while inﬁnite FCN and CNNVEC outperform their respective ﬁnite networks, inﬁnite CNN-GAP networks under-perform their ﬁnite-width counterparts, i.e. there exists a ﬁnite optimal width.
One notable margin, of about 15% accuracy on CIFAR10, was shown in Novak et al.
(2018) for the case of CNN-GAP.
It was further pointed out there, that the NNGPs associated with CNN-VEC, coincide with those of the corresponding Locally Connected Networks (LCNs), namely CNNs without weight sharing between spatial locations.
Furthermore, the performance of SGD-trained
LCNs was found to be on par with that of their NNGPs.
We argue that our framework can account for this observation.
The priors P 0 [ f ] of a LCN and CNN-VEC agree on their second cumulant (the covariance), which is the only one not vanishing as N → ∞ , but they need not agree on their higher order cumulants, which come into play at ﬁnite N .
In App.
I we show that U appearing in our leading FWC, already differentiates between CNNs and LCNs.
Common practice strongly suggests that the prior over functions induced by CNNs is better suited than that of LCNs for classiﬁcation of natural images.
As a result we expect that the test loss of a ﬁnite-width CNN trained using our protocol will initially decrease with N but then increase beyond some optimal width N opt , tending towards the loss of the corresponding GP as N → ∞ .
This is in contrast to SGD behavior reported in some works where the CNN performance seems to saturate as a function of N , to some value better than the NNGP (Novak et al., 2018; Neyshabur et al., 2018).
Notably those works used maximum over architecture scans, high learning rates, and early stopping, all of which are absent from our training protocol.
To test the above conjecture we trained, according to our protocol, a CNN with six convolutional layers and two fully connected layers on CIFAR10, and used CNN-VEC for the readout.
We used MSE loss with a one-hot encoding into a 10 dimensional vector of the categorical label; further details and additional settings are given in App. G. Fig.
3 demonstrates that, using our training protocol, a ﬁnite CNN can outperform its corresponding GP and approaches its GP as the number of channels increases.
This phenomenon was observed in previous studies under realistic training settings (Novak et al., 2018), and here we show that it appears also under our training protocol.
We note that a similar yet more pronounced trend in performance appears here also when one considers the averaged MSE loss rather the the MSE loss of the average outputs.
C ONCLUSION
In this work we presented a correspondence between ﬁnite-width DNNs trained using Langevin dynamics (i.e. using small learning rates, weight-decay and noisy gradients) and inference on a stochastic-process (the NNSP), which approaches the NNGP as N → ∞ .
We derived ﬁnite width corrections, that improve upon the accuracy of the NNGP approximation for predicting the DNN outputs on unseen test points, as well as the expected ﬂuctuations around these.
In the limit of a large number of training points n → ∞ , explicit expressions for the DNNs’ outputs were given, involving no costly matrix inversions.
In this regime, the FWC can be written in terms of the discrepancy of GP predictions, so that when GP has a small test error the FWC will be small, and vice versa.
In the small n regime, the FWC is small but grows with n , which implies that at large N , FWCs are only important at intermediate values of n .
For no-pooling CNNs, we build on an observation made by Novak et al.
(2018) that ﬁnite CNNs outperform their corresponding NNGPs, and show that this is because the leading FWCs reﬂect the weight-sharing property of CNNs which is ignored at the level of the NNGP.
This constitutes one real-world example where the FWC is well suited to the structure of the data distribution, and thus improves performance relative to the corresponding GP.
In a future study, it would be very interesting to consider well controlled toy models that can elucidate under what conditions on the architecture and data distribution does the FWC improve performance relative to GP.
A E DGEWORTH SERIES
The Central Limit Theorem (CLT) tells us that the distribution of a sum of N independent RVs will tend to a Gaussian as N → ∞ .
Its relevancy for wide fully-connected DNNs (or CNNs with many channels) comes from the fact that every pre-activation averages over N uncorrelated random variables thereby generating a Gaussian distribution at large N (Cho & Saul, 2009), augmented by higher order cumulants which decay as 1 /N r/ 2 − 1 , where r is the order of the cumulant.
When higher order cumulants are small, an Edgeworth series (see e.g. Mccullagh (2017);
Sellentin et al.
(2017)) is a useful practical tool for obtaining the probability distribution from these cumulants.
Having the probability distribution and interpreting its logarithm as our action, places us closer to standard ﬁeld-theory formalism.
For simplicity we focus on a 2-layer network, but the derivation generalizes straightforwardly to networks of any depth.
We are interested in the ﬁnite N corrections to the prior distribution
P
[ f ] , i.e. the distribution of the DNN output f ( x ) = (cid:80)
Ni =
1 a i φ ( w T i x ) , with a i ∼ N (0 , ς 2 a N ) and w i ∼ N ( 0 , ς 2 w d I ) .
Because a has zero mean and a variance that scales as 1 /N , all odd cumulants are zero and the 2 r ’th cumulant scales as 1 /N r − 1 .
This holds true for any DNN having a fully-connected last layer with variance scaling as 1 /N .
The derivation of the multivariate Edgeworth series can be found in e.g. Mccullagh (2017); Sellentin et al.
and our case is similar where instead of a vector-valued RV we have the functional RV f ( x ) , so the cumulants become "functional tensors" i.e. multivariate functions of the input x .
Thus, the leading FWC to the prior P
[ f ] is
P
[ f ] = 1 Z e −
[Equation]
where S GP [ f ] is as in the main text Eq.
8 and the 4th Hermite functional tensor is
H
[Equation]
where by the integers in [ · ] we mean all possible combinations of this form, e
[Equation]
g.
[Equation]
H [ f ] is the functional analogue of the fourth Hermite polynomial: H 4 ( x ) =
x 4 − 6 x 2 + 3 , which appears in the scalar Edgeworth series expanded about a standard Gaussian.
B F IRST ORDER CORRECTION TO POSTERIOR MEAN AND VARIANCE B. P OSTERIOR MEAN
The posterior mean with the leading FWC action is given by
[Equation]
where
S [ f ] =
S GP [ f ] + S Data [ f ] + S U [ f ];
[Equation]
where the O (1 /N 2 ) implies that we only treat the ﬁrst order Taylor expansion of S [ f ] , and where
S GP [ f ] , S U [ f ] are as in the main text Eq. 8.
The general strategy is to bring the path integral (cid:82) D f to the front, so that we will get just correlation functions w.r.t.
the Gaussian theory (including the data term S Data [ f ] ) (cid:104)· · · (cid:105) 0 , namely the well known results (Rasmussen & Williams, 2005) for
¯ f GP ( x ∗ ) =
(cid:104) f
( x ∗ ) (cid:105) 0 and Σ GP ( x ∗ ) =
(cid:10) ( δf ( x ∗ )) 2 (cid:11) 0 , and then ﬁnally perform the integrals over input space.
Expanding both the numerator and the denominator of Eq.
B.1, the leading ﬁnite width correction for the posterior mean reads
[Equation]
This, as standard in ﬁeld theory, amounts to omitting all terms corresponding to bubble diagrams, namely we keep only terms with a factor of (cid:104) f
( x ∗ ) f
( x (cid:48)
α ) (cid:105)
and ignore terms with a factor of (cid:104)
f ( x ∗ ) (cid:105) 0 , since these will cancel out.
This is a standard result in perturbative ﬁeld theory (see e.g.
Zee (2003)).
We now write down the contributions of the quartic, quadratic and constant terms in H [ f ] : 1.
For the quartic term in H [ f ] , we have
[Equation]
We dub these terms by ¯ f ΣΣ ∗ and ¯ f ¯ f ¯ f Σ ∗ to be referenced shortly.
We mention here that they are the source of the linear and cubic terms in the target y appearing in Eq.
11 in the main text.
For the quadratic term in H [ f ] , we have
[Equation]
we note in passing that these cancel out exactly together with similar but opposite sign terms/diagrams in the quartic contribution, which is a reﬂection of measure invariance.
This is elaborated on in §B.3.
For the constant terms in H [ f ] , we will be left only with bubble diagram terms ∝ (cid:82)
D f f ( x ∗ ) which will cancel out in the leading order of 1 /N .
B. P OSTERIOR VARIANCE
The posterior variance is given by
[Equation]
Following similar steps as for the posterior mean, the leading ﬁnite width correction for the posterior second moment at x ∗ reads
[Equation]
As for the posterior mean
[Equation]
the constant terms in H [ f ] cancel out and the contributions of the quartic and quadratic terms are
[Equation]
and quadratic
[Equation]
B. M EASURE INVARIANCE OF THE RESULT
The expressions derived above may seem formidable, since they contain many terms and involve integrals over input space which seemingly depend on the measure µ ( x ) .
Here we show how they may in fact be simpliﬁed to the compact expressions in the main text Eq.
11 which involve only discrete sums over the training set and no integrals, and are thus manifestly measure-invariant.
For simplicity, we show here the derivation for the FWC of the mean ¯ f U ( x ∗ ) , and a similar derivation can be done for Σ U ( x ∗ ) .
In the following, we carry out the x integrals, by plugging in the expressions from Eq.
and coupling them to U .
As in the main text, we use the Einstein summation notation, i.e. repeated indices are summed over the training set.
The contribution of the quadratic terms is
[Equation]
where we deﬁned
[Equation]
Fortunately, this seemingly measure-dependent expression will cancel out with one of the terms coming from the ¯ f ΣΣ ∗ contribution of the quartic terms in H [ f ] .
This is not a coincidence and is a general feature of the Hermite polynomials appearing in the Edgeworth series, thus for any order in 1 /N in the Edgeworth series we will always be left only with measure invariant terms.
Collecting all terms that survive we have
[Equation]
where we deﬁned
[Equation] [Equation]
This is a more explicit form of the result reported in the main text, Eq. 11.
C F INITE WIDTH CORRECTIONS FOR MORE THAN ONE HIDDEN LAYER
For simplicity, consider a fully connected network with two hidden layers both of width N , and no biases, thus the pre-activations h ( x ) and output z ( x ) are given by
[Equation]
We want to ﬁnd the 2nd and 4th cumulants of z ( x ) .
Recall that we found that the leading order
Edgeworth expansion for the functional distribution of h is
[Equation]
where K − 1 ( x (cid:48) 1 , x (cid:48) 2 ) and U ( x (cid:48) 1 , x (cid:48) 2 , x (cid:48) 3 , x (cid:48) 4 ) are known from the previous layer.
So we are looking for two maps:
[Equation]
so that the mapping between the ﬁrst two cumulants K and U of two consequent layers is (assuming no biases)
[Equation]
where the starting point is the ﬁrst layer ( N (0) ≡ d
[Equation]
The important point to note, is that these functional integrals can be reduced to ordinary ﬁnite dimensional integrals.
For example, for the second layer, denote
[Equation]
we ﬁnd for K (2)
[Equation]
and for U (2) we denote h
[Equation]
so that
[Equation]
This iterative process can be repeated for an arbitrary number of layers.
D F OURTH CUMULANT FOR THRESHOLD POWER - LAW ACTIVATION FUNCTIONS D. F OURTH CUMULANT FOR R E LU ACTIVATION FUNCTION
The U ’s appearing in our FWC results can be derived for several activations functions, and in our numerical experiments we use a quadratic activation φ ( z ) =
z 2 and ReLU.
Here we give the result for ReLU, which is similar for any other threshold power law activation (see derivation in App. D.2), and give the result for quadratic activation in App.
E.
For simplicity, in this section we focus on the case of a 2-layer FCN with no biases, input dimension d and N neurons in the hidden layer, such that φ iα :=
φ ( w ( i ) · x α ) is the activation at the i th hidden unit with input x
α sampled with a uniform measure from S d − 1 ( √ d ) , where w ( i ) is a vector of weights of the ﬁrst layer.
This can be generalized to the more realistic settings of deeper nets and un-normalized inputs, where in the former the linear kernel L is replaced by the kernel of the layer preceding the output, and the latter amounts to introducing some scaling factors.
For φ = ReLU , (Cho & Saul, 2009) give a closed form expression for the kernel which corresponds to the GP.
Here we ﬁnd U corresponding to the leading FWC by ﬁrst ﬁnding the fourth moment of the hidden layer µ 4 := (cid:104) φ 1 φ 2 φ 3 φ 4 (cid:105) (see Eq. 9), taking for simplicity ς 2 w =µ 4 = (cid:112) det( L − 1 ) (2 π )
[Equation]
where L − 1 above corresponds to the matrix inverse of the 4 × 4 matrix with elements L αβ =
( x α · x β ) /d
which is the kernel of the previous layer (the linear kernel in the 2-layer case) evaluated on two random points.
In App.
D.2 we follow the derivation in Moran (1948), which yields (with a slight modiﬁcation noted therein) the following series in the off-diagonal elements of the matrix L
[Equation]
where the coefﬁcients A (cid:96)mnpqr are
[Equation]
For ReLU activation, these G ’s read
[Equation]
and similar expressions can be derived for other threshold power-law activations of the form φ ( z ) =
Θ( z ) z ν .
The series Eq.
D.2 is expected to converge for sufﬁciently large input dimension d since the overlap between random normalized inputs scales as O (1 / √ d ) and consequently L ( x, x (cid:48) ) ∼
O (1 / √ d ) for two random points from the data sets.
However, when we sum over U α 1 ...α 4 we also have terms with repeating indices and so L αβ ’s are equal to 1 .
The above Taylor expansion diverges whenever the 4 × 4 matrix L αβ − δ αβ has eigenvalues larger than 1 .
Notably this divergence does not reﬂect a true divergence of U , but rather the failure of representing it using the above expansion.
Therefore at large n , one can opt to neglect elements of U with repeating indices, since there are much fewer of these.
Alternatively this can be dealt with by a re-parameterization of the z ’s leading to a similar but slightly more involved Taylor series.
D. D ERIVATION OF THE PREVIOUS SUBSECTION
In this section we derive the expression for the fourth moment (cid:104)
f 1 f 2 f 3 f 4 (cid:105) of a two-layer fully connected network with threshold-power law activations with exponent ν : φ ( z ) =
Θ( z ) z ν ; ν =corresponds to a step function, ν = 1 corresponds to ReLU, ν = 2 corresponds to ReQU (rectiﬁed quadratic unit) and so forth.
When the inputs are normalized to lie on the hypersphere, the matrix L is
[Equation]
where the off diagonal elements here have L αβ
=
O (cid:16) 1 / √ d (cid:17) .
We follow the derivation in Ref.
Moran (1948), which computes the probability mass of the positive orthant for a quadrivariate
Gaussian distribution with covariance matrix L :
[Equation]
The characteristic function (Fourier transform) of this distribution is ϕ ( t 1 , t 2 , t 3 , t 4 ) =
exp (cid:18) −2 t T L t (cid:19) =
exp (cid:32) − 4 (cid:88)
α
[Equation]
t 2 α (cid:33) exp   − (cid:88)
α<β
L αβ t α t β 
 =
exp (cid:32) − 4 (cid:88)
α
= t 2 α (cid:33) ∞ (cid:88) (cid:96),m,n,p,q,r = ( − ) (cid:96) + m + n + p + q + r L (cid:96) 12 L m 13 L n 14 L p
L q 24 L r(cid:96) !
m !
n !
p !
q !
r !
t
[Equation]
+ m + nt (cid:96) + p
[Equation]
q
[Equation]
t m + p + rt n + q + r(D.7)
Performing an inverse Fourier transform, we may now write the positive orthant probability as
[Equation]
where the coefﬁcients A (cid:96)mnpqr are
[Equation]
and the one dimensional integral is
[Equation]
We can evaluate the integral over t to get
[Equation]
and performing the integral over z yields
[Equation]
We can now obtain the result for any integer ν by inserting z ν inside the z integral:
[Equation]
[Equation]
Using integration by parts we arrive at the result Eq.
D.4 reported in the main text
[Equation]
Similar expressions can be derived for other threshold power-law activations of the form φ ( z ) =
Θ( z ) z ν for arbitrary integer ν .
In a more realistic setting, the inputs x may not be perfectly normalized, in which case the diagonal elements of L are not unity.
It amounts to introducing a scaling factor for each of the four z ’s and makes the expressions a little less neat but poses no real obstacle.
E F OURTH CUMULANT FOR QUADRATIC ACTIVATION FUNCTION
For a two-layer network, we may write U , the 4th cumulant of the output f ( x ) = (cid:80) Ni =
1 a i φ ( w T i x ) , with a i ∼ N (0 , ς 2 a /N ) and w i ∼ N ( 0 , ( ς 2 w /d )
I ) for a general activation function φ as
[Equation] [Equation]
[Equation] [Equation]
with
[Equation]
For the case of a quadratic activation function φ ( z ) =
z 2 the V ’s read
[Equation]
where the linear kernel from the ﬁrst layer is L ( x, x (cid:48) )
=
ς
[Equation]
x · x (cid:48) .
Notice that we distinguish between the scaled and non-scaled variances:
[Equation]
These formulae were used when comparing the outputs of the empirical two-layer network with our
FWC theory Eq.
One can generalize them straightforwardly to a network with M layers by recursively computing K ( M − 1) the kernel in the ( M − 1) th layer (see e.g. Cho & Saul (2009)), and replacing L with K ( M − 1) .
F A UTO - CORRELATION TIME AND ERGODICITY
As mentioned in the main text, the network outputs ¯ f DNN ( x ∗ ) are a result of averaging across many realizations (seeds) of initial conditions and the noisy training dynamics, and across time (epochs) after the training loss levels off.
Our NNSP correspondence relies on the fact that our stochastic training dynamics are ergodic, namely that averages across time equal ensemble averages.
Actually, for our purposes it sufﬁces that the dynamics are ergodic in the mean , namely that the time-average estimate of the mean obtained from a single sample realization of the process converges in both the mean and in the mean-square sense to the ensemble mean:
[Equation]
where µ ( x ∗ ) is the ensemble mean on the test point x ∗ and the time-average estimate of the mean over a time window ˜ T is
[Equation]
This is hard to prove rigorously but we can do a numerical consistency check using the following procedure: Consider the time series of the network output on the test point x ∗ for the i ’th realization as a row vector and stack these row vectors for all different realizations into a matrix F , such that F ij = f DNN i ( x ∗ ; t j ) .
(1) Divide the time series data in the matrix F into non-overlapping sub-matrices, each of dimension n seeds × n epochs .
For each of these sub-matrices, ﬁnd ˆ f ( x ∗ ) i.e. the empirical dynamical average across that time window and across the chosen seeds; (2) Find the empirical variance σ 2emp ( x ∗ ) across these ˆ f ( x ∗ ) ; (4) Repeat (1)-(3) for other combinations of n epochs , n seeds .
If ergodicity holds, we should expect to see the following relation
[Equation]
where τ is the auto-correlation time of the outputs and σ 2 m is the macroscopic variance.
The results of this procedure are shown in Fig.
F.1, where we plot on a log-log scale the empirical variance σ 2emp vs. the number of epochs n epochs used for time averaging in each set (and using all 500 seeds in this case).
Performing a linear ﬁt on the average across test points (black x’s in the ﬁgure) yields a slope of approximately − 1 , which is strong evidence for ergodic dynamics.
G N UMERICAL EXPERIMENT DETAILS G. FCN EXPERIMENT DETAILS
We trained a 2 -layer FCN on a quadratic target y ( x ) = x T Ax where the x ’s are sampled with a uniform measure from the hyper-sphere S d − 1 ( √ d ) , with d = 16 and the matrix elements are sampled as A ij ∼ N (0 , 1) and ﬁxed for all x ’s. For both activation functions, we used a training noise level of σ 2 = 0 .
2 , training set of size n = 110 and a weight decay of the ﬁrst layer γ w = 0 .
Notice that for any activation φ , K scales linearly with ς 2 a = σ 2 a N =
( T/γ a ) · N , thus in order to keep K constant as we vary N we need to scale the weight decay of the last layer as γ a ∼ O ( N ) .
This is done in order to keep the prior distribution in accord with the typical values of the target as N varies, so that the comparison is fair.
We ran each experiment for 2 · 10 6 epochs, which includes the time it takes for the training loss to level off, which is usually on the order of 10 4 epochs.
In the main text we showed GP and FWC results for a learning rate of dt = 0 .
Here we report in Fig.
G.1 the results using dt ∈ { 0 . 003 , 0 .
For a learning rate of dt = 0 .
003 and width N ≥ 1000 the dynamics become unstable and strongly oscillate, thus the general trend is broken, as seen in the blue markers in Fig.
G.1.
The dynamics with the smaller learning rates are stable, and we see that there is a convergence to very similar values up to an expected statistical error.
G. CNN EXPERIMENT DETAILS AND ADDITIONAL SETTINGS
The CNN experiment reported in the main text was carried as follows.
Dataset: In the main text Fig.
3 we used a random sample of 10 train-points and 2000 test points from the CIFAR10 dataset, and in App.
H we report results on 1000 train-points and 1000 test points, balanced in terms of labels.
To use MSE loss, the ten categorical labels were one-hot encoded into vector of zeros and one.
Architecture: we used 6 convolutional layers with ReLU non-linearity, kernel of size 5 × 5 , stride of 1, no-padding, no-pooling.
The number of input channels was 3 for the input layer and C for the subsequent 5 CNN layers.
We then vectorized the outputs of the ﬁnal layer and fed it into an ReLU activated fully-connected layer with 25 C outputs, which were fed into a linear layer with 10 outputs corresponding to the ten categories.
The loss we used was MSE loss.
Training: Training was carried using full-batch SGD (GD) at varying learning-rates around 5 · 10 − 4 ,
Gaussian white noise was added to the gradients to generate σ 2 = 0 .
2 in the NNGP-correspondence, layer-dependant weight decay and bias decay which implies a (normalized by width) weight variance and bias variance of σ 2 w = 2 and σ 2 b = 1 respectively, when trained with no-data.
During training we saved, every 1000 epochs, the outputs of the CNN on every test point.
We note in passing that the standard deviation of the test outputs around their training-time-averaged value was about 0.1 per CNN output.
Training was carried for around half a million epochs which enabled us to reach a statistical error of about 2 · 10 − 4 , in estimating the Mean-Squared-Discrepancy between the training-time-averaged CNN outputs and our NNGP predictions.
Notably our best agreement between the DNN and GP occurred at 112 channels where the MSE was about 7 · 10 − 3 .
Notably the variance of the CNN (the average of its outputs squared) with no data, was about 25 .
Statistics.
To train our CNN within the regime of the NNSP correspondence, sufﬁcient training time (namely, epochs) was needed to get estimates of the average outputs ¯ f E ( x α ) =
¯ f ( x α ) + δf α since where
Notably, apart from just random noise when estimating the relative MSE between the averaged CNN outputs and the GP, a bias term appears equal to the variance of δf α averaged over all α ’s as indeed
(cid:88)
[Equation]
In all our experiments this bias was the dominant source of statistical error.
One can estimate it roughly given the number of uncorrelated samples taken into ¯ f E ( x α ) and correct the estimator.
We did not do so in the main text to make the data analysis more transparent.
Since the relative MSEs go down to 7 · 10 − 3 and the ﬂuctuations of the outputs quantiﬁed by Σ α = ( δf α ) 2 are of the order 0 .
1 2 , the amount of uncorrelated samples of CNN outputs we require should be much larger than 0 .
1 2 / (7 · 10 − 3 )
≈ 1 .
To estimate this bias in practice we repeated the experiment with 3different initialization seeds and deduced the bias from the variance of the results.
For comparison with NNGP (our DNN − GP plots)
the error bars were proportional to the variance of δf α .
For comparison with the target, we took much larger error bars equal to the uncertainty in estimating the expected loss from a test set of size 1000 .
These latter error bars where estimated empirically by measuring the variance across ten smaller test sets of size 100 .
Lastly we discarded the initial “burn-in" epochs, where the network has not yet reached equilibrium.
We took this burn-in time to be the time it takes the train-loss to reach within 5% of its stationary value at large times.
We estimated the stationary values by waiting until the DNNs train loss remained constant (up to trends much smaller than the ﬂuctuations) for about 5 · 10 5 epochs.
This also coincided well with having more or less stationary test loss.
Learning rate.
To be in the regime of the NNSP correspondence, the learning rate must be taken small enough such that discrepancy resulting from having discretization correction to the continuum
Langevin dynamics falls well below those coming from ﬁnite-width.
We ﬁnd that higher C require lower learning rates, potentially due to the weight decay term being large at large width.
In Fig.
G.2.
we report the relative MSE between the NNGP and CNN at learning rates of 0 .
001 , 0 .and
C = 48 showing good convergence already at 0 .
Following this we used learning rates of 0 .
0005 for C ≤ 48 and 0 .
00025 for C > 48 , in the main ﬁgure.
Comparison with the NNGP.
Following Novak et al.
(2018), we obtained the Kernel of our CNN.
Notably, since we did not have pooling layers this can be done straightforwardly without any approximations.
The NNGP predictions were then obtained in a standard manner (Rasmussen &
Williams, 2005).
H F URTHER NUMERICAL RESULTS ON CNN S
Here we report two additional numerical results following the CNN experiment we carried (for details see App. G).
Fig.
H.3b is the same as Fig.
H.3a apart from the fact that we subtracted our estimate of the statistical bias of our MSE estimator described in App.
G.
[Figure]
Concerning the experiment with 10 training points.
Here we used the same CNN as in the previous experiment.
The noise level was again the same and led to an effective σ 2 = 0 .
1 for the GP.
The weight decay on the biases was taken to be ten times larger leading to σ 2 b = 0 .
1 instead of σ b = 1 .as before
For C ≤ 80 we used a learning rate of dt
= 5 · 10 − 5 after verifying that reducing it further had no appreciable effect.
For C ≤ 80 we used dt = 2 . 5 · 10 − 5 .
For c ≤ 80 we used 6 · 10 +5 training epochs and we averaged over 4 different initialization seeds.
For C > 80 we used between 10 −different initialization seeds.
We reduced the aforementioned statistical bias in estimating the MSE from all our MSEs.
This bias, equal to the variance of the averaged outputs, was estimated based on our different seeds.
The error bars equal this estimated variance which was the dominant source of error.
I T HE FOURTH CUMULANT CAN DIFFERENTIATE CNN S FROM LCN S
Here we show that while the NNGP kernel K of a CNN without pooling cannot distinguish a CNN from an LCN, the fourth cumulant, U , can.
For simplicity let us consider the simplest CNN without pooling consisting of the following parts: (1) A 1D image with one color/channel ( X i ) as input i ∈ { 0 , . . .
, L − 1 } ; (2) A single convolutional layer with some activation φ acting with strideand no-padding using the conv-kernel T cx where c ∈ { 1 , . . .
, C } is a channel number index and x ∈ { 0 , . . .
, 2 l } is the relative position in the image.
Notably, in an LCN this conv-kernel will receive an additional dependence on ˜ x , the location on X i on which the kernel acts.
A vectorizing operation taking the C outputs of each convolutional around a point ˜ x ∈ { l, . . .
, L − l } , into a single index y ∈ { 0 , . . .
, C ( L − 2 l ) } .
A linear fully connected layer with weights W oc ˜ x where o ∈ { 0 , . . .
, #outputs } are the output indices.
Consider ﬁrst the NNGP of such a random DNN with weights chosen according to some iid Gaussian distribution P 0 ( w ) , with w including both W oc ˜ x and T cx .
Denoting by z o ( x ) the o ’th output of the
CNN, for an input x we have (where we denote in this section (cid:104)· · · (cid:105) :
= (cid:104)· · · (cid:105) P 0 ( w ) )
K oo (cid:48) ( x, x (cid:48) ) ≡ (cid:104)
z o ( x ) z o (cid:48)
( x (cid:48) ) (cid:105)
=
δ oo (cid:48) (cid:88)
[Equation] [Equation]
˜
[Equation]
˜
[Equation]
(cid:104) W
[Equation]
W o
[Equation]
[Table]
(I.1)
The NNGP kernel of an LCN is the same as that of a CNN.
This stems from the fact that (cid:104) W oc ˜ x W oc (cid:48)
˜ x (cid:48) (cid:105) yields a Kronecker delta function on the ˜ x, ˜ x (cid:48) indices.
Consequently, the difference between LCN and CNN, which amounts to whether T cx (˜ x ) is the same (CNN) or a different (LCN) random variable than T cx (cid:48) (cid:54) =
x (˜ x (cid:48) ) , becomes irrelevant as the these two are never averaged together.
For simplicity, we turn to the fourth cumulant of the same output, given by
[Equation]
with the second term on the LHS implying all pair-wise averages of z o ( x 1 ) ..
z o ( x 4 ) .
Note that the ﬁrst term on the LHS is not directly related to the kernel, thus it has a chance of differentiating a
CNN from an LCN.
Explicitly, it reads
[Equation]
The average over the four W ’s yields non-zero terms of the type W oc ˜ x W oc ˜ x W oc (cid:48)
˜ x (cid:48) W oc (cid:48)
˜ x (cid:48) with either ˜ x
= ˜ x (cid:48) (type 1), ˜ x (cid:54) =
˜ x (cid:48) and c (cid:54) =
c (cid:48) (type 2), or ˜ x (cid:54) =
˜ x (cid:48) and c = c (cid:48) (type 3).
The type 1 contribution cannot differentiate an LCN form a CNN since, as in the NNGP case, they always involve only one ˜ x .
The type 2 contribution also cannot differentiate since it yields
[Equation]
(I.4)
Examining the average involving the four T ’s, one ﬁnds that since T cx (˜ x ) is uncorrelated with T c (cid:48)
x (cid:48) (˜ x (cid:48) ) for both LCNs and CNNs, it splits into (cid:88)
[Equation]
where as in the NNGP, two T ’s with different ˜ x are never averaged together and we only get a contribution proportional to products of two K ’s. We note in passing that these type 2 terms yield a contribution that largely cancels that of K ( x α , x β ) K ( x γ , x δ )[3] , apart from a “diagonal" contribution ( ˜ x = ˜ x
(cid:48) )
[Equation]
We turn our attention to the type 3 term given by
[Equation]
Examining the average involving the four T ’s, one now ﬁnds a sharp difference between an LCN and a CNN.
For an LCN, this average would split into a product of two K ’s since T cx (˜ x ) would be uncorrelated with T c x (˜ x (cid:48) ) .
For a CNN however, T c x (˜ x ) is the same random variable as T c x (˜ x (cid:48) ) and therefore the average does not split giving rise to a distinct contribution that differentiates a CNN from an LCN.
Notably, it is small by a factor of 1 /C
owing to the fact that it contains a redundant summation over one c -index while the averages over the four W ’s contain a 1 /C 2 factor when properly normalized.
J C ORRECTIONS TO EK
Here we derive ﬁnite- N correction to the Equivalent Kernel result.
Using the tools developed by
Cohen et al.
(2019), the replicated partition function relevant for estimating the predictions of the network ( f ( x ∗ ) ) averaged ( (cid:104)· · · (cid:105) n ) over all draws of datasets of size n (cid:48) with n (cid:48) taken from a Poisson distribution with mean n is given by
[Equation]
with S GP [ f ] and S U [ f ] given in Eq. 8.
We comment that the above expression is only valid for obtaining the leading order asymptotics in n .
Enabling generic n requires introducing replicas explicitly (see Cohen et al.
Notably, the above expression coincides with that used for a ﬁnite dataset, with two main differences: all the sums over the training set have been replaced by integrals with respect to the measure, µ x , from which data points are drawn.
Furthermore σis now accompanied by n .
Following this, all the diagrammatic and combinatorial aspects shown in the derivation for a ﬁnite dataset hold here as well.
For instance, let us examine a speciﬁc contribution coming from the quartic term in H [ f ] : U x 1 ..x 4 K − 1 x 1 x (cid:48) 1 · · · K − 1 x 4 x (cid:48) 4 f
( x (cid:48) 1 ) · · · f ( x (cid:48) 4 ) , and from the diagram/Wick-contraction where we take the expectation value of 3 out of the 4 f ’s in this quartic term, to arrive at an expression which is ultimately cubic in the targets y
[Equation]
where we recall that (cid:104) f
( x ) (cid:105) ∞
[Equation]
and Σ
[Equation]
K
[Equation]
x 2 being the posterior covariance in the EK limit, where ˜ K xx (cid:48)
f ( x (cid:48) )
=
K xx (cid:48)
f ( x (cid:48) )
+ ( σ 2 /n )
f ( x ) .
Using the fact that K − 1 xx (cid:48)
K x (cid:48) x
(cid:48)(cid:48) gives a delta function w.r.t.
the measure, the integrals against K − 1 x α x (cid:48)
α can be easily carried out yielding
[Equation]
Introducing the discrepancy operator ˜
[Equation]
xx
[Equation]
xx
[Equation]
− K xx
[Equation]
xx
[Equation]
, we can write a more compact expression
[Equation]
This with the additional 1 / 4!
factor times the combinatorial factor of 4 related to choosing the "partner" of f ( x ∗ ) in the Wick contraction, yields an overall factor of 1 / 6 as in the main text, Eq.