Predicting the Outputs of Finite Networks Trained with Noisy Gradients
Authors: Gadi Naveh
Oded Ben-David
Haim Sompolinsky
Zohar Ringel
Venue:  (2021)
Keywords: Gaussian process
deep learning theory
finite DNNs
statistical mechanics

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 parameters maps to the so-called 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.

Under review as a conference paper at ICLR 2021
PREDICTING THE OUTPUTS OF FINITE NETWORKS TRAINED WITH NOISY GRADIENTS
Anonymous authors Paper under double-blind review
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 GP governed by the Neural Tangent Kernel (NTK), whereas earlier works showed that DNN with an i.i.d. prior over its weights maps to the so- called 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 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 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 ((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 image classification tasks.
INTRODUCTION
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 - oo 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 the Neural Tangent Kernel (NTK) (Jacot 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 Basni et al. 2019) of DNNs. Roughly speaking, highly over- parameterized DNNs generalize well because they have a strong implicit bias 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 mnder-performs, often by large margin, convolutional neural networks (CNNs) trained with SGD (Arora et al. (019). (c) Deriving explicit finite width corrections (FWCs) is challenging, it requires solving a set of coupled ODEs (Dyer & Gur-Ari 2020; Huang & Yau 2019). Thus, there is a need for extended theory for end-to-end trained deep networks which is valid for finite width DNNs.

---

Under review as a conference paper at ICLR 2021
Our contribution is three-fold. First, we prove correspondence between DNN trained with noisy gradients and a Stochastic Process (SP) which at N + oo 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 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 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 width corrections are important for intermediate values of n (Fig.
Third, we propose an explanation for why finite CNNs trained on image classification tasks can outperform their infinite- 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
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 paeerraratttiiiiioonn, stochas- ticity, 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.
1.1 RELATED 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 al., 2017 Teh et al.. 2016| et al. 2019 Ye et al.) [017]. However, to the best of our knowledge, a correspondence with concrete SP or a non-pan non-parametric model was not established nor was 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. Hanni& & Nica (219)) study the NTK of finite DNNs, but where the depth scales together with width, whereas we keep the depth fixed. Dyer Gur-An (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 -ontrretureeeive 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- non-Gaussian priors induced by finite-N DNNs. Unlike here, there was no attempt to establish correspondence with trained DNNs. The formulation presented here has the conceptual advantage of representing 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 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 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-Planne type. In contrast, in our framework are interested in the distribution over function
2

---

Under review as a conference paper at ICLR 2021
space equilibrium, i.. for cc. (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, showed in section 4.2.
2 THE NNSP CORRESPONDENCE
In this section we show that finite-width DNNs, trained in specific manner, correspond to Bayesian inference using a non-pan parametric model which tends to the NNGP as N + Dc. We first give a short review of Langevin dynamics in weight space described by Neal al. (2011), Welling & Teh (201T), 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 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
(1) where Wy the vector of all network weights at time step t, 7 is the strength of the weight decay, C(zw) is the loss as a function of the output Z T' is the temperature (the magnitude of noise), dt the learning rate and & (O, I). As dt - these discrete-time dynamics converge to the
Aw, == W++1 - Wi = - (rw; + V. (zw)) dt + d2T
continuous-time Langevin equation given by do (L) = - v. (3 [1(((1)))) 2 + & (z.)) + V (1) with (t),,, (t)) 6,,6 (t f'), so that - oo the weights will be sampled from the equilibrium
distribution in weight space, given by (Risken & Frank 1996)
(2 11 + + ((..)) = exp (-(22
L (Z (zw)))
P (w) x exp -
|00l2 + 20
(2)
20
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 (0, a I). Thus we can map the hyper- paremeearreeerss of the training to those of the BNN: x2 = T/y and o² T/2. Notice that a sensible scaling for the weight variance at layer & is The ~ (1/Ne-1), thus the weight decay needs to scale as 7 (NN-1).
transition from weight space to function space- We aim to move from a distribution over weight space Eq. 2 to one over function space. Namely, we consider the distribution of 2 (x) implied by the above (ww) where for concreteness we consider a DNN with single scalar output zn (x) E TR on a regression task with data { {x Na) )a) Rd R. Denoting by P[J] the induced measure on function space we formally write
Pt - dwols - =nm (w: e 202 / (dwe 20' llwll 8lf zw]
where f fd denotes an integral over all weights and we denote by 8[f Zw a delta-function in function-space. As common in path-integrall or field-theory formalism (Schulman! 2012). such a delta function is understood as limit procedure where one chooses suitable basis for function space, trims it to a finite subset, treats oʃ - zn|] 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[S]
Po[S] where the
prior over functions is Po[s] x f 2 || 8|ʃ - z. The integration over weights now obtains clear meaning: it yields the distribution over functions induced by DNN with i.i.d. random weights z w|12 Thus, we can relate any correlation function in function space and weight space, for instance (Df is the integration measure over function space)
chosen according to the prior (o(w)
f (FR(L)(()))))))) - DS dddP = dw (ooPPo()zo zm am
zw
As noted by Cho & Saul (2009), for highly over parmmerraeteri DNNs the r.h.s. of equals the kernel of the NNGP associated with this DNN, K K(x,x) Moreover Po[S] tends Gaussian and can be
3

---

Under review as a conference paper at ICLR 2021
written
Polf] x exp / (u(x)))(())(())) (ax ((,)6)))) (1(/N)
(5)
where ua(z) the measure of the imput space, and the (LL/N) scaling of the finite-NV correction
will be explained in $3 If we now plug 5 in 3 take the loss to be the total square erro [ c[f]
Ex=1 (ya- - f (xa)) and take N + so have that the posterior P[f] is that of GP. Assuming ergodicity, one finds that training- time averaged output of the DNN is given by the posterior mean of GP, with measurement moiserequss
o² 222 and kernel given by the NNGP of that DNN.
We refer to the above expressions for Pa[ and P[ 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 Po[/] at finite N, one can predict the DNN's outputs at finite N. Henceforth, we refer to Pol.) as the prior distribution, as it the prior distribution of a DNN with random weights drawn from Po(w).
Evidence supporting ergodicity Our derivation relies on the ergodicity of the dynamics. Ergodicity 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 ceff-consst 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 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-crgodic dynamics). However, the empirical ACT gives 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 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 solution to constraints using far more variables (roughly M 22 where M is the number of layers). From a different angle, in 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).
INFERENCE ON THE RESULTING NNSP
Having mapped the analyze the predictions of this NNSP in the case where N large but finite, such that the NNSP is only weakly non- non-Gaussian (i.e. its deviation from a GP is (1/NN).. The main result of this section a derivation of leading FWCs to the standard GP regression results for the posterior mean for (a.) and variance (xp(x.) on unseen test point T., given training set ((x., Va) Va)) C Rd x R. namely (Rasmussen & Williams. 2005)
outputs of DNN to inference on the above NNSP, we turn to
For(x.) Σ Ya K K K3;
Ear(x.) K'* E KKaKd a
(6)
a.
where Kag== == K(Fa, o sagi
K; K(I,,za): KM K(I,,
We take the total error, i.c. we don't 't divide by so that [(f) becomes more dominant for larger n. Here o² is property of the training protocol and of the data itself, our prior it.
4

---

Under review as a conference paper at ICLR 2021
3.1 EDGEWORTH EXPANSION AND PERTURBATION THEORY
Our first task is to find how P[/] changes compared to the Gaussian (N + oo) scenario. As the data-dependent part e --Iifl/20 is independent of the DNNs, this amounts to obtaining finite width corrections to the prior f. One way to characterize this is to perform an Edgeworth expan- sion of Po[s] (Mecullagh/ 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 { and assume WLOG (2,) 0, (Z?) = 1, with higher cumulants x² for 2 3. Now consider their normalized sum Yv = I 2i Ze From additivity and homogeneity of cumulants Nu
VN
we have R,22 == K222 = (VN) = N-/2-1- Now, let (py) = (2T) 1/2e v/ The charac-
teristic function of Y is f(t) = F(f(v)) = exp (Er= Mi (6t)" = exp (Er= Aor (86) r! ((t). Taking the inverse Fourier transform F-1 has the effect of mapping it + -0, thus we get f(v) = exp E (-8,)" (()) = (v() (1 Dr= H+))) where Hy(y) is the rth Her- = y
y!
If we were to consider vector-valued RVs, then the r'th order cumulant becomes 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, are continuously indexed by the inputs Ze.
mite polynomial, e.g.
This is especially convenient here since for all DNNs with the last layer being fully-comnected, all odd cumulants vanish and the th cumulant scales as 1/N*-1, Consequently, at large N we can char- acterize Po[f]] up to (2N-2) by its second and fourth cumulants, K (I1, X2) and U ((11, T2, T3. I4). respectively. Thus the leading order correction Polf] reads
Polf] Sar([] (1 1
Su (A) + (1/N2)
where the GP action SeP and the first FWC action Sy are given by
1 SulfI = 41 /app/dpplll dpiiaUz aa H m
/
Scr = 2 dp1:25z, Krt Kagzzfzzi
(8)
Here, H is the 4th functional Hermite polynomial (see App. A). U is the 4th order functional cumulant of the NN output which depends on the choice of the activation function & U U21 s/ (18)),)))))))))) (d.0)) (d)))s)) + other perins. of (a,8,7,0 7, (1,...,44
4}
(9)
where &a == -(2f-1(xa)) and the pre-activations are 2f(x) - of Σ Ny W2 (
(x)). Here
we distinguished between the scaled and non-scaled weight variances: a 2 /NN, where a are the weights of the last layer. Our shorthand notation for the integration measure over inputs means e.g.
du144 := dp(x1) ... ..p(z))))
Using perturbation theory, in App. we compute the leading FWC to the posterior mean f(x,) and variance (())(x=))²) on a test point T.
f(x.) (x.) + N [v(a +T( ((6f(x.))²) = Eor(z, +N- EN 0(N-)) with Ex( ((((xx)))) - (x. and
(10)
To( U+ Uiiozas
(Jon Hazaaa
3k S
(11)
((((x.)))) Üre
(Nnnoos Noz
Kaloz)
where all repeating indices are summed over the training set (i.c. range over (l,..., n)), denoting: V. R- us. and defining
U Úniazas = U U
Kajazajaa K Kayk3 K3 Ü (diazas) )
anazas
(12)
Ü
Uata2
U UVjazas
k R-1 K
Here we take U ~ O(1) to emphasize the scaling with N in Eqs.
5

---

Under review as a conference paper at ICLR 2021
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 fu (2+) has a linear term in y, which can be viewed as correction to the GP kernel, but also cubic term in y, unlike for (ar.) which is purely linear. The correction to the predictive variance Ev ) has quartic and quadratic terms in y, unlike [xp(x.) which is j-independent. Ü*
terms of GP regression: if we consider the indices 01, Q22 Otg as aa has a clear interpretation in can be thought of as the ground truth value of target function (analogous to y.), and the second term on the r.h.s. K a.. K is then the GP prediction of U Vajagas with the kernel K, where at runs on the training set (compare to far (a.) in Eq. (). Thus U 02 is the discrepancy in predicting U Uajagaga using a GP with kernel K. In 3.2 we study the behavior of (x.) as a function of n.
then U U2razas a20s
Un Varazagas
The posterior variance Exx) (5f (x)))² has a clear interpretation our correspondence: is
decrease the test loss by averaging. Our procedure for generating empirical network outputs involves time-averaging over the training dynamics after reaching equilib- rium 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 E(E.). Following this, one finds that the expected MSE test loss after training saturates is n,
measure of how much
Ea=i ((F(()) - y (za))² + E(xa)) where n,, is the size of the test set.
3.2 FINITE WIDTH CORRECTIONS FOR SMALL AND LARGE DATA SETS
The expressions in Eqs. 6 II 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 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; Soolich & Williams 2004). The EK approximates the GP predictions at large nt, after averaging on all draws of (roughly) n. training points representing the target function. Even if one is interested in 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. aree replaced by integrals over all input space, which together with spectral decomposition of the kernel function K(I, x²) = E, X (z)))s(z') yields the well known result
S dp(ae)) Σ >xtt(((.))xx(")))
EK JP
(a..)
\ 02/02 g(x')
(13)
Here we develop an extension of Eq. 13 for the NNSPs we find large but finite N. In particular, we find the leading non-linear correction to the EK result, i.e. the "EK analogue" of Eq. [T To this end, consider the average predictions of NNSP trained on ensemble of data sets of size n', corresponding to n' independent draws from a distribution j(() over all possible inputs x. Following the steps in App. D we find
13
3n²
FEK (I.) = Da U2i
g(x')) σf dm Be, (((1))))))
(14)
where an integral f du(zz) is implicit for every pair of repeated x coordinates. We introduced the discrepancy operator 8zz. which acts on some function to as f
== (x'' = (xx) - (EK (zr). Essentially, Eq. i4 derived from Eq. 11 by replacing each K by (n/o²) )) and noticing that in this regime Ü+ Uz in Eq. 12 becomes 5,, U2, "02.03,04" Interestingly, FEK (xx) is written explicitly in terms of meaningful quantities: "=m'(((")) and 6 U #2200,24
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. (011)). At N - co, we obtain the standard EK result. Eq. [3] It is basically high-pass, -ashs linear filter which filters out features of g that have support on eigenfunctions th associated with eigenvalues A; that are small relative to o² /nn. We stress that the the X, 's are independent of any particular size n dataset but rather

---

Under review as a conference paper at ICLR 2021
property of the average dataset. In particular, computationally costly data dependent matrix inversion is needed to evaluate Eq.
13
Turning to our FWC result, Eq. 14 it depends on g(z) only via the discrepancy operator 8wn.. Thus these FWCs would be proportional to the error of the DNN, at N + co. In particular, perfect performance at N + oo, implies no FWC. Second, the DNN's average predictions act as a linear transformation the target function combined with cubic non-linearity. Third, for g(x) having support only on some finite set of cigenfunctions a of K, òzz'g(z") would scale as o² /n 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 δegz, also scales as o² /m, so that the whole expression in the f-..} has no scaling with n. In addition, some decreasing behavior with n is expected due to the 0g. U21.22.25.0 factor which can be viewed as the discrepancy in predicting d2 at fixed 22; 33, I44 based on random samples (I 's) of Una 02 #2.4 In Fig. Te 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 only important at intermediate values of n.
2
Loading FNC
2
2-
-3
2-
-3003
0.0003
2-18
2
220003
2
2
27
2'
2'
fram set Size
tram set size
train set size
FWC vs. n for d
(b) FWC vs. n for
(c) GP RMSE vs. for d
Figure 1: Leading FWC to the mean [x.)). (Eq. 11) and GP discrepancy (RMSE) as a function of train set size n for varying training noise o². The target is quadratic g(x) x Ax = (1)) with x E Sa-1(Vd) 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 [zx() increases linearly for small n (dashed-dotted lines in (a)) before decays (best illustrated for larger d and o²). For sufficiently large n, both the GP discrepancy and Ifu (x.) scale as 1/n (diagonal dashed black lines in (c), This verifies our prediction for the scaling of FWCs with n. Eq. 14jin the large regime. Notably, it implies that at large N' FWCs are only important intermediate values of n.
NUMERICAL EXPERIMENTS
In this section we numerically test our analytical results. We first demonstrate that in the limit N + do the outputs of FCN trained in the regime of the NNSP correspondence converge to 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 f Eq. I1 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.
4.1 TOY EXAMPLE: FULLY CONNECTED NETWORKS ON SYNTHETIC DATA
We trained 2-layer FCN f(x) MN agp(w()) 2r) on a quadratic target v(x: = TT Aar where the I's are sampled with a uniform measure from the hyper-sphere S see App. G. 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) (x.) is much closer to the GP prediction fcr(x.) than to the ground truth ys. Otherwise, the convergence of the network output to the corresponding NNGP as N grows (shown in Fig. [cp would be trivial, since all reasonable estimators would be close to the target and hence close to each other.
7

---

Under review as a conference paper at ICLR 2021
In Fig. 2c we plot in log-log scale (with base 10) the MSE (normalized by (fDNN)²) between the predictions of the network JDNN and the corresponding GP and FWC predictions for quadratic and ReLU activations. We find that indeed for sufficiently large widths (N 2 500) the slope of the GP-DNN MSE approaches - 2.0 (for both ReI.U and quadratic), which is expected from our theory, since the leading FWC scales 1/N. For smaller widths, higher order terms (in 1 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 regime where the GP and FWC MSEs intersect at N E 100, below which our FWC actually increases the MSE, which suggests scale of how large N needs be for our leading FWC theory to hold.
Choak beu
(a) Predictions
(b) Auto-correlation functions
(c) MSE scaling
Figure 2: Fully connected 2-layer network trained on a regression task. (a) Network outputs on a test point fx..,)) t) vs. normalized time: the time-averaged DNN output f(x.) (dashed line) is much closer to the GP prediction far(2, (dotted line) than to the ground truth y. (dashed-dotted line). (b) ACFs of the time series of the 1st and 2nd layer weights, and of the outputs: the output converges to equilibrium faster than the weights. (c) Relative MSE between the network outputs and the labels y (triangles), GP predictions fG.(2.) Eq. (dots), and FWC predictions Eq. To (x's), shown vs. width for quadratic (blue) and ReLU (red) activations. For sufficiently large widths (N 2 500) the slope of the GP-DNN MSE approaches 2.0 and the FWC-DNN MSE is further improved by more than order of magnitude.
4.2 PERFORMANCE GAP BETWEEN FINITE CNNS AND THEIR NNGP
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 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 mindeninnsional vector (CNN-VEL VEC) applying global average pooling to the spatial dimensions CENN-GAP). It was observed that while infinite FCN and CNN- VEC 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 CIFAR10, was shown in Novak al. (2188) 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 be on par with that of their NNGPs. We argue that our framework can account for this observation. The priors [f of a LCN and CNN-VEC agree on their second cumulant (the covariance), which the only one not vanishing as N - co, but they need not agree on their higher order cumulants, which come into play at finite N. In App. [w 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 + co. This is in contrast to SGD behavior reported in some works where the CNN performance seems saturate function of N, to some value better than
8

---

Under review as a conference paper at ICLR 2021
the NNGP (Nvakk et al. 2018 Neyshabur et al. 2018). Notably those works used maximum architecture scans, high learning rates, and early stopping, all which are absent from our training protocol.
To test the above conjecture we trained, according to protocol, CNN with six convolutional lay- ers and two fully connected layers on CIFARIO, and
1.010
used CNN-VEC for the readout. We used MSE loss with one-hot encoding into a 10 dimensional vector of the categorical label; further details and additional settings are given in App. G Fig. demonstrates that,
1.005
MSE/TGP y29
MSEWaryy
1.000
2
using our training protocol, a finite CNN can outper- form its corresponding GP and approaches its GP as the number of channels increases. This phenomenon was observed in previous studies under realistic train- ing settings (Novak et al. 2018), and here we show that it appears also under training protocol. We note that a similar yet more pronounced trend in per- formance appears here also when one considers the averaged MSE loss rather the the MSE loss of the average outputs.
0.990
GPY
0.980
2
2 #Channels
Figure 3: DNN-GP MSE demonstrates con- vergence to slope of -2.0, validating the theoretically expected scaling. DNN-ground truth (Y) MSE shows finite CNN outper- form corresponding GP.
CONCLUSION
In this work 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 + co. 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 large number of training points + oo, explicit expressions for the DNNs' outputs were given, involving no costly matrix inversions. that when GP has small
this regime,
the FWC can be written in terms of the discrepancy of GP predictions, test сгrоr the FWC will be small, and vice versa. In the small n regime, the FWC is small but grows with n, which implies that large N, FWCs only important intermediate values of 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 the structure of the data distribution, and thus improves performance relative to the corresponding GP. In future study, it would be very interesting 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.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On Exact Computation with an Infinitely Wide Neural Net. arXiv e-prints, art. 1rXiv:1904.11955, Apr 2019.
Ronen Basri, David Jacobs, Yoni Kasten, and Shira Kritchman. The Convergence Rate of Neural Networks for Learned Functions of Different Frequencies. arXiv e-prints, art. arXiv:1906.00425, Jun 2019.
Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang, networks: Non-asymptotic rates and generalization bounds. arXiv preprint 0rXiv:2002.04022 2020.
analysis of two-layer neural
Youngmin Cho and Lawrence Saul. Kernel methods for deep learning. In Proceedings of the 22Nd International Conference Neural Information Processing Systems, NIPS 09, pp. 342-350, USA, 2009. Curran Associates Inc. ISBN 978-1-61567-911-9, URL http: /al acm. org/
citation wtt..ccmmmm==242m00000 2984112
Omry Cohen, Or Malka, and Zohar Ringel. Learning Curves for Deep Neural Networks: A Gaussian Field Theory Perspective. arXiv e-prints, art. arXiv: 19X,::190.00501.,, Jun 2019.
A. Daniely, R. Frostig. and Y. Singer. Toward Deeper Understanding of Neural Networks: The Power of Initialization and Dual View on Expressivity. ArXiv e-prints, February 2016.
Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020. URL httpps:///penreeieew net/ forum?
E. Gardner and B. Derrida. Optimal storage properties of neural network models. Journal of Physics Mathematical General, 21:271-284, January 1988. doi: 10. 00.0888/0005-44722//1001...
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. arXiv preprint 1XXiv:1909005009 2019.
Jiaoyang Huang and Horng- Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. arXiv preprint arXiv: 19t::0999900550 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernl:: Convergence and Generalization in Neural Networks. arXiv e-prints, art. arXiv:1806.07572, Jun 2018
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04430 2016.
Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In International Conference on Learn- ing Representations, 2018. URL
www.: 2a
Jachoon Lee, Samuel S Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. arXiv preprint arXiv:2007.15801, 2020.
Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A Simple Baseline for Bayesian Uncertainty in Deep Learning. arXiv e-prints, art. arXiv:1902.02476 Feb 2019.
Stephan Mandt, Matthew D. Hoffiman, and David M. Blei. Stochastic Gradient Descent as Approxi- mate Bayesian Inference. arXiv e-prints, art arXiv: 104.01704.04289, Apr 2017.
Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint 1rXiv:1804.11277 2018.
Peter Mccullagh. Tensor Methods in Statistics. Dover Books on Mathematics, 2017.
10

---

Under review as a conference paper at ICLR 2021
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences. 115(33) 2018. ISSN 0027-8424. doi: 1111077/ppaa..8805799990 URL https: www. www.ppaas.... org/
content/115/333
P A P Moran. Rank Correlation and Product-Moment Correlation. Biometrika, 35(1):203-206, 1948.
Radford M Neal et al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards in Generalization of Neural Networks, arXiv
Understanding the Role of e-prints, art, arXiv: 18K5v:1805.12001 12076, May 2018.
Roman Novak, Lechao Xiao, Jachoon Lee, Yasaman Bahri, Greg Yang, Daniel A. Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes. arXiv e-prints, art.
October 2018.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred A. Hamprecht, Yoshua Bengio, and Aaron Courville. On the Spectral Bias of Neural Networks. arXiv e-prints, art. 1XXiv:100.008344, Jun 2018.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press, 2005. ISBN 026218253X.
H. Risken and T. Frank. The Fokker-Planee Equation: Methods of Solution and Applications. Springer Series in Synergetics. Springer Berlin Heidelberg, 1996. ISBN 9783540615309. URL
google i1 htll///ooksss wtt.ooook.s...MM222//S
Lawrence S Schulman. Techniques and applications of path integration. Courier Corporation, 2012.
Elena Sellentin, Andrew H Jaffe, and Alan F Heavens. On the of the edgeworth expansion in cosmology how to foresee and evade its pitfalls. arXiv preprint (eX9.0/709.034: 2017.
Peter Sollich and Christopher KI Williams. Understanding gaussian process regression using the equivalent kernel. International Workshop on Deterministic and Statistical Methods in Machine Learning, pp. 211-228. Springer, 2004.
Yee Whye Teh, Alexandre H. Thiery, and Sebastian J. Vollmer. Consistency and fluctuations for stochastic gradient langevin dynamics. J. Much. Learn, Res., 17(1):193-225, January 2016. ISSN 1532-4435.
Belinda Tzen and Maxim Raginsky. A mean-field theory of lazy training in two-layer neural nets: entropic regularization and controlled mekean-vlasov dynamics. arXiv preprint arXiv:2002.01987, 2020.
Max Welling and Yee whye Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML'11, pp. 681 6888, USA, 2011. Omnipress. ISBN 978-1-4503--619-5. URL
http://dl. acm. o ci tat ion c fmm?: d= 04482.44... 3104661
Sho Yaida. Non Non-gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pp. 165-192, PMLR, 2020.
N. Ye, Z. Zhu, and R. K. Mantiuk. Langevin Dynamics with Continuous Tempering for Training Deep Neural Networks. ArXiv e-prints, March 2017.
Zee. Quantum Field Theory in a Nutshell. Nutshell handbook. Princeton Univ. Press, Princeton, NJ, 2003. URL https: : cds. cern. ch. wwt.eeeorrr706225555555
11

---

Under review as a conference paper at ICLR 2021
A EDGEWORTH 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 + oc. 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 Gaussian distribution at large N (Cho & Saul 2009), augmented by higher order cumulants which decay as 1 1/N*/2-1, where is the order of the cumulant. When higher order cumulants are small, an Edgeworth series (see e.g. Mccullagh (2017); |Sellentin et al. (2017) 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 the prior distribution Po[S], i.e the distribution of the DNN output f(x) = Ei (w x), with a; No0,, 3 and
wy ~N d. I). Because has mean and variance that scales 1/N, all odd cumulants zero and the 2r'th cumulant scales as 1 1/N This holds true for any DNN having 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 al. (2017), and our case is similar where instead of vector-valued RV we have the functional RV f(r). so the cumulants become "functional tensors" i.e. multivariate functions of the input σ. Thus, the leading FWC to the prior Polf] is
Po [f]
1 Sarlf] Z
1
1
/ du (21) ...))..aee du (x1) U (21, T3, 1 [f: 3 J2, 13, aAl 41
+O(I/N3)
(A.1)
where SGP [S] is in the main text Eq. and the 4th Hermite functional tensor is
H If1 = / du (x')) ... du (x'4) K (01, 21) ... K- (x1, 2'4) (z1) ... (a44)
K (Fa: 23) / du (z/H) du (x/)) K-1 (x)) K-1 (Ix, 2()) f (x/)) f (x(() [6] (A.2)
+ K (Ir, 28) K (I [3]
where by the integers in [-] we mean all possible combinations of this form, c.g.
K K j0r K Ki3 Kz + K K23
(A.3)
H[f] is the functional analogue of the fourth Hermite polynomial: H. (x) x² 6x + +3, which
appears in the scalar Edgeworth series expanded about a standard Gaussian.
B FIRST ORDER CORRECTION TO POSTERIOR MEAN AND VARIANCE
B. POSTERIOR MEAN
The posterior mean with the leading FWC action is given by
(x.))) = DDfe-Sifff S10 (x.)
j DDfe + O(1/N3)
(B.1)
where
a (f (xa) ya)
S[s] = Scr[[] + [Datal/] + SS
S [
(B.2)
20
where the (1 implies that we only treat the first order Taylor expansion of S(f], and where Sce[f), [f] are as in the main text Eq. g The general strategy is to bring the path integral f Df to the front, so that we will get just correlation functions w.r.t. the Gaussian theory (including the data term [f)) ( namely the well known results (Rasmussen & Williams) 2005) for
12

---

Under review as a conference paper at ICLR 2021
for (2.) (xx)) and Ear(=.)= (z)f(z)))) 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
(Aaraau (0), 22, va, 03, v4) (x.) H (11) - (S (#a))o / dpraa (x), 22, 23, v.) (H (f)))) fa. = 4!
1
(B.3)
This, as standard in field theory, amounts to omitting all terms corresponding to bubble diagrams, namely we keep only terms with factor of (f (xx) f (a%a)) and ignore terms with a factor of
(f (x, Di since these will cancel out. This is 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
1. For the quartic term in H [f], we have
(f.) (2 (a/) (c1))) (f (x.)) (f ( E (x,, (1)) f (2/) [1]] E (x., (1/) f (2.) (2)) [4]
(B.4)
We dub these terms by fEE, and fffz.. 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. I1 in the main text.
2. For the quadratic term in H [f]. have
(f (x.) (z/)) f (2.)) (f (f
- (Ix, fa()) [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 reflection of measure invariance. This is elaborated on $B.3
3. For the constant terms in H [f], we will be left only with bubble diagram terms DD f (xx) which will cancel out in the leading order of I/N.
B.2 POSTERIOR VARIANCE
The posterior variance is given by
E(x.) (S (x.) (x.))
(f.) (f Par 2ferfu + +01//33) Exr(x.) + (S ( - 25r 0(1N22)
(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.))
1
[auiau (1, xA) (x.) (x.)) (S) ( (((.) / duurs (11, 22, uz, x3) (H (S1))
4!
(B.7)
As for the posterior mean, the constant terms in H[S] cancel out and the contributions of the quartic and quadratic terms are
quartic terms E.. E. [12]+ + E.g [12]
(B.8)
and
quadratic terms Ean
(B.9)
13

---

Under review as a conference paper at ICLR 2021
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 p(xx). Here we show how they may in fact be simplified to the compact expressions in the main text Eq. T which involve only discrete sums over the training set and no integrals, and are thus manifestly measure-invariam
For simplicity, we show here the derivation for the FWC of the mean fu (2+). and a similar derivation can be done for Ey (x.). In the following, we carry out the I 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
Aa K a. U8,
K=1 Aaiaz Ka Kalg, Kaye, Möi K32,*
(B.10)
where we defined
A (13, == // du(z1)dp(z2)) (01, 23, T4) K=1 (x1,
(B.11)
Fortunately, this seemingly measure-dependern expression will cancel out with one of the terms coming from the SE,. contribution the quartic terms in H[ʃ]. This not a coincidence and is general feature of the Hermite polynomials appearing in the Edgeworth series, thus for any order in I/N in the Edgeworth series we will always be left only with measure invariant terms. Collecting all terms that survive we have
1
AU aazga Kaje, Kaya, Kayl, Kayg, 9e, - 12Ù mogos Kayo Kaj 3,
Uo.
(B.12)
4!
where we defined
U aaagas == U aiagas -
Uni K K K2
(B.13)
This is more explicit form of the result reported in the main text, Eq. 11
C FINITE 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 (I:) and output (1:) are given by
h (z)
w(2)
Own ( Vd
VN
Vd
(C.1)
(x)
a
(x)))())
VN"
We want to find the 2nd and 4th cumulants of 2 (z). Recall that we found that the leading order Edgeworth expansion for the functional distribution of is
Pk. [b] x - 11(4)) (1) (1 N 26 [h; 2/,, 22 ang. a'y)
(C.2)
where K-1 (z/,
and U (xi, -²ا 2/g, 244) are known from the previous layer. So we are looking for x 1-1
two maps:
K& UK (I, x²) = (o (h (x)) Φ (h (x Prc [A] Up (K, (K, (I1, 22, I3, x4) = (ø (h (z1)) o (h (xz)) o (h (Ir)) & (h (æ4))) Pr.0[M
(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²) Kis (())) (x,
(i++1) (2), 22, X3, 22, 23, x2)
Up (K10) (K(),0()) (o), 22, 23, x34)
(C.4)
0 ((111))
- Kis (((0)
(ay Kod (K10, v(()) (Ing: Kag) [3]
14

---

Under review as a conference paper at ICLR 2021
where the starting point is the first layer (N()) d)
0² (1 x x'
k()) (x, x')
U (21, 02, 23, vz) = 0
(C.5)
NO
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
h1 h2
(c()) (1) K = K()) (21,22)
K()) (#1, 222) K()) (x2, x22)
h ==
(C.6)
find for K222)
K(22) (2))
= 1 dhe 3h kh'k,a) (x11) (h2)
=f
(C.7)
0 (3)
and for (22)
we denote
h1
K(1) (Ii, x1)
K()) (I); [2)
(()) (x) Fz)
K(I) (x1,
h
h2
K
K(1) (2)
K()) (I2; 222)
K (L)) (x2, 23)
K(1) (12, I4))
hs
K()) (33))
K(1) (122 233)
k()) ( 23))
K()) (I3, I4)
h4
K(1) (41,
K(I) (I2,
K() (zz, 244)
K()) (I4, T4)
(C.8)
so that
Up (KW)ww)) uw) (0), 22, 23, x2) =
/ dhe tn tb'K,bb o (hi) Φ (12) @ (ha) o (ha)
(C.9)
This iterative process can be repeated for an arbitrary number of layers.
D FOURTH CUMULANT FOR THRESHOLD POWER-LAW ACTIVATION FUNCTIONS
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) = 2² and ReIU.. 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 2-layer FCN with no biases, input dimension d and N neurons in the hidden layer, such that ø', of ( Ia) is the activation at the ith hidden unit with input Da sampled with uniform measure from (v-i(vd), where (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 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 = (222224)4) (see Eq. 9p. taking for simplicity 2
#4 V)e((L-T) (2T)²
dze L 22
(D.1)
where --1 above corresponds to the matrix inverse of the 4 x 4 matrix with elements Lag = (Ia 2d)/d /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 (948). which yields (with slight modification noted therein) the following series the off off-diagonal elements of the matrix L
P4
Atmpngr
Lt L I L L Lg L
(D.2)
15

---

Under review as a conference paper at ICLR 2021
where the coefficients Atmppgr
-
e
Getm+n
(D.3)
lmm!nllggr!!
For ReLU activation, these G's read
V2T
GReLU
(D.4)
and odd
(-) (26)
2k +
à: = 0, 0,
/2#2*k!
and similar expressions can be derived for other threshold power-law activations of the form ((2) (=))=". 'he 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/Vd) and consequently (xxxx')
O(IVVd) for two random points from the data sets. However, when sum Uni...04 also have terms with repeating indices and so Tais's ' are equal to 1. The above Taylor expansion diverges whenever the 4 matrix Log 8a, has cigenvalues larger than Notably this divergence does not reflect a true divergence of U, but rather the failure of representing it using the above expansion. Therefore 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 re- parameterrization of the z's leading similar but slightly more involved Taylor series.
D.2 DERIVATION OF THE PREVIOUS SUBSECTION
In this section derive the expression for the fourth moment (fif2fafa) of two-layer fully connected network with threshold-power law activations with exponent v: 6(2) = ( v corresponds to step function, y corresponds to ReLU, y = corresponds ReQU (rectified quadratic unit) and so forth.
When the inputs normalized lie the hypersphere, the matrix T. is
L T
L
L12
L23
L24
L
=
L13
L23
L34
(D.5)
L14
L2a
L34
1
where the off diagonal elements here have Lag O (1/Va). We follow the derivation in Ref.
Moran (9488), which computes the probability mass of the positive orthant for a quadrivariate Gaussian distribution with covariance matrix L:
(det(I.-I) (2T)²
/
dze 1-'
P+
(D.6)
The characteristic function (Fourier transform) of this distribution is
sp (f1, t2, 13, t4)
1
T
exp
2
2
Lagtato asB
Σ
exp
exp
(-)etmtntptatt T T. Lp T. T.2,
2
cxp
Alm!nlplqir!
(D.7)
Performing inverse Fourier transform, we may write the positive orthant probability
16

---

Under review as a conference paper at ICLR 2021
1
| dz / dt pp (41, 12,,,,, (3,4)) 14) e- Eam1
P+
(2T)
JRL
00
(-)etmtatptett 11 I L LiA
Σ
lmm!nlplgr!!
E.m.np-a.-rm
(D.8)
{ dt ( 2f w
dz
Jz
(2T)*
Σ
Atmapor L T 14 T T.
12
where the coefficients Atmapgr
(-yttmaatpttt Ge+m+n GrannemmttttttneGnee
Atmuppq
(D.9)
lmm!n!plg!!!
and the one dimensional integral is
1 [-
G(v=0)
dz
t' exp
itz dt
(D.10)
We can evaluate the integral over t to get
((v==)
/
d
22
(D.11)
(-1)² (2m)
1/2
dz
and performing the integral over z yields
= 0
G(u=0)
seven and s 2 2
(D.12)
(2k)
2k k= 0,1,,,
i(2m) 1/22 (/2m)
We can now obtain the result for any integer v by inserting z" inside the z integral:
(
1
dz dz
ʃ ts exp
(-)p itz dt (-i)" (2x) 1/2
[
2"
d
c
22
2T
dz
(D.13)
Using integration by parts we arrive the result Eq. D.4 reported in the main text
0
1
GReLt G((=1)=
(D.14)
2 3 and odd
(2))
2k
0,1,
2
Similar expressions can be derived for other threshold power-law activations of the form 6(2) = (2))2" for arbitrary integer In more realistic setting, the inputs x may not be perfectly normalized, in which case the diagonal elements L are not unity. It amounts to introducing a scaling factor for each of the four z's and makes the expressions little less neat but poses real obstacle.
17

---

Under review as a conference paper at ICLR 2021
E FOURTH CUMULANT FOR QUADRATIC ACTIVATION FUNCTION
For two-layer network, we may write U. the 4th cumulant of the output f(x) Ei= a
with 04 ~ (0 S//N) and w, (O (d)) for general activation function & as
U 02100200300
N (Man, + (9).00)) + (o1,,e)
(E.1)
with
V(a (2)) = 00 For the case of quadratic activation function x(2) z2 the V's read
60)))) (0 (600644)
(E.2)
V(naa))aan = In Tas (Iz2a)² Iu T.44 (Izz3) + Lm2 (Ina)² + Leztas (I13)² )t. {{ (L1a)) (Lea)² + (L14) (L23) ) (8 + L 113 + Lagb1241n + Lnbipbiabzs)tt.
16 (L12241a + L12 + L1311a (E.3)
2
where the linear kernel from the first layer is L(x, x') = x x'. the scaled and non-scaled variances:
Notice that we distinguish between
2
S
(E.4)
N
d
These formulae were used when comparing the outputs of the empirical two-layer network with our FWC theory Eq. [I One can generalize them straightforwardly to 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 with ((M-1).
F AUTO-CORRELATION TIME AND ERGODICITY
As mentioned in the main text, the network outputs JDNN (xr.) are 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 the process converges in both the mean and in the mcan-square sense to the ensemble mean:
lim E [[KDNN (x.; (e,;t)) - (xxx)] =
7-400
(R.1)
lim E (NN (Ix: - T-2
where μ(x,) is the ensemble mean on the test point v, and the time-average estimate of the mean over time window T is
1 7 JNN 7 (ONN(NN: (z.:t)dt= x
1'
1 :=
(DNN (I.it)))))
Σ DNN (Ixitz)
(F2)
7
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 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 Fg f.NN (z.: (1) Divide the time series data in the matrix F into non-overlapping sub-matrices, each of dimension nsodd X neporhs (2) For each these sub matrices, find f (x.) i.e. the empirical dynamical average across that time window and across the chosen seeds; (2) Find
18

---

Under review as a conference paper at ICLR 2021
the empirical variance Cemp (xx) across these f(x.); (4) Repeat (1)-(3 for other combinations of Nepochs, needds If ergodicity holds, we should expect to see the following relation
Cemp (I+)
(F.3)
RepochsReeedd
where 7 the auto-correlation time of the outputs and o the macroscopic variance. The results of this procedure are shown in Fig. F.1 where we plot on log-log scale the empirical variance oe vs. the number epochs Repochs 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 in the figure) yields a slope of approximately -, which is strong evidence for ergodic dynamics.
var vs. u_ n_epochs
slope -0.984
intercept=1.502
-250
-3.00
logi polemp.no var)
-3.25
-3.50
-3.75
4.00
4.25
4.4
4.6
4.8
5.0
5.2
log(n_epochs)
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 are averages across these.
Figure F.1: Ergodicity check. Empirical variance cmmp (x..)
G NUMERICAL EXPERIMENT DETAILS
G.1 FCN EXPERIMENT DETAILS
We trained 2-layer FCN on a quadratic target y(z) = xT Ar where the x's are sampled with uniform measure from the hyper-sphere 1V-r(Vd), with d = 16 and the matrix elements are sampled as A, -0, 1) and fixed for all z's. For both activation functions, we used a training noise level of x² 0.2, training set of size n 110 and weight decay of the first layer Tw 0.05. Notice that for any activation é, K scales linearly with x² a,. = (T/ra) N, thus in order to keep K constant we vary N' we need to scale the weight decay of the last layer as %o O(N). This done in order to keep the prior distribution in accord with the typical values of the target N varies, so that the comparison fair.
We ran each experiment for 2- 10 epochs, which includes the time it takes for the training loss to level off, which is usually on the order 10 epochs. the main text we showed GP and FWC results learning rate of dt 0.001. Here report in Fig. G.1 the results using dt {0.003, 0.001, 0.0005). For learning rate dt 0.003 and width N 2 1000 the dynamics become unstable and strongly oscillate, thus the general trend is broken, as seen in the blue markers in Fig. G. The dynamics with the smaller learning rates are stable, and we see that there is a convergence to very similar values up expected statistical error.
19

---

Under review as a conference paper at ICLR 2021
GP-net & FWC-net MSES vs width: seed averaged
logiGP-net and FWC--et MSE)
X
x
-3
0.003
0.003 FWC
0.001 GP
X
0.001 FWC
0.0005 GP
X
0.0005 FWC
1.8 2.0 2.2
2.4
2.6
2.8
3.0
3.2
log(width)
Figure G.1: Regression task with fully connected network: (un-normalized) MSE vs. width on log-log scale (base 10) for quadratic activation and dif- ferent leaning rates. The learning rates dt 0.001, 0.0005 converge to very similar values (recall this is log scale), demonstrating that the learning rate sufficiently small that the discrete-time dynamics is good approximation the continuous matinuuuus-time dynamics. For a learning rate of dt = 0.003 (blue) and width N 2 1000 the dynamics become unstable, thus the general trend is broken, so one cannot take the di to be too large.
G.2 CNN EXPERIMENT DETAILS AND ADDITIONAL SETTINGS
The CNN experiment reported in the main text was carried as follows.
used random sample of 10 train-points and 2000 test points from the CIFARIO 10 datasct, and in App. H we report results on 1000 train-points and 1000 test points, balanced in terms of labels. To MSE loss, the categorical labels were one-hot encoded into vector of zeros and one.
Dataset: In the main text Fig.
Architecture: used 6 convolutional layers with ReLU non-linearity, kernel of size 5 x 5, stride of 1, modpadding, no-pooling. The number of input channels was for the input layer and C' for the subsequent CNN layers. We then vectorized the outputs of the final layer and fed it into ReLU activated fully- cully-connected layer with 25C outputs, which were fed into linear layer with 10 outputs corresponding to the 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 the gradients to generate o² 0.2 the NNGP-corresoondenc; layer-dependant weight decay and bias decay which implies (normalized by width) weight variance and bias variance of o 2 and of we saved, every 1000 epochs, the outputs of the CNN every test point. We note in passing that the standard deviation of the test outputs around their praning-time--veragee value was about 0.1 per CNN output. Training was carried for around half million epochs which enabled to reach a statistical error of about 10-4, in estimating the Mean-Squared Seee--Suuard-DDisrrpaang between the training-tinn sinting--imm-vverageee CNN outputs and our NNGP predictions. Notably
respectively, when trained with no-data. During training
best agreement between the DNN and GP occurred 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 get estimates of the average outputs fF(xa) ) f(xa) + of. since
20

---

Under review as a conference paper at ICLR 2021
the estimators' fluctuations, ffa, scale as (r (7/ttraining) 1/// where is cuto-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 of averaged over all a's i' indeed
E (fe(Ea)- -ataaaaaa))--(taaaaaaaaaaae
((f(ra)
Bennt
Σ (fe(sa)-fcr(taa
Σ (ofa)²
(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(()) 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 Σ. (ōfa)² )² are of the order 0.12, the amount of uncorrelated samples of CNN outputs we require should be much larger than 0.13/(7. . 10-3) R 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 ofa- 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 be the time it takes the train-loss 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- 10 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 wiith. 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.
0.07
0.20
0.6
0.508
485
0.006
0.44
0.84
0.000
0.50
Figure G.2: MSE between 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).
Novokk
H FURTHER NUMERICAL RESULTS ON CNNS
Here report two additional numerical results following the CNN experiment carried (for details see App. G Fig. H.3b is the same as Fig. H.3a from the fact that we subtracted our estimate of the statistical bias of our MSR estimator described in App.
21

---

Under review as a conference paper at ICLR 2021
100
300
038
0.18
034
ASC/UOP -02
500000000
0.00
0.50
0.03
002 0.50
0.00
Athannets
#Chareess
(a) 1000
(b) subtract bias of MSE estimator
Figure H.3: CNNs trained on CIFARIO in the regime of the NNSP corre- spondence 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 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 s GP and the trained DNN on the log-log scale was -1.77, placing us close to the perturbartive regime where slope of -20 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 an effective 0² 0.1 for the GP. The weight decay on the biases was taken to be ten times larger leading to of 0.1 instead of a; 1.0
as before. For C S 80 we used a learning rate of dl = 5 10 after verifying that reducing it further had no appreciable effect. For C < 80 we used di 2.5 20- For c S 80 we used 6 10+5 training epochs and we averaged over 4 different initialization seeds. For C > S0 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.
I THE FOURTH CUMULANT CAN DIFFERENTIATE CNNS FROM LCNS
Here we show that while the NNGP kernel K of 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 ID image with one color/channel (X,) as input i E (O.... L - 1): (2) A single convolutional layer with some activation & acting with stride I and no-padding using the conv-kernel T. where E {1.....C) C} is a channel number index and {O..... 2/} } is the relative position in the image. Notably, in LCN this conv-kernel will receive an additional dependence on J. the location on X, on which the kernel acts. (3) A vectorizing operation taking the C outputs of each convolutional around point ă (Z...., L - 1), into single index E (O..... C'LL } (4) A linear fully connected layer with weights W where E {0..... }outputs} are the output indices.
Consider first the NNGP of such a random DNN with weights chosen according to some iid Gaussian distribution Po(w), with including both Wo and TG. Denoting by =²(x) the o'th output of the CNN, for an input we have (where we denote in this section (...) == (.. )o(w)
and (I,x') (2 (x)2" (a')) 800
Σ 0,6' ,5,5"
(i')X++8
(L.1)
22

---

Under review as a conference paper at ICLR 2021
The NNGP kernel of an LCN is the same that of a CNN. This stems from the fact that (Wa W c' yields a Kronecker delta function on the 2, z' indices. Consequently, the difference between LCN and CNN, which amounts whether (2) is the same (CNN) or different (LCN) random variable than T T (7'), becomes irrelevant as the these two are never averaged together.
For simplicity, we turn to the fourth cumulant of the output, given by
... = (z"(x)
(2
- (x) ='(x4)) K (Ioa Fg)) ) (2n, 23)) 2x){3
(1.2)
with the second term on the LHS implying all pair-wise averages of 2 (21))_22"(x Note that the first term on the LHS is not directly related to the kernel, thus it has a chance of differentiating CNN from an LCN. Explicitly, it reads
Σ 01-66411...4
(Wa
W Wa a (6( (51) (2:)x=:+8.))) -1) O (2))
(I.3)
The average over the four W's yields non-zero terms of the type Wo W We'z W22 with either ã = &' (type 1), 2 H' and с + c' (type 2), or ā o X' and c = è' (type 3).
The type contribution cannot differentiate an LCN form a CNN since, as in the NNGP case, they always involve only one T. The type 2 contribution also cannot differentiate since it yields
(Wo WaW
(
espc'
(I.4)
Examining the average involving the four T's, one finds that since (8) is uncorrelated with T (i') for both LCNs and CNNs, it splits into
Σ CF
(Wo We Wa W XOTT
(L.5) where in the NNGP, two T's with different 3 are never averaged together and we only get 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 (( ( ((ry [3). apart from "diagonal" contribution
(a .').
We turn attention to the type term given by
(Wo Wo (Wo car
W
(I.6) Examining the average involving the four s, one now finds sharp difference between an LCN and CNN. For an LCN, this average would split into product of two K's since ( would be uncorrelated with (z'). For CNN however, (ā()) the same random variable as (T')") and therefore the average does not split giving rise to distinct contribution that differentiates a CNN from an L.N.. Notably, small by a factor of 1/C owing to the fact that contains a redundant summation over c-index while the averages over the four W's contain 1/C² factor when properly normalized.
J CORRECTIONS TO EK
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 (z. )) averaged ((... . )n) over all draws of datasets of size n' with n'' taken from a Poisson distribution with mean is given by
Zn = Dfe Serlfl Sn dptp))(() (1 Stlf)) O(1NN3)
f
(J.I)
with Sc and [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
23

---

Under review as a conference paper at ICLR 2021
finite dataset, with two main differences: all the sums the training set have been replaced by integrals with respect to the measure, Ma. from which data points are drawn. Furthermore o² is now accompanied by 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 [f]: U2 K=1 æiي K aaz! f(a() ... f(44), and from the diagram Weeee/Wickkoooooeeeee contraction where we take the expectation value out of the 4 f's in this quartic term, to arrive at expression which is ultimately cubic in the targets y
Un aa 2 (f(z))) _0
22d, (f(22)) T3gs (f(aa)) K Eso (10
(J.2)
where we recall that (f(z))
K Kay K Ko" being the posterior covariance the EK limit, where Kny f(x') Ke ((x))+
KZ and (x1
K21,22
(o² (z). Using the fact that K Kr gives delta function w.r.t. the measure, the integrals against K 7ang can be easily carried out yielding
(Un Un ang R-1
"aina K Keisss)) R R-y Kiint Kazins "agin v (((()((()))
(I.3)
Introducing the discrepancy operator 8gwr == Sgg' KE K Rwls o² K oz's compact expression
we can write a more
71 (
(J.4)
This with the additional 1/4! factor times the combinatorial factor of 4 related to choosing the "partner" of (x.) in the Wick contraction, yields an overall factor of 1/6 as in the main text, Eq. 14 The other term therein, which linear in y, is result of following similar steps with the fE2. contributions that do not get canceled by the quadratic part in [f].
24


Reviews:
[{"review_id": "jSUZUzWLka", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "It has been show that practical performance of finite width DNNs deviates from the infinite limiting cases of both the NTK and GP. The reason behind this is still not well understood. This papers develops theory trying to explain this by showing the finite width correction (FWC) for DNNs. The results look very interesting. However, as an educational guest on this topic, I find the paper written in a way that causes a lot of confusions, making me not fully understand the results.\n\nFirst of all, the paper is talking about the FWC. Although it has been used in previous work and I can get some sense of the meaning, I think a formal definition is needed in order to make the paper more readable.\n\nRegarding the technical details:\n\n1. What is Df in eq.4? And why can it be eliminated? And it is not clear to me about the sentence below eq.4: why the rhs of eq.4 equals the kernel of the NNGP only in highly over-parameterized DNNs?\n2. I am not sure how eq.5 is obtained?\n3. I think the description below eq.5 can be made more formal because this gives the definition of NNSP. Based on the current description, I still cannot get why NNSP is special? Is it still a GP and why?\n4. Second line below section 3.1, it says e^{-L[f]/2\\sigma^2} is independent of the DNNs, but why? I think f corresponds to the DNN?\n5. In Section 3, the paper derives complicated formulations for the posterior mean and variance. I am not sure how these results are useful? It seems to say that these results describe how the finite width DNNs are deviated from the NTK and NNGP, but by looking at these formulas, it is hard for me to figure out why NNSP is better than NTK and NNGP?\n6. Paragraph above Figure 1: it says \"the above cubic term would lose its explicit dependence on n\", but I think there is still a quadratic term, which still depends on n. How can it say that the FWC is negligible in the large n regime?\n\nOverall, I found this paper addressing an important and interesting problem, but the current presentation cannot convince me a pass.", "extra_scores": {}, "extra_text": {}}, {"review_id": "bisK0AMRkG0", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": " In this paper, the authors propose an analytical framework for DNN training with noisy gradient descent and weight decay. For infinite width, the authors have established a correspondence between DNNs trained with\nnoisy gradients and the NNGP, and for finite width, the authors introduce a finite width correction (FWC). Overall, the paper is quite interesting and well organized. However, I still have several concerns.\n\nWhile I agree with the authors that this is an analytical framework different from the Neural Tangent Kernel, my major concern is that this paper doesn't show the advantage of the new framework because this paper hasn't provided any theoretical result for convergence or generalization. \n\nThe authors have commented in the introduction that for NTK \"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.\" However, as far as I know, there are plenty of good works that have considered NTK trained with SGD([1], [2], [3]). The authors may want to comment on these papers. In the infinite width limit, the authors establish a correspondence between DNNs trained with noisy gradients and the NNGP. Actually, there is a line of work that directly studies the infinite wide NN trained with noisy gradients and weight decay regularization in the so-called mean-field regime([4], [5], [6]).  It would be much better if the authors can characterize the difference between NNGP and those papers I have mentioned.\n\nMy other minor concerns are as follows:\n\n1. In the eq 2 of this paper, it seems that the probability measure of the parameter is very close to the gaussian with \\sigma_{w}^2 variance. Would it be better for the NN to initialize at \\sigma_{w}^2 too? If so, whether the parameters will still only make a small change in the distribution space?\n\n2. The analysis in this paper highly depends on the ergodicity of the dynamics. The authors mention that a small ACT of the dynamics implies this key property, especially for the output. It would be better if the authors can explain more about why the ACT can imply ergodicity.\n3. There is no explicit theorem, lemma, or proof in this paper, which is quite hard to read or follow. I suppose the finite width correction would be the key theorem of this paper?\n\nReferences:\n\n[1] Allen-Zhu, Z., Li, Y., & Song, Z. (2019, May). A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning (pp. 242-252). PMLR.\n\n[2] Zou, D., Cao, Y., Zhou, D., & Gu, Q. (2020). Gradient descent optimizes over-parameterized deep ReLU networks. Machine Learning, 109(3), 467-492.\n\n[3] Daniely, A. (2017). SGD learns the conjugate kernel class of the network. In Advances in Neural Information Processing Systems (pp. 2422-2430).\n\n[4] Mei, S., Misiakiewicz, T., & Montanari, A. (2019). Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. arXiv preprint arXiv:1902.06015.\n\n[5] Chen, Z., Cao, Y., Gu, Q., & Zhang, T. (2020). Mean-Field Analysis of Two-Layer Neural Networks: Non-Asymptotic Rates and Generalization Bounds. arXiv preprint arXiv:2002.04026.\n\n[6] Tzen, B., & Raginsky, M. (2020). A mean-field theory of lazy training in two-layer neural nets: entropic regularization and controlled McKean-Vlasov dynamics. arXiv preprint arXiv:2002.01987.\n\n\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "80j4_zdUzeO", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies equivalence between neural networks and Gaussian process. The authors make three main contributions: (1) they prove that infinite-width DNNs trained with noisy gradients are corresponded to neural network Gaussian process (NNGP), not neural tangent kernel (NTK); (2) they prove that the finite width correction (FWC) for DNNs with arbitrary activation functions and depth has a general analytical form which can help predict the outputs of empirical finite networks with high accuracy; and (3) they show how these FWCs can improve the performance of finite convolutional neural networks (CNNs) on image classification tasks.\n\nOverall, I vote for accepting. The paper establishes solid correspondence between neural networks and Gaussian process which may have a real impact to applications. The authors also give empirical results to support their arguments.\n\nPros:\n\n+ The paper takes one of the most important problem in deep learning theory – the correspondence between neural networks and Gaussian kernels.\n\n+ The paper makes real contributions compared with existing works on NNGP (in NNGP, weights were determined by the distribution of the DNN weights at initialization; and in this paper, the weights are sampled across the stochastic training dynamics).\n\n+ Give results for finite-width networks, while many works only give asymptotical analysis.\n\n+ Explain why finite CNNs outperforms infinite CNNs and quantify the difference.\n\nCons:\n\n- The presentation can be improved. For example, Sections 2 and 3 would better if the contents can be further organized into smaller parts.\n\nQuestions: \n\n* How different is the techniques used in the proves comparing to existing results? I have gone through the proves and the techniques wherein seem similar to existing works for me. Could you please give detailed proof skeletons in the manuscript and discussion the differences with related works?", "extra_scores": {}, "extra_text": {}}, {"review_id": "55PU86fd6UA", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThis paper is written in a physics style which makes it hard to read. Many notations are defined in words instead of mathematical formulas. This is quite annoying since it is hard to follow what are scalers, vectors, matrices, functions, random variables, and what are the arguments of the functions.\nFor the first claimed contribution, I feel that this is a simple observation rather than a significant contribution: the authors just pointed out that the stationery distribution of Langevin dynamics is a Gibbs measure.\nFor the third claimed contribution, I feel that the authors over-claimed what their results can suggest. The authors showed that the FWC contains the information of the weight sharing structure in CNN, and then cited a 'common practice' to argue that this weight sharing structure will lead to better test error. This doesn't give a convincing argument that finite width CNN will outperform its corresponding GPs.\n\nThe derivation of the finite width correction of the NNGP could indeed be a concrete contribution, but the result is presented in a very ambiguous way. For example, the $U_{x1, x2, x3, x4}$ quantity seems to be a very important quantity in defining this finite width correction. However, its definition in Eq. (9) is quite ambiguous. Especially, U_{x1, x2, x3, x4} depend on $\\phi_\\alpha = \\phi(z_i^{l-1}(x_\\alpha))$, and I didn't see where does the subscript $i$ come from. The authors used the Einstein convention in subscripts but I don't know what are ranges of \\alpha, \\beta, \\gamma, and \\delta to be summed over (is it 1 to n or 1 to 4?). There are many other places where the quantities are not concretely defined and I can only guess what do these notations mean in order to follow the results.\n\nAlthough the finite width correction of NNGP (if properly written) could be an interesting result, due to the ambiguous writing style of the paper, I don't think it achieves the goal of clearly presenting the results. I will suggest a rejection.\n\nAfter reading the rebuttal: \nI increased my score to 5. I still feel that the writing style is hard to follow. Besides the examples that I wrote in my original review, there are many other places where the notations and definitions are not clearly written. ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper develops an interesting new angle on the behavior of large-width neural networks by elucidating the connection between the NNGP and noisy gradient descent and by examining finite-width corrections through an Edgeworth expansion. While these contributions are important, the paper would better serve the community if its presentation were significantly improved before publication. The main issue is not one of presentation style -- papers with physics-style prose are welcomed and appreciated at ICLR -- but rather one of presentation substance. In addition to the various specific points raised by the reviewers, I would add that the figures and captions are difficult to interpret, the experiments need a more in-depth discussion, and the notations should all be defined at the time of their introduction, among other things. For these reasons, I cannot recommend accepting the paper in its current form, but I hope to see a more polished version of the manuscript at a subsequent conference.
Decision: Reject