How does Gradient Descent Learn Features --- A Local Analysis for Regularized Two-Layer Neural Networks
Authors: Mo Zhou
Rong Ge
Venue: NeurIPS 2024 poster (2024)
Keywords: deep learning theory
feature learning
optimization
Primary area: optimization_for_deep_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.

How does Gradient Descent Learn Features - A Local Analysis for Regularized Two-Layer Neural Networks
Mo Zhou" University of Washington mozhou170cs washington, edu
Rong Ge Duke University ronggedcs duke. edu
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 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 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 certain threshold, gradient descent with 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 also occur towards the end of training.
Introduction
Feature learning has long been considered to be major advantage of neural networks. However, how gradient- based training algorithms can learn useful features not well-understood. In particular, the most widely applied analysis for overparametrized neural networks is the neural tangent kernel (NTK) (Jacot et al., 2018: 8; Du et al., 2019; Allen-Zhu et al., 2019b). this setting, the neurons don't t move far from their initialization and the features are determined by the network architecture and random initialization (Chizat et al., 2019).
While there 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, 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 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 the ground-truth function.
Recently, growing line of works (Daniely and Malach, 2020; Damian et al., 2022; Abbe et al., 2021,
2022, 2023; Ychudai and Shamir, 2019; Shi et al., 2022: Ba et al., 2022; Mousavi-Hossee et al., 2023; Barak et al., 2022; Dandi et al., 2023; Wang et al., 2024; Nichani et al., 2024a,b) showed that early stages gradient training (either one/a few steps of gradient descent or a small amount time of gradient flow)
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
Work done Duke University.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 fw, as long as the width of student network m is at least some quantity my that only depends on fa, a variant of gradient descent algorithm (Algorithm I, roughly gradient descent with decreasing weight decay) can recover the larget 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-overparamettiie 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 kernellaay 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; Mousavi- Hosseini 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 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 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/multt-index models with neural networks Single/Multi-ineex 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 singlc-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-trute 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 convex nor satisfies PI. condition. Chizat (2022) considered square loss with l2 regularization similar to our setup and showed the local loss landscape 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 & regularization instead of &g 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, focus more natural but technically challenging case that second-laye 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 {I,..., For vector w, we use |||||2 for its 2-norm and w = w/ ||0|| as its normalized version. For two vectors w, we use L(w, v) = arccos( w v p|/(|//l2 || E [0, T .7/2 the angle between them (up to a sign). matrix A let ||A|| be its Frobenius norm. We use standard O, θ to hide constants and O S2, θ 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 ~ NO, Ia). The goal is to learn the teacher network of size m,
(x = Σ of (ww x) + wi x bū:
where (x) == max(0, x is ReLU activation, S, := span(w, wi...... w } is the target subspace.
Without loss of generality, we will assume ||w/ l| due to the homogeneity of ReLU.
Following the recent line of works learning fingl//multiiinde models (Ba et al., 2022; Damian et al., 2022), assume the target network has low dimensional structure.
Assumption 1. Teacher neurons form low dimensional subspace in Rd, that is
dim(S.) dim(span ........ w }) < d.
We will also assume the teacher neurons are non-degenerate in the following sense:
Assumption 2. Teucher neurons are A-separated, that angle L(w], ,wz) 2 A for all + j.
Assumption 3. H: Em w T is non-degenerat in target subspace S., i.e., rank H) r. Denote = (H))
Assumption 2 simply requires all teacher neurons pointing to different directions, which 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 at most 2 due Assumption 3. Indeed, one can show [/ ha)))(( x)) ā2 He, where h2(x)) is the 2nd-order normalized Hermite polynomial and 6₂ 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. There- fore, we will use Q., O, to hide poly(r, m., A, |a.l...., jom. k)), 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:
f(x; = Σ (w, x) + a + BTxT. i=1
(I)
where a = (0.,...,am,) a E RM, W- = (w) ... ...)) Wm)) E Baxd and 0 = (a, W, a, B).
3

---

Loss and algorithm Consider the square loss function with & regularization under Gaussian imput
(0) = (o
m) + |all
W 1
We will use to denote the square loss for simplicity. The l regularization the same as the commonly used weight decay in practice. Our goal is to find the minima of unregularized problem (A 0) to recover teacher network fa. 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 preprocessed data (a,, 17) in the loss function as Damian al. (2022). Specifically, given any (æ, 3) with = (x) denote a, = Ex [v] and 3_ = En [yæ]], we get
(x) = y - a. B xc
This preprocessing process essentially removes the O-th and 1-st order term in the Hermite expansion of a.. See Appendix A for a brief introduction of Hermite polynomials and Claim B.I.
Our algorithm shown in Algorithm 1. weight decay A that goes to O. 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.
It is roughly the standard GD following given schedule of
We will use symmetric initialization that a; dm-m//2 w; m with a, Unif w; Unif(()/ (1/Vm)8=-1) a 0, B 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
Input:
initialization g(0), weight decay A and stepsize т)
Data preprocess: get (x, y) according to (3)
Stage I: one step gradient update
((1) t ((0) 77 VeL2. (0(0))
Stage 2: norm adjustment by convex program
a (T)) a (T2), (Ta)
(T))) + mina,a,9 L(a, W()) a, 3) + +xE
w.
Balancing norm between two layers s.t. Jas| || 122 for all Stage 3: local convergence
for k < K do
// for each epoch, run GD until convergence
for T3,k-1 < t < T3 do
((++l) +- f(t) 1, VoL (0())
Output: (Ts,k) (a (T, (T,,x) a (a,kk) ((Txx))
3
Main results
In this section, we give our main result that shows training student network using Algorithm can recover the target network within polynomial time. We will foeus the case that d Q (1)1) when the complexity of target function small.
Theorem 2 (Main result). Under Assumption 1, 2, 3, consider Algorithum on loss (2). There exists a
schedule of weight decay X₂ and step size T) such that given m 2 mg = (1) (1/eg) 0(r) neurons with small enough Eo (1), with high probability we will recover the target network L(0) S within time T = O, (1/n=²) where = poly 1/m).
Moreover, when + every student neuron w, either aligns with of teacher neuron w as L(w;, w 0 or vanishes |a.| || 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 mo only depends on the complexity of target function f- (only quantities that are related to fa, not student network f or error E), so any mildly overparametrized networks can learn fo 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 Eo 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 (see Section 4.1 for more details). Intuitively, these random features need to span the whole space of the target function class 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/e)0(r) neurons to achieve desired accuracy e. In contrast, continue training both layer at the last phase of training allows us learn not only subspace but also exactly the ground-truth directions. Moreover, we only use (1(E))0(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.
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-wide width networks)
C =L,,
min (EM(3A-1)
Lx, (P),
where (( is the set of measures the sphere S0-1, As example, if p= as || 6w: then [x( recovers (f) when linear term G, are perfectly fitted and norms are balanced |a:1 |l: We defer the precise definition of [x( to (6) in appendix.
4.1
Stage 1
For Stage I, 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 with Ao = 1)0 = and
m 2 ma = O.() (1/Eo) with any eo = (1). After first step, with probability 1 - 5 we have
(i) for every teacher neuron wi, there exists at least one student neuron w, L(w, w,) S so-
(ii)
w || ||, (1),, |a 1 s O.(1/ (m) all E [mm.,, 01 and Br 0.
The key observation here is similar to Damian et al. (2022) that w!! ~ 2 (0) (o? given H is non-degenerate in target subspace S, we essentially sample w (1) from the target subspace. It is then natural expect that the neurons form an cn-net in the target subspace given mg neurons.
Hw:) so that
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 low loss solution in Stage 2.
Lemma 4 (Stage 2). Under Assumption 1,2,3, consider Algorithm 1 with A = / Given Stage 1 in Lemma we have Stage 2 ends within time T2 (1(/ee)) such that optimality gap Çr2 O. (Eo).
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 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 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.
5

---

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 multi-epoch gradient descent with decreas- ing weight decay À at every epoch. Note that this result only requires the initial optimality gap is small and width m 2 m, (target network width, not mo).
Lemma 5 (Stage 3). Under Assumption 1,2,3, consider Algorithm I on loss (2). Given Stage 19/5), in Lemma 4, if the initial optimality gap S3,0 S Oz(23,0), (,00)) weight decay X follows the schedule of
initial value λ (1), and k-th epoch A3,k = 2-, and stepsize Tak =1 S O d 3) for all Ts S t S Ts,++1 in epoch k. then within K = O. (1o)(()EE) epochs and total T3 T2
O. (A5.6' - 2) time we recover the ground-trudt network L(O0) S &.
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 Lojasiewicz property (Lojasiewiez, 1963). This means GD can always make progress until the optimality gap Ç is small.
Lemma 6 (Gradient lower bound). When Qt (AP) 5.50 (.(2//5 and X S O.()), we have
Vo (5
Note that this generalizes previous result Zhou 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 (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.
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 (-u, 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 '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 & regular- ization Σ, |a,12 + ||w;|l would become the effective l regularization 2 E, |asl ||w,|| 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 gd-1) based on their distance to the closest teacher neurons: denote T, = (w : L(w, wx{" S /ww, (ww,w)) for any j # i} (break the tie arbitrarily) so that U,T, -d We will also use 6, to denote L(w,, wx)) for j Te.
As described above, weight decay can always lead descent direction when norms are not balanced or norm cancellation happens (see Lemma F.15 and Lemma F. F.1 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 balanced and no norm cancellation happens. there exists properly chosen qij 2 and ΣjeT; aj9ı = of such that
Σ
Σ
V wi T,x, w. - 9is ; = =(())
SE SE(ma) JET

---

Teacher neuron
Student neuron
Descent direction
Figure Illustration of descent direction
In words, this descent direction is the following: move neuron wy Ti toward either ground-trun truth direction w or 0 depending on whether it is in the neighborhood teacher neuron w.. Specifically, move far-away neurons towards 0 (and thus setting 9i 0) and move close-by neurons towards its "closest' minima Øz w (the fraction of wn that neuron w should target to approximate). See Figure for an illustration.
The proof of the above lemma requires a dedicated characterization the low loss solution's structure, which we describe in Section 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, compute the inner product between gradient and constructed descent direction. We can lower bound by (assuming norms are balanced)
+ Σ Σ En [R(x)),))) 20(o((w)) (wi Tx) - o' duww] x))),
VE 3E(m.Jje JJET
where R(x)= = f(x) - (xx) is the residual. Thus, in order to get a lower bound, the goal is to show second term above is small than Ç. As the inner product between R(x) and h(xx)= Exe(m.] İjET, ajfty " x(o'(ww) x) o'(w] x)).
see, this term is quite complicated and can be viewed as
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 denote v, Ejer. ayw; as the average neuron. Intuitively, this average neuron v; stands for an idcalize where all neurons belong to T, (closer to w than other wz) collapse into a single neuron.
We decompose the residual R(x) f(x) Fx()) into the 3 terms below: denote é; U, w
1
Rí (x:)
Σ
6. a sign(w; a)), R2(x))
Σ
a,w. [siign(w) x)) sign(w;; x)),
ig(mml
sElm.lJjT.
Rg(x:)
at || Σ [E[m..]
as ||uc:ll2 +a aa-a-t +a -3] TT
v 2m
(E(m)
Rj can be thought as the exact- pxacttparameeriization setting (use neurons to learn m. neurons), where the average neurons {u,) ) are the effective neurons. The difference between this exact- parametrization and overparametrization setting is then characterized by the term Rg, which captures loss landscape degenerate in overparametrized case and slows down the convergence (Zhou et al., 2021; Xu and Du, 2023). Overall, this residual decomposition similar to Zhou al. (2021), with additional modification of R3 to deal with ReLU activation and linear term
difference nonlinear activation pattern. This term fact suggests

---

To some extent, our residual decomposition can be viewed as a kind of "bias-variance" decomposition in the sense that the 'bias' term R captures the overall average contribution of all neurons, and the 'variance' term R2 captures the individual contributions of each neuron that are not reflected in Ri.
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, hi).
High-level proof plan of Lemma 7
1. Show that neurons mostly concentrate around ground-truth directions.
2. Show that average neuron v; is close teacher neuron wo for all E [ma]
3. Use above structure to bound (R,, ht.. Specifically, bounding (R1, h) relies on the fact that average neuron is close to teacher neuron (step 2); a bound on (R2, h) follows from far-away neurons are small (step 1); third term (Ra, lt) 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 and (iii) corresponding to step above).
Lemma (Informal). Suppose the optimality gap is S, we have
(i) Total norm of far-away neurons is small: Eie(m.] ESET, 1a511|05122 - -O((5/)), where angle 8,- L(wj, w;) for w w that j E T.
(ii) For every wi there exists at least one close-by neuron w s.t. L(w, wF) S Eclose = (2
(iii) Average neuron is close to teach neurons: we have ||0s - w S (CC/A)9)/1
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 (6, = 0) or norm becomes 0 (Jyyl ||w,|| || = 0). This can be understood from the & regularized loss (cquivalent to lp 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 develop for the proof.
6.2 Neurons concentrate around teacher neurons: dual certificate analysis and test function
We focus on Lemma 8())(i)) here. We will use dual certificate technique similar to Poon et al. (2023) to prove Lemma 8(), and a more general construction of test function to prove Lemma 8(ii). In below, we consider a relaxed version original optimization problem (2) by allowing infinite number of neurons, i.e., distribution of neurons, with (1)) = ReLUU()) 1/v2m x/2 instead of ReLU:
min (x( = 02.) + A/m41,
(5)
-)
where 23 is the minimizer. We use σ22 activation because this is the effective activation when linear terms a, 3 are perfectly fitted (remove Oth and 1st order Hermite expansion of ReLU, see Claim B.1 and (6) in appendix).
This is the loss function would have in the idealized setting: (1) linear term a, B reach their global minima (this is easy to achieve as loss is convex in them); (2) use & regularization instead of l2 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 its Hermite expansion is well-defined, generalizing Zhou al. (2021) that can only handle absolute/ReLU activation. In below we will focus on the activation 022 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 al., 2023). One common way is study its dual problem. The dual solution (a (maps TRd to R) of (5) when À = 0 satisfies E, dual problem can be found in e.g., Poon et al. (2023)). Here 7(w) = E, [==
T x]] E 1 (more detailed discussions x)) is often called dual certificate, as serves as a certificate of whether a solution μ is optimal. Its meaning will be clear in the discussions below.
this
8

---

We now introduce the notion of non non-degenerate dual certificate, motivated by Poon et al. (2023). Note that the condition n(w)) E |s implies that n(w;) sign(a;) and || S 1. The following definition is a slightly stronger version of the above implications as it requires y to decay at least quadratic when moves away from u.j.
Definition 1 (Non-degenerate dual certificate). n(w) is called non-degenerate dual certificate if there exists p(œ) such that n)ww) (=()))))22( for w gd-1 and
(i) n(w?) sign(a;) for i 1, I,...,,I... m.,
y(00)
(ii) |n(w)) < 0,, w if E Tr where "(w,w,) (wx, wf)
The existence and construction of the non-degenerate dual certificate deferred to Appendix G. We focus on the implications of such non- Figure 2: Dual certificate degenerate dual certificate below.
Roughly speaking, the dual certificate only focuses on the position of ground-truth directions wi as it decays fast when moving away from these directions (Figure 2). Thus, if μu exactly recovers ground-truti ,t, then we have M..) = |p//1. The gap between (n, (1) and |x/2 large when μ is away from pxx. Therefore, can be viewed as certificate to test the optimality of μ.. The lemma below makes it more precise.
Lemma 9, Given a non non-degenerate dual certificate n, then
(i) (n, (²) = 1eth and (n - -0)1 S lllll2 v L(p).
(ii) For any measure μ E M(S 167, S 1442 Py Exelm.1 Jy; 6(w, w dl (lw))
In the finite width case, we have Eie(m. .r. ,(w, w dlw((w) = Σ. |a,] |m 67. 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 14 The proof for general j with optimality gap ( is similar (Lemma F:5).
Claim (Lemma 8()) for 23)) Eie(m.] JT. 5(w, )² dlp)((w) S (A)
Proof. It is not hard show |,ū1 S |n^|2 (Lemma F.3) we have 1nth 1p011 (n.14 n²) - 14 n²))
Using Lemma 9 and the fact L(H)) = (X²) from Lemma F.3,
LHS =110/1 (n. MN Py Σ
/ 5(w, ((w,w;)) ddp)l(w) RHS < llpl22 v [(x))= = Oz()).
6E -e.mm.]
We have Exe(m.) ST. 8(w, uo; = OA(A).
Test function The idea of using test function is to identify certain properties of the target func- tion/distribution that 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 specific case of test function: the correlation between dual certificate 1) and distribution of neurons f is large (reach |p[1) only when μ z p4.
In below, (Lemma 8(i)). Denote (,())= == {j : L(wj, w.) S 6) T.T, as the neurons that are 8-close to w." Lemma 10 (Lemma 8(i), informal). Given the optimality gap S, we have the total mass near each target direction is large, i.e., (T)( sign(a;) Z 1 for all E [m.] and any 5 O. (33.
this test function idea to show that every ground-trui -ruth direction has close-by neuron
Note that although the results in the dual certificate part (Lemma 9(i)) can imply that there are neurons close teacher neurons, the bound we get here using carefully designed test function are sharper 1/3 vs. <2/4). This is in fact important the descent direction construction (Lemma 7).

---

In the proof, we view the residual R(x) fu(x) - (xx as the target function and construct test function that will only have large correlation if there is teacher neuron that have no close student neurons. Specifically, the test function 9 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). looks like a single spike dual certificate 7, but in fact decays much faster than when moving away. It is more flexible to choose test function than dual certificate, so test function g can focus only on a local region of one ground ground-truth direction and give better guarantee than dual certificate analysis.
a(m)
Figure 3: Test function
6.3 Average neuron is close to teacher neuron: residual decomposition and average neuron
We give the proof idea for Lemma 8(ii) that shows average neuron is close to teacher neuron wi using the residual decomposition R R1 + R2 + Rg.
The key is to observe that R1 is an analogue to pact-parametrizztiion case where loss is often w I|-- Then the goal is upper bound R |
strongly-convex, so we have Rı 1/ (1) E, ||::
Given the decomposition R = Ri + R2 + Ra. it is easy to bound | < || + ||]]]]]] || + Rall-
We focus on || as the other two are not hard to bound (loss is small in local regime). R2
in fact closely related with the total weighted norm bound in Lemma 8: we show | =
O.((C/A)//2). Thus, get bound for l|u, w See 3/2
(1) EjET, layl ||ee,, |l2 03)
Appendix F.1.4 for details.
Conclusion
In this paper we showed that gradient descent converges in a large local region depending 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.
Acknowledgement
Rong Ge and Mo Zhou are supported by NSF Award DMS-2031849 and CCF-1845171 (CAREER).
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, 2--29999-27002. 2021.
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning 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 neural networks: lapp complexity and saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory, pages 2.22- 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 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.
10

---

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019a.
Zeyuan Allen Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. 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. High- dimensional 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 Rexearch, (8(19)11-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-2117 217644, 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-indes 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-555 509, 2006.
Lenaic Chizat. Sparse optimization on measures with over-parameterized gradient descent. Mathe- matical Programming. 194(1):487-532 2022.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for perapearamteriieed 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 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 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 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 I. Donoho. Compressed sensing. IEEE Transactions on information theory, 1294):128- 2006.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes prer-earameterized neural networks. In International Conference on Learning Representations, 2019.
11

---

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 over- parametrized tensor decomposition. Advances in Neural Information Processing Systems, 34: 1299-1311, 2021.
Margalit Glasgow. SGD finds then tunes features 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- parammetrrie two-layer neural networks beyond ntk. In Conference on learning theory, pages 2613-2683 PMLR, 2020.
Stanislaw Lojasiewicz. Une propriété topologique des sous-ensembles analytiques réels. Les équations aux dérivées partielles, 1 11::8789,, 1963.
Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: mean-fiell 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 two- layer 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-Hossein 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 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' O'Donnell. Analysis of boolean functions. arXiv preprint arXiv:2105. 1000::2005.1086, 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- prramammetttinnoooonnn 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 imputs 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, (8(1):267-288, 1996.
12

---

Zihao Wang, Eshaan Nichani, and Jason D. Lee. Learning hierarchical polynomials with three-layer neural networks. 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 overpaarranttrizzaainn exponentially slows down gradient descent for learning single neuron. In The Thirty Sixth Annual Conference Learning Theory, pages 1155-1198. PMLR, 2023.
Gilad Yehudai and Shamir Ohad. Learning 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- wer-parameterized two-layer neural network. In Conference on Learning Theory, pages 4577-4632. PMLR, 2021.
13

---

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 more complete discussion on Hermite polynomials. Let Hy be the probabilists' Hermite polynomial where
Hu = (-1)* 02 d
dx (e 12-
and h4p - Vk! Hx be the normalized Hermite polynomials.
Hermite polynomials are classical orthogonal polynomials, which means No. (o. 1 if and otherwise O. Given function d, call (xx) Σks (x))(())) the Hermit expansion of a and 0 Ez=N(0.1) a (x)hx(r) as the k-th Hermite coefficient of a.
The following is useful property of Hermite polynomial.
Claim A.1 ((O (O''Dnnell,, 2021), Section 11.2). Let (z,y) be p-correlated standard normal variables (that is, both I. huve marginal distribution N(0, 1) and Elry] = p). Then, n Smn where 6mn = 1 if m and otherwise 0.
=
The following lemma gives the Hermite coefficients for absolute value function and ReL.U. Lemma A.1. Let o₂ N (1 be the Hermite coefficient of a. For a is ReLU or absolute function, we have 1ōx1 = 6/44)
Proof. From Goel et (2020); Zhou et al. (2021) we have
odd
2/
Ōaba,k
(k-2)!
even and 2
#
VR12t/2 (k/2-11) 1)1
0
k odd and 2 3
11/2T
k 0
relu,k
1/2
k
(-1)
(K- 2-2 2x Vk/nh/22211 (1/22-111)
A is and
Using Stirling's formula, we get lāabs,el, @relu,kl 1|| = Ok:
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 a, B corresponding to Oth and 1st order of Hermite expansion. The effective activation is in fact 022 as defined below.
Claim B.1. Denote -(1//2m
Σn 0i | 12g 3 -(122) E. a We have square loss
2
a/2 2 + B L( = la -
Em (1/2(())) - [.()))
where (22(2;; 0) = Eie(n 2,022 ( x) and 022(2)) = ( 1 1/V x122 is the activation that after removing Oth and 1st order term in Hermite expansion.
As a result, when B are perfectly fitted and norms are balanced we have
LA (0) = F.g [x((f2((z)) - F=(=))
Σ lal [0lll
[E(m)
14

---

Proof. Following Ge et al. (2018), we can write the loss L(Ø) as sum of tensor decomposition problem using Hermite expansion as in Section A (recall |wo l0 and preprocessing procedure removes the 0-th and 1-st order term in the Hermite expansion of a):
L(O) =En
Σ i [m]
a;
l2
Σ
Tghi W. x))
11 x) - Σ ag | Ea x) (e(m-)
k20
k22
2
a + ōo
Σ 0i //w
+
3 +
δ1
xE(m)
ic(m)
Σ
Σ sE(m]
ag |1,,,)
112
w
80
Σ
|0
1| w 6
se/m.]
Note that 60 1/v2n σ 1/2 as in Lemma A. we get the result.
The proof of main result Theorem is simply 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 on loss (2). There exists a
schedule of weight decay X₂ and step size Th such that given m 2 = Ó.()) (1/eo) Oor) neurons
with small enough Eo O. (1)), with high probability we will recover the target network L(0) <
within time T O. 1/)r²)) where = poly(s, 1/d, 1/m).
Moreover, when + every student neuron w, either aligns with of teacher neuron w3 as L(w,, wf) 0 or vanishes |ai| || 0.
Proof. Combine Lemma (Stage 1), Lemma (Stage 2) and Lemma 5 (Stage 3) together and follow the choice of X and we get the result.
For the student neurons' alignment, it is direct corollary from Lemma F.6 and Lemma F.5.
C Stage 1: first gradient step
In this section, show that after the first gradient update the first layer weights Wi...., Wm form en-ne of the target subspace S.. given (1/eo))ir) Oii) neurons. The proof is deferred to Section C.
Lemma 3 (Stage 1), Under Assumption 1,2.3, consider Algorithm with Ao Yo and mi = Ō.) (1/eo))(r) with any Eo = O.(1). After first step, with probability 1 - we have
(i) for every teacher neuron wi, there exists at least one student neuron w, s.t. L(w,, wj) S sp.
(ii)
w 0 (1),
(0 s (.(1/00) all E [m..,, on and B 0.
The proof relies the following lemma from Damian al. (2022) that shows after the first step update ws's 's are located at positions as 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 & norm sense
w; (1)
-1o V ()) (()) 2m)a (0) (8) ( + (,H 20
where or == Ey
is the k-th Hermite polynomial coefficient.
C.1 Proofs in Section
We now ready give the proof of Lemma 3.
15

---

Lemma 3 (Stage 1), Under Assumption 1,2.3, consider Algorithm with λo тo and mo = O.()) (I/eo) with any Eo = (L). After first step, with probability 1 we have
(i) for every teacher neuron wi there exists at least one student neuron w, L(w,, w,) S Eu- (ii) un " 0.(1), la s (0(1//m) all [m.], 0 and 3r 0.
Proof. We show them one by one.
Part (i) From Lemma C.I and the fact that क samples uniformly from unit sphere, know the probability <(w(") w) for any given w is at least (x()). Applying union bound we get the desired result.
Part (ii) We have
w (1) 170 V I(a w()) a;
(0) E. [f. x)))
For the norm bound, using Lemma C. we know
Va (1 H (0)
o(vr))
Ō vr
| S
vd Vd
((( + +
(0) initializes from Gaussian distribution, we know the desired bound hold. Similarly, one Since : can bound (.
Since we use symmetric initialization and have preprocessed the data, it easy to see a, B remains at O.
D Stage 2: reaching low loss
In Stage 2, we show that given the features learned in Stage one can adjust the norms on top of it to reach low loss that enters the local convergence regime in Stage
Procedure We first specify the procedure to solve mina min, 8 L( + + weil2 la:|- For current point, we first solve the inner optimization problem, which is linear regression on a, B. From Claim B. know the global minima is (á. 3). For simplicity of the proof, we just directly set
(a. 3) = (á, 3). Then given the o. B, the outer optimization is convex optimization for a, which also be solved efficiently. Specifically, perform 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
(a)(aa) = Ezl(a 022 g) +
Σ || |2allla las)-
The following lemma shows that after Stage we reach low loss solution given the first layer features learned after first gradient step. The proof requires 1j to be small enough that depends on 1 1/m, mostly due 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 Algorithum with X = Veg. Given Stage in Lemma 3, we have Stage 2 ends within time T2 (1/nEo) such that optimality gap ÇT, O+
16

---

Proof. Denote a, as the minima of L1.x. Then, we have
| aalt1) a.
a a.
2n V I (aa" ) - a.) n² Val
1
(1()) a.
2n((s [á,)(a))) + (m)
( a. 2
2n(Enx((a)) [ā=)( + 7
where (a) we use idea loss L1, is convex in
lterating the above inequality over all t have
a aa a. 2
2
- - nITEu/2.
2n
2
1T
which means
||al)) - ā.11 l12
min (() (ā.) 1 ((ns(aaa)) - EST
min
+ 80/2.
n/
It is easy to see |a 1|ā=1 = (1). Thus, when 2 (1(n/))) we know L (a()(a(())) (..) & 3c//4.
This suggests the optimality gap after balancing the norm (so that (x(TTa)) = ((aa(a(5))))
ST² (O(TT)))
min -EM(BN-1)
(2(1))
=a -I,x(ā.)) + [n,)(ā.)
min Lx(e). HEM(B4 1)
For L1 (T)) - (ā.)) we just show above that it is less than 350/4.
For (ā.) minpEMiSs Lx(H), we have
[I,((ā))-
min Ex(1)= (T.))(a-)- min Lx(n)
(EM(3d-1)
MEM(Sa-1)
<0.(.)) + + +|lla,ll1 A14i/1 O,
where in the last inequality we use Lemma F.3 and μ arg min mi Lx(p). Here à, is rescaled version of a, and is constructed as: for every teacher neuron w choose the closest neuron w s.t. L(w,, wx) S Eo and set à.. = a |,, Set other å. =.
Together with above calculations, we have Cra < (s()))
E Stage 3: local convergence for regularized 2-layer neural networks
In this section we show the local convergence that loss eventually goes within polynomial time and recovers teacher neurons' direction.
The results in this section only need the width 2 m. as long as its initial loss is small.
Lemma 5 (Stage 3). Under Assumption 1,2,3, consider Algorithum I on loss (2). Given Stage 9/5 weight decay follows the schedule of initial value λ3,0 (1), and k-th epoch λ3 1/ and stepsize Vaλ 1] O. O -3) for all T3,k < t < T3,4+1 in epoch k. then within K O, (log(1/=)) epochs and total T3 T2
in Lemma 4, if the initial optimality gap Ça.0 < (x(2a))
O. O272.n 1 time we recover the ground-trua truth network L( <
The goal of each epoch is to minimize the loss LA with a fix X. The lemma below shows that long as the initial optimality gap is O. (29/5), then at the end of each epoch, L., could decrease to (X?).. Therefore, using a slow decay of weight decay parameter À for each epoch we could stay the local convergence regime for each epoch and eventually recovers the target network.
17

---

Lemma E.1 (Loss improve within one epoch). Suppose |a
w (00
for all i [m]. If
2
So S. (19/5) and S (1) and y - (a(A then within O+ Ox (\+ time the optimality gap becomes L2 (x(x5) -.
The above result relies on the following characterization of local landscape of regularized loss. We show the gradient is large whenever the optimality gap large. This is the main contribution of this paper, see Section F for detailed proofs.
Lemma 6 (Gradient lower bound). When Ω, (A2) S 5. 0.(2%/) and X s 0.(1), we have
VoLAll 2 (2(
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) allow the convergence analysis.
Lemma E.2 (Smoothness). Suppose |ail S ||w;ll2 and [[
2 x)
O.(d)) for all
e [m] n = (L/dd), then
L2 VoLo)) & (0) n +0.
E.1 Proofs in Section E
We now are ready to show the convergence of Stage by using Lemma E. to show the loss makes progress every epoch.
Lemma 5 (Stage 3). Under Assumption 1,2,3, consider Algorithm I on loss (2). Given Stage
in Lemma if the initial optimality gap Ça,o S O= 9/5 weight decay À follows the schedule of initial value X3,0 = (1),,, and le-th epoch A3,k = 1/2 and stepsize qskk = 1 S O. (A 12 d-33 for all Ts S < T3,441 in epoch ki then within K = O, (log(1/s)) epochs and total T3 T2 O. 0.72bb1 E 2) time we recover the ground-truth network L( S E.
Proof. Since la!0
(0) for all i [m] at the beginning of Stage 3, from Lemma E.3 we know w
they will remain hold for all epoch and all time t.
From Lemma we know for epoch it finishes within O, (A time and achieves Lx LA, (1x))) O.(A)) To proceed next epoch + 1, we only need show the solution at the end of epoch (()) gives the optimality gap Ç = O. 9/5 (2841) for the next Ak+11 We have
A6+1
Ak+1
(0(k) (0(*))
Dxwes (1 (T.((())) - I(N +
(k)
+
ARta Aarllreesl
2
2
(n)
50.(A + AR+1
[[(())
2
d
W 1|²
d
2
(b)
SO.(A2) + Ak+1
(O.(22) + I((3)) -(((())))
d
(0)
50.(1 S (x9//)
where (a) due to Lemma F.4 that (()) )) is small; (b) the optimality gap the end of epoch Is is O,(2) and 1rt, - 1, = (A²)) from Lemma F.3; (c) due to Lemma F.3 that L(x3)) is small. In this way, we can apply Lemma E. again for epoch k + 1.
From Lemma F.4 we know at the end of epoch the square loss T((O()) = Ox Thus, to reach E square loss, we need A = (e222)) which means we need to take O, (log(1/e)) epoch. Since epoch i it finishes within O. (A 4)-1) time, we know the total time is at most O- (2,) 1 time.
18

---

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).
(0) 1g for all i [m]. Lemma E.1 (Loss improve within epoch). Suppose |a: m
Su S (A%/5) and & (1) and y I Ox(A'ad33), then within O, "n "n=)) time the optimality gap becomes L (x(p)) = O.
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
Proof. Since |a; (0)
(0) :
Lx - VeLLs) & (o) V I1 + (,
Recall Çe = (0(0))) (x Using gradient lower bound Lemma 6 and consider the time before Se reach (A²) we have
S++1 S Se - +
S 22 (2C2//222))
where we use 0. ( to small enough.
The above recursion implies that
4 O + 2)
Thus, within (1/A/)) the optimality gap Çe reaches 0, (22).
The lemma below shows regularity condition the norm between two layers.
Lemma E.3. If we start at la; S all i E [m.] and all time t.
w and r = O1(1), then we have la 10
(000 [ for
Proof. Denote R(x) f(x) - (zx). Assume 1a
w (8)
s 0 we show it remains t 1.
We have
la;
((+1)
w
((+1) 2
-la!0
(0)
1 LA 0
(1) 1/
w0 V/ Lx((t))
-|a,00
101
w
(0)
+ 1
V (0())) - 172 V (1(0))))
(9)
1 2F [( (*
(T) x)) TTx²) Xa 101 172 1| (a (0 (' x))=]] + Xao 1
We first focus on the last 2 terms. We have
F (w (+) m)] da; 12) 2F (((r) (t) o' (w) (t) x) + Xun 1
w 122 |2 (R(z) T 1) >2 1a! 12 2E_ (R())o) =))=) X2 w
(
la:
[a42 100)2-1au (no
x2 - 2E R((z)' T |=: |13)
where (a) due to |2E, (w (+)
2)1 < 2E [R()) (w) (t) x)
Therefore, plug it back to the above equation, we have
a 12 w (++) 2 < |0 wo (11) 11+r222 - 77 2E (Rl[A(x)'') (w!!) x))0|]]]]]) 2
0,
19

---

where (a) due to 1= w
(4)
I 0 and we use [
O.(d) from =||
Lemma E.4 and is small enough.
Therefore, we can see that la (2)
I 0 remains for all t.
This lemma shows the smoothness of loss function. The proof requires a careful calculations to bound the error terms.
Lemma E.2 (Smoothness). Suppose |a.] S ||w,l22 and ||.n (0) x) O.(d) for all
i EE [m] If
(1/d/)), then
[ vVo..) S (0) ( //Vooll 0.(1))/2
Proof. Denote Ro(x: fo (x) fa) to denote the dependency on 0. For simplicity, we will
use Vo- VoL and same for others. Since (_ x)=l = Ox(d), we know
Va (1) 1illp and ا w. (-laaa,d)
We have
Lx 1) Vo) (x +
=0 Vo) (0) (V. Ve)
-E Bo vo (I À
a V. |
>
w + Vw
Ex [x(
X
2 |aal
X
2 W
2
Σ Egl Ro w. V.,1
Σ
Eg (w)) a) /... E. (Ro (x) - E Ro(()) //3l
sE/ml
kE[m]
Maa, Va) A,W, )
= Exl( &o (z) - (
(1)
+ Ex
RØ(
Rot5o (x) -Ro(x))
Σ
(w, x
x)
Σ
ago
V..
-aa x xTV
sE/m]
VE m)]
(IS)
2
+
w
2
2
The last line is easy on O. (1 using norm bound in Lemma F.12, so in below we going to bound (I) and (II) one by one. The goal is to show they are small in the sense of on order o(1)--
Bound (I) For (I), we can write out the expression
27
En (R (2) Ro(x))) =E,
Σ (e(m)
(a: + V Va.)o((w; + Tw) a,0 (w. + V. +
V
27
<2 Ex
Σ (a: Vac)o((w; + Vw)))x) ao x)
1E/m]
(1.4)
Va + x [//) 27
+ 2 Eg
(1.i))
20

---

For (Ii), we can split into 2 terms as
27
Ex
Σ 1m(m]
(at + Vayoo((w, + (w)) - (ws x)
27
27
Σ
Val oaw, =x))
2
Σ
ai (ww; vw))=)) a, x))
EE mm)
iE/m!
27
Va, (wo; + Vw,)
21
Σ ic(ml
as
V
vo c0
We then can bound them separately as
2 wi + V w: 1
2
(I.) 20(1)
01)
Σ 1E(ml
lail
A
112
a.
(E(m]
2
2
(b)
(0,(dP)
Σ
17
11w; + +
12: 12
+ (dd(
Σ
na
xE(m]
ic(ml
dr)
where (a) use Lemma E.5; (b) recall [a,l -n(n |1w l d) and
V ws
(-la,,dd);; (c)
|all [W] WilF Eiemm la,| w1llz (1(1) from Lemma F.12 and Lemma F.4, well 17 is small
enough.
For (I.i), we have
En V
(Va+ + x: V.)
<O(V.. | 0.(1))))
where we use Lemma F.4.
Combine (I.i) and (Lii) we know (1)=
Bound (II) For (II), we have
E,
Ro(
(0)5.(00) Ro(w=) -
Σ
(w. x)
Σ
0go (w xx) V m V, x TV
2E|m)
kc [m]
=E,
(z)
Σ (E(m)
(as + Va,)o((w; + Vw.) (w, x) oww V)V. a,o'( (w, x)x² T A
I(())
<
|
||
(E(m)
focus on bound || below. The goal is to show o(1). For [xxx), we have
|
F
(a: + V Va + V T (w, x) oun m) V aso x.)... "V. V.n.)
a. (a((w) + Vw:) x) - olw: ((wT))))² 2 + (a:o[(ww + v..)) Tx) - o(w! x) - o'(w, x)x vw.)))
<E. m
2
<2 Ex
[ a,
2 P
x 2
+2a Ex [ww - V..)) T x 2 ((w; + Tw) o'(,, x))
wi
(11.1)
(11.1))
21

---

For (II,), recall /Va,1 On 11w:l12 d) and \
O. (1)aaddd) have
E
Vas
2 V 1/
12 | 0.(1/aas W d2
For (II.ii, we have
[0; + V w. 2/ (o'(1s + V
o' owww 2²))
=
1e.e
(w, + Vw,) ) x0 sig(((w "aa
at) a)
<
Wi + V mn
where = L(w;+ + V u. wi) is the angle between w;+ + V w, and w; Since
V w
(nlasldd
(1 [1:ll d), we know 0 / given 1) (1/d) be small enough.
Combine (II.i) and (II.ii) we have
|1:122 S (, a Io (2 + wi V... 5 {// 50.11 1w:112dd)
Since |Roll = O.)1, this implies
(mn) < Σ 0,(11/2 1E/m]
111 //2) = O, 0,
Combine (IIII Finally, combing (I) and (II) we have
Lx - - LA (0) + v
Going back 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.
Lemma E.4. We have (R()) ( 2) O.(d)
Proof. It is easy to see given |R| = Ox(1).
Lemma E.5 (Lemma D.4 in Zhou et al. (2021)). Consider as; E Rd for E [n]. We have
n
27
2
Ex-N(0.)
Σ
a
S co Σ ||as||
where og is a constant.
Local landscape of population loss
In this section, we are going to show Lemma that characterizing the population local landscape with fixed A by giving the lower bound of gradient.
22

---

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 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 closely-related problem in Section F.1 (rewrite (5)):
min REM(SA-1
+ Nuali == (m.s((f(())) - D) + A/4l1
(6)
=E [0,
p o222 a:dd
pe
+ A/111,
(7)
where () is the measure space over unit sphere gd-1, ps Eie(m.] aj dw; and 022(2)
1e-11////////// is the activation that after removing Oth and 1st order term in Hermite expansion.
Note that when represents a finite-wdith network, we have μ = Eie[m as ||w:|l2 â,, is a empirical
measure over the neurons. In particular, when μr Mrs, model fp recovers the target f..
We call (5) the ideal loss because the original problem (2) would become the above (5) when balance the norms (I/w;ll 2|2 = la:D: perfectly fit a; B and relax the lnite-width constraints to allow wndhh (see Claim B.1). This is why we slightly abused the notation 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 a, 3 are not fitted well.
Notation
Denote the optimality gap between the loss at μ and the optimal distribution μ² as
((p) == (x) - Lx (PQ),
where u² is the optimal measure that minimize (5). For simplicity denote ū; a, || I|ا so that
[p|1 = ||āll when μ = Eie(m) as ||l -म. Often we use 4 = ((m) to denote the optimality gap
time just Ç for simplicity. We slightly abuse the notation to also use < = (0) La (1()). Finally denote μ² Eie(m.) of dws; (assuming || ll 1) that fa- (a:) Enowye [a22 u x)).
F1 Structure of the ideal loss solution
In this section, we will focus on the structure of approximated solution for the Eι regularized regression problem (5).
In the rest of this section, we will first introduce the idea of non-degenerate dual certificate and then use as tool to characterize the structure of the solutions. The proofs are deferred to Section H.
F111 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). v)(w) is called a non-degenerate dual certificate if there exists p(x) such that ()w) = Em (z(r()))22((1)) for u E 3d-1 and
(i) n(w?)) sign(a;) for 1, ,,..., ,., (ii) |n(w) < 1 - p,,(w, )² if E T, where v(w,w;) = L(w, wf).
We first show that there exist such non-degenerate dual certificate. More discussion and detailed proof are deferred to Section G.
Lemma F.1. There exists non-degenerate dual certificate 17 E,, (p( x)] with P., = +(1) and [|pl|2 S poly(m....)) A)
The following lemma (restate of Lemma 9) gives the properties that will be used in the later proofs: the non-eegeeeeate dual certificate v allows us to capture the gap between the current position y and the target μ²"
23

---

Lemma F2. Given non-degenerate dual certificate 1), then
(i) (n.1²) = 1p"
(ii) For any measure μ E MM(S-11), 167, S 144/ Pn Eielm.) ST. 6ww, wt) dl dlw)) Then K0. 1²)) s
(iii) (n, 1 1²)
Sp - Sur) where fp(x) = Eeo 122 [T)] x)).
llll|2 v L(n).
F.1.2 Properties of μ
Given the non-degenerate dual certificate 7, we now are ready to identify several useful properties of Mx The lemma below essentially says that μ is similar to μ² in the sense that most of the norm are concentrated in the ground-trum direction and the square loss small. The proof relies on comparing H with f using the optimality conditions. Lemma F.3. We have the following hold
(1) |1-11 A 2lpll < |xX| S |u*1 - 1="1==|aa"1]
(ii) M(x)) S X2 -
(iii) Eie[m.] .T. 5(w, w dlm}l(w) S X 1llll/// /p, O.(A)
F.1.3 Properties of with optimality gap
We now characterize the structure of μ when the optimality gap is Ç. The proof mostly relies 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
Lemma F4.. Recall the optimality gap Ç = =x((1= - Lx(x)). Then, the following holds:
(i) ((p) & 5X2 1pll + 4Ç- = O.(A2 + C).
(ii) ifC S 11 and X < |/411/ 11p12. then 1A/2 S 3/ 311
The following lemma characterize the structure of μu 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 S (14 Tx(1)) Then, we have
Σ E 8(w, dl & (C/A+ 2A IPlB)/P, = O.(C/A+ A). lE(m.]" T.
In particular, when P = Eie(m) ai |l00: Jm, represents finite number of neurons, we have
Σ
lay1 122 03 2 O.(5/++
KE(m.. JET:
where 5, = (wn, wx)) for j T.
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 Lemma F.6. Under Lemma 6, if the Hermite coefficient of a decays as O with some constant co > O, then the total mass near each target direction is large, i.e., u(T((O)) sign(a;) Z with large enough hidden constant.
handle almost all activation.
122l/2for for all i E [m.] and any delose 2 Ωi (
(4/))3/(ac 3/((2c0
In particular, for is ReLU absolute function, Selose 2 5i (2(1 ((4))) Here amin = min (as] is the smallest entry of a2 in absolute value.
As a corollary, if the optimality gap Ç Lx(11) [x(3)), then deoose 2 й. ((S + 12)1/(460-2)
and for ReLU absolute delose 2 Ωr (IS 22)1/3).
24

---

E11.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:) [x) (x) Ri (2) + Ra( + Ra( (this can be directly
verified noticing that (2)(2)) = | 1/ 2TT,
T
1 Σ
(i(x)
Σ JET,
ayw, w
sign(w, TT x),
2
iE(m.)
(2(x))
1 Σ
Σ
ayw
x (iig((w) x) sign(w; x)),
(8)
2
kE[m.]jcT:
Ra(x)
Σ icfm=
of 2 ll2
a 1| Σ TE ic(ml
22T
In the following we characterize R1, R2, R3 separately. In Lemma F.7 we relate Ry with the average neuron. In Lemma F.8 and Lemma F.9 bound R2 and R3 respectively.
Lemma E.7 (Zhou et al. (2021), Lemma 11), | 2 = 2243/m/2 Eie(m.) [EeT, asw, mn!
Lemma E8, Under Lemma 6, recall the optimality gap Ç = L2(H) [x(13). Then
| O.(1S/++ 2)3/2).
Lemma F9. Under Lemma and recall the optimality gap Ç = La(n) - La(e)). Ifōo = and σk > O with some k (11/A2) [a>ll;)), then
|| RRa 22 =Ō.(( + X2)1// + (C/A+ X) +
Now we are ready to bound the difference between average neuron with its corresponding ground -ruth neuron.
Lemma F10. Under Lemma 6, recall the optimality gap = Lx(R) - Lx(H3). Then for any i [m.). Ç = (2(22) and ç. < 1/ polylo,,,,.) A, 11a.)11)
1/2
Σ JET.
a un
Σ
Σ JET
-.((//))3/))
ajw;
ie
F22 From ideal loss solution to real loss solution
In previous section, we consider the ideal loss solution that assumes the norms are perfectly balanced (la.) = I| (1 and a, 3 are perfectly fitted. However, during the training we are not able to guarantee achieve these exactly but only approximately. This section devoted to show that the results in previous section still hold though the conditions are only approximately satisfied. Recall that the original loss
A A 1a12++ 2
Lx = L(0) +
IWI
so that when norm are balanced and ax, 33 are perfectly fitted, (x LA (n).
L(0)+ + + |a: |w:1/2 =
The lemma below shows that the properties of ideal loss solution previous section still hold for the solution of original loss, when a. f are approximately fitted.
Lemma F.11. Given any (a, W. a, (3) satisfying a12 O(S), ||g- O()),
=((),, where (11/22h En a; ||wlll22 and & -(112) Enh a,w, Let its corresponding balanced
à
25

---

version Obal = (Abal: Wbat: Obbal Bbal) as Obal,i = sign(a:) (a:)
w 2
Wball
W. /la
w
Olbal and Boot B. Then, we have
[x(0) (x(bbat) la - ar2 | 3-
(asll 1w: 0.
Moreover, let the optimality gap Ç (0) L2(1X), 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 (x) with the change of R3 in (8) as
Ra(
Σ #E/m.]
112
Σ xE(m)
+ - a (3 B) x. as a
2
The following lemma shows the norm remains bounded.
Lemma F.12. Under Lemma 6, suppose optimality gap Ç Lx(0) Lx(x)) Then || 3 ||a.ll1.
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 (())). We will assume Ç 52(x²) in this section for simplicity.
We first show gradient is always large whenever o, G are not fitted well. This is a direct corollary of Claim B.I. Lemma F.13 (Descent direction, a and (3). We have
T a))-) |V 4 a1lB
Before proceeding to the following descent direction, we first make simplification assumption that Assumption F.1. For every T., for all neuron w, E Tú assume w w 2 O.
This is because due to the linear term 3, the effective activation is symmetry )=)(1)) = d-2(~z) (-r)..
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 set of parameter (by flipping the sign of neurons) that satisfy the assumption and gives almost same gradient norm. Thus, making such assumption will cause any issue when a, 33 are perfectly fitted.
Lemma F.14. Suppose (a - a)2 | - =
| S to be small enough and ||a||2 |/ = O.(1).
Then, given any parameter there exists another set of parameter that satisfies Assumption F.I such that fo fg and ||oL>l| || S O.(vT).
Proof. Denote 0 (a, wT. Wm, a, 3). We first construct 0 = (ā, Wi,..., Wm. ā. 3).
Let à a. For W; there exists such sign vector s = (81....,$m) {+1] m so that by flipping
the sign of neurons have w. s,m;; satisfies Assumption F.I. Let + 0;30;-
and
One can verify that fo fø. Moreover, for the gradient of a, B we have
V.lA= VaLA, VoL, VaLx.
For gradient of a, wi, when s; = we know they are the same. When 8; 11, note that
V LA - Va, L = 2E (oz) x) - o x 2(3 - - ww; V w. LA V w. LA F., [R(x) (w, æx) + o )a) 2a, (33 (3)
Therefore, we get the desired result by noting the norm bound.
26

---

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 decrease the loss due the regularization term.
Lemma E.15 (Descent direction, norm balance). We have
ΣΣ JET;
(Vo, La, -a,) + (V LA, wyyl =>
Σ
w. 12:11
1E Lmm
max
Al ||
|W
Σ (laal VE[m_]
|| w. 212)
Lemma F.16 (Descent direction, norm cancellation). Under Lemma 6 and Assumption F.I, suppose the optimality gap Lx Lx(H)), For any w., consider òsign such that detose < òsign ((2/51/) with small enough hidden constant (dclose defined in Lemma F.6), then
E E ΣEA
02 VogLAN ΣjeT, (Saiin)) [ajl |105lll
Vw, Las EjET, (haiia))
w
(((2)).
|aj 1100
+E(+i-) (JETT)))ee
where T;, (8siga) = [j T; 8 w S bsign, sign (a; Ti.- (2.-(6sig)) e T.
5(w,, w
S x sign(a;) # sign(a;)) are the set of neurons that close to w with/without same
sign of a.
As a result,
I V 21²
Σ
layl 1|5ll2
-ETi.-
Now given the above lemmas, it suffices to consider the remaining case that B 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. Thus, could leverage these detailed characterization of the solution (far-away neurons are small and average neuron close to corresponding ground-truth neuron) construct descent direction.
Lemma F.17 (Descent direction). Under Lemma 6 and Assumption F.1, suppose the optimality gap < (0) Lx(H))) Suppose
norms are (almost) balanced: || - <// Elsim 1|0; Ox(52/2²)
(ii) (almost) no norm cancellation: consider all neurons w that are Saign close w.F.t. teacher neuron w. but has a different sign, ie., sign(a;) * sign(a²) with ōzign = OA(C(//2) we have Ejer, (Bsin) layl [| ||g I = (//6////))) / with small enough hidden constant, where (6) defined in Lemma F.16.
(ii)) a, B well fitted: Ja à -
O.),),
| 1 O.(s) with small enough hidden factor.
Then, we can construct the following descent direction
(a+ (a+aa)Vala (V + Σ Eivi, I. L,,, wy 26y ()) = S(())
xE(ma]JET.
where Oij satisfy the following conditions with delose Ösign and detase O.((1/3): (1) Σjer. ajıij a;; (2) qi 2 0: (3) di O when sign(a;) * sign(a;) or 8, Sclose- (4)
Eie(m.] Ejet, 2 (1).
F4 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.
27

---

Lemma 6 (Gradient lower bound). When Ω2 (AP) 5. and X S O1(1). we have
|eLall 2 /x²(
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 Eielm.)] |0 -1,,1,, -2/ by Lemma F.15 we know
(V, -a. + LA Lx, 22,
JET:
With Lemma F.12, this implies
+ 0(11a~11)2 VIVoL>11 11 | ||
= -2,(2/A)),
which means
VOLAII 2 | + V L 2 (5((2)
For assumption (ii) (norm cancellation) Lemma F.17, whenever does not hold, by Lemma F.16 we know
VoLal 2 Vaball V 21²
Σ
(5
JET; (B.ig))
For assumption (iii) (a. (3 in Lemma F.17, whenever it does not hold, by Lemma F.13 we know
VaLa22 (a a)² = (2(())) |V>4ا112 = 4 1 B 2 = 92 (C²)
which implies
V V. VoLA1l2 (22
Thus, the remaining case is the one that all assumption (i)-(iii) in Lemma F.17 hold and also we choose
Eielm.]
af
|| 12 OC
a,47 L (Sstor))
if TT,
46
otherwise
so that condition (1)-(4) on di all hold: condition (1)-(3 are easy to check, Lemma H.4 shows condition (4) holds. Now know from Lemma F.17 that
(a+a.) + (VaLs, B,) + Σ
Σ
(Vw, LA, wi () = S2(C).
xE[m.] JET
Note that
(a + a..aA (Vgl,, Σ
V Wm, LAx W0 - -gi
iE [m..) JET,
V Tex
(a aa+ 113 B.112
Σ
|wo 2:
[m_.] JET.
and
la + a.. < |ál + |a.1 O.((5) 8 0.11)
||3 + 1xll2 - | + 1|3.112 + +O((S) S O1(1)
(b)
(
Σ Σ ||m u 11 2 Σ Σ
1 12
2
|
$ O.(1),
vE(m.]JET.
ve(m.) JET:
28

---

where (a)(b) by Lemma F.4; (c) we use Lemma F.12 and condition (4) on 9ij.
Therefore, get
|V |V VaLLll VaLAll |V - 12n(5)).
Combine all cases above, we know
|Vaaall2 VW (min 20/0,, 12(())12))
as long as Ç = 0 poly m., A, |a+ Amin)).
We now use Lemma 1.14 to show when Assumption F.1 is not true, we can get similar gradient lower bound. Denote the above gradient lower bound as To = Q2 (5* />²). Let = 7o/2.
When (a - à)² 2 r or 3 3 | 2 T, from Lemma F.13 we know / 2T.
When (a (a-a)² 3 2 T, using Lemma F. 14 we know there exists @ such that | 2
and - S VT. Thus, know || 2 0.1r.
Therefore, combine above we can show /VoLlle = S2,((\>2))
G Non-degenerate dual certificate
In this section, show that there indeed exists non-degenerate dual certificate that satisfies Definition and therefore proving Lemma F.1.
Lemma F.1. There exists non-degenerate dual certificate 1j Ex ( æ)] with p. = ((1) and l|pl|2 & poly(m.,,)) A)
Recall that we want to use the dual certificate 7) to characterize the (approximate) solution for the following regression problem:
min (A( = E. (f((x) - . + A/p11 = Ew (( x)d jddd p (_.) + Alp|1, 1)
27
where 022 is the ReLU activation after removing Oth and 1st order (corresponding to and 8 terms) and pa Lie(m.] at 8w; the ground-truth.
Notation We need to first introduce few notations before proceeding to the proof. Denote the kemell K2r(w, u) Ee |0 wh x 2) x))
as the kernel induced by activation 02r(2), where (2)(1) Enze δxh (z. Zr = || - Exze = ((433/4 is the normalizing factor, he (z) is the normalized k-th (probabilistic) Hermite polynomial and o₂ is the corresponding Hermite coefficient. We will specify the value 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 E):
n)w)=
Σ jE(m.)
on K (wo
w) +
Σ V VKKKW, w)
jE[m.]
such that it satisfies
n sign and Vn)(w)) for all [m.].
Here 0i (aI.... am.) ) E m__ .22 (a (.... ..... a_mm) )T mm. are the parameters that
are going to solve and V, means the gradient w.r.t. i-th variable (for example, V,K (xx, y) means
gradient with respect to x).
One can rewrite the above constraints (9) into the matrix form:
[x1
r
b.
(10)
022
29

---

where b (agg(aj)... sign(am..). ). mm) E lmm(d+t), r Beh()))<=)]]
(m.(d+1)xmm(dd 7(x)
(727 w .),..., 22 wm. x) V Vwo wi x)
V T) E Rm)(d+))
Here V T Ze x)
Pa Pz ( x)) R', where Pw; is the projection matrix defined
below.
Notions the unit sphere As could see, the kernel K is invariant under the change of norms, so it suffices to focus on the input on the unit sphere gd-1, On the unite sphere, we could compute the gradient and hessian of function (ww) on the sphere (e.g., Absil et al. (2013))
grad ((w) = Pw V f(w),
H [[(w) = P.. "VV [(w) " ()))))= for all tangent vector that z =T w = 0,
where P = www is the projection matrix.
Then, we could define the derivative as Poon et al. (2023); Absil et al. (2008); for tangent vectors 2, z'
Do f(w) = f(w)
D1 f[w)[=] = (z. grad f(w)) = =x T V
D2 ([w))==) = (H f(w) x²) = P (V f(w) w - (f(w))))
and their associated norms
[Di f(w)ll == sup
D1 [[w)[z] P f(w))l2
|==|,.=
[D f(w)ll
sup D2 D2f(w) x11 Pw H f(w P 1
where ||=|| = P
For simplicity, we will use K(6)) (w, u) to denote v V (w,, u). One can check that this is in fact the same as the one defined Poon et (2023) under our specific kernel K, i + j < 3 and i, S Let
( k (w, u)
KLO u)
(1)
sup
w.u
- IL -
||_ =1,
2 Vpē(i).4Elil
where 2 applies to the dimension corresponding to w and similarly 2u. for u.
Before solving (10), we first present Section for the proofs). The lemma below shows that kernel K K(w, u.) is non-degenerate in the sense that it decays at least quadratic at each ground-truth direction (w Rs wx) and contributes almost nothing when w, u are away.
useful propricties of kernel K that will be used later (see
Lemma G.1 (Non-degeneracy of kernel K). For any 0, let 0( - (log/=/ kernel Ket deggeereee in the sense that there exists r ( following hold:
((1). P2 O()) such that
(i) K(w, u) s - for all ō(w, u) == L(w, u) 2 r.
(ii) k<20) u)(= x] S -P2 1| for tangent vector that z w 0 and 8(w, u)
(iii) ([ (un). wn ,00; S h/m2 for (i E {0, 1} x (0. 1,2)
The following lemma shows that K and its derivatives are bounded.
Lemma G.2 (Regularity conditions kernel K). Let Bij = supp.. || u) ul)l and
Bo Boo + Bio + 1, B2 B20 + B21 + 1. We have Boo 0(1), Big 0 Bn O()),
B20 O(), B21 ( and therefore Bo O(21/2), Bg 0(13/2).
30

---

The following lemma from Poon et al. (2023) connects the non-degeneracy of kernel K to the dual certificate 7) that we are interested in.
Lemma G.3 (Lemma 2, Poon al. (2023), adapted in our setting). Let {+1). Suppose that for B-1/2 we have: for all 6(w, wo) and Rdd with zTw O, holds that (02) (Wo, w)[z 2, 2] > p 21/ and [|]] (wa, w) S B. Let r be smooth function. If 17( wo) = (n)))) 0 and a D2 (ww) (2) (wo, w) I T for all 6(w, wo) S r with 2/2, then we have (n(w)
some p 0, B 0 and
((p 2T) ())6(()) wo)? for all 6(w, wo)
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), can show that a Rs sa..(a.) and Or Ra O and therefore we can transfer the non-degeneracy kernel K the dual certificate 7) with Lemma G.3.
Lemma F.1. There exists non-degenerate dual certificate 1, 22 x)] with pn = + and ||p||2 S poly(mu, A
Proof. Note that Y SDYDS, where
Im.
Im.
Pwt
(Zo I
D
S
Pwi.. wwi..
(Zo
are block diagonal matrices, T- (x(((()))(())) R. mm (
7(x) (ozz( ....... . Oze(w' (Zo { o2.k Zo o ((4-1/4) is the normalizing factor so that the diagonal of T are all 1.
(w)) x)x
/Z jozz (wi,, x) 7) E Rm-(d+1),
Thus, solve (10), is sufficient solve the following: denote K DYD
b,
(11)
and a1 = Q1. a22 )Z to get the solution of (10).
In the following, we are going to first show that K R: DD because all off-diagonal terms T small due to Lemma G.1 (iii) (we choose h to be small enough, and will choose it later). Specifically, we have
K -DD = sup
1 (K - DD)
|=||==
sup
Σ i.
21, (wj w3)21. + +2(Zo/Lo)
31, V1 ([K(w,,))) 22.3 Σ
|
i,j
(Zo/2or)) )2
Σ
2, V V2K V2K(w, wo) 22.j
[(w] w1i1 113-1/2) (x(10) (w0
(w, | w/)l + + ( (1)))
S 2h,
m 207
where (zī TT) 21 (21,11..... 21.m. T and 2 (2 .............. 22. has the same block structure as (ai1 (2) and we use Lemma G.I and Lemma G.2 in the last line.
Note that DD has exactly m.d eigenvalues of and m. eigenvalues of and K also has m. eigen- values 0. By Weyl's inequality, we know hy-11 < 2h where K EiEi df Yo v. is its eigende-
composition. Here v. Ta 0 for all TT V1 (o w. 0....,00 0)T (O...., ... 0, w "_m) ]T
31

---

in the null space of D. Since Tr = 0 for all Ta E V1. we have
α1
Ktb-
Σ u. b
(7 - 1),,0 b+ E a
sEmmad]
iE iEmm.d]
iE|m.d]
(7
10, b.
d
Therefore,
(% - 1)
S max his 111mm, =
=(hmmmm) =1(.-mmmm00//
=: h².
2
HE(m,d) d]
This implies ||o sign (a.) = |ā - sign(a.)ll S Ri |a = || l1 < 1 + / and ||e2l|2 = (Z./ ||ă2,ll2 I ((h'E
Now, given the a1, a2 O and O(A log(m./ A)) so that the condition in Lemma G.1 holds.
show the corresponding 7) is non-degenerate. Choosing h
Consider w Tr, when o(w, w;) ( using Lemma G.1 and Lemma G.2 we have
[n(w)]
Σ
anjK(wj, w)+
Σ
02 V K (woj, w)
jE
jE(m.)
J01, [(w3, w)) +
| m0 || VVKK(w w)) wj
jE(m.]
5(1 h1)(1 P1 h) O 1/2 (B10 h < 1 - p1/2 < ((p)) (p)))((w, (w,
where we choose h O 1/2) be small enough.
When 6(w, w S (12)) again using Lemma G.1 and Lemma G.2 we have
at D2 n(w) k(02)) ((02)(w; w)
01 K
(202) (02) (w) w) K(02) (w0 w]
Σ
|n K (02) w) w))l + Σ
|
k(12) w)
10
jE(m.]
20
<h' Buz + 0 + h) S P2/16,
where again due choice of small h. Using Lemma G.3 we know that In(w)] I 1 (pz/4)6(w, )
Combine the above two cases, have In(w) 1-)())(( w and n)w)) E (( xx)] with
p(x)=
1
Σ
az x)+
Σ
22
T (I w/ (w, x))
zz
jE(m_]
jE(m.]
We have ||p|| = O(33/4 m, + m. m. 3(A23/2
H Proofs in Section F
In this section, we give the omitted proofs in Section
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
32

---

Lemma F2. Given non-degenerate dual certificate then
(i) (n.H²) = 1p
(ii) For any measure μr E MM(S 16, S 1u/1 - Py Eie(m.) ST. 8(ur, w dl (l(w).
(iii) (n. F 16)) (p, - fur), where fu(x) = Beoo 122 x)}. Then Km. 1²)) s
lll12 ((u).
Proof. We show the results one by one.
Part (i)ii) We have
10) S ( [[lw))
Σ / |n(w)) |el1 Pn Σ [m.]]
/ o(w, wf)² [lw)..)).
iE(m.]
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 (iit)
We have
(n.4 - 11) ( v((w) d(u ((a))d(("""))) (w) = Jeu Ex ( (x))) d)) - 1 (w)
=Fn p((x) / g/-- 222(10) (TT))) d(u - ,² (wv)
=B (p():)(() Sn- (x=))
Note that L(n) = I1fa - fa- 112- this leads to (17, 1 - n²)l s lllll2 v L(p).
Given the above lemma and the optimality of px we are able to characterize the structure of n as below: norm is bounded, square loss is small and far-away neurons are small.
Lemma F.3. We have the following hold
(1) 11 A 1 - 1p11 $ 14011 = =Iaa"li
(ii) M(x & X2 1p 0.(22)
(iii) Eie(m.] FT. 8(w, w d|uil(w) &
O.(A)
Proof. We show the results one by one.
Part (i) Due to the optimality of μ² have
n(xE) + Alp5l1 = LA(U)) s Lx(u') L(x²) + Alp"11-
Rearranging the terms, have
AlRSl1 - A111111 & T(p²) - T(13) = -I((() s O.
For the lower bound, with Lemma F.2 we have
0 < || - ||0) (n. 1² n²) < |1t11 - |,=1 + Illlla / (((3).
Using part (ii) get the desired lower bound.
33

---

Part (ii) We first have the following inequality due to the optimality of μe and adding A17.13 -1²) on both side:
T(332) + X(1x/h - 11111)) - X(n, 1z - -1) S I((²) - Xmn, - 10))
For (I), have
(1) = A(1231 - (n. x2)) +
- -11m"1""))-)" 2 0,
where we use Lemma F.2 in the last inequality. Therefore, the above inequality leads to
L(u)) s L(p") N/n,u n²) S l|lll2 v ((MR),
where we again use Lemma F.2. This further leads to L(13) S X² 11
Part (iii) Using part (i) we have
|eil - 1A (n. -²) S -m - )()).
With Lemma F.2, LHS and RHS become
LHS =Ip|| (n Pn xE(m.]
Σ
/ 8(w, wx)² dl
RHS < Ipllz v L(N)).
Then using part (ii) we have the desired result.
We are now ready characterize the approximated solution by comparing μ and p Lemma F4. Recall the optimality gap Ç Lx(x) - Lx(U))) Then, the following holds:
(i) L(R) < 5X² 1pll + 4Ç- = -0. + C)-
(ii) ifC A1p²11 and 1/*/1/ [1ll22 then 1u/2 S 3/u"11 3|
Proof. We show the results one by one.
Part (i) By the definition of the optimality gap S and adding - -X(n, pa n²) on both side, we have
L(p) X(14) - (mXh1) A(n. - n²) S L(N)) + - A(11 2 n²))
Note that on LHS.
A(11 1x(1) M07 μ μ²) A(1m) (n. + 1111==1~-1Mx11 2 0,
where we use Lemma F.2 and Lemma F.3. Therefore, with Lemma F.2 and Lemma F.3 we get
L(p) S M(x + X(1, u n²) s X2 1pp + + lllllz / VL()).
Solving the above inequality on L(μ) gives L(n) < 5X2 ||p| + 45.
Part (ii) Again from the definition of the optimality gap Ç, we have
Alad1 L(xZ) + A1pT1 + + L(x) < x3 |p| Alp*l +
where we use Lemma F.3. Thus, |4|: S ||r1 + |p*1 + </> < 311-
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 μr and nz and give bound for it.
34

---

Lemma F.5. Recall the optimality gap Ç = L2(1) - Ex(1))- Then, we have
Σ Σ
O.(C/A+ A). J 8(w, (0; d(u(((w))) & (C/A+ + 22 l, =
xe(m.]" T
In particular, when J = Eve(m) a4 Ilwaall2 ,w. represents finite number of neurons, we have
Σ
Σ
layl |105|12 s 2x 0- -O.C/AA
ve(m.) JET:
where 55 = (ung, wx) for j E T.
Proof. By the definition of the optimality gap C. have
L(p)= + Alpl1 L(n + A|a + 5.
Rearranging the terms and adding - (n n²²) both side, get
|x| |لä| - (n, / - , = x (1(x))) L())
+ + =())+
(11 u p¹).
For LHS, with Lemma F.2 and Lemma F.3 we have
Σ /
LHS |x/| (n p)) |xā1a + | Pm 6E|m.]] T.
6(w, wi) (lw)
For RHS, with Lemma F.2 and Lemma F.3 we have
1
L(H)
RHS S I (AP 1e11 L(x) + +) + llll22 VL(u) = A A11e1
A lllll v L(u).
When L(u) > X2 1pll we have RHS < </> + A |Iel When L(M) s X3 I1ll2- we have RHS <
</> + 2A 1112. Thus, in summary RHS S Ç/> + 2X ||112
Combine the bounds on LHS and RHS we have
Pn Σ
/
8(w, w dlw((w) S C/A+ 2X
iE [E/mm.)"
The following lemma shows that every teacher neuron must have least one close-by student neuron within angle O. (ç1/3).. This generalize and greatly simplify the previous results Lemma 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 a decays as lōxl fxl==((k-0) with some constant Co > O, then the total mass near each target direction large, i.e., (x(x(5)) sign(a;) Z with large enough hidden constant.
|atl/2 for all E [m.)] and any Sclose" Z 5.
)2/(4co-2) >//44)))//
a min
In particular, for a is ReLU or absolute function, 8cosse 2 á the smallest entry of a. in absolute value.
(5((()) (44))11))) Here amin min |a is
As corollary, if the optimality gap Ç Lx(1) - Lx(3)), then de 2 σ. ((S + (2)c/((40-2)) for ReLU or absolute delose 2 Ω. (IS + X2)1/33.
Proof. Assume towards contradiction that there exists some [m.] with some Setose 2 Φ ( with large enough hidden constant such that μ v(T(()) sign(a;) < |af1/2. For simplicity, we will use 5 for detose in the following.
Let 9(x) Ercke2 sign(a;) sign(0&) (ogn@t)he(w""" T æ) be a test function, where (x( is the k-th normalized probabilistic Hermite polynomial and f will be chosen later.
35

---

Denote R(x)= = fu(x) - fu- (x) so that | L(μ). We have
v L(m) 119112 2(-R,g)
apo(un) x:) Sra olw ((w")) du((w) g(ar)
= Ra
Σ Js6i
aujolwiTan x) (gd-1(1(6)) aww T x)) du((un) g(ar)
Rez
Recall the Hermite expansion of oxx) Exz0 (x))((r) and its property in Claim A.1. For the first term, it becomes
[a;][aul - - TT() 1& sign(a)) (w du(ww) 2 ag
Σ l
ESk<2e
For the second term, it becomes
Σ
Σ
Cajlāsl (ign(a;)(w)) w/)x -B--1
|ōt sign(a)) "w) (w;)) du(w)
1SK<28
<(1aa1), + (1111)
1ōkl max (w T wx)
<(w )28
ASk$2e
<(Ia"1] + (m/1))
Σ
1ōx1(1 82/ /5)e
&
<4||a"|| 1I, (1 82
|ox| S [ Σ lōx1.
tkk<2
ESk<2E
where (i) in the third line cos 5 < 1 82/5 for
[0, x0, and (ii) in the last line
Lemma F.4 and choose = [ log( 16 ||a||l111111// /|a,1))
Thus, given |ōv| O=)- we have
1 |aš1
LI()Ve = ((()) |1all2 2 4
|ō
4 |0;l
Σ
la
PSk<28
With the choice of = e(1 we have 5 = δ
L(a) 10
1/(10x-2)
Since 5 2
Lud 2/ with a large enough hidden constant, we know this is a contradiction.
As corollary, with Lemma F.4 that L(u) 4C + 5X ||pl Ω (#C+5*PIpPlB 22/(cc0
we have 6 >
4nmin
For the activation a is ReLU or absolute function, by Lemma A.1 we know Co = 5/4, which gives the desired result.
The lemma below bounds R2 using the fact that it is spiky (has small non-zero support). Lemma F88. Under Lemma 6, recall the optimality gap Ç = Lx(u) Lx(1)). Then
|R (1(//A++ + 2)3/2).
Proof. Using the same calculation as in Lemma 12 in Zhou et al. (2021), we have
1/2
3/2
R2ll (m..)
Σ
Σ
layl
w 11a 82 lasl
TE(m.]
JET
TET
With Lemma E.4 and Lemma F.5, have Ra 12
Om 1 + 2)3/22.
36

---

The following lemma bounds R3. In fact, in the view of expressing the loss as sum of tensor decomposition problem, R3 corresponds to the 0-th order term in the expansion. would become small when high-order terms become small, as shown in the proof below.
Lemma E.9. Under Lemma 6 and recall the optimality gap < Lx(n) La(U)). 1f 0 and ōk > 0 with some k (((1/ log(S/ |a.ll,)), then
|| Rs l2 -0.(( + 22) + (c/>+ + +))
Proof. As shown in Ge et al. (2018); Li et al. (2020), we can write the loss L(p) sum of tensor decomposition problem (recall | 1 = 1):
L(n)
Σ %
J w dp(w)-
Σ of
||w0
l12 w
&k
k20
iE[m-]
F
Thus, know for any k > 1,
w ok du(ww)
Σ
at |1 20
S L(4)/66
-
iE im(m.)
F
Given any w and even k, we have
S du(w)-
Σ of
|ww w oot
/wssa-.
xE(m.)
/ wegga-
Σ az ||wn w0
w @ du(()); w 8K
[iEmma]
az || 1l2 S (ws, w0s dpp())
Σ az /1willl (wo w))h Joe 1VTs Knf
(w, wgyh du(w)
a) ||000 112 S,p apleo) - -11 du(ww) f (un, woih du((w))
Σ of | 2| (ww w7)h f 18A-1 (w, with du(w)
We show the last terms are small.
For the second term on RHS, we have
Vg du((w) 1 (w,, w dp(w)
(2) / 1- 1-(1- 10-1-1100 -d(w m33)//2 dlp)( S [ (1 (w, wn)" dlu)(w) S
T5
T,
(b)
O(k) 8(w, w dlu((w)+ + f J 17..00m.11333311 dlu((w)
<O(k) / 5(w, w (d),,))) dlu((w))
where (a) cos 8 2 1 - 82/2 for 8 E [0, T/, (b) (1 - -x)* 2 1 kr for x E [0. 1].
For the third term on RHS, we have
o || ll2 (w, w2)² { Jo (w, w² du(w)
< <(1|a.11 11: + +|ml1)
max
(w T wx
/wm,0))22/2
(0)
(||a, 1l, + (1 - -2/110/ < 0()),
37

---

where (a) cos 2 1 82/5 5 for 5 E /0, (b) we choose A = @((1/3)) (|xxl1)) and Lemma F.4.
Therefore, we have
w w ((()))
Σ af || wo
ic(ma)
F
f u(w))
- -O()) / 8(w, (or,eww)))) (u)))
2
T
This implies that
m. (Tu))/x
2
Σ
aj ||011122
L. u(w))
-O(k) Σ
/ 8(ur, (1))(((())))) O(m,s) jE(ma]]
jE(m.]
{
Σ [E(m-.]
a
|
ll2
p(w)
-O(((/+ - - O(m.S);
J84
where we use Lemma F5, Rearranging the terms and recalling L(n) = O+(S + X2) 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 R1, R2, R3. respectively.
Lemma F.10. Under Lemma 6, recall the optimality gap Lx(R) Ex(H)). Then for any
i (m.). < 22(2) and 5, X & 1/ poly(m., ||a. 11,
1/2
Σ JET:
ajw, w s
Σ
Σ JET
w
w
=O_((//A)/4)
(E(m.]]
Proof. With the relation of residual decomposition, Lemma F.7, Lemma F.8 and Lemma F.9, we have for any i [m.]
2 1/2
2/43/2 0(A3/2/m222)
Σ
ajmn
w.
[ [ + Rallz + || Rall2
:E|m.]
JET.
=O+(( + 22)11/2 (S/A~A)3/7 + ++(((+
+
+ 23)1/2 + (c)A+
Rearranging the terms, we get the result.
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 a, f perfectly their target values only decrease the optimality gap.
| a| =0(),
Lemma F.1.. Given any 0- = (a, W, a,, satisfying la a/2 O((),
O()), where = -(1/V2m) Em a; |luzall2 and = - (1/2) 2 a:w; Let its corresponding balanced version 0bat = (Abal Woat: Olbal, Boat) as Obal,i = sign(a; (o:) v [ail I|: Tly Whal.; w. /Tay 1 O'bal = and Bool G. Then, we have
(+) I2 (fbat) la ài² + ||3
Σ (las) l10s112)2 2 0.
ie(m]
38

---

Moreover, let the optimality gap Ç (x Lx( we have results in Lemma F.4, Lemmna F.5, Lemma F.6. Lemma F.7, Lemma F.8, Lemma F.9 and Lemma F.10 still hold for (x) with the change of R3 in (8) as
R2(
1
Σ ic(m.]
of 1kwil2
Σ as 11ws ic(mm
ta a (3 ()
22T
Proof. Recall in Claim B.1 we have
L(0) = ---a=2 al2
Σ
a
20,
Σ a W 112 w0
00k
R22
rE[m]
ie[m.]
Note that [a:| 1w:ll2 = |abat,al Wbbat, i ill2 so that T((0) T.(Obat) + 1ax 0r/2 |.3
3
2
We then
have
X
A
(0) (2(bl)) =la -6
a 222
|
|
11
W bat
2
X
=/a à:2 +
3
(1as) 11277 100:l122
ic/ml
Therefore, we have the optimality gap that Obal corresponds to network that has perfect balanced norms and fitted a, B, thus all results in Lemma F4, Lemma F.5, Lemma F.6, Lemma F.7, Lemma F.8. Lemma F.9 and Lemma F.10 hold for Obat. Since Ç 2 Sbal. |a:1 ||/ll2 = josat.il Wbot,sll2 and L(0) = L(Obai) + O()), we can easily check that all of them also hold for For the bound on R3, note that
Lx L LA(ut) 2 Lx LA Lx(x)) Shat- Note
Rall2 S
Σ af 1wtll2 AE(m..]
Σ allslll xE(m]
+ +la - -al||+|/B - -011 -
V2T
so that the same bound still hold for R3.
Lemma F.12. Under Lemma 6, suppose optimality gap Ç (x(0) Then + S
3 |lalli--
Proof. We have
||a
W
= AKbl L(0) s < X² 11ll2 A/pil1
where we use Lemma F.3. Rearranging the terms, we get the result by noting that | S ||ax 11
H.3 Omitted proofs Section F.3
In this section, we give the omitted proofs in Section F.3. We will consider them case case. The lemma below says that one can always decrease the loss if norms are balanced. Lemma F.15 (Descent direction, norm balance). have
ΣΣ (Va Lx, -a,) + (V LA wo =x
Σ of
JET:
BE mm.
max
>/ A|||a|/
Σ
(1a.1 - /|w:lbb) 11))
Ve (m_)
39

---

Proof. We have
Vay Tax; aa)) +V L,, w,n
vE[m]
Σ
20 (z) f.(x))aa, a)] - Xa + + (((()) - [a)))) æ]] + X A| 1
vE[m]
=X
a3 ww
Note that |ad
+ 1 1| 2 ||as| - || we get the result.
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.I, suppose the optimality gap Ç = (f) (x(()) For any w,, consider ōzign such that òelose < Ösign = ((2/ç1 with small enough hidden constant (ōclose defined in Lemma F.6), then
E Σ
Vo, LAi Ejet, [a [w 1l
0j
V Ww, LAs EjET, |a 20,
w
Q(A)..
where Tú (5siga) {j T; : 5( 8(w;, (,) daigu: sign(a,) sign(a,) }. Tú (8sigu) {j T3 5(w,, w S Osign, sign(a;) # sign(a;)) are the set of neurons that close to w with/without same
sign of a.
As a resull,
V IaLa Lw 212
Σ
JET. (duian)
Proof. Denote R(x) f(x) - (x)).
We have
Σ
Σ
aj Vaglix E EjET,
Vw,, Lan0 EjET, lagl |100
w
lay 12
|ayl |15l1 aj |20
1
Aa
Σ
Σ
2Bx(R(x)) x)] +
Ejet,
Ejeri.(Buan |0,1 120/l1 ll2
<E(+,-) JJET:
ag || sE(+,) (EE. ESETi layl
EE
Σ
Σ
1
2E (((r)oo) x)] EjeT. (Brien) |ay)
All20y 11
w1yl2
We split the above into two terms (depending on square loss or regularization). WLOG, assume sign(a;) 1. For the first term that depends gradient on square loss
||w llz
44 Σ
Σ
(x( x)
(SET.. |a
|0,122
sE(+.-) (6ain))
layl |100
Σ
Eg [ x)]
EJET, (Brin)) lay1 |,00
lagl 20, 11g
4
Ex (x x)]
Ejer, (drna))
layl
1w5ll2
jET. (Snign)
layl |205112
-4 Σ
E, [[(( T x))]
ΣJeT, (Boiga)
layl |w0112
jET, (6xign)
lagl |0,,, l
-4
Ex (OT a a
x))
Ejer, (Done))
lajl
[w 1l
jET. (Ssign)
40

---

Since w is ósign close to w; and R = L(0), we have
(()) (ōsign) || O. (1 where we use Lemma F.11 that L( = (c)..
For the second term that depends regularization, we have
EjeT, oz + /wy112 2A + 2A -A1
(1)) Σ SE - Σjer. (Boign)
[ayl W5ll2
Therefore, when (I) S 2A, ie., deign =0x(A/51/2) we have
ΣE
Σ
E
VogLi1 EjeT.o ag |2001
sign(a,)lag)
Vw, Lx,
w
EjeT. (Faiaa)) la 200
sE(+,-) JET (6xign)
EjeTe. (Brgg))))) az 1105112
Σ
EjeThh (drisn)) [gyl 11il22
+E(+,-)
We compute a upper bound for LHS. Note that
Σ
Σ
Vos EjeT, (Miia)))
02
Vws Tax
w
ag
2 11
ZjeT, (daian)
|a 1202122
Σ Σ
V w. Lx } (Va, +
Σ
0
+
w
EjET (Biin)))
|رyا |w, 1|,
(E(t.-)JET,.oe
NE(+.- JETra (Gsign)
EjET,, (driia)) a² + |20 11
VoL 12 \
Σ
JE: la w l1 )2
sE(+--
ΣjET, (6 (om
0z + [20) 11
SVIVoLAl2 V
Σ
ESET. (Fmina)
lay1
|2,, ||,
v E E |ayl 11w
xE(+.-)
where the last line we use Lemma F.6: ESET. 12y1 [0yl22 < EJer, (Bmgga) (Urla)) lay 1e0 because (T) = EjeT((Engn) aj ||00 > 0.
Combine with the above descent direction, we have
V V LW
Σ
EJETr. ag + [
1
Σ ΣJeT, layl 10yll v EEET, (Gmgo)) lay1 ||
sE(+.-} }
A
Ejetho
a² || 112
Σ
2
E)er. layl 2 l1
sE(+.-)
which implies
2X²
Σ
layl ||1,
JET (deign)
The lemma below shows that when all previous cases are not hold, then there is 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).
41

---

Lemma E.17 (Descent direction). Under Lemma 6 and Assumption F.I, suppose the optimality gap
Ç (0) Lx(r)) Suppose
(i) norms (almost) balanced: [[]]] ||ll21 S </> Eieia |1, 1 0.(C(/29)
(ii) (almost) norm cancellation: consider all neurons w that are òsign- close W.T.f. teacher neuron W but has a different sign, ie., sign(a;) + sign(a?) with ōzign (2/C122) have Σjet, - (6mig)) |ayl [|w, 112 S = O, /X) with small enough hidden constant, where T (6) defined in Lemma F.16.
(iii) a, are well fitted: a-- | 0 =O(() with small enough hidden factor.
Then, we can construct the following descent direction
(Val) + Σ Σ
vv, La, (6500,) = 22()),
[E(m.]JET.
where 9iş satisfy the following conditions with delose ōeign and òeiose Ou((1/3): (1) ΣJET. asqj a; (2) quj 2 0: (3) qtj when sign(a,) + sign(a;) or 8g delose- (4)
Eie(m. EjeT, =
Proof. Recall residual R(x) f(x) - (z)). We have
aa+ + (Vgs, Σ EVvu, (),w7)
iE(m.]JETi
(9) ZE [[(x)
2kx 33) 3.)) xc]
Σ
[R[[()) - Σ Σ E. (1 x))
iE [m JET.
iE|m- JET.
2 Σ
[x( (w,; x) - o' o' x))]
:E|m.] JET.
Σ |w,, 10,
Σ
Σ
W w
xEfm)
iE [-m JET.
(b)
2|
2
+ X
W Pr x Σ
Σ
w :
selm. JET
+2 Σ Σ
En (a
al (w T x)) - o'(nn] x))
1JET.
(c)
+ 2
-
Σ Σ no |w5
tElm.]jET,
+2 Σ Σ
Eg (((x)) w x=(o (wo x) - d' [w) x))]
(12)
where (a) we plug in the gradient expression and add and minus the term 2 Eie((m. ΣjEt, Ex R( æ)]; (b) rearranging the terms; (c) using (x
||R||2 + (A/2) W F (A/2) |aal = Lx(H3) +
42

---

For the first line on RHS of (12), we have
1x(33 + + 2 W -1aa12)) - Σ
Σ Di 10y12
i JE(m,)JETs
@ 2//2 + I(M3) + A|g/l - A Σ
Σ
Øi |20y122
E mm. JET
(b) 2C/2 + A143l A|la-111 +1
Σ
Σ
co([ay] - (2,)12)
sEfm.]jET.
1/2
1/2 (d)
3 2 2/22-0(13)- -
ΣΣ
Σ FE(m]
- -1aj1-1wssly 122
2 C/4,
jE|m-]JET
where (a) due to assumption that norms are balanced; (b) we ignore ((2)) and add and minus ||: (c) due to Lemma F3; (d) due to assumption that norms are balanced and the choice of quz-
In the following, we will lower bound the last term of (12) to show it is smaller than -C/8 so that we get the desired lower bound. Recall the residual decomposition (8) that R(x)) R1 (2) R2 (x) + (x), we have
Σ
Σ
Ex [R(x) wi alo'(w) x) - d'ww, x))]
ielm. JET:
Σ
Σ
Ex
T wlo x) - o' [w) x))]
iE/m.
Σ
Ee ( (wo x) - d'wj] x))]
E EE
(I
E.g ( of al (w)Taa) x) - o' d'(w] x))]
E JET.
(111)
Bound (I) For (I), recall Rí (1/2) Eie(m) v. sign(w; T ax), where w; EjET, aywy - w is the difference between average neuron and corresponding ground-truth and (Eie(m)) |/1 11/12//22 ((S/A)8//) from Lemma F.10 and Lemma F.11. We have
Σ
Σ
E.n Ra (z: ( aaa'(ww" Tx) o' dmww] (w]]))]
[m. JET:
(8)
2
Σ
Eee (o al|ayayy||w0? x0 sigu(w) #signn(w; x))
(E(m-] JET, kE(m-)
(b)
Σ
lagass 11 E& [ कl [wo Ø sign(w)] () ()ssiin(( a))
tE (E(m.JJET kc[m..
Σ
la5ası| | 1, 1 sign((w) T) siignn
[E[m.]] JET: kE[m.]]
Σ Σ
Σ
lagassl |/ 0(53)
+E(m-. JET RE(m-)
(_ 52 (
Σ
Σ
la54yo -..((/A))3// 22
:EIm.
43

---

where in (a) we plug in the definition of R1 and using the fact that wt (w;
x) - [w x))
|ws æ: Il sign(u]] s)ssigg(( az) (b) क is 3-dimensional Gaussian since the expectation only depends on Uk₁ w wy; (c) w z
dj || - when sign(w; Ø) sign(w;" 3); (d) direct calculation
bound as Lemma H.2; (e) definition of qij-
Bound (II) For (II, recall
(2( - Σ Σ
ayw mssign(w) æ)) signns)) x))
Σ
Σ
ag /w0 ae sign(w] a) siggnni mn))
ic(m.) JET:
iE[m.] JET:
For each term in (II) with j E Ti, we can split it into two terms that corresponding to Ti and other Tg's.
Bu ( w [(o'(w"" x) - d'(w) (w]xx) x))]
Σ
Em [as] [a ml signtw] "wig a) Oy aro'(un? x)) - o' d'ww] x))
kE(mm] TETr
Σ asdgThy B.n mm a0 Il sign(w? x) (w w) / xc sigu(wo; w))ssign(() mn)
KE [m.] RET
araj Be 12 x lwo x sigu(w] a) #sigu(( w) sigul( [) m))
ET
(11.6)
aca F. wu a we// x sign(w,; (wo Ta) 1 sign(to; T Trr (w)]) (13)
Right TET
(11.58)
For (II.i,, we further split neurons into Ti (6siga) and others:
(II.)) Σ
ara E.
[ur a al|l0, ax| sign(w, #)ssign(w; a) aigutt;; ()ssign(( mn)
(ETi(
Bml a: al|w; sign a)sssggn(w mm) (ing ()siign(( m))
(14)
Consider the first line of (14), from the choice of dzy know aydyya 2 0. For E T (Fsign)
know sag)(a;) sign(a;), (a),, which implies aęa qtj 2 0 for these terms. We thus only need to deal with neurons in T7, (8aign). we have the first line bounded
Σ
ara,y.y E.. Ee ae| w x sign(w] u)) #sigu (w;; w) sign(w;? (siguu( æe)
Σ
agag Em (1w, x=||0 ar sign(w) siggu(w;; To Tw) sign(w;" =))gnnn(( m))
-la,asl
aael [elll2 Exe 20 x /wo T signts #)ssigg(( Ta) Asign(w; T ()ssigg
(ETgg-(8mgn)
(6)
Σ
lael [ewll2 E.d
sigu( ssiguu(( sigga(w; ()
LETi,- (Brign)
(
Σ
lael 1ewlll 06)55))
EET, (Brim))
la
44

---

where (a) æ is a 3-dimensional Gaussian since the expectation only depends on we, wy, w (b) J Ø0 - de |اتت| when sign( "n x) A sign(w? 3:) and ww, ã
- 6g وا|| when ( a)) o
sign(w) (w) æ); (c) direct calculation in Lemma H.2; (d) assumption that norm cancellation is small.
For the second term of (14), similar above, have
Σ
a [ a: w x: sign(wn] m) m T)) sigg(("?" m) =)/sign((1) (
(
2 2laz9isl
awell2 lael
l2 Eg 2 [ellw T 2 sig (w))) p)sssign(w T (8) sign(w;; T ()slsign(( [
(ETATT(Sagg))
(b)
2
2|a,9is1
lael [well ded || Lsign(w: #) my)
(T
2la,9i5l0(5)
E
agl Welll dg
(ETg)T
2/a,giz| (6.66.222)))) (62tose
sign
where (a) Æ is 3-dimensional Gaussian vector since the expectation only depends on wy, Wy, wt: (b) | Ø < 6. ||ق||| when sign(w;; z) f sign(w, 2') and |wo Øi| < 6g ||قق|| when sign(w; x) of
sign( W Ø): (c) a direct calculation as in Lemma H.2; (d) choice of qs and Lemma F.5 and
Lemma E.1l that far-away neurons are small.
Thus, for (II.) we have
(11.i) -21a:qg11O.Tddinnn Odose Setose SA S0--ss1--00 Osign)
For (II.ii,, we have
((T..))
| |/ ||un] a: /w a0 sign(w, a)ssiig((( aign(wo; (ws a))
A BET
l welll 6, 2 1 sign([] () a) 1 leign(w? T si
A RET
[aellaşgo| |/well 6,0, | wt T Iw: [[,
(6)
lael [ewlll dc O(866//A)
Ryll BET
( 2/a 10, (82
class - -A-),
where (a)(b) is 4-dimensional Gaussian vector, | x S 8 | when sign(w; T)) + Swg x) and lww TT) Ø < 8, ||0|| when sign(w;; x) * sign(w) क); (c) by Lemma H.1; (d) choice of gtz and Lemma F.5 and Lemma F.1l that far-away neurons are small.
Combine (ILi) (II,i), we have for (13)
[2 (w)a) a<o' (w," æ) o' (w]x))] x))] -2|a, O( & + Ez C 12-11 Ssign)
This further gives the lower bound on (II):
Σ
Σ
F. (Ra()))) aloo(nn x)) -oa--o'nn o'(10] x))]
iE [m_] JET.
> 2 Σ
O 82 close + S2 1151
xE(m. JET
O, (r
+82 22 C doiga)
-1811)
45

---

Bound (III) For (II), recall Ra(x)
V2 E [ie(m.) of |,ofl2 Eielm as 1100ll12)
(B TTax. a. We have
Σ
Σ
Eun (u T a= (wn - d' (un x))]
8E/m JET
(a)
2 0,(5/2)
(g Re 1|o a sign(w) J a)) #aign((
1E(m.] JET
Σ Σ iElm.JJEE.
1yyby
ExlI( B] TTTZi "=llww x sigu(w] s)ssign(w
[6]
O.(C/A)
Σ
laggs ((65)
iE(m.]
JET.
- 13
Σ
Σ
la
10 Ez ||
sign(w) #)sssggn(w; x))
iE|m.JJET,
( (C/A))
Σ
Σ
1a
JE(m.]JET.
_ 0 0. (82osee/A))
where (a) plugging in the expression of R3 and using Lemma F.9 and Lemma F.11; (b) using Lemma H.3 and the fact that z is 3-dimensional Gaussian vector and | x I dj @2122 when sign(w; w Ø) # sign(w) w/ T); (c) Lemma H.2; (d) choice of qij-
Combine bounds
Combine (I) (II) (III) we now get the last term of (12)
Σ kE(m.]JET,
Σ
E. (R( ( (w) x) - o' (w x))] 2 -O(((/A)3// close + Tosiga Pataa + 8etoso
52-10-1111111
sign
From Lemma F.6 we can choose deosse 0. ( and from Lemma F.16 we can choose daign (2((() Also with T = O(55/6 /A), we finally get
Es (x([(z)a w x(d(ww))) d'ww;" x))] 2 C/8. VE(m.JJET.
Σ
as long as <= 0(295 poly(r, m., A, ||laalll, (min)) with small enough hidden constant.
Thus, we eventually get the lower bound of (12)
aa+ (V B.) +
Σ
(Vv, (Vw,Lx Lx, 9:5 w)) C/4 C/8 C/8.
(E(m.)JET
H.4 Technical Lemma
In this section, we collect several technical lemmas that are useful in the proof. Lemma H.1. Consider a, a R4 with á L(ar, (3) [0, and ||ax||2 |1B11, and
(O, Then, for any 0 < 81, 82 & é we have
[1 a =|5h 11 |
0 62/ sin
Proof. We first consider the case when at least one of 81, 82 2 cΦ for fixed small enough constant. WLOG, suppose 5n 2 cω. In this case, it suffices show bound ((,)). We have
Ex ||||0:|] 1 1l la (611)=1) 1 | S Ez |||20|| 1 |ox #|56 O(61).
46

---

Then, we focus on the case when 82.02 I cœ for fixed small enough constant. WLOG, assume a 01 3 (cos o. sin Φ, 0, 0) and @ E [0, π Then we have
Bm Bn 2 11 1
(2T)² ( r 12 dr
sin 0g d02 d0 /" 2
J0SOIST, c000
sin S0. cos Or cos dit Og cos Og sin ei Lón
d0g
/ Josa 1050:ST//0000:155
=0(1)
sin 01
sin 0g c00 d01
J0S0ST,
-by-eon cne
Scos Og
fo- cond ain 01 sin
win 21 ain
J Josnise 1560
82
sin² 01. 0
d01
sin 01 sin
8182 sin
=0
Lemma H.2 (Lemma C.9 in Zhou al. (2021)). Consider a. E IR3 with L(xx, B3 and
a B 2 0. We have
En [|a==|1 1 sigu(a #) (BT))) [T)) = 0(ø).
Lemma H.3. Consider B E IRd with ((a, o, ||x|2 || 1 and a 2T 0. We have
[a x signn( (t))siig = O(
Proof. It suffices to consider a, B, ER?. WLOG, assume (1, and B (cos Φb sin Φ) We have
/
cos 1 Asigu(eon s +)) d0 re- 12 dr /
2
Ez T x Asign(a ()) 20)1
2n
=0))²).
Lemma H.4. Under Lemma let
a,a;
ifj E To,tooe
Qij
Ee (Fetone)
0
otherwise
Eie(m.) o3 12:1::ll3
s amin/ then Eielm.) EjeT. 22 O(la..11)
Proof. We have
a2
a 22
Σ
Σ Σ
EjeTe, (Betone) az
EjeT, (Betare) ag
jE
sE (E(m- JETh) (Betone)
iE [mmm
In the following, aim to lower bound
ΣjET, (Bstone) az. Given Ejet (Betore) laz ||00,121
|a 1/, we have
2 Σ
2
Σ
af+ + ||,0 1at1/222
Σ
|a51 1a.1/2 2 10/1/2,
jET, (betoaa)
(
jET,, (Batore)
where the last inequality is due Lemma F.6: Σjet, (8ature) la 1011120112
EyET, Og ||w/llll 2 |//1/2. Thus, we have Eie(m.] ΣjeT, d O1laa)11))
47

---

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 0, let ((A=>\oogmmm//ha))) kernel Ke is deneeeeeee in the sense that there exists r ((E-1/3) P1 = ((1). P2 =e(E) such that following hold:
(i) (ww, u) S 1 pr for all 6(w, u) / (wx, u)
(ii) k u2|2, 2) S -p ||=11²2 for tangent vector z that =T w = and 8(w, u) S
(iit) | w m0 S h/m for (i.5) E {0, 13 x (0,1,22
Proof. With the property of Hermite polynomials in Claim A.1, we have
K(w, u) =Bx
J21
20
x
1I
x] =
cos 0,
Z k
1
1
K (100) (w,
(2, u)
cos
k-1 0
(I
ww )
Zz
||w||2
k2t
k(111) (ws, u.)
cosk-22
0
[ |1112 (I ww )uw (I uu.
Za
k>t
1
k c-1 0
(I www )I wu
z
|| 2 1
k2t
k2(20) (w, u.
(w,,
1
Σ
1) cas
k-
wup
Yuu (I uaw
Z?
k2
0ākk cos 8-1 8
w
77(I
Z3 k28
||00|2
2
((21) u): - (20)
(w, u)
1
kk 1)(k - 2) cos 3 0
e. (I wu
) (I
ww Yuu (I www
23 k2
||0 1 |ull2
kk 1) *- k-2 0
(I
www 1) ((I wu
"ee uue (I wu )) (I ww
72 R2F
|/
|u|] 2
kāk 1)
k-2
e, T(
uu
w 7 we
72
k2!
|
[
2
1
cos
1-1
w
(I
wu
)e,(I
ww
Z3
||00
2
k2f
2
(15)
where 0 arccos( w u))
Part (i) Given that r r (1/VV() with small enough hidden constant, we know for 5(ur u) 2
K(w, u)
Σ G cos 0 5-
Σ 0 (1-
=c<l, -
Z2
72
kz
k2t
where is constant less than 1. Thus, P1 +()).
48

---

Part (ii)
For tangent vector z that x -ww = 0, we have (1/2|12 = ||u|l2 = 1, 1, u) r)
k(20) (w; [2,
(k 1) 20 0- (u 2)²
ok cos
0- 1
Z3 k2t
|| = h
Σ
ok - 1) c)s 22 (u 2)
Σ
ok cos k 1-1
-1 0 w a
A28
kzt
|
og - 1) cos k- 0 0
Σ
0 cos 0
Z
A28
For the first term, we have
Σ
ok - cos* 20
k2
Σ oR - 1) (11kk) +
Σ
O(k 0(k-1/3))2
k21/12
ESKSI/r2
Σ
0l138/2) + 0()) = (r)),
121/1²
where we use Lemma I.1 and 6% = 0k 5/2) Lemma A.I.
For the second term, we have
Σ
okk 0 2 1((-112)(1-13))
PSA-228
Given that (11/V)) with a small enough hidden constant, we know
||=12² O (006-1/)=-000
K ((20) (w. 2]
= - (()) ||112,
z3
since Z = -
Part (iii) Recall that 6(w,,, wx) 2 A for i j. It suffices to bound || (w, u)]l2 S h/m2, for
0 = 5(w, u) 2 A. Given that 2 ( 22 llog(m, with large enough hidden constant, from (15) we have for |! |u||
1
1
K(w, u)
S
01
22/5) < hh/m².
Zz
kT10 (w, ui
<
1
Σ
- a k- 0sin 0 S -0(0)11-22/5)))) A2/5)--1 S hh/m²
72
k2?
K k111) u1)
Σ oz - 1) cos 2 21 30 22 +
Σ
ogk cos 0z] 22
sup
22
1
zg u= bZy 1|1 -||
k28
||
1
(k 1) k 2 0
Σ
0 k cosk-1 0
z2 RZ8
z3 kz8
<0(2)22)
Σ 1(kk1122((
13/5)-1- sh/m²
S
49

---

k (1,0 u)
sup
EoRk(k - 1) cos -0 ø1 ū 22
Σ
o
cos k-1 0. u
u. 2,
x, w= w==
122 11ll2
k28
||2 122
1 Σ
Lik(k 1) cos sin? 2 0++
og cos
23 122
<0(2/3)
- -11=-13(1-127877 (1)
Σ @(K
22/5)--1 1 hh/m..
12
K (21)
u
w.u
2R 1)(k 2) cos* 3 Σ q: (I .uu" )w 21 u 22
sup
7 1:2
2 w
20 =2211 - 1=12=1
| 1=ill=11
Σ
oR. 1) cos 0
Σ liz) (I uu )ei u T 22 Σ 9:22 (I uu )e: 21
ZZ
12!
1
okk(k 1) 2 0
q.ee (I Eq.eT
tu w
u 21 22
Z2 k2!
1
Σ
of *
- 1 Σ q: (I uu 'ee 27 22
Za
k2!
$ Zz
Σ
og - 1) - 2) cos sin ++
2 Σ
okk(( cos 0in 0 +
Σ
of
kos
cos * s
z3 12t
z3
A28
k28
(a)
S
<h/m²
where we use 0% ((k-5/2)) in Lemma A.I and (a) the last two terms bound similarly in K (20)
and first term zg Σ k2 of 11) 2) k- 3 0 sin ((8322 Σxze ((11/2((1 22/5//////
h /
Lemma G.2 (Regularity conditions on kernel K). Let Bij = supw ( u) and
Bu Buo + B10 + 1. B2 = Bzo + Br + 1. We have Buo = 01), B = O Bn = O
B20 = ((), B11 = ( and therefore Bo = (2) B = 0
Proof. We compute Bij one by one from (15) (see part (iii) proof in Lemma G. Using Lemma I.1
have
1
Boo = sup
11
Z
k2
1
B10
Jlk cos
k- 0sin0 S 3(13/22)
O 0(2-5/2)6 k
0(e12)) VI
ZZ k2l
k2t
B1i <
1
Σ
kk cos k-2
x0s 2-2 sin? 0 +
6 k cos k-10 0
72
22
kzt
k2t
(0(33/2)
Σ
O(k (/k)5/22) 2/2 + (2222) Σ
O 5/2 - O()),
k
k2!
50

---

B20 5 2 kzt
kk 1) cos k2 0 ssin 0+
Σ
oRk cos 0
Z3
126
<0((3/2)
( 5/2 + 0
Σ
O 6/2 O(),,
k2!
k28
1 B21 $
(k: 1)(k 2) cos A 0
dkk((k 1) 0sin +
1
o cos* s
- 0
27
Z
23 k22
128
k2!
<0(63/2)
((11122(11
/2/55/--0// + (
-1
(kk-1/2)((
k2t
For first term above Ekzs (kk 0² /5) 503 using Lemma 1.2 have
Σ
(3 8²/ 15) 38 <
0
Vln(1/(1 02)) (1 - 82/ 02/5)///2-3/ 383
12
k2
(1 0 15) 0())
0 (1 =
8²
R28
For second term above Σnzr ((k-1/2) 02 12/55*116 have
/ x- 112 /5) < ((0)0
-2
/In(1/(1 - 02) 0(1)
R28
Therefore, we have B21 ((312))
L.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
max cos 8 0 - @(1/VK), max cos sin" 0 @(1/k), max cos* 0 sin 0 (11/13/2)
Proof. We only compute the first max cos s 0 1/vk. Others are similar.
We compute the gradient of f(0) cos sin 0 and get f'( = cos 0(cos 0 - kin 0). We only need to consider 0 [0, 2r). So the maximum is achieved either at boundary = 0, æ or (0) 0. Then one can verify that the bound is true.
Lemma 1.2. For 3 1 and k O, we have 122 gh/2 < (21(())))
Proof. Let f(k) = g//2 We have fk( -1/2 + R 1/ 34/2 Set f'(ko) = 0
we have ko 2In(2/7) It is easy to max f(k) f(ko) - 221n(2/8)
J Notes Sample Complexity
The current paper focuses on the analysis on population loss, which is already highly non-trivial and requires new ideas that developed in the paper. The finite-sample analysis is not focus, omit it in the current paper.
For sample complexity, we believe the following strategy would work to get polynomial sample complexity. We can break down the analysis into parts: early-stage feature learning (Stage 1 and 2) and final-stage feature learning (Stage 3).
51

---

Stage I 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 O. where is the number of sample and d is input dimension.
Stage 3: In local convergence regime, all weights have norms bounded in (+(1) due l2 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.
52

---

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 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 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 long 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 theory paper so all assumptions are clearly listed and discussed. The limitations, for example Stage in Algorithm 1, are clearly discussed the paper.
Guidelines:
The answer NA means that the paper has 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 (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn 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 the scope of the claims made, e.g., if the approach was only tested on few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
53

---

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 include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide 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 ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless 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 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, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or 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 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 possible for other researchers to have some path reproducing verifying the results. 5. Open access to data and code
54

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [NA]
Justification:
No experiments in the paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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 NeurlPS code and data submission guidelines (https:
wttps.cc/ppblccc//cccccccpcppppp geiic//gguee...ee CoodeSummii agionPolicy 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. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, 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 level of detail that is necessary to appreciate the results and make sense of them.
The full details can 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, confi- dence 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.)
55

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation 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 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 emrr 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 well 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 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: : http///eurrpp..e/
Answer: [Yes]
Justification:
The paper follows the NeurlPS 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations 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 theory paper and has no foresee direct societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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.
56

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This is theory paper.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
57

---

If this information is not available online, the authors are encouraged to reach out to the asset's '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/moodl 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 details about compensation (if any)?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, 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 involve crowdsourcing 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 NeurlPS 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.
58


Reviews:
[{"review_id": "HL4OTNQoBt", "reviewer": "Reviewer_WTVj", "summary": "The 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.", "questions": "please see weaknesses", "limitations": "please see weaknesses", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "Comments:\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer's efforts to provide detailed feedback. We will try to address reviewer's concerns below and incorporate the suggestions in the revision.\n\n>1. In the abstract, ...\n\nTheorem 2 is a combination of early-stage feature learning (Stage 1 and 2) and final-stage feature learning (Stage 3). This sentence refers to the local convergence result (Stage 3), and the threshold is $\\varepsilon_0$ in Theorem 2. We show at the end of Stage 2 (Lemma 4) the loss/optimality gap is below $\\varepsilon_0$ so that we enter the local convergence regime in Stage 3 (Lemma 5). We agree that this should have been made clearer and we will address this issue in the revision.\n\n> 2&3. The relation between assumption 3 and information exponent should be made explicit ...\n\nThanks for the suggestion. We will make it more explicit in the revision. The information exponent of teacher network is indeed 1. Assumption 3 we made here is the same as in Damian et al. (2022), and we use the same preprocessing procedure to remove the linear part. After this preprocessing step, the information exponent becomes 2 due to Assumption 3.\n\nOur results hold when information exponent is 2 after preprocessing. If it's larger, the current analysis does not work, and it would require a generalization of Damian et al. (2022) to make it work (Dandi et al. (2023) explored a generalized version of this problem and left it as a conjecture).\n\nNeural networks can learn representations with gradient descent, Damian et al., 2022\n    \nHow Two-Layer Neural Networks Learn, One (Giant) Step at a Time, Dandi et al., 2023\n\n>4. Why is the head a and the back layer W ... \n- Regularization on $a$ and $W$: \n\nThis is the same as using weight decay on both layers. At the minimum we have $\\|\\|w_i\\|\\|_2^2+|a_i|^2 = 2\\|\\|w_i\\|\\|_2|a_i|$, so $\\ell_2$ regularization $\\sum_i\\|\\|w_i\\|\\|_2^2+|a_i|^2$ becomes $2\\sum_i\\|\\|w_i\\|\\|_2|a_i|$ that is effectively $\\ell_1$ regularization over the norm $\\|\\|w_i\\|\\|_2|a_i|$ of neurons. Such a regularization favors sparse solution and helps us to recover the ground-truth directions and remove irrelevant directions.\n\n- Analyze the problem of $\\lambda=0$ or $\\lambda\\to0$:\n    \nOur goal is to minimize the loss with $\\lambda=0$. However, directly analyzing the unregularized case is challenging so we choose to analyze the case where $\\lambda$ is positive but gradually decreasing to 0 ($\\lambda\\to 0$). In this way, the solution we get will eventually converge to the minima with $\\lambda=0$. In Theorem 2 we show that the unregularized loss ($\\lambda=0$) is small at the end of training.\n\n>5. The analysis is for gradient descent ran on expected loss ...\n\nPlease refer to the general response where we address the questions regarding sample complexity.\n\n>6. The algorithm analyzed in this paper ...\n\n- Non-standard algorithm: \n    \nRecent feature learning literature often focuses on early-stage learning, using layer-wise training (training 1st-layer weights $W$ first, then 2nd-layer weights $a$), which is also not standard in practice. As discussed in the paper, analyzing the entire training dynamics of standard gradient descent, especially the middle stage (Stage 2), is technically challenging and remains an open problem. We hope our results contribute to understanding standard training methods.\n    \n- Role of Stage 2: \n    \nIn our algorithm, the Stage 1 and 2 essentially follow the similar layer-wise training procedure that are common in early-stage feature learning literature. The role of Stage 2 is to do a regression on top of the learned feature after Stage 1. This makes the loss small so that we enter the local convergence regime in Stage 3.\n\n- Regularization and balancing:\n    \nWe assume `normalizing' mentioned by reviewer refers to the regularization or the norm balance. We'd be happy to further clarify if this does not answer the original question.\n\nThe regularization is essentially the same as $\\ell_2$ regularization or weight decay used in practice, as $\\|\\|w_i\\|\\|_2^2 + |a_i|^2 \\ge 2\\|\\|w_i\\|\\|_2|a_i|$, with equality at the minimum of $\\ell_2$ regularized loss. The reason of using weight decay is further discussed in Section 5: it induces an effective $\\ell_1$ regularization over the neurons and helps to reduce norm cancellation between neurons.\n\nThe norm balancing step is just a preparation step for Stage 3 for technical convenience to enter the local convergence regime.\n\n>7. In the discussion below Theorem 2, ...\n\nWe will make it clear in the revision. To be more specific, we could look at Damian et al. (2022) (although the discussion also applies to other works that rely on (large) first-step gradient, e.g., Ba et al. (2022), Abbe et al. (2022)). It is shown that after the first-step gradient update neurons $w_i^{(1)}\\approx c Hw_i^{(0)}$, where $c$ is a scalar and $w_i^{(0)}$ is neuron $w_i$ at random initialization. Given the assumption that matrix $H$ has full rank within the target subspace, $w_i^{(1)}$ is essentially sampled randomly from target subspace. This suggests neural networks in fact only learn the target subspace and do random feature within it.\n\n>8. Instead of the lengthly discussion on the construction of the dual certificate, ....\n\nThanks for the suggestion. We do believe the construction of the dual certificate is the core technical contribution. However We will clarify general take-aways more clearly in the revision. Our main finding is that training both layers of 2-layer networks to convergence can achieve 0 loss and recover ground-truth directions. This finding highlights a strong form of feature learning that complements the early-stage feature learning literature, where the first-layer weights are typically fixed after one or a few steps of gradient descent.\n\n>9. Missing citations and discussion of prior work...\n\nThanks for the references. We will add discussions about these works in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XrPPDEWeKH", "reviewer": "Reviewer_2mhB", "summary": "This 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.", "questions": "- 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.", "limitations": "The limitations are addressed in the submission.", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The 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.", "weaknesses": "This 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer's efforts to provide detailed feedback. We will incorporate suggestions on the presentation of the paper and discuss more related works in the revision. Below we try to address reviewer's concerns.\n\n> In equation (3) you preprocess the target network to remove its first Hermite coefficient ...\n\nRemoving the first Hermite direction is important for the current analysis and Damian et al. (2022). If there is a non-vanishing first Hermite direction, then it will become the dominate term in the first-step gradient. To be more specific, after first-step gradient update neurons $w_i^{(1)}\\approx c_1 \\beta +c_2 Hw_i^{(0)}$, where $c_1,c_2$ are scalars, $\\beta$ is the first Hermite direction and matrix $H$ has full rank within the target subspace due to Assumption 3. When first Hermite direction $\\beta\\neq 0$, $\\beta$ dominates so $w_i^{(1)}\\approx c_1 \\beta$. This makes all neurons collapse into a single direction.\n\nWe would also like to note that this issue is likely a technical difficulty introduced by the setting we consider. In experiments, we can still achieve strong feature learning ability with proper choice of hyperparameters. Providing rigorous analysis beyond current setting is an interesting and challenging problem to consider in the future.\n    \n> Could the author be more precise on how they would extend the findings to a polynomial number of samples after Theorem 2 ...\n\nPlease refer to the general response where we address the questions regarding sample complexity.\n\n    \n> What do the authors mean when saying \"complexity of $f_*$\" on Page 4? It would be nice to mathematically formalize this concept.\n\nThanks for the suggestion. We will make it clear in the revision. Here by ``complexity of $f_*$\", we mean it only depends on the quantities of $f_*$, such as teacher neurons' norm $|a_i^*|,\\|\\|w_i^*\\|\\|$, target subspace dimension $r$ and angle separation $\\Delta$.\n    \n> What is $\\bar{w}_i$ at page 5?\n\nIt is the normalized version of neuron $w_i$, i.e., $\\bar{w}_i = w_i/\\|\\|w_i\\|\\|_2$.\n\n> What is meant for $\\varepsilon_0$-net at page 5?\n\nWe mean an $\\epsilon$-net with $\\epsilon=\\varepsilon_0$. That is, neurons cover the whole target subspace well in the sense that for every direction $v$ in the target subspace, there exists a neuron $w$ such that the angle between them is small $\\angle(w,v)\\le \\varepsilon_0$.\n\n\n> Presentation of paper and discussion of prior work.\n\nWe will improve the presentation of paper based on the suggestions and discuss more related works in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pyEGr76llf", "reviewer": "Reviewer_6HCx", "summary": "The present manuscript concerns the study how gradient descent achieves feature learning in a certain class of two-layer neural networks.\nThe 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.", "questions": "-", "limitations": "See weaknesses", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewers for the feedback. For the concern about sample complexity, please refer to the general response where we address the questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We appreciate the reviewers' detailed feedback and their efforts in providing constructive comments. One common question raised by reviewers is about sample complexity. We try to address it below.

First, we would like to emphasize that even the analysis on population loss is 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 = \Theta(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 $\ell_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. 


In this way, we could get a polynomial sample complexity result. 

Below we address more detailed questions on sample complexity:

- Better sample complexity

    It is indeed not clear how to get a better sample complexity bound like the one reviewer asked for ($n=\Theta(d)$). A simple parameter counting approach would suggest likely $\Theta(dr)$ samples are needed. As discussed in Damian et al. (2022), they showed at least $\Theta(d)$ sample is required in this setting and only provided a sample complexity depending on $\Theta(d^2)$. Moreover, if we focus on the first-step gradient, Dandi et al. (2023) showed that with only $n=\Theta(d)$ samples only 1 direction could be learned.  In contrast, $n = \Theta(d^2)$ is essential for neurons to learn multiple relevant directions of the target with a single gradient step. This suggests more than $\Theta(d)$ samples might be necessary to learn multi-index models like the 2-layer nets we consider in the paper. Improving sample complexity, especially in early-stage feature learning (e.g., Abbe et al. (2023), Damian et al. (2024), Arnaboldi et al. (2024)), is an active and interesting research topic, which we leave for future work.

- Impact of information exponent

    The information exponent indeed would matter for sample complexity. Due to Assumption 3, the information exponent after preprocessing step is 2 in our setting. In the outline above, $d^2$ sample complexity appears in Stage 1 and 2. This is because in fact we need an accurate estimation of degree-2 Hermite polynomial (matrix $H$), as discussed in Damian et al. (2022). When the target function is harder (information exponent is larger), we expect the sample complexity would increase accordingly.




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

Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics, Abbe et al., 2023
    
How Two-Layer Neural Networks Learn, One (Giant) Step at a Time, Dandi et al., 2023

Computational-Statistical Gaps in Gaussian Single-Index Models, Damian et al., 2024

Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions, Arnaboldi et al., 2024

Meta-review: This paper provides a convergence analysis of gradient descent to learn a two layer neural network in a teacher-student setting. Under a bit strong assumption, it is shown that the dynamics decreases the loss to 0 and it can be divided into three stages. The computational complexity is polynomial.   

One main drawback of this paper is that the dynamics is analyzed only for the population risk minimization. A typical approach in this direction is to show the sample complexity as well as its polynomial time computational complexity and discuss its optimality (referring an SQ or CSQ lower bound). This drawback is somehow resolved in the author rebuttal but it requires substantial revision to include the finite sample analysis.   
On the other hand, the analysis itself is not trivial. The authors provided an interesting contribution even though the loss is of population.   

As a summary, I think the result in this paper is informative in the literature and worth publication in NeurIPS. Hence, I recommend acceptance of this paper.
Decision: Accept (poster)