Paper: A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The assumptions for the activation function are unclear. I had a hard time understanding Condition 2.3 and Condition 2.4. Could the authors directly tell us which activation functions satisfy these conditions?

Author response:
2- Since the Hermite polynomials form an orthonormal basis of $L^2(\mu)$, where $\mu$ is the standard Gaussian measure on $\mathbb{R}$, any activation in $L^2(\mu)$ admits the Hermite expansion.
In Condition 2.3, we additionally require that $c_1 \neq 0$ and a certain decay rate.
In the limitation section, we discuss why this assumption is necessary.
Many popular activation functions including (centered) ReLU, tanh, sigmoid satisfy Condition 2.3.
We added this discussion in the revision.

Reviewer comment (Theory):
I noticed that the theoretical results regarding the training loss are only provided for the cases where $\ell$ equals 1 and 2. While the authors mentioned that results for the general case of $\ell$ can be found in the Appendix, I believe it is essential to include and discuss these results in the main content.

Author response:
4- We moved the general $\ell$ result to the main text in the revision.
We note that due to space limitations, it was not possible to include the full formula for the limiting error.

Reviewer comment (Evaluation):
Show the decrease in training loss may not be surprising. Instead, the focus should be on generalization and test error.

Author response:
5- We added a limitations section to the end of this paper discussing this point. In Section 4 of the paper, we wanted to demonstrate that the emergence of spikes in the spectrum of the feature matrix will indeed translate to going beyond linear models for large enough step size through the analysis of training loss. Unlike the training loss, the test error does not allow a simple expression such as Lemma K.2. Specifically, the bias term in the test error will contain a term involving two resolvent matrices with the covariance matrix of the learned features multiplied in the middle (see equation C.10 in Ba et al. (2022)). The strategy that we use to simplify such expressions is to use Woodbury formula to take out the spike terms from the resolvent matrices and study terms appearing in the expansion. However, for the case of this complicated formula, one needs to study various forms of interactions between the spike terms. We find that this would require a much more laborious calculation, and since our paper is already quite long (cca 50 pages total), we believe it is better suited for follow-up work.

Reviewer comment (Theory):
Emphasizing the technical challenges involved in deriving the proofs would make it easier to appreciate the technical novelty. Currently, it's challenging to discern these technical aspects when looking at the LONG proofs in the Appendix.

Author response:
6- We have added *proof idea* sections after the theorems to briefly emphasize the techniques and challenges involved in proving each result.

Reviewer comment (Experiments):
We need some experiments of higher-order feature learning to justify the theoretical findings in Theorem 4.1.

Author response:
7- We repeated the experiment in Section 5 with a target function with cubic term.
See Figure 4 in the revised manuscript.
We observe that the cubic part is learned with a large enough step size.

Reviewer comment (Novelty):
The setup of this paper and Zhenmei et al. 2022 seems significantly different from my POV, but both are counted as feature learnings. It would be better to clarify the field of feature learning and how they are connected.

Author response:
1- In Zhenmei et al. (2022), the input is defined as a linear combination of $D$ patterns, and feature learning in this setting means identifying the patterns and their relation to the label from training data.
In our problem, the feature means the direction $\boldsymbol{\beta}_{\star}$ in the single-index model.
We believe both perspectives are valid and provide independent insights to feature learning in neural networks.
We added details to the paper,
please see Section 1.1 and Appendix A in the revision.

Reviewer comment (Evaluation):
The values of training error and test error are quite confusing. It would be helpful if the authors included some baseline measures to assess the algorithm's performance. Specifically, it's challenging to see the significance of studying an algorithm with a large test error. While it may be caused by the scaling issue, the authors should consider making adjustments to avoid any misunderstanding.

Author response:
2- Setting 1 serves as a baseline in our experiment.  We see that for small step size, in setting 2, the higher order component is treated as noise and the resulting loss is equal to the loss in setting 1. When we increase the step size further, higher order components starts to be learned in setting 2 and we see an improvement compared to the baseline (setting 1).

Reviewer comment (Experiments):
How does Figure 2 change as the number of iterations increases?

Author response:
3- We conducted the same experiment with five gradient steps (see Figure 7 in the revised manuscript). We see that multiple spikes emerge in the spectrum. However, this goes beyond the scope of this paper and our theory for one step of gradient update. We leave the study of multiple steps for future work.

Reviewer comment (Theory):
The required condition on the learning rate (scaling with the number of samples) is not scalable. I never see a step size grows with the sample size in practice, which will lead to unreasonably large learning rate when learning on large-scale dataset. I understand the authors need a way to precisely characterize the benefit of large learning rates, but this condition is not realistic itself.

Author response:
Giant step sizes have recently gained a lot of attention in the deep learning theory community. See e.g., (Dandi et al. 2023). One can heuristically think of one giant step as an approximation to taking multiple small steps.

We would like to emphasize that the main purpose of our work is to gain a deeper understanding of feature learning through analysis of a simple toy model.
Modern deep learning operates based on vague intuitions and works for mysterious reasons.  It is a grand challenge understand its behavior, and any progress in the area is valuable in our view.
Thus, we believe fundamental research is necessary for future advances in the field. The use of simple theoretical models to demystify various phenomenon in deep learning (such as double descent, etc.) has been a very active and successful area of research (see e.g., Mei \& Montanari, (2022), etc.).

Reviewer comment (Experiments):
In Figure 3 left/middle: is the total number of training steps fixed? I.e. more iterations for small LR, and fewer iterations for large LR? This is important for a fair comparison between small and large learning rates.

Author response:
In all the experiments in the paper, we stick to the one gradient step update and only change the step size.

Reviewer comment (Theory):
My main concern is that unlike (Ba et al. 2022), the theoretical results in the current submission does not translate to learning guarantees for the studied single-index teacher. As a result, it is unclear if a larger learning rate provides any statistical benefits, so the claim that *"for large enough step sizes, the model can learn non-linear components of the teacher function"* is not supported.

Author response:
Unlike the training loss, the test error does not allow a simple expression such as Lemma K.2. Specifically, the bias term in the test error will contain a term involving two resolvent matrices with the covariance matrix of the learned features multiplied in the middle (see equation C.10 in Ba et al. (2022)). The strategy that we use to simplify such expressions is to use Woodbury formula to take out the spike terms from the resolvent matrices and study terms appearing in the expansion. However, for the case of this complicated formula, one needs to study various forms of interactions between the spike terms. We find that this would require a much more laborious calculation, and since our paper is already quite long (cca 50 pages total), we believe it is better suited for follow-up work. 

In the nonlinear random matrix theory community, it is common to only derive the limiting training loss  after the study of the spectrum of the feature matrix (see e.g., Pennington \& Worah (2017); Louart et al. (2017), etc.). Deriving the precise asymptotics for the test loss is a very interesting problem and we leave it as a future work.


We have added this to a limitation section at the end of the paper in the revised manuscript.
Also, we changed out statement to "Moreover, for large enough step sizes, 
the model can fit non-linear components of the teacher function."

Reviewer comment (Theory):
In Figure 3 the authors plotted the test error which exhibits improvement due to the learning of the quadratic component, but such improvement is not proved, and the experimental setting is only for target function with information exponent $s=1$. In fact, by inspecting the formulae for $\Delta$, it appears that the test error cannot improve for $s>1$.

Author response:
We thank the reviewer for pointing this out. It is true that when learning single index functions, with one gradient step, in the proportional regime, the information exponent needs to be one in order for the gradient step to be beneficial. This is also true for the results of Ba et al. (2022).  Recently, Dandi et al. (2023) showed that for single index models with information exponent $\kappa$, there are hard directions such that learning them will require $\Theta(d^\kappa)$ samples. 
They also show that with one gradient step, and a sample size $\Theta(d)$, only a single direction of a multi-index target function can be learned. In the present work, we study the problem of learning nonlinear components of a single-index target function with $\kappa = 1$. We clarified these points in the related works, and the limitations section that can be found in the revised manuscript.

Reviewer comment (Presentation):
This limitation needs to be explicitly mentioned in the main text.

Author response:
We have added this to a limitation section at the end of the paper in the revised manuscript.

Reviewer comment (Presentation):
In Figure 2, do the crosses represent the theoretical predictions of the spike location? If so, how are these values obtained?

Author response:
In this plot, the crosses are showing the location of the singular values, based on simulations.

Reviewer comment (Evaluation):
In Figure 3, do the solid lines correspond to the analytic predictions based on Theorems 4.4 and 4.5? If so, why do we observe fluctuations in the curve?

Author response:
In this plot, the the curves are drawn using simulation results. Note that the convergence can be very slow in this problem. For example note that when $\ell = 2$, we need to have a large enough $n$, $N$, $d$ such that terms like $n^{1/2}/\eta^2$ to go to zero with $\eta = n^{\alpha}$ with $\alpha \in (1/4, 1/3)$  (Section K.2.1, analysis of terms (4,2) and (2,4)).
For example, if we choose $\alpha = 0.3$, making $n^{1/2}/\eta^2 = 0.01$ requires $n = 10^{20}$.
In other words, with small $n, N, d \approx 10^{3} \sim 10^4$ that can be handled with a personal computer and academic cluster that we have access to, we cannot expect the simulated result to closely match the asymptotics.
We believe some prior works in the polynomial scaling settings might have experienced the same scalability issue, and therefore they could only provide cartoons demonstrating the staircase plot (see e.g.,  Ghorbani et al. (2020), Misiakiewicz (2022), etc.).

Reviewer comment (Theory):
A limitations section is missing. In the conclusion section, the authors should state the limitations of the assumptions and the results. The authors only proved the improvement of training loss in this two-stage training process for neural networks (NNs). There is a lack of analysis for generalization errors, although I understand there may be some difficulty with this kind of theoretical result. There should be some remark or discussion on this, or providing some conjectures related to the generalization error.

Author response:
1- We thank the reviewer for the suggestion to write a limitations section. We have written a limitation section at the end of the paper in the revised manuscript. 

In Section 4 of the paper, we wanted to demonstrate that the emergence of spikes in the spectrum of the feature matrix will indeed translate to going beyond linear models for large enough step size through the analysis of training loss. Unlike the training loss, the test error does not allow a simple expression such as Lemma K.2. Specifically, the bias term in the test error will contain a term involving two resolvent matrices with the covariance matrix of the learned features multiplied in the middle (see equation C.10 in Ba et al. (2022)). The strategy that we use to simplify such expressions is to use Woodbury formula to take out the spike terms from the resolvent matrices and study terms appearing in the expansion. However, for the case of this complicated formula, one needs to study various forms of interactions between the spike terms. We find that this would require a much more laborious calculation, and since our paper is already quite long (cca 50 pages total), we believe it is better suited for follow-up work.

In the nonlinear random matrix theory community, it is common to only derive the limiting training loss  after the study of the spectrum of the feature matrix (see e.g., Pennington \& Worah (2017); Louart et al. (2017), etc.). Deriving the precise asymptotics for the test loss is an interesting problem and we leave it as a future work.

Reviewer comment (Experiments):
Additional simulations are needed. In Section 5, there are only cases for linear and quadratic target functions. It would be better to provide more simulations for training and testing errors with more complicated target functions to show the feature learning when $\eta$ is sufficiently large. There is no empirical simulation for the staircase phenomenon in Figure 3 (Right).

Author response:
2- Regarding your point on conducting experiments for the staircase plot in Figure~3 (Right), for example note that when $\ell = 2$, we need to have a large enough $n$, $N$, $d$ such that terms like $n^{1/2}/\eta^2$ to go to zero with $\eta = n^{\alpha}$ with $\alpha \in (1/4, 1/3)$  (Section K.2.1, analysis of terms (4,2) and (2,4)).
For example, if we choose $\alpha = 0.3$, making $n^{1/2}/\eta^2 = 0.01$ requires $n = 10^{20}$.
In other words, with small $n, N, d \approx 10^{3} \sim 10^4$ that can be handled with our personal computers and academic clusters at the moment, we cannot expect the simulated result to closely match the asymptotics.
We believe some prior works in the polynomial scaling settings might have experienced the same scalability issue, and therefore they could only provide cartoons demonstrating the staircase plot (see e.g.,  Ghorbani et al. (2020), Misiakiewicz (2022), etc.). 

We repeated the experiment in Section 5 with a target function with cubic term.
See Figure 4 in the revised manuscript.
We observe that the cubic part is learned with a large enough step size.

Reviewer comment (Theory):
Theorems 4.4 and 4.5 rely on Conjecture 4.3. However, this conjecture is not well stated in the main text. It would be also better to explain the difficulty of the proof and why this conjecture cannot be proved by previous results like Hu&Lu, (2023) and Ba et al. (2022).

Author response:
3- Our results in Section 4.2 rely on a Gaussian equivalence
conjecture for the untrained features $\mathbf{F}_0$ (and not the trained features $\mathbf{F}$). In the proofs of Section 4.2, by the application of Woodbury's formula, and taking out the spike terms in $\mathbf{F}$ from the resolvent matrix $\mathbf{R}$ to reduce the problem to a problem involving  $\mathbf{F}_0$ and spikes, and we only need a Gaussian equivalence theorem to linearize $\mathbf{F}_0$; as it is common in the literature of random features models. The Gaussian equivalence conjecture we use, despite being very similar to the results discussed in Section J, does not directly follow from prior work. One can try to use the Lindeberg exchange idea used by Hu \& Lu (2023) and Bosch et al. (2023) to try to prove this specific conjecture. Although this approach might work, due to the complication form of the terms under consideration (listed in Section I),  we have not yet been able to prove this specific form of Gaussian equivalence conjecture, despite some serious attempts.
We have also discussed this in the limitations section.

Reviewer comment (Theory):
In Section 2.1, the scaling of the neural network is different from Ba et al. (2022). In Ba et al. (2022), they used a mean-field regime with learning rate $\eta\sqrt{N}$ and there is an extra $1/\sqrt{N}$ for the second layer $\mathbf{a}$. Is this regime the same as the setting of this paper?

Author response:
1- Yes. The two settings are exactly identical. The scaling difference of the second layer is the reason behind the extra $\sqrt{N}$ term in the steps sizes in Ba et al. (2022) and this paper.

Reviewer comment (Theory):
The initialization of $\mathbf{W}_0$ is sampled from a uniform distribution on the unit sphere, which is different from the Gaussian initialization of Ba et al. (2022). I guess this initialization will make the analysis simpler, e.g. Lemma B.1 can be applied directly. This should be mentioned somewhere in the paper and explain why you use this initialization.

Author response:
2- Yes, as pointed out by the reviewer, this choice of initialization will make the analysis simpler; e.g., it will enable us to use Lemma C.1 without the need for a normalization. However, many arguments can be shown to hold if we switch from the uniform distribution over sphere to Gaussian. For example, please see Section M.5. The choice of uniform over sphere distribution for weights is commonly used in the papers that study random features regression (see e.g. Mei \& Montanari (2022)). We mentioned this in the revision in Section 2.1.

Reviewer comment (Theory):
Condition 2.3 assumes that $\sigma$ has bounded first three derivatives but this won't be true if you consider the general polynomial activation function, which is set in Theorem 3.3. For Theorems 3.4, 4.1-4.2, and 4.4-4.5, do you only consider $\sigma$ as a polynomial or center ReLU function? I am confused why Theorem 3.3 needs polynomial activation functions which may contradict with Condition 2.3.

Author response:
3- Thanks for pointing out this error. In the revision, we added a decay condition to Condition 2.3 and modified the theorem statements and proof in Appendix F accordingly.
This condition is satisfied by activation functions such as (shifted) ReLU, hyperbolic tangent, and sigmoid.

Reviewer comment (Theory):
For Figure 2, are the locations of the spikes and the alignments empirically simulated or can you predict them from your theory? From Theorem 3.4, these alignments should converge to one, right?

Author response:
4- Locations of the spikes in Figure 2 are not from the theoretical predictions but are empirically generated, see Section 5 for the experimental detail.
For the $\ell = 1$ case, the alignment between $\tilde{\mathbf{X}} \boldsymbol{\beta}$ and the first left singular vector converges to one by Theorem 3.4. However, for $\ell > 1$, since $[(\tilde{\mathbf{X}} \boldsymbol{\beta})^{\circ k}]_{k = 1}^\ell$ is not necessarily orthogonal, we can only argue about the subspace distance between their span and the span of top-$\ell$ left singular vectors.

Reviewer comment (Novelty):
[1] and [2] also studied the initial feature matrix $\mathbf{F}_{0}$. And [1] also presented the limit of training error for random feature ridge regression but with a slightly different definition than yours.

Author response:
6- We thank the reviewer for pointing out these two papers. We have added them (alongside other nonlinear random matrix theory papers) to the related works section.

Reviewer comment (Writing):
Above Theorem 3.4, vector $\mathbf{w}_i$ is not defined.

Author response:
7- Thanks for pointing this out. We defined them in the revised manuscript when defining the updated weight matrix $\mathbf{W}$.

Reviewer comment (Theory):
For Theorems 4.4 and 4.3, can you say something about some extreme cases? For instance, $n\gg N,d$ or $N\gg n,d$.

Author response:
9- In the extreme case, the expressions for $m_1, m_2$ will further simplify. For example, if $n \gg N, d$, then $m_1, m_2  \to \phi / \lambda \psi$. Note that in this limit  $L_{tr}(\mathbf{F_0}) = \sigma_\epsilon^2 + c_\star^2 = \sigma_\epsilon^2 + c_{\star,1}^2 + 2 c_{\star,2}^2 + c_{\star,>2}^2$. Plugging the expressions in Corollary 4.5, we see that for example when $\ell = 2$, we have $\mathcal{L}(\mathbf{F}) \to \sigma_\epsilon^2 + \frac{2 c_{\star,2}^2}{3} + c_{\star,>2}^2$. Note that the term corresponding to the linear component of the teacher function in $\mathcal{L}(\mathbf{F_0})$ completely canceled out with the corresponding term in $\Delta_2$. 

Note that the training loss improvement becomes independent of $\sigma_\varepsilon$. This is intuitive since large sample size can make the gradient more accurate regardless of the noise level.

We have added this to appendix $M$ of the revised manuscript.

Reviewer comment (Theory):
In Section 4.2, why not present the theory of training loss for general $\ell$ like Appendix L? Can Appendix L directly cover Theorems 4.4 and 4.3? Besides, using Appendix L, can you plot Figure 3 (Right) and show that the training loss is always decreasing?

Author response:
10, 16- We thank the reviewer for this suggestion. As the reviewer suggested, we now present the theorem for the limiting training loss for general $\ell$ first and state the special cases $\ell = 1$ and $\ell = 2$ as special cases. 

As the reviewer points out, there was a minor mistake in the final expression of limiting training loss with general $\ell$ in the previous version. This is now fixed. Please see section L in the revised manuscript. The result for general $\ell$ does indeed cover the results for $\ell = 1$ and $\ell = 2$ and gives the expressions in Corollary 4.5. We thank the reviewer for pointing this out.


According to the expression of $\Delta_\ell$ in Theorem L.1, we can see that for any $\ell$, there is always an improvement due to the gradient update compared to random features regression. However, it is not guaranteed that the loss will decrease at every threshold. Looking at the special cases of $\ell$ in Corollary 4.5, we see that when going from 
$\ell = 1$ to $\ell = 2$, the loss will indeed decrease.

Reviewer comment (Writing):
In Appendix A, a typo for the definition of $\mathbf{R}_{0}$.

Author response:
8, 12, 13- Thanks for pointing these out. These typos are now fixed in the revised manuscript.

Reviewer comment (Theory):
Lemma B.1, do you need to require $\|a\|=\|b\|=1$?

Author response:
10, 16- We thank the reviewer for this suggestion. As the reviewer suggested, we now present the theorem for the limiting training loss for general $\ell$ first and state the special cases $\ell = 1$ and $\ell = 2$ as special cases. 

As the reviewer points out, there was a minor mistake in the final expression of limiting training loss with general $\ell$ in the previous version. This is now fixed. Please see section L in the revised manuscript. The result for general $\ell$ does indeed cover the results for $\ell = 1$ and $\ell = 2$ and gives the expressions in Corollary 4.5. We thank the reviewer for pointing this out.


According to the expression of $\Delta_\ell$ in Theorem L.1, we can see that for any $\ell$, there is always an improvement due to the gradient update compared to random features regression. However, it is not guaranteed that the loss will decrease at every threshold. Looking at the special cases of $\ell$ in Corollary 4.5, we see that when going from 
$\ell = 1$ to $\ell = 2$, the loss will indeed decrease.

Reviewer comment (Theory):
In the proof of proposition 3.1, how do you use Lemma J.1 to derive the limit of $\boldsymbol{\beta}^\top\boldsymbol{\beta}_{\*}$?

Author response:
14- We have clarified how we used Lemma K.1 in the proof of Proposition 3.1 in the revised manuscript.