# Gradient-Based Feature Learning under Structured Data

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

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 66 references that are in the map.
- First posted on arXiv 7 Sep 2023.

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

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [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)
- [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)
- [When Do Neural Networks Outperform Kernel Methods?](https://mlsanity.com/paper/2006.13409/) (2020)
- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (2023)
- [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)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [On the Power and Limitations of Random Features for Understanding Neural Networks](https://mlsanity.com/paper/1904.00687/) (2019)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [Limitations of Lazy Training of Two-layers Neural Networks](https://mlsanity.com/paper/1906.08899/) (2019)
- [Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed](https://mlsanity.com/paper/2102.11742/) (2021)
- [Neural Networks Efficiently Learn Low-Dimensional Representations with SGD](https://mlsanity.com/paper/2209.14863/) (2022)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (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)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Learning ReLUs via Gradient Descent](https://mlsanity.com/paper/1705.04591/) (2017)
- [Learning a Single Neuron with Gradient Methods](https://mlsanity.com/paper/2001.05205/) (2020)
- [Agnostic Learning of a Single Neuron with Gradient Descent](https://mlsanity.com/paper/2005.14426/) (2020)
- [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)
- [Learning Polynomials of Few Relevant Dimensions](https://mlsanity.com/paper/2004.13748/) (2020)
- [Time/Accuracy Tradeoffs for Learning a ReLU with respect to Gaussian Marginals](https://mlsanity.com/paper/1911.01462/) (2019)
- [How rotational invariance of common kernels prevents generalization in high dimensions](https://mlsanity.com/paper/2104.04244/) (2021)

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

- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [Repetita Iuvant: Data Repetition Allows SGD to Learn High-Dimensional Multi-Index Functions](https://mlsanity.com/paper/2405.15459/) (2024)
- [Asymptotics of feature learning in two-layer networks after one gradient-step](https://mlsanity.com/paper/2402.04980/) (2024)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [Nonlinear spiked covariance matrices and signal propagation in deep neural networks](https://mlsanity.com/paper/2402.10127/) (2024)
- [Survey on Algorithms for multi-index models](https://mlsanity.com/paper/2504.05426/) (2025)
- [Learning Multi-Index Models with Neural Networks via Mean-Field Langevin Dynamics](https://mlsanity.com/paper/2408.07254/) (2024)
- [Sliding down the stairs: how correlated latent variables accelerate learning with neural networks](https://mlsanity.com/paper/2404.08602/) (2024)
- [On the Mechanisms of Weak-to-Strong Generalization: A Theoretical Perspective](https://mlsanity.com/paper/2505.18346/) (2025)
- [Kernel ridge regression under power-law data: spectrum and generalization](https://mlsanity.com/paper/2510.04780/) (2025)
- [Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks](https://mlsanity.com/paper/2506.02651/) (2025)
- [From Spikes to Heavy Tails: Unveiling the Spectral Evolution of Neural Networks](https://mlsanity.com/paper/2406.04657/) (2024)
- [Low-dimensional Functions are Efficiently Learnable under Randomly Biased Distributions](https://mlsanity.com/paper/2502.06443/) (2025)
- [Provable Benefits of Unsupervised Pre-training and Transfer Learning via Single-Index Models](https://mlsanity.com/paper/2502.16849/) (2025)
- [Robust Feature Learning for Multi-Index Models in High Dimensions](https://mlsanity.com/paper/2410.16449/) (2024)
- [Low Rank Gradients and Where to Find Them](https://mlsanity.com/paper/2510.01303/) (2025)
- [Asymptotic Analysis of Two-Layer Neural Networks after One Gradient Step under Gaussian Mixtures Data with Structure](https://mlsanity.com/paper/2503.00856/) (2025)
- [Learning a Single Index Model from Anisotropic Data with vanilla Stochastic Gradient Descent](https://mlsanity.com/paper/2503.23642/) (2025)
- [Online Learning and Information Exponents: On The Importance of Batch size, and Time/Complexity Tradeoffs](https://mlsanity.com/paper/2406.02157/) (2024)
- [Universality of high-dimensional scaling limits of stochastic gradient descent](https://mlsanity.com/paper/2512.13634/) (2025)
- [How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability](https://mlsanity.com/paper/2601.19208/) (2026)
- [Positive Distribution Shift as a Framework for Understanding Tractable Learning](https://mlsanity.com/paper/2602.08907/) (2026)
- [Input-Label Correlation Governs a Linear-to-Nonlinear Transition in Random Features under Spiked Covariance](https://mlsanity.com/paper/2409.20250/) (2024)
- [Guided by the Experts: Provable Feature Learning Dynamic of Soft-Routed Mixture-of-Experts](https://mlsanity.com/paper/2510.07205/) (2025)
- [From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD](https://mlsanity.com/paper/2510.21020/) (2025)
- [Learning Hierarchical Polynomials of Multiple Nonlinear Features with Three-Layer Networks](https://mlsanity.com/paper/2411.17201/) (2024)
- [Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning](https://mlsanity.com/paper/2510.12026/) (2025)
- [The Mechanism of Weak-to-Strong Generalization: Feature Elicitation from Latent Knowledge](https://mlsanity.com/paper/2605.12908/) (2026)
- [Learning with Shallow Neural Networks on Cluster-Structured Features](https://mlsanity.com/paper/2605.14927/) (2026)
- [Optimal Representation Size: High-Dimensional Analysis of Pretraining and Linear Probing](https://mlsanity.com/paper/2605.20105/) (2026)

## Abstract (arXiv)

Recent works have demonstrated that the sample complexity of gradient-based learning of single index models, i.e. functions that depend on a 1-dimensional projection of the input data, is governed by their information exponent. However, these results are only concerned with isotropic data, while in practice the input often contains additional structure which can implicitly guide the algorithm. In this work, we investigate the effect of a spiked covariance structure and reveal several interesting phenomena. First, we show that in the anisotropic setting, the commonly used spherical gradient dynamics may fail to recover the true direction, even when the spike is perfectly aligned with the target direction. Next, we show that appropriate weight normalization that is reminiscent of batch normalization can alleviate this issue. Further, by exploiting the alignment between the (spiked) input covariance and the target, we obtain improved sample complexity compared to the isotropic case. In particular, under the spiked model with a suitably large spike, the sample complexity of gradient-based training can be made independent of the information exponent while also outperforming lower bounds for rotationally invariant kernel methods.