# SGD Finds then Tunes Features in Two-Layer Neural Networks with near-Optimal Sample Complexity: A Case Study in the XOR problem

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

## Key facts

- Cited by 39 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 46 references that are in the map.
- First posted on arXiv 26 Sep 2023.

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

- [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)
- [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)
- [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)
- [Linearized two-layers neural networks in high dimension](https://mlsanity.com/paper/1904.12191/) (2019)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [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)
- [Learning Single-Index Models with Shallow Neural Networks](https://mlsanity.com/paper/2210.15651/) (2022)
- [The Landscape of Empirical Risk for Non-convex Losses](https://mlsanity.com/paper/1607.06534/) (2016)
- [Saddle-to-Saddle Dynamics in Deep Linear Networks: Small Initialization Training, Symmetry, and Sparsity](https://mlsanity.com/paper/2106.15933/) (2021)
- [Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias](https://mlsanity.com/paper/2110.13905/) (2021)
- [Learning One-hidden-layer Neural Networks with Landscape Design](https://mlsanity.com/paper/1711.00501/) (2017)
- [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)
- [High-dimensional limit theorems for SGD: Effective dynamics and critical scaling](https://mlsanity.com/paper/2206.04030/) (2022)
- [Classifying high-dimensional Gaussian mixtures: Where kernel methods fail and neural networks succeed](https://mlsanity.com/paper/2102.11742/) (2021)
- [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)
- [The staircase property: How hierarchical structure can guide deep learning](https://mlsanity.com/paper/2108.10573/) (2021)
- [Gradient Descent Quantizes ReLU Network Features](https://mlsanity.com/paper/1803.08367/) (2018)
- [Saddle-to-Saddle Dynamics in Diagonal Linear Networks](https://mlsanity.com/paper/2304.00488/) (2023)
- [Towards Understanding Hierarchical Learning: Benefits of Neural Representations](https://mlsanity.com/paper/2006.13436/) (2020)

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

- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)
- [Neural network learns low-dimensional polynomials with SGD near the information-theoretic limit](https://mlsanity.com/paper/2406.01581/) (2024)
- [RL for Reasoning by Adaptively Revealing Rationales](https://mlsanity.com/paper/2506.18110/) (2025)
- [On the Complexity of Learning Sparse Functions with Statistical and Gradient Queries](https://mlsanity.com/paper/2407.05622/) (2024)
- [Survey on Algorithms for multi-index models](https://mlsanity.com/paper/2504.05426/) (2025)
- [Saddle-to-Saddle Dynamics Explains A Simplicity Bias Across Neural Network Architectures](https://mlsanity.com/paper/2512.20607/) (2025)
- [Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks](https://mlsanity.com/paper/2506.06489/) (2025)
- [Complexity Matters: Dynamics of Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2403.03375/) (2024)
- [Shallow diffusion networks provably learn hidden low-dimensional structure](https://mlsanity.com/paper/2410.11275/) (2024)
- [Learning Gaussian Multi-Index Models with Gradient Flow: Time Complexity and Directional Convergence](https://mlsanity.com/paper/2411.08798/) (2024)
- [Sharper Guarantees for Learning Neural Network Classifiers with Gradient Methods](https://mlsanity.com/paper/2410.10024/) (2024)
- [Robust Feature Learning for Multi-Index Models in High Dimensions](https://mlsanity.com/paper/2410.16449/) (2024)
- [When Are Bias-Free ReLU Networks Effectively Linear Networks?](https://mlsanity.com/paper/2406.12615/) (2024)
- [Learning a Single Index Model from Anisotropic Data with vanilla Stochastic Gradient Descent](https://mlsanity.com/paper/2503.23642/) (2025)
- [How Does Label Noise Gradient Descent Improve Generalization in the Low SNR Regime?](https://mlsanity.com/paper/2510.17526/) (2025)
- [Transformers Provably Learn Sparse XOR with Polylogarithmic Parameters](https://mlsanity.com/paper/2502.07553/) (2025)
- [On the Theory of Continual Learning with Gradient Descent for Neural Networks](https://mlsanity.com/paper/2510.05573/) (2025)
- [Universality of high-dimensional scaling limits of stochastic gradient descent](https://mlsanity.com/paper/2512.13634/) (2025)
- [Positive Distribution Shift as a Framework for Understanding Tractable Learning](https://mlsanity.com/paper/2602.08907/) (2026)
- [Neural Networks Provably Learn Spectral Representations for Group Composition](https://mlsanity.com/paper/2606.02993/) (2026)
- [Linearly Separable Features in Shallow Nonlinear Networks: Width Scales Polynomially with Intrinsic Data Dimension](https://mlsanity.com/paper/2501.02364/) (2025)
- [Convergence of Shallow ReLU Networks on Weakly Interacting Data](https://mlsanity.com/paper/2502.16977/) (2025)
- [Favorability of Loss Landscape with Weight Decay Requires Both Large Overparametrization and Initialization](https://mlsanity.com/paper/2505.22578/) (2025)
- [Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning](https://mlsanity.com/paper/2510.12026/) (2025)
- [The two clocks and the innovation window: When and how generative models learn rules](https://mlsanity.com/paper/2605.10019/) (2026)
- [The Benefits of Temporal Correlations: SGD Learns k-Juntas from Random Walks Efficiently](https://mlsanity.com/paper/2605.10237/) (2026)
- [Learning with Shallow Neural Networks on Cluster-Structured Features](https://mlsanity.com/paper/2605.14927/) (2026)
- [A Function-Space Dichotomy for Compositional Learning: Exponential Sub-Optimality of the Neural Tangent Kernel](https://mlsanity.com/paper/2607.06382/) (2026)
- [A Derandomization Framework for Structure Discovery: Applications in Neural Networks and Beyond](https://mlsanity.com/paper/2510.19382/) (2025)
- [Shallow Neural Networks Learn Low-Degree Spherical Polynomials with Feature Learning by Learnable Channel Attention](https://mlsanity.com/paper/2512.20562/) (2025)

## Abstract (arXiv)

In this work, we consider the optimization process of minibatch stochastic gradient descent (SGD) on a 2-layer neural network with data separated by a quadratic ground truth function. We prove that with data drawn from the d-dimensional Boolean hypercube labeled by the quadratic ``XOR'' function y = -xixj, it is possible to train to a population error o(1) with d :polylog(d) samples. Our result considers simultaneously training both layers of the two-layer-neural network with ReLU activations via standard minibatch SGD on the logistic loss. To our knowledge, this work is the first to give a sample complexity of tildeO(d) for efficiently learning the XOR function on isotropic data on a standard neural network with standard training. Our main technique is showing that the network evolves in two phases: a textitsignal-finding phase where the network is small and many of the neurons evolve independently to find features, and a textitsignal-heavy phase, where SGD maintains and balances the features. We leverage the simultaneous training of the layers to show that it is sufficient for only a small fraction of the neurons to learn features, since those neurons will be amplified by the simultaneous growth of their second layer weights.