# High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation

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

## Key facts

- Cited by 150 papers in the map (data through 14 Oct 2026).
- Builds on 68 of its 95 references that are in the map.
- First posted on arXiv 3 May 2022.

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

- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Rich feature hierarchies for accurate object detection and semantic segmentation](https://mlsanity.com/paper/1311.2524/) (2013)
- [Reconciling modern machine learning practice and the bias-variance trade-off](https://mlsanity.com/paper/1812.11118/) (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)
- [On Exact Computation with an Infinitely Wide Neural Net](https://mlsanity.com/paper/1904.11955/) (2019)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [Nonparametric regression using deep neural networks with ReLU activation function](https://mlsanity.com/paper/1708.06633/) (2017)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [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)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Breaking the Curse of Dimensionality with Convex Neural Networks](https://mlsanity.com/paper/1412.8690/) (2014)
- [Kernel and Rich Regimes in Overparametrized Models](https://mlsanity.com/paper/2002.09277/) (2020)
- [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020)
- [Deep learning: a statistical viewpoint](https://mlsanity.com/paper/2103.09177/) (2021)
- [The large learning rate phase of deep learning: the catapult mechanism](https://mlsanity.com/paper/2003.02218/) (2020)
- [Just Interpolate: Kernel "Ridgeless" Regression Can Generalize](https://mlsanity.com/paper/1808.00387/) (2018)
- [Feature Learning in Infinite-Width Neural Networks](https://mlsanity.com/paper/2011.14522/) (2020)
- [Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the Neural Tangent Kernel](https://mlsanity.com/paper/2010.15110/) (2020)
- [Adaptivity of deep ReLU network for learning in Besov and mixed smooth Besov spaces: optimal rate and curse of dimensionality](https://mlsanity.com/paper/1810.08033/) (2018)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [Linearized two-layers neural networks in high dimension](https://mlsanity.com/paper/1904.12191/) (2019)
- [The Break-Even Point on Optimization Trajectories of Deep Neural Networks](https://mlsanity.com/paper/2002.09572/) (2020)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [High-Dimensional Asymptotics of Prediction: Ridge Regression and Classification](https://mlsanity.com/paper/1507.03003/) (2015)
- [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)
- [Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernel](https://mlsanity.com/paper/1810.05369/) (2018)
- [Generalisation error in learning with random features and the hidden manifold model](https://mlsanity.com/paper/2002.09339/) (2020)
- [Learning curves of generic features maps for realistic datasets with a teacher-student model](https://mlsanity.com/paper/2102.08127/) (2021)

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

- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [Dataset Distillation using Neural Feature Regression](https://mlsanity.com/paper/2206.00719/) (2022)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [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)
- [How Deep Neural Networks Learn Compositional Data: The Random Hierarchy Model](https://mlsanity.com/paper/2307.02129/) (2023)
- [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)
- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)
- [How Feature Learning Can Improve Neural Scaling Laws](https://mlsanity.com/paper/2409.17858/) (2024)
- [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)
- [Learning threshold neurons via the "edge of stability"](https://mlsanity.com/paper/2212.07469/) (2022)
- [A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks](https://mlsanity.com/paper/2310.07891/) (2023)
- [From Lazy to Rich: Exact Learning Dynamics in Deep Linear Networks](https://mlsanity.com/paper/2409.14623/) (2024)
- [The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit](https://mlsanity.com/paper/2306.17759/) (2023)
- [Muon Outperforms Adam in Tail-End Associative Memory Learning](https://mlsanity.com/paper/2509.26030/) (2025)
- [Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning](https://mlsanity.com/paper/2406.06158/) (2024)
- [The Neural Covariance SDE: Shaped Infinite Depth-and-Width Networks at Initialization](https://mlsanity.com/paper/2206.02768/) (2022)
- [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)
- [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)
- [Gradient-Based Feature Learning under Structured Data](https://mlsanity.com/paper/2309.03843/) (2023)
- [Spectral Evolution and Invariance in Linear-width Neural Networks](https://mlsanity.com/paper/2211.06506/) (2022)
- [There Will Be a Scientific Theory of Deep Learning](https://mlsanity.com/paper/2604.21691/) (2026)
- [Are Gaussian data all you need? Extents and limits of universality in high-dimensional generalized linear estimation](https://mlsanity.com/paper/2302.08923/) (2023)
- [Nonlinear spiked covariance matrices and signal propagation in deep neural networks](https://mlsanity.com/paper/2402.10127/) (2024)
- [The Onset of Variance-Limited Behavior for Networks in the Lazy and Rich Regimes](https://mlsanity.com/paper/2212.12147/) (2022)

## Abstract (arXiv)

We study the first gradient descent step on the first-layer parameters boldsymbolW in a two-layer neural network: f(boldsymbolx) = frac1√Nboldsymbola^topσ(boldsymbolW^topboldsymbolx), where boldsymbolW∈ℝd× N, boldsymbola∈ℝN are randomly initialized, and the training objective is the empirical MSE loss: frac1nsumi=1n (f(boldsymbolxi)-yi)2. In the proportional asymptotic limit where n,d,N→∞ at the same rate, and an idealized student-teacher setting, we show that the first gradient update contains a rank-1 "spike", which results in an alignment between the first-layer weights and the linear component of the teacher model f^*. To characterize the impact of this alignment, we compute the prediction risk of ridge regression on the conjugate kernel after one gradient step on boldsymbolW with learning rate η, when f^* is a single-index model. We consider two scalings of the first step learning rate η. For small η, we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random features model, but cannot defeat the best linear model on the input. Whereas for sufficiently large η, we prove that for certain f^*, the same ridge estimator on trained features can go beyond this "linear regime" and outperform a wide range of random features and rotationally invariant kernels. Our results demonstrate that even one gradient step can lead to a considerable advantage over random features, and highlight the role of learning rate scaling in the initial phase of training.