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

We would like to thank reviewers for the feedback. For the concern about sample complexity, please refer to the general response where we address the questions.

We would like to thank reviewer's efforts to provide detailed feedback. We will incorporate suggestions on the presentation of the paper and discuss more related works in the revision. Below we try to address reviewer's concerns.

In equation (3) you preprocess the target network to remove its first Hermite coefficient ...

Removing the first Hermite direction is important for the current analysis and Damian et al. (2022). If there is a non-vanishing first Hermite direction, then it will become the dominate term in the first-step gradient. To be more specific, after first-step gradient update neurons $w_i^{(1)}\approx c_1 \beta +c_2 Hw_i^{(0)}$, where $c_1,c_2$ are scalars, $\beta$ is the first Hermite direction and matrix $H$ has full rank within the target subspace due to Assumption 3. When first Hermite direction $\beta\neq 0$, $\beta$ dominates so $w_i^{(1)}\approx c_1 \beta$. This makes all neurons collapse into a single direction.

We would also like to note that this issue is likely a technical difficulty introduced by the setting we consider. In experiments, we can still achieve strong feature learning ability with proper choice of hyperparameters. Providing rigorous analysis beyond current setting is an interesting and challenging problem to consider in the future.

Could the author be more precise on how they would extend the findings to a polynomial number of samples after Theorem 2 ...

Please refer to the general response where we address the questions regarding sample complexity.

What do the authors mean when saying "complexity of $f_*$" on Page 4? It would be nice to mathematically formalize this concept.

Thanks for the suggestion. We will make it clear in the revision. Here by ``complexity of $f_*$", we mean it only depends on the quantities of $f_*$, such as teacher neurons' norm $|a_i^*|,\|\|w_i^*\|\|$, target subspace dimension $r$ and angle separation $\Delta$.

What is $\bar{w}_i$ at page 5?

It is the normalized version of neuron $w_i$, i.e., $\bar{w}_i = w_i/\|\|w_i\|\|_2$.

What is meant for $\varepsilon_0$-net at page 5?

We mean an $\epsilon$-net with $\epsilon=\varepsilon_0$. That is, neurons cover the whole target subspace well in the sense that for every direction $v$ in the target subspace, there exists a neuron $w$ such that the angle between them is small $\angle(w,v)\le \varepsilon_0$.

Presentation of paper and discussion of prior work.

We will improve the presentation of paper based on the suggestions and discuss more related works in the revision.

We would like to thank reviewer's efforts to provide detailed feedback. We will try to address reviewer's concerns below and incorporate the suggestions in the revision.

1. In the abstract, ...

Theorem 2 is a combination of early-stage feature learning (Stage 1 and 2) and final-stage feature learning (Stage 3). This sentence refers to the local convergence result (Stage 3), and the threshold is $\varepsilon_0$ in Theorem 2. We show at the end of Stage 2 (Lemma 4) the loss/optimality gap is below $\varepsilon_0$ so that we enter the local convergence regime in Stage 3 (Lemma 5). We agree that this should have been made clearer and we will address this issue in the revision.

2&3. The relation between assumption 3 and information exponent should be made explicit ...

Thanks for the suggestion. We will make it more explicit in the revision. The information exponent of teacher network is indeed 1. Assumption 3 we made here is the same as in Damian et al. (2022), and we use the same preprocessing procedure to remove the linear part. After this preprocessing step, the information exponent becomes 2 due to Assumption 3.

Our results hold when information exponent is 2 after preprocessing. If it's larger, the current analysis does not work, and it would require a generalization of Damian et al. (2022) to make it work (Dandi et al. (2023) explored a generalized version of this problem and left it as a conjecture).

Neural networks can learn representations with gradient descent, Damian et al., 2022

How Two-Layer Neural Networks Learn, One (Giant) Step at a Time, Dandi et al., 2023

4. Why is the head a and the back layer W ...

• Regularization on $a$ and $W$:

This is the same as using weight decay on both layers. At the minimum we have $\|\|w_i\|\|_2^2+|a_i|^2 = 2\|\|w_i\|\|_2|a_i|$, so $\ell_2$ regularization $\sum_i\|\|w_i\|\|_2^2+|a_i|^2$ becomes $2\sum_i\|\|w_i\|\|_2|a_i|$ that is effectively $\ell_1$ regularization over the norm $\|\|w_i\|\|_2|a_i|$ of neurons. Such a regularization favors sparse solution and helps us to recover the ground-truth directions and remove irrelevant directions.

• Analyze the problem of $\lambda=0$ or $\lambda\to0$:

Our goal is to minimize the loss with $\lambda=0$. However, directly analyzing the unregularized case is challenging so we choose to analyze the case where $\lambda$ is positive but gradually decreasing to 0 ($\lambda\to 0$). In this way, the solution we get will eventually converge to the minima with $\lambda=0$. In Theorem 2 we show that the unregularized loss ($\lambda=0$) is small at the end of training.

5. The analysis is for gradient descent ran on expected loss ...

Please refer to the general response where we address the questions regarding sample complexity.

6. The algorithm analyzed in this paper ...

• Non-standard algorithm:

Recent feature learning literature often focuses on early-stage learning, using layer-wise training (training 1st-layer weights $W$ first, then 2nd-layer weights $a$), which is also not standard in practice. As discussed in the paper, analyzing the entire training dynamics of standard gradient descent, especially the middle stage (Stage 2), is technically challenging and remains an open problem. We hope our results contribute to understanding standard training methods.

• Role of Stage 2:

In our algorithm, the Stage 1 and 2 essentially follow the similar layer-wise training procedure that are common in early-stage feature learning literature. The role of Stage 2 is to do a regression on top of the learned feature after Stage 1. This makes the loss small so that we enter the local convergence regime in Stage 3.

• Regularization and balancing:

We assume `normalizing' mentioned by reviewer refers to the regularization or the norm balance. We'd be happy to further clarify if this does not answer the original question.

The regularization is essentially the same as $\ell_2$ regularization or weight decay used in practice, as $\|\|w_i\|\|_2^2 + |a_i|^2 \ge 2\|\|w_i\|\|_2|a_i|$, with equality at the minimum of $\ell_2$ regularized loss. The reason of using weight decay is further discussed in Section 5: it induces an effective $\ell_1$ regularization over the neurons and helps to reduce norm cancellation between neurons.

The norm balancing step is just a preparation step for Stage 3 for technical convenience to enter the local convergence regime.

7. In the discussion below Theorem 2, ...

We will make it clear in the revision. To be more specific, we could look at Damian et al. (2022) (although the discussion also applies to other works that rely on (large) first-step gradient, e.g., Ba et al. (2022), Abbe et al. (2022)). It is shown that after the first-step gradient update neurons $w_i^{(1)}\approx c Hw_i^{(0)}$, where $c$ is a scalar and $w_i^{(0)}$ is neuron $w_i$ at random initialization. Given the assumption that matrix $H$ has full rank within the target subspace, $w_i^{(1)}$ is essentially sampled randomly from target subspace. This suggests neural networks in fact only learn the target subspace and do random feature within it.

8. Instead of the lengthly discussion on the construction of the dual certificate, ....

Thanks for the suggestion. We do believe the construction of the dual certificate is the core technical contribution. However We will clarify general take-aways more clearly in the revision. Our main finding is that training both layers of 2-layer networks to convergence can achieve 0 loss and recover ground-truth directions. This finding highlights a strong form of feature learning that complements the early-stage feature learning literature, where the first-layer weights are typically fixed after one or a few steps of gradient descent.

9. Missing citations and discussion of prior work...

Thanks for the references. We will add discussions about these works in the revision.