Learning Orthogonal Multi-Index Models: A Fine-Grained Information Exponent Analysis
Authors: Yunwei Ren
Jason D. Lee
Venue: NeurIPS 2025 poster ()
Keywords: multi-index model
online stochastic gradient descent
information exponent
sample complexity
Primary area: theory

Abstract: The information exponent (Ben Arous et al. [2021]) and its extensions --- which are equivalent to the lowest 
  degree in the Hermite expansion of the link function (after a potential label transform) for Gaussian single-index 
  models --- have played an important role in predicting the sample complexity of online stochastic gradient descent 
  (SGD) in various learning tasks. In this work, we demonstrate that, for multi-index models, focusing solely on the 
  lowest degree can miss key structural details of the model and result in suboptimal rates. 

  Specifically, we consider the task of learning target functions of form $f_*(x) = \sum_{k=1}^{P} \phi(v_k^* \cdot x)$, 
  where $P \le d$, the ground-truth directions $\\{ v_k^* \\}_{k=1}^P$ are orthonormal, and the information exponent of 
  $\phi$ is $L$. Based on the theory of information exponent, when $L = 2$, only the relevant subspace (not the exact 
  directions) can be recovered due to the rotational invariance of the second-order terms, and when $L > 2$, 
  recovering the directions using online SGD require $\tilde{O}(P d^{L-1})$ samples. In this work, we show that by 
  considering both second- and higher-order terms, we can first learn the relevant space using the second-order
  terms, and then the exact directions using the higher-order terms, and the overall sample and complexity of online 
  SGD is $\tilde{O}( d P^{L-1} )$.

Learning Orthogonal Multi-Index Models: A FraieGr Information Exponent Analysis
Yunwei Ren Princeton University yunwei www.princeton.ccommmmmmmm edu
Jason D. Lee Princeton University j jasonlee@princetton. edu
Abstract
The information exponent ([BAG21I) and its extensions which equivalent to the lowest degree in the Hermite expansion of the link function (after a potential label transform) for Gaussian single-index models have played important role in predicting the sample complexity of online stochastic gradient descent (SGD) various learning tasks. In this work, we demonstrate that, for multi-index models, focusing solely on the lowest degree can miss key structural details of the model and result in suboptimal rates.
Specifically, we consider the task of learning target functions of form f.:)
Σe-n ob(1) (0) x). where P < d. the ground-truth directions {uīje- are orthonor- mal, and the information exponent of is L. Based the theory of information exponent, when L= 2, only the relevant subspace (not the exact directions) can recovered due to the rotational invariance of the second-order terms, and when L > 2, recovering the directions using online SGD require Ó ((Pd'-1) samples. In this work, we show that by considering both second- and higher-order terms, we can first learn the relevant space using the second-order terms, and then the exact directions using the higher-order terms, and the overall sample and complexity of online SGD is O(dPl-1).
Introduction
In many learning problems, the target function exhibits or is assumed exhibit a low-dimensional structure. A classical model of this type is the multi-index model, where the target function depends only on P-dimensional subspace of the ambient space R. with typically much smaller than When the relevant dimension P = 1, the model is known as the single-index model, which dates back least [Ich93]. Both single- and multi-index models have been widely studied, especially in the context of neural network and stochastic gradient descent (SGD) in recent years, sometimes under the name "feature learning" [BAGJ21, BBSS22, DI.S22, AAM22, AAM23, DKPS24, DPVLB24, OSSW24, UTA+244)
In [BAGJ21], the authors show that for single-index models, the behavior of online SGD can be split into two phases: an initial "searching" phase, where most of the samples are used to boost the correlation with the relevant (one-dimensional) subspace phase, where the correlation further increases to 1. They introduce the concept of the information exponent (IE), defined as the index of the first nonzero coefficient in the Taylor expansion of the population loss around 0, which also corresponds to the lowest degree in the Hermite expansion of the link function in Gaussian single-index models. They prove that the sample complexity of online SGD TE- After that, various lower and upper bounds have been established for single-index models in [BBSS22, DNGL23, DPVLB24]. Similar results for certain multi-index models have also been derived in [AAM22, AAM23, BBPV23, OSSW24). In all cases, the sample complexity of online SGD scales with L when IE 2 3.
constant, and subsequent "descending"
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Later, it was realized that the notion of information exponent is not stable under modifications of the algorithm. In particular, the information exponent of a link function can be greatly reduced by reusing batches or applying a suitable label transformation [ADK* 24, 24A+24, LOSW24, DPVLB24].. For example, the IE of any fixed degree polynomial can be reduced to at most 2 via monomial transformations. This observation leads to the notion of generative exponent (GE) (PPVLB24) which is defined as the lowest information exponent among all L transform of the original link function. It yields bounds that match the previous results for non-gradient--iae -baed methods [CM20, TDD*24 24, BKM 19J. Despite the improvement over the vanilla information exponent, in the framework of generative exponents. still only the lowest order is considered. As we will discuss later, this makes it suffer from the same issue of information exponent in the context of multi-inder models.
Consider multi-index models of form f.()) = =E p x), where {of are orthonormal vectors. In this setting, there are two types of recovery: recovering each direction Vi2 and recovering the subspace spanned by {ug The former notion is stronger, and once the directions are learned, the learning task essentially reduces to learning the me-dimensional ø R + R. However, directional recovery is not always possible. To see this, consider the case ((2) = hz(2), where h; is the I-th (normalized) Hermite polynomial. One can show that this corresponds decomposing the projection matrix (a second-order tensor) of the subspace span (gg k- Hence, recovering the directions is impossible due to the rotational invariance (see Section 3.1 for more discussion). In other words, in any framework that considers only the lowest order (IE or GE), if the lowest order is 2, we cannot get guarantees beyond subspace recovery due to the existence of the Φ = h22 example.
On the other hand, if & contains some higher-order terms, e.g., Φ = h22 + h4, then one should expect that directional convergence is possible, even though IE(o) is still 2, because of identifiability property of (higher-order) orthogonal tensor decomposition problem [GLM18, LMZ20, GRWZ21]. In addition, we should be able to first recover the subspace using the second-orde terms, which should require Ö(d poly P) samples, and then recover the directions using the higher-order terms. Moreover, since we have learned the relevant subspace, the number of samples needed in the second step should be much smaller than what is needed if there were no second-order terms.
In this work, we formally prove the above conjecture and show that the overall sample complexity OddPL-I), where I. the (lowest) order of the higher-order terms. Note that this bound scales linearly with the ambient dimension d, and it matches the sample complexity of separately learning P independent single-index models with the relevant subspace known but the noise scales with d, up to potential logarithmic terms. More formally,
prove the following theorem.
Theorem 1.1 (Informal version of Theorem 2.1). Suppose that the target function is given by [a = Ekm oot x) where & = 22h2 + Σi= Srhy, with L 2 3, ø2 87 > O, and ) are orthonormal, and the input x follows the standard Gaussian distribution N (0, Ia)- Then, we can use online SGD (followed by a ridge regression step) to train a two-layer network of width Ö(P) to learn (with high probability) this target function using ((dPL-1) samples and steps.
Organization The rest of the paper is organized as follows. First, we review the related works and summarize our contributions. Then, we describe the detailed setting and state the formal version of the main theorem in Section In Section 3, we discuss the easier case where the training algorithm is population gradient flow. Then, in Section 4, we show how to convert the gradient flow analysis to an online SGD one. Finally, we conclude and discuss the limitations in Section 5. The proofs, simulation results (Appendix F), and table of contents can be found in the appendix.
1.1 Related work
In this subsection, we discuss works that are directly related to ours or were not covered earlier in the
introduction.
Along the line of information exponent, the paper most related to ours [OSSW24]. They show that for near orthogonal multi-index models, the sample complexity of recovering all ground-truth directions using online SGD is ((-))- when IE 2 3. Their results do not apply to the case
'The IE = case is particularly important as the information exponent of many functions, including all fixed degree polynomials, can be reduced to at most 2 by a suitable label transform [DPVLB24], and the information exponent of any even functions is at least 2.
2

---

IE = 2 for the reason we have discussed earlier. They propose first removing the second-orde terms using the technique in [DLS22. DKPS24), which requires d' samples. Our result considers the situation where both the second and I-th order terms are present and show that in this case, the sample complexity of online SGD (without any preprocessing) is (ddPL-1).
Another recent related work is [BAGP24]. Our main results are not directly comparable since the settings are different. They run SGD on the Stiefel manifold, which automatically prevents the model from collapsing
single direction, but allow the target model have condition number larger than 1. In addition, only the lowest degree is considered in their work. However, they also show (in their setting) that when the second order term is isotropic, the subspace and only the subspace can be recovered. A similar idea is also used in our analysis of Stage 1.1 (cf. Section 3.1).
Another related line of research is learning two-layer networks in the teacher- setting ([ZSJ+ 17, LY17, Tia17, LMZ20, ZGJ21, GRWZ21J. Among them, the ones most relevant to this work are [LMZ200 and the follow-up [GRWZ21], both of which consider orthogonal models similar to and use similar ideas in the analysis of the population process. However, they do not assume a low-dimensional structure and only provide very crude poly(d)-sty) sample complexity bounds.
1.2 Our contributions
We summarize our contributions as follows:
We demonstrate that information/ggeneraiive exponent alone is insufficient to characterize certain structures in the learning task and show that for specific orthogonal multi-index model, if we consider both the lower- and higher-order terms, the sample complexity of directional recovery using online SGD can be greatly improved over the vanilla information exponent- based analysis.
As a by-product, we derive collection of user-friendly technical lemmas to analyze the difference between noisy one-dimensional processes and their deterministic counterparts, which may be of independent interests (cf. Section 4.1 and Appendix E).
2 Setup and main result
In this section, we describe the setting of our learning task and the training algorithm, and then formally state main result. We will also convert the problem an orthogonal tensor decomposition task using the standard Hermite argument as in [GLM18].
often drop the subscript and simply write |-||. For a, b, ö R. = 6 means |a - b] S (ö] and V b maxfa, b} and A b- = min b). Beside the standard asymptotic (big O) notations, we also use the notation fa = Og(g). which means there exists a constant Cs > that can depend only on such that fa < Caga for all large enough d. Sometimes we also write fa So Id for fa Op(9a) The actual value Co can vary between lines.
Notations We use II-IIp denote the p-norm of a vector. When p
2.1 Input and target function
We assume the input x follows the standard Gaussian distribution N (0, Ia) and the target function has form fx = Ep_ ( x),, where log d < P S d for large universal constant C' > {ut) are orthonormal and I + R is the link function. In addition, we assume & satisfies the following.
Assumption (Assumptions on the link function). Let h& denote the degree-k normalized Hermite polynomial and & = Eiso dehy denote the Hermite expansion of o E L2 (0, Ia))
(a) (IE structure) For some constant L > 2, ((2) = )221((2) + 6zh((2) - Eist,
(b) EE regularity) &2. &1 = 211 and !| = -ie 107 s CJ for some constant Co > 0.
(c) (Polynomial growth) There exists universal constants C, q 0 such that (d(x)) V 10 (x)l C(1 + 22)0/22 for all all x E R.
3

---

Our target model and algorithm will all be invariant under rotation. Hence, we will assume w.l.o.g. that 2 e₁ where (ex is the standard basis of Rd
2.2
Learner model and the training algorithm
Our learner model is width-m two-layer network
m
f()) = f(x;
a, a) = Σ a.o(v; x),
where = (aI,..., am) Rn and V (U), (UI...., Um) E (gd-1)m are the trainable parameters. We call {v; Jelm) the first-layer neurons. We measure the difference between the learner and the target model using the correlation loss. Given sample (a, f. (2))), we define the per-sample and population MSE loss as
- (((()--(a;a,,))))) v))? LMSE(a, V) = E ELMSE (2; a, V). 1 2 (f.(x) (M) == ((:; a, V) V): =
Now, we describe the training algorithm. First, we initialize each output weight a, to be Then, we symmetrically initialize the first layer neurons. That is, for i E [m/2], we initialize Unif(S independently and set / vv; for the other half of the neurons. After the initialization, we fix the output weights a and train the first-layer weight correlation loss (corr(
using online (spherical) SGD with the -f ( (xx and step size 7) > for T iterations. Then, we fix the first-layer weights and use ridge regression to train the output weights
Let {( (z: be our samples where {æ,} are i..d. standard Gaussian vectors, and let V. (I ve T V., denote the spherical gradient. Then, we can formally describe the training procedure as follows:
Initialization:
00.1 = 00.5 LLd. Un
00.m/2+1 Do,i
Vi E [m/2];
0.41, = Uz.i + nf (xxe Vo, 61u xx),
Stage 1:
Vi [m], t [T];
U141,1 uit1s/ /
Stage
a argmin
Σ MøTTnia' a', Vr) + + 2|a1a'1 n=1
2N
(1)
Here, the hyperparameters are the network width m > 0, step size 17 > 0, time horizon > 0, the number of samples N in Stage 2, and the regularization strength À > 0.
We will show that after the first stage, for each ground truth direction oh there will some neurons 0, that has converged to that direction. As result, in the second stage, we can use ridge regression to pick out those neurons and use them to fit the target function. The analysis of this second stage is standard and has been done in [DLS22, AAM22, BES 22, LOSW24, OSSW24). Hence, we will not further discuss this stage in the main text and defer the proofs of this stage to Appendix C.
For the gradient update in Stage 1, we have the following lemma on its expectation and tail. The proof of this lemma is rather standard and can be found in, for example, [GLM18, OSSW24). We also provide proof in Appendix A.1 1 for completeness.
Lemma 2.1 (First-layer gradients). Consider the setting described above. Suppose that & satisfies Assumption and ag 1 for all E [m] and {vt} are orthonormal. Then, for each [m], we have
E Eff.(z) x)] 2 Σ
(vt v.) 2z
>V l0i (v vi) 1T.k=l
wi)² vi.
(2)
Then, for each fixed neuron (a, v) and direction 5d-1 that is independent of x ~ N (0, Ia), have
E (ff((x) 0(0; x), So P.
(f.( 2), x),u) So /2/ log 2(1+g) log(m/6p) with probability at least 1 8p.

---

Remark.
We say a random variable X is - [VGNA20, KC22] if
P[[X]
M) S exp ((MM/6)1/0) WA 0.
(3)
Hence, this lemma implies that V,, æ), is (P. 1/ (egggg)))l)iibblnee subweibull.
2.3 Main result
The following is our main result. The proof of it can be found in Appendix D.
Theorem 2.1 (Main Theorem). Consider the setting and algorithm described above. Let C > O be a large universal constant. Suppose that log d I I d and {oz hE((P) are orthonormal. Let óp rexp(-log" d), and E. be given. Suppose that choose ag T, satisfying
p
2,00
PPL/2-1)
m O(P), N =
1 6
PdPL.2-1)
=
2.0
Then, there exists some À O such that at the end of training, we have LMSE(a, V) S 8. with probability least (())
Remark. Note that N < T. Hence, the total number of samples needed is T (dPL--)) which matches the sample complexity of separately learning P single-index models with the relevant subspace known priori and the noise scales with the ambient dimension d.
3 The gradient flow analysis
In this section, we consider the situation where the training algorithm in Stage I is gradient flow over the population correlation loss instead of online SGD. The discussion here is non-rigorous and our formal proof does not rely anything in this section. Nevertheless, this gradient flow analysis will provide valuable intuition on the behavior of online SGD.
For notational simplicity, we will assume w.l.o.g. that of = ex. In addition, we will assume = fo + Jy with L for case of presentation. Let be an arbitrary first-layer neuron. By Lemma 2.1, the dynamics of v are controlled by²
e, 2 E uk (I vo ee + of wo Teek-
The second term on the RHS comes from the normalized/projection. For each k: [d], we have
dr 2 = 21 P) (2 Luf-2) -2 2 [ vz.
(4)
We further split Stage into two substages. In Stage 1.1, the second-order terms dominate and || | grows from (PP/d) to O(1). In Stage 1.2, e converges one ground-truth direction relying on the signal from the higher-order terms.
The direction which e will converge depends the index of the largest v at the beginning argument, one can show that max V is least + times larger than the second-largest Vf for small constant >> 0 with probability at least 1/ poly(P) (P) the initialization (of Stage 1.1). Hence, as long as this gap can be preserved throughout Stage 1, we can choose m poly ()) to ensure all ground-truth directions can be found after Stage 1.2.
of Stage 1.2. With some standard pnncemtration/anii-cocceee
3.1 Stage 1.1: learning the subspace and preservation of the gap
In this substage, we track || || 21 and o///0t /u3 where p. q E [P] are arbitrary. The goal is to
show that || pll will grow to constant while /// stay close to its initial value.
In the main text, we use to index the time in this continuous time process (as t has been used to index the steps in the discrete-ting process) and will often omit it when is clear from the context.
5

---

For the norm ratio, by (4), we have
d |0ge|12 d
|ugell
| 2 d
||
4 A|
2L 11u1
dr | dr ||>r||
dr |0sp|1
+
| pll
||>e|1
|1
2 (211ugp11 + L llaeatt Tosell
2 (2222lte (2||0gp|1
+ L 11 ||
||uee|l
||
|,>p|| 1
In particular, note that the terms coming from normalization cancel with each other. Moreover, this
implies 4 |Psell" dr || 2 4 ||o and therefore, it takes only at most ++o(1) log(d/P) (log(d/P))
||
amount of time for the ratio to grow from ((/dd) to ( If we choose small step size so
that online SGD closely tracks the gradient flow, then the number of steps one should expect
Ollgg(d/P) P)/n)-
Meanwhile, for any p.g E [P], we have
d 22 2 + Lup 2)
22
2 (2110gr11 + + 11l0sp1t
dr 1/
2
(2 + Log 2 -2 (2110gp11 + (1borllt)) 2L (u/-2 p
(0p 1/g -
Note that not only those terms coming from normalization cancel with each other, but also the second-order terms. In particular, this also implies that we cannot learn the directions using only the second-order terms. At initialization, with high probability vf = ((1/d) for all E [P]. Hence, d
we assume the induction hypothesis ox (L1PP), then above will become di og/ o2 S P -e /u/ As result, p //. < (1+ + o(1)) /ö., for any & < ((log(d/))) as long as P 2 poly log d.
3.2 Stage 1.2: learning the directions
Let e be first-layer neuron with of 2 (1 + c) maxgsk< of for some small constant c >> at initialization. By our previous discussion, we know at the end of Stage 1.1, the above bound still holds with a potentially smaller constant > 0. In addition, since ||u<p| = ( we also have of 2 (2(1P)) P) at the end of Stage 1.1. We claim that v will converge to e₁. The argument here is similar to the proofs in [LMZ20] and [GRWZ21].
Again, by (4), we have
d oî = (2 -
2| + Lof-² L -1lesp1li) of 2 2L (ot-2 (os of
dr
Assume the induction hypothesis uf 2 (1+0) +c) max of and write
of-2 -f-² - -1l-lllpll - (1-221110))) vf) - (1 (10<p|1- - of
2/
L-2
of
k=
Note that the summation is a weighted average of (ot "J22 and therefore can be upper bounded by (oi/(L e)) 4/2-1 (1 -cr)))u-22 for some constant Cr. 0 that can only depend L. Thus, we have
d dr
2T 1 (1 (1 2
(1-)))o) -10
(1 -
os 2cr (1 vf) of
When of < 3/4, this implies duf of Z1 of. As result, it takes at most Oz Oz (PL/2-1) amount of time for of to grow from ((1/P)
(P)() to 3/4 under gradient flow. It important that vz 52(1 P) instead of ((1/d) at the start of Stage 1.2, since otherwise the time needed will be O (d(--11) After of reaches
3/4, we have d - uf) Se (1 vf) Thus, of will converge lincarly to afterwards.

---

From gradient flow to online SGD
In this section, we discuss how to convert the previous gradient flow analysis to an online SGD one. Our actual proof will be based directly on the online SGD analysis, but the overall idea is still proving that the online SGD dynamics of certain important quantities closely track their population gradient descent (GD) counterparts. Our choice of learning rate nl will be much smaller than what needed for GD to track GF, so the bottleneck comes from the GD-to-SGD conversion, not the GF-to-GD one. Provided that SGD tracks GD well, the number of steps/samples it needs to finish each substage is roughly the amount of time GF needs, divided by the step size
The rest of this section is organized as follows. In Section 4.1, we collect few useful lemmas for controlling the difference between noisy dynamics and their deterministic counterparts. The idea behind them has appeared in [BAGJ21] and is also used in [AAM22]. Here, we simplify and slightly generalize their argument and provide a user-friendly interface. When used properly, it reduces the GD-to-SGD proof to routine calculus. Then, in Section 4.2, discuss how to apply those general results to analyze the dynamics of online SGD in our setting.
4.1 Technical lemmas for analyzing general noisy dynamics
We start with the lemma that will be used to analyze ||wgr| 1// /|/s>r/l The formal proofs of it and all other lemmas in this subsection can be found Section
Lemma 4.1 (Stochastic Gronwall's lemma). Suppose that (Xt)e satisfies
X (1 + a)X, + E141 + Ze+1, Xo = 20 > O,
(5)
where the signal growth rate a > 0 and initialization xg > 0 are given, (Et)e is an adapted process, and (Ze); is martingale difference sequence. Define x; (1 xx-
Let T and 8y E (0. 1) be given. Suppose that there exists some ōF. E (0, 1) and E,
oz > 0 such that for every t 2 0, Xf - (1 0.5) then we have |5+41] (1 aa' with probability at least
1 - -7, and Ze41 is conditionally ((1 a))'2,,22 0) subbeebiil Then, if
z
E and oz
(6)
Tlog (T
we have X; = (1 (1 0.5).z for all t e [T] with probability at least 1 - 8q -6.ze
Condition (6). One may interpret Z141 as those terms coming from the difference between the population and mini-batch gradients, whose variance is typically quadratic in and Er+1 as the higher-order error terms. a is usually small. In our case, it is proportional to the step size n. T is usually the time needed for X₂ to grow from a small ro > 0 to @(1), which is roughly a log(1 1 / Since the LHS' of (6) are (17)) while the RHS' are (()). (6) can be alternatively viewed as condition on 17.
Stochastic induction. One important feature of this lemma is that only requires the bounds |s++| < (1 + a and E[22+1 Fi) S (1 + a)'az to hold when X, (1 + 0.5)x; This can be viewed as form of induction and is particularly useful when considering the dynamics of, say, ei. Similar to how the RHS of d 2 20,, depends on U, the size of E++1, Ze+1 will usually depend on X,, Hence, we will not be able to bound them without suitable induction hypotheses.
Remark on the subweibull condition. We assume the martingale difference terms (Ze+1), ) are conditionally subweibull. This allows us to get poly-logarithmic dependence on 8z, which is important as we will eventually take union bound over poly events. One may replace this condition with the weaker condition E[ Fi] < (1 a)'d 22 This will lead to a linear dependence on öp.
Proof sketch of Lemma 4.1. For the ease of presentation, we assume that |5++1| < (1 + aa with probability at least ōr.g and E[ Fi]] S (1 + a) always hold. This step can be made formal using stopping time argument. See Section E for details. Then, Unroll (5) to obtain
7

---

X++1 = a1+ zo- + Ei=1( (1 + a)) + Ei=1 (1 -aaa and replace + 1 with t. Then, the above becomes
Divide both sides with
a) (1+a)(+1
(1+
X1(1 + a) = zu + E1= +5 + E'= (1++)
The second term is bounded by TE (uniformly over S T) with probability at least 1 - Tõp,ç. Note that (1 + = Z is still a martingale difference sequence. Hence, by Doob's su-submari inequality, the third term is bounded by 2o/4 with probability at least 160 (axgg). Thus, when (6) holds, the RHS is (1 + 0.5)zz with probability at least - T óp. Multiply both sides with (I+a)*, and we complete the proof. To improve the dependence on óp from linear to poly-logarihmmi. it suffices to replace Doob's s L'=submartinggle inequality with a variant of Freedman's 's inequality that works with subweibull variables (ef. Appendix E).
Using the same strategy, one can prove a similar lemma that deals with the case a = 0, which will be used show the preservation of the gap Stage 1.1, Another interesting case is where the growth is not linear but polynomial. This is the case of Stage 1.2 in our setting. For this case, we have the following lemma.3
Lemma 4.2. Let (Xi)e be a non-negative stochastic process satisfying
X 2 Xt + a.X Z+11 + Ee+1, Xo 30 >
(7)
where a > 0, (Z4+3)e is martingale difference sequence, and (Et)s an adapted process. Let te be the solution to the deterministic recurrence relationship Ri+1 = E2 + aci, 2o za/2.
Let &p E (0, be given and T := inf { S (pa(xx// (2222)2211 12)--1)- X. 2 1 Suppose that there exists
E, oz > 0 and og,g E (0, 1) such that if X, 2 & and t S T, we have |k,l S EX, with probability at least 1 ε and Zu+1 1 is conditionally (o2 O)-suweeibull Then, if
a a - /p, = 02 S pazd poly log(T/5r),
(8)
we have X, 2 8. for all t S T and X, 2 with probability least Tör,s 85.
The proof of this lemma can be found in Section F. It is similar to the previous proof in spirit: we replace (1 + a)* with IT= + axp-1) and unroll the recurrence. However, unlike the linear case, it is generally difficult to upper bound the difference between X, and ét, as this type of polynomial systems exhibit sharp transitions and blow up in finite time, Consequently, & S EX, and B[Z Fil S x2 do not directly imply that & S E&, and E| Fil S 2 and this makes the analysis tricky as the RHS' are not deterministic. To handle this issue, we use the following recoupling argument: whenever X, 2 45², we replace the current &, with X./2. Clear that this can only increase ét, and it ensures X, I 4., always holds. Moreover, after each recoupling, dt will at least double. As a result, the conditions we need absorb the noises will also become weaker.
4.2 Sample complexity of online SGD
In this subsection, we demonstrate how to use the previous results to obtain results for online SGD and discuss why the sample complexity is (dPL-11 instead of (dd--1) even though we are relying on the L-th order terms learn the directions.
4.2.1 A simplified version of Stage 1.1
As an example, we consider the dynamics of Poz/(du2) where p S P and > P and assume both of up and By are small and (du²) < 1. This can be viewed as simplified version of the analysis of || || 2 in Stage 1.1. The analysis of other quantities/stages is essentially the same we rewrite the update rule to single out martingale difference terms and the higher-order error terms, and apply a suitable lemma from the previous subsection (or Section E) to complete the proof.
For the ease of presentation, in this subsection, we ignore the higher-order terms. In particular, we
assume the approximation 1, & R Ut.k + 21 (1k k < P) - [|0rrll") + +10)+++2200 1, for all E [d],
In an early version of this manuscript, we did not relax the conditions on the noises when X grows as in Lemma 4.1. We thank Eshaan Nichani for pointing out that this would result in suboptimal rate.
8

---

where I(+ represents the difference between the population and mini-batch gradients. Then, we compute
1111 (1+ 4n
v + 24 CLT V zz).
Here, the last term is the higher-order term and will eventually be included in &. For simplicity, will also ignore them in the following discussion. The second term is the martingale difference term. Its (conditional) variance depend on Uk. and this necessitates the induction-style conditions in Lemma 4.1. Note that 1 1.p / / 1.9
1p/ /0/2 /62 Hence, we have
1 4n (1 [(sr)(?))
2
of + 2pp ui + 21/6
" 1+1,
(1- 4nl 4n|10sell))
for any 0 Repeatedly use the clementary identity - (1- (1-atz) xe (1- ( and small > O, we can rewrite the above equation as
PUi+Le se (1 41+4n)
P
2nu 2Pmju d
dvi+1.q
de
Suppose that op og initialization and assume the induction hypothesis Prp/ (du)) (1 + 0.5)(1 + 4n)² Pup (dvā.q). Then, by Lemma 2.1, the conditional variance of the martingale difference terms (the last two terms) is bounded by Oz((1 + 4n))',² p2p /d). Using the language of Lemma 4.1, this means 02 < (n² p2 /d). Meanwhile, by our gradient flow analysis, the number steps Stage 1.1 needs is roughly log d/n). Hence, in order for (the second condition of) (6) to hold, suffices to choose 1) = (1/d). One can also show that for the higher-order terms to small, suffices to choose 1j = (1)(P))) As a result, for Stage 1.1, the sample complexity is Ö(dP).
4.2.2 The improved sample complexity for Stage 1.2
To why the existence of the sccond-order terms reduce the sample complexity from d to poly(P), first note that after Stage 1.1, maxpe[r] op will be 12(1/P). Also note that the conditions in Lemma 4.2 depend the initial value. With the initial value being 2(1/P) instead of ((1/d), the largest possible step size we can choose will be (Pap/Pp/p//--)) which is much larger than the usual ((1((Pd by gradient flow analysis, we know the number of iterations needed is O(p4/2-1 /n). Combine these and we obtain the ÓddPT-1) sample complexity.
requirement from the vanilla information exponent argument. Meanwhile,
5 Conclusion and limitations
In this work, we study the task of learning multi-index models form f.(x) = E d( x) with P < d, {ug be orthogonal and = Φ2hn + Ei=, gohy. By considering both the lower- and higher-order terms, we prove an (o bound on the sample complex for strong recovery of directions using online SGD, which improves the results one can obtain using vanilla information exponent-base bapsed analysis.
The main limitation of this work is the orthogonality condition. This can potentially be relaxed to near-orthogonality as in [OSSW24). Extending this result beyond near-orthogonal teacher neurons is an interesting but challenging future direction, as when the teacher neurons are not ncar-orthogonal, this task is hard in general. However, we conjecture that when the target model has hierarchical structure across different orders, online SGD can gradually learn the directions using those terms of different order sequentially.
Another limitation of this work is the assumption that the signal strengths are isotropic. When this is not true, training with the second-order terms and correlation loss will make all neurons collapse to the largest direction or require d² samples if we perform only one gradient step [DLS22, DKPS24). That being said, it is still reasonable to expect the overall sample complexity to be improved if we can leverage the second-order terms properly. Formally establishing this is also potential future direction.

---

Acknowledgements
JDL acknowledges support of NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994. We thank Eshaan Nichani for pointing out an earlier version of Lemma 4.2 suboptimal, and anonymous reviewer 4V3d for helpful discussions on relaxing Assumption 1.
References
[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairenn property: necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks. In Proceedings of Thirty Fifth Conference on Learning Theory, pages 1782-48877 PMLR, June 2022. ISSN: 2640-3498.
[AAM23] Emmanuel Abbe, Enric Boix Adserå, and Theodor Misiakiewicz. SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics. In Proceedings of Thirty Sixth Conference on Learning Theory, pages 2552-2622 PMLR, July 2023.
ISSN: 2640-3498.
[ADK*24] Luca Amaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita luvant: Data Repetition Allows SGD to Learn High- High-Dimensional Multi-Index Functions. June 2024.
[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from figh-dimensional inference. Journal of Machine Learning Research, 22(106):1-51, 2021.
[BAGP24] Gérard Ben Arous, Cédric Gerbelot, and Vanessa Piccolo. High-dimensions opti- mization for multi-spiked tensor PCA, August 2024, arXiv:2408.064M0 06400 [cs, math, stat].
[BBPV23] Alberto Bietti, Joan Bruna, and Loucas Pillaud- Vivien. Learning Gaussian Multi- index Models with Gradient Flow, November 2023. arXiv:2310.19793,
[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single- index models with shallow neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances Neural Information Processing Systems, 2022.
[BES 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, December 2022.
[BKM 19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal emrors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, Publisher: Proceedings of the National Academy of Sciences.
--64500, March 2019.
[CM20] Sitan Chen and Raghu Meka. Learning Polynomials in Few Relevant Dimensions. In Proceedings of Thirty Third Conference Learning Theory, pages 1161-1227. PMLR, July 2020. ISSN: 2640-3498.
[DKPS24] Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at a time. J. Mach. Learn. Res., 25(1), January 2024. Publisher: JMLR.org.
[DLS222] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural Networks can Learn Representations with Gradient Descent. Proceedings of Thirty Fifth Conference on Learning Theory, pages 5413-5452. PMLR, June 2022. ISSN: 2640-3498.
10

---

[DNGL23] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models. In Advances in Neural Information Processing Systems, November 2023.
[DPVLB24] Alex Damian, Loucas Pillaud- Vivien, Jason D. Lee, and Joan Bruna. Computational- Statistical Gaps in Gaussian Single-Index Models, March 2024. arXiv:2403.05529 [es, stat].
[DTA 24] Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborova, and Florent Krzakala. The Benefits of Reusing Batches for Gradient Descent in Two-Layer Networks: Breaking the Curse of Information and Leap Exponents. In Proceedings of the st International Conference on Machine Learning, pages 9991-10016, PMLR, July 2024. ISSN: 2640-3498.
[GLM18] Rong Ge, Jason D. Lee, and Tengyu Ma. Learning the-hidden-lee Neural Networks with Landscape Design. In International Conference on Learning Representations, 2018.
[GRWZ21] Rong Ge, Yunwei Ren, Xiang Wang, and Mo Zhou. Understanding Deflation Process in Over-parametrized Tensor Decomposition, October 2021. arXiv:2106.06573 [es. stat].
[1ch93] Hidehiko Ichimura. Semiparametric least squares (SL.S) and weighted SL.S estimation of single-index models. Journal of Econometrics, 58(1):11-120,, July 1993.
[KC22] Arun Kumar Kuchibhotla and Abhishek Chakrabortty. Moving beyond sub- uub--Gassianity in high-dimensional statistics: applications in covariance estimation and linear regres- sion. Information and Inference: A Journal of the IMA, 11(4):189-1146, December 2022.
[LMZ20] Yuanzhi Li, Tengyu Ma, and Hongyang R. Zhang. Learning Layer Neural Networks beyond NTK. In Jacob Abernethy and Shivani Agarwal, editors, Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Muchine Learning Research, pages 2613-2682. PMLR, July 2020.
Two-
[LOSW24] Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit, June 2024. arXiv:2406.01581 [cs, stat] version: 1.
[LY17] Yuanzhi Li and Yang Yuan. Convergence Analysis of Two-layer Neural Networks with ReLU Activation. Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[MZ03] Ron Meir and Tong Zhang. Generalization Error Bounds for Bayesian Mixture Algo- rithms. Journal of Machine Learning Research, 4(Occt::39 860, 2003.
[O'D14] Ryan O'Donnell. Analysis of Boolean Functions. Cambridge University Press, 1 edition, June 2014.
[OSSW24] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning of diverse features: computational hardness and efficient gradient-based training for ridge com- binations. In Proceedings of Thirty Seventh Conference Learning Theory, pages PMLR, June 2024. ISSN: 2640-3498.
[Pin20] Iosif Pinelis. Concentration and anti-concentration of gap between largest and second largest value in gaussian iid sample. MathOverflow, 2020. PRLLhtps//maatooerfooveee///////
(version: 2020-12-25).
[TDD+24] Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborová, Bruno Loureiro, and Florent Krzakala. Fundamental computational limits of weak learnability in high- dimensional multi-index models, October 2024. arXiv:2405.15480.
11

---

[Tial7] Yuandong Tian. An Analytical Formula of Population Gradient for two-layered ReLU network and its Applications in Convergence and Critical Point Analysis. In Proceed- ings of the 34th International Conference on Machine Learning, pages 3404-3413. PMLR, July 2017. ISSN: 2640-3498.
[Ver18] Roman Vershynin. Diaen-Dimesional Probability: An Introduction with Applications in Data Science. Cambridge University Press, 1 edition, September 2018.
[VGNA20] Mariia Vladimirova, Stéphane Girard, Hien Nguyen, and Julyan Arbel. Sub-Weibull distributions: Generalizing sub-Gaussian and sub-Exponential properties to heavier tailed distributions. Stat, 91))e311, January 2020.
[Wai19] Martin J. Wainwright. High--iimnssooal Statistics: Non-Asymptotic Viewpoint. Cambridge University Press, edition, February 2019.
[ZGJ21] Mo Zhou, Rong Ge, and Chi Jin. A Local Convergence Theory for Mildly Over- Parameterized Two-Layer Neural Network. In Proceedings of Thirty Fourth Conference Learning Theory, pages 4577-4632. PMLR, July 2021. ISSN: 2640-3498.
[ZS] 17] Kai Zhong, Zhao Song, Prateek Jain, Peter L. Bartlett, and Inderjit S. Dhillon. Recov- ery Guarantees for Onc-hidden-layer Neural Networks. In Proceedings of the 34th International Conference Machine Learning, pages 4140-4101 4149. PMLR, July 2017. ISSN: 2640-3498.
12

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: See Theorem 2.1.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as 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:
We discuss the limitations in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
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 use 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?
Answer: [Yes]
13

---

Justification:
See Section 2 and the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in 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 to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should 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: [Yes]
Justification: See Appendix F.
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, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all 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 to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [Yes]
Justification:
See the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nipp 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 (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /Codeeuomiisseeeeeeee
/nips. wwtpccccpppcccaecccccccccccccclt /onee ssisssolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
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 understand the results?
Question:
Answer: [Yes]
Justification:
See Appendix and the code.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, 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: [No]
Justification: Our experiments are very simple simulations with small variance across different runs.
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, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures 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 to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / /ttp///eerppp.//// cc htte/public//ee auublic/Ehhccsuuudeeleen. cwwuideluulliiie ines?
Answer: [Yes]
Justification: This is theoretical result and has no potential negative societal impacts.
Guidelines:
The answer NA means that the authors have 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., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is theory paper considering theoretical model.
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 socictal 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.
16

---

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 of 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
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?
Question:
Answer: [NA]
Justification: The simulations use synthetic data.
Guidelines:
The answer NA means that the paper does not 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.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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. Crowdsoureing 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: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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 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 the institution conducting the review.
16. Declaration of LLM usage
18

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs are only used check the grammars and polish the sentences.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https //tp:///eeu.pppeeeeee 22//e//eeeeee for what should or should not be described.
19

---

Contents
1
Introduction
1.1
Related work
1.2 Our contributions
2
Setup and main result
2.1
Input and target function
2.2
Learner model and the training algorithm
2.3 Main result
3
The gradient flow analysis
3.1
Stage 1.1: learning the subspace and preservation of the gap
3.2 Stage 1,2: learning the directions
From gradient flow to online SGD
4.1
Technical lemmas for analyzing general noisy dynamics
4.2
Sample complexity of online SGD
8
4.2.1
A simplified version of Stage 1.1
4.2.2
The improved sample complexity for Stage
5
Conclusion and limitations
Preliminaries
A.I Population and per-sample gradients
A.2 Typical structure at initialization
22
A.3
Concentration and anti-concentration of Gaussian
24
Stage 1: recovery of the subspace and directions
27
B.1
Stage 1.1: recovery of the subspace and preservation of the gap
28
B.1.1
Learning the subspace
28
B.1.2
Preservation of the gap
32
B.1.3
Other induction hypotheses
34
B.2
Stage 1.2: recovery of the directions
35
B.3 Deferred proofs in this section
39
C Stage 2: training the second layer
40
Proof of the main theorem
42
E Stochastic Induction
42
E.1 Deferred proofs
48
F Simulation
51
20

---

A Preliminaries
A.1 Population and per-sample gradients
In this subsection, we show that the task of learning the multi-index target function f+(x) Σ á z) can be reduced to tensor decomposition and prove the tail bounds in Lemma 2.1. For the first goal, will need the following classical result Hermite polynomials (cf. Chapter 1 11.2 of (14]) and correlated Gaussian variables.
Lemma A.1 (Proposition 11.31 of [O'D14]). For & N20 denote the normalized Hermite polyno- mials. Let p E [- and 2, be p-correlated standard Gaussian variables. Then, we have E
(hx(z))zz=')) = 1ik = jbp*.
2,8'
Lemma A.2. Under the setting described in Section 2, we have
p
E 2
[f.x())V x)] = ΣΣ
l (w w)--1 vz.
Proof. Let o == he be the Hermite expansion of & where the convergence is in L sense. For any p E [-1,1] and p-correlated standard Gaussian variables 2, 2', we have
E (0(=))(2)))
Σ
δuăi E (ha(2))u(22))
k.1=
where the first equality comes from the Dominated Convergence Theorem and the second from Lemma A.1. Therefore, we have
8
E [l.((z) x²)
E [o(o
2))(v x)]
o (vt, 4)
Then, we compute
P
P
E Vl
v (U w/
=
1ot (ui,
2x
k=1/=1
Now, we consider the per-sample gradient. The goal here is to prove variance and tail bounds for (f. (z) æ), u) where u E gd-1 is an arbitrary direction that is independent of First, we upper bound f.. To this end, we will use need the following concentration inequality for the sum of independent subweibull variables. It is a consequence of Theorem 3.2 of [KC22] and the discussion after Definition 2.3 of the same paper.
Lemma A.3 ([KC22)). Let (I:) exp(x*) and 11- denote the corresponding Orlicz norm. Let a S 1 and X1....,X, be i.i.d. mean zero random variables with variance o² and [| [X1ll,a < c0. Then, for any 6p E 10, we have
Xx
Vna log 12 (1/6=) + V log//(n//6>), with probability least 8p.
Lemma A.4. Suppose that Assumption holds and (vitn h are orthonormal. Then, for any ôp E (0, 1), we have
[++(x)) S (P/be, with probability at least 1 -8p.
Proof. Write Y == p0 x). By the orthonormality of {wi)s, (Ye)x are independent variables.
For any p 2 1, we have
[1 + =2ymn/3] 1/P
Dillze 5 (NL
(0P)(z)
c
E
x0.10.1)
< Cp"
E [11 + 22)9 S p².
21

---

where the first inequality in the second line comes from the Gaussian hypercontractivity. This implies that pl S and Var Y S Thus, by Lemma A.3, we have with probability at least öp
that
[f(x)) & VP Plgg 1/2 12 (1/6p) VPlog" (P/03) & VPlog"
Lemma A.5. Suppose that Assumption 1 holds and {vi are orthonormal. Then, we have
(f x), u²² Ss P.
Vol x), u) So p1/2 log2 21++) log(m/ ó with probability at least 1 8p, where q is the degree of @ if it is a polynomial and Q = 0 if ø Lipschitz.
Proof. Note that V x) u) d'(v x) (z. u). First, for the variance, we have
(xf'o x) (u, x))² 2 f E(o' (v x))² + xu,
where the second inequality comes from Assumption I and the hypercontractivity of Gaussian. This implies E (Vf(z), u)² < P. For the tail bound, first recall from the previous lemma that [f) 2 (Pllo"(P/6p with probability at least - ôp. The proof it also implies jo'(v (v(vx)) S log0((/0p) with probability at least - 8p. Finally, since (æe, u) is I-subgaussian, we have | ul S log 1/2 (1/62) with probability least óz. Combine these bounds, take the union bound over m learner neurons.
and we complete the proof.
A.2 Typical structure at initialization
In this subsection, initialization. Recall that we initialize v, with Und(() independently. Meanwhile, note that for U we have v d Z / Z|| where Z N' Ka).
the results in Section A.3 to analyze the structure of Ui.... Um at
We start with a lemma on the largest coordinate.
Lemma A.6 (Largest coordinate). Let v ~ Unif(S For any 2 1, we have
4v2K log
4
max
with probability at least 1 dk"
EEId
Vd
As a corollary, for any 8p E (0, 1), at initialization, we have
4/2 (log(4m/pp)
max sE|m]
|elloo
with probability at least 6r.
Vd
Proof. Let Z N (0. Ia). Recall that Z/ Z follows the uniform distribution over the sphere. By Lemma A.11, we have || |Z| 2 Vd/2 with probability at least 2 exp((-d/18) Then, by
Lemma A.13, with probability at least 1 - 2e - 2c -2-22/2 we have
max;eld (4 [z|
log d +
2 d
2s
+
vd/2
Vd
Vd
Let K 2 be arbitrary. Choose s
/2K log and the above becomes
maxie(4)
4/2K log
with probability at least 1
vd
dk
For the corollary, use union bound and choose K = log(4m/6p)) log d, we have
4//2log(Am/0) Vd
4m
max 0 10/1la0
5
with probability least
5g.
x6 m
dk
22

---

Suppose that we only have higher-order terms. Then, for a neuron v E Sd to converge to ground -ruth direction e₁ in a reasonable amount of time, we need E to be the largest among all of and there gap between it and the second largest vf. The following lemma ensures that when m is large, for every ground-truth direction (e*Ike]p)) there will be at least one neuron satisfying the above property. Note that in our case, we only need to ensure ei is the largest among all {oî he([)] instead of [v]] (e(d)) as the second-order term will help us identify the correct subspace.
Lemma A.7 (Existence of good neurons). Let óp E (e-log" 2P log(P/Ep) (P). Then, at initialization, with probability at least 1 öp, we have
log ,,1) be given. Suppose that m Z
vp E P| E [m]] such that
max U.. 200 200 1 1(1.
Vip
21+
Proof. Let do be a parameter to chosen later and let ōр,0 be the probability that ( og is smaller than 1 + do times the second largest oyg. For each p [P]. let Bp be the event Vk E [m]] 22.p S (1 do) max }
}. To bound P[Bp], we write
F[Bp] =
P D-Unni((3d))
2 up S (1 +
max
(34-1)
(P(P))))
PP
2 p
+ max P.
+ 1p $ (1 + c)
2 2
Up
max
P
Up
2
DE
PMMM
4
P)]
1
60,0
m
By Corollary A.12, if we choose 8p,0 1/2, then choose
80 = 200 log(48P)
With the above choices parameters, we have
1 )p S P (- (2p)
1 m
IP
B [E[
Bp
S P 1
2P
For the last term to be bounded by ôp, it suffices to choose 2 2P 2ogo (log((//5 P/8p).
Lemma A.8 (Typical structure at initialization). Let 8p E (ce-log log d, 1) be given and cg > 0 be a small constant. Suppose that {wg] ~ Unif(gd-1) independently with
m θ (Plog())6))
Then, with probability at least 1 35g, we have
Vp E [P] e [m] such that
Us.
1+
((1)
(9)
maxgelPMMp) [ūgg]
log P
20 Covlog(P/) (P/Op)
Vi [m].
| 5
Vd
VP
|gepl
3vP
Wi E [m]]
3
|
Vd
Proof. The first two bounds comes directly from Lemma A.6 and Lemma A.7 and the fact that we use symmetric initialization. By Lemma 1I we have
P
(11
Z - E
2
vd/2) < 2e -d/8
P -
>
vP/2) < 2e -P/8
23

---

As result, for any v a~ Unif(S" 1), we have with probability at least - 4e" P/8 that
Jusell
Z<pl
E zrlll + vP/2 [1/3,3] x V d'
P
||l||
|
E Z + dd/2
Since we assume P 2 log C' d for large C², we have 4e 4 S dz dz/m. This gives the third bound.
A.3
Concentration and anti-concentraatin of Gaussian
Lemma A.9. Let Zi,..., Za be independent N (0, 1) variables. Let Y1, Y2 be the largest and second largest of [2 b/.b.... IZa). For any óp E (0, 1), have
Y
6p
2
1+
with probability at least - ôp.
Y2
12 log (12d/6p)
Proof. The following proof is adapted from this MathOverflow answer [Pin20]. Let and F denote the PDF and CDF of Z with Z ~ N (0, 1). respectively. We have f(y) = 2 2 0} and F(y) = (rrf(g//2)1 2 0}, where is the PDF of V (0, 1) and erf is the error function. We will the following formula for the joint PDF of two order statistics:
fy Ya 32) d(d 1)F (oo)) < y2 < y1}.
Consider small O. We compute
Y1
z
1
1y 2 (1+ s)y2) (yn 12) dyzdy1
fo fy ((1 + s)y2 r. 32) dyzdr
=d(d- d(d-1) Fr (2=2(22)((ve) (%, S((1 s)32 r) (dr) dye
d(d dd(d F F(22( (1 F((1 + s)uz)) dy2-
Let G = F-1, With the change-of-variables u F(y2), = G(u), we can rewrite the above as
Y1
/ u d-2 21 F(11 + f((()) dG(u)
12 2 + = d(d - d(d-1)
- 1) / n=22 11-211F((14 ())) F'(((w) du
= =d(d-1)
d(d 1) d- F((1 + (G++)) du.
Now, we analyze the last integral. We will use the following expansion of the (complementary) error function:
v-2212
1 1- F = - erf (1/v2) = 3//7/2
1
2
dr
V
yy
For notational simplicity, put = G(u). Then, we have
F(1 + s)w) (1 s)w x/2
(1++)))w222
1
222-22 dr
V
(1tt)
exp (-(s + 82 22/2)=?)
w2222
1 0oo
r²e r2 -2 dr
1
1
T/2
V
(1-s)w
exp (-(s + 82 2/22ww²)
1 1-
vā (1
1
22 -2 dr
(1 - (1-((w))
1 1
1 1 22 dr
va Jw
24

---

Note that the last factor is at least 1 and F(w) = F(G(u)) = F(FF((F11)) (u)) u. Therefore,
exp (- 82/2)ur ()-u).
- F((1 + s)wo
1+s
As result, we have
Y1 P 2 11+
2 d(d d(d-1) S/"' 2 u)
exp (-(s + s² 52/2)G² (u))
du
Y2
1
2 d(d 1) /
exp (- + s² /2/22/ 22) (u) du,
u dd-2(1 u)
1
where E 0 is a parameter to be chosen later. By the next lemma, when a S 1 - E, we have (u) - 2 (2g(( Therefore.
Y 2 1+
/-¹ " 21-2(1- - uddu
exp (- 52) log(2/s))
2 d(d ((d-1)
22
1
2 d(d - 1) /
p-=
1
4a
u) du
+ s
2
de
= (1 5)d 14 1-8 1
2
Let 83 E () be our target failure probability. We choose
de
6p
5y
(1 =) =)d
1
2
E
-
3
6d
With this choice of E, compute
4s
8p
48
62
8p
4s 4s
2
8 5
3
12d
3
12d
3
12 log
12d
dp
Note that for 8 satisfying this condition, we automatically have 1/(1 + 8) 2 1 - 0p/33. Thus, we have
Y1
ôp
with probability least 1 8p.
2
12 log
hp
Lemma A.10. Let F(y) = erf(w/V2)1 2 0} be the CDF of I (0, 1) respectively. Let 1 d), then G(u) I (2oo(22/e).
Proof. Note that G(u) S M iff a S F(M) iff -
F(M) iff 1 u 2 P(ZZ M). In
other words, our goal here is to find the smallest M such that P(Z Z M) < E. By the standard
Gaussian concentration, we have P([Z) Z M) cxxp(- 22)2). For the RHS be upper bounded by e, suffices to choose M 2 / log
Lemma A.11. Let Zi...., Zm be independent N (0, Ia) random vectors. Then, for any we have
Vk E [m]. P
1
2 1 - -2me -5/333
Proof. is well-known that any 1 1-Lipschitz function of / (0, I&) is 1-subgaussian (see, for example, Theorem 5.2.2 of [Ver18]). Hence, for any O, we have m
I
max kE[m]
||| ||
E
z
P 2xll
E
zx
2 S 2me -2/2
Set
|| E zx|
and the above becomes
||
IP
Vk [m],
S 2
2me -
E æ1
To complete the proof, it suffices to note that |||]] || follows the Xa distribution, and therefore we have E || 1ill 2 vd(1 - 2/d).
25

---

Corollary A.12. Let v ~ Unif(s and wi and w2 denote the largest and second largest of |oi|...., /vpl. Suppose that log d 2 8f
(P/8)) and óy (e-log" log," a 111. Then, we have
w1 21+ 100 (24P/6pp)
Oz
with probability at least Ög.
W2
Proof. Let N (0, Ia) vectors. Note that v ōp E (0, 1), that
2 =/||=||- By Lemma A.11, we have, for any
1og(2/00) d
with probability at least 6p.
Suppose that |zk, and |24,| are the largest and second largest of [2r|,..., /zpl. By Lemma A.9 (with d replaced by P), we have
Fkl
87
21+
with probability at least ōp.
12
12log(12P/6g)
Write
3log(2P/6))
W1
1
8p
1
1+
8p 12 log (2PP/pp)
3V
3 log (2d/2/)))
W2
12 log log
(2P/
1+
(2gg(2/
In order to merge the last term into the second last term, it suffices to require
8p
3 (log(2)/6p)
log 1og2(P/6)
61
12 log
d
log d
of
Then, with probability at least
26p, we have
8p
+
W22
24 1olog122//g) (12P/pp)
Replace &p with 8p/2 and we complete the proof.
Lemma A.13 (Upper tail for the maximum). Let Zn...., Za ~ N (0, 1) be independent. have the upper tail
Z 2 22 d + 8) 20
max
Vs Ys2 0.
rEld
Proof. For notational simplicity, put Z max Zr By union bound and the Chernoff bound, we have for each s, 0 > 0,
d P(Z 2 =P
EC
V 2
<dP(Z 2 8) S d
de de 1/2 12
c
Choose 0 s to minimize the RHS, and we obtain Piz > 8) < 22 /2 log + 52 and this becomes
Replace s with
P pz 2 22loog + < <p [2' 2 + o2 S 60
Use the fact min Z, maxield Zi and we complete the proof.
26

---

B Stage I: recovery of the subspace and directions
In this section, consider the stage where the second layer is fixed to be small value and the first layer is trained using online spherical SGD. Let e be a first-layer neuron that is good in the sense of (9). Assume w.l.o.g. that is the largest. Our goal in this section is to show will converge to close to e1 with probability at least 6ạ at the end of Stage 1.
For notational simplicity, let leorr and Leory denote the per-sample and population correlation loss, respectively. By Lemma 2.1, we can write its update rule as
0z+1
8 = vz + 1/V Lcorr + n
1++
where Ze+1= = (I - v0 ( (Vrboorr( (x)
VoLcor)) and, by Lemma 2.1,
-oooo = -(I vo /V
P
P
= 22
Uk us - - vo )ek
Jekk +
1- (I ww Jes.
k=1
12 ==
In particular, for each & E [d], we have
1++1, ve + 2n 11k P} - ||0grp1?) uk + Lnoi (1(k < Plut- (|0rpl8) 24
Σ loj (1 < p)ut- ||0gr(l1) oo 7
+ 11
b
= 04 + 1/1 (ks S P} 223 + uf + Σ 189
1/ 203 /gpll2 + Ldz ||uppll + Σ
100 |105
Un + 11
1>L
For notation simplicity, we define
:= 2 |spll2 + Lo / +
Σ
109 ||<pll
(10)
1>L
Note that p is independent of the coordinate k and we can write
111 0r + 11lkk < P}
2 L6int L-2 + Σ
10fok 1-2
nppu + 17
(11)
CST.
For the martingale difference term Z. note that by Lemma 2.1, for any u Sd-1, (Z u) is (M2, 0) subweibull variable with Mz p1/2 and 1/0 2(1 + Q). In particular, this implies
| [Z u2) So Mz 20+2()))) log(d/5) =: Mz, with probability least 1 6 / (12) where C' > 0 is any fixed constant.
In addition, have the following lemma the dynamics of vz. The proof is routine calculation and is deferred to the end of this section (ef. Section B.3).
Lemma B.1 (Dynamics of vg). For any first-layer neuron v and k E [d], we have
0f = UR + 22 + 21T Z1t1, + E1+1,k
where Th = 1kk < P} 2 1 Eis1 10f
is Fe-meauurable random
variable with mr,al < 20 and (Er+1)ee(T) is (uniformly) bounded by (n(r2 with probability least
8p. 8p.
To proceed, we split Stage into two substages. In Stage 1.1, we rely on the second-order terms to learn the relevant subspace. We will also show that the gap between largest and second-larg -lagest coordinates. which can guaranteed with certain probability initialization, is preserved throughout Stage 1.1. These give Stage 1.2 a nice starting point. Then, we show that in Stage 1.2, online spherical SGD can recover the directions using the L-th order terms.
'ee will often drop the subscript i when it clear from the context.
27

---

B.1 Stage 1.1: recovery of the subspace and preservation of the gap
In this subsection, first we show that the ratio ||_rl||* /|>>>|l will grow from R(P/d) to (1) within O(dP) iterations. We will rely the second-order terms and bound the influence of higher- order terms. This leads to the desired complexity. The next goal to show the initial randomness can be preserved. In case, only need the gap between the largest and the second-largest coordinate to be preserved, which will ensure that the neurons will not collapse to one single direction. Formally, have the following lemma.
Lemma B.2 (Stage 1.1). Let E g be an arbitrary first--l first-layer neuron satisfying || I log² d/d.
||u_ell ||>p|l 2 P/d, and u = (1 + 5u) ge( Di at initialization. Let öy be given.
Suppose that
%
P
8 dP
P p log d and 1 So aP logg d
= 0,
M2
M3 log
0+1 (d/op)
Then, with probability at least ((p)), we have
|V<P 2 2
within T 1 o(1) log
= (dP) iterations.
||eee||
4021
P
In addition, at the end of Stage 1.1, we have 2i = (1 + 6u/2) gggmax,e(PMMPe 4.
Proof. It suffices combine Lemma B.4 and Lemma B.6.
To prove this lemma, we will use stochastic induction (cf. Section E), in particular, Lemma 4. 1 and Lemma E.3. For example, to analyze the dynamics of ||rlll /// 1|" it suffices to write down the update rule of || | and decompose it into a signal growth term, higher-order error term, and a martingale difference term as in Lemma 4.1. Then, we bound the higher-order error terms, and estimate the covariance of the martingale difference terms, assuming the induction hypotheses.
The induction hypotheses will maintain in this substage the following:
| |
10, / 3
log" d P
(11) 402)))
2 s
(13)
They are established in Lemma B.4 and Lemma B.7.
B.1.1
Learning the subspace
Now, we derive formulas for the dynamics of the ratio ||/r|| / rll. As we have mentioned earlier, the goal here separate the signal terms, martingale difference terms, and higher-order error terms.
Lemma B.3 (Dynamics of the norm ratio). Assume the induction hypotheses (13) at time S T. Suppose that (dnrz) Let be arbitrary first-layer neuron. Then, time have
|14115r112
||rr|l (1 + +noz + 2nE.) + He41 + Ett1.
|1141,1111
||
where E. = Eizl 107 ||0gell| || where (He+1): is martingale difference sequence that is conditionally (Os((L+ 402n)² 2), subweibull, and (Ex), is an adapted process with 1s+1l So (1+ '1 MM MM P for all t E [T' with probability at least - Sp.
Proof. First, recall from Lemma B. that 1 022 21 a + +1 23, 77 2222111 + 2n U Zz+1.
28

---

where The 1[k I P} (203 + L 2 11>110ffk
1070/
p is a Je-meaurable random
variable with 17c - 20% First, for @لrاl we have
||0+1, (1+ 4n02 - 2mp) ||ura|1 + 2m Σ
|| + 2n (USP, Zsp)
12L
1 180 ||gpp11 + 21)² 1z
Similarly, for |/0s) pll, we have
| |/011.2P1 ||02||2 - 2np |1 222 p
p) =80,n² |usp1l 21 + 2nj² 1 pl1
For notational simplicity, we also write E, Ep Eizr l0f || || Note that by Assump- tion and the fact that ||vgrll// /|oss// 2 11 S C. Since || ///5/ = /| @0s we have
|041,5p112 |0+1,>P11
(1 (1+44) 4n 2np + 21T50)
||ur||
(1 (1-
2n (U>P, Z>p)
E>P
(1 - 2np) || U |>>p1
|/++1,,1|1*
|0++1,111?
+ 2n (VSP, Zzp)
ESP
111+,,>2111
|0(41,>p!1 2
1 + 4n03 - 22p + 21TEE
||
(1- - 2np) |0 pl
+ +n103 20p + 2//
||vgr||
2n (Usp, Zsp)
2n (USUpP. Z<p)
(1 (1-2np) 2n1p) ||5> ||
|0+1,,1||
|1+1,
+ 4no2 2mp + 2nE
||ur|||
ESP
ExP
(1 2np)
|
|011121111
10++1,2111
|1141.P11
Tı
T2
+ Ta
|
where each T; represents one line. Note that, up to some higher order terms, T1 contains the signal terms and contains the martingale difference terms. Now, our goal is to factor out those higher order terms.
For Ti, recall that |p| Că and use the fact that
=
222,
1/2 $ 1/2,
(14)
It
to obtain
1141,PP11
|u<p 12
Ti
+ 4100 2np + 2ne.)) (1+ 2np + 41)pp)
|141,211|
| 1
| |2
(1+ 4no2 2ne, + 30C
|upp|' 2
Now, we consider
||gp||
+ 4n0 - 2mp + 2ne 21 (U>P Zsp)
2n (0Sp, Zep) |0 1
T2
||
1 2np
|
29

---

First, we estimate the 1/ ||111, We write
1
|
(1 2mp) ||0rr
+ 2n (vyr, r. Zyr) + ExP,+11
1 2n (U>P, Z>p) + +
(1 2mp) || 1
10++1,>11
By (12), we have with probability at least 8p that
4>r- Z>rl A |ugr Zerl A max
[ So Mz. kE|d
Note that the above conditions also imply
Ikspl - 80 uusel 1| + 22 So PM
|s>1] S " /
1 + 2r/ddNzz So ddd
By our definition of Stage 1.1, we have |
1/2. Then, have
1
1
1 + 4n + 320
||+ (1 - 221pp) |0>p|
1
[ + (s(Mi)))
(1 - -2mpp) |/0ي,|1 2
Using the above two estimations of 1/110/11,>>11 2 we can rewrite T2
T2 = | 2 + 4n03
2mp
27cr
2n (usp, Zyp)
(1 (1
||u>p|1
1 1 -2mp
(1-2np) ||
2n7 (u<p, IsP)
+ (1 - 201p) |/0>p11'
|usel| 1|
+ 4no2
2np + 2nE
2n (U> Z>p)
2n (u<p, Zsp)
+
|
2np
(1 - 21p)) ||>5p1
(1 - 20p)) |00ppl
|erlll
||rr||
20
(s + 4
nAl (s(AAiz)
|0>p||
|
||up||
2 +
4n00
2np0 + 27T.,
2n (U> Zxp)
21 (v<p, Zzp)
+
| [
2mp
(1 2mp) ||>م|11
(1 2mp) |0yp|1
||><p|| ||
+ O
1/2
Finally, consider the third term
11 + 4n08 -2pp + 211
|grl]
ESP
E<P
T3 = By our previous bounds on &, we have
(1-211) ||
|
|
1| Car2
|
Tal s 64
+ 32
PM & n
PV
| 2
|u>p| |0++1,2|1'
30

---

Combine the above bounds, and we get
|011,<p11*
|ugel"
(1
4n12 + 21/e,) 12
Us> pl
Uugp
4m02
21p + 21je.
21/ (U> p, Zsp)
2n (0gp, Zep)
12
1 1-270
(1 (1-2mp) pl
(1 2np) |0ypl
|uspll
30'
10000
17 12103 V ndd(z)
|
" 10 pl
U PV
+ 10
2
10>
Let H++1 denote the second line and f++1 denote the last two lines. Recall our induction hypothesis
|| |
(1)(11+403))
|| ||
4(1++403n) P/d). Meanwhile,
mgr
note that M>
Then, we compute
P d P " + (1+40) So (1 + 402))*
M2 (IVndMZz) + n2 M2 + 4d2n)*r
Ki+1l
d
VV p
So (1 (1+ 17 M2 (VPdmMZ V VP)
So (1 40 N)
Then, consider H141- We have
|uep|l 2
4n02 2np + 2n, 2n (V>P, Zyp)
2n (0< P. Zcp)
He+1s <
+
||>p|11
1 2mp
(1-) || /
(1 - 21pp) |05 ml
P 11 2 (1 422n)²
P. (1+ Z>p)++ +
1/2
(U>P.
d)
7 UUEP Zep)l-
d
Since both (0>p, Z>p) and KTEP, ZepYi are conditionally (P,,0)suuwweibul,, M4+1 is condition- ally (O&((1 + 483n)² ( subweibull.
L
With the above formula, can now use Lemma 4.1 to analyze the dynamics of the ratio of the norms.
Lemma B.4 (Learning the subspace). Suppose that
p2
P
P > à log" d and n E dP
=
=
M2
M3 log
0+1 (d/ōz)
dP
Then, throughout Stage 1.1, we have (1 + 40%n)*
|
|
3(1 + 403n)²
|00, p
5
|
|
and Stage 1.1 takes most (1 + (41)) log (d/P) Ös (dP) iterations. For this result to hold the P good neurons (satisfyyng (9)), it suffices to replace oy with Sp/P.
Proof. First, by Lemma B.3, we have for any ST
||
(1 4n02 + 27T.)) + Het1 + E1t1,
||141,>1|1
||aع,|1
where E. Σizl 107 ||ugell/ // where (H++))s is martingale difference sequence that is conditionally (0s((1 + 4033)) 2),0) subweibull, and (Et), is an adapted process with
31

---

15tt1l So (1 + n for all t E [T] with probability at least óp. By our induction
hypothesis oz < log d/ we have
0g log 22-2 C2log p1/2-1 =: =:.
- E. = ||
U
S
lot
|ugell 1-2 &
12t
12L
In particular, note that 8, does not depend and is o(1). For notational simplicity, let X₃ | el0= ||>p|l 2t (1 44) and 2 + (1+4n(1 + 6,)) xt will serve as the lower
and upper bounds for the deterministic counterpart of X, since
(1+ 40221) X + Ee+11 + He+1 & Xe+1 < (1 + 453n(11 + 5.)) Xt + E1t1 + Hi-1-
Moreover, note that for any t I T, we have
zi
1 + 403m(1 5c)
(1144n(1 +5,)) (1-403) +
2t
1 + 4031
(1 402n0, =00(1²))
S exp (0 (6, + n))
Since T So log d/n, the above implies
1 S exp (1+(1) log d (6, + +)) S 1 51+0 log (6. + 11) + 0(),
Zt2
where the last (approximate) identity holds whenever
C2 log
8, logd
P> log² d.
p1./2-1
logd
In particular, this implies that the (multiplicative) difference between ot and oe is small. Now, apply Lemma 4.1 to X₁. In our case, we have
P 0 20 n M2.
So o
4(1 + o(1))22n and X, (PP/d). Recall that T Zo log d/n. Hence, meet the conditions of Lemma 4.1, it suffices to choose
Xo T
1/² PM3
So
+
So
dM3 log d'
P 11² ā M2 So
Tlog #+1 (T
dM3 log d log (T/
Then, by Lemma 4.1, we have, with probability at least ((6)) 0.5x, S X, I 1.5z.". Since o (1 + o(1))=7. this implies 0.5xr, S X, I 2x,, To complete the proof, suffices to note that for to grow from ((P/d) to I, the number of iterations needed is bounded by (1 (4d)))))))) log (d/P).
B.1.2 Preservation of the gap
Now, we show that the gap between the largest coordinate and the second-largest coordinate can be preserved in Stage 1.1. Let p = argmax;e[p] v0) and consider the ratio v / where ạ E [P] arbitrary. The proof is conceptually very similar to the previous one, except that we will use Lemma E.3 instead of Lemma 4.1.
Lemma B.5, For p = e (0) and any g E [P], we have
v2 + Hes1 + Ec-1,
where (He+1)) is martingale difference sequence that is conditionally (Oo(n' d.))), -ubbeebibl, and (Ee), is an adapted process that is uniformly bounded by (s(rddM2 ) with probability at least 1 8p.
32

---

Proof. For notational simplicity. define Pt.g/p := / Our goal is to upper bound Pt.a/r- Recall from Lemma B.1 that for any i f P. we have
&+1AA= = - + 2 + 2 + V + +1² 2222 + v Vt, 1
where Th,t 203 + ToOl 1 + Eisl 10 p. Then, we compute
(1 2)))) v²., + 211 - + Ee+1,g
Pt+1.//p -
(1 + 2)p(t.p)
1 + p
ELtLP
(1 2)))t
(1+ 21))) u2. 2 21+1
2910s,0 Zt-1.9
+
(1 + k))
ut.pp
(1 + 21)tt.p))) uL.p
(1 + 217 of
+ f++1.
(1 + 21y)t.g) u E1+1
21p0 Z
27 Z Eit1.
1241,
(1 + /t,)) ui
22
(1 + m),p
of
09 1i+1.
=: (P((L111.//))) + T2 Tp(ptt../p):
where I contains the first term (signal term), T2 contains the next terms (approximate martingale difference terms), and T3 contains the last line (higher order error terms).
First, for the first term, compute
Ti = Pt.g/p +1+ (227))))) 11-21/ 11
+ 2077.
Pt.a/p (1 22772) (1 20771. = =5)("(;p)
Pt.f/ (1+ 21 (7t.g Yt.p) + 20m2 (nÊ.p V (igg)
Recall that hr,u < 20 and note that 71.4 - M1 = Σi>l |07 1 Σizl 107up < 0. Hence,
Ti S Pt.a/p (1 + 80C
Now, consider the martingale difference term
1 1 21, 210t,pZt11 +
21 Zitie
To = Pt.a/p + 277
1/t1.p
+1+ 217 oi
We rewrite the denominator as
11+ ui 2f 2n + Eltlp
1
2m Z141.p + ELTL.
"t utp "i + 20
o2t1,p
By (12), with probability least 1 - 50/d0, we have
kp.t_1 So 12 + np MM + 1/ Ie MMz So 172 M2Z
Therefore,
1 ndlz
- #0
02 1+1.p
r2,p 211 12
Vtpp utp
Then, we can rewrite T2
1 + 21/Yt. 21pUt. 7.t1,
27y0', 7t
Nr2
T2 -tt,// 2171.p "²,p 21
+ Os
112
(1+ 21y ) 12
H,t1 + 0 (1²"
Note that H4t1 is martingale difference term with
1/ Z
HH41 So
So mVd (1741 17e+1
Ut.p
2
33

---

Since ZitIp and Z++1. are both conditionally (Os(rfdM2), 0) subbeebbbll Finally, consider
(M2, -usbweebbll H4+1 is conditionally
f++1
(1 + 21) vi.q EL+1.
2 21-11g 2yputpZZ41, + Ee+1.
(1 + 2)))tp)) vi 02 s1i4
(1 + 21/ vř,p
0/t,.p
Since |5,,411 V (p,,411 | $ 1 M2 and 1etlppl V 1/11-1.l S Mz, we have
(Tsl So rddMz..
Combining the above bounds, we get
Ptt1.//p S Pt.g/p +1+ 800 + H4t1 + Os (npdM13) = =Pt.ap + Het1 + Os (1fdM13)
Lemma B.6 (Preservation of the gap). Consider de E (0,1), [rgmax,e(P] (0) and any
q E [P]. Suppose that
8
P
P
8e
S dPlog d
A
= =
M2
M2 log
0+1 (T/6p)
dP
Then, we have
sup E<T
1.g
Vū.q
S
6
with probability least 6y.
V
Ep
Proof. By Lemma B.5, we have
vi.g He+1 + E1+1.
2
1.4
<
2
+
1tp.
where (He41), )r is martingale difference sequence that is conditionally (0+())dM2), 0) -ubbweibull, and (Ex) is an adapted process that is uniformly bounded by (177ddM))) with probability at least
1 dp. Hence, by Lemma E.3, we have
vo.q
So Tnp TrRdM + nPdM3T 9+1 (
sup tT
0.
So 1 2 log d + VndM3 0+1 (T) logd
For the RHS to be bounded by 5, E (0, 1), it suffices require
ndyli log d So de + So
dMz log d'
(T/6p) V /ndM3 1 log d. So de #
82
So
dM3 log
0+1
(T (T
logd
B.1.3 Other induction hypotheses
In this subsection, we verify the induction hypothesis: 0z < log² d/P for all p E [P]. This condition is used to ensure the influence of the higher-order term is small compared to the influence of the second-order terms.
Lemma B.7 (Upper bound on v.) Suppose that
log d dP
P M2
P
So
^
M2 log (d/8p)
Then, throughout Stage I, we have o < log d/P.
34

---

Proof. First, by Lemma B.1, for any p S P. we have
1t1pp S up.pp + +2n
202 +
1oput
Vi.p" 21 Lttip
12L
+ 1. "t + 17 2
Tt. Ut.p 1
=
C² log 6-2 d log--2
$ Upp 4020
1+
0/ + 27p0 Z Ee41 203 PL12-1
where 7p.t == 20- + -
1-2 + Ei>L
10i kk 1-2 is F.-measurable random variable with
httpl s 20 By (12), with probability at least 5p we have
ketl So v2 + 1) M2 + 1/² Mr,p Mr Eo 1j M3
We maintain the induction hypothesis " s 2 + 4 8.)))² log d/d. Under this induction hypothesis, we have
2n S // (1 + 4337(1 + 5.)) log² Z
and therefore, is (Os (17211 + ( 5x)) 13, -sbweibbull Using the language of Lemma 4.1, we have
log d M2.
Eo 772 M and 2 So 7 402n 8,))
Therefore, as long as
log 2 d/d
log d dM3
M2 So
4=
11 Zs
T
2 log
x
log d
M2 Zs
So
d
Tlog
8+1
(T/
dM3 log 8+1 (T/
we have o2,p $ 2(1 + 453n((
+ 6, log² d/d throughout Stage 1. In particular, by Lemma B.4, this
implies
log² d log" d
1 Z +5. (403nT)
B.2 Stage 1.2: recovery of the directions
Let be an arbitrary first-layer neuron. Assume w.l.o.g. that of is the largest at initialization and / maxagkg "i 2 + dg. By Lemma B.2, we know this gap can be approximately preserved in the sense that / maxg< v/.x 1 + 6 holds. For notational simplicity, factor 1/2 in the sequel. Moreover, since we symmetric initialization, can further assume that V1 > 0. In this subsection, we show that of will grow from 2(1/P) to 3/4 and then to close to 1. Formally, prove the following lemma.
will drop the
Lemma B.8 (Stage 1.2). Let E gá-1 be an arbitrary first-layer neuron satisfying vf 2 c/P and / max2çksp 1 x 2 1 + e ffror some small universal constant c <> 0. Let og (0, 1) and
e. 0 be given. Suppose that we choose PPA
80
P
d
P
e.dP
n So JPL/2
M2
M2
log
0+1
(d//p)
dP
M2 log(1/e.)
M3 log d log 201 (d/67)
Then, with probability at least 1 1-0 (8Fr), we have of 2 1 1 within Og ((PL/2-1+ +log( (og(1/e.)) /n) iterations.
35

---

Proof. It suffices to combine Lemma B.10 and Lemma B.11.
Lemma B,9 (Dynamics vz). We have
1+1.1
2
2 U1,1
1 2n
Σ 12..
10
2n
Σ 12L.
loi
|er H1t1 Ei41,
where He41 is martingale difference term that is conditionally (f, O)-subeeibull and S1+1 is bounded by Do(T)dMM3ui, uniformly over E [T'] with probability least 1 òp.
Proof. Recall from Lemma B.1 that
1+1, vi, 227 2 2 Ec+1,
where Thet 1(k - P) 203 + LØ2 u
1-2 E1st L0f 1-2 i F--meauurable random
16f
variable with lyc,l S 20% and (Es+1)eetry (uniformly) bounded by Os(1" M))) with probability at least 1 - 6p. Sum over E [d] and we get
0++1 12 11 11+ 2n
Lik S P}
20 2 Σ
1ōfut Elofut
-0)
vi + 2m (Ue, Zit1) E4+
12b
= =1+ 2m
202 |1,ppll? + Σ
1ot |urgell - p |10011?
+ 2n (ve, Zet1) + El+1
12L
$ 1 + 21
2202 - P) + 2m
Σ
lot
|engell +2n (vr. Ze-1) + E1+1.
12
=:N3
where E is bounded by (M()) Recall from (12) that (ve, Zc+1) | So. Mr and choose S (dM22 -1 As a result,
2n (v,, Z(41)
Si+1
(1-
2n (1,, Zet1) + S1+1 )
16+11
N2
N3
N3 + 21/ (vee Ze+1) + Ec+1
1
2n (ve, Z++1)
+ +o((?ad13))
NZ
N
N
Meanwhile, have
6241, = + 21 (203 - p) uf + +2
Σ
l0f + 2 + E4111
121
where 16+11,1 S6 172 172. Therefore.
2n (Ur, Z(+1)
1, 2 12
+ +0001r
N2
NZ
N3
67711.1
02 2n (vi, Ze+1)
+ 0n(nndddz3) d2)3
N2
N?
N3
0f 2n
2203 p)
v2 2n Eizl
N2
v2 + 21
20 p) v/.1 + 2m Eizz 102f10ff 2n (ve, Z
21fu,,1 Z441,1
N2 1 #0 (s² 1fdN22 20(1
N2
N2
T Ti (,²+1,1) + To (vi+1,1) + Os (1PaN13321))
36

---

where we have used the fact that of,1 2 1/P to merge enror terms of form M/ into Os (rPdMZZu))) Meanwhile, for T2, have [Tal So by Zenial + [meg, (ve, Z.+1)/. There- fore, is (0 uz M2), 0) subbeeibull For the signal term Tı, by (14), we have
of,1 + 2n
2 P) uf.1 + 2n E12r 107of
Ti
1+ 27
203 P) + 2n Eizl 10, ||0,
21,1
1 21
202 p) + +2 Σ
10701-1
1 11-27 (2Ø3 - p)
(2
2n
Σ
lo7
01 |r + (17)
12L.
127
1 2n
10f 1-2
2n
Σ
1ot
|e (1²)
12L
12L
Combine the above estimations, set H4+1 = T2, and we complete the proof.
Lemma B.10 (Weak reocovery of directions). Suppose that choose
60
P
d
1
dPL/2
M
M2
log
0+1 (d/oz)
Then with probability at least (69)) we will have of 2 3/4 within the following number of iterations:
PL/2-1
PdpL-2 (
= Os
Mr
M3 log
911 logh*1 (d/ōp)
Remark. Note that when M3, M3 = (P). then the above is roughly P x (dPL-2). The dPL-2 the usual bound for online SGD when the noise has order d instead of The first P comes from the fact that there are P directions.
Proof. By Lemma B.9, we have
14111 2 oi,1
1+ 27
Σ 12L
187u1-2
2n
Σ 12L
1si |errll H441 E+t1,
where Het1 is martingale difference term that is conditionally (1) M2 vf, 1), 0))subbeibulll and E++1 is bounded by dNf33N for all t E [T] with probability at least 8p. For the signal term, we write
of |00l 1 of
Σ k=2
1/ 2 of (|0gr1122
-z
1-2
Uk
ugp
uf
k=
Note that the last term is weighted average of - Similar to the proof in Section B.1.2, one show that the induction hypothesis of/ max25kgp 2i 2 1 50/2 remains true, which gives
P
2/2-1
L/2-1
2 I-2
Σ
1-2 k
<
max 25k5P
2 S
<
of
1 0
1 80/2
'The only difference is that now the L-th order terms cannot be simply ignored as we no longer have the induction hypothesis v < log² d/P. To handle them, it suffices note that if of 2 oi then those L-th order terms of of are also larger, which will even lead to amplification of the gap. In fact, this is why we can recover the directions using them.
37

---

Therefore,
o11-2
2 - -1|0l1 2 of 21 os) (10gp112-1t) -r9)
1 50/2
1 - (1 do(1 24t) - |1grl(²)
60 y (1-1f)
= 2
1+ 6
As result, for the signal term, we have
of + 2m
|0 01 1-2
21
Σ
10i ||mel||
2 2 L0 (1 (1 nou
12L
12L
In partieular, when uf $ 3/4, have
Løz donor
0f 2i +
Het1 + E1+1-
Thus, using the notations of Lemma E.5, have
Lot dOn,
E 3 dM?
o of
p= L/2,
zo 0(1/P).
To meet the conditions of Lemma E.5, suffices to choose
a S 2
p-1 00
20
do
E S and p-1
11 S
dM3 2
21/2-1
0 S0
n S
log
(oo((//o)/(az///)
5p)
M PL/2
log 0+ (d/6p)
Combine the above and we get the condition
50
P
d
1
dPL/2
M2
MZ log
(d/62)
Finally, we apply Lemma E.5 to complete the proof.
Lemma B.11 (Strong recovery of directions). Let u E Sd-1 be an arbitrary first-layer neuron. Let ôp and e. be given. Suppose that we choose
E. n S6
P
E.dP
dP M2 log(1/e.)
M3 log dl log 01 (10++(dd//))
(d/5p)
Then, with probability at least 1 O(p),, we have of 2 1 - E. within (log((/EE.)) iterations.
Proof. Again, By Lemma B.9, have
1 2 of
1+ 2n Σ 10fot-2
2n
Σ l07 |0 H141 E
12L
where Het1 is a martingale difference term that is conditionally (O& M2 -)suuweiibull, and Ei+1 is bounded by Os(n d for E [T] with probability at least őp. Meanwhile, by the proof of the previous lemma, we have
2
1
1 2m
10 1-2
2n
Σ 107 |eslll 22
2 or + 2T.
Cg.b (1 vf) not
1 Cg.
12
2L
(3 (3)
2 r + = + Cg.b
(1 of)
4
=: up + 7c (1 - v?)
38

---

for some constant cg. 0 that depends only on Cg.L and Φ. Thus,
1-
2 uit,, 2 (1 - vs) nCg.L (1 -2î) -H++1 E(+1.
In the language of Lemma 4.1, we have
a = L,, r/I' So log(L/E.), To meet the conditions of Lemma 4.1, it suffices choose
22 So M2
E So rf did...
Ex ES
E dM3 log(1/e.)
11 So
T
2
E?
2
7
Tlog 9+1 (pp)
M log dlog 0+1 (d/ōp)
Under the above conditions, by Lemma 4.1, we have of 2 E, within T (s iterations with probability at least - ((p).
B.3 Deferred proofs in this section
Proof of Lemma B.I. First, recall from (11) that
81+1x= = Ve, + n1 - P}
20 + L6
1- + Σ
107 )
Uk mppkk + 11
121.
Therefore,
0z 1,k v2. + 21puE
[1 P} 2022 LΦiuk
1-2
+
Σ
101
1-2
11k ppok Z
131
2
7? 11k 1kk PP 20 LØIut
2-2
+ Σ
10,
k
1-2
pux + Ze+1
1>b
=: 22 + TT (b)+1,) + +T2 (1it1.))
For the first term, rewrite it as
T = 2
1[k < PJ
2 + L 2-2 +
Σ
10t
1-2
-1)
+ 22yu1. 2 I
1>L
Consider the second term. For notation simplicity, put
The == 1 (k s P)
203 + I. vt-2 Σ
187
p.
12L
Note that Th.t is me-easurable and by Assumption 1, we have
202 + LØiut- Σ
of
2
Loi
Σ
l01 S C2.
12L
p= 202 0-pll²2 + LQz |u
Σ
103 ||<Pl 22d + Lo7 + Σ
101 - C2.
1>I.
1>1.
and therefore (Sr.ll S 202 p.-
Then, we compute
T2
= (7 + 2+ = "2, + 7241, + +2,,, 2
2
Combine the above two bounds and we get
02+1 of + 2U,,, + 2270r
+ T2 v1,6 + +12 22 + 2n12 U2 Ze+1.
Now, consider the last three terms. By (12), we have 1z+1,x| So Mz with probability at least 1 8p for all E [T]. Thus,
10 "Z, + 17 241, + 2 u Z So M2.
WWhen a is negative, it suffices replace 2o with our target E..
39

---

C Stage 2: training the second layer
Lemma C.1. Suppose that for each p [P] there exists a first-layer neuron vi, with Vip.p 2
VI-eu Ev for some small positive Ee = O()P)), then we can choose a, E Rmm with ||ax|l = vP such that
C(a.. V) = E (f=(a) - f(æ; ax V)) - 200 E
Proof. Choose one U for each p E [P]. Then, we set the ip-th entries of a to be 1 and all other
entries 0. Then, we write
P -d(v x))
(f+(z) - f(x; ax, v))
Σ (ō(xx) - (
(Ō(IN) o(vr, - x)) (6(2)) Ø(vi; - -x))
k,1=1
By expanding & in the Hermite basis, for any v gd-1, have
E
क (2, = Σ of (',0')' 10(0 x))((0 x)] - Σ
1=N(0,1)
i=0
i=2
Hence, for ls = 1, we have
E (Ø(xz) - olvi . xx)² (x)( + EEO x)) xz) 2E -2 ((N))(() x2)
20 2
oj 1 (Ek. (z)))
S 20
Meanwhile, for k # I, we have
E (é(aa) oee x)) (d(z)) - olva x))
Ed(x))((z)) + Edlo, 2) x) Eω(In)((0 - x) -EΦ(v,, x))(2))
Σ &i va)) wha - vlui))
Note that oi Uind S En and
(Diy, va)) 32 S 20 + 2 (vix - ex, 02,)²
0,, 3 20, + 2/r, -xxl== 28, 41 Vir,k) S 68..
As result,
E (Ø(xx)) o(v,, x)) (d(x)) - Φ(v; x)) s 100
Combining these two cases, we obtain
C = E (f+(x)) f(x; a, v))² < 20
Now, ready to prove the following generalization bound for Stage 2. The proof of it is adapted from Section B.8 of [OSSW24], which in turn is based on (DDLS22, AAM22, BES+ +2]).
Lemma C.2. Suppose that for each p [P), there exists a first-layer neuron ns with up u 2 Ex for some small positive E. = (1)P). Then, there exists some such that the ridge estimator obtain in Stage 2 satisfies
81/a.11 vm
II((á, V) fallau
IOLPP
VNo
with probability at least 1 28p.
40

---

Proof. For notational simplicity, let D = (0, 1) and D = 1 distribution of the samples we use in Stage 2. addition, we write fa for f(:aa V) where V the first-layer weights we have obtained in Stage I and X (#T+n) )
N Σ ser-n denote the empirical
Let a, Rm denote the second-layer weights we constructed in Lemma C.1 and á E Ihm denote the ridge estimator obtained via minimizing + II -af.-flliaa ( + A ||all By the equivalence such that
between norm-constrained linear regression and ridge regression, there exists
fl falliz()) & Ilff - fa. [[; and ||||| = ||a=|l
Choose this À and let F = (f(:a)= ||a|| s ||a. ||} be our hypothesis class. Note that fa E F. Moreover, we have
Ilfa follar (llfa (ooffllla Ilfa felle=(6)) + lIfa - follea(D)
< sup
(11a- - (alL - 115a - -bllv())) + -15a foll
a:
(15a - (blle2D) - 11fa -b)llu(b)) + I15a, - Sallez(0))
sup
|:
where we used the fact that I|ʃa Allm(( $ 115a - -lllz((()) ll... /bll in the last line.
Now, we bound the first term. Let a == (o,) i.i.d. Rademacher variables that are also independent of everything else. By symmetrization and Theorem 7 of [MZ03], we have
(Ilfa folloo Ilfa fblle:()))
E
sup
x
a: |aa Slaa.11 Il
< 2 E
sup
(a(((Tt)) (a
X. a: ||||||||a_||
N
t=1
2 E
sup
oe (Ja(eT+n) ( (a(T4n))
X.o ||||||||.|||
N
2 E
1
Σ t=1
(( +2 E oij. (T (aT
sup
N
X
a
|
X.oN-
1=1
Note that the first term is two times the Rademacher complexity Rad (F) of F (see, for example, Chapter 4 of [Wail9]). By (the proof of) Lemma 48 of [DLS22], we have
|a.1]
||a.||
Rad (F)
E m>.1
lVo(1))
E N0.1()
0²(uk x)
VN
|a. vm
E
(2(11)
VN
12.1((0.1)
2 2la m
VN
In other words, we have
41la.ll
E sup
(1lfa- follL Ilfa fAllo(()) 5
a: a:lla|ss |S|aa_
VN
Hence, for any óp E (0, 1). by Markov's incquality, we have
(Il.a - fallenD) - Ilfa Sullo S
4 4||a_1 vm
sup a: a:[|a||
VN8g
with probability at least 1 - öp. Apply the same argument to Ilfa. - fellea(6) and recall from Lemma C. that Ilaa fall(A(D) I 10LP3 and we obtain
llfa - -f-llza(D) _ 8 ||| vm
VNop + --VV100 11OLPE
41

---

with probability at least - 20p.
D Proof of the main theorem
Theorem 2.1 (Main Theorem). Consider the setting and algorithm described above. Let C be a large universal constant. Suppose that log d S P S d and {ai ]E(PP are orthonormal. Let óp (exp(- d), 1) and E. > 0 be given. Suppose that we choose ao T,,,,
. N satisfying
p2
7 6 (
2200
ph./2-1
Ō(P). N θ
(PdPLIZT))
T O
ne
Then, there exists some O such that at the end of training, we have LMSE(a, V) S E with probability at least O(6)).
Proof. First, by Lemma A.8, we should choose θ (Plog(P/6e) and the in Lemma B.2 and Lemma B.8 can be chosen to be ( log P). Meanwhile, by Lemma C.2, to achieve target L -eror
E, with probability at least ((6r), we need
Pm 0
p lPg(P/))) 220z
Eu = Og
p2
By Lemma 2.1, have Mz So p1/2 and Mz So pl/2 log" (og((/)7) where 0 1/(2(1++ Then, to meet the conditions of Lemma B.2 and Lemma B.8 (uniformly over those P good neurons), it suffices choose
7
log
20 (d/ópp
dDPL22
P Plog(1/E.)
Then, by Lemma B.2 and Lemma B.8, the numbers of iterations needed for Stage 1.1 and Stage 1.2 are Og ()/P)/n) and Os ((pL/2-1 + log(1/s.)) /n), respectively. Thus, the total number of iterations is bounded by
logd p1/2-1 + log( (/e()) = -,s (dpl-ly p1/2 log log(1/E.)
T Og
2
E Stochastic Induction
Our proof is essentially large induction: When certain properties hold, we know how to analyze the dynamics and can show certain quantities are bounded with high probability. Meanwhile, certain properties hold as long as those quantities are still well-controlled. In the deterministic setting, this seemingly looped argument or the continuity argument (in continuous time). In this subsection, we show the same can also be done in the presence of randomness and derive stochastic version of Gronwall's lemma and its generalizations.
be made formal by either mathematical induction (in discrete time)
We start with an example where Doob's submartingale inequality can be directly used. Let (2, (Fi)..
our filtered probability space and (Z,); h be a martingale difference sequence. Suppose that E EE[Z,,, Fi] is uniformly bounded by oz. Then, Doob's submartingale inequality, for any M > 0 and T O, we have
Toz M
IP
sup KTT
Z
M
-2
Z.
In particular, this implics that when M w(oz we have sup E probability.
Z, S M with high
Note that there is no need any kind of "induction" in the above example because of the unconditional uniform bound on [ Fi). However, things become subtle if instead of assuming
42

---

2| Fi] is always bounded by o we assume it to be bounded by o when sup, .i Er= Zr1 S
M. Intuitively. since M is chosen that suP.ST Σ Z, M holds with high probability, the bounds E|Z Fil < o should also hold with high probability and we can still use Doob's submartingale inequality as before. Now, we formalize this argument.
Lemma E.1. Let (Ze); be martingale difference sequence. Suppose that there exists M, M, > such that if supas: D 7al S M, then we have B B|22+1 1 [i) S oz. Then, we have
T
P
sup
Z
> M
S
M
Note that this bound is the same as the one we obtained with the assumption that E| Fi] S oz always holds.
E.=1 Z. > M). By definition, we have sup&gt 1S Zel < M for all & < T. Then, we define Y = Zr-1 1[ < r). Note that (Y)) is
Proof. Consider the stopping time == in 2 0
a martingale difference sequence with E[ Fil S oz. As result, by Doob's submartingale
inequality, we have IP SuP:ST
Σ ===1 Y
M
s Tož/M /M2. To relate it to (Ze)., we compute
sup P
7
M
P
sup tT
Σ
7
>
MAT T
P
Z
MA ST
T
P
Y,
MAT
ST
To
M
where the first and second identities comes from the definition of r and the third from the fact Zr Y for all t S T.
Now, we consider more complicated case, where the process of interest is not a pure martingale. Suppose that the process (X,), satisfies
X = (1 + a)X, + Ei+1 + Ze+1. Xo = 2o > O,
where the signal growth rate 0 and initialization up > 0 given and fixed, (Er), is an adapted process, and ( (Z,), is a martingale difference sequence. In most cases, (Ex); will represent the
higher-order error terms.
Our goal control the difference between X, and deterministic counterpart x; = (1 + a)) To this end, we recursively expand the RHS to obtain
X = (1 + 2 + (1 + ta) + E+t1 + + (1+ + +7+1
= (1 + a)) xD
(1+ S.+1 + Σ (1 a) 'Z.41.
1=1
8=1
Divide both sides with (1 + a) and replace t + 1 with f. Then, the above becomes
X.(1 + at =x0 E(++)) + Σ 11+a) (1+a)=%%
Note that ((1 + -Z;) is still martingale difference sequence. Ideally, [Eel should be small represents the higher-order error terms, and we have bounds on the conditional variance of Ze so that we can apply Doob's 's submartingale inequality the last term. Unfortunately, in many cases, since E1+1 and Z441, particularly their maximum and (conditional) variance, can potentially depend on (X,)est, may only be able to assume [s+1l I (1 + a² with probability at least 1 ōg.g (for each t) and B|2 Fi] < (1 + a)'o2 for some g Eat and of when, say, X, (1 + 0.5) Still,
43

---

we can use the previous argument to estimate the probability that X, & (1 + 0,5)zz for some t S T. We now formalize this argument. In addition, instead of Doob's L submartingale inequality, we will use the following extension of Freedman's 's inequality, which allows us to improve the dependence on failure probability from linear to poly-logarithmic. The proof of this lemma is deferred the end of this section.
Lemma E.2 (Freedman's 's incquality with subweibull variables). Let {Ze)t be martingale difference sequence that is conditionally (o² (o",,))subveiiil -)suuvveillll, i.e.
P|ZZ|] 2
Fi-1 s Coxp
( (-(M/o)VO).
A 0,
for some universal constant C >> 0. Then, for any öp E (0, 1). we have
Σ Ze So oy Tlog*+1 (T/6e), with probability least 1 öp.
Lemma 4.1 (Stochastic Gronwall's lemma). Suppose that (Xt), satisfies
X++1 (1 a)xx Eet1 + Ze+1, Xo 20 > O,
(5)
where the signal growth rate O and initialization zn 0 are given, (Ee)a adapted process. and (Ze),, is martingale difference sequence. Define xạ = (1 + xx0
Let T 0 and óp E (0, 1) be given. Suppose that there exists some 8 E (0, 1) and E, oz > 0 such that for every & 2 0, X, = (1 = 0.5) then we have |5++1] S (1+ a) with probability at least
1 - -7 and Zt41 is conditionally ((1 + a) 1000 -usswebeiblll. Then, if
E 20
z
and o
(6)
Tlog 0+1 (T
we have X, = 0.5)a, for all i E [T] with probability at least 1 85 T
This lemma can be easily generalized to cases where we have multiple induction hy- potheses. For example, if we have another process X (1 + a') + E(+1 + Z'41 and we need both X, (1 + 0.5)., and X (1 + 0.5)ar, for the bounds on K1t1/ Kit1l. E[Z? Fil, E[ (44 1) Fi] to hold. In this case, the final failure probability will be bounded by T( + )p.g) )++ + 20p.
Remark.
Remark. If the recurrence relationship is X S (1 + aa)x + E1+1 + Ze+1, and we only want an upper bound, then we can replace zo with any oo 2 zo in (6) and the definition of the deterministic process (zz).
Proof. Let 7 == inf (t 20200 0: X, 61+)) and set E == f4+1 S r} and 2441 Zi+1 1st < r}. Clear that 7 is a stopping time, & is adapted, and is still a martingale differ- ence sequence. Moreover, by our hypotheses, we have É& < (1 + a) with probability at least 1 -rr, and 2t+1 conditionally ((1 02, 0) -usbwebelllll. As result,
1 (1 + a) -
Σ s=1
< Et <
TE with probability at least Tôp.g.
and by Lemma E.2,
sup 1E(2)
(1++a) Z, s sozv Tlog 0+1 Tlog*+ (T/Ep)
with probability at least 1 - 5p.
Hence, for any öp E (0, 1), if we assume
E and
2 < T
0+1 (T/5p)
44

---

then with probability at least 1 - 8p T we have
Σ (1 (1+ s=1
3 a)
vt E [7].
#=1
Recall that
X₁ = (1)) + (1+ + Σ (1+ (1+
and It (1 + ax
s=1
s=1
Then, we compute
P P [7), X2 & 0.5) 0.5)a [T], (1 AT
= PP XX, $ (1 + 0. ATST]
P
a1+ E. + Σ
(1 (1+a) =Z,
2 0..5% ATST
s=1
+ E E(+a)
PP
2,
2 0.5z0 ATST
$=1
- 0p + TÖp
The above lemmas will be used in Stage 1.1 to estimate the growth rate of the signals. The next lemma considers the case where a and will used to show the gap between the largest and second-largest coordinates can be preserved during Stage 1.1.
Lemma E.3. Suppose that (Xi), satisfies
X1+1 S Xx Ex+1 + Ze+1,
Xo = Z0 > 0,
where the signal growth rate process, and (Ze), is a martingale difference sequence.
and initialization æn 0 are given and fixed, (Er), an adapted
Let T > 0 and 62 E (0, 1) be given. Suppose that there exists some ōβ,5 (0, 1) and E, 0 such that for every t S T. l5el I with probability at least - ōz and Zi41 is conditionally (az, --ssbwweillll. Then, we have
sup Xx mal TE V Tlog (T/00)) with probability at least Tör 87. SST
Proof. Recursively expand the RHS, and we obtain
X, S 20 +
Ex
Z..
Clear that
sup tST s=1
Σ &
S TE
with probability at least Tör,g-
Meanwhile, by Lemma E.2, we have
Σ 2, S <oz Tlog" 0+1 (T/86) with probability at least 1 8p.
sup 8ST
Combine the above bounds and we complete the proof.
Now, we consider the case where the signal grows at a polynomial instead of linear rate. This lemma will used in Stage 1.2, where the L-thh order terms dominate. We will need the following estimations on the corresponding deterministic process. Its proof is deferred the end of this section.
45

---

Lemma E.4. Consider the process 2++1 2t + az, where 20, are small positive real mumbers and p 1. Let T be the time 24 first goes above We have
T-1
T (p 1)000
and Σ t=
Ft
ifa S x0 -pp
poap
2- paaro
This lemma provides upper bounds on the time needed for 2i to grow from zn = (1) to 1 and the sum 24 this process. Note that the second upper bound is essentially Twg. Intuitively, this is because due to the sharp transition behavior of this polynomial system, I2 RJ xo for most of the time.
Remark.
Lemma E.5. Let (Xt)r be non-negative stochastic process satisfying
X1+1 2 X, + Zu+1 + Ee+1, Xo 30 >
(15)
where (Z1+1); is martingale difference sequence, and (Et); is an adapted process. Let tt
be the solution to the deterministic recurrence relationship R++1 = ãg + aai; 20 = zo/2.
Let op E (0. 1) be given and T == inf (t 2 0 X; 1211 21.. Suppose that there exists E,02 > 0 and g 10 such that X; 2 8; and S T, we have [Eel S EX, with probability at least 8p and Zut1 conditionally (,, O-subweibblll Then,
ax
a a 20 /P,
E S parth
02 S0
log 0+1
55))
then with probability at least - / (c(xx/2)pp-) åR, we have T S (po(20/2)pp-)) and X, 2 for all S T.
-1
Proof. Note that we can rewrite (15) as X 2 X₁(1 axp- + & + Ze and view it as the linear recurrence relationship in Lemma 4.1 with a non non-constant growth rate. This suggests defining the counterpart of (1 + a)* as
axp-1), t>s.
Pa..
t=
Then, we can unroll (15) as
X1 2 Xy
+ axx-1
+ + +Z,
X2 2 No
[ 1xx011111 1) + + 21) (1 (1+ +
2 X 1 11xx%-11 1) (1 (1
+1+1 (5) + Z1) +
= XpPo. + P1, (91 + 71) +
Xa 2 X (1+aX + + 2z
2 (XoPo,2 + P1,2 (f1 + Z1) + & + Z2) [1+ (1+aX] axx- + + Zs
= XoPo.3 + P1.3 (f1 + Z - P2.3 (E2 + Z2) + E1 + Zz.
Continue the above expansion, and eventually we obtain
1-1 X, 2 X Poo.t + P.tLA (5.+1 + Zot1),
Vt 2 4p 2 0.
a=to
Since X is non-negnennn non-negative, we have ps. 2 1 > O. Hence, we can rewrite the above as
t-1
P. 2 Xt,
Σ
P. (Eo+1 + Zo+1), Wt > 2 0.
46

---

We wish the repeat the argument in the proof of Lemma 4.1, showing that the last term is smaller than Fo /2. Unfortunately, this approach will not work directly. We have only assumed | I EX; and ([2441/Fi I oz X,. Since X, can be much larger than ét, we cannot directly use our assumption to control the size of noises. On the other hand, note that if X, > Lt, the induction hypothesis will less likely be violated, so in principle, X; > Ig should help us. To "enforce" the X, Z 3h condition, we consider the following recoupling strategy: whenever X; 2 4f,, we restart St at X,/2. This recoupling will only increase the value of 2,, and it ensures X, S dt always hold.
We now formalize the above argument. To this end, let o : R550 + R>o be the flow map of the recurrence relationship 2++1 24 + az. That is, (2(2) is the value of x, if (r,)s is generated by It-1 = x with zg z. Then, we inductively define the following sequences of "deterministic" processes and stopping times:
( (x(oo/2),
( = inf {+ 0 X1 2 4.. }
# + (X,u))
(k+1) inf { (k) X, 2 48 (1)) s: Ak 11
In words, (k) is the time we switch to the kth coupling. By construction, By construction, a. in
non-decreasing in both L and k; 0 =: (0) <... ... < (a) < ... and ((()) x ,.0/²2 2 24 (k-1)
((k) z
25 1) 2 2 - In particular, the last property implies that there are only finitely many couplings before at reaches any fixed constant.
Then, we abuse notations, redefining
(h) (t< (10
k=0
Clear that at each t, only one summand is nonzero. By construction, we always have X, S 42,, Since this & is no smaller than the original one, it suffices to bound the probability that X, < 8g for some t < T. Note that X, S &y if and only if there exists some k E N20 with E [c(0), 1) such that
X,00 +
Pat. s4 st1 (Es+1 + 2s+1) S P ( a
(0)=0
In addition, note that if X, 2 S. for all s < t, then we have P-1 2(1) Therefore,
S
((0) (3k)
X20
X, 2 Res <t
38 X1 S B4
3K E N [(()), ((+D)) s.t.
Σ ami(k)
t-1
Pidl 0- (Eut1 + 7x41) S -2
-Z(00-
In other words, it suffices to upper bound the probability that RHS happens before L. To this end, detine T := inf (t 2 0 X. S fe). Ê = Ertilft < r}. and Ze+1 Ze+1 (t 7}. Then, we can further rewrite the above as
36 < T, < &t
(1AT)-1
(AT)) Σ
3 E N20, L E ((00),ttt))) s.t.
Σ
Est1 +
Prj
2s41
2 (k)
(10),
((k)
az((k)
s=c(k)
We now estimate the last term as follows. First, for (Er)e, we have 1ép1l S EX, S 43x(() t (1(4),,(0+3)) (k++)). Therefore,
(MTT)-1
(1AT)-1
Σ
Pol E
P
4
Σ
(B)
=
#+1
(8) 2ーー
pa
S=((t)
(k)
47

---

where the second inequality comes from Lemma E.4. For the RHS to be smaller than ((0) to require
(00) it suffices
E S pa
(K) (00) (00) p--
= S paca P-1
Also, by Lemma E.4, when the induction hypothesis is true, we have T S (axp=1)
( Thus, the
above implies that with probability at least 1 - /oxg/(axx the total contribution of (Et)e is small, as long as 3 X pazz 1
Then, we consider the martingale difference terms. Note that P,th) (00), , Ze+1 is a martingale difference sequence that conditionally (403 (k) (kis 0) -ssubweiblll.. Hence, for each k, by Lemma E.2, we have
log 0+1
(ala(())- 10, )
(ATT)-
Σ
Pat (0,),
2ot1 So \
03 ((6)
al2/(1) 2(4)
|o-1
amc(k)
log"+1
(1/(azg 54,22)
So oz
with probability at least probability at least 1 require
Se,2. For the RHS to smaller than zo, it suffices to
oa
02 So
logo+1
(1/(ar) 8r,2))
Recall that we recouple most (loog[L/zo) times. Hence, it suffices to replace 85,2 with óe (og((/)))) to ensure the total contribution of (Zi), is small.
E.1 Deferred proofs
Proof of Lemma E.2. First, consider the where 1. Let M 2 1 be parameter be chosen
later. Then, define Z, 1,1 [1 < M) and write
T
T
T + Σ
Σ
2 1=1
2, - E
Ze Fr
+ Σ 2
Z,1 (12
M) =: T1 + Te + T3.
1=1
Since Ze conditionally (1, -ub)-uubbiillll, we have
P(Ts + 0) < P (3t E [7], 2 M) S CT exp (-M!/0)
For the last term to be bounded by ōp, it suffices to choose
M log" (CT/6p).
Then, we consider T2. Since E|Z,
F 0, we have
E S/4 F, E 7. 2, F E M) [F--).
For the last term, using the layer cake representation, we obtain
E [Z,1 > M] Fi-1l < || ž M] F(-1]
P(Z 1 ( Z M) 8 Fe-1) ds
P(Zz MV F.-1) ds
= MMP(ZIZ2MILL. 1 P(Z 2 F(-1) ds.
48

---

Therefore, for each summand in T2, we have
2 Fi-1 - cxp
(-M-/0) + / s
Coxp (-st10) ds
= CM exp (-M'/0) + Co
/ e 806002 do.
1M1100
Note that if s/ log 2 2(0 - 1), we have M such that
S = Hence, if we choose
M1/0
M
2( 1)
log (M1/8) 2(0 1
log M
then we have
C S
(-) ---1110
c
ds S S
e-112 ds 2C 2C'xxp
M1/0
M110
As result, have
T2 CT
M [Mexp (M11/0)
2exp M < 4CMT exp (-
Finally, consider Tj. Note that (2, B|2, Fe-il) is martingale difference that is bounded by 2.M and has conditional variance bounded by Cg for some Co > Therefore, by Bernstein's 's inequality, we have
(K/VT)²
P( 2 K) 2exp
C. 2M
For the RHS to be bounded by dz, it suffices require
K> TT (Co 2M) log (2/62)
Finally, combining the above analysis, obtain
Ze <CCM exp (-M\/012) VTV 2M) log (2/dp),
with probability at least 28p, where M 2 log" (CT)8p) and M/ log" 2 (2(0-1) 0 simplify the RHS follows. Note that
Now,
4CMT exp [-12 < VT 2M
/2M log (2/6²)
1
log M - 1/1/0)
1
/2 (2/8p)
exp
2
2
4CvT
M
M > 4" log" (
8C2T
+
29,
Tog (2/
log
M
In other words, we can choose
M = 0, (Tg2((//b))
and obtain
T
Ze So V T Tlog*+1
9+1 (T/6p).
Finally, for general 0, it suffices note that if X is (a², -ubbuwiiibull, then X/a is (1,0)- subweibull.
49

---

Proof of Lemma E.4. First, we consider the upper bound on T. We compute
11+1
3t
1+1
2++1
3t
a
dy
11 11
141
1+1
T++1
p1+ 1
1
dy
tt
yp
Ft
2
p-1
In addition, note that /4/1/a/ = (1 aar axpp- 1)
oar
< a) S pp. Therefore,
eap
1
1)a
a
1
x
p--1
2/2
cap
Sum both sides from 0 to - and we get
-
Upp 1)xx
1
pop
St
10
p-1
op (p
1)ad
In particular, this implies
1
T S
2/0
p-1
(p 1)a
Now, we consider the upper bound on E, Ft. Let (#x), be the solution the continuous-ting time ODE de Ti = dh with in = zg- Note that increasing and therefore
i(41ja Tea + S
f'aaat Flato dr Jea caa
Hence, by induction, we have Zea 2 x, for all In addition, Tn has the closed-form formula:
Th =
- (p - 1)h
Thus,
T-1
T-1
(t+1)a
Ta
<
Tra S a
Σ
Fg ds
Ty dh
1=0
0=
to
-
1
(p
1)h
dh.
p p- -
0
When 2, we have
T-1
1 Ta
(/0 (1 (1-h)
-1
1
1 log
2
Ft
0
J0
xo
dh log
a
1 To
a
( 1- + z1
as long as a S 2n so that ZT. 0(1). When > 2, to have TTa S 2, it suffices to have
1
oD 1
1
(p 1)aT >
av
$ 2/-1/p /p.
2/ p-
2p-1
1-1
2
50

---

Let T be the time on reaches 2. We have T 2
and
(0-
1)
T-1
Σ 84 S
1
1
(p 1)h
7-1
-
0'
0
20
-
1
1
(p 1
a
2
20
1-
-
1
X
(p
1)T
-
(p 1)T
1
20
(p-
2)a
p-2"
F
Simulation
We include simulation results for Stage I in this section. The goal here is to provide empirical evidence that (i) we have both the second- and L-th order terms, then the sample complexity of online SGD scales linearly with d and (ii) without the higher-order terms, online SGD cannot recovery the exact directions.
The setting is the same as the one we have described in Section 2. We choose the hyperparameters roughly according Theorem 2.1. 1. To reduce the demand of computational resources, we choose m O( (2) instead of Ω2( p$). Note that by the Coupon Collector problem, we need m S2 log P) to ensure that for each p E [P], there exists at least one neuron a with v 2 max ui? Since we are mostly interested in the dependence on d, for the learning rate, we choose tunable constant that independent of d can depend on everything else. T is chosen according to Theorem 2.1 and we early-stop the training when for all p E [P]. there exists a neuron with ug 2 0.95 (in the moving average sense).
c/d, where c
All experiments are performed on the authors' laptop without using GPUs, and it takes less than one day to complete the experiments.
51

---

= h2 ha
+ + nl (train bath layers)
V5.
0.6
0.6
may
0.4
d=20
d=20
d-30
d-100
d-100
d=203
d0200
0.0
3000 4000 5000 6000
2000
4000
8000
Iterations (x1000)
tterations (x1000)
d h2 hg
i = abs
1.0
1.0
0.8
max
0.4
0.4
d
d-
0.2
d-50
0.2
d-50
d=100
d= 100
d=200
d=200
1000
2000
3000
4000
5000
6000
0
5000
10000
15000
20000
fterations (x1000)
(x1000)
The above plots show the evolution of the correlation with each of the ground-truth directions. We fix the relevant dimension P 5 and vary the ambient dimension d. Different colors represent different d. For each color, one curve represents max, v² for one p E [P]- In the first row, the link function is & h2 h4. In the left plot, the right plot, we train both layers simultaneously. The second row contains simulation results for other link functions.
Figure 1: Recovery of directions.
the algoríthm (1), while in
d= b2 VS d = hg he
fu
0.75
max ///
0.65
0.60
0.4
0.50
0.40
10000
20000
30000
40000
50000
10000
20000
30000
40000
50000
Iterations (x1000)
(x1000)
Figure 2: Necessity of the higher order terms. In these two figures, we choose P= 10 and = 100.
The left plot shows the maximum correlation each of the ground-truth directions (also Figure I). We can see that in the isotropic case, whether online SGI) can recover the ground-truth directions is determined by the presence/absence of the higher-order terms. The right plot shows the change of max, / ||'p'l for each p E [P] Stage I when the link function is h2g One can observe that they are almost unchanged throughout training. This, together with the left plot, shows that the increase of the correlation is caused by learning the subspace instead of the actual directions.
52


Reviews:
[{"review_id": "4nV1NIzpEj", "reviewer": "Reviewer_4V3d", "summary": "The paper analyzes the sample complexity of learning multi-index models of the form $\\sum_{k=1}^P \\phi(v^\\star_k . x)$ with orthogonal $v^\\star_k$  through online-SGD on a two-layer neural network model. Unlike existing works focusing on the recovery of the target subspace and subsequent learning of the target non-linearity, the main results show that with sufficient overparameterization, the neurons of the trained model exactly recover individual target directions $v^\\star_k$ with $\\mathcal{O}(P dP^{L-2})$ sample-complexity, where $L$ denotes the information exponent. Crucially, the sample-complexity for recovering individual directions remains linear in the dimension. The exact recovery of directions proceeds through higher-order terms in the gradient dynamics, which amplify the neurons' initial difference in overlaps across different directions.", "questions": "* What factors does the polynomial in the Poly(P) sample-complexity depend on? Can it be well-specified?\n* Can assumption 1(c) be extended to targets with polynomially-bounded derivatives?\n* where does the diverging width requirement $ \\log d^c < P$ arise from?\n* what would happen if Assumption 1(a) is relaxed to allow a linear-component in the target non-linearity?\n\nIncluding the above clarifications would strengthen the paper.", "limitations": "The work is primarily theoretical and does not pose direct societal risks. The paper would benefit from a more detailed discussion of the limitations of their model listed under \"weaknesses\".", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths:\n- The paper is well-written and organized with adequate discussion of prior work.\n- The analysis is insightful. In particular, the insights into how higher-order information in the gradients leads to the exact recovery of individual directions are expected to carry over to general setups, such as while jointly training all the layers, or for deeper networks.\n- The proofs are simple and the paper provides readable proof-sketches explaining the key ideas.\n\n# Weaknesses\n- The width-complexity or overparameterization in the main result is specified as Poly(P). A precise description of this polynomial dependence would greatly strengthen the results.\n- The paper suffers from the limitations of using correlation loss (which prevents interactions between neurons). However, since this simplifying assumption is used by several related works in the literature, I consider it a justifiable starting point.\n- Another related limitation is the absence of a linear component in the target non-linearity. With correlation loss, this presumably leads to a spike from the linear components dominating the dynamics. A discussion of possible directions to relax such limitations would strengthen the work.", "was_revised": true, "final_justification": "I recommend acceptance based on the strengths highlighted in my review, in particular the insights gained from the analysis.", "rebuttal": "Thank you for the reviews and suggestions! We'd like to answer your questions as follows.\n\n* **Width complexity.** We realized that the width can be as small as $O(P \\log P)$ after submitting the main text. This\n  is tight up to the logarithmic term and matches the coupon collector bound. We have included this in the appendix. \n  See Lemma A.7 and the statement of the main theorem in Appendix D.\n* **Relaxing Assumption 1(c).** Thank you this question! It makes us realize that polynomial boundedness is sufficient.\n  Assumption 1(c) is only used to get a tail bound on the per-sample gradient, which amounts to controlling the tail of \n  $f_* = \\sum_p \\phi(v_p^* \\cdot x)$ (times the derivative of a single neuron). We can use the polynomial boundedness \n  to bound the $L^p$ norm of each $\\phi(v_p^* \\cdot x)$ and then show that they are sub-Weibull. \n  Since $v_p^\\*$ are orthogonal, $\\phi(v_p^\\* \\cdot x)$ are independent, and therefore, the sum of them is still \n  sub-Weibull and this leads to the desired tail bound. We will add this into the revision!\n* **The diverging width assumption.** The condition $P \\ge \\log^C d$ is mostly an artifact of our analysis. First, we \n  want to ensure certain concentration properties hold in the relevant subspace (such as the norm in the relevant \n  subspace). We do not want to explicitly track the failure probability of these events, so we assume $P \\ge \\log^C d$,\n  so that the failure is inverse super-polynomially small. In addition, this also creates a cleaner separation between \n  the two stages. For example, at the end of the first stage, the 2nd order terms have order $\\| v\\_{\\le P} \\|\\_2^2 \\approx 1$, \n  while the higher order terms have order $\\| v\\_{\\le P} \\|\\_L^L$. The higher order terms are much smaller than the 2nd \n  order terms when $P \\gg 1$, and therefore can be safely ignored. We believe that this condition can be removed at the \n  cost of a messier proof.\n* **Allowing linear terms.** First, we wish to stress that we consider the case IE=2 since it is the lowest possible when\n  the link function is even, or equivalently, both $v_k^\\*$ and $-v_k^\\*$ appear in the teacher network. Now, back to\n  the question of the influence of linear components. You are right that there will be a spike (which is the all 1 \n  direction if we assume w.l.o.g. that the targets are $\\{ e_p \\}_p$) and the all neurons will be attracted by that \n  direction. \n\n  To fix this issue, the easiest approach is to first train one neuron to fit that direction to remove it. \n  Due to the linearity, learning the linear term is roughly equivalent to learning a single-index model with IE=1, \n  and it costs only $O(d)$ samples. \n\n  The more natural but technically more challenging approach is to use MSE loss. Similar to the 2nd vs 4th order terms\n  argument, one should be able to show that the first order terms will get learned before the higher terms become large.\n  The tricky thing here is that after learning the first order term, the learner distribution will be shifted to be \n  centered at the all-one point (assuming $v_k^* = e_p$ for simplicity), instead of one single neuron learning that \n  direction with all other neurons staying close to the initialization, as the gradient of the linear term depends only \n  on the average of the neurons and not on any individual neuron. The analysis become tricky after that, since we can\n  no longer ignore the interaction between learner neurons, though in principle, a finite-width mean-field type analysis\n  similar to the one in [LMZ20] can potentially be used to analyze the dynamics of neurons in those later stages."}, {"review_id": "8OsmMbV3k9", "reviewer": "Reviewer_x9Fa", "summary": "This paper studied the problem of learning orthogonal multi-index models under Gaussian distribution. The paper assumes the multi-index model takes the specific form of $f^\\ast(x) = \\sum_{k=1}^P \\phi(v_k^\\ast\\cdot x)$, where $\\phi$ are non-linear activations that have the form of $\\phi=he_2 +\\sum_{l\\geq L}\\hat{\\phi}_l he_l$. The authors showed that, using online sgd algorithm, using $dP^{O(L)}$ samples suffices to learn the target $f^\\ast$ to $\\epsilon$-$L_2^2$ error. The authors provided a set of frameworks and tools for the analysis of sgd dynamics on such objectives, and provided a white-box, detailed analysis of the learning algorithm.", "questions": "See weakness", "limitations": "no negative limitations", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. I think the paper conducted a thorough study of an interesting special case of multi-index models. The paper is solid in techniques. \n2. I think the characterization of the sample complexity, in general, is not surprising, as similar (qualitatively) results have been derived in prior literature like [DLS22]. However, I think it is still an interesting contribution to obtain such a result for the special case.\n\nWeakness\n1. I think the condition that $P\\geq \\log^C(d)$ is a weird condition that seemly most likely to be a side-effect of the analysis. Is it possible to remove this by refining your arguments?\n2. In [DLS22], the authors studied learning MIMs of the form $g(v^\\ast_1\\cdot x,\\dots,v^\\ast_p\\cdot x)$ where $g$ is a polynomial of bounded degree, using two layer neural nets. Though [DLS22] has sample complexity $d^2 P + d P^L$ (consider $g$ being a degree $L$ polynomial), this paper considers a much simpler form of the multi-index model. I wonder if it is possible to relax the orthogonality condition to the condition of [DLS22], i.e., $E[\\nabla^2 f^*(x)]$ has rank $P$. Otherwise, it seems the major contribution of this paper is to improve the degree of $d^2$ to $d$, with the sacrifice of model generality.", "was_revised": true, "final_justification": "I think the authors cleared my concersn. I maintain my evaluation.", "rebuttal": "Thank you for your feedback! We'd like to address your concerns as follows. \n\n* **The condition on $P$.** The condition $P \\ge \\log^C d$ is mostly an artifact of our analysis. First, we \n  want to ensure certain concentration properties hold in the relevant subspace (such as the norm in the relevant \n  subspace). We do not want to explicitly track the failure probability of these events, so we assume $P \\ge \\log^C d$,\n  so that the failure is inverse super-polynomially small. In addition, this also creates a cleaner separation between \n  the two stages. For example, at the end of the first stage, the 2nd order terms have order $\\| v\\_{\\le P} \\|\\_2^2 \\approx 1$, \n  while the higher order terms have order $\\| v\\_{\\le P} \\|\\_L^L$. The higher order terms are much smaller than the 2nd \n  order terms when $P \\gg 1$, and therefore can be safely ignored. We believe that this condition can be removed at the \n  cost of a messier proof.\n* **Comparison with [DLS22].** First, we wish to stress that our sample complexity of the second stage is also tighter \n  than the one in [DLS22]. They use a random feature argument in Stage 2, which leads to the sample complexity $d r^q$,\n  where $q$ is the degree of the polynomial. In our case, we show that SGD can do feature learning, and the sample \n  complexity is $d r^L$, where $L$ is the information exponent, which can be much smaller than the degree. In addition, \n  our results can be extended to any link functions with a polynomially bounded derivative (see the response to \n  Reviewer 4V3d), while their results only deal with polynomials. Moreover, their width complexity is $P^q$ and ours is \n  $P \\log P$ (cf. the statement of Theorem 2.1 in Appendix D). In some sense, [DLS22] shows that SGD can learn the \n  subspace using $d^2$ samples, while our results show that, at least in a more restricted but still non-trivial \n  setting, SGD can automatically do *hierarchical* feature learning, which yields sample/width complexities that are \n  (almost) tight, in the sense that they match the complexities of separately learning the subspace, and then \n  independently learning $P$ single-index models in the relevant subspace. \n* **Relaxing the orthonormality condition.** First, we believe that our results can be extended to the near orthogonal \n  case. However, we expect the non-orthogonal case to be hard, as it is known that non-orthogonal tensor decomposition \n  is hard in general, whence we cannot expect SGD to learn the directions with polynomial samples/iterations, \n  even when the relevant subspace is given. \n\n  Things are trickier if we consider only the 2nd order terms (which is always \"orthogonal\" because of the spectral \n  theorem), but do not assume isotropicity. In this case, correlation loss will fail, since all neurons will simply \n  converge to the largest eigendirection (similar to what will happen in the matrix power iteration). This can be \n  fixed by using the MSE loss and an initialization whose scale matches the target's. Our analysis does not directly \n  apply to this case due to the interaction between the learner neurons, but the behaviors should be similar.\n  That is, the 2nd order terms drive the distribution of learner neurons to (close to) the Gaussian in the relevant \n  subspace, with the correct covariance matrix, and then the higher order terms become non-negligible and allow the \n  neurons to recover the exact directions."}, {"review_id": "KrP9fFX3Ns", "reviewer": "Reviewer_uJL9", "summary": "The authors propose to understand the sample complexity of learning a multivariate function of the form $x \\to \\sum_{i = 1}^P \\phi(v_k \\cdot x)$. The input $x \\sim N(0,1)$ and the learning is performed via a modified version of \"SGD\" with a two-layer neural network.\n\nIn the article, the information exponent is $2$, but the aim of the paper is to show that while this first order term $h_2$ has a rotational invariance that prevent the $v_k$ to be retrieved (only the span of it), higher order terms can help having a perfect recovery, without paying any cost.\n\nThe main theorem is presented in Section 2, while the sketch of the proof is divided between the two subsequent parts : the gradient flow study and then the discretization via SGD.", "questions": "The aim of the paper is to show that the network is able to retrieve the $v_k$, and not only the span of it. In the theorem, it is not clear to me that it is the case since it is stated that the population loss is small. I understand that it could be that since there are higher order terms, this implies that the $v_k$ are retrieved but this needs to be clarified\n\nI have trouble with the claim that *the information/generative exponent is insufficient* ? For me it still appears as the integer driving the complexity of the dynamics?", "limitations": "Already discussed", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- the paper seems technically very solid, with clear sketch of proofs and a concrete goal. Analyzing sample complexities in such tasks if often quite technical\n- the aim of the paper is clear and it is overall well written\n\nWeaknesses:\n- the scope of the paper is quite narrow : it deals with a specific problem in the multi-index model (a separable sum with the same link function), and the link function has a very specific structure : $h_2 + h_L + \\dots$.\n- every thing is sort of tailored for the result to hold : the neural network is adapted to the problem in the sense that the link function which could be a priori unknown is the non-linearity of the Neural Network. The algorithm is also very specific : freezing layers after the other ones. I understand these assumptions and algorithms are abundant in the literature, yet I have to signal this as a weakness.", "was_revised": true, "final_justification": "I have read the answer of the authors. It is concise and precise and I understand the goal of the paper perfectly. Due to the limitation of the scope of the result, I am not willing to raise my score, however, I tend to a 4+.", "rebuttal": "Thank you for the reviews! We'd like to address your concerns as follows. \n\n* **Recovering the ground truth directions.** We show in Lemma B.8 that at the end of Stage 1.2, every ground truth \n  direction will be recovered by a neuron, in the sense that the correlation is at least $1 - \\epsilon$. The only \n  thing Stage 2 does is picking out these neurons, or, more accurately, find a combination of the neurons that is \n  at least as good as the one consists of these neurons. If one really wants to recover the directions explicitly, \n  one may add a $l^1$ regularizer $\\lambda \\sum_{k=1}^m |a_k|$ and do lasso regression, which returns a sparse solution. \n  Alternatively, this step can also be done with $l^1$-regularized SGD via a local convergence argument similar to \n  the one in [Chi20].\n* **On the scope of the paper.** Our results can be extended to the misspecified case, where the link functions along \n  different directions are different and unknown to the learner, as long as the activation of the learner has some \n  universal approximation property, and the target link functions still have isotropic 2nd order terms, and the \n  condition number of the $L$-th order terms is constant. This is similar to the situation in [BBSS22]. Recovering \n  the subspace/directions does not require the activation to be exactly the same as the target link function, and \n  after learning the directions, fitting the target link function can be done with a one-dimensional random feature \n  argument.\n\n  Our results can be (partially) extended to other settings as well (polynomially bounded link functions, \n  including linear terms, and/or non-isotropic). You could check the responses to Reviewer 4V3d and x9Fa for details. \n* **Issue of IE/GE.** We are not claiming that the IE/GE theory is wrong. Instead, our point is that \n  the vanilla IE/GE theory, which considers only the lowest order, is insufficient to characterize the richer behavior\n  of multi-index models. For example, we cannot distinguish between $h_2$ and $h_2 + h_4$ when looking only at the \n  lowest order. On the other hand, SGD's behaviors for these target functions are different, as directional convergence \n  is possible only in the 2nd case. We propose to look beyond the lowest order and show that, in the setting considered\n  in this work, this leads to a sharper characterization. \n  \n  In addition, we believe this point of view might shed some light on the much harder non-orthogonal setting. \n  In this case, the minimizers of the correlation loss, truncated at any finite order, are not aligned with the \n  ground truth directions. However, as the order increases, they will become localized around the ground-truth \n  directions, as the influence decay faster (for example, the \"infinite-order\" terms are nonzero only if the learner \n  neuron is fully aligned with a ground truth direction). Intuitively, one might be able to use the lower order terms \n  first to obtain warm starts, and then gradually move to higher order terms to move toward the ground truth \n  directions. Formalizing this or finding conditions under which this is possible is challenging, and the orthogonal \n  $h_2 + h_L + \\cdots$ setting considered in this paper provides a simple example where this type of argument is \n  doable. \n\n\n[Chi20] Lenaic Chizat. Sparse Optimization on Measures with Over-parameterized Gradient Descent. 2020.\n\n[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, Min Jae Song. Learning Single-Index Models with Shallow Neural Networks. 2022."}, {"review_id": "Em79fS7pm0", "reviewer": "Reviewer_9mhx", "summary": "PAPER SUMMARY:\nThe information exponent (IE) of a multi-index model $f(x) = \\sum_i \\phi(v_i \\cdot x)$ is the order of the first nonzero term in the Hermite expansion of $\\phi$ around $0$.\nThe IE and its generalization, the generative exponent (GE), provide characterizations for the sample complexity of learning multi-index models via SGD. In particular, roughly $d^{IE -1}$ samples are needed if $IE \\geq 3$.\n\nWhen $IE = 2$, rotational invariance makes it impossible to recover the exact orthonormal vectors $\\{v_i\\}_i$ and only $Span(\\{v_i\\}_i)$ can be recovered. This paper investigates the statistical complexity of recovering all the orthonormal vectors $\\{v_i\\}_i$ as a function of $L$, where $IE = 2$, and $L$ is the the next non-zero term in the Hermite expansion of $\\phi$.\n\nThis work proves that $O(PdP^{L-2})$ samples are sufficient for recovering all vectors $\\{v_i\\}_i$.\nHere is the catch: even if it does not allow recovery of all the vectors $\\{v_i\\}_i$, having a non-zero second-order term is advantageous, as $d^{L-1} \\gg PdP^{L-2}$ samples would be required if the second-order term was null. In particular, the IE does not tell the full statistical complexity story when $IE = 2$.", "questions": "I probably missed something, but I am not sure why you write $PdP^{L-2}$ rather than $dP^{L-1}$.", "limitations": "Yes.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "I believe that the authors could have done a better job at articulating why this particular problem should be interesting for ML theorists. Overall, it seems like the studied phenomenon manifests itself in the particular case $IE = 2$, when we want to recover the exact vectors $\\{v_i\\}_i$, and does not add much to the bigger picture of multi-index models.\n\nThis work claims that \"focusing solely on the lowest degree can miss key structural details of the model and result in suboptimal rates\".\nIt seems to me that the main message of this work is that, for the particular case of $IE = 2$, looking at the next non-zero term is informative. This does not detract from the way more general statistical bound for $IE \\geq 3$, and does not add much depth to that line of work.", "was_revised": true, "final_justification": "I confirm my evaluation and recommend acceptance.", "rebuttal": "Thank you for the reviews! We'd like to address your concerns as follows. \n\n* **On the scope of the paper.** We focus mainly on the case where $\\mathrm{IE}=2$ because $\\mathrm{IE}=2$ is special \n  in the following ways: (1) the IE of even functions is at least $2$, (2) most link functions/activations have IE \n  at most $2$, and (2') the IE of many functions, including all finite degree polynomials, can be reduced to $2$ by a \n  proper label transformation. \n\n  Moreover, we believe our *hierarchical* feature learning analysis can potentially shed some light on the much harder non-orthogonal \n  version of this problem.  In the non-orthogonal case, the minimizers of the correlation loss, truncated at any finite order, are not aligned with the ground truth directions. \n  However, as the order increases, they will become localized around the ground-truth directions, as the influence decay \n  faster (for example, the \"infinite-order\" terms are nonzero only if the learner neuron is fully aligned with a ground \n  truth direction). Intuitively, one might be able to use the lower order terms \n  first to obtain warm starts, and then gradually move to higher order terms to move toward the ground truth \n  directions. Formalizing this or finding conditions under which this is possible is challenging, and the orthogonal \n  $h_2 + h_L + \\cdots$ setting considered in this paper provides a simple example where this type of argument is \n  doable. \n\n  Finally, our results can be (partially) extended to other settings as well (polynomially bounded link functions, \n  including linear terms, non-isotropic, and/or misspecified). You could check the responses to Reviewer 4V3d, x9Fa and uJL9 for details. \n* **$P d P^{L-2}$ vs $d P^{L-1}$.** We write $P d P^{L-2}$ instead of $d P^{L-1}$ to emphasize the different sources of \n  $P$ and $P^{L-1}$. The first $P$ comes from the fact we have $P$ different directions, while the second part \n  $d P^{L-1}$ comes from learning a single-index model in a $P$-dimensional subspace with noises scaling with $d$."}]

Author rebuttal:

Meta-review: The paper analyzes the sample complexity of learning multi-index models of the form    $\sum_{k=1}^P\phi(v_k^{*}\cdot x)$    
 
with orthogonal directions  $v_k^{*}$,  through online-SGD on a two-layer neural network model. The authors demonstrate that information/generative exponent alone is insufficient to characterize certain structures in the learning task and show that, in this specific orthogonal setting, by considering both second and higher-order terms, the sample complexity of directional recovery using online SGD can be improved over the vanilla information exponent-based analysis.
 
The paper has ratings: Accept, Accept, Borderline Accept, Borderline Accept.

Strengths: Reviewers all appreciate the technical contributions of the paper and agree that the presented theoretical analysis is insightful.

Weaknesses: Reviewer uJL9 is concerned that the scope of the paper is limited, since it deals with a specific problem in the multi-index model (a separable sum with the same link function), and the link function has a very specific structure. Nevertheless, Reviewer uJL9 still leans towards acceptance.

Base on these reasons, the final recommendation is: Accept

Note: I agree with Reviewer 9mhx that writing $PdP^{L-2}$ instead of $dP^{L-1}$ can be potentially confusing to readers.
Decision: Accept (poster)