Neural Networks Learn Generic Multi-Index Models Near Information-Theoretic Limit
Authors: Bohan Zhang
Zihao Wang
Hengyu Fu
Jason D. Lee
Venue: ICLR 2026 Poster (2026)
Keywords: Feature Learning
Multi-Index Models
Two-Layer Network
Gradient Descent
Sample Complexity
Primary area: learning theory
TL;DR: We prove that neural networks trained via GD can efficiently learn generic Multi-Index models near the information-theoretic limit.

Abstract: In deep learning, a central issue is to understand how neural networks efficiently learn high-dimensional features. To this end, we explore the gradient descent learning of a general Gaussian Multi-index model $f(\boldsymbol{x})=g(\boldsymbol{U}\boldsymbol{x})$ with hidden subspace $\boldsymbol{U}\in \mathbb{R}^{r\times d}$, which is the canonical setup to study representation learning. We prove that under generic non-degenerate assumptions on the link function, a standard two-layer neural network trained via layer-wise gradient descent can agnostically learn the target with $o_d(1)$ test error using $\widetilde{\mathcal{O}}(d)$ samples and $\widetilde{\mathcal{O}}(d^2)$ time. The sample and time complexity both align with the information-theoretic limit up to leading order and are therefore optimal. During the first stage of gradient descent learning, the proof proceeds via showing that the inner weights can perform a power-iteration process. This process implicitly mimics a spectral start for the whole span of the hidden subspace and eventually eliminates finite-sample noise and recovers this span. It surprisingly indicates that optimal results can only be achieved if the first layer is trained for more than $\mathcal{O}(1)$ steps. This work demonstrates the ability of neural networks to effectively learn hierarchical functions with respect to both sample and time efficiency.

Published as a conference paper at ICLR 2026
NEURAL NETWORKS LEARN GENERIC MULTI-INDEX MODELS NEAR INFORMATION-THEORTII LIMIT
Bohan Zhang
Zihao Wang"
Peking University
Stanford University
2200010903@stu.ppam www.:/ppwccuueuue.oooooooooo
z zihaow stanford..ce.oooo edu
Hengyu Fu UC Berkeley
Jason D. Lee
UC Berkeley j jasondlee@berkeleyy.
hengyu.feberkeley.ed
ABSTRACT
In deep learning, central issue is understand how neural networks efficiently learn high- -igh-dimensional features. To this end, we explore the gradient descent learning of general Gaussian Multi-inde model f(z) g(Ux) with hidden subspace U E IRrxd, which is the canonical setup to study representation learn- ing. We prove that under generic non-degenerate assumptions on the link fune- tion, standard two-layer neural network trained via layer-wis, wise gradient descent agnostically learn the target with (1( test error using O(d) samples and (dd)) time. The sample and time complexity both align with the information- theoretic limit up to leading order and are therefore optimal. During the first stage of gradient descent learning, the proof proceeds via showing that the inner weights can perform power-iteration process. This process implicitly mimics spectral start for the whole span of the hidden subspace and eventually eliminates finite- sample noise and recovers this span. It surprisingly indicates that optimal results can only be achieved if the first layer is trained for more than O(1) steps. This work demonstrates the ability of neural networks to effectively learn hierarchical functions with respect to both sample and time efficiency.
INTRODUCTION
Modern large language models powered by deep neural networks Krizhevsky et al. (2012); He al. 2016); Vaswani et al.] (2023) have shown strong empirical success across many domains. It is be- lieved that much of the effectiveness of the pretraining-fime fineingnniinnninnggggg paradigm comes from their ability to extract useful representations from data. In practice, neural networks often first capture important hidden features, which helps them to learn complex tasks more efficiently. This adaptability offers a clear performance advantage over conventional fixed feature methods, such as kernel techniques Wei al. 2020); Allen-ZZi Zhu & Li (2020) [Bii & Lee (002)).
Despite these empirical successes, the theoretical understanding of why neural networks can ef- ficiently and adaptively learn important features from data remains largely incomplete. Significant effort has been devoted in this direction, with a series of works Damian et al. (2022): Ba et al. (2022); Berthier et al. (2024); [ee et al. (2024); Montanari Urbani (2025); Montanari & Wang (2026) in- vestigating how neural networks can uncover hidden features and subsequently learn single-index and -uuli--nndex models, where the labels depend on the d-dimensional data inputs only through their projection onto one-dimensi low-dimensional subspace, respectively. The study of learning single-img index models via neural networks is relatively well developed. Notably, [ee et al.] (2024)
show that neural networks can successfully learn single-index models with ((d) sample complex- ity, which is optimal up leading order, while Montanari & Urbani (2025) provides sharp DMFT theory analyzing the learning dynamics of single-index models using two-layer neural network when the sample size and input dimension are proportional.
"Equal contribution. The order determined by coin flip.

---

Published as a conference paper at ICLR 2026
However, a single-index model involves only one hidden feature, whereas in practice often expect multiple distinct features, motivating the study of neural networks learning of multi-index models. On the contrary, positive results on neural networks learning of multi-indes models remain rather limited. Note first that the information-theorriic limit for learning multi-index models is still (d)) samples, and there exist polynomial time algorithms Chen & Meka (2020); (Troiani al.] (2025); Defilippis et al. (2025); Kovačević et al. (2025); Damian et al. (2025) that achieve this for generic generativeeepponnn--oo multi-index models which covers almost all canonical examples. Never-
theless, in general, it is not clear whether neural networks can achieve this algorithmic limit and, if so, how. The earliest work on neural networks learning of multi-inde models Damian et al. (2022)
showed that standard two-layer neural networks can efficiently learn large subclass of multi-inde
models with sub-optimal ((d)) sample size. Subsequently, Abbe et al. 2022 Abbe al. (023a)
studied the problem of neural network learning multi-index models with a staircase structure. How- ever, their results are far from optimal and are restricted to special classes of functions. The closest attempt to the general problem is Montanari & Wang (002)), which shows that a two-layer network trained via gradient descent can weakly recover some of the features in a multi-index model using (d)) samples. Nevertheless, since the features are only weakly recovered and are not guaranteed to cover the entire span, their results do not establish that these features can be directly leveraged to successfully learn the targets in later stages of training. Therefore, the following remains open.
Can neural networks trained via standard optimizers (e.g.. gradient descent) achieve efficient, learning of generic multi-index models?
1.1 MAIN CONTRIBUTIONS
In this work, we work on this and conclude it positively for generative-staircase multi-index models. Specifically, we consider learning the target function fxx) = =g(xx) over the standard d-dimensional Gaussian input distribution, where U E Tkrxd is the hidden subspace and g : IR" R is polynomial of constant degree. This corresponds to the standard multi-index model with a polynomial link function, which generalizes the single-index model as the special case 1.
non-
Our main result, Theorem I shows that under generic assumptions, a two-layer neural network trained with layer-wise gradient descent can agnostically learn the target function up to (1) test error using (d) samples and O(d²) training time These sample and time complexities are optimal up to leading order, thereby demonstrating the adaptivity and ef- fectiveness of neural network learning. Moreover, our analysis applies to generic activation func- tions and broad subclass of loss functions, including but not limited to square loss, l loss, the Huber loss, and the pseudo- Huber loss. We emphasize that the assumption on the link function in our work is rather mild and generic, as all polynomials possess a generative leap exponent of two Chen & Meka (220)) Damian et al. (2025) and the loss function can be interpreted as a form of data preprocessing. The idea of data preprocessing has a rich literature in spectral methods and so on, but rarely considered in theoretical analysis of neural networks learning. The only substantial assump- tion in this work is that the target function has no generative-staircan structure. We acknowledge that this assumption is important and it is not clear whether and how to extend our results to the case with generative-staircas structure within a neural network framework. For further discussion, please see Section 2.1
Our proof begins by observing that near initialization, the gradient descent dynamics is nearly in- dependent across different neurons, with each neuron evolving similarly to power method iterations applied to the local Hessian. The local Hessian exhibits large eigenvalues corresponding to the sig- nal, namely the hidden subspace, as well as smaller eigenvalues arising from the noise because of the finite sample size. The key insight is that if this power iteration approximation works for too long, all features will align with the directions of the largest eigenvalues. that is, the most dominant directions, while ignoring the others, leading to failure of learning. However, if the number of train- ing steps too small, those noisy directions cannot be sufficiently eliminated, particularly when the sample size is nearly proportional to the input dimension. Therefore, the optimal strategy is to stop training the first layer at an appropriately chosen intermediate time. The features will still span the entire subspace while the noise is simultaneously eliminated.
We use Õ(-) to hide those (1)(1) factors.
2

---

Published as a conference paper at ICLR 2026
We note that using a spectral initialization or a hot start as in Mignacco et al. (202I); Chen & Shen (2025). The neural networks can only learn the right features from and only from gradient descent, trained for divering number of steps. This time scale requirement rule out all the approaches based on DMFT state evolution, such as Agoritsas al. (20T8); [Ba et al. (2022); Celentano et al. (2025); Gerbelot iet al.] (2023); Berthier et al. (2024) which requires either information/generaa hot start/spectral initialization,
standard initialization scheme for the neural network parameters, without
exponent one or guarantee the learning timescale can be made constant. We also highlight that the power iteration approximation does not work if the learning rate is too small. It important to choose a moderately large learning rate. For a more detailed discussion, please see Section 4
We defer the detailed discussion of related works to Appendix
1.2 NOTATIONS
We shall use X S Y to denote X < CY for some absolute positive constant C' and X 2 Y is defined analogously. We also use the standard big-O notations: (-)), O(-) and ((-)) to only hide absolute positive constants. In addition, we use Ô θ and Q to hide all the dda(1) terms throughout the paper. Let [k] {1, 21. k} for k E N+ and [a; b] = {a, ...,...,,,. bb for a, E N+ For a vector v, denote by || = (E |0g (2)/// the norm. When p = 2, we omit the subscript for simplicity. For matrix A, let ||A aap and || r be the spectral norm and Frobenius norm, respectively. We use Amax (-) and (-in(-) to denote the maximal and the minimal eigenvalue of a real symmetric matrix.
2
PROBLEM SETUP
Our purpose is to learn the target f R2 + R. where Rd is the input domain equipped with the standard Gaussian distribution N(0, Ia). Our target follows a multi-index model f"(x) g" (Ux) with hidden directions U E Rexd and the link function g" R" + R. We assume U has orthogonal unit rows without loss of generality.
Assumption 1. g g-)+) is a polynomial of degree p. Furthermore, 100 < C where C is an absolute constant.
We treat r; p as absolute constants throughout the paper. Next, we consider the model as a standard two-layer neural network
Se fe( (x: = -ajoow.a bz)
Here we use θ ; {aa, by. wy } to denote all the parameters. We make the following assumption on the activation function.
Assumption 2. o(.) E (R) is smooth activation function satisfying d'(0) 0, o² 1 and (g(3))) (-) M where M is absolute constant.
Let D= = {(x;, ))) be one dataset, where æ; Rd are drawn i.i.d. from distribution N(0, Ia) and y; f^(22 (xx). The empirical and population loss is given respectively by
(o) = Σ ((e(..) vi),
(D( = =R[
(2)
where / - E(t. y) RxR + R is a general scalar loss function. We make the following assumptions on the loss function.
Assumption 3. f(.-) (IR") l(t, u)|t=y = 0 and that it is convex in its first argument. Further- more, we assume
i
82
lt.
v)
L
and
e(t, v)
SL
(3)
Øt
82t
for all y.y R where I is positive absolute constant.
3

---

Published as a conference paper at ICLR 2026
We also consider a slightly relaxed version of Assumption for some specific activation functions. This assumption includes square loss as a special case.
Assumption 3'. We weaken the assumption in Assumption by allowing the first derivative to grow linearly with the inpul
a ltt. v)
S L + [8] lul) for all t, e R.
(4)
0t
Denote g.) = S0 ltt, y/ 1=) - as li.
Assumption 4. We assume E[f;] = O.
We define
Σ 1 71
(5)
and also its population counterpart Σg == (0, y) ax Recall that & is function of yi and that ye is itself function of w, Therefore, f; is function of U and we conclude that rank (E)) S via invoking Stein's 's Lemma and performing integration by parts.
Assumption 5. The rank is full, i.e., rank( =
Furthermore, denote the non-zero eigenvalues A S ... S. A1 We will assume X 0 and set A1 = 1 without loss of generality.
In our setting, we assume Xy > to ensure that Eg not degenerate in its support and provides a meaningful directional signal during optimization. We will denote and track the condition number = d but will still assume +x(1).
A,
2.1 DISCUSSION OF ASSUMPTIONS
We note that the first four assumptions are quite standard in theory literature. Assumption T is fairly generic. In Assumption 2 the key requirement are o0( O, ensuring that the activation behaves locally quadratic. The smoothness assumptions can be relaxed with more refined analysis. For Assumption 3 or 3 convexity in the first argument of the loss is essential to guarantee that the second layer weights converge its minimizer during the second stage of training, We note that this condition holds for a large subclass of commonly used losses such as square loss, Ll loss and Huber loss. The bounds condition under a more careful analysis, since most of our data remain in a bounded regime and the outliers are rare. Finally, Assumption is also standard and could be removed with either more refined argument or an additional data pre- pre-processing step.
the derivatives not critical and can be replaced by polynomial growth
The central assumption in our analysis is Assumption 5 Informally, this assumption requires that the second-order information of the data remain non-degenerate after applying the pre-processing map y + ('( y). We admit that for one fixed loss function f, it is straightforward to construct target task such that the second-order information of E' (((0.)) becomes degenerate. In practice, however, the target task is typically fixed, and one has the flexibility choose among different loss functions to learn it. Recent studies on generative exponents and learning thresholds Troiani et al. (2025);
Damian et al. (2024); Kovačević et al. (2025); Damian et al. (2025); Defilippis et al.) (2025) show that for a large subclass of common tasks, including all polynomial targets with non-generative- staircase structure, there exists pre-processing function that ensures the second-order information is non-degenerate, and importantly, such function need not be highly specialized. Consequently, we expect that (O( y) exhibits non-degenerate second-order information for at least some those commonly used loss functions in practice, and at least this already includes all targets in Damian et al. 2022). We will expand on this in Section
2.2 TRAINING ALGORITHM
We introduce layer-wise gradient descent algorithm.
4

---

Published as a conference paper at ICLR 2026
At initialization, set by zero and impose symmetric structure aj and wy Specifically, let the width m be even, and for all j E [m] we set
w (0)
0 am--j
(6)
m
For the first half of indices, , let a, be i.i.d. Rademacher variables: aj lid. Uniform +1}. For the first half of the features wj, we initialize them independently and uniformly on the sphere of radius 60.
The network is trained via layer-wise gradient descent with sample splitting, a common approach in theoretical analysis such as Wang al. 2023); Nichani et al. (2025). start with, we generate two independent datasets D, where algorithm proceeds as follows. we first train W via gradient descent for Ti steps on the empirical loss LD, (0), then reinitialize b and train a via gradient descent for T2 steps on the empirical loss CD. ))) The pscudocode of this procedure is given below.
1, 2 and each consisting of n independent samples. The training
Algorithm wase Gradient Descent
Require: Learning rates 1/2 f22, weight decay B1, B2, initialization level co, number of steps Tı, T2-
1: Initialize a, b, W.
2: Train W on dataset D
3: for t 1 to T₁ do
4: ( - (t m1 ( WW
+ ( (Q!t-1))
5: end for
6: ((T)) (T c0 V Lp, (0 + 3Ic0 BIcoW((T-1))
7: Re-initialize
8: Tf and Unif([ (--, 3)3, [m]
9: (TT)) (0)
10: (a W(T)) b(T)))
11: end
12: Train on dataset D3
13: for = + 1 to Tı Ty do
14: (!!) +- (t-11) - 72l 72 (0)+ + Bra
15: end for
16: Note T + T2
17: return Prediction function for (+) + (T( (O(T),o(Wir) (WW(T)) bbT,)))
MAIN RESULTS
Our goal to show that the network defined trained via Algorithm can efficiently learn multi-index models near the information-theoreti limit.
We use to denote an arbitrarily large absolute constant and, for technical convenience, assume n. m S d//4 throughout the paper. Before stating the main results, we further assume below that finite linear combinations of the activation can approximate any one-dimensional monomial to arbitrary precision.
Assumption 6. We assume the activation function has the monomial approximation property with exponent 3 2 0. That is, for every integer 0 and every p >> there exists a bounded weight
function (p) {+1} x 3, 3] + R with
1 11x S VE.P Sa.k p-
where Sa,k hides absolute constants depending solely on k and activations such that
sup |
Ea, [1k (a, b) (aa + b)] 2 S
(8)
where a is Rademacher RV i.e. P(a = +1) = P(a -1) = and b ~ Unifl 3, 3] independent
from
5

---

Published as a conference paper at ICLR 2026
We treat 3 as absolute constant throughout. We remark that this assumption is rather mild, which only requires that a two-layer network with one-dimensional input can approximate any monomial to arbitrary precision using expectation form. Moreover, conjecture that for most activations one can take 3 = and further set = 0, as suggested by Barron 1993); E et al. (202T). For concrete examples, Wang al. (2023) shows that have B 0 for ReLU activations in their Appendix D.6 and we also show in our Appendix E.3 that for the following activation
{ 2/4 12
M 2
(t)
(9)
|t|
<
we provably have 9 0 and thus can set p well.
We also require a small initialization scale, which is our last assumption.
Assumption 7. We set the initialization level en S C.mmlragr7 (3) To where Ce universal constans.
We note that is only required to be polynomially small provided T1 o(log d). With our as- sumptions in place, we are ready state our main theorem.
Theorem 1. Under assumptions 2 and assumptions 7
under the additional assumption that B we have the following result: For any first stage training time
T1 o(log d) and any sample size
2 d log d + d x²
(10)
there exists a proper choice of the hyperparameters and second stage training time T2 such that with probability at least 1 Old D/2), the Algorithm returns
+ Vn
c(O(T) S 4pT (log d) 4p+1
(II)
learn the target up the vanishing test with the width m O(1) and the sample size n 2 d'log 2,T + dd+ We specify those hyperparameters later in the statement of Theorem 2 Appendix [ We further remark that the sample size require- ment dllog 2TT, + d 11+T/T, x guarantees that the first-stage training learns the correct
In short, Theorem I states that
B is contributed purely by representations, whereas the eπror term <ATTi (log (Iogd)4p+11 v vn
the second-state stage training. We note that under the first set assumptions, the error bound can be refined 22pTT (log d))p+1((())))))))))))111
Via balancing two terms in our sample complexity requirement
こ d log dT? i + d' x²
(12)
the optimal dependence shall be derived from setting T1 θ and Ω to hide all the doa(1) terms.
log [
Note that below we use Ō
log
Corollary 3.1. Consider T,
log log
Under this, we further require
m = (log d))pp(++)++1 Sp((++)
Tng
and n = Ck2 (log d)² 22 log
(13)
where C a large universal constant.
Then, there exists a proper choice hypepparmnetrss that the Algorithm returns ((O( = oa(1) with T, @(1) and probability 1 (d D/2).
We note that the overall time complexity of Algorithm is ((nmdT). Under the setup of Corol- lary 3.1 have T ( m ((1), and (dd), yielding overall time of ( This also optimal up to leading order, since any algorithm must spend at least Q(nd) = 2(d2) time to read the full dataset.

---

Published as a conference paper at ICLR 2026
Compared with Damian et al. (2022), , note that our results exactly and completely recover the results in Damian et al. (2022) when we set T1 = 1. In this regime, successful learning is guaranteed only with a larger sample size π. = (dd)), and Dandi et al. (2025) further indicates that n = (cl)) necessary when T2 Our results improve upon Damian et al. (2022) and all previous works by managing to execute careful analysis of the multi-step fradient-dsseett dynamics on the empirical
loss over divering time horizon, We provide the details in Section
OUTLINE OF THE PROOF
4.1 GRADIENT DESCENT CAN MIMIC POWER ITERATIONS
To begin, we compute the gradient with respect to the feature wy as follows.
V w (0(0)
= n 1=1
(t)T Σ el (ferro U:) ayo w
zi x;
(14)
Our initialization ensures that the network output is zero at the beginning of the training. Therefore, within moderate large time horizon, we expect fei (a) remain near zero. Furthermore, recall
that || (0) = od()). For small ; 2; we have
(t)T
(w (w (H) x (0))) (TT x; + oa(1)
(15)
using d'(0) = 0 and substituting these yields
V Ven En (e (t)
(0, ((0,)))))"")))u (TT a, + smaller order terms.
(16)
Letting 4 := (0(0,, vi), under proper learning rate and weight decay, the dynamics can be expressed
as
1 n Σ Lara
((+1) w x
(17)
We emphasize that the above approximation works whenever fe (æx) and w0 (t) both remain small. Under those approximations, the neurons decouple from each other.
Under these approximations, the dynamics can be viewed power method iteration with respect to Σg. Recall that Ex the empirical version of Eg, satisfying |Σ3 Eellop O ((ddn)). Moreover, from the assumptions, the population one admits the eigendecomposition
Σ
(18)
where {u:}__] are orthonormal unit vectors spanning U, with eigenvalues ordered as 1 = X > 2 X, = 1/K, in accordance with our assumptions. We further assume a constant eigenvalue gap only in this section for convenience, namely min;zi [A, >yy 2 c for some universal constant e; We shall emphasize that we manage to avoid this assumption in our final results via a more fine-grained analysis (see Appendix C
Eigenvalues and eigenvectors of ΣT. tion
The empirical covariance Σg then admits the decomposi-
Σ = I
(19)
where {ū,) forms an orthonormal unit basis of TR.. We note that we have 1 - >11 = = (d/nn)
and || - uill = ((/d/)) for S r. For the others we have [Ael = ((vd/n) This decomposi- tion makes it straightforward to analyze powers of ES
d
E = Σ X ū u.!
(20)
1=1
7

---

Published as a conference paper at ICLR 2026
where ' denotes the number of gradient descent steps in the first stage. For random initialization " (0) the contribution from the noise components is
d
Σ Xmp++++
X 0
w) | (vVd/n)
T
(0) |10)
(21)
while the signal contribution along direction 2z
X ū w
00
(22)
Va
To ensure that the signal is at least comparable with the noise, require
K n-T
vd 2 (Vd/n)
(23)
which corresponds to the condition n Z d=+1/T R2
Assuming the noise term is negligible, for &o E [r] we now focus on the signal strength along direction Ug instead üs
u T. (0)
X
Σ
(24)
IsiSr 1
Here we require the first term to dominate the second so that the signal strength along the direction ue is well aligned with that along üg For a typical w the magnitude of the first term is at least
θ N //T w (0) )). while we have
Σ Fu.ū,u, =o 01 /vr)
(25)
(1103) V1/d). Therefore, suffices to take n 2 d log 2T We emphasize and highlight that it is quite interesting to observe that two conflicting factors involved the same time. If T is too small, the noise is not sufficiently suppressed; whereas if is too large, the weakest signal direction vanishes exponentially and becomes undetectable. Therefore, we have balance these two effects find the optimal T, which is interestingly found to be of the order O (Volggd).
using the standard estimates |u] = ((dd/)) and |ū] |0 e
Therefore, assuming n 2 dllog xT + d! +x² the learned feature approximately lies in and successfully spans the whole U subspace. We can then conclude that the second stage is able to learn any multi-index model with hidden subspace U from the learned features via using a standard kernel method argument. The complexity of this learning scales independent with d and polynomial with TT_ as some compensation is required to learn the weakest direction, which is still Õ(1) when T ( /log d). For concrete explanation, we refer to the proof roadmap of the second stage in Appendix D
4.2 COMPARISON WITH THE EXISTING APPROACHES
First of all, we note that our results cannot be recovered using a DMFT or state evolution framework, such as Agoritsas et al. (2018); Ba et al. (022);; Celentano et al. 20025;; ; Gerbelot et al. (2023); Berthier et al. (2024) and many more. Essentially, ly, in order to learn those generative-exponent two directions under the sample size @(d), one needs a diverging time horizon to escape the initial saddle point if there is no hot start explicit spectral initialization. Under the DMFT/state evolution propotional limit, the neurons shall not escape the initial stationary point.
Second, methods differ from the standard power iteration in matrix/tensor PCA problems & Richard (2014). If we enforce the matrix power iteration approx-
Balzano et al. (201T);[Mnnte imation for long time, the neurons shall collapse to single direction, which prevents successful
8

---

Published as a conference paper at ICLR 2026
learning. Instead, we force early stop (but still longer than constant) to the power iteration to ensure the neurons span the entire subspace. In practice, we believe neural networks can adaptively find the optimal stopping time when the second layer is trained together with the first layer. We provide numerical validations of our main results in Section
5 FURTHER DISCUSSIONS
We here further explain Assumption 5 which essentially requires target function has gen- erative exponent two without generative-staircat structure. To start with, we briefly introduce the concept of generative exponents. The so-called generative exponents has long standing literature (see related works) and is introduced in its most general form in recent work Damian et al. 20025) which sharply characterize the complexity of efficiently learning Gaussian multi-index models up to leading order. In the following, we use h2 to denote degree k Hermite polynomials and Hey to denote degree k Hermite tensors. A short introduction to these objects provided in Appendix EE.1
We denote = = x ~ N(O, I as the hidden directions and y as the model output. Denote further as subspace of the U subspace, and let zs E denote the orthogonal projection of onto S. The generative exponent given a subspace S, kg = kg (2, y), is defined as the largest l such that E[T 2s) p=)=)) = 0 for any measurable function T and any polynomial p of degree lower than k. This indicates that the tensor E[ 2ş) Hex, (=s )] * 0 and we will learn some new directions Snew- Therefore, we get new learned subspace S' = U Seew- Repeating this process, we will get an increasing sequence of subspace = So Ç S C ... C St = span(U), and the overall generative leap exponent is then formally defined as k ;= max; kis.
The generative exponent important since Damian et al. (2025) proves that the complexity of ef- ficiently learning a multi-index model is Odd via providing both the LDP lower bound and a matching spectral algorithm upper bound. Further, it can be shown that almost all common multi- index models have the generative leap exponent k* I 2. For instance, it works when the link function is 1) all the polynomials, 2) piecewise linear functions, 3) intersection of halfspace. There- fore, we can always extract information from computing EJ T 2s. (Zs,)))) with a proper pre- processing function T. In addition, Lemma F.14 of Damian et al. (2024) indicates that T need not to be highly specific. In fact, they prove that even random polynomials suffice in the single-index setting.
We emphasize that this observation motivates Assumption In addition to generative leap exponent two, we are essentially requiring that there exists one pre-processing function T which is able to extract all the spanned directions in the subspace U from the second order information. Our assumption is weaker than the previous literature using neural networks to learn multi-index models up to vanishing test error. For instance, via using Huber loss, we can deal with the link function g(z) Σ hz,),) which not covered previous papers like Damian et al. 2022)).
6 CONCLUSIONS
In this paper, we demonstrate that under Gaussian input, a two-layer neural network trained with layer-wise gradient descent can learn generic multi-index models up to oa(1) test error using Ö(d) sample and (dd)) time complexity, which matches the ceammatiionthhoorete threshold up to the leading order. We emphasize that our analysis does not rely on any particular choice of activation, loss, or link function; rather, the assumptions imposed are generic and hold for a broad class of functions. An important and somewhat unexpected insight from the proof is that the inner layer must be trained for more than a constant number of steps to successfully learn the hidden features, in order to recover the hidden features most efficiently and thus achieve optimal performance. This observation impossible to be observed via a DMFT or state evolution framework over bounded time horizon, and has the potential to extend to other settings. Overall, our results demonstrate that neural networks are capable of effective representation learning and can fully learn the target function this canonical task with near-optimal efficiency.

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
The authors thank anonymous reviewers for their valuable feedback and suggestions. Zihao Wang also thanks Alex Damian, Andrea Montanari, Eshaan Nichani, Fan Nie, Yunwei Ren, Rui Sun, Zixuan Wang, Denny Wu, Jiatong Yu for their insightful discussions and valuable feedback during the completion of this work. Jason D. Lee acknowledges support of NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 2144994.
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-stairease property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782 4887. PMLR, 2022.
Emmanuel Abbe, Enric Boix Adserå, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to-saddiy dynamics. In The Thirty Sixth Annual Conference on Learn- ing Theory, pp. 2552-26666622 2623. PMLR, 2023a.
Emmanuel Abbe, Enric Boix-Adsera and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saddle-to -oaddle dynamics, 2023b.
Emmanuel Abbe, Enric Boix-Adsera, and Theodor Misiakiewicz. The merged-staircantt property: a necessary and nearly sufficient condition for sgd learning of sparse functions two-layer neural networks, 2024.
Elisabeth Agoritsas, Giulio Biroli, Pierfrancesco Urbani, and Francesco Zamponi. Out-of- equilibrium dynamical mean-field equations for the pereeptron model. Journal of Physics A: Mathematical and Theoretical, 5 (8))008500,,, January 2018. ISSN 1751-8121.
Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels?, 2020.
Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows sgd to learn high-dimensional multi-index functions, 2025.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference, 2021.
Benjamin Aubin, Antoine Maillard, Jean Barbier, Florent Krzakala, Nicolas Macris, and Lenka Zdeborová. The committee machine: computational to statistical gaps in learning a two-layers neural network. Journal of Statistical Mechanics: Theory and Experiment, 2010(1)) December 2019, ISSN 1742-5468.
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, 2022.
Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks, 2020.
Laura Balzano, Robert Nowak, and Benjamin Recht. Online identification and tracking of subspaces from highly incomplete information, 2011.
Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and I Lenka Zéeborová. Optimal errors and phase transitions high-dimensional generalized linear models. Proceedings of the National Academy of Sciences, 16((1))5451-4460 March 2019. ISSN 1091-6490.
Andrew Barron. Barron, a.e.: Universal approximation bounds for superpositions of a sigmoidal function. ieee trans. on information theory 39, 930-945. Information Theory, IEEE Transactions on, 39:930 945, 06 1993.
William Beckner. Soboley inequalities, the poisson semigroup, and analysis on the sphere s" Pro- ceedings of the American Mathematical Society, 118(2)519-5228 1992.
10

---

Published as a conference paper at ICLR 2026
Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. Foundations of Computational Mathematics, 25(5):622--1710, August 2024. ISSN 1615-3383.
Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
Joan Bruna and Daniel Hsu. Survey algorithms for multi-inder models, 2025.
Michael Celentano, Chen Cheng, and Andrea Montanari. The high-dimensional asymptotics of first order methods with random data, 2025.
Sitan Chen and Raghu Meka. Learning polynomials of few relevant dimensions, 2020.
Yuchen Chen and Yandi Shen. Learning single index model with gradient descent: spectral initial- ization and precise asymptotics, 2025.
Alex Damian, Jason D. Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent, 2022.
Alex Damian, Loucas Pillaud- -iviie, Jason D. Lee, and Joan Bruna. Computationall-satisttti gaps in gaussian single-index models, 2024.
Alex Damian, Jason D. Lee, and Joan Bruna. The generative leap: Sharp sample complexity for efficiently learning gaussian multi-index models, 2025.
Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krza- kala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents, 2024.
Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. How two-layer neural networks learn, one (giant) step at time, 2025.
Leonardo Defilippis, Yatin Dandi, Pierre Mergny, Florent Krzakala, and Bruno Loureiro. Optimal spectral transitions in high-dimensional multi-index models, 2025.
Ilias Diakonikolas, Giannis lakovidis, Daniel M. Kane, and Lisheng Ren. Algorithms and sq lower bounds for robustly learning vaalev multi-index models, 2025.
Weinan E, Chao Ma, and Lei Wu. The barron space and the flow-induced function spaces for neural network models, 2021.
Hengyu Fu, Zihao Wang. Eshaan Nichani, and Jason D. Lee. Learning hierarchical polynomials of multiple nonlinear features with threc-layer networks, 2024.
Cedric Gerbelot, Emanuele Troiani, Francesca Mignacco, Florent Krzakala, and Lenka Zdeborova. Rigorous dynamical mean field theory for stochastic gradient descent methods, 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Filip Kovačević, Yihan Zhang, and Marco Mondelli. Spectral estimators for multi-index models: Precise asymptotics and optimal weak recovery, 2025.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet classification with deep convo- lutional neural networks. In Advances in Neural Information Processing Systems, 2012.
Jason D. Lee, Kazusato Oko, Taiji Suzuki, and Denny Wu. Neural network learns low-dimensional polynomials with sgd near the information-theoree limit, 2024.
Yue M. Lu and Gen Li. Phase transitions of spectral initialization for high-dimensional nonconvex estimation, 2019.
Wangyu Luo, Wael Alghamdi, and Yue M. Lu. Optimal spectral initialization for signal recovery with applications to phase retrieval. IEEE Transactions on Signal Processing, 67(9):2347-2356 May 2019. ISSN 1941-0476.
11

---

Published as a conference paper at ICLR 2026
Antoine Maillard, Florent Krzakala, Yue M. Lu, and Lenka Zdeborová. Construction of optimal spectral methods in phase retrieval, 2021.
Francesca Mignacco, Pierfrancesco Urbani, and Lenka Zdeborová. Stochasticity helps to navigate rough landscapes: comparing gradient-desccet-baaed algorithms in the phase retrieval problem. Machine Learning: Science and Technology, 2(3):035029, July 2021. ISSN 2632-2153.
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of Machine Learning. MIT Press, edition, 2018.
Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval, 2018.
Andrea Montanari and Emile Richard. A statistical model for tensor pea, 2014.
Andrea Montanari and Pierfrancesco Urbani. Dynamical decoupling of generalization and overfit- ting in large two-layer networks, 2025.
Andrea Montanari and Zihao Wang. Phase transitions for feature learning in neural networks, 2026.
Eshaan Nichani, Alex Damian, and Jason D Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. arXiv preprint arXiv:2305.0600 066966 2023.
Eshaan Nichani, Alex Damian, and Jason D. Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks, 2025.
Ryan O' O'Donnell. Analysis of Boolean Functions. Cambridge University Press, 2014.
Giuseppe Da Prato and Luciano Tubaro. Wick powers stochastic pdes: an introduction. 2007.
Yunwei Ren, Mo Zhou, and Rong Ge. Depth separation with multilayer mean-field networks. arXiv preprint X3Xx:22344.0063,, 2023.
Emanuele Troiani, Yatin Dandi, Leonardo Defilippis, Lenka Zdeborová, Bruno Loureiro, and Flo- rent Krzakala. Fundamental computational limits of weak learnability in high- diggeniiinsional multi- index models, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018
Zihao Wang, Eshaan Nichani, and Jason D. Lee. Learning hicrarchical polynomials with three-layer neural networks, 2023.
Colin Wei, Jason D. Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and optimization of neural nets v.s. their induced kernel, 2020.
12

---

Published as a conference paper at ICLR 2026
Appendix
Related works
13
B Numerical illustrations
15
Results
15
B.2 Proper choice of loss function
15
Proof of Lemma C.1: Feature learning
17
Proof of Lemma C.1
19
Proof of supporting lemmas in Appendix
D
Proof of Theorem
26
D.I Proof roadmap of Theorem
26
D.2 Events of high probabilities
27
D.3 Proof of T'heorem
28
D.4 Proof under 3 0 and Assumption
36
D.5
Proof of supporting lemmas in Appendix D.3
39
Technical background
57
E.
Hermite polynomials
57
E.2
Uniform generalization bounds
58
E.3
Univariate approximation
59
E,4
Convex optimization
61
Sphere lemmas
E.6
Supporting technical lemmas
62
A RELATED WORKS
Learning multi-index models. In multi-index model, the output depends on a low-dimensional subspace of the input data. This framework has a rich history in statistics and ML, where it used to model -ow-iimensional structures with high- sich-dimensionll data. A comprehensive survey of current algorithms for learning multi-index models provided in Bruna Hsu (2025). In particular, is known that under generic conditions, the information-theored learn such targets is proportional to the input dimension, as rigorously established in Barbier et al. 2019); Aubin et al. (2019). Nevertheless, it remains largely unclear whether and how widely used polynomial algorithms can efficiently recover the hidden subspace.
sample complexity required to
To start with, Barbier et al. (2019): Mondelli & Montanari (2018); Lu & Li 2019 Luo et al. 2019); Maillard et al. (2021) studies generic single-inder models and sharply derives the learning threshold. In particular, Mondelli & Montanari (2018), : Lu & Li (2019); Luo et al. (2019): Maillard ct al. (202T) show that for broad class of single-index models, spectral methods with suitable data pre-processing recover the hidden direction the optimal threshold, matching the performance
13

---

Published as a conference paper at ICLR 2026
of optimal approximate message passing (AMP) Barbier al. (2019). In contrast, Arous et al. (2021) considers a more restrictive setting and examines the learning of single-indes models when using online SGD the optimizer without data preprocessing, and introduces the concept of an information exponent, which characterizes the learning complexity in their setup.
A substantial line of works studies how to extend these results to multi-index models such Chen & Meka 2020); AAbee et al. (2022 20 Troiani et al. (2025): Damian et al. (2025): Kovačević et al. (2025 Defilippis al. (2025); Diakonikolas al. (2025). On one hand, Abbe et al. [2024] 2023b) introduces the notion of leap complexity, showing that hierarchicallstaircass structure in multi-index
Damian et al. (2024) 2025) extend this line of work by formalizing the generative exponent for single-index models and the generative leap exponent for multi-index models, extending the previous results to higher order and incorporating the benefits of data pre-pne mre-processing. In particular, we highlight Damian et al. (2025) for its two key contributions.
models can facilitate learning. On the other hand, Damian
For multi-index models with generative exponent k", the authors establish up leading order lower bound [2(dk /2) under the low-degree polynomial (LDP) framework and matching upper bound O(d'" /2) achieved by a spectral polynomial-time aammeealggiite algorithm for the sample complexity.
Furthermore, it shows that the generative exponent is most two for almost all common multi-index models, indicating that ((d) samples are sufficient to efficiently recover the hidden subspace.
These papers are in general not tight in terms of leading order constants. We also note that Di- akonikolas et al. (2025) introduces a related notion, i.c., the m-well-behavee multi-index model, which coincides with the generative exponent in many cases. While their sample complexity ex- hibits a worse dependence on d compared to Damian et al. (2025), their work offers two advantages: it targets agnostic PAC learning rather than subspace recovery, and provides uniform sample com- plexity guarantees that apply agnostically to all multi-index models within certain given class.
Subsequently, for the precise (weak recovery) learning threshold for multi-index models, TTroiani et al. (2025) non-rigorously derives the sharp optimal sample size required for weakly learning peteeraiiee-cap-xppeeeette
multi-index models by performing a local analysis Bayes AMP, extending the results of Barbier et al. (2019); Mondelli & Montanari (2018) for single-index mod- els when the generative exponent the target equals two. Moreover, as two follow-up works, : Kovačević et al. (2025); Defllipiis et al. (2025) study learning multi-inded models using optimal
spectral method, although their sample thresholds in Kovačević et al. (2025) are in general not al- setup. We emphasize that the above works depend on the special structure of multi-index models and therefore do not use general- algorithms such as learning with standard neural networks trained via gradient descent.
ways optimal, and none of those works can deal with the generative-staire
Feature learning in neural networks. The term feature learning refers to the following phe- nomenon. In practice, neural networks trained with local gradient-based algorithms such as gradient descent can adapt to the low-dimensional structure of the tasks and efficiently learn useful hid- den features, which in turn facilitates successful learning and enables effective transfer learning to downstream tasks. Among the first several works in this direction, as we have mentioned previously, Damian et al. (2022) demonstrates that neural networks can learn generie multi-index models at sub-optimal sample size of = ((d2) under some additional assumptions. Similarly, Abbe et al. 2024 2023b) establish a sub-optimal rate of
Ö(d²) for neural network learning, where L denotes the leap complexity. [B et al. (2022) studies the problem under a much more restrictive as- sumption of an information exponent one, and in this setting shows that neural networks can succeed with (dd) samples.
Furthermore, as noted earlier, Lee et al. (2024) shows that neural networks can learn generic single- index models with generative exponent at most two with sample complexity O(d) through data re-using while Montanari & Urbani (2025) provides sharp DMFT theory analyzing the learning dynamics of single-index index models using standard two-layer neural network when the sample size
and input dimension are proportional. Together, these two works essentially conclude the studies of learning single-index models with neural networks. Next, Dandi et al. (2024) demonstrates that
14

---

Published as a conference paper at ICLR 2026
neural networks can learn multi-index models with generative leap exponent using = (d)) samples in constant time, leveraging DMFT theory. Finally, as noted previously, Arnaboldi et al. 2025) provides loose result, showing that neural networks with ld d) samples can recover the hidden subspace of multi-index models but in weak recovery sense when the generative leap exponent of the target equals two. Besides linear features, we note that there is also a substantial line of work on learning nonlinear features with deep neural networks, including Nichani et al. (2023) Ren et al. ) Wang et al. (2023); Fu et (2024).
We aware of and highlight this recent concurrent work Montanari & Wang 2026). Our works differ the following aspects. Montanari & Wang (2026) studies learning a multi-index model us- ing two-layer network trained with standard gradient descent instead of modified gradient-based algorithms. Their arguments are performed in the proportional regime, that is, n/d - (0, +0o), leveraging toolboxes from discrete time DMFT and random matrix theory. Besides, their results are sharp in constants, and they derive the sharp phase transition threshold at which the trained neu- ral network can start weakly recover the first hard direction (with generative exponent two). In contrast, Montanari & Wang (2026) does not establish full subspace recovery and therefore cannot provide an end-to-end analysis of the problem. Instead our work provides an end-to-end learning guarantee via somehow different mechanism, albeit at the cost of slightly worse sample com- plexity and the use of a modified layer-wise training scheme.
B NUMERICAL ILLUSTRATIONS
We conduct preliminary numerical experiments to validate our theoretical predictions.
We consider the following target function (a) = + /x3)/ /5/2)//5/22/ 5/2 where E F2 and x; denotes the i-th coordinate. This is a multi-index model with r = 2 hidden directions. We use two-layer neural network with width m = and quadratic activation function (t)) = t². Both layers use standard Kaiming uniform initialization for weights and biases. For each dimension d and sample size exponent a, generate
|dd] training samples from N(O, Ia). We train the network using Adam optimizer with learning rate 0.005 for up to 1000 epochs with batch size 32. We vary the dimension d logarithmically and search the sample size exponent in the range [1.1, 1.8] with step size 0.01. For each (d, a) pair, we run 5 independent trials with different random seeds and report the average minimal sample size exponent a required achieve test error below threshold
B.1 RESULTS
Figure shows the aanimal-sammle-siiee exponent a required to achieve test error below threshold converges to one as d increases for different thresholds E {1.0, 0.1, 0.01). Each point represents the average over 5 random seeds. These experiments validate that neural networks can indeed learn multi-index models with near-optimal sample complexity. The decreasing trend in dimension grows provides empirical support for our theoretical finding that Ö(d) samples suffice for successful learning, achieving the information-theorrtiic limit to leading order.
B.2 PROPER CHOICE OF LOSS FUNCTION
We demonstrate that the choice of loss function critically impacts learning performance for subset
of targets. We consider a challenging target (z12)) + (x where h4(2) = /24 (2+ -624-62- 62 + 3) is the degree-4 Hermite polynomial. This target has generative exponent 2 but information exponent 4 which is larger than generative exponent and requires a properly chosen loss function rather than square loss to ensure the second-order information is non-degenerate. To this end, we train a two- layer neural network with width m 4 and cosine activation (t) = cos(t) using two different loss functions: Mean Squared Error (MSE) and Huber loss. We measure the learning process by the alignment metric cos(best), defined as the maximum cosine similarity between any first-layer weight vector and the true directions e₁ or e2. A value close to one indicates successful recovery of the hidden directions. For dimensions d E {200, 300, 500), we vary the sample size n and plot the alignment cos(best) function of sample complexity ratio n/d, For each (d, n) pair, 10 independent trials with different random seeds and report the median along with 30th and 70th percentiles.
15

---

Published as a conference paper at ICLR 2026
011
Minimal (wheee d²)
1.4
133
log d
Figure 1: Minimal sample size exponent (where d²) required to achieve test error S function dimension d. Three thresholds are shown: e = 1.0 (light blue), c = 0.1 (medium blue), and 0.01 (dark blue). More stringent thresholds require larger (more samples), but all curves exhibit similar downward trends.
0.8
cue(best) (Final Aligumeat)
MBE (0x20)
0.6-
MSE (suaao
(0=200)
Hober 60- 200
Haler 61-300
(0-5000
0.0
a/d (Sample Complexity
Figure 2: MSE and Huber loss for learning (x1) + h1(22) across different dimensions d e {200, 300, 500). The z-axis shows n/d, and the y-axis shows the cos similarity cos(best). Lines show median with error bars indicating 30th to 70th percentiles over 10 random seeds. Blue shades represent MSE and red/orange shades represent Huber loss.
Figure 2 reveals striking difference between the two loss functions. MSE loss consistently fails to learn target function, with alignment remaining near zero even when n/d = 80. stark contrast, Huber loss exhibits clear phase transition: the alignment rapidly increases from near ( to near once the sample size exceeds a critical threshold around n/d = 20. This threshold consistent across all three dimensions. This experiment validates our theoretical assumption (Assumption) 5) that the pre-processing function l' (0, y) must ensure non non-degenerate second-order information. For this particular target, MSE loss corresponds to degenerate pre-processing (since E(yxx J| 0), while Huber loss provides the necessary non-degeneracy. This demonstrates that: (i) the choice of loss function is merely matter of optimization convenience but fundamentally influences
16

---

Published as a conference paper at ICLR 2026
whether the thforration-theoretii limit be achieved, and (ii) framework correctly predicts that with proper loss selection, neural networks can learn targets beyond previous theoretical guar- antees with linear sample complexity.
C PROOF LEMMA C.1 FEATURE LEARNING
We will use 7) to refer to the learning rate 7/1 and use 3 to refer 31 throughout this section. We recall that we use D to refer to large absolute constant.
We will first introduce several more notations to help to do the proof. In the first stage, expect the prediction fe (z;) to remain small, since the norm of the features should remain small. We apply first-order Taylor expansion of the first argument of e(()) around zero
E' (0, mi) ('( feece (x:)
for some
(c
S (e (3x)
(26)
Denote the remainder term by
= ("(5,") (t ye)
(27)
which satisfies the bound ( S L| (ax) and should be small in principle. Substituting the expansion into the gradient expression for the j-th neuron, we get
V wy (0(00) = 1 Σ (0) +
(0) x;)x,,
(28)
Grouping terms, we have
V Ww, ()))) = aj Σ E' (0, jo' (w
(1) x ax
Σ A o (w)) TT ,,)
(29)
7eroth-oddr term
Perturhation
To analyze the gradient dynamics, we further expand the derivative the activation function using second-order Taylor approximation around zero. Since a d0( = obtain
d((z) = 00))= + 20 (E)== for some
K <
(30)
Substituting the expansion of d'(-) into Vw, Σ, we obtain
n
aj n
V =
aj
Lo (E:)(w) a: Pa,
Σ
A (4)
o'
(unJar;)an
a: a.)
2n
n 1=1
1=1
Perturbation
-aoo(0))
Qj
Σ
Lo (3) (E)) 20
(31)
2n
12
+20"(0)
Σ
[w)
x az
A O (E)
The first term captures the leading-order linear dynamics that should be dominating, while the sec- ond term is higher-order perturbation term that will be controlled using norm bounds on un5) and also bounds on A (0)
Under gradient descent with proper weight decay (recall that we set }) and the learning rate as y) the update rule for w becomes
w ((+1) =n nVv,, E (1) -njaaaaaaaa (0)) + R
(1)
(32)
17

---

Published as a conference paper at ICLR 2026
where the residual term R2.
(()
is defined by
n
10 =
1 2n
Σ
4o (3) (g())) (w0) x;)
(t)T
n
(1) (w)
(t)T
24 ja, + A (E, (t) xx) x1)²x,
+2
(33)
i=1
Since (w (T) 22)22; (w)
( x; ) w (0)
we can write Rio as RR))) ma, ((02 where
Q(t) E Rddd is a symmetric matrix defined by
Q10 2n
1
Σ 4o (c.0
(t)T x,)x,
n
1
x Σ / (c()) (0) x.)),aa (
+2
(34)
i=1
Now the entire update can be expressed
w w ((+1)
= - (St+ (00 (0)
(35)
Define
M (0)
Σ. + Q so that
w w
(1+1) na, (1))
(36)
Through all the calculations above, we obtain a compact form of the update rule with clear linear structure. Then by repeatedly following the updates for Tr iterations, we achieve
Ti-1 II
(T) 1 (-a,) 60
Ti
TT
Ti
(1 1=
(37)
This shows that w lies in the image of a composition of linear operators, each consisting of a dominant direction determined by Ee, plus time-dependent perturbation ((1),
The general picture is that Q is negligible when = o(log d). In this case, we have w) (1+1) o E and the dynamics of GD will be similar to a power method iteration. For now note that our power method iteration is with respect to E, not Σ, since the sample size is finite.
This section is devoted to a complete proof to bound the difference between w update
(Tı) wT(T) and the ideal
W (T)) = (-a,) T' ΣT. 3 (0)
(38)
Note that we have
w
(T))
w
(T))
- T
M IT.
(0)
(39)
(II.. M (()
101
and therefore we just need to bound
Lemma C.1 (Total Perturbation Error). With probability at least 1 - (d assuming Z dT? and m < D/4 for all j E [m], the total perturbation error is bounded by
[
TT-1
log d + To where 7= CV d
II
(0)
M (0)
ScoTiy
(40)
d
where C a universal constant.
The proof of Lemma C.1 is provided in Appendix C.1
18

---

Published as a conference paper at ICLR 2026
C.I PROOF OF LEMMA CA
We first give sharp characterization of Σx, We will bound the difference of the empirical and the population. Define
He Σ La B[(ax""
(41)
Lemma C.2 (Operator Norm Bound for Hz). With probability O(d-2D) we have Hi llop Va
From this lemma, we expect sharp concentration once n > d. This is standard in literature. Next, bound the magnitude of (a).
Lemma C.3. Let S (a)) -ー=1 (w (1) x)) be the oulput of a two-layer neural nelwork at iteration t, where each w R& satisfies ||u S and the input IR satisfies ||>||| 22d..
Then, for all such æ, the network output satisfies the bound
4 [x)) S 2e M de/2
Me
(42)
Next, we start to control the operator norm of the matrix Q
Lemma C.4 (Operator Norm Bound for (()). Suppose the following holds. Each input satisfies ||لح||| < 2/a. Each weight vector satisfies | || S c. The network prediction at iteration t satisfies (Sece (a:)) & Cy-
Then the operator norm of Q(t) satisfies
|q [1 S 4/d (LM LC + 4LC,d
(43)
Plugging in the concrete estimation for the constant Cy, also note that without loss of generality we assume I 2 1 and M 1. have the following.
Corollary C.5. Suppose the following holds. Each input satisfies || S 2vd. Each weight vector satisfies |ml (0) c.
Then (()) satisfies
40 3
Q ( lop S 4LMed + &Lmed + 2. 2/ + LM s M2de (44)
With all the lemmas and corollaries above, define the following high probability events. First, define the event εи samile noise bound) as follows:
Eи = Hellop S y where 7 = C
(45)
By Lemma C.2 this holds with probability at least 1 - Odd-2D) Furthermore, provided n 2 dTT know that 7 also 7 2T, Second, define the event &D (truncated norm) follows:
ED [|<=||| < 2vd for i [n]]-
(46)
By Lemma E.18 this holds with probability at least - exp(-cd), thus at least 1 - Od
Third, we define the event E as e8¹ = @ op - } Note that under event &p. provided
I 12 we have el happens. notice that el¹ has no independent randomness,
is defined just for notation convenience.
In the following we analyze the gradient update conditioned on the event EH n Ep, which holds with high probability. First bound the norm of un (1)
19

---

Published as a conference paper at ICLR 2026
Lemma C.6 (Norm Bound after T₃ Steps). Suppose the event Eh NED holds, and let the learning rate be set as
1 1/(2T))
7 Cn
(47)
rt2
for a sufficiently large but fixed constant Cg where & 4D log Suppose
11 | =
(48)
128mL2M2d/22
Then for all 0 Sts
have
|/00 S Vdeoo
and ||/ 128mL2
1
(49)
Next, remove Q out of (d = Eg Q(t) and compute the perturbation emror.
Lemma C.7. Suppose the event D holds. Define (MT M ET 10 Assume that ED S 125mL2M=072 Then the perturbation error is bounded by
(0)
5 T1-1 4
| S T
22
(50)
Next, we deal with the difference between Eg and Et Lemma C.8. Let Uni)())-((Eo)) be a random independent vector. Then with probability at least 1 - O(d ) we have
log d
|1EEe Hz) w DT'wil w|l $ ca (
(l + He (Hellop)*
-1Eellop
d
col He llop-
(51)
Corollary C.9. Under the setup of Lemma C.8 and note assumed |He llop y with y
1/(2T1). Then with probability at least 1 O
2D), have
[ + Ho) To E/'ru w/ Sco
log Trmp
d + , )
(52)
Proof. From the binomial expansion,
Ti (T)
Ti7 (1 + )z - - (1= Σ -
Σ j-2
(53)
Since (T= < (T (( Ti>) 2(T < Tiy, get (1 + TT - 1 2T,y. On the other hand, as in Lemma C.8 one shows under the same event that
(Ee + H.) Ti E w S co (1+ 1/ d t
d
(54)
Combining the two bounds yields the claimed result.
By putting all the things together, obtain the final bound of Lemma C.1
Proof of Lemma C.7 We first prove another lemma. We aim first show that under the stated condition on En. the perturbation error from the vire-varying iterates smaller than the leading- order perturbation induced by Hi.
20

---

Published as a conference paper at ICLR 2026
Lemma C.10. Suppose the event Еи NED holds. Recall the error from time-varying perturbations.
(1T.-0 M
Eg 10)
denote
Then, provided that 6q mLM
log d (4/5)T:
have
A S T logg
(55)
Proof. Recall the bounds established in previous lemmas. We have the one-step expansion error from iterated time-varying perturbations
T1-1
||||
S
Ti
(56)
We plug in one eg
5
T1-1
|
S T
md
2/2
4
5
T1
4
Ti
log d
S T1
4
7/2
E0
mdd/2
3
'log d
2 TTnV
E0
(57)
and the proof complete.
Finally we decompose the total perturbation error into two parts
T 1
Ti-1
II 1=0
(
ET
w/ =
II t=0
M (1)
TT
w/0 + STD +ETT E w
log
logd
$ Tiny
co + cu
Trny
d + tx
d
Llogd + tT
logd d
$ 60
Trn
(58)
The conclusion is then straightforward if we adapt the previous two lemmas to bound these two parts separately.
C.2 PROOF SUPPORTING LEMMAS IN APPENDIX
C.2.1 PROOF OF LEMMA C
Proof. We analyze the operator norm via a net argument.
Recall that Hellop [uPl|//=|elpp u Hiv. Let M/4 1/44 be two --nets of the unit sphere in Rd Then
Hellop 2
u HEv
(59)
WEMVA
Fix any u, u, and define
7: := (u (x) x,),, Fl(( (x) (T)(v<))( (=Tr) m:){]
(60)
Note that we have u T H.v = + - Z; Since æ; is standard Gaussian and f E [-L, L), we have
|| Ii & L. Therefore we can invoke Bernstei'ss inequality and do union bound over the net to
obtain
enl2
P( H lop 2 t) < 81d exp
(61)
I,2
Choosing t= = CLy d for sufficiently large C, we obtain the desired conclusion.
21

---

Published as a conference paper at ICLR 2026
C.2.2
PROOF OF LEMMA C
Proof. Using the Taylor expansion of (-)) at with d'(0) 0 and o² 1, we have
1 2
<(2) = (0) +
(c)= for some
[ S |=|
(62)
Then
(T) x few (x) a w
Σ ag
(0)
w
(
x)
2
1
(3)
((0)
(OT
x)3
(63)
Since Σg aj 0, the constant term cancels and
fem (x)
(1)
x
1
aoo
(3) (20
(t) x)3
(64)
6
We then bound each term. We have
(t)
S [m0
(1) |||||||| S 2evd..
(65)
Therefore (w) (t) x)² s 4 and (w)) T
x) I 8/2 Then
(w)) (()T x²)² S d,
Σ ajo (5)) (un) ax) < <
(66)
and the following is straightforward
Isem (x)) - (Ame2d) 1
(8mMc)d//2) 2 _MM20/22
(67)
The proof is complete.
C.2.3
PROOF LEMMA C.4
We upper bound the operator norm of (()) by treating the three additive components sepa- rately.
Proof.
Let
1 n
Σ Lo (c( (Ew) (w] )
(68)
2n
We bound each component |wo xi (((())
I [2 |||=ا||||| S 2evd, |o = ||x;1| <
4d,
< LM. Thus,
(10
(8) lop - (2V 4L2
(69)
2n
Let
@
(t)
1 A z
Σ
(70)
71 1=1
Each term satisfies | < LCy, | llop S 4d, so
|1 lo < (4d)
(t)
(71)
Let
@ =
2n Σ A (t) σ (3)
(72)
We again use 14 (3) (e( ) < M,, w æ; < 2e [ llop S 4d, to obtain
1 |
Q f lop -
ALC,
(73)
2n
We combine the operator norm bounds
|| < [|/ llop + // llop + 11/ & (MM + 4LCdd.
(74)
The proof is complete.
22

---

Published as a conference paper at ICLR 2026
C.2.4 PROOF OF LEMMA 6.6
Proof. We first prove the following supporting lemma.
Lemma C.11 (One-Step Norm Control). Suppose the event Eи n &D EQ holds, and let the learning rate set as
1/(21))
1 71
(75)
On
rt
for sufficiently large constant Cy where = 4D log d. Then the gradient update satisfies
| s d2 10 |
(76)
Proof. Under the event E n @ we have
M Tlop
EEllop + Hellop + |
I 1 +
+
1
3
(77)
4
Then,
[ [1 = v/ M 2 S 7/ M
loop
|1
11
3 2
ア
|
11
(78)
Plug the formula for 7 into the above one and we obtain the desired conclusion.
We prove by induction on t E {0, I,.., Ti } that ||w,/1| S (2(T)) where we assume
6o 1287 The t 0 case is trivial. Suppose for some T1, we have
| s (221)))
(79)
Plugging 6o in, |
I Maa thus the event EH n &D n &Q holds, apply the previous
e(t)
lemma C.11 we get
( s d |10 |10 ((+1)/(2T))
(80)
Repeating this induction, the result is straightforwwrd. In particular, for all 0 SISTi 1, we have
| (1) S Vdep.
C.2.5 PROOF LEMMA CV7
We begin by expanding the product of perturbed operators using a standard multiplicative expansion
Proof.
II 2=0
M
ST
Σ
Σ
(TT) T
(81)
a=
OST,
<T
Subtracting ET from both sides yields
T--1
II
M -
ST Σ Σ
(T) ST
ST)
(82)
We now apply this expansion to un) (0) and estimate the resulting norm. For each term in the expan- sion with s perturbation matrices, we use the operator norm bound
(T--+)-1(())) Q (T) ST. w
(0) < 1T TT-. II
(1))))
(0) llop
1|0
(83)
A=1
Note that the event EH EDE holds, Σgllop < Σ Eellop
Hellop < 4- Combining with
60 5 128mL2
(84)
23

---

Published as a conference paper at ICLR 2026
by Lemma C.6 for all 0 S -1 ||u0] 2 S Vdeo- Moreover, using Corollary C.5 to know that
it is easy
0 S 32LM m
(85)
up
Hence, each term with perturbation matrices is bounded by
(32L. 2co co-
(86)
Summing over all such terms (total of (Ti) per order 8), we have
T
II
M
ET
w (0)
T
5
T1-9
(32L 2 (20
(87)
4
t=
This sum can be rewritten as
5 + 32L M
T1
5
T1
6o
md
Eo
(88)
Applying the first-order binomial expansion which is valid for the small enough in our assumption and we conclude
5
T1-1
|
Tı
2
(89)
The proof is complete.
C.2.6 PROOF OF LEMMA C.8
Proof. We first introduce several standard concentration results.
Lemma C.12. Suppose A is matrix in ddd,, and w Un Then with probability at least 1 - O(d-2D), we have
log for all s E [r].
(Aw)s $ ||A| op
(90)
Proof. Fix [r], and define the function f(w) (Aw), Aw.
This function is Lipschitz with
1/ = | A| < A|lpp
(91)
By the standard concentration results for Lipschitz function on the sphere for instance Theorem 5.1.4 in Vershynin (0118), and noting that [[](w)] O, have the following concentration
edl2
P(]((w) 2
2exp
(92)
|||
Set t -|||A|lop
D log
Then we obtain
D Diog d cd
IP
(Aw)al 2 A||
< 22
(93)
Applying a union bound over all the coordinates and the desired result follows.
Lemma C.13. Let w ~ Unif(Sd-1 (co)) Let Hę E Raxd be a fixed matrix. Let II e gr denote the projection onto the first coordinates, that is (II*v), U, for E [r).
Then with probability at least 1 - D(d-2D) we have
rlogd for all k 0, 1...., Tı.
II' H will S co || llop
(94)
24

---

Published as a conference paper at ICLR 2026
Proof. Fix k (0.1,..., Tı}, and define the matrix A H. By Lemma C.12 for each fixed E [r], with probability at least 1 O(d-2D) we have
4D log d
(Aw)). = w)&l - H flop
(95)
ed
for some small absolute constant c. Applying a union bound over all bound holds
[r]. we get the following
4D log d ed
vs E (H S coll H Flop
(96)
Hence
((HH))),) r @ (oll 1lop |IT H w|l
4D log d
(97)
ed
and the conclusion follows with a square root calculation and a union bound over time.
To prove Lemma
C.8 we invoke the following auxiliary results.
Operator norm bound. Hellopp S V with probability at least 1 - O(d-2D)
iii)
Projection bound. For any integer 0 S l < T1, we have by Lemma C.13
log d
II' H/w|| S 6o 1
with probability at least 1 - ODd
(98)
Under those two events, we expand
Ex H. w D..um w = T1
Σ Mx(w),
(99)
where Ms(w) collects all degree-T, monomials with exactly occurrences of He. Among them we have two kinds of terms.
First, the pure H2 term. There exactly one such term: | T w|l S col Hellop
(100)
Second, the mixed terms. Each M&(w) consists of expressions of the form
T(w) = AΣT H/w, where A consists of a product of T1 1 m Ee's and k l He's. Because Σg rank-r, and leftmost D acts as projection to top-r eigenspace, we may write
(101)
T(w) AΣ II" H)
(102)
Now, we estimate each term.
A op S ||Er
<|> Ti
H
1-- lop
(103)
E lp
|Ee|or
(104)
rlog d
II' (H:w))
He l
(105)
d
and therefore
r rlog d
[[T(w)]] S en Σx T2
H flop
V
(106)
d
The number of such terms is at most (f)). and summing over < k: < T
T
rlogd
E ITT H &
60
k
(107)
z cu
(1) (() + He Hellop) )T
Ex 1Tp
Hell(
V
log
d
25

---

Published as a conference paper at ICLR 2026
Adding the pure He term back, conclude the final result with 1 - probability
r Flogd
|| + He) wil S ((EEllo + |HHellpp) To [Eellap)
+ col (l H2 llop
d
(108)
The proof is complete.
D PROOF OF THEOREM
D.1 PROOF ROADMAP OF THEORRM
General roadmap. For notational convenience, given parameters we denote the population loss and its empirical counterpart on the second-stage training set as
(( == Re ( and ((0) = 1 Σ e (fe
(109)
To prove Theorem i.c., upper bounding the population loss ((O(T)) where ( returned by Algorithm 1 we first construct an "ideal" parameter candidate O based on our feature learning analysis (Appendix C.T. and analysis standard random feature models as transition term, which lies in proper parameter class O* E F with well-controlled Then, decompose the error ((O(T))
|| O and yields small empirical error.
((O(T)) Lemma csbble issue)
((0e)
) - [(0") + _(( -
(110)
Lemmaa D.18
Lemmn D.
Next, by choosing a proper weight decay hyperparameter which relates control the complexity of F and leveraging the convexity of loss function with respect to a, we show in Lemma D.18 that with (1) second-stage training steps, the second term can be bounded appropriately. Then, given standard Rademacher complexity analysis in Lemma D.19 with proper
that (T) E F, truncation arguments to bound the generalization error (the third term) via uniform convergence arguments. Now suffices to construct the ideal O* with a controllable norm and small error.
Subtle issues in constructing O¹. An intricate problem arises when constructing O and estimat- ing (( We note that each neuron w (T) after training depends simultaneously all neurons initialization in a complicated manner. Consequently, the network after time cannot be expressed as average of conditionally independent terms, which hinders direct Monte Carlo sampling analysis.
To overcome this, we introduce the so-called decoupled parameter group ē = (a, b, W). In this construction, each row W, preserves its dependence on the corresponding initialization wn but independent of the other neurons at initialization. Then, further decompose the error (()))
as
Z(00)= Z(O
-
(111)
Leman DB
Lemana D.16
For the first term, we invoke standard Monte Carlo concentration arguments (Lemma D.14 4) to bound the squared loss, and then use the Lipschitz continuity
to extend the squared loss to the general loss [(@*) (Lemma D.15). Moreover, we show that the empirical risks of O* and O differ only by negligible amount (Lemma (.16), which is well-controlled by the Lipschitz continuity of the loss function and the small distance between [T) (T) (Appendix C.I ensuring that θ and O are sufficiently close.
Now by combining all the above analysis, we can give an upper bound for ((O (Proposition D.20p. Last, we manage to get a balance among the error terms with undetermined coefficients related to the monomial approximation property of the target function in Lemma D.21 and compute the time complexity explicitly in Lemma D.26 concluding our formal proof of Theorem
The formal version stated in Theorem where we explicitly specify all the hyperparameters.
26

---

Published as a conference paper at ICLR 2026
D.2 EVENTS OF HIGH PROBABILITIES
We will first define the following events.
Briefly, we define & and Eg events that the features ensured be well learned in the first stage, which has been proved in Appendix C.T, we then define Enj and & to control the range of the learned features and the second-stage training data for convenience of learning the outer polynomial link function in the second training stage.
First, we define the events solely related to our first dataset. Recall from Lemma C.2 that the con-
centration error He admits the high-probability operator-noma norm bound Hellop S 7 == Cvd/r where C' is absolute constant. We note that for sufficiently but moderately large (n 2 d'T2) the perturbation bound satisfies S 2T,
We will assume this regime throughout.
We define
(n iE(n]
{
s
E = Hellop S U
- | 1 S
(112)
and we know from Lemma C.2 and Lemma E. 18 that F(E.) > 1- - O(d-2D). Next, for a fixed
perturbation H&, the projection control event (Tor w drawn uniformly from the sphere sdgd-1 1) for single feature is
[H,WW)
log d d
Eg
H lop
for all l e TT
(113)
and we know from Lemma C.13 that P(Es He) 1 O(d Next, consider the events related to the sampling of our features at initialization. Note that because we use symmetric ini- tialization, only the first m/2 weight vectors independent. For fixed perturbation H; and
1 J S m/2, the projection control event E is
II (H, (0)
log d for all s T3 }
H Iolp
(114)
60
and we know from Lemma the intersection Eg-
C.13 that P(E He) 2 1 Od-22D). Next, define the desired event
Applying union bound gives
(PE H.) 2 - 0(md
(115)
From here, in this section we will write 7h r from time to time. Next, conditioning on the first dataset and the features, after the first stage of training, we define new events over the randomness of the second dataset. We need to introduce more notations first. From the previous section, we have the following formula for our final feature
w
(-a,)
Ti
II 1=0
(
(0) w
(116)
6n
and recall the ideal feature w = (-a,) Ti, Th E w define the approximate feature to remove the influence of the interaction between neurons.
For technical convenience, we next
(Tr) (-a,) Tin
Ti
(
un (0)
(117)
To help the proof, we further define
1
h (ur) =
ST.
w and h(w) = Σ u
(118)
6o
6p
and the residual r(w) = (w) - h(w) for notational convenience. These approximate features useful for our future analysis since Σ, does not depend {wn ym Therefore, the approximate
27

---

Published as a conference paper at ICLR 2026
dynamics then will decouple across neurons. We can bound the difference between the approximate features and true features afterwards with some additional efforts.
We next define the following event over the randomness of the second dataset conditioning on the first dataset for one feature
E = { (1hn(w) I [h,(ww)] for all iiisi 2m
(119)
where 4D log d. Via invoking the Gaussian tail bound plus union bound over dataset, we have Hr, w/ 2 (nd- Then we define
E { hhn (w)) (0) T (w0 ( 2 for all n + +1si 2n
(120)
where again 4D log d for each individual neuron. We have ( P( [He, (w)
)==0) )
1
(nd 2D) due to the same reasons as above. Define the desired event as the intersection Eyj := mi cl. Applying union bound over these indices j [m/2] gives P(E,) HH,, fw) ((mdd-2D)
(0) 1m 1==0) 2
Last, also need to define another event depending fær (m+)> that is, &
{ | 11 s d) NkETa {1(0:+n)s) 2 Vlogd for all S rr.. By Lemma E.18 and
Niemn
Lemma E.19 we have that (EE;) 2 Od--))
Finally, we denote Eμx := & &g Ne, N& from here. We will condition on this event throughout this section unless we particularly specify.
D.3 PROOF OF THEOREM 1
We first recall some basic facts from the results of the first stage. We have |elll S y under E which gives rise to
Eellop = E + Help S Dellop + Hellop S 1 +
(121)
and taking the product over t yields
S 1l (1 (1+19 T.
On those events obtain the crude uniform bound
T w
(1 + TT+ |m
hn (w)ll
(1 + +
(122)
E0
6g
that we might use later. Note that this estimation is uniform for all the features. We denote this upper bound by V (1+
Furthermore, we know by Corollary C.9 for w drawn uniformly from the sphere e under Ea n Eg, the residual term is dominated by
D log d TTV
|((w)) s
+ +g
(123)
where a universal constant. For notational convenience, we denote the right side as K. We note that this is also uniformly correct for all features under the full event E.
Next, we need to straightforward since need to take the expectation of w.
following moment estimation for the residual term. The below is not
Lemma D.1. Under the high probubility event Eco for any integer S j S 4p we have
(Ew [T'))))) 1/3 S K.
(124)
The next lemma computes the sample size needed to control K.
28

---

Published as a conference paper at ICLR 2026
Lemma D.2. For each absolute constant in order bound K as dr K2 I c² it just sufficient to require
2 Tid dk d d
(125)
Lemma D.3. Given n 2 Tid T + d x², under the high probability event of & n we have
[w)(w))] S logg d Vd
(126)
The proofs of the three lemmas above are provided in Appendices D.5.1 D.5.2 and D.5.3 tively. Then the following corollary holds.
respec-
Corollary D.4. Let Eg U Given
2 Tod dx d x², under the high
probability event of & n Egg En we have
log d Vd
max je 1,2nn
| (w) 20
(127)
Proof of Corollary Eg n E,, we have
D.4 From previous proof we can see under the high probability event of & n
max jE
(w) xayl S ||,, (w))ll log d.
(128)
Also by Lemma D.3 as & Eg happens, have (w))l Vlog d tt. Plug this in and we get our conclusion.
Before proving the results below, we introduce the notation used throughout. First, let S* span(e e C Rd be the subspace spanned by the first standard basis vectors, and recall II Rd - S* denote the orthogonal projector onto this subspace (i.e. II' preserves the first coordinates and removes other coordinates). Then, for a symmetric tensor T E Sym?* (onn""("" (IR)), write Mat(T) for its canonical matricization as a linear operator on Sym (R*); equivalently, Mat(T) is the unique linear operator satisfying
u @k Mat(TT) T(u vok) for all u. U Rd
(129)
The next lemma will be important for our proof. We can now show that the obtained features h. (w) are sufficiently expressive to allow us to efficiently represent any polynomial of degree restricted to the principal subspace S*.
Lemma D.5. Under the high probability event of Eas given the sample size 2 Tid 2T +
t141/Ti x2, for any k S p, we have
d
*
Mat (Ew (II' h ()) 628)))
sym
(130)
2T:
where (ym denotes the projection onto symmetric k tensors restricted to S*.
In this lemma we will use monomial basis to construct the following approximations for (T, x Sk) for any le S and any & tensor T supported on S = fee......es e Lemma D.6. Assume the sample size 2 Tid dr²T, d2+1/Tin2 Then under the high probability event E.,, for each &o S and any symmetric k tensor T supported on St, there exists 2r(wr) such that
Em [2T(w) (n,) (w) (w)]] = (T,x:8
(131)
and we have the bounds
d
Esp [2r(w)?] $
|T and |
|| 11,
(132)
A,Ti
29

---

Published as a conference paper at ICLR 2026
The proofs provided in Appendices D.5.4 and D.5.5
In the following, we show that bias reinitialization together with our learned features can be used to transform wide class of activations into (or uniformly approximate by) the monomial ø², thereby making it able to approximate low-dimensional polynomials.
Definition D.7. Let be a Rademacher RV i.e. P(a +1) P(a -1) - and let Unif[-3,3 3), independent from a.
We say the activation function a()) has the monomial approximation property with exponent for every integer k 2 0 and every positive there exists a bounded weight function uk {=1}} x
[-3,3] + with
|| 1100 S Vx.p S (k)p
(133)
such that
(0 b)o(az + - =² S Ea [uk |
sup
(134)
The uniform approximation property in Definition D.7 expected to hold for many standard activa- tions, with different growth behavior of the required weights Vk,p- For example, we provably have 3 = 0 for the activation function below.
Definition D.8. The locally quadratic activation is defined as
1 - 1,
|t 2 1,
o(1) - 12 | <1
(135)
In Secttion we show that /3 = O for locally quadratic function.
We first prove the following lemma in preparation for the proof that follows.
Lemma D.9. Under the sample-size condition Z lod dx + kR Then under the
high probability event Ep. we have for all features E [m]
(T1) w
S1 and
max 1,22m
(10) (T1))T maa $ and | (T1)
$
(136)
Note that we already constructed the approximations for (T, x Sk) using monomial basis, now we will utilize those bounds above to construct the following approximations for (T, for any k S and any h tensor T supported on S = [ei...... } using activation function a.
Lemma D.10. Under the sample-size condition n 2 Tfdlog d + d'+1/Tin also Ti = o(log d). Then under the high probability event Ea n Consider
for (x) = Eng,b,w [NT(a, b, (T
(137)
where (T')) = (-nja) (w()). Here
lu((a, b, w) = uk
(ew) (a, (w)
1 TT (l,,, (1))l II IT
72
(138)
(-n)AT,
where Ug (ex) is constructed from Definition D.7 and 2T (Tw) is the same from Corollary D.6 We
further assume c 2 123
123
Then, for any training point æ E {x,)?-nt1r it holds thats
(x) - xT D S d |T
(139)
where
A
Ck TTTT
En (2)( S E HT: log T|
(140)
Additionally we have the following moment estimations where Vk.ea is from Definition D.7
Elbig] I V 2T: log² d) |)
| S Venc (N: log" d) [
(141)
30

---

Published as a conference paper at ICLR 2026
The proofs of Lemmas D.9 and D.10 are provided in Appendices D.5.6 and D.5.7 respectively.
The previous lemmas establish the pointwise approximation guarantee for an arbitrary data point a under the event &a Eg. The next corollary will bound the empirical squared loss over the dataset
{a:
) 12
Corollary D.11. Under the sample-size condition n Z Tfd dr + d'+1/TiR also T1 ollog d). Then under the event &a n Ec, for any k S and any k tensor T supported on S (er,.... erb, we have the empirical squared loss bound
2n
Σ
(ai)
(T, - A + d D/4 |T
(142)
\=++1
where fho is defined in Lemma D.10
Proof of Corollary D.I1
Let
Jur (2,) -T xT By Lemma D.10 have, for every
training point 2;, 72 + < 2n,
led D + d ||
(143)
where
A
Ek RT;
Eu [eT(W)] S 4Ts log [||
(144)
Squaring and averaging, we get the following bound as desired.
2n
Σ (Ax + d (T))) s 24 + 2d + |
(145)
i=++1
The proof is complete.
Definition D.12. For accuracy parameters tsn define
Vp.c = max VR,cs
(146)
OskSp
where VkLer is the uniform bound (1o(()) |l from Definition D.7
Notice that we already have empirical squared loss bounds for all T supported on S* in the Corollary D.II Next we will derive empirical squared loss bounds for (x))
Lemma D.13. We assume the event En n &c, the sample-size condition n 2 Tid log dx + +1//TT and T, = o(logd) Further, let h(a, b, w) = Eesp hT, (a, w.), where {Tu is defined from tensor expansion of J introduced in Lemma E.21
(x))) Exsp (Te, x3k) and
hr, (+) is given by Lemma D.10 accuracy parameters {c& )
Then we have the following empirical squared loss bound for f (x)
2
( x
Σ
A ASP
+ d-1/1
(147)
Where fx(x)) is defined as
fn(x): = Ea,bo [b(a, b, w (T +
(148)
Additionally we have the following moment estimation
E|h21 S p.c (k T1) ())g
|| S
2p
(149)
The proof is provided Appendix D.5.8 We next introduce the following lemma using Monte-
Carlo Approximation.
31

---

Published as a conference paper at ICLR 2026
Lemma D.14. Under the sample-size condition 2 Tidlog dR d x² also T₁
ollog d). Consider aj == 1 h(a,, bj. w (0)) where lu()) is defined Lemma D.13 and 0g bj. w (0) initialized in Algorithm / Denote θ (a", ",,bb, b(T1), W(T.)).
is
Then under the event Ep. we have
2 Ess kSP
2n Σ
V2 (KT) (log d) d 4/1 +
(150)
7
m
1=++1
and ||a*11 m V2.
(K
T2) (log cl) d))tp.
The proof in Appendix D.3.9 For notational simplicity, given the parameter group @ = (a, b, W), we denote the population loss and its empirical counterpart on the second-stage training set as
n
(()) == (f)
and Z(0)
= n
((e(a))
(151)
respectively. Also write ((0) for Ö(fe).
In this lemma, we bound our general loss using the squared loss bound. Lemma D.15. Under the sample-size condition n. Z Tfdlog dx2T, 2T, + d2+1/TTx also T₁ = o(log d) and under the event Ey. have
(n (Aogg ddap+2
2(0') LL
Σ
A +
(152)
m
kSP
where A₂ is defined from Lemma D.10
The proof is presented in Appendix D.5.10
Next, we consider the gap between the real features we learn W(T)) and the approximate features (T1), Recall θ+ = (a", (((T), W(T)) and @' = (a", ,((T)), W(T.)).
Lemma D.16. Under the sample-size condition n Z Tfdlog 2T, + d2+1/TTs also T₁ = o(log d) and under the event Ey, we have
(I (log dyw
- SL
(153)
m
We provide the proof in Appendix D.5.11
Combining the above two lemmas, we directly get the bound for C(fe.). Lemma D.17. Under the same condition as the last lemma have
V2 (1TT. T A (log d'
L(fe.) SL
Σ ASP
(154)
The proof this bound is provided in D.5.12
The preceding lemmas construct a parameter e* with small empirical risk Z(fer). We now analyze the parameter output by the second-stage procedure,
(T) (T (T)) W(T))),
T Ti T2
(155)
where (T) denotes the iterate after T iterations. We recall that only a is updated in the second stage while b and W remain at their values time TT. The next lemma gives upper bound on the empirical risk (1)) of this output.
For notation convenience, we define = L (Ensp Do +
(n Fi (4p d
and
m
(N (log djsp
32

---

Published as a conference paper at ICLR 2026
Lemma D.18. We consider doing gradient descent the ridge-regularized empirical loss F(a)= = [((a, (T)) w(T)))
32
(156)
with the hyperparameters B2 = thr and 1/2
62
LymU+42J
Then for any Te 2 L' log( (m) also under the event En, our second stage algorithm returns a such that for ( = (T( ((T)), W(T)))
A (T) S
and ||a x U.
(157)
The proof is provided in Appendix D.5.13 loss and population loss for (T). Lemma D.19. Under the event E, where & is defined in Defination D.37 and the sample-size condi- tion n 2 Tod d + d d also T1 o(log d) we have
The following lemma bounds the gap between empirical
V (K 4 (log d)
[(Som) ((Jem)) )
LV
(158)
n
The proof is provided in Appendix D.5.14 Then the following proposition combines the previous lemmas and gives a quick proof map of bounding the excess population risk ((fer)) Proposition D.20. On the event & and the cample-size condition Z Tid d 2T: + d also T = o(log d) the excess population risk satisfies
V2. VA (log dd
pie W2TT
(log d) 4p+2
((fem))<((") LL
ΣΦ + SSP
(159)
n
m
Proof of Proposition D.20 Note that to upper bound the population loss (()(T)) of the parameter group (T) returned by Algorithm|] we first construct an "ideal" parameter group O which lies in proper parameter class @* F (small
|[+||] and yields small empirical error.
So the first several lemmas this section give an upper bound for fee-) through θ+ which is introduced to decouple the dependence
ffe-
fe A + Afa.) - As-)
D.I6
D.I5
V2 (I: (n d))
V2 (n: (lpg d) 4p22
$ L
+L
Σ A
(160)
m
m
ASP
( 4
< L
Σ Ax
per
m
ASP
And recall that we denote : Ensp A V v2 ( (log d) 1p42
m
Also note that we can decompose the error L(O in the equation that follows
((fem) - _(fer) (C(omm - [(S@m)) + -(fem) )]
Lemma D.19
Lemma
2 (n ( (log dl 4p+1
$ LV
V2 kApT, (log dyy++1
V2 MA d) 4p42
$
Σ AB + V
+
m
6Sp
(161)
33

---

Published as a conference paper at ICLR 2026
Combining this two yields
V2
pie (K T: (4pT:))4p(()))) (log d) dp
V2 (1(: ) (log d'
-
Σ RSP
A +
(162)
which is precisely the claimed bound. This completes the proof.
Back to the main proof, in the previous proposition, we have the generalization bound
V2. (1 ( (log 40
V2
(..T. (log d) 4p+2
[(eem) L
Σ As + KSP
+
(163)
m
We note that the bound is depending on both {Ax) and Vyees which are left unspecified. Fixing O, notice that A has the bound relying c which is
4
Ck TTT
Ew
(2,, (w)]
(164)
S RT
log" d||Talrr S C log" d.
On the other hand, we observe that Vk.c. S S.. Thus we need to choose proper {<=] to keep the balance of ViLez and Ag. For notation convenience, let
C. = V.
C. = V
= n. TT (log d) 20+1
(165)
m
Thus the e-dependent contribution is upper bounded (up to absolute constants) by
o((c)us)) A(
< ck (Cm + Vn.)) kSP
(166)
we know that the approximation precision p can be zero, which indicates that = 0 for all [< and also Vp,e S 1. Thus
Specifically, when 3 0, from Definition D.7
S A((Cm + V..)) A(Cm + Cr))
(167)
Furthermore, for B > 0, plugging in Ve,r S cg -8
we have
ক ((en)Res) = A( Σ <x + (Cm + Cn)
max (a-8) kSp &
(168)
The following lemma gives the minimum value of Φ.
Lemma D.21. Fix > 0. Consider as a function of {cK Hezo attained at the symmetric point
Then the minimizer is unique and
6 = = [8(Cm +Cn)
Cn)
8-1
(169)
p+1
Moreover, it holds that
la (p 1) (Cm + Cn) S A(Cm + Ch
min + A (c>>0)
(170)
The proof is provided in Appendix D.5.15 The following corollary is direct from the previous statement.
Corollary D.22. When 3 = 0, we have inf & = A(C.. Cn).
Now we state the main theorem. For the theorem stated below, we recall that we assume T o(log d), dog 2T, dl+1/Ti,, in addition the assumptions in Section and Section
34

---

Published as a conference paper at ICLR 2026
Theorem (Main Theorem).
Under the assumptions above, there exists set of hyperparameters
d
1 m
U
m C. 4D2 log d such that for any T2 2 Lm log(Um) the population risk satisfies
B1
1/2
LmmL² 2J
B2
(171)
(V=+VV
3+1
Em (F((( S k 2 ( d))xt1
(172)
under the high probability event which happens with probability at least 1 D(2-D/2), Cn is an absolute constant from Lemma C.6 and J, are quantities defined in D.3
where
further estimate sample and time complexity in the next lemmas, showing that they are both nearly optimal.
Lemma D.23. Under the setting Tı = IV
log log
have
Tfd TT d2+1/TT x² z k²d
2
Inne Tng
(log d)²
(173)
The proof is provided Apendix 1.5.16 Note that from Theorem for n Z T9dlog ds2T,, + g++1/Ti,² x² the main result applies. Since our sample-complexity upper bound coincides with that threshold, the corollary follows immediately.
Corollary D.24 (Sample Complexity). The sample complexity is @(x2d' subpolynoiil
where d Jog
Next first estimate T2 under particular m, choices and other hyperparameters. Lemma D.25, Fix the set of hyperparameters
logd log
J
U
L
T1
(V
B2
U
1/2
mm/ 2J
T2 J
log(Um)
(174)
where J and U is defined as
V2 (x:T1 ( (log 4p+2
V2 (KT1) (logg d)
J:=L E
and U
(175)
m
m
kSP
{ hesp is defined in Lemma D.10 and V is defined in Defination fixed in Lemmal D.21 conclusively and are also fixed.
D.12 Note that {Ex He was
loR (log d 4p(++1) 1)+ Ap(3+1)
Furthermore, fix the sample size and net width
log
Cn2(log 2)2 222 where is a large universal constant. Then under the above setting.
lean
T2 S 3 Plog 8p8-+18+2222 108 load
(176)
By ignoring the logarithmic terms, get T2 (1p8018)))1133)))
The proof is provided in Appendix D.5.17
We note that the total time complexity is upper bound by nmd- T up to absolute constants, so plugging and the next lemma follows directly.
Lemma D.26 (Time complexity). Under the same condition as the last lemma, we have T Tcomplexity S x.5++ d) 12p0+179+5 (100))700005570120001111117111
log
The proof is provided in Appendix D.5.18
35

---

Published as a conference paper at ICLR 2026
D.4 PROOF UNDER B = 0 AND ASSUMPTION 3
We now extend our analysis to a more general loss class (Assumption 31 under an additional as- sumption / 0. For clarity, we provide complete proof here.
Under Assumption 3 and 3 = 0, Lemma D.5.1 to Lemma D.14 still hold. In particular, given that = 0 we can conclude that 4 0 for all S p and Vp, bounded by universal constant. We
restate this as the following lemma, which is parallel to Lemma D.14 the previous section.
Lemma D.27. Under the sample-size condition n 2 Tod dx + d+1/TT x² also Tí ollog d). Consider a3 == 1 Lh(a,, bj, un (0)) where is(-) defined Lemma D. and 0g by, w) (0)
initialized as in Algorithm| / Denote θ- = (a", (D(T)), W(T))).
Then under the event Ep. we have
2n
(10) (w:)) <d-D/4 (Ni (P d) 4022
(177)
m
and ||a"||
< (K ( dlo
We sequentially present the parallel versions of the other theorems in the previous section. In the fol- lowing lemma, we bound our general loss using the squared loss bound. For notation convenience, we denote Laug := Ly (log 2r, It can be observed that the original conclusion remains almost unchanged, except that L replaced by the new Laug.
Lemma D.28. Under n Z Tfd log dx + d2+1/TT x², also T2 o(log d) and under the event Eyp, we have
A (log d) 4p-2
((07) S Laue
(178)
m
Proof. Under Assumption 3' for any 2, 2 E R,
(1(z,) a(=', v) s L(1 + +1|++11++ [x)) =//
(179)
Combining with the above Lemma, we have
2(0²)
Σ
n
8 (z.,
_
n Σ
n
n
Σ
L (1
|6 (x: + (a:))) ( - (xx))
(180)
Furthermore, by the same argument as Lemma D.9 we have (ww
S for all j. thus we
have
of W (T)) + b) i vm
(181)
1 (ax:) = a' w (T)) + b) $ U Um
(182)
Also note that S T1 (log d) and 1f (x:)) $ (log d)p Thus by Jensen's inequality and the Lipschitz property of l, we have
m
( < Σ L(1 11 (2x)) (m:))+15(<0 | (2x.) l (xx))
(183)
n
RAPTI (log d) 40+2
< L(** T 1 (log c) 20
Σ (xx) f² (o) Max)) $ Lauz
(184)
n
m
The proof is complete.
36

---

Published as a conference paper at ICLR 2026
Next, consider the gap between the real features learn W(T)) and the approximate features W(T),, Recall O = (a", b(T:), W(T)) and O' (a", b(TT), W(T))). Lemma D.29. Under the sample-size condition n 2 lid dx + d x², also T₁ ollog d) and under the event Ey. have
(N: A dytr
(10))-C(03 2 Laug
(185)
m
We omit the proof here, since the only change is to change L to Laug Combining the above two lemmas, we directly get the bound for [(
Lemma
D.30. Under the same condition as Lemma D.28 we have [(fe-)
24
T1 (log d m
Laug
We now analyze the parameter output by the second-stage procedure, θ (T) = (a (T) ((Ti), ((
T +T2
(186)
in the new setting. Define Laug V similar to the above section, but simpler under the new setting.
(x (logg d and U
(A Ti (log dydp
This is
Lemma D.31. consider doing gradient descent on the ridge-regularized empirical loss
82
F(a): = -[((a, W(T,)))
|a
(187)
with the hyperparameters B2 fiz and 1/2
1
LmU"++
Then for any T2 ~ log log(mn² log d) also under the event Eno our second stage algorithm returns (T)) such that for ( (a(1), ((Tx), w(x)))
()( S J' and |a S U".
(188)
Proof. Consider the training objective F(a) = 2 ((a, b(Ti), w(T:))) + Ballaa]]² with 82 T
Let
(00) = arg min F(a) ac
(189)
be the minimizer of F.
We further denote Eo is the prescribed optimization tolerance (a(T)a((oo) s Eg. Furthermore, let Tus
By optimality of a (c0) and the constructed comparator a" we have
32 _ ||a1|1 F(a( =(a')) [(a")
(190)
Here only use that a* is feasible point with ||a*|| U" and E(a²) S J' (Lemma D.30 On the other hand,
F(a Z(a(00)) + B2 l
3/
82
(00
(191)
2
combining with 331 gives
2J
2J'
(00) < 1|a + A(a") $ 122
= B2
12 +
Ju S U22
(192)
B2
Thus
|||00|||| S..
(193)
37

---

Published as a conference paper at ICLR 2026
and also
|a || S
(|(
[|a!T)
a - c +U' SU
(194)
Moreover, from (331 and 334
82
13 2 J' 82 1 z J'
[(a = F(aa)))
(195)
2
Thus we have
IF(a F(a F(a
(Laug a (T) dW(()) x, + ((T))
(00)
}
sup
*E 1,2n]
(2(aaa + 1)(am -
2 aug |a -a (x))
sup
+
iefn+1,2m
+ (2((la(T)) ||a (00) - (
(196)
By previous proof ( + b ((T) loe be bounded by absolute constant for all under
event E. Thus can have loose e norm bound and
(F(a(T) - F(a (()() z Laug (a -a |l
||o b (T))
sup
sE(n+1,20)
(2(aa (100 (T) a (00)))
28 Laug aa (T) a (00) l (2U" a (00)
2 (Laug V m + B²U )0 J
(197)
Combining with the same argument as above,
((a S J.
(198)
By Assumption 3' on the loss function f, the function F(a) Eq ||||| 2 convex and satisfies the gradient Lipschitz condition
By the similar arguments as in the proof of Lemma D.18 the lip constant for F satisfy Lup z L. + /2
Therefore, by Lemma E.13 if the step size 172 S
gradient descent will approximate a (0o) to
Lap
arbitrary accuracy in
log [Ia(Ti) ( -alo))a(o00100 -122)112/6))
T2 T232
(199)
steps. Note that when the step size 7/2 Lm+J'//uz
1
(200)
1232
Using the assumption that S d and ||
(20) || < U, [T))]
S can further simplify:
1 |la (T:) - (00)
(Ia((T))
log
Iog(U'mLaug)) x logd
m
log(mu log d)
(201)
32
Thus, we conclude that gradient descent will approximate a (x) to arbitrary accuracy in at most
Te 2
log log(mn" log d)
(202)
steps. The proof is complete.
38

---

Published as a conference paper at ICLR 2026
The following lemma bounds the gap between empirical loss and population loss.
Lemma D.32. Under the event E, where & is defined in Defination D.37 and the sample-size condi- also Ti = o(log d) we have
tion n 2 Tidlog d + d
(n Am (log d) 40
[(fom) -(fem)) < Laug
(203)
n
Proof. Use the same argument as the proof of Lemma D.19 truncation, we can get
Appendix D.5.14 and use the same
log d
sup
- CM()) Laug U
(204)
SEF
n2
We notice that under the event ε M holds for all æ,, i [n + 1, 2nl, thus (M( = Z(f).
For the other term, we have
(M C(f) -
MB]
|100--+|[[(]++|1/|]]]00001100111 & Mi]
[ MP + Mi))
<L(VE[f(a)PP(= MD} + (=) & M)])
(205)
Recall that it has been assumed that the activation function has bounded third derivatives as in Assumption] Therefore, there exists an absolute constant C, > 0 such that
[o(2) -
(u(*) for all u E R.
(206)
Via a crude moment estimation bound (plus some hypercontractivity arguments), we can see the right side is S L Thus, we have the following final estimation
m lolog
sup B( - -() Trug U' FEF
+ D/8
(207)
Since we have assume n. S d/44, we have sup [()) L()) 2 Laug U' Plug in the formula for U and get our desired result.
log +LU' -- +
We then have the final proposition.
Proposition D.33, On the event E and the sample-siz size condition 2 Todlogdn" 2T d also T₁ = o(log d) the excess population risk satisfies
((Jem) - L(5*) S Traugh 2pT: (log d) 2p+1
V
(208)
assumption is not essential and can be removed via careful analysis. To be concrete, one can keep the Aş and Vp.e terms in final generalization bound under the weakened assumption 3") Next, we perform another round of balancing those terms. The minimizer shall be slightly different and bit technical troublesome, but the final result remains the same up to logarithmic factors.
Roadmap remove 3 0 assumption. We note that
D.5 PROOF OF SUPPORTING LEMMAS IN APPENDIX D.3
D.5.1 PROOF LEMMA D.
Proof. We split the expectation over w into the good event Eg and its complement
r r = E
E +
P1
(209)
39

---

Published as a conference paper at ICLR 2026
We first deal with the first term.
By the high-proode probbbbbbbbiitty bound, we have [[rr(w) S K for all w, so
E 1
S K) - K'.
(210)
Next for the second one, we only need invoke the trivial operator norm bound. By definition we have - - (1++)) +1+ Thus
|r(w)]
S 1 EdlT = 2M
(211)
Hence
E
II rl 1
< (2M) $ (2M)rTTid
(212)
-Ea
Therefore we have
1/5
2M j
E II' rlp
SK' (2M)rTidd-
and [B 1/1 K
1+
rTid D
K
(213)
To bound the right side, further note
2M
2(1+ + 7)
2(1+ + 7)
Y 1/(2T)
(214)
K
/r log d/d
Tin Dr log d/d
Since 71 < e Tin S Ve, we get
2M
Boog TiCh V rogog d X Vn
(215)
K
Tn
Thus via choosing large enough universal constant D, have for all 4p
2M
rT, < 2d 1
(216)
K
and we finally obtain the desired conclusion via combining two terms.
D.5.2 PROOF OF LEMMA D.2
Proof. We want to ensure equivalent to
K X-
4 < 2. Taking square roots on both sides, we note that this is
Vd X.. KSc and K
AT
(217)
Vd
To get a sufficient condition, we bound the two terms in K separately. Recall that we set
rDlog d TT
K K Trry
(218)
therefore we need
rD Dlog d
XTi 2/d
XTi 2vd
CTrMV
and CyT. <
(219)
d
AT
Under these, we get K <
as desired. Note that
vd
d and CCHTIV
r log d
r rD log d
y= C
CH
< T
(220)
n which is equivalent to
Then, to satisfy the first one, we require T1
rDlogd
Dir
Thr Dd Dol 2T.
(221)
For the second one, we require
Ti
1/T;
d (Vd)usT
which is equivalent to
,
(222)
X
Combining the two parts above, the final conclusion directly follows.
n
40

---

Published as a conference paper at ICLR 2026
D.5.3 PROOF OF LEMMA D.3
A
Proof. From Lemma D.2 we know that K X
Then on the intersection of event & n &g the
residual satisfies
||(w)]] S K
x
Vlog d/d.
(223)
vd
Moreover he, (w) h(w) + +(w) and h(w) E Under Eg. we have ||h(w)]] v llg d/d. Therefore, under the event & n Eg we have
[h,(w)) - [k(w)]] + ||[(w)]] & Vlogd/d.
(224)
D.5.4 PROOF LEMMA D.3
Proof. Note that since every vector in span (at ((I' (w)))22))) a vectorized symmetric k tensor, it suffices to show that
d
-
II n (w)) @2k
2 (T.T)
(225)
x 2T,
for all symmetric tensors T supported on with |T = 1.
For notation convenience, we denote w, := w/cn where we know w. Unif(S(-1). Recalling that II n, (w) = Σ w. + r["r(w) and the tensor binomial expansion, we have
E we @(k-i)
II hn (w)) 88
X ((w) 84
(226)
1=0
Thus for any symmetric k-tensor T, we have
(T. (II' (w)) 80M) = (T, T. (D. )anr
88
+ 6(w)
(227)
where
k (T, 1TT w.
@(k-1)
8(w) =
(IIr)(W)) 81)
(228)
Note that by Young's inequality,
E (T. (II* h (w))
27
> E
-
1
T.
w.T' w.
)ar) - E [d(
@k 27
(229)
2
Therefore we need to upper bound (l(((u)) and lower bound EE (T) [E) un, @kk2). Intuitively the 5(w) term should be negligible since r(w) is of small norm.
Note that there crude upper bound for 6(w). Contracting T in its first k slots and applying the Frobenius
inequality gives
(T.
x w.
@(k-a)
(I')(()))* 88)
$ T
((77 w.
olk-1)
II r(w)
(230)
Hence
b
E w.:w w. )
1(w))
i T
[['r(")]1'
(x we
) [['r(w)]]'
T
(231)
41

---

Published as a conference paper at ICLR 2026
Via Cauchy- Sechar-- again, have
Blo(w) S
E ElIT(( w.
w. [r)
2 BIIT(E w. (0.)) [I
(232)
where the second line is derived via the Gaussian hypercontractivity lemma that we will state at Lemma E.22 in Section [.6
Next, let T be the symmetric & tensor defined by T(vi.... Uz) T(E)
01
2 E
vx). For
each i and each unit-Frobenius nomm tensor V supported on S we have
(T w. ()). V) = i(w (0 (t)
(233)
i.e. by moving Σ/ from the arguments of T onto the test tensor one applies (E. each slot of V. By Cauchy- Cauchy-Schwarz again,
T [ (w/co) (0(M-1) ((--1))
v S iT 00k MF](
(234)
Since bound
S 1/2 for VIIF 1, taking the supremum over unit V gives the pointwise
T w ((--1)
|e S 1/1 [Tw, )(k-0)]|
(235)
Squaring and taking expectation over u yields
E T w. (.)k
(w. (k-- 113
(236)
2
Combining with previous calculation, we have the estimation below
E EITT((E) w. 0(K S
El
T(w. s(k-1)
X 22T.
ET E w. w..k 20
(237)
where the second inequality is deduced from the lemma below. Lemma D.34. Let u ~ Uni and each tensor T and each k: S p we have
* I| z d d=-*E(T w 22
(238)
The proof Lemma D.34 is provided in Appendix
We also know from Lemma D.I that for any integer $ < 4p we have under the high probability (d- that
event of &,, which happens with probability least
E II
r
r
1/1 $ K
(239)
Therefore, have
B[ S E(T, w.) @k &
(240)
2T;
1=1
Because we have assumed Lemma that
Do 2T 2T: +
we have the following from d'+1/Tin? x2,
K2 c²
(241)
2
with some prechosen absolute constant into the previous equation we have
Thus via picking small c and plugging the above bound
E. [o P1 ETT ( w. )0M
(242)
42

---

Published as a conference paper at ICLR 2026
Combining everything gives
E (II len 2 (T, (2) T w.) ) - Ew (6(w)
B[T w. )
(243)
d
(244)
where the third line is via taking the expectation over the detailed proof for this can be also found in Appendix E,3
The proof is finished.
D.5.5
PROOF OF LEMMA D.6
Proof. Let
2r (w) = MattTT (ee))) (w)
(245)
Note that Vec(T) E span (ataa (w) )2b1) by Lemma D.5 Therefore,
(wnlz((w)(hmm (w) x²)*1 Eu (eo Mat (w) 02k Vec(hnn(u @k (ec(h(h)( (k) V
= (T, x @k).
(246)
For the bounds on zT(w) we have the following estimation
En (w)) = E. (Vc((T) Mat (w) 'eee)))/VV((( Vec(h,, (w)) (t) (Mahh, 62k1) (2)) Vec(T))
VVe((T) (a,,,,(,, (2))
( Voc(T)
S
|
(247)
2T
and
Vec(h,, (w) d
|zr(W)] Veo(T) (B/ (w))))))
||T|| | hn (w)]l
2T
(248)
The proof is complete.
D.5.6 PROOF OF LEMMA D.9
Proof. Recall that from Lemma D.3 and D.4 under the high probability event of Ea n Eg O Egy which happens with probability at least
Od P) the following bounds hold
log" and [w)] < log
max (w) jE[n+1122g]
(249)
d
Those bounds will be helpful for us to construct the low-dimensional approximation the second stage.
1/(2T))
where the universal constant is chosen sufficiently large we can absorb the implicit constants hidden in the S bounds. This means choose C' to make sure on the same high-probabilie event the following controls are valid
Recall that we have set = m = 1/ log'2
T [
and max (
Im) To ten (w)Tx/
xz
(250)
which indicates that
(T))
1 (-a;)
Ti,
Tr, w
(251)
43

---

Published as a conference paper at ICLR 2026
Note that w (T())) J c0
(II:20 T-a
(0) w and by Lemma C.I
we have
log d
|| T K 1
(252)
d
Thus (/w s 2 $ and we complete our proof.
D.5.7 PROOF OF LEMMA D.10
Proof. For any training point and given initialization w. on events & N&C N ne, we have the following control
= (-n) T. h. (w) E [-1,1].
(253)
By our previous assumption on the activation function, have the following low dimensional approximation result uniformly for z E [-1,1]
Ba [o) (cx)) (o (a, b)o(az + b)] -0]- 2² S Ek.
(254)
Divide both sides by (-1))aT TT, and substitute the features back we obtain
(-n))<1 KT, D k (a, b)o((-aa))) n, (w) x + 6)] (ww xx)* S ck
(255)
Recall that from Corollary D.6 under the high probability event & we have
E. pr w) (w, 2) (T, x*).
(256)
And also recall from Lemma D and the statement before the lemma know that under the high probability event Ea N n &g Eg
In To he (w)) [x S
(257)
17 Ti hn 1
(258)
thus
(x)
(T, zk)
6xc
Ew
(wl2z(()) (w)) Eunene 2i,1 T
Bw (xT(W)] (
KT
1
Ew(2r(w)))
(ex) (a, ()o((a)))" (w) + b)1 en ennnennnn
(259)
<
6
+ hT (w (.ne
+ 17
KTi En (w) (a.b)o((-na)) h (w) b) (.
We bound the third term first. Denote
X(a, w) = TT, ok (rk) (a, (((n()) h. (w)
+b.
(260)
By Cauchy-Schwar
Ea, [(a, b, w) (. S Be (X,, w))l P(E NE (Es )))))) (261)
We bound the two factors on the right. Recall that has been assumed that the activation function has bounded third derivatives as in Assumption 2 Therefore, there exists constant C, 0 such that
[a(u) Ca (w(3) for all R.
(262)
44

---

Published as a conference paper at ICLR 2026
and also
((, w)1 I (eeeler(()))) (1 In h (w) x 8/8).
(263)
We recall that the following crude bound works for all w
[w)n(w)]| (1 + 7 2
(264)
where y is defined in via directly bounding the operator norm of I1 M Notice that under the event &g. we have ||a||| - 2vd.
Thus via another crude bound we have In) hn (w))(()) x + bJ S and
[X] S Vk,co |=r(w)ld?
(265)
In addition, from Definition D.7 we also have Vi,ce ||0 | 11 S x 2 d/122. Also invoking
[= S x [T from
D.6 note that for Ti = o(logd) d), we have T & d, thus it shows
that we obtain
E[X b, w) S Vhne d V El S dD/12+3//2+3 TT
(266)
Second, Lemma
D.3 and D.4 give
P N (Eg (Eg ( S TP((E) En) z d/22
(267)
when 1 2D log Hence
Ea [a b,
/ 2 d (
(268)
for sufficiently large universal constant D.
The above argument shows that the third term is negligible.
Then we bound the second term. Using essentially the same method and argument as above, we obtain
Eno (()((()))))))))) 1 Eo
2
T V Bul
227
P(E, N& n (Eg ))
TT
I 3h/2 T pd 13/4
x d 1
(269)
For the first one, we define the approximation error
A
TT Eu (2)(
(270)
By Cauchy-Schwarz,
EN:
Em 2T (2)
(271)
KT
and Substitute the second-moment moment bound mentioned in Corollary D.6 and we obtain
d
Ew (r
|
(272)
thus
1
d
4/2
Au S Ex
IT |F S A log" d|[]|F F.
(273)
KT,
2Tg
For the other moments, from Definition
D.7
we have 11
(
1lo0
S Vkiee and
Viee
1
B <
E[zT ([zT(w)²1
Viek 2T
1
2Tg
|
(274)
Substitute r) back and we can get
E[ha] & V22 (NT log d) 1|
(275)
Similarly, using [2r(w)) & the definition of hrr (a, b, w), we have
1 |T (h (w)) again and noticing the truncation terms in
|| S Vanes TT log? |
(276)
and the lemma is proved.
45

---

Published as a conference paper at ICLR 2026
D.5.8 PROOF OF LEMMAD.L.
Proof. Since h(a,b, b, w) Ensp Nu,, (a, w), combining with linearity of expectation, have
fx(())
Σ
furr (a)
(277)
RSP
For each k. define the empirical error vector O R by
O (i == fun, (z:) (Tk 2; 2,
i= 1,...,21 2n.
(278)
Then for the aggregated function we have, for all æ; from the second stage
fx:1))) - (x,) Eo,(t,
(279)
kSp
For now throughout this proof we equip R" with the empirical & norm, defined for any vector (v(1). v(n)) E R" by
Σ o(1)?
== n i=1
(280)
Hence write the empirical squared loss by
Σ (fu(2a)) - (xx)))
Σ Ok
(281)
72
ASP
Note that we have
2
Σ
Ok
S
Σ
1lollIn
(282)
ASP
kSP
By Corollary D.11 each single-term error satisfies | 1
< ~ A + d d/
Thus
||o In E A
+
d D/8 [
(283)
Notice that by Lemma E.21 we have ||T S f(0-k)/2 k3/4 x 1. Therefore
2
2
Σ 11ol11
2
Σ D-8
|T&|lr
ASP
ASp
PSP
V2
Σ Ae + D/D
KSp
2
+ (a-D)
k
ASp
Ax
+ dd-D/4
(284)
S
where the last step absorbs r,p into the $ constant. This proves equation 147
For the moment bounds, we have by Schwarz
27
E E E
Σ ln,
ElbT
E
(285)
K.ESP
ASP
Substituting the second-moment bound mentioned Lemma D.10 we obtain for each k.
E
V2 (T. log" d)* |
(286)
46

---

Published as a conference paper at ICLR 2026
hence
S (p.c( log 1 Talls.
(287)
Using the bound |T
p=k
and summing over &s - p, we obtain
Σ
E
( d) x Viee Tx (
(288)
kSp
where the last step absorbs the polynomial degree into the X constant. Further, squaring both sides yields
E|h²| - Vi. (log 3p
(289)
which is the first inequality of Equation 293 For the sup-norm, substitute the sup-norm bound mentioned in Lemma D.10 we know, for each
|/rrylls & Vr.ca (k2Th log" Txlle.
(290)
Combining the bounds together we have
|11l
Σ
|
1o S
EVec Vi,ca (N2 (log (1ogd)3)$( S ( (log d)
(291)
RSP
kSp
This proves the second one of Equation 293
concluding our proof.
D.5.9 PROOF OF LEMMA D.14
Proof. We first bound ||a") where (aī,..... a. Since
1
aj h(aj, bj. wzz
(292)
and via Lemma D.13 we have the following moment estimation
||/|loo & Vp.e T1 (2 (log dp
(293)
Thus directly get as desired
V2 V2. (1 dy ||a" 12
(294)
m
To bound the square loss for finite width, take the basic decomposition
(f((z;) - f*,)))) <2 - Σ (fz - fx:)))) +2- Σ (f((z)) 2
i=+11
(295)
We next bound these two terms separately.
For the second term, by Lemma D.13 we have
2
2n Σ
2
Σ As ASP
+ d D/4 D/4
(296)
n
It thus suffices to control the Monte Carlo fluctuation term (the first term). Note that the parameters are initialized symmetrically and this symmetry is preserved by first stage training, that is, the first half neurons are free. Therefore we should group neurons into symmetric pairs to help the proof.
For j S m/2, define
Zz(x) = a= (Ti))T (b;)bb))aaaa ( (T)))T x +bm-)) w (Ti) = (-na,) fen (w.
(297)
so that fz. - fn(x)) = Ejsm/2 (Z (z: E|Z, (x)) Introduce the truncated version
Zz(z) = Zz (x)1
| (T (T
(298)
47

---

Published as a conference paper at ICLR 2026
By the concentration event Ens which occurs with probability at least 1 O(d O we have 7 (2,) Z (2:) for all 1 S i I 2n and all j m/2. Hence, E, have the follow- ing decomposition
- fn(x))) -
(Z)( Σ (Z)(z,) ((2,)))EZ((33)) + [EZ (2,)) 35m/2
(299)
We bound the second term B, first.
Recall
m B: = 25:2(0 (1E, (2)) = m - (aa
(300)
because Z, Z3 on E. By Cauchy-Schware
[ -
E (a:: ( + "E "
2
$ 2
EZ, (x,)² + EZ, EZ,
V (Eg
<m (Z,((;;)) (E N n (Eg (E,))))))
(301)
The final inequality follows from the fact that truncation does not increase magnitude, that is, [my( $ (3,) For symmetric pair
(T:)T x; bb) + =0 (W0m)) (T))T I, + bm-))
Zz (,,
(302)
Recall that it has been assumed that the activation function has bounded third derivatives in Assumption 2 Therefore, there exists an absolute constant C. > 0 such that
[o(u)] < C. (1 + |ul3) for all R.
(303)
Then, for any fixed i.j. (())(a)))) (1 h. ()) x; tby by) (mm-)) hn (w J + bm 133
(304)
Notice that under the event Eç,, we have æ; is bounded for all i [n + 1,2n] by || S 2 Thus via another crude bound, since we have ( ( x, byl S d and 1y h. (w (0) ) xz +
5m- S d and we have
V [Z,) S d [(03)2+ (a (a
S 2 (I (log d)Ap
(305)
m
In addition, from Definition D.7 we also have Vkkes = 1| ((((e)) Iloo S 6 S d
thus
V2 I d n.e
Also note that for T, = o(logd). we have T2 $ d, thus it shows that
4+ P
V EZ, (2,(x;)² $
(306)
This gives the intermediate bound
d d++1p+
B
S m
(((E) S d++40+B f /P((EE
(307)
m
By the concentration estimates proved above we have Pr((E, E,) S d- (union bound over the m/2 independent draws yields the extra factor m), hence
B. d'+1p+ Pa d D/1 < D/18
(308)
48

---

Published as a conference paper at ICLR 2026
for sufficiently large universal constant D.
Fix any æ E {a: The variables {Z)(z) - EZ, (x) 1smm are i.idd and centered, so we can use Berstein bound the sum. Notice that (w) ) Jx and w m-j truncated by the defination of (z). we have ((ww)T)) + by) and ((w (T1) To bm-z) also bounded by a universal constant.
(Ti)T
Combining with az m hhaaaa by. wj), we have for a universal constant C
(z - (x)) - |
(309)
Lemma D.35 (Hoeffding's inequality). Let X1, X2........ ... Xn be i.dd. random variables with mean Ia and X2 [a, b]. Then,
(-1o- o]}
2nt?
X; p 2t s exx
(310)
11
Due to the fact that [2,(2)) x /l/oo, Hoeffding inequality then yields with probability at least
1 - Od D)
log d |1 S
Σ EZz(
(311)
Vm
55m/2
Recall from Lemma D.13 have |11|o0 S Vn (1 (log 2r Plugging these bounds into the above one, it shows that
Vi jap (log d)2p+1
Σ (Z)(z) 1EZ)(a+)) 2
(312)
Vm
35m/2
Therefore
(z:) fn((x;) S
Σ (xx;) EZ, (xx1)
+ 111 EZ, (z.) EZ,((;))
2
j iSm/2
V2 (NCT' (4p df) 4p+2
+ d D/8
(313)
m
Squaring both sides and applying a union bound over the n training points, we obtain the following.
1
Σ
(J.) )xx)))) x
12
n.e
(n: TerTiee (oooo 4p-2
d DD4
(314)
m
By the initial decomposition and Lemma D.13
2
2n
V2 (x d) 4pt2
Σ
(z.) <
Σ
A
+ d 2 +
(315)
m
1=++1
ASP
which is 177
This completes the proof.
49

---

Published as a conference paper at ICLR 2026
D.5.10 PROOF OF LEMMA D.15
Proof. From Lemma D.14 and Lemma D.13 we know that with high probability least
Odd D/2) there exists a" E I such that for θ = (a", b(Ti), W(T.)),
2
2n Σ i=n+
V2 (1.T. (.TT))d(( (log d) 1)4+2 + dd-D/ D/4
- fxxx))
x
Σ
m
ASP
2
V2
Die (16 To ) (log d)
A
(316)
kSp
The last inequality is because m I dD/4, so d D/4 2 1 thus d D/4 can be absorbed by ve, (xc T (log d
m
By Jensen's inequality and the Lipschitz property of l, we have
((f)) - -(() = 1
Σ
( (2. (EL)))))
7
S
( (x) - - (xctt))
11
(317)
S
( (2it+n) / (x)++)))
V2
P.e (NCF ((cF)))p((le (log of 4p+2
X L
Σ ASp
As +
17
The proof is complete.
D.5.11 PROOF OF LEMMA D.16
Proof. Note that
1 Σ [[(f (2:+n), - (2s+n..))
[((fg)) [(fe-)
n i=1
S L sup iE(n]
f (2i+n) - fe. (di+m))
( (T:))))))))))))) TTi + b,) - ((w) witn + bb))
Lm sup
xinn
ve(n]JE[m]
Lm||a*I sein]jj[mm]
sup
(( (T)) TT + bb) (((w,((1)) (T)))TTanl ((t))TTzae Tiin by)
(318)
(0))) = and [+(3)() < M. So we can bound the gap of a ((w, (Tı) T + by) ((w + by) by the following lemma.
By Assumption 2 we know that d0()) = 0,
Lemma D.36. If an activation fincttion a satisfies a' (0) = 0, o²(0) and [x()(()) S M, then we have for all x, R:
(a(x) av))l [b((cl] bll) NY. (|0r2 + 1|/ |00 v--
(319)
50

---

Published as a conference paper at ICLR 2026
Proof of Lemma D.36 Define r(t) == ott) - 1/2²² -0(). Since o0) = 1 and ||0 M,
have that o" M -ipischhitz Hence
t = =1( (d"(s) (((0(()) - ds - Md It/
(320)
By the mean value theorem,
(n(x) - r(y)] < sup {Elm.y]
|r(( - |r - y/ - M max |z], 1a132 y/.
(321)
Therefore,
((x))o(v)) S E11²² -1 ( S [[xx]] + |w)||> vl + 1 max | 1013 v/. (322)
Finally, since max 6)2² v a² + 62, we get the conclusion.
Back to the proof of Lemma D.I6 under event En- from Lemma D.9 notice that for all in {a; p we have (wo) (1) x and w x are truncated, so ((w0)) ( x + by) and
(T+)T
((w ( T bm- are bounded by universal constant. Respectively, we have (T) Tx and (T))T w m-j x are truncated, so ((wi (TT))T + bz) and ((wm- ((wm))((1)) (Ti) ) JT. + bm-)) are also bounded by universal constant. Thus we have
[( [(fe-)) s Lm||a"|l sup ic(n],E[(m
( (Ti) T aitn + +bú) - (1w, (T))) Tiitn + b,,
Lm||a"|l sup sE|n].E[[m]
(wj (Ti)) Litn - (w, Iiun|
x Lmlla*||va
sup iE(m].JE(m]
(Ti)
wy (Taill
1
2 Lm- 1aal-aallvvin T. 60
(323)
where is defined in Lemma C.7
and from this lemma, we have the bound
| 2 (9) TT-1 Md E
(324)
By Assumption
know that cp C. 7 1772 V2 (x " (log d)
877a (1) where C. is universal constant. More-
over, have ||a ||a+1|2 1
Thus, it holds that
((...) [(fe-) - &
x Lmmlaa*| (IdTT, (4) TT-1-_1 Mdd. 60
V2 (N:T (log d)dp
S T.mp
1 log d.M
m
C.m/ndã
V22 (K6T) (1 d)40
1
V2 (4 d)
5 L
log d
S L Vn
(325)
The proof is complete.
D.5.12
PROOF OF LEMMA D.17
Proof. Lemma
D.15
implies that
Vi )ap (log dd
LL
Σs +
(326)
k:P
51

---

Published as a conference paper at ICLR 2026
Lemma D.16 further suggests that
V (nTT) (lo ) d)4p
((.6)) Z(se.) & LL
(327)
m
By combining these two bounds, we easily obtain
V2.
(A.T. (I:Ts (4 d) 4p+2
(afo-)--(())) LL
Σ A+
(328)
RSP
The proof is complete.
D.5.13
PROOF LEMMA D.I8
Proof. From Lemma D.17 we have the following bound
V2
(re (1TT (dp d))p+2 J
E(fe. L
Σ A +
(329)
m
&SP
Consider the training objective F((a) = ((a, ((T)) w(T)))) + 32/1a111 with B3 = tz. Let
J
a (00) == arg min F(a) a
(330)
be the minimizer of F.
We further denote 6o is the prescribed optimization tolerance ||a!T) (((T)a((op) [(00)|| S 6g. Furthermore, let
mL.
By optimality of a (00) and the constructed comparator a* we have
32 a|a01 112.
=(a') -((a=)()))))))))))))">)) +
(331)
Here we only use that a" is feasible point with ||a*|| 5 U and ī(a") S J (Lemma D.17).
On the other hand,
82
32
F(a [(aloo)) +
(2())))
>
a 11)
(332)
2
so combining with (31)) gives
2 D(a") $ SU² +
2J
2J
aa 100 < || +
= 12 +
J/U² S U22
S
(333)
B2
B2
Thus
||a (00) x U.
(334)
and also
|a S ||a [|)()))|
+|||| (T) a (00) 5o+
(335)
- lla (x0 F(a') S /+ & 42
E(a (x)) = (()
2
(336)
Thus we have
(F(a()) F(a
x; + b(7T)) - (o)) o(W(()) +b
sup
iE[n-1,2m]
+ (2(aa (T) + ||||| ||]| a (T) (|م0)
(T) - :c|=+1.20
(o( + +b(T)]
sup
" + (2) _
(337)
52

---

Published as a conference paper at ICLR 2026
By previous proof [o(W(T)) x; + b(T)) loc be bounded by absolute constant for all under event &. Thus have loose g norm bound and |F(a F(a F(a )| x Lila
|
(W x, b(T))
sup 32(1la (2) ||(00)))\\a0 1
2 (Llla (T) - a (0)) Vm + f2 (T) - a (x0)
S
(338)
Denote Copt (I-/ (2 J
LVM+OLU It is easy know that copt
Lym++&U
m.
v UmL E J. Thus, we obtain
((((T)) S.
(339)
Denote
:= ((W(T) ( E Rm.
i= +,.,,.... .n...........
(340)
82 2 [|a|| is convex and satisfies the By Assumption on the loss function the function F(a) gradient Lipschitz condition
-
Σ
(e(aT C' (a p.)) Ø:
n
i=n+1
2n
(aaa 01) - (a 0:) |61l|
-=++1
(341)
L
(a - az) Tbi ||:
2n
=
[10:11?) ||a| [2/1-
i=n+1
By previous arguments, ([a(((T));; x; + DT)) Ilo can be bounded by an absolute constant for all i [n 1, 2n] under event E. In addition, have E, || Tr( Σ, 0.0 I sup; ||
Thus the Lip constant for F satisfies Trpp T. J/U². Therefore, by Lemma E.13 if the step size 112
T.ip' gradient descent will approximate a (oc) to arbitrary accuracy in
1 log - T2 2 7232
(342)
steps.
Note that when the step size f)2 Lm/+
Lm + J/U
U
I J
(343)
12
32
Using the assumption that m S d and ||o
|| S U, (T(T)) S 1, can further simplify:
|a (T1) -
-a(o)) 12
U
log
log (mm)
(344)
T2
Thus, we conclude that gradient descent will approximate a (x) to arbitrary accuracy in at most
T2 2 J log(Um)
(345)
steps. The proof is complete.
n
53

---

Published as a conference paper at ICLR 2026
D.5.14 PROOF OF JEMMAD.19
Proof. Let F(b(T), W(T) {f(x) a (W b(T)), θ : ||a|| S U}. Note that under & we have ferm E F. We first define a new event for a ~ (0, I) under fixed W(T) and D(T)
M2 [h... (wn, TTa h ( ( || 4Dlogd
(346)
Define the finite-width event for the initialization as the intersection Ma == N (0)
M!! It is
straightforward to have Pe, {w, yy 2 1 - Odd-). We also define another event for a order to bound its first r components
M3 := sup (x)) S 4D log
(347)
ke[r]
Applying a union bound gives immediately P(Ms) 2 - O(d Denote M .Ma MM. Define the truncated loss
(M (L.)))
(348)
n1 8 (J(assa).Uttt))
(M()
(349)
By the symmetrization lemma,
sup LM - CM) 2 FEF
sup Σ Eee (f Vitun) [EF
vita)
(350)
where (6:)-1 are i..d. Rademacher signs independent of {Eiinn - Denote S ((re+n Wi+n) ) For fixed (2 Wetn), the mapping + 8M (1, yitn) has Lipschitz constant
in satisfying
(eM (s. (e(t, v) ((s. 1fx E M) s L|t s|.
(351)
Thus, by the Ledoux-Taiee -Talagrnd contraction inequality Lemma E.5 applied coordinatewise to the maps (M (tt))
Eg sup
Eee (.eM(f (f(#tt))) Vita) I LB
sup
(f(+))))))) = LRad,, (F).
(352)
Recalling Rad, (F) E E[Rad, (F)), we obtain the in in-expectation bound
(M( - cmcns < 2T. Rad,
Es
sup fEr
(353)
Furthermore, for under the event M. by the same argument as in Lemma (w Ta| $ for all j, thus we have
D.9 we have
(WW $ Vm
(354)
[f(x) = a (W x + TT) s U.m
(355)
(M (M((z). (((
$ <LIf(m) M)
M}
x LUVm LWV/m+L(loggd L(log < LU Vm.
(356)
We now give the following two high probability events in order to bound the gap between empirical loss and population loss.
54

---

Published as a conference paper at ICLR 2026
Definition D.37 (Event &, and E). We say the event E, holds if
log
sup fer
LM (f))
[LRad, (F) LU my
(357)
Using the standard immirrical-ooppoplaaai deviation bound by Lemma probability at least (O
E.4 we have E, holds with
Furthermore, we define event
E E
(358)
Hence, by the preceding probability bounds (and a union bound), & occurs with probability at least 1 Old -
By Lemma E.6 we have Rad,, (F)
S SUV T Thus
- CMM m log d
sup SEF
(359)
We notice that under the event Egs M holds for all a; E in + 1, 2m). thus (M( Z()). For the other term, we have
(M - c(f)
E((lf()),))( MY
[R[[ M)]
[E([[[(())1 Mi] {z) & MF))
ST. (VB(s(())))( & My] + E + ((( & MI])
(360)
Recall that it has been assumed that the activation function has bounded third derivatives as in Assumption Therefore, there exists an absolute constant C, > 0 such that
(a(4) < (1 + )u
for all E R.
(361)
Via a crude moment estimation bound, we can see the right side is S LUd P Thus, we have the following final estimation
log d
sup -(())-(()) LU FEF
LUd-D/S
(362)
Since we have assume n < dD/4, we have sup -(D) S LU formula for U and get our desired result.
m log
Plug in the
D.5.15 PROOF LEMMA D.21
Proof. Let 6min == minogksp &x- Because x + =- 8 is strictly decreasing on (0, oc), we have maxosk<p 6g -8 <min Also note that Ep EN 2 (p + 1)mmin with equality if and only if all ep are equal. Consequently,
((cx)) 2 A((p + 1)cmin + (Cm.. + C k min
(363)
and equality holds precisely symmetric choices eo Ep. Therefore the problem reduces to minimizing the strictly convex one-variable function
s(c) == (p + 1)c (Cm + C (c> 0).
(364)
Since
(c( (p + 1) B(Cm Cn)e
(365)
55

---

Published as a conference paper at ICLR 2026
and
g"()) B(B 1)(C + C.)c-(8+2) there is unique critical point, which is the global minimizer
0,
(366)
B(Cm + Cn) øों
=
(367)
p+1
Substituting back and using (Cm. + C. (E" ( - = pt1 E" yields 3
min g (3 + 1)8 - (p + 1) of (Cm Cn) Multiplying by LA gives the stated expression for min db
(368)
Finally, as is fixed, its contribution can be absorbed into the S constant, also note that $ (3 + 1)33 ササं
thus we get
= (3 1)5 * (p + 1) (Cm + Cn) Cz)
min (.+00)
2 A(Cm + Cr,) 35r
S 2DT log tog2p+14 d
(V +
VV=+V
1
341
(369)
The proof is complete.
D.5.16 PROOF OF LEMMA D.23
Proof. Let
log log
so Ti [T] E [T, + 1]. Then
2T exp(2T, logn) E L2 log
lox 22d2
log
(370)
and therefore
Inne Rid d 2T, S x² (log 222
ToR
(371)
Moreover, since 1, <
log logd have
11/T:Np2 S +
logd
(372)
and combining the two bounds yields
lid 2T + d1t//T, x² S x² log(d)d
222 lage
S dl 4l
Inga
(log d)
(373)
The proof is complete.
D.5.17 PROOF OF LEMMA D.25
Proof. Plugging in U and we get
T2 L.L
log(Um)
v2 (n (1)TT " d) Lm
log (v (ITT)) d) S V2 k5511 (logd)5P (374)
V22 (1cPa 1 (log d1
56

---

Published as a conference paper at ICLR 2026
((()++Ca)) गfr
Note that we have Vk,en S & by Defination D.7 and also note that &k
form
Lemma D.21 thus
T2 H 1/2 P.c. boT: (log d) m E %% bo (log d
24
+1
TT (log d) pm
3( 3(V +
I min (m, n) महमा v TpT, (log d) m S x 5pT: (log d) mm?
(375)
Moreover plugging = (log d) (p( +)+1 Ap(0+1) tet and Ti IV
logd
in, we get
log
T2 x spI (log d) mm² 2 xmx "log d) 8p3 13p+2 13ppp (8p2+11pp) The proof is complete.
Inga lord
(376)
D.5.18 PROOF OF LEMMA D.26
Proof. By Lemma D.25 we have
T2 S N (log d) 2
Ing
(377)
Plugging m (log d) 4p(8+1)+1 A((+1)) Iuen
Ck2 (log 2 in we get
amd . T = nmd (Ti T2)
(log d) (p(0+1)+1 ( S k² (log d)² d 2
lue
Iuen plog d) (100
(808+113))
S k 50+2 2log d) 1228+17046 (1228 / flos
(378)
(379)
The proof is complete.
E TECHNICAL BACKGROUND
E. HERMITE POLYNOMIALS
We briefly introduce Hermite Polynomials and Hermite Tensors. Definition E.1 (ID Hermite polynomials). The k-th normalized probabilist's Hermite polynomial, ha : + R is the degree polynomial defined as
(-1) d'
(x)
hx(x))
d=
(380)
Vkl
(x)
where ug(x) = exp(-12/2)/ 2T is the density of the standard Gaussian.
The first such Hermite polynomials
22 -
23 -32
(2) 1,
(2) = = h22
.2((2)
(381)
V6
Denote 3 = (0, 1) to the standard Gaussian ID. A key fact is that the normalized Hermite polynomials form orthonormal basis of L (3)): that is (hy
The multidimensional analogs the Hermite polynomials are Hermite tensors.
Definition E.2 (Hermite tensors). The k-th Hermite tensor Hep Ed (IR)) is
Her(x: = (-1)* v (x) Vk!
(382)
Hr (x)
where Hoy (x)) exp(- - exx((=b|ax(3)) (22 is the density of the d-dimensional Gaussian.
57

---

Published as a conference paper at ICLR 2026
The Hermite tensors form orthonormal basis of (2));; that is, for any (3(r), write the Hermite expansion
f(x) Σ Her(z)
where C+(f) := [aoD====eo)])())He(e (R*)**.
(383)
E.2
UNIFORM GENERALIZATION BOUNDS
Definition E.3, The empirical Rademacher complexity of a function class F on finite samples is defined as
Rad
Reg sup Σ ()(())) fEF t=1
(384)
where E1, En, are i...d Rademacher random variables PP(E: = 1) P(G = -1) 2/2 Let
(F) (FRa be the population Rademacher complexity.
Then recall the uniform law of large number via Rademacher complexity, which can be found in Theorem 3.3 in Mohri et al. (2018)
Lemma E.4. Assume that ranges in [0, R) for all F. For any (0,1) 1), with probability at least 1 - 8 over the i.idd training set S = (X1..... Xn) we have
log log(4/5)
sup JEF
Σ fx,) -E(()) < Rad, RV 1=1
(385)
n
Then we recall the contraction Lemma in /Vershynin 2018 Exercise 6.7.7) compute Rademacher complexity.
Lemma E.5 (Contraction Lemma). Let ys R + R with 1,..., n B-Lispehitz continuous. Then
Ee sup n fEF
Σ Exp o I (I:) < (R)
(386)
The next Lemma estimates the Rademacher complexity.
Lemma E.6. Let f(x) a a'o(Wa x + b) be our two layer neural network with our previous assumptions the activation functions in Assumption 2 Let
F F=o ||a|| S Ba, ||æ S Ibbll S SC}
(387)
where C' is a universal constant. Then we have the following estimation Rad, (F) S
Baym
vn
Proof.
Rad, (F)
Ex,
sup JEF
Σ Ci a ( x; + b))
n
Ba En
Σ Soo (W +
sup /EF
n
(388)
B
Ex E
sup feF
Σ Coo(Wo + b)
B
B Baym m vn
P
Σ alu x + b,)2
n
58

---

Published as a conference paper at ICLR 2026
E.3 UNIVARIATE APPROXIMATION
We will consider the following activation function
| - 1,
[t| 2 1, |1| < 1.
(t) =
(389)
11 12
Lemma E.7. There exists (a b), supported on {+1} [2, 3], such that for any |=| S
Eas(vo(a, ()z + b)] = 1, sup (o b)155 S
(390)
a.b
Proof. Let b)) = 12. 1a=1 2). 1og a(b)
Then, since 2 + 2 1,
{ (b 2 jo(2 + b)db
5
Ea,olvo(a, + b)] = 6
(391)
=6 / (b 2) )22 2b 1 1)
(392)
/ (b )db + 6
3
5.
3
5
2 6
(b (22 - 1)db
(393)
2 2
2
2
=1.
(394)
The proof is complete.
Lemma E.8. There exists in (a. b), supported on {+1} x (2, 3), such that for any |2| 1 1
Eac//i (az + +b]] = 2, sup ((a $1151
(395)
a.b
Proof. Let bo(a,b) = 1-=1 (-24 + +1)
bepa m(b)
Then, since z + +b> 1,
/ 3
b)) + +b]
+ +1)00 + bjdb
(396)
2
3 (-24b + 62) + 2b 1)db
(397)
J2
(-246 + 61)db + 2 /'''
(-24b + 61) - l)db
(398)
2
2.
(399)
The proof is complete.
Lemma E.9. There exists 12 b), supported on {=1} x [-2, 3), such that for any |2| S 1
Ea,lvz(a, + b)] = =2,
sup lu(a 6) $
(400)
a.b
Proof. First, see that
S a(2 + al2++))))
fate o(b)db
2+z
+b)lb
(401)
_242
( (-2b - 1)db + 1'/ 12 +
/ 2+
(2b 1)db
(402)
= =1-12 b)/22_ +
+ 1b2 112+z
(403)
3
2
= - 2)2 + (2 - 2) +
(2 + 2)² (2+2) 2)
(404)
3
14
= 2 +
(405)
3
59

---

Published as a conference paper at ICLR 2026
Let (2) la=1
12.1 (()) 30o(a,b) b) Then
1 ²
{
7
Eatl/u((a, ()a + +b] - -
a(
b)db
(406)
= + 7 7 =
(407)
The proof is complete.
Lemma E.10. Let u(b) = bk(k 1)(k: 2) - b)k-3 loc(o.1) Then p(b)
Eo(vx(b)o(z + b)] = = 1,>00 k(k - 1) 22
kz 1.
(408)
2
Proof. Plugging in vn (b) and applying integration by parts yields
/'' k(k - 1)(k - 2)(1 3)6-90(22 + b)db
E (b)or( bj] = jn
(409)
2
/- *k(k
= -1gh(h 1)(1 - 6)k - 20
6)1d
1)(1 b)* 6)--22 20' + bbdb
(410)
1 k(k 12)(2) + 1/ +1212t k(1 6)--111- d'(2
S'" k(1 6)* 10)--)/"(" 00"(2 + bjdb
1 + 6)9b
(411)
- 1))(2))
1 ko') f k(1 - b)* '1 11
(412)
When 1 2 >> 0, we have
1 k(1 (0)--ba-- -))-11 2+0 = f/r k(1 b)k- db [11 -))11= = =--11 -1.
1-z
(413)
When -1 S 2 S.5 0, we have
J. k(1 - b)h 124b/// = /-' k(1 - b)b = 11
11
(414)
.
Since 2 E [-1, we have that o(2) = z² and 0'(2) = 22. Therefore for E [-1,1]
Ex(vx(b)o + b)] = *² 1220 k(k
1) 2 kz 1.
(415)
2
The proof is complete.
Lemma E.11. There exists en (a, b), supported on {+11 x 1-2,3], such that for any |=| S
Eu,lvu(a, b)o(az + b]]] = sup
yax (a. b)] S poly(k).
(416)
a.b
Proof. We focus & 2 3. We have that
k(k k(k-1) - 1)
[b(ve(b)o(z + b)] = =,,, 1220
22 kz 1.
(417)
2
k(k 11)
1)
Es(vx(()((=z+)))) b)] (-=)** 12x0
2
4 kz
(418)
2
Therefore if is even
Ea ())
b)]
= - k(k 1)z² 2.
(419)
Let (a b) = (b) + kkk Then
Eu,lvvl(a, ()z + b)] (o(vx(b)=(++))) + (x b)] + k(k - -1)²² 2 2 (420)
If k is odd,
(a(u(b)o(z -b)) b)] =² -k
(421)
Let un b) = (b) (b) 2av (a, Then
E., (a, (az + 6)] = Eg Eolux(bo0( b) - nn (b)((( - 6)] + 2kz = 2/
(422)
The proof is complete.
60

---

Published as a conference paper at ICLR 2026
E.4
CONVEX OPTIMIZATION
Denote (a) as a C² function defined in IR.. Assume that
There exists m such that ((x)
m ||
convex.
[V[()) V f fz)) - L| y/1-
The following result standard and can found in most convex optimization textbooks like Boyd & Vandenberghe (2004).
Lemma E.12. There exists unique x² such that f(x") = inf f f(x). And if we start at the point x and do gradient descent with learning rate 17, in S m. then we will get
||ee x²1 S 1= - x²1
(423)
where e -1 m
We compute the minimal number of the required iterations.
Corollary E.13. Lere E (0, 10= ==) be target squared-distance accuracy. In order to have
|| x 12 e² is enough have
k E m
log || - x= x²12
(424)
MML
The above corollary can be verified from a direct computation.
E.5 SPHERE LEMMAS
Lemma E.14. Let x(d). Then,
k-1
2 [I d(d+ 2). ...(d. 2h: 2) (dd*).
(425)
Proof. Write 2 (d) with density proportional to x 1--12-=/2 on x > function identity,
Using the Gamma
E/ E[( = (A
k-1 = (d+
(426)
r(
==0
Since I (d + 2f) = d for fixed &s (e.g., by Stirling or monotonicity of the ratio), we obtain (ct)). Lemma E.15. Let Unif((d-1)) Then,
E N.) 2k
E
22
(427)
E Exnx(d) (d) 222k]
Proof. This follows from w = 1/10 with v ~ x(d), w ~ Unif independent.
Corollary E.16. By Wick/Isserlis' 's formula for Gaussian moments,
Ew-NNO, lw 82k (2k 1)!lSym(I
(428)
Meanwhile, Lemma E.14 shows E Ev)y(d)) 2k (ddh). Hence overall we have
w (Smm @k).
(429)
Corollary E.17. Let T be symmetric p-tensor with dim (T)) 1-2d
With probability at least
r (log d)*
T T(
S
[Tlla]]
(430)
db
61

---

Published as a conference paper at ICLR 2026
Proof. Denote Y [T(w F- We first bound the second moment.
Flatten T matrix
M Rd xx so that T( 8k
M (w Then
u || (M, MCL) = -IMO - | B[ = E M
(431)
E S the d df-dimensional k-mode space. Next, from Lemma E.15 and Wick/Isserlis' formula for Gaussian even moments, Cs = cs,aSym( with
where Cs
E
(w
T
C,,d x Restricting to the r-dimensional span of (on each paired contraction) yields an extra factor [k/2] from pairing counts. Hence
LLA/2] de
rlA/2/ de
BY S
|M
(432)
We then bound Y with high probability. The map + T( is homogeneous polynomial of degree & restricted to gá-1. By (spherical) hypercontractivity for degree-k polynomials on sd-1,
Lemma E.22 for all q 2 2,
|Y (84-11 S (CVV))"lX
(433)
Combining with equation (432) gives M I (CVi)*|IT(l) inequality 222g with g x Dlog d yields
k/2l dt
and applying Markov's
r Lk/22 d)b
P Y
S 2d
(434)
dk
which concludes the proof.
E.5.1 PROOF OF LEMMA D.34
Proof.
As the previous proof, flatten T into M E Rxd"
kxd Then
E T > 1
| 1/2 k
}/ s |M C
=
E i 02k
Cy.aSym(Iet)
(435)
so || = Ck.d = k
On the other hand,
B(T, w Sp)2 (T, E w 2
T with Cppt x P
(436)
Therefore
E T || S P S d d cp.al |T| "T="E(TTw"" k
(437)
which proves the claim.
E.6 SUPPORTING TECHNICAL LEMMAS
Lemma E.18 (Norm Concentration for Gaussian Samples). Let f1,..., Ba N(O, Ia) be i.i.d. Gaussian vectors. Then with probability at least 1 - nexpp -ed), we have
| < 2vd Vie
(438)
where c is some positive universal constant.
Proof. Note that ||x||² follows a chi-squured distribution with freedom and standard concentration for the chi-squared distribution, for instance Theorem 2.8.2 in Vershynin 2018 gives us
P(J 2 2Vd) s exp(-cd)
(439)
with some universal positive constant The result follows with a union bound argument.
Lemma E.19 (Coordinate concentration parameterized by D). Let Fi,..., En ~ r E [d]. For any > O, with probability at least 2nrd-D.
Ltd.
N(0, Ia) and fix
mAx
[(x:)al & 2 Dlogd.
(440)
VE REi
62

---

Published as a conference paper at ICLR 2026
Proof. For Z N(0. 1), P([Z] 2 t) < 2/- Setting v4D log d gives P(|Z] v4Dlog d) I 2d 2D Then we do a union bound over all [n] and S r.
Let the Hermite expansion of our target be f(x) = Ei= nn (Ck, (e)())) where C₁ (RR)) is the k-tensor defined by Cg = E V [(x)] and supported on S Here He, the unnormalized k-th Hermite tensor.
Lemma E.20 (Parseval's 's Identity).
= Ex(f(x)
=
C
(441)
k!
As immediate consequence of Lemma E.20 we have |C S k!. Note that f(x) also has the following tensor expansion
f(x) = Σ
(Te. x x*h)
(442)
kSP
The following lemma shows |T F S r S
Lemma E.21. There exist To, .... Tp such that
f(x)
Σ Tr
(443)
kSP
and [Tal|e Sr < 1 1 Sp.
Proof. Note that from the Taylor expansion of fax we have
V [(0)
CHe((He;(0))
(21 - 11C5))1))
T = k!
Σ
(444)
klj!
M!(2))!
2j k
We recall from the consequence of Lemma E.20 that |Ciil I k! and therefore,
|Tale $
Σ
|C22+t (T+))| S S - 2
(445)
25Sp-k k
shall use the next Gaussian hypercontractivity lemma from Theorem 4.3, Prato Tubaro 2007
Lemma F.22. For any é N and f [²((r) be a degree polynomial where the input distribu-
tion & = N(0, Ia), for any 2 2, we have
Ez~r [s(=)9 C..t (Bzny [)(=)²))012
(446)
where we use Cg to denote some universal positive constant that only depends on E,
We also need this Spherical hypercontractivity lemma from Beckner 1992); O D'onelll (2014).
Lemma E.23. Let Y., be a spherical harmonic of degree m Sd-1 Then for all < p S q 00 we have
- m/2
Y LL(0--1)
I |YmllLP(BB-1)-
(447)
In particular, for any polynomial P of degree at most m and 2 2 we have P| 1) S (g
||
63


Reviews:
[{"review_id": "F21LPLICDX", "reviewer": "Reviewer_Quxi", "summary": "The paper studies the theoretical ability of two-layer neural networks trained with gradient descent to learn so-called multi-index models of the form f(x) = g(Ux). The authors assume that the link function g is a polynomial function of a few linear projections of the input. They analyze a specific layer-wise training algorithm: first, the inner weights are updated by a few small gradient-descent steps (the “feature-learning” phase), and then the outer weights are optimized on a fresh dataset (the “readout” phase). The main theoretical result (Theorem 1) shows that, under certain regularity and non-degeneracy assumptions on the loss and activation, the resulting network achieves vanishing test error using O(d log d) samples—thus matching the information-theoretic limit of O(d) up to logarithmic factors.\n\nThe central idea is to replace the usual spectral initialization by a short sequence of small gradient-descent steps followed by one large step, and to prove that with sufficiently many samples this procedure achieves perfect learning. In other words, the first phase acts as a substitute for spectral initialization, while the final large gradient step ensures convergence to the correct subspace. This approach avoids relying on an explicit spectral start, which is the standard technique in related problems.\n\nWhile this is an interesting and elegant way to re-establish an expected result, the paper in its current form contains several overstated claims, imprecise statements made to artificially broaden its generality, and misleading references to prior work. The contribution is promising and potentially valuable, but the presentation requires revision and clarification (see detailed comments below). \n\nThis issue is the main reason for my low overall score: the technical results are solid and interesting, but the framing, scope, and literature discussion are overstated. Importantly, these problems should be relatively easy for the authors to correct through a careful revision that accurately reflects what is proved and properly situates the work within the existing literature.", "questions": "To strengthen the paper and make its contribution both accurate and impactful, I suggest the authors consider the following revisions:\n\n* Reframe the scope and claims: The current title and framing promise a general solution to “learning multi-index models near the information-theoretic limit,” while the actual analysis applies only to non-staircase, generative-exponent-2 polynomial functions. The paper would benefit from a more precise and honest statement of scope. A few possible titles that would better reflect the contribution: “Learning Low-Dimensional Polynomial Features via Layer-Wise Gradient Descent”\n\n* Clarify conceptual novelty: The key insight is that gradient descent can mimic spectral initialization—an elegant technical observation. The authors should emphasize this as the main conceptual contribution, rather than claiming a fundamentally new learning phenomenon. It would be useful to explicitly contrast their method with the spectral algorithms, clarifying what is gained by expressing them in the language of neural-network training.\n\n* Specify assumptions and limitations clearly. The paper should explicitly list the restrictions implied by Assumptions 1–5, particularly that (i) the link function must be a low-degree polynomial, (ii) the activation is smooth and quadratic near zero, and (iii) the loss must be specially chosen to ensure non-degeneracy. The authors could also discuss whether these conditions are essential or merely technical conveniences.\n\n* Correct and balance the related-work section. The discussion should be revised to include and properly attribute earlier foundational results. Correcting the chronology and expanding the citations will make the contribution clearer and situate it fairly within the literature. It would be useful to comment explicitly on how this analysis relates to the established dynamical mean-field (DMFT) or “state evolution” approaches. Are the authors’ results a finite-time counterpart of those dynamical equations? A short discussion would make the connection to the broader learning-dynamics literature much clearer.\n\n* Address anonymity issues: The reference to “Andrea Montanari and Zihao Wang, in preparation (2025)” should be removed or replaced by an anonymized placeholder consistent with the ICLR double-blind policy.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "It was already known from Chen & Meka (2020) and related works that low-dimensional polynomial functions (i.e., generative-exponent-2 multi-index models) can be learned efficiently, near the information-theoretic limit. The present paper’s contribution is to re-establish this result within a neural-network training framework, showing that a standard two-layer network trained by gradient descent can implicitly reproduce the behavior of these earlier polynomial-learning algorithms. Replacing an explicit spectral initialization by a short sequence of small gradient steps is an elegant idea that bridges classical spectral methods with early gradient-descent dynamics.\n\nOverall, this is a good and solid piece of work. It addresses a well-motivated and long-standing open question—whether standard gradient-descent training of neural networks can, on its own, achieve near–information-theoretic learning of structured high-dimensional targets. The main statements are interesting and relevant to the current theoretical understanding of representation learning, and the approach provides a clear and rigorous connection between classical spectral algorithms and neural-network optimization dynamics. Even though the setting is somewhat restricted, the paper gives a precise and constructive answer to this important question, and the results are likely to be of interest to researchers working on the theory of learning dynamics and high-dimensional inference.", "weaknesses": "I think the main result is genuinely interesting—essentially showing that, in a specific regime, the network dynamics are equivalent to a power iteration on a certain matrix. But the framing is off: the paper presents this as a broad statement about neural networks learning generic multi-index models, whereas in reality it addresses a much narrower, well-understood setting. A more honest and focused framing would make the contribution stronger and easier to appreciate. As of now, the paper currently overstates its contributions, blurs the limits of its applicability, and misrepresents parts of the existing literature.\n\n1. Statements and Overclaimed Results\n\n• Overstated main claim.\nThe paper does not fully align with either its title or its main motivating question. It is therefore misleading to describe the work as proving that “neural networks learn multi-index models near the information-theoretic limit.” In practice, the analysis concerns only non-staircase, generative-exponent-2 polynomial functions, a regime already known to be learnable efficiently since the work of Chen & Meka (2020). The authors simply replace the spectral start used in those earlier algorithms by a short gradient-descent phase with a (largely unspecified) modified loss that mimics the spectral initialization. This is a nice idea and could justify publication, but it is essential to correctly reflect the actual scope and contribution of the work.\n\nFrom the title and framing, readers are led to believe that the authors have proved that “neural networks learn multi-index models near the information-theoretic limit.” In reality, it is not even clear which specific class of models this applies to (see Assumption 5). The claim is clearly overstated; even for single-index models, such a general statement would not be justified.\n\nIndeed, the spectral-start strategy (or its gradient-descent surrogate) can only access functions that are not staircase, or more precisely, not generative-staircase in the terminology of Troiani et al. Using the rigorous interpolation method of Aubin et al., those authors showed that learning general multi-index models requires an iterative staircase mechanism—learning one direction at a time and restarting the spectral step. All these regimes are excluded by construction here. While such an extension might be possible, it is far from trivial and should not be implicitly claimed.\n\n• On Assumptions 4 and 5.\nThese assumptions effectively linearize the loss and amount to introducing a spectral start. The text states:\n\n“Recent studies on generative exponents Damian et al. (2024; 2025) show that for almost all common tasks … there exists a pre-processing function that ensures the second-order information is non-degenerate.”\n\nFor generative exponent 2, however, this result is in fact due to Yue Lu (2017) or Mondelli & Montanari (2018). Reading Damian et al. (2024, 2025) reveals no genuinely new results for exponent 2—they mainly restate existing work and extend it to higher exponents. It is therefore misleading to credit them for ideas that are much older. The fact that the paper never references the original works on spectral methods—single-index (Lu 2017; Mondelli & Montanari 2018; Maillard et al. 2020) or multi-index (Troiani et al. 2025; Defilippis et al. 2025; Kovačević et al. 2025)—is surprising. The key matrix Σ already appears in these earlier analyses. (See, for instance, Lu 2017, arXiv:1702.06435; Maillard et al., arXiv:2012.04524; Aubin et al., arXiv:2010.03460, none of which are cited despite predating most of the works referenced here.)\n\n\n2. Imprecise Statements\n\nThe paper relies on a modified loss, special activation functions, and other restrictive assumptions, which make it unclear what it truly means for these models to be “learned by neural networks.” For instance, using a standard ℓ₂ loss would not work, nor would a ReLU activation. This is not a problem in itself, but given the ambitious title and claims, it is misleading.\n\n• Choice of loss function.\nThe authors do not use the square loss but rather a specially chosen loss that ensures non-degeneracy. Ass.4 and 5 are so vague that it is unclear which functions are actually covered. Given the limitations above, the authors should clearly specify the excluded cases (e.g., all staircase structures). Even if such a loss can be generic, this remains a limitation that should be clearly acknowledged. Saying that “neural networks learn” is somewhat misleading, since the success of the algorithm depends on a modified loss. Similar connections have already been discussed in the literature—for instance arXiv:2403.02418, arXiv:2510.18435, and, in the DMFT context, arXiv:2509.23527 and arXiv:2103.04902—none of which are cited. (Notably, Yue Lu’s 2017 paper was itself titled “Spectral Start.”)\n\n\n3. Discussion of Related Work\n\nThe discussion of related work is incomplete and historically inaccurate, which substantially weakens the scholarly framing of the paper. In fact, reading it one might believe that the relevant literature began in 2024. Several key theoretical developments that underpin the results are misrepresented or omitted. More generally, the section over-emphasizes extremely recent U.S.-based works while under-citing the earlier theoretical literature that established (i) Bayes-optimal errors and phase diagrams for GLMs, (ii) spectral thresholds (BBP-type) for pre-processed covariance methods, and (iii) AMP algorithmic limits—all directly relevant to the present theorems.\n\n(1) Information-theoretic foundations.\nThe paper attributes the identification of the information-theoretic scaling n = Θ(d) for single- and multi-index models to Dudeja & Hsu (2024) and Damian et al. (2024). In fact, these results were rigorously established much earlier by Barbier, Krzakala, Macris, Miolane & Zdeborová (2017, PNAS), who derived the exact Bayes-optimal risk and phase transitions in high-dimensional GLMs using the Guerra/Talagrand interpolation method. The multi-index generalization appeared the following year in Aubin et al. (NeurIPS 2018, “The Committee Machine: Computational to Statistical Gaps in Learning a Two-Layer Neural Network”)—not to mention the non-rigorous statistical-physics literature from the 1980s that already predicted these results.\n\n(2) Spectral and AMP results for single-index models.\nThe manuscript credits Mondelli & Montanari (2018) for “generic single-index models and sharply deriving the learning threshold.” Yet this threshold was obtained earlier in Barbier et al. (2017, PNAS), where Eq. (11) gives the celebrated “generative-exponent 2” transition derived via AMP. The contribution of Mondelli & Montanari was to employ pre-processing within a spectral method—but even this approach was introduced earlier by Yue M. Lu and collaborators (arXiv:1702.06435; arXiv:1811.04420), who analyzed spectral initialization and its BBP-type phase transitions for a wide class of models.\n\nThe quoted papers by Damian et al. on generative-exponent 2 for single-index models do not introduce new results; they merely restate those earlier works. Their genuine (and important) contribution was the extension to higher exponents, which is not relevant to the present paper, as it remains confined to exponent 2.\n\n(3) Causality and scope mismatches: The following sentence is incorrect: “Subsequently, Troiani et al. (2025) … recovers the results of Mondelli & Montanari (2018).” In fact:\n- Causality is reversed: The Bayes-AMP theory (Barbier et al., 2017–2019) came first; Mondelli & Montanari (2018) recovered its single-index specialization.\n- Scope mismatch: Mondelli & Montanari treated single-index models only. The generalization to multi-index models followed Troiani et al. (2025), which inspired two independent recent spectral studies:\n(i) Kovačević, Zhang & Mondelli (2025, COLT, arXiv:2502.01583)\n(ii) Defilippis, Dandi, Mergny, Krzakala & Loureiro (2025, NeurIPS, arXiv:2502.02545)\n\nBoth derived equivalent results, and the latter appeared slightly earlier (as acknowledged in the former’s conclusion). Only the first is cited here. Because these spectral starts are central to the present paper, such omissions and reversed chronology substantially weaken the discussion.\n\n(4) On staircase and generative-exponent structures.\nThe authors emphasize the staircase or leap-complexity results of Abbe et al. but ignore the more relevant generative-exponent framework. Unfortunately, no gradient-descent results are known for that case, and the strategy used here cannot address it. These limitations are never discussed, even though the paper is restricted to polynomial targets (see Assumption 1). This omission is problematic given the claimed generality of the title (see Section 1 of this review). \n\n(5) DMFT and dynamical analyses.\nNo discussion of DMFT-based results appears apart from the recent Montanari & Urbani (2025). Rigorous DMFT studies have long shown that, with an appropriate “hot start,” two-layer networks can be learned efficiently (see e.g. arXiv:1710.04894; arXiv:2006.06098; arxiv:2112.07572 arXiv:2210.06591, rXiv:2303.00055). Instead, the authors cite an unpublished and identifiable work, “Andrea Montanari and Zihao Wang, Spiked Model for the Hessian of Two-Layer Neural Networks, In preparation, 2025.”\n\nWe note that this citation violates the ICLR anonymity policy: it (i) names specific, non-anonymous researchers; (ii) refers to unverifiable, unpublished work; and (iii) could reveal the submitting authors’ identities through known collaborations. Even if the reference is genuinely independent, such a citation breaks the spirit of the double-blind process and should be removed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YkX25oiNow", "reviewer": "Reviewer_t3JK", "summary": "This paper studies learning multi-index models with two-layer neural nets using gradient descent (GD). The authors prove that for polynomials with leap exponent at most 2, $\\Theta(\\sqrt{\\log d})$ steps of GD on the first layer with $\\tilde{O}(d)$ samples learns the latent representation, which is almost information-theoretically optimal, with a computational complexity of $\\tilde{O}(nd)$ which is also nearly optimal. The main intuition is that with a super-constant number of steps, GD mimics power iteration on a certain matrix and act as a spectral method.", "questions": "4. The authors mention the possibility of extending their analysis to non-Gaussian inputs. One relevant work here is [1], where the authors prove that GD can learn generic functions with near information-theoretically optimal sample complexity, at the expense of a potentially super-polynomial computational complexity. One interesting aspect of such result is that it can also accommodate non-isotropy in the input. In fact, I'm curious if a sample complexity better than $\\tilde{O}(d)$ could be achieved if we assume a structured covariance matrix for the input that reveals some information about the target direction.\n\n5. Wouldn't you need some assumption on the loss that enforces the prediction and ground-truth label to be close (e.g. by enforcing that the loss grows as $t$ and $y$ diverge)? Currently the loss seems too general, but for the network to learn the correct representation (which is how the proof works) I suspect that such an assumption might be necessary?\n\n6. I might've missed this but I couldn't find the assumption $r=\\mathcal{O}(1)$ written in the paper. I think one either needs this, or explicit dependence on $r$ in the bounds.\n\n\nReferences:\n\n[1] A. Mousavi-Hosseini, D. Wu, M. A. Erdogdu. \"Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics.\" ICLR 2025.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The finding that a super-constant number of steps can improve the sample complexity of feature learning under the multi-index model is interesting, and the proof sketch seems intuitive and can be used beyond this work.", "weaknesses": "1. In comparison with the recent literature, Assumption 5 is a bit restrictive. It would be interesting if one could show that gradient descent would automatically reduce leap complexity down to at most 2, similar to Lee et al., 2024 for the single-index case. I agree however that it would be technically challenging to achieve such result, and it's perhaps an interesting direction for future research in this area.\n\n2. The Lipschitz assumption on loss (Assumption 3) rules out squared loss, and I'm not sure if it's really necessary.\n\n3. Similar to other papers in this area, the authors have to consider a layer-wise training dynamics. Having an analysis of training both layers at the same time however is highly non-trivial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YkK1KQOa24", "reviewer": "Reviewer_jpVz", "summary": "The authors investigate the problem of learning a polynomial Gaussian multi-index model using a two-layer neural network trained via gradient descent. Through a layer-wise training procedure, and under mild assumptions on the loss function and activation, they prove that there exists a suitable choice of hyperparameters for which the network efficiently learns the hidden subspace and achieves a small generalization error, close to the information-theoretic sample complexity.", "questions": "1. Following point 1 in Weaknesses, could you offer some insights on which assumptions regarding initialization and the training scheme (e.g., sample splitting, bias reinitialization) could be relaxed, under a more careful analysis, without changing the qualitative results? Which ones do you believe are necessary for the theoretical guarantees to hold?\n\n2. At the end of Sections 1.1 and 4, you mention that choosing $T_1 = o(\\log d)$ is necessary to prevent the weights from collapsing onto the top eigenvector of $\\hat{\\Sigma}_\\ell$. However, Appendix A also states that the power-iteration approximation remains accurate only up to $o(\\log d)$ steps. Could you clarify the relationship between these two statements? In particular:\n - does the breakdown of the power-iteration approximation beyond $o(\\log d)$ still imply that part of the signal is suppressed, or the analysis no longer applies?\n - more generally, could you provide some intuition on what happens to subspace recovery if the first-layer training continues for longer times, when corrections to the power-iteration dynamics become non-negligible?\n\n3. The discussion of Assumption 5 could be expanded. In particular, which polynomial target tasks would violate it for standard choices of the loss function ($\\ell_1$, Huber, squared loss, etc.)? Conversely, could you give examples of reasonably natural losses for which the assumption would be satisfied for most polynomial targets?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper presents a novel theoretical analysis of the ability of shallow neural networks to learn multi-index models. The assumptions on the loss function and activation are fairly general and their necessity is discussed transparently. The connection between the early-stage training dynamics and power iteration provides insight into the trade-off between training time and full subspace recovery. The bound in Theorem 1 improves upon previous results, showing that, in principle, two-layer neural networks can generalize polynomial multi-index models with $n = \\widetilde{O}(d)$, approaching the information-theoretic limit. The manuscript is well structured, with Section 4 (and B.1) offering clear intuition behind the main steps of the proof.", "weaknesses": "1. The analysis strongly depends on specific initialization assumptions and on the particular layer-wise training scheme. The manuscript lacks a discussion of their necessity or the validity of the claims beyond these settings.\n\n2. Several references are missing or misplaced:\n\n- Line 95: [1] and [2] are two other relevant references on learning thresholds of single-index models.\n\n- Line 116: spectral methods for multi-index models (with generative exponent 2) achieving the optimal learning threshold have been studied in [3], concurrent with [Kovacevic et al. (2025)]. Both works should also be cited on line 45, together with [Damian et al. (2025)].\n\n- [Troiani et al. (2025)]: this work shows that many multi-index models of interest can be learned with $\\Theta(d)$ samples; therefore, it should be included in the discussions on lines 110 and 437. It could also be mentioned on line 72.\n\n- [Lee et al. (2024)], [Arnaboldi et al. (2025)]: both works cite each other as concurrent and show equivalent results. Moreover, could the authors clarify the meaning of the statement “the recovered features are learned only in a quite weak sense” (lines 52 and 139)?\n\n3. The paper is purely theoretical and lacks numerical validation. While not a weakness per se, a few illustrative examples would strengthen the claims and improve the overall presentation.\n\n4. The bibliography should be harmonized, and missing journal or conference names should be added when necessary.\n\n\n\n[1] Barbier et al. \"Optimal Errors and Phase Transitions in High-Dimensional Generalized Linear Models\"\n\n[2] Lu, Li \"Phase Transitions of Spectral Initialization for High-Dimensional Nonconvex Estimation\"\n\n[3] Defilippis et al. \" Optimal Spectral Transitions in High-Dimensional Multi-Index Models\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1F5dHRRvAR", "reviewer": "Reviewer_5Tcn", "summary": "The paper studies the problem of learning multi-index models with a two layer network using gradient descent. In particular the authors show that first training the first layer for a few step, and then freezing it an only training the second one it's possible to achieve vanishing test error with linearly many samples in the input dimension and quadratically many iterations. The key insight is that with a carefully chosen initialization (first layer with infinitesimally small weights, second layer either +1 or -1) GD is essentially performing a fixed point iteration on the empirical covariance of the data, which means that after $O(log(d))$ iterations the learned features are in the subspace of the multi-index model directions. One can then train the second layer to have good generalization.", "questions": "1. I am puzzled by the need to train the first layer for no more than $O(log(d))$ iterations. I find this very surprisingly low, could you comment on this? \n2. More on 1: is that something that is a direct consequence of the layer-wise training? Beyond what you can prove, how much of a role do you expect layer-wise training to play in your results?\n3. In Troiani et al. '25 the authors show that AMP undergoes a saddle to saddle dynamic. Do you expect GD to undergo similar dynamics? Would something like the \"grand staircase\" in chapter 5 of their paper happen for GD?\n4. Would the results change if $g$ would be a generic leap exponent 2 function instead of a polynomial?\n5. I find the total absence of numerical experiments slightly disappointing. I would greatly appreciate having a set of experiments, especially investigating theeffect of changing $T_1$.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is extremely well written and does an incredibly good job at explaining the general idea of the paper and its proof directly in the main text. I believe this work to be a significant contribution to the theoretical understanding of machine learning, extending the literature on multi-index models significantly.", "weaknesses": "Both the data model and the training procedure is extremely idealized and ad-hoc, and it's unclear to me how these results would generalize in more realistic settings. Even within the theory landscape, the target multi-index model is just a polynomial, in contrast with what stated in the abstract. In particular training the layers one at a time is not standard practice, and this is made even worse by the need to control the number of iterations when training the first layer. Additionally the paper lacks any experimental validation or numerical check of the results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)