[Paper Review] A Mean Field View of the Landscape of Two-Layers Neural Networks

Mei Song, Andrea Montanari+1 more|arXiv (Cornell University)|Apr 18, 2018

Model Reduction and Neural Networks23 references44 citations

TL;DR

The paper derives a mean-field scaling limit for stochastic gradient descent on two-layer neural networks, formulating a nonlinear PDE (distributional dynamics) that describes SGD as a gradient flow in Wasserstein space, and proves convergence results indicating SGD can reach near-optimal generalization in several settings.

ABSTRACT

Multi-layer neural networks are among the most powerful models in machine learning, yet the fundamental reasons for this success defy mathematical understanding. Learning a neural network requires to optimize a non-convex high-dimensional objective (risk function), a problem which is usually attacked using stochastic gradient descent (SGD). Does SGD converge to a global optimum of the risk or only to a local optimum? In the first case, does this happen because local minima are absent, or because SGD somehow avoids them? In the second, why do local minima reached by SGD have good generalization properties? In this paper we consider a simple case, namely two-layers neural networks, and prove that -in a suitable scaling limit- SGD dynamics is captured by a certain non-linear partial differential equation (PDE) that we call distributional dynamics (DD). We then consider several specific examples, and show how DD can be used to prove convergence of SGD to networks with nearly ideal generalization error. This description allows to 'average-out' some of the complexities of the landscape of neural networks, and can be used to prove a general convergence result for noisy SGD.

Motivation & Objective

• Motivate and analyze learning in two-layer neural networks under a one-pass SGD regime.
• Introduce a distributional dynamics PDE that describes SGD in the limit N→∞, ε→0.
• Show how the PDE exploits symmetries and simplifies landscape analysis.
• Demonstrate convergence to near-optimal generalization in representative data/models.
• Provide finite-N and noisy-SGD extensions with convergence guarantees.

Proposed method

• Represent the population risk as R_N(θ) = R# + 2∫V(θ)ρ(dθ) + ∫∫U(θ,θ′)ρ(dθ)ρ(dθ′).
• Derive the distributional dynamics PDE: ∂tρ_t = 2ξ(t) ∇·(ρ_t ∇Ψ(θ;ρ_t)) with Ψ = V + ∫U(θ,θ′)ρ(dθ′).
• Show the connection to Wasserstein gradient flow for the infinite-N limit.
• Extend to noisy SGD yielding a diffusion-augmented PDE: ∂tρ_t = 2ξ(t)∇·(ρ_t ∇Ψ_λ(θ;ρ_t)) + 2ξ(t)/β Δθρ_t.
• Prove propagation of chaos: empirical distribution from SGD converges to ρ_t under specified scaling.
• Provide non-asymptotic bounds linking R_N(θ^k) and R(ρ_t).
• Apply the framework to isotropic/anisotropic Gaussian data and ReLU activations to illustrate convergence and failure modes.

Experimental results

Research questions

• RQ1 Does SGD on two-layer networks converge to a global optimum or do local minima persist under typical data distributions?
• RQ2 Can a mean-field PDE accurately describe SGD dynamics in the large-N limit and what are the implications for generalization?
• RQ3 How do data distributions with symmetry (isotropic/anisotropic Gaussians) affect the limiting dynamics and convergence?
• RQ4 What finite-N and noisy-SGD guarantees can be established within the distributional dynamics framework?
• RQ5 Under what conditions can SGD escape poor local minima and achieve near-ideal generalization?

Key findings

• SGD dynamics for two-layer networks are captured by a nonlinear PDE (distributional dynamics) in the scaling limit (N→∞, ε→0).
• DD acts as a gradient flow in the Wasserstein space, minimizing an asymptotic risk R(ρ) with local mass conservation.
• For noisy SGD, the dynamics converge to the minimizer of a entropy-regularized free energy, yielding global convergence in many steps independent of N.
• In several constructed examples (centered isotropic and anisotropic Gaussians, with varied activations), SGD converges to networks with near-ideal generalization, and finite-N behavior closely matches the PDE predictions.
• The theory provides non-asymptotic error bounds linking finite-N risk to the limiting risk and describes fixed points and stability properties of the DD and diffusion DD.
• Numerical experiments corroborate the DD predictions for both statics (minimizers) and dynamics (convergence trajectories).