Lune

NeurIPS2024Top-tier venue

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

Jason D. Lee, Kazusato Oko, Taiji Suzuki, Denny Wu

2024Year
49Citations
34Top-tier citations

Abstract

We study the problem of gradient descent learning of a single-index target function f∗(x)=σ∗(⟨x,θ⟩)f_*(\boldsymbol{x}) = \textstyle\sigma_*\left(\langle\boldsymbol{x},\boldsymbol{\theta}\rangle\right) under isotropic Gaussian data in Rd\mathbb{R}^d, where the unknown link function σ∗:R→R\sigma_*:\mathbb{R}\to\mathbb{R} has information exponent pp (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)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∗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 ⁣⋅ ⁣polylogd)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≳d(p∗−1)∨1n\gtrsim d^{(p_*-1)\vee 1} samples are sufficient to achieve low generalization error, where p∗≤pp_* \le 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.

Ask about this paper

Your agent reads all of it.

Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.

Questions to start from

Your agent calls

Luneget_paper_fulltext

Ask in Lune

Free to start. No credit card required.

lune papers fulltext 753279e7-e1cd-4eac-8fb0-d350af0dbaa1

Cited by top-tier papers34

Ask how each one uses it

Builds on9

Related papers

Dusk over the sea between two cliffs drawn in fine vertical lines