Simplicity Bias and Optimization Threshold in Two-Layer ReLU Networks

We sincerely appreciate the reviewer’s insightful feedback.

The gap between the theoretical setting in this work and the empirically observed phenomenon with in-context learning and diffusion modes is large, and it is not yet clear whether and how the properties of the training dynamics identified here are related to what actually happens in those practical situations. (E.g. it might be that case that the stability issues studied by Qiao et al. in NeurIPS 2024 are at play there to a greater extent than the neuron alignment in early training.)

We agree that it remains unclear whether empirical observations in more complex architectures are better explained by our theoretical results or those of Qiao et al. (2024). Evaluating this in real-world settings would require extensive experiments beyond the scope of our primarily theoretical contribution, as also noted by the reviewer.

That said, Section A.5 in the appendix suggests that in our toy model—while not necessarily representative of more complex architectures—non-convergence appears to be driven by early alignment rather than stability issues.

We view our work as an alternative perspective to that of Qiao et al. (2024) in explaining non-convergence. In practice, it is likely that both phenomena contribute to the observed behavior in real-world settings.

In the References, check whether papers cited as on arXiv have been published, and consider providing a clickable link for each item.

We will update the references accordingly.

We sincerely appreciate the reviewer’s insightful feedback.

the theorems and propositions could be rephrased in a more accessible and less mathematically rigorous manner to improve clarity

We we will enhance the clarity of our results in the revised version. Specifically, we will include an informal statement of Theorem 4.1 at the end of the introduction. However, we do not find it feasible to do the same for Section 3, as it requires prior knowledge of extremal vectors. More precisely, we will make the following revisions regarding the clarity and organization of the paper:

• Clarification of key results: We will include an informal statement of Theorem 4.1 at the end of the introduction.
• Emphasis on the goal of Section 3: This section aims to provide a simple characterization of extremal vectors (up to $\sqrt{\frac{d}{n}}$ terms) when the number of training samples is large. This serves two purposes: (1) forming the first step in the proof of Theorem 4.1 and (2) offering a general framework that may be useful in future research, as noted by Reviewer TN6U.
• Additional clarification of key results: We will add an explanation of how Proposition 3.1 relates to and complements Theorem 3.1.
• Discussion of main implications: We will briefly summarize the key implications of Theorem 4.1 at the beginning of the discussion section.
• Expanded discussion in the appendix: We will elaborate on our work’s connections to related literature, particularly regarding the double descent phenomenon. Additionally, we will compare the NTK/lazy training regime with the feature learning/mean-field regime.

We appreciate the reviewers' valuable suggestions and look forward to refining our work accordingly

We sincerely appreciate the reviewer’s insightful feedback.

In lines 72 and 73, is there any reference to support this claim?

Yes, this phenomenon is extensively discussed by Raventos et al. (2024). In particular, Figure 4 and the discussion at the bottom of page 7 suggest that the model does not interpolate, even when trained indefinitely.

For diffusion models, Kadkhodaie et al. (2023) indicate that models are trained until convergence, yet interpolation ceases beyond a certain number of training samples. However, their experimental details do not specify the exact duration of training.

The connection to real-world large-scale models remains abstract

Assessing the causes of non-convergence in real-world setups—beyond the analysis in Section A.5 for our toy model—would require extensive experiments that warrant a separate study. As reviewer TN6U noted, our primary contribution is theoretical, making such empirical investigations beyond our current scope.

That said, Raventos et al. (2024) provide strong empirical evidence supporting the relevance of our work to ICL. Their experiments suggest that after a certain number of task examples, the model stops reaching the ERM solution (modeled as a Nadaraya-Watson kernel with poor generalization) and instead converges to a spurious local minimum that generalizes well, aligning with the optimal estimator. Furthermore, their findings indicate that prolonged training does not improve performance, reinforcing our claim.

A similar pattern is observed in Kadkhodaie et al. (2023) for diffusion models: the training loss remains low with few samples but increases beyond a certain threshold, while the test loss continues to improve.

The organization of the paper is somewhat unclear.

In response to the different reviewers’ feedback, we will make the following revisions regarding the clarity and organization of the paper:

• we will include an informal statement of Theorem 4.1 at the end of the introduction
• emphasis on the goal of Section 3 (see more below)
• we will add an explanation of how Proposition 3.1 relates to and complements Theorem 3.1
• we will summarize the key implications of Theorem 4.1 at the beginning of the discussion section
• we will elaborate on our work’s connections to related literature, particularly regarding the double descent phenomenon. Additionally, we will compare the NTK/lazy training regime with the feature learning/mean-field regime.

We appreciate the reviewers' valuable suggestions and look forward to refining our work accordingly.

the mixed use of $\mathbb{E}_ {X,y}$ and $\mathbb{E}_{\mu}$ is confusing.

We appreciate this feedback. The notation will be clarified (and unified) in the revised version.

Is Proposition 3.1 a consequence of Theorem 3.1?

As stated after Proposition 3.1, “Proposition 3.1 relies on the tail bound version of Theorem 3.1 and continuity arguments.” While Theorem 3.1 alone does not directly imply Proposition 3.1, its tail bound version plays a crucial role. This distinction is why we did not present Proposition 3.1 as a corollary. We will add further clarification in the revised version.

for what choices of $\mu$ does the extremal vector in Equation (6) exist?

Extremal vectors exist for any distribution $\mu$ since they correspond to critical directions of the continuous function $G$ (defined on line 196). However, if $G$ is non-differentiable, the definition of extremal vectors should be adapted to account for subgradients. Equation (6) assumes differentiability, which holds for distributions $\mu$ that are continuous with respect to the Lebesgue measure (or, more specifically, when the marginal of $x$ is continuous).

What is the relationship between Section 3 and Section 4?

The first step in the proof of Theorem 4.1 relies on the tail bound version of Theorem 3.1. Specifically, Section 4 requires a characterization of extremal vectors in the finite-data setting, which is facilitated by first studying their infinite-data counterparts (Equation (6)) and then applying the tail bound version of Theorem 3.1 to establish their proximity.

While Proposition 3.1 is not used in Section 4, it supports the idea that early alignment behavior in the infinite-data case is not significantly different from the large but finite $n$ case. Theorem 3.1 bounds $|D_n-D|$, but additional assumptions are needed to directly conclude that the extremal vectors of $G_n$ are close to those of $G$. Proposition 3.1 addresses this gap. We believe that Theorem 3.1 (and its tail bound version) and Proposition 3.1 are broadly applicable and may be useful in future work, as suggested by Reviewer TN6U. The paper is structured accordingly: Section 3 presents a general theoretical result that can be used beyond this work, while Section 4 applies it to a specific setting. We will clarify this connection in the revised version.

We sincerely appreciate the reviewer’s insightful feedback.

Atanasov et al. (2021) seems particularly relevant since they also study early alignment and how it relates to feature learning.

Thank you for pointing out this reference. We will include it in the revised version, as it provides valuable insights into the early alignment phenomenon.

It might be helpful to readers to have some high-level explanations of the theoretical parts as well as some visuals to help with the intuition behind the theory. There could also be more emphasis and discussion on the simplicity bias. I suggest bringing another one of the simulations from the supplementary into the main text for more experimental support of the theory. It could also be interesting to link this to double descent and “data-double descent” (Henighan et al., 2023) in a discussion. In general I suggest having a more extensive discussion relating the results of this work to existing literature.

We appreciate these suggestions. Our discussion section already connects our results to the literature on convergence to global minima, implicit/simplicity bias, and benign & tempered overfitting. We thank the reviewer for pointing out Henighan et al. (2023), which we will discuss in the revised version. Our setting does not exhibit a data double descent phenomenon, as the observed test loss consistently decreases with the number of samples, whereas double descent is characterized by an intermediate bump in the loss curve. However, the toy experiments in Henighan et al. (2023) illustrate a similar phenomenon: for a sufficiently large number of training points, the training loss remains high while the model learns optimal features. It remains unclear though whether this high training loss stems from an underparameterized regime (i.e., the model lacks sufficient capacity to memorize the data) or if optimization fails to reach the ERM in their setup.

Given the breadth of related topics, we initially omitted certain discussions (e.g., NTK and double descent) that we felt were less central to our study. However, based on the reviewer’s feedback, we will add a section in the appendix relating our work to these topics.

Additionally, we will move Figure 3 (unless another figure is preferred by the reviewers) from the appendix to the main text, utilizing the allowed extra page.

Other work studying early alignment has done so in the context of feature learning and the neural tangent kernel (Jacot et al. 2019). Does data interpolation correspond to the “lazy” learning regime, or would you consider this to be something separate? Can you discuss the connection of this work to feature learning?

There might be some confusion here. We distinguish between feature learning and the NTK/lazy regime, as they involve fundamentally different training dynamics (see Chizat et al., On Lazy Training in Differentiable Programming for an in-depth discussion). Our study specifically focuses on the feature learning regime with small initialization, as indicated by our initialization choice (Equation (2)), where both the inner and outer layers scale as $\frac{1}{\sqrt{m}}$.

In contrast, in the NTK/lazy regime (corresponding to large initialization scales), theory predicts that interpolation should occur at convergence, which is contrary to our main result. However, empirically demonstrating this interpolation in our toy model (with large $n$) is computationally challenging, as it would require an extremely large number of parameters. We will clarify this distinction and add a discussion on NTK vs. feature learning regimes in the appendix.

Do other (smooth) activations exhibit early alignment and do they learn generalizing solutions, or are they constrained to data interpolation?

Our experiments in the appendix include results with the smooth GeLU activation, where we observe similar early alignment behavior. While theory predicts that with an infinite number of neurons, interpolation should eventually occur, it remains unclear how many neurons are required in practice. Our experiments suggest that this number is quite large.