Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit
Authors: Jason D. Lee
Kazusato Oko
Taiji Suzuki
Denny Wu
Venue: NeurIPS 2024 poster (2024)
Keywords: single-index model
feature learning
statistical query
SGD
Primary area: learning_theory

Abstract: We study the problem of gradient descent learning of a single-index target function $f_*(\boldsymbol{x}) = \textstyle\sigma_*\left(\langle\boldsymbol{x},\boldsymbol{\theta}\rangle\right)$ under isotropic Gaussian data in $\mathbb{R}^d$, 
where the unknown link function $\sigma_*:\mathbb{R}\to\mathbb{R}$ has information exponent $p$ (defined as the lowest degree in the Hermite expansion). Prior works showed that gradient-based training of neural networks can learn this target with $n\gtrsim d^{\Theta(p)}$ samples, and such complexity is predicted to be necessary by the correlational statistical query lower bound. 
Surprisingly, we prove that a two-layer neural network optimized by an SGD-based algorithm (on the squared loss) learns $f_*$ with a complexity that is not governed by the information exponent. Specifically, for arbitrary polynomial single-index models, we establish a sample and runtime complexity of $n \simeq T = \Theta(d\cdot\mathrm{polylog} d)$, where $\Theta(\cdot)$ hides a constant only depending on the degree of $\sigma_*$; this dimension dependence matches the information theoretic limit up to polylogarithmic factors. More generally, we show that $n\gtrsim d^{(p_*-1)\vee 1}$ samples are sufficient to achieve low generalization error, where $p_* \le p$ is the \textit{generative exponent} of the link function. Core to our analysis is the reuse of minibatch in the gradient computation, which gives rise to higher-order information beyond correlational queries.

Neural network learns low-dimensional polynomials with SGD near the information- limit
Jason D. Lee', Kazusato Oko Taiji Suzuki Denny Wu
Princeton University, University of California, Berkeley, University of Tokyo "RIKEN AIP, New York University, Flatiron Institute
rasonlee@princeton edu, oko@berkeley. edu, taiji@mist.. ac. dennymu@nyu.com edu
Abstract
We study the problem of gradient descent learning of a single-index target function (x = a_( ((xx 0)) under isotropic Gaussian data Ed where the unknown link function 0, R + E has information exponent (defined as the lowest degree in the Hermite expansion). Prior works showed that gradient-base training of neural networks can learn this target with n Z dB(p) samples, and such complexity is pre- dicted to be necessary by the correlational statistical query lower bound. Surpris- ingly, we prove that two-layer neural network optimized by an SGD-based algo- rithm (on the squared loss) learns f. with complexity that is not governed by the information exponent. Specifically, for arbitrary polynomial single-indes mod- els, we establish a sample and runtime complexity of n 2 T O(d- polylogd). where (()) hides a constant only depending on the degree of a.; this dimension dependence matches the information theoretic limit up to polylogarithmic factors. More generally, we show that n Z d(p- (()1- samples are sufficient to achieve low generalization error, where p. S p is the generative exponent of the link function. Core to our analysis is the reuse of minibatch in the gradient computation, which gives rise to higher-order information beyond correlational queries.
Introduction
Single-index models are classical class of functions that capture low-dimensional structure in the learning problem. To efficiently estimate such functions, the learning algorithm should extract the relevant (one-dimensional) subspace from high-dimensional observations: hence this problem set- ting has been extensively studied in deep learning theory [BL20, BES 22, BBSS22, MHPG 23, MZD+22 23, WWF24], to examine the adaptivity tation learning in neural networks (NNs) optimized by gradient descent (GD). In this work we study the learning of single-index target function under isotropic Gaussian data:
low-dimensional targets and benefit of represen-
= fa(x)) + f. a. ((x; 0)), 2;
Lid. N/O,
(1.1)
where is i.i.d. label noise, 0 E Rd is the direction of index features, and we assume the link function o, : R + R. has information exponent p e N+ defined as the index of the first non-zero coefficient in the Hermite expansion (see Definition 1).
Equation (1.1) requires the estimation of the one-dimensional link function a. and the relevant direc- tion 0; it is known that learning information theoretically possible with d training examples [DH24, DPVLB24].. Indeed, when a. is polynomial, such statistical complexity can be achieved up to logarithmic factors by tailored algorithm that exploit the structure of low-dimensional target [CM20]. On the other hand, for gradient-based training of two-layer NNs, existing works estab- lished a sample complexity of n (r) [BAGJ21, BBSS22, DNGL23], which presents a gap
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

1.4
10
0.08
10a
sample size n
1.0
sizee
10
0.04
0.6
0.02
n n=dd
10a
=
0.2
10 dimensionality d (a) Online SGD (weak recovery).
10a 8)²
20a dimensionality d
10a Em
(b) Same-batch GD (generalization error).
Figure I:
We train a ReLU NN (3.1) with N 1024 neurons using SGD (squared loss) with step size 7)= 1/d learn a single-index target f- (x) = Hes(( 0)); heatmaps are values averaged over 10 runs. (a) online SGD with batch size B = 8; (b) GD on the same batch of size n for T 2 steps. For online SGD we only report weak recovery (i.e., averaged overlap between neuron and target 0) since the test error does not drop.
between the information theoretic limit and what is computationally achievable by (S)GD. Such gap is also predicted by the correlational statistical query (CSQ) lower bound [DLS22, AAM23), which roughly states that for CSQ algorithm to learn (isotropic) Gaussian single-index models using less than exponential compute, sample size of n 2 dp/2 is necessary.
Although CSQ lower bounds are frequently cited to imply fundamental barrier of learning via SGD) (with the squared/correlation loss), strictly speaking, the CSQ model does not include empirical risk minimization with gradient descent, due to the non-adversarial noise and existence of non- correlational terms in the gradient computation. Very recently, [DTA 224] exploited higher-order terms in the gradient update arising from the reuse of the same training data, and showed that for certain link functions with high information exponent (p > 2). two-layer NNs may still achieve weak recovery (i.e., nontrivial overlap with 0) after two GD steps with ((d) batch size. While this presents evidence that GD-trained NNs can learn f. beyond the sample complexity suggested by the CSQ lower bound, the weak recovery statement in [DTA+24) may not translate to statistical guarantees; moreover, the class of functions where SGD can achieve vanishing generalization error is not fully characterized, as only few specific examples of link functions are discussed.
Given the existence of (non-NN) algorithms that learn any single-index polynomials in n Ö(d) samples [CM20] regardless of the information exponent p. and more generally, non-CSQ algo- rithms with sample complexity surpassing the CSQ lower bound [DPVLB24], is natural ask gradient-basee based training of NNs can achieve similar statistical efficiency for this function class. Mo- tivated by observations in [DTA + 24] that SGD with reused data may break the "curse of information exponent", we aim to address the question:
Can NN optimized by SGD with reused batch learn single-index f. beyond the CSQ lower bound? And for polynomial o,, can learning succeed near the information-theorees limit d?
Empirically, the separation between one-pass (online) and multi-pass SGD is clearly observed in Figure 1, where we trained the same two-layer ReLU neural network to learn a single-index poly- nomial with information exponent low test error using roughly n ~ samples, whereas online SGD fails to achieve even weak recov- ery with much larger sample size = l(dd). Our main contribution is to establish this improved statistical complexity for two-layer NNs trained by a variant of SGD with reused training data.
3. We see that SGD with reused data (Figure 1(b)) reaches
1.1 Our Contributions
We answer the above question in the affirmative by showing that SGD training (with the squared loss) on a natural class of shallow NNs can achieve small generalization eтror using polynomial compute and a sample complexity that is not governed by the information exponent, if we employ a layer-wise optimization procedure (analogous to that [BES+ DLS22, AAM23J) and reuse of the same minibatch. The core insight is that SGD can implement a full statistical query (SQ) algorithm that goes beyond CSQ, despite the correlational structure of the squared loss. Our main finding is summarized by the following theorem.

---

SGD + batch reuse [Thss work)
Smoothed SGD [DNGL23] CSQ lower bound [DLS22]
One-pass SGD
Kernel methods [GMMM21]
SQ algorithm [CM20]
[BAGJ21]
Information
theoretic limit
d (d)
((dp/2)
(
(dd9)
Figure 2: Complexity of learning single-index model where the link function a. is degree-q polynomial with information exponent For the CSQ lower bound, we translate the tolerance to sample complexity using the i.i.d. concentration heuristic T ss n 11/2, restrict ourselves to algorithms using polynomial compute; this excludes the sphere-covering procedure in [DPVLB24]] or exponential-width neural network in [Bac17, TS24].
Theorem (informal). A shallow NN with N O,(1) neurons can learn arbitrary single-inder models up to small population loss: E. algorithm (with reused training data) to minimize the squared loss objective, with a sample and runtime complexity of n, T (2(d(x--1))) where p., the generative exponent of the link a..
oar(1), if we employ an SGD-bas based
Note that the generative exponent [DPVLB24] is defined as the minimum information exponent of the link function O% after arbitrary L transformation, and hence by definition p= I p (equality is achieved by the identity transformation). We make the following remarks our main result.
We know that p. < 2 for arbitrary polynomial link functions. Therefore, the theorem suggests that NN SGD with reused batch can learn single-index polynomials with a sample complexity n Óa(d) which is information theoretically optimal up to polylogarithmic factors, hence matching the efficiency of SQ algorithms tailored for now-dimessional polynomial regression [CM20].
For non-polynomial a= with high generative exponent p. > 2, our sample complexity Z d can be interpreted as an SQ version of the online SGD result in [BAGJ21]. Since the information exponent can be arbitrarily larger than the generative exponent pa, our main theorem disproves a conjecture in [AAM23] stating that n x dp/2 is the optimal sample complexity for empirical risk minimization with SGD on the squared loss correlation loss.
A key observation in our analysis is that with suitable activation function, SGD with reused batch can go beyond correlational queries and implement (a subclass of) SQ algorithms. This enables polynomial transformations to the labels that reduce the information exponent, and therefore op- timization can escape the high-entropy "equator" at initialization in polylogarithmic time.
Upon completion of this work, we became aware of the preprint [ADK* + 24] showing weak recovery (for polynomial targets with p. < 2) with similar sample complexity, also by exploiting the reuse of training data. Our work was conducted independently and simultaneously.
2 Problem Setting and Prior Works
Notations. ||=|| denotes the & norm for vectors and the lg - Ey operator norm for matrices. O2(-) and oa(-) stand for the Oig-O and little-o notations, where the subscript highlights the asymptotic variable d and suppresses dependence on p, q: we write Ō(-) when
factors are ignored. Oa.(--) (resp. (-) represents big-O (resp. little-o) in probability as d - xc. (2.)), (-)) are defined analogously. 7 is the standard Gaussian distribution in R. We denote the - of function f with respect the data distribution (which will be specified) [l For g E R, we denote gf as its i-th exponentiation, and g() the i-th derivative. We say an event happens with high probability when the failure probability is bounded by exp(-C log d) for large constant C.
2.1 Complexity of Learning Single-index Models
We aim to learn single-index model (1.1) where the link function a, .RR. I has information exponent p defined as follows [DHI8, BAGJ2I].
Definition 1 (Information exponent). Let 1o denote the normalized Hermite polynomials. The information exponent of g E L²()),, denoted by IE(g) = p N4, the index of the first non-zero Hermite coefficient of g. that is, given g(z) Di (2)), := min i > 0 a #0).
By definition, when a, a degree-q polynomial, we always have p S q. Note that f. contains (d) parameters to be estimated, and hence information theoretically Z d samples are both suf-
3

---

ficient and necessary for learning [MM18, BKM* 19, DPVLB24]; however, the sample complexity achieved by different (polynomial time) algorithms depends on structure of the link function.
Kernel Methods. Rotationally invariant kernels cannot adapt the www-dimensional structure of single-index f. and hence suffer from the curse of dimensionality [YS19, GMMM21, DWY21, BES+ By standard dimension argument [KMS20, HSSVG21, AAM22], we know that in the isotropic data setting, kernel methods (including neural networks in the lazy regime [JGH18, 8, COB19) require n Z d9 samples learn degree-q polynomials in R4.
Gradient-based Training NNs. While NNs can easily approximate a single-index model [Bac17] 17,, the sample complexity of gradient-based learning established in prior works typically scales as n Z d in the well-specified setting, [BAGJ2I] proved sample complexity of n ((d"-1) for online SGD, which is later improved to () by smoothed objective [DNGL23]; for the misspecified setting, [BBSS22. DKL UXL+223 233] showed that n 2 dh samples suffice, and in some cases a (dd-11 complexity is achievable [AAM23, OSSW24a]. Consequently, at the information-theonr
limit (n x d), existing results can only cover the learning of low information exponent targets [AAM22, BMZ23, BES+23), 23]. This exponential dependence on p also appears in the CSQ lower bounds [DLS22, AAM22J, which is often considered to be indicative of the performance of SGD learning with the squared loss (see Section 2.2).
Statistical Query Learners. If do not restrict ourselves correlational queries, the sample complexity of learning (1.1) can be drastically improved. Specifically, for polynomial d., [CM20] gave an SQ algorithm that achieves low generalization error in ı = ((d) samples, which is near the informitioo--hheortiic limit; the key ingredient is construct nonlinear transformations to the labels that lowers the information exponent to 2; similar preprocessing also appeared in context of phase retrieval [MM18, BKM + 19]. Such transformations do not belong to CSQ, but can be utilized by full SQ learner to enhance the statistical efficiency. Recently, [DPVLB24] introduced the generative exponent which governs the complexity of SQ algorithms.
Definition 2 (Generative exponent). The generative exponent (GE) of g E L2(r)) defined as the lowest information exponent (IE) after arbitrary L transformation, that is,
P= =: GE(g) = inf IE(T g)).
(P))(P,)
The generative exponent the smallest information exponent obtained by all possible label transfor- mations. By definition we always have p" S p. and the gap between the two indices can be arbitrarily large; for example. for the Hermite polynomials we have IE(H),)
whereas GE(He;) S
[DPVLB24] established a sample complexity lower bound of n = ((d for full SQ learners with polynomial compute (assuming T = n- 1-122,, ), and obtained matching upper bound by a tensor partial-trace algorithm. Our goal is to show that SGD training of two-layer neural network can also achieve sample and runtime complexity that scales with ~ d((-)) where the dimension dependence governed by the generative exponent pa instead of the information exponent p.
2.2 Can Gradient Descent Go Beyond Correlational Queries?
Correlational statistical query. statistical query (SQ) learner (Kea98, Rey20] accesses the tar- get f. through noisy queries é with error tolerance т: - E y)l S T. Lower bound on the performance of SQ algorithm is classical measure of computational hardness. In the context of gradient- based optimization, an often-studied subclass of SQ is the correlational statistical query (CSQ) [BF02] where the query is restricted to (noisy version of) Emn,y (o(x))). To see the connection between CSQ and SGD, consider the gradient of expected squared loss for one neuron fao (æ:):
V En. (fu)) y)² x -- En.y V (x)) +
E, [f V (x))
correlational query
be evaluated without
One can see that information of the target function encoded in the correlation term in the gradi- ent. To infer the statistical efficiency of GD in the empirical risk minimization setting, replace the population gradient with the empirical average ( Σm (fw
) 1 y²)²), and heuristically equate the CSQ tolerance 7 with the scale of i.i.d. concentration error n -1/2

---

For the Gaussian single-index model class with information exponent p. [DLS22] proved lower bound stating that a CSQ learner either has access to queries with tolerance X d or exponen- tially many queries are needed to learn fa with small population loss. Using the heuristic r = n 1/2, this suggests a sample complexity lower bound n Z dp/2 for polynomial time CSQ algorithm. This lower bound can be achieved by landscape smoothing procedure [DNGL23] (in the well-specified setting), and is conjectured to be optimal for empirical risk minimization with SGD [AAM23].
SGD with reused data. As previously discussed, the gap between SQ and CSQ algorithms primar- ily stems from the existence of label transformations that decrease the information exponent. While such transformation cannot be utilized by a CSQ learner, [DTA + 24] argued that they may arise from two consecutive gradient updates using the same minibatch. For illustrative purposes, consider one neuron (x) = o((x, (az, w)) updated by two GD steps using the same data point (x. y), starting from zero initialization wt 0 (we focus on the correlational term in the loss for simplicity):
w² = w 2 + r7 yo' ((xr, w w=)))= no' (0) y y. +1 ny (1jo' (0) y).
(2.1)
CSQ
non CSQ
Under appropriate learning rate scaling n |xe 1 (1), one can see that in the second gradient step, the label y is transformed by the nonlinearity a', even though the loss function itself is not modified. Based on this observation, [DTA 224] showed that the non-CSQ term (2.1) reduces the information exponent to I, then weak recovery (i.c., nontrivial overlap between the first-layer parameters w and index features 0) can be achieved after two GD steps with n = (dd) samples.
2.3 Challenges in Establishing Statistical Guarantees
Importantly, the analysis in [DTA+ does not lead to concrete learnability guarantees for the class of single-indes polynomials for the following reasons: (i) it is not clear if an appropriate nonlinear transformation that lowers the information exponent can always be extracted from SGD with reused data, and (ti) the weak recovery guarantee may not translate sample complexity for the trained NN to achieve small generalization error. We elaborate these technical challenges below.
SGD decreases information exponent. To show weak recovery, [DTA+24 Definition 3.1] as- sumed that the student activation a can reduce the information exponent of the labels to 1; while few examples are given, the existence of such transformations in SGD is not guaranteed:
The label transformation employed in prior SQ algorithms [CM20] is based on thresholding, which reduces the information exponent to 2 for any polynomial aı; however, isolating such function from SGD updates on the squared loss is challenging. Instead, we make use of monomial transformation which can be extracted from SGD via Taylor expansion.
If the link function satisfies 2 2. its information exponent after arbitrary nonlinear transforma- tion is least 2; such functions are predicted not be not learnable by SGD in the n d regime [DTA+ + 24]. To handle this setting, we analyze the SGD update up to poly(d) time, at which a non- trivial overlap can be established by a Grönwall -type argument similar to [BAGJ21]. For p. 2, this recovers results on phase retrieval when a (z) z² which requires n M(d d) samples.
From weak recovery to sample complexity. Note that weak recovery (i.c., (w, 0) > for some small constant ε > 0) generally insufficient to establish low generalization error of the trained NN. Therefore, we need to show that starting from a nontrivial overlap, subsequent gradient steps can achieve strong recovery of the index features (i.e., (wx, 0) misspecification. After the first-layer parameters align with the target function, we train the second- layer parameters with SGD learn the link function o, with the aid of random bias units [DLS22].
E), despite the link
3 Learning Polynomial f. in Linear Sample Complexity
We first consider the setting where o. is polynomial with degree q specified as follows. Assumption 1. The target function is given as fx) ((x, 0)), where the link function d. IR + I. admits the Hermite decomposition o. (z) = Σi a, (e)
For single-index polynomials, we not expect a computational-et -osttttsstttt gap under the SQ class [CM20] indeed, we will establish learning guarantees near the information theoretic limit x d.
5

---

Algorithm 1: Gradient-based training of two-layer neural network
Input : :Step sizes n/: momentum parameters ç' training time Ti, T2; l2 regularization A.
Initialize w/ ~ 1 (1), aj )nif(
Phase I: normalized SGD on first-layer parameters
for t = to Ti do
if L even then
x ~ (0, Ia): y y= (.(r) +s;
1/ Draw i.i.d. data (m,
w/ w w + (w) w)²²) (when t 0);
// Interpolation step
w/ w/ |w0||:
// Normalization
end
w + w vw - y)? (J=1....,N)) end nitialize by Unil([-Ca, Ct)).
// SGD step
Phase II: SGD second-layer parameters
argmin T2
1 D (fe y:)² + X||a|1?
// Ridge regression
Output: Prediction function fe (x) with è (āj, ( 6,)3
3.1
Training Algorithm
We train the following two-layer network with N neurons using SGD to minimize the squared loss:
N az, w)) + bb))
1 N
fa(x)=
(3.1)
where C = (w., ag. by) are trainable parameters, and oj R + R is the activation function
defined as the sum of Hermite polynomials up to degree C.: (2) == E.e He, where C, only depends on the degree of link function d.. Note that allow each neuron to have different nonlinearity as indicated by the subscript in oy: this subscript is omitted when we focus on the dynamics of one single neuron. Our SGD training procedure is described in Algorithm 1, and below we outline the key ingredients of the algorithm.
Algorithm employs layer- wise training strategy common in the recent feature learning theory literature [DLS222 BES 22, BBSS22, AAM23, MHWSE23), where in the first stage, we optimize the first-layer parameters (w; } with normalized SGD to learn the low-dimensional latent rep- resentation (index features 0), and in the second phase, we train the second-layer {az {aا} to fit the unknown link function d..
The most crucial part in Phase of Algorithm is the reuse of the same minibatch in the gradient computation. Specifically, we sample a fresh batch of training examples in every two GD steps; this enables to extract non-CSQ terms from two consecutive gradient updates outlined in (2.1).
We introduce an interpolation step between the current and previous iterates with hyperparameter & to stabilize the training dynamics: this resembles a negative momentum often seen in optimiza- tion algorithms [AZ18, ZLBHI9]; the role of this interpolation is discussed in Section 4.2. We use a projected gradient update (w = (I " V ((w) for steps 2t and 2t+11 1, where V w is the Euclidean gradient; similar use of projection also appeared [DNGL23, AAM23).
3.2 Convergence and Sample Complexity
Weak Recovery Guarantee. We first consider the "search phase" of SGD, and show that after running Phase of Algorithm I for T = polylog(d) steps, subset of parameters achieve non- trivial overlap with the target direction We denote H (g: as the j-th Hermite coefficient of some g E L7()). Our main theorems handle polynomial activations satisfying the following condition.
Assumption 2. We require the activation function to a polynomial (2)) E (,He,(2) and its degree Co to be sufficiently large so that C. 2 C holds (Cg is defined in Proposition 6). For all 2 < l. & Co and k = 0, 1, we assume that H a (f) (1))-11;)) > 0.

---

As discussed in Appendix B.I, for a given Gx, the above assumption only needs to be met for one pair of (k. f). Appendix B.1.3 states that H (o(1)) ;k) * 0 also suffices if we set the momentum parameter & differently. Now we verify this condition for wide range of polynomial activations. Lemma 3. Given f 2 2 and & 2 O. For C. 2 2 we choose { { where Br is randomly drawn from some non-emne non-empty interval [a,, be]. then H (o k) + with probability 1.
The next theorem states that n = (dd) samples are sufficient for SGD to achieve weak recovery. Theorem 1. Under Assumptions and 2, for suitable choices of hyperparameteers of' =dd(Nd and 1 - 5 oa(1), there exists constant C(g) such that after Phase I of Algorithm / is run for 2T1,1 C(q) dpolylog(d) steps, with high probability, there exists subset of neurons w E W with W| = (N) such that (wn] 0) > cor some с 2 1/polylog(d).
Recall that at random initialization we have (w, d-1/2 with high probability. The theorem hence implies that SGD "escapes from mediocrity" after seeing n O(d) samples, analogous to the information exponent p 2 setting studied in [BAGJ21]. We remark that due to the small second-layer initialization, the squared loss dominated by the correlation loss, which allows us to track the evolution of each neuron independently; similar use of vanishing initialization also appeared in [BES 22, AAM23].
Strong recovery and sample complexity. After weak recovery is achieved, we continue Phase I to amplify the alignment. Due to the nontrivial overlap between w and 9, the objective is no longer dominated by the lowest degree in the Hermite expansion. Therefore, to establish strong recovery ((w, 0) > 1 E), we place an additional assumption on the activation function.
Assumption 3. Given the Hermite expansions a (2) = Ei=p He, 0,(2) Ei= (e,), assume the coefficients satisfy 0r 2 0 p S S q.
This assumption is easily verified in the well well-specified setting o, = a [BAGJ21] since a; = 3,, and under link misspecification, it has been directly assumed in prior work [MHWSE23). We follow [OSSW24a] and show that by randomizing the Hermite coefficients of the activation function, a subset of neurons satisfy the above assumption for any degree-q polynomial link function o..
Lemma 4. If we set a)(=) = Eir He,()), where for each neuron we sample B3g i.ddd )nr)((+r:)) with appropriate constant rú, then Assumption and 3 are satisfied in exp(-O(a))- fraction of neurons.
~
Note that in our construction of activation functions for both assumptions, we do not exploit knowl- edge of the link function a, other than its degree q which decides the constant Co; see Appendix B.1 for more discussion of Assumption 3 and Lemma 4. The next theorem shows that by running Phase I for (dd) more steps, a subset of neurons achieves sufficiently large overlap with the index features. Theorem 2. For student neurons satisfying Assumptions 2, 3 and parameter w starting from non- trivial overlap > 0 specified in Theorem 1, if Phase I of Algorithm continues for 2,2 (a(de-2) steps with hyperparameters 1f = =a(Nd-le), gL = 1, we achieve (un, 2(TT,1 O)
1 1-8 with high probability.
The following proposition shows that after strong recovery, training the second- -ayd-llaer parameters in Phase II is sufficient for the NN model (3.1) to achieve small generalization error.
Proposition 5. After Phase terminates, for suitable X 0, the output of Phase Il satisfies
Ex (x) - fz f=
with probability 1 as d + co, if we set To C(g) C polylog(d)= N = =CCl)pooyloo(())=" for some constant Cgg) depending on the target degree q.
Putting things together. Combining the above theorems, we conclude that in order for two-layer NN (3.1) trained by Algorithm to achieve s population squared loss, it is sufficient to set
N = C()) - - polylog(d), where constant C(g) only depends on the target degree q (although exponentially). Hence we may set -1 x polylogd to conclude an almost-linear sample and computational complexity for learning arbitrary single-index polynomials up to 0a(1) population error.
n = Ti + T2 = C(q) (de-2 e - ppolylooog)))))
7

---

4
Proof Sketch
In this section we outline the high-level ideas and key steps in our derivation.
4.1 Monomial Transformation Reduces Information Exponent
To prove the main theorem, first establish the existence of nonlinear label transformation that (t) reduces the information exponent, and (it) can be casily extracted from SGD updates. If we igmore desideratum (ii), then for polynomial link functions, transformations that decrease the information exponent to at most 2 have been constructed in [CM20, Section 2.1]. However, prior results are based on the thresholding function, and it not clear if such function naturally arises from SGD with batch reuse. The following proposition shows that the effect of thresholding can also be achieved by a simple monomial transformation where the required degree can be uniformly upper bounded.
Proposition Let g : > R be any polynomial with degree up to p and I19li;>() then
(i) There exists some i C. N+ such that IE(g") S 2, where constant Cgg only depends on q.
(ii) Let godd RR + R be the odd part of g with (lgodd (12)))) 2 p O. Then there exists some & Cg.p E M4 such that (F(g')) 1, where constant Ca.p only depends on q and p.
The proof can be found in Appendix A. We make the following remarks.
The proposition implies that for any polynomial link function that is not even, there exists some i E N+ only depending on the degree of a, such that raising the function to the i-th power reduces the information exponent to (this implies the generative exponent p. = 1). For even o., the information exponent after arbitrary transformation at least 2 (p. 2), which can also be attained by monomial transformation. Furthermore, we provide uniform apper-bound on the required degree of transformation i via compactness argument.
The advantage of working with monomial transformations is that they can be obtained from two GD steps on the same training example, by Taylor expanding the activation a'. Section 4.2, we build upon this observation to show that Phase I of Algorithm achieves weak recovery using n 2 d polylog(d) samples.
Intuition behind the analysis. Our proof is inspired by [CM20] which introduced a (non- polynomial) label transformation that reduces the information exponent of any degree-q polynomial to at most 2. To prove the existence of monomial transformation for the same purpose, we first show that for fixed link function o., there exists some í such that the i-th power of the link function has information exponent 2, which mirrors the transformation used in [CM20]. Then, we make use of the compactness of the space of link functions define a test function and obtain uniform bound on i. As for the polynomial transformation for non-even functions, we exploit the asymmetry of a, further reduce the information exponent 1.
4.2 SGD with Batch Reuse Implements Polynomial Transformation
Now we present a more formal discussion of (2.1) to illustrate how polynomial transformation can be utilized in batch reuse SGD. We let n = 17. When neuron fw((x) (x) = a(x,, w)) is updated by
two GD steps using the same sample (x, y), starting from w == w, the alignment with @ becomes
0, w (0. [w- 17- ((z, w )=)) = (0, w)+
n yww x:) (0, x)) + -Co-1 (n/|==|1) (()) "o' x)) (l)))) (0, (4.1)
We take 17 I c with a small constant C, so that 7/ æ|| < with high probability. Crucially, the strength of each term in (4.1) can vary depending on properties of the unknown link function Hence careful analysis is required to ensure that a suitable monomial transformation is always singled out from the gradient. We establish the following lemma on the evolution of alignment.
Lemma 7. Under the assumptions per Theorem 1, the following holds for p. = 1,
(0, 22 2(+11) > (0, 2² + cxeccod 2 Vx (x2 + Cr,ced" 2 v1 2t
8

---

See Lemma 16 for the formal version. For p. = 1, taking expectation immediately yields that weak
recovery within (17(1 - {))) -1 = O(d) steps. For p. 2. (0, w =: x² can be approximated dx'
=xx exp((((( E)-( x by differential equation dr r11 S)x''. Solving this yields x² = *x*
d i exp(n(1 - E)t), and weak recovery is obtained within & S (n<(1-)7) E)-) log d = Og d)
-
steps, similar the analysis in [BAGJ2I].
Why interpolation is needed. In our setting, the signal strength may not dominate the error from discarding the effect of normalization. In prior analyses for online SGD, given the gradient 9 and projection Pw Ia - ww the spherical gradient changes the alignment as (0, w(+1) T| (0, corresponds to the signal, and - v sufficiently small, the normalization error shrinks faster than the signal. However, in our case the signal shrinks at the rate of cf (recall that 1) = G,d='), improve the signal-to-non nonise ratio when the degree of transformation I is large. The interpolation step in Algorithm reduces the effect of normalization without shrinking the signal too much. In particular, by setting = ij. we see that the signal is affected by factor of ij whereas the normalization error shrinks by rf:; this allows us to boost the signal-to-noise ratio by taking ij small.
(0, -=" P +g/ 2 (0, +,00) - 11"119]] w²) + (h.o.t.), see [BAGJ2I, DNGL23J. Here
w²) comes from normalization. Thus, taking
=cnd ) and hence taking a smaller step may not
4.3 Analysis of Phase II and Statistical Guarantees
Once strong recovery is achieved for the first-layer parameters, we turn to Phase II and optimize the second-layer with & regularization. Since the objective is strongly convex, gradient- based op- timization can efficiently minimize the empirical loss. In Appendix B.6, the learnability guarantee follows from standard analysis analogous to that in [AAM22, DLS22, BES BES+22], where we construct a "certificate" second-layer a" E R that achieves small loss and small norm:
Ex (f((z) - + Σ (oos((w)), x) + b;) < ||"||
from which the population loss of the regularized empirical risk minimizer can be bounded via stan- dard Rademacher complexity argument. To construct such a certificate, we make use of the random bias units {b; ازرر to approximate the link function o, as done in [DLS22, BBSS22, OSSW24a).
5 Beyond Polynomial Link Functions
Thus far we have shown that for polynomial single-index target functions (which satisfy p. 2), SGD with data reuse can implement a polynomial transformation to the labels that reduces the information exponent at most 2: consequently, the trained two-layer neural network can achieve small generalization error with n = d polylog(d) samples. However, as shown in [DPVLB24] there exists (non moommmolynomial) o, with generative exponent p. (i.c., label transformations cannot lower the information exponent to 2) and thus not learnable by SQ algorithms in linear sample complexity.
Nevertheless, for a single-indes model with generative exponent p.. we know there exists an "opti- mal" label transformation that reduces the information exponent p.. If SGD transformation, then from the arguments in [BAGJ21], it is natural to conjecture that a sample size of d is sufficient. In this section we confirm this intuition by proving that SGD with data reuse (A (Algorithm 1) indeed matches this complexity. The following lemma is an analogue of Proposition 6 stating that polynomial transformations are sufficient lower the information exponent.
make use of such
Lemma 8. Given d, with generative exponent p. E N4- Suppose we can take an orthonormal polynomial basis {Ørbk for the space (Py) with inner product (f,g) = [x)= Then there exists some degree of transformation I E N, such that IE(o!) = pa.
We outline the differences and additional technical challenges to handle the GE(o.) > 2 setting.
For general 2 link functions σ., we can no longer make use of the compactness argument (see proof of Proposition 6) to upper bound the degree of monomial transformation. Hence in Lemma we do not state a uniform upper bound on the required degree I, unlike the polynomial setting.
Any link function with p. > 2 cannot be polynomial, and hence we cannot achieve low gener- alization error using a neural network with polynomial nonlinearity. We therefore need to use an activation function with universal function approximation ability.
9

---

5.1 Sample Complexity for Weak Recovery
We first show that Algorithm achieves weak recovery (i.e., nontrivial overlap with the ground truth 0) with complexity governed by the generative exponent of the link function p. GE( Similar to Section 3.2, we make use of randomized activation functions to ensure the desired label transformation is encoded we defer the conditions on the student activation to Appendix B.1.2.
Proposition 9. Suppose the link function o. has generative exponent Pa, and let I E N4 be the smallest degree of monomial transformation that lowers the information exponent to p. (i.e., IE(a!.) = p). We can find a student activation function a depending only on P.P. and I, such that if we take 17th, n12 = G,Nd' (add)))((())))
(2(t-1) = cqd tr- 1/) for small cyp. Cg = od(1), and set
(if p. = =1)
Ti,1 = CE
d(log d) dp.-1
(i 2)
(ifp. 2 3),
then if the initial alignment (w 2c,
4/2, there exists T. 1 Ti, such that for all 2 To.
2T 0 2 @(1) with probability oa()).
Proposition 9 is generalization of Theorem beyond polynomial a, (the proof of both results are presented in Appendix B.3,B.4), and can be interpreted an SQ counterpart to [BAGJ21]: we es- tablish sufficient sample size of n d<v--1)v1 for Algorithm I to exit the search phase, which is parallel the d(e-1)v1 rate for one-pass SGD (note that our rates slightly sharper due to logarithmic factors removed, since cE can grow arbitrarily slowly with d). For high generative ex- ponent with p. > 2, we no longer match the information theoretically optimal sample complexity n d, which is consistent with the computationalltosstaiisittt gap observed [DPVLB24).
5.2 Generalization Error Guarantee
After Phase I of Algorithm 1, we learn the unknown link function の, via training the second-layed To approximate non-polynomial functions, introduce ReLU component in the student nonlin- earity a (see Lemma 12 for discussions), and make use of the approximation result for the (uni- variate) ReLU kernel in [BBSS22), which handles general a, whose second derivative has bounded 4th moment. Combining the above, we arrive at the following end-to-end guarantee for learning single-index models with arbitrary generative exponent using SGD training of neural network.
Proposition 10 (Informal). Suppose the link function a, has generative exponent p- E N. and satisfies o o² E L (7),). For appropriately chosen activation function a (see Appendix B.1.2), neural network (3.1) with N = (1) neurons optimized by Algorithm I achieves small population loss [ff (x) - f. {x=)) (1,) with sample complexity of n ((d(r--1))))
See Appendix B.6 for the full statement with e dependence. This proposition confirms that the sample complexity for weak recovery (Proposition 9) is the bottleneck in single-index learning, as the total sample size required for Algorithm to achieve low test error also scales with -Ivv-
6 Conclusion and Future Directions
We showed that two-layer neural network (3.1) trained by SGD with reused batch learn single- index model (with generative exponent p.) using when the link function is polynomial, we established sample complexity of n (de--2) to achieve & population loss, which almost information theoretically optimal. Our analysis is based on the observation that by reusing the same training data twice in the gradient computation, non- correlational term arises in the SGD update that transforms the labels (despite the loss function not modified). We proved that monomial transformations that lower the information exponent of can be extracted by Taylor-expanding the SGD update; then we showed via careful analysis of the trajectory that strong recovery and low population loss is achieved under suitable activation function.
samples and compute; in particular,
Interesting future directions include extension to multi-index models [BAGJ22, BBPV23, CWPPS23), hierarchical target functions [AZL19, NDL23], and in-context learning [OSSW24b). Also, the SGD algorithm that we employ requires a layer-wise training procedure and a specific batch reuse schedule; one may therefore ask if standard multi-pass SGD training of all parameters simultancously [Gla23] (as reported in Figure 1) also achieves the same statistical efficiency.
10

---

Acknowledgements
The authors thank Gerard Ben Arous, Joan Bruna, Alex Damian, Marco Mondelli, and Eshaan Nichani for the discussions and feedback on the manuscript. JDL acknowledges support of the ARO under MURI Award
NSF CCF 2002272, NSF IIS 2107304, NSF CIF 2212262, ONR Young Investigator Award, and NSF CAREER Award 21 2144994, KO was partially supported by JST ACT-X (JPMJAX23C4), TS was partially supported by JSPS KAKENHI (24K02905) and JST CREST (JPMJCR2015). This research is unrelated to DW's work at xAI.
References
[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged- staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
[AAM23] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. SGD learning on neural networks: leap complexity and muddle-to-saddee dynamics. In The Thirty Sixth Annual Conference Learning Theory, pages 2552-2623, PMLR, 2023.
[ADK*24] Luca Arnaooli,, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows sgd learn high-dimensional multi-index functions. arXiv preprint arXiv:2405. 15459, 2024.
[AZ18] Zeyuan Allen-Zhu. Katyusha: The first direct acceleration of stochastic gradient methods. Journal of Machine Learning Research, 18(221):1-51, 2018.
[AZL19] Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? Advances in Neural Information Processing Systems, 32, 2019.
[Bac17] Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, (8(1):629-688) 2017.
[BAGJ21] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradi- ent descent on non-convex losses from high-dimensional inference. The Journal of Machine Learning Research, (:488- 4838, 2021.
[BAGJ22] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. digh-ddmensiioal limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Infor- mation Processing Systems, 11-25349-25666 2022.
[BBPV23] Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien On learning Gaussian multi- index models with gradient flow. arXiv preprint arXiv:2310.19993, 2023.
[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning single- index models with shallow neural networks. Advances in Neural Information Pro- cessing Systems, 35:9768 17733,8-7783, 2022.
[BES+22] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. Advances in Neural Information Processing Systems, 35:37932-37946, 2022.
[BES+23 Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, and Denny Wu. Learning in the presence of low-dinensional structure: A spiked random matrix perspective. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[BF02] Nader H Bshouty and Vitaly Feldman. On using extended statistical queries to avoid membership queries. Journal of Machine Learning Research, 2(Feb):359-33) -959999999999 2002.
[BKM+19] Jean Barbier, Florent Krzakala, Nicolas Macris, Léo Miolane, and Lenka Zdeborová. Optimal errors and phase transitions in high- nigh-dimensional generalized linear models. Proceedings of the National Academy of Sciences, (16(1):5451-546) 2019.
11

---

[BL20] Yu Bai and Jason D. Lee. Beyond linearization: On quadratic and higher-order ap- proximation of wide neural networks. In International Conference Learning Rep- resentations, 2020.
[BMZ23] Raphaël Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00 300.00053000 2023.
[CCMI1] Seok-Ho Chang, Pamela Cosman, and Laurence Milstein. Chemoff-type bounds for the Gaussian error function. IEEE Transactions Communications, 59(11):2939 2944, 2011.
[CM20] Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference Learning Theooy, pages 1161-1227, PMLR, 2020.
[COB19] Lenaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, 32, 2019.
[CWPPS23] Elizabeth Collins- -Woodfn, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the nigh-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. arXiv preprint arXiv: 089777 2023.
[DH18] Rishabh Dudeja and Daniel Hsu. Learning single-index models in gaussian space. In Conference On Learning Theory, pages 1887-1930. PMLR, 2018.
[DH24] Rishabh Dudeja and Daniel Hsu. Statisicall-oomputttinaaal trade-offs in tensor pea and related problems via communication complexity. The Annals of Statisties, 52(1):131-156, 2024.
[D3L+23] Yatin Dandi, Florent Krzakala, Bruno Loureiro, Luca Pesce, and Ludovic Stephan. Learning two-layer neural networks, one (giant) step at a time. arXiv preprint 1.Xi::2305.18270, 2023.
[DLS22] Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn representations with gradient descent. Conference on Learning Theory, pages 5452. PMLR, 2022.
[DNGL23] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index mod- els. In Thirty-seventh Conference Neural Information Processing Systems, 2023.
[DVLB224] Alex Damian, Loucas Pillaud- Pillaud-Vivien, Jason D Lee, and Joan Bruna. The com- putational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529 2024.
[4] Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborová, and Florent Krzakala. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. arXiv preprint arXiv:2402.03220. 2024.
[DWY21] Konstantin Donhauser, Mingqi Wu, and Fanny Yang. How rotational invariance of common kernels prevents generalization in high dimensions. International Con- ference on Machine Learning, pages 2804-2814, PMLR, 2021.
[Gl223] Margalit Glasgow. Sgd finds then tunes features in two-layer neural networks with near-optimal sample complexity: A case study in the xor problem. arXiv preprint arXiv:2309.15111 2023.
[GMMM21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Lin- carized two-layers neural networks in high dimension. The Annals of Statistics, 49(2):1029-100) 2021.
[HSSVG21] Daniel Hsu, Clayton H Sanford, Rocco Servedio, and Emmanouil Vasileios Vlatakis- Gkaragkounis. On the approximation power of two-layer networks of random relus. In Conference on Learning Theory, pages 2423-2461. PMLR, 2021.
12

---

[JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Con- vergence and generalization in neural networks. In Advances in neural information processing systems, pages 3571-8580 2018.
[Kea98] Michael Kearns. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM), 45(6):983-1006, 1998.
[KMS20) Pritish Kamath, Omar Montasser, and Nathan Srebro. Approximate is good enough: Probabilistic variants of dimensional and margin complexity. In Conference on Learning Theory, pages 2236-2262. PMLR, 2020.
23] Alireza MlussviiHosseenng Sejun Park, Manuela Girotti, Ioannis Mitliagkas, and Mu- rat A Erdogdu. Neural networks efficiently learn low-dimensional representations with SGD. In The Eleventh International Conference on Learning Representations, 2023.
[MHWSE23] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A. Erdogdu. Gradient-based feature learning under structured data. In shirth-seventh Conference on Neural Information Processing Systems (NeurIPS 2023), 2023.
[MM18] Marco Mondelli and Andrea Montanari. Fundamental limits of weak recovery with applications to phase retrieval. In Conference On Learning Theory, pages 1445-1450. PMLR, 2018
[M2D+23] Arvind Mahankali, Haochen Zhang, Kefan Dong, Margalit Glasgow, and Tengyu Ma. Beyond ntk with vanilla gradient descent: A mean-field analysis neural networks with polynomial width, samples, and time. Advances in Neural Information Process- ing Systems, 36, 2023.
[NDL223] Eshaan Nichani, Alex Damian, and Jason D Lee. Provable guarantees for nonlin- car feature learning in three-layer neural networks. Advances in Neural Information Processing Systems, 36, 2023.
[O'D14] Ryan O'Donnell. Analysis of Boolean Functions. Cambridge University Press, 2014.
[OSSW24a] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Learning sum of diverse features: computational hardness and efficient gradient-based training for ridge com- binations. Conference on Learning Theory. PMLR, 2024.
[OSSW24b] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Pretrained trans- former efficiently learns low-dimensional target functions in-context. arXiv preprint arXiv:2411.02544, 2024.
[Rey20] Lev Reyzin. Statistical queries and statistical algorithms: Foundations and applica- tions. arXiv preprint 20Xi::200400055,, 2020.
[Sch80] Jacob Schwartz. Fast probabilistic algorithms for verification of polynomial iden- tities. Journal of the ACM (JACM), 27(4):701-717, 1980.
[TS24] Shokichi Takakura and Taiji Suzuki. Mean-field analysis on two-layer neural net- works from a kernel perspective. arXiv preprint arXiv:2403.14911 2024.
[WWF24] Zhichao Wang, Denny Wu, and Zhou Fan. Nonlinear spiked covariance matrices and signal propagation in deep neural networks. In Conference on Learning Theory. PMLR, 2024.
[YS19] Gilad Yehudai and Ohad Shamir. On the power and limitations of random features for understanding neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[ZLBH19] Michael Zhang, James Lucas, Jimmy Ba, and Geoffrey E Hinton. Lookahead op- timizer: k steps forward, step back. Advances in neural information processing systems, 32, 2019.
13

---

Table of Contents
1
Introduction
1.1
Our Contributions
2
Problem Setting and Prior Works
2.1
Complexity of Learning Single-index Models
2.2
Can Gradient Descent Go Beyond Correlational Queries?
2.3
Challenges in Establishing Statistical Guarantees
3
Learning Polynomial in Linear Sample Complexity
3.1
Training Algorithm
6
3.2
Convergence and Sample Complexity
4
Proof Sketch
4.1
Monomial Transformation Reduces Information Exponent
4.2
SGD with Batch Reuse Implements Polynomial Transformation
4.3
Analysis of Phase and Statistical Guarantees
5
Beyond Polynomial Link Functions
5.1
Sample Complexity for Weak Recovery
10
5.2
Generalization Error Guarantee
10
6
Conclusion and Future Directions
10
A Polynomial Transformation
15
A.1 Proof for Even Functions (i)
15
A.2 Proof for Non-even Functions (it)
16
A.3 Proof Non-Polynomial Functions
17
SGD with Reused Batch
17
B.1
Assumptions on Link Function
18
B.2 Initialization..
22
B.3 Weak Recovery: Population Update
22
B.4 Weak Recovery: Stochastic Update
26
B.5 From Weak Recovery to Strong Recovery
29
B.6 Second Layer Training
14

---

A Polynomial Transformation
Proof of Proposition 6.
We use a thresholding and compactness argument inspired by [CM20].
A.1 Proof for Even Functions (i)
We divide the analysis into the following steps.
(i-1): Monomials reducing the information exponent. Define r(f) = max- -2r52 lf()), This entails that if [f((] rff), then we have |t| > 2.
Consider the following expectation:
f(t)
(12 -1-1)
(A.1)
2T(
We evaluate the case when is even. (A.1) can be lower bounded as
2T( ) (t2 (f- (1- 11
[(1)
(A.1) E.
1ll((t) 2 2T())
1(1) ;
)
+ E
[[T(f) s 11((1) < 27(11]
2T(f) 2 (t2 -1)
f(t)
+ E.wN(..1)
(12 -1)
(T(T))
(2(1))
2 B ((1(0)122(())
(22 -1)
2T(
(2((5) ) (2 -1)
S(I)
[(r()) s <0(00)0(0(00)) 22(00]
1
<
(<(0))++0000))) (T()
(2((f) (02 (02-1)
+ 1~
2 3P ((.,).1.)))) 2 2t( 2-
Note that (P[[(()) 2 2T( is positive (since f is polynomial) and independent of i, while 2- decays to 0 as increases. Therefore, for sufficiently large E N, (A.1) is positive and hence IE(f") 2. The subsequent analysis aims to provide upper bound on i.
(i-2): Construction of test function. We introduce the notation H (=:5) which takes any function
(in ( and returns its j-th Hermite coefficient. We consider the following test function:
DC
H(f::2)
2
(A.2)
(2
(i-3): Lower bound of test function via compactness. Let F,, be of polynomials with degree
up to g with unit 2 norm. Because HH (f) is positive for any f E Fy. (ff:)) is continuous with respect to f, and Fg is a compact set, inf feJ, (f) admits a minimum value Mo which is positive.
(i-4): Conclusion via hypercontracttvitt, Because f is polynomial with degree most g, Gaus- sian hypercontractivity [O'D14] yields that
2H( < ( [((())) & (2i 1)t2 (EL~ [F())]]* = (2i 1)1..
Therefore, for all polynomials in Fg. partial sum of (A.2) uniformly bounded by
H(f';2
2 2-111 = - oo).
27 -
Combining this with the fact that (f) 2 Ao > 0, we know that there exists some C, - 1 + 1og2(". such that
C,
H(
Σ
Ho 0.
12
for all polynomials in Fg.. This means that there is at least one i Cg such that H ( + 0.
15

---

A.2 Proof for Non-even F Functions (ii)
(ii-1): Monomials reducing the information exponent. We prove that some exponentiation of
g == f has non-zero first Hermite coefficient. Denote godd as the odd part of g. and similarly g Let v(g)= E R4 be the value at which the followings hold:
godd (1) for all 2 vg) and (t) for all t S -g)).
(b) g oven (t) ggod (t) for all 2 v(g) and S -gg)..
(c) For for all 2 v(g) and -v(g), g(s) gtt (as an equation of s) only has two real-valued solutions with opposing signs.
Such threshold v()) exists because the tail of f² is dominated by the highest degree which is even. Then, we let T(g) (()xs()))(ss((g) (9())). Consider the following expectation:
(
(gg(t)
g(t) (T(())
[.
(A.3)
(A.3) is decomposed as
(A.3) E, N(..1)
[llg(t)] 3r(f)]
(r(G) ) i
g(t)
$)
g(t)
+ EmA((.1)
1[27(g) s 1o(0) < <3r(0)]
4
3r(
+ EuAA((.))
[[lo(0)] < 2T(1
g(t)
1
(A.4)
3r(g)
We first evaluate the first term. Because of (c), g(t) 3r 3f( has two real-valued solutions < O B. Because of (a) and (b), ((3) geven (3) + godd (B) 3r( > g evee + godd
godd (-19). Because lim godd (1) t00, and is the only solution in i 0. have < 33 Moreover, for all t >> 8, we have g(t) (e)en(t) + godd(t) > geven (-t) + godd(1)) godd(-t).
Combining the above, the first term of (A.4) is bounded as
g(t)
E. EL
[llo(t)] 3r(f)]
3r(g)
g(()
g(() $ 1lt -a]
= -Ew((o)))
1[3 <t -al
t + IE (L
3r(g) t
3r(
3r(g) )²
g())
+ E
1lt 1lt s a]
1 -al
-0(((())
(((0))(((((-)))
-(((to) (9602) ) {
E.-((O.1)) [1 < 5 -a](] E
(3((g)
SP <t -a].
Following the exact same reasoning, we know that the second term of (A.4) is positive. Finally, the third term which is bounded by
g(t) 3r(g) ) t] 2 -Emwio,)) ( (T
2
E [1llo(t]] [All((t)))lll( (r)
Putting things together,
( 2+(0)]|0 (2)
2
(A.4) > OT> 1 sss -a]
The first term is independent of and positive, while the second term goes to zero as i grows. Therefore, there exists some such that IE(g":11 1) = 1.
16

---

(ii-2): Construction of test function. This time we consider the following function:
:(f (2
H(
(ii-3): Lower bound of test function via compactness. Let F.g be a set of unit n.rm polynomi- als with degree up to and Et-W( fyffodd (t)²] | 2 c. Since H
(f) (f) is always positive for Fu. (f) continuous with respect to f. and F, is a compact set, inf FEF, (J) has the minimum value Ho. that is positive. Note that .H6 (f) might depends on
(ii-4): Conclusion via hypercontractivity. Using the same argument as in (i), we conclude that there exists some Cg.e such that
1( 2 22(2i-11 1)
Σ
O.
Because Hy depends on e, Cg.e depends on as well as q.
A.3 Proof for Non-Polynomial Functions
For non-polynomial link functions, note that similar (DPVLB24), the existence of polynomial basis is needed to exclude extreme cases, and we cannot upper bound the required degree because general link functions are not included in a compact space.
Proof of Lemma 8. The derivation is analogous to [DPVLB24, Lemma F.14). Let ~ NO. 1) and y o_(2). We define Sp. (v) E v
1 Hep. (2) y) and its basis expansion Sp. (9)
Ex=o EN=OUEOR URōR. Let K be smallest integer such that Un * 0. Then, there exists an integer with I IS K such that IE(y') P.. Indeed,
El (l)x()) (z)] = E, 1 (1)
(HH,H (z)[vil
Ey Φ
(arbe(u) = uK + 0,
k==
which means that at least one y, y² k yields non-zero p--th Hermite coefficient.
B SGD with Reused Batch
In this section we show that Algorithm learns single-index models in ((d samples with high probability. The algorithm trains the first layer for Ti SGD steps, where sample data point in every two steps. The first layer training is further divided into two phases: weak recovery (w 0 2 1) and strong recovery (Ilw- 0|l S z). Then, we learn the second layer parameters.
Specifically, Section B.2 shows that at initialization, (nearly) constant fraction of neurons has alignment w TO beyond certain threshold. We focus on such neurons in the first phase of train- ing. Section B.3 lower bounds the expected update of alignment wt O of two gradient steps, and Section B.4 establishes that the neurons achieve weak recovery within 2Ti, ((d2v(p--1)) steps. Section B.5 discusses how convert weak recovery strong recovery using 2T,,2 Ö(de more steps. We let Ti = 2T,,1 + 2T1,2- Finally, Section B.6 analyzes second layer training and concludes the proof.
In the following proofs, mically. Specifically, asymptotic strength of the constants ordered as follows.
several constants, which depends d at most most polylogarith-
cl C2 poly(c,
c
88- "oll((c,))) c poly(ci S Scy"
1 co GS
C2
S polylog(d) = C4.
17

---

Here, Cyy and & should satisfy lims-pa Cyj = lim 8 = O, but the convergence can be arbitrarily slow, (e.g.. as slow as 1/ log log log log d). This requirement comes from the fact that we do not know the exact value of P.). To ensure that one signal term (from the Taylor series) isolated, taking 7 = d-1 with sufficiently small constant is insufficient but v = G,d-1 with arbitrarily slow Gy suffices. Also, to guarantee that the failure probability is oa(1), we require 6 to be (1). Cg can also decay arbitrarily slowly. as long as it satisfies cg A époly(c, )). C4 = polylog(d) will be used to represent any polylogarithmic factor that comes from high probability bounds.
For the first-layer training,
reduce the argument into training of one neuron using the corre-
lation loss as follows. At each step, the gradient update (Line 8 of Algorithm 1) is written as
w + w "7 v ((fe(z)- 7)
=w म V.
(w) x)) + 2n
Σ ==1
aaj x)
y
N
==1
w
2ny c
1
Σ
N
(w) T x) (v (wo T x)) +
Ca
N y (V ( x)) (B.1)
N
While the second term scales with 17' n/c2..N' N the third term scales with ry'c.NN1. Thus, by setting on sufficiently small, we can ignore the interaction between neurons. We will show that the strength of the signal in the direction of 0 is at least (x)))=-11 2 d Paj' (up to a polylogarithmic factor, and p. = )). On the other hand, we can easily see that 0 ( o (,(3) x)) (v 1f æ)) is bounded by Õ(1) with high probability. Therefore, by simply letting Cu = O "
) we can ignore the effect of the second term in (B.1). Moreover, for simplicity, we will reparameterize 2n'c.
as nf below. Consequently, we may analyze the following update
w
+ w + 17 'V (ya, x)).
instead of Line 8 of Algorithm 1. Since there is no interaction between neurons now, we omit the subscript j when the context is clear.
B.1 Assumptions Link Function
The analysis consists of three different phases: weak recovery and strong recovery of the first-layer weights, and approximation of the link function (ridge regression of the second-layer). Each phase requires different assumptions on the activation functions, depending on the link function. Before starting the analysis, we decompose Assumptions 2 and 3 and clarify which conditions are needed in each phase. We prove that instead of using specific activation function tailored different link functions, a randomized activation function satisfies all required assumptions with probability (2(1).
In the following, we write the student activation function as
(s)
By. (s)
with coefficients ) (sometimes the subscript j. which is the index of the neurons, is omitted).
B.1.1 For polynomial link functions
In the following, we summarize the precise conditions to be satisfied by the activation functions (these conditions are weaker than Assumptions 2 and 3). For polynomial link functions, we focus on polynomial activation functions (with bounded degree) for simplicity, but non- non-polynomial activation functions would not change the proof significantly.
Let p and q be the minimum and maximum degree of non-zero Hermite coefficients of a.. Note that GE(a.) = 1 or 2 holds (see Proposition 6). Let I < C2 (according to Proposition 6) be the smallest integer such that IE(o!) ) = GE(a.) = p. and C₂2 be the degree of the activation function.
(I) If + IE(o.)) = GE(aa) P.- Weak recovery: ap. 3p. (covered by Assumption 3).
18

---

Strong recovery: Ej=p. ر > 0 for all > 0 (covered by Assumption 2).
Approximation (ridge regression): B; + for some 2 q (covered by Assumption 3).
(II) 2 <== {min i [I(([)) = GE(o..) = p.) - Co.
p. 1) > 0 (covered by Assumption 3). Strong recovery: Σ) 1 1 > 0 for all > (covered by Assumption 2).
Weak recovery: (((()))
Approximation: B, # 0 for some 2 q (covered by Assumption 3).
Note that it is difficult to construct a deterministic activation function that satisfies all of the as- sumptions for any link function d. (the simplest counterexample is to consider -0 which flips the Hermite coefficients). Instead, we show the existence of randomized construction of such an activation function that satisfies all of the assumptions with constant probability, which entails that a subset of neurons can achieve strong recovery. The construction does not depend on properties of the link function itself except for its degree q.
the activation function simultaneously
Lemma 11. There exists a randomized activation function sampled from a discrete set such that the above conditions hold with constant probability.
Proof. Let e be a sufficiently small constant only used in this proof and Co be the minimum odd 1, + 2,3), where Cg introduced in Proposition 6. With proba- bility 1. we let 3 ~ Unif([+1]) and Bj Un))((=c]) for S C.. With probability 1. let Bj ~ +nif({+t)} for 1 S j < C - 2 and Bc.-1 = Sc. ~ (nni((+1)))
integer with C₂ 2 max
We first consider (I), When 31 ~ Unif([=1]), and B, ~ Unif({=e}) for 2 j S Co. it is easy to see sign(a;) = 3ig)(3,) for all j = 1, q hold with probability at least 2 which is sufficient
to satisfy (1).
We then consider (II). First focus on the case when p. and is even. When 3 ~ Unif((+1])
and Bj Unif((=c)) for 2 S j S Co. by taking c sufficiently small, we have
(000 0) 11,,(
(B.2)
When is even, by adjusting the sign of B1, (O 0) non-zero and has the same sign (((a))(;1) with probability z.2 Note that the sign of 31 is independent from whether holds. This holds with probability at least 2 2-++1 Thus
Elay,
> 0 for all we verified (II) for p. 1 and even T.
For = 1 and odd I, consider B, ~ Unif((+c)) for S j < C. 2 and Bc.-1 Sc, ~ Unif((+1)). Note that - fla,,, > 0 for all > 0 (this is the condition for strong recovery) and the condition for ridge regression also holds. Furthermore, the sign of H((Hec, + Hec,- + (1) 0) is #1 with equiprobability, independent of S2..... ... Therefore, by taking sufficiently small, we can obtain the desired sign of (())
Thus we proved (II) for p. = and odd I.
Regarding (II) for p. 2 and even I, when 31 Unif((+1)) and B3 Unir((+te) for 2 م Co, we have
(0 = (I+ 1)48r41 (1)))
Thus, similar to (II) with = 1 and even I, we get (II) for p. = and even I.
Finally, consider (II) for p. = 2 and odd 1. When 33 ~ Unif((+t)) for 1 < < Co - 2 and Sc,-1= = 3c. ~ Unif({ the sign of H( + ( +1 with cquiprobability when / is odd, and this term dominates the others in I/ (( (o()) Thus, (1I) for p. and odd / holds similarly to (II) for p. and odd I. Now have obtained the assertion for all cases.
+ Hec,,-1 (1)1-11; 1) is
B.1.2 For general link functions
Now we consider non-polynomial link functions with potentially large generative exponent p. GF(a.) 2 For weak and strong recovery to succeed, the conditions on the activation function are essentially the same those for polynomial link functions:
19

---

(I) I + IE(o,) = GE(O.)) = P..
Weak recovery: ap, Bp. > 0.
Strong recovery: Ej=p. 1101885s-11 for all 0,
(II) If I-= {min i IB(a,)) = GE((o.) = =
:p. 1->1>00 0. Strong recovery: -) 3,,, 1- 0 for all 0.
Weak recovery:
Due to the proof strategy (which uses Taylor expansion), we also require that all differentials and sum of expectations appearing in the following proofs are well-defined and bounded.
To approximate non-polynomial o. we introduce the following condition on the activation func- tion. We sample Jj from a discrete set (with bounded cardinality). Let J be an index set such that the coefficients of o; (j J) satisfy the conditions above. Because we are selecting σj from a discrete set, [J] = N holds. We introduce the following condition, which states that the target single-inder index model can be well-approximated by a linear combination of student neurons.
Assumption 4. When bj Unif([-Cb, CtJ) where (C, polylog(d)) and æ1, 3T2 (0, Ia). there exists a set of coefficients 01, ......aaaaa such that
(
- To Σ JEJ
Σ
a,, xb) bb) o.((Tai)) -..00 SE?
T2
holds with coefficients of reasonable magnitudes Eje, az = =(([I) with high probability (w.r.. the randomness of bj and wi). Moreover, E. (o, x b s polylog(d) for all j with high probability (w.r.. the randomness of bz)-
The following lemma states that we can design randomized activation function that satisfies all of the above assumptions with probability (2(1), long as the link function satisfies Assumption for a = ReLU. In other words, we are able to cover the class of link functions that can be efficiently approximated by two-layer ReLU network. Since the general link functions are not included in a compact space, we do not have an upper bound of exponent to obtain IE(o!) = GE(a,) as we had Cg in the polynomial case. Consequently, our student activation not entirely agnostic to the link function o., as we require knowledge of p (information exponent), P. and I.
Lemma 12. Suppose the target link function a, satisfies Assumption 4 for oj = ReLU. There exists randomized activation sampled from discrete set such that the above conditions hold with constant probability.
Before we sketch the design of activation function, we present the following approximation result from [BBSS22], which establishes that Assumption 4 with oj ReLU is satisfied for broad class of functions, according Lemma 4.4 and 4.5 of [BBSS22]. Specifically, taking 1/2 and À = -1 yields that E.r [ ΣjE, a + b,) a. (0 )² < - Although they sample by from o.(0'
Gaussian (00, 2). the result translates to uniform sampling of biases from [-Cb, C&J by introducing additional logarithmic factor.
Lemma 13 (Lemma 4.4, 4.5 of [BBSS22)). Suppose that Ex-. E (o) 00. Then, Assumption 4 with oj ReLU holds with N and C. = vlogd.
Proof of Lemma 12. We show the existence of suitable a two steps: first we construct a randomized polynomial activation function that satisfies conditions (I(I) with constant probability; then we add a small ReLU perturbation so that the activation can approximate non non-polynomial o...
Recall p E N4 is the information exponent of o.. We first show that there exists a randomized polynomial activation that satisfies the conditions for weak and strong recovery with probability 2(1). Note that the issue of differentiability and bounded moment avoided when we focus on the polynomial activation functions. We specify the following two distributions. With probability 2' let 31 Unif([-1,
By ~ Unif([- c)) for j = 1.... ,P. + 7 and 3, = otherwise, sufficiently small constant. With probability 1 let 81 Unif(f-1, 1)), B2 2.
where c > Unif([-c,e)) c)), B = Unir([_e2, (e)) for all 2 S j S (p- + T) V p for sufficiently small constant > 0, and B, = 0 otherwise.
20

---

Regarding (I), consider the case when the coefficients are sampled from the first distribution, and |3,1 except for p.. Then, -j=p. 3,,) = p.lap, Sp. s", Choosing the sign of Sp., we have that the assumption holds with probability ((1).
Regarding (II) with even L, consider coefficients sampled from the first distribution, and (3)) (ign(a;) for (p- + I 1) V p. Then, 2ijpp+1 Bla, for all 0. Also, similarly to (B.2),
(0 p. 1) =0111+.. -11))+9-1(6111 (81))-+0((0))
x
By flipping the sign of 31. we can change the sign of H I is satisfied by a randomized choice of 31.
1). Thus, (II) for even
For (II) with odd 1, consider coefficients sampled from the second distribution, and Sign(3,) = Sign(a,) for j S (p. + +1) Vp. Then, D3 p. +1 3ر for all > 0.
H (0)0) p. 1) (p= - 1)!
Blo (1) (o Hep.
1
(p. - 1)! E[(I 11)113(ppppp + Hep. H)p (B, (B) -1l -1]+ Occc).
2(1 p. 11((11))) -2 (p. + II! +0(c)).
(p. - 1)!
x
By flipping the sign of 31, can change the sign of H (((d))()) :p_ 1). Thus, (II) for odd I is satisfied by a randomized choice of 31.
Therefore, we have constructed a randomized polynomial activation a that satisfies all of the con- ditions for weak and strong recovery. Now we provide sketch of reasoning that when the link function a, well-approximated by ReLU as specified in Assumption 4, we can find some that additionally satisfies Assumption 4 by introducing small ReLU component. Specifically, we add en ReLU to the activation function with probability 2. with sufficiently small CR 2(11, e.g..
en (logd) -C for some C > O. When a two-layer ReLU network approximates a. that satisfies Assumption 4, by using the neurons with added ReLU component, d, can be approximated up to some polynomial residual with degree (p. I) V p. And by using the remaining polynomial neu- rons, we can approximate the additional polynomial terms in G. (see Lemma 22,23). Subtracting the latter from the former, we obtain the desired approximation result. When eR is sufficiently small, this additional term does not impact the conditions for weak and strong recovery and the moment calculations; similarly, since CR <> 1 we may discard this non-smooth term before Taylor expansion without affecting the analysis of optimization dynamics. We remark that to avoid such unnatural de- sign of activation function, we can also train the first-layer parameters using a polynomial activation specified above, and then perturb it before the second-layer training to enhance the approximation ability- - such strategy has also been employed in prior layer- training analysis [AAM22]. O
B.1.3 More Discussion on Assumption
Assumption 2 requires ((o()(((0)(1 1) is not zero and has the same sign H(o! p.)-
We remark that if we allow negative momentum parameter larger than 1, i.e., setting g2(++1)
we can negate the opposite sign of H ((())(()))))) the subsequent analysis still holds. Therefore, what we essentially need is H (o())-- kk) # 0. Lemma 3 confirms that it is satisfied by almost all polynomials:
1 + cge
(2+
1) (see Lemma 16), and
Proof of Lemma 3. We note that H( (o())==1 k) = (o is polynomial of (3, ( This polynomial is not identically equal to zero. To confirm this consider = C= +
-C Because (00) (1))) expanded as a sum of (i(CC 3) < 1 < i(C, - 2) + 1 with positive coefficients and each x' is sum of Her, He:-2 with positive coefficients, R (1)) has all positive Hermite coefficients for degree 0, 1,.... O,,, - 2) 1. If i(C. 2) + 1
21

---

this choice of a yields H( ((0(()()))) > O, which confirms that (1 as polynomial of { is not identically equal to zero. Hence the assertion follows from so-called Schwartz-Zipper Lemma [Sch80], or the fact that zeros of non-zero polynomial form a measure- zero set.
B.2 Initialization
We first consider the initial alignment. In the following sections, we focus on the neurons that satisfy x3 0 w 2 2c,y1 ' the initialization. The following lemma states that roughly constant portion of the neurons satisfy the initial alignment condition upon random initialization. In d)=
particular, if we take Cy (2(( log ), the fraction of neurons that satisfy the initial alignment condition at least e -16e, = (11.. Let us write C2 = Ciy for simplicity in the following.
Lemma 14. At the time of initialization, ki = @ w0 satisfies the the following:
P(e 2 2C2d = P[xx s 202d 2 1602 52(1).
We make use of the following lemma. Lemma 15 (Theorem 2 of [CCM11]). For any B > 1 and R, have
/ /2eBB - 1)
de3
00
e 2 dd
28
2t
Because x²' = vT Tw d cl ||l
Proof of Lemma 14.
where g ~ NoO,
Plaj 2 2C2d b1 = P gr(((
ei g 2 4C2 A || - 2d
2 B gr.\((....)
e 2 4C₂ - PgnN((.La))
||9|| 2 2d
22 -1)
8303
((d))
23
where we used Lemma 15 for the final inequality. By letting B 2. we have that P[N2] Czd $1 1603. Because of the symmetry, P(x]] < 2C2d $] P(x.]] 2 2O2d $).
B.3
Weak Recovery: Population Update
We divide the first layer training into the first phase (weak recovery) and the second phase (strong recovery). We first evaluate the expected update of two gradient steps with the same training exam- ple.
Lemma 16. Let 2 .1 2,2111 = - G,d g2((+1) - ced 120-2)1
Suppose that the link function satisfies TE(o!) GE(a.) p. (we choose the smallest such I) and activation functions satisfy all of the assumptions in Section B.I for weak recovery. Then, for 2 with c
c-ld S
0 2t < cpp the alignment 0" 2(++1) can be evaluated as,
0 w 2(++1)
0 2 Tw + df (12 + cr.ced 4= -1
Here Co = =paaap. Sp. (when IE(a,) = GE(o.)) or co = Ma (a(0)) ip (other- wise), and 2tt a mean-zero sub-exponential random variable.
2(1-131
Proof. The expected alignment 0 (t++1) after two gradient steps from 2 using the same
sample (z, y). step size 1 2 c,d- and momentum parameter e2(c+1)
22

---

1- cgd is evaluated as follows. With a projection matrix Pw = ww., the first step updates the weight as
w 2++1 + 2 + V wo (w x) = + Tyo' (wTx))) Pw
(B.3)
and the next gradient step with the same sample is computed as
V ww 2 x) yo' vw'ww" 2 x)xx:
= yo ((w + +nyo (w T x)) Pwix) Ta) Pw
= yo yo'(wl x + 7 x p d'(w] (e)aa)y))) Pwx.
(B.4)
here we used the notation ||0]|2 0 AO for a vector 0 E R. and positive symmetric matric A E Rdddd From (B.3) and (B.4), the parameter after the two steps obtained as
w 2(++l)
w 21+1
+ vw 2 x))
w nyo' Pw) + nyo' (w x + ع|l d'(w" (wwTT)y)) Pwx.
= w + ng" 21
where
-²t = = (w T P x + yo (w" x + n|| o' dww (wwTTa))) x)y)) Pwx.
Finally, the normalization step yields
- (w) (w+ 2²t
+ وو
w + Cy Cn 3 V2
w 2(41))
[ - ( 2
||ا + nggク'
|| + Cyreed - 2/ v2
Therefore, by writing 0 2 2t the update of the alignment
2(+1) -T 2(++1)
22 + c.reet 411 aT g2t
|| + c. 4 v1
v1 2
2 22 + cod v 0 T g2t
k 2 2
22 c
32 v3 vs10 """"
(B.5)
2
We can easily see that 2 S d and E o
2| S d. which implies that the expec- tation of the last two terms of (B.5) is bounded by S x²1 2dd v 2d 1)v2 <
dd 1p.-1 21
Now we bound F[ by 2 C) d/-1g.-11, Let C. be the maximum degree of the activation
function with non-zero coefficients of Hermite expansion, which may be infinity when we consider general link functions, and there appear some infinite sums. For these cases we simply assume the sums converge we discuss the validity of this condition in Section 1.1.2. We omit the subscript 2t in the following for simplicity. We divide the analysis into the two cases.
(1 T + IF(..)) GF((o)) Pa-
For the first term of E(0" T g), we have
C.
0 E Elyo'(as (w T Pw = P
He (0 x)
j8, He, - (w T xx) )z)
jmp.
80 =0 Pw Σ
,la,3, (0T) 1- w-) 0+ (j + 2) 30 (Tw)/w
j=p-
Co
Σ 3la,3, w)) -- 0 P.0
j=P.
P. 3p, p1
11 + O(PP.).
(B.6)
23

---

For the second term of E g). the following decomposition can be made.
Elyo'
(w x + +n|= o' d'(w x))y))))))) P
C
(i)) 10)000 0
P. you (
(w))) 20) ( yللر سw x;)² x0 001 D'Rlyo'(u (w=)) P
C.
(11)~1100 Pu | P 1ا (litt) (w))) (o())( (w)) '
(B.7)
+ pp Sp. 8 + O(KP.).
(B.8)
We evaluate each term in the summation. We need to show that although || is a function of x and w x, it is mostly independent from the two quantities. To verify this, let e be the orthogonal component of 0 to w. Then, we have that
0-10 wiw
0 wjwi
| x (I ww ce ) + (e (eTxx²²
10 x (0 ww) To
x' (I ww ee T )ar
+
100 (0 w)w/|
~ independent
and
With Pu.0 = La ww T ee (B.7) expanded as
-1 8
(++))
(B.7)
i!! 0 (0 12
i=1
1=0 1=0
PuR E.
(x: T ) it (0. x (OT w) ())(()))==))))) x)) (o (w)))(() ==))' (w))))
C.-1 2
(-) B [2
ΣEΣΣ
i - (0 w/
E (o.(0Tx)))
-www (1+) (w x.) (al T ())) x x- o ww ax) (B,9)
For a general differentiable function g(x), we have (()e(()))((z)) E Bid g(x)]. If g(x) is polynomial (with a bounded coefficients) of 21 and u x and its degree with respect to 21 is at most #(S t), (Z) )x)) S |ua] 1-= because differentiation of g(x) = g(x1, u x)) is taken with respect to the first variable at most 8 times. Each term of (B.9) is an expectation of (a. (oT)))))) x))), multiplied by the polynomial of 0² x and w x, where its degree with respect to 0 ar is at most 2j k. Thus each term of (B.9) is evaluated as (here we omit the constants)
n (-1)²E
(0 m)) m) "ww o (w a)) (o (T)) (0T x 0 @www ww)) xr)
E
0 (0 w | 21
1F2P.
degree w.t.. 8 most 2j
S S
-2j+k-1) (---j----o z cnd Sg(l. - 1)- S co -d/x2)) S c
The lower bound follows in the same fashion. Therefore,
(B.9)] S c.p
Now, E g] can be evaluated as 0 T g] (B.6) + (B.7) (B.8) 2p. Pp.lap..pppht kp- + Ole + kP)).
24

---

(I) If / = {min i IE(a')) GE( = p-} 2 2. Note that a, 0 for all j S from the
assumption. Following (B.6), the first term of E g) is evaluated as
0 Elyo (w Pw
,a,,
P =
(B.10)
j=p
For the second term of E[O" g). similarly (B.9), the following decomposition can be made.
θ T Elyo' (w
x + n|| d'(w"
z)y) P
(OT) (-1)
J x2_ [2t
ΣEΣ = j=0 100 Am
0 w 2j
E (o. (TT) (0T
2)2 T z)) ((++)) (w x) (o! x)) x))'((1 (0T a -www x))B.111
Each of (B.11) (omitting constants) is evaluated as
n''hh TReBB 2 1 (0' (0" 20 - (wo o (Tt)) (0T x o ww ww)x)) x))
(B.12)
cg d B
(o. (0 (o.()x))))
(0 2)) ( (w (w"x) (o( (0
(0T o ww Tw)
P+ (1 2 1)
degree 0 aost
IE2
p. + 1 < 1)
S (BB))
(B.13)
When I and = c,d we have / < and (0T(o,(0Tx)) 2 p. + Thus
(B.13) X x
When 2 1, 1E(((((0 xx)²')))) 2 and we get
(B.13) S c, x pー
Now the case of = I - 1. When = -- and * 0, and using the assumption that k S Cy.
(B.13) $ c kP -1 (m /d) S c
J-1
When = 1, = 0, and + 0,
(B.13) S c d-1
When i i= - 1,j 0, 0, and 1 1-
(B.13) S c
1-1 1-1
Therefore, except for = 1 0, k = 0, and < - 1, we can bound (B.13) by S of RP kP. The lower bound follows in the same way. Finally, consider the case of = 1. j = 0, k 0, and = I.
(B.12) = - 1 (o. (0 x:) ((+1)) (w) x)) (w Ta)) 1-1
10T x '0
0 (ww
C,1
11
2 1-1
1)(1 n:3 m-
X
25

---

= -pld (1-1) smx p
1)(11 - x2 + 0(ep-118))
Putting it all together (recovering the constants omitted in (B.12) again), (B.11)
cn J-1 Pold (1-1) 1) x~~ (
H(o (( -1) OcG + +x(²):
(1
and
E g] = -B.10) +
1-1 (1-1d((-11)B_...
H(o..p.) (00 11) Occypei- xP.).
"
(T 1)!
(())
Combining (i) and (ii). we have
Ei0 g) 2 Conpo + +ccc kPt)
for positive constant Co = ((1) Here Co > satisfies 2co 2p_lap, Sp. (for (i)) or
2cr ( -1)-1))
(for (ii)). Going back to (B.5), by setting 22 = (0 g2
(1-1))
E I), we have
((21) 2 22 2c V VV gdd + credd (O g 210Tg20 + Ole -1] -1]v1
2x 22 7 v co (121)pp0 + creed " vi
(p. (p.1)V (1. + 1+l d+tced 4 (n + creed 4 V1 V (n
When cg < of and cnjy d S s cp terms in the big-O notation is smaller than cleed & co (1221)p0 and we have
(1++) 2h + ef 2 (x=*tp- + cyrced vi,,
straightforward to check 21 has sub- Weibull tail.
B.4 Weak Recovery: Stochastic Update
This subsection proves weak recovery using the results on population update from the previous section. Specifically, from the previous section, we know that
2 ((++)) 0 u +
+ (x 21
with the mean-zero sub- luibuli random variable 2² and positive Co @(1). For notational simplicity we write c/ = C1. The following lemma is detailed version Proposition
c,d (2(++1) = = 1- ced (Pa-2)+ Suppose that the link
Lemma 17. Take 112t n7
function satisfies IE(o!) GE(o.) = p. (we choose the smallest such 1) and activation funcctions satisfy all of the assumptions in Section B. for weak recovery. Let
(ifp. GE(a..) 1)
Tı =.1=Ca Cac5
d(logd) d
(else pf = GE(a.)) =2)
(else GE(o.) 2 3),
and take eg S 6poly(c,), 02 Z poly(cn), and C3 = G If " 2 20, dd 4, there exists some
S STi, such that
2r. > 202
with probability at least 1 - 8, and 2r 2 20g for all S Ti,1 with high probability.
26

---

We may take 5 = (1) with arbitrarily slow decay. The proof is adapted from [BAGI2I], but our bound on Ti, is slightly sharper (by a log factor for p. and by (log d) factor for P- 2 3). For p. = 2, this is because of a trick that we carefully "restart" the dynamics, whose failure probability exponentially decays.
Proof. We divide the proof into the following cases.
(i) When p. = 1. Note that (Ex=0 1²=]r is Martingale with B(((22)))
S 1. By Doob's maximal
inequality and Markov's 's inequality, with probability 6, we have
2
T
T
max SSTST
Σ ==0
2a
8
Σ
C18
(B.14)
==0
==0
for any fixed T 2 O, with a sufficiently large constant C₁ = (1). In the following we consider the case when (B.14) holds for T = = - id
If of tt S 2d S o for all & = 0, 1, ................................ T, we have
2(++)))
2 2 + ared + cared
> 2 + ccq((T 1)d -
Σ 22 ==0
(B.15)
Now, applying (B.14) to get
1++ 2 (B.15) 2 20 d (Tc + -) + 8 2d
when
1 c c
d - 1. By letting cg S c C C 8, we have o7 d s cn - c o
c
and
2 2 Gg -d + +
which verifies 1d 1 < 20 for = + 1. Thus, there exists some T. Ci cg t such that
ne 2r 2 402,
for e S tch, with probability 1 1 8.
Now we prove that x² 2 20g holds for all T. St Ti,1 = C3c Because 2/ are mean-zero
sub- Weibull random variables, we also have that Σ T44' 2 < CIVi'' for all 0 < T. < s=T
Ti.1 with high probability. Also, because n/ < d and || - Ld < 1, we can easily see that (x+((+ x O(d=) for all T 0,1, Tí, 1, with high probability. Thus, when there exists 2 J+ such that ((( > 4c2 and 2² < 4c2, we have x²² 3c2 with high probability. Moreover, following the above argument, we can inductively show that
x² 2((++)) 2 30, + aoce'dd
CaV7
(7' S Sc 22 202 2T)
2 3c, + ci di
02
(1ceT'd
(7' 2 c ci)
2 202,
for r' S Ti, ==,1= d or until n ((++') > 4c2 holds again. By repeating this argument (if there
are multiple such T), we see that 2t 2 202 holds for all T. S L S Tı,1 = Cac.'d with high probability.
(ii) When p. p.=2 2. We define to = 0,41 log 2)(4),12 2 2lo 1y(4).... We show that, for each if n.2 2 22-11 1 1, have 2(((+)) 2n.". with probability at least 1 - 84- or there exists some t (4 < S 4+1) with 2! > ch.
Assume that the above statement holds until some i 2 (we do not need assume anything for
i = 0). Then, we have 2 2 2 2 2c, 2 Similarly to (B.15), if c, - S 2t S c, for all
f = + 1, , T, we have
((++) > 24 + ciced
Σ
20 c,ced
Σ 22
8=4
27

---

Applying (B.14) with d = 8/4" and T = ic
(6/4 - to get
((+1) 2 2 + areddt Σn - - v
x + crced
Σ xza =
a
when S
15 / (0/4 (x/4
2 - 1, which verifies 77t S 2,2m S x2t for
1 T +
This implies that, with probability 1 5/4", we have
1 a0 crcedt
2
2 2"
for all r 2 (6"4) 2 ) 1, which is equivalent to
2e; x²² 2 (1 + ciced )- k
for all 4,4 + 1,...,,, ,4 + to 20 (. 2d. c1,, (8/4())))(1122 we have 2 (x- 2 log log(++ared-11(4) and we get
By taking og
21 2 2k. 21
with probability 1 - 5/4 (or there exists t S 4++ such that 2 > cl))
Thus, by induction, for all i, we have that
24 2'x
(B.16)
or that there exists some t S 4 such that 2 larger than c/ with probability 1- 8.
The LHS of (B. 16) becomes larger than of for some S log d. Because 4 O(ic ce 2d) within -1 (c aed log d) steps, there exists at least one T. = ( cg log such that 2T. 2 4c² for 02 1o with probability 5.
Once such т. is obtained, following the last paragraph of (i), we can see that x2 2 2c2 holds until
t= Ti, with high probability.
(iii) When p. 2 3. We apply (B.14) with T= = -i 11 (1² -(0- 2) to obtain that
-1 c
Groed= 4
22
c ci C to 4
(B.17)
for all r = 0, 1,. T - 1, with probability 1 8-
We take cg < C7 2 - _ _ that (B.17) is bounded by o ---- Then,
hrd
x2(++1) 2 x + c 4
Σ (k
p + croed F4
22
2 7 + creed
p1
5=0
is easy to see that (((++ is lower bounded by a+1 where a = 1cid 1 and a+1 = a² +
cieed # (1 By applying Lemma 18, we have
k0
2 (1 ---ieed= (p. (1 2)1) p.-Z
28

---

Thus, until < (ciced (p, ()0))(0-))) <T+1 < du.1,, with probability at least 1 1 6, there exists some T. such that
2 2 402 2 %
when 62 S J/
Once such T, obtained, following the last paragraph of (i), we can see that x22 2 20g holds until t = Ti,1 with high probability.
In the above proof we used the (discrete version of) Bihari-LaSalle inequality from [BAGJ22].
Lemma 18. For 2 3 and > O, consider positive sequence (a* (22) such that
r+ = a' + c
Then, have
a
a'
(1 - c(p (11))()-))))=- pー
Proof. From definition, we have
a ++| a²
1
(at))-1
tmat xP-1 p
[at)p=2
(at-1)p-2
Taking the summation and re-arranging the terms yield
(a')-(0-2) < (ao) (P-2) clp 2)t,
a
a.a'
(o) =tz"
which gives the lower bound.
B.5 From Weak Recovery Strong Recovery
In the previous subsection, we proved that after = 271, = O(d) steps, with probability Ộ((1) 2TT.1 > 202. This subsection
the randomness of initialization, we obtain nontrivial alignment Kij discusses how to convert the weak recovery into the strong recovery.
Lemma 19. Suppose the neuron satisfies 2T2 2 202- Take = = an 2/2011 O.
= g2(++1) 0 for all t 2 T1,1, where En S poly(ci). If the activation functions satisfy all of the assumptions Section B.1 for strong recovery, then we have
0 2(T., 21-
with high probability, where T. S Ti.2 C,de= Moreover, U w ( > T. < < T,,2 -2 with high probability.
E for all
Proof. Consider the Hermite expansions of a. and Let p be the smallest degree that both 0,, and a have non-zero coefficients. First we compute the population gradient (of the correlation term) as
V Tuwyoww" 21 T x) = E
A w
ag He, He, m)) (0Tx))
Σ He, (w xt)
flay 2 0 (j 2)1a, 20 21 (B.18)
Applying Passi, we have
E P V Vwyo(w22 x)] 10 - 2 T 0)w2t
Σ
jla 300 205-1 (
(B.19)
j=p
29

---

Thus, the update of the alignment n.2 = 0 w is
11++1 2 2r + +n0² g 2 n: 2t
2 n""1"" 3 e allg|ll
where
g Pa 2t yw
(w
x x)x.
From (B.18), the expectation of 19) is bounded by
DC
B]n2(e+1]] 2 2r + n1 - (xx2)
Nlay - C/d(((( - -n) Σ
j=p
2 2t + 1)(1 - 22)
11 - - + 1))-
By letting /) S & = when x2 < 1 6, we have
E 2(+2111
2 21
1 11
Σ
11,,3,(12 2 2n² 2² + nee.
j=p
It is easy to see that the noise 2 has sub- weibulibull tail and we obtain that
1
K ((++)
2 2t
+
n
31 (1 + ?pp2 2 2t + nEet + np
(B.20)
2
p=p
Suppose that 20² S 2(11144) (B.20), we have
S 1 - for all t 0,1,...,, 1. By taking the summation of
T1.17-1
R 2TL1++)
x + nete +1
Σ
2 2 20² + NETCL CanVT,
(B.21)
8=T1,
with high probability. The third term is bounded by C < 02 when S CgT 11 11- c2 0 222 and by IeTe when 4s 2c 20 C2, we can bound (B.21) by
46 2 2p CC..
Because
o2 C
67 2-- 2
(TT1177 202 + 2 nETel
(B.22)
which verifies 202 SA" 2x2(TI.147
Therefore, by induction, until x 2 1 we have the lower bound (B.22), whose RHS exceeds
1 --8 6 when т 2 - -ep" < 22 Thus, there exists T_ < Ti, = Cads-2 such that x2((1,1+T.) 2 1 - E, with high probability.
Now, what remains is to prove that x 2((,11+77 > 1 - 35 holds for all T. S t < Ti.2 C3ds Because 222 are mean-zero sub- weibul random variables, we have that |Σ 1 2,2 < CAV7'' for all 0 A T, T' < Ti,1 with high probability. Also, because r7' < Ed we can easily see that (x+((++)) - 2- = O for all 7 = 0, ,i.1 - 1, with high probability. Thus, when there exists 2 т, such that - 2 11- - and n +T) < 1 - E, we have (T2(T.,++)) 2 - 2 with high probability, Moreover, following the above argument, we can inductively show that
k 2(T1,1+T+T) 2 1 2: + nET'd CanVT'
2 2s nET'd { (r' S S
m (7' (
2 35,
for < T1,2 until 2. () > & holds again. Note that the last inequality follows from t 2 =C By repeating this argument (if there are multiple such r), we can see that ((T1)+))) 21 1 - holds for all T. < 1 < 7ı,2 = C2de with high probability.
Adjusting hidden constants to remove a factor of 3 from 3e yields the desired result.
30

---

o(1) B.6 Second Layer Training
From the previous analysis, know that at least (2(1) portion of the neurons will satisfy the weak and strong recovery conditions (Appendix B.I), at least 22(1) portion of the neurons (independent from the choice of σj) satisfy initial alignment conditions (Appendix B.2), and at least fraction of them achieves strong recovery. This subsection proves a generalization error bound after second-layer training. Let fa (x) fe( for (w) aj. b where IR and (w,, b are the parameters trained in the first stage. Let a" E IR be the "certificate" with ||a || O(N)
that is shown exist Lemma 22.
The following lemma is a complete version of Proposition 5. Lemma 20. There exists 4q-th order polynomial Q b, g¹) of R maxy ا b = (b,) (ww b, for some óg > 0, the ridge estimator satisfies
Polynomial Link Functions.
such that, if we X @(1 Tzho Ne
1
QR +2)))(((()))) + 3\
Ilfa f.1l2A(P.) & (N-2 +E²)+ +
a (|²(B.23)
T2Aōn
2
with probability 1 - 86. Hence taking T2 = Ō(E-1) we have
+ (1) and N =
[x((fa(x)- - fz(x))
Proof. Let Pr, be the empirical distribution of the second stage: Pra == 1/ E/e oz. Let vxx)
(a((a, (xy)) + b,)) so that fa(x) (a, (xx)))
Part (1).
We first bound the term IIʃa falle2(P))) Since L(Sa) A|lal1 s ((Ja)) ) A11a*112
have
I1fa s_112>(p))) + A||āl12
(B.24)
$ Ilfa. - Sall?a(P5)))) T2 (fa (z.) facc)-fa(x;))e fa(x;))e + 11
Now, by the Cauchy-Schward inequality, we have
T2 (a,) fa(x:))e = (a" - -à) 2 T'
2
(xx)e,
T2
T2 1=1
Eis EEE (xE) (xx,)
S 2 - -ally
T3
By applying Markov's inequality to the right hand side, it can be further bounded by
Ex [1P)
> 2 11 + +
Ex (x=))
la'- -ally
ally
I 2
||all
T201
T28 A
with probability 1- 81. Thus, by combining with (B.24), we arrive at
flfa - f.112 + |lal|2 & Ilfa- - -f + Ex (l
37
+ 2
||a
T581A
Here, by using the evaluation can be further bounded by
|a. = = + 8) in Lemma 22, the right hand side
X
Eg b=ll1(()))
31
Ifa f- f1 + 2 1|all < O(N-2+
+
||a1| 1|2
T
2
31

---

Part (2). Next we lower bound llfa fx fxllia(r))) by noticing that
llfa fxlli2(P7)))
= =Ilfa- - - - 11fa
ful(22(22) + l1fa - follER(P2)
T2
= =I.fallia(p5)) - 11.all22 - T2 i=1
Σ
La Ia( - EIaa (B)
+ +I (I
Ilf 11fe112((.)) + llfa - fx11E2(P.)
(B.25)
The first terms of Eq. (B.25) can be bounded by
E 22 -i-l (x)))(
Il -lfallar
a
T2
S ||āl1
sup
Il (22())))
1lfall22
a:||a||s1
The standard Rademacher complexity bound yields that
|| [[all12((P)) - 11 (1
T3 (a:)
sup aERM:
=||a||51
T2 xz;)
22B
(2
T2
sup [|RR:||a||S|
T2
1 T2
22
E sup
T (2x:)))
(0: (=)
DERM:||al|s!
T3
2
E
T2
T2
(ax(()))11
1
2
E.
(1))1
T2
where (o the i.idd Rademacher sequence independent of (z:) Hence, Markov's 's inequal- ity yields
[1lfalli((P))) faa (P))
2/ V T262 E. (x))
with probability 1 1 82.
The third term in Eq. (B.25) can be evaluated as
La (m.)). (x)]
T
à
T
(x::s.. (x(((a))
T2
Ts
T2
< ||ل|
- [ete) (Ea(
1
|ally
Emn (1l)bbb))))
T253
< |ally
E.
[1]
1 (1)
32

---

1 |1ala1 + XT203 Bm [|<|(=))])** + (x))1))
with probability 1 - d3 where we used Markov's incquality again in the second inequality. Finally, the fourth and fifth term in Eq. (B.25) can be bounded as
||1f-1B2 - 111-1122(p)) = V (ISAl1BE(PP)) - 11
Em ((P* II (1 (1..)112(()))))
V T2d4
Em (x))))))
V
T264
with probability 1 - 54 where used Markov's inequality in the last inequality. Combining these inequalities, finally arrive
I1 - sal(2((2)) + (7 V
T2d2 En
|lal?
(
[1)
E.
Ea [=(((())))))
- zŌ(N~2 + x² + Tz1
+
la
81
8z
dz
2
with probability 1 ド-E 8z. Hence, by setting
>> 8v
Ee (Ilv(()))))))))))))))) have that
T202
flfa - --1(22(P2)
E. [+)(()11 Ex [+
Ex (=)())))))
3A | |a11.
S SŌ(N-2+67)
+ 2
T₂A
83
5g
When the activation function a is a polynomial, then each ()(()) = o wr) + b,) is an order q-polynomial of Gaussian random variable (ar, wy) which has mean 0 and variance E[(x,, w2
||,,|| Ô(1). Then, if we let R := max, ||w,||| Ö(1), the term maxj max _m (1)((( (x)) Em [x])()))] can be bounded by 4q-th order polynomial of R. and b, which can be denoted by (Ru,, 4q).
Part (3). By combining evaluations of (1) and (2) together, if we let 8V for some & > O, (by ignoring polylogarithmic factors) we obtain that
2 T200
Ex [+))
31
+ llfa -00-ff12z00 ((ZAA)) & (N-2 +22) + T2A50 (2N2Q b, (,4') + Exl(f(()))))
|aa
2
with probability 1 45g. Thus, since ||a" Ö(N), by setting T2 @((N"Q b, ([f(()))[)) 4), and N Ô(e-1) we obtain that (B.23) S
Higher Generative Exponent Functions. For general link functions, under Assumption 4 and the bounded fourth moment of the link function, we have the following counterpart of Lemma 20, which provides the formal statement of Proposition 10.
Lemma 21. Suppose that El x) < co and Assumption 4 hold. Then, by setting X
N2 Taho
for some dg > O, the ridge estimator à satisfies
1
Ia fAl12A())) S VT200
+ (=((())*) + VT200 ||a*12 CN
with probability 5g. By taking T2 = @((N+ + N2))--),, we have
(el(fa())-f(1e1) - s²
33

---

Furthermore, applying Lemma 12 and 13 yields that, when o, a satisfies E and (l 2)d1 are bounded, with a properly designed randomized activation in Lemma 12, by taking N =) and T2 6(= -2), Algorithm 1 yields
Em -
with probability od(1).
Proof. The proof is identical that of Lemma 20, with the difference being that we replace the bounded moment assumptions with E( x> do or Assumption
Approximation Guarantee. Note that for non-polynomial link function with generative exponent P- 2 3, the approximation error is already controlled in Assumption based on [BBSS22, Lemma 4.4, 4.5] (using activation function with ReLU component). If is degree-q polynomial, have the following approximation result using polynomial activation, which follows Lemmas and 30 of [OSSW24a].
Lemma 22. Suppose that there exist at least N' = Ô(N) neurons that satisfy ||00] 2Th 0 S & and a is a polynomial link function with degree at least q. Let bj Unif([-Cb, Cb]) with C, = ((),, and consider approximation of ridge function h(O x) with its degree at most q. Then, there exists a1..... .....aa aN such that
N E ajo, (w)
27: b,) h(O" x)
=ÓNN- +
with high probability, where (x, y) is random sample, and omit dependence on the degree q in the big-O notation. Moreover, we have Σソ= a ÓNN).
Lemma 22 can be established from the following result in [OSSW24a].
Lemma 23. Suppose that C2 2 For any polynomial h(()) with its degree at most g, there exists bbb h) with (6(b; h3] S Co such that for all 8,
hlo(b; + 6)] h(s).
34

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, ad- dressing issues of reprodueibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes]. [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evalu- ation. While "[Yes] is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] provided a proper justification is given (e.g.. "error
are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification elaborate, All supporting evidence can appear either the main paper the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
The paper's main contribution is studying the computational complexity of learning single-index models with polynomial link functions by training neural networks with stochastic gradient descent, which is what we claim in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
35

---

Justification:
We discussed the limitations of our results in Section 4.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions required in our work are stated in Assumptions 1,2, and 3. All proofs are presented in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [NA]
Justification: This paper theoretical and experiments are only used for the purpose of illustration in the Introduction.
36

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), relcasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via
While NeurlPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [NA]
Justification: This paper is theoretical and experiments are only used for the purpose of illustration in the Introduction.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips. catlpps.cc/uubbcee httpliicgguide//e/e Code/oodeSubbt ssi aesonnoolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
37

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification:
This is theory paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [NA]
Justification:
This is theory paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK report I-sigma error bars, but should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors verified.
For asymmetric distributions, the authors should careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification:
This is theory paper.
Guidelines:
The answer NA means that the paper does not include experiments.
38

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g.. preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http::/herrriss./ppbitt cetoolii//thiiaudeeelllllll ines?
Answer: [Yes]
Justification:
The authors have read and followed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper carries out a theoretical study, and we believe our work does not have specific societal impacts that require discussion.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
39

---

Justification:
This paper is theoretical and the experiment is only about two-layer neural network and synthetic data, which we do have such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safery filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: This paper is theoretical and does not use any assets.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, paperswi thhoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: This paper is theoretical and does not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
40

---

Question: For crowdsoureing experiments and research with human subjects, does the pa- per include the full text of instructions given participants and screenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: This paper is theoretical and does not include crowdsourcing experiments nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper is theoretical and does not include crowdsourcing experiments nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva-
lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
41


Reviews:
[{"review_id": "nQ8Hnz9u44", "reviewer": "Reviewer_p2Vt", "summary": "This manuscript studies the learning properties of two-layer networks trained with SGD reusing the batch. The authors show that this simple modification allows SGD to surpass the limits of CSQ algorithms and learn single-index functions efficiently. The submission considers both recovery of the target features and generalization properties of two-layer networks. The claims are supported by rigorously proven theorems.", "questions": "- The idea of label transformation implemented by [CM20] could be reported. The contrast with what SGD (with batch reusing)  is implementing would give an idea to the reader of the strength of the claims. \n- Naively, one might think that sample complexity guarantees might come easily from weak recovery plus [BAGJ21]. Maybe the authors could comment more on the technicalities that arise. \n- What is the role of the hyperparameters, e.g. the interpolation one $\\xi$, the learning rate, and all the quantities appearing in the non-CSQ transformation? Are they randomly drawn (and the theorems hold with high probability? \n- The authors mention that interpolation is required and correctly dedicate a paragraph to it. However, is not clear to me if they believe is only necessary for the technicalities of the proof or if they believe it is valid in general. \n- The authors do an amazing job in introducing the CSQ/non-CSQ parallel when reusing the batch in section 2.2. However, to even more enhance this subsection I think it would be great to state more clearly details over the SQ class. Although SQ is formally defined, the authors could refer in the submission to the lower bounds achievable by SQ. More precisely, it was reported that $n \\simeq$ d is both sufficient and necessary for learning by citing works using AMP-type algorithms. Are these algorithms belonging to SQ? \n- Closely related to the above question. If AMP-type algorithms belong to SQ, could the author comment more on the link with [Gla23] and non-adversarial noise and the possibility of going beyond SQ? Of course, the Information Theoretic barrier cannot be broken, but I think more clarity on these points would be more than welcome. I think the insights presented are crucial and interesting and deserve more space in the main body.", "limitations": "The limitations are addressed in the manuscript.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The strength of this submission resides in the strong theoretical claims. The questions addressed are of great interest to the theoretical machine learning community.", "weaknesses": "This submission has no outstanding weaknesses but the presentation could be enhanced. I will detail in the section below some suggestions to improve the manuscript which sometimes is a bit obscure for a non-expert reader.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.  \n\n---\n\n*\"The idea of label transformation implemented by [CM20] could be reported\"*\n\nIn the current manuscript, the difference between our label transformation and that in [CM20] is discussed in the paragraph starting line 157 and line 250. Specifically, prior SQ algorithms are typically based on thresholding, whereas we use monomial transformations which can be easily extracted from SGD update.\n\n---\n\n*\"one might think that sample complexity guarantees might come easily from weak recovery plus [BAGJ21]\"*\n\nThe difficulty is discussed in the paragraph starting line 162. In particular, since the link function $\\sigma_*$ is unknown, we cannot directly employ the argument in [BAGJ21] which assumed a well-specified model, i.e. $\\sigma_* = \\sigma$. Instead, we make use of the randomized Hermite coefficients of the student activation function to translate weak recovery to strong recovery. \n\n---\n\n*\"What is the role of the hyperparameters ... and all the quantities appearing in the non-CSQ transformation?\"*\n\nMost hyperparameters in Algorithm 1 are deterministic, with the only exceptions being the sign of student activation function (see Lemma 2) and the momentum parameter $\\xi$ randomized over student neurons.\nThis aims to guarantee that, for any target function, there exists some student neurons achieving strong recovery (Theorem 2). \nWe realize that there is a typo in Theorem 1: \"with probability\" should be replaced by \"with *high* probability\". We apologize for the confusion this may cause.  \nWe will clarify this in the revision. \n\n---\n\n*The role and necessity of interpolation*\n\nWe use an interpolation step to improve the signal-to-noise ratio in the gradient; this is crucial in the generative exponent 2 setting where the signal and noise are almost at the same magnitude (see Section 3.2 and paragraph starting line 307 for details). It is possible that other modifications or hyperparameter choices of the learning algorithm also achieve a similar effect, but we do not pursue these in the current submission. \n\nBelow we recap the intuition provided in Section 3.2 on the failure of the standard online SGD analysis. \nWhen we analyze the training dynamics of learning single-index model, we characterize the progress via the projection of $\\boldsymbol{w}^t$ onto $\\boldsymbol{\\theta}$, which we refer to as the alignment $\\kappa^t = \\langle \\boldsymbol{\\theta}, \\boldsymbol{w}^t \\rangle$\n(e.g., [DNG+23]).\nWe want to show that $\\mathbb{E}[\\kappa^{t+1}]$ increases from $\\kappa^t$.\nGiven the gradient $\\boldsymbol{g}^t$ (assumed to be orthogonal to $\\boldsymbol{w}^t$ for simplicity) and step size $\\eta$, the update of alignment is given as \n$$\\kappa^{t+1} = \\left\\langle \\boldsymbol{\\theta}, \\frac{\\boldsymbol{w}^t + \\eta \\boldsymbol{g}^t}{||\\boldsymbol{w}^t + \\eta \\boldsymbol{g}^t||} \\right\\rangle \\gtrsim \\underbrace{\\left\\langle \\boldsymbol{\\theta}, \\boldsymbol{w}^t \\right\\rangle}_{=\\kappa^t}+ \\eta \\left\\langle \\boldsymbol{\\theta}, \\boldsymbol{g}^t \\right\\rangle - \\frac{1}{2} \\eta^2 ||\\boldsymbol{g}^t||^2 \\kappa^t  + \\text{(mean-zero noise)}.$$\n\nOne sees that for the expectation to be larger than $\\kappa^t$, $\\eta \\mathbb{E}[\\langle\\boldsymbol{\\theta}, \\boldsymbol{g}^t\\rangle]$ should be larger than $\\frac12\\eta^2 ||\\boldsymbol{g}^t||^2\\kappa^t$.\nTo achieve this, we can simply let $\\eta^t$ sufficiently small ($\\Theta(d^{-1})$ for $\\mathrm{IE}=2$); because the signal term linearly depends on $\\eta$, while the noise term quadratically. \n\nHowever, in our case, the signal comes from non-CSQ term.\nFor example, when $\\mathrm{IE}(y^I)=2$, the signal term is proportional to $(\\eta^t)^Id^{I-1}\\kappa^t$ (under $\\eta^t \\lesssim d^{-1}$).\nTherefore, decreasing $\\eta^t$ does not improve the SNR. The interpolation step provides a remedy by preventing the parameters from changing too fast and reducing the projection error.\n\n ---\n\n*\"it would be great to state more clearly details over the SQ class\"*\n\nThank you for the valuable suggestion. We will discuss the SQ complexity and generative exponent in [DPVLB24] in more details. Indeed our analysis is based on the observation that SGD with reused batch can implement SQ, and hence $\\tilde{O}(d)$ samples are sufficient to learn target functions with generative exponent at most 2; this sample complexity is also achieved by AMP algorithms (typically assuming knowledge of the link function to construct the optimal preprocessing), and is consistent with the SQ lower bound. \n\n---\n\n*\"could the author comment more on the link with [Gla23] and non-adversarial noise and the possibility of going beyond SQ?\"*\n\nThank you for the close reading. We realize that this remark on future direction is misleading. We initially thought that for the $k$-parity problem, the SQ lower bound suggests that $n\\gtrsim d^k$ samples are required for rotationally invariant algorithms with polynomial compute (as opposed to $n\\gtrsim d^{k/2}$ in the Gaussian case). However this is not the case, and the total computation in [Glasgow 23] only matches that suggested by the SQ lower bound $q/\\tau^2\\asymp d^2$. \nWe will include appropriate references on the gap between statistical and adversarial noise, such as [[Dudeja and Hsu 20]](https://jmlr.csail.mit.edu/papers/volume22/20-837/20-837.pdf) [[Abbe and Sandon 20]](https://onlinelibrary.wiley.com/doi/full/10.1002/cpa.22121). \n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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": "QCTqYCgWHC", "reviewer": "Reviewer_hk3C", "summary": "This paper addresses the problem of learning single-index targets with polynomial link functions under Gaussian inputs. The authors demonstrate that using SGD on a two-layer fully connected network with a specific activation function can learn such targets with O(d poly(log(d))) samples. The analysis involves the reuse of samples, which allows to improve previous bounds obtained for online single-pass SGD.", "questions": "Q1) Can you clarify how many hidden neurons N are needed for the main result to hold?\n\nQ2) How does the bound depend on the degree of the target q?\n\nQ3) Would you expect the same result to hold if the bias weights are trained?\n\nQ4) Do you have any high-level intuition on why the even polynomials are harder than the odd ones? Do you believe that the poly(log(d)) terms are needed for the even ones?\n\nQ5) Do you expect a similar analysis to hold for other losses? For example, L1 loss.\n\nQ6) Do you have an intuition for what could be an optimal mini-batch re-use schedule?\n\nTypos/suggestions:\nProposition 4ii): Can you formally define 'the odd part of f'?\nLine 211: missing the word 'high'.\nLine 326: 'high' -> 'weak'.\nLine 175: missing the word 'be'.\nLine 163: typo 'not be not'.", "limitations": "The authors adequately addressed the limitations of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "S1) The paper advances the analysis of the complexity of learning Gaussian single-index models, a currently very popular model for the theoretical study of neural networks. This contribution is thus significant for the deep learning theory community.\n\nS2) The technical contributions are novel and well presented.\n\nS3) The paper builds on previous work that demonstrated the benefits of re-using batches for learning single-index models, providing concrete evidence of strong learnability of the target by SGD on a shallow network.", "weaknesses": "W1) The authors provide minimal empirical validation, with only one experiment demonstrating their claims. They do not address more standard SGD practices, such as training both layers simultaneously, using standard activations/initializations, or employing larger learning rates.\n\nW2) The analysis relies on several theoretical assumptions and is limited to a very structured data distribution, which is common in deep learning theory proofs. While the assumptions are well stated, there is little discussion on whether or how these assumptions could be relaxed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below. \n\n---\n\n*Standard SGD practices, relaxing assumptions*\n\nWe agree that Algorithm 1 deviates from the most standard training procedure in practice. Note that layer-wise training and fixed bias units are fairly standard algorithmic modifications when proving end-to-end sample complexity guarantees, as seen in many prior works on feature learning [DLS22] [AAM23]. Without such modifications, stronger target assumptions  are generally needed, such as matching link function [BAGJ21]. \nBelow we discuss the possible extensions to more general settings. \n\n- **Simultaneous training**. We believe that the statistical efficiency of Algorithm 1 can be achieved by simultaneous training of all parameters, under appropriate choice of hyperparameters. For instance, the layer-wise procedure may be approximated by a two-timescale dynamics [BBPV23], and the required ``diversity'' of bias units may be met when the learning rate of the bias weights is sufficiently small. As mentioned in the Conclusion section, it is also interesting to consider a more standard multi-pass algorithm instead of the currently employed data-reuse schedule. \n\n- **Different losses**. \nWe focus on the squared loss as it is the most standard objective for regression tasks. Note that the restriction to correlational information is a feature of the squared/correlation loss -- see Section 2.2. If we employ a different loss such as L1 loss, it is possible that SGD can implement non-correlational queries just from the loss function itself -- such direction is tangential to our analysis, since our goal is to show that such non-CSQ component naturally arises from reuse of training data. \n\n---\n\n*Number of neurons and dependence on target degree $q$* \n\nThe required width of the neural network is almost dimension-free. In particular, we only need to set $N=\\text{polylog}(d)$ to achieve $o_d(1)$ population error -- see line 245 for details. \n\nOn the other hand, our big-$O$ notation hides a constant that might depend exponentially on the target degree $q$. This is due to the sign of Hermite coefficients required in Assumption 3, and the compactness argument to uniformly upper-bound $C_q$ in Proposition 4. \nNote that similar dependence on $q$ is also present in tailored algorithms\nfor learning low-dimensional polynomials [CM20].\n\n---\n\n*\"why the even polynomials are harder than the odd ones\"*\n\nIntuitively speaking, the neural network is initialized at an (approximate) saddle point when $f_*$ is even, since the expectation $\\mathbb{E}[\\mathcal{T}(y)\\langle\\boldsymbol{x},\\boldsymbol{\\theta}\\rangle]=0$ for any $\\mathcal{T}\\in L^2(y)$. See [BAGJ21] for more discussions. \n\nTo elaborate, let us consider one-pass SGD and evaluate the scale of population gradient of one neuron $\\sigma(\\langle \\boldsymbol{w},\\boldsymbol{x}\\rangle)$ at random initialization. \nLet $\\sigma_*=\\sum_{i=0}^q c_i He_i$, we have \n$\n\\mathbb{E}[\\nabla_{\\boldsymbol{w}}\\sigma_*(\\langle \\boldsymbol{\\theta},\\boldsymbol{x}\\rangle)\\sigma(\\langle \\boldsymbol{w},\\boldsymbol{x}\\rangle)]\n\\approx \\sum_{i=1}^q \\mathbb{E}_{t\\sim \\mathcal{N}(0,1)}[\\sigma^{(i)}(t)] c_i\\langle \\boldsymbol{\\theta},\\boldsymbol{w}\\rangle^{i-1}\\boldsymbol{\\theta}.\n$\nTherefore, when $c_1 \\ne 0$, the scale of the population gradient is $\\Theta(1)$, while when $c_1=0$ and $c_2\\ne 0$, it is $\\langle \\boldsymbol{\\theta},\\boldsymbol{w}\\rangle \\simeq d^{-\\frac12}$  with high probability.\n\nSimilarly for reuse-batch SGD, the population gradient is a linear combination of $\\mathbb{E}[\\nabla_{\\boldsymbol{w}}\\sigma_*(\\langle \\boldsymbol{\\theta},\\boldsymbol{x}\\rangle)^i\\sigma^{(i-1)}(\\langle \\boldsymbol{w},\\boldsymbol{x}\\rangle)(\\sigma^{(1)}(\\langle \\boldsymbol{w},\\boldsymbol{x}\\rangle))^{i-1}]$.\nWhen $\\sigma_*$ is even, one can similarly see that the scale of each is $O(d^{-\\frac12})$.\n\n---\n\n\n*\"what could be an optimal mini-batch re-use schedule?\"*\n\nWe intuitively expect that any mini-batch size $b$ between $1$ to $d$ with $\\eta=\\tilde{\\Theta}( d^{-1}b)$ would achieve similar sample complexity. But for simplicity, we employed $b=1$, which excludes the correlation between different samples.  \nInvestigating the benefit of more intricate mini-batch schedule is an interesting direction for future work. \n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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": "RY6Ub5acBS", "reviewer": "Reviewer_Mv1X", "summary": "This paper studied the problem of learning single index models under Isotropic Gaussian distribution. The target model $f^*(x) = \\sigma(\\theta^\\top x)$ is a polynomial function $\\sigma$ composed with a one-dimensional structure $\\theta^\\top x$, where the polynomial $\\sigma$ is of degree at most $q$ and has information component (i.e., the order of the first non-zero Hermite polynomial expansion) $p$. The paper studied the sample and computational complexity of learning such single index models with 2-layer neural networks, using gradient descent type methods. The critical aspect of the algorithm is reusing the same batch of samples every two iterations, which frees the algorithm from CSQ constraints and becomes an SQ algorithm. A critical observation is that by reusing samples at each iteration, the algorithm induced a monomial transformation of the labels, which effectively reduced the information component from $p$ to less than 2. Hence, the algorithm can achieve near-optimal sample complexity $\\tilde{O}(d)$.", "questions": "1. Do ReLUs and sigmoids satisfy the assumption 2?\n2. I think there are typos on line 593 to 595. What is 'x' on the right-hand side of line 593 and 595?\n3. Since only $exp(-q)$ neurons satisfy assumptions 2 and 3, does it imply that the width of the network is at least $exp(q)$?\n4. Since Theorem 2 relies on neurons that satisfy assumptions 2 and 3, does it imply that having only an $exp(-q)$ fraction of good neurons  (neurons with $w_j^\\top \\theta>1 - \\epsilon$) is enough to achieve small $L_2^2$ error? What is the intuition behind this?", "limitations": "The authors addressed the limitations of the paper and provided inspiring future directions.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is clearly written with intuitions and useful explanations.\n2. This paper provides new perspectives on designing SQ algorithms to learn single-index models using neural networks. Though the idea of reusing samples has already appeared in prior works ([DTA+24]), this work shows that resuing samples can achieve strong recovery of the hidden direction and provides a well-rounded analysis of the sample and computational complexity. Importantly, the authors showed that by reusing the mini-batches, one can learn the target model with $\\tilde{O}(d)$ samples, which is near the information-theoretic limit.\n3. This paper provides a very interesting intuition on reducing the information exponent of the link function $\\sigma$ using monomial transformations, which could be of independent interest for future works.", "weaknesses": "1. Though the authors claimed that they were using neural networks to learn the single index model, the activation of each neuron turns out to be a combination of polynomials. Hence, the neural network $\\sum a_j\\sigma(w_j^\\top x + b_j)$ is essentially a linear combination of Hermite polynomials. In this case, I am wondering what the differences are between using the 'neural network' to learn the single index models and using polynomials to learn the single index models, which is already done in [CM20]. Of course, [CM20] requires a warm start procedure, which is not a gradient descent type algorithm, but I think it would be more interesting if that analysis is carried out on conentional neural networks like ReLU networks.\n\n2. The authors hide many constants in the big-O notations. However, I am skeptical that all those parameters are independent of the dimension $d$. For example, in the proof of Proposition 4, the upper bound on $C_q$ is $1+\\log_2(H_0^{-1})$. However, there is no actual lower bound on $H_0$ other than that being non-zero. Therefore, I am wondering if it is possible that $H_0$ can be as small as $2^{-d}$? I think the paper will be more theoretically sounded if the authors can explicitly present the dependence on the parameters $C_q$, $C_\\sigma$ etc. in the final bounds on the sample complexity and iteration complexity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below.\n\n---\n\n*Polynomial activation function*\n\nWe make the following clarifications. \n\n1. Note that all square-integrable activation functions (ReLU, sigmoid, etc.) can be written as a linear combination of Hermite polynomials. The key differences between our algorithm and [CM20] are as follows, \n\n- [CM20] employed a label transformation (based on thresholding) prior to running SGD, whereas we use the squared loss without preprocessing and extract the transformation from reusing training data. \n\n- [CM20] considered optimization jointly over the low-dimensional subspace (finding index features) and the coefficients of the polynomials. In contrast, in our setting the coefficients of the polynomial activation are fixed, and we optimize the parameters of the neural network. \n\n2. We restrict ourselves to polynomial activation because it is easy to construct coefficients that satisfy Assumption 3 (required for strong recovery). For strong recovery, ReLU or sigmoid may not suffice, and the use of well-specified or polynomial activations is common in the literature, e.g., [AGJ21][AAM22][AAM23][DNG+23].  \nAs for weak recovery, Assumption 2 is satisfied with probability 1 by a shifted ReLU/sigmoid (e.g., see Lemma 15 in [BES+23]). \nTherefore, we can establish weak recovery using standard choices of activation. \nWe will comment on this in the revised manuscript.\n\n---\n\n*Dimension dependence in constants* \n\nAll constants in our theorems are dimension free. Specifically, the lower bound on $H_0$ does not depend on dimensionality, since Proposition 4 is for univariate functions $f:\\mathbb{R}\\to\\mathbb{R}$. We notice that there is a typo in Appendix A: the expectation in A.1 should be with respect to $\\mathcal{N}(0,1)$ since $f$ is a scalar function. We apologize for the confusion this may cause. \n\n---\n\n*\"does it imply that the width of the network is at least $\\text{exp}(q)$?\"*\n\nYes, the required student width is exponential in the target degree $q$, which is treated as a constant the big-$O$ notation. \nAlthough only a small fraction of neurons can achieve strong recovery, the entire neural network can achieve small $L^2$ error because these ``good'' neurons can be singled out in the second-layer training, as shown in Proposition 3.  \nNote that similar dependence on $q$ is also present in tailored algorithms\nfor learning low-dimensional polynomials [CM20]. \n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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": "VktodmpomP", "reviewer": "Reviewer_6kdk", "summary": "The paper studies the classical problem of the single index model over Gaussian inputs, i.e. $x\\sim \\mathcal{N}(0,I_d)$ and $f_*(x)=\\sigma_*(\\langle \\theta_*,x \\rangle)$ for an unknown direction $\\theta_*$. Information theoretically, one needs $\\Omega(d)$ samples to learn this function class. The paper shows that reusing the batches with a certain SGD-based training on a 2-layer neural network achieves vanishing $L_2$ error with $O(d \\text{polylog}(d))$ samples--that is this SGD algorithm by some reused batches learns this function class at nearly the information-theoretic limit. \n\nThe problem has been studied extensively. CSQ lower bound in terms of ``information exponent'' $(IE=p)$, the lowest degree of a non-zero Hermite coefficient of the link $\\sigma_*$, was considered the correct complexity measure for SGD. That is $\\Theta(d^{p/2})$ samples are necessary. However, when one reuses the batch, this can be seen as a non-correlation query on the same example and hence CSQ lower bound is breached, allowing us to learn at $\\tilde{O}(d)$ sample complexity, irrespective of the information exponent.", "questions": "1. Can the use of momentum be avoided in the first layer of training? If \"no\" then, in what way this is important in the current analysis?\n2. Why cite [Gla'23] (on line 338 page 9) for the point of adversarial noise in SQ vs non-adversarial noise in GD?\n3. Do the authors believe what was said in lines 337-338: ``It is also possible that SGD can achieve a statistical complexity beyond the SQ lower bound`` \nfor the generative exponent more than two? While I understand this is not ruled out (and it is indeed important to make the point of non-adversarial noise), saying that ``it is possible\" sounds slightly strong to me as if the authors are indicating their belief about the situation. Unless there is any strong experimental evidence for this or the authors truly believe this is possible, I would encourage the authors to reword this part.\n4. Should the abbreviation on line 31 page 1 be CSQ instead of SQ?", "limitations": "None.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The paper addresses an important question in our understanding of the complexity of gradient descent-type algorithms on regular neural networks. There have been considerable efforts devoted to understanding this. This paper goes beyond most (if not all) of these works by analyzing an SGD with reused batches, over vanilla batch SGD. While a recent paper of [DTA+24] provided the first evidence of the benefit of reusing the batches and that CSQ bound can be escaped, the paper goes beyond this and in an important way as follows. \n2. This paper considers strong recovery which is more satisfying, and technically much more challenging (in contrast to [DTA+24] that only considered weak recovery). To achieve this, there were important pieces to be figured out, which the authors successfully did. This paper provides a clear end-to-end analysis and establishes the learning guarantee in contrast to [DTA+24].", "weaknesses": "I do not see any major weaknesses. The training procedure is slightly non-standard, but it is completely understandable from a technical point of view. The layer-wise training and the use of the projected gradients are completely standard in theoretical research. However, I could not see a clear motivation/need for the momentum for the first-layer training. I was wondering if it can be avoided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thoughtful comment and constructive feedback. We address the technical concerns below. \n\n---\n\n*\"Can the use of momentum be avoided in the first layer of training?\"*\n\nWe use an interpolation step to improve the signal-to-noise ratio in the gradient; this is crucial in the generative exponent 2 setting where the signal and noise are almost at the same magnitude -- please see our response to Reviewer p2Vt for more details and a recap of our explanation in Section 3.2.  \nIt is possible that other modifications or hyperparameter choices of the learning algorithm also achieve a similar effect, but we do not pursue these in the current submission. \n\n---\n\n*[Gla'23] and the possibility of going beyond SQ*\n\nThank you for the close reading. We realize that this remark on future direction is misleading. We initially thought that for the $k$-parity problem, the SQ lower bound suggests that $n\\gtrsim d^k$ samples are required for rotationally invariant algorithms with polynomial compute (as opposed to $n\\gtrsim d^{k/2}$ in the Gaussian case). However this is not the case, and the total computation in [Glasgow 23] only matches that suggested by the SQ lower bound $q/\\tau^2\\asymp d^2$. \nWe will include appropriate references on the gap between statistical and adversarial noise, such as [[Dudeja and Hsu 20]](https://jmlr.csail.mit.edu/papers/volume22/20-837/20-837.pdf) [[Abbe and Sandon 20]](https://onlinelibrary.wiley.com/doi/full/10.1002/cpa.22121) (albeit with highly nonstandard architecture).  \nWe also agree with the reviewer that the statement needs to be reworded in the absence of empirical evidence. \n\n---\n\nWe would be happy to clarify any further concerns/questions in the discussion period.", "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:

Meta-review: This paper studies the problem of learning single index models under Isotropic Gaussian distribution. It shows that reusing the batches with SGD-based training on a 2-layer neural network achieves the sample complexity that matches the information theoretic limit up to polylogarithmic factor, which is a significant improvement over existing theoretical results. 

All reviewers appreciate the novelty and significant contributions of this paper. Therefore, I recommend accepting it.
Decision: Accept (poster)