# Online stochastic gradient descent on non-convex losses from high-dimensional inference

arXiv: https://arxiv.org/abs/2003.10409 · HTML: https://mlsanity.com/paper/2003.10409/

## Key facts

- Cited by 119 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 23 Mar 2020.

## Built on by (most cited first, 30 of 119)

- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [Grokking as the Transition from Lazy to Rich Training Dynamics](https://mlsanity.com/paper/2310.06110/) (2023)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [High-dimensional limit theorems for SGD: Effective dynamics and critical scaling](https://mlsanity.com/paper/2206.04030/) (2022)
- [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://mlsanity.com/paper/2305.18270/) (2023)
- [Neural Networks Efficiently Learn Low-Dimensional Representations with SGD](https://mlsanity.com/paper/2209.14863/) (2022)
- [Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models](https://mlsanity.com/paper/2305.10633/) (2023)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [The Benefits of Reusing Batches for Gradient Descent in Two-Layer Networks: Breaking the Curse of Information and Leap Exponents](https://mlsanity.com/paper/2402.03220/) (2024)
- [Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks](https://mlsanity.com/paper/2202.00293/) (2022)
- [Grokking as a First Order Phase Transition in Two Layer Networks](https://mlsanity.com/paper/2310.03789/) (2023)
- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks](https://mlsanity.com/paper/2302.05882/) (2023)
- [SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem](https://mlsanity.com/paper/2309.15111/) (2023)
- [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://mlsanity.com/paper/2405.15459/) (2024)
- [How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad](https://mlsanity.com/paper/2406.06467/) (2024)
- [Gradient-Based Feature Learning under Structured Data](https://mlsanity.com/paper/2309.03843/) (2023)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [A distributional simplicity bias in the learning dynamics of transformers](https://mlsanity.com/paper/2410.19637/) (2024)
- [Bayes-optimal learning of an extensive-width neural network from quadratically many samples](https://mlsanity.com/paper/2408.03733/) (2024)
- [Provable Advantage of Curriculum Learning on Parity Targets with Mixed Inputs](https://mlsanity.com/paper/2306.16921/) (2023)
- [Provably learning a multi-head attention layer](https://mlsanity.com/paper/2402.04084/) (2024)
- [On the Cryptographic Hardness of Learning Single Periodic Neurons](https://mlsanity.com/paper/2106.10744/) (2021)
- [Six Lectures on Linearized Neural Networks](https://mlsanity.com/paper/2308.13431/) (2023)
- [Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks](https://mlsanity.com/paper/2305.06986/) (2023)
- [Fundamental computational limits of weak learnability in high-dimensional multi-index models](https://mlsanity.com/paper/2405.15480/) (2024)
- [On Single Index Models beyond Gaussian Data](https://mlsanity.com/paper/2307.15804/) (2023)
- [A Random Matrix Theory Perspective on the Spectrum of Learned Features and Asymptotic Generalization Capabilities](https://mlsanity.com/paper/2410.18938/) (2024)
- [Escaping mediocrity: how two-layer networks learn hard generalized linear models with SGD](https://mlsanity.com/paper/2305.18502/) (2023)

## Abstract (arXiv)

Stochastic gradient descent (SGD) is a popular algorithm for optimization problems arising in high-dimensional inference tasks. Here one produces an estimator of an unknown parameter from independent samples of data by iteratively optimizing a loss function. This loss function is random and often non-convex. We study the performance of the simplest version of SGD, namely online SGD, from a random start in the setting where the parameter space is high-dimensional. We develop nearly sharp thresholds for the number of samples needed for consistent estimation as one varies the dimension. Our thresholds depend only on an intrinsic property of the population loss which we call the information exponent. In particular, our results do not assume uniform control on the loss itself, such as convexity or uniform derivative bounds. The thresholds we obtain are polynomial in the dimension and the precise exponent depends explicitly on the information exponent. As a consequence of our results, we find that except for the simplest tasks, almost all of the data is used simply in the initial search phase to obtain non-trivial correlation with the ground truth. Upon attaining non-trivial correlation, the descent is rapid and exhibits law of large numbers type behavior. We illustrate our approach by applying it to a wide set of inference tasks such as phase retrieval, and parameter estimation for generalized linear models, online PCA, and spiked tensor models, as well as to supervised learning for single-layer networks with general activation functions.
