# Feature selection with gradient descent on two-layer networks in low-rotation regimes

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 50 references that are in the map.
- First posted on arXiv 4 Aug 2022.

## Builds on (most cited first, 26 of 26)

- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networks](https://mlsanity.com/paper/1901.08584/) (2019)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [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)
- [Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data](https://mlsanity.com/paper/1808.01204/) (2018)
- [Taming the Monster: A Fast and Simple Algorithm for Contextual Bandits](https://mlsanity.com/paper/1402.0555/) (2014)
- [Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit](https://mlsanity.com/paper/2207.08799/) (2022)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Directional convergence and alignment in deep learning](https://mlsanity.com/paper/2006.06657/) (2020)
- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (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)
- [Regularization Matters: Generalization and Optimization of Neural Nets v.s. their Induced Kernel](https://mlsanity.com/paper/1810.05369/) (2018)
- [Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks](https://mlsanity.com/paper/1909.12292/) (2019)
- [Risk and parameter convergence of logistic regression](https://mlsanity.com/paper/1803.07300/) (2018)
- [How Much Over-parameterization Is Sufficient to Learn Deep ReLU Networks?](https://mlsanity.com/paper/1911.12360/) (2019)
- [Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias](https://mlsanity.com/paper/2110.13905/) (2021)
- [Beyond Linearization: On Quadratic and Higher-Order Approximation of Wide Neural Networks](https://mlsanity.com/paper/1910.01619/) (2019)
- [Learning Parities with Neural Networks](https://mlsanity.com/paper/2002.07400/) (2020)
- [A Theoretical Analysis on Feature Learning in Neural Networks: Emergence from Inputs and Advantage over Fixed Features](https://mlsanity.com/paper/2206.01717/) (2022)
- [Margins, Shrinkage, and Boosting](https://mlsanity.com/paper/1303.4172/) (2013)
- [On Margin Maximization in Linear and ReLU Networks](https://mlsanity.com/paper/2110.02732/) (2021)
- [Neural tangent kernels, transportation mappings, and universal approximation](https://mlsanity.com/paper/1910.06956/) (2019)

## Built on by (most cited first, 14 of 14)

- [SGD learning on neural networks: leap complexity and saddle-to-saddle dynamics](https://mlsanity.com/paper/2302.11055/) (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)
- [Linguistic Collapse: Neural Collapse in (Large) Language Models](https://mlsanity.com/paper/2405.17767/) (2024)
- [Why Do You Grok? A Theoretical Analysis of Grokking Modular Addition](https://mlsanity.com/paper/2407.12332/) (2024)
- [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)
- [Provable Multi-Task Representation Learning by Two-Layer ReLU Neural Networks](https://mlsanity.com/paper/2307.06887/) (2023)
- [Butterfly Effects of SGD Noise: Error Amplification in Behavior Cloning and Autoregression](https://mlsanity.com/paper/2310.11428/) (2023)
- [Provable Guarantees for Neural Networks via Gradient Feature Learning](https://mlsanity.com/paper/2310.12408/) (2023)
- [Pareto Frontiers in Neural Feature Learning: Data, Compute, Width, and Luck](https://mlsanity.com/paper/2309.03800/) (2023)
- [Benign Overfitting in Two-Layer ReLU Convolutional Neural Networks for XOR Data](https://mlsanity.com/paper/2310.01975/) (2023)
- [LoRanPAC: Low-rank Random Features and Pre-trained Models for Bridging Theory and Practice in Continual Learning](https://mlsanity.com/paper/2410.00645/) (2024)
- [Pruning Before Training May Improve Generalization, Provably](https://mlsanity.com/paper/2301.00335/) (2023)
- [Convergence of Gradient Descent for General Neural Network Architectures Beyond the NTK Regime](https://mlsanity.com/paper/2606.23364/) (2026)
- [Data Uniformity Improves Training Efficiency and More, with a Convergence Framework Beyond the NTK Regime](https://mlsanity.com/paper/2506.24120/) (2025)

## Abstract (arXiv)

This work establishes low test error of gradient flow (GF) and stochastic gradient descent (SGD) on two-layer ReLU networks with standard initialization, in three regimes where key sets of weights rotate little (either naturally due to GF and SGD, or due to an artificial constraint), and making use of margins as the core analytic technique. The first regime is near initialization, specifically until the weights have moved by O(√ m), where m denotes the network width, which is in sharp contrast to the O(1) weight motion allowed by the Neural Tangent Kernel (NTK); here it is shown that GF and SGD only need a network width and number of samples inversely proportional to the NTK margin, and moreover that GF attains at least the NTK margin itself, which suffices to establish escape from bad KKT points of the margin objective, whereas prior work could only establish nondecreasing but arbitrarily small margins. The second regime is the Neural Collapse (NC) setting, where data lies in extremely-well-separated groups, and the sample complexity scales with the number of groups; here the contribution over prior work is an analysis of the entire GF trajectory from initialization. Lastly, if the inner layer weights are constrained to change in norm only and can not rotate, then GF with large widths achieves globally maximal margins, and its sample complexity scales with their inverse; this is in contrast to prior work, which required infinite width and a tricky dual convergence assumption. As purely technical contributions, this work develops a variety of potential functions and other tools which will hopefully aid future work.