Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper studies the exact asymptotics of  feature learning in a two-layer neural network. The training procedure includes two steps: 1) fix the second layer and train the frist layer using gradient descent; 2) fix the first layer and optimize the second layer. The data samples used in the two steps are independent. Under the Gaussian assumption on the initialized weights and data input, a Gaussian equivalence result is established, based on which the exact formula of the generalization error is obtained. Specifically, it is rigorously shown that 1) after one small gradient step, the generalization error is smaller than the random feature model; 2) after one large step, it is possible to obtain a better estimate than the linear model. These results demonstrate the benefits of feature learning in a precise way.
strengths_and_weaknesses: Strengths: The exact asymptotics of learning in the linear (or NTK) regime has been studied in several works before. However, what is much less unknown is the exact chacracterization beyond the linear regime, where feature learning is possible. This is also one of the key reasons why neural network can outperform traditional methods such as kernel method. This paper take an initial step towards this direction and the results are quite new and interesting. For example, a quantitative result (Theorem 5) is given on the improvement after one gradient step. Besides, the presentation is clear and the proof parts are well-written.

Weaknesses: The model (single-index) and training algorithm (data-spliting) considered in this paper is relatively special and it will be interesting to further extend to more general settings.
questions: 1. Should $\boldsymbol{W}_1$ in Eq. (2.1) be $\boldsymbol{W}_0$?

2. Does the upper bound on $\delta$ (line 292) directly follow from the form of $\delta(\cdot)$ in Theorem 5? If so, it is better to make an explicit statement.

3. Can the authors give some intuition on how the specific model in Lemma 6 is constructed and where Eq. (5.1) come from?

4. Can the also comment on the impact of the choice of $\sigma(x)$? For example, which $\sigma(x)$ can maximize $\delta(\cdot)$ in Theorem 5 and what is a good $\sigma(x)$ for large gradient step? 

5. Can the authors comment on the exact asymptotics after multiple gradient steps?
limitations: NA
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: The authors analyze the behavior of a two-layer network in the proportional asymptotic regime. The authors consider a training procedure whereby a single gradient step is taken for the weights of the initial layer, which then fixes the features for a kernel ridge regressor. The authors analyze two distinct regimes: 1) a small step size regime (step size $O(1)$), and 2) a “large” step size regime (step size $O(\sqrt{N})$). In the small step size regime, the authors leverage a perturbative argument to derive a gaussian equivalence theorem, through which they are able to give a detailed analysis of the prediction risk. In the large step size regime, the authors establish an upper bound on the prediction risk, which for some problems can be contrasted with known lower bounds for random features model to demonstrate a separation for the 1-step training procedure.
strengths_and_weaknesses: Overall, I found the paper well-written and easy to follow. The paper presents interesting insights in bridging between recent advances in characterizing random features models and more accurate models of neural network training. The results obtained by the authors are as sharp as can be reasonably expected given the problem under consideration.

In my view, the main weakness of the theory presented in this paper is that the analysis performed by the authors requires independence between the data used for the gradient step in $W$ and the data used to estimate $a$. In particular, the obvious to obtain this empirically is to split the data, causing an effective change in the ratio $N / d$ by a factor of 2. In the proportional asymptotic, this may qualitatively change the interpretation of the obtained result.


questions: 1) Given that the assumption of independent training sets for $W$ and $a$ is one of the main limitations of the theoretic results, would it be possible to provide empirical or heuristic evidence as to its impact (or lack thereof)? In other words - is this limitation due to the analysis or is the behavior fundamentally different?

2) Is it possible to carry out the analysis of theorem 5 by taking different sample sizes for the estimation of $W$ and $a$? This could provide interesting insights into scenarios inspired by e.g. pre-training where the number of samples for $W$ may be much larger than that for $a$. This could also enable one, in the context of applying the current result to a bona-fide sample using data splitting, to tune the ratio to optimize the prediction risk.

limitations: The limitations of the work are well addressed by the authors. I do not believe this work has any particular negative social impact.
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper characterizes the asymptotic risk of a single hidden layer neural network after one gradient descent step. The asymptotic theoretical analysis reveals two regimes of operation depending on the scaling of the learning rate: a regime where the scaling is of order 1 where the learning is better than random weights (random features) but worse than the best linear classifier. In the second regime where the rate is $\mathcal{0}(\sqrt{N})$ with $N$ the number of neurons, the learning not only outperforms a random feature but also any linear classifier. Theoretical analysis is performed using tools from random matrix theory and opens the perspective of a better understanding of the feature learning process in a neural network at least after one learning stage. Some insights from the theoretical formulas are also derived to explain the interplay between the number of samples and the number of neurons in the hidden layer.
strengths_and_weaknesses: *Strengths*

1- Theoretical understanding of the learning features process of neural network is a very important and challenging issue in machine learning and is one of the motivations of this work.

2- The paper is very well written and the authors have done a substantial job in making it pedagogical despite the technical difficulties inherent to the tackled problem and the heaviness of the theoretical analysis.

3 - Although the model studied is simple, the insights derived are non-trivial and the technical approach is original. 

*Weaknesses*

1- In an empirical or intuitive approach to explaining neural networks, one would be tempted to think that the strong interdependence of the joint optimization of the hidden layer weights $W$ and the regression vector $a$ is the key factor that explains the high performance of neural networks. It is tempting to say that this work, by breaking this dependency (by taking another independent set of data for the second learning phase) misses a fundamental understanding of the learning process. It would have been very interesting to compare, at least empirically, the strong intuitions that are lost when this kind of hypothesis is made, for example by adding a global learning (using the same dataset) after a gradient descent step in an experiment that would not break this dependence, in order to at least quantify empirically the impact of this strong hypothesis.

2- Although the experiments conducted on synthetic data are very well done, it would have been interesting to run simulations on some real data sets to test the robustness of the conclusions in a real application context.


questions: 1- In Equation (2.1) line 119, the gradient should be taken in the initial weights $W_0$ and not $W_1$. In other words, $G_0$ should be function of $W_0$ and not $W_1$?

2- In the line 123 about the discussion of the ridge regression for the 2nd layer, the authors might have discussed based on one experiment how the assumption of using fresh dataset for learning $a$ affects the insights drawn in the paper. The explanation of the theoretical motivation of such assumption is great but it needs a better study of the impact of such assumption.

3- Line 226 whence --> when ?

4 - As mentioned earlier, some experiments on real data would have been welcome (at least not for comparing the close fit between theory and practice which will likely be not good but at least the intuitions drawn).
limitations: 1- The authors have discussed quite thoroughly the limitations of their works. But I think applications to real data might give more strengths to the theoretical analysis. Also some experiments about assessment of the strong assumption of using different dataset for learning $W$ and $a$ would have been a great addition to the experimental part.

2- The setup of the experiments should also be thoroughly discussed in the main paper for the sake of reproducibility and assessment of the pertinence of the model used.
review_title: null

Meta-review: metareview: The paper studies a two layer neural network in the setting of one step of gradient descent on the first layer (after random init) and then freezing the first layer and training the last layer. The study is in proportinal asymptotic limits of parameters going to infinity. They identify two regimes depending on the learning rate for the single step on the first layer, in the first regime the step can improve over the random features model but stays below the best linear predictor and in the higher learning rate regime it can improve over the best linear predictor. The latter is established for a class of f* (student teacher model assumed). The meta takeaway is a neat study how feature learning can be happening. Although the precise setting is somewhat restricted, the results are strong and can lead to further research in this important area. 

Overall all the reviewers felt that the paper is a strong contribution to Neurips community. I am happy to recommend acceptance. 
confidence: Certain
Decision: Accept