# Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias

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

## Key facts

- Cited by 83 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 56 references that are in the map.
- First posted on arXiv 26 Oct 2021.

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

- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Spectrally-normalized margin bounds for neural networks](https://mlsanity.com/paper/1706.08498/) (2017)
- [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)
- [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)
- [Fantastic Generalization Measures and Where to Find Them](https://mlsanity.com/paper/1912.02178/) (2019)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [The Pitfalls of Simplicity Bias in Neural Networks](https://mlsanity.com/paper/2006.07710/) (2020)
- [Learning and Generalization in Overparameterized Neural Networks, Going Beyond Two Layers](https://mlsanity.com/paper/1811.04918/) (2018)
- [Gradient Descent Maximizes the Margin of Homogeneous Neural Networks](https://mlsanity.com/paper/1906.05890/) (2019)
- [A PAC-Bayesian Approach to Spectrally-Normalized Margin Bounds for Neural Networks](https://mlsanity.com/paper/1707.09564/) (2017)
- [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)
- [Characterizing Implicit Bias in Terms of Optimization Geometry](https://mlsanity.com/paper/1802.08246/) (2018)
- [Implicit Bias of Gradient Descent on Linear Convolutional Networks](https://mlsanity.com/paper/1806.00468/) (2018)
- [SGD on Neural Networks Learns Functions of Increasing Complexity](https://mlsanity.com/paper/1905.11604/) (2019)
- [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)
- [Algorithmic Regularization in Learning Deep Homogeneous Models: Layers are Automatically Balanced](https://mlsanity.com/paper/1806.00900/) (2018)
- [Towards Resolving the Implicit Bias of Gradient Descent for Matrix Factorization: Greedy Low-Rank Learning](https://mlsanity.com/paper/2012.09839/) (2020)
- [Stochastic Gradient Descent Optimizes Over-parameterized Deep ReLU Networks](https://mlsanity.com/paper/1811.08888/) (2018)
- [Stochastic subgradient method converges on tame functions](https://mlsanity.com/paper/1804.07795/) (2018)
- [Implicit Regularization of Discrete Gradient Dynamics in Linear Neural Networks](https://mlsanity.com/paper/1904.13262/) (2019)
- [Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow ReLU networks](https://mlsanity.com/paper/1909.12292/) (2019)
- [Convergence of Gradient Descent on Separable Data](https://mlsanity.com/paper/1803.01905/) (2018)

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

- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)
- [On the Implicit Bias in Deep-Learning Algorithms](https://mlsanity.com/paper/2208.12591/) (2022)
- [Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models](https://mlsanity.com/paper/2210.14199/) (2022)
- [Gradient flow dynamics of shallow ReLU networks for square loss and orthogonal inputs](https://mlsanity.com/paper/2206.00939/) (2022)
- [Implicit Bias of AdamW: ℓ_∞ Norm Constrained Optimization](https://mlsanity.com/paper/2404.04454/) (2024)
- [Implicit Regularization Towards Rank Minimization in ReLU Networks](https://mlsanity.com/paper/2201.12760/) (2022)
- [Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Data](https://mlsanity.com/paper/2210.07082/) (2022)
- [Neural Redshift: Random Networks are not Random Functions](https://mlsanity.com/paper/2403.02241/) (2024)
- [Understanding Incremental Learning of Gradient Descent: A Fine-grained Analysis of Matrix Sensing](https://mlsanity.com/paper/2301.11500/) (2023)
- [Get rich quick: exact solutions reveal how unbalanced initializations promote rapid feature learning](https://mlsanity.com/paper/2406.06158/) (2024)
- [Implicit Bias of Gradient Descent on Reparametrized Models: On Equivalence to Mirror Descent](https://mlsanity.com/paper/2207.04036/) (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)
- [Random Feature Amplification: Feature Learning and Generalization in Neural Networks](https://mlsanity.com/paper/2202.07626/) (2022)
- [Feature emergence via margin maximization: case studies in algebraic tasks](https://mlsanity.com/paper/2311.07568/) (2023)
- [Which Features are Learnt by Contrastive Learning? On the Role of Simplicity Bias in Class Collapse and Feature Suppression](https://mlsanity.com/paper/2305.16536/) (2023)
- [Benign Overfitting in Linear Classifiers and Leaky ReLU Networks from KKT Conditions for Margin Maximization](https://mlsanity.com/paper/2303.01462/) (2023)
- [Early Neuron Alignment in Two-layer ReLU Networks with Small Initialization](https://mlsanity.com/paper/2307.12851/) (2023)
- [Implicit Regularization in Hierarchical Tensor Factorization and Deep Convolutional Neural Networks](https://mlsanity.com/paper/2201.11729/) (2022)
- [Simplicity Bias in 1-Hidden Layer Neural Networks](https://mlsanity.com/paper/2302.00457/) (2023)
- [The Asymmetric Maximum Margin Bias of Quasi-Homogeneous Neural Networks](https://mlsanity.com/paper/2210.03820/) (2022)
- [The Double-Edged Sword of Implicit Bias: Generalization vs. Robustness in ReLU Networks](https://mlsanity.com/paper/2303.01456/) (2023)
- [Empirical Phase Diagram for Three-layer Neural Networks with Infinite Width](https://mlsanity.com/paper/2205.12101/) (2022)
- [An overview of condensation phenomenon in deep learning](https://mlsanity.com/paper/2504.09484/) (2025)
- [Phase Diagram of Initial Condensation for Two-layer Neural Networks](https://mlsanity.com/paper/2303.06561/) (2023)
- [Learning a Neuron by a Shallow ReLU Network: Dynamics and Implicit Bias for Correlated Inputs](https://mlsanity.com/paper/2306.06479/) (2023)
- [Saddle-to-Saddle Dynamics Explains A Simplicity Bias Across Neural Network Architectures](https://mlsanity.com/paper/2512.20607/) (2025)
- [Feature selection with gradient descent on two-layer networks in low-rotation regimes](https://mlsanity.com/paper/2208.02789/) (2022)
- [Provable Guarantees for Neural Networks via Gradient Feature Learning](https://mlsanity.com/paper/2310.12408/) (2023)
- [Transformers Learn Low Sensitivity Functions: Investigations and Implications](https://mlsanity.com/paper/2403.06925/) (2024)

## Abstract (arXiv)

The generalization mystery of overparametrized deep nets has motivated efforts to understand how gradient descent (GD) converges to low-loss solutions that generalize well. Real-life neural networks are initialized from small random values and trained with cross-entropy loss for classification (unlike the "lazy" or "NTK" regime of training where analysis was more successful), and a recent sequence of results (Lyu and Li, 2020; Chizat and Bach, 2020; Ji and Telgarsky, 2020) provide theoretical evidence that GD may converge to the "max-margin" solution with zero loss, which presumably generalizes well. However, the global optimality of margin is proved only in some settings where neural nets are infinitely or exponentially wide. The current paper is able to establish this global optimality for two-layer Leaky ReLU nets trained with gradient flow on linearly separable and symmetric data, regardless of the width. The analysis also gives some theoretical justification for recent empirical findings (Kalimeris et al., 2019) on the so-called simplicity bias of GD towards linear or other "simple" classes of solutions, especially early in training. On the pessimistic side, the paper suggests that such results are fragile. A simple data manipulation can make gradient flow converge to a linear classifier with suboptimal margin.