Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Rigorous Framework for the Mean Field Limit of Multilayer Neural Networks

Phan-Minh Nguyen, Huy Tuan Pham

TL;DR

This work develops a mathematically rigorous mean-field framework for multilayer neural networks via a neuronal embedding that embeds arbitrary widths into a non-evolving probabilistic space. It proves well-posed MF ODEs and a precise coupling between finite networks trained by SGD and their MF limits, with quantitative width-dependent error bounds. The authors unveil degeneracy under IID initializations in deeper networks, and remedy it with bidirectional diversity and correlated initializations, obtaining global convergence results for two-, three-, and deeper-layer architectures under non-convex losses and Morse-Sard conditions. The framework, which applies to broad architectures and initialization schemes, provides a principled non-convex optimization theory for infinite-width neural networks and connects theoretical insights with potential practical implications for feature learning and initialization design.

Abstract

We develop a mathematically rigorous framework for multilayer neural networks in the mean field regime. As the network's widths increase, the network's learning trajectory is shown to be well captured by a meaningful and dynamically nonlinear limit (the \textit{mean field} limit), which is characterized by a system of ODEs. Our framework applies to a broad range of network architectures, learning dynamics and network initializations. Central to the framework is the new idea of a \textit{neuronal embedding}, which comprises of a non-evolving probability space that allows to embed neural networks of arbitrary widths. Using our framework, we prove several properties of large-width multilayer neural networks. Firstly we show that independent and identically distributed initializations cause strong degeneracy effects on the network's learning trajectory when the network's depth is at least four. Secondly we obtain several global convergence guarantees for feedforward multilayer networks under a number of different setups. These include two-layer and three-layer networks with independent and identically distributed initializations, and multilayer networks of arbitrary depths with a special type of correlated initializations that is motivated by the new concept of \textit{bidirectional diversity}. Unlike previous works that rely on convexity, our results admit non-convex losses and hinge on a certain universal approximation property, which is a distinctive feature of infinite-width neural networks and is shown to hold throughout the training process. Aside from being the first known results for global convergence of multilayer networks in the mean field regime, they demonstrate flexibility of our framework and incorporate several new ideas and insights that depart from the conventional convex optimization wisdom.

A Rigorous Framework for the Mean Field Limit of Multilayer Neural Networks

TL;DR

This work develops a mathematically rigorous mean-field framework for multilayer neural networks via a neuronal embedding that embeds arbitrary widths into a non-evolving probabilistic space. It proves well-posed MF ODEs and a precise coupling between finite networks trained by SGD and their MF limits, with quantitative width-dependent error bounds. The authors unveil degeneracy under IID initializations in deeper networks, and remedy it with bidirectional diversity and correlated initializations, obtaining global convergence results for two-, three-, and deeper-layer architectures under non-convex losses and Morse-Sard conditions. The framework, which applies to broad architectures and initialization schemes, provides a principled non-convex optimization theory for infinite-width neural networks and connects theoretical insights with potential practical implications for feature learning and initialization design.

Abstract

We develop a mathematically rigorous framework for multilayer neural networks in the mean field regime. As the network's widths increase, the network's learning trajectory is shown to be well captured by a meaningful and dynamically nonlinear limit (the \textit{mean field} limit), which is characterized by a system of ODEs. Our framework applies to a broad range of network architectures, learning dynamics and network initializations. Central to the framework is the new idea of a \textit{neuronal embedding}, which comprises of a non-evolving probability space that allows to embed neural networks of arbitrary widths. Using our framework, we prove several properties of large-width multilayer neural networks. Firstly we show that independent and identically distributed initializations cause strong degeneracy effects on the network's learning trajectory when the network's depth is at least four. Secondly we obtain several global convergence guarantees for feedforward multilayer networks under a number of different setups. These include two-layer and three-layer networks with independent and identically distributed initializations, and multilayer networks of arbitrary depths with a special type of correlated initializations that is motivated by the new concept of \textit{bidirectional diversity}. Unlike previous works that rely on convexity, our results admit non-convex losses and hinge on a certain universal approximation property, which is a distinctive feature of infinite-width neural networks and is shown to hold throughout the training process. Aside from being the first known results for global convergence of multilayer networks in the mean field regime, they demonstrate flexibility of our framework and incorporate several new ideas and insights that depart from the conventional convex optimization wisdom.

Paper Structure

This paper contains 70 sections, 45 theorems, 596 equations.

Table of Contents

  1. Introduction
  2. Two-layer MF network: a brief overview
  3. Multilayer MF network: the challenge and our contributions
  4. Notations
  5. A General Framework
  6. Multilayer neural network and generalized learning dynamics
  7. Mean field limit
  8. Preliminaries
  9. Existence and Uniqueness of the Solution of the MF ODEs
  10. Main Result: Connection between Neural Network and MF Limit
  11. Neuronal Embedding and the Coupling Procedure
  12. Main Theorem
  13. Proof of Theorem \ref{['thm:gradient descent coupling']}
  14. Proof of Theorems \ref{['thm:particle coupling']} and \ref{['thm:gradient descent']}
  15. Simplifications under Independent and Identically Distributed Initialization
  16. ...and 55 more sections

Key Result

Theorem 1

Under suitable regularity conditions, for any finite constant $T$, as $n\to\infty$, $\mu_{t}^{n}\to\mu_{t}$ weakly and uniformly over $t\in\left[0,T\right]$.

Theorems & Definitions (92)

  • Theorem 1: Two-layer MF network with distributional representation, $n\to\infty$, informal and simplified
  • Theorem 2: Two-layer MF network with distributional representation, $t\to\infty$, informal and simplified
  • Example 3: Fully-connected networks
  • Example 4: Convolutional networks
  • Remark 5
  • Remark 6
  • Theorem 7
  • Lemma 8
  • Lemma 9
  • Lemma 10
  • ...and 82 more