A self consistent theory of Gaussian Processes captures feature learning effects in finite CNNs
Authors: Gadi Naveh
Zohar Ringel
Venue: NeurIPS 2021 Poster (2021)
Keywords: Statistical mechanics
Mean field theory
Gaussian Processes
feature learning
finite DNNs

Abstract: Deep neural networks (DNNs) in the infinite width/channel limit have received much attention recently, as they provide a clear analytical window to deep learning via mappings to Gaussian Processes (GPs). Despite its theoretical appeal, this viewpoint lacks a crucial ingredient of deep learning in finite DNNs, laying at the heart of their success --- \textit{feature learning}. Here we consider DNNs trained with noisy gradient descent on a large training set and derive a self-consistent Gaussian Process theory accounting for \textit{strong} finite-DNN and feature learning effects. Applying this to a toy model of a two-layer linear convolutional neural network (CNN) shows good agreement with experiments. We further identify, both analytically and numerically, a sharp transition between a feature learning regime and a lazy learning regime in this model. Strong finite-DNN effects are also derived for a non-linear two-layer fully connected network. We have numerical evidence demonstrating that the assumptions required for our theory hold true in more realistic settings (Myrtle5 CNN trained on CIFAR-10).
Our self-consistent theory provides a rich and versatile analytical framework for studying strong finite-DNN effects, most notably - feature learning.

A self-consistent theory of Gaussian Processes captures feature learning effects in finite CNNs
Gadi Naveh
Zohar Ringel
Racah Institute of Physics Hebrew University Jerusalem, Israel 91904
Racah Institute of Physics
Hebrew University
Jerusalem, Israel 91904
gad, mintmmmailml....mmmmmmmmmmmmm hujuuii..... auuji.ac.ii ac.
zohar. wawgerrineellln moaggee@maal...ie www.hh................ ac. il
Abstract
Deep neural networks (DNNs) in the infinite width/channel limit have received much attention recently, they provide clear analytical window to deep learning via mappings to Gaussian Processes (GPs). Despite its theoretical appeal, this viewpoint lacks crucial ingredient of deep learning in finite DNNs, laying the heart of their success - feature learning. Here we consider DNNs trained with noisy gradient descent Gaussian Process theory accounting for strong finite-DNN and feature learning effects. Applying this to a toy model of two-layer linear convolutional neural network (CNN) shows good agreement with experiments. We further identify, both analytically and numerically, sharp transition between feature learning regime and a lazy learning regime in this model. Strong finite-DNN effects are also derived for non-linear two-layer fully connected network. We have numerical evidence demonstrating that the assumptions required for our theory hold true in more realistic settings (Myrtle5 CNN trained on CIFAR-10), Our self-consistent theory provides rich and versatile analytical framework for studying strong finite-DNN effects, most notably- feature learning.
large training and derive a self-consistent
Introduction
The correspondence between Gaussian Processes (GPs) and deep neural networks (DNNs) has been instrumental in advancing our understanding of these complex algorithms. Early results related randomly initialized strongly wer-parraeteriiee DNNs with GP priors [33, 24, 30]. More recent results considered training using gradient flow (or noisy gradients), where DNNs, potentially following some ensembling, map to Bayesian inference on GPs governed by the neural tangent kernel [21. 25, 18] (or the NNGP kernel [32]). These correspondences carry to wide variety of architcctures, going beyond fully connected networks (FCNs) to convolutional neural networks (CNNs) [4, 36), recurrent neural networks (RNNs) [2] and even attention networks [20]. They provide us with closed analytical expressions for the outputs of strongly wer-paaaaeterizze trained DNNs, which have been used to make accurate predictions for DNN learning curves [11, 9, 8].
Despite their theoretical appeal, GPs are unable to capture feature learning [47, 46], which is well-observed key property of trained DNNs. Indeed, it was noticed [21] that as the width tends to infinity, the neural tangent kernel (NTK) tends to a constant kernel that does not evolve during training and the weights in hidden layers change infinitesimally from their initialization values. This regime of training was thus dubbed lazy training [10]. Other studies showed that for CNNs trained on image classification tasks, the feature learning regime generally tends to outperform the lazy regime [15, 14, 23]. Clearly, working in the feature learning regime is also crucial for performing transfer learning [45, 46].
35th Conference on Neural Information Processing Systems (NcurIPS 2021),

---

It is therefore desirable to have a theoretical approach to deep learning which enjoys the generality and analytical power of GPs while capturing feature learning effects in finite DNNs. Here we make several contributions towards this goal:
1. We show that the mean predictor of finite DNN trained on large data set with noisy gradients, weight decay and MSE loss, can be obtained from GP regression on a shifted target ($3). Central to our approach is non-linear self-consistent equation involving the higher cumulants of the finite DNN (at initialization) which predicts this target shift.
2. Using this machinery on a toy model of a two-layer linear CNN in a teacher-student setting, we derive explicit analytical predictions which are in very good agreement with experiments even well away from the GP/lazy-learning regime (large number of channels, C) thus accounting for strong finite-DNN corrections (64.1). The match empirical values of our theory in this toy model is shown to be clearly superior to two alternative theoretical predictions, one of them being a purely perturbative theory. Similarly strong corrections to GPs, yielding qualitative improvements in performance, are demonstrated for the quadratic two-layer fully connected model of Ref. [28].
3. We show how framework can be used to study statistical properties of weights in hidden layers. In particular, in the CNN toy model, we identify, both analytically and numerically, sharp transition between feature learning phase and lazy learning phase ($4.1.4). We define the feature learning phase as the regime where the features of the teacher network leave a clear signature in the spectrum of the student's hidden weights posterior covariance matrix. In essence, this phase transition is analogous to the transition associated with the recovery of a low-rank signal matrix from noisy matrix taken from the Wishart ensemble, when varying the strength of the low-rank component [6].
1.1 Additional related work
Several previous papers derived leading order finite-DNN corrections to the GP results [32, 44, 13, 39]. While these results are in principle extendable to any order in perturbation theory. such high order expansions have not been studied much, perhaps due to their complexity. These previous perturbative approaches are expected to be satisfactory only when the GP limit already gives a reasonable approximation to the DNN behavior. In contrast, we develop analytically tractable non-perturbative approach which we find crucial for obtaining non-negligible feature learning and associated performance enhancement effects.
Previous works [15, 14, 43] studied how the behavior of infinite DNNs depends on the scaling of the top layer weights with its width. In [45] it is shown that the standard and NTK parameterizations of a neural network do not admit an infinite-width limit that can learn features, and instead suggest an alternative parameterization which can learn features in this limit. While unifying various viewpoints on infinite DNNs, this approach does not immediately lend itself to analytical analysis of the kind proposed here. Also, in our work, feature learning as a purely finite-width effect.
Several works [38, 27, 16, 3] show that finite width models can generalize either better or worse than their infinite width counterparts, and provide examples where the relative performance depends on the optimization details, the DNN architecture and the statistics of the data. Another study [1] considered cases where the enhanced flexibility of a finite DNN relative to its infinite width counterpart provides the former with superior performance. Here we demonstrate analytically that finite DNNs outperform their GP counterparts when the latter have a prior that lacks some constraint found in the data (e.g. positive-definitenee [28] or translation invariance / cquivariance [36]).
Deep linear networks (FCNs and CNNs) similar to our CNN toy example have been studied in the literature [5. 40, 22, 17]. These studies use different approaches and assumptions and do not discuss the target shift mechanism which applies also for non-linear CNNs. In addition, their analytical results hinge strongly on linearity whereas our approach could be useful whenever several leading cumulants of the DNN output are known or can be approximated, and we give an example of this in a non-linear setting in $4.2. While there are some similarities between the phase transition we discuss in $4.1.4 44 and the one appearing in Ref. [22], there are several important differences, to mention one of them: that study considers the statistics of a noisy teacher, whereas we consider the statisties of the first layer of the trained student network.
2

---

Two concurrent works [48, 34] derived exact expressions for the output priors of finite FCNs induced by Gaussian priors over their weights. However, these results only apply to the limited case of prior over single training point and only for FCN. In contrast, our approach applies to the setting of large training set, it not restricted to FCNs and yields results for the posterior predictions, not the prior. Focusing on deep linear fully connected DNNs, recent work [26] derived analytical finite-width renormalization results for the GP kernel, by sequentially integrating out the weights of the DNN, starting from the output layer and working backwards towards the input. Our analytical approach, its scope, and the models studied here differ substantially from that work.
Preliminaries
We consider fixed set of n training inputs | C d and a single test point x, over which we wish to model the distribution of the outputs of DNN. We consider generic DNN architecture where for simplicity we assume a scalar output f(x) E R. The learnable parameters of the DNN that determine its output, training set and on the test point into vector f (f ...,,...,, (x,). ) (x++1)) R=+1, where we denoted the test point as x, = Xn+1. We train the DNNs using fulng full-batch gradient decent with weight decay and external white Gaussian noise. The discrete dynamics of the parameters are thus
collected into single vector 0. We pack the outputs evaluated on the
0c+1- 02 = (70, + VoL (Jo)) n + 20 Tic
(1)
where Oe is the vector of all network parameters at time step y is the strength of the weight decay, ((f)) is the loss as a function of the DNN output fo (where we have emphasized the dependence on the parameters 0), a is the magnitude of noise, 77 is the learning rate and & (N I). As - 0 this discrete-time dynamics converge to the continuous-time Langevin equation given by à (t) V (3 (1)()) + (fa) 2n5 (t) with (t), 6,,ō (t t'), so that as oo the DNN parameters 0 will be sampled from the equilibrium Gibbs distribution P(0).
As shown in [32], the parameter distribution P(0) induces a posterior distribution over the trained
DNN outputs P(f) with the following partition function:
z - afPo (i) exp (-20e (1Jm)) (ow))m)))
(J)=
n+1
+
JuSR
(2)
#=1
Here Po( is the prior generated by the finite-DNN with 0 drawn from (N, 2002/02)) where the weight decay may be layer-dependent, {- are the training targets and J are source terms used to calculate the statistics of f. We keep the loss function L arbitrary at this point, committing to a specific choice in the next section. As standard [19], to calculate the posterior mean any of the training points or the test point x,+1 from this partition function one uses
(fa) 0s, log Z (J))
Vu E fI...., +1
(3)
J=
3 A self-consistent theory for the posterior mean and covariance
In this section we show that for a large training set, the posterior mean predictor (Eq. 3) amounts to
GP regression on shifted target (g,, ~ 9p Ag This shift to the target (Agn) is determined by solving certain self-consistent equations involving the cumulants of the prior Po( f). For concreteness, we focus here on the MSE loss L = Σp= (fa g,,)² and comment on extensions to other losses, e.g. the cross entropy, in App. C. To this end, consider first the prior of the output of a finite DNN. Using standard manipulations (see App. A), it can be expressed follows
a+1 Σ Ulps +
Po (f)
di d'exp exx
Σ AL...,
Kp ilps ilp
2x11
where Ky.,.... is the r'th multivariate cumulant of Po( [31]. The second term in the exponent is the cumulant generating function (CGF), denoted by (()), corresponding to Po. As discussed in
3

---

App. B and Ref. [32], for standard initialization protocols the r'th cumulant will scale as 1/C(r/2-1), where C' controls the prrepaaaaettiiiiionnn e.g. number of neurons / channels in each layer for
FCNs/ CNNs, respectively. The second (r = 2) cumulant which is C-independent, describes the NNGP kernel of the finite DNN and is denoted by [(xp) Xpa) = -/
Consider first the case of C' + oo [24, 30, 33] where all > 2 cumulants vanish. Here one can explicitly perform the integration in Eq. 4 to obtain the standard GP prior Po (f) x
exp 2 Emtpp=i Hap fpz) Plugging this prior into Eq. 2 with MSE loss, one recovers
standard GP regression formulas [37]. In particular, the predictive mean x. is: (f (x.))
nu Iv where K2 K (x+; and Kur K (x,, x,) + o 0 Another set of quantities shall find useful the discrepancies in GP prediction, which for the training set read
Emv=1 KK
Aμ E m., (ōgu) = 9p (x,)) = 94
Kay K Kw v
(5)
Saddle-point approximation for the mean predictor. For a DNN with finite C, the prior (f()) will no longer be Gaussian and cumulants with would contribute. This renders the partition function in Eq. 2 intractable and so some approximation is needed to make progress. To this end we note that can be integrated out (see App. A. to yield a partition function of the form
2 (J)
/
dii ......dlyy .....ddee 2
dlne d
(6)
where S(T, J) the action whose exact form given Eq. A.14. Interestingly, the il, variables appearing above are closely related to the discrepancies ōgn. in particular (it,)) = (ógm)/
To proceed analytically we adopt the saddle-point (SP) approximation [12] which often relies having a partition function of the form Z f dde-nS(t) where is a large number and S is 0(1). In our settings we cannot simply extract such a large factor from the action and make 0(1). Nonetheless, we argue in App. A.4 that the saddle-point is still good approximation for large (training set size). This relies on the fact that the non-linear terms in the action comprise of many it,.' ss Given that this sum is dominated by collective effects coming from all data points (as opposed to only a selected few), expanding S(T, J) around the saddle-point yields terms with increasingly negative powers of n.
sum
For the training points μ E (1....,n)) ou, S (E.J) Tーū
n). taking the saddle-point approximation amounts setting O. This yields set of equations that has preciscly the form of Eq. 5, but where
the target is shifted as 9r + gv~ Agr and the target shift is determined self celf-consistently by
Ag. Σ (T 1)!
20
1
T
Kepp (o (8mm)) (..p)..)((())) (o-00000)))) 28pp....)
M
Equation 7 is thus an implicit equation for Ag, involving all training points, and it holds for the training set and the test point v E (I...., 1}. Once solved, either analytically or numerically, one calculates the predictions on the test point via
(f-) Ag. + Σ
Ki (g. Agu)
(8)
Equation 5 with 9r - 9v Ag, along with Eqs. and are the first main result of this paper. Viewed as an algorithm, the procedure to predict the finite DNN's output on test point x. is as follows: we shift the target in Eq. as + g Ag with Ag as in Eq. 7, arriving at a closed equation for the average discrepancies (6gz) on the training set. For some models, the cumulants Ke......... can be computed for any order and it can be possible to sum the entire series, while for other models several leading cumulants might already give reasonable approximation due to their 1/Cr/2-1

---

scaling, The resulting coupled non-linear equations can then be solved numerically, to obtain Agp from which predictions on the test point are calculated using Eq. 8.
Notwithstanding. solving such equations analytically is challenging and one of our main goals here to provide concrete analytical insights. Thus, in $4.1.2 propose an additional approximation wherein to leading order we replace all summations over data-points with integrals over the measure from which the data-set is drawn. This approximation, taken in some cases beyond leading order as in Ref. [1I], will yield analytically tractable equations which we solve for two simple toy models, one of linear CNN and the other of non-linead FCN.
Saddle-point plus Gaussian Iluctuations for the posterior covariance. The SP approximation can be extended compute the predictor variance by expanding the action S to quadratic order in the deviation from the SP value ot, = Ey - (see App. A.3). Due to the saddle-pons being an En., 6t jt, This leaves the previous SP approximation
extremum this leads to S = SsP for the posterior mean the training set unaffected (since the mean and maximizer of a Gaussian coincide), but is necessary to get sensible results for the posterior covariance. Using the standard Gaussian integration formula, one finds that Auw is the covariance matrix of it,.. Performing such an expansion one finds
Anu (0" + KE + AKWw)
AKμ Di. Oic, (itt,....., itn)
where the it's on the r.h.s. are those of the saddle-point and Ć is the CGF C without the second cumulant (see App. A.1). This gives an expression for the posterior covariance matrix on the training set:
Ene (faft) (fr) -o' - AK jer 0
(10)
where the r.h.s. coincides with the posterior covariance a GP with kernel equal to K + AK [37]. The variance the test point is given by (repeating indices are summed over the training set)
E.. K.. K + (02,, Cli + ((9g.K7tt)) - (Ao.) (K2il)) + Var (Ag.)
(11)
where here Ag. is in Eq. 7 but where the (o 20-28 's replaced the it, 's that have Gaussian fluctuations. The first two terms in Eq. 11 yield the standard result for the GP posterior covariance matrix on a test point [37], for the case of AK 0 (see Eq. 9). The rest of the terms can be evaluated by the SP plus Gaussian fluctuations approximation, where the details would depend on the model at hand.
4
Experiments
4.1
The two layer linear CNN
4.1.1
Setting of the model and its properties
Here we define teacher-student toy model showing several qualitative real-world aspects of feature learning and analyze it via our self-consistent shifted target approach. Concretely, we consider the simplest student CNN (xx,, having one hidden layer with linear activation, and a corresponding teacher CNN, g(x)
N C
N
c
fxx = ΣV ai,. X,
(xx)
Σ
a
w x,
(12)
i=1 aml
i=1
c=l
This describes a CNN that performs 1-dimensional convolution where the convolutional weights for each channel are we E R.S. These are dotted with convolutional window of the in-
put x; = (Ds(i-1)4)...... (NN.)., (K) xx) N.S. Namely, the input dimension is d NSS where N is
E I and there are no overlaps between them so that x
5

---

the number of (non-overlappinng) convolutional windows, S is the stride of the conv-kernel and it is also the length of the conv-kernel, hence there is no overlap between the strides. The inputs x are sampled from (0,, Ta).
Despite its simplicity, this model distils several key differences between feature learning models and lazy learning or GP models. Due to the lack of pooling layers, the GP associated with the student fails to take advantage of the weight sharing property of the underlying CNN [36]. In fact, here it coincides with a GP of a fully-connected DNN (see Eq. 13) which is quite inappropriate for the task. We thus expect that the finite network will have good performance already for n C* + S) whereas the GP will need n of order of the dimension (N to learn well [11]. Thus, for N + S & NS there should be a broad regime in the value of where the finite network substantially outperforms the corresponding GP. We later show ($4.1.4) 1.4) that this performance boost over GP is due to feature learning, as one may expect.
Conveniently, the cumulants of the student DNN of any order can be worked out exactly. Assuming ? a 5 Is] and 6² of the noisy GD training are chosen such that ai ~ N (0, o /CN), W. ~ N (0,
(and similarly for the teacher DNN) the covariance function for the associated GP reads
0200
0g Txx
K (x,x')=
(13)
NS
NS
o7 0² N 9
Denoting A the even cumulant of arbitrary order 2m is (see App. F):
N
K2m (X1,--
X2m)
(z)) (oim- im (im-1
(2m 1)!
Cm-
(14)
while all odd cumulants vanish due to the sign flip symmetry of the last layer. In this notation, we
mean that the e's stand for integers in {1..... 2m) } and e.g. (1i 24,) == (R), -x2 and the bracket
notation (2m 1)!] stands for the number of ways to pair the integers {1, 2m3 into the above
form. This result can then be plugged in to obtain the self-consistent (saddle-point) equations on the
training set. See App. A.4 for a convergence criterion for the saddle-point, supporting its application
here.
4.1.2 Self-consistent equation in the limit of large training set
In $3 our description of the self-consistent equations was for finite and fixed training set. Further analytical insight can be gained we consider the limit of a large training set, known in the GP literature the Equivalent Kernel (EK) limit [37, 42]. For a short review of this topic, see App. D. In essence, in the EK limit we replace the discrete sums over a specific draw of training set, as in Eqs. 5, 7, 8. with integrals over the entire input distribution μuxx). Given a kernel that admits spectral decomposition in terms of its eigenvalues and eigenfunctions: K (x, x') = Σ. Aot (x) v/ (x'), the standard result for the GP posterior mean test point is approximated by [37]
(f (x.)) - f du (x), (x);
f
h (x., x) = Σ
A
v. (x (x)
(15)
A + 0²/n"
This has several advantages, already at the level of GP analysis. From a theoretical point of view, the integral expressions retain the symmetries of the kernel K [xxx') unlike the discrete sums that ruin these symmetries. Also, Eq. 15 does not involve computing the inverse matrix K-1 which is costly for large matrices.
In the context of our theory, the EK limit allows for derivation of a simple analytical form for the self-consistent equations. As shown in App. E, in our toy CNN we can write both Ag and ôg in terms of the target g using corresponding proportionality factors. Thus the self-consistent equations can be reduced to a single equation governing the proportionality factor a between ôg and g (ôg = ag). Notice that a itself is governed by an equation that is non-linear in g. which means that Ag and ôg do not scale linearly with g, only that we can trade the function-valle valued self-consistent
'Generically this requires C dependent and layer dependent weight decay.

---

equation for a ecalar-valued one. Thus, starting from the general self-consistent equations, 5, 5,7, 8, taking their EK limit, and plugging in the general cumulant for our toy model (14) we arrive the following equation for a
02/n a= A+ A x
+ (1- (1-9)A, + (1NTorTm -1
(
_/ C
Xa
a 02/n
[
(
2 C
a
27 {
(16)
o2
Setting for simplicity 0² = =1 où we have À = 1/ (NS) and we also introduced the constant
q = (1 - (G))() + +=2/n) where åGr computed using the empirical GP predictions on either the training set or test set: åGP
( fo 9.. (En, oi.) or analytically in the perturbation 1/n correction to the
theory approach developed in [11]. The quantity q has an interpretation EK approximation [11] but here can be considered as a fitting parameter. It is non-negative and is typically O(1); for more details and analytical estimates of g see App. E.2.
Equation 16 is the second main analytical result of this work. It simplifies the highly non-linear inference problem to single equation that embodies strong non-linear finite-DNN effect and feature learning (see also $4.1.4). In practice, to compute Otest we numerically solve 16 using Qraai for the training set to get Ohraann and then set a Otrain in the r.h.s. of 16 but use can also be used to bound a analytically on both the training set and test point, given the reasonable assumption that a changes continuously with C. Indeed, at large C' the pole in this equation lays at Opole = (o² (n²/)) >> whereas R o&P < Opllle- As C' diminishes, continuity implies that a must remain smaller than Opole- The latter decays o² DCVCNS/m n implying that the amount of data required for good performance scales as VCNS rather than as NS in the GP case.
quest. Equation 16
4.1.3 Numerical verification
In this section we numerically verify the predictions of the self-consistent theory of Sec. $4.1.2, by training linear shallow student CNNs on a teacher with C* 1 in Eq. 12, using noisy gradients as in Eq. 1, and averaging their outputs across noise realizations and across dynamics after reaching equilibrium.
For simplicity we used N = S and n E {62, 200, 650) S E {15, 30, 60) so that n & S1.7 The latter scaling places us in the poorly performing regime of the associated GP while allowing good performance of the CNN. Indeed, as aforementioned, the GP here requires n on the scale of X-1 NS ( for good performance [11], while the CNN requires the scale of the
number of parameters (C(N - S) = O(S))
62
200
e50
E
200
650
(A)
number crametic
(B)
empercal ChN bast
(C)
number chantats
Figure 1: (A) The CNNs' cosine distance a, defined by (1)=1111)) a)g between the ensemble- averaged prediction (f) and ground truth g plotted vs. number of channels C for the test set (for the train set, see App. H.1). As n increases, the solution of the self-consistent equation 16 (solid line) yields an increasingly accurate prediction of these empirical values (dots). (B) Same data as in (A), presented as empirical a vs. predicted a. As n grows, the two converge to the identity line (dashed black line). Solid lines connecting the dots here are merely for visualization purposes. (C) The theoretical predictions of the self-consistent theory but using only the fourth cumulant rather than all cumulants (dotted lines), and the predictions of perturbation theory to order 1/C' (dashed lines, truncated to avoid large negative values in the figure).
The results are shown in Fig. 1 where we compare the theoretical predictions given by the solution of the self-consistent equation (16) to the empirical values of a obtained by training actual CNNs and averaging their outputs across the ensemble. We can see that as n increases, the predictions of our theory match the empirical data more closely (panels A and B), as can be expected from our SP approximation which is valid for large π.
7

---

In panel C, we compare the empirical values to two alternative theoretical predictions: (i) our self- consistent theory but using only the fourth cumulant, equivalent to equation Eq. 16 while ignoring the [.... [...]_] term, rather than summing the geometric series of cumulants of all orders which gives rise to this term (dotted lines); (ii) purely perturbative theory (see e-g. [32, 44, 391) keeping only a leading ( (C) term on top of the GP predictions (dashed lines, truncated to avoid large negative values in the figure). We see that the predictions of our full self-consistent theory (panels A and B) are clearly better than those of either (i) or (ii) above, at least for sufficiently small C' or sufficiently large nt e.g. the predictions of (ii) for n E {62, 200} match the empirical data nicely for C' 2 2a but diverge rapidly for smaller values of C, so rapidly from the empirical data and always gives 0 ك a S 1, but are also quite poor for C' 2 2 This shows that in these settings knowing all the cumulants provides much more accurate theory. For sufficiently large C, all predictions, including our own theory, coincide as they all converge to the GP limit.
would expect. The predictions of (i) don't 't diverge
4.1.4 Feature learning phase transition in the CNN model
At this point there is evidence that our self-consistent shifted target approach works well within the feature learning regime of the toy model. Indeed, GP is sub-optimal here, since it does not represent the CNN's s weight sharing present in the teacher network. Weight sharing is intimately tied with feature learning in the first layer, since it aggregates the information coming from all convolutional windows refine a single set of repeating convolution-file filteers. Empirically, we observed a large performance gap of finite-C' CNNs over the infinite-C (GP) limit, which was also observed previously in more realistic settings [23, 15, 36]. Taken together with the existence of a clear feature in the teacher, a natural explanation for this performance gap is that feature learning, which is completely absent in GPs, plays a major role in the behavior of finite C CNNs.
To analyze this we wish to track how the feature of the teacher w* is reflected in the student network's s first layer weights w. across training time (after reaching equilibrium) and across training realizations. However, as our formalism deals with ensembles of DNNs, computing averages of w, with respect to these ensembles would simply give zero. Indeed, the chance of DNN with specific parameters 0 = {are. } appearing is the same as that of -6. Consequently, to detect feature learning the first reasonable object to examine is the empirical covariance matrix Σw = E SWT, where the matrix W E R has w. as its c'th column. This Ew is invariant under such a change of signs and provides important information on the statistics of we.
As shown in App. G, using our field-theory or function-space formulation, we find that to leading order in L/C the ensemble average of the empirical covariance matrix, for teacher with a single feature w", is
(Ewloor)> = 1
(
n 0z
6as +
2 A C (A + o² z w.w. - 0(1/C2)
(17)
A first conclusion that could be drawn here, is that given access to ensemble of such trained CNNs, feature learning happens for any finite C' as a statistical property. We turn to discuss the more common setting where one wishes to use the features learned by a specific randomly chosen CNN from this ensemble.
To this end, we follow Ref. [29] and model Σw Wishart matrix with a rank-one perturbation. The variance of the matrix and details of the rank one perturbation are then determined by the above equation. Consequently the eigenvalue distribution is expected to follow spiked Marchenko- -astuur (MP) model, which was studied extensively in [7]. To test this modeling assumption, for each snapshot of training time (after reaching equilibrium) and noise realization we compute Ew 's eigenvalues and aggregate these across the ensemble. In Fig. 2 we plot the resulting empirical spectral distribution for varying values of C' while keeping S fixed. Note that, differently from the usual spiked-MP model, varying C here changes both the distribution of the MP bulk (which is determined by the ratio S/C) as well as the strength of the low-rank perturbation.
Our main finding is a phase transition (analogous to (7)) between two regimes which becomes sharp as one takes S + co. In the regime of large C' the eigenvalue distribution of Σw is indistinguishable from the MP distribution, whereas in the regime of small C an outlier eigenvalue An departs from the support of the bulk MP distribution and the associated top eigenvector has non-zero overlap
8

---

with w", see Fig. 2. We refer to the latter as the feature-learning regime, since the feature w" is manifested in the spectrum of the students weights, whereas the former is the non-feature learning regime. We use the quantity Q = E' as surrogate for Am. as it is valid on both sides of the transition. Having established the correspondence to the MP plus low rank model, we can use the results of [7] to find the exact location of the phase transition, which occurs at the critical value Corit given by
Cerit = (o2 ( 1+
S5² + و2)
n2 )
+ +0
(1+ 1+
1 n A
) o²
(18)
where we assumed for simplicity N = S so that À = 5-2
n 650 S 60
650 S = 60
24-
= 256
Tww"
C 512
10
density.
1.0
1024
h_ and v "sow"
theory
8.8
04-
94"
00
3.3
2.0
810
1000
(A)
Iw eigenvalues
(B)
C
Figure 2: (A) Aggregated histograms of Σw cigenvalues where Ew = & is the normalized empirical covariance matrix of the hidden layer weights during training. Different colors indicate varying number of channels, C. Solid smooth lines indicate the corresponding Marchenko-Pann Pastur
(MP) distributions with support X4] where: } = VS/C) The quantity Q
Ewww"" which correlates with the strength of rank-1 component of the feature w", is represented by thick short bars. For large C, Q remains within the MP bulk whereas for small C it pops out. (B) The theoretical X_ curve and interpolated curve of Q intersect very close to the theoretically predicted value given in Eq. 18, here given by Cerit = 473 (dashed vertical line).
4.2 Two-layer FCN with average pooling and quadratic activations
Another setting where GPs are expected to under-perform finite-DNNs is the case of quadratic fully connected teacher and student DNNs where the teacher is rank-1, 11 also known as the phase retrieval
problem [28]. Here we consider some positive target of the form g(x) (w. x)2 o ||| where
W., X E IRd and student DNN given by fxx = Em=1 (Wm x)² - o |xx We consider training p=1 using noisy GD training with weight decay 7 = 2Mo² /o?
this DNN on n train points {
n
Similarly to the previous toy model, here too the GP associated with the student at large M (and finite o²) overlooks a qualitative feature of the finite DNN - the fact that the first term in f(x) is non-negative. Interestingly, this feature provides a strong performance boost [28] in the o² 0 limit compared to the associated GP. Namely the DNN, even at large M, performs well for > 2d [28] whereas the associated GP expected to work well only for n = ((d2) [11].
We wish to solve for the predictions of this model with our self-consistent GP based approach. As shown in App. the cumulants of this model can be obtained from the following cumulant generating function
M Tr log I- 1-2M-11 Σ it x 2
n+1
Σ
it ot |x
(19)
The associated GP kernel is given by K ((xx, (.y.xx,) x,) shift equation, at the saddle-point level, appears as
2M -1 1o (x₁, )² Following this, the target
8gn
8g x X
Ag = - - K (Xx; xx)
σ² + o
- 2M 11
2 Σ
x, - -,-2,110011 (20)
o2
P

---

In App. I, we solve these equations numerically for o² = 10 and show that our approach captures the correct n 2d threshold value. An analytic solution of these equations low o² using EK or other continuum approximations is left for future work (see Refs. [11, 8, 9] for potential approaches). As first step towards this goal, in App. we consider the simpler case of a² = 1 and derive the asymptotics of the learning curves which deviate strongly from those of GP for M < d.
4.3 Validity of the saddle-point approximation in realistic settings
In this subsection we complement the above results on controlled toy models by some preliminary results in more realistic settings. The approximation underlying our analysis was that the fluctuations of the integrand in the partition function of Eq. 6 are near- Gaus-Gassin.n A strong indication of how adequate this approximation is in real-worlde settings can be given by performing normality tests on an ensemble of the outputs of trained DNNs on the training set. Indeed, this is the ensemble described by our partition function. App. H.3 we report some preliminary results we obtained by training the Myrtle-5 deep CNN [41] using our protocol on tiny subsets (n E {16, 32, 64)) of CIFAR- 10, and inspecting the normality of output fluctuations, as measured by the 4th cumulant of the outputs. In the experiments with small enough C, although the CNN predictions already deviated strongly from those at infinite width, we haven't 't found a measurable signal of non non-nonme non-normality, thus our Gaussian fluctuation assumption is consistent with these findings.
Discussion
In this work we presented correspondence between ensembles of finite DNNs trained with noisy gradients and GPs trained on a shifted target. The shift in the target can be found by solving a set of self-consistent equations for which we give a general form. We found explicit expressions for these equations for the case of a 2-layer linear CNN and non non-linear FCN, and solved them analytically and numerically. For the former model, we performed numerical experiments on CNNs that agree well with our theory both in the GP regime and well away from it, i.e. for small number of channels C, thus accounting for strong finite effects. For the latter model, the numerical solution of these equations captures remarkable and subtle effect in these DNNs which the GP approach completely overlooks the n = 2d threshold value.
Considering feature learning in the CNN model, we found that averaging over ensembles of such networks always leads form of feature learning. Namely, the teacher always leaves a signature on the statistics of the student's 's weights. However, feature learning is usually considered at the level of a single DNN instance rather than an ensemble of DNNs. Focusing on this case, we show numerically that the cigenvalues of Ew, the student hidden weights covariance matrix, follow a Marchenko-Paasee Pastur distribution plus a rank-1 perturbation. We then use our approach derive the critical number of channels Cerit below which the student is in feature learning regime.
There are many directions for future research. Our toy models possible in order to demonstrate the essence of our theory on problems where lazy learning grossly under-performs finite-DNNs. Even within this setting, various extensions are interesting consider such as adding more features to the teacher CNN (e.g. biases or a subset of linear functions which are more favorable). studying linear CNNs with overlapping convolutional windows, or deeper linear CNNs. As for non- non-linear CNNs, we believe it is possible to find the exact cumulants of any order for a variety of toy CNNs involving, for example, quadratic activation functions. For other cases may be useful to develop methods for characterizing and approximating the cumulants.
chosen be simple as
advocated here physics-style methodology using approximations, self- consistency checks, and experimental tests. As DNNs are very complex experimental systems, we believe this mode of research both appropriate and necessary. Nonetheless hope the insights gained by our approach would help generate richer and more relevant set of toy models on which mathematical proofs could be made.
More generally,
Acknowledgements
We would like to thank Haim Sompolinsky for useful discussions and Jonathan Kadmon, Oded Ben- Ren-David, Dar Gilboa and Inbar Seroussi for comments the manuscript. GN was partially
10

---

supported by the Gatsby Charitable Foundation, the Swartz Foundation, the National Institutes of Health (Grant No. 1U19NS104653) and the MAFAT Center for Deep Learning.
References
[T] L. Aitchison. Why bigger is not always better: on finite and infinite neural networks. arXiv preprint arXiv:1910.08013, 2019.
[2] S. Alemohammad, Z. Wang. R. Balestriero, and R. Baraniuk. The recurrent neural tangent kernel. arXiv preprini urXiv:2006. 10246, 2020.
[3] A. Andreassen and E. Dyer. Asymptotics wide convolutional neural networks. arXiv preprint arXiv:2008.08675, 2020.
[4] S. Arora, S. S. Du, W. Hu, Z. Li, R. Salakhutdinov, and R. Wang. On Exact Computation with an Infinitely Wide Neural Net. arXiv e-prints, page arXiv:1904.11955 Apr 2019.
[5] P. Baldi and K. Hornik. Neural networks and principal component analysis: Learning from examples without local minima. Neural networks, 2(1):53-58, 1989.
[6] Bengyeh-Georges and R. R. Nadakuditi. The eigenvalues and eigenvectors of finite, low rank perturbations of large random matrices. Advances in Mathematics, 227(1):494-521 2011.
[7] F. Benogeh-Georges and R. R. Nadakuditi. The singular values and vectors of low rank perturbations of large rectangular random matrices. Journal of Multivariate Analysis, 111:120- 135, 2012.
[8] B. Bordelon, A. Canatar, and C. Pehlevan. Spectrum dependent learning curves in kernel regression and wide neural networks, 2020.
[9] A. Canatar,
B. Bordelon, and C. Pehlevan. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks. Nature Communications, 12(1), May 2021.
[10] I. Chizat, E. Oyallon, and F. Bach. On lazy training in differenniable programming. In Advances in Neural Information Processing Systems, pages 2937-2947, 2019..
[11] O Cohen, O. Malka, and Z. Ringel. Learning Curves for Deep Neural Networks: A Gaussian Field Theory Perspective. arXiv e-prints, page urXiv:1906.0530. Jun 2019.
[12] H. Daniels. Saddlepoint Approximations in Statistics. The Annals of Mathematical Statistics, 25(4):631 650, 1954.
[13] E. Dyer and G. Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.
[14] M. Geiger, I.. Petrini, and M. Wyart. Landscape and training regimes in deep learning. Physics Reports, 2021.
[15] M. Geiger, S. Spigler, A. Jacot, and M. Wyart. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020(1 I1): 1):111301.. 2020.
[16] B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. When do neural networks outperform kernel methods? arXiv preprint arXiv:2006.1134 13409, 2020.
[17] S. Gunasekar, J. Lee, D. Soudry, and N. Srebro. Implicit bias of gradient descent linear convolutional networks. arXiv preprint arXiv: 2018.
[18] B. He, B. Lakshminarayanan, and Y. W. Tch. Bayesian deep ensembles via the neural tangent kernel. arXiv preprint 2000:200000006664 2020.
[19] M. Helias and D. Dahmen. Statistical field theory for neural networks. arXiv preprint 1x11v:1901.10416. 2019.

---

[20] J. Hron, Y. Bahri, J. Sohl-Dickstein, and R. Novak. Infinite attention: Nogp and ntk for deep attention networks. In International Conference on Machine Learning, pages 4376 4336. PMLR, 2020.
[21] A. Jacot, F. Gabriel, and C. Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. arXiv e-prints, page arXiv: 180.078,, Jun 2018.
[22] A. K. Lampinen and S. Ganguli. An analytic theory of generalization dynamics and transfer learning in deep linear networks. arXiv preprint arXiv:1809.1033 2018.
[23] Lee, S. Schoenholz, J. Pennington, Adlam, L. Xiao, R. Novak, and Sohl- Sich-Dickstein. Finite versus infinite neural networks: an empirical study. arXiv preprint arXiv:2007.15800 2020.
[24] J. Lee, J. Sohl-dickstein, J. Pennington, R. Novak, S. Schoenholz, and Y. Bahri. Deep neural networks gaussian processes. In International Conference Learning Representations, 2018.
[25] J. Lee, I.. Xiao, S. S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. arXiv preprint arXiv: /902.060000000000 2019.
[26] Q. Li and H. Sompolinsky. Statistical mechanics of deep linear neural networks: The back- propagating renormalization group. arXiv preprint arXiv:20/2.040300 2020.
[27] E. Malach, P. Kamath. E. Abbe, and N. Srebro. Quantifying the benefit of using differentiable learning over tangent kernels. arXiv preprint arXiv:2103.01210 2021.
[28] S. S. Mannelli, E. Vanden- and L. Zdeborová. Optimization and generalization of shallow neural networks with quadratic activation functions. arXiv preprint arXi::2006.15455 2020.
[29] C. H. Martin and M. W. Mahoney. Implicit Seff-Regularizztiio in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning. arXiv e-prints, page 1xXiv:111000100 Oct. 2018.
[30] A. G. d. G. Matthews, M. Rowland, J. Hron, R. E. Turner, and Z. Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv:1804. 1271, 2018
[31] P. Mccullagh. Tensor Methods in Statistics. Dover Books on Mathematics, 2017.
[32] G. Naveh, O. Ben-David, H. Sompolinsky, and Z. Ringel. Predicting the outputs of finite networks trained with noisy gradients. arXiv preprint arXiv:2004.01190 2020.
[33] R. M. Neal. Priors for infinite networks. In Bayesian Learning for Neural Networks, pages 29-53. Springer, 1996.
[34] L. Noci, G. Bachmann, K. Roth, S. Nowozin, and T. Hofmann. Precise characterization of the prior predictive distribution of deep relu networks. arXiv preprint arXiv:2106.066115 2021.
[35] R. Novak, L. Xiao, J. Hron, J. Lee, A. A. Alemi, J. Sohl-Dickstein, and S. S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python, arXiv preprint arXiv: 2019.
[36] R. Novak, L. Xiao, J. Lee, Y. Bahri, G. Yang, D. A. Abolafia, J. Pennington, and J. Sohl- Dickstein. Bayesian Deep Convolutional Networks with Many Channels are Gaussian Processes. arXiv e-prints, page arXiv:1810.05111 Oct. 2018
[37] C. Rasmussen and C. K. I. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press, 2005.
[38] M. Refinetti, S. Goldt, F. Krzakala, and I.. Zdeborová. Classifying high-dimensional gaus- sian mixtures: Where kernel methods fail and neural networks succeed. arXiv preprint 1724 2021.
12

---

[39] D. A. Roberts, S. Yaida, and B. Hanin. The principles of deep learning theory. arXiv preprint arXiv:2106.100 10165, 2021.
[40] A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint
2013.
[41] V. Shankar, A. Fang, W. Guo, S. Kridovich-Keil, J. Ragan-Kelley, L. Schmidt, and B. Recht. Neural kernels without tangents. In International Conference on Machine Learning, pages 8614-8623. PMLR, 2020.
[42] Sollich and C. K. Williams. Understanding gaussian process regression using the equiva- lent kernel. In International Workshop on Deterministic and Statistical Methods in Machine Learning, pages 21 211-228. Springer, 2004.
[43] B. Woodworth, Gunasekar, D. Lee, Moroshko, P. Savarese, I. Golan, D. Soudry, and N. Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673. PMLR, 2020.
[44] S. Yaida. Non- gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning, pages 165-192. PMLR, 2020.
[45] G. Yang and E. Hu. Feature learning in infinite-width neural networks. arXiv preprint arXiv:2011. 14522, 2020.
[46] J. Yosinski, J. Clune, Y. Bengio, and H. Lipson. How transferable are features in deep neural networks? Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
[47] D. Yu, M. Seltzer, Li, J.-T. Huang, and Scide. Feature learning in deep neural networks studies speech recognition. In International Conference Learning Representations, May 2013.
[48] J. A. Zavatone- Veth and C. Pehlevan. Exact priors of finite neural networks. arXiv preprint a04iv:2/04. /1734, 2021.
13


Reviews:
[{"review_id": "hdqAOv1Wt5c", "reviewer": "Reviewer_P6G1", "summary": "In this work, the authors solve for the infinite time mean predictor of a model with weight decay trained via noisy full batch gradient descent. The authors develop a general formalism appropriate for model weights initialized with mean zero and variance related to the weight decay and gradient noise parameters and push this formalism to explicit characterization of the model output for the case of one-hidden-layer linear teacher student CNNs trained with MSE loss and quadratic teacher student fully connected DNNs. In the former case, the authors observe a phase transition between the lazy and feature learning regimes as the model width decreases.", "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": "The formalism and solvable model presented are interesting and have the potential to provide important insight into more realistic deep learning systems. At present, however, the manuscript falls short in a few key areas. I am open to increasing my score if these concerns are addressed (potentially using the additional page).\n\n**Relevance to realistic DL models**\n\nThe theoretical and empirical setup presented here is interesting, but quite limited. Even if the formalism requires restrictive assumptions, the paper would be much stronger if the qualitative phenomena exhibited by the solvable models could be explored empirically in more realistic settings. For example, can one see evidence of the spiked-MP phase transition in the spectrum of the weights for realistic models trained with a teacher target? Does the same transition also empirically explain models trained with categorical labels? Do the solvable models make additional predictions that can be investigated empirically in a wider context?\n\n**Relationship to existing literature**\n\nAnother area where the paper could be significantly strengthened is casting the present work in light of previous work. The authors have a fairly inclusive related work section, however there are two qualitative discussions which are important to include.\ni) The discussion of how student weights learn teacher weights presented in Section 4.4 seems quite similar to the analysis in 1809.10374. Both discuss this process in terms of low rank perturbation of a noise matrix and the associated phase transition, following Benaych-Georges & Nadakuditi. It is true that the authors of 1809.10374 were not concerned with connecting this behavior to the lazy / feature transition, however more discussion of the connection to their analysis would be welcomed.\nii) Various mechanisms underlying the lazy/feature learning transition have been discussed in the literature. For example a transition as a function of model normalization was discussed in 1812.07956 and 1906.08034, a transition as a function of learning rate was discussed in 2003.02218, and a transition as a function of weight decay strength is implicit in 2006.08643. How does the transition here relate to those mechanisms? If different, are there measurements in realistic models we can perform to disentangle the key factors underlying the transition in relevant setups? \n\n**Additional questions**\n\nFundamental limitations of formalism v convenient simplification -- It would be nice to emphasize a little more clearly which simplifications are essential for the analysis and which are used for simplicity of presentation. For example, the requirement that the weight initialization is related to the weight decay and gradient noise is essential, however the MSE loss is less essential. Is that correct? Is it correct even if one wants to understand the lazy/feature transition? Perhaps this could be emphasized more explicitly for more of the modeling assumptions?\n\nDo these results hold for both C > n and C < n? If C <n, is the measure in 2 still unconstrained? If \\sigma were 0, \\tilde{K} would not be invertible for C<n so expressions such as 5 would be singular with non-zero sigma, is there any avatar of this discontinuity?\n\nDo these results shed light on the \\sigma-> 0, finite \\gamma limit?\n\n**Minor questions / comments**\n\n- The authors provide a consistency condition for the saddle point approximation, but comment that one cannot reformulate the action with an explicit pre-factor. Can one introduce auxiliary fields (e.g. as in hep-th/0306133) to accomplish this?\n\n- 2106.10165 -- which came out post submission -- might, none-the-less, be nice to reference in a final version, as they also attempt to use corrections to large width to understand the feature learning / lazy transition.\n\n-----------------------------------\n\nAfter the response from the authors and discussion with other reviewers. I have decided to revise my score up.\n\nI appreciate the time the authors took in responding to my questions! I am only partially appeased by my main question -- what aspects of this solvable system are applicable (at least approximately) in realistic systems? The authors point to [27] (\"Implicit Self-Regularization in Deep Neural Networks...\") as a potential example of the spiked MP behavior they are analyzing in realistic setups. Though [27] does exhibit a variety of empirical distributions for spectra, it is not clear to me that they are realizing a version of the lazy feature transition discussed by the authors. That being said, I do think a model that can be analyzed through the non-perturbative lazy/feature transition is attractive. I also appreciated the additional discussion of Lampinen and Ganguli.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Some limitations, and encouragement for the authors to discuss the limitations, are discussed in the review above. The authors have adequately discussed societal impact.", "review_text": "", "extra_scores": {"ethical_concerns": "I do not see any ethical concerns. ", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4.5"}, "extra_text": {}}, {"review_id": "h0Y22yAP1zJ", "reviewer": "Reviewer_9Tuv", "summary": "This paper introduces a new perspective on feature learning in wide neural networks by demonstrating that DNNs trained with noisy gradients and large training data converge to solutions with mean predictions dictated by GP regression with shifted targets when using MSE loss. The authors then consider two toy CNN models, one linear and one non-linear, where they are able to make more analytic progress with their theory, and demonstrate the benefits of of feature learning of wide enough but finite NNs compared to their corresponding infinite-width GPs in terms of sample complextiy on student-teacher tasks. Finally, they also show a phase transition (marked by a critical value of width/channels) in the distribution of eigenvalues empirical covariance of weights in the first layer of a CNN instantiation, between a feature learning regime (where there is an outlier eigenvalue that is associated to the teacher's features) and a non-feature learning regime. Experimental evidence is provided to support these claims. ", "questions": "", "limitations": "", "rating": 7, "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": "*originality* The view of trained NNs as GP regression with modified targets is novel as far as I can tell, as are the insights developed on toy CNN models. The paper is clear in how it uses and extends previous work in my opinion.\n\n*quality* I enjoyed reading this paper, and believe it is of high quality. The approach is non-rigorous but I believe technically sound, and the insights that the authors obtain (e.g. the main equations 8 and 16) are interesting. The choice of toy models is also made well, in that they are simple enough to obtain analytic results but also complicated enough to lead to profound results of benefits of feature learning (in terms of sample complexity).\n\n*clarity* the paper is well written and organised, particularly with regards to relation to previous work.\n\n*significance* The paper examines a key question which is of feature learning in wide NNs, which will be of interest to the general NeurIPS community. The new results about trained NNs being GP regression with shifted targets is certainly surprising to me, and I can see these techniques being built upon in future work. \n\nMinor points:\n- The comment on line 20 about 'DNNs map to Bayesian inference on GPs governed by the NTK' is incorrect. https://arxiv.org/abs/2007.05864 shows you need to add in extra randomness to the initialised NN to obtain a Bayesian interpretation to trained NNs in the NTK regime.\n- where -> were in line 332\n- There is another concurrent paper to [41] in line 77, https://arxiv.org/abs/2106.06615, which provides similar results.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "- I'm a bit unsure what the setup is in Figure 1. E.g. what is C*, what is the input dataset? Also isn't it concerning that as you increase the number of channels \\alpha gets larger (so the discrepancy with the ground truth increases)?\n- This isn't really a limitation, but it's quite hard for someone (like myself) who is not well versed in the techniques used in statistical physics to understand all the details of the author's arguments.", "review_text": "", "extra_scores": {"ethical_concerns": "N/a", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "_ipd6w_N9v6", "reviewer": "Reviewer_ZQro", "summary": "The authors develop a self-consistent Gaussian Processes (GPs) theory to account for finite size effects in Deep neural networks (DNNs). The equivalence between GPs and DNNs is a well known result that holds in the so called ‘lazy learning’ regime in the infinite width/channel limit, where the DNNs’ parameters don’t change during training. In practical contexts, however, DNNs are finitely sized and operate in a different regime, called ‘feature learning’, where parameters change substantially during training. This paper provides a theoretical framework to account for finite size effects, and so it allows to investigate the feature learning regime.\n\nTo accomplish this, the authors consider the partition function $Z$ associated with the DNN posterior distribution $P(f)$, in the case of training via noisy (full batch) gradient flow with MSE loss. $Z$ is given by an integral over the prior distribution $P_0(f)$, which can be expressed in terms of cumulants. Following their notation, $C$ is the hyper-parameter controlling the over-parametrization. In the infinite-$C$ limit, only the first two cumulants survive, and the standard GP-DNN equivalence is recovered, with the GP kernel given by the second cumulant.\nFor $C$ finite, all cumulants contribute, and the authors have to rely on a saddle point approximation to make progress. By doing so, they arrive at a self-consistent set of equations for the target shift that describes the predictive mean of the DNN. In fact, the predictive mean is given by the same GP-like expression found for the infinite-$C$ limit, though with shifted targets. Next, the authors consider a second order expansion of the action so as to evaluate the posterior covariance.\n\nIn order to test their results, the authors consider two examples of DNN architectures that over-perform their associated GPs. Of the two, they mostly focus on the case of a linear two-layer CNN, in a teacher-student set up.  For this simple architecture, the cumulants can be worked out analytically. However, in order to deal with the series appearing in the expression of $P_0(f)$, the authors consider the limit of infinite number of training data, also known as Equivalent Kernel limit. This leads to an analytical expression for $\\alpha$, the proportionality factor between the discrepancy in prediction and the target value. Numerical verifications confirm their theoretical predictions (showing mismatches decreasing with the number of training samples).\nBy modelling the empirical weight covariance matrix as a Wishart matrix with a rank-one perturbation, the authors are able to detect a feature learning transition for this model. This is marked by the appearance of outliers in the associated spiked Marchenko-Pastur eigenvalue spectrum.\n\nFinally, the authors consider the case of a two-layer FCN with average pooling and quadratic activations, with a rank-1 teacher. They provide the expression of the cumulant generating function, as well as the equation for the targets shift, and leave the analysis to Appendix I.", "questions": "", "limitations": "", "rating": 7, "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": "The overall quality of the paper is good: I found it technically solid, clearly written, original (at least to my knowledge), and all claims are well supported. The organisation of the paper is fine but could be improved (see my comments below).\n\nI have the following comments for the authors:\n\n1. Other works from the literature (see e.g. https://arxiv.org/abs/1906.08034) distinguish lazy and feature learning regimes in the infinite-$C$ limit, according to the scaling of the network's parameters with $C$. In this paper, feature learning appears to be solely a finite size effect (the parameters’ scaling with $C$ gives the lazy learning regime in the infinite-$C$ limit). I wonder if the self consistent theory developed here would work for different scalings of the parameters with $C$, and in which case, as I presume the cumulants too will scale differently with $C$. I would appreciate if the authors could comment on this point.\n\n2. The authors use a saddle point (SP) method in a somewhat non standard way, as explained in the first paragraph of Appendix A.4. I think the paper would gain in clarity with a sentence in the main text that spells out this point. I would also suggest to state more clearly that the approximation assumes $n$ large, with $n$ the training set size, in the paragraph where the SP method is introduced. A further improvement would be to comment on the error induced by this approximation when $n$ is finite, as e.g. to explain the mismatches observed in figure 1 (increasing as $n$ decreases).\n\n3. With the SP method, the authors are able to find the expression for the posterior covariance matrix by expanding the action to quadratic orders. This coincides with the posterior covariance of a GP with a kernel given by $K + \\Delta K$ (with $\\Delta K$ properly defined in eq. (9)). As far as I understand, the term $\\Delta K$ makes this kernel different from the one appearing in the predictive mean with shifted targets. If this is the case, the predictions for finitely sized DNNs cannot be really described as equivalent to a GP regression with shifted targets (while the equivalence in the infinite-$C$ limit is exact). If so, this should be stated clearly in the paper, as to avoid misinterpretations.\n\n4. Although section 4.4 provides a valid and interesting investigation of the feature learning transition in the CNN model, it is not completely clear to me how the proposed analysis fits with the rest of the paper, i.e. if/how the self-consistent shifted target approach is relevant for the task. I would recommend the authors to make this clearer in the main text; as of now, it appears to be more of a stand alone section (also, it is not immediately obvious why the feature-to-lazy learning transition observed in this specific model is of interest, nor is it clear if the same method could be applied to different architectures).\n\n5. The discussion of the FCN architecture presented in the main text is limited to the equation for the target shift, while the rest of it is left to Appendix I. Even there, there is no figure showing a comparison between theoretical predictions and empirical results. I would have expected to see a plot of this kind in the main text (i.e. one similar to fig. 1), as to support the general applicability of the proposed method. Basically, the content of the FCN discussion doesn't seem to be enough for a separate section in the main text.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors do not put great effort in explaining the limitations of their analysis. In particular, the obstacles for their theory to be applied to different DNN architectures should be discussed with more details.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "16"}, "extra_text": {}}, {"review_id": "-26w2MH7kV", "reviewer": "Reviewer_ZAE8", "summary": "This work addresses the question of how finite-width DNNs differ from their infinite-width GP limit. This is important since, while the GP limit is a powerful theoretical tool for studying DNNs, it's unable to capture _feature learning_, a key behavior of DNNs. The authors show that the mean predictor of a finite DNN trained with noisy gradients and weight decay corresponds to GP regression on a shifted target, and give an approximation for the shift based on the cumulants of the prior. They apply the framework to two toy models. In particular, for a two-layer linear CNN they derive a large-training-set approximation for the cumulants, and show that as the width increases, there is a phase transition out of the feature-learning regime.  ", "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": "My impression is that this paper makes important contributions – particularly the characterization of the finite-width DNN as shifted GP regression and the the demonstration of feature learning. However, my familiarity with the topic is limited and the paper is written in a fairly inaccessibly way, which makes it hard for me to fully evaluate the contribution. \n\n\n * I felt that the experimental sections could be made more convincing, since as written, I found the discussion to be imprecise and hard to follow. For example:\n\t * The discussion of how large the training set must be for the CNN vs the GP to \"perform well\" is somewhat vague.\n\t * In section 4.3, the idea is to show that the prediction of $\\alpha$ improves as $n$ grows. But it was unclear if the existence of such $\\alpha$ depends on linearity occurring in the limit $C\\rightarrow \\infty$ anyways, (eqn. E.3)? And there is the fact that the linearity only results from the EK large $n$ limit, so it's confusing to say that the prediction of $\\alpha$ is improving as $n\\rightarrow\\infty$. And there is the approximation from finding $\\alpha_{test}$ using $\\alpha_{train}$. In essence, I didn't follow the discussion well enough to feel confident that what is demonstrated indicates an improved prediction.\n\t * In section two 4.4, I'm uncertain what is meant by saying the phase transition \"becomes sharp as one takes $n, S  \\rightarrow  \\infty$.\" Is this demonstrated? \n\t \n * It would be nice to see some comparison of the predictions of this approach compared to those of a leading-order finite width correction, e.g. [30]. Are these other approaches incapable of exhibiting feature learning?\n * A number of terms and notation are introduced without explanation, and I'm unsure which are standard to someone more familiar with the topic. Perhaps adding additional references to background material could aid the uninitiated. For example, $C$ is not defined when first introduced, SRN is undefined, and I still don't fully understand what \"self-consistent\" means.\n * Minor: \"self consistent\" should be hyphenated everywhere.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper generated a lot of lively discussion! Across the board the reviewers thought the theoretical development was interesting and insightful. However, several of the referees thought the exposition could be made more approachable for non-experts. On the other hand, several of the reviewers were concerned that the teacher-student setting in this paper was a bit too idealized; they thought the authors did not do a good enough job synthesizing the results to come up with practical lessons for researchers that were backed up by experimental evidence on more realistic networks. Although I recommend acceptance since the theory seems quite interesting, I do think the impact of the work would be significantly improved if the authors take some time to run experiments on realistic systems before the camera ready deadline to see if vestiges of their theory remain.
Decision: Accept (Poster)