Consider a single-layer neural network with the Xavier initialization:

where , and ) : . The number of hidden units is N and the output is scaled by a factor (the widely-used Xavier initialization, see [2]). The objective function is

where the data (, and the parameters = (. For notational convenience, we may refer to ) as ) in our analysis below.

The model parameters are trained using stochastic gradient descent. The parameter updates are given by:

for k = 0, 1, . . ., T N where T > 0. is the learning rate (which may depend upon N). The data samples are () are i.i.d. samples from a distribution ). We impose the following assumption.

• The activation function (R), i.e. is twice continuously differentiable and bounded.

• The randomly initialized parameters () are i.i.d., mean-zero random variables with a distribution ).

• The random variable has compact support and .

• The sequence of data samples () is i.i.d. from the probability distribution ).

• There is a fixed dataset of M data samples (and therefore ) = 1M

Note that the last assumption also implies that ) has compact support.

We will study the limiting behavior of the network output (x) for as the number of hidden units N and stochastic gradient descent steps T N simultaneously become large. The

network output converges in distribution to the solution of a random ODE as .

1.1 Main Results

Define the empirical measure

Note that the neural network output can be written as the inner-product

Due to Assumption 1.1, as and for ,

where is a Gaussian random variable. We also of course have that

Define the scaled processes

We will study convergence in distribution of the random process () as in the space ([0, T ]) where ([0, T ]) is the Skorokhod space and M(S) is the space of probability measures on S.

The main contribution of our work is a rigorous proof that a neural network with the Xavier initialization and trained with stochastic gradient descent converges in distribution to a random ODE as the number of units and training steps become large. In addition, our convergence analysis will also address several interesting questions:

• Our results provide a rigorous convergence guarantee for the Xavier initialization (i.e., the nor- malization factor), which is almost universally used in deep learning models. A priori it is unclear if the neural network (x) will converge as since, for k > 0, the ) is correlated with ) and therefore a limit may not exist. If a limit did not exist, this would imply that the neural network model could have poor numerical behavior for large N. We prove that a limit does exist.

• Although the pre-limit problem of optimizing a neural network is non-convex (and therefore the neural network may converge to a local minimum), the limit equation minimizes a quadratic objective function.

• We show that the matrix in the limiting quadratic objective function is positive definite, and therefore the neural network (in the limit) will converge to a global minimum with zero loss on the training set.

Convergence to a global minimum for a neural network has been recently proven in [3], [4], and [5]. Our work contributes to this growing literature by showing that convergence to a global minimum is a simple consequence of the mean-field limit for neural networks. A detailed discussion of these papers and other related literature is provided in Section 1.2.

Our main results are presented below.

Theorem 1.2. The process () converges in distribution in the space ([0, T ]) as to () which satisfies, for every ), the random ODE

Proof. See Sections 3, 4, and 5.

Recall that is a Gaussian random variable; see equation (1.2). In addition, note that ¯in the limit equation (1.3) is a constant, i.e. for [0, T ]. Therefore, (1.3) reduces to

Since (1.4) is a linear equation in ([0, T ]), the solution is unique.

To better understand (1.4), define the matrix where

where is finite-dimensional since we fixed a training set of size M in the beginning. Then, (1.4) becomes

where ˆY = (). Therefore, is the solution to a continuous-time gradient descent algorithm which minimizes a quadratic objective function.

Therefore, even though the pre-limit optimization problem is non-convex, the neural network’s limit will minimize a quadratic objective function.

An interesting question is whether ˆY as . That is, in the limit of large numbers of hidden units and many training steps, does the neural network model converge to a global minimum with zero training error. Theorem 1.3 shows that ˆY as if A is positive definite. Corollary 1.4 proves that, under reasonable hyperparameter choices and if the data samples are in distinct directions (see [1]), A will be positive definite.

Theorem 1.3. If A is positive definite, then

Proof. Consider the transformation ˜ˆY . Then,

Then, ˜0 (and consequently ˆY ) as if A is positive definite.

Corollary 1.4. Assume Assumption 1.1. A sufficient condition for A to be positive definite is ) is nonpolynomial and slowly increasing (i.e., lim= 0 for every a > 0), is positive when evaluated on sets of positive Lebesgue measure and the data samples are in distinct directions (as defined on page 192 of [1]).

Proof. See Section 6.

Examples of activation units ) satisfying the conditions in Corollary 1.4 include sigmoid functions and hyperbolic tangent functions. Using a normal distribution for the initialization of the parameters in the neural network is a common choice in practice (covered by the requirements of Corollary 1.4).

Remark 1.5. For presentation purposes we have not explicitly denoted the bias term in the model. However, it is clear that this can be handled by requiring the first component of the vector x to be equal to one for example. This would result in the neural network taking the form ) = ). We leave the rest of the details to the interested reader.

1.2 Literature Review

[11], [12], [13], and [15] study the asymptotics of single-layer neural networks with a normalization; that is, ) = ). [14] studies the asymptotics of deep (i.e., multi-layer) neural networks with a normalization in each hidden layer. In the single layer case, the limit for the neural network satisfies a partial differential equation. As discussed in [11], it is not necessarily true that the limiting equation (a PDE in this case) will converge to the global minimum of an objective function with zero training error.

The normalization studied in [11], [12], [13], and [15] is convenient since the single-layer neural network is then in a traditional mean-field framework where it can be described via an empirical measure of the parameters. However, the normalization that we study in this paper is more widely-used in practice (it is referred to as the Xavier initialization and was first introduced in [2]). The normalization requires different analysis than the standard mean-field analysis , and it produces a completely different limit. Importantly, under reasonable conditions, the limit equation converges to a global minimum with zero training error. In addition, for the limit to hold, we show that the normalization requires the effective learning rate for the parameters to be of the order .

Convergence to a global minimum for a neural network has been recently proven in [3], [4], and [5]. Although it has been long understood that neural networks have universal approximation properties (see [8], [9], and [10]), it has until recently been commonly believed that training algorithms for neural networks (e.g., gradient descent) may converge to a local minimum (and not a global minimum) since neural networks are non-convex. [3], [4], and [5] showed that neural networks (under suitable conditions) will converge to a global minimum during training. This result is quite remarkable considering the optimization problem is non-convex, and it provides an important mathematical guarantee for the field of deep learning.

[3], [4], and [5] do not study the mean-field limit of a neural network with the Xavier initialization, which is the focus of our paper. Once the mean field limit is established, we show that convergence to a global minimum is a simple consequence of the limit equation. There are also some differences between our assumptions and the assumptions required for the theorems of [3], [4], and [5]. [3] and [4] study gradient descent while our paper studies stochastic gradient descent, which introduces additional technical challenges due to the stochastic dynamics. [5] studies stochastic gradient descent for a framework where the neural network’s output layer parameters are not trained. In their paper, the parameters are randomly generated and then frozen (i.e., they do not change during training). In practice, all of the parameters in the neural network, including the output layer parameters, are trained with stochastic gradient descent and therefore it is worthwhile to consider the more general case. [5] also imposes an assumption that the loss function vanishes at infinity. [3], [4], and [5] all require that every data sample has the same magnitude, i.e.= 1 for every i = 1, . . . , M. We do not require this assumption.

[6] proved a limit for neural networks with a Xavier initialization when they are trained with continuous-time gradient descent. Our paper proves a limit for neural networks trained with the (standard) discrete-time stochastic gradient descent algorithm which is used in practice. Our method of proof is also different than the approach of [6]. Whereas [6] begins their analysis in continuous time (due to their framework being continuous-time gradient descent), our paper rigorously passes from discrete time (where the stochastic gradient descent updates evolve) to continuous time through weak convergence analysis of appropriate stochastic processes and measure-valued processes. In [6], the authors directly study the evolution of the derivatives of the output with respect to the parameters, while we address the limiting behavior of the underlying associated stochastic processes and measure-valued processes.

1.3 Organization of Paper

Section 2 derives equations describing the evolution of the pre-limit process (). Relative compactness of the family of processes () is proven in Section 3. Section 4 proves that any limit point of the process must satisfy the equation (1.3). These results are collected together in Section 5 to prove that () converges in distribution to the solution of equation (1.3). Corollary 1.4 is proven in Section 6.

We begin by analyzing evolution of the network output (x). Using a Taylor expansion,

for points and in the line segment connecting the points and . Let . Substi- tuting (1.1) into (2.1) yields

Using (2.3), we can write the evolution of for [0, T ] as

Next, we decompose the summations into a drift and martingale component.

For convenience, we define the martingale terms (the third and fourth terms in the equation above) as and , respectively. The equation for can be re-written in terms of a Riemann integral and the scaled measure , yielding

In addition, using conditional independence of the terms in the series for and as well as the bounds from Lemmas 3.2 and 3.1, we have that

We can also analyze the evolution of the empirical measure in terms of test functions (). Using a Taylor expansion, we find that

for points ( ¯), ( ˆ) and ( ˜) in the segments connecting with and with , respectively.

Substituting (1.1) into (2.5) yields

Therefore,

In this section we prove that the family of processes is relatively compact. Section 3.1 proves compact containment. Section 3.2 proves regularity. Section 3.3 combines these results to prove relative compactness.

3.1 Compact Containment

Proof. The unimportant finite constant may change from line to line. We first observe that

where the last inequality follows from the definition of (x) and the uniform boundedness assumption on ). Then, we subsequently obtain that

This implies that

Let us now define =

have that

By the discrete Gronwall lemma and using ,

where the last inequality follows from the random variables taking values in a compact set. Now, we turn to the bound for . We start with the bound (using Young’s inequality)

for a constant that may change from line to line. Taking an expectation, using Assumption 1.1, the bound (3.2), and using the fact that , we obtain

for all and all k such that , concluding the proof of the lemma.

Using the bounds from Lemma 3.1, we can now establish a bound for (x) for .

Lemma 3.2. For all , all k such that , and any ,

Proof. Recall equation (2.2), which describes the evolution of (x).

This leads to the bound

We now square both sides of the above inequality.

Then, using a telescoping series,

Taking expectations,

Taking advantage of the fact that is sampled from a fixed dataset D of M data samples,

and therefore

Recall that

Combining this bound with the bound (3.4) and using the discrete Gronwall lemma yields, for any 0 ,

Substituting this bound into equation (3.3) produces the desired bound

for any 0 .

Proof. For each L > 0, define = [. Then, we have that is a compact subset of , and for each 0 and ,

where we have used Markov’s inequality and the bounds from Lemma 3.1. We define the compact subsets of )

and we observe that

Given that lim= 0, we have that, for each 0, there exists a compact set ˆsuch that

Due to Lemma 3.2 and Markov’s inequality, we also know that, for each 0, there exists a compact set U = [such that

Therefore, for each 0, there exists a compact set ˆsuch that

3.2 Regularity

Proof. We start by noticing that a Taylor expansion gives for 0

for points ( ¯¯) and ( ˆˆ) in the segments connecting with and with , respectively.

Let’s now establish a bound on for with 0 1.

where is some unimportant constant. Then, the statement of the Lemma follows.

We next establish regularity of the process in ([0, T ]). For the purposes of the following lemma, let the function ) = minwhere and .

Lemma 3.5. For any (0, 1), there is a constant such that for 0 1, 0 , [0, T ],

Proof. Recall that

Therefore,

This yields the bound

where we have used the boundedness of ) (from Assumption 1.1) and the bounds from Lemma 3.1. Taking expectations,

Using the bounds (3.6) and (3.7),

The bound (3.7) holds for each . Therefore,

The statement of the Lemma then follows.

3.3 Combining our results to prove relative compactness

Lemma 3.6. The family of processes is relatively compact in ([0, T ]).

Proof. Combining Lemmas 3.3 and 3.4, and Theorem 8.6 of Chapter 3 of [7] proves that is relatively compact in ([0, T ]). (See also Remark 8.7 B of Chapter 3 of [7] regarding replacing supwith limin the regularity condition B of Theorem 8.6.) Similarly, combining Lemmas 3.3 and 3.5 proves that is relatively compact in ([0, T ]). Since relative compactness is equivalent to tightness, we have that the probability measures of the family

of processes are tight. Similarly, we have that the probability measures of the family of process

are tight. Therefore, is tight. Then, is also relatively compact.

Let be the probability measure of a convergent subsequence of. Each takes values in the set of probability measures ([0, T ]). Relative compactness, proven in Section 3, implies that there is a subsequence which weakly converges. We must prove that any limit point of a convergent subsequence will satisfy the evolution equation (1.3).

Lemma 4.1. Let be a convergent subsequence with a limit point . Then, is a Dirac measure concentrated on (([0, T ]) and () satisfies equation (1.3).

Proof. We define a map ) : ([0, T ]) for each [0, T ], (), ), and 0 .

Then, using equations (2.4) and (2.6), we obtain

Since ) is continuous and ) is uniformly bounded (due to the uniform boundedness results of Section 3),

Since this holds for each [0, T ], () and ), () satisfies the evolution equation (1.3).

We now combine the previous results of Sections 3 and 4 to prove Theorem 1.2. Let be the probability measure corresponding to (). Each takes values in the set of probability measures ([0, T ]). Relative compactness, proven in Section 3, implies that every subsequence has a further sub-sequence which weakly converges. Section 4 proves that any limit point of will satisfy the evolution equation (1.3). Equation (1.3) is a finite-dimensional, linear equation and therefore has a unique solution. Therefore, by Prokhorov’s Theorem, weakly converges to , where is the distribution of (), the unique solution of (1.3). That is, () converges in distribution to ().

This section proves that under reasonable hyperparameter choices, the matrix A in the limit equation will be positive definite.

Proof. We first show that A is equivalent to the covariance matrix of the random variables U = ,

which are defined as

To prove that A is positive definite, we need to show that 0 for every non-zero .

The functions ) are linearly independent since the are in district directions (see Remark3.1 of [1]). Therefore, for each non-zero z, there exists a point such that

for some 0 such that for (

Then,

and A is positive definite, concluding the proof of the Corollary.

[1] Yoshifusa Ito. Nonlinearity creates linear independence. Advances in Computational Mathematics, 5: 189-203, 1996.

[2] X. Glorot and Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. 2010.

[3] S. Du, J. Lee, H. Li, L. Wang, and X. Zhai. Gradient Descent Finds Global Minima of Deep Neural Net- works. Proceedings of the 36th International Conference on Machine Learning, Long Beach, California, PMLR 97, 2019.

[4] S. Du, X. Zhai, B. Poczos, and A. Singh. Gradient Descent Provably Optimizes Over-Parameterized Neural Networks. ICLR, 2019.

[5] D. Zou, Y. Cao, D. Zhou, and Q. Gu. Stochastic Gradient Descent Optimizes Over-parameterized Deep ReLU Networks. arXiv: 1811.08888, 2018.

[6] A. Jacot, F. Gabriel, and C. Hongler. Neural Tangent Kernel: Convergence and Generalization in Neural Networks. 32nd Conference on Neural Information Processing Systems (NeurIPS 2018), Montreal, Canada.

[7] S. Ethier and T. Kurtz. Markov Processes: Characterization and Convergence. 1986, Wiley, New York, MR0838085.

[8] K. Hornik, M. Stinchcombe, and H. White. Multilayer feedforward networks are universal approximators. Neural Networks, 2(5), 359-366, 1989.

[9] K. Hornik. Approximation capabilities of multilayer feedforward networks. Neural Networks, 4(2), 251-257, 1991.

[10] C. Kuan and K. Hornik. Convergence of learning algorithms with constant learning rates. IEEE Transactions on Neural Networks, 2(5), 484-489, 1991.

[11] S. Mei, A. Montanari, and P. Nguyen. A mean field view of the landscape of two-layer neural networks Proceedings of the National Academy of Sciences, 115 (33) E7665-E767, 2018.

[12] J. Sirignano and K. Spiliopoulos. Mean Field Analysis of Neural Networks: A Law of Large Numbers. SIAM Journal on Applied Mathematics, Vol. 80, Issue 2, pp. 725–752, 2020.

[13] J. Sirignano and K. Spiliopoulos. Mean Field Analysis of Neural Networks: A Central Limit Theorem. Stochastic Processes and their Applications, Volume 130, Issue 3, March 2020, pp. 1820-1852, 2020.

[14] J. Sirignano and K. Spiliopoulos. Mean Field Analysis of Deep Neural Networks. Mathematics of Operations Research, 2021, to appear.

[15] G. M. Rotskoff and E. Vanden-Eijnden. Neural Networks as Interacting Particle Systems: Asymptotic Convexity of the Loss Landscape and Universal Scaling of the Approximation Error. arXiv:1805.00915, 2018.