Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit
Using batch-reused SGD on two-layer neural networks enables learning low-dimensional polynomials near the information-theoretic limit, surpassing classical p-dependent complexity.
Key Findings
Methodology
This work introduces a batch-reuse SGD framework combined with Hermite polynomial analysis to activate high-order statistical information. By employing a layered training strategy—initially achieving weak feature recovery, then refining to strong recovery—the approach leverages nonlinear label transformations to lower the target's Hermite information exponent p. The core mechanism involves reusing minibatches in gradient computations, which uncovers higher-order moments beyond correlational queries. Theoretical proofs establish that for polynomial targets, the sample complexity approaches n≈Θ(d·polylog d), nearly matching the information-theoretic lower bound. The analysis hinges on the activation functions' Hermite coefficients and the interplay between layer-wise training and high-order statistical information, enabling the neural network to learn efficiently even when p is large.
Key Results
- Theoretical analysis shows that two-layer neural networks trained with batch-reused SGD can learn polynomial single-index models with sample complexity n≈Θ(d·polylog d), significantly below the classical p-dependent bounds, and close to the information limit.
- For general target functions with generative exponent p_*, the sample complexity scales as n≳d^{(p_*−1)∨1}, surpassing the correlational statistical query (CSQ) lower bounds, demonstrating the high-order information utilization capacity of SGD.
- Empirical experiments validate that batch-reused SGD achieves low test error with sample sizes proportional to dimension d, outperforming online SGD which requires orders of magnitude more samples, confirming the theoretical predictions.
Significance
This research fundamentally advances the understanding of neural network training efficiency, showing that high-order statistical information can be exploited via simple SGD modifications. It demonstrates that neural networks can approach the information-theoretic limits for low-dimensional polynomial functions, even in high-dimensional settings. The findings bridge the gap between statistical optimality and computational feasibility, opening pathways for efficient learning in complex, high-dimensional data scenarios with limited samples. Such insights could influence the design of future algorithms and deepen theoretical comprehension of neural network capabilities.
Technical Contribution
The paper provides the first rigorous proof that batch-reuse SGD can implement a form of high-order statistical query (SQ) algorithm, breaking the traditional p-dependent complexity barrier. It introduces a novel label transformation mechanism via Hermite polynomial powers, enabling the reduction of the target's information exponent. The layered training approach, combined with a detailed Hermite coefficient analysis, establishes near-optimal sample complexity bounds. These contributions significantly extend the theoretical framework of neural network learning, linking high-order statistical information extraction with practical training algorithms.
Novelty
This work is the first to demonstrate that simple modifications to SGD—specifically, batch data reuse—can unlock high-order statistical information, enabling neural networks to learn polynomial targets near the information-theoretic limit. Unlike prior studies limited by correlational query bounds, this approach exploits the nonlinear structure of label transformations to surpass these barriers. The layered training strategy and Hermite-based analysis provide a new theoretical lens, establishing a fundamental link between high-order moments and neural network training efficiency, marking a significant leap in understanding neural learning dynamics.
Limitations
- The theoretical guarantees rely on assumptions about activation functions' Hermite coefficients, which may not directly translate to practical neural architectures or real-world data distributions.
- The analysis presumes high-dimensional Gaussian data, limiting immediate applicability to more complex, structured datasets.
- Computational costs associated with the layer-wise training and Hermite polynomial evaluations may be high in practice, requiring further optimization for large-scale deployment.
Future Work
Future research will explore extending these high-order information exploitation techniques to more general data distributions and nonlinear activation functions. Developing adaptive algorithms that automatically identify beneficial label transformations and optimize training dynamics remains an open challenge. Additionally, integrating these insights into scalable architectures and real-world applications such as image and speech recognition could significantly impact industry. Investigating robustness, convergence speed, and generalization in more realistic settings will be key directions.
AI Executive Summary
This paper tackles a fundamental challenge in deep learning: efficiently learning low-dimensional polynomial functions in high-dimensional spaces. Traditional bounds suggest that the sample complexity for such tasks scales exponentially with the information exponent p, making high-p targets computationally infeasible. The authors propose a novel approach that leverages batch data reuse within a two-layer neural network trained by stochastic gradient descent (SGD). By incorporating nonlinear label transformations based on Hermite polynomial expansions, the method effectively reduces the target's Hermite information exponent, unlocking higher-order statistical information. This enables the neural network to learn polynomial single-index models with a sample complexity of roughly n≈Θ(d·polylog d), approaching the information-theoretic limit. The core innovation lies in the realization that SGD, with batch reuse, can implement a form of high-order statistical query (SQ) algorithm, surpassing the classical correlation-based bounds. Theoretical proofs confirm that for polynomial targets, the method achieves near-optimal sample and runtime complexity, significantly outperforming traditional gradient-based algorithms that depend on p. Empirical experiments validate the theory, showing that the proposed approach requires far fewer samples than online SGD, especially in high p regimes. This breakthrough not only advances the theoretical understanding of neural network learning but also opens new avenues for practical, efficient training in high-dimensional, low-sample scenarios. Future work will focus on extending these techniques to broader data distributions, optimizing computational costs, and applying them to real-world tasks such as image and speech recognition, promising a substantial impact on both academia and industry.
Deep Dive
Abstract
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.