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