How Does Gradient Descent Learn Features -- A Local Analysis for Regularized Two-Layer Neural Networks

Summary

The present manuscript concerns the study how gradient descent achieves feature learning in a certain class of two-layer neural networks. The main idea of the manuscript, which builds heavily on previous line of works, is to consider how the population loss is minimized during the first steps of gradient descent. Using that after these steps the student network spans the feature space of the teacher network, the authors show that the successive steps of gradient descent are essential to fully recover the teacher network.

Strengths

The results of the paper show that learning both layers of 2-layer neural networks can be lead to a perfect recovery of the teacher network. This is a stronger result than what was previously known in the literature, where the first layer of the network was trained only for few initial steps and then fixed, leading to a particular class of random feature models.

Weaknesses

As far as I understand the paper concerns mainly the minimization of the population loss and there are no statements about the empirical loss and how sample complexity enters in the results that are presented.

Questions

-

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

See weaknesses

Summary

This paper studies the learning properties of networks trained with gradient descent. More precisely, the authors focus on the late stages of the dynamics where the algorithm learns the ground truth directions. These findings extend the usual ones in the literature that are focused on the early stages of the dynamics. The main theoretical result shows through a local landscape analysis the presence of a strong notion of feature learning, i.e., matching the ground-truth teacher directions in a non-simplified setting where second-layer weights are allowed to have negative values.

Strengths

The main strength of this submission is the nice theoretical contribution. The results are proved through a challenging local landscape analysis that significantly extends previous contributions.

Weaknesses

This submission has no strong weaknesses. However, the presentation could be improved in some parts of the manuscript. I suggest in the following section some possible changes to enhance to quality of the presentation for non-expert readers.

Questions

- I would emphasize more the local loss landscape analysis. As of now, it appears in a small paragraph at the end of Page 2, but introducing the challenges that the authors face for this type of analysis would help to grasp the quality of the contribution. - The authors correctly highlight many related works that focus on feature learning in the early stages of gradient descent dynamics and how they surpass kernel methods/random features. I believe it would help the non-expert reader to mention [1,2] that showed an equivalent non-linear feature map for networks trained with 1 step. This contrasts the noisy linear equivalent map of random features emerging through Gaussian equivalence. - The nice theoretical characterization of this work looks at a stronger feature learning metric, in contrast with closely related works that focus only on weak recovery. Are there other works that dealt with this matter in the context of gradient descent learning? I believe it would be nice to mention more generally previous works that made this "feature learning distinction" even in different contexts. For example, see [3] for Bayes-optimal learning of two-layer networks (specialization phase transition). - In equation (3) you preprocess the target network to remove its first Hermite coefficient. As correctly mentioned by the authors, this is reminiscent of the procedure done by Damian et al. (2022); however, it would be nice to describe what would happen if such a pre-processing could not be done. Are the strong feature learning capabilities of this network lost due to the presence of a non-vanishing first Hermite direction? - Could the author be more precise on how they would extend the findings to a polynomial number of samples after Theorem 2? At that point, the hardness of the target function (e.g. information exponent) would matter? - What do the authors mean when saying "complexity of $f_*$" on Page 4? It would be nice to mathematically formalize this concept. - The key passage at the end of page 4 is a bit obscure to me and I would suggest rephrasing it more clearly, e.g., remind $\varepsilon_0$. - What is $\bar{w}_i$ at page 5? - What is meant for $\varepsilon_0$-net at page 5? - A schematic drawing of the descent direction after the description on page 6 would help the reader to grasp intuitively the concepts. [1] A theory of non-linear feature learning with one gradient step in two-layer neural networks. Moniri et al. ICML 2024 [2] Asymptotics of feature learning in two-layer networks after one gradient-step. Cui et al. ICML 2024 [3] The committee machine: Computational to statistical gaps in learning a two-layers neural network. Aubin et al. NeurIPS 2018.

Rating

6

Confidence

2

Soundness

2

Presentation

2

Contribution

3

Limitations

The limitations are addressed in the submission.

Summary

The present paper studies feature learning in the end phase of training. The authors show that when the loss is small, gradient steps capture relevant directions.

Strengths

- The problem of feature learning studied in the paper is important. - From a technical point of view, the analysis of phase 3 of the algorithm that shows that the local landscape is benign is interesting and can be used in other problems.

Weaknesses

Comments: 1. In the abstract, the authors write "We show that once the loss is below a certain threshold, gradient descent with a carefully regularized objective will capture ground-truth directions". Where is this threshold? Looking at theorem 2, its hard to connect the description in the abstract to the statement proved here. 2. The relation between assumption 3 and information exponent should be made explicit. Based on my understanding, the information exponent of the teacher function is always one (because it is assumed that it has a linear part). 3. What happens if information exponent is larger than 1? 4. Why is the head a and the back layer W have the same regularization parameter \lambda? Also, in line 128, it seems that you are not analyzing the problem at \lambda = 0, but you study the problem in the ridgeless case where \lambda \to 0. 5. The analysis is for gradient descent ran on expected loss (Eq. 2). This problem is not very realistic and a finite sample analysis should be performed. What is the sample complexity here? I don't think a concentration type argument is possible here; at least in the realistic setting where the dimenion of the covariates and the number of samples are roughly in the same order. 6. The algorithm analyzed in this paper is not close to typical training methods used in practice. What is the role of phase 2? Why are we normalizing in this particular way? 7. In the discussion below Theorem 2, the authors write " In these works, neural networks only learn the target subspace and do random features within it". Is this correct? Specifically what result are the authors pointing towards? Can the authors be more formal here? This discussion is very vague. 8. Instead of the lengthly discussion on the construction of the dual certificate, I think the authors should have discussed in more detail the general take-aways of this result. 9. Missing citations and discussion of prior work. The following (highly relevant) paper have not been discussed. Papers [1] and [2] are missing in the discussion of feature learning in the early phase of training. The authors should also discuss other approaches to analyze feature learning [3], [4], etc. [1] B Moniri, D Lee, H Hassani, E Dobriban. A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks. [2] H Cui, L Pesce, Y Dandi, F Krzakala, YM Lu, L Zdeborová, B Loureiro, Asymptotics of feature learning in two-layer networks after one gradient-step. [3] A Radhakrishnan, D Beaglehole, P Pandit, M Belkin, Mechanism for feature learning in neural networks and backpropagation-free machine learning models. [4] D Beaglehole, I Mitliagkas, A Agarwala, Gradient descent induces alignment between weights and the empirical NTK for deep non-linear networks.

Questions

please see weaknesses

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

please see weaknesses

Thank you for the response

I sincerely thank the authors for their rebuttal, I have no further concerns to discuss and I believe the proposed changes will improve the submission. After carefully reading the response along with other reviewers’ comments I would like to keep my score as in the original review.