Title: Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit
Venue: NeurIPS 2024
Accepted: yes
Keywords: single-index model

feature learning

statistical query

SGD
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-theoretic limit

Abstract

Westudy the problem of gradient descent learning of a single-index target function f ∗ ( x ) = σ ∗ ( ⟨ x , θ ⟩ ) under isotropic Gaussian data in R d , where the unknown link function σ ∗ : R → 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 ≳ d Θ( 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 ≃ T = Θ( d · polylog d ) , where Θ( · ) hides a constant only depending on the degree of σ ∗ ; this dimension dependence matches the information theoretic limit up to polylogarithmic factors. More generally, we show that n ≳ d ( p ∗ -1) ∨ 1 samples are sufficient to achieve low generalization error, where p ∗ ≤ 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.

1 Introduction

Single-index models are a 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 setting has been extensively studied in deep learning theory [BL20, BES + 22, BBSS22, MHPG + 23, MZD + 23, WWF24], to examine the adaptivity to low-dimensional targets and benefit of representation learning in neural networks (NNs) optimized by gradient descent (GD). In this work we study the learning of a single-index target function under isotropic Gaussian data:

where ς i is i.i.d. label noise, θ ∈ R d is the direction of index features, and we assume the link function σ ∗ : R → R has information exponent p ∈ 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 σ ∗ and the relevant direction θ ; it is known that learning is information theoretically possible with n ≳ d training examples [DH24, DPVLB24]. Indeed, when σ ∗ is polynomial, such statistical complexity can be achieved up to logarithmic factors by a 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 established a sample complexity of n ≳ d Θ( p ) [BAGJ21, BBSS22, DNGL23], which presents a gap

10 2

10 3

dimensionality

d

10 4

10 5

sample size

n

n

d

2

w

,

2

0.02

0.04

0.06

0.08

(a) Online SGD (weak recovery).

10 2

10 3

dimensionality

d

10 3

10 4

sample size

n

n

d

0.2

0.4

0.6

0.8

1.0

1.2

1.4

(b) Same-batch GD (generalization error).

between the information theoretic limit and what is computationally achievable by (S)GD. Such a gap is also predicted by the correlational statistical query (CSQ) lower bound [DLS22, AAM23], which roughly states that for a CSQ algorithm to learn (isotropic) Gaussian single-index models using less than exponential compute, a sample size of n ≳ d p/ 2 is necessary.

Although CSQ lower bounds are frequently cited to imply a 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 noncorrelational terms in the gradient computation. Very recently, [DTA + 24] 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 θ ) 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 a few specific examples of link functions are discussed.

Given the existence of (non-NN) algorithms that learn any single-index polynomials in n = ˜ O ( d ) samples [CM20] regardless of the information exponent p , and more generally, non-CSQ algorithms with a sample complexity surpassing the CSQ lower bound [DPVLB24], it is natural to ask if gradient-based training of NNs can achieve similar statistical efficiency for this function class. Motivated 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 σ ∗ , can learning succeed near the information-theoretic limit n ≃ 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 polynomial with information exponent p = 3 . We see that SGD with reused data (Figure 1(b)) reaches low test error using roughly n ≃ d samples, whereas online SGD fails to achieve even weak recovery with much larger sample size n = Ω( d 2 ) . Our main contribution is to establish this improved statistical complexity for two-layer NNs trained by a variant of SGD with reused training data.

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 error 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 in [BES + 22, DLS22, AAM23]) 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.

Information

theoretic limit

SGD + batch reuse [

This work

]

SQ algorithm [CM20]

Smoothed SGD [DNGL23]

CSQ lower bound [DLS22]

One-pass SGD

[BAGJ21]

Kernel methods

[GMMM21]

d

˜

Θ(

d

)

˜

Θ(

d

p/

2

)

˜

Θ(

d

p

-

1

)

Θ(

d

q

)

Theorem (informal) . A shallow NN with N = ˜ Θ d (1) neurons can learn arbitrary single-index models up to small population loss: E x [( f Θ ( x ) -f ∗ ( x )) 2 ] = o d, P (1) , if we employ an SGD-based algorithm (with reused training data) to minimize the squared loss objective, with a sample and runtime complexity of n, T = ˜ Θ d ( d ( p ∗ -1) ∨ 1 ) , where p ∗ is the generative exponent of the link σ ∗ .

Note that the generative exponent [DPVLB24] is defined as the minimum information exponent of the link function σ ∗ after arbitrary L 2 transformation, and hence by definition p ∗ ≤ p (equality is achieved by the identity transformation). We make the following remarks on our main result.

Weknow 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 = ˜ O d ( d ) which is information theoretically optimal up to polylogarithmic factors, hence matching the efficiency of SQ algorithms tailored for low-dimensional polynomial regression [CM20].

For non-polynomial σ ∗ with high generative exponent p ∗ > 2 , our sample complexity n ≳ d p ∗ -1 can be interpreted as an SQ version of the online SGD result in [BAGJ21]. Since the information exponent p can be arbitrarily larger than the generative exponent p ∗ , our main theorem disproves a conjecture in [AAM23] stating that n ≍ d p/ 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 optimization 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 ℓ 2 norm for vectors and the ℓ 2 → ℓ 2 operator norm for matrices. O d ( · ) and o d ( · ) stand for the big-O and little-o notations, where the subscript highlights the asymptotic variable d and suppresses dependence on p, q ; we write ˜ O ( · ) when (poly-)logarithmic factors are ignored. O d, P ( · ) (resp. o d, P ( · ) ) represents big-O (resp. little-o) in probability as d →∞ . Ω( · ) , Θ( · ) are defined analogously. γ is the standard Gaussian distribution in R . We denote the L 2 -norm of a function f with respect to the data distribution (which will be specified) as ∥ f ∥ L 2 . For g : R → R , we denote g i as its i -th exponentiation, and g ( i ) is 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 a single-index model (1.1) where the link function σ ∗ : R → R has information exponent p defined as follows [DH18, BAGJ21].

̸

Definition 1 (Information exponent) . Let { He j } ∞ j =0 denote the normalized Hermite polynomials. The information exponent of g ∈ L 2 ( γ ) , denoted by IE( g ) := p ∈ N + , is the index of the first non-zero Hermite coefficient of g , that is, given g ( z ) = ∑ ∞ i =0 α i He i ( z ) , p := min { i> 0 : α i =0 } .

By definition, when σ ∗ is a degreeq polynomial, we always have p ≤ q . Note that f ∗ contains Θ( d ) parameters to be estimated, and hence information theoretically n ≳ d samples are both suf-

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 to the low-dimensional structure of single-index f ∗ and hence suffer from the curse of dimensionality [YS19, GMMM21, DWY21, BES + 22]. By a standard dimension argument [KMS20, HSSVG21, AAM22], we know that in the isotropic data setting, kernel methods (including neural networks in the lazy regime [JGH18, COB19]) require n ≳ d q samples to learn degreeq polynomials in R d .

Gradient-based Training of NNs. While NNs can easily approximate a single-index model [Bac17], the sample complexity of gradient-based learning established in prior works typically scales as n ≳ d Θ( p ) : in the well-specified setting, [BAGJ21] proved a sample complexity of n = ˜ Θ( d p -1 ) for online SGD, which is later improved to ˜ Θ( d p/ 2 ) by a smoothed objective [DNGL23]; as for the misspecified setting, [BBSS22, DKL + 23] showed that n ≳ d p samples suffice, and in some cases a ˜ Θ( d p -1 ) complexity is achievable [AAM23, OSSW24a]. Consequently, at the information-theoretic limit ( n ≍ d ), existing results can only cover the learning of low information exponent targets [AAM22, BMZ23, BES + 23]. This exponential dependence on p also appears in the CSQ lower bounds [DLS22, AAM22], 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 we do not restrict ourselves to correlational queries, the sample complexity of learning (1.1) can be drastically improved. Specifically, for polynomial σ ∗ , [CM20] gave an SQ algorithm that achieves low generalization error in n = ˜ O ( d ) samples, which is near the information-theoretic limit; the key ingredient is to 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 a 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 ∈ L 2 ( γ ) is defined as the lowest information exponent (IE) after arbitrary L 2 transformation, that is,

The generative exponent is the smallest information exponent obtained by all possible label transformations. By definition we always have p ∗ ≤ p , and the gap between the two indices can be arbitrarily large; for example, for the Hermite polynomials we have IE( He k ) = k whereas GE( He k ) ≤ 2 .

[DPVLB24] established a sample complexity lower bound of n = Ω( d p ∗ / 2 ∨ 1 ) for full SQ learners with polynomial compute (assuming τ ≈ n -1 / 2 ), 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 a sample and runtime complexity that scales with n ≃ d Θ( p ∗ ) , where the dimension dependence is governed by the generative exponent p ∗ instead of the information exponent p .

2.2 Can Gradient Descent Go Beyond Correlational Queries?

Correlational statistical query. A statistical query (SQ) learner [Kea98, Rey20] accesses the target f ∗ through noisy queries ˜ ϕ with error tolerance τ : | ˜ ϕ -E x ,y [ ϕ ( x , y )] | ≤ τ . Lower bound on the performance of SQ algorithm is a 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) E x ,y [ ϕ ( x ) y ] . To see the connection between CSQ and SGD, consider the gradient of expected squared loss for one neuron f w ( x ) :

One can see that information of the target function is encoded in the correlation term in the gradient. To infer the statistical efficiency of GD in the empirical risk minimization setting, we replace the population gradient with the empirical average ∇ w ( 1 n ∑ n i =1 ( f w ( x i ) -y i ) 2 ) , and heuristically equate the CSQ tolerance τ 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 a lower bound stating that a CSQ learner either has access to queries with tolerance τ ≲ d -p/ 4 , or exponentially many queries are needed to learn f ∗ with small population loss. Using the heuristic τ ≈ n -1 / 2 , this suggests a sample complexity lower bound n ≳ d p/ 2 for polynomial time CSQ algorithm. This lower bound can be achieved by a landscape smoothing procedure [DNGL23] (in the well-specified setting), and is conjectured to be optimal for empirical risk minimization with SGD [AAM23].

SGDwithreused data. As previously discussed, the gap between SQ and CSQ algorithms primarily 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 f w ( x ) = σ ( ⟨ x , w ⟩ ) updated by two GD steps using the same data point ( x , y ) , starting from zero initialization w 0 = 0 (we focus on the correlational term in the loss for simplicity):

Under appropriate learning rate scaling η · ∥ x ∥ 2 = Θ(1) , one can see that in the second gradient step, the label y is transformed by the nonlinearity σ ′ , even though the loss function itself is not modified. Based on this observation, [DTA + 24] showed that if the non-CSQ term in (2.1) reduces the information exponent to 1 , then weak recovery (i.e., nontrivial overlap between the first-layer parameters w and index features θ ) can be achieved after two GD steps with n = Θ( d ) samples.

2.3 Challenges in Establishing Statistical Guarantees

Importantly, the analysis in [DTA + 24] does not lead to concrete learnability guarantees for the class of single-index 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 ( ii ) the weak recovery guarantee may not translate to a 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] assumed that the student activation σ can reduce the information exponent of the labels to 1 ; while a 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 σ ∗ ; 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 p ∗ ≥ 2 , its information exponent after arbitrary nonlinear transformation is at 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 nontrivial overlap can be established by a Gr¨ onwall-type argument similar to [BAGJ21]. For p ∗ = 2 , this recovers results on phase retrieval when σ ∗ ( z ) = z 2 which requires n = Ω( d log d ) samples.

From weak recovery to sample complexity. Note that weak recovery (i.e., |⟨ w , θ ⟩| > ε for some small constant ε > 0 ) is 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., |⟨ w , θ ⟩| > 1 -ε ), despite the link misspecification. After the first-layer parameters align with the target function, we train the secondlayer parameters with SGD to learn the link function σ ∗ with the aid of random bias units [DLS22].

3 Learning Polynomial f ∗ in Linear Sample Complexity

We first consider the setting where σ ∗ is polynomial with degree q specified as follows.

Assumption 1. The target function is given as f ∗ ( x ) = σ ∗ ( ⟨ x , θ ⟩ ) , where the link function σ ∗ : R → R admits the Hermite decomposition σ ∗ ( z ) = ∑ q i = p α i He i ( z ) .

For single-index polynomials, we do not expect a computational-to-statistical gap under the SQ class [CM20] - indeed, we will establish learning guarantees near the information theoretic limit n ≍ d .

3.1 Training Algorithm

We train the following two-layer network with N neurons using SGD to minimize the squared loss:

where Θ = ( w j , a j , b j ) N j =1 are trainable parameters, and σ j : R → R is the activation function defined as the sum of Hermite polynomials up to degree C σ : σ j ( z ) := ∑ C σ i =0 β j,i He i ( z ) , where C σ only depends on the degree of link function σ ∗ . Note that we allow each neuron to have a different nonlinearity as indicated by the subscript in σ j ; 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 1 employs a layer-wise training strategy common in the recent feature learning theory literature [DLS22, BES + 22, BBSS22, AAM23, MHWSE23], where in the first stage, we optimize the first-layer parameters { w j } N j =1 with normalized SGD to learn the low-dimensional latent representation (index features θ ), and in the second phase, we train the second-layer { a j } N j =1 to fit the unknown link function σ ∗ .

The most crucial part in Phase I of Algorithm 1 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 us 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 optimization algorithms [AZ18, ZLBH19]; the role of this interpolation is discussed in Section 4.2. We use a projected gradient update ˜ ∇ w L ( w ) = ( I d -w 2 t w 2 t ⊤ ) ∇ w L ( w ) for steps 2 t and 2 t +1 , where ∇ w is the Euclidean gradient; similar use of projection also appeared in [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 I of Algorithm 1 for T = polylog( d ) steps, a subset of parameters w achieve nontrivial overlap with the target direction θ . We denote H ( g ; j ) as the j -th Hermite coefficient of some g ∈ L 2 ( γ ) . Our main theorems handle polynomial activations satisfying the following condition.

Assumption 2. We require the activation function to be a polynomial σ ( z ) = ∑ C σ i =0 β i He i ( z ) and its degree C σ to be sufficiently large so that C σ ≥ C q holds ( C q is defined in Proposition 6). For all 2 ≤ ℓ ≤ C σ and k = 0 , 1 , we assume that H ( σ ( ℓ ) ( σ (1) ) ℓ -1 ; k ) > 0 .

̸

As discussed in Appendix B.1, for a given σ ∗ , the above assumption only needs to be met for one pair of ( k, ℓ ) . Appendix B.1.3 states that H ( σ ( ℓ ) ( σ (1) ) ℓ -1 ; k ) = 0 also suffices if we set the momentum parameter ξ differently. Now we verify this condition for a wide range of polynomial activations.

̸

Lemma 3. Given ℓ ≥ 2 and k ≥ 0 . For C σ ≥ 2 ℓ + k -1 ℓ , if we choose { β i } C σ i =0 where β i is randomly drawn from some non-empty interval [ a i , b i ] , then H ( σ ( ℓ ) ( σ (1) ) ℓ -1 ; k ) = 0 with probability 1.

The next theorem states that n = ˜ Θ( d ) samples are sufficient for SGD to achieve weak recovery.

Theorem 1. Under Assumptions 1 and 2, for suitable choices of hyperparameters η t = ˜ O d ( Nd -1 ) and 1 -ξ t = o d (1) , there exists constant C ( q ) such that after Phase I of Algorithm 1 is run for 2 T 1 , 1 = C ( q ) · d polylog( d ) steps, with high probability, there exists a subset of neurons w 2 T 1 j ∈ W with |W| = ˜ Θ( N ) such that ∣ ∣ ⟨ w 2 T 1 j , θ ⟩ ∣ ∣ > c for some c ≳ 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 is 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 θ , the objective is no longer dominated by the lowest degree in the Hermite expansion. Therefore, to establish strong recovery ( ⟨ w , θ ⟩ > 1 -ε ), we place an additional assumption on the activation function.

Assumption 3. Given the Hermite expansions σ ∗ ( z ) = ∑ q i = p α i He i ( z ) , σ j ( z ) = ∑ C σ i =0 β j,i He i ( z ) , we assume the coefficients satisfy α i β j,i ≥ 0 for p ≤ i ≤ q .

This assumption is easily verified in the well-specified setting σ ∗ = σ [BAGJ21] since α i = β i , 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 degreeq polynomial link function σ ∗ .

Lemma 4. If we set σ j ( z ) = ∑ C σ i =0 β j,i He i ( z ) , where for each neuron we sample β j,i i . i . d . ∼ Unif( {± r i } ) with appropriate constant r i , then Assumption 2 and 3 are satisfied in exp( -Θ( q )) -fraction of neurons.

Note that in our construction of activation functions for both assumptions, we do not exploit knowledge of the link function σ ∗ other than its degree q which decides the constant C σ ; see Appendix B.1 for more discussion of Assumption 3 and Lemma 4. The next theorem shows that by running Phase I for ˜ Θ( d ) 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 j starting from nontrivial overlap c > 0 specified in Theorem 1, if Phase I of Algorithm 1 continues for 2 T 1 , 2 = ˜ Θ d ( dε -2 ) steps with hyperparameters η t = ˜ O d ( Nd -1 ε ) , ξ t = 1 , we achieve 〈 w 2( T 1 , 1 + T 1 , 2 ) j , θ 〉 > 1 -ε with high probability.

The following proposition shows that after strong recovery, training the second-layer parameters in Phase II is sufficient for the NN model (3.1) to achieve small generalization error.

Proposition 5. After Phase I terminates, for suitable λ > 0 , the output of Phase II satisfies

with probability 1 as d →∞ , if we set T 2 = C ( q ) N 4 polylog( d ) ε -4 , N = C ( q )polylog( d ) ε -1 for some constant C ( q ) 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 1 to achieve ε population squared loss, it is sufficient to set

where constant C ( q ) only depends on the target degree q (although exponentially). Hence we may set ε -1 ≍ polylog d to conclude an almost-linear sample and computational complexity for learning arbitrary single-index polynomials up to o d (1) population error.

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, we first establish the existence of nonlinear label transformation that ( i ) reduces the information exponent, and ( ii ) can be easily extracted from SGD updates. If we ignore 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 is 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.

(i) There exists some i ≤ C q ∈ N + such that IE( g i ) ≤ 2 , where constant C q only depends on q .

(ii) Let g odd : R → R be the odd part of g with ∥ g odd ∥ 2 L 2 ( γ ) ≥ ρ > 0 . Then there exists some i ≤ C q,ρ ∈ N + such that IE( g i ) = 1 , where constant C q,ρ only depends on q and ρ .

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 ∈ N + only depending on the degree of σ ∗ such that raising the function to the i -th power reduces the information exponent to 1 (this implies the generative exponent p ∗ = 1 ). For even σ ∗ , the information exponent after arbitrary transformation is at least 2 ( p ∗ = 2 ), which can also be attained by monomial transformation. Furthermore, we provide a uniform upper-bound on the required degree of transformation i via a 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 σ ′ . In Section 4.2, we build upon this observation to show that Phase I of Algorithm 1 achieves weak recovery using n ≳ d polylog( d ) samples.

Intuition behind the analysis. Our proof is inspired by [CM20] which introduced a (nonpolynomial) label transformation that reduces the information exponent of any degreeq polynomial to at most 2 . To prove the existence of monomial transformation for the same purpose, we first show that for a fixed link function σ ∗ , there exists some i 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 to define a test function and obtain a uniform bound on i . As for the polynomial transformation for non-even functions, we exploit the asymmetry of σ ∗ to further reduce the information exponent to 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 η t ≡ η . When one neuron f w ( x ) = σ ( ⟨ x , w ⟩ ) is updated by two GD steps using the same sample ( x , y ) , starting from w 0 := ω , the alignment with θ becomes

We take η ≤ c η d -1 with a small constant c η so that η ∥ x ∥ 2 ≪ 1 with high probability. Crucially, the strength of each term in (4.1) can vary depending on properties of the unknown link function σ ∗ . Hence a 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 , 2 :

See Lemma 16 for the formal version. For p ∗ = 1 , taking expectation immediately yields that weak recovery within ( η (1 -ξ ) γ ) -1 = O ( d ) steps. For p ∗ = 2 , ⟨ θ , w 2 t j ⟩ =: κ t can be approximated by a differential equation d κ t d t = η (1 -ξ ) γκ t . Solving this yields κ t = κ 0 exp( η (1 -ξ ) γt ) ≈ d -1 2 exp( η (1 -ξ ) γt ) , and weak recovery is obtained within t ≲ ( η (1 -ξ ) γ ) -1 · log d = O ( d log d ) steps, similar to the analysis in [BAGJ21].

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 -g and projection P w = I d -ww ⊤ , the spherical gradient changes the alignment as ⟨ θ , w t +1 ⟩ = 〈 θ , w t + ηP w g ∥ w t + ηP w g ∥ 〉 ≥ ⟨ θ , w t ⟩ + η ⟨ θ , g ⟩ -1 2 η 2 ∥ g ∥ 2 ⟨ θ , w t ⟩ + (h.o.t.), see [BAGJ21, DNGL23]. Here η ⟨ θ , g ⟩ corresponds to the signal, and -1 2 η 2 ∥ g ∥ 2 ⟨ θ , w t ⟩ comes from normalization. Thus, taking η sufficiently small, the normalization error shrinks faster than the signal. However, in our case the signal shrinks at the rate of c I η (recall that η = c η d -1 ), and hence taking a smaller step may not improve the signal-to-noise ratio when the degree of transformation I is large. The interpolation step in Algorithm 1 reduces the effect of normalization without shrinking the signal too much. In particular, by setting ξ = 1 -˜ η , we see that the signal is affected by a factor of ˜ η whereas the normalization error shrinks by ˜ η 2 ; this allows us to boost the signal-to-noise ratio by taking ˜ η small.

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 ℓ 2 regularization. Since the objective is strongly convex, gradient-based optimization can efficiently minimize the empirical loss. In Appendix B.6, the learnability guarantee follows from standard analysis analogous to that in [AAM22, DLS22, BES + 22], where we construct a 'certificate' second-layer a ∗ ∈ R N that achieves small loss and small norm:

from which the population loss of the regularized empirical risk minimizer can be bounded via standard Rademacher complexity argument. To construct such a certificate, we make use of the random bias units { b j } N j =1 to approximate the link function σ ∗ 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 to 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-polynomial) σ ∗ with generative exponent p ∗ > 2 (i.e., label transformations cannot lower the information exponent to 2 ) and thus not learnable by SQ algorithms in linear sample complexity.

Nevertheless, for a single-index model with generative exponent p ∗ , we know there exists an 'optimal' label transformation that reduces the information exponent to p ∗ . If SGD can make use of such transformation, then from the arguments in [BAGJ21], it is natural to conjecture that a sample size of n ≃ d p ∗ -1 is sufficient. In this section we confirm this intuition by proving that SGD with data reuse (Algorithm 1) indeed matches this complexity. The following lemma is an analogue of Proposition 6 stating that polynomial transformations are sufficient to lower the information exponent.

Lemma 8. Given σ ∗ with generative exponent p ∗ ∈ N + . Suppose we can take an orthonormal polynomial basis { ϕ k } k for the space L 2 ( P y ) with inner product ⟨ f, g ⟩ = E y = σ ∗ ( z ) [ f ( y ) g ( y )] . Then there exists some degree of transformation I ∈ N ∗ such that IE( σ I ∗ ) = p ∗ .

We outline the differences and additional technical challenges to handle the GE( σ ∗ ) > 2 setting.

For general L 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 8 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 generalization error using a neural network with polynomial nonlinearity. We therefore need to use an activation function with universal function approximation ability.

5.1 Sample Complexity for Weak Recovery

We first show that Algorithm 1 achieves weak recovery (i.e., nontrivial overlap with the ground truth θ ) with a 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 σ ∗ has generative exponent p ∗ , and let I ∈ N + be the smallest degree of monomial transformation that lowers the information exponent to p ∗ (i.e., IE( σ I ∗ ) = p ). We can find a student activation function σ depending only on p, p ∗ and I , such that if we take η 2 t , η 2 t +1 = c η Nd -1 , ξ 2( t +1) = 1 -c ξ d -( p ∗-2) + / 2 for small c η , c ξ = o d (1) , and set

then if the initial alignment ⟨ w 0 , θ ⟩ ≥ 2 c -1 d -1 / 2 , there exists τ ≤ T such that for all τ ≥ τ ,

η ∗ 1 , 1 ∗ .

Proposition 9 is a generalization of Theorem 1 beyond polynomial σ ∗ (the proof of both results are presented in Appendix B.3,B.4), and can be interpreted as an SQ counterpart to [BAGJ21]: we establish a sufficient sample size of n ≃ d ( p ∗ -1) ∨ 1 for Algorithm 1 to exit the search phase, which is parallel to the n ≃ d ( p -1) ∨ 1 rate for one-pass SGD (note that our rates are slightly sharper due to logarithmic factors removed, since c -1 ξ can grow arbitrarily slowly with d ). For high generative exponent σ ∗ with p ∗ > 2 , we no longer match the information theoretically optimal sample complexity n ≍ d , which is consistent with the computational-to-statistical gap observed in [DPVLB24].

5.2 Generalization Error Guarantee

After Phase I of Algorithm 1, we learn the unknown link function σ ∗ via training the second-layer. To approximate non-polynomial functions, we introduce a ReLU component in the student nonlinearity σ (see Lemma 12 for discussions), and make use of the approximation result for the (univariate) ReLU kernel in [BBSS22], which handles general σ ∗ 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 σ ∗ has generative exponent p ∗ ∈ N ∗ and satisfies σ ∗ , σ ′′ ∗ ∈ L 4 ( γ ) . For appropriately chosen activation function σ (see Appendix B.1.2), a neural network (3.1) with N = ˜ Θ(1) neurons optimized by Algorithm 1 achieves small population loss E x [( f ˆ Θ ( x ) -f ∗ ( x )) 2 ] = o d, P (1) with a sample complexity of n = ˜ Θ( d ( p ∗ -1) ∨ 1 ) .

See Appendix B.6 for the full statement with ε 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 1 to achieve low test error also scales with d ( p ∗ -1) ∨ 1 .

6 Conclusion and Future Directions

We showed that a two-layer neural network (3.1) trained by SGD with reused batch can learn singleindex model (with generative exponent p ∗ ) using n ≃ d ( p ∗ -1) ∨ 1 samples and compute; in particular, when the link function σ ∗ is polynomial, we established a sample complexity of n = ˜ O ( dε -2 ) to achieve ε population loss, which is almost information theoretically optimal. Our analysis is based on the observation that by reusing the same training data twice in the gradient computation, a noncorrelational 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.

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 simultaneously [Gla23] (as reported in Figure 1) also achieves the same statistical efficiency.

References

[AAM22] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The mergedstaircase 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 saddle-to-saddle dynamics. In The Thirty Sixth Annual Conference on Learning Theory , pages 2552-2623. PMLR, 2023.

[ADK + 24] Luca Arnaboldi, Yatin Dandi, Florent Krzakala, Luca Pesce, and Ludovic Stephan. Repetita iuvant: Data repetition allows sgd to 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 , 18(1):629-681, 2017.

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

[BAGJ22] Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems , 35:25349-25362, 2022.

[BBPV23] Alberto Bietti, Joan Bruna, and Loucas Pillaud-Vivien. On learning Gaussian multiindex models with gradient flow. arXiv preprint arXiv:2310.19793 , 2023.

[BBSS22] Alberto Bietti, Joan Bruna, Clayton Sanford, and Min Jae Song. Learning singleindex models with shallow neural networks. Advances in Neural Information Processing Systems , 35:9768-9783, 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-dimensional 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-395, 2002.

[BKM + 19] Jean Barbier, Florent Krzakala, Nicolas Macris, L´ eo Miolane, and Lenka Zdeborov´ a. Optimal errors and phase transitions in high-dimensional generalized linear models. Proceedings of the National Academy of Sciences , 116(12):5451-5460, 2019.

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

[BMZ23] Rapha¨ el Berthier, Andrea Montanari, and Kangjie Zhou. Learning time-scales in two-layers neural networks. arXiv preprint arXiv:2303.00055 , 2023.

[CCM11] Seok-Ho Chang, Pamela C Cosman, and Laurence B Milstein. Chernoff-type bounds for the Gaussian error function. IEEE Transactions on Communications , 59(11):2939-2944, 2011.

[CM20] Sitan Chen and Raghu Meka. Learning polynomials in few relevant dimensions. In Conference on Learning Theory , 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-Woodfin, Courtney Paquette, Elliot Paquette, and Inbar Seroussi. Hitting the high-dimensional notes: An ode for sgd learning dynamics on glms and multi-index models. arXiv preprint arXiv:2308.08977 , 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. Statistical-computational trade-offs in tensor pca and related problems via communication complexity. The Annals of Statistics , 52(1):131-156, 2024.

[DKL + 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 arXiv:2305.18270 , 2023.

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

[DNGL23] Alex Damian, Eshaan Nichani, Rong Ge, and Jason D. Lee. Smoothing the landscape boosts the signal for SGD: Optimal sample complexity for learning single index models. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

[DPVLB24] Alex Damian, Loucas Pillaud-Vivien, Jason D Lee, and Joan Bruna. The computational complexity of learning gaussian single-index models. arXiv preprint arXiv:2403.05529 , 2024.

[DTA + 24] Yatin Dandi, Emanuele Troiani, Luca Arnaboldi, Luca Pesce, Lenka Zdeborov´ a, 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. In International Conference on Machine Learning , pages 2804-2814. PMLR, 2021.

[Gla23] 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. Linearized two-layers neural networks in high dimension. The Annals of Statistics , 49(2):1029-1054, 2021.

[HSSVG21] Daniel Hsu, Clayton H Sanford, Rocco Servedio, and Emmanouil Vasileios VlatakisGkaragkounis. On the approximation power of two-layer networks of random relus. In Conference on Learning Theory , pages 2423-2461. PMLR, 2021.

[JGH18] Arthur Jacot, Franck Gabriel, and Cl´ ement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems , pages 8571-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.

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

[MHWSE23] Alireza Mousavi-Hosseini, Denny Wu, Taiji Suzuki, and Murat A. Erdogdu. Gradient-based feature learning under structured data. In Thirty-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.

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

[NDL23] Eshaan Nichani, Alex Damian, and Jason D Lee. Provable guarantees for nonlinear feature learning in three-layer neural networks. Advances in Neural Information Processing Systems , 36, 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 combinations. In Conference on Learning Theory . PMLR, 2024.

[OSSW24b] Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu. Pretrained transformer efficiently learns low-dimensional target functions in-context. arXiv preprint arXiv:2411.02544 , 2024.

[Rey20] Lev Reyzin. Statistical queries and statistical algorithms: Foundations and applications. arXiv preprint arXiv:2004.00557 , 2020.

[Sch80] Jacob T Schwartz. Fast probabilistic algorithms for verification of polynomial identities. Journal of the ACM (JACM) , 27(4):701-717, 1980.

[TS24] Shokichi Takakura and Taiji Suzuki. Mean-field analysis on two-layer neural networks from a kernel perspective. arXiv preprint arXiv:2403.14917 , 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 optimizer: k steps forward, 1 step back. Advances in neural information processing systems , 32, 2019.

Table of Contents

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 τ ( f ) = max -2 ≤ t ≤ 2 | f ( t ) | . This entails that if | f ( t ) | ≥ τ ( f ) , then we have | t | > 2 . Consider the following expectation:

We evaluate the case when i is even. (A.1) can be lower bounded as

Note that P [ | f ( t ) | ≥ 2 τ ( f )] is positive (since f is polynomial) and independent of i , while 2 -i decays to 0 as i increases. Therefore, for sufficiently large i ∈ N , (A.1) is positive and hence IE( f i ) ≤ 2 . The subsequent analysis aims to provide an upper bound on i .

(i-2): Construction of test function. We introduce the notation H ( · ; j ) which takes any function (in L 1 ) and returns its j -th Hermite coefficient. We consider the following test function:

(i-3): Lower bound of test function via compactness. Let F q be a set of polynomials with degree up to q with unit L 2 norm. Because H ( f ) is positive for any f ∈ F q , H ( f i ; 2) is continuous with respect to f , and F q is a compact set, inf f ∈F q H ( f ) admits a minimum value H 0 which is positive.

(i-4): Conclusion via hypercontractivity. Because f is a polynomial with degree at most q , Gaussian hypercontractivity [O'D14] yields that

Therefore, for all polynomials in F q , a partial sum of (A.2) is uniformly bounded by

Combining this with the fact that H ( f ) ≥ H 0 > 0 , we know that there exists some C q ≤ 1 + log 2 ( H -1 0 ) such that

for all polynomials in F q . This means that there is at least one i ≤ C q such that H ( f i ; 2) = 0 .

̸

A.2 Proof for Non-even Functions ( ii )

(ii-1): Monomials reducing the information exponent. We prove that some exponentiation of g := f 2 has non-zero first Hermite coefficient. Denote g odd as the odd part of g , and similarly g even . Let υ ( g ) ∈ R + be the value at which the followings hold:

(a) g odd ( t ) > 0 for all t ≥ υ ( g ) and g odd ( t ) < 0 for all t ≤ -υ ( g ) .

(b) g even ( t ) > | g odd ( t ) | for all t ≥ υ ( g ) and t ≤ -υ ( g ) .

(c) For for all t ≥ υ ( g ) and t ≤ -υ ( g ) , g ( s ) = g ( t ) (as an equation of s ) only has two real-valued solutions with opposing signs.

Such threshold υ ( g ) exists because the tail of g = f 2 is dominated by the highest degree which is even. Then, we let τ ( g ) = max -υ ( g ) ≤ t ≤ υ ( g ) | g ( t ) | .

Consider the following expectation:

(A.3) is decomposed as

We first evaluate the first term. Because of (c), g ( t ) = 3 τ ( f ) has two real-valued solutions α < 0 < β . Because of (a) and (b), g ( β ) = g even ( β ) + g odd ( β ) = 3 τ ( f ) > g even ( -β ) + g odd ( -β ) = g odd ( -β ) . Because lim t →-∞ g odd ( t ) = + ∞ , and α is the only solution in t < 0 , we have α < -β . Moreover, for all t > β , we have g ( t ) = g even ( t ) + g odd ( t ) > g even ( -t ) + g odd ( -t ) = g odd ( -t ) . Combining the above, the first term of (A.4) is bounded as

Following the exact same reasoning, we know that the second term of (A.4) is positive. Finally, the third term which is bounded by

Putting things together,

The first term is independent of i and positive, while the second term goes to zero as i grows. Therefore, there exists some i such that IE( g i ; 1) = 1 .

(ii-2): Construction of test function. This time we consider the following function:

(ii-3): Lower bound of test function via compactness. Let F q be a set of unit L 2 -norm polynomials with degree up to q and E t ∼N (0 , 1) [ f odd ( t ) 2 ] ≥ c . Since H ( f ) is always positive for F q , H ( f ) is continuous with respect to f , and F q is a compact set, inf f ∈F q H ( f ) has the minimum value H 0 that is positive. Note that H ( f ) might depends on c .

(ii-4): Conclusion via hypercontractivity. Using the same argument as in (i), we conclude that there exists some C q,c such that

Because H 0 depends on c , C q,c depends on c as well as q .

A.3 Proof for Non-Polynomial Functions

For non-polynomial link functions, we note that similar to [DPVLB24], the existence of polynomial basis is needed to exclude extreme cases, and we cannot upper bound the required degree I 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 z ∼ N (0 , 1) and y = σ ∗ ( z ) . We define ζ p ∗ ( y ) = E [ 1 √ p ∗ ! He p ∗ ( z ) | y ] and its basis expansion ζ p ∗ ( y ) = ∑ ∞ k =0 v k ϕ k . Let K be a smallest integer such that v k = 0 . Then, there exists an integer with I ≤ K such that IE( y I ) = p ∗ . Indeed,

̸

which means that at least one of y, y 2 , · · · , y K yields a non-zero p ∗ -th Hermite coefficient.

B SGD with Reused Batch

In this section we show that Algorithm 1 learns single-index models in ˜ O ( d 1 ∨ ( p ∗ -1) ) samples with high probability. The algorithm trains the first layer for T 1 SGD steps, where we sample a new data point in every two steps. The first layer training is further divided into two phases: weak recovery ( w ⊤ θ ≳ 1 ) and strong recovery ( ∥ w -θ ∥ ≲ ε ). Then, we learn the second layer parameters.

Specifically, Section B.2 shows that at initialization, a (nearly) constant fraction of neurons has alignment w ⊤ θ beyond a certain threshold. We focus on such neurons in the first phase of training. Section B.3 lower bounds the expected update of alignment w ⊤ θ of two gradient steps, and Section B.4 establishes that the neurons achieve weak recovery within 2 T 1 , 1 = ˜ O ( d 1 ∨ ( p ∗ -1) ) steps. Section B.5 discusses how to convert weak recovery to strong recovery using 2 T 1 , 2 = ˜ O ( dε -2 ) more steps. We let T 1 = 2 T 1 , 1 + 2 T 1 , 2 . Finally, Section B.6 analyzes second layer training and concludes the proof.

In the following proofs, we use several constants, which depends on d at most at most polylogarithmically. Specifically, asymptotic strength of the constants is ordered as follows.

Here, c η and δ should satisfy lim d →∞ c η = lim d →∞ δ = 0 , 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 H ( σ I ∗ ; p ∗ ) . To ensure that one signal term (from the Taylor series) is isolated, taking η ≍ d -1 with a sufficiently small constant is insufficient but η ≍ c η d -1 with arbitrarily slow c η suffices. Also, to guarantee that the failure probability is o d (1) , we require δ to be o d (1) . c ξ can also decay arbitrarily slowly, as long as it satisfies c ξ ≲ δ poly( c -1 η ) . C 4 = polylog( d ) will be used to represent any polylogarithmic factor that comes from high probability bounds.

For the first-layer training, we can reduce the argument into training of one neuron using the correlation loss as follows. At each step, the gradient update (Line 8 of Algorithm 1) is written as

While the second term scales with η t c 2 a N -1 , the third term scales with η t c a N -1 . Thus, by setting c a sufficiently small, we can ignore the interaction between neurons. We will show that the strength of the signal in the direction of θ is at least ( κ t j ) p ∗ -1 ≳ d -p ∗-1 2 (up to a polylogarithmic factor, and p ∗ = GE( σ ∗ ) ). On the other hand, we can easily see that θ ⊤ ( 1 N ∑ N j =1 σ j ( w t j ⊤ x ) )( ˜ ∇ w σ j ( w t j ⊤ x ) ) is bounded by ˜ O (1) with high probability. Therefore, by simply letting c a = ˜ Θ( d -p ∗-1 2 ) , we can ignore the effect of the second term in (B.1). Moreover, for simplicity, we will reparameterize 2 η t c a N as η t below. Consequently, we may analyze the following update

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 on 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 a specific activation function tailored to different link functions, a randomized activation function satisfies all required assumptions with probability Ω(1) .

In the following, we write the student activation function as

with coefficients { β j,i } C σ i =0 (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-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 σ ∗ . Note that GE( σ ∗ ) = 1 or 2 holds (see Proposition 6). Let I ≤ C q (according to Proposition 6) be the smallest integer such that IE( σ I ∗ ) = GE( σ ∗ ) = p ∗ and C σ be the degree of the activation function.

Weak recovery: α p ∗ β p ∗ > 0 (covered by Assumption 3).

q

̸

Strong recovery: ∑ j = p ∗ j ! α j β j s j -1 > 0 for all s > 0 (covered by Assumption 2). Approximation (ridge regression): β i = 0 for some i ≥ q (covered by Assumption 3). (II) If 2 ≤ I = { min i | IE( σ I ∗ ) = GE( σ ∗ ) = p ∗ } ≤ C σ .

̸

Weak recovery: H (( σ ∗ ) I ; p ∗ ) H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) > 0 (covered by Assumption 3). Strong recovery: ∑ q j = p ∗ +1 j ! α j β j s j -1 > 0 for all s > 0 (covered by Assumption 2). Approximation: β i = 0 for some i ≥ q (covered by Assumption 3).

Note that it is difficult to construct a deterministic activation function that satisfies all of the assumptions for any link function σ ∗ (the simplest counterexample is to consider -σ ∗ which flips the Hermite coefficients). Instead, we show the existence of randomized construction of such an activation function that satisfies all of the assumptions on the activation function simultaneously 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 .

Lemma 11. There exists a randomized activation function sampled from a discrete set such that the above conditions hold with constant probability.

Proof. Let c be a sufficiently small constant only used in this proof and C σ be the minimum odd integer with C σ ≥ max { C q +1 , q +2 , 3 } , where C q was introduced in Proposition 6. With probability 1 2 , we let β 1 ∼ Unif( {± 1 } ) , and β j ∼ Unif( {± c } ) for 2 ≤ j ≤ C σ . With probability 1 2 , we let β j ∼ Unif( {± c } ) for 1 ≤ j ≤ C σ -2 and β C σ -1 = β C σ ∼ Unif( {± 1 } ) .

We first consider (I). When β 1 ∼ Unif( {± 1 } ) , and β j ∼ Unif( {± c } ) for 2 ≤ j ≤ C σ , it is easy to see sign( α j ) = sign( β j ) for all j = 1 , · · · , q hold with probability at least 2 -q , which is sufficient to satisfy (I).

We then consider (II). First focus on the case when p ∗ = 1 and I is even. When β 1 ∼ Unif( {± 1 } ) and β j ∼ Unif( {± c } ) for 2 ≤ j ≤ C σ , by taking c sufficiently small, we have

When I is even, by adjusting the sign of β 1 , H ( σ ( I ) ( σ (1) ) I -1 ; 0) is non-zero and has the same sign as H (( σ ∗ ) I ; 1) with probability 1 2 . Note that the sign of β 1 is independent from whether ∑ q j =2 j ! α j β j s j -1 > 0 for all s > 0 holds. This holds with probability at least 2 -q +1 . Thus we verified (II) for p ∗ = 1 and even I .

For p ∗ = 1 and odd I , consider β j ∼ Unif( {± c } ) for 1 ≤ j ≤ C σ -2 and β C σ -1 = β C σ ∼ Unif( {± 1 } ) . Note that ∑ q j =2 j ! α j β j s j -1 > 0 for all s > 0 (this is the condition for strong recovery) and the condition for ridge regression also holds. Furthermore, the sign of H (( He C σ + He C σ -1 ) ( I ) (( He C σ + He C σ -1 ) (1) ) I -1 ; 0) is ± 1 with equiprobability, independent of β 2 , . . . , β q . Therefore, by taking c sufficiently small, we can obtain the desired sign of H ( σ ( I ) ( σ (1) ) I -1 ; 0) . Thus we proved (II) for p ∗ = 1 and odd I .

Regarding (II) for p ∗ = 2 and even I , when β 1 ∼ Unif( {± 1 } ) and β j ∼ Unif( {± c } ) for 2 ≤ j ≤ C σ , we have

Thus, similar to (II) with p ∗ = 1 and even I , we get (II) for p ∗ = 2 and even I .

Finally, consider (II) for p ∗ = 2 and odd I . When β j ∼ Unif( {± c } ) for 1 ≤ j ≤ C σ -2 and β C σ -1 = β C σ ∼ Unif( {± 1 } ) , the sign of H (( He C σ + He C σ -1 ) ( I ) (( He C σ + He C σ -1 ) (1) ) I -1 ; 1) is ± 1 with equiprobability when I is odd, and this term dominates the others in H ( σ ( I ) ( σ (1) ) I -1 ; 1) . Thus, (II) for p ∗ = 2 and odd I holds similarly to (II) for p ∗ = 1 and odd I .

Now we have obtained the assertion for all cases.

B.1.2 For general link functions

Now we consider non-polynomial link functions with potentially large generative exponent p ∗ = GE( σ ∗ ) ≥ 2 . For weak and strong recovery to succeed, the conditions on the activation function are essentially the same as those for polynomial link functions:

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 a non-polynomial σ ∗ , we introduce the following condition on the activation function. We sample σ j from a discrete set (with bounded cardinality). Let J be an index set such that the coefficients of σ j ( 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-index model can be well-approximated by a linear combination of student neurons.

Assumption 4. When b j ∼ Unif([ -C b , C b ]) where ( C b = polylog( d )) and x 1 , . . . , x T 2 ∼ N (0 , I d ) , there exists a set of coefficients a 1 , . . . , a | J | such that

holds with coefficients of reasonable magnitudes ∑ j ∈ J a 2 j = Θ( | J | ) with high probability (w.r.t. the randomness of b j and x i ). Moreover, E x [ σ j ( θ ⊤ j x + b j ) 4 ] ≤ polylog( d ) for all j with high probability (w.r.t. the randomness of b j ).

The following lemma states that we can design a randomized activation function that satisfies all of the above assumptions with probability Ω(1) , as long as the link function σ satisfies Assumption 4 for σ = ReLU . In other words, we are able to cover the class of link functions σ that can be efficiently approximated by a 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( σ I ∗ ) = GE( σ ∗ ) as we had C q in the polynomial case. Consequently, our student activation is not entirely agnostic to the link function σ ∗ , as we require knowledge of p (information exponent), p ∗ and I .

Lemma 12. Suppose the target link function σ ∗ satisfies Assumption 4 for σ j = ReLU . There exists a randomized activation sampled from a 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 σ j = ReLU is satisfied for broad class of functions, according to Lemma 4.4 and 4.5 of [BBSS22]. Specifically, taking τ = 1 / 2 and λ = N -1 yields that E x [( 1 | J | ∑ j ∈ J a j σ j ( θ ⊤ x + b j ) -σ ∗ ( θ ⊤ x )) 2 ] ≤ N -2 7 . Although they sample b j from Gaussian N (0 , 2) , the result translates to uniform sampling of biases from [ -C b , C b ] by introducing additional logarithmic factor.

Lemma 13 (Lemma 4.4, 4.5 of [BBSS22]) . Suppose that E z ∼N (0 , 1) [ σ ∗ ( z ) 4 ] , E z ∼N (0 , 1) [ σ ′′ ∗ ( z ) 4 ] < ∞ . Then, Assumption 4 with σ j = ReLU holds with ε = N -1 7 and C b ≃ √ log d .

Proof of Lemma 12. We show the existence of suitable σ in two steps: first we construct a randomized polynomial activation function that satisfies conditions (I)(II) with constant probability; then we add a small ReLU perturbation so that the activation can approximate non-polynomial σ ∗ .

Recall p ∈ N + is the information exponent of σ ∗ . We first show that there exists a randomized polynomial activation that satisfies the conditions for weak and strong recovery with probability Ω(1) . Note that the issue of differentiability and bounded moment is avoided when we focus on the polynomial activation functions. We specify the following two distributions. With probability 1 2 , let β 1 ∼ Unif( {-1 , 1 } ) , β j ∼ Unif( {-c, c } ) for j = 1 , · · · , p ∗ + I -1 and β j = 0 otherwise, where c > 0 is a sufficiently small constant. With probability 1 2 , let β 1 = Unif( {-1 , 1 } ) , β 2 = Unif( {-c, c } ) , β j = Unif( {-c 2 , c 2 } ) for all 2 ≤ j ≤ ( p ∗ + I ) ∨ p for a sufficiently small constant c > 0 , and β j = 0 otherwise.

Regarding (I), consider the case when the coefficients are sampled from the first distribution, and | β j |≪ 1 except for j = p ∗ . Then, ∑ ∞ j = p ∗ j ! α j β j s j -1 ≈ p ∗ ! α p ∗ β p ∗ s p ∗ . Choosing the sign of β p ∗ , we have that the assumption holds with probability Ω(1) .

Regarding (II) with even I , consider coefficients sampled from the first distribution, and Sign( β j ) = Sign( α j ) for j ≤ ( p ∗ + I -1) ∨ p . Then, ∑ ∞ j = p ∗ +1 j ! α j β j s j -1 > 0 for all s > 0 . Also, similarly to (B.2),

By flipping the sign of β 1 , we can change the sign of H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) . Thus, (II) for even I is satisfied by a randomized choice of β 1 .

For (II) with odd I , consider coefficients sampled from the second distribution, and Sign( β j ) = Sign( α j ) for j ≤ ( p ∗ + I ) ∨ p . Then, ∑ ∞ j = p ∗ +1 j ! α j β j s j -1 > 0 for all s > 0 .

By flipping the sign of β 1 , we can change the sign of H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) . Thus, (II) for odd I is satisfied by a randomized choice of β 1 .

Therefore, we have constructed a randomized polynomial activation σ that satisfies all of the conditions for weak and strong recovery. Now we provide a sketch of reasoning that when the link function σ ∗ is well-approximated by ReLU as specified in Assumption 4, we can find some σ that additionally satisfies Assumption 4 by introducing a small ReLU component. Specifically, we add c R · ReLU to the activation function with probability 1 2 , with a sufficiently small c R = ˜ Ω(1) , e.g., c R = ( log d ) -C for some C > 0 . When a two-layer ReLU network approximates σ ∗ that satisfies Assumption 4, by using the neurons with added ReLU component, σ ∗ can be approximated up to some polynomial residual with degree ( p ∗ + I ) ∨ p . And by using the remaining polynomial neurons, we can approximate the additional polynomial terms in σ ∗ (see Lemma 22,23). Subtracting the latter from the former, we obtain the desired approximation result. When c R is sufficiently small, this additional term does not impact the conditions for weak and strong recovery and the moment calculations; similarly, since c R ≪ 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 design 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-wise training analysis [AAM22].

B.1.3 More Discussion on Assumption 2

Assumption 2 requires H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) is not zero and has the same sign as H ( σ I ∗ ; p ∗ ) . We remark that if we allow a negative momentum parameter larger than 1 , i.e., setting ξ 2( t +1) = 1+ c ξ d -( p ∗-2) + 2 , we can negate the opposite sign of H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) (see Lemma 16), and the subsequent analysis still holds. Therefore, what we essentially need is H ( σ ( i ) ( σ (1) ) i -1 ; k ) = 0 . Lemma 3 confirms that it is satisfied by almost all polynomials:

̸

Proof of Lemma 3. We note that H ( σ ( i ) ( σ (1) ) i -1 ; k ) = E [ σ ( i ) ( σ (1) ) i -1 He k ] is a polynomial of { β j } C σ j =0 . This polynomial is not identically equal to zero. To confirm this, consider σ = x C σ + x C σ -1 . Because σ ( i ) ( σ (1) ) i -1 is expanded as a sum of x l ( i ( C σ -3) ≤ l ≤ i ( C σ -2) + 1 with positive coefficients and each x l is a sum of He l , He l -2 · · · with positive coefficients, σ ( i ) ( σ (1) ) i -1 has all positive Hermite coefficients for degree 0 , 1 , · · · , i ( C σ -2) + 1 . If k ≤ i ( C σ -2) + 1 ,

this choice of σ yields H ( σ ( i ) ( σ (1) ) i -1 ; k ) > 0 , which confirms that H ( σ ( i ) ( σ (1) ) i -1 ; k ) as a polynomial of { β j } C σ j =0 is not identically equal to zero. Hence the assertion follows from so-called Schwartz-Zippel Lemma [Sch80], or the fact that zeros of a non-zero polynomial form a measurezero set.

B.2 Initialization

We first consider the initial alignment. In the following sections, we focus on the neurons that satisfy κ 0 j = θ ⊤ w 0 j ≥ 2 c -1 η d -1 2 at the initialization. The following lemma states that roughly a constant portion of the neurons satisfy the initial alignment condition upon random initialization. In particular, if we take c η = Ω((log log d ) -1 2 ) , the fraction of neurons that satisfy the initial alignment condition is at least e -16 c -2 η = ˜ Ω(1) . Let us write C 2 = c -1 η for simplicity in the following.

Lemma 14. At the time of initialization, κ 0 j = θ ⊤ w 0 satisfies the following:

We make use of the following lemma.

Lemma 15 (Theorem 2 of [CCM11]) . For any β > 1 and s ∈ R , we have

Proof of Lemma 14. Because κ 0 = v ⊤ w d = e ⊤ 1 g ∥ g ∥ , where g ∼ N (0 , I d ) ,

where we used Lemma 15 for the final inequality. By letting β = 2 , we have that P [ κ 0 j ≥ C 2 d -1 2 ] ≳ e -16 C 2 2 . Because of the symmetry, P [ κ 0 j ≤ 2 C 2 d -1 2 ] = P [ κ 0 j ≥ 2 C 2 d -1 2 ] .

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 example.

Lemma 16. Let η 2 t , η 2 t +1 = η = c η d -1 , ξ 2( t +1) = ξ = 1 -c ξ d -( p ∗-2) + 2 . Suppose that the link function satisfies IE( σ I ∗ ) = GE( σ ∗ ) = p ∗ (we choose the smallest such I ) and activation functions satisfy all of the assumptions in Section B.1 for weak recovery. Then, for w 2 t with c -1 η d -1 2 ≤ θ ⊤ w 2 t ≤ c I η , the alignment θ ⊤ w 2( t +1) can be evaluated as,

Here c σ = p ∗ ! α p ∗ β p ∗ (when IE( σ ∗ ) = GE( σ ∗ ) ) or c σ = p ∗ ! H ( σ I ∗ ; p ∗ ) H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) 2( I -1)! (otherwise), and ν 2 t is a mean-zero sub-exponential random variable.

Proof. The expected alignment θ ⊤ w 2( t +1) after two gradient steps from w 2 t = ω using the same sample ( x , y ) , step size η 2 t = η 2 t +1 = η = c η d -1 and momentum parameter ξ 2( t +1) = ξ =

1 -c ξ d -( p ∗-2) + 2 is evaluated as follows. With a projection matrix P ω = I -ωω ⊤ , the first step updates the weight as

and the next gradient step with the same sample is computed as

here we used the notation ∥ θ ∥ 2 A = θ ⊤ A θ for a vector θ ∈ R d and a positive symmetric matric A ∈ R d × d . From (B.3) and (B.4), the parameter after the two steps is obtained as

where

Finally, the normalization step yields

Therefore, by writing θ ⊤ w 2 t = κ 2 t , the update of the alignment is

We can easily see that E [ ∥ g 2 t ∥ 2 ] ≲ d and E [ | θ ⊤ g 2 t |∥ g 2 t ∥ 2 ] ≲ d , which implies that the expectation of the last two terms of (B.5) is bounded by ≲ κ 2 t c 2 η c 2 ξ d -( p ∗ -1) ∨ 1 ∨ c 3 η c 3 ξ d -( 3 p ∗ 2 -1) ∨ 2 ≤ c 2 η c 2 ξ d -( p ∗ -1) ∨ 1 κ 2 t .

Now we bound E [ θ ⊤ g 2 t ] by ≳ c I -1 η κ p ∗ -1 . 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 B.1.2. We omit the subscript 2 t in the following for simplicity. We divide the analysis into the two cases.

(I) If I = 1 ⇔ IE( σ ∗ ) = GE( σ ∗ ) = p ∗ . For the first term of E [ θ ⊤ g ] , we have

For the second term of E [ θ ⊤ g ] , the following decomposition can be made.

We evaluate each term in the summation. We need to show that although ∥ x ∥ 2 P ω is a function of θ ⊤ x and ω ⊤ x , it is mostly independent from the two quantities. To verify this, let e = θ -( θ ⊤ ω ) ω ∥ θ -( θ ⊤ ω ) ω ∥ be the orthogonal component of θ to ω . Then, we have that

With P ω , θ = I d -ωω ⊤ -ee ⊤ , (B.7) is expanded as

For a general differentiable function g ( x ) , we have E [ He t ( x 1 ) g ( x )] = E [ d t d x t 1 g ( x )] . If g ( x ) is a polynomial (with a bounded coefficients) of x 1 and u ⊤ x and its degree with respect to x 1 is at most s ( ≤ t ) , | E [ He t ( x 1 ) g ( x )] | ≲ | u 1 | t -s , because differentiation of g ( x ) = ¯ g ( x 1 , u ⊤ x ) is taken with respect to the first variable at most s times. Each term of (B.9) is an expectation of ( σ ∗ ( θ ⊤ x )) l , multiplied by the polynomial of θ ⊤ x and ω ⊤ x , where its degree with respect to θ ⊤ x is at most 2 j -k . Thus each term of (B.9) is evaluated as (here we omit the constants)

The lower bound follows in the same fashion. Therefore,

Now, E [ θ ⊤ g ] can be evaluated as

(II) If I = { min i | IE( σ i ∗ ) = GE( σ ∗ ) = p ∗ } ≥ 2 . Note that α j = 0 for all j ≤ p ∗ from the assumption. Following (B.6), the first term of E [ θ ⊤ g ] is evaluated as

For the second term of E [ θ ⊤ g ] , similarly to (B.9), the following decomposition can be made.

Each term of (B.11) (omitting constants) is evaluated as

When i ≤ I -2 and η = c η d -1 , we have l ≤ i +1 < I and IE(( σ ∗ ( θ ⊤ x )) l ) ≥ p ∗ +1 . Thus (B.13) ≲ κ p ∗

When i ≥ I , IE(( σ ∗ ( θ ⊤ x )) l ) ≥ p ∗ and we get

̸

Now the case of i = I -1 . When i = I -1 and j = 0 , and using the assumption that κ ≤ c η ,

̸

When i = I -1 , j = 0 , and k = 0 ,

When i = I -1 , j = 0 , k = 0 , and l ≤ I -1 ,

Therefore, except for i = I -1 , j = 0 , k = 0 , and l ≤ I -1 , we can bound (B.13) by ≲ c I η κ p ∗ -1 + κ p ∗ . The lower bound follows in the same way. Finally, consider the case of i = I -1 , j = 0 , k = 0 , and l = I .

Putting it all together (recovering the constants omitted in (B.12) again),

(B.11)

and

Combining (i) and (ii), we have

for a positive constant c σ = Θ(1) . Here c σ > 0 satisfies 2 c σ = 2 p ∗ ! α p ∗ β p ∗ (for (i)) or 2 c σ = p ∗ ! H ( σ I ∗ ; p ∗ ) H ( σ ( I ) ( σ (1) ) I -1 ; p ∗ -1) ( I -1)! (for (ii)). Going back to (B.5), by setting ν 2 t = ( θ ⊤ g 2 t -E [ θ ⊤ g 2 t ]) , we have

When c ξ ≤ c I η and c -1 η d -1 2 ≤ κ ≤ c I η , terms in the bigO notation is smaller than c I η c ξ d -p ∗ 2 ∨ 1 c σ ( κ 2 t ) p ∗ -1 and we have

It is straightforward to check ν 2 t 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

with the mean-zero sub-Weibull random variable ν 2 t and a positive c σ = Θ(1) . For notational simplicity we write c I η c σ = c 1 . The following lemma is a detailed version of Proposition 9.

Lemma 17. Take η 2 t , η 2 t +1 = η = c η d -1 , ξ 2( t +1) = ξ = 1 -c ξ d -( p ∗-2) + 2 . Suppose that the link function satisfies IE( σ I ∗ ) = GE( σ ∗ ) = p ∗ (we choose the smallest such I ) and activation functions satisfy all of the assumptions in Section B.1 for weak recovery. Let

and take c ξ ≲ δ poly( c η ) , c 2 ≳ poly( c η ) , and C 3 ≃ c -1 1 . If κ 0 ≥ 2 c -1 η d -1 2 , there exists some τ ∗ ≤ T 1 , 1 such that

with probability at least 1 -δ , and κ 2 τ ≥ 2 c 2 for all τ ∗ ≤ τ ≤ T 1 , 1 , with high probability.

We may take δ = o d (1) with arbitrarily slow decay. The proof is adapted from [BAGJ21], but our bound on T 1 , 1 is slightly sharper (by a log d factor for p ∗ = 2 and by a (log d ) 2 factor for p ∗ ≥ 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 { ∑ τ s =0 ν 2 s } τ is Martingale with E [( ν 2 s ) 2 ] ≲ 1 . By Doob's maximal inequality and Markov's inequality, with probability 1 -δ , we have

for any fixed T ≥ 0 , with a sufficiently large constant C 1 = Θ(1) . In the following we consider the case when (B.14) holds for T = c -1 1 c -1 ξ d -1 .

Now, applying (B.14) to get

when τ ≤ c -1 1 c -1 ξ d -1 . By letting c ξ ≤ c -4 η c 1 C -1 1 δ , we have c -1 η d -1 2 ≤ c η c 1 2 ξ c -1 2 1 C 1 2 1 δ -1 2 d -1 2 , and

which verifies c -1 η d -1 2 ≤ κ 2 t for t = τ +1 . Thus, there exists some τ ∗ ≤ c -1 1 c -1 ξ d such that

for c 1 ≤ 1 4 c I η , with probability 1 -δ .

Now we prove that κ 2 t ≥ 2 c 2 holds for all τ ∗ ≤ t ≤ T 1 , 1 = C 3 c -1 ξ d . Because ν 2 t are mean-zero sub-Weibull random variables, we also have that | ∑ τ + τ ′ -1 s = τ ν 2 s | ≤ C 4 √ τ ′ for all 0 ≤ τ, τ ′ ≤ T 1 , 1 with high probability. Also, because η t ≪ d -1 and | 1 -ξ t | ≪ 1 , we can easily see that | κ 2( τ +1) -κ 2 τ | = ˜ O ( d -1 ) for all τ = 0 , 1 , · · · , T 1 , 1 -1 , with high probability. Thus, when there exists τ ≥ τ ∗ such that κ 2( τ -1) ≥ 4 c 2 and κ 2 τ < 4 c 2 , we have κ 2 τ ≥ 3 c 2 with high probability. Moreover, following the above argument, we can inductively show that

for τ ′ ≤ T 1 , 1 = C 3 c -1 ξ d or until κ 2( τ + τ ′ ) ≥ 4 c 2 holds again. By repeating this argument (if there are multiple such τ ), we see that κ 2 t ≥ 2 c 2 holds for all τ ∗ ≤ t ≤ T 1 , 1 = C 3 c -1 ξ d with high probability.

(ii) When p ∗ = 2 . We define ι 0 = 0 , ι 1 = log (1+ c 1 c ξ d -1 ) (4) , ι 2 = 2log (1+ c 1 c ξ d -1 ) (4) , . . . . We show that, for each i , if κ 2 ι i ≥ 2 c -1 η d -1 2 , we have κ 2( ι i +1 ) ≥ 2 κ 2 ι i , with probability at least 1 -δ 4 -i , or there exists some t ( ι i < t ≤ ι i +1 ) with κ 2 t > c I η .

Assume that the above statement holds until some i -1 ≥ 0 (we do not need to assume anything for i = 0 ). Then, we have κ 2 ι i ≥ 2 i κ 0 ≥ 2 c -1 η d -1 2 . Similarly to (B.15), if c -1 η d -1 2 ≤ κ 2 t ≤ c I η for all t = ι i , ι i +1 , · · · , τ , we have

Applying (B.14) with δ = δ/ 4 i and T = 1 4 c -2 η c -2 ξ C -1 ( δ/ 4 i )( κ 2 ι i ) 2 d 2 -1 to get

when τ ≤ ι i + 1 4 c -2 η c -2 ξ C -1 ( δ/ 4 i )( κ 2 ι i ) 2 d 2 -1 , which verifies c -1 η d -1 2 ≤ 1 2 κ 2 ι i ≤ κ 2 t for t = τ +1 .

This implies that, with probability 1 -δ/ 4 i , we have

for all τ = 1 4 c -2 η c -2 ξ C -1 ( δ/ 4 i )( κ 2 ι i ) 2 d 2 -1 , which is equivalent to

for all τ = ι i , ι i + 1 , · · · , ι i + 1 4 c -2 η c -2 ξ C -1 ( δ/ 4 i )( κ 2 ι i ) 2 d 2 . By taking c ξ ≪ c 1 c -2 η C -1 ( δ/ 4 i )( κ 2 ι i ) 2 d , we have 1 4 c -2 η c -2 ξ C -1 ( δ/ 4 i )( κ 2 ι i ) 2 d 2 ≥ log (1+ c 1 c ξ d -1 ) (4) , and we get κ 2 ι i +1 ≥ 2 κ 2 ι i with probability 1 -δ/ 4 i (or there exists t ≤ ι i +1 such that κ 2 t > c I η ). Thus, by induction, for all i , we have that

or that there exists some t ≤ ι i such that κ 2 t is larger than c I η , with probability 1 -δ .

The LHS of (B.16) becomes larger than c I η for some i ≤ log d . Because ι i = Θ( ic -1 1 c -1 ξ d ) , within O ( c -1 1 c ξ d log d ) steps, there exists at least one τ ∗ = O ( c -1 1 c -1 ξ d log d ) such that κ 2 τ ∗ ≥ 4 c 2 for c 2 ≤ 1 4 c I η , with probability 1 -δ .

Once such τ ∗ is obtained, following the last paragraph of (i), we can see that κ 2 t ≥ 2 c 2 holds until t = T 1 , 1 with high probability.

(iii) When p ∗ ≥ 3 . We apply (B.14) with T = 1 p ∗ -2 c -1 1 c -1 ξ d p ∗ 2 ( κ 0 ) -( p ∗ -2) to obtain that

for all τ = 0 , 1 , · · · , T -1 , with probability 1 -δ .

We take c ξ ≪ c -2 η c 1 C -1 δd p ∗ 4 ( κ 0 ) p ∗ 2 so that (B.17) is bounded by c -1 1 d -1 2 . Then,

It is easy to see that κ 2( τ +1) is lower bounded by a τ +1 , where a 0 = c -1 η d -1 2 and a τ +1 = a τ + c 1 c ξ d -p ∗ 2 ( a τ ) p ∗ -1 . By applying Lemma 18, we have

Thus, until τ ≤ ( c 1 c ξ d -p ∗ 2 ( p ∗ -2)( κ 0 ) ( p ∗ -2) ) -1 ≤ T +1 ≪ d p ∗ -1 , with probability at least 1 -δ , there exists some τ ∗ such that

when c 2 ≤ 1 4 c I η .

Once such τ ∗ is obtained, following the last paragraph of (i), we can see that κ 2 t ≥ 2 c 2 holds until t = T 1 , 1 with high probability.

In the above proof we used the (discrete version of) Bihari-LaSalle inequality from [BAGJ22].

Lemma 18. For p ≥ 3 and c > 0 , consider a positive sequence ( a t ) t ≥ 0 such that

Proof. From definition, we have

Taking the summation and re-arranging the terms yield

which gives the lower bound.

B.5 From Weak Recovery to Strong Recovery

In the previous subsection, we proved that after t = 2 T 1 , 1 = ˜ Θ( d ) steps, with probability ˜ Ω(1) over the randomness of initialization, we obtain nontrivial alignment κ 2 T 1 , 1 j ≥ 2 c 2 . This subsection discusses how to convert the weak recovery into the strong recovery.

Lemma 19. Suppose the neuron satisfies κ 2 T 1 , 1 ≥ 2 c 2 . Take η 2 t = η = ¯ c η εd -1 , η 2 t +1 = 0 , ξ 2( t +1) = 0 for all t ≥ T 1 , 1 , where ¯ c η ≲ poly( c 1 ) . If the activation functions satisfy all of the assumptions in Section B.1 for strong recovery, then we have

with high probability, where τ ∗ ≤ T 1 , 2 = C 3 dε -2 . Moreover, θ ⊤ w 2( T 1 , 1 + t ) ≥ 1 -ε for all τ ∗ ≤ t ≤ T 1 , 2 = C 3 dε -2 , with high probability.

Proof. Consider the Hermite expansions of σ ∗ and σ . Let p be the smallest degree that both σ ∗ and σ have non-zero coefficients. First we compute the population gradient (of the correlation term) as

Applying P w 2 t , we have

Then, we have

Thus, the update of the alignment κ 2 t = θ ⊤ w 2 t is

where

From (B.18), the expectation of (B.19) is bounded by

By letting η ≤ c p -1 1 εd -1 , when κ 2 t ≤ 1 -ε , we have

It is easy to see that the noise ν 2 t has sub-Weibull tail and we obtain that

Suppose that 2 c 2 ≤ κ 2( T 1 , 1+ τ ) ≤ 1 -ε for all t = 0 , 1 , . . . , τ -1 . By taking the summation of (B.20), we have

with high probability. The third term is bounded by C 4 η √ τ ≤ c 2 when τ ≤ c 2 2 C -2 4 η -2 = c 2 2 C -2 4 ¯ c -2 η ε -2 d 2 and by 1 2 ηετc p 1 when τ ≥ 4 ε -2 c -2 p 1 C 2 4 . Because c 2 2 C -2 4 ¯ c -2 η ε -2 d 2 ≥ 4 ε -2 c -2 p 1 C 2 4 , we can bound (B.21) by

which verifies 2 c 2 ≤ κ 2( T 1 , 1+ τ ) .

Therefore, by induction, until κ 2 t ≥ 1 -ε , we have the lower bound (B.22), whose RHS exceeds 1 -ε when τ ≥ 2 η -1 ε -1 c -p 1 ≤ C 3 dε -2 . Thus, there exists τ ∗ ≤ T 1 , 2 = C 3 dε -2 such that κ 2( T 1 , 1+ τ ∗ ) ≥ 1 -ε , with high probability.

Now, what remains is to prove that κ 2( T 1 , 1+ τ ) ≥ 1 -3 ε holds for all τ ∗ ≤ t ≤ T 1 , 2 = C 3 dε -2 . Because ν 2 t are mean-zero sub-Weibull random variables, we have that | ∑ τ + τ ′ -1 s = τ ν 2 s | ≤ C 4 √ τ ′ for all 0 ≤ τ, τ ′ ≤ T 1 , 1 with high probability. Also, because η t ≪ εd -1 , we can easily see that | κ 2( τ +1) -κ 2 τ | = ˜ O ( εd -1 ) for all τ = 0 , 1 , · · · , T 1 , 1 -1 , with high probability. Thus, when there exists τ ≥ τ ∗ such that κ 2( T 1 , 1 + τ -1) ≥ 1 -ε and κ 2( T 1 , 1 + τ ) < 1 -ε , we have κ 2( T 1 , 1 + τ ) ≥ 1 -2 ε with high probability. Moreover, following the above argument, we can inductively show that

for τ ′ ≤ T 1 , 2 or until κ 2( T 1 , 1+ τ + τ ′ ) ≥ 1 -ε holds again. Note that the last inequality follows from ¯ c -2 η C -2 4 d 2 ≥ ε -2 C 2 4 c -2 p 1 . By repeating this argument (if there are multiple such τ ), we can see that κ 2( T 1 , 1+ t ) ≥ 1 -ε holds for all τ ∗ ≤ t ≤ T 1 , 2 = C 3 dε -2 with high probability.

Adjusting hidden constants to remove a factor of 3 from 3 ε yields the desired result.

B.6 Second Layer Training

From the previous analysis, we know that at least Ω(1) portion of the neurons will satisfy the weak and strong recovery conditions (Appendix B.1), at least ˜ Ω(1) portion of the neurons (independent from the choice of σ j ) satisfy initial alignment conditions (Appendix B.2), and at least 1 -o (1) fraction of them achieves strong recovery. This subsection proves a generalization error bound after second-layer training. Let f a ( x ) = f Θ ( x ) for Θ = ( ˆ w j , a j , ˆ b j ) N j =1 where a ∈ R N and ( ˆ w j , ˆ b j ) N j =1 are the parameters trained in the first stage. Let a ∗ ∈ R N be the 'certificate' with ∥ a ∗ ∥ 2 = ˜ O ( N ) that is shown to exist in Lemma 22.

Polynomial Link Functions. The following lemma is a complete version of Proposition 5.

Lemma 20. There exists a 4 q -th order polynomial Q ( R w , b, q ′ ) of R w = max j ∥ w j ∥ , b = ( b j ) N j =1 such that, if we set λ = Θ (√ 2 T 2 δ 0 N 2 Q ( R w , b, q ′ ) ) for some δ 0 > 0 , the ridge estimator ˆ a satisfies

with probability 1 -δ 0 . Hence taking T 2 = ˜ Θ(( N 4 Q 2 ( R w , b, q ′ ) + E [ f ∗ ( x ) 4 ] 2 ) ε -4 ) and N = ˜ Θ( ε -1 ) , we have

Proof. Let P T 2 be the empirical distribution of the second stage: P T 2 := 1 T 2 ∑ T 2 i =1 δ x i . Let ψ ( x ) = ( σ ( ⟨ x , ̂ w j ) ⟩ + b j )) N j =1 so that f a ( x ) = ⟨ a , ψ ( x ) ⟩ .

Part (1). We first bound the term ∥ f ˆ a -f ∗ ∥ L 2 ( P T 2 ) . Since ˆ L ( f ˆ a ) + λ ∥ ˆ a ∥ 2 ≤ ˆ L ( f a ∗ ) + λ ∥ a ∗ ∥ 2 , we have

Now, by the Cauchy-Schwarz inequality, we have

By applying Markov's inequality to the right hand side, it can be further bounded by

with probability 1 -δ 1 . Thus, by combining with (B.24), we arrive at

Here, by using the evaluation ∥ f a ∗ -f ∗ ∥ L 2 ( P T 2 ) = ˜ O ( N -1 + ε ) in Lemma 22, the right hand side can be further bounded by

Part (2). Next we lower bound ∥ f ˆ a -f ∗ ∥ 2 L 2 ( P T 2 ) by noticing that

The first two terms of Eq. (B.25) can be bounded by

The standard Rademacher complexity bound yields that

where ( σ i ) T 2 i =1 is the i.i.d. Rademacher sequence independent of ( x i ) T 2 i =1 . Hence, Markov's inequality yields

with probability 1 -δ 2 .

The third term in Eq. (B.25) can be evaluated as

with probability 1 -δ 3 where we used Markov's inequality again in the second inequality. Finally, the fourth and fifth term in Eq. (B.25) can be bounded as

with probability 1 -δ 4 where we used Markov's inequality in the last inequality. Combining these inequalities, we finally arrive at

with probability 1 -∑ 4 j =1 δ j . Hence, by setting λ ≥ 8 √ 2 T 2 δ 2 E x [ ∥ ψ ( x ) ∥ 4 ] , we have that

When the activation function σ is a polynomial, then each ψ j ( x ) = σ ( ⟨ x , w j ⟩ + b j ) is an order q -polynomial of a Gaussian random variable ⟨ x , w j ⟩ which has mean 0 and variance E [ ⟨ x , w j ⟩ 2 ] = ∥ w j ∥ 2 = ˜ O (1) . Then, if we let R w := max j ∥ w j ∥ = ˜ O (1) , the term max j max { E x [ ψ ( x ) 2 j ] , E x [ ψ ( x ) 4 j ] } can be bounded by a 4 q -th order polynomial of R w and b , which can be denoted by Q ( R w , b, 4 q ) .

Part (3). By combining evaluations of (1) and (2) together, if we let λ = 8 √ 2 T 2 δ 0 E x [ ∥ ψ ( x ) ∥ 4 ] for some δ 0 > 0 , (by ignoring polylogarithmic factors) we obtain that

with probability 1 -4 δ 0 . Thus, since ∥ a ∗ ∥ 2 = ˜ O ( N ) , by setting T 2 = ˜ Θ(( N 4 Q 2 ( R w , b, q ′ ) + E [ f ∗ ( x ) 4 ] 2 ) ε -4 ) , and N = ˜ Θ( ε -1 ) , we obtain that (B.23) ≲ ε 2 .

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 E [ σ ∗ ( θ ⊤ x ) 4 ] < ∞ and Assumption 4 hold. Then, by setting λ = ˜ Θ (√ N 2 T 2 δ 0 ) for some δ 0 > 0 , the ridge estimator ˆ a satisfies

with probability 1 -δ 0 . By taking T 2 = ˜ Θ(( N 4 + N 2 ) ε -4 ) , we have

Furthermore, applying Lemma 12 and 13 yields that, when σ ∗ = ∑ ∞ j =0 α j He j satisfies ∑ ∞ j =0 j 2 j ! α 2 j and E [ σ ∗ ( θ ⊤ x ) 4 ] are bounded, with a properly designed randomized activation in Lemma 12, by taking N = ˜ Θ( ε -7 ) and T 2 = ˜ Θ( ε -32 ) , Algorithm 1 yields

with probability 1 -o d (1) .

Proof. The proof is identical to that of Lemma 20, with the difference being that we replace the bounded moment assumptions with E [ σ ∗ ( θ ⊤ x ) 4 ] < ∞ or Assumption 4.

Approximation Guarantee. Note that for non-polynomial link function with generative exponent p ∗ ≥ 3 , the approximation error is already controlled in Assumption 4 based on [BBSS22, Lemma 4.4, 4.5] (using activation function with a ReLU component). If σ ∗ is a degreeq polynomial, we have the following approximation result using polynomial activation, which follows Lemmas 29 and 30 of [OSSW24a].

Lemma 22. Suppose that there exist at least N ′ = ˜ Θ( N ) neurons that satisfy ∥ w 2 T 1 j -θ ∥ ≤ ε and σ is a polynomial link function with degree at least q . Let b j ∼ Unif([ -C b , C b ]) with C b = ˜ O (1) , and consider approximation of a ridge function h ( θ ⊤ x ) with its degree at most q . Then, there exists a 1 , . . . , a N such that

with high probability, where ( x , y ) is a random sample, and we omit dependence on the degree q in the bigO notation. Moreover, we have ∑ N j =1 a 2 j = ˜ O ( N ) .

Lemma 22 can be established from the following result in [OSSW24a].

Lemma 23. Suppose that C b ≥ q . For any polynomial h ( s ) with its degree at most q , there exists ¯ v ( b ; h ) with | ¯ v ( b ; h ) | ≲ C b such that for all s ,

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, 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 a 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 it (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 evaluation. While '[Yes] ' is generally preferable to '[No] ', it is perfectly acceptable to answer '[No] ' provided a proper justification is given (e.g., 'error bars 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 a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a 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

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

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 in the paper.

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

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

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

2. Limitations

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

Answer: [Yes]

Justification: 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 are encouraged to create a separate 'Limitations' section in their paper.

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

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

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

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

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

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

3. Theory Assumptions and Proofs

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

Answer: [Yes]

Justification: All assumptions 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 crossreferenced.

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

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

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

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

4. Experimental Result Reproducibility

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

Answer: [NA]

Justification: This paper is theoretical and experiments are only used for the purpose of illustration in the Introduction.

Guidelines:

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

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

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

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

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

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

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

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

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

5. Open access to data and code

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

Answer: [NA]

Justification: 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/ public/guides/CodeSubmissionPolicy ) for more details.

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

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

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

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

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

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

6. Experimental Setting/Details

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

Answer: [NA]

Justification: This is a 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 a level of detail that is necessary to appreciate the results and make sense of them.

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

7. Experiment Statistical Significance

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

Answer: [NA]

Justification: This is a 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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

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

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

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

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

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

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

8. Experiments Compute Resources

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

Answer: [NA]

Justification: This is a theory paper.

Guidelines:

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

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

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

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

9. Code Of Ethics

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

Answer: [Yes]

Justification: The authors have read and followed the NeurIPS Code of Ethics.

Guidelines:

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

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

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

10. Broader Impacts

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

Answer: [NA]

Justification: This paper carries out a theoretical study, and we believe our work does not have specific societal impacts that require a discussion.

Guidelines:

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

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

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

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

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

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

11. Safeguards

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

Answer: [NA]

Justification: This paper is theoretical and the experiment is only about two-layer neural network and synthetic data, which we do not have such risks.

Guidelines:

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

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

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

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

12. Licenses for existing assets

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

Answer: [NA]

Justification: This paper is theoretical and does not use any assets.

Guidelines:

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

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

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

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

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

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

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

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

13. New Assets

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

Answer: [NA]

Justification: 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 of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

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

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

14. Crowdsourcing and Research with Human Subjects

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

Answer: [NA]

Justification: 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 is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

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

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

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

Answer: [NA]

Justification: 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.

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

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

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

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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].\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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)$.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors addressed the limitations of the paper and provided inspiring future directions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nS1) 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.\n\nWEAKNESSES:\nW1) 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.\n\nQUESTIONS:\nQ1) 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'.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe strength of this submission resides in the strong theoretical claims. The questions addressed are of great interest to the theoretical machine learning community.\n\nWEAKNESSES:\nThis 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe limitations are addressed in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}