Generalization for Least Squares Regression with Simple Spiked Covariances

The authors claim l.083 that Moniri et al. (2023) do not quantify the test error after one gradient step. To the best of my understanding, they do provide such a tight characterization (Theorem 4.5). Could the authors clarify their claim, and emphasize how their work is positioned with respect to Moniri et al. (2023)?

We should have been more precise in our claim. While Theorem 4.5 in Moniri et al. 2023 shows that the difference between the test error using $F_1$ (features after 1 step) and $F_0$ (initial features) converges to a constant, there are important differences in our approach:

1. Moniri et al. 2023 have expressions requiring solutions to fixed point equations (see equations (5) in their paper). These only hold asymptotically with hard-to-quantify approximation rates.

2. In contrast, we provide:

   • Closed form expressions for the risk itself (not just differences as is the case in Moniri et al. 2023)

   • Better control on approximation error rates, enabling analysis of finite matrices

   • The above two allow better control in understanding the relationship between the bulk and spike.

I find the discussion l.332-355 somewhat confusing

We apologize for the unclear notation. We use the Vinogradov notation where $f \ll g$ means $f = O(g)$. Therefore, setting $\theta^2 = \tau^2 n$ is consistent with our assumptions.

I believe more intuition about the different scaling considered would help solidify the intuition for the spn case

You raise an excellent point. Let's clarify the scaling relationships:

1. For the bulk term: $a_i^T \beta_* \approx \tau_A \|\beta_*\|$

2. For the spike term: $z_i^T \beta_* \approx \theta\|\beta_*\|$

Using our scaling $\theta = \tau \sqrt{n}$, the signal part is always larger. However, when the bulk also grows (i.e., $\tau_A = \Theta(d)$), the spike's effect becomes invisible. Specifically, we need:

• $a_i^T \beta_* = \Theta(1)$

• $z_i^T \beta_* = \Theta(\sqrt{n})$

(assuming $\|\beta_*\|= \Theta(1)$) for the spike to have a detectable effect.

In 3.3, more discussion in the main text about why only the unregularized case is considered for the spn case, while generic is considered for the signal-only model, would be helpful for intuition, whether it is for technical reasons or because it is not interesting.

This limitation is primarily technical. The regularized case for the signal-plus-noise model introduces many additional terms whose mean and variance would need to be bounded, significantly complicating the analysis.

typos

Thank you for identifying these issues. We will correct all typos in the revision.

I thank the authors for taking the time to provide all the detailed clarifications and answers to my interrogations. I think the paper is scientifically sound, and thus increase slightly my score, although I have not checked the proofs.

We thank the reviewer for the feedback.

The paper is motivated by the spiked covariance from the one-step gradient feature learning in neural networks (Section 1). However, it did not show how the results can be applied to the feature learning scenario. I question the amount of contribution this paper provides.

Prior work (Moniri et al. 2023 and Ba et al. 2022) established the existence of spikes for specific target types $y$ (single index models). Our work focuses on understanding the spike's effect through several novel contributions:

1. We analyze various targets and alignments between spikes and targets, providing rigorous proofs of generalization bounds

2. We demonstrate how asymptotic results may not capture behavior in finite networks

3. We provide precise quantification of how spikes affect generalization in both finite and asymptotic regimes

The assumption in line 2221-222 and 253-255 is too strong. The analysis breaks down if there is dependency in the cross term. However, the paper did not show how big the difference the predicted result would be when there is dependence in the cross term. It is questionable if the result in this paper is applicable in realistic machine learning settings.

We have now removed this assumption. When considering the dependence structure from Moniri et al. 2023, the analysis actually simplifies. A key term in our analysis is the projection of the spike direction on the bulk eigenvectors. Due to the dependence, this term becomes more tractable. Please see the general response.

What novel results could the authors conclude in the feature learning setting in neural networks using the main theorems 3,4?

Our analysis reveals several important insights for feature learning:

1. Relationship between targets and spike size:

   • For targets depending on the bulk (input data), large spikes are crucial

   • The required spike size scales with input dimension and dataset size

2. Importance of spike-target alignment:

   • The alignment between spike direction and targets significantly affects generalization

   • This alignment term exhibits its own double descent behavior

   • Small alignment improvements can yield large generalization gains

3. Double descent characteristics:

   • Peak location depends on bulk variance and regularization strength

   • Suggests weight decay regularization primarily affects the bulk, not learned features (spike)

While some of these phenomena have been observed before, we provide simplified, quantitative connections between them.

How could the authors show the assumption on the dependence does not affect the result? Is there any experimental validation?

We provide new theoretical results that explicitly handle the dependence structure. Please see the general response.

From Figure 4.1, we can see that the effect of the spike correction term is small when is large. Is the main theorem still useful to explain the phenomenon we see from feature learning?

Yes - prior work shows that the spike represents the learned feature. Our results allow for larger spikes than previously considered in works like Hastie et al. 2022. However, this shows that for the spike to effect the generalization, we need even bigger spikes.

This paper has problems with the wordings, even in main theorems. This makes the reading difficult. For instance:

Thank you for identifying these issues. We will fix all typos and improve clarity in the revised version.

We thank the reviewer for their comments.

They reference the work of Moniri et al., but this work is unrelated to neural networks or gradient descent; it addresses a purely linear regression problem for data with simple spiked covariances.

We respectfully disagree. Our work is directly motivated by and connected to neural networks through the following chain of reasoning:

1. Ba et al. (2022) and Moniri et al. (2023) show that after one gradient step, the feature matrix $F_1$ can be written as $F_0 + P$, where $P$ is a rank-$ell$ matrix.
2. This creates a spiked covariance structure in $F_1^TF_1$.
3. To understand the generalization error of such networks, we need to analyze least squares regression with $F_1$ as the feature matrix.
4. Our work studies this exact setting, though in a simplified form, to make the analysis tractable.
5. In the general rebuttal, we removed many of our simplifications. This further strengthens the connection

They do not account for the generalization ability of neural networks after a single gradient step, as they bypass the gradient step entirely by assuming the W1 matrix directly, which does not reflect the full process of neural network training.

We agree with the reviewer. Building on the results from Ba et al. 2022 and Moniri et al. 2023, our new results take us towards understanding the generalization error for two-layer networks.

Our analysis provides valuable insights:

• It shows how spikes affect generalization in finite vs asymptotic regimes
• It demonstrates the importance of the alignment between the spike direction and the target function

We never meant our paper to claim that we understood the generalization error for two-layer networks, as reflected in our title's focus on least squares regression.

Could you provide a reference for the statement, 'It has been shown that to understand the generalization...' on line 39?

In addition to the RMt paper cited in the paper, see [1] for an empirical result on more realistic networks.

[1] Martin and Mahoney, "Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning", JMLR 2021.

Is your generalization analysis very different from the work of Li & Sonthalia (2024)?

Compared to Li and Sonthalia 2024. Only one of their two models allows for an eigenvalue to diverge. This model is closely related to the Signal only model in this paper. However, we have output noise $\varepsilon$. This creates many new dependencies requiring novel analysis techniques.

1. Consider the proof sketch on page 10. Line 491 shows that for the signal-only problem, our solution is the solution from the Li and Sonthalia 2024 paper plus an extra term. The $\hat{\varepsilon}$ is not isotropic in the ridge regularied version. Hence, we do not immediately have free independence between $\hat{Z}+\hat{A}$ and $\hat{\varepsilon}$. Hence, we need to be very careful about the alignment between the two.

2. Hence, we get terms that are cubic and quartic in eigenvalues (vs quadratic in prior work). These required the development of new concentration bounds for these higher-order terms. See Lemmas 17 through 22.

The Li and Sonthalia 2024 paper does not consider the signal plus noise case, which introduces further terms that need bounding. Finally, the Li and Sonthalia 2024 only consider the case when $\tau_A = 1$ or more broadly $\tau_A = \Theta(1)$, we allow $\tau_A = \Theta(\sqrt{d})$. This is also significant.

We thank the reviewer for their detailed feedback. Let us address the key points:

Limited contribution/novelty... Most of the results in this paper are trivial extensions of the results by Hastie et al. (2022) and Li & Sonthalia (2024)

We respectfully disagree. Our contributions extend beyond prior work in several important ways:

1. Finite vs Asymptotic Analysis: Prior work focused on asymptotic results. We provide finite-sample corrections that reveal how spikes affect generalization. We show when these corrections matter (small bulk variance) and when they don't (large bulk variance)

2. Technical Novelty:

   • As the reviewers point out, we allow one eigenvalue to diverge compared to Hastie et al. 2022. This is a significant difference.

   • Compared to Li and Sonthalia 2024. Only one of their two models allows for an eigenvalue to diverge. This model is closely related to the Signal only model in this paper. However, we have output noise $\varepsilon$. This creates many new dependencies requiring novel analysis techniques

     • Consider the proof sketch on page 10. Line 491 shows that for the signal-only problem, our solution is the solution from the Li and Sonthalia 2024 paper plus an extra term. The $\hat{\varepsilon}$ is not isotropic in the ridge regularied version. Hence, we do not immediately have free independence between $\hat{Z}+\hat{A}$ and $\hat{\varepsilon}$. Hence, we need to be very careful about the alignment between the two.
     • Hence, we get terms that are cubic and quartic in eigenvalues (vs quadratic in prior work). These required the development of new concentration bounds for these higher-order terms. See Lemmas 17 through 22.

     The Li and Sonthalia 2024 paper does not consider the signal plus noise case, which introduces further terms that need bounding. Finally, the Li and Sonthalia 2024 only consider the case when $\tau_A = 1$ or more broadly $\tau_A = \Theta(1)$, we allow $\tau_A = \Theta(\sqrt{d})$. This is also significant.

Although the paper is motivated by Moniri et al. (2023), there are significant discrepancies.

We agree. Please see the general response on how we can handle the dependency.

Additionally, we think of the difference in the targets as a strength of the paper, as we can show things as the following.

1. Relationship between targets and spike size:

   • For targets depending on the bulk (input data), large spikes are crucial
   • The required spike size scales with input dimension and dataset size

2. Importance of spike-target alignment:

   • The alignment between spike direction and targets significantly affects generalization
   • This alignment term exhibits its own double descent behavior
   • Small alignment improvements can yield large generalization gains

3. Double descent characteristics:

   • Peak location depends on bulk variance and regularization strength
   • Suggests weight decay regularization primarily affects the bulk, not learned features (spike)

There exists a related work (Cui et al., 2024)... Dandi et al. (2024)

Thank you for bringing these to our attention. The second is quite new and, as the reviewer points out, was only posted well after the submission deadline. Hence, we believe that it is concurrent work and should not affect the review of our paper. We shall nonetheless discuss the two papers in the revision.

There are many differences to Cui et al. 2024

1. We work in a more restricted setting but provide rigorous proofs.

2. We simplify expressions. For example, $\zeta$ in equation 17 in Cui et al. 2024 is exactly $\xi - 1$ in our paper (see Lemma 13 in the appendix for a definition). The expression in Cui et al. 2024 is left in terms of $\zeta$. This is because the results are a product of dependent terms. Hence, simplification is not easy. However we a. Compute the expectations and variances of each of the terms b. Compute the expectations of the products. c. Greatly simplify expressions These simplifications are a major strength of our work, as the expressions are interpretable without numerical computations.

In Line 178, denotes the case with a single spike (as shown by Moniri et al., 2023)...

We only consider the single spike case. The analysis here is similar to Sonthalia and Nadakuditi 2023. Kaushik et al. 2024 extend Sonthalia and Nadakuditi 2023 to the higher rank version. We would need to do the same to extend to multiple spikes. As mentioned before, the difficulty in the analysis was a. Compute the expectations and variances of each of the terms b. Compute the expectations of the products. c. Greatly simplify expressions

These are currently scalar expressions, so we can use commutativity. For multiple spikes, we have matrix expressions, so we no longer have commutativity. The analysis is possible, but it is just quite tedious.

We ignore the $o(\sqrt{n})$ we shall highlight this as another discrepensacy

We thank the reviewer for the feedback and comments. Key differences between our work and important prior research are that we (1) provide finite matrix correction terms and (2) offer simplified closed-form expressions.

Could the authors comment on the link between their results and (Ba et al. 2022) in the context of Gaussian Universality (see e.g. [1]) ?

Yes, the problem we and prior work are interested in understanding is the generalization error for the following. First, we solve a regression problem $\beta\_{LS} = argmin ||y - \beta^T F||\_F^2 + \lambda ||\beta||\_2^2$ Then, we are interested in the generalization performance of $\beta\_{LS}$. Let's call this risk $R(F)$ to highlight the dependence on $F$. The difference between setups lies in the $F$ term. There are three different $F$'s considered.

1. $F_{CK} := \sigma(WX)$. For Gaussian $X$ and $W$, after taking a step of GD. This is from Ba et al. 2022
2. $F_{CE} := \theta_1 WX + \theta_2 A$ where $A$ has IID standard Gaussian entries independent of $W,X$ and $W$ is after taking gradient step of GD
3. $F_{SP} := A + \theta uv^T$, where $A$ has IID standard Gaussian entries.

The Ba et al. 2022 paper shows that in the small learning rate regime, $R(F_{CK}) = R(F_{CE})$ asymptotically.

However, we do things differently:

1. We allow spikes from the large learning rate limit. Hence, the result from Ba et al. 2022 does not apply. Equation 3.1 in Ba et al. shows that for small learning rate, the size of the spike is $\Theta(1)$, where as for large learning rates, it is $\Theta(\sqrt{d})$ (note for the rank one spike the Frobenius norm is equal to the spectral norm). We are interested in the case when the spike is large. The idea behind the large step size is that we are in a regime in which the Gaussian Equivalence Property is no longer true.

2. We provide more precise correction terms for finite matrices. While the prior work is purely asymptotical.

3. In our rebuttal, we also generalize to models closer to that from Moniri et al.

One additional weakness of this submission is the related works coverage.

We thank the reviewer for pointing us to these works. We shall add these references.

[5] characterizes the risk for the setting from Moniri et al. 2023 using, as the reviewer and the paper says, using the non-rigorous replica symmetry method. The differences are three-fold:

1. Our results in the paper are for a restricted setting; however, we provide proof.
2. We simplify expressions. For example, $\zeta$ in equation 17 in Cui et al. 2024 is exactly $\xi - 1$ in our paper (see Lemma 13 in the appendix for a definition). The expression in Cui et al. 2024 is left in terms of $\zeta$. This is because the results are a product of dependent terms. Hence, simplification is not easy. However we a. Compute the expectations and variances of each of the terms b. Compute the expectations of the products. c. Greatly simplify expressions

We believe these are the main challenges we overcome in our proof.

Could the authors elaborate on the connection with Moniri et al. 2024?

While Theorem 4.5 in Moniri et al. 2023 shows that the difference between the test error using $F_1$ (features after 1 step) and $F_0$ (initial features) converges to a constant, there are important differences in our approach:

1. Moniri et al. 2023 has expressions requiring solutions to fixed point equations (see equations (5) in their paper). These only hold asymptotically with hard-to-quantify approximation rates.

2. In contrast, we provide:

   • Closed form expressions for the risk itself (not just differences as is the case in Moniri et al. 2023)

   • Better control on approximation error rates, enabling analysis of finite matrices

   • The above two allow better control on understanding the relationship between the bulk and spike.

What is the bottleneck for analyzing multiple spikes?

The analysis here is similar to Sonthalia and Nadakuditi 2023. Kaushik et al. 2024 extend Sonthalia and Nadakuditi 2023 to the higher rank version. We would need to do the same to extend to multiple spikes. As mentioned before, the difficulty in the analysis was bounding variances. These are currently scalar expressions hence we can use commutativity. For multiple spikes we have matrix expressions. Hence no longer have commutativity. The analysis is possible, but it is just quite tedious.

maximal scaling regime

We don't know this regime. Does the reviewer mean when the step size is too small and we do not see a spike or the regime from [1] where the spectrum becomes heavy-tailed? In the heavy-tailed situation analysis similar to [2] can be used.

[1] Martin and Mahoney JMLR 2021 - Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning

[2] Wang et al. 2024 AISTATS - Near-interpolators: Rapid norm growth and the trade-off between interpolation and generalization

I warmly thank the authors for addressing my concerns. After reading carefully the other reviewers' comments, I believe that the paper still heavily relies on previously published works, and I would like therefore to keep my original score. On the writing side,I believe the authors should expand the discussion about the connection with previous works not belonging only to the Random Matrix Theory literature (e.g. see the suggested references [1,5]).