A self-consistent theory of Gaussian Processes captures feature learning effects in finite CNNs
Year: 2021

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 — 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 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.

heading: 1 Introduction
text: 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 over-parameterized 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 over to a wide variety of architectures, 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 over-parameterized 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 a 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 (NeurIPS 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 a finite DNN trained on a 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 a 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 (§4.1). The match to 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 our 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, a sharp transition between a feature learning phase and a 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 a noisy matrix taken from the Wishart ensemble, when varying the strength of the low-rank component [6].

heading: 1.1 Additional related work
text: 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 an 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-definiteness [28] or translation invariance / equivariance [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 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 statistics of the first layer of the trained student network.
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 a prior over a single training point and only for a FCN. In contrast, our approach applies to the setting of a large training set, it is 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.

heading: 2 Preliminaries
text: We consider a fixed set of n training inputs {xµ}nµ=1 ⊂ R d and a single test point x∗ over which we wish to model the distribution of the outputs of a DNN. We consider a generic DNN architecture where for simplicity we assume a scalar output f(x) ∈ R. The learnable parameters of the DNN that determine its output, are collected into a single vector θ. We pack the outputs evaluated on the training set and on the test point into a vector ~f ≡ (f (x1) , . . . , f (xn) , f (xn+1)) ∈ Rn+1, where we denoted the test point as x∗ = xn+1. We train the DNNs using full-batch gradient decent with weight decay and external white Gaussian noise. The discrete dynamics of the parameters are thus
θt+1 − θt = − (γθt +∇θL (fθ)) η + 2σ √ ηξt (1)
where θt is the vector of all network parameters at time step t, γ is the strength of the weight decay, L(fθ) is the loss as a function of the DNN output fθ (where we have emphasized the dependence on the parameters θ), σ is the magnitude of noise, η is the learning rate and ξt ∼ N (0, I). As η → 0 this discrete-time dynamics converge to the continuous-time Langevin equation given by θ̇ (t) = −∇θ ( γ 2 ‖θ(t)‖ 2 + L (fθ) ) + 2σξ (t) with 〈ξi(t)ξj(t′)〉 = δijδ (t− t′), so that as t → ∞
the DNN parameters θ will be sampled from the equilibrium Gibbs distribution P (θ).
As shown in [32], the parameter distribution P (θ) induces a posterior distribution over the trained DNN outputs P (~f) with the following partition function:
Z ( ~J ) = ∫ d~fP0 ( ~f ) exp ( − 1 2σ2 L ( {fµ}nµ=1 , {gµ} n µ=1 ) + n+1∑ µ=1 Jµfµ ) (2)
Here P0(~f) is the prior generated by the finite-DNN with θ drawn from N (0, 2σ2/γ) where the weight decay γ may be layer-dependent, {gµ}nµ=1 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 at any of the training points or the test point xn+1 from this partition function one uses
∀µ ∈ {1, . . . , n+ 1} : 〈fµ〉 = ∂Jµ logZ ( ~J ) ∣∣∣∣
~J=~0
(3)

heading: 3 A self-consistent theory for the posterior mean and covariance
text: In this section we show that for a large training set, the posterior mean predictor (Eq. 3) amounts to GP regression on a shifted target (gµ → gµ −∆gµ). This shift to the target (∆gµ) is determined by solving certain self-consistent equations involving the cumulants of the prior P0(~f). For concreteness, we focus here on the MSE loss L = ∑n µ=1 (fµ − gµ)
2 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 as follows
P0 ( ~f ) ∝ ∫ Rn+1 d~t exp ( − n+1∑ µ=1 itµfµ + ∞∑ r=2 1 r! n+1∑ µ1,...,µr=1 κµ1,...,µr itµ1 · · · itµr ) (4)
where κµ1,...,µr is the r’th multivariate cumulant of P0(~f) [31]. The second term in the exponent is the cumulant generating function (CGF), denoted by C(~t), corresponding to P0. As discussed in
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 over-parameterization, 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 K(xµ1 ,xµ2) = κµ1,µ2 .
Consider first the case of C → ∞ [24, 30, 33] where all r > 2 cumulants vanish. Here one can explicitly perform the integration in Eq. 4 to obtain the standard GP prior P0 ( ~f ) ∝
exp ( − 12 ∑n+1 µ1,µ2=1 κµ1,µ2fµ1fµ2 ) . Plugging this prior into Eq. 2 with MSE loss, one recovers standard GP regression formulas [37]. In particular, the predictive mean at x∗ is: 〈f (x∗)〉 =∑n µ,ν=1K ∗ µK̃ −1 µν gν where K ∗ µ = K (x∗,xµ) and K̃µν = K (xµ,xν) + σ
2δµν . Another set of quantities we shall find useful are the discrepancies in GP prediction, which for the training set read
∀µ ∈ {1, . . . , n} : 〈δ̂gµ〉 ≡ gµ − 〈f (xµ)〉 = gµ − n∑
ν,ν′=1
Kµν′K̃ −1 ν′,νgν (5)
Saddle-point approximation for the mean predictor. For a DNN with finite C, the prior P0(~f) will no longer be Gaussian and cumulants with r > 2 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 f can be integrated out (see App. A.1) to yield a partition function of the form
Z ( ~J ) ∝ ∫ Rn dt1 · · · dtne−S( ~t, ~J) (6)
where S(~t, ~J) is the action whose exact form is given in Eq. A.14. Interestingly, the itµ variables appearing above are closely related to the discrepancies δ̂gµ, in particular 〈itµ〉 = 〈δ̂gµ〉/σ2. To proceed analytically we adopt the saddle-point (SP) approximation [12] which often relies on having a partition function of the form Z = ∫ dte−nS(t) where n is a large number and S is O(1). In our settings we cannot simply extract such a large factor from the action and make it O(1). Nonetheless, we argue in App. A.4 that the saddle-point is still a good approximation for large n (training set size). This relies on the fact that the non-linear terms in the action comprise of a sum of many itµ’s. 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.
For the training points µ ∈ {1, . . . , n}, taking the saddle-point approximation amounts to setting ∂itµS ( ~t, ~J ) ∣∣∣
~J=~0 = 0. This yields a set of equations that has precisely the form of Eq. 5, but where
the target is shifted as gν → gν −∆gν and the target shift is determined self-consistently by
∆gν = ∞∑ r=3
1
(r − 1)! n∑ µ1,...,µr−1=1 κν,µ1,...,µr−1 〈 σ−2δ̂gµ1 〉 · · · 〈 σ−2δ̂gµr−1 〉 (7)
Equation 7 is thus an implicit equation for ∆gν involving all training points, and it holds for the training set and the test point ν ∈ {1, . . . , n+ 1}. Once solved, either analytically or numerically, one calculates the predictions on the test point via
〈f∗〉 = ∆g∗ + n∑
µ,ν=1
K∗µK̃ −1 µν (gν −∆gν) (8)
Equation 5 with gν → gν − ∆gν along with Eqs. 7 and 8 are the first main result of this paper. Viewed as an algorithm, the procedure to predict the finite DNN’s output on a test point x∗ is as follows: we shift the target in Eq. 5 as g → g−∆g with ∆g as in Eq. 7, arriving at a closed equation for the average discrepancies 〈δ̂gµ〉 on the training set. For some models, the cumulants κν,µ2,...,µr can be computed for any order r and it can be possible to sum the entire series, while for other models several leading cumulants might already give a reasonable approximation due to their 1/Cr/2−1
scaling. The resulting coupled non-linear equations can then be solved numerically, to obtain ∆gµ 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 is to provide concrete analytical insights. Thus, in §4.1.2 we 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. [11], will yield analytically tractable equations which we solve for two simple toy models, one of a linear CNN and the other of a non-linear FCN.
Saddle-point plus Gaussian fluctuations for the posterior covariance. The SP approximation can be extended to compute the predictor variance by expanding the action S to quadratic order in the deviation from the SP value δtµ ≡ tµ − tSPµ (see App. A.3). Due to the saddle-point being an extremum this leads to S ≈ SSP + 12 ∑ µ,ν δtµA −1 µν δtν This leaves the previous SP approximation for the posterior mean on 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 Aµν is the covariance matrix of itµ. Performing such an expansion one finds
A−1µν = − ( σ2δµν +Kµν + ∆Kµν ) (9)
∆Kµν = ∂itµ∂itν C̃ (it1, . . . , itn)
where the it’s on the r.h.s. are those of the saddle-point and C̃ 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:
Σµν ≡ 〈fµfν〉 − 〈fµ〉 〈fν〉 = −σ4 [ σ2I +K + ∆K ]−1 µν + σ2δµν (10)
where the r.h.s. coincides with the posterior covariance of a GP with a kernel equal to K + ∆K [37]. The variance on the test point is given by (repeating indices are summed over the training set)
Σ∗∗ = K∗∗ −K∗µA−1µνK∗ν + 〈 ∂2it∗ C̃|it∗=0 〉 + 2 (〈 ∆g∗K ∗ µitµ 〉 − 〈∆g∗〉 〈 K∗µitµ 〉) + Var (∆g∗)
(11) where here ∆g∗ is as in Eq. 7 but where the 〈 σ−2δ̂gµ 〉 ’s are 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 ∆K = 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.

heading: 4 Experiments

heading: 4.1 The two layer linear CNN

heading: 4.1.1 Setting of the model and its properties
text: Here we define a 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 f(x), having one hidden layer with linear activation, and a corresponding teacher CNN, g(x)
f (x) = N∑ i=1 C∑ c=1 ai,cwc · x̃i g (x) = N∑ i=1 C∗∑ c=1 a∗i,cw ∗ c · x̃i (12)
This describes a CNN that performs 1-dimensional convolution where the convolutional weights for each channel are wc ∈ RS . These are dotted with a convolutional window of the input x̃i = ( xS(i−1)+1, . . . , xS·i )T ∈ RS and there are no overlaps between them so that x = (x1, . . . , xN ·S) T = (x̃1, . . . , x̃N ) T ∈ RN ·S . Namely, the input dimension is d = NS, where N is
the number of (non-overlapping) 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 N (0, Id). 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∗(N + S) whereas the GP will need n of order of the dimension (NS) to learn well [11]. Thus, for N + S NS there should be a broad regime in the value of n where the finite network substantially outperforms the corresponding GP. We later show (§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 γ and σ2 of the noisy GD training are chosen such that1 ai,c ∼ N ( 0, σ2a/CN ) , wc ∼ N ( 0, σ2w S IS ) (and similarly for the teacher DNN) the covariance function for the associated GP reads
K (x,x′) = σ2aσ 2 w
NS N∑ i=1 x̃Ti x̃ ′ i = σ2aσ 2 w NS xTx′ (13)
Denoting λ := σ 2 a N σ2w S , the even cumulant of arbitrary order 2m is (see App. F):
κ2m (x1, . . . ,x2m) = λm
Cm−1 N∑ i1,...,im=1 (•i1 , •i2) · · · ( •im−2 , •im−1 ) ( •im−1 , •im ) · · · [(2m− 1)!]
(14) while all odd cumulants vanish due to the sign flip symmetry of the last layer. In this notation, we mean that the •’s stand for integers in {1, . . . , 2m} and e.g. (1i1 , 2i2) ≡ ( x̃1i1 · x̃ 2 i2 ) and the bracket notation [(2m− 1)!] stands for the number of ways to pair the integers {1, ..., 2m} into the above form. This result can then be plugged in 7 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.

heading: 4.1.2 Self-consistent equation in the limit of a large training set
text: In §3 our description of the self-consistent equations was for a finite and fixed training set. Further analytical insight can be gained if we consider the limit of a large training set, known in the GP literature as 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 µ(x). Given a kernel that admits a spectral decomposition in terms of its eigenvalues and eigenfunctions: K (x,x′) = ∑ s λsψs (x)ψs (x
′), the standard result for the GP posterior mean at a test point is approximated by [37]
〈f (x∗)〉 = ∫ dµ (x)h (x∗,x) g (x) ; h (x∗,x) = ∑ s λs λs + σ2/n ψs (x∗)ψs (x) (15)
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(x,x′) 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 a derivation of a simple analytical form for the self-consistent equations. As shown in App. E.1, in our toy CNN we can write both ∆g 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 α between δ̂g and g (δ̂g = αg). Notice that α itself is governed by an equation that is non-linear in g, which means that ∆g and δ̂g do not scale linearly with g, only that we can trade the function-valued self-consistent
1Generically this requires C dependent and layer dependent weight decay.
equation for a scalar-valued one. Thus, starting from the general self-consistent equations, 5, 7, 8, taking their EK limit, and plugging in the general cumulant for our toy model (14) we arrive at the following equation for α
α = σ2/n
λ+ σ2/n + (1− q)λ λ+ σ2/n +
( q
λ λ+ σ2/n − 1 ) λ2 C ( α σ2/n )3 [ 1− λ C ( α σ2/n )2]−1 (16)
Setting for simplicity σ2a = 1 = σ 2 w we have λ = 1/ (NS) and we also introduced the constant q ≡ λ−1(1− α̂GP)(λ+ σ2/n) where α̂GP is computed using the empirical GP predictions on either the training set or test set: α̂GP ≡ 1− (∑ µ f GP µ gµ ) / (∑ µ g 2 µ ) , or analytically in the perturbation
theory approach developed in [11]. The quantity q has an interpretation as a 1/n correction to the 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 q see App. E.2.
Equation 16 is the second main analytical result of this work. It simplifies the highly non-linear inference problem to a single equation that embodies strong non-linear finite-DNN effect and feature learning (see also §4.1.4). In practice, to compute αtest we numerically solve 16 using qtrain for the training set to get αtrain, and then set α = αtrain in the r.h.s. of 16 but use q = qtest. Equation 16 can also be used to bound α analytically on both the training set and test point, given the reasonable assumption that α changes continuously with C. Indeed, at large C the pole in this equation lays at αpole = (σ2/n)(C/λ)1/2 1 whereas α ≈ αGP < αpole. As C diminishes, continuity implies that α must remain smaller than αpole. The latter decays as σ2 √ CNS/n implying that the amount of data required for good performance scales as √ CNS rather than as NS in the GP case.

heading: 4.1.3 Numerical verification
text: 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 as 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 ∈ {62, 200, 650} , S ∈ {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 λ−1 = NS = O(S2) for good performance [11], while the CNN requires n on the scale of the number of parameters (C(N + S) = O(S)).
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 α 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 n.
In panel C, we compare the empirical values to two alternative theoretical predictions: (i) our selfconsistent theory but using only the fourth cumulant, equivalent to equation Eq. 16 while ignoring the [· · · ]−1 term, rather than summing the geometric series of cumulants of all orders which gives rise to this term (dotted lines); (ii) A purely perturbative theory (see e.g. [32, 44, 39]) keeping only a leading O(1/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 n: e.g. the predictions of (ii) for n ∈ {62, 200} match the empirical data nicely for C ≥ 24 but diverge rapidly for smaller values of C, as one would expect. The predictions of (i) don’t diverge so rapidly from the empirical data and always gives 0 ≤ α ≤ 1, but are also quite poor for C . 25. This shows that in these settings knowing all the cumulants provides a much more accurate theory. For sufficiently large C, all predictions, including our own theory, coincide as they all converge to the GP limit.

heading: 4.1.4 Feature learning phase transition in the CNN model
text: 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 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 to refine a single set of repeating convolution-filters. 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 first layer weights wc across training time (after reaching equilibrium) and across training realizations. However, as our formalism deals with ensembles of DNNs, computing averages of wc with respect to these ensembles would simply give zero. Indeed, the chance of a DNN with specific parameters θ = {ai,c, wc} appearing is the same as that of −θ. Consequently, to detect feature learning the first reasonable object to examine is the empirical covariance matrix ΣW ≡ SCWW
T, where the matrix W ∈ RS×C has wc as its c’th column. This ΣW is invariant under such a change of signs and provides important information on the statistics of wc.
As shown in App. G, using our field-theory or function-space formulation, we find that to leading order in 1/C the ensemble average of the empirical covariance matrix, for a teacher with a single feature w∗, is
〈[ΣW ]ss′〉 =
( 1 + ( 1
λ +
n
σ2
)−1) δss′ + 2
C
λ
(λ+ σ2/n) 2w ∗ sw ∗ s′ +O(1/C 2) (17)
A first conclusion that could be drawn here, is that given access to an 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 as a 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 a spiked Marchenko-Pastur (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 ΣW ’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 n, S →∞. 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 λm departs from the support of the bulk MP distribution and the associated top eigenvector has a non-zero overlap
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 quantityQ ≡ w∗TΣWw∗ as a surrogate for λm, 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 Ccrit given by
Ccrit = 4
S (S−1 + (σ2/n)S) 4
( 1 + ( S2 + n
σ2
)−1) +O ( 1 + ( 1
λ +
n
σ2
)−1) (18)
where we assumed for simplicity N = S so that λ = S−2.
(MP) distributions with support on [λ−, λ+] where: λ± =
(
1±
√
S/C
)2
. The quantity Q ≡
w∗TΣWw ∗, 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 λ+ curve and interpolated curve of Q intersect very close to the theoretically predicted value given in Eq. 18, here given by Ccrit = 473 (dashed vertical line).

heading: 4.2 Two-layer FCN with average pooling and quadratic activations
text: 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, also known as the phase retrieval problem [28]. Here we consider some positive target of the form g(x) = (w∗ · x)2 − σ2w‖x‖
2 where w∗,x ∈ Rd and a student DNN given by f(x) = ∑M m=1(wm · x)2 − σ2w‖x‖
2. We consider training this DNN on n train points {xµ}nµ=1 using noisy GD training with weight decay γ = 2Mσ 2/σ2w.
Similarly to the previous toy model, here too the GP associated with the student at large M (and finite σ2) 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 σ2 → 0 limit compared to the associated GP. Namely the DNN, even at large M , performs well for n > 2d [28] whereas the associated GP is expected to work well only for n = O(d2) [11].
We wish to solve for the predictions of this model with our self-consistent GP based approach. As shown in App. I, the cumulants of this model can be obtained from the following cumulant generating function
C(t1, ..., tn+1) = − M
2 Tr
( log [ I − 2M−1σ2w ∑ µ itµxµx T µ ]) − n+1∑ µ=1 itµσ 2 w‖xµ‖ 2 (19)
The associated GP kernel is given by K(xµ,xν) = 2M−1σ4w(xµ · xν)2. Following this, the target shift equation, at the saddle-point level, appears as
∆gν = − ∑ µ K(xν ,xµ) δ̂gµ σ2 + σ2wx T ν [ I − 2M−1σ2w ∑ µ δ̂gµ σ2 xµx T µ ]−1 xν − σ2w‖xν‖ 2 (20)
In App. I, we solve these equations numerically for σ2 = 10−5 and show that our approach captures the correct n = 2d threshold value. An analytic solution of these equations at low σ2 using EK or other continuum approximations is left for future work (see Refs. [11, 8, 9] for potential approaches). As a first step towards this goal, in App. I we consider the simpler case of σ2 = 1 and derive the asymptotics of the learning curves which deviate strongly from those of GP for M d.

heading: 4.3 Validity of the saddle-point approximation in realistic settings
text: 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-Gaussian. A strong indication of how adequate this approximation is in real-world 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. In 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 ∈ {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 found a measurable signal of non-normality, thus our Gaussian fluctuation assumption is consistent with these findings.

heading: 5 Discussion
text: In this work we presented a 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 a 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 C effects. For the latter model, the numerical solution of these equations captures a 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 to a form of feature learning. Namely, the teacher always leaves a signature on the statistics of the student’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 eigenvalues of ΣW , the student hidden weights covariance matrix, follow a Marchenko–Pastur distribution plus a rank-1 perturbation. We then use our approach to derive the critical number of channels Ccrit below which the student is in a feature learning regime.
There are many directions for future research. Our toy models were chosen to be as simple as 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 to 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-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 it may be useful to develop methods for characterizing and approximating the cumulants.
More generally, we advocated here a physics-style methodology using approximations, selfconsistency checks, and experimental tests. As DNNs are very complex experimental systems, we believe this mode of research is both appropriate and necessary. Nonetheless we hope the insights gained by our approach would help generate a richer and more relevant set of toy models on which mathematical proofs could be made.

heading: Acknowledgements
text: We would like to thank Haim Sompolinsky for useful discussions and Jonathan Kadmon, Oded Ben-David, Dar Gilboa and Inbar Seroussi for comments on the manuscript. GN was partially
supported by the Gatsby Charitable Foundation, the Swartz Foundation, the National Institutes of Health (Grant No. 1U19NS104653) and the MAFAT Center for Deep Learning.

Review summaries:
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.

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.  

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.  To 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. For $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.  In 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). By 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.  Finally, 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.

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.