[Paper Review] On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport

Lénaïc Chizat, Francis Bach|arXiv (Cornell University)|May 23, 2018

Stochastic Gradient Optimization TechniquesComputer Science1 references187 citations

TL;DR

The paper studies gradient flows on discretized measures (particles) and proves that, in the many-particle limit, the associated Wasserstein gradient flow converges to global minimizers under certain homogeneity and initialization separation conditions. This provides a qualitative consistency principle for over-parameterized, non-convex models.

ABSTRACT

Many tasks in machine learning and signal processing can be solved by minimizing a convex function of a measure. This includes sparse spikes deconvolution or training a neural network with a single hidden layer. For these problems, we study a simple minimization method: the unknown measure is discretized into a mixture of particles and a continuous-time gradient descent is performed on their weights and positions. This is an idealization of the usual way to train neural networks with a large hidden layer. We show that, when initialized correctly and in the many-particle limit, this gradient flow, although non-convex, converges to global minimizers. The proof involves Wasserstein gradient flows, a by-product of optimal transport theory. Numerical experiments show that this asymptotic behavior is already at play for a reasonable number of particles, even in high dimension.

Motivation & Objective

• Explain when and why non-convex particle gradient flows converge to global minima in over-parameterized settings.
• Link finite-particle gradient flows to an infinite-dimensional Wasserstein gradient flow framework.
• Show how homogeneity and structured initialization lead to global optimality in lifted formulations relevant to neural networks and sparse deconvolution.
• Provide guidance by applying results to single-hidden-layer neural networks and sparse spikes deconvolution.

Proposed method

• Lift the unknown measure into a convex functional F(mu) that separates a smooth loss R and a potential V via a lifted representation Phi and domain Omega.
• Discretize mu as a mixture of m particles and study the particle gradient flow of F_m, with velocity given by a projected negative gradient (Eq. 5).
• Move from particle dynamics to a Wasserstein gradient flow of F on probability measures, governed by the continuity equation with velocity in the Wasserstein subdifferential (Definition 2.4).
• Establish a general many-particle limit: mu_{m,t} converges to a Wasserstein gradient flow mu_t solving the limit dynamics (Theorem 2.6).
• Analyze global convergence in two homogeneous settings: 2-homogeneous (covering ReLU, lifted problems) and partially 1-homogeneous (bounded Phi, e.g., sparse deconvolution and sigmoid networks) cases (Theorems 3.3 and 3.5).
• Apply results to sparse deconvolution and neural networks with a single hidden layer, detailing initialization and boundary/Sard-type regularity assumptions (Section 4).

Experimental results

Research questions

• RQ1 Under what structural conditions (homogeneity) and initialization patterns do gradient flows over particle representations converge to global minimizers of the lifted functional F?
• RQ2 Can the finite-particle gradient flow dynamics for over-parameterized models be described by a Wasserstein gradient flow in the mean-field limit?
• RQ3 How do lifted, homogeneous representations apply to practical architectures like sparse deconvolution and single-hidden-layer neural networks with sigmoid or ReLU activations?
• RQ4 What role do initialization separation properties and Sard-type regularity play in ensuring convergence to global minima?
• RQ5 Do empirical (finite m) particle systems exhibit the asymptotic global convergence predicted by the Wasserstein framework?

Key findings

• In the many-particle limit, the discrete particle gradient flow converges to a unique Wasserstein gradient flow for F with appropriate initial conditions (Theorem 2.6).
• If the Wasserstein gradient flow converges under 2-homogeneity or partial 1-homogeneity, the limit is a global minimizer of F (Theorems 3.3 and 3.5).
• For lifted problems corresponding to neural networks and sparse deconvolution, the analysis prescribes initialization patterns that separate certain spheres in parameter space to ensure global convergence (assumptions in Theorems 3.3, 3.5).
• The results cover ReLU and sigmoid activations in single-hidden-layer networks, with explicit discussion of boundary conditions and Sard-type regularity to guarantee convergence (Section 4).
• Numerical experiments indicate that the asymptotic regime is observable with a reasonable number of particles, and particle gradient flow can outperform optimization on a fixed large particle grid (Section 4.3).