# A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks

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

## Key facts

- Cited by 49 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 86 references that are in the map.
- First posted on arXiv 11 Oct 2023.

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

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [The generalization error of random features regression: Precise asymptotics and double descent curve](https://mlsanity.com/paper/1908.05355/) (2019)
- [Adversarial Feature Learning](https://mlsanity.com/paper/1605.09782/) (2016)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit](https://mlsanity.com/paper/1902.06015/) (2019)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [High-Dimensional Asymptotics of Prediction: Ridge Regression and Classification](https://mlsanity.com/paper/1507.03003/) (2015)
- [Learning curves of generic features maps for realistic datasets with a teacher-student model](https://mlsanity.com/paper/2102.08127/) (2021)
- [Spectral Methods for Data Science: A Statistical Perspective](https://mlsanity.com/paper/2012.08496/) (2020)
- [Mean Field Analysis of Neural Networks: A Central Limit Theorem](https://mlsanity.com/paper/1808.09372/) (2018)
- [The Neural Tangent Kernel in High Dimensions: Triple Descent and a Multi-Scale Theory of Generalization](https://mlsanity.com/paper/2008.06786/) (2020)
- [Dynamics of Deep Neural Networks and Neural Tangent Hierarchy](https://mlsanity.com/paper/1909.08156/) (2019)
- [Understanding Double Descent Requires a Fine-Grained Bias-Variance Decomposition](https://mlsanity.com/paper/2011.03321/) (2020)
- [Universality of empirical risk minimization](https://mlsanity.com/paper/2202.08832/) (2022)
- [How Two-Layer Neural Networks Learn, One (Giant) Step at a Time](https://mlsanity.com/paper/2305.18270/) (2023)
- [A Theoretical Analysis on Feature Learning in Neural Networks: Emergence from Inputs and Advantage over Fixed Features](https://mlsanity.com/paper/2206.01717/) (2022)
- [Spectrum of inner-product kernel matrices in the polynomial regime and multiple descent phenomenon in kernel ridge regression](https://mlsanity.com/paper/2204.10425/) (2022)
- [Moving Beyond Sub-Gaussianity in High-Dimensional Statistics: Applications in Covariance Estimation and Linear Regression](https://mlsanity.com/paper/1804.02605/) (2018)
- [The curse of overparametrization in adversarial training: Precise analysis of robust generalization for random features regression](https://mlsanity.com/paper/2201.05149/) (2022)
- [Bayes-optimal Learning of Deep Random Networks of Extensive-width](https://mlsanity.com/paper/2302.00375/) (2023)
- [Spectral Evolution and Invariance in Linear-width Neural Networks](https://mlsanity.com/paper/2211.06506/) (2022)
- [What causes the test error? Going beyond bias-variance via ANOVA](https://mlsanity.com/paper/2010.05170/) (2020)
- [Concentration Inequalities for Statistical Inference](https://mlsanity.com/paper/2011.02258/) (2020)
- [Precise Asymptotic Analysis of Deep Random Feature Models](https://mlsanity.com/paper/2302.06210/) (2023)
- [A Random Matrix Perspective on Mixtures of Nonlinearities for Deep Learning](https://mlsanity.com/paper/1912.00827/) (2019)
- [Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks](https://mlsanity.com/paper/2305.06986/) (2023)
- [On double-descent in uncertainty quantification in overparametrized models](https://mlsanity.com/paper/2210.12760/) (2022)

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

- [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)
- [Asymptotics of feature learning in two-layer networks after one gradient-step](https://mlsanity.com/paper/2402.04980/) (2024)
- [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://mlsanity.com/paper/2405.15459/) (2024)
- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)
- [Nonlinear spiked covariance matrices and signal propagation in deep neural networks](https://mlsanity.com/paper/2402.10127/) (2024)
- [Dimension-free deterministic equivalents and scaling laws for random feature regression](https://mlsanity.com/paper/2405.15699/) (2024)
- [Scaling Laws and Spectra of Shallow Neural Networks in the Feature Learning Regime](https://mlsanity.com/paper/2509.24882/) (2025)
- [Fundamental computational limits of weak learnability in high-dimensional multi-index models](https://mlsanity.com/paper/2405.15480/) (2024)
- [A Random Matrix Theory Perspective on the Spectrum of Learned Features and Asymptotic Generalization Capabilities](https://mlsanity.com/paper/2410.18938/) (2024)
- [Provable Guarantees for Neural Networks via Gradient Feature Learning](https://mlsanity.com/paper/2310.12408/) (2023)
- [Asymptotics of Learning with Deep Structured (Random) Features](https://mlsanity.com/paper/2402.13999/) (2024)
- [High-dimensional learning of narrow neural networks](https://mlsanity.com/paper/2409.13904/) (2024)
- [On Regularization via Early Stopping for Least Squares Regression](https://mlsanity.com/paper/2406.04425/) (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)
- [Spurious Correlations in High Dimensional Regression: The Roles of Regularization, Simplicity Bias and Over-Parameterization](https://mlsanity.com/paper/2502.01347/) (2025)
- [Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks](https://mlsanity.com/paper/2506.02651/) (2025)
- [On the Mechanisms of Weak-to-Strong Generalization: A Theoretical Perspective](https://mlsanity.com/paper/2505.18346/) (2025)
- [Gradient Descent with Large Step Sizes: Chaos and Fractal Convergence Region](https://mlsanity.com/paper/2509.25351/) (2025)
- [Understanding the Mechanisms of Fast Hyperparameter Transfer](https://mlsanity.com/paper/2512.22768/) (2025)
- [Spectral Dynamics in Deep Networks: Feature Learning, Outlier Escape, and Learning Rate Transfer](https://mlsanity.com/paper/2605.07870/) (2026)
- [Towards Understanding the Word Sensitivity of Attention Layers: A Study via Random Features](https://mlsanity.com/paper/2402.02969/) (2024)
- [Robust Feature Learning for Multi-Index Models in High Dimensions](https://mlsanity.com/paper/2410.16449/) (2024)
- [Generalization Analysis of Online Stochastic Gradient Descent for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/2407.07670/) (2024)
- [Asymptotic Analysis of Two-Layer Neural Networks after One Gradient Step under Gaussian Mixtures Data with Structure](https://mlsanity.com/paper/2503.00856/) (2025)
- [Low Rank Gradients and Where to Find Them](https://mlsanity.com/paper/2510.01303/) (2025)
- [Online Learning and Information Exponents: On The Importance of Batch size, and Time/Complexity Tradeoffs](https://mlsanity.com/paper/2406.02157/) (2024)
- [Bayes optimal learning of attention-indexed models](https://mlsanity.com/paper/2506.01582/) (2025)
- [A Law of Data Reconstruction for Random Features (and Beyond)](https://mlsanity.com/paper/2509.22214/) (2025)
- [Learning Multi-Index Models with Hyper-Kernel Ridge Regression](https://mlsanity.com/paper/2510.02532/) (2025)

## Abstract (arXiv)

Feature learning is thought to be one of the fundamental reasons for the success of deep neural networks. It is rigorously known that in two-layer fully-connected neural networks under certain conditions, one step of gradient descent on the first layer can lead to feature learning; characterized by the appearance of a separated rank-one component -- spike -- in the spectrum of the feature matrix. However, with a constant gradient descent step size, this spike only carries information from the linear component of the target function and therefore learning non-linear components is impossible. We show that with a learning rate that grows with the sample size, such training in fact introduces multiple rank-one components, each corresponding to a specific polynomial feature. We further prove that the limiting large-dimensional and large sample training and test errors of the updated neural networks are fully characterized by these spikes. By precisely analyzing the improvement in the training and test errors, we demonstrate that these non-linear features can enhance learning.