Mean Field Analysis of Deep Neural Networks

TL;DR

This paper provides a rigorous mathematical analysis of multi-layer neural networks in the large-size and long-training limit, revealing their asymptotic behavior and convergence properties.

Contribution

It introduces a novel limit analysis for deep neural networks, characterizing their behavior via deterministic equations and proving convergence to global minima.

Findings

Limit neural network behavior described by integro-differential equations

Convergence to global minimum under certain conditions

Applicable to networks with any number of hidden layers

Abstract

We analyze multi-layer neural networks in the asymptotic regime of simultaneously (A) large network sizes and (B) large numbers of stochastic gradient descent training iterations. We rigorously establish the limiting behavior of the multi-layer neural network output. The limit procedure is valid for any number of hidden layers and it naturally also describes the limiting behavior of the training loss. The ideas that we explore are to (a) take the limits of each hidden layer sequentially and (b) characterize the evolution of parameters in terms of their initialization. The limit satisfies a system of deterministic integro-differential equations. The proof uses methods from weak convergence and stochastic analysis. We show that, under suitable assumptions on the activation functions and the behavior for large times, the limit neural network recovers a global minimum (with zero loss for the objective function).

Equations

$$g_{\theta }^{N_1,N_2}(x)=\frac{1}{N_2}\sum _{i=1}^{N_2}{C}^i\sigma \left(\frac{1}{N_1}\sum _{j=1}^{N_1}{W}^{2,i,j}\sigma \left(W^{1,j}\cdot x\right)\right).$$

$$H^{1,j}(x)$$

$$Z^{2,i}(x)$$

$$H^{2,i}(x)$$

$$g_{\theta }^{N_1,N_2}(x)$$

$$\theta =(C^1,\dots ,C^{N_2},W^{2,1,1}\dots ,W^{2,N_1,N_2},W^{1,1},\dots W^{1,N_1}),$$

$$\displaystyle L^{N_{1},N_{2}}(\theta)=\frac{1}{2}\mathbb{E}_{Y,X}\bigg{[}(Y-g^{N_{1},N_{2}}_{\theta}(X))^{2}\bigg{]},$$

$$C_{k+1}^i$$

$$W_{k+1}^{1,j}$$

$$W_{k+1}^{2,i,j}$$

$$H_k^{1,i}(x_k)$$

$$Z_k^{2,i}(x_k)$$

$$H_k^{2,i}(x_k)$$

$$g_{{\theta }_k}^{N_1,N_2}(x_k)$$

$${\alpha }_C^{N_1,N_2}$$

$${\tilde{\gamma }}_k^{N_1,N_2}:=\frac{1}{N_1}\sum _{j=1}^{N_1}{\delta }_{W_k^{1,j},W_k^{2,1,j},\dots ,W_k^{2,N_2,j},C_k^1,\dots ,C_k^{N_2}}.$$

$$\displaystyle g_{\theta_{k}}^{N_{1},N_{2}}(x)=\frac{1}{N_{2}}\sum_{i=1}^{N_{2}}\left\langle c_{i},\tilde{\gamma}^{N_{1},N_{2}}_{k}\right\rangle\sigma\bigg{(}\left\langle w^{2,i}\sigma(w^{1}\cdot x),\tilde{\gamma}^{N_{1},N_{2}}_{k}\right\rangle\bigg{)}.$$

$${\gamma }_t^{N_1,N_2}:={\tilde{\gamma }}_{\lfloor N_{1}t\rfloor }^{N_1,N_2},$$

$$\displaystyle g_{t}^{N_{1},N_{2}}(x)\vcentcolon=g_{\theta_{\left\lfloor N_{1}t\right\rfloor}}^{N_{1},N_{2}}(x)=\frac{1}{N_{2}}\sum_{i=1}^{N_{2}}\left\langle c_{i},\gamma^{N_{1},N_{2}}_{t}\right\rangle\sigma\bigg{(}\left\langle w^{2,i}\sigma(w^{1}\cdot x),\gamma^{N_{1},N_{2}}_{t}\right\rangle\bigg{)}.$$

$$⟨f,{\gamma }_t^{N_2}⟩-⟨f,{\gamma }_0^{N_2}⟩$$

$$\displaystyle+\int_{0}^{t}\int_{\mathcal{X}\times\mathcal{Y}}\big{(}y-g_{s}^{N_{2}}(x)\big{)}\left\langle\sigma(w^{1}\cdot x)(\sigma^{\prime}(Z_{s}(x))\odot c)\cdot\nabla_{w^{2}}f,\gamma^{N_{2}}_{s}\right\rangle\pi(dx,dy)ds$$

$$\displaystyle+\int_{0}^{t}\int_{\mathcal{X}\times\mathcal{Y}}\big{(}y-g_{s}^{N_{2}}(x)\big{)}\frac{1}{N_{2}}\sum_{i=1}^{N_{2}}\left\langle c_{i}\sigma^{\prime}(Z_{s}^{i,N_{2}}(x))w^{2,i}\sigma^{\prime}(w^{1}\cdot x)x\cdot\nabla_{w^{1}}f,\gamma^{N_{2}}_{s}\right\rangle\pi(dx,dy)ds,$$

$$Z_s^{i,N_2}(x)$$

$$H_s^{i,N_2}(x)$$

$$g_s^{N_2}(x)$$

$${\gamma }_0^{N_2}(d{w}^1,d{w}^2,dc)$$

$$\underset{N_1\to \infty }{\lim }{g}_t^{N_1,N_2}(x)=g_t^{N_2}(x),$$

$$\underset{N_2\to \infty }{\lim }\underset{N_1\to \infty }{\lim }{g}_t^{N_1,N_2}(x)=g_t(x),$$

$$g_t(x)$$

$$d{\tilde{C}}_t^c$$

$$d{\tilde{W}}_t^{1,w}$$

$$d{\tilde{W}}_t^{2,c,w,u}$$

$${\tilde{H}}_t^{1,w}(x)$$

$${\tilde{Z}}_t^c(x)$$

$${\tilde{H}}_t^{2,c}(x)$$

$$V_t^w(x)$$

$$\displaystyle\sup_{x\in\mathcal{X}}\mathbb{E}\bigg{[}|g_{t}^{N_{2}}(x)-g_{t}(x)|\bigg{]}\leq KN_{2}^{-1/2},$$

$$g_t(x)$$

$${\int }_{\mathcal{X}}h(x)\sigma (w\cdot x)\pi (dx)=0\text{ for all }w\in {\mathbb{R}}^d,$$

$$d{\tilde{C}}_t^c$$

Full text

1 Introduction

Machine learning, and in particular deep learning, has achieved immense success, revolutionizing fields such as image, text, and speech recognition. It is also increasingly being used in engineering, medicine, and finance. However, despite their success in practice, there is currently limited mathematical understanding of deep neural networks. This has motivated recent mathematical research on multi-layer learning models such as [41], [42], [43], [22], [23], [44], [51], [52], [45], and [50].

Neural networks are nonlinear statistical models whose parameters are estimated from data using stochastic gradient descent (SGD) methods. Deep learning uses neural networks with many layers (i.e., “deep” neural networks), which produces a highly flexible, powerful and effective model in practice. Typically, a neural network with multiple layers between the input and the output layer is called a “deep” neural network, see for example [26]. We analyze multi-layer neural networks that have a fixed number of layers between the input and output layer, and where the number of hidden units in each layer becomes large.

Applications of deep learning include image recognition (see [37] and [26]), facial recognition [61], driverless cars [8], speech recognition (see [37], [6], [38], and [62]), and text recognition (see [64] and [59]). Neural networks have also been applied in engineering, robotics, medicine, and finance (see [39], [40], [60], [28], [49], [5], [53], [54], [55], and [56]).

In this paper we characterize multi-layer neural networks in the asymptotic regime of large network sizes and large numbers of stochastic gradient descent iterations. We rigorously prove the limit of the neural network output as the number of hidden units increases to infinity. The proof relies upon weak convergence analysis for stochastic processes. The result can be considered a “law of large numbers” for the neural network’s output when both the network size and the number of stochastic gradient descent steps grow to infinity. We show that the neural network output in the large hidden-units and large SGD-iterates limit depends on paths of representative weights that go from input to output layer. This result is then used to show that, under suitable assumptions, the limit neural network seeks to minimize the limit objective function and achieve zero loss.

Recently, law of large numbers and central limit theorems have been established for neural networks with a single hidden layer [12, 32, 45, 50, 51, 52]. For a single hidden layer, one can directly study the weak convergence of the empirical measure of the parameters. However, in a neural network with multiple layers, there is a closure problem when studying the empirical measure of the parameters (which is explained in Section 4.3). Consequently, the law of large numbers for a multi-layer network is not a straightforward extension of the single-layer network result and the analysis involves unique challenges which require new approaches. In this paper we establish the limiting behavior of the output of the neural network.

To illustrate the idea, we consider a multi-layer neural network with two hidden layers:

[TABLE]

As we will see in Section 4.2, the limit procedure can be extended to neural networks with three layers and subsequently to neural networks with any fixed number of hidden layers.

Notice now that (1.1) can be also written as

[TABLE]

where $C^{i},W^{2,i,j}\in\mathbb{R}$ and $x,W^{1,j}\in\mathbb{R}^{d}$. The neural network model has parameters

[TABLE]

which must be estimated from data. The number of hidden units in the first layer is $N_{1}$ and the number of hidden units in the second layer is $N_{2}$. The multi-layer neural network (1.2) includes a normalization factor of $\frac{1}{N_{1}}$ in the first hidden layer and $\frac{1}{N_{2}}$ in the second hidden layer.

The loss function that we focus on in this paper is the mean squared error

[TABLE]

where the data $(X,Y)\sim\pi(dx,dy)$. The goal is to estimate a set of parameters $\theta$ which minimizes the objective function (1.3).

The literature frequently refers to minimizing the mean squared error loss as regression. Our results also hold for a more general class of error functions. In particular, we could have also considered the error function $\mathbb{E}_{Y,X}\Psi(Y-g^{N_{1},N_{2}}_{\theta}(X))$ for a function $\Psi$ that is smooth, convex and satisfies $\displaystyle\min_{x\in\mathbb{R}}\Psi(x)=\Psi(0)=0$. However, for the purposes of simplicity, we will focus on the standard regression task (1.3).

The stochastic gradient descent (SGD) algorithm for estimating the parameters $\theta$ is, for $k\in\mathbb{N}$,

[TABLE]

where $\alpha_{C}^{N_{1},N_{2}}$, $\alpha_{W,1}^{N_{1},N_{2}}$, and $\alpha_{W,2}^{N_{1},N_{2}}$ are the learning rates. The learning rates may depend upon $N_{1}$ and $N_{2}$. The parameters at step $k$ are $\theta_{k}=(C^{1}_{k},\ldots,C^{N_{2}}_{k},W^{2,1,1}_{k}\ldots,W^{2,N_{1},N_{2}}_{k},W^{1,1}_{k},\ldots W^{1,N_{1}}_{k})$. $(x_{k},y_{k})$ are samples of the random variables $(X,Y)$.

The goal of this paper is to characterize the limit of an appropriate rescaling of the multi-layer neural network output $g^{N_{1},N_{2}}_{\theta_{k}}(x)$ as both the number of hidden units $(N_{1},N_{2})$ and the stochastic gradient descent iterates $k$ become large. This is the topic of Theorem 2.3. The idea is to first take $N_{1}\rightarrow\infty$ with $N_{2}$ fixed. In Lemma 2.2, we prove that the empirical measure of the parameters converges to a limit measure as $N_{1}\rightarrow\infty$ (with $N_{2}$ fixed) which satisfies a measure evolution equation. This naturally implies a limit for the neural network output $g^{N_{1},N_{2}}$ as $N_{1}\rightarrow\infty$. The next step is to take $N_{2}\rightarrow\infty$. Theorem 2.3 proves that the limiting distribution can be represented via a system of ODEs.

As previously discussed, related limiting results for the single-layer neural network case have been investigated in [12, 32, 45, 50, 51, 52]. In those papers, it is proven that as the number of hidden units and stochastic gradient descent steps, in the appropriate scaling, diverge to infinity, the empirical distribution of the neural network parameters converges to the weak solution of a non-local PDE. This non-local PDE turns out to be a gradient flow for the limiting objective function in the space of probability measures endowed with the Wassenstein metric (this result is analogous to our Theorem 3.4). More results with non-asymptotic bounds, again for the single-layer neural network, can be found in [46].

Let us also mention that after our current paper appeared as a preprint on arXiv, [4] studied the limit of a multi-layer neural network, but under some important differences as compared to our paper. [47] also derives asymptotics for multi-layer neural networks. In [4], the weights in the first and last hidden layers are held fixed throughout training, and the number of hidden units in the rest of the layers is sent to infinity. In our paper, we train all parameters in all layers of the neural network, which introduces additional technical challenges. There is also the related papers [11, 33], whose authors look again at the limit as the units per layer go to infinity, but under a $\frac{1}{\sqrt{N}}$ normalization instead of our $\frac{1}{N}$ normalization. In the $\frac{1}{\sqrt{N}}$ normalization, a completely different limit equation will appear. The $\frac{1}{\sqrt{N}}$ case results in a perturbation around the network’s randomized initialization, leading to a kernel-type regression limit.

We address the multi-layer neural network case in the mean field scaling of $\frac{1}{N}$. In particular, we study the behavior of the neural network’s output (a) as the number of hidden units in each layer go to infinity one by one, and (b) as the number of stochastic gradient descent iterates (i.e., during training of the network) goes to infinity at the speed of the number of the hidden units of the first layer. In this asymptotic regime, we are able to obtain a well-defined limit in Theorem 2.3. Consequences of this result along with a simulation study are presented in Sections 3 and 4.

The rest of the paper is organized as follows. Our main result, which characterizes the asymptotic behavior of a neural network with two hidden layers when the number of hidden units becomes large, is presented in Section 2. The result can be easily extended to an arbitrary number of hidden layers. Section 3 is devoted to the global convergence arguments. In particular, we show that under the proper assumption the limiting problem derived in Section 2 seeks to minimize the limiting objective function, recovering the global minimum. Section 4 discusses further the theoretical results, includes a numerical study to showcase some of the theoretical implications, and, as an example, presents the limit for a three-layer neural network. The proof of the convergence theorem is in Section 5. The uniqueness of a solution to the limiting system is established in Section 6. The proof of the limit of the first layer (i.e., the proof of Lemma 2.2) and a few other results are provided in the Appendix. Section 7 has concluding remarks and possible directions for future work.

2 Main Results

Let us start by presenting our assumptions, which will hold throughout the paper. We shall work on a filtered probability space $(\Omega,\mathcal{F},\mathbb{P})$ on which all the random variables are defined. The probability space is equipped with a filtration $\mathfrak{F}_{t}$ that is right continuous and $\mathfrak{F}_{0}$ contains all $\mathbb{P}$-negligible sets.

Assumption 2.1**.**

We assume the following conditions throughout the paper.

• •

$\sigma(\cdot)\in C^{2}_{b}$, i.e., it is twice continuously differentiable and bounded.

• •

The distribution $\pi(dx,dy)$ has compact support, i.e. the data $(x_{k},y_{k})$ takes values in the compact set $\mathcal{X}\times\mathcal{Y}$.

• •

The random initialization of the parameters, i.e. $\{C^{i}_{\circ}\}_{i},\{W^{2,i,j}_{\circ}\}_{i,j},\{W_{\circ}^{1,j}\}_{j}$, are i.i.d. and take values in compact sets $\mathcal{C},\mathcal{W}^{1}$, and $\mathcal{W}^{2}$.

• •

The probability distributions of the initial parameters $(C^{i}_{\circ},W^{2,i,j}_{\circ},W^{1,j}_{\circ})_{i,j}$ admit continuous probability density functions.

We denote by $\mu_{c}(dc)$, $\mu_{W^{2}}(du)$, and $\mu_{W^{1}}(dw)$ the probability distributions of $\{C^{i}_{\circ}\}_{i}$, $\{W^{2,i,j}_{\circ}\}_{i,j}$, and $\{W_{\circ}^{1,j}\}_{j}$ respectively.

For reasons that will become clearer later on, we shall choose the learning rates to be

[TABLE]

Note that the weights in the second layer are trained faster than the other parameters. This choice of learning rates is necessary for convergence to a non-trivial limit as $N_{1},N_{2}\rightarrow\infty$. If the parameters in all the layers are trained with the same learning rate, it can be mathematically shown that the network will not train as $N_{1},N_{2}$ become large. We further explore this interesting fact in Section 4.1.

Define the empirical measure

[TABLE]

If $f$ is an appropriate test function on some space $X$ and $\gamma$ is a finite measure on $X$, then we denote the inner product $\left\langle f,\gamma\right\rangle=\int_{X}f(x)\gamma(dx)$. Using this inner product, the neural network’s output can be re-written in terms of the empirical measure:

[TABLE]

Let us next define the time-scaled empirical measure

[TABLE]

and the corresponding time-scaled neural network output is

[TABLE]

At any time $t$, $\gamma^{N_{1},N_{2}}_{t}$ is measure-valued. The scaled empirical measure $(\gamma^{N_{1},N_{2}}_{t})_{0\leq t\leq 1}$ is a random element of $D_{E}([0,1])$333$D_{S}([0,1])$ is the set of maps from $[0,1]$ into $S$ which are right-continuous and which have left-hand limits. with $E=\mathcal{M}(\mathbb{R}^{d+2N_{2}})$.

We study convergence using iterated limits. We first let $N_{1}\rightarrow\infty$ where the number of units in the first layer is $N_{1}$ and the number of stochastic gradient descent steps is $\left\lfloor N_{1}\right\rfloor$. Then, we let the number of units in the second layer $N_{2}\rightarrow\infty$.

We begin by letting the number of hidden units in the first layer $N_{1}\rightarrow\infty$.

Lemma 2.2**.**

The process $\gamma^{N_{1},N_{2}}\vcentcolon=(\gamma_{t}^{N_{1},N_{2}})_{0\leq t\leq 1}$ converges in distribution to the measure valued process $\gamma^{N_{2}}$ that takes values in $D_{E}([0,1])$ as $N_{1}\rightarrow\infty$. For every $f\in C^{2}_{b}(\mathbb{R}^{d+2N_{2}})$, $\gamma^{N_{2}}$ satisfies the measure evolution equation

[TABLE]

where

[TABLE]

Proof.

The proof of this lemma is related to the limit in the first layer as the number of hidden units in the first layer grows with the number of hidden units in the second layer is held fixed. The proof is analogous to the proof in [51] and the details are presented for completeness in the Appendix A. ∎

Lemma 2.2 studies the limit of the empirical measure $\gamma_{t}^{N_{1},N_{2}}$ as $N_{1}\rightarrow\infty$ with $N_{2}$ fixed. The limit is characterized by the stochastic evolution equation (2.3)-(2.4). Notice that Lemma 2.2 immediately implies that

[TABLE]

in probability, as $N_{1}\rightarrow\infty$

The next step is to study the limit as $N_{2}\rightarrow\infty$. To do so, we study the limit of the random ODE as $N_{2}\rightarrow\infty$ whose law is characterized by (2.3)-(2.4). Our main goal is the characterization of the limit neural network output $g_{t}^{N_{1},N_{2}}(x)$. The following convergence result characterizes the neural network output $g_{t}^{N_{1},N_{2}}(x)$ for large $N_{1}$ and $N_{2}$.

Theorem 2.3**.**

For any $t\in[0,1]$ and $x\in\mathcal{X}$,

[TABLE]

in probability444$\displaystyle\lim_{N_{2}\rightarrow\infty}\lim_{N_{1}\rightarrow\infty}X^{N_{1},N_{2}}=X$ in probability if, for all $\epsilon>0$, \displaystyle\lim_{N_{2}\rightarrow\infty}\lim_{N_{1}\rightarrow\infty}\mathbb{P}\bigg{[}\left\lVert X^{N_{1},N_{2}}-X\right\rVert>\epsilon\bigg{]}=0., where we have that

[TABLE]

with

[TABLE]

The system in (2.6) has a unique solution. In addition, letting $g_{t}^{N_{2}}(x)$ defined through Lemma 2.2 we have the following rate of convergence

[TABLE]

for some constant $K<\infty$.

Notice that we can also write that $g_{t}(x)$ satisfies

[TABLE]

The proof of Theorem 2.3 is given in Section 5. Theorem 2.3 indicates that the neural network output in the large hidden units and large SGD-iterates limit depends on paths of representative weights that connect the input layer to the output layer. Even though, we restricted the statement of Theorem 2.3 in the interval $t\in[0,1]$, its proof makes it clear that the statement is true for $t\in[0,T]$ for any $0<T<\infty$. In addition, even though this does not mean that we can take $T=\infty$, we can still examine what happens to the limit problem as $t$ grows. In particular, in Section 3 we show that under the proper assumptions, one does expect to recover the global minimum as $t\rightarrow\infty$.

Section 4 discusses some further consequences of Theorem 2.3 as well as challenges that come up in the study of the limiting behavior of multi-layer neural networks as the number of the hidden units grows.

3 On global convergence

The goal of this section is to demonstrate that, under appropriate assumptions, it can be expected that the global minimum is recovered as $t\rightarrow\infty$. Let the target data $y$ be a function $f(x)$ that we seek to learn, i.e. $y=f(x)$.

Assumption 3.1**.**

The activation function $\sigma(\cdot)$ is real analytic, bounded, and $\sigma^{\prime}(\cdot)>0$. Let $\mathcal{X}=\{1\}\times D$ where $D\subset\mathbb{R}^{d-1}$, which is equivalent to including bias weights in the first layer of the neural network. $\pi$ is positive on sets in $D$ with positive Lebesgue measure and its first marginal is a Dirac measure at $x=1$.

Notice that by [31], the fact that $\sigma$ is bounded and non-constant by Assumption 3.1 implies that $\sigma$ is also discriminatory in the sense of [13, 31]. Namely, if we have that

[TABLE]

then $h(x)=0$ for $x\in\mathcal{X}$ .

Remark 3.2*.*

An example of a real analytic, bounded activation function where $\sigma^{\prime}(\cdot)>0$ (and therefore discriminatory too), i.e. that satisfies Assumption 3.1, is the sigmoid function $\sigma(z)=\frac{e^{z}}{1+e^{z}}$.

Theorem 2 of [25] shows that for activation functions $\sigma$ that are non-constant, bounded and monotone (in which case Assumption 3.1 holds), then two layer neural networks, such as $g_{\theta}^{N_{1},N_{2}}(x)$, (and more generally multilayer neural networks) are dense in $\mathcal{C}(\mathcal{X})$. This result basically implies that for large enough $N_{1},N_{2}$ functions of the form $g_{\theta}^{N_{1},N_{2}}(x)$ approximate to arbitrary accuracy functions in $\mathcal{C}(\mathcal{X})$. Training the weights with stochastic grading descent and sending $N_{1},N_{2}$ to infinity, leads then, by Theorem 2.3, to function approximator $g_{t}(x)$, which is given through a system of integro-differential equations of the form

[TABLE]

where the neural network prediction is

[TABLE]

We establish in this section that under reasonable assumptions $f(x)$ can indeed be recovered by $g_{t}(x)$ as $t\rightarrow\infty$.

Let us denote $\Theta_{t}(c,w,u)=(\tilde{C}^{c}_{t},\tilde{W}^{1,w}_{t},\tilde{W}^{2,c,w,u}_{t})$ for the components of the ODE in (3.1). For notational convenience, we shall often write $\theta=(c,w,u)$. Then, we obviously have that $g_{t}(x)$ depends on $t$ only through $[\Theta_{t}]=\left\{\Theta_{t}(\theta)\right\}_{\theta\in\mathcal{C}\times\mathcal{W}^{1}\times\mathcal{W}^{2}}$. In order to emphasize that we shall write $g_{t}(x)=g(x;[\Theta_{t}])$.

Analogously, we will denote $[\Theta_{t}(c,\cdot)]=\left\{\Theta_{t}(c,w,u)\right\}_{(w,u)\in\mathcal{W}^{1}\times\mathcal{W}^{2}}$ and $[\Theta_{t}(\cdot,w,\cdot)]=\left\{\Theta_{t}(c,w,u)\right\}_{(c,u)\in\mathcal{C}\times\mathcal{W}^{2}}$, which then leads to the notation $\tilde{Z}^{c}_{t}(x)=\tilde{Z}(x;[\Theta_{t}(c,\cdot,\cdot)])$ and $V^{w}_{t}(x)=V(x;[\Theta_{t}(\cdot,w,\cdot)])$. For notational convenience let us also denote $h(x;[\Theta_{t}])=(f(x)-g_{t}(x))$.

With these definitions in place, and for $(c,w)\in\mathcal{C}\times\mathcal{W}^{1}$, let us also define

[TABLE]

and let us set

[TABLE]

The purpose of the notation above is to make it clear that the functions $H(\cdot)$ depends on $[\Theta_{t}]$, on $\Theta_{t}(\theta)$ and on $\theta$ separately. With these identifications we can write that the ODE system in (3.1) can be written in the form

[TABLE]

Next we investigate the limiting objective function

[TABLE]

where $g(x;[\Theta_{t}])=g_{t}(x)$ is given by (3.2).

Notice that the function $\Theta\mapsto\bar{L}(\Theta)$ is always non-negative and becomes zero only at a global minimum. In particular, our goal is to demonstrate that as $t\rightarrow\infty$, there are sufficient conditions that guarantee that the global minimum is realized in the sense that

[TABLE]

We notice that $\bar{L}$ acts as a Lyapunov function for the dynamical system (3.4) in that

[TABLE]

Indeed, we calculate

[TABLE]

[TABLE]

[TABLE]

The fact that $\frac{d}{dt}\bar{L}(\Theta_{t})\leq 0$ from (3.5) means that $\bar{L}(\Theta_{t})$ is decreasing in the gradient direction of the paths governing the limiting behavior of the weights.

Let us define $\eta_{0}$ to be the joint measure for the random initialization of the parameters, i.e. for $\{C^{i}_{\circ}\}_{i}$, $\{W_{\circ}^{1,j}\}_{j}$, $\{W^{2,i,j}_{\circ}\}_{i,j}$. By Assumption 2.1 we have that this is the product measure $\eta_{0}=\mu_{c}\times\mu_{W^{1}}\times\mu_{W^{2}}$. Let us make a convenient assumption for the support of $\eta_{0}$. In particular,

Assumption 3.3**.**

We have that $\textrm{support}(\eta_{0})=\mathcal{C}\times\mathcal{W}^{1}\times\mathcal{W}^{2}$.

Notice now that we can write

[TABLE]

As a matter of fact, more is true. We investigate how the objective function changes with perturbation from $\Theta$ to $\hat{\Theta}$. In particular, we have for $\epsilon=(\epsilon_{1},\epsilon_{2},\epsilon_{3})$

[TABLE]

which, in combination with (3.6) and Assumption 3.3, then says that a minimizer for $\bar{L}(\Theta)$, say $\Theta^{*}(\theta)$, will satisfy for almost all $(c,w)\in\mathcal{C}\times\mathcal{W}^{1}$ the relations

[TABLE]

By analogy, let us now define $\eta_{t}$ to be the measure at time $t$ of the random vector $\Theta_{t}=(\tilde{C}_{t},\tilde{W}^{1}_{t},\tilde{W}^{2}_{t})$ as governed by the solution to the random ODE system (3.1). Then, $\eta_{t}$ is the pushforward of $\eta_{0}$ under $\Theta_{t}$ given by (3.4) (see for example Chapter 8 in [3]), i.e.,

[TABLE]

Theorem 3.4**.**

Let assumptions 3.1 and 3.3 hold. If $\eta_{t}\rightarrow\eta^{*}$, weakly, where $\eta^{*}$ is a non-degenerate measure that admits a density with finite first moments, then we have that $\eta^{\ast}$ is a global minimum with zero loss.

Before proving Theorem 3.4 we discuss its assumptions in the remarks that follow.

Remark 3.5*.*

It is interesting to note that Theorem 3.4 actually shows that any stationary point is a global minimum with zero loss. In addition, even though we do not show this here, one can see from the proof of Theorem 3.4 that the assumption on analyticity of the activation function can be replaced by an assumption on $\eta^{*}$ having full support.

Remark 3.6*.*

The assumption on convergence of $\eta_{t}$ as $t\rightarrow\infty$ is a required assumption for convergence in the limit as $t\rightarrow\infty$ to occur. For completeness we mention here that verifying this assumption for the (commonly used in practice) unregularized case that we study here is still an open problem even for the corresponding unregularized case in the single layer case, see Appendix C.5 in [12] or [45]. Even though we do not show this here, one way to guarantee that such an assumption holds is to add an appropriate regularization term in the loss function, or to add an appropriate non-degenerate stochastic perturbation in the training algorithm (1.4) for the evolution of each one of $C^{i}_{k+1},W^{1,j}_{k+1},W^{2,i,j}_{k+1}$ (noisy SGD); see [32, 50] for the corresponding arguments in the single layer case. See also [10] for conditions under which the assumption on convergence of $\eta_{t}$ as $t\rightarrow\infty$ holds for non-degenerate sparse problems that include an appropriate regularizing term in the loss function.

Remark 3.7*.*

It is clear that convergence of $[\Theta_{t}]$ as $t\rightarrow\infty$ is linked to the convergence of the measure $\eta_{t}$ as $t\rightarrow\infty$. At this point, we also mention that convergence of the trajectories $[\Theta_{t}]$ as $t\rightarrow\infty$ is linked to the study of Łojasiewicz inequalities [9, 29]. To the best of our knowledge no general results in this direction are currently known neither for our case of interest, nor for problems corresponding to the single-layer case (see Appendix C.5 of [12]) (i.e. no general results currently exist for the non-geodesically convex case that we are dealing with here).

Proof of Theorem 3.4.

We will use the notation $\tilde{\mu}_{c}(c)$, $\tilde{\mu}_{W^{1}}(w)$, and $\tilde{\mu}_{W^{2}}(w)$ for the densities of the measure $\mu_{c}(dc)$, $\mu_{W^{1}}(dw)$, and $\mu_{W^{2}}(dw)$ respectively. From equations (3.7), we have that any stationary solution must satisfy

[TABLE]

Equation (3.9) and Assumption 3.3 imply that

[TABLE]

for almost every $w\in\mathcal{W}^{1}$ and $c\in\mathcal{C}$. Recalling now that

[TABLE]

we obtain

[TABLE]

where $\tilde{\eta}^{\ast}_{W^{2}}(u)$ denotes the marginal density of the limiting measure $\eta^{*}$ with respect to $\tilde{W}^{2}$. Analogously we shall define the marginal densities $\tilde{\eta}^{\ast}_{c}(c)$ and $\tilde{\eta}^{\ast}_{W^{1}}(w)$ with respect to the variables $\mathcal{C}$ and $\tilde{W}^{1}$ respectively. Similar to before, $\tilde{\eta}^{\ast}(c,w,u)$ is the density of the measure $\eta^{\ast}(dc,dw,du)$.

We must therefore have that $\displaystyle\sup_{x\in\mathcal{X}}|Z^{\ast,c}(x)|\tilde{\mu}_{c}(c)<\infty$ for almost every $c\in\mathcal{C}$. In addition, since $\eta^{\ast}$ has a density, we have that there exists a set $B=\{c\in\mathbb{R}:|c-c_{0}|<\delta\}$ with $c_{0}\neq 0$ and $0<\delta<|c_{0}|$ such that

[TABLE]

Therefore, there must be a set $K\subset\mathbb{R}$ with Lesbegue measure $\lambda(K)>0$ such that, for almost every $c\in K$,

[TABLE]

which in turn implies that $C^{\ast,c}\neq 0$ and $\displaystyle\sup_{x\in\mathcal{X}}|Z^{\ast,c}(x)|<\infty$ for almost every $c\in K$. Due to equations (3.11) and (3.10), for almost every $(c,w)\in K\times\mathcal{W}^{1}$,

[TABLE]

Thus, there exists a $c_{0}\in K$ such that

[TABLE]

for almost every $w\in\mathcal{W}^{1}$. This of course implies that, for almost every $w\in\mathcal{W}^{1}$,

[TABLE]

Consequently, we have that

[TABLE]

Therefore, since $\eta^{\ast}_{W^{1}}$ has a density, there exists a compact set $A\subset\mathbb{R}^{d}$ such that, for every $w\in A$,

[TABLE]

Since $\sigma(\cdot)$ is analytic, $\Gamma(w)$ is analytic (see Lemma 3.8). Therefore, by the identity theorem for real-analytic functions, $\Gamma(w)=0$ for every $w\in\mathbb{R}^{d}$. It then follows, since $\sigma(\cdot)$ is discriminatory,

[TABLE]

for every $x\in\mathcal{X}$. Finally, since $\sigma^{\prime}(\tilde{Z}^{\ast,c_{0}}(x))>0$ due to equation (3.12),

[TABLE]

for every $x\in\mathcal{X}$, concluding the proof of the theorem. ∎

We conclude with the the result on analyticity of $\Gamma(w)$ defined in (3.13).

Lemma 3.8**.**

Assume $\mathcal{X}$ is a compact set and the activation unit $\sigma(\cdot)$ is real analytic and bounded. Then,

[TABLE]

defined in (3.13) is a real analytic function as well.

Proof of Lemma 3.8.

The partial derivative of $\sigma(w\cdot x)$ is

[TABLE]

Before we prove analyticity of $\Gamma(w)$, we show that $g(x)$ is finite. Recall that

[TABLE]

Therefore, due to $\sigma(\cdot)$ being bounded and $\eta^{\ast}$ having marginals with finite first moments,

[TABLE]

Next,

[TABLE]

Due to the compactness of $\mathcal{X}$, $\displaystyle\sup_{x\in\mathcal{X}}|\sigma^{\prime}(\tilde{Z}^{\ast,c_{0}}(x))|<\infty$, and the fact that $\sigma$ is a real analytic function, Proposition 2.2.10 of [34] shows there exist constants $C_{1},C_{2}>0$ such that

[TABLE]

where $R_{2}=\frac{1}{C_{2}}$ and $R=R_{1}R_{2}$. It follows then again from Proposition 2.2.10 of [34] that $\Gamma(w)$ is a real analytic function. This concludes the proof of the lemma. ∎

4 Discussion on the limiting results and extensions to multi-layer networks with greater depth

In Subsection 4.1, we discuss some of the implications of our theoretical convergence results and presents related numerical results. In Section 4.2, we show that the procedure can be extended to treat deep neural networks with more than two hidden layers. General challenges in the study of multi-layer neural networks are explored Subsection 4.3.

4.1 Discussion on the limiting results

It is instructive to notice that the results of this paper recover the results of [51] (see also [45, 50]) if we restrict attention to the one-layer case. Indeed, let us set $N_{2}=1$, $N_{1}=N$, $C^{i}=1$ and $H^{2,i}=Z^{2,i}$ in (1.1)-(1.2), and we get the single-layer neural network

[TABLE]

with the corresponding empirical measure of the parameters becoming

[TABLE]

In that case notice that we can simply write

[TABLE]

Then, it is relatively straightforward to notice that the result of Lemma 2.2 boils down to the one layer convergence results of [51], see also [45, 50]. Namely, if we write $\gamma_{t}$ for the limit in probability of $\gamma^{N}_{t}$ we get that

[TABLE]

It is useful to compare the limits of the neural network output in the one layer and two layer cases, (4.2) and (2.7) respectively.

Some general remarks of interest follow.

• •

It is clear that the two layer case (and more generally the multi-layer case) is more complex than the single-layer case, which provides some intuition for the increased complexity of deep neural networks when compared to shallow neural networks. For instance, in the single layer case the neural network output is given explicitly by (4.2). On the other hand, in the multi-layer case, the neural network’s output is the solution to an integral equation as given by (2.6)-(2.7).

• •

In contrast to the single-layer case, in the multi-layer case the asymptotic weight distribution is layer dependent and weights do not necessarily become independent in the large hidden units limit. In fact, as our result demonstrates, see (2.5)-(2.6), the neural network output in the large hidden units limit and large SGD-iterates does not depend on the individual weights, but on paths of weights that connect the input layer to the output layer. One could regard these paths, i.e. the solution to the random ODE’s in (2.6), as typical representative paths of the weights in the limit connecting the different layers of the “limiting” neural network.

• •

In the multi-layer case, units at a given layer receive an aggregate signal from units of the previous and/or of the next layer.

• •

The law of large numbers for a single-layer network indicates that the network will converge in probability to a deterministic limit. That is, after a certain point, adding more hidden units will not increase the accuracy. Our main result Theorem 2.3 suggests that the same conclusion is true for multi-layer neural networks as well.

• •

Theorem 3.4 combined with Theorem 2.3 characterizes the limiting behavior of the objective function $L^{N_{1},N_{2}}(\theta)$ from (1.3). Section 3 shows that, under the proper assumptions, the limit objective function decreases in the gradient direction of the paths governing the limiting behavior of the weights. The limit ODEs seek to minimize the limit objective function and the global minimum is expected to be recovered.

One practical consequence of our analysis is that the parametrization of the learning rates, see (2.1), indicates that one should use larger learning rates for the weights that connect the different layers (e.g., $W^{2,i,j}$) as compared to the weights for the input or output layers (e.g., $C^{i}$ and $W^{1,j}$). Notice that this is also the case for the three layer case outlined in Subsection 4.2 below.

As will be explained in Section 4.2, the law of large numbers can be extended to multi-layer neural networks with an arbitrary number of layers. The law of large numbers will only hold under a certain choice of the learning rates. The learning rates need to be scaled with the number of hidden units in each layer. For a multi-layer network with $L$ layers, the learning rates are

[TABLE]

where $N_{\ell}$ is the number of hidden units in the $\ell$-th layer.

If the learning rates are constant in the number of hidden units $N_{1},\ldots,N_{L}$, it turns out that the network will not train as $N_{1},\ldots,N_{L}\rightarrow\infty$ (i.e., in the limit, the network parameters will remain at their initial conditions).

The necessity of scaling the learning rates in the asymptotic regime of large numbers of hidden units (i.e., wide layers) is one of the interesting products of the mean-field limit analysis. A numerical example is presented in Figure 1 below. A deep neural network is trained to classify images for the CIFAR10 dataset [35]. The CIFAR10 dataset contains $60,000$ color images in $10$ classes (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, and truck). The dataset is divided into $50,000$ training images and $10,000$ test images. Each image has $32\times 32\times 3$ pixels. The goal is to train a neural network to correctly classify each image based solely on the image pixels as an input. The neural network we use has the mean-field normalizations $\frac{1}{N_{\ell}}$ in each layer $\ell$. There are $8$ convolution layers which each have $64$ channels. This is followed by two fully-connected layers which each have $500$ units. We first train the neural network using the scaled learning rates. Then, we also train the neural network with the standard stochastic gradient descent algorithm (no scaling of the learning rates). Using the scaled learning rates, we achieve a high test accuracy. However, without the scalings, the neural network does not train (i.e., it remains at a very low accuracy).

4.2 Extension to multi-layer neural networks with more layers

The procedure developed in this paper naturally extends to multi-layer neural networks with more layers than two layers. For brevity, let us present the result in the case of three layers. The situation for more layers is the same, albeit with more complicated algebra. A multi-layer neural network with three layers takes the form

[TABLE]

which can be also written as

[TABLE]

where $C^{i},W^{2,j,\nu},W^{3,i,j}\in\mathbb{R}$ and $x,W^{1,\nu}\in\mathbb{R}^{d}$. The neural network model has parameters

[TABLE]

which must be estimated from data. The number of hidden units in the first layer is $N_{1}$, the number of hidden units in the second layer is $N_{2}$, and the number of hidden units in the third layer is $N_{3}$. Naturally, the loss function now becomes

[TABLE]

where the data $(X,Y)\sim\pi(dx,dy)$.

The stochastic gradient descent (SGD) algorithm for estimating the parameters $\theta$ is, for $k\in\mathbb{N}$, $\nu=1,\ldots,N_{1}$, $i=1,\ldots,N_{3}$ and $j=1,\ldots,N_{2}$ is

[TABLE]

where

[TABLE]

where $\alpha_{C}^{N_{1},N_{2},N_{3}}$, $\alpha_{W,1}^{N_{1},N_{2},N_{3}}$, $\alpha_{W,2}^{N_{1},N_{2},N_{3}}$ and $\alpha_{W,3}^{N_{1},N_{2},N_{3}}$ are the learning rates. The parameters at step $k$ are

[TABLE]

$(x_{k},y_{k})$ are samples of the random variables $(X,Y)$. We assume a condition analogous to Assumption 2.1.

Let us now choose the learning rates to be

[TABLE]

Similar to before, define the empirical measure

[TABLE]

The time-scaled empirical measure is

[TABLE]

and the corresponding time-scaled neural network output is $g_{t}^{N_{1},N_{2},N_{3}}(x)=g_{\theta_{\left\lfloor N_{1}t\right\rfloor}}^{N_{1},N_{2},N_{3}}(x)$.

Following the same procedure as for the two layer case one expects to get the following limit that describes $g_{t}(x)$ (the iterated limit of $g_{t}^{N_{1},N_{2},N_{3}}(x)$ as first $N_{1}\rightarrow\infty$, then $N_{2}\rightarrow\infty$, and then $N_{3}\rightarrow\infty$):

[TABLE]

In other words, one expects to be able to write that the neural network’s output is

[TABLE]

A computation analogous to the one in Section 3 shows that, as expected, the limit ODEs in (4.8) seek to minimize the corresponding limit objective function as $t\rightarrow\infty$. We leave the rigorous proof for the form of $g_{t}(x)$ in the three layer neural network case to the interested reader.

4.3 Challenges in the analysis of multi-layer neural networks

Challenges arise in the study of the asymptotics of multi-layer neural networks. [51], [45], and [50] leveraged traditional approaches from the mean-field and interacting particle system literature to characterize the asymptotics of single-layer neural networks. However, it turns out that the traditional mean-field approach cannot be used for multi-layer neural networks.

A standard approach for analyzing (1.4) as $N_{1},N_{2}\rightarrow\infty$ would be to construct an empirical measure $\rho_{k}^{N_{1},N_{2}}$ of the parameters $\theta_{k}$ at training step $k$, as for example $\tilde{\gamma}_{k}^{N_{1},N_{2}}$. Then, we could study the behavior of $\rho_{k}^{N_{1},N_{2}}$ as $N_{1},N_{2}\rightarrow\infty$. This empirical measure $\rho_{k}^{N_{1},N_{2}}$ needs to be designed such that the dynamics of $\rho_{k}^{N_{1},N_{2}}$ can be written in terms of $\rho_{k}^{N_{1},N_{2}}$ itself and the data $(x_{k},y_{k})$ (plus martingale and remainder terms). That is, the dynamics of $\rho_{k}^{N_{1},N_{2}}$ are closed.

This is straightforward for single-layer neural networks (see [51]), but it is challenging to do for multi-layer neural networks. In the case of single-layer networks, the empirical measure is simply given by (4.1) and its analysis has been successfully carried on in [51]. One is tempted to do the same thing for the multi-layer case, i.e. study the limit of the empirical measure defined in (2.2). The problem that one faces with this formulation is that, in contrast to the single-layer case, the dimension of the space on which the empirical measure takes values increases with $N_{2}$. Therefore, the problem cannot be studied using such a empirical measure.

An alternative approach, which is also natural in this case, is to try and create “nested measures”, sometimes called multi-level measure valued processes in mathematical biology, see [14, 16, 17, 18, 19, 20, 21] and the review paper [15].

Let us explain how to construct a nested empirical measure for (1.4) and why it will not work. In order to simplify notation, let $N_{1}=N_{2}=N$ and $\rho_{k}^{N_{1},N_{2}}=\rho_{k}^{N}$ in the following example. Considering $N_{1}\neq N_{2}$ and taking subsequent limits does not alter the conclusions below.

• •

Let’s first examine the parameter $W^{1,j}$ in the first layer. The $j$-th unit in the first layer (i.e., $H^{1,j}=\sigma(W^{1,j}\cdot x)$) is connected to all of the hidden units in the second layer (i.e., $H^{2,i}$) via the weights $W^{2,i,j}$. Therefore, there is a measure associated with each $W^{1,j}$ which must track $\{W^{2,i,j},Z^{i},C^{i}\}_{i=1}^{N}$. $W^{2,i,j}$ and $Z^{i}$ are required for calculating the SGD update for $W^{1,j}$ (see 1.4). This measure is $\nu^{N,j}_{k}=\frac{1}{N}\sum_{i=1}^{N}\delta_{W^{2,i,j}_{k},Z^{i}_{k},C^{i}_{k}}$.

• •

Let’s next examine the parameter $C^{i}$ in the second layer. The $i$-th unit in the second layer is connected to all of the hidden units in the first layer via the weights $W^{2,i,j}$. Therefore, for each $C^{i}$, we must track $\{W^{2,i,j},W^{1,j}\}_{j=1}^{N}$ in order to calculate the SGD update for $C^{i}$ (see 1.4). Furthermore, updating $C^{i}$ requires tracking $W^{1,j}$, and updating $W^{1,j}$ requires tracking $\nu^{j}$. Therefore, updates to $C^{i}$ require the empirical measure $\mu^{N,i}_{k}=\frac{1}{N}\sum_{j=1}^{N}\delta_{W^{2,i,j}_{k},W^{1,j}_{k},\nu^{N,j}_{k}}\in\mathcal{M}(\mathbb{R}\times\mathbb{R}^{d}\times\mathcal{M}(\mathbb{R}^{3}))$.

• •

Finally, the entire network at iteration $k$ is specified by the empirical measure

[TABLE]

where $\mathcal{M}(E)$ is the space of measures on the metric space $E$. Notice that the process (4.9) involves nested measures (sometimes called “multi-level processes”). The process $\rho_{k}$ takes values in a space of nested measures $\mathcal{M}(\mathcal{M}(\mathcal{M}(\cdots)))$.

Careful inspection of $\rho_{k}^{N}$ identifies a crucial problem: its dynamics are not closed. The evolution of $\nu^{N,j}_{k}$ (the innermost measure in the nested measures) cannot be written in terms of $\rho_{k}^{N}$. In particular, updating $\nu^{N,j}_{k}$ requires also updating $(W^{2,i,j},Z^{i}_{k},C^{i}_{k})_{i=1}^{N}$. This would in fact require knowledge of $(W^{2,i,j},Z^{i}_{k},C^{i}_{k},\mu^{N,i}_{k})_{i=1}^{N}$, i.e. we would have to re-define $\nu^{N,j}_{k}$ as $\frac{1}{N}\sum_{i=1}^{N}\delta_{W^{2,i,j}_{k},Z^{i}_{k},C^{i}_{k},\eta^{N,i}_{k}}$ where $\eta^{N,i}_{k}=\frac{1}{N}\sum_{m=1}^{N}\delta_{W^{2,i,m}_{k},W^{m}_{k}}$. This leads to re-defining $\rho_{k}^{N}$ as taking values in $\mathcal{M}(\mathbb{R}\times\mathcal{M}(\mathbb{R}\times\mathbb{R}^{d}\times\mathcal{M}(\mathbb{R}^{3}\times\mathcal{M}(\mathbb{R}\times\mathbb{R}^{d}))))$, i.e. the space of $4$-nested measures (before it was $3$-nested measures). However, the closure problem remains, since the evolution of $\eta^{N,i}_{k}$ cannot be written in terms $\rho_{k}^{N}$. In fact, there does not seem to exist any finite number of nested measures for which the empirical measure $\rho_{k}^{N}$’s dynamics are closed and despite our best efforts we have not managed to find one. This closure problem then leads to non-trivial issues associated with establishing a well define limit.

For completeness, we also remark that a third alternative is to define the candidate empirical measure as an appropriately normalized double sum over both indices corresponding to the two hidden layers and then consider the limit of this measure as $N_{1},N_{2}\rightarrow\infty$. However, this approach also suffers a closure problem, similar to the situation described above.

The discussion in this section highlights some of the challenges that must be addressed in the analysis of multi-layer neural networks. Such problems are not present in the analysis of single-layer neural networks. Therefore, a different approach is required for the asymptotic analysis of multi-layer neural networks and this paper is a first step in this direction.

The problems that we describe above led us to the approach used in this paper. In particular, the approach in Section 2 first studies the limit of the empirical measure as the number of hidden units in the first layer grows to infinity. This is similar to [51]. The limit is a solution to an evolution equation and it is the law of a system of random ODEs. We then make the crucial observation that one can characterize the resulting system in terms of the initialization for the stochastic gradient descent iterates. This means that we can reformulate the limiting system of the first layer into an equivalent system of random ODEs and then consider the limit of the second layer. This allows us to obtain the limit of the output of the neural network as the number of hidden units of all layers grow to infinity by studying the limit of the random ODE in Theorem 2.3.

5 Proof of Theorem 2.3-Characterization of the limit

In preparation for the proof of Theorem 2.3, we first re-express the result from Lemma 2.2.

Corollary 5.1**.**

Consider the particle system:

[TABLE]

Let $\nu_{t,(c_{1},\ldots,c_{N_{2}})}$ be the conditional Law of $(W^{1}_{t},W^{2,1}_{t},\ldots W^{2,N_{2}}_{t},C_{t}^{1},\ldots,C_{t}^{N_{2}})_{0\leq t\leq T}$ given $(C^{1}_{0},\ldots,C^{N_{2}}_{0})=(c_{1},\ldots,c_{N_{2}})$. Then, $\nu_{t,(C_{\circ}^{1},\ldots,C_{\circ}^{N_{2}})}$ is the unique solution to the evolution equation (2.3)-(2.4).

Proof.

See Appendix A. ∎

Due to exchangeability properties and without loss of generality, we can transform (5.1) into the equivalent particle system:

[TABLE]

Since for all $i=1,\cdots,N_{2}$, $C^{i}_{\circ}$ have probability density function as described in Assumption 2.1, we have that

[TABLE]

This allows us to substitute the variable names

[TABLE]

for $i=1,\ldots,N_{2}$.

This produces the system:

[TABLE]

The system (5.3) is exactly the same system as (5.1). Notice also that $\hat{C}_{t}^{i}$ in (5.3) depends also on $\{C_{\circ}^{i}\}_{i=1}^{N_{2}}$ in a symmetric way via $g_{t}^{N_{2}}(x)$. Similarly $\hat{W}_{t}^{1}$ and $\hat{W}_{t}^{2,i}$ depend also on $\{C_{\circ}^{i}\}_{i=1}^{N_{2}}$ and on $\{W_{0}^{2,i}\}_{i=1}^{N_{2}}$ symmetrically via $g_{t}^{N_{2}}(x)$ and $V_{t}^{N_{2},W_{0}}(x)$. This is also the situation for (5.1). Then independence and identical distribution of the initial conditions together with the aforementioned exchangeability property imply that (5.1) and (5.2) are equivalent.

5.1 Limiting System

The goal is to prove that for any $t\in[0,1]$ and $x\in\mathcal{X}$,

[TABLE]

in $L^{1}$, where $g_{t}^{N_{2}}(x)$ is from (5.2) and the limit $g_{t}(x)$ is given by

[TABLE]

where,

[TABLE]

Before presenting the proof of this result, let us define a quantity that will be of central interest in the sequel. In particular, for $(c,w,u)\in\{(C_{\circ}^{i},W^{1}_{0},W^{2,i}_{0}),i=1,\cdots N_{2}\}$ and $x\in\mathcal{X}$, let’s define the error function

[TABLE]

Note that we certainly have,

[TABLE]

We will show below that $\mathbb{E}\left[E_{t}^{N_{2}}(C^{i}_{\circ},W_{0},W_{0}^{2,i},x)\right]$ appropriately converges to zero as $N_{2}\rightarrow\infty$ which will then imply that $g_{t}^{N_{2}}(x)$ converges to $g_{t}(x)$ as indicated.

5.2 A Priori Bounds

Let us first establish uniform bounds on the processes $C_{t}^{c},W_{t}^{1,w},W_{t}^{2,c,w,u},$ and $g_{t}^{N_{2}}(x)$ for the system (5.2). For any $t\in[0,T]$ and any $N_{2}\in\mathbb{N}$,

[TABLE]

$|H_{s}^{2,C^{i}_{\circ}}(x)|<K$ since $\sigma(\cdot)\in C_{b}^{2}$. Then, using the fact that $\mathcal{X}\times\mathcal{Y}$ is compact,

[TABLE]

By Gronwall’s inequality, we have that for any $N_{2}\in\mathbb{N}$ and for any $t\in[0,T]$,

[TABLE]

Using the same approach, we can establish uniform bounds on the other processes $W_{t}^{1,w}$ and $W_{t}^{2,c,w,u}$. Therefore, for any $N_{2}\in\mathbb{N}$, $t\in[0,1]$, $x\in\mathcal{X}$, and $(c,w,u)\in\{(C_{\circ}^{i},W^{1}_{0},W^{2,i}_{0}),i=1,\cdots N_{2}\}$,

[TABLE]

5.3 Bound for $\mathbb{E}|V_{t}^{N_{2},W_{0}}(x)-V_{t}^{W_{0}}(x)|^{2}$.

We have:

[TABLE]

The first term in (5.6) can be studied using a decomposition:

[TABLE]

where the uniform bounds from (5.5) were used. In addition, we also have for some constant $K<\infty$

[TABLE]

where we used the assumed independence of $C^{i}_{\circ}$, $W_{0}$, and $W_{0}^{2,i}$ as well as the a-priori bound from (5.5).

Hence, we obtain that for some unimportant constant $K<\infty$

[TABLE]

5.4 Bound for $|g_{t}^{N_{2}}(x)-g_{t}(x)|$.

We can write

[TABLE]

Let’s analyze the first term in (5.8). Using the uniform bounds from (5.5) we have for some unimportant constant $K<\infty$

[TABLE]

The second term in (5.8) is bounded, as follows,

[TABLE]

where the independence of $C^{i}_{\circ}$ was used. Hence, putting things together we get

[TABLE]

5.5 Bound for \mathbb{E}\bigg{[}(C_{t}^{C_{\circ}^{i}}-\tilde{C}_{t}^{C_{\circ}^{i}})^{2}+(W_{t}^{1,W_{0}}-\tilde{W}_{t}^{1,W_{0}})^{2}+(W_{t}^{2,C^{i}_{\circ},W_{0},W_{0}^{2,i}}-\tilde{W}_{t}^{2,C^{i}_{\circ},W_{0},W_{0}^{2,i}})^{2}\bigg{]}.

Let us write for notational convenience $c={C_{\circ}^{i}}$. We have

[TABLE]

Using Young’s inequality and the uniform bounds from (5.5),

[TABLE]

For the term $\displaystyle|C_{s}^{c}-\tilde{C}_{s}^{c}|\Gamma_{s}^{g,1}(x)$, where we recall $c=C_{\circ}^{i}$, we have

[TABLE]

Therefore, using (5.9) and (LABEL:Cbound0011A) we obtain

[TABLE]

Using similar arguments, we can also show, using (5.7), that

[TABLE]

Therefore, we overall get that

[TABLE]

5.6 Bound for \mathbb{E}\bigg{[}(Z_{t}^{C_{\circ}^{i}}(x)-\tilde{Z}_{t}^{C_{\circ}^{i}}(x))^{2}\bigg{]}.

We next consider $(Z_{t}^{C_{\circ}^{i}}(x)-\tilde{Z}_{t}^{C_{\circ}^{i}}(x))^{2}$. For the following calculations, we define $\mathcal{F}_{C}^{N_{2}}=(C^{1}_{\circ},C_{\circ}^{2},\ldots,C^{N_{2}}_{\circ})$.

For $i=1,2,\ldots,N_{2}$, we have

[TABLE]

where the uniform bounds from (5.5) were used together with the compactness of the state space assumption from Assumption 2.1.

Therefore, using iterated expectations, we have that

[TABLE]

where (5.11) was used.

Finally, using the assumption that $\sigma(\cdot)$ is globally Lipschitz,

[TABLE]

for possibly different, but finite constants $K_{1},K_{2},K_{3}<\infty$.

5.7 Bound for $\mathbb{E}\left[E_{t}^{N_{2}}(C^{i}_{\circ},W_{0},W_{0}^{2,i},x)\right]$

Collecting our results from (5.12), (5.13) and (5.14), together with the definition of the error function $E_{t}^{N_{2}}$, we have, for $i=1,\cdots,N_{2}$, the bound

[TABLE]

Therefore, by averaging over all $i=1,\cdots,N_{2}$ and then using Gronwall’s inequality, we get for any $0\leq t\leq T$,

[TABLE]

for an appropriate constant $K<\infty$. Combining the last two displays we naturally get, again using Gronwall’s inequality, that for $i=1,\cdots,N_{2}$

[TABLE]

where the constant $K<\infty$ does not depend on $i$ or $N_{2}$.

5.8 Convergence of neural network prediction

The bound (5.15) of course proves the (uniform) convergence in probability of the neural network output $g_{t}^{N_{2}}(x)\rightarrow g_{t}(x)$. Recall, by (5.8), that

[TABLE]

where

[TABLE]

Then, using the Cauchy-Schwartz inequality and (5.15),

[TABLE]

Therefore, for $0\leq t\leq T$, and for some unimportant constant $K<\infty$

[TABLE]

concluding the identification of the limit in Theorem 2.3.

6 Proof of Theorem 2.3-Uniqueness of the limit

Lemma 6.1**.**

The solution to the limiting system (2.6) is unique in $C([0,1],\mathcal{C}\times\mathcal{W}^{1}\times\mathcal{W}^{2}\times\mathcal{X})$.

Proof.

Suppose there are two solutions to (2.6). Let’s denote the first solution as $(\hat{W}_{t}^{1,w},\hat{W}^{2,c,w,u}_{t},\hat{C}_{t}^{c},\hat{Z}_{t}^{c}(x),\hat{V}_{t}^{w}(x),\hat{g}_{t}(x))$ and the second solution as $(\bar{W}_{t}^{1,w},\bar{W}^{2,c,w,u}_{t},\bar{C}_{t}^{c},\bar{Z}_{t}^{c}(x),\bar{V}_{t}^{w}(x),\bar{g}_{t}(x))$. Define the function

[TABLE]

Note that

[TABLE]

We next study the evolution of $Q_{t}$ for $t>0$.

Using the same approach as in Section 5.2, we can show that any $\tilde{W}^{2,c,w,u}_{t},\tilde{C}_{t}^{c},\tilde{W}_{t}^{1,w},\tilde{H}_{t}^{1,w}(x),\tilde{Z}_{t}^{c},V_{t}^{w},$ and $g_{t}(x)$ which solve (2.6) are uniformly bounded on $\mathcal{C}\times\mathcal{W}^{1}\times\mathcal{W}^{2}\times\mathcal{X}\times[0,1]$.

We can then prove the inequality

[TABLE]

where we have also used Young’s inequality and the fact that $\mathcal{X}\times\mathcal{Y}$ is compact. Therefore,

[TABLE]

Similarly,

[TABLE]

Next, using the Cauchy-Schwarz inequality and Young’s inequality,

[TABLE]

Using a similar approach,

[TABLE]

Therefore,

[TABLE]

Finally,

[TABLE]

Consequently,

[TABLE]

Collecting our results,

[TABLE]

Therefore, by Gronwall’s inequality,

[TABLE]

for all $t\in[0,1]$, completing the proof of the lemma.

∎

7 Conclusions and future work

In this paper, we have developed an approach that allows us to obtain the limiting behavior of multi-layer neural networks in the mean-field scaling as the number of units per layer and stochastic gradient steps grow to infinity. We have demonstrated that the limit is characterized by paths of weights that connecting the input layer to the output layer. The limit procedure shows that the main hyperparameters (the learning rates) must be chosen using a specific scaling (as a function of the number of hidden units per layer) in order to get a well-defined limit. A numerical study demonstrates that the mean-field scaling of the learning rates has a significant effect on how well the network can be trained. In addition, we have shown that the limit neural network seeks to minimize the limit objective function.

There are of course a number of open questions. The discussion in Section 4.3 shows that a genuine mean field formulation in terms of convergence of empirical distributions of trained parameters has yet to be understood. Secondly, our methodology works for convex error functions like the regression task. It is of interest to characterize the limit in other cases such as Wassenstein distances or regularized entropies. Thirdly, the universal approximation theorems say that a neural network approximator in the infinite unit per layer limit can approximate any continuous function [7, 25, 30, 31]. It would be of interest to investigate whether the limit procedure developed in this paper could lead to more specific information on convergence rates. Fourthly, the, commonly used in practice, unregularized algorithm and its finer properties (such as more detailed behavior as time grows to infinity) need to be better understood and we hope that the results in this paper build towards this direction.

Appendix A Limit of First Layer, Lemma 2.2

In this appendix we prove Lemma 2.2, which is about the limit of the first layer. The proof is analogous to [51]. Hence, instead of repeating all the arguments we will only present the general outline emphasizing the differences and refer the interested reader to [51] when the arguments are the same.

We let $N_{1}\rightarrow\infty$ (with $N_{2}$ fixed). We want to prove that for each $N_{2}\in\mathbb{N}$, $\tilde{\gamma}_{\cdot}^{N_{1},N_{2}}\overset{d}{\rightarrow}\gamma_{\cdot}^{N_{2}}$ in $D_{E}([0,1])$. $\gamma_{t}^{N_{2}}$ is a random measure-valued process which, for every $f\in C^{2}_{b}(\mathbb{R}^{d+2N_{2}})$, $\gamma^{N_{2}}$, satisfies the evolution equation (2.3)-(2.4).

A.1 Relative Compactness

Let’s first establish a bound on $C_{k}^{i}$. Recall that $\sigma(\cdot)$ is bounded and $\alpha_{C}^{N_{1},N_{2}}=\frac{N_{2}}{N_{1}}$. In the following calculations, the unimportant constants, $K,K_{0},K_{1},$ and $K_{2}$ may change from line to line.

For $k=0,1,\ldots,\left\lfloor N_{1}\right\rfloor$,

[TABLE]

Since $\sigma(\cdot)$ is bounded, $|H^{2,i}_{k}|<K$. Therefore,

[TABLE]

This yields

[TABLE]

This implies that

[TABLE]

By Gronwall’s inequality, for $k\leq\left\lfloor N_{1}\right\rfloor$,

[TABLE]

Then, we also have that

[TABLE]

This immediately yields that for $k\leq\left\lfloor N_{1}\right\rfloor$

[TABLE]

Let’s now address the parameters $W_{k}^{2,i,j}$. Recall that $\alpha_{W,2}^{N_{1},N_{2}}=N_{2}$. Then,

[TABLE]

Therefore, for $k\leq\left\lfloor N_{1}\right\rfloor$,

[TABLE]

Similarly, we have

[TABLE]

Therefore, we obtain

[TABLE]

Collecting our results, for all $k/N_{1}\leq 1$ and for all $j=1,\cdots,N_{2}$, we have the uniform bound

[TABLE]

We now prove relative compactness of the family $\{\gamma^{N_{1},N_{2}}\}_{N_{1}\in\mathbb{N}}$ in $D_{E}([0,1])$ where $E=\mathcal{M}(\mathbb{R}^{d+2N_{2}})$. It is sufficient to show compact containment and regularity of the $\gamma^{N_{1},N_{2}}$’s (see for example Chapter 3 of [24]).

Lemma A.1**.**

For each $\eta>0$ and $t\geq 0$, there is a compact subset $\mathcal{K}$ of E such that

[TABLE]

Proof.

This uniform bound (A.1) actually implies the stronger statement of compact support. In particular, notice that the set $[-K,K]^{d+2N_{2}}$ is compact, and define

[TABLE]

Then $\mathcal{K}\subset M(\mathbb{R}^{d+2N_{2}})$, and $\mathbb{P}$-a.s. $\gamma_{t}^{N_{1},N_{2}}\in\mathcal{K}$ for all $N_{1}\in\mathbb{N}$ and $t\in[0,1]$. This concludes the proof. ∎

We now establish regularity of the $\gamma^{N_{1},N_{2}}$’s. Define the function $q(z_{1},z_{2})=\min\{|z_{1}-z_{2}|,1\}$ where $z_{1},z_{2}\in\mathbb{R}$.

Lemma A.2**.**

There is a constant $K<\infty$ such that for $0\leq u\leq\delta$, $0\leq v\leq\delta\wedge t$, $t\in[0,1]$,

[TABLE]

Proof.

Let $0\leq s<t<\leq T$ and $\delta<1$, such that $t-s<\delta<1$. We then have

[TABLE]

Using the same approach, we can establish similar bounds for the other parameters:

[TABLE]

The desired result then follows. ∎

We conclude this section now with the required relative compactness of the sequence $\{\gamma^{N_{1},N_{2}}\}_{N_{1}\in\mathbb{N}}$. This implies that for each fixed $N_{2}$, every subsequence $\gamma^{N_{1},N_{2}}$’s has a convergent sub-subsequence as $N_{1}\rightarrow\infty$.

Lemma A.3**.**

The sequence of probability measures $\{\gamma^{N_{1},N_{2}}\}_{N_{1}\in\mathbb{N}}$ is relatively compact in $D_{E}([0,1])$.

Proof.

Given Lemmas A.1 and A.2, Theorem 8.6 and Remark 8.7 B of Chapter 3 of [24], gives the statement of the lemma. ∎

A.2 Identification of the Limit

We consider the evolution of the empirical measure $\gamma^{N_{1},N_{2}}_{t}$ via test functions $f\in C^{2}_{b}(\mathbb{R}^{d+2N_{2}})$. Using a Taylor expansion based argument similarly to [51], we can show that the scaled empirical measure satisfies

[TABLE]

where $M^{N_{1},N_{2}}(t)$ is a martingale term,

[TABLE]

and we have defined $H_{t}^{N_{1},N_{2}}$ as the vector $(H_{t}^{1,N_{1},N_{2}},\ldots,H_{t}^{N_{2},N_{1},N_{2}})$, $c$ as the vector $(c_{1},\ldots,c_{N_{2}})$, and $Z_{t}^{N_{1},N_{2}}$ as the vector $(Z_{t}^{1,N_{1},N_{2}},\ldots,Z_{t}^{N_{2},N_{1},N_{2}})$. The martingale term $M^{N_{1},N_{2}}(t)$ converges to [math] in $L^{2}$ as $N_{1}\rightarrow\infty$. That is,

[TABLE]

The proof for (A.3) follows as in Lemma 3.1 of [51] and thus it is omitted. The limit point of $\gamma^{N_{1},N_{2}}$, as $N_{1}\rightarrow\infty$ and for a fixed $N_{2}$, will satisfy the evolution equation

[TABLE]

where

[TABLE]

and

[TABLE]

Let $\pi^{N_{1},N_{2}}$ be the probability measure of $\left(\gamma_{t}^{N_{1},N_{2}}\right)_{0\leq t\leq 1}$. Each $\pi^{N_{1},N_{2}}$ takes values in the set of probability measures \mathcal{M}\big{(}D_{E}([0,1])\big{)}. Relative compactness, proven in Section A.1, implies that there is a subsequence $\pi^{N_{1_{k}},N_{2}}$ which weakly converges. We must prove that any limit point $\pi^{N_{2}}$ of a convergent subsequence $\pi^{N_{1_{k}},N_{2}}$ will satisfy the evolution equation (A.4).

Lemma A.4**.**

Let $\pi^{N_{1_{k}},N_{2}}$ be a convergent subsequence with a limit point $\pi^{N_{2}}$. Then, $\pi^{N_{2}}$ is a Dirac measure concentrated on $\gamma^{N_{2}}\in D_{E}([0,1])$ and $\gamma^{N_{2}}$ satisfies the measure evolution equation (A.4).

Proof.

We define a map $F(\gamma):D_{E}([0,1])\rightarrow\mathbb{R}_{+}$ for each $t\in[0,T]$, $f\in C^{2}_{b}(\mathbb{R}^{d+2N_{2}})$, $g_{1},\cdots,g_{p}\in C_{b}(\mathbb{R}^{d+2N_{2}})$ and $0\leq s_{1}<\cdots<s_{p}\leq t$.

[TABLE]

Then,

[TABLE]

Therefore,

[TABLE]

Since $F(\cdot)$ is continuous and $F(\gamma^{N_{1},N_{2}})$ is uniformly bounded (due to the uniform boundedness results of Section A.1),

[TABLE]

Since this holds for each $t\in[0,T]$, $f\in C^{2}_{b}(\mathbb{R}^{d+2N_{2}})$ and $g_{1},\cdots,g_{p}\in C_{b}(\mathbb{R}^{d+2N_{2}})$, $\gamma^{N_{2}}$ satisfies the evolution equation (A.4). ∎

It remains to prove that the evolution equation (A.4) has a unique solution. This is the content of Section A.3.

A.3 Uniqueness

Lemma A.5**.**

There exists a unique solution to the evolution equation (A.4).

Proof.

We only sketch the proof as the argument is similar to the uniqueness argument of [51]. Consider the particle system:

[TABLE]

Let $\nu_{t,(c_{1},\ldots,c_{N_{2}})}$ be the conditional law of $(W_{t}^{1},W^{2,1}_{t},\ldots W^{2,N_{2}}_{t},C_{t}^{1},\ldots,C_{t}^{N_{2}})_{0\leq t\leq T}$ given $(C^{1}_{\circ},\ldots,C^{N_{2}}_{\circ})=(c_{1},\ldots,c_{N_{2}})$. Similarly to the general results of [36], $\nu_{t,(C^{1}_{\circ},\ldots,C^{N_{2}}_{\circ})}=\gamma^{N_{2}}_{t}$ is a solution to the evolution equation (A.4) and conversely, any solution $\nu_{t,(C^{1}_{\circ},\ldots,C^{N_{2}}_{\circ})}$ to the evolution equation (A.4) must also be the law of the solution to (A.6).

Let us next prove that we can write Z_{t}^{i}(x)=\mathbb{E}\bigg{[}W^{2,i}_{t}H^{1}_{t}(x)\bigg{|}C^{1}_{\circ},\ldots,C^{N_{2}}_{\circ}\bigg{]}. Recall that $\mathcal{F}_{C}^{N_{2}}=(C^{1}_{\circ},\ldots,C^{N_{2}}_{\circ})$ and let us similarly define $\mathcal{F}_{W^{2}}^{N_{2}}=(W^{2,1}_{0},\ldots,W^{2,N_{2}}_{0})$. Inspecting (A.6) it becomes clear that the random variables $(W_{t}^{1},W^{2,1}_{t},\ldots W^{2,N_{2}}_{t},C_{t}^{1},\ldots,C_{t}^{N_{2}})_{0\leq t\leq T}$ depend, in addition to their own initial conditions, also on $\mathcal{F}_{C}^{N_{2}}$ and $\mathcal{F}_{W^{2}}^{N_{2}}$ in a symmetric way through the terms $g^{N_{2}}_{t}(x)$ and $\frac{1}{N_{2}}\sum_{i=1}^{N_{2}}C_{t}^{i}\sigma^{\prime}(Z_{t}^{i}(x))W^{2,i}_{t}$. In order to make this dependence specific in the notation, we write in particular that

[TABLE]

This, then leads to the following calculations

[TABLE]

The second to the last equality is due to the assumed independence of the initial conditions via Assumption 2.1 together with the fact that the marginal of $\gamma_{0}^{N_{2}}$ with respect to $C^{1}_{\circ},\ldots,C^{N_{2}}_{\circ}$ is a product of delta Dirac distributions.

It remains to show that the solution to (A.6) is unique. We do this via a fixed point argument, completely analogous to Section 4 of [51]. The details are omitted due to the similarity of the argument.

∎

A.4 Convergence Result for First Layer

Let $\pi^{N_{1},N_{2}}$ be the probability measure corresponding to $\gamma^{N_{1},N_{2}}$. Each $\pi^{N_{1},N_{2}}$ takes values in the set of probability measures \mathcal{M}\big{(}D_{E}([0,1])\big{)}. Relative compactness, proven in Section A.1, implies that every subsequence $\pi^{N_{1_{k}},N_{2}}$ has a further sub-sequence $\pi^{N_{1_{k_{m}}},N_{2}}$ which weakly converges. Section A.2 proves that any limit point $\pi^{N_{2}}$ of $\pi^{N_{1_{k_{m}}},N_{2}}$ will satisfy the evolution equation (A.4). Section A.3 proves that the solution of the evolution equation (A.4) is unique. Therefore, by Prokhorov’s Theorem, $\pi^{N_{1},N_{2}}$ weakly converges to $\pi^{N_{2}}$, where $\pi^{N_{2}}$ is the distribution of $\gamma^{N_{2}}$, the unique solution of (A.4). That is, $\gamma^{N_{1},N_{2}}$ converges in distribution to $\gamma^{N_{2}}$.