Year: 2020

Abstract:
A recent line of works studied wide deep neural networks (DNNs) by approximating them as Gaussian Processes (GPs). A DNN trained with gradient flow 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 finite 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 finite width. Our contribution is three-fold: (i) In the infinite 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 finite width correction (FWC) for DNNs with arbitrary activation functions and depth and use it to predict the outputs of empirical finite networks with high accuracy. Analyzing the FWC behavior as a function of n, the training set size, we find 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 flesh-out algebraically how these FWCs can improve the performance of finite convolutional neural networks (CNNs) relative to their GP counterparts on image classification tasks.

heading: 1 INTRODUCTION
text: 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 flow, namely with vanishing learning rate and involving no stochasticity.
The NTK result has provided the first 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 finite 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 finite 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 finite-width DNNs, significantly 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 verification for our aforementioned ergodicity assumption. In the regime where the GP RMSE error scales as 1/n, we find that the leading FWC are a decaying function of n, and thus overall negligible. In the small n regime we find that the FWC is small and grows with n. We thus conclude that finite-width corrections are important for intermediate values of n (Fig. 1).
Third, we propose an explanation for why finite CNNs trained on image classification tasks can outperform their infinite-width counterparts, as observed by Novak et al. (2018). The key difference is that in finite CNNs weight sharing is beneficial. Our theory, which accounts for the finite width, quantifies 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 finite over-parameterization, stochasticity, and depth. We note that there are several factors that make finite 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 quantifies the contribution of finite-width effects to this difference, distilling it from the contribution of these other factors.

heading: 1.1 RELATED WORK
text: 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 finite DNNs, but where the depth scales together with width, whereas we keep the depth fixed. Dyer & Gur-Ari (2020) obtained a finite 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 finite N corrections whereas finite 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 finite-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 specific 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 field 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 field regime one is interested in the dynamics (finite 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 field 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.

heading: 2 THE NNSP CORRESPONDENCE
text: In this section we show that finite-width DNNs, trained in a specific manner, correspond to Bayesian inference using a non-parametric model which tends to the NNGP as N → ∞. We first give a short review of Langevin dynamics in weight space as described by Neal et al. (2011), 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
∆wt := wt+1 − wt = − (γwt +∇wL (zw)) dt+ √ 2Tdtξt (1)
where wt is the vector of all network weights at time step t, γ is the strength of the weight decay, L(zw) is the loss as a function of the output zw, 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 ( γ 2 ||w(t)|| 2 + L (zw) ) + √
2Tξ (t) with 〈ξi(t)ξj(t′)〉 = δijδ (t− t′), so that as t → ∞ the weights will be sampled from the equilibrium distribution in weight space, given by (Risken & Frank, 1996)
P (w) ∝ exp ( − 1 T (γ 2 ||w||2 + L (zw) )) = exp ( − ( 1 2σ2w ||w||2 + 1 2σ2 L (zw) )) (2)
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, σ2wI). Thus we can map the hyper-parameters of the training to those of the BNN: σ2w = T/γ and σ
2 = T/2. Notice that a sensible scaling for the weight variance at layer ` is σ2w,` ∼ O(1/N`−1), thus the weight decay needs to scale as γ` ∼ O(N`−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 zw(x) implied by the above P (w) where for concreteness we consider a DNN with a single scalar output zw(x) ∈ R on a regression task with data {(xα, yα)}nα=1 ⊂ Rd × R. Denoting by P [f ] the induced measure on function space we formally write
P [f ] = ∫ dwδ[f − zw]P (w) ∝ e− 1 2σ2 L[f ] ∫ dwe − 1 2σ2w ||w||2 δ[f − zw] (3)
where ∫ dw denotes an integral over all weights and we denote by δ[f − zw] a delta-function in function-space. As common in path-integrals or field-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 finite subset, treats δ[f − zw] as a product of regular delta-functions, and at the end of the computation takes the size of the subset to infinity.
To proceed we decompose the posterior over functions Eq. 3 as P [f ] ∝ e− 1 2σ2 L[f ]P0[f ] where the prior over functions is P0[f ] ∝ ∫ dwe − 1 2σ2w ||w||2
δ[f − zw]. 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 P0(w) ∝ e − 1 2σ2w ||w||2
. Thus, we can relate any correlation function in function space and weight space, for instance (Df is the integration measure over function space)∫ DfP0[f ]f(x)f(x′) = ∫ Df ∫ dwP0(w)δ[f−zw]f(x)f(x′) = ∫ dwP0(w)zw(x)zw(x ′) (4)
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′). Moreover P0[f ] tends to a Gaussian and can be
written as
P0[f ] ∝ exp ( −1
2
∫ dµ(x)dµ(x′)f(x)K−1(x, x′)f(x′) ) +O (1/N) (5)
where µ(x) is the measure of the input space, and the O(1/N) scaling of the finite-N correction will be explained in §3. If we now plug 5 in 3, take the loss to be the total square error1 L[f ] =∑n
α=1 (yα − f (xα)) 2, and take N →∞ we have that the posterior P [f ] is that of a GP. Assuming ergodicity, one finds that training-time averaged output of the DNN is given by the posterior mean of a GP, with measurement noise2 equal to σ2 = T/2 and a kernel given by the NNGP of that DNN.
We refer to the above expressions for P0[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 P0[f ] at finite N , one can predict the DNN’s outputs at finite N . Henceforth, we refer to P0[f ] as the prior distribution, as it is the prior distribution of a DNN with random weights drawn from P0(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 find 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 sufficient that the dynamics of the outputs zw 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 sufficient 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 find a solution to n constraints using far more variables (roughly M · N2 where M is the number of layers). From a different angle, in a statistical mechanical description of satisfiability 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).

heading: 3 INFERENCE ON THE RESULTING NNSP
text: 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 finite, 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 ⊂ Rd × R, namely (Rasmussen & Williams, 2005)
f̄GP(x∗) = ∑ α,β yαK̃ −1 αβK ∗ β ; ΣGP(x∗) = K ∗∗ − ∑ α,β K∗αK̃ −1 αβK ∗ β (6)
where K̃αβ := K(xα, xβ) + σ2δαβ ; K∗α := K(x∗, xα); K ∗∗ := K(x∗, x∗).
1We take the total error, i.e. we don’t divide by n so that L[f ] becomes more dominant for larger n. 2Here σ2 is a property of the training protocol and not of the data itself, or our prior on it.

heading: 3.1 EDGEWORTH EXPANSION AND PERTURBATION THEORY
text: Our first task is to find 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 finite width corrections to the prior P0[f ]. One way to characterize this is to perform an Edgeworth expansion of P0[f ] (Mccullagh, 2017; Sellentin et al., 2017). We give a short recap of the Edgeworth expansion to elucidate our derivation, beginning with a scalar valued RV. Consider continuous iid RVs {Zi} and assume WLOG 〈Zi〉 = 0, 〈 Z2i 〉
= 1, with higher cumulants κZr for r ≥ 3. Now consider their normalized sum YN = 1√N ∑N i=1 Zi. From additivity and homogeneity of cumulants we have κr≥2 := κYr≥2 = NκZr ( √ N)r = κZr Nr/2−1 . Now, let ϕ(y) := (2π)−1/2e−y 2/2. The charac-
teristic function of Y is f̂(t) := F [f(y)] = exp (∑∞
r=1 κr (it)r r!
) = exp (∑∞ r=3 κr (it)r r! ) ϕ̂(t).
Taking the inverse Fourier transform F−1 has the effect of mapping it 7→ −∂y thus we get f(y) = exp (∑∞ r=3 κr (−∂y)r r! ) ϕ(y) = ϕ(y) ( 1 + ∑∞ r=3 κr r!Hr(y) ) where Hr(y) is the rth Hermite polynomial, e.g. H4(y) = y4− 6y2 + 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 defined 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 2rth cumulant scales as 1/Nr−1. Consequently, at large N we can characterize P0[f ] up to O(N−2) by its second and fourth cumulants, K(x1, x2) and U(x1, x2, x3, x4), respectively. Thus the leading order correction to P0[f ] reads
P0 [f ] ∝ e−SGP[f ] (
1− 1 N SU [f ]
) +O ( 1/N2 ) (7)
where the GP action SGP and the first FWC action SU are given by
SGP[f ] = 1
2
∫ dµ1:2fx1K −1 x1,x2fx2 ; SU [f ] = − 1
4!
∫ dµ1:4Ux1,x2,x3,x4Hx1,x2,x3,x4 [f ] (8)
Here, H is the 4th functional Hermite polynomial (see App. A), U is the 4th order functional cumulant of the NN output3, which depends on the choice of the activation function φ Ux1,x2,x3,x4 = ς 4 a (〈φαφβφγφδ〉 − 〈φαφβ〉 〈φγφδ〉) + 2 other perms. of (α, β, γ, δ) ∈ {1, . . . , 4} (9) where φα := φ(z`−1i (xα)) and the pre-activations are z ` i (x) = b ` i + ∑N` j=1Wijφ(z `−1 j (x)). Here we distinguished between the scaled and non-scaled weight variances: σ2a = ς 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µ(x1) · · · dµ(x4). Using perturbation theory, in App. B we compute the leading FWC to the posterior mean f̄(x∗) and variance 〈 (δf(x∗)) 2 〉
on a test point x∗ f̄(x∗) = f̄GP(x∗) +N
−1f̄U (x∗) +O(N−2)〈 (δf(x∗)) 2 〉 = ΣGP(x∗) +N −1ΣU (x∗) +O(N−2)
(10)
with ΣU (x∗) = 〈 (f(x∗)) 2 〉 U − 2f̄GP(x∗)f̄U (x∗) and
f̄U (x∗) = 1
6 Ũ∗α1α2α3
( ỹα1 ỹα2 ỹα3 − 3K̃−1α1α2 ỹα3 ) 〈 (f(x∗)) 2 〉 U = 1 2 Ũ∗∗α1α2 ( ỹα1 ỹα2 − K̃−1α1α2
) (11) where all repeating indices are summed over the training set (i.e. range over {1, . . . , n}), denoting: ỹα := K̃ −1 αβ yβ , and defining
Ũ∗α1α2α3 := U ∗ α1α2α3 − Uα1α2α3α4K̃ −1 α4β K∗β
Ũ∗∗α1α2 := U ∗∗ α1α2 − ( U∗α1α2α3 + Ũ ∗ α1α2α3 ) K̃−1α3βK ∗ β
(12)
3Here we take U ∼ O(1) to emphasize the scaling with N in Eqs. 7, 10.
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α2α3 has a clear interpretation in terms of GP regression: if we consider the indices α1, α2, α3 as fixed, 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α1α2α3α4K̃ −1 α4β K∗β is then the GP prediction of U ∗ α1α2α3 with the kernel K, where α4 runs on the training set (compare to f̄GP(x∗) in Eq. 6). Thus Ũ∗α1α2α3 is the discrepancy in predicting Uα1α2α3α4 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) = 〈 (δf (x)) 2 〉
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 fluctuations that will scale with Σ(xα). Following this, one finds that the expected MSE test loss after training saturates is n−1∗ ∑n∗ α=1 (〈( f̄ (xα)− y (xα) )2〉 + Σ(xα) ) where n∗ is the size of the test set.

heading: 3.2 FINITE WIDTH CORRECTIONS FOR SMALL AND LARGE DATA SETS
text: 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 specific 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′) = ∑ i λiψi(x)ψi(x ′) yields the well known result
f̄EKGP (x∗) =
∫ dµ(x′) ∑ i λiψi(x∗)ψi(x ′) λi + σ2/n g(x′) (13)
Here we develop an extension of Eq. 13 for the NNSPs we find at large but finite N . In particular, we find the leading non-linear correction to the EK result, i.e. the "EK analogue" of Eq. 11. To this end, we consider the average predictions of an NNSP trained on an ensemble of data sets of size n′, corresponding to n′ independent draws from a distribution µ(x) over all possible inputs x. Following the steps in App. J we find
f̄EKU (x∗) = 1
6 δ̃x∗x1Ux1,x2,x3,x4
{ n3
σ6 δ̃x2x′2g(x
′ 2)δ̃x3x′3g(x ′ 3)δ̃x4x′4g(x ′ 4)−
3n2
σ4 δ̃x2,x3 δ̃x4,x′4g(x
′ 4) } (14)
where an integral ∫ dµ(x) is implicit for every pair of repeated x coordinates. We introduced the
discrepancy operator δ̃xx′ which acts on some function ϕ as ∫ dµ(x′)δ̃xx′ϕ(x ′) := δ̃xx′ϕ(x ′) = ϕ(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̃∗x2,x3,x4 in Eq. 12 becomes δ̃x∗x1Ux1,x2,x3,x4 . Interestingly, f̄ EK U (x∗) is written explicitly in terms of meaningful quantities: δ̃xx′g(x′) and δ̃x∗x1Ux1,x2,x3,x4 .
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. (2019). At N →∞, we obtain the standard EK result, Eq. 13. It is basically a high-pass linear filter which filters 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′ . 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 finite set of eigenfunctions ψi of K, δ̃xx′g(x′) 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 δ̃x2x3 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∗x1Ux1,x2,x3,x4 factor which can be viewed as the discrepancy in predicting Ux,x2,x3,x4 , at fixed x2, x3, x4, based on n random samples (xα’s) of Uxα,x2,x3,x4 . In Fig. 1 we illustrate this behavior at large n and also find 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.

heading: 4 NUMERICAL EXPERIMENTS
text: In this section we numerically test our analytical results. We first 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/N2 which is the scaling of the leading FWC squared. Second, we show that introducing the leading FWC term N−1f̄U (x∗), Eq. 11, further reduces this MSE by more than an order of magnitude. Third, we study the performance gap between finite CNNs and their corresponding NNGPs on CIFAR-10.

heading: 4.1 TOY EXAMPLE: FULLY CONNECTED NETWORKS ON SYNTHETIC DATA
text: We trained a 2-layer FCN f(x) = ∑N i=1 aiφ(w
(i) · x) on a quadratic target y(x) = xTAx where the x’s are sampled with a uniform measure from the hyper-sphere Sd−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 find that indeed for sufficiently large widths (N & 500) the slope of the GP-DNN MSE approaches −2.0 (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 find 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 atN . 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.

heading: 4.2 PERFORMANCE GAP BETWEEN FINITE CNNS AND THEIR NNGP
text: Several papers have shown that the performance on image classification tasks of SGD-trained finite 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 final readout layer: flattening 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 infinite FCN and CNNVEC outperform their respective finite networks, infinite CNN-GAP networks under-perform their finite-width counterparts, i.e. there exists a finite 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 P0[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 finite 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 classification of natural images. As a result we expect that the test loss of a finite-width CNN trained using our protocol will initially decrease with N but then increase beyond some optimal width Nopt, 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 finite 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.

heading: 5 CONCLUSION
text: In this work we presented a correspondence between finite-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 finite 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 fluctuations 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 finite CNNs outperform their corresponding NNGPs, and show that this is because the leading FWCs reflect 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.

heading: A EDGEWORTH SERIES
text: 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/Nr/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 field-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 finite N corrections to the prior distribution P0[f ], i.e. the distribution of the DNN output f(x) = ∑N i=1 aiφ(w T i x), with ai ∼ N (0, ς2a N ) and wi ∼ 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 2r’th cumulant scales as 1/Nr−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. (2017), 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 P0[f ] is
P0 [f ] = 1
Z e−SGP[f ]
[ 1 + 1
4!
∫ dµ (x1) · · · dµ (x4)U (x1, x2, x3, x4)H [f ; x1, x2, x3, x4] ] +O(1/N2)
(A.1)
where SGP[f ] is as in the main text Eq. 8 and the 4th Hermite functional tensor is
H [f ] = ∫ dµ (x′1) · · · dµ (x′4)K−1 (x1, x′1) · · ·K−1 (x4, x′4) f (x′1) · · · f (x′4)
−K−1 (xα, xβ) ∫ dµ ( x′µ ) dµ (x′ν)K −1 (xµ, x′µ)K−1 (xν , x′ν) f (x′µ) f (x′ν) [6] (A.2) +K−1 (xα, xβ)K −1 (xµ, xν) [3]
where by the integers in [·] we mean all possible combinations of this form, e.g.
K−1αβK −1 µν = K −1 12 K −1 34 +K −1 13 K −1 24 +K −1 14 K −1 23 (A.3)
H[f ] is the functional analogue of the fourth Hermite polynomial: H4 (x) = x4 − 6x2 + 3, which appears in the scalar Edgeworth series expanded about a standard Gaussian.

heading: B FIRST ORDER CORRECTION TO POSTERIOR MEAN AND VARIANCE
text: B.1 POSTERIOR MEAN
The posterior mean with the leading FWC action is given by 〈f (x∗)〉 = ∫ Dfe−S[f ]f (x∗)∫ Dfe−S[f ] +O(1/N2) (B.1)
where
S[f ] = SGP[f ] + SData[f ] + SU [f ]; SData[f ] = 1
2σ2 n∑ α=1 (f (xα)− yα)2 (B.2)
where the O(1/N2) implies that we only treat the first order Taylor expansion of S[f ], and where SGP[f ], SU [f ] are as in the main text Eq. 8. The general strategy is to bring the path integral ∫ Df to the front, so that we will get just correlation functions w.r.t. the Gaussian theory (including the data term SData[f ]) 〈· · · 〉0, namely the well known results (Rasmussen & Williams, 2005) for
f̄GP(x∗) = 〈f(x∗)〉0 and ΣGP(x∗) = 〈 (δf(x∗)) 2 〉 0 , and then finally perform the integrals over input space. Expanding both the numerator and the denominator of Eq. B.1, the leading finite width correction for the posterior mean reads
f̄U (x∗) = 1
4!
(∫ dµ1:4U (x1, x2, x3, x4) 〈f (x∗)H [f ]〉0 − 〈f (x∗)〉0 ∫ dµ1:4U (x1, x2, x3, x4) 〈H [f ]〉0 ) (B.3)
This, as standard in field theory, amounts to omitting all terms corresponding to bubble diagrams, namely we keep only terms with a factor of 〈f (x∗) f (x′α)〉0 and ignore terms with a factor of 〈f (x∗)〉0 , since these will cancel out. This is a standard result in perturbative field 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
〈f (x∗) f (x′1) f (x′2) f (x′3) f (x′4)〉0 − 〈f (x∗)〉0 〈f (x ′ 1) f (x ′ 2) f (x ′ 3) f (x ′ 4)〉0 = Σ (x∗, x ′ α) Σ ( x′β , x ′ γ ) f̄ (x′δ) [12] + Σ (x∗, x ′ α) f̄ ( x′β ) f̄ ( x′γ ) f̄ (x′δ) [4] (B.4)
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.
2. For the quadratic term in H [f ], we have〈 f (x∗) f ( x′µ ) f (x′ν) 〉 0 − 〈f (x∗)〉0 〈 f ( x′µ ) f (x′ν) 〉 0 = Σ ( x∗, x ′ µ ) f̄ (x′ν) [2] (B.5)
we note in passing that these cancel out exactly together with similar but opposite sign terms/diagrams in the quartic contribution, which is a reflection of measure invariance. This is elaborated on in §B.3. 3. For the constant terms in H [f ], we will be left only with bubble diagram terms ∝∫ Df f (x∗) which will cancel out in the leading order of 1/N .
B.2 POSTERIOR VARIANCE
The posterior variance is given by
Σ(x∗) = 〈f (x∗) f (x∗)〉 − f̄2
= 〈f (x∗) f (x∗)〉0 + 〈f (x∗) f (x∗)〉U − f̄ 2 GP − 2f̄GPf̄U +O(1/N2) = ΣGP(x∗) + 〈f (x∗) f (x∗)〉U − 2f̄GPf̄U +O(1/N 2)
(B.6)
Following similar steps as for the posterior mean, the leading finite width correction for the posterior second moment at x∗ reads
〈f (x∗) f (x∗)〉U = 1
4!
(∫ dµ1:4U (x1, x2, x3, x4) 〈f (x∗) f (x∗)H [f ]〉0 − 〈f (x∗) f (x∗)〉0 ∫ dµ1:4U (x1, x2, x3, x4) 〈H [f ]〉0 ) (B.7)
As for the posterior mean, the constant terms in H[f ] cancel out and the contributions of the quartic and quadratic terms are
quartic terms = Σ∗αΣ∗β f̄γ f̄δ [12] + Σ∗αΣ∗βΣγδ [12] (B.8)
and quadratic terms = Σ∗µΣ∗ν [2] (B.9)
B.3 MEASURE 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 simplified 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. 6 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
Aα1,∗K̃ −1 α1β1 yβ1 −Aα1α2K̃−1α1β1K̃ −1 α2β2 yβ1Kβ2,∗ (B.10)
where we defined
A (x3, x4) := ∫∫ dµ(x1)dµ(x2)U (x1, x2, x3, x4)K −1 (x1, x2) (B.11)
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
1
4!
{ 4Ũ∗α1α2α3K̃ −1 α1β1 K̃−1α2β2K̃ −1 α3β3 yβ1yβ2yβ3 − 12Ũ∗α1α2α3K̃ −1 α2β2 K̃−1α1β1yβ1 } (B.12)
where we defined Ũ∗α1α2α3 := U ∗ α1α2α3 − Uα1α2α3α4K̃ −1 α4β4
K∗β4 (B.13) This is a more explicit form of the result reported in the main text, Eq. 11.

heading: C FINITE WIDTH CORRECTIONS FOR MORE THAN ONE HIDDEN LAYER
text: 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
h (x) = σw2√ N W (2)φ ( σw1√ d W (1)x ) z (x) =
σa√ N aTφ
( h(2) (x) ) (C.1) We want to find 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
PK,U [h] ∝ e− 1 2h(x ′ 1)K −1(x′1,x ′ 2)h(x
′ 2) ( 1 + 1
N U (x′1, x ′ 2, x ′ 3, x ′ 4)H [h;x ′ 1, x ′ 2, x ′ 3, x ′ 4]
) (C.2)
where K−1 (x′1, x ′ 2) and U (x ′ 1, x ′ 2, x ′ 3, x ′ 4) are known from the previous layer. So we are looking for two maps: Kφ (K,U) (x, x′) = 〈φ (h (x))φ (h (x′))〉PK,U [h]
Uφ (K,U) (x1, x2, x3, x4) = 〈φ (h (x1))φ (h (x2))φ (h (x3))φ (h (x4))〉PK,U [h] (C.3)
so that the mapping between the first two cumulants K and U of two consequent layers is (assuming no biases)
K(`+1) (x, x′)
σ2 w(`+1)
= Kφ ( K(`), U (`) ) (x, x′)
U (`+1) (x1, x2, x3, x4)
σ4 w(`+1)
= Uφ ( K(`), U (`) ) (x1, x2, x3, x4)
−Kφ ( K(`), U (`) ) (xα1 , xα2)Kφ ( K(`), U (`) ) (xα3 , xα4) [3]
(C.4)
where the starting point is the first layer (N (0) ≡ d)
K(1) (x, x′) = σ2 w(1)
N (0) x · x′ U (1) (x1, x2, x3, x4) = 0 (C.5)
The important point to note, is that these functional integrals can be reduced to ordinary finite dimensional integrals. For example, for the second layer, denote
h := ( h1 h2 ) K(1) = ( K(1) (x1, x1) K (1) (x1, x2) K(1) (x1, x2) K (1) (x2, x2) ) (C.6)
we find for K(2) K(2) (x1, x2)
σ2 w(2)
= ∫ dhe − 12h TK−1 (1) h φ (h1)φ (h2) (C.7)
and for U (2) we denote
h :=  h1h2h3 h4  K(1) =  K(1) (x1, x1) K (1) (x1, x2) K (1) (x1, x3) K (1) (x1, x4) K(1) (x1, x2) K (1) (x2, x2) K (1) (x2, x3) K (1) (x2, x4) K(1) (x1, x3) K (1) (x2, x3) K (1) (x3, x3) K (1) (x3, x4) K(1) (x1, x4) K (1) (x2, x4) K (1) (x3, x4) K (1) (x4, x4)  (C.8)
so that Uφ ( K(1), U (1) ) (x1, x2, x3, x4) = ∫ dhe − 12h TK−1 (1) h φ (h1)φ (h2)φ (h3)φ (h4) (C.9)
This iterative process can be repeated for an arbitrary number of layers.

heading: D FOURTH CUMULANT FOR THRESHOLD POWER-LAW ACTIVATION FUNCTIONS
text: D.1 FOURTH CUMULANT FOR RELU 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) = z2 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 ith hidden unit with input xα sampled with a uniform measure from Sd−1( √ d), where w(i) is a vector of weights of the first 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 find U corresponding to the leading FWC by first finding the fourth moment of the hidden layer µ4 := 〈φ1φ2φ3φ4〉 (see Eq. 9), taking for simplicity ς2w = 1
µ4 =
√ det(L−1)
(2π) 2
∞∫ 0 dze− 1 2z TL−1zz1z2z3z4 (D.1)
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 modification noted therein) the following series in the off-diagonal elements of the matrix L
µ4 = ∞∑ `,m,n,p,q,r=0 A`mnpqrL ` 12L m 13L n 14L p 23L q 24L r 34 (D.2)
where the coefficients A`mnpqr are
(−)`+m+n+p+q+r G`+m+nG`+p+qGm+p+rGn+q+r `!m!n!p!q!r!
(D.3)
For ReLU activation, these G’s read
GReLUs =  1√ 2π s = 0 −i 2 s = 1
0 s ≥ 3 and odd (−)k(2k)!√
2π2kk! s = 2k + 2 k = 0, 1, 2, ...
(D.4)
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 sufficiently large input dimension d since the overlap between random normalized inputs scales as O(1/ √ d) and consequently L(x, x′) ∼ 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 reflect 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.2 DERIVATION OF THE PREVIOUS SUBSECTION
In this section we derive the expression for the fourth moment 〈f1f2f3f4〉 of a two-layer fully connected network with threshold-power law activations with exponent ν: φ(z) = Θ(z)zν ; ν = 0 corresponds to a step function, ν = 1 corresponds to ReLU, ν = 2 corresponds to ReQU (rectified quadratic unit) and so forth.
When the inputs are normalized to lie on the hypersphere, the matrix L is
L =  1 L12 L13 L14L12 1 L23 L24L13 L23 1 L34 L14 L24 L34 1  (D.5) where the off diagonal elements here have Lαβ = O ( 1/ √ d )
. 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:
P+ =
√ det(L−1)
(2π) 2
∞∫ 0 dze− 1 2z TL−1z (D.6)
The characteristic function (Fourier transform) of this distribution is
ϕ (t1, t2, t3, t4)
= exp ( −1
2 tTLt
)
= exp ( −1
2 4∑ α=1 t2α
) exp −∑ α<β Lαβtαtβ  = exp ( −1
2 4∑ α=1 t2α ) ∞∑ `,m,n,p,q,r=0 (−)`+m+n+p+q+r L`12Lm13Ln14L p 23L q 24L r 34 `!m!n!p!q!r! t`+m+n1 t `+p+q 2 t m+p+r 3 t n+q+r 4
(D.7)
Performing an inverse Fourier transform, we may now write the positive orthant probability as
P+ = 1
(2π) 4 ∫ R4+ dz ∫ R4 dt ϕ (t1, t2, t3, t4) e −i ∑4 α=1 zαtα
= ∞∑ `,m,n,p,q,r=0 (−)`+m+n+p+q+r L`12Lm13Ln14L p 23L q 24L r 34 `!m!n!p!q!r! × · · ·
× 1 (2π) 4 ∫ R4+ dz ∫ R4 dt e ∑4 α=1(− 12 t 2 α−izαtα)t`+m+n1 t `+p+q 2 t m+p+r 3 t n+q+r 4
= ∞∑ `,m,n,p,q,r=0 A`mnpqrL ` 12L m 13L n 14L p 23L q 24L r 34
(D.8)
where the coefficients A`mnpqr are
A`mnpqr = (−)`+m+n+p+q+r G`+m+nG`+p+qGm+p+rGn+q+r
`!m!n!p!q!r! (D.9)
and the one dimensional integral is
G(ν=0)s = 1
2π ∞∫ 0 dz ∞∫ −∞ ts exp ( −1 2 t2 − itz ) dt (D.10)
We can evaluate the integral over t to get
G(ν=0)s = 1
(−i)s (2π)1/2 ∞∫ 0 ( d dz )s e−z 2/2dz (D.11)
and performing the integral over z yields
G(ν=0)s =  1 2 s = 0 0 s even and s ≥ 2 (2k)!
i(2π)1/22kk! s = 2k + 1 k = 0, 1, 2, ...
(D.12)
We can now obtain the result for any integer ν by inserting zν inside the z integral:
G(ν)s = 1
2π ∞∫ 0 dz zν ∞∫ −∞ ts exp ( −1 2 t2 − itz ) dt =
1
(−i)s (2π)1/2 ∞∫ 0 zν ( d dz )s e−z 2/2dz
(D.13)
Using integration by parts we arrive at the result Eq. D.4 reported in the main text
GReLUs = G (ν=1) s =  1√ 2π s = 0 −i 2 s = 1
0 s ≥ 3 and odd (−)k(2k)!√
2π2kk! s = 2k + 2 k = 0, 1, 2, ...
(D.14)
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 xmay 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.

heading: E FOURTH CUMULANT FOR QUADRATIC ACTIVATION FUNCTION
text: For a two-layer network, we may write U , the 4th cumulant of the output f(x) = ∑N i=1 aiφ(w T i x), with ai ∼ N (0, ς2a/N) and wi ∼ N (0, (ς2w/d)I) for a general activation function φ as
Uα1,α2,α3,α4 = ς4a N ( V(α1,α2),(α3,α4) + V(α1,α3),(α2,α4) + V(α1,α4),(α2,α3) ) (E.1)
with V(α1,α2),(α3,α4) = 〈φ α1φα2φα3φα4〉w − 〈φ α1φα2〉w 〈φ α3φα4〉w (E.2)
For the case of a quadratic activation function φ(z) = z2 the V ’s read
V(α1,α2),(α3,α4) = 2 { L11L33 (L24) 2 + L11L44 (L23) 2 + L22L33 (L14) 2 + L22L44 (L13) 2 } +...
4 { (L13) 2 (L24) 2 + (L14) 2 (L23) 2 } +8 (L11L23L34L24 + L22L34L14L13 + L33L12L14L24 + L44L12L13L23)+...
16 (L12L13L24L34 + L12L14L23L34 + L13L14L23L24) (E.3)
where the linear kernel from the first layer is L(x, x′) = ς 2 w
d x · x ′. Notice that we distinguish between
the scaled and non-scaled variances:
σ2a = ς2a N ; σ2w = ς2w d
(E.4)
These formulae were used when comparing the outputs of the empirical two-layer network with our FWC theory Eq. 11. 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).

heading: F AUTO-CORRELATION TIME AND ERGODICITY
text: 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 suffices 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:
lim T̃→∞
E [〈 fDNN(x∗; t) 〉 T̃ − µ(x∗) ] = 0
lim T̃→∞
E [(〈 fDNN(x∗; t) 〉 T̃ − µ(x∗) )2] = 0
(F.1)
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
〈 fDNN(x∗; t) 〉 T̃ := 1
T̃ ∫ T̃ 0 fDNN(x∗; t)dt ≈ 1 T̃ tj=T̃∑ tj=0 fDNN(x∗; tj) (F.2)
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 Fij = fDNNi (x∗; tj). (1) Divide the time series data in the matrix F into non-overlapping sub-matrices, each of dimension nseeds × nepochs. (2) For each of these sub-matrices, find 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 nepochs, nseeds. If ergodicity holds, we should expect to see the following relation
σ2emp(x∗) = σ 2 m
τ
nepochsnseeds (F.3)
where τ is the auto-correlation time of the outputs and σ2m 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 nepochs used for time averaging in each set (and using all 500 seeds in this case). Performing a linear fit on the average across test points (black x’s in the figure) yields a slope of approximately −1, which is strong evidence for ergodic dynamics.
Figure F.1: Ergodicity check. Empirical variance σ2emp(x∗) vs. the number of epochs used for time averaging on a (base 10) log-log scale, with dt = 0.003 and N = 200. The colored circles represent different test points x∗ and the black x’s are averages across these.

heading: G NUMERICAL EXPERIMENT DETAILS
text: G.1 FCN EXPERIMENT DETAILS
We trained a 2-layer FCN on a quadratic target y(x) = xTAx where the x’s are sampled with a uniform measure from the hyper-sphere Sd−1( √ d), with d = 16 and the matrix elements are sampled as Aij ∼ N (0, 1) and fixed 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 first layer γw = 0.05. Notice that for any activation φ, K scales linearly with ς2a = σ 2 aN = (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 ·106 epochs, which includes the time it takes for the training loss to level off, which is usually on the order of 104 epochs. In the main text we showed GP and FWC results for a learning rate of dt = 0.001. Here we report in Fig. G.1 the results using dt ∈ {0.003, 0.001, 0.0005}. 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.
Figure G.1: Regression task with fully connected network: (un-normalized) MSE vs. width on log-log scale (base 10) for quadratic activation and different leaning rates. The learning rates dt = 0.001, 0.0005 converge to very similar values (recall this is a log scale), demonstrating that the learning rate is sufficiently small so that the discrete-time dynamics is a good approximation of the continuous-time dynamics. For a learning rate of dt = 0.003 (blue) and width N ≥ 1000 the dynamics become unstable, thus the general trend is broken, so one cannot take the dt to be too large.
G.2 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 final layer and fed it into an ReLU activated fully-connected layer with 25C 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 σ2w = 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, sufficient training time (namely, epochs) was needed to get estimates of the average outputs f̄E(xα) = f̄(xα) + δfα since
the estimators’ fluctuations, δfα, scale as (τ/ttraining)−1/2, where τ is an auto-correlation time scale. 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 ntest∑ α=1 (f̄E(xα)−fGP (xα))2 = ntest∑ α=1 (f̄(xα)−fGP (xα))2−2 ntest∑ α=1 (f̄E(xα)−fGP (xα))δfα+ ntest∑ α=1 (δfα) 2 (G.1) 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 fluctuations of the outputs quantified by Σα = (δfα)2 are of the order 0.12, the amount of uncorrelated samples of CNN outputs we require should be much larger than 0.12/(7 · 10−3) ≈ 1.43. To estimate this bias in practice we repeated the experiment with 3-7 different 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 fluctuations) for about 5 ·105 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 finite-width. We find 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.002, 0.001, 0.0005 and C = 48 showing good convergence already at 0.001. Following this we used learning rates of 0.0005 for C ≤ 48 and 0.00025 for C > 48, in the main figure.
Figure G.2: MSE between our CNN with C = 48 and its NNGP as a function of three learning rates.
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).

heading: H FURTHER NUMERICAL RESULTS ON CNNS
text: 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.
(a) n = 1000 (b) subtract bias of MSE estimator
Figure H.3: CNNs trained on CIFAR10 in the regime of the NNSP correspondence compared with NNGPs MSE test loss normalized by target variance of a deep CNN (solid green) and its associated NNGP (dashed green) along with the MSE between the NNGP’s predictions and CNN outputs normalized by the NNGP’s MSE test loss (solid blue, and on a different scale). We used balanced training and test sets of size 1000 each. For the largest number of channels we reached, the slope of the discrepancy between the CNN’s GP and the trained DNN on the log-log scale was−1.77, placing us close to the perturbartive regime where a slope of −2.0 is expected. Error bars here reflect statistical errors related only to output averaging and not due to the random choice of a test-set. The performance deteriorates at large N = #Channels as the NNSP associated with the CNN approaches an NNGP.
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 σ2b = 0.1 instead of σb = 1.0 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− 16 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.

heading: I THE FOURTH CUMULANT CAN DIFFERENTIATE CNNS FROM LCNS
text: 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 (Xi) as input i ∈ {0, . . . , L − 1}; (2) A single convolutional layer with some activation φ acting with stride 1 and no-padding using the conv-kernel T cx where c ∈ {1, . . . , C} is a channel number index and x ∈ {0, . . . , 2l} is the relative position in the image. Notably, in an LCN this conv-kernel will receive an additional dependence on x̃, the location on Xi on which the kernel acts. (3) 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 − 2l)}. (4) A linear fully connected layer with weights W ocx̃ where o ∈ {0, . . . ,#outputs} are the output indices. Consider first the NNGP of such a random DNN with weights chosen according to some iid Gaussian distribution P0(w), with w including both W ocx̃ and T c x . Denoting by z
o(x) the o’th output of the CNN, for an input x we have (where we denote in this section 〈· · · 〉 := 〈· · · 〉P0(w))
Koo ′ (x, x′) ≡ 〈zo(x)zo ′ (x′)〉 = δoo′ ∑ c,c′,x̃,x̃′ 〈W ocx̃W o ′ c′x̃′〉〈φ(T cx(x̃)Xx+x̃−l)φ(T c ′ x (x̃ ′)Xx+x̃′−l)〉
(I.1)
The NNGP kernel of an LCN is the same as that of a CNN. This stems from the fact that 〈W ocx̃W oc′x̃′〉 yields a Kronecker delta function on the x̃, x̃′ 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′ 6=x(x̃ ′), becomes irrelevant as the these two are never averaged together.
For simplicity, we turn to the fourth cumulant of the same output, given by
〈zo(x1) · · · zo(x4)〉−〈zo(xα)zo(xβ)〉〈zo(xγ)zo(xδ)〉[3] = 〈zo(x1) · · · zo(x4)〉−K(xα, xβ)K(xγ , xδ)[3] (I.2) with the second term on the LHS implying all pair-wise averages of zo(x1)..zo(x4). Note that the first 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∑
c1..c4x̃1..x̃4
〈W oc1x̃1 · · ·W o c4x̃′4 〉〈φ(T c1x1 (x̃1)Xx1+x̃1−l) · · ·φ(T c4 x4 (x̃4)Xx4+x̃′4−l)〉 (I.3)
The average over the four W ’s yields non-zero terms of the type W ocx̃W o cx̃W o c′x̃′W o c′x̃′ with either x̃ = x̃′ (type 1), x̃ 6= x̃′ and c 6= c′ (type 2), or x̃ 6= x̃′ and c = c′ (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∑ c 6=c′;x̃ 6=x̃′ 〈W ocx̃W ocx̃〉〈W oc′x̃′W oc′x̃′〉〈φ(T cx(x̃)Xx+x̃−l)φ(T cx(x̃)Xx+x̃−l)φ(T c ′ x′ (x̃ ′)Xx′+x̃′−l)φ(T c′ x′ (x̃ ′)Xx′+x̃′−l)〉 (I.4) Examining the average involving the four T ’s, one finds that since T cx(x̃) is uncorrelated with T c′ x′ (x̃ ′) for both LCNs and CNNs, it splits into∑ c 6=c′;x̃ 6=x̃′ 〈W ocx̃W ocx̃〉〈W oc′x̃′W oc′x̃′〉〈φ(T cx(x̃)Xx+x̃−l)φ(T cx(x̃)Xx+x̃−l)〉〈φ(T c ′ x′ (x̃ ′)Xx′+x̃′−l)φ(T c′ x′ (x̃ ′)Xx′+x̃′−l)〉 (I.5) 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 ofK(xα, xβ)K(xγ , xδ)[3], apart from a “diagonal" contribution (x̃ = x̃′).
We turn our attention to the type 3 term given by∑ c;x̃ 6=x̃′ 〈W ocx̃W ocx̃〉〈W ocx̃′W ocx̃′〉〈φ(T cx(x̃)Xx+x̃−l)φ(T cx(x̃)Xx+x̃−l)φ(T cx′(x̃′)Xx′+x̃′−l)φ(T cx′(x̃′)Xx′+x̃′−l)〉 (I.6) Examining the average involving the four T ’s, one now finds 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 cx(x̃ ′). For a CNN however, T cx(x̃) is the same random variable as T c x(x̃ ′) 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/C2 factor when properly normalized.

heading: J CORRECTIONS TO EK
text: Here we derive finite-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 (〈· · · 〉n) over all draws of datasets of size n′ with n′ taken from a Poisson distribution with mean n is given by
Zn =
∫ Dfe−SGP[f ]− n 2σ2 ∫ dµx(f(x)−y(x))2(1 + SU [f ]) +O(1/N2) (J.1)
with SGP[f ] and SU [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. (2019)). Notably, the above expression coincides with that used for
a finite 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 σ2 is now accompanied by n. Following this, all the diagrammatic and combinatorial aspects shown in the derivation for a finite dataset hold here as well. For instance, let us examine a specific contribution coming from the quartic term in H[f ]: Ux1..x4K
−1 x1x′1 · · ·K−1x4x′4f(x ′ 1) · · · f(x′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
Ux1,x2,x3,x4K −1 x1x′1 〈f(x′1)〉∞K−1x2x′2〈f(x ′ 2)〉∞K−1x3x′3〈f(x ′ 3)〉∞K−1x4x′4Σ∞(x ′ 4, x∗) (J.2)
where we recall that 〈f(x)〉∞ = Kxx′K̃−1x′x′′y(x′′) and Σ∞(x1, x2) = Kx1,x2 − Kx1,x′K̃ −1 x′,x′′Kx′′,x2 being the posterior covariance in the EK limit, where K̃xx′f(x ′) = Kxx′f(x ′)+ (σ2/n)f(x). Using the fact that K−1xx′Kx′x′′ gives a delta function w.r.t. the measure, the integrals against K−1xαx′α can be easily carried out yielding(
Ux1,x2,x3,x∗ − Ux1,x2,x3,x4K̃−1x4,x′4Kx′4,x∗ ) K̃−1x1,x′1 K̃−1x2,x′2 K̃−1x3,x′3 y(x′1)y(x ′ 2)y(x ′ 3) (J.3)
Introducing the discrepancy operator δ̃xx′′ := δxx′′ −Kxx′K̃−1x′x′′ = σ2 n K̃ −1 xx′′ , we can write a more compact expression( n σ2 )3 δ̃x∗,x4Ux1,x2,x3,x4 δ̃x1,x′1 δ̃x2,x′2 δ̃x3,x′3y(x ′ 1)(x ′ 2)y(x ′ 3) (J.4)
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. 14. The other term therein, which is linear in y, is a result of following similar steps with the f̄ΣΣ∗ contributions that do not get canceled by the quadratic part in H[f ].

Review summaries:
This paper shows a correspondence between deep neural networks (DNN) trained with noisy gradients and NNGP. It provides a general analytical form for the finite width correction (FWC) for NNSP expanding around NNGP. Finally, it argues that this FWC can be used to explain why finite width CNNs can improve the performance relative to their GP counterparts on image classification tasks.