Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.  
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. 


 * 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:
	 * The discussion of how large the training set must be for the CNN vs the GP to "perform well" is somewhat vague.
	 * 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.
	 * 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? 
	 
 * 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?
 * 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.
 * Minor: "self consistent" should be hyphenated everywhere.

limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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).

**Relevance to realistic DL models**

The 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?

**Relationship to existing literature**

Another 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.
i) 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.
ii) 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? 

**Additional questions**

Fundamental 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?

Do 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?

Do these results shed light on the \sigma-> 0, finite \gamma limit?

**Minor questions / comments**

- 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?

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

-----------------------------------

After the response from the authors and discussion with other reviewers. I have decided to revise my score up.

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

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_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.

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

I have the following comments for the authors:

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

2. 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).

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

4. 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).

5. 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.
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_title: null

final_score: rating: 7: Good paper, accept
review_content: 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. 
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.

*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).

*clarity* the paper is well written and organised, particularly with regards to relation to previous work.

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

Minor points:
- 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.
- where -> were in line 332
- There is another concurrent paper to [41] in line 77, https://arxiv.org/abs/2106.06615, which provides similar results.
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)?
- 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_title: null

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)