Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit

We study the problem of gradient descent learning of a single-index target function $f_*(\boldsymbol{x}) = \textstyleσ_*\left(\langle\boldsymbol{x},\boldsymbolθ\rangle\right)$ under isotropic Gaussian data in $\mathbb{R}^d$, where the unknown link function $σ_*:\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^{Θ(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 = Θ(d\!\cdot\! \mathrm{polylog} d)$, where $Θ(\cdot)$ 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\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.

Paper

References (54)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer 6kdk7/10 · confidence 4/52024-07-07

Summary

The paper studies the classical problem of the single index model over Gaussian inputs, i.e. $x\sim \mathcal{N}(0,I_d)$ and $f_*(x)=\sigma_*(\langle \theta_*,x \rangle)$ for an unknown direction $\theta_*$. Information theoretically, one needs $\Omega(d)$ samples to learn this function class. The paper shows that reusing the batches with a certain SGD-based training on a 2-layer neural network achieves vanishing $L_2$ error with $O(d \text{polylog}(d))$ samples--that is this SGD algorithm by some reused batches learns this function class at nearly the information-theoretic limit. The 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.

Strengths

1. The paper addresses an important question in our understanding of the complexity of gradient descent-type algorithms on regular neural networks. There have been considerable efforts devoted to understanding this. This paper goes beyond most (if not all) of these works by analyzing an SGD with reused batches, over vanilla batch SGD. While a recent paper of [DTA+24] provided the first evidence of the benefit of reusing the batches and that CSQ bound can be escaped, the paper goes beyond this and in an important way as follows. 2. This paper considers strong recovery which is more satisfying, and technically much more challenging (in contrast to [DTA+24] that only considered weak recovery). To achieve this, there were important pieces to be figured out, which the authors successfully did. This paper provides a clear end-to-end analysis and establishes the learning guarantee in contrast to [DTA+24].

Weaknesses

I do not see any major weaknesses. The training procedure is slightly non-standard, but it is completely understandable from a technical point of view. The layer-wise training and the use of the projected gradients are completely standard in theoretical research. However, I could not see a clear motivation/need for the momentum for the first-layer training. I was wondering if it can be avoided.

Questions

1. Can the use of momentum be avoided in the first layer of training? If "no" then, in what way this is important in the current analysis? 2. Why cite [Gla'23] (on line 338 page 9) for the point of adversarial noise in SQ vs non-adversarial noise in GD? 3. 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`` for 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. 4. Should the abbreviation on line 31 page 1 be CSQ instead of SQ?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

None.

Reviewer Mv1X6/10 · confidence 3/52024-07-10

Summary

This paper studied the problem of learning single index models under Isotropic Gaussian distribution. The target model $f^*(x) = \sigma(\theta^\top x)$ is a polynomial function $\sigma$ composed with a one-dimensional structure $\theta^\top x$, where the polynomial $\sigma$ is of degree at most $q$ and has information component (i.e., the order of the first non-zero Hermite polynomial expansion) $p$. The paper studied the sample and computational complexity of learning such single index models with 2-layer neural networks, using gradient descent type methods. The critical aspect of the algorithm is reusing the same batch of samples every two iterations, which frees the algorithm from CSQ constraints and becomes an SQ algorithm. A critical observation is that by reusing samples at each iteration, the algorithm induced a monomial transformation of the labels, which effectively reduced the information component from $p$ to less than 2. Hence, the algorithm can achieve near-optimal sample complexity $\tilde{O}(d)$.

Strengths

1. This paper is clearly written with intuitions and useful explanations. 2. 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. 3. This paper provides a very interesting intuition on reducing the information exponent of the link function $\sigma$ using monomial transformations, which could be of independent interest for future works.

Weaknesses

1. Though the authors claimed that they were using neural networks to learn the single index model, the activation of each neuron turns out to be a combination of polynomials. Hence, the neural network $\sum a_j\sigma(w_j^\top x + b_j)$ is essentially a linear combination of Hermite polynomials. In this case, I am wondering what the differences are between using the 'neural network' to learn the single index models and using polynomials to learn the single index models, which is already done in [CM20]. Of course, [CM20] requires a warm start procedure, which is not a gradient descent type algorithm, but I think it would be more interesting if that analysis is carried out on conentional neural networks like ReLU networks. 2. 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.

Questions

1. Do ReLUs and sigmoids satisfy the assumption 2? 2. I think there are typos on line 593 to 595. What is 'x' on the right-hand side of line 593 and 595? 3. Since only $exp(-q)$ neurons satisfy assumptions 2 and 3, does it imply that the width of the network is at least $exp(q)$? 4. 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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors addressed the limitations of the paper and provided inspiring future directions.

Area Chair rw5H2024-08-12

Please engage with the authors' response: have they addressed your concerns adequately, and how has your score been affected? Best, your AC

Reviewer Mv1X2024-08-13

I thank the authors for their detailed response. I would like to keep my score unchanged.

Reviewer hk3C7/10 · confidence 3/52024-07-11

Summary

This paper addresses the problem of learning single-index targets with polynomial link functions under Gaussian inputs. The authors demonstrate that using SGD on a two-layer fully connected network with a specific activation function can learn such targets with O(d poly(log(d))) samples. The analysis involves the reuse of samples, which allows to improve previous bounds obtained for online single-pass SGD.

Strengths

S1) The paper advances the analysis of the complexity of learning Gaussian single-index models, a currently very popular model for the theoretical study of neural networks. This contribution is thus significant for the deep learning theory community. S2) The technical contributions are novel and well presented. S3) The paper builds on previous work that demonstrated the benefits of re-using batches for learning single-index models, providing concrete evidence of strong learnability of the target by SGD on a shallow network.

Weaknesses

W1) The authors provide minimal empirical validation, with only one experiment demonstrating their claims. They do not address more standard SGD practices, such as training both layers simultaneously, using standard activations/initializations, or employing larger learning rates. W2) 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.

Questions

Q1) Can you clarify how many hidden neurons N are needed for the main result to hold? Q2) How does the bound depend on the degree of the target q? Q3) Would you expect the same result to hold if the bias weights are trained? Q4) 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? Q5) Do you expect a similar analysis to hold for other losses? For example, L1 loss. Q6) Do you have an intuition for what could be an optimal mini-batch re-use schedule? Typos/suggestions: Proposition 4ii): Can you formally define 'the odd part of f'? Line 211: missing the word 'high'. Line 326: 'high' -> 'weak'. Line 175: missing the word 'be'. Line 163: typo 'not be not'.

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors adequately addressed the limitations of their work.

Reviewer p2Vt6/10 · confidence 2/52024-07-12

Summary

This manuscript studies the learning properties of two-layer networks trained with SGD reusing the batch. The authors show that this simple modification allows SGD to surpass the limits of CSQ algorithms and learn single-index functions efficiently. The submission considers both recovery of the target features and generalization properties of two-layer networks. The claims are supported by rigorously proven theorems.

Strengths

The strength of this submission resides in the strong theoretical claims. The questions addressed are of great interest to the theoretical machine learning community.

Weaknesses

This submission has no outstanding weaknesses but the presentation could be enhanced. I will detail in the section below some suggestions to improve the manuscript which sometimes is a bit obscure for a non-expert reader.

Questions

- The idea of label transformation implemented by [CM20] could be reported. The contrast with what SGD (with batch reusing) is implementing would give an idea to the reader of the strength of the claims. - 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. - 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? - 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. - 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? - 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.

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

The limitations are addressed in the manuscript.

Area Chair rw5H2024-08-09

Dear reviewers, Thank you for taking the time to review the paper. Notice that the discussion period has begun, and will last until August 13 (4 more days). During this time your active participation and engagement with the authors is very important, and highly appreciated. Specifically, please _read the responses_, _respond to them early on in the discussion_, and **discuss points of disagreement**. Thank you for your continued contributions to NeurIPS 2024. Best, Your Area Chair.

Reviewer p2Vt2024-08-10

Thank you for the rebuttal

I thank the authors for their rebuttal that clarified my concerns. After carefully reading it along with other reviewers’ comments I would like to keep my score as in the original review.

Authorsrebuttal2024-08-11

Thank you for the update. We would appreciate knowing if there are any outstanding concerns that may have led to the reviewer's decision to maintain the current score. We would be more than happy to provide further clarifications during the discussion period.

Reviewer hk3C2024-08-12

Thanks

I thank the authors for their response, which addressed all my concerns. I will keep my score.

Authorsrebuttal2024-08-14

Dear Reviewers and Area Chair, We appreciate your continued time and effort in providing detailed feedback on our paper. As the rebuttal period comes to a close, we would like to briefly summarize our responses and revisions, which we believe have addressed all of the reviewers' concerns. - Reviewers hk3C and 6kdk inquired about the role of the momentum step. We explained that its purpose is to improve the signal-to-noise ratio, which is crucial in the $\mathrm{GE}=2$ setting, where the signal and noise are nearly at the same scale. - Reviewers Mv1X and hk3C raised questions on the constants in our big-$O$ notation. We clarified the dependence on the target degree $q$, and addressed Reviewer Mv1X's misunderstanding regarding the dimension dependence in the constants. - We discussed the motivation behind the layer-wise training procedure and the possibility simultaneous training and the use of different loss functions, in response to Reviewer hk3C's concern about standard SGD training. - Reviewers Mv1X and p2Vt asked about the differences between our analysis and that of [CM20]. We highlighted that our work focuses on SGD training of a two-layer neural network, without label preprocessing (with thresholding) or optimization of the polynomial coefficients. --- To recap the main findings of our submission, we showed that by simple data-reuse, two-layer neural network can learn *arbitrary* polynomial single-index models near the information-theoretically optimal sample complexity (beyond that suggested by the CSQ lower bound). We built upon the insight is that SGD with reused data can implement *non-correlational* queries (despite the squared loss), and established end-to-end learning guarantees that goes beyond weak recovery. We believe that our results will be of interest to the NeurIPS community. Best regards, Authors of submission 12755

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC