# Analysis of a Two-Layer Neural Network via Displacement Convexity

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

## Key facts

- Cited by 27 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 5 Jan 2019.

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

- [Neural Networks can Learn Representations with Gradient Descent](https://mlsanity.com/paper/2206.15144/) (2022)
- [Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction](https://mlsanity.com/paper/2106.15013/) (2021)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Trajectory Inference via Mean-field Langevin in Path Space](https://mlsanity.com/paper/2205.07146/) (2022)
- [Birth-death dynamics for sampling: Global convergence, approximations and their asymptotics](https://mlsanity.com/paper/2211.00450/) (2022)
- [Stochastic gradient descent with noise of machine learning type. Part II: Continuous time analysis](https://mlsanity.com/paper/2106.02588/) (2021)
- [Generalization and Stability of Interpolating Neural Networks with Minimal Width](https://mlsanity.com/paper/2302.09235/) (2023)
- [Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction](https://mlsanity.com/paper/2306.07221/) (2023)
- [Global Convergence of Three-layer Neural Networks in the Mean Field Regime](https://mlsanity.com/paper/2105.05228/) (2021)
- [Analysis of feature learning in weight-tied autoencoders via the mean field lens](https://mlsanity.com/paper/2102.08373/) (2021)
- [Mean-field Analysis of Piecewise Linear Solutions for Wide ReLU Networks](https://mlsanity.com/paper/2111.02278/) (2021)
- [On Feature Learning in Neural Networks with Global Convergence Guarantees](https://mlsanity.com/paper/2204.10782/) (2022)
- [Fundamental limits of overparametrized shallow neural networks for supervised learning](https://mlsanity.com/paper/2307.05635/) (2023)
- [Partially Observed Trajectory Inference using Optimal Transport and a Dynamics Prior](https://mlsanity.com/paper/2406.07475/) (2024)
- [Doubly Regularized Entropic Wasserstein Barycenters](https://mlsanity.com/paper/2303.11844/) (2023)
- [Understanding the training of infinitely deep and wide ResNets with Conditional Optimal Transport](https://mlsanity.com/paper/2403.12887/) (2024)
- [Wasserstein Flow Meets Replicator Dynamics: A Mean-Field Analysis of Representation Learning in Actor-Critic](https://mlsanity.com/paper/2112.13530/) (2021)
- [Stochastic Modified Flows, Mean-Field Limits and Dynamics of Stochastic Gradient Descent](https://mlsanity.com/paper/2302.07125/) (2023)
- [Kernel Approximation of Fisher-Rao Gradient Flows](https://mlsanity.com/paper/2410.20622/) (2024)
- [A blob method for inhomogeneous diffusion with applications to multi-agent control and sampling](https://mlsanity.com/paper/2202.12927/) (2022)
- [Information-theoretic reduction of deep neural networks to linear models in the overparametrized proportional regime](https://mlsanity.com/paper/2505.03577/) (2025)
- [Limiting fluctuation and trajectorial stability of multilayer neural networks with mean field training](https://mlsanity.com/paper/2110.15954/) (2021)
- [Efficient displacement convex optimization with particle gradient descent](https://mlsanity.com/paper/2302.04753/) (2023)
- [Towards Understanding the Word Sensitivity of Attention Layers: A Study via Random Features](https://mlsanity.com/paper/2402.02969/) (2024)
- [Mean-field analysis for heavy ball methods: Dropout-stability, connectivity, and global convergence](https://mlsanity.com/paper/2210.06819/) (2022)
- [A Functional-Space Mean-Field Theory of Partially-Trained Three-Layer Neural Networks](https://mlsanity.com/paper/2210.16286/) (2022)
- [Sharp asymptotics on the compression of two-layer neural networks](https://mlsanity.com/paper/2205.08199/) (2022)

## Abstract (arXiv)

Fitting a function by using linear combinations of a large number N of `simple' components is one of the most fruitful ideas in statistical learning. This idea lies at the core of a variety of methods, from two-layer neural networks to kernel regression, to boosting. In general, the resulting risk minimization problem is non-convex and is solved by gradient descent or its variants. Unfortunately, little is known about global convergence properties of these approaches. Here we consider the problem of learning a concave function f on a compact convex domain Omegasubseteq mathbb Rd, using linear combinations of `bump-like' components (neurons). The parameters to be fitted are the centers of N bumps, and the resulting empirical risk minimization problem is highly non-convex. We prove that, in the limit in which the number of neurons diverges, the evolution of gradient descent converges to a Wasserstein gradient flow in the space of probability distributions over Omega. Further, when the bump width δ tends to 0, this gradient flow has a limit which is a viscous porous medium equation. Remarkably, the cost function optimized by this gradient flow exhibits a special property known as displacement convexity, which implies exponential convergence rates for N→∞, δ→ 0. Surprisingly, this asymptotic theory appears to capture well the behavior for moderate values of δ, N. Explaining this phenomenon, and understanding the dependence on δ,N in a quantitative manner remains an outstanding challenge.