A self consistent theory of Gaussian Processes captures feature learning effects in finite CNNs

1 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 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 ] . 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 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 . 2 Preliminaries . 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 ) 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 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 can not 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 .

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.