Title: How does Gradient Descent Learn Features --- A Local Analysis for Regularized Two-Layer Neural Networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: deep learning theory

feature learning

optimization
Abstract: The ability of learning useful features is one of the major advantages of neural networks. Although recent works show that neural network can operate in a neural tangent kernel (NTK) regime that does not allow feature learning, many works also demonstrate the potential for neural networks to go beyond NTK regime and perform feature learning. Recently, a line of work highlighted the feature learning capabilities of the early stages of gradient-based training. In this paper we consider another mechanism for feature learning via gradient descent through a local convergence analysis. We show that once the loss is below a certain threshold, gradient descent with a carefully regularized objective will capture ground-truth directions. We further strengthen this local convergence analysis by incorporating early-stage feature learning analysis. Our results demonstrate that feature learning not only happens at the initial gradient steps, but can also occur towards the end of training.

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

Abstract

The ability of learning useful features is one of the major advantages of neural networks. Although recent works show that neural network can operate in a neural tangent kernel (NTK) regime that does not allow feature learning, many works also demonstrate the potential for neural networks to go beyond NTK regime and perform feature learning. Recently, a line of work highlighted the feature learning capabilities of the early stages of gradient-based training. In this paper we consider another mechanism for feature learning via gradient descent through a local convergence analysis. We show that once the loss is below a certain threshold, gradient descent with a carefully regularized objective will capture ground-truth directions. We further strengthen this local convergence analysis by incorporating early-stage feature learning analysis. Our results demonstrate that feature learning not only happens at the initial gradient steps, but can also occur towards the end of training.

1 Introduction

Feature learning has long been considered to be a major advantage of neural networks. However, how gradient-based training algorithms can learn useful features is not well-understood. In particular, the most widely applied analysis for overparametrized neural networks is the neural tangent kernel (NTK) (Jacot et al., 2018; Du et al., 2019; Allen-Zhu et al., 2019b). In this setting, the neurons don't move far from their initialization and the features are determined by the network architecture and random initialization (Chizat et al., 2019).

While there are empirical and theoretical evidence on the limitation of NTK regime (Chizat et al., 2019; Arora et al., 2019), extending the analysis beyond the NTK regime has been challenging. For 2-layer networks, an alternative framework for analyzing overparametrized neural networks called mean-field analysis was introduced. Earlier mean-field analysis (e.g., Chizat and Bach, 2018; Mei et al., 2018) require either infinite or exponentially many neurons. Later works (e.g., Li et al., 2020; Ge et al., 2021; Bietti et al., 2022; Mahankali et al., 2024) can analyze the training dynamics of mildly overparametrized networks with polynomially many neurons with stronger assumptions on the ground-truth function.

Recently, a growing line of works (Daniely and Malach, 2020; Damian et al., 2022; Abbe et al., 2021, 2022, 2023; Yehudai and Shamir, 2019; Shi et al., 2022; Ba et al., 2022; Mousavi-Hosseini et al., 2023; Barak et al., 2022; Dandi et al., 2023; Wang et al., 2024; Nichani et al., 2024a,b) showed that early stages of gradient training (either one/a few steps of gradient descent or a small amount of time of gradient flow) can be useful in feature learning. These works show that after the early stages of gradient training, the first layer in a 2-layer neural network already captures useful features (usually in the form of a low dimensional subspace), and continuing training the second layer weights will

Rong Ge

give performance guarantees that are stronger than any kernel or random feature based models. In this work, we consider the natural follow-up question:

Does feature learning only happen in the early stages of gradient training?

We show that this is not the case by demonstrating feature learning capability for the final stage of gradient training - local convergence. In particular, we prove the following result:

Theorem 1 (Informal) . If the data is generated by a 2-layer teacher network f ∗ , as long as the width of student network m is at least some quantity m 0 that only depends on f ∗ , a variant of gradient descent algorithm (Algorithm 1, roughly gradient descent with decreasing weight decay) can recover the target network within polynomial time. Moreover, the student neurons align with the teacher neurons at the end.

Our result highlights the different mechanisms of feature learning: previous works show that in the early stages of gradient descent, the network learns the subspace spanned by the neurons in the teacher network. Our local convergence result shows that at later stages, gradient descent is able to learn the exact directions of the teacher neurons, which are much more informative compared to the subspace and lead to stronger guarantees.

Analyzing the entire training dynamics is still challenging so in our algorithm (see Algorithm 1) we use a convex second stage to 'fast-forward' to the local analysis. Our technique for local convergence is similar to the earlier work (Zhou et al., 2021), however we consider a more complicated setting with ReLU activation and allow second-layer weights to be both positive or negative. This change requires additional regularization in the form of standard weight decay and new dual certificate analysis.

1.1 Related works

Neural Tangent Kernel Early works often studied neural network optimization using NTK theory (Jacot et al., 2018; Allen-Zhu et al., 2019b; Du et al., 2019). It is shown that highly-overparametrized neural nets are essentially kernel methods under certain initialization scale. However, NTK theory cannot explain the performance of neural nets in practice (Arora et al., 2019) and leads to lazy training dynamics that neurons stay close to their initialization (Chizat et al., 2019). Hence, later research efforts (e.g., Allen-Zhu et al., 2019a; Bai and Lee, 2020; Li et al., 2020), as well as current paper, focus on feature learning regime where neural nets can learn features and outperform kernel methods.

Early stage feature learning Researchers have recently tried to understand how neural networks trained with gradient descent (GD) can learn features, going beyond the kernel/lazy regime (Jacot et al., 2018; Chizat et al., 2019). A typical setup is to use 2-layer neural networks to learn certain target function, often equipped with low-dimensional structure. Examples include learning polynomials (Yehudai and Shamir, 2019; Damian et al., 2022), single-index models (Ba et al., 2022; MousaviHosseini et al., 2023; Moniri et al., 2024; Cui et al., 2024), multi-index models (Dandi et al., 2023), sparse boolean functions (Abbe et al., 2021, 2022, 2023), sparse parity functions (Daniely and Malach, 2020; Shi et al., 2022; Barak et al., 2022) and causal graph (Nichani et al., 2024b). Also, few works use 3-layer networks as learner model (Nichani et al., 2024a; Wang et al., 2024). These works essentially showed that feature learning happens in the early stage of training. Specifically, they often use 2-stage layer-wise training procedure: first-layer weights/features are only trained with one or few steps of gradient descent/flow and only update the second-layer afterwards. Our results give a complementary view that feature learning can also happen in the final stage training that leading student neurons eventually align with ground-truth directions. This cannot be achieved if first-layer weights are fixed after few steps.

Learning single/multi-index models with neural networks Single/Multi-index models are the functions that only depend on one or few directions of the high dimensional input. Many recent works have studied the problem of using 2-layer networks to learn single-index models (Soltanolkotabi, 2017; Yehudai and Ohad, 2020; Frei et al., 2020; Wu, 2022; Bietti et al., 2022; Xu and Du, 2023; Berthier et al., 2023; Mahankali et al., 2024) and multi-index models (Damian et al., 2022; Bietti et al., 2023; Suzuki et al., 2024; Glasgow, 2024). These works show the advantages of feature learning over fixed random features in various settings. In this paper, we consider target multi-index function that can be represented by a small 2-layer network, and show a variant of GD with weight decay can learn it and, moreover, recover the ground-truth directions.

Local loss landscape Safran et al. (2021) showed that in the overparametrized case with orthogonal teacher neurons, even around the local region of global minima, the landscape neither is convex nor satisfies PL condition. Chizat (2022) considered square loss with ℓ 2 regularization similar to our setup and showed the local loss landscape is strongly-convex under certain non-degenerate assumptions. However, it is not known when such assumptions actually hold and the proof cannot handle ReLU. Later Akiyama and Suzuki (2021) gives a result for ReLU, but the non-degeneracy assumption is still required (and also focus on effective ℓ 1 regularization instead of ℓ 2 regularization). Zhou et al. (2021) studies a similar local convergence setting but restricts second-layer weights to be positive and uses absolute activation. In this paper, we focus on a more natural but technically challenging case that second-layer can be positive and negative and using ReLU activation. We develop new techniques to overcome the above challenges (additional assumption, ReLU, standard second-layer, etc).

2 Preliminary

Notation Let [ n ] be set { 1 , . . . , n } . For vector w , we use ∥ w ∥ 2 for its 2-norm and w = w / ∥ w ∥ 2 as its normalized version. For two vectors w , v we use ∠ ( w , v ) = arccos( | w ⊤ v | / ( ∥ w ∥ 2 ∥ v ∥ 2 )] ∈ [0 , π/ 2] as the angle between them (up to a sign).For matrix A let ∥ A ∥ F be its Frobenius norm. We use standard O, Ω , Θ to hide constants and ˜ O, ˜ Ω , ˜ Θ to hide polylog factors. We use O ∗ , Ω ∗ , Θ ∗ to hide problem dependent parameters that only depend on the target network (see paragraph above (1)).

Teacher-student setup We will consider the teacher-student setup for two-layer neural networks with Gaussian input x ∼ N ( 0 , I d ) . The goal is to learn the teacher network of size m ∗

where σ ( x ) := max { 0 , x } is ReLU activation, S ∗ := span { w ∗ 1 , . . . , w ∗ m ∗ } is the target subspace. Without loss of generality, we will assume ∥ w ∗ i ∥ 2 = 1 due to the homogeneity of ReLU.

Following the recent line of works in learning single/multi-index models (Ba et al., 2022; Damian et al., 2022), we assume the target network has low dimensional structure.

Assumption 1. Teacher neurons form a low dimensional subspace in R d , that is

We will also assume the teacher neurons are non-degenerate in the following sense:

̸

Assumption 2. Teacher neurons are ∆ -separated, that is angle ∠ ( w ∗ i , w ∗ j ) ≥ ∆ for all i = j .

Assumption 3. H := ∑ m ∗ i =1 a ∗ i w ∗ i w ∗⊤ i is non-degenerate in target subspace S ∗ , i.e., rank ( H ) = r . Denote κ := | λ r ( H ) | .

Assumption 2 simply requires all teacher neurons pointing to different directions, which is crucial for identifiability (Zhou et al., 2021).

Assumption 3 says the target network contains low-order (second-order) information, which is related with the notion of information exponent (Arous et al., 2021). In our setting, the information exponent is at most 2 due to Assumption 3. Indeed, one can show E x [ f ∗ ( x ) h 2 ( v ⊤ x )] = ˆ σ 2 v ⊤ Hv , where h 2 ( x ) is the 2nd-order normalized Hermite polynomial and ˆ σ 2 is the 2nd Hermite coefficient of ReLU. See Appendix A for more details. Many previous works also rely on same or similar assumption to show neural networks can learn features to perform better than kernels (Damian et al., 2022; Abbe et al., 2022; Ba et al., 2022).

In this paper, we are interested in the case where the complexity of target network is small. Therefore, we will use O ∗ , Ω ∗ , Θ ∗ to hide poly( r, m ∗ , ∆ , | a 1 | , . . . , | a m ∗ | , κ ) , which is the polynomial dependency on relevant parameters of target f ∗ (does not depend on student network).

We will use the following overparametrized student network:

where a = ( a 1 , . . . , a m ) ⊤ ∈ R m , W = ( w 1 · · · w m ) ⊤ ∈ R m × d and θ = ( a , W , α, β ) .

Loss and algorithm Consider the square loss function with ℓ 2 regularization under Gaussian input

We will use L to denote the square loss for simplicity. The ℓ 2 regularization is the same as the commonly used weight decay in practice. Our goal is to find the minima of unregularized problem ( λ = 0 ) to recover teacher network f ∗ . However, directly analyzing the unregularized problem is challenging so instead we choose to analyze the regularized problem and will gradually let λ → 0 .

In above, we use preprocessed data ( x, ˜ y ) in the loss function as in Damian et al. (2022). Specifically, given any ( x , y ) with y = f ∗ ( x ) , denote α ∗ = E x [ y ] and β ∗ = E x [ y x ] , we get

This preprocessing process essentially removes the 0-th and 1-st order term in the Hermite expansion of σ . See Appendix A for a brief introduction of Hermite polynomials and Claim B.1.

Our algorithm is shown in Algorithm 1. It is roughly the standard GD following a given schedule of weight decay λ t that goes to 0. Due to the difficulty in analyzing gradient descent training beyond early and final stage, we choose to only train the norms in Stage 2 as a tractable way to reach the local convergence regime.

We will use symmetric initialization that a i = -a i + m/ 2 , w i = w i + m/ 2 with a i ∼ Unif {± √ d } , w i ∼ Unif((1 / √ m ) S d -1 ) , α = 0 , β = 0 . Our analysis is not sensitive to the initialization scale we choose here. The choice is just for the simplicity of the proof.

Algorithm 1: Learning 2-layer neural networks

3 Main results

In this section, we give our main result that shows training student network using Algorithm 1 can recover the target network within polynomial time. We will focus on the case that d ≥ Ω ∗ (1) when the complexity of target function is small.

Theorem 2 (Main result) . Under Assumption 1, 2, 3, consider Algorithm 1 on loss (2) . There exists a schedule of weight decay λ t and step size η t such that given m ≥ m 0 = ˜ O ∗ (1) · (1 /ε 0 ) O ( r ) neurons with small enough ε 0 = Θ ∗ (1) , with high probability we will recover the target network L ( θ ) ≤ ε within time T = O ∗ (1 /ηε 2 ) where η = poly( ε, 1 /d, 1 /m ) .

Moreover, when ε → 0 every student neuron w i either aligns with one of teacher neuron w ∗ j as ∠ ( w i , w ∗ j ) = 0 or vanishes as | a i | = ∥ w i ∥ = 0 .

Note that our results can be extended to only have access to polynomial number of samples by using standard concentration tools. We omit the sample complexity for simplicity. See more discussion in Appendix J. We emphasize that the required width m 0 only depends on the complexity of target function f ∗ (only quantities that are related to f ∗ , not student network f or error ε ), so any mildly overparametrized networks can learn f ∗ efficiently to arbitrary small error.

The analysis consists of three stages: early-stage feature learning (Stage 1 and 2) and final-stage feature learning/local convergence (Stage 3). It will be clear in the later section that ε 0 is in fact the threshold to enter the local convergence regime. See Section 4 for more details.

Our result improves the previous works that only train the first layer weight with small number of gradient steps at the beginning (Damian et al., 2022; Ba et al., 2022; Abbe et al., 2021, 2022, 2023). In these works, neural networks only learn the target subspace and do random features within it (see Section 4.1 for more details). Intuitively, these random features need to span the whole space of the target function class to perform well, which means its number (the width) should be on the order of the dimension of target function class. For 2-layer networks, random features in the target subspace need (1 /ε ) O ( r ) neurons to achieve desired accuracy ε . In contrast, continue training both layer at the last phase of training allows us to learn not only subspace but also exactly the ground-truth directions. Moreover, we only use (1 /ε 0 ) O ( r ) neurons that only depends on the complexity of target network. This highlights the benefit of continue training first layer weights instead of fixing them after first step.

4 Proof overview

In this section, we give the proof overview of these three stages separately.

Denote the optimality gap ζ at time t as the difference between current loss and the best loss one could achieve with networks of any size (including infinite-width networks)

where M ( S d -1 ) is the set of measures on the sphere S d -1 . As an example, if µ = ∑ i a i ∥ w i ∥ δ w i , then L λ ( µ ) recovers L λ ( θ ) when linear term α, β are perfectly fitted and norms are balanced | a i | = ∥ w i ∥ . We defer the precise definition of L λ ( µ ) to (6) in appendix.

4.1 Stage 1

For Stage 1, we show in the lemma below that the first step of gradient descent identifies the target subspace and ensures there always exists student neuron that is close to every teacher neuron.

Lemma 3 (Stage 1) . Under Assumption 1,2,3, consider Algorithm 1 with λ 0 = η 0 = 1 and m ≥ m 0 = ˜ O ∗ (1) · (1 /ε 0 ) O ( r ) with any ε 0 = Θ ∗ (1) . After first step, with probability 1 -δ we have

(i) for every teacher neuron w ∗ i , there exists at least one student neuron w j s.t. ∠ ( w ∗ i , w j ) ≤ ε 0 .

The key observation here is similar to Damian et al. (2022) that w (1) i ≈ -2 η 0 a (0) i ( ˆ σ 2 2 Hw i ) so that given H is non-degenerate in target subspace S ∗ we essentially sample w (1) i from the target subspace. It is then natural to expect that the neurons form an ε 0 -net in the target subspace given m 0 neurons.

4.2 Stage 2

Given the learned features (first-layer weights) in Stage 1, we now perform least squares to adjust the norms and reach a low loss solution in Stage 2.

Lemma 4 (Stage 2) . Under Assumption 1,2,3, consider Algorithm 1 with λ t = √ ε 0 . Given Stage 1 in Lemma 3, we have Stage 2 ends within time T 2 = ˜ O ∗ (1 /ηε 0 ) such that optimality gap ζ T 2 = O ∗ ( ε 0 ) .

It remains an open problem to prove the convergence when training both layers simultaneously beyond early and final stage. To overcome this technical challenge, we choose to use a simple least square for Stage 2. We use the simple (sub)gradient descent to optimize this loss. There exist many other algorithms that can solve this Lasso-type problem, but we omit it for simplicity as this is not the main focus of this paper.

Note that the regularization in Algorithm 1 is the same as standard weight decay when we train both layers. This regularization leads to several desired properties at the end of Stage 2: (1) prevent norm cancellation between neurons: neurons with similar direction but different sign of second layer weights cancel with each other; (2) neurons mostly concentrate around ground-truth directions. As we will see later, these nice properties continue to hold in Stage 3, thanks to the regularization.

4.3 Stage 3

After Stage 2 we are in the local convergence regime. The following lemma shows that we could recover the target network within polynomial time using a multi-epoch gradient descent with decreasing weight decay λ at every epoch. Note that this result only requires the initial optimality gap is small and width m ≥ m ∗ (target network width, not m 0 ).

Lemma 5 (Stage 3) . Under Assumption 1,2,3, consider Algorithm 1 on loss (2) . Given Stage 2 in Lemma 4, if the initial optimality gap ζ 3 , 0 ≤ O ∗ ( λ 9 / 5 3 , 0 ) , weight decay λ follows the schedule of initial value λ 3 , 0 = O ∗ (1) , and k -th epoch λ 3 ,k = λ 3 ,k -1 / 2 and stepsize η 3 k = η ≤ O ∗ ( λ 12 3 ,k d -3 ) for all T 3 ,k ≤ t ≤ T 3 ,k +1 in epoch k , then within K = O ∗ (log(1 /ε )) epochs and total T 3 -T 2 = O ∗ ( λ -4 3 , 0 η -1 ε -2 ) time we recover the ground-truth network L ( θ ) ≤ ε .

The lemma above relies on the following result that shows the local landscape is benign in the sense that it satisfies a special case of Łojasiewicz property (Lojasiewicz, 1963). This means GD can always make progress until the optimality gap ζ is small.

Lemma 6 (Gradient lower bound) . When Ω ∗ ( λ 2 ) ≤ ζ ≤ O ∗ ( λ 9 / 5 ) and λ ≤ O ∗ (1) , we have

Note that this generalizes previous result in Zhou et al. (2021) that only focuses on 2-layer networks with positive second layer weights. This turns out to be technically challenging as two neurons with different signs can cancel each other. We discuss how to deal with this challenge in the next section.

5 Descent direction in local convergence (Stage 3): the benefit of weight decay

In this section, we give the high-level proof ideas for the most technical challenging part of our results -characterize the local landscape in Stage 3 (Lemma 6).

The key idea is to construct descent direction - a direction that has positive correlation with the gradient direction. The gradient lower bound follows from the existence of such descent direction.

It turns out that the existence of both positive and negative second-layer weights introduces significant challenge for the analysis: there might exist neurons with similar directions (e.g., ( a, w ) and ( -a, w ) ) that can cancel with each other to have no effect on the output of network. Intuitively, we would hope all of them to move towards 0, but they have no incentive to do so. Moreover, if they are not exactly symmetric it's hard to characterize which directions these neurons will move.

We use standard weight decay to address the above challenge. Specifically, weight decay helps us to

Balance norm between neurons . When norm between two layers are balanced, the ℓ 2 regularization ∑ i | a i | 2 + ∥ w i ∥ 2 would become the effective ℓ 1 regularization 2 ∑ i | a i | ∥ w i ∥ over the distribution of neurons. Such sparsity penalty ensures most neurons concentrate around the ground-truth directions, especially preventing norm cancellation between far-away neurons.

Reduce cancellation between close-by neurons . For close-by neurons, weight decay helps to reduce the norm of neurons with the 'incorrect' sign (different sign with the ground-truth neuron). This is because weight decay prefers low norm solutions, and reducing cancellations between neurons can reduce total norm (regularization term) while keeping the square loss same.

̸

We will group the neurons (i.e., partitioning S d -1 ) based on their distance to the closest teacher neurons: denote T i = { w : ∠ ( w , w ∗ i ) ≤ ∠ ( w , w ∗ j ) for any j = i } (break the tie arbitrarily) so that ∪ i T i = S d -1 . We will also use δ j to denote ∠ ( w j , w ∗ i ) for j ∈ T i .

As described above, weight decay can always lead to descent direction when norms are not balanced or norm cancellation happens (see Lemma F.15 and Lemma F.16). The following lemma shows that in other scenarios we can always improve features towards the ground-truth directions.

Lemma 7 (Feature improvement descent direction, informal) . When norms are balanced and no norm cancellation happens, there exists properly chosen q ij ≥ 0 and ∑ j ∈T i a j q ij = a ∗ i such that

Teacherneuron

Studentneuron

Descentdirection

In words, this descent direction is the following: we move neuron w j ∈ T i toward either ground-truth direction w ∗ i or 0 depending on whether it is in the neighborhood of teacher neuron w ∗ i . Specifically, we move far-away neurons towards 0 (and thus setting q ij = 0 ) and move close-by neurons towards its 'closest' minima q ij w ∗ i (the fraction of w ∗ i that neuron w j should target to approximate). See Figure 1 for an illustration.

The proof of the above lemma requires a dedicated characterization of the low loss solution's structure, which we describe in Section 6.

6 Structure of (approximated) minima

In this section, we first highlight the importance of understanding local geometry by showing the challenges in proving the existence of descent direction (Lemma 7). Then after presenting the main result of this section to show the structure of (approximated) minima (Lemma 8), we discuss several proof ideas such as dual certificate analysis in the remaining part.

6.1 Constructing descent direction requires better understanding of local geometry

To show the existence of descent direction in Lemma 7, we compute the inner product between gradient and constructed descent direction. We can lower bound it by (assuming norms are balanced)

where R ( x ) = f ( x ) -˜ f ∗ ( x ) is the residual. Thus, in order to get a lower bound, the goal is to show second term above is small than ζ . As we can see, this term is quite complicated and can be viewed as the inner product between R ( x ) and h ( x ) = ∑ i ∈ [ m ∗ ] ∑ j ∈T i a j q ij w ∗⊤ i x ( σ ′ ( w ∗⊤ i x ) -σ ′ ( w ⊤ j x )) .

Average neuron and residual decomposition To deal with above challenge, we use the idea of average neuron and residual decomposition. For each teacher neuron w ∗ i , denote v i = ∑ j ∈T i a j w j as the average neuron. Intuitively, this average neuron v i stands for an idealize case where all neurons belong to T i (closer to w ∗ i than other w ∗ j ) collapse into a single neuron.

We decompose the residual R ( x ) = f ( x ) -˜ f ∗ ( x ) into the 3 terms below: denote ˆ v i = v i -w ∗ i

R 1 can be thought as the exact-parametrization setting (use m ∗ neurons to learn m ∗ neurons), where the average neurons { v i } m ∗ i =1 are the effective neurons. The difference between this exactparametrization and overparametrization setting is then characterized by the term R 2 , which captures the difference in nonlinear activation pattern. This term in fact suggests the loss landscape is degenerate in overparametrized case and slows down the convergence (Zhou et al., 2021; Xu and Du, 2023). Overall, this residual decomposition is similar to Zhou et al. (2021), with additional modification of R 3 to deal with ReLU activation and linear term α, β .

To some extent, our residual decomposition can be viewed as a kind of 'bias-variance' decomposition in the sense that the 'bias' term R 1 captures the overall average contribution of all neurons, and the 'variance' term R 2 captures the individual contributions of each neuron that are not reflected in R 1 .

High-level proof plan of Lemma 7 We now are ready to give a proof plan for Lemma 7. The key is to show properties of minima that can help us to bound ⟨ R,h ⟩ .

Show that neurons mostly concentrate around ground-truth directions.

Show that average neuron v i is close to teacher neuron w ∗ i for all i ∈ [ m ] .

Use above structure to bound ⟨ R i , h ⟩ . Specifically, bounding ⟨ R 1 , h ⟩ relies on the fact that average neuron is close to teacher neuron (step 2); a bound on ⟨ R 2 , h ⟩ follows from far-away neurons are small (step 1); third term ⟨ R 3 , h ⟩ can be directly bounded using the loss. Detailed calculations are deferred into Appendix H.3.

We give main result of this section that shows the desired local geometry properties more precisely ((i)(ii) corresponding to step 1 and (iii) corresponding to step 2 above).

Lemma 8 (Informal) . Suppose the optimality gap is ζ , we have

(i) Total norm of far-away neurons is small: ∑ i ∈ [ m ∗ ] ∑ j ∈T i | a j | ∥ w j ∥ 2 δ 2 j = O ∗ ( ζ/λ ) , where angle δ j = ∠ ( w j , w ∗ i ) for w j that j ∈ T i .

(ii) For every w ∗ i , there exists at least one close-by neuron w s.t. ∠ ( w , w ∗ i ) ≤ δ close = O ∗ ( ζ 1 / 3 ) .

(iii) Average neuron is close to teach neurons: we have ∥ v i -w ∗ i ∥ 2 ≤ O ∗ (( ζ/λ ) 3 / 4 ) .

These properties give us a sense of what the network should look like when loss is small: neurons have large norm only if they are around the ground-truth directions. Moreover, when ζ/λ → 0 , student neuron must align with one of teacher neurons ( δ j = 0 ) or norm becomes 0 ( | a j | ∥ w j ∥ = 0 ). This can be understood from the ℓ 1 regularized loss (equivalent to ℓ 2 regularization on both layers) that promotes the sparsity over the distribution of neurons. In the rest of this section, we discuss new techniques such as dual certificate that we develop for the proof.

6.2 Neurons concentrate around teacher neurons: dual certificate analysis and test function

We focus on Lemma 8(i)(ii) here. We will use a dual certificate technique similar to Poon et al. (2023) to prove Lemma 8(i), and a more general construction of test function to prove Lemma 8(ii). In below, we consider a relaxed version of original optimization problem (2) by allowing infinite number of neurons, i.e., distribution of neurons, with σ ≥ 2 ( x ) = ReLU ( x ) -1 / √ 2 π -x/ 2 instead of ReLU:

where µ ∗ λ is the minimizer. We use σ ≥ 2 activation because this is the effective activation when linear terms α, β are perfectly fitted (remove 0th and 1st order Hermite expansion of ReLU, see Claim B.1 and (6) in appendix).

This is the loss function we would have in the idealized setting: (1) linear term α, β reach their global minima (this is easy to achieve as loss is convex in them); (2) use ℓ 1 regularization instead of ℓ 2 regularization, since this is the case when the first and second layer norm are balanced (weight decay encourages this to happen). Note that the results in this part can handle almost all activation as long as its Hermite expansion is well-defined, generalizing Zhou et al. (2021) that can only handle absolute/ReLU activation. In below we will focus on the activation σ ≥ 2 for simplicity.

Dual certificate This optimization problem (5) can be viewed as a natural extension of the classical compressed sensing problem (Donoho, 2006; Candès et al., 2006) and Lasso-type problem (Tibshirani, 1996) in the infinite dimensional space, which has been studied in recent years (Bach, 2017; Poon et al., 2023). One common way is to study its dual problem. The dual solution p 0 ( x ) (maps R d to R ) of (5) when λ = 0 satisfies E x [ p 0 ( x ) σ ≥ 2 ( w ⊤ x )] ∈ ∂ | µ ∗ | ( S d -1 ) (more detailed discussions on this dual problem can be found in e.g., Poon et al. (2023)). Here η ( w ) = E x [ p ( x ) σ ≥ 2 ( w ⊤ x )] is often called dual certificate, as it serves as a certificate of whether a solution µ is optimal. Its meaning will be clear in the discussions below.

We now introduce the notion of non-degenerate dual certificate, motivated by Poon et al. (2023). Note that the condition η ( w ) ∈ ∂ | µ ∗ | ( S d -1 ) implies that η ( w ∗ i ) = sign( a ∗ i ) and ∥ η ∥ ∞ ≤ 1 . The following definition is a slightly stronger version of the above implications as it requires η to decay at least quadratic when moves away from w ∗ i .

Definition 1 (Non-degenerate dual certificate) . η ( w ) is called a non-degenerate dual certificate if there exists p ( x ) such that η ( w ) = E x [ p ( x ) σ ≥ 2 ( w ⊤ x )] for w ∈ S d -1 and

The existence and construction of the non-degenerate dual certificate is deferred to Appendix G. We focus on the implications of such nondegenerate dual certificate below.

+1

0

η

(

w

)

w

∗

1

w

∗

2

w

∗

3

-1

Roughly speaking, the dual certificate only focuses on the position of ground-truth directions w ∗ i as it decays fast when moving away from these directions (Figure 2). Thus, if µ exactly recovers ground-truth µ ∗ , then we have ⟨ η, µ ∗ ⟩ = | µ ∗ | 1 . The gap between ⟨ η, µ ⟩ and | µ | 1 is large when µ is away from µ ∗ . Therefore, η can be viewed as a certificate to test the optimality of µ . The lemma below makes it more precise.

Lemma 9. Given a non-degenerate dual certificate η , then

In the finite width case, we have ∑ i ∈ [ m ∗ ] ∫ T i δ ( w , w ∗ i ) 2 d | µ | ( w ) = ∑ i | a i | ∥ w i ∥ δ 2 i . This is exactly the quantity that we are interested in Lemma 8.

To see the usefulness of Lemma 9, we show a proof for total norm bound of the optimal solution µ ∗ λ . The proof for general µ with optimality gap ζ is similar (Lemma F.5).

Proof. It is not hard to show | µ ∗ λ | 1 ≤ | µ ∗ | 1 (Lemma F.3) so we have

Using Lemma 9 and the fact L ( µ ∗ λ ) = O ∗ ( λ 2 ) from Lemma F.3,

Test function The idea of using test function is to identify certain properties of the target function/distribution that we are interested in. Specifically, we construct test function so that it only correlates well with the target function that has the desired property. Generally speaking, the dual certificate above can be consider as a specific case of a test function: the correlation between dual certificate η and distribution of neurons µ is large (reach | µ | 1 ) only when µ ≈ µ ∗ .

In below, we use this test function idea to show that every ground-truth direction has close-by neuron (Lemma 8(ii)). Denote T i ( δ ) := { j : ∠ ( w j , w i ) ≤ δ } ∩ T i as the neurons that are δ -close to w ∗ i .

Lemma 10 (Lemma 8(ii), informal) . Given the optimality gap ζ , we have the total mass near each target direction is large, i.e., µ ( T i ( δ )) sign( a ∗ i ) ≥ | a ∗ i | / 2 for all i ∈ [ m ∗ ] and any δ ≥ Θ ∗ ( ζ 1 / 3 ) .

Note that although the results in the dual certificate part (Lemma 9(ii)) can imply that there are neurons close to teacher neurons, the bound we get here using carefully designed test function are sharper ( ζ 1 / 3 vs. ζ 1 / 4 ). This is in fact important to the descent direction construction (Lemma 7).

In the proof, we view the residual R ( x ) = f µ ( x ) -f ∗ ( x ) as the target function and construct test function that will only have large correlation if there is a teacher neuron that have no close student neurons. Specifically, the test function g only consists of high-order Hermite polynomial such that it is large around the ground-truth direction and decays fast when moving away (Figure 3). It looks like a single spike in dual certificate η , but in fact decays much faster than η when moving away. It is more

0

g

(

x

)

w

∗

1

w

∗

2

w

∗

3

flexible to choose test function than dual certificate, so test function g can focus only on a local region of one ground-truth direction and give a better guarantee than dual certificate analysis.

6.3 Average neuron is close to teacher neuron: residual decomposition and average neuron

We give the proof idea for Lemma 8(iii) that shows average neuron v i is close to teacher neuron w ∗ i using the residual decomposition R = R 1 + R 2 + R 3 .

. .

The key is to observe that R 1 is an analogue to exact-parametrization case where loss is often strongly-convex, so we have ∥ R 1 ∥ 2 2 = Ω ∗ (1) ∑ i ∥ v i -w ∗ i ∥ 2 2 . Then the goal is to upper bound ∥ R 1 ∥ Given the decomposition R = R 1 + R 2 + R 3 , it is easy to bound ∥ R 1 ∥ ≤ ∥ R ∥ + ∥ R 2 ∥ + ∥ R 3 ∥ We focus on ∥ R 2 ∥ as the other two are not hard to bound (loss is small in local regime). R 2 is in fact closely related with the total weighted norm bound in Lemma 8: we show ∥ R 2 ∥ = O ∗ (1) ( ∑ j ∈T i | a j | ∥ w j ∥ 2 δ 2 j ) 3 / 2 = O ∗ (( ζ/λ ) 3 / 2 ) . Thus, we get a bound for ∥ v i -w ∗ i ∥ . See Appendix F.1.4 for details.

7 Conclusion

In this paper we showed that gradient descent converges in a large local region depending on the complexity of the teacher network, and the local convergence allows 2-layer networks to perform a strong notion of feature learning (matching the directions of ground-truth teacher networks). We hope our result gives a better understanding of why gradient-based training is important for feature learning in neural networks. Our results rely on adding standard weight decay and new constructions of dual certificate and test functions, which can be helpful in understanding local optimization landscape in other problems. A natural but challenging next step is to understand whether the intermediate steps are also important for feature learning.

References

Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems , 34:26989-27002, 2021.

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory , pages 4782-4887. PMLR, 2022.

Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory , pages 2552-2623. PMLR, 2023.

P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds . Princeton University Press, 2008.

P-A Absil, Robert Mahony, and Jochen Trumpf. An extrinsic look at the riemannian hessian. In International conference on geometric science of information , pages 361-368. Springer, 2013.

Shunta Akiyama and Taiji Suzuki. On learnability via gradient method for two-layer relu neural networks in teacher-student setting. In International Conference on Machine Learning , pages 152-162. PMLR, 2021.

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in neural information processing systems , 32, 2019a.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In International conference on machine learning , pages 242-252, 2019b.

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems , 32, 2019.

Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research , 22(1):4788-4838, 2021.

Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. Highdimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems , 35:37932-37946, 2022.

Francis Bach. Breaking the curse of dimensionality with convex neural networks. Journal of Machine Learning Research , 18(19):1-53, 2017.

Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations , 2020.

Boaz Barak, Benjamin Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems , 35:21750-21764, 2022.

Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055 , 2023.

Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single-index models with shallow neural networks. Advances in Neural Information Processing Systems , 35:9768-9783, 2022.

Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. On learning gaussian multi-index models with gradient flow. arXiv preprint arXiv:2310.19793 , 2023.

Emmanuel J Candès, Justin Romberg, and Terence Tao. Robust uncertainty principles: Exact signal reconstruction from highly incomplete frequency information. IEEE Transactions on information theory , 52(2):489-509, 2006.

Lenaic Chizat. Sparse optimization on measures with over-parameterized gradient descent. Mathematical Programming , 194(1):487-532, 2022.

Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in neural information processing systems , pages 3036-3046, 2018.

Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems , pages 2933-2943, 2019.

Hugo Cui, Luca Pesce, Yatin Dandi, Florent Krzakala, Yue Lu, Lenka Zdeborova, and Bruno Loureiro. Asymptotics of feature learning in two-layer networks after one gradient-step. In Forty-first International Conference on Machine Learning , 2024.

Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference on Learning Theory , pages 5413-5452. PMLR, 2022.

Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023.

Amit Daniely and Eran Malach. Learning parities with neural networks. Advances in Neural Information Processing Systems , 33:20356-20365, 2020.

David L Donoho. Compressed sensing. IEEE Transactions on information theory , 52(4):1289-1306, 2006.

Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2019.

Spencer Frei, Yuan Cao, and Quanquan Gu. Agnostic learning of a single neuron with gradient descent. Advances in Neural Information Processing Systems , 33:5417-5428, 2020.

Rong Ge, Jason D Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design. In International Conference on Learning Representations , 2018.

Rong Ge, Yunwei Ren, Xiang Wang, and Mo Zhou. Understanding deflation process in overparametrized tensor decomposition. Advances in Neural Information Processing Systems , 34: 1299-1311, 2021.

Margalit Glasgow. SGD finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the XOR problem. In The Twelfth International Conference on Learning Representations , 2024.

Surbhi Goel, Aravind Gollakota, Zhihan Jin, Sushrut Karmalkar, and Adam Klivans. Superpolynomial lower bounds for learning one-layer neural networks using gradient descent. In International Conference on Machine Learning , pages 3587-3596. PMLR, 2020.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems , pages 8571-8580, 2018.

Yuanzhi Li, Tengyu Ma, and Hongyang R Zhang. Learning over-parametrized two-layer neural networks beyond ntk. In Conference on learning theory , pages 2613-2682. PMLR, 2020.

Stanislaw Lojasiewicz. Une propriété topologique des sous-ensembles analytiques réels. Les équations aux dérivées partielles , 117:87-89, 1963.

Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis of neural networks with polynomial width, samples, and time. Advances in Neural Information Processing Systems , 36, 2024.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of twolayer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 2018.

Behrad Moniri, Donghwan Lee, Hamed Hassani, and Edgar Dobriban. A theory of non-linear feature learning with one gradient step in two-layer neural networks. In Forty-first International Conference on Machine Learning , 2024.

Alireza Mousavi-Hosseini, Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Murat A Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In The Eleventh International Conference on Learning Representations , 2023.

Eshaan Nichani, Alex Damian, and Jason D Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. Advances in Neural Information Processing Systems , 36, 2024a.

Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint arXiv:2402.14735 , 2024b.

Ryan O'Donnell. Analysis of boolean functions. arXiv preprint arXiv:2105.10386 , 2021.

Clarice Poon, Nicolas Keriven, and Gabriel Peyré. The geometry of off-the-grid compressed sensing. Foundations of Computational Mathematics , 23(1):241-327, 2023.

Itay M Safran, Gilad Yehudai, and Ohad Shamir. The effects of mild over-parameterization on the optimization landscape of shallow relu neural networks. In Conference on Learning Theory , pages 3889-3934. PMLR, 2021.

Zhenmei Shi, Junyi Wei, and Yingyu Liang. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations , 2022.

Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural information processing systems , 30, 2017.

Taiji Suzuki, Denny Wu, Kazusato Oko, and Atsushi Nitanda. Feature learning via mean-field langevin dynamics: classifying sparse parities and beyond. Advances in Neural Information Processing Systems , 36, 2024.

Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society Series B: Statistical Methodology , 58(1):267-288, 1996.

Zihao Wang, Eshaan Nichani, and Jason D. Lee. Learning hierarchical polynomials with three-layer neural networks. In The Twelfth International Conference on Learning Representations , 2024.

Lei Wu. Learning a single neuron for non-monotonic activation functions. In International Conference on Artificial Intelligence and Statistics , pages 4178-4197. PMLR, 2022.

Weihang Xu and Simon Du. Over-parameterization exponentially slows down gradient descent for learning a single neuron. In The Thirty Sixth Annual Conference on Learning Theory , pages 1155-1198. PMLR, 2023.

Gilad Yehudai and Shamir Ohad. Learning a single neuron with gradient methods. In Conference on Learning Theory , pages 3756-3786. PMLR, 2020.

Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems , 32, 2019.

Mo Zhou, Rong Ge, and Chi Jin. A local convergence theory for mildly over-parameterized two-layer neural network. In Conference on Learning Theory , pages 4577-4632. PMLR, 2021.

A Some properties of Hermite polynomials

In this section, we give several properties of Hermite polynomials that are useful in our analysis. See O'Donnell (2021) for a more complete discussion on Hermite polynomials. Let H k be the probabilists' Hermite polynomial where

and h k = 1 √ k ! H k be the normalized Hermite polynomials.

Hermite polynomials are classical orthogonal polynomials, which means E x ∼ N (0 , 1) [ h m ( x ) h n ( x )] = 1 if m = n and otherwise 0. Given a function σ , we call σ ( x ) = ∑ ∞ k =0 ˆ σ k h k ( x ) as the Hermit expansion of σ and ˆ σ k = E x ∼ N (0 , 1) [ σ ( x ) h k ( x )] as the k -th Hermite coefficient of σ .

The following is a useful property of Hermite polynomial.

Claim A.1 ((O'Donnell, 2021), Section 11.2) . Let ( x, y ) be ρ -correlated standard normal variables (that is, both x, y have marginal distribution N (0 , 1) and E [ xy ] = ρ ). Then, E [ h m ( x ) h n ( y )] = ρ n δ mn , where δ mn = 1 if m = n and otherwise 0.

The following lemma gives the Hermite coefficients for absolute value function and ReLU.

Lemma A.1. Let ˆ σ k = E x ∼ N (0 , 1) [ σ ( x ) h k ( x )] be the Hermite coefficient of σ . For σ is ReLU or absolute function, we have | ˆ σ k | = Θ( k -5 / 4 ) .

Proof. From Goel et al. (2020); Zhou et al. (2021) we have

Using Stirling's formula, we get | ˆ σ abs,k | , | ˆ σ relu,k | = Θ( k -5 / 4 ) .

B Useful facts and proof of Theorem 2

In this section we provide several useful facts and present the proof of Theorem 2.

The following claim shows that the square loss can be decomposed into 3 terms, where α, β are corresponding to 0th and 1st order of Hermite expansion. The effective activation is in fact σ ≥ 2 as defined below.

Claim B.1. Denote ˆ α = -(1 / √ 2 π ) ∑ m i =1 a i ∥ w i ∥ 2 , ˆ β = -(1 / 2) ∑ m i =1 a i w i .We have square loss

where f ≥ 2 ( x ; θ ) = ∑ i ∈ [ m ] a i σ ≥ 2 ( w ⊤ i x ) and σ ≥ 2 ( x ) = σ ( x ) -1 / √ 2 π -x/ 2 is the activation that after removing 0th and 1st order term in Hermite expansion.

As a result, when α, β are perfectly fitted and norms are balanced we have

Proof. Following Ge et al. (2018), we can write the loss L ( θ ) as a sum of tensor decomposition problem using Hermite expansion as in Section A (recall ∥ w ∗ i ∥ 2 = 1 and preprocessing procedure removes the 0-th and 1-st order term in the Hermite expansion of σ ):

Note that ˆ σ 0 = 1 / √ 2 π , ˆ σ 1 = 1 / 2 as in Lemma A.1, we get the result.

The proof of main result Theorem 2 is simply a combination of few lemmas appear in other sections. We refer the detailed proof and discussion to their corresponding sections.

Theorem 2 (Main result) . Under Assumption 1, 2, 3, consider Algorithm 1 on loss (2) . There exists a schedule of weight decay λ t and step size η t such that given m ≥ m 0 = ˜ O ∗ (1) · (1 /ε 0 ) O ( r ) neurons with small enough ε 0 = Θ ∗ (1) , with high probability we will recover the target network L ( θ ) ≤ ε within time T = O ∗ (1 /ηε 2 ) where η = poly( ε, 1 /d, 1 /m ) .

Moreover, when ε → 0 every student neuron w i either aligns with one of teacher neuron w ∗ j as ∠ ( w i , w ∗ j ) = 0 or vanishes as | a i | = ∥ w i ∥ = 0 .

Proof. Combine Lemma 3 (Stage 1), Lemma 4 (Stage 2) and Lemma 5 (Stage 3) together and follow the choice of λ t and η t we get the result.

For the student neurons' alignment, it is a direct corollary from Lemma F.6 and Lemma F.5.

C Stage 1: first gradient step

In this section, we show that after the first gradient update the first layer weights w 1 , . . . , w m form a ε 0 -net of the target subspace S ∗ , given m = (1 /ε 0 ) O ( r ) neurons. The proof is deferred to Section C.1. Lemma 3 (Stage 1) . Under Assumption 1,2,3, consider Algorithm 1 with λ 0 = η 0 = 1 and m ≥ m 0 = ˜ O ∗ (1) · (1 /ε 0 ) O ( r ) with any ε 0 = Θ ∗ (1) . After first step, with probability 1 -δ we have

(i) for every teacher neuron w ∗ i , there exists at least one student neuron w j s.t. ∠ ( w ∗ i , w j ) ≤ ε 0 .

The proof relies on the following lemma from Damian et al. (2022) that shows after the first step update w i 's are located at positions as if they are sampled within the target subspace S ∗ .

Lemma C.1 (Lemma 4, Damian et al. (2022)) . Under Assumption 3, we have with high probability in the ℓ 2 norm sense

where ˆ σ k := E x [ σ ( x ) h k ( x )] is the k -th Hermite polynomial coefficient.

C.1 Proofs in Section C

We now are ready to give the proof of Lemma 3.

Lemma 3 (Stage 1) . Under Assumption 1,2,3, consider Algorithm 1 with λ 0 = η 0 = 1 and m ≥ m 0 = ˜ O ∗ (1) · (1 /ε 0 ) O ( r ) with any ε 0 = Θ ∗ (1) . After first step, with probability 1 -δ we have

(i) for every teacher neuron w ∗ i , there exists at least one student neuron w j s.t. ∠ ( w ∗ i , w j ) ≤ ε 0 .

Proof. We show them one by one.

Part (i) From Lemma C.1 and the fact that w (0) i samples uniformly from unit sphere, we know the probability of ∠ ( w (1) i , w ) for any given w is at least Ω ∗ ( ε r 0 ) . Applying union bound we get the desired result.

Part (ii) We have

For the norm bound, using Lemma C.1 we know

Since w (0) i initializes from Gaussian distribution, we know the desired bound hold. Similarly, one can bound | a (1) i | .

Since we use a symmetric initialization and have preprocessed the data, it is easy to see α, β remains at 0.

D Stage 2: reaching low loss

In Stage 2, we show that given the features learned in Stage 1 one can adjust the norms on top of it to reach low loss that enters the local convergence regime in Stage 3.

Procedure We first specify the procedure to solve min a min α, β L ( θ ) + λ ∑ i ∥ w i ∥ 2 | a i | . For a at current point, we first solve the inner optimization problem, which is a linear regression on α, β . From Claim B.1 we know the global minima is (ˆ α, ˆ β ) . For simplicity of the proof, we just directly set ( α, β ) = (ˆ α, ˆ β ) . Then given the α, β , the outer optimization is a convex optimization for a , which can also be solved efficiently. Specifically, we perform 1 step of (sub)gradient on the loss function. We repeat the above 2 steps until convergence.

From Claim B.1 we know the actual objective that we optimize is

The following lemma shows that after Stage 2 we reach a low loss solution given the first layer features learned after first gradient step. The proof requires η to be small enough that depends on 1 /m , mostly due to the large gradient norm. We believe using more advance algorithm for this type of problem can alleviate this issue. However, as this is not the focus of this paper, we omit it for simplicity.

Lemma 4 (Stage 2) . Under Assumption 1,2,3, consider Algorithm 1 with λ t = √ ε 0 . Given Stage 1 in Lemma 3, we have Stage 2 ends within time T 2 = ˜ O ∗ (1 /ηε 0 ) such that optimality gap ζ T 2 = O ∗ ( ε 0 ) .

Proof. Denote ˜ a ∗ as the minima of ˜ L 1 ,λ . Then, we have

where (a) we use idea loss ˜ L 1 ,λ is convex in a .

Iterating the above inequality over all t we have

which means

It is easy to see ∥ ∥ a (1) ∥ ∥ 2 , ∥ ˜ a ∗ ∥ 1 = O ∗ (1) . Thus, when T ≥ O ∗ (1 /ηε 0 ) we know ˜ L 1 ,λ ( a ( T 2 ) ) -˜ L 1 ,λ ( ˜ a ∗ ) ≤ 3 ε 0 / 4 .

This suggests the optimality gap after balancing the norm (so that L λ ( θ ( T 2 ) ) = ˜ L 1 ,λ ( a ( T 2 ) ) )

For ˜ L 1 ,λ ( a ( T 2 ) ) -˜ L 1 ,λ ( ˜ a ∗ ) , we just show above that it is less than 3 ε 0 / 4 .

For ˜ L 1 ,λ ( ˜ a ∗ ) -min µ ∈M ( S d -1 ) L λ ( µ ) , we have

where in the last inequality we use Lemma F.3 and µ ∗ λ = arg min µ ∈M ( S d -1 ) L λ ( µ ) . Here ˆ a ∗ is a rescaled version of a ∗ and is constructed as: for every teacher neuron w ∗ i choose the closest neuron w j s.t. ∠ ( w j , w ∗ i ) ≤ ε 0 and set ˆ a ∗ ,j = a ∗ i / ∥ w j ∥ 2 . Set all other ˆ a ∗ ,k = 0 .

Together with above calculations, we have ζ T 2 ≤ O ∗ ( ε 0 ) .

E Stage 3: local convergence for regularized 2-layer neural networks

In this section we show the local convergence that loss eventually goes to 0 within polynomial time and recovers teacher neurons' direction.

The results in this section only need the width m ≥ m ∗ as long as its initial loss is small.

Lemma 5 (Stage 3) . Under Assumption 1,2,3, consider Algorithm 1 on loss (2) . Given Stage 2 in Lemma 4, if the initial optimality gap ζ 3 , 0 ≤ O ∗ ( λ 9 / 5 3 , 0 ) , weight decay λ follows the schedule of initial value λ 3 , 0 = O ∗ (1) , and k -th epoch λ 3 ,k = λ 3 ,k -1 / 2 and stepsize η 3 k = η ≤ O ∗ ( λ 12 3 ,k d -3 ) for all T 3 ,k ≤ t ≤ T 3 ,k +1 in epoch k , then within K = O ∗ (log(1 /ε )) epochs and total T 3 -T 2 = O ∗ ( λ -4 3 , 0 η -1 ε -2 ) time we recover the ground-truth network L ( θ ) ≤ ε .

The goal of each epoch is to minimize the loss L λ with a fix λ . The lemma below shows that as long as the initial optimality gap is O ∗ ( λ 9 / 5 ) , then at the end of each epoch, L λ could decrease to O ∗ ( λ 2 ) . Therefore, using a slow decay of weight decay parameter λ for each epoch we could stay in the local convergence regime for each epoch and eventually recovers the target network.

Lemma E.1 (Loss improve within one epoch) . Suppose | a (0) i | ≤ ∥ ∥ ∥ w (0) i ∥ ∥ ∥ 2 for all i ∈ [ m ] . If ζ 0 ≤ O ∗ ( λ 9 / 5 ) and λ ≤ O ∗ (1) and η ≤ O ∗ ( λ 12 d -3 ) , then within O ∗ ( λ -4 η -1 ) time the optimality gap becomes L λ -L λ ( µ ∗ λ ) = O ∗ ( λ 2 ) .

The above result relies on the following characterization of local landscape of regularized loss. We show the gradient is large whenever the optimality gap is large. This is the main contribution of this paper, see Section F for detailed proofs.

Lemma 6 (Gradient lower bound) . When Ω ∗ ( λ 2 ) ≤ ζ ≤ O ∗ ( λ 9 / 5 ) and λ ≤ O ∗ (1) , we have

In order to use the above landscape result with standard descent lemma, we also need certain smoothness condition on the loss function. We show below that this regularized loss indeed satisfies certain smoothness condition (though weaker than standard smoothness condition) to allow the convergence analysis.

E.1 Proofs in Section E

We now are ready to show the convergence of Stage 3 by using Lemma E.1 to show the loss makes progress every epoch.

Lemma 5 (Stage 3) . Under Assumption 1,2,3, consider Algorithm 1 on loss (2) . Given Stage 2 in Lemma 4, if the initial optimality gap ζ 3 , 0 ≤ O ∗ ( λ 9 / 5 3 , 0 ) , weight decay λ follows the schedule of initial value λ 3 , 0 = O ∗ (1) , and k -th epoch λ 3 ,k = λ 3 ,k -1 / 2 and stepsize η 3 k = η ≤ O ∗ ( λ 12 3 ,k d -3 ) for all T 3 ,k ≤ t ≤ T 3 ,k +1 in epoch k , then within K = O ∗ (log(1 /ε )) epochs and total T 3 -T 2 = O ∗ ( λ -4 3 , 0 η -1 ε -2 ) time we recover the ground-truth network L ( θ ) ≤ ε .

Proof. Since | a (0) i | ≤ ∥ ∥ ∥ w (0) i ∥ ∥ ∥ 2 for all i ∈ [ m ] at the beginning of Stage 3, from Lemma E.3 we know they will remain hold for all epoch and all time t .

From Lemma E.1 we know for epoch k it finishes within O ∗ ( λ -4 k η -1 ) time and achieves L λ k -L λ k ( µ ∗ λ k ) = O ∗ ( λ 2 k ) . To proceed to next epoch k +1 , we only need to show the solution at the end of epoch k θ ( k ) gives the optimality gap ζ = O ∗ ( λ 9 / 5 k +1 ) for the next λ k +1 . We have

where (a) due to Lemma F.4 that L ( θ ( k ) ) is small; (b) the optimality gap at the end of epoch k is O ∗ ( λ 2 k ) and | µ ∗ λ k | 1 -| µ ∗ λ k +1 | 1 = O ∗ ( λ k ) from Lemma F.3; (c) due to Lemma F.3 that L ( µ ∗ λ k ) is small. In this way, we can apply Lemma E.1 again for epoch k +1 .

From Lemma F.4 we know at the end of epoch k the square loss L ( θ ( k ) ) = O ∗ ( λ 2 k ) . Thus, to reach ε square loss, we need λ k = O ∗ ( ε 1 / 2 ) , which means we need to take O ∗ (log(1 /ε )) epoch. Since epoch k it finishes within O ∗ ( λ -4 k η -1 ) time, we know the total time is at most O ∗ ( λ -4 0 η -1 ε -2 ) time.

To show the lemma below that loss makes progress within every epoch, we rely on the gradient lower bound (Lemma 6) and smoothness condition of loss function (Lemma E.2).

Lemma E.1 (Loss improve within one epoch) . Suppose | a (0) i | ≤ ∥ ∥ ∥ w (0) i ∥ ∥ ∥ 2 for all i ∈ [ m ] . If ζ 0 ≤ O ∗ ( λ 9 / 5 ) and λ ≤ O ∗ (1) and η ≤ O ∗ ( λ 12 d -3 ) , then within O ∗ ( λ -4 η -1 ) time the optimality gap becomes L λ -L λ ( µ ∗ λ ) = O ∗ ( λ 2 ) .

Proof. Since | a (0) i | ≤ ∥ ∥ ∥ w (0) i ∥ ∥ ∥ 2 for all i ∈ [ m ] at the beginning of current epoch, from Lemma E.3 we know they will remain hold for all time t . Then combine Lemma E.4 and Lemma E.2 we know

Recall ζ t = L λ ( θ ( t ) ) -L λ ( µ ∗ λ ) . Using gradient lower bound Lemma 6 and consider the time before ζ t reach O ∗ ( λ 2 ) we have

where we use η = O ∗ ( λ 12 d -3 ) to be small enough.

The above recursion implies that

Thus, within O ∗ (1 /λ 4 ) the optimality gap ζ t reaches O ∗ ( λ 2 ) .

The lemma below shows a regularity condition on the norm between two layers.

Lemma E.3. If we start at | a (0) i | ≤ ∥ ∥ ∥ w (0) i ∥ ∥ ∥ 2 and η = O ∗ (1) , then we have | a ( t ) i | 2 ≤ ∥ ∥ ∥ w ( t ) i ∥ ∥ ∥ 2 2 for all i ∈ [ m ∗ ] and all time t .

Proof. Denote R ( x ) = f ( x ) -f ∗ ( x ) . Assume | a ( t ) i | 2 -∥ ∥ ∥ w ( t ) i ∥ ∥ ∥ 2 2 ≤ 0 we show it remains at t +1 . We have

We first focus on the last 2 terms. We have

where (a) due to | 2 E x [ R ( x ) σ ( w ( t ) ⊤ i x )] | 2 ≤ ∥ ∥ ∥ 2 E x [ R ( x ) σ ′ ( w ( t ) ⊤ i x ) x ] ∥ ∥ ∥ 2 2 .

Therefore, plug it back to the above equation, we have

where (a) due to | a ( t ) i | 2 -∥ ∥ ∥ w ( t ) i ∥ ∥ ∥ 2 2 ≤ 0 and we use ∥ ∥ ∥ 2 E x [ R ( x ) σ ′ ( w ( t ) ⊤ i x )] x ∥ ∥ ∥ 2 2 = O ∗ ( d ) from Lemma E.4 and η is small enough.

This lemma shows the smoothness of loss function. The proof requires a careful calculations to bound the error terms.

The last line is easy to see on O ∗ ( η 2 d 2 ) using norm bound in Lemma F.12, so in below we are going to bound (I) and (II) one by one. The goal is to show they are small in the sense of on order o ( η ) .

Bound (I) For (I), we can write out the expression as

For (I.i), we can split into 2 terms as

We then can bound them separately as

where (a) we use Lemma E.5; (b) recall | ˜ ∇ a i | = O ∗ ( η ∥ w i ∥ 2 d ) and ∥ ∥ ∥ ˜ ∇ w i ∥ ∥ ∥ 2 = O ∗ ( η | a i | d ) ; (c) ∥ a ∥ , ∥ W ∥ F , ∑ i ∈ [ m ] | a i | ∥ w i ∥ 2 = O ∗ (1) from Lemma F.12 and Lemma F.4, as well as η is small enough.

For (I.ii), we have

Combine (I.i) and (I.ii) we know (I)= O ∗ ( η 2 d 2 ) .

Bound (II) For (II), we have

We focus on bound ∥ I i ∥ below. The goal is to show it is o ( η ) . For I i ( x ) , we have

For (II.i), recall | ˜ ∇ a i | = O ∗ ( η ∥ w i ∥ 2 d ) and ∥ ∥ ∥ ˜ ∇ w i ∥ ∥ ∥ = O ∗ ( η | a i | d )

For (II.ii), we have

̸

where δ = ∠ ( w i + ˜ ∇ w i , w i ) is the angle between w i + ˜ ∇ w i and w i . Since ∥ ∥ ∥ ˜ ∇ w i ∥ ∥ ∥ = O ∗ ( η | a i | d ) =

O ∗ ( η ∥ w i ∥ 2 d ) , we know δ = O ( ∥ ∥ ∥ ˜ ∇ w i ∥ ∥ ∥ ) given η = O ∗ (1 /d )

Combine (II.i) and (II.ii) we have

Since ∥ R θ ∥ = O ∗ (1) , this implies

Combine (I)(II) Finally, combing (I) and (II) we have

Going back to the beginning of this proof, we get the desired result.

E.2 Technical Lemma

We present technical lemmas that are used in the proof of this section. They mostly follow from direct calculations.

Proof. It is easy to see given ∥ R ∥ = O ∗ (1) .

Lemma E.5 (Lemma D.4 in Zhou et al. (2021)) . Consider α i ∈ R d for i ∈ [ n ] . We have

2 to be small enough.

where c 0 is a constant.

F Local landscape of population loss

In this section, we are going to show Lemma 6 that characterizing the population local landscape with a fixed λ by giving the lower bound of gradient.

Outline We generally follow the high-level proof plan that outlines in Section 6. In Section F.1 and Section F.2, we characterize the local geometry as in Lemma 8. Then, we use it to construct descent direction in Section F.3. Finally we give the proof of Lemma 6 in Section F.4.

We start by identifying the structure of (approximated) solution of a closely-related problem in Section F.1 (rewrite (5)):

where M ( S d -1 ) is the measure space over unit sphere S d -1 , µ ∗ = ∑ i ∈ [ m ∗ ] a ∗ i δ w ∗ i and σ ≥ 2 ( x ) = σ ( x ) -1 / √ 2 π -x/ 2 is the activation that after removing 0th and 1st order term in Hermite expansion. Note that when µ represents a finite-wdith network, we have µ = ∑ i ∈ [ m ] a i ∥ w i ∥ 2 δ w i is a empirical measure over the neurons. In particular, when µ = µ ∗ , model f µ recovers the target ˜ f ∗ .

We call (5) as the ideal loss because the original problem (2) would become the above (5) when we balance the norms ( ∥ w i ∥ 2 = | a i | ), perfectly fit α, β and relax the finite-width constraints to allow infinite-width (see Claim B.1). This is why we slightly abused the notation to use L λ in both (2) and (5).

In Section F.3 we will use the solution structure to construct descent direction that are positively correlated with gradient and also handle the case when norms are not balanced or α, β are not fitted well.

Notation Denote the optimality gap between the loss at µ and the optimal distribution µ ∗ λ as

where µ ∗ λ is the optimal measure that minimize (5). For simplicity denote ˜ a i = a i ∥ w i ∥ 2 so that | µ | 1 = ∥ ˜ a ∥ 1 when µ = ∑ i ∈ [ m ] a i ∥ w i ∥ 2 δ w i . Often we use ζ t = ζ ( µ t ) to denote the optimality gap at time t and just ζ for simplicity. We slightly abuse the notation to also use ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . Finally denote µ ∗ = ∑ i ∈ [ m ∗ ] a ∗ i δ w ∗ i (assuming ∥ w ∗ i ∥ 2 = 1 ) so that f µ ∗ ( x ) = E w ∼ µ ∗ [ σ ≥ 2 ( w ⊤ x )] .

F.1 Structure of the ideal loss solution

In this section, we will focus on the structure of approximated solution for the ℓ 1 regularized regression problem (5).

In the rest of this section, we will first introduce the idea of non-degenerate dual certificate and then use it as a tool to characterize the structure of the solutions. The proofs are deferred to Section H.

F.1.1 Non-degenerate dual certificate

We first recall the definition of non-degenerate dual certificate, which is similar as in (Poon et al., 2023) but slightly adapted for fit our need.

Definition 1 (Non-degenerate dual certificate) . η ( w ) is called a non-degenerate dual certificate if there exists p ( x ) such that η ( w ) = E x [ p ( x ) σ ≥ 2 ( w ⊤ x )] for w ∈ S d -1 and

We first show that there exist such non-degenerate dual certificate. More discussion and a detailed proof are deferred to Section G.

LemmaF.1. There exists a non-degenerate dual certificate η = E x [ p ( x ) σ ≥ 2 ( w ⊤ x )] with ρ η = Θ(1) and ∥ p ∥ 2 ≤ poly( m ∗ , ∆)

The following lemma (restate of Lemma 9) gives the properties that will be used in the later proofs: the non-degenerate dual certificate η allows us to capture the gap between the current position µ and the target µ ∗ .

Lemma F.2. Given a non-degenerate dual certificate η , then

F.1.2 Properties of µ ∗ λ

Given the non-degenerate dual certificate η , we now are ready to identify several useful properties of µ ∗ λ . The lemma below essentially says that µ ∗ λ is similar to µ ∗ in the sense that most of the norm are concentrated in the ground-truth direction and the square loss is small. The proof relies on comparing µ ∗ λ with µ ∗ using the optimality conditions.

Lemma F.3. We have the following hold

F.1.3 Properties of µ with optimality gap ζ

We now characterize the structure of µ when the optimality gap is ζ . The proof mostly relies on comparing µ with µ ∗ λ and the structure of µ ∗ λ in previous section.

The following lemma shows the square loss is bounded by the optimality gap and norms are always bounded. Note that the conditions are true under Lemma 6.

Lemma F.4. Recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then, the following holds:

The following two lemma characterize the structure of µ using the fact that the square loss is small in previous lemma. The lemma below says that the total norm of far away neuron is small.

Lemma F.5. Recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then, we have

In particular, when µ = ∑ i ∈ [ m ] a i ∥ w i ∥ 2 δ w i represents finite number of neurons, we have

where δ j = ∠ ( w j , w ∗ i ) for j ∈ T i .

The lemma below shows there are neurons close to the teacher neurons once the gap is small. The proof idea is similar to Section 5.3 in Zhou et al. (2021) that use test function to lower bound the loss, but now we can handle almost all activation.

Lemma F.6. Under Lemma 6, if the Hermite coefficient of σ decays as | ˆ σ k | = Θ( k -c σ ) with some constant c σ > 0 , then the total mass near each target direction is large, i.e., µ ( T i ( δ )) sign( a ∗ i ) ≥ | a ∗ i | / 2 for all i ∈ [ m ∗ ] and any δ close ≥ ˜ Ω ( ( L ( µ ) a 2 min ) 1 / (4 c σ -2) ) with large enough hidden constant.

In particular, for σ is ReLU or absolute function, δ close ≥ ˜ Ω ( ( L ( µ ) a 2 min ) 1 / 3 ) . Here a min = min | a i | is the smallest entry of a ∗ in absolute value.

As a corollary, if the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) , then δ close ≥ ˜ Ω ∗ ( ( ζ + λ 2 ) 1 / (4 c σ -2) ) and for ReLU or absolute δ close ≥ ˜ Ω ∗ ( ( ζ + λ 2 ) 1 / 3 ) .

F.1.4 Residual decomposition and average neuron

In this section, we introduce the residual decomposition and average neuron as in (Zhou et al., 2021) that will be used when proving the existence of descent direction.

Denote the decomposition R ( x ) = f µ ( x ) -f µ ∗ ( x ) = R 1 ( x ) + R 2 ( x ) + R 3 ( x ) (this can be directly verified noticing that σ ≥ 2 ( x ) = | x | / 2 -1 / √ 2 π ),

In the following we characterize R 1 , R 2 , R 3 separately. In Lemma F.7 we relate R 1 with the average neuron. In Lemma F.8 and Lemma F.9 we bound R 2 and R 3 respectively.

Lemma F.7 (Zhou et al. (2021), Lemma 11) . ∥ R 1 ∥ 2 2 = Ω(∆ 3 /m 3 ∗ ) ∑ i ∈ [ m ∗ ] ∥ ∥ ∥ ∑ j ∈T i a j w j -w ∗ i ∥ ∥ ∥ 2 2 .

Lemma F.8. Under Lemma 6, recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then

Lemma F.9. Under Lemma 6 and recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . If ˆ σ 0 = 0 and ˆ σ k > 0 with some k = Θ((1 / ∆ 2 ) log( ζ/ ∥ a ∗ ∥ 1 )) , then

Nowweare ready to bound the difference between average neuron with its corresponding ground-truth neuron.

Lemma F.10. Under Lemma 6, recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then for any i ∈ [ m ∗ ] , ζ = Ω( λ 2 ) and ζ, λ ≤ 1 / poly( m ∗ , ∆ , ∥ a ∗ ∥ 1 )

F.2 From ideal loss solution to real loss solution

In previous section, we consider the ideal loss solution that assumes the norms are perfectly balanced ( | a i | = ∥ w i ∥ 2 ) and α, β are perfectly fitted. However, during the training we are not able to guarantee achieve these exactly but only approximately. This section is devoted to show that the results in previous section still hold though the conditions are only approximately satisfied. Recall that the original loss

so that when norm are balanced and α, β are perfectly fitted, L λ ( θ ) = L ( θ ) + λ ∑ i | a i | ∥ w i ∥ 2 = L λ ( µ ) .

The lemma below shows that the properties of ideal loss solution in previous section still hold for the solution of original loss, when α, β are approximately fitted.

version θ bal = ( a bal , W bal , α bal , β bal ) as a bal,i = sign( a i ) √ | a i | ∥ w i ∥ 2 , w bal,i = w i √ | a i | ∥ w i ∥ 2 , α bal = ˆ α and β bal = ˆ β . Then, we have

Moreover, let the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) , we have results in Lemma F.4, Lemma F.5, Lemma F.6, Lemma F.7, Lemma F.8, Lemma F.9 and Lemma F.10 still hold for L λ ( θ ) , with the change of R 3 in (8) as

The following lemma shows the norm remains bounded.

LemmaF.12. Under Lemma 6, suppose optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . Then ∥ a ∥ 2 2 + ∥ W ∥ 2 F ≤ 3 ∥ a ∗ ∥ 1 .

F.3 Descent direction

In this section, we show that there is a descent direction as long as the optimality gap is small until it reaches O ( λ 2 ) . We will assume ζ = Ω( λ 2 ) in this section for simplicity.

We first show gradient is always large whenever α, β are not fitted well. This is a direct corollary of Claim B.1.

Lemma F.13 (Descent direction, α and β ) . We have

Before proceeding to the following descent direction, we first make a simplification assumption that Assumption F.1. For every T i , for all neuron w j ∈ T i , assume w ⊤ j w ∗ i ≥ 0 .

This is because due to the linear term β , the effective activation is symmetry σ ≥ 2 ( x ) = σ ≥ 2 ( -x ) . This introduce the ambiguity of the sign of neurons. Such assumption clarifies the ambiguity of neurons' direction.

As the lemma below shows, there always exists a set of parameter (by flipping the sign of neurons) that satisfy the assumption and gives almost same gradient norm. Thus, making such assumption will not cause any issue when α, β are perfectly fitted.

Lemma F.14. Suppose ( α -ˆ α ) 2 , ∥ ∥ ∥ β -ˆ β ∥ ∥ ∥ 2 2 ≤ τ to be small enough and ∥ a ∥ 2 , ∥ W ∥ F = O ∗ (1) .

Then, given any parameter θ , there exists another set of parameter ˜ θ that satisfies Assumption F .1 such that f θ = f ˜ θ and | ∥∇ θ L λ ∥ -∥ ∥ ∇ ˜ θ L λ ∥ ∥ F | ≤ O ∗ ( √ τ ) .

Proof. Denote θ = ( a , w 1 , . . . , w m , α, β ) . We first construct ˜ θ = ( ˜ a , ˜ w 1 , . . . , ˜ w m , ˜ α, ˜ β ) .

Let ˜ a = a . For ˜ w i , there exists such sign vector s = ( s 1 , . . . , s m ) ∈ {± 1 } m so that by flipping the sign of neurons we have ˜ w i = s i w i satisfies Assumption F.1. Let ˜ α = α and ˜ β = β + + ∑ i : s i = -1 a i w i .

One can verify that f θ = f ˜ θ . Moreover, for the gradient of α, β we have

For gradient of a , w i , when s i = 1 we know they are the same. When s i = -1 , note that

Therefore, we get the desired result by noting the norm bound.

We then show that if norms are not balanced or norm cancellation happens for neurons with similar direction, then one can always adjust the norm to decrease the loss due to the regularization term.

Lemma F.15 (Descent direction, norm balance) . We have

Lemma F.16 (Descent direction, norm cancellation) . Under Lemma 6 and Assumption F.1, suppose the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . For any w ∗ i , consider δ sign such that δ close < δ sign = O ( λ/ζ 1 / 2 ) with small enough hidden constant ( δ close defined in Lemma F.6), then

̸

Now given the above lemmas, it suffices to consider the remaining case that α, β are well fitted, norms are balanced and no cancellation. In this case, the loss landscape is roughly the same as the ideal loss (5) from Lemma F.11. Thus, we could leverage these detailed characterization of the solution (far-away neurons are small and average neuron is close to corresponding ground-truth neuron) to construct descent direction.

Lemma F.17 (Descent direction) . Under Lemma 6 and Assumption F.1, suppose the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . Suppose

(i) norms are (almost) balanced: | ∥ W ∥ 2 F -∥ a ∥ 2 2 | ≤ ζ/λ , ∑ i ∈ [ m ] ( | a j | - ∥ w j ∥ 2 ) 2 = O ∗ ( ζ 2 /λ 2 )

̸

(ii) (almost) no norm cancellation: consider all neurons w j that are δ sign -close w.r.t. teacher neuron w ∗ i but has a different sign, i.e., sign( a j ) = sign( a ∗ i ) with δ sign = Θ ∗ ( λ/ζ 1 / 2 ) , we have ∑ j ∈ T i, -( δ sign ) | a j | ∥ w j ∥ 2 ≤ τ = O ∗ ( ζ 5 / 6 /λ ) with small enough hidden constant, where T i, -( δ ) defined in Lemma F.16.

(iii) α, β are well fitted: | α -ˆ α | 2 = O ∗ ( ζ ) , ∥ ∥ ∥ β -ˆ β ∥ ∥ ∥ 2 2 = O ∗ ( ζ ) with small enough hidden factor.

Then, we can construct the following descent direction

̸

where q ij satisfy the following conditions with δ close < δ sign and δ close = O ∗ ( ζ 1 / 3 ) : (1) ∑ j ∈T i a j q ij = a ∗ i ; (2) q ij ≥ 0 ; (3) q ij = 0 when sign( a j ) = sign( a ∗ i ) or δ j > δ close . (4) ∑ i ∈ [ m ∗ ] ∑ j ∈T i q 2 ij = O ∗ (1) .

F.4 Proof of Lemma 6

Now we are ready to prove the gradient lower bound (Lemma 6) by combining all descent direction lemma in the previous section together.

where T i, + ( δ sign ) = { j ∈ T i : δ ( w j , w ∗ i ) ≤ δ sign , sign( a j ) = sign( a ∗ i ) } , T i, -( δ sign ) = { j ∈ T i : δ ( w j , w ∗ i ) ≤ δ sign , sign( a j ) = sign( a ∗ i ) } are the set of neurons that close to w ∗ i with/without same sign of a ∗ i .

As a result,

Lemma 6 (Gradient lower bound) . When Ω ∗ ( λ 2 ) ≤ ζ ≤ O ∗ ( λ 9 / 5 ) and λ ≤ O ∗ (1) , we have ∥∇ θ L λ ∥ 2 ≥ Ω ∗ ( ζ 4 /λ 2 ) .

Proof. We check the assumption of Lemma F.17 one by one. We first assume Assumption F.1 holds to get a gradient lower bound.

For assumption (i) (norm balance) in Lemma F.17, whenever ∑ i ∈ [ m ∗ ] ∣ ∣ ∣ a 2 i -∥ w i ∥ 2 2 ∣ ∣ ∣ = Ω ∗ ( ζ 2 /λ 2 ) , by Lemma F.15 we know

With Lemma F.12, this implies √ ∥∇ a L λ ∥ 2 2 + ∥∇ W L λ ∥ 2 F · O ( ∥ a ∗ ∥ 1 ) ≥ √ ∥∇ a L λ ∥ 2 2 + ∥∇ W L λ ∥ 2 F √ ∥ a ∥ 2 2 + ∥ W ∥ 2 F = Ω ∗ ( ζ 2 /λ ) , which means

For assumption (ii) (norm cancellation) in Lemma F.17, whenever it does not hold, by Lemma F.16 we know

For assumption (iii) ( α, β ) in Lemma F.17, whenever it does not hold, by Lemma F.13 we know

which implies

Thus, the remaining case is the one that all assumption (i)-(iii) in Lemma F.17 hold and also ∑ i ∈ [ m ∗ ] ∣ ∣ ∣ a 2 i -∥ w i ∥ 2 2 ∣ ∣ ∣ = O ∗ ( ζ 2 /λ 2 ) , we choose

so that condition (1)-(4) on q ij all hold: condition (1)-(3) are easy to check, Lemma H.4 shows condition (4) holds. Now we know from Lemma F.17 that

Note that

and

where (a)(b) by Lemma F.4; (c) we use Lemma F.12 and condition (4) on q ij .

Therefore, we get

Combine all cases above, we know

We now use Lemma F.14 to show when Assumption F.1 is not true, we can get similar gradient lower bound. Denote the above gradient lower bound as τ 0 = Ω ∗ ( ζ 4 /λ 2 ) . Let τ = τ 0 / 2 .

When ( α -ˆ α ) 2 , ∥ ∥ ∥ β -ˆ β ∥ ∥ ∥ 2 2 ≤ τ , using Lemma F.14 we know there exists ˜ θ such that ∥ ∥ ∇ ˜ θ L λ ∥ ∥ 2 F ≥ τ 0 and | ∥ ∥ ∇ ˜ θ L λ ∥ ∥ F -∥∇ θ L λ ∥ F | ≤ √ τ . Thus, we know ∥∇ θ L λ ∥ 2 F ≥ 0 . 1 τ .

Therefore, combine above we can show ∥∇ θ L λ ∥ 2 F = Ω ∗ ( ζ 4 /λ 2 ) .

G Non-degenerate dual certificate

In this section, we show that there indeed exists a non-degenerate dual certificate that satisfies Definition 1 and therefore proving Lemma F.1.

LemmaF.1. There exists a non-degenerate dual certificate η = E x [ p ( x ) σ ≥ 2 ( w ⊤ x )] with ρ η = Θ(1) and ∥ p ∥ 2 ≤ poly( m ∗ , ∆)

Recall that we want to use the dual certificate η to characterize the (approximate) solution for the following regression problem:

where σ ≥ 2 is the ReLU activation after removing 0th and 1st order (corresponding to α and β terms) and µ ∗ = ∑ i ∈ [ m ∗ ] a ∗ i δ w ∗ i is the ground-truth.

Notation We need to first introduce few notations before proceeding to the proof. Denote the kernel K ≥ ℓ ( w , u ) = E x ∼ N (0 , I ) [ σ ≥ ℓ ( w ⊤ x ) σ ≥ ℓ ( u ⊤ x )] as the kernel induced by activation σ ≥ ℓ ( x ) , where σ ≥ ℓ ( x ) = ∑ k ≥ ℓ ˆ σ k h k ( x ) /Z σ , Z σ = ∥ σ ≥ ℓ ∥ 2 = √ ∑ k ≥ ℓ ˆ σ 2 k = Θ( ℓ -3 / 4 ) is the normalizing factor, h k ( x ) is the normalized k -th (probabilistic) Hermite polynomial and ˆ σ k is the corresponding Hermite coefficient. We will specify the value of ℓ later and use K instead of K ≥ ℓ for simplicity.

We will construct the dual certificate η following the proof strategy in Poon et al. (2023) with the form below (the difference is that we now only keep high order terms that are at least ℓ ):

such that it satisfies

Here α 1 = ( α 1 , . . . , α m ∗ ) ⊤ ∈ R m ∗ , α 2 = ( α ⊤ 2 , 1 , . . . , α ⊤ 2 ,m ∗ ) ⊤ ∈ R m ∗ d are the parameters that we are going to solve and ∇ i means the gradient w.r.t. i -th variable (for example, ∇ 1 K ( x , y ) means gradient with respect to x ).

One can rewrite the above constraints (9) into the matrix form:

where b = (sign( a ∗ 1 ) , . . . , sign( a ∗ m ∗ ) , 0 ⊤ m ∗ d ) ⊤ ∈ R m ∗ ( d +1) , Υ = E x [ γ ( x ) γ ( x ) ⊤ ] ∈ R m ∗ ( d +1) × m ∗ ( d +1) ,

Here ∇ w σ ≥ ℓ ( w ∗⊤ i x ) = P w ∗ i σ ≥ ℓ ′ ( w ∗⊤ i x ) x ∈ R d , where P w ∗ i is the projection matrix defined below.

Notions on the unit sphere As we could see, the kernel K is invariant under the change of norms, so it suffices to focus on the input on the unit sphere S d -1 . On the unite sphere, we could compute the gradient and hessian of a function f ( w ) on the sphere (e.g., Absil et al. (2013))

where P w = I -ww ⊤ is the projection matrix.

Then, we could define the derivative as in Poon et al. (2023); Absil et al. (2008): for tangent vectors z , z ′

and their associated norms

For simplicity, we will use K ( ij ) ( w , u ) to denote ∇ i 1 ∇ j 2 K ( w , u ) . One can check that this is in fact the same as the one defined Poon et al. (2023) under our specific kernel K , i + j ≤ 3 and i, j ≤ 2 . Let

where z ( p ) w applies to the dimension corresponding to w and similarly z ( q ) u for u .

Before solving (10), we first present some useful proprieties of kernel K that will be used later (see Section I for the proofs). The lemma below shows that kernel K ( w , u ) is non-degenerate in the sense that it decays at least quadratic at each ground-truth direction ( w ≈ u ≈ w ∗ i ) and contributes almost nothing when w , u are away.

Lemma G.1 (Non-degeneracy of kernel K ) . For any h > 0 , let ℓ ≥ Θ(∆ -2 log( m ∗ ℓ/h ∆)) , kernel K ≥ ℓ is non-degenerate in the sense that there exists r = Θ( ℓ -1 / 2 ) , ρ 1 = Θ(1) , ρ 2 = Θ( ℓ ) such that following hold:

The following lemma shows that K and its derivatives are bounded.

Lemma G.2 (Regularity conditions on kernel K ) . Let B ij := sup w , u ∥ ∥ K ( ij ) ( w , u ) ∥ ∥ w , u and B 0 = B 00 + B 10 +1 , B 2 = B 20 + B 21 +1 . We have B 00 = O (1) , B 10 = O ( ℓ 1 / 2 ) , B 11 = O ( ℓ ) , B 20 = O ( ℓ ) , B 21 = O ( ℓ 3 / 2 ) , and therefore B 0 = O ( ℓ 1 / 2 ) , B 2 = O ( ℓ 3 / 2 ) .

The following lemma from Poon et al. (2023) connects the non-degeneracy of kernel K to the dual certificate η that we are interested in.

Lemma G.3 (Lemma 2, Poon et al. (2023), adapted in our setting) . Let a ∈ {± 1 } . Suppose that for some ρ > 0 , B > 0 and 0 < r ≤ B -1 / 2 we have: for all δ ( w , w 0 ) and z ∈ R d with z ⊤ w = 0 , it holds that -K (02) ( w 0 , w )[ z , z ] > ρ ∥ z ∥ 2 2 and ∥ ∥ K (02) ( w 0 , w ) ∥ ∥ w ≤ B . Let η be a smooth function. If η ( w 0 ) = a , ∇ η ( w 0 ) = 0 and ∥ ∥ a D 2 η ( w ) -K (02) ( w 0 , w ) ∥ ∥ w ≤ τ for all δ ( w , w 0 ) ≤ r with τ < ρ/ 2 , then we have | η ( w ) | ≤ 1 -(( ρ -2 τ ) / 2) δ ( w , w 0 ) 2 for all δ ( w , w 0 ) ≤ r .

We now are ready to proof the main result in this section Lemma F.1 that shows the non-degenerate dual certificate exists. Roughly speaking, following the same proof as in Poon et al. (2023), we can show that α ≈ sign( a ∗ ) and α 2 ≈ 0 and therefore we can transfer the non-degeneracy of kernel K to the dual certificate η with Lemma G.3.

LemmaF.1. There exists a non-degenerate dual certificate η = E x [ p ( x ) σ ≥ 2 ( w ⊤ x )] with ρ η = Θ(1) and ∥ p ∥ 2 ≤ poly( m ∗ , ∆)

Proof. Note that Υ = SD ˜ Υ DS , where

are block diagonal matrices, ˜ Υ = E x [ ˜ γ ( x ) ˜ γ ( x ) ⊤ ] ∈ R m ∗ ( d +1) × m ∗ ( d +1) ,

Z σ ′ = √ ∑ k ≥ ℓ ˆ σ 2 k k = Θ( ℓ -1 / 4 ) is the normalizing factor so that the diagonal of ˜ Υ are all 1.

Thus, to solve (10), it is sufficient to solve the following: denote ˜ K = D ˜ Υ D

and let α 1 = ˜ α 1 , α 2 ,i = ( Z σ /Z σ ′ ) ˜ α 2 ,i to get the solution of (10).

In the following, we are going to first show that ˜ K ≈ DD because all the off-diagonal terms of ˜ Υ are small due to Lemma G.1 (iii) (we can choose h to be small enough, and we will choose it later). Specifically, we have

where z = ( z ⊤ 1 , z ⊤ 2 ) ⊤ , z 1 = ( z 1 , 1 , . . . , z 1 ,m ∗ ) ⊤ and z 2 = ( z ⊤ 2 , 1 , . . . , z ⊤ 2 ,m ∗ ) ⊤ has the same block structure as ( α 1 , α 2 ) and we use Lemma G.1 and Lemma G.2 in the last line.

Note that DD has exactly m ∗ d eigenvalues of 1 and m ∗ eigenvalues of 0, and ˜ K also has m ∗ eigenvalues of 0. By Weyl's inequality, we know | γ i -1 | ≤ 2 h where ˜ K = ∑ i ∈ [ m ∗ d ] γ i v i v ⊤ i is its eigendecomposition. Here v ⊤ i v ⊥ = 0 for all v ⊥ ∈ V ⊥ = span { ( 0 , w ∗ 1 , 0 , . . . , 0 ) ⊤ , . . . ( 0 , . . . , 0 , w ∗ m ∗ ) ⊤ }

in the null space of D . Since b ⊤ v ⊥ = 0 for all v ⊥ ∈ V ⊥ , we have

Therefore,

This implies ∥ α 1 -sign( a ∗ ) ∥ ∞ = ∥ ˜ α 1 -sign( a ∗ ) ∥ ∞ ≤ h ′ , ∥ α 1 ∥ ∞ = ∥ ˜ α 1 ∥ ∞ ≤ 1 + h ′ and ∥ α 2 ∥ 2 = ( Z σ /Z σ ′ ) ∥ ˜ α 2 ,i ∥ 2 ≤ Θ( h ′ ℓ -1 / 2 ) .

Now, given the α 1 , α 2 , we can show the corresponding η is non-degenerate. Choosing h = O ( m -1 / 2 ∗ ) and ℓ = Θ(∆ -2 log( m ∗ / ∆)) so that the condition in Lemma G.1 holds.

Consider w ∈ T i , when δ ( w , w ∗ i ) ≥ r = Θ( ℓ -1 / 2 ) , using Lemma G.1 and Lemma G.2 we have

where we choose h = O ( m -1 / 2 ∗ ) to be small enough.

When δ ( w , w ∗ i ) ≤ r = Θ( ℓ -1 / 2 ) , again using Lemma G.1 and Lemma G.2 we have

̸

where again due to our choice of small h . Using Lemma G.3 we know that | η ( w ) | ≤ 1 -( ρ 2 / 4) δ ( w , w ∗ i ) 2 .

Combine the above two cases, we have | η ( w ) | ≤ 1 -Θ(1) δ ( w , w ∗ i ) 2 and η ( w ) = E x [ p ( x ) σ ( w ⊤ x )] with

H Proofs in Section F

In this section, we give the omitted proofs in Section F.

H.1 Omitted proofs in Section F.1

We give the proofs for these results that characterize the structure of ideal loss solution. The following proof follows from the definition of non-degenerate dual certificate η .

Lemma F.2. Given a non-degenerate dual certificate η , then

(i) ⟨ η, µ ∗ ⟩ = | µ ∗ | 1

Proof. We show the results one by one.

Part (i)(ii) We have

where the last inequality follows the property of non-degenerate dual certificate (Definition 1). The other part then follows directly by the definition of µ ∗ .

Part (iii) We have

Given the above lemma and the optimality of µ ∗ λ , we are able to characterize the structure of µ ∗ λ as below: norm is bounded, square loss is small and far-away neurons are small.

Lemma F.3. We have the following hold

Proof. We show the results one by one.

Part (i) Due to the optimality of µ ∗ λ , we have

Rearranging the terms, we have

For the lower bound, with Lemma F.2 we have

Using part (ii) we get the desired lower bound.

Part (ii) We first have the following inequality due to the optimality of µ ∗ λ and adding λ ⟨ η, µ ∗ λ -µ ∗ ⟩ on both side:

For ( I ) , we have

where we use Lemma F.2 in the last inequality.

Therefore, the above inequality leads to

where we again use Lemma F.2. This further leads to L ( µ ∗ λ ) ≤ λ 2 ∥ p ∥ 2 2 .

Part (iii) Using part (i) we have

With Lemma F.2, LHS and RHS become

Then using part (ii) we have the desired result.

We are now ready to characterize the approximated solution by comparing µ and µ ∗ λ . Lemma F.4. Recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then, the following holds:

Proof. We show the results one by one.

Part (i) By the definition of the optimality gap ζ and adding -λ ⟨ η, µ -µ ∗ ⟩ on both side, we have

Note that on LHS,

where we use Lemma F.2 and Lemma F.3.

Therefore, with Lemma F.2 and Lemma F.3 we get

Solving the above inequality on L ( µ ) gives L ( µ ) ≤ 5 λ 2 ∥ p ∥ 2 2 +4 ζ .

Part (ii) Again from the definition of the optimality gap ζ , we have

where we use Lemma F.3. Thus, | µ | 1 ≤ λ ∥ p ∥ 2 2 + | µ ∗ | 1 + ζ/λ ≤ 3 | µ ∗ | 1 .

The lemma below shows that far-away neurons are still small even for the approximated solution. Intutively, we use the non-degenerate dual certificate to certify the gap between µ and µ ∗ λ and give a bound for it.

Lemma F.5. Recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then, we have

In particular, when µ = ∑ i ∈ [ m ] a i ∥ w i ∥ 2 δ w i represents finite number of neurons, we have

where δ j = ∠ ( w j , w ∗ i ) for j ∈ T i .

Proof. By the definition of the optimality gap ζ , we have

Rearranging the terms and adding -⟨ η, µ -µ ∗ ⟩ on both side, we get

For LHS, with Lemma F.2 and Lemma F.3 we have

For RHS, with Lemma F.2 and Lemma F.3 we have

When L ( µ ) ≥ λ 2 ∥ p ∥ 2 2 , we have RHS ≤ ζ/λ + λ ∥ p ∥ 2 2 . When L ( µ ) ≤ λ 2 ∥ p ∥ 2 2 , we have RHS ≤ ζ/λ +2 λ ∥ p ∥ 2 2 . Thus, in summary RHS ≤ ζ/λ +2 λ ∥ p ∥ 2 2 .

Combine the bounds on LHS and RHS we have

The following lemma shows that every teacher neuron must have at least one close-by student neuron within angle O ∗ ( ζ 1 / 3 ) . This generalize and greatly simplify the previous results Lemma 9 in Zhou et al. (2021). In particular, we design a new test function using the Hermite expansion to achieve this.

Lemma F.6. Under Lemma 6, if the Hermite coefficient of σ decays as | ˆ σ k | = Θ( k -c σ ) with some constant c σ > 0 , then the total mass near each target direction is large, i.e., µ ( T i ( δ )) sign( a ∗ i ) ≥ | a ∗ i | / 2 for all i ∈ [ m ∗ ] and any δ close ≥ ˜ Ω ( ( L ( µ ) a 2 min ) 1 / (4 c σ -2) ) with large enough hidden constant.

In particular, for σ is ReLU or absolute function, δ close ≥ ˜ Ω ( ( L ( µ ) a 2 min ) 1 / 3 ) . Here a min = min | a i | is the smallest entry of a ∗ in absolute value.

As a corollary, if the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) , then δ close ≥ ˜ Ω ∗ ( ( ζ + λ 2 ) 1 / (4 c σ -2) ) and for ReLU or absolute δ close ≥ ˜ Ω ∗ ( ( ζ + λ 2 ) 1 / 3 ) .

Proof. Assume towards contradiction that there exists some i ∈ [ m ∗ ] with some δ close ≥ ˜ Ω ( ( L ( µ ) a 2 min ) 1 / (4 c σ -2) ) with large enough hidden constant such that µ ( T i ( δ )) sign( a ∗ i ) ≤ | a ∗ i | / 2 . For simplicity, we will use δ for δ close in the following.

Let g ( x ) = ∑ ℓ ≤ k< 2 ℓ sign( a ∗ i ) sign(ˆ σ k ) h k ( w ∗⊤ i x ) be a test function, where h k ( x ) is the k -th normalized probabilistic Hermite polynomial and ℓ will be chosen later.

̸

Recall the Hermite expansion of σ ( x ) = ∑ k ≥ 0 ˆ σ k h k ( x ) and its property in Claim A.1. For the first term, it becomes

For the second term, it becomes

̸

where (i) in the third line we use cos δ ≤ 1 -δ 2 / 5 for δ ∈ [0 , π/ 2] and (ii) in the last line we use Lemma F.4 and choose ℓ = ⌈ (5 /δ 2 ) log(16 ∥ a ∗ ∥ 1 / | a ∗ i | ) ⌉ .

Thus, given | ˆ σ k | = Θ( k -c σ ) we have √ L ( µ ) √ ℓ = √ L ( µ ) ∥ g ∥ 2 ≥ 1 4 | a ∗ i | ∑ ℓ ≤ k< 2 ℓ | ˆ σ k | = 1 4 | a ∗ i | ∑ ℓ ≤ k< 2 ℓ Θ( k -c σ ) = | a ∗ i | Θ( ℓ 1 -c σ ) .

With the choice of ℓ = ˜ Θ(1 /δ 2 ) , we have δ = ˜ O ( ( L ( µ ) | a ∗ i | 2 ) 1 / (4 c σ -2) ) . Since δ ≥ ˜ Ω ( ( L ( µ ) a 2 min ) 1 / (4 c σ -2) ) with a large enough hidden constant, we know this is a contradiction.

As a corollary, with Lemma F.4 that L ( µ ) = 4 ζ + 5 λ 2 ∥ p ∥ 2 2 , we have δ ≥ ˜ Ω ( ( 4 ζ +5 λ 2 ∥ p ∥ 2 2 a 2 min ) 1 / (4 c σ -2) ) .

For the activation σ is ReLU or absolute function, by Lemma A.1 we know c σ = 5 / 4 , which gives the desired result.

The lemma below bounds R 2 using the fact that it is spiky (has small non-zero support).

Lemma F.8. Under Lemma 6, recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then

Proof. Using the same calculation as in Lemma 12 in Zhou et al. (2021), we have

With Lemma F.4 and Lemma F.5, we have

∥

R

∥

2

2

=

O

(

m

2

|

µ

∗

|

1

/

2

(

ζ/λ

+

λ

)

3

/

2

)

.

2

∗

The following lemma bounds R 3 . In fact, in the view of expressing the loss as a sum of tensor decomposition problem, R 3 corresponds to the 0-th order term in the expansion. It would become small when high-order terms become small, as shown in the proof below.

Lemma F.9. Under Lemma 6 and recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . If ˆ σ 0 = 0 and ˆ σ k > 0 with some k = Θ((1 / ∆ 2 ) log( ζ/ ∥ a ∗ ∥ 1 )) , then

Proof. As shown in Ge et al. (2018); Li et al. (2020), we can write the loss L ( µ ) as sum of tensor decomposition problem (recall ∥ w ∗ i ∥ 2 = 1 ):

Thus, we know for any k ≥ 1 ,

Given any w ∗ j and even k , we have

̸

We show the last 2 terms are small.

For the second term on RHS, we have

where (a) cos δ ≥ 1 -δ 2 / 2 for δ ∈ [0 , π/ 2] ; (b) (1 -x ) k ≥ 1 -kx for x ∈ [0 , 1] . For the third term on RHS, we have

̸

̸

where (a) cos δ ≤ 1 -δ 2 / 5 for δ ∈ [0 , π/ 2] ; (b) we choose k = Θ((1 / ∆ 2 ) log( ζ/ ∥ a ∗ ∥ 1 )) and Lemma F.4.

Therefore, we have

This implies that

where we use Lemma F.5. Rearranging the terms and recalling L ( µ ) = O ∗ ( ζ + λ 2 ) from Lemma F.4, we get the bound.

The following lemma gives the bound on the average neuron to its corresponding teacher neuron. It follows directly from the residual decomposition and previous lemmas that characterize R 1 , R 2 , R 3 respectively.

Lemma F.10. Under Lemma 6, recall the optimality gap ζ = L λ ( µ ) -L λ ( µ ∗ λ ) . Then for any i ∈ [ m ∗ ] , ζ = Ω( λ 2 ) and ζ, λ ≤ 1 / poly( m ∗ , ∆ , ∥ a ∗ ∥ 1 )

Proof. With the relation of residual decomposition, Lemma F.7, Lemma F.8 and Lemma F.9, we have for any i ∈ [ m ∗ ]

H.2 Omitted proofs in Section F.2

In this section, we give the omitted proofs in Section F.2. The key observation used in the proofs is that balancing the norm and setting α, β perfectly to their target values only decrease the optimality gap.

Lemma F.11. Given any θ = ( a , W , α, β ) satisfying | α -ˆ α | 2 = O ( ζ ) , ∥ ∥ ∥ β -ˆ β ∥ ∥ ∥ 2 2 = O ( ζ ) , where ˆ α = -(1 / √ 2 π ) ∑ m i =1 a i ∥ w i ∥ 2 and ˆ β = -(1 / 2) ∑ m i =1 a i w i . Let its corresponding balanced version θ bal = ( a bal , W bal , α bal , β bal ) as a bal,i = sign( a i ) √ | a i | ∥ w i ∥ 2 , w bal,i = w i √ | a i | ∥ w i ∥ 2 , α bal = ˆ α and β bal = ˆ β . Then, we have

Moreover, let the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) , we have results in Lemma F.4, Lemma F.5, Lemma F.6, Lemma F.7, Lemma F.8, Lemma F.9 and Lemma F.10 still hold for L λ ( θ ) , with the change of R 3 in (8) as

Proof. Recall in Claim B.1 we have

Note that | a i | ∥ w i ∥ 2 = | a bal,i | ∥ w bal,i ∥ 2 so that L ( θ ) = L ( θ bal ) + | α -ˆ α | 2 + ∥ ∥ ∥ β -ˆ β ∥ ∥ ∥ 2 2 . We then have

Therefore, we have the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) ≥ L λ ( θ bal ) -L λ ( µ ∗ λ ) = ζ bal . Note that θ bal corresponds to a network that has perfect balanced norms and fitted α, β , thus all results in Lemma F.4, Lemma F.5, Lemma F.6, Lemma F.7, Lemma F.8, Lemma F.9 and Lemma F.10 hold for θ bal . Since ζ ≥ ζ bal , | a i | ∥ w i ∥ 2 = | a bal,i | ∥ w bal,i ∥ 2 and L ( θ ) = L ( θ bal ) + O ( ζ ) , we can easily check that all of them also hold for θ . For the bound on R 3 , note that

so that the same bound still hold for R 3 .

LemmaF.12. Under Lemma 6, suppose optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . Then ∥ a ∥ 2 2 + ∥ W ∥ 2 F ≤ 3 ∥ a ∗ ∥ 1 .

Proof. We have

where we use Lemma F.3. Rearranging the terms, we get the result by noting that | µ ∗ λ | 1 ≤ ∥ a ∗ ∥ 1 .

H.3 Omitted proofs in Section F.3

In this section, we give the omitted proofs in Section F.3. We will consider them case by case.

The lemma below says that one can always decrease the loss if norms are not balanced.

Lemma F.15 (Descent direction, norm balance) . We have

i i 2 i i 2

The following lemma shows that one can always decrease the loss if there are close-by neurons that cancels with others. Intuitively, reducing such norm cancellation decrease the regularization term while keeping the square loss term, which decreasing the total loss as a whole.

Lemma F.16 (Descent direction, norm cancellation) . Under Lemma 6 and Assumption F.1, suppose the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . For any w ∗ i , consider δ sign such that δ close < δ sign = O ( λ/ζ 1 / 2 ) with small enough hidden constant ( δ close defined in Lemma F.6), then

̸

We split the above into two terms (depending on square loss or regularization). WLOG, assume sign( a ∗ i ) = 1 . For the first term that depends on gradient on square loss,

where T i, + ( δ sign ) = { j ∈ T i : δ ( w j , w ∗ i ) ≤ δ sign , sign( a j ) = sign( a ∗ i ) } , T i, -( δ sign ) = { j ∈ T i : δ ( w j , w ∗ i ) ≤ δ sign , sign( a j ) = sign( a ∗ i ) } are the set of neurons that close to w ∗ i with/without same sign of a ∗ i .

As a result,

Since w j is δ sign -close to w ∗ i and ∥ R ∥ 2 2 = L ( θ ) , we have

where we use Lemma F.11 that L ( θ ) = O ∗ ( ζ ) .

For the second term that depends on regularization, we have

Therefore, when ( I ) ≤ 2 λ , i.e., δ sign = O ∗ ( λ/ζ 1 / 2 ) , we have

We compute a upper bound for LHS. Note that

where the last line we use Lemma F.6: ∑ j ∈ T i, -( δ sign ) | a j | ∥ w j ∥ 2 < ∑ j ∈ T i, + ( δ sign ) | a j | ∥ w j ∥ 2 because µ ( T i ( δ )) = ∑ j ∈ T i ( δ sign ) a j ∥ w j ∥ 2 > 0 .

Combine with the above descent direction, we have

which implies

The lemma below shows that when all previous cases are not hold, then there is a descent direction that move all close-by neurons towards their corresponding teacher neuron. The proof relies on calculations that generalize Lemma 8 in Zhou et al. (2021).

Lemma F.17 (Descent direction) . Under Lemma 6 and Assumption F.1, suppose the optimality gap ζ = L λ ( θ ) -L λ ( µ ∗ λ ) . Suppose

(i) norms are (almost) balanced: | ∥ W ∥ 2 F -∥ a ∥ 2 2 | ≤ ζ/λ , ∑ i ∈ [ m ] ( | a j | - ∥ w j ∥ 2 ) 2 = O ∗ ( ζ 2 /λ 2 )

̸

(ii) (almost) no norm cancellation: consider all neurons w j that are δ sign -close w.r.t. teacher neuron w ∗ i but has a different sign, i.e., sign( a j ) = sign( a ∗ i ) with δ sign = Θ ∗ ( λ/ζ 1 / 2 ) , we have ∑ j ∈ T i, -( δ sign ) | a j | ∥ w j ∥ 2 ≤ τ = O ∗ ( ζ 5 / 6 /λ ) with small enough hidden constant, where T i, -( δ ) defined in Lemma F.16.

(iii) α, β are well fitted: | α -ˆ α | 2 = O ∗ ( ζ ) , ∥ ∥ ∥ β -ˆ β ∥ ∥ ∥ 2 2 = O ∗ ( ζ ) with small enough hidden factor.

Then, we can construct the following descent direction

̸

where q ij satisfy the following conditions with δ close < δ sign and δ close = O ∗ ( ζ 1 / 3 ) : (1) ∑ j ∈T i a j q ij = a ∗ i ; (2) q ij ≥ 0 ; (3) q ij = 0 when sign( a j ) = sign( a ∗ i ) or δ j > δ close . (4) ∑ i ∈ [ m ∗ ] ∑ j ∈T i q 2 ij = O ∗ (1) .

Proof. Recall residual R ( x ) = f ( x ) -˜ f ∗ ( x ) . We have

where (a) we plug in the gradient expression and add and minus the term 2 ∑ i ∈ [ m ∗ ] ∑ j ∈T i E x [ R ( x ) a j q ij σ ( w ∗⊤ i x )] ; (b) rearranging the terms; (c) using L λ ( θ ) = ∥ R ∥ 2 2 +( λ/ 2) ∥ W ∥ 2 F +( λ/ 2) ∥ a ∥ 2 2 = L λ ( µ ∗ λ ) + ζ .

For the first line on RHS of (12), we have

where (a) due to assumption that norms are balanced; (b) we ignore L ( µ ∗ λ ) and add and minus λ ∥ a ∗ ∥ 1 ; (c) due to Lemma F.3; (d) due to assumption that norms are balanced and the choice of q ij . In the following, we will lower bound the last term of (12) to show it is no smaller than -ζ/ 8 so that we get the desired lower bound. Recall the residual decomposition (8) that R ( x ) = R 1 ( x ) + R 2 ( x ) + R 3 ( x ) , we have

Bound (I) For (I), recall R 1 ( x ) = (1 / 2) ∑ i ∈ [ m ∗ ] v ⊤ i x sign( w ∗⊤ i x ) , where v i = ∑ j ∈T i a j w j -w ∗ i is the difference between average neuron and corresponding ground-truth and ( ∑ i ∈ [ m ∗ ] ∥ v i ∥ 2 2 ) 1 / 2 = O ∗ (( ζ/λ ) 3 / 4 ) from Lemma F.10 and Lemma F.11. We have

̸

̸

̸

̸

where in (a) we plug in the definition of R 1 and using the fact that w ∗⊤ i x ( σ ′ ( w ∗⊤ i x ) -σ ′ ( w ⊤ j x )) = | w ∗⊤ i x | ✶ sign( w ⊤ j x ) =sign( w ∗⊤ i x ) ; (b) ˜ x is a 3-dimensional Gaussian since the expectation only depends on v k , w ∗ i , w j ; (c) | w ∗⊤ i ˜ x | ≤ δ j ∥ ˜ x ∥ 2 when sign( w ⊤ j ˜ x ) = sign( w ∗⊤ i ˜ x ) ; (d) a direct calculation bound as Lemma H.2; (e) definition of q ij .

Bound (II) For (II), recall

̸

For each term in (II) with j ∈ T i , we can split it into two terms that corresponding to T i and other T k 's.

̸

̸

For (II.i), we further split neurons into T i ( δ sign ) and others:

̸

̸

̸

Consider the first line of (14), from the choice of q ij we know a j q ij a ∗ i ≥ 0 . For ℓ ∈ T i, + ( δ sign ) , we know sign( a ℓ ) = sign( a ∗ i ) , which implies a ℓ a j q ij ≥ 0 for these terms. We thus only need to deal with neurons in T i, -( δ sign ) , we have the first line is bounded as

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

where (a) ˜ x is a 3-dimensional Gaussian since the expectation only depends on w ℓ , w j , w ∗ i ; (b) | w ⊤ ℓ ˜ x | ≤ δ ℓ ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ ℓ ˜ x ) and | w ∗⊤ i ˜ x | ≤ δ j ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ j ˜ x ) ; (c) a direct calculation as in Lemma H.2; (d) assumption that norm cancellation is small.

̸

For the second term of (14), similar as above, we have

̸

̸

̸

where (a) ˜ x is 3-dimensional Gaussian vector since the expectation only depends on w ℓ , w j , w ∗ i ; (b) | w ⊤ ℓ ˜ x | ≤ δ ℓ ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ ℓ ˜ x ) and | w ∗⊤ i ˜ x | ≤ δ j ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ j ˜ x ) ; (c) a direct calculation as in Lemma H.2; (d) choice of q ij and Lemma F.5 and Lemma F.11 that far-away neurons are small.

Thus, for (II.i) we have

For (II.ii), we have

̸

̸

̸

where (a)(b) ˜ x is a 4-dimensional Gaussian vector, | w ⊤ ℓ ˜ x | ≤ δ ℓ ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ ℓ ˜ x ) and | w ∗⊤ i ˜ x | ≤ δ j ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ j ˜ x ) ; (c) by Lemma H.1; (d) choice of q ij and Lemma F.5 and Lemma F.11 that far-away neurons are small.

Combine (II.i) (II.ii), we have for (13)

This further gives the lower bound on (II):

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

where (a) plugging in the expression of R 3 and using Lemma F.9 and Lemma F.11; (b) using Lemma H.3 and the fact that ˜ x is a 3-dimensional Gaussian vector and | w ∗⊤ i ˜ x | ≤ δ j ∥ ˜ x ∥ 2 when sign( w ∗⊤ i ˜ x ) = sign( w ⊤ j ˜ x ) ; (c) Lemma H.2; (d) choice of q ij .

Combine all bounds Combine (I) (II) (III) we now get the last term of (12)

From Lemma F.6 we can choose δ close = O ∗ ( ζ 1 / 3 ) and from Lemma F.16 we can choose δ sign = Θ ∗ ( λ/ζ 1 / 2 ) . Also with τ = O ( ζ 5 / 6 /λ ) , we finally get

as long as ζ = O ( λ 9 / 5 / poly( r, m ∗ , ∆ , ∥ a ∗ ∥ , a min )) with small enough hidden constant.

1

Thus, we eventually get the lower bound of (12)

H.4 Technical Lemma

In this section, we collect several technical lemmas that are useful in the proof.

Lemma H.1. Consider α , β ∈ R 4 with ϕ = ∠ ( α , β ) ∈ [0 , π ] and ∥ α ∥ 2 = ∥ β ∥ 2 = 1 and x ∼ N ( 0 , I ) . Then, for any 0 < δ 1 , δ 2 ≤ ϕ we have

✶ Proof. We first consider the case when at least one of δ 1 , δ 2 ≥ cϕ for a fixed small enough constant. WLOG, suppose δ 2 ≥ cϕ . In this case, it suffices to show a bound O ( δ 1 ) . We have

̸

̸

Then, we focus on the case when δ 1 , δ 2 ≤ cϕ for a fixed small enough constant. WLOG, assume α = (1 , 0 , 0 , 0) ⊤ , β = (cos ϕ, sin ϕ, 0 , 0) and ϕ ∈ [0 , π/ 2] . Then we have

Lemma H.2 (Lemma C.9 in Zhou et al. (2021)) . Consider α , β ∈ R 3 with ∠ ( α , β ) = ϕ and α ⊤ β ≥ 0 . We have

̸

̸

✶ Lemma H.3. Consider α , β ∈ R d with ∠ ( α , β ) = ϕ , ∥ α ∥ 2 = ∥ β ∥ 2 = 1 and α ⊤ β ≥ 0 . We have E x [ | α ⊤ x | sign( α ⊤ x ) =sign( β ⊤ x ) ] = O ( ϕ 2 ) .

✶ Proof. It suffices to consider α , β , x ∈ R 2 . WLOG, assume α = (1 , 0) ⊤ and β = (cos ϕ, sin ϕ ) ⊤ We have

̸

Proof. We have

In the following, we aim to lower bound ∑ j ∈ T i, + ( δ close ) a 2 j . Given ∑ j ∈ T i, + ( δ close ) | a 2 j -∥ w j ∥ 2 2 | ≤ | a ∗ i | / 2 , we have

where the last inequality is due to Lemma F.6: ∑ j ∈ T i, + ( δ close ) | a j | ∥ w j ∥ 2 ≥ | ∑ j ∈T i ( δ close ) a j ∥ w j ∥ 2 | ≥ | a ∗ i | / 2 . Thus, we have ∑ i ∈ [ m ∗ ] ∑ j ∈T i q 2 ij = O ( ∥ a ∗ ∥ 1 ) .

̸

I Proofs in Section G (non-degenerate dual certificate)

In this section, we give the omitted proofs in Section G. The proofs are mostly direct computations with the properties of Hermite polynomials in Claim A.1.

Lemma G.1 (Non-degeneracy of kernel K ) . For any h > 0 , let ℓ ≥ Θ(∆ -2 log( m ∗ ℓ/h ∆)) , kernel K ≥ ℓ is non-degenerate in the sense that there exists r = Θ( ℓ -1 / 2 ) , ρ 1 = Θ(1) , ρ 2 = Θ( ℓ ) such that following hold:

(i) K ( w , u ) ≤ 1 -ρ 1 for all δ ( w , u ) := ∠ ( w , u ) ≥ r .

Proof. With the property of Hermite polynomials in Claim A.1, we have

where θ = arccos( w ⊤ u ) .

Part (i) Given that r = Θ(1 / √ ℓ ) with a small enough hidden constant, we know for δ ( w , u ) ≥ r

where c is a constant less than 1. Thus, ρ 1 = Θ(1) .

Part (ii) For tangent vector z that z ⊤ w = 0 , we have ( ∥ w ∥ 2 = ∥ u ∥ 2 = 1 , δ ( w , u ) ≤ r )

For the first term, we have

where we use Lemma I.1 and ˆ σ 2 k = Θ( k -5 / 2 ) in Lemma A.1. For the second term, we have

Given that r = Θ(1 / √ ℓ ) with a small enough hidden constant, we know

since Z 2 σ = Θ( ℓ -3 / 2 ) .

Part (iii) Recall that δ ( w ∗ i , w ∗ j ) ≥ ∆ for i = j . It suffices to bound ∥ ∥ K ( ij ) ( w , u ) ∥ ∥ 2 ≤ h/m 2 ∗ for θ = δ ( w , u ) ≥ ∆ . Given that ℓ ≥ Θ(∆ -2 log( m ∗ ℓ/h ∆)) with large enough hidden constant, from (15) we have for ∥ w ∥ = ∥ u ∥ = 1

̸

where we use ˆ σ 2 k = Θ( k -5 / 2 ) in Lemma A.1 and (a) the last two terms bound similarly as in K (20) and first term 1 Z 2 σ ∑ k ≥ ℓ ˆ σ 2 k k ( k -1)( k -2) cos k -3 θ sin 3 θ ≤ Θ( ℓ 3 / 2 ) ∑ k ≥ ℓ Θ( k 1 / 2 )(1 -∆ 2 / 5) k ≤ h/ 3 m 2 ∗ .

Lemma G.2 (Regularity conditions on kernel K ) . Let B ij := sup w , u ∥ ∥ K ( ij ) ( w , u ) ∥ ∥ w , u and B 0 = B 00 + B 10 +1 , B 2 = B 20 + B 21 +1 . We have B 00 = O (1) , B 10 = O ( ℓ 1 / 2 ) , B 11 = O ( ℓ ) , B 20 = O ( ℓ ) , B 21 = O ( ℓ 3 / 2 ) , and therefore B 0 = O ( ℓ 1 / 2 ) , B 2 = O ( ℓ 3 / 2 ) .

Proof. We compute B ij one by one from (15) (see part (iii) proof in Lemma G.1). Using Lemma I.1 we have

For first term above ∑ k ≥ ℓ Θ( k 1 / 2 )(1 -θ 2 / 5) k -3 θ 3 , using Lemma I.2 we have

For second term above ∑ k ≥ ℓ Θ( k -1 / 2 )(1 -θ 2 / 5) k -1 θ we have

Therefore, we have B 21 = O ( ℓ 3 / 2 ) .

I.1 Technical lemma

We collect few lemma here used in the proof. They mostly rely on direct calculations.

Lemma I.1. For large enough integer k , we have

Proof. We only compute the first one max | cos k θ sin θ | = 1 / √ k . Others are similar.

We compute the gradient of f ( θ ) = cos k θ sin θ and get f ′ ( θ ) = cos k -1 θ (cos 2 θ -k sin 2 θ ) . We only need to consider θ ∈ [0 , 2 π ] . So the maximum is achieved either at boundary θ = 0 , π or f ′ ( θ ) = 0 . Then one can verify that the bound is true.

J Notes on Sample Complexity

The current paper focuses on the analysis on population loss, which is already highly non-trivial and requires new ideas that we developed in the paper. The finite-sample analysis is not our focus, so we omit it in the current paper.

For sample complexity, we believe the following strategy would work to get a polynomial sample complexity. We can break down the analysis into 2 parts: early-stage feature learning (Stage 1 and 2) and final-stage feature learning (Stage 3).

Stage 1 and 2: This should follow the results in Damian et al. (2022). The most important step is to show the concentration of first-step gradient (Stage 1). As shown in Damian et al. (2022), using concentration tools we can get sample complexity n = Θ ∗ ( d 2 ) , where n is the number of sample and d is input dimension.

Stage 3: In local convergence regime, all weights have norms bounded in O ∗ (1) due to ℓ 2 regularization we have. Thus, we can apply standard concentration tools to show the empirical gradients are close to population gradients given a large enough polynomial number of samples.

Achieving a tight sample complexity is an interesting and challenging open problem that is beyond the scope of current work.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The main contribution is our Theorem 2, which matches the claims in abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: This is a theory paper so all assumptions are clearly listed and discussed. The limitations, for example Stage 2 in Algorithm 1, are clearly discussed in the paper.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All assumptions are clearly listed in the main text and full proofs are given in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [NA]

Justification: No experiments in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [NA]

Justification: No experiments in the paper.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [NA]

Justification: No experiments in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [NA]

Justification: No experiments in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [NA]

Justification: No experiments in the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper follows the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This is a theory paper and has no foresee direct societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This is a theory paper.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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:
-

LIMITATIONS:
See weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThis 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.\n\nQUESTIONS:\n- 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.\n- 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. \n-  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).\n- 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? \n- 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? \n- What do the authors mean when saying \"complexity of $f_*$\" on Page 4? It would be nice to mathematically formalize this concept.\n- 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$. \n- What is $\\bar{w}_i$ at page 5?\n- What is meant for $\\varepsilon_0$-net at page 5? \n- A schematic drawing of the descent direction after the description on page 6 would help the reader to grasp intuitively the concepts.\n\n\n[1] A theory of non-linear feature learning with one gradient step in two-layer neural networks. Moniri et al. ICML 2024\n[2] Asymptotics of feature learning in two-layer networks after one gradient-step. Cui et al. ICML 2024\n[3] The committee machine: Computational to statistical gaps in learning a two-layers neural network. Aubin et al. NeurIPS 2018.\n\nLIMITATIONS:\nThe limitations are addressed in the submission.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- The problem of feature learning studied in the paper is important.\n- 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.\n\nWEAKNESSES:\nComments:\n\n\n1. 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.\n\n\n2. 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).\n\n3. What happens if information exponent is larger than 1?\n\n4. 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.\n\n5. 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. \n\n6. 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?\n\n7. 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.\n\n8. 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.\n\n9. 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.\n\n\n\n\n[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.\n\n[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.\n\n[3] A Radhakrishnan, D Beaglehole, P Pandit, M Belkin, Mechanism for feature learning in neural networks and backpropagation-free machine learning models.\n\n[4] D Beaglehole, I Mitliagkas, A Agarwala, Gradient descent induces alignment between weights and the empirical NTK for deep non-linear networks.\n\nQUESTIONS:\nplease see weaknesses\n\nLIMITATIONS:\nplease see weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}