# Neural Networks Efficiently Learn Low-Dimensional Representations with SGD

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

## Key facts

- Cited by 63 papers in the map (data through 14 Oct 2026).
- Builds on 66 of its 94 references that are in the map.
- First posted on arXiv 29 Sep 2022.

## Builds on (most cited first, 30 of 66)

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [The generalization error of random features regression: Precise asymptotics and double descent curve](https://mlsanity.com/paper/1908.05355/) (2019)
- [In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning](https://mlsanity.com/paper/1412.6614/) (2014)
- [Stronger generalization bounds for deep nets via a compression approach](https://mlsanity.com/paper/1802.05296/) (2018)
- [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020)
- [Implicit Bias of Gradient Descent on Linear Convolutional Networks](https://mlsanity.com/paper/1806.00468/) (2018)
- [SGD: General Analysis and Improved Rates](https://mlsanity.com/paper/1901.09401/) (2019)
- [Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit](https://mlsanity.com/paper/1902.06015/) (2019)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [Towards moderate overparameterization: global convergence guarantees for training shallow neural networks](https://mlsanity.com/paper/1902.04674/) (2019)
- [Stability of Stochastic Gradient Descent on Nonsmooth Convex Losses](https://mlsanity.com/paper/2006.06914/) (2020)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)
- [Convergence Analysis of Two-layer Neural Networks with ReLU Activation](https://mlsanity.com/paper/1705.09886/) (2017)
- [Implicit Regularization of Discrete Gradient Dynamics in Linear Neural Networks](https://mlsanity.com/paper/1904.13262/) (2019)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup](https://mlsanity.com/paper/1906.08632/) (2019)
- [Implicit Bias of SGD for Diagonal Linear Networks: a Provable Benefit of Stochasticity](https://mlsanity.com/paper/2106.09524/) (2021)
- [On the Power and Limitations of Random Features for Understanding Neural Networks](https://mlsanity.com/paper/1904.00687/) (2019)
- [Recovery Guarantees for One-hidden-layer Neural Networks](https://mlsanity.com/paper/1706.03175/) (2017)
- [Sharper bounds for uniformly stable algorithms](https://mlsanity.com/paper/1910.07833/) (2019)

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

- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [Transformers learn through gradual rank increase](https://mlsanity.com/paper/2306.07042/) (2023)
- [Grokking as a First Order Phase Transition in Two Layer Networks](https://mlsanity.com/paper/2310.03789/) (2023)
- [The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit](https://mlsanity.com/paper/2306.17759/) (2023)
- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [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)
- [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)
- [Emergence in non-neural models: grokking modular arithmetic via average gradient outer product](https://mlsanity.com/paper/2407.20199/) (2024)
- [Mechanism of feature learning in deep fully connected networks and kernel machines that recursively learn features](https://mlsanity.com/paper/2212.13881/) (2022)
- [Beyond NTK with Vanilla Gradient Descent: A Mean-Field Analysis of Neural Networks with Polynomial Width, Samples, and Time](https://mlsanity.com/paper/2306.16361/) (2023)
- [Six Lectures on Linearized Neural Networks](https://mlsanity.com/paper/2308.13431/) (2023)
- [PoLAR: Polar-Decomposed Low-Rank Adapter Representation](https://mlsanity.com/paper/2506.03133/) (2025)
- [Variation Spaces for Multi-Output Neural Networks: Insights on Multi-Task Learning and Network Compression](https://mlsanity.com/paper/2305.16534/) (2023)
- [Mechanism of feature learning in convolutional neural networks](https://mlsanity.com/paper/2309.00570/) (2023)
- [Provable Multi-Task Representation Learning by Two-Layer ReLU Neural Networks](https://mlsanity.com/paper/2307.06887/) (2023)
- [Linear Recursive Feature Machines provably recover low-rank matrices](https://mlsanity.com/paper/2401.04553/) (2024)
- [Towards a theory of model distillation](https://mlsanity.com/paper/2403.09053/) (2024)
- [On the Complexity of Learning Sparse Functions with Statistical and Gradient Queries](https://mlsanity.com/paper/2407.05622/) (2024)
- [Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics](https://mlsanity.com/paper/2408.07254/) (2024)
- [Survey on Algorithms for multi-index models](https://mlsanity.com/paper/2504.05426/) (2025)
- [Provable Guarantees for Neural Networks via Gradient Feature Learning](https://mlsanity.com/paper/2310.12408/) (2023)
- [Autoencoders for discovering manifold dimension and coordinates in data from complex dynamical systems](https://mlsanity.com/paper/2305.01090/) (2023)
- [Mean-field Analysis on Two-layer Neural Networks from a Kernel Perspective](https://mlsanity.com/paper/2403.14917/) (2024)
- [Agnostic Active Learning of Single Index Models with Linear Sample Complexity](https://mlsanity.com/paper/2405.09312/) (2024)
- [Feature learning as alignment: a structural property of gradient descent in non-linear neural networks](https://mlsanity.com/paper/2402.05271/) (2024)
- [From Spikes to Heavy Tails: Unveiling the Spectral Evolution of Neural Networks](https://mlsanity.com/paper/2406.04657/) (2024)
- [On the symmetries in the dynamics of wide two-layer neural networks](https://mlsanity.com/paper/2211.08771/) (2022)
- [ReLU Neural Networks with Linear Layers are Biased Towards Single- and Multi-Index Models](https://mlsanity.com/paper/2305.15598/) (2023)

## Abstract (arXiv)

We study the problem of training a two-layer neural network (NN) of arbitrary width using stochastic gradient descent (SGD) where the input boldsymbolx∈ ℝd is Gaussian and the target y ∈ ℝ follows a multiple-index model, i.e., y=g(langleboldsymbolu1,boldsymbolxrangle,...,langleboldsymboluk,boldsymbolxrangle) with a noisy link function g. We prove that the first-layer weights of the NN converge to the k-dimensional principal subspace spanned by the vectors boldsymbolu1,...,boldsymboluk of the true model, when online SGD with weight decay is used for training. This phenomenon has several important consequences when k ll d. First, by employing uniform convergence on this smaller subspace, we establish a generalization error bound of O(√kd/T) after T iterations of SGD, which is independent of the width of the NN. We further demonstrate that, SGD-trained ReLU NNs can learn a single-index target of the form y=f(langleboldsymbolu,boldsymbolxrangle) + ε by recovering the principal direction, with a sample complexity linear in d (up to log factors), where f is a monotonic function with at most polynomial growth, and ε is the noise. This is in contrast to the known dOmega(p) sample requirement to learn any degree p polynomial in the kernel regime, and it shows that NNs trained with SGD can outperform the neural tangent kernel at initialization. Finally, we also provide compressibility guarantees for NNs using the approximate low-rank structure produced by SGD.