A Mean Field View of the Landscape of Two-Layers Neural Networks (1804.06561v2)

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.

• The paper demonstrates that SGD dynamics in two-layer networks can be approximated by a non-linear PDE using a mean field approach.
• It shows that the empirical parameter distribution converges via a Wasserstein gradient flow, simplifying the risk landscape.
• Numerical experiments on Gaussian data validate the framework, highlighting robust convergence toward global minima.

A Mean Field View of the Landscape of Two-Layers Neural Networks

Introduction to Concepts and Methodology

The paper explores the optimization landscape of two-layer neural networks using a mean field approach. This perspective addresses classical questions regarding the behavior of stochastic gradient descent (SGD) in non-convex scenarios, such as whether SGD converges to a local or global minimum, and how it avoids poor minima with complex landscapes.

The authors propose that the dynamics of SGD can be approximated by a non-linear partial differential equation (PDE), termed distributional dynamics (DD). This PDE effectively describes the evolution of the empirical distribution of network parameters, allowing for an 'averaging-out' of the intricate details of the optimization landscape inherent in neural networks.

Implementation of Distributional Dynamics (DD)

The key insight is realizing that in the large $N$ limit—where $N$ is the number of neurons—the empirical distribution of network parameters converges to a probability measure $\rho$ that solves the DD PDE:

\partial_t \rho_t = 2 \xi(t) \nabla_\theta \cdot \left( \rho_t \nabla_\theta \Psi(\theta; \rho_t) \right)

This equation employs the Wasserstein gradient flow framework in the space of probability measures, which is crucial for analyzing convergence properties under SGD dynamics.

Steps for Implementation:

1. Define the Framework:
   • Start by specifying the architecture of a two-layer neural network and the associated risk (or loss) function.
   • Select activation functions and parameterize the network.
2. Empirical Distribution Initialization:
   • Initialize parameters $(\theta_i)_{i \le N} \sim_{iid} \rho_0$.
   • Define the empirical distribution of parameters as $\rho = \frac{1}{N}\sum_{i=1}^{N} \delta_{\theta_i}$.
3. Stochastic Gradient Descent Setup:
   • Implement standard SGD iterations using the risk function.
   • Ensure each training example is visited only once (one-pass assumption).
4. Numerical PDE Solution:
   • Use discretization techniques to simulate the PDE over time, which models large-scale network parameter updates.
   • Apply techniques such as multiple-deltas approximation for efficient computation.

Analysis and Expected Outcomes

The analysis introduced various conditions under which DD would converge to minimizers effectively removing unwanted local minima:

• Investigation into $R(\rho)$ where $R$ is the risk function defined for distribution $\rho \in P(\mathbb R^D)$.
• Focus on convexity in large $N$ limits, demonstrating that the landscape effectively simplifies.

Numerical results, such as those shown in figures involving risk plots and dynamics, suggest that under structured random initialization, SGD pathologically avoids poor minima and converges to solutions near the global minimum.

Examples and Case Studies

Cases of Gaussian data distributions were examined to illustrate convergence phenomena. In experiments comparing isotropic and anisotropic Gaussian models:

• Isotropic Gaussian: Demonstrated rapid convergence with simpler optimization surfaces.
• Anisotropic Gaussian: Showed neural networks successfully identifying relevant subspace features.

These numerical examples validate theoretical predictions, typically using high-dimensional PDE approximations.

Conclusion and Implementation Considerations

This framework broadly addresses optimization challenges in neural networks, potentially guiding hybrid combinatory algorithms that leverage mean-field landscapes for faster convergence. Given the validation across various numerical scenarios, the approach provides promising insights into designing neural architectures resistant to hill phenomena.

Implementation Considerations:

• Assess computational cost as a function of neuron count $N$.
• Make use of efficient numerical solvers and grid methods for PDE evaluations.
• Ensure validation against theoretical risk landscape simplifications in applications involving massive data.

Deploying this methodology involves exploiting symmetry and SDP techniques for high-dimensional feature space analysis, making it applicable across supervised learning tasks in both structured and unstructured data forms.