Global Convergence of Three-layer Neural Networks in the Mean Field Regime
Authors: Huy Tuan Pham
Phan-Minh Nguyen
Venue: ICLR 2021 Oral (2021)
Keywords: deep learning theory

Abstract: In the mean field regime, neural networks are appropriately scaled so that as the width tends to infinity, the learning dynamics tends to a nonlinear and nontrivial dynamical limit, known as the mean field limit. This lends a way to study large-width neural networks via analyzing the mean field limit. Recent works have successfully applied such analysis to two-layer networks and provided global convergence guarantees. The extension to multilayer ones however has been a highly challenging puzzle, and little is known about the optimization efficiency in the mean field regime when there are more than two layers.

In this work, we prove a global convergence result for unregularized feedforward three-layer networks in the mean field regime. We first develop a rigorous framework to establish the mean field limit of three-layer networks under stochastic gradient descent training. To that end, we propose the idea of a neuronal embedding, which comprises of a fixed probability space that encapsulates neural networks of arbitrary sizes. The identified mean field limit is then used to prove a global convergence guarantee under suitable regularity and convergence mode assumptions, which – unlike previous works on two-layer networks – does not rely critically on convexity. Underlying the result is a universal approximation property, natural of neural networks, which importantly is shown to hold at any finite training time (not necessarily at convergence) via an algebraic topology argument.

Published as a conference paper at ICLR 2021
GLOBAL CONVERGENCE OF THREE-LAYER NEURAL NETWORKS IN THE MEAN FIELD REGIME
Huy Tuan Pham and Phan-Minh Nguyen
ABSTRACT
In the mean field regime, neural networks are appropriately scaled so that the width tends to infinity, the learning dynamics tends to a nonlinear and nontriv- ial dynamical limit, known the mean field limit. This lends way to study large-width neural networks via analyzing the mean field limit. Recent works have successfully applied such analysis two-layer networks and provided global con- vergence guarantees. The extension to multilayer ones however has been a highly challenging puzzle, and little is known about the optimization efficiency in the mean field regime when there are more than two layers.
In this work, we prove global convergence result for unregularized feedforward three-layer networks in the mean field regime. We first develop a rigorous frame- work to establish the mean field limit of three-layer networks under stochastic gra- dient descent training. To that end, we propose the idea of a neuronal embedding. which comprises of fixed probability space that encapsulates neural networks of arbitrary sizes. The identified mean field limit is then used to prove a global con- vergence guarantee under suitable regularity and convergence mode assumptions, which unlike previous works on two-layer networks does not rely critically on convexity. Underlying the result a universal approximation property, natural of neural networks, which importantly is shown to hold at any finite training time (not necessarily at convergence) via an algebraic topology argument.
INTRODUCTION
Interests in the theoretical understanding of the training of neural networks have led to the recent discovery of new operating regime: the neural network and its learning rates are scaled appro- priately, such that as the width tends to infinity, the network admits a limiting learning dynamics in which all parameters evolve nonlinearly with time!. This is known as the mean field (MF) limit (Mei et al. (2018) 8); Chizat & Bach (2018 8); Rotskoff & Vanden-Eijnden (2018); Sirignano & Spiliopoulos (2018); Nguyen (2019); Araújo et al. (2019); Sirignano & Spiliopoulos (2019)). The four works Mei et al. (2018 8); Chizat & Bach (2018 8); Rotskoff & Vanden-Eijnden (2018) 8); Sirignano & Spiliopoulos (2018) led the first wave of efforts in 2018 and analyzed two-layer neural networks. They estab- lished a connection between the network under training and its MF limit. They then used the MF limit to prove that two-layer networks could be trained to find (near) global optima using variants of gradient descent, despite non-convexity (Mei et al. (2018) 2); Chizat & Bach (2018)). The MF limit identified by these works assumes the form of gradient flows in the measure space, which factors the invariance from the action of symmetry group on the model. Interestingly, by lifting to the measure space, with a convex loss function (e.g. squared loss), one obtains limiting optimiza- tion problem that is convex (Bengio et al. (2006); Bach (2017)). The analyses of Mei et (2018);
"This paper is conference submission. We refer to the work Nguyen & Pham (2020) and its companion note Pham & Nguyen (2020) for generalizations well other conditions for global convergence in the case of multilayer neural networks.
Department of Mathematies, Stanford University. This work was done in parts while H. T. Pham was the University of Cambridge.
The Volcon Group. This work was done while P.M. Nguyen was at Stanford University.
Thee author ordering randomized.
'This is to be contrasted with another major operating regime (the NTK regime) where parameters essen- tially do not evolve and the model behaves like kernel method (Jacot al. (2018); Chizat et al. (2019); Du al. (2019); Allen-Zhu et al. (2019); Zou et al. (2018); Lee al. (2019)).

---

Published as a conference paper at ICLR 2021
Chizat & Bach (2018) utilize convexity, although the mechanisms to attain global convergence in these works are more sophisticated than the usual convex optimization setup in Euclidean spaces.
The extension to multilayer networks has enjoyed much less progresses. The works Nguyen (2019); Araújo et al. (2019); Sirignano & Spiliopoulos (2019) argued, heuristically or rigorously, for the existence of MF limiting behavior under gradient descent training with different assumptions. In fact, it has been argued that the difficulty is not simply technical, but rather conceptual (Nguyen (2019)): for instance, the presence of intermediate layers exhibits multiple symmetry groups with intertwined actions on the model. Convergence to the global optimum of the model under gradient- based optimization has not been established when there are more than two layers.
In this work, we prove a global convergence guarantee for feedforward three-layer networks trained with unregularized stochastic gradient descent (SGD) in the MF regime. After an introduction of the three-layer setup and its MF limit in Section 2, our development proceeds in two main steps:
Step (Theorem in Section 3): We first develop a rigorous framework that describes the MF limit and establishes its connection with large-width SGD- three-layer network. Here propose the new idea of a neuronal embeddingg which comprises of an appropriate non-evolving probability space that encapsulates neural networks of arbitrary sizes. This probability space is in general abstract and is constructed according to the (not necessarily i.i.d.) initialization scheme of the neural network. This idea addresses directly the intertwined action of multiple symmetry groups, which is the aforementioned conceptual obstacle (Nguyen (2019)), thereby covering setups that cannot be handled by formulations Araújo et al. (2019); Sirignano & Spiliopoulos (2019) (see also Section 5 for a comparison). Our analysis follows the technique from Sznitman (1991); Mei et al. particular, the MF limit yields good approximation
(2018) and gives a quantitative statement: of the neural network as long as "min log mamx < 1 independent of the data dimension, where Ammi and max are the minimum and maximum of the widths.
Step 2 (Theorem 8 in Section 4): We prove that the MF limit, given by our framework, con- verges to the global optimum under suitable regularity and convergence mode assumptions. Several elements of our proof are inspired by Chizat & Bach (2018); the technique in their work however does not generalize to our three-layer setup. Unlike previous two-layer analyses, we do not exploit convexity; instead we make use of a new element: a universal approximation property. The result turns out to be conceptually new: global convergence can be achieved even when the loss function is non-convex. An important crux of the proof is to show that the universal approximation property holds at any finite training time (but not necessarily at convergence, i.c. at infinite time, since the property may not realistically hold at convergence).
Together these two results imply a positive statement on the optimization efficiency of SGD-trained unregularized feedforward three-layer networks (Corollary 10). Our results can be extended to the general multilayer case with new ideas on top and significantly more technical works or used to obtain new global convergence guarantees in the two-layer case (Nguyen & Pham (2020); Pham & Nguyen (2020)). We choose to keep the current paper concise with the three-layer case being pro- totypical setup that conveys several of the basic ideas. Complete proofs are presented in appendices.
Notations. K denotes generic constant that may change from line to line. |-| denotes the absolute value for scalar and the Euclidean norm for a vector. For an integer n, we let [r]] = n..
2 THREE-LAYER NEURAL NETWORKS AND THE MEAN FIELD LIMIT
2.1 THREE-LAYER NEURAL NETWORK
We consider the following three-layer network at time k E Nso that takes as input x E Rd;
y (z; W (k:) = PP (H; W (A))),
(I)
H3 (z; W (k)) = 1 - Σ
n2
wa (k. P2 (H2 (z, jo; W (k))),
112
j==1
2

---

Published as a conference paper at ICLR 2021
1 11
a1
H2 (2, W (l:)
w2 (k, sp ((w1 (k; ,0))
in which W (k) (w1 (k: ,,, W2 -,.,,,), Wa (k. :) consists of the weights W1 (1. E R W2 (k, (1. j2) E R and Wa (k, 12) E R. Here Φi : R + R. 92 : R + R and pa R + R are the activation functions, and the network has widths (n1,
We train the network with SGD w.r.t the loss L: R R20. We assume that at each time k we draw independently a fresh sample 2 (l) = (x ( E R x R from a training distribution P. Given an initialization W (0), we update W (k) according
W3 (As + 1, W3 (2) - <5s (le) Grads (= (k)) W (k)),
W2 (A: + 1, J1: W2 j1. - < (Rt) Gradz (a (k) f2; W (A)),
(A: + 1,51) wi (k.ji) - E (At)) Gradi (= (k) .Ji; W (k)),
in which J1 = 1,, = =1, -.....n.g. 22, e E Rso is the learning rate, & Rgo + Rgo is the learning rate schedule for wa, and for z (ar, y), we define
Grady (2, W (k)) OzC (y. ÿ (z; W (k))) & (Hs (2; W (k))) (H2 (2. W (K))), Grad2 (2,51.32: W (k)) A! (2, W (k)) P1 ((W1 (k. (2))))))) 2)),
1
Gradi (2,56 W (k)) 112
Σ
AN H (z,j2; W (k)) w2 (k.J)..2)) s ((wi (k; x)) x,
A (2, W (K)) = O2L & (z; W (k)) & (Ha (z; W (k)) wa (k, j2) vp (H2 (r, W (k:)).
We note that this setup follows the same scaling w.r.t. n, and n2, which is applied to both the forward pass and the learning rates in the backward pass, as Nguyen (2019).
2.2 MEAN FIELD LIMIT
The MF limit is a continuous-time infinite- width analog of the neural network under training. To describe it, we first introduce the concept of a neuronal ensemble. Given product probability space (,.FFP) P) = (S) x 122, x F1, Pi x P2), we independently sample C; ~ P., i = 1,2. In the following, we use Ec, to denote the expectation w.r.t. the random variable C, ~ P; and e; to denote an arbitrary point o; E Ωg The space (,, P) is referred to as a neuronal ensemble.
Given a neuronal ensemble F2, P), the MF limit is described by a time-evolving system with state/parameter W (D), where the time 1 E R2 and W (t) (wi (1,-), w2 (t, ; .) wa (1, ()) with IR and wys Rza x O2 - IR. It entails the
w1= R2o x Ω + Rdd, W2 R2 x 22 x 92 - quantities:
9 ( W (1)) 4P (H3 (I; W (())),
W (O))), Ha (z; W (1)) = EcN [wa (1, C2) 92 (H2 (z, W
H1 (2, 02 W (t)) = Ec. ww2 (t, C1 02) pi (w1 (1, (1,C)),00 x>)]].
Here for each t E Rz0. W1 (t, is (II, meeaurablee and similar for we (t, .),, ws (t.-). The MF limit evolves according to a continuous-time dynamics, described by a system of ODEs, which (wn (0,-), w2 (0, :)), wa (0, )), the dynamics solves:
refer to as the MF ODEs. Specifically, given initialization W (0)
Owwz (2) = -Ea (1) (02; W (1))
D.22tt,c, - (f) 22 (c1, ozi W (t)), Dewi(t,c) = -Ei (t) A (c: W (t))
Here C1 E Q2, C2 E f2 Ez denotes the expectation w.r.t the data Z (X,Y) P, and for
2 = (z, v), we define
Ag (o2; W (t)) = Ez |026 (V, (X;W (O)) &a (H& (X; W (1)) 62 (H2 (x,oz oz W (0)).
To absorb first layer's hias to Wi, we assume the input a to have appended to the last entry.
3

---

Published as a conference paper at ICLR 2021
22 (2) 02 W (t)) = Ez - (Z, c2 W (t) 41 ((wr (t,
A (cr; W (t)) = Ez [Ec, [ (Z, C2: W (t)) wy (t,c1, C2)] pi ((w) (t, (1) .X))
A 02; W (t)) = 026 (v. 9 (z; W (t))) & (Ha (z; W (t))) wy (t, 62) da (H₂ (z, o2 W (t))).
In Appendix B, we show well pell-poses of MF ODEs under the following regularity conditions. Assumption 1 (Regularity). We assume that 91 and p2 are K-bounded, s 4/ and se's are K- bounded and K K-piischitz. 12/2 and 4p/3 are non-zero everywhere, 0zL(+.) is K K-Lipschitz in the second variable and K -ouudddd,, and X| S K with probability 1. Furthermore En, &2 and Es are K-bounded and K K-Lipschitz.
Theorem 1. Under Assumption 1, given any neuronal ensemble and an initialization W (0) such that ess-sup [uz (0, C1, C2)1, ess-sup (wa (0, C2) S K, there exists a unique solution W to the MF ODEs on t E (o).
An example of suitable setup is 91 92 tanh, (P3 is the identity, L is the Huber loss, although non-convex sufficiently smooth loss function suffices. In fact, all of our developments can be easily modified to treat the squared loss with an additional assumption |Y| I K with probability 1.
So far, given an arbitrary neuronal ensemble F. for each initialization W (0), we have de- fined MF limit W (t). The connection with the neural network's s dynamics W (It) is established in the next section.
3 CONNECTION BETWEEN NEURAL NETWORK AND ITS MEAN FIELD LIMIT
3.1 NEURONAL EMBEDDING AND THE COUPLING PROCEDURE
To formalize connection between the neural network and its MF limit, we consider their initializa- tions. In practical scenarios, to set the initial parameters W (0) of the neural network, one typically randomizes W (0) according to some distributional law p. We note that since the neural network defined w.r.t. a set of finite integers (21,n2)) so is p. We consider a family Init of initialization laws, each of which is indexed by the set {M1, n2
Init {p: p is the initialization law of neural network of size (n1,11)) n1, n2 N
This is helpful when one is to take a limit that sends n1, 12 - oo, in which case the size of this family |lnit] is infinite. More generally we allow |nit] < co (for example, Init contains a single law P of a network of size {ni, and hence |Init| 1). We make the following crucial definition.
Definition 2. Given family of initialization laws Init, we call (0,F, P. (w?) 2,13)) neuronal embedding of Init if the following holds:
1. (2, P. (Ω) x 92 Fi x F2, P1 P2) a product measurable space. As reminder, we call it a neuronal ensemble.
2. The deterministic functions wf Ω + Rd, ud::: Ωi x l2 + R and wb Q2 + IR are such that, for each index n2) of Init and the law p of this index, i with an abuse of notations - we independently sample {C, (id)) ~ P, ii.d. for each i = 1, 2, then
Law (we (Ci (i)),, w (Ci(i), C2 (iz)), wg (C2(j2)), j. E [m]] i = 1,2) = p.
To proceed, given Init and (mi, n2) in its index set, we perform the following coupling procedure:
1. Let (2,F, P, { w 2 be neuronal embedding of Init.
2. We form the MF limit W (t) (for RRz)) associated with the neuronal ensemble
(9, F, P) by setting the initialization W (0) to (0,-) = w (.), wa (0,,,) = wg (s:) and W3 (0, .) wgg (.) and running the MF ODEs described in Section 2.2.
We recall the definition of ess-sup Appendix A.
4

---

Published as a conference paper at ICLR 2021
3. We independently sample C₁ (ji) P. for i 1,2 and j. 1, ..., B2. We then form
the neural network initialization W (0) with w1 (0, = w (C1 (ji)), W2 (0, (1) wº (Ci (5i), C2 (fz)) and wa (2) the network's 's trajectory W (k) for l E Ngo as in Section 2.1, with the data 2 (R) generated independently of {Ci (i));=1,2 and hence W (0).
wg (C2 (1z)) for
[n]] J2 [nz]. We obtain
We can then define a measure of closeness between W (Lt/el) and W (1) for t E [0,T]: DT (W, W) = sup { [w ([t/e] - w1 (t, C1 (fi)), we2 (Lt/e] (1. - w2 (t, C1 Wa [Lt/c], - w's (t, C2 (fa)) 1 T, J1 S m,, f2 S n2)
(2)
Note that W (1) is a deterministic trajectory independent of (11,722) whereas W (k) is random for all & Nzo due the randomness of {C, (i:))=1,2 and the generation of the training data = (k).
Similarly Pr (W, W) is a random quantity.
The idea of the coupling procedure is closely related to the coupling argument in Sznitman (1991); Mei et al. (2018). Here, instead of playing the role of a proof technique, the coupling serves as a vehicle to establish the connection between W and W on the basis of the neuronal embedding. This connection shown in Theorem 3 below, which gives an upper bound on 2r (W, W).
We note that the coupling procedure can be carried out to provide connection between W and W as long as there exists a neuronal embedding for Init. Later in Section 4.1, we show that for common initialization scheme (in particular, i.i.d. initialization) for Init, there exists neuronal embedding. Theorem applies to, but not restricted to, this initialization scheme.
3.2
MAIN RESULT: APPROXIMATION BY THE MF LIMIT
Assumption 2 (Initialization of second and third layers). We assume that ess-sup Mu (C1, C2) ess-sup wug (C2) S K, where w9 and w9 are as described in Definition 2.
Theorem 3. Given a family Init of initialization laws and a tuple (n1:22) that is in the index set of Init, perform the coupling procedure as described in Section 3.1. Fix terminal time T eNioo- Under Assumptions and 2. for e S 1, we have with probability at least 1 28,
2T (W, W) S KT (11 + vt
log 1/2
3 (T + 1) m
te) erra,r (c, n1, n2),
VTimin
in which Mminn = min {ni, n2). 7max = max n2)), and Kr = K (1 +++70 TT).
The theorem gives connection between W (Lt/eJ), which is defined upon finite widths n, and n2g, and the MF limit W (t), whose description is independent of nı and n2. It lends a way to extract properties of the neural network in the large-width regime.
Corollary 4, Under the same setting as Theorem 3, consider any test function sb : R x R + R which is K K-Lipschitz in the second variable uniformly in the first variable (an example of vp the
loss L). For any 8 >> O, with probability at least - 36,
sup Ez [bb (Y, (X; W (Lt/el))) - Ez w (X: W (+)) < eTerra (c, n2)
These bounds hold for any n2 and n2, similar to Mei et al. (2018); Araújo et al. (2019), in contrast with non-quantitative results in Chizat & Bach (2018); Sirignano & Spiliopoulos (2019). These bounds suggest that nı and n2 can be chosen independent of the data dimension d. This agrees with the experiments in Nguyen (2019), which found width s 1000 to be typically sufficient observe MF behaviors in networks trained with real-life high-dimensional data.
We observe that the MF trajectory W (1) is defined as per the choice of the neuronal embedding (R, [w]] ),12))) which may not be unique. On the other hand, the neural network's trajectory
W (k) depends on the randomization of the initial parameters W (0) according to an initialization law from the family Init (as well as the data z (k)) and hence is independent of this choice. Another
corollary of Theorem is that given the same family Init, the law the MF trajectory is insensitive
to the choice of the neuronal embedding of Init.
5

---

Published as a conference paper at ICLR 2021
Corollary 5. Consider a famity Init of initialization laws, indexed by set of tuples {m1, m2) that contains a sequence of indices {m1 (m), me (m) m E N} in which as m + oo, min {m (m). m2 (m)} 1 log (max {m (m)) m2 (m)}) + O. Let W (t) and W (t) be two MF P. (1 322123))- The following statement holds for any T 2 0 and any two positive integers
trajectories associated with two choices of neuronal embeddings of Init, (Q,F, { {w] 11 and
nn and n2: if we independently sample C₁ (it) ~ P. and C2 (j:) ~ P. for js E [ma], i = 1.2, then Law (W n n2) Law(W (11, T)), where we define W (2,, T) as the below collection w.r.t W (1), and similarly define W (m1,712, T) w.r.l W (I):
W (11, n2. T) = { [w1 (t,, C1 (ft)), W2 (t, C1 (hi) C2 (f2)), wy (t. C2 (f2)):
ji E [mi], j2 E [12], 1 E [..
The proof's are deferred to Appendix C.
4 CONVERGENCE TO GLOBAL OPTIMA
In this section, we prove global convergence guarantee for three-layer neural networks via the MF limit. We consider a common class of initialization: i.idd initialization.
4.1 LLD. INITIALJZATION
Definition 6. An initialization law P for neural network of size {mi, n22 is called (p",p², ".pppp))) i..d. initialization (or i.idd initialization, for brevity), where p², p² and g² are probability mea- sures over R&,, I and respectively, if {w1 (0,f1))ssemm are generated i.idd according to p², {w2 (0, (1))))) } in e(ma), Jat[mal are generated i..d. according to p² and {wa (0, 220|=2] are gen- erated i.i.d. according to p², and W1, W2 and W3 are independent.
Observe that given (p'. p²), one can build family Init of i.i.d. initialization laws that contains any index set {πI, n2). Furthermore i.d.d initializations are supported by our framework, as stated in the following proposition and proven in Appendix D.
Proposition 7. There exists a neuronal embedding P. initialization laws, which are p², -ii.i..
for any family Init of
4.2 MAIN RESULT: GLOBAL CONVERGENCE
To measure the learning quality, we consider the loss averaged over the data Z P:
(V) Ez (XX))
where V = (U1, U2 03) a set of three measurable functions U1 Ω2 + Rd. U2 Qn x l2 + R. U3 G2= RR
Assumption 3. Consider a neuronal embedding (wL) (im12))) of the (p) - -i.id.
initialization, and the associated MF limit with initialization W (0) such that (0, = w0 (.), W2 (0,,,' .) = =wh (c,)) and Wa (0.-) = wg (-). Assume:
1. Support: The support of pl is Rdd.
2. Convergence mode: There exist limits Wh W2 and Wy such that as t + xc.
E [11 + (irz(C2)) |wa (C2) ( (Ci, C2) -
(3)
E (1 (C2((())) (C2)l) ( (C2) [t2 Ci - +
(4)
E [11 + [wx(C2)) |ws C2) m2(C2)] +
(5)
ess-supRe, [10 (1, C, C2]
0
(6)

---

Published as a conference paper at ICLR 2021
3. Universal approximation: {s1 ((,, [ has dense span in L2 (Px) (the space of square integrable functions w.r.t. Px the distribution of the input X).
Assumption 3 is inspired by the work Chizat & Bach (2018) on two-layer networks, with certain differences. Assumptions 3.1 and 3.3 are natural in neural network learning (Cybenko (1989): Chen & Chen (1995)), while we note Chizat Bach (2018) does not utilize universal approximation. Similar to Chizat & Bach (2018). Assumption 3.2 is technical and does not seem removable. Note that this assumption specifies the mode of convergence and not an assumption on the limits Wh, t2 and de Specifically conditions (3)-(5) are similar to the convergence assumption in Chizat & Bach (2018). We differ from Chizat & Bach (2018) fundamentally in the essential supremum condition (6). On one hand, this condition helps avoid the Morse-Sard type condition in Chizat & Bach (2018), which is difficult to verify in general and not simple to generalize to the three-layer case. On the other hand, it turns out to be natural assumption to make, in light of Remark 9 below.
We now state the main result of the section. The proof is in Appendix D.
Theorem 8. Consider neuronal embedding (2, { (w]) ) of (p' .').p"")-iiddd initial- ization. Consider the MF limit corresponding to the network (1), such that they are coupled together by the coupling procedure in Section 3.1, under Assumptions 2 and 3. For simplicity, assume & (.) &2 1. Further assume either:
(untrained third layer) Es (-) = 0 and w/ (C2) + 0 with positive probability, or
(trained third layer) Eg (.) 11 and L (wf, wh, wd) Ez [E, (0)))-
Then the following hold:
Case 1 (convex loss): If L is convex in the second variable, then
lim L ((W (4) inf V L (V)
inf Ez (C
9: RR
Case 2 (generic non-negative loss): Suppose that 0z (y. 9) = 0 implies C (y.0) = 0. If cc.
y - y(x is a function of x, then L (W (1))
Remarkably here the theorem allows for non-convex losses. A further inspection of the proof shows that no convexity-based property is used in Case (see, for instance, the high-level proof sketch in Section 4.3); in Case 1, the key steps in the proof are the same, and the convexity of the loss function serves as a convenient technical assumption to handle the arbitrary extra randomness of Y conditional on X. We also remark that the same proof of global convergence should extend beyond the specific fully-connected architecture considered here. Similar to previous results on SGD-trained two-layer networks Mei et al. (2018); Chizat & Bach (2018), our current result in the three-layer case non-quantitative.
Remark 9. Interestingly there is a converse relation between global convergence and the essential supremum condition (6): under the same setting, global convergence is unattainable if condition (6) does not hold. A similar observation was made in Wojtowytsch (2020) for two-layer ReLU networks. A precise statement and its proof can be found in Appendix E.
The following result straightforward from Theorem and Corollary 4, establishing the optimiza- tion efficiency of the neural network with SGD.
Corollary 10. Consider the neural network (1). Under the same selting as Theorem 8, in Case 1,
lim lim lim Ez (Y, (X; W ((((()))] inf L (S1, fz) inf E [C (Y, 9 (X))
t-100 m2 +0
fu.J2.Ja
probability, where the limit of the widths is such that min {nı, n22-1 log (max (ni, n2)) + 0. In Case 2, the same holds with the right-hand side being O.
4.3 HIGH-LEVEL IDEA OF THE PROOF
We give a high-level discussion of the proof. This is meant to provide intuitions and explain the technical crux, our discussion may simplify and deviate from the actual proof.
7

---

This task is the more major technical step of the proof. To that end, we first show that there exists a mapping (t, u) + M (t, u) that maps from (t, W1 (0, c1)) (t. u) to W1 (t, (1) via careful measurability argument. This argument rests on a scheme that exploits the symmetry in the network Published as a conference paper at ICLR 2021
Our first insight is to look at the second layer's weight u2- At convergence time co, expect to have zero movement and hence, denoting W (oo) = (wi, W2, Wz):
A2 (c1, C2: W (oo)) = Ez a (Z, c2: W (oo)) pi ((w2
= 0,
for P-almost every C1 Oz- Suppose for the moment that are allowed to make an additional (strong) assumption on the limit wn: supp (0), (C1)) = R4. It implies that the universal ap- proximation property, described in Assumption 3, holds at l co; more specifically, it implies {s1 ((0) ()),,))= <1 E { has dense span in 1,2 (Px). This thus yields
Ex - (Z, 62; W (xo) X = =x =0
for P-almost every x. Recalling the definition of A, one can then easily show that
Ez (Y, i (X; W (oo))) X = =|=0.
Global convergence follows immediately; for example, in Case 2 of Theorem 8, this is equivalent to that 0,C (y (x). 1 (z; W (oo))) = O and hence C (y (z), 9 (æ; W (oo))) = for P-almost every x. In short, the gradient flow structure of the dynamics of w2 provides scamless way to obtain global convergence. Furthermore there is no critical reliance on convexity.
However this plan of attack has potential flaw in the strong assumption that supp (w1 (Ct) R", Le. the universal approximation property holds at convergence time. Indeed there are setups where it is desirable that supp (w, (C,)) * Re (Mei et al. (2018); Chizat (2019)); for instance, is the case where the neural network is to learn some "sparse and spiky" solution, and hence convergence time, if successfully trained, cannot have full support. On the other hand, one can entirely expect that if supp (w) (0, C1)) = IR initially at f = O, then supp (w) (t, C1)) = RRd any finite t 2 O. The crux of our proof is to show the latter without assuming supp (w] (Cz)) Rd.
the weight distribution
evolution. Furthermore the map M is shown to be continuous. The desired conclusion then follows from an algebraic topology argument that the map M preserves a homotopic structure through time.
DISCUSSION
The MF literature is fairly recent. A long line of works (Nitanda & Suzuki (2017); Mei et al. (2018); Chizat & Bach (2018); Rotskoff & Vanden- et al. (2019); Javanmard et al. (2019); Mei et al. (2019); Shevchenko & Mondelli (2019); Woj- towytsch (2020)) have focused mainly on two-layer neural networks, taking interacting particle system approach to describe the MF limiting dynamics as Wasserstein gradient flows. The three works Nguyen (2019); Araūjo et al. (2019); Sirignano & Spiliopoulos (2019) independently develop different formulations for the MF limit in multilayer neural networks, under different assumptions. These works take perspectives that are different from ours. In particular, while the central object in Nguyen (2019) is a new abstract representation of each individual neuron, our neuronal embed- ding idea instead takes keen view on a whole ensemble of neurons. Likewise our idea is also distant from Araújo et al. (2019); Sirignano & Spiliopoulos (2019): the central objects in Araújo al. (2019) are paths over the weights across layers; those in Sirignano & Spiliopoulos (2019) are dime-dependent functions of the initialization, which are simplified upon i.i.d. initializations.
(2018); Sirignano & Spiliopoulos (2018); 8); Wei
The result of our perspective is neuronal embedding framework that allows one describe the MF limit in a clean and rigorous manner. In particular, it avoids extra assumptions made in Araújo et al. (2019); Sirignano Spiliopoulos (2019): unlike our work, Araújo et al. (2019) assumes untrained first and last layers and requires non-trivial technical tools; Sirignano & Spiliopoulos (2019) takes an unnatural sequential limit n + oo before n2 + do and proves a non-quantitative result, unlike Theorem 3 which only requires sufficiently large min 11).. We note that Theorem 3 can be extended to general multilayer networks using the neuronal embedding idea. The advantages of framework come from the fact that while MF formulations in Araújo et al. (2019): Sirignano & Spiliopoulos (2019) are specific to and exploit i.dd initializations, our formulation does not. Remarkably shown in Araújo et al. (2019), when there more than three layers and no biases,
8

---

Published as a conference paper at ICLR 2021
Li.d. initializations lead to certain simplifying effect the MF limit. On the other hand, our framework supports non-ii.d. initializations which avoid the simplifying effect, as long as there exist suitable neuronal embeddings (Nguyen & Pham (2020)). Although our global convergence result in Theorem 8 is proven in the context of i.i.d. initializations for three-layer networks, in the general multilayer case, it turns out that the of a special type of non-i.i.d. initialization allows one to prove a global convergence guarantee (Pham & Nguyen (2020)).
In this aspect, our framework follows closely the spirit of the work Nguyen (2019), whose MF formulation is also not specific to i.i.d. initializations. Yer though similar in the spirit, Nguyen (2019) develops heuristic formalism and does not prove global convergence.
Global convergence in the two-layer case with convex losses has enjoyed multiple efforts with lot of new and interesting results (Mei al. (2018); Chizat & Bach (2018); Javammard al. (2019); Rotskoff et al. (2019); Wei et al. (2019)). Our work is the first to establish a global convergence guar- antee for SGD-trained three-layer networks in the MF regime. Our proof sends a new message that the crucial factor is not necessarily convexity, but rather that the whole learning trajectory maintains the universal approximation property of the function class represented by the first layer's neurons, together with the gradient flow structure of the second layer's 's weights. As remark, our approach can also be applied to prove a similar global convergence guarantee for two-layer networks, remov- ing the convex loss assumption in previous works (Nguyen & Pham (2020)). The recent work Lu et al. (2020) on a MF resnet model (a composition of many two-layer MF networks) and a recent update of Sirignano & Spiliopoulos (2019) essentially establish conditions of stationary points be global optima. They however require strong assumptions on the support of the limit point. As ex- plained in Section 4.3, we analyze the training dynamics without such assumption and in fact allow it to be violated.
Our global convergence result non--uanittative. An important, highly challenging future direction is to develop a quantitative version of global convergence; previous works on two-layer networks Javanmard et al. (2019); Wei al. (2019); Rotskoff et al. (2019); Chizat (2019) have done so under sophisticated modifications of the architecture and training algorithms.
Finally we remark that our insights here can be applied to prove similar global convergence guaran- tees and derive other sufficient conditions for global convergence of two-layer or multilayer networks (Nguyen & Pham (2020); Pham Nguyen (2020)).
ACKNOWLEDGEMENT
H. T. Pham would like to thank Jan Vondrak for many helpful discussions and in particular for the shorter proof Lemma 19. We would like to thank Andrea Montanari for the succinct description of the difficulty in extending the mean field formulation to the multilayer case, in that there multiple symmetry group actions in a multilayer network.
REFERENCES
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 242-252, 2019.
Dyego Araújo, Roberto I Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint 19Xiv:1906.00/93, 2019. 1, 5,2,,
Francis Bach. Breaking the curse of dimensionality with convex neural networks. The Journal of Machine Learning Research, 18(1):629-611, 2017.
Yoshua Bengio, Nicolas L Roux, Pascal Vincent, Olivier Delalleau, and Patrice Marcotte. Convex neural networks. Advances in neural information processing systems, pp. 123-130, 2006.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions Newral Networks, 6(4):911-91 1995. 4.2

---

Published as a conference paper at ICLR 2021
Lenaic Chizat. Sparse optimization measures with wer-parametteizee gradient descent. arXiv preprint urXiv1907.103000 2019. 4.3,5
Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. In Advances in Neural Information Processing Systems, pp. 3040-3050. 2018. 1, 1, 3.2, 4.2, 4.2, 5, D.2
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems 32, pp. 2937-2947. 2019.
George Cybenko. Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989. 4.2
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference Learning Representations, 2019. URL https: httpps:///ppaaee /pwp///veeeeeeeeee//eee/ee
Vitaly Feldman and Jan Vondrak. Generalization bounds for uniformly stable algorithms. In Ad- vances in Neural Information Processing Systems, pp. 9747-97517 2018. F
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kemel: Convergence and gen- cralization in neural networks. In Advances in Neural Information Processing Systems 31, pp. 8580-8589.8 2018. 1
Adel Javammard, Marco Mondelli, and Andrea Montanari. Analysis of a two-layer neural network via displacement convexity. arXiv preprint 190100390...135.. 2019. 5
Jachoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. In Advances in Neural Information Processing Systems 32, pp. 8572- 8583. 2019.
Yiping Lu, Chao Ma, Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean-field analysis of deep resnet and beyond: Towards provable optimization via overparameterivztion from depth. arXiv preprint arXiv:2003.05508, 2020. 5
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layers neural networks. In Proceedings of the National Academy of Sciences, volume 115, pp- 7665-7671, 2018. 1,11, 3.2, 4.2, 5
Song Mei, Theodor Misiakiewiez, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. arXiv preprint arXiv:1902.0660 2019.
Phan-Minh Nguyen. Mean field limit of the learning dynamics of multilayer neural networks. arXiv preprint
02880, 2019. 1, 1, 2.1, 3.2.
Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multi-
layer neural networks. arXiv preprint arXiv:2001. 2020. 1,5
Atsushi Nitanda and Taiji Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint arXiv: /7//:
2017. 5
Huy Tuan Pham and Phan-Minh Nguyen. A note the global convergence of multilayer neural networks the mean field regime. arXiv preprint arXiv:2006.09355, 2020. *, 1,5
Iosif Pinclis. Optimum bounds for the distributions of martingales in banach spaces. The Annals of Probability, 22(4):1679-1706 1994. F
Grant Rotskoff and Eric Vanden-Eijnden Parameters interacting particles: long time convergence and asymptotic eиror scaling of neural networks. In Advances in Neural Information Processing Systems 3/, pp. 7146-7155. 2018. 1,5
Grant Rotskoff, Samy Jelassi, Joan Bruna, and Eric Vanden-Eijnden Neuron birth-death dynamics accelerates gradient descent and converges asymptotically. In Proceedings of the 36th Interna- fional Conference on Machine Learning, volume 97, pp. 5508-5517, 2019. 5
10

---

Published as a conference paper at ICLR 2021
Alexander Shevchenko and Marco Mondelli. Landscape connectivity and dropout stability of sgd solutions for over-parameterizan neural networks. arXiv preprint arXiv:1912.10995, 2019.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks. arXiv
preprint 1.K5:0800000005 2018. 1.5
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of deep neural networks. arXiv preprint arXiv:1903.044: 2019. 1, 3.2,
Alain-Sol Sznitman. Topics in propagation of chaos. In Ecole d' été de probabilités de Saint-Flour XIX-1989, pp. 165-251. Springer, 1991. 1, 3.1
Colin Wei, Jason D Lee, Qiang Liu, and Tengyu Ma. Regularization matters: Generalization and op- timization of neural nets v.s. their induced kernel. In Advances in Neural Information Processing Systems 32, pp. 9712-9724. 2019.
Stephan Wojtowytsch. On the convergence of gradient descent training for two-layer relu neluonetworks in the mean field regime. arXiv preprint urXiv:2005. 13530, 2020. 9, 5
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Stochastic gradient descent optimizes over-parameterized deep relu networks. arXiv preprint 1X1v:18111.0888, 2018. 1

---

Published as a conference paper at ICLR 2021
A NOTATIONAL PRELIMINARIES
For real-valued random variable Z defined on probability space (9, P), we recall
ess-supZ inf e R: P (Z = 0}
We also introduce some convenient definitions which throughout the appendices. For set of neural network's 's parameter W, we define
sup w2 (Lt/c) (1, max sup wa (Lt/c) }. Ja5 KT
Will max
max Jismi, 1 ITT
Similarly for set of MF parameters W, we define:
|| || = max (sss-sup sup wu2 (4, C1, C2)1, ess-supsup wua (4, C2)1
1ST
1ST
For two sets of neural network's parameters W W", we define their distance:
|W lr = sup
[ (Lt/e] - w" (Lt/e] , ww's (Lt/el .1. w (Lt/e] 12,12)1, w/ (L4/e] (2) - w (L4/e] :2)] LE [o, ji E [mi], 12 E [r2] }.
Similarly for two sets of MF parameters W W" define their distance:
||W w IT ess-sup
|wy] (1, C1) - wn (1. C1),, w/ (1, C1, C2) - wyy (1. C1, C2)1 1E/0,71
sup
Iu's (t, C2) w/y (t, Ca)l
B EXISTENCE AND UNIQUENESS OF THE SOLUTION MF ODES
We first collect some a priori estimates.
Lemma 11. Under Assumption I, consider solution W to the MF ODEs with initialization W (0) such that [|lo < co. If this solution exists, it satisfies the following a priori bounds, for any T 1 0:
ess-sup sup wws (1, S I|[ Ill +KT [|[ lllo + Ko,3 (T),
tST
ess-sup sup Iu2 (t, C1, C2) < W Ilo + KTK.o, (T) = || W Iloo + Ko, (7),
ITT
and consequently, || S 1 + max (Ko (T). Ka,3 (T))
Proof. The bounds can be obtained easily by bounding the respective initializations and update quantities separately. In particular,
ess-sup sup Jwa (1, C2) < ess-sup wwa (0, C2) + T Tess-sup sup tTT
W3 (4, C2) < || llo +KT,
tT
a
ess-sup sup 1u (t, C1, Cz) < ess-sup W2 (0, (O, KT
cus- sup (ST
W2 (t, C1, C2)
s ess-sup W2 (0, C2) + KTess-
|Wa (t, C2)
15T
W + (T))
Inspired by the a priori bounds in Lemma 11, given an arbitrary terminal time T and the initialization W (0), let consider:
12

---

Published as a conference paper at ICLR 2021
for a tuple (a. b) E Bo. space Wr (a, of w
(W' (t)s
(wi (t, wug (1,),) ;) w/ (t, lisr such that
sup-sup wug (t, C2) S b
IST
sup lu (1, C1, C2) S a,
EST
R2 x Q f2 + RR wpg Rzo U2 + R. for a tuple (a, b) E R2 and W (0), a space WT (a, b, W (0)) of W E Wr (a, b) such that W" (0) = W (0) additionally (and hence every W in this space shares the same initialization W (0))
where w/ R22 S2 +RRd u²:
W 11. It is easy to that both spaces are complete. Note that Lemma 11 implies, under Assumption and ||| [| Illo < 0o, we have any MF solution W, if exists, is in Wr ([][W ||oo Ko.2 (T)) ; Illo K,3 (T)). For the proof of Theorem 1, we work mainly with Wi (l]I Illo +o. (T') ( |I| |lo K0, (2)) W (0)), although several intermediate lemmas are proven in more generality for other uses.
We equip the spaces with the metrie [W"-W"
Lemma 12. Under Assumption 1, for T 2 O, any W', W" E WT (a, b) and almost every ~ P:
uss-supsup
A (z, C2; W" (t)) & Ka,b:
'5T
ess-supsup sup H2 (z. Cz; W' (1) H2 (I,, Cz W" (1)) S Ka, IIW" W"
ST sup (z; W (t)) - Ha (z; W" (f)) S Ka, I|W - W"Il::
1ST
sup 1ST
W (f))) - 826 (v.0 (z; W (f)) - Ka,b IW W C2;W" A (z, C2; W (1) - (z. Cz; W" (0)) & Ka IW' - W" W
ess-supsup
IST
where Ka,b 2 1 is generic constant that grows polynomially with and b.
Proof.
The first bound is casy to see:
ess-sup sup
A (z, Cz; W' (4)) S sup lu's (1, C2) s b. tST
H
C2;
tST
We prove the second bound, invoking Assumption
HHz (z, C2; W (t)) - H2 (z, Cz: W" (t))
< K lu (1, C1 (2) ((w) (1, - iP (ww" (1, x)))
K Iw's (t,Ci, w/' C2)
< K (lw's (t, C1, C2)l + |W W'llz.
which yields by the fact W E Wr (a,b): ess-sup sup 'T
H2 (x. C2; W (t)) H2 (z. (r,C2:W" (f) ((a+1) -
Consequently, we have:
Mt (z; W' (t) - Ha W" (0) s K (4.C)))l (f)) 12 (1, C2 W" (f))
K w/ ug
S K Iw's (1, C2)l H2 (z, Cz; W (1) H2 (z;
C2; (1))
K ||W' - W
102C (z: W (f))) 0,C (8)) < K ( (z; W' (t)) - [2; W" (f) K H3 (z; W' (1)) Hy W" (1)),
13

---

Published as a conference paper at ICLR 2021
which then yield the third and fourth bounds by the fact W W E Wr (a, b). Using these bounds, we obtain the last bound:
| (z, C2: W (1)) - A/ (z, C2: C2;
$ K luls (t, C2)1 |02_ (y. 0 (x; (1; (f)) - 02C (v.0 (x; (t))
+ H3 (c; W (t)) H3 (z; W" (0)) + H2 (z, C2; W (0) - H2 (2., C2; C2: (t)))
+ K w/g (t. C2) - w/ (t. C2)1
from which the last bound follows.
To prove Theorem 1, for given W (0), we define a mapping Fw(o) that maps from W (w/, w/y w()) E Wr b) to Fw (W") W = (wo W72 w/y), defined by W (0) W (0)
and
8
"/ (1,02) - (1) As (c2: W (I)),
Ø1
0
ōt w/2 (t, = fe (t) 42 (c, c2: W (t)),
/ / (t,ct) = - (t) A (c: W (t)) 0t
Notice that the right- hand sides do not involve W Note that the MF ODEs' solution, initialized at W (0), is a fixed point of this mapping.
We establish the following estimates for this mapping.
Lemma 13. Under Assumption I, for T 2 0, any initialization W (0) and any W", W" Wr (a,
.51 ess-supsup A (C2i W ())) (5)) - As (C2; (s)) (s)) Kab I - l De C21 W (s)) A2 (C1. C2; W (s)) - Ka,b W - W"
ess-sup sup 1
sup-
A (C1; W (s)) - A (CI: W" (s)) S Kub I - W"
15t
and consequently, if in addition W" (0) = W" (0) (not necessarily equal W (0)) then
w'g (l) 35.' (1, C2) S Kabb L I|W' W I|. ds,
ess-sup sup
EST
ess-sup sup isT
Th (t,C,, (2) - (t, C1 C2) S Ka,b
IIW - w ds.
ess-supsup
|o| (1) - w/ (t, Ct) & Ka,b
I| W ds,
in which W' = (wh, W22y w5a) = o) (W'), W = (w". wh wh = FwIO) (W") and Ka.b 2 1 is a generic constant that grows polynomially with a and b.
Proof. From Assumption I and the fact W W E WT (), we get:
(X0_ Aa (C2; W (s)) - Aa (C2; W" (s)) < KEZ (X) (s))
+ KEZ H3 (X; (xx
+ KEz | H2 (2, (x) - (X,C;; (x))))
AA (C1, C2; (+)) - Ae (C1, (""), < Ka, ww (x, (1) - w/ (s,
+ K Ez [ (2 (s)) - A (Z, W2" (s))
A (C;;W' (s)) (C;, (s)) < K2 A (7, (')))) - (7,C%) W" (+)]
14

---

Published as a conference paper at ICLR 2021
+ Ka,b lug (s, C2) w2 (s. C1, C2)
+ Ka,b lun (s, Ci) - w/ (s, Ci)I:
from which the first three estimates then follow, in light of Lemma 12. The last three estimates then follow from the fact that W (0) = W" (0) and Assumption 1; for instance,
ess-sup sup juns (1, C2) w6 (1, C2)1 < / esT
Ø w C2) 0I.
ess-sup
w'3 (0, C2) ds
0t
I K
ess-sup As (C2: W (s)) As (C2: W" (s)) ds.
We now ready prove Theorem
Proof Theorem 1. We will use Picard-type iteration. To lighten notations:
W W W Illa + (T), |||W Illo Ko, (T), W (0)),
F= Fw(()-
Since |||Wl| < K by assumption, we have |||W |llo KK, (T) + Ko,3 (T) Kr. Recall that Wi complete. For arbitrary T O, consider W W WT- Lemma 13 yields:
IF (W") - F (W")l, Kr [
I|W" Wl ds.
Note that F maps to WT under Assumption I by the same argument as Lemma 11. Hence we are allowed iterating this inequality and get, for an arbitrary T >
s KT /' |EII-D (W") - (() (W'")
p()) (W")- (1)) (Wn))
dTz
T2
11
5 K2
J
p(k-2) (W") p( (W"")
1[s S T) dTzdT2
Ta
T -12
< Kx Jo Jo
I| W" (1x S .... S T2 TT) dl k +++1...Tf
1
< k!
k I| - l'
By substituting W = F (W"), we have:
F (k+1)) (w") - F(t) (W =
Σ ((()) (W") - ((() (w") ||
D
K
I - -
k=1
0C.
Hence k oc, p(k (W') converges a limit in Wir, which is fixed point of F. The unique- ness of a fixed point follows from the above estimate, since if W and W" are fixed points then
1 K IW W || = F(() (W') - (()) (W")) | S k!
W
while one take i arbitrarily large. This proves that the solution exists and is unique on t [0, T]. Since T arbitrary, we have existence and uniqueness of the solution on the time interval
[0, oo).
15

---

Published as a conference paper at ICLR 2021
C CONNECTION BETWEEN THE NEURAL NET AND ITS MF LIMIT: PROOFS FOR SECTION
C.I PROOF OF THEOREM 3
We construct auxiliary trajectory, which call the particle ODEs:
i
1i (1. -E3 (1) Ez 02h Y Xx W (4)) 5/ Ha (x:(()) 92 HH2 (x, W (1))]
01
0
W2 (1. (t) (t) Ez
A
ZZ W (t) fi ((001 (t, (t.)
0t
i wi (1,j1) - (l)Ez
1
AN
Z, W (1) W2 (L. vi ((w) (1))),
01
112
$2=1
22 (t, (t.)). (t. (t. and E Rzo- We specify the initialization W (0): 10 (0, w (C1 (fi)), U2 (0.ji..j2) = wh (C (2 (j2)) and Wi (s. w (C2 (f2)). That is, it shares the same initialization with the neural network one W (0), and hence is coupled with the neural network and the MF ODEs. Roughly speaking, the particle ODEs are continuous-time trajectories of finitely many neurons, averaged over the data distribution. We note that W (t) is random for all i E Rzo due to the randomness of (C; (id)),
in which ji 1, ..... 11. L,..., n2, W (t) (w)
The existence and uniqueness of the solution to the particle ODEs follows from the same proof as in Theorem 1, which we shall not repeat here. We equip W (t) with the norm
max sup lws (t.j2)) }
|| max
max (iS Jo S IST
sup U62 (t,31,52)l, janng IST
One can also define the measures 2r
W) and Dr (W, W
similar Eq. (2):
Ør (w,W) sup { |mm (t, C1 (fi) - w1 (t. C1 (f1))., Iw2 (f, C1 (f1), C2 (jz) - و2 (t. C1 (t, (z (z)),,
Iut (t, C2 (f2)) - dis (1, C2 (f2)) TST, j S m,, 52 S n2)
Ør (,, w) sup { [w (It/c] - @b (t:C(( |w2 (Lt/c) (1..je) 02 (t, C1 (11) (2) (z)))
Ww' (It/cl ,52) - -'s (t, (fz)) ts j1 S m. f/ $ n2).
We have the following results: Theorem 14. Under the same setting as Theorem for any 8 0, with probability least 6,
log" 1/2 (T + 1) m + e Kr
2r (,, (W,W)
Vemiu
5
in which Amin = min {ni, n Tmax = max nz), and KT = K (1+ TTK)
Theorem 15. Under the same setting as Theorem 3, for any d > and S 1, with probability at least 1 8,
Dr (w, (W.w)=
$1elog
/clg
22 + +e )eettr
in which Kr = K (1 + (K)
Proof of Theorem 3. Using the fact
gr (W. W) S 2T Ww. + 2r (W:w).
the thesis is immediate from Theorems 14 and 15.
16

---

Published as a conference paper at ICLR 2021
C.2 PROOF OF THEOREMS 14 AND 15
Proof of Theorem 14. In the following, let K₁ denote an generic positive constant that may change from line line and takes the form
K2 = K (1+t*),
such that K₂ 2 1 and Kr S KT for all t S T. We first note that at initialization, Do (W. W) 0. Since W| I K, I|Wlll I Kr by Lemma 11. Furthermore
easy to see that |||W N || S K almost surely. By the same argument as in Lemma 11, || S Kr almost surely.
We shall use all above bounds repeatedly in the proof. We decompose the proof into several steps.
Step 1 Main proof.
Let us define, for brevity
q (t, H3
x; W (t) H3 (x; (t)),
L2 12, 02)
H2 f2; W (1) H2 (2, 02; W (t)),
ds (1, = - Z, j2; W (t) W2 (1, - A (z, 02; W (1) un (2)..
Consider 1. 2 0. We first bound the difference in the updates between W and W. Let us start with uy and Wi.. By Assumption I, we have:
i
a ws (1, Co (o)) < KEz (1e3 (1, lea (1, X,,, C2 (1a))]
3a (2)
a1
o
Similarly, for w2 and WZ2
0
Hi "1 (1, (1), (22)
a
22 (1:31:
at.
$ KE ||
ZZ W (1) - (z, (z): [
+ K Iuss (1. C2 (f2)) (n (1)))
< K,Ez 1lys (1, (22 (1,,,,
+ Kr (lw (1,51) - wi
(ūa (t) - ws (t, C2 (i2))) C2
S K2Ez (l9a (t, X)] + 102 (2 + K. (w, W)
and for wi and W1, by Lemma 12,
Ot 11 (t,
w1 (t, C1 (it)
Ōt
1
< KKz
Ec, lوs (t, j2,,,,,,, Ci1), Ca) (G)
712 J==1
+ EEc2
[| (Z, Cz W (1)) 102 (4. (fi) Ca) W1 (1,51) wi (4, (1))
1
< KEZ
Ec, lus 2,J.,,1,2 C1 (G)
122
+ K.Pt (W,
To further the bounding, we now make the following two claims:
Claim 1: For any & >>0 0,
mAX fz
W3 (t + &, C2 (12)
w (1, "1, (1,22 S Kerss
Ot
17

---

Published as a conference paper at ICLR 2021
0 01
max S... Jeing
22 (t
+ E,
Ci
- - w2 (t, (1) (2), (i2)
3 Ke4eS,
max S1
w1 (t + 5, C1 (i1)
" (t, (it)
$ Ketg5,
0t
at
and similarly,
max S25
Li (t + 5,
i (2)
I Ke+55:
Öt
Öt
0
max ja Sm1, Sma
10
(1 + 52)
2/ (1, j1.
$ Kete5:
Ot
ōt
8
max Si
Wi + 5,51)
in (t,
I Ke+e5.
Öt
Öt
Claim 2: For any 71, 72, 73 > 0 and 2 0.
max
max N25
Ez [/02 (1, X, j2, C2 (f2))]I, Ez 112a (1, X)I,
1 R2
max Ez S15m1
Ec, (9s (t, Z. C2
T12
32
2 K. 2 (W. W
+ + 72 + 73)
with probability most
ni exp (-Ko
n2n2 Kr
112
1172
122
+
exp
+ exp
72
Kr
73
K
Combining these claims with the previous bounds, taking a union bound over t E {0, E, 25,..... [T/E] {} for some (0, 1), we obtain that
w' (t, C2 (2
max
max f=
th (t,
Ot
ōt
i U2 (t.
a
u2 (t, C1 (i)) Co(()))
max
51 11.. J25ng
Øt
at
wr (t, C (m))
max SA1
th (t,
Ōt
Øt
< Kr
Lr
(w, + 71 + 7/ + 73 + +E).
v [0,7].
with probability at least
T ni exp (-
nusn Kr
n2
(-
1172
12n3 Kr
]
1-
+
exp
cxp
92
Kr
73
The above event turn implies
< Kr 1'' (w. + + + 23 + E) ds,
2. (W.W)
SKr
and hence by Gronwall's lemma and the fact 2y (W, w) 0. we get
2r W. W
S (n + 72 + 73 cr.
The theorem then follows from the choice
KT
1og!/2
log 1/2
3 (T' + 1) n
+ +e). 71
+
KT
log 1/2
(3
(T' + 1) 1 to)
/mmax
n1
8
/n2
8
We are left with proving the claims.
18

---

Published as a conference paper at ICLR 2021
Step 2- 2- Proof of Claim 1. We have from Assumption 1,
J.
0 W' (s, C2) ds Öt
ess-sup Iwy (t +. C2) w3 (t, C2)1 S K
ess-sup
K,
ess-sup /uz (1 + E, C1, C2) w2 (1, C1 C2)1 S K
ess-sup
W2 (s, C1, C2) ds
Ot
K
ess-sup lwo (8, C2) ds
S Ki+ef,
0 w1 (s, C1) ds ess-sup Ot
ess-sup tml E.
w1 (t, C1)1 S K
$ K !.
ess-sup ww's (s, C2) 22 (s, C1, C2) ds
- Ke+s5.
By Lemma 12, we then obtain that
ess-supBzz [|H2 (X,C2;
C2; W (L+ - (X, (2, W (()W())) (I)] V Kiteš, Ez [|2 (X; W (t+ (X;W (t)] I KYtE5,
ess-supEz [| (Z, C2: W (t + E)) - - (Z.C2; (0)] S Kit55.
Using these estimates, we thus have, by Assumption 1.
i Wy (t + C2 (f2)) 8 ws (t, C2 (dz)) Ot
max S25mng
Ot
< Kite5 + KEz [|H3 (X;W (L+ - Ha (X; W (I)I
H2 + Kess-supBz [|// (X, W (O)I
< KitES,
0 का w2 (t, (1) Co (2))
max J2Sm2
W2 w2
(t)
0l
< Kitcf Kess-supEz - (L++) -
AN (Z. Cz. W W()lI
+ Kess-sup Wy (t, C2) [ - wn ( (t.
< KetcE,
① ot un (t, C1 (1)))
max fi
Øt W1 (t + E,+
Ki+55 + Kess-supEz [Ec, |
(t. C1, Ct. Cb]]]]
- - w2 w2
+ K Kess-supEc, [/ve (1, C2) w2 (1 ( C2,
+ Kess-supEc, [|w (1, C2) w2 (1, C1 C2) [m] (t (1+ wi (1,
< KrtE5.
The proof of the rest of the claim is similar.
Step 3- - Proof of Claim 2. We recall the definitions of gs. Q2 and g3. Let us decompose them as follows. We start with q2:
1g2 (1,,,12 Co (f2))
1 ny We (1, f2) 1p ((0r (1,1), xx) Ec, (wz (1, 2n ((wn (,))),a)))
n1 n=1
max
|ū2 (1, s ((0) ((.5)..)) we (t, C. (1), (z), 9 (ww) (t, C1 (f)) .x))
Si
19

---

Published as a conference paper at ICLR 2021
n1 Σ 22 (1, Ci (12) 91 ((w) C1 (fi)), x)) Ec, [w2 C1,
C2 (1)) i1 (ww (1) x))1]
J1
Q2 (2) + 22 (2)
Similarly, we have for q3
lgs x))
1 12
H2 (z. W (t)) Ec, [wn (t, C2) P2 (H2 C2; (+)))
iy (t, 82
n2
max jaSnz
Wi (2, 822
H2
x. J2i W (t)) Wis (t. C2 (1e) 62 H2 (z. (t))
1
Wy (t, C2 (22)) P2 H2 (z, (i2); (0)) - Ec, (wy (t, C2) 92 (H2 (z. C2; W (t)))
122
Q3.1 (z) + Q3.2 (x).
Finally we have for qu:
1
Ec, 9s (t, Ci (ii), C2)
12
A
2, J2: W (t) W2 (t.ji. A (z, C2 ()2): (t) (t
Snnng
+
A (z, C2 (fz) (t. (tGG))) - Ec, | (z, C2; W (+) 22 (t, (1), C2)
112
1=
= Qu11 (2, + Q1. (1,
Now let us analyze each of the terms.
We start with Q2,1. We have from Assumption 1,
max Ez 1Q2,1 (X,j2)
fz5m2
< K
max
102 (t. j2) w2 (t, C1 (2), (i))
+ max
Im2 (t, (11), (i))) (6 (t.ji) - wi (t. C (i1)
ji S1Si J25ng
< K, D. W, W)
To bound Q2.2, let us write:
Ze (x, 02, = wa (1, pi ((wi (1, x))).
Recall that C₁ (ji) and C2 (jz) are independent. We thus have:
E 22 (X, C (i)) C. (fz)) Co (iz) = Ec, [Ze are independent, conditional on C2 (j2). We
(i))))-
Furthermore (Zg (C, (11), C2 also have, almost surely
C2 Zz (X, Ci (1),, (jz)) < Kr.
by Assumption
1. Then by Lemma 19,
P (Ez [Q2,2 (X,jz) 2 KiT2) < (1/12)xxp (-n1 K.))
20

---

Published as a conference paper at ICLR 2021
To bound Q3,1, we have from Assumption 1,
Ez (Qa,1 (X)] S max (K |i03 (1, wa (t, C2 (52)) + K,Ez 11 (t, J2, C2 (12))
S=5
< K
(,,W)=
+ Ke max fasna
Ez 1l2e (4, X2, C2 (1a))))-
To bound Q3,2, noticing that almost surely
Iwa (t, C2 (fz)) p2 (H2 C2, (f)) - K,
by Assumption 1, obtain
P ( [Qa,2 (X)) 2 Krys) S (1/73) exp (-n2n3/ ki) similar to the treatment of Q2.2.
To bound Q1.1, using Assumption 1,
(w2 C1 (2)) Ez || (Z. W (t) AN (Z, C2 (12 w [
Ez (Q11 (Z, f1)1 S K max j2Sm2
ال f2) u2 (t, C1 (f1), C2 (f2)) Ez | (Z.52: W [0]
K max JuSn2
S K max Juy (1, Ci (11) C2 (iz)) [igg (1, - wg (1, C2 (f2)) JuSm2
+ + (t. C2 (ia)) Ez [l (t, + 1 (t. X, Cs (12))1)
+ K max |022 (2, - 12 (t. C1 (h) C2 (1z) 00 (1, juSnz
9 (W,W)= + Ez 100 (1, X)I + max (42 X2,, C2 (2))
S Ke
To bound Q1.2. we note that almost surely
A (z, C2 (jz); W (t)) u2 (t, Ci (1). C2 (2))
C. s K wua (t, C2 (L
< Kr.
Then similar to the bounding of Q2,2, we get:
P (Ez [Q1,2 (Z,j1) 2 Kmn) S (1/7) exp (-niit/K) Ki))
Finally, combining all of these bounds together, applying suitably the union bound over ji E [m]] and j2 E [122], we obtain the claim.
Proof of Theorem 15. We consider t < T, for a given terminal time T E eN2 We again reuse the notation K, from the proof of Theorem 3. Note that K₁ < Kr for all f < T. We also note that at
initialization, Po (w, W)) = 0. We also recall from the proof of Theorem 3 that ||Wlll S Kr almost surely.
For brevity, let us define several quantities that relate to the difference in the gradient updates be- tween W and W:
93 (kc 24, = O2 (y. y (z; W (I))) & (Hs (c; W (k)) P2 (H2 (I. f2; W (k)))
- -df J (z; (ke)
2 Hg (;; W (ke)
322 H2 (7, f2i W (ke) (ke))) 90 Hg (z; W (ke)
T3 2, = Es (KE) 3z6 y 2; W (kE)
80 H2 (z: W (ke))
Es (kc) Ez [2 (Y,8
x; W (kc)) &6 H3
Xx W (Ae)))
82 H2 (x, W (kc)) )]
22 (k. 2, (1, = = (z, W (k)) 9e ((w1 (z))
21

---

Published as a conference paper at ICLR 2021
42 2, W (ke)) 1o ((02 (he, ; x)),
r2 (1,5) = E2 (kt) A
22; W (kt)
sn (fir (1),
- E2 (kt)
Ez
A
Z: f2: W (ke) si ((00 (ke, J1) xx)]]
1
n2
qi 2,,,,
(z, j2; W (k)) w2 (k, 12) si ((wi (k, (1),),,)) a
j
1
j2 W (ke)) ú2 Jh, (2) 4 ((00 (ke, (i), ;;)) st
12
ja
r1 2k)) = E1 (ke)
HH
22 W (ke) is (ke, j1. vo' ((w1 (ke, ji), ,2))
122
52=1
22
Ei (kt) Ez
A
Z,,2; W (ke)) 22 (ke, v ((rr (ke, x X
112
j==1
Let us also define:
43 (k, = H3 (1; W (k)) H3 ax; W (ke)):
c (k, = H2 j2 W (k:) H2 f2: W (ke))
We proceed in several steps.
Step 1: Decomposition. As shown in the proof of Theorem 3:
i
a
max 22
W3 + E2)
Wz (2)
$ KitEš,
Ōt
θ
max Ji Smi. Sno
12 (t
+
22 (t. (1))))
$ Kitsf,
Øt
Øt
i
max Si5m
wi (t (,5))
w1 (t,
S Ki4SE.
Ōt
for any 2 0 and { 2 0. These time iimeiiterpollation estimates, along with Assumption allow to derive the following. We first have:
max |w' (Lt/e] .32) - -ma) (t, JuSma
[t/c|-1
S max
6
Σ K=
Ea (ke) Ez (ua (k, = (k), [2] + KKKE
< K max
[Qa, (L1/e] + Qa.2 (L4/e] .ja) 1Kr5,
S2522
where we define
1t/c|
Qx (L1)e),5)) =6
lgs (k, = (k)) (k)))
K=0
We/e-11
Q2 ([t/el,52)
e
Ta (k, 2 (k)
(Here -:4el-1 1/cl-1
if [t/c] 0.) We have similarly:
max Ww (Lt/c) .i) - we (t.hi) < K max (Q1, (Lt/c) + Q1. (Lt/c), + ttKic Si
22

---

Published as a conference paper at ICLR 2021
max S n2S in which
W2
ú2 (t,f1. s K
max Si n25ng
[Q2,1 ,2. + Q2,2 (It/c] f2)1 + tKec,
Lt/el-1
Q1.1 (Lt/ee)
lai (k, 2 (k), ii),
Q1 (Lt/el/)),,1)
é
T1 (k, 2 (k)
k=
|t/c|-1
22,1 [L4/e] (2, =
Σ
I/2 (k, = (k) -1,
k=0
Lt/e--1
22.2 [Lt/e] .f1,22) = e
Σ
12 (k, (k), .1.
The task is now bound Q1, Q1, Q2,1, Q2, Q3,1 and Q3.2-
Step 2: Bounding the terms. Before we proceed, let give some bounds for o and of which hold for any x E Rd:
1 (k, (k, & K max (/ib2 (ke, (1,,)) W1 (kt, ui (kc, + [w2 (k:, (1.22 w2 (kc, (1, Jisn
S Kk. (W, W)
15 (k, S K max JaSn2
(Iwa (ki, - wa (k), + Vãa (Ae ja)] 105 (k, (k)),
S Kee W W)
With these, we have the following:
Let us bound Qa By Assumption 1,
1ga (ks = (K) 2 (At) ,2)1 S K (og (kt x (4),52) + +155 (k, (k))).
We then get:
[t/c|-1
max Q3,1 (Lt/e) .f2) S Kie S25
Σ
D (,, W)
k=0
Similarly to Qa,1. consider Q2,1
122 (ks, (k) (A) .21.52) < K A (2 (k) j2: W (k)) - A (= jos W (ke)
+ K A 2 (k) W (kc) wwa (kr.1) - wr (kc,
< KR (/05 (k, (k) .2)] + 1 (k, (k))])
+ K |w3 (k,j2) is (ke, j2)l Kr. |w1 (k.j1) To (ke,ji))
which yields
1t/e|11
max Q2 (lt/el,fi.da) s Kie Σ ja Smn So2
Lke W W)
km=0
Again we get a similar bound for Q1,1:
K D2
Iqi (k, (k) z (k) 10ii))
j (ke, J., f2))
A (=(k). jo: W (k)) -A (=(4), W (ke)
72
23

---

Published as a conference paper at ICLR 2021
K
ng
A 2 (k)
./2; :,;; W (kre) |w2 (1:)) ú0 (Nte; (1.)))))))
12/2
12=1
K
f2
Li (ke, J1. A (a (k) j2: w (ke)) Iw2 (R:.1) wr (ke,
n22
j2
112 (k, :2)) + 10 (ks, (k)))) + K Ohe (W, W.). JaSnz
S Kk
max
which yields
[t/e]-1
max Q1, (Lt/e] .hi) S K,E Sn
Σ R=0
The (.. W
Let us bound Q3,2. Let us define:
Is (2)
Σ (K, (k) (k))))))
Ls (0.j2) 0.
t=
Let Fk be the sigma-algebra generated by {= (f) E (o...., k 1}}. Note that {rs (2) is a martingale adapted to {Fa Furthermore, for & S T/e, the mar- tingale difference is bounded: Ira (ki (kk, S K by Assumption Therefore, by
Theorem 20 and the union bound, we have:
< 2nz exp (- (-K
max Su5nz BE (O.1....TTT
max
QI (l,52) ((.5))2
K (T (T++))e
The bounding of Q2,2 is similar: |ra (k; = (k), (1, f2)1 s Kir almost surely by Assumption I, and thus
Q2.2 (L,j1, 26) S 2n exp (- Kr (T + 1)e
IP
max jismi, S25 (o. T/e)
max
Again the bounding of Q1,2 is also similar: Irr (k, (1:) ,131 I Kice almost surely by Assumption 1, and thus
2 < 2n1 exp (-Kr (T + 1)e
P max
max Sm (/e..
Q1, (C.51) (c.51)
Step 3: Putting everything together.
All the above results give us
We/l-1
Dusejo (W. W) < KyC
Σ
Dre
(. (W.W) w) + TKre
S
which hold with probability at least
E2 (Trr
2n 72 exp
KT (T + 1)c
The above event implies, by Gronwall's lemma,
2r ( w) $ (E + c) Kr.
Choosing { = Kr (T + eog (2n completes the proof.
C.3 PROOFS OF COROLLARIES 4 AND 5
Proof of Corollary 4. By the assumption on v and Assumption 1, we have:
BE (1) (Y. (X; W (LteE))) (Y, (X; (()))
24

---

Published as a conference paper at ICLR 2021
KE Hg (X; W ([t/el)) (X; (+))
KEz [Haa (X; (t) - H3 (X; W [Lt/e] e)) + K |H (X; W (Lt/el)) - Ma ( (Lt/eJ))))
An inspection of the proof of Theorem 3 (in particular, the proofs of Theorems 14 and 15) reveals that firstly,
sup Ez [|4 (X; W (t)) - Ha (x;W (Lt/e - KTE, ES
and secondly,
sup Ez [1 (X; W (L(/e - Ha (x; W (L(feD) RT
I KT DT (W, W)
102
1/2
Vimin
with probability at least - 5. Together with Theorem 3, we obtain the claim.
Proof of Corollary 5. Observe that for each index {N1, N₂3 of Init, one obtains a neural network
initialization (0)) with law p by setting
W1 (1) C((i), W2 (1) wy (0, C1 (f), (fz)),
W3 (3, = w3 (0, C2 (22)). ji E [Ni], j2 E [N2].
We consider the evolution W (k) starting from W (0), which is independent of W. Note that W (k) is deterministic function of its initialization W (0) and the data {= (i) Similarly, we
consider the counterpart for W: the evolution W (k) as a function of the initialization W (0) and the data {8 (i))zsk- Due to sharing the same distribution for both the initialization and the data, these evolutions have the same law; to be specific, W (n1, n2, T) and W (n1, n2. T) has the same distribution for any n1, f2 and T, where we define
W (n1 n2, T)
{wi (k, wa (k,J1, j2), w3 (2,
j1 [m1], J2 E [m2], ke S [T/c]]
and similar definition for W (n1, n2, ,,,T), T). In other words,
W
W. w
inf coupling of [))
E
max JiS. JaSna { Ww (k:,1) - w1 (k:,1)),
W2 (1.
W2 (ks, (1. |w3 (k: ws (k. (22, =0.
Theorem 3 implies that for any tuple { n2) such that S Nı and n2 S N2. with probability at least 1 -28,
Φ m2)) (W, W) = max
sup (ST, fišm J2S
W2 (2.5).2)) - w2 (t, C1 (1) C2 (i2)),
sup
Ww (t.jz) ws (t. C2 (f2))),
1ST, J25na
|w] (t, - w1 (t. C1 (iDil s ÖST (e, N1,
sup ST,, Si
where ŌS. N2) + 0 0 and N log Numax 0 with Nmin min NN2, and Nunx max (N1, N2). We also have similar result for W and W. As such, with probability at least 1 -6,
inf W (W, (w, coupling of (W.w)
{ Iu (t, (1, (fi) -r (1))
E
sup
Sim. jo Snz
1m2 (1, O2 (5z))
12 tt, 11C10in (ii,, O2 (62)
(t, (t,C (J2)) 43 (1,22(b))) (6z)
I 9 ml (W, W) of m2). (W.W) + w W. W
$ 200T N1,
By fixing the tuple [n,, while letting claim.
+ 0, -11 log Nmax + and + O, we obtain the min
25

---

Published as a conference paper at ICLR 2021
D GLOBAL CONVERGENCE: PROOFS FOR SECTION 4
D. PROOF OF PROPOSITION 7
Proof of Proposition 7. Consider a probability space (A. Po) with random processes -alluee pi (01), RR-alued p2 (01, 02) and IR-valued pa (02), which are indexed by (01, 02, E [0, x [o, 4], such that the following holds. Let my and m2 be two arbitrary finite positive integers and,
with these integers, let (o(h)) [0, 1] 4 E [mL], i == 1.2] be arbitrary collection. For each
i = 1, let S be the set of unique elements in {o!t)) hi E [mal]. Similarly, let R2 be the set k1 E [mm]], k2 E [mel] We have that {pr (01): 01 S1). {ps (82) 82 ES2) Jand {p2 (01, 62) (01, 02) Ry} } are all mutually independent. In addition, also have Law (p) (8)) = p², Law (Ps (82) = p² and Law (P2 (8', 82)) = p² for any 0 E Si, O2 S2 and (2) E R2. Such space exists by Kolmogorov's extension theorem.
of unique pairs in { {(f(t)) (0(()))
We now construct the desired neuronal embedding. For i = 1.2, consider 22 = A x [0. 1] and F. 9 x B ((0, 1)), equipped with the product measure Po Unif ((O, 1]) in which Unif ((0, 1)) is the uniform measure over [0, 1] equipped with the Borel sigma-algebra B ([0, 1]). We construct
Ω = f2 x S2 and F = Fi x F2, equipped with the product measure P = (Po x Unil ((0, ID)).
Define the deterministic functions w/ Ω + Rdd wh- On x f2 - IR and wb : Q2 + R:
wf (A1,0:)) P1 (01) (A1),
w ((A), 01) (22, 02)) = p2 (0,, 02) (22). wg ((22, 02)) = Pa (02) (A2). It is easy to check that this construction yields the desired neuronal embedding.
D.2 PROOF OF THEOREM 8
We first present a measurability argument, which is crucial to showing that a certain universal ap- proximation property holds throughout the course of training.
Lemma 16 (Measurability argument). Consider family Init of initialization laws, which are (p>,p",p") ) such that p²-almost surely |w2l S K and p²-almost surely (wal S K. There
exists a neuronal embedding (2, F,, { (w) -1,2,3) of Init such that there exist Borel Borel functions wi
and A
** for which P-almost surely, for all 1 2 0,
Wi (t, (1) = w (t, w (Ci)).
A (z, Cz; W (f)) = A (t, u. (C2)),
where W (1) is the MF dynamics formed under the coupling procedure with this neuronal embedding as described in Section 3.1. Furthermore,
θ Ot w (t, u1) = -c) (t) / Ez A (t, 2, us) u2y'1 ((wi (t, x)) X] p² (duz) p² (dus)
/'' Ez E2, I / (t, Z, us) A (s, (s, us) pa (dus)
+ t (t)
x pi (ww) (s, X')) of ((wi (t. u1) X)) x ds.
with initialization w (0, u,) = in for all in E supp (p)) and 2 0. where Z' an independent copy of Z.
Proof. We denote by K2 a constant that may depend on t and is finite with finite t. By Proposition 7, there exists neuronal embedding that accommodates Init. We recall its construction and reuse the notations from the proof of Proposition 7; in particular:
wof ((11, 0:)) P1 (01) (A1),
26

---

Published as a conference paper at ICLR 2021
2 ((A1, (22, 02)) P2 (01, 82) (A2).
w3 ((22, 02)) pa (02) (A2).
Let S1, S3 and S2 denote the sigma-algebras generated by w (C1), wz (C2) and (we (Ci), w (Ci, C2), w (C2)) respectively. Let S13 denote the sigma-algebra generated by Si and Sy. We also let S to denote the sigma-algebra generated by S and the sigma-algebra of the data Z. We define similarly for S and Sg.
Step 1: Reduced dynamics. Given the MF dynamics W (t). let us define Ås (2) = E| (C2; W (4)]Sa] (c2),
2 (2) = E [ (C1, C2; W (t)] (c1, 02), A (1) = E [Ar (C1; W (I))
We recall from the proof of Theorem 3 that for any 8,8 2
ess-sup [3a (t, C2) - w3 (s, C2) _ K It s
ess-sup (22 (t, C2) u2 (s, C2) I Kivs |t - 8)
ess-sup |wi (1) - wi (8, Ci) S Kivs |1 888
Then by Lemma 13,
E
As C2) As (s. C2)
Kivs It s|?
E
| C1 A2 )s, C1, Ca)l²]
Kiva || s
E
A (1. - A (s, Gi) Itvslt-8h
Therefore, by Kolmogorov continuity theorem, there exist continuous modifications of the (time- indexed) processes ÂI. 2₂ and Åg. We thus replace them with their continuous modifications, written by the same notations.
Given these continuous modifications, consider the following reduced dynamics:
i
Øt W3 (t) Ea (t) A (t,
8
(t, C1, 02) =2 (t) AS 02),
Ot
0
(1,4)) - (1) A (1,c1), Øt
in which:
11 R20 x 92 + Rd W2 Rz0 x 0 x 22 + Wy Rz0 x 2gg + I.
W (t) { i. (t, .) U2 (t, ......)....... w3 (t, }} is the collection of reduced parameters at time t,
the initialization is w7 (0.-) wf (-), T (0.,,)) wg (:. and W3 (0. = wg (-), i..
W (0) = W (0).
Step 2: Measurability of the reduced dynamics. It is easy to see that W3 (t. C2) is ma-mesuraable by its construction and the fact W3 (0, C2) wg (Cb) is me-maauuaale. Similarly, U2 (t, C1 C2) is Sy-measurable and w (t, C1) is -m-maaurrale..
Notice that there exist Borel functions wi, wy and vy for which P-almost surely,
T (t, C1) wi (t, w/ (Ci)).
W2 (1, C1, C2) wi (1, w (Ci) wi (C1, (2), mn (C2)),
W3 (1, C2) wi wg (C2)).
Indeed, since W2 (t, Ci, C2) is Sy-measuaalle, there exists function W5 (f,-) for each rational t such that the desired identity holds for P-almost every (C1, C2) and for all rational 2 0. Since V continuous in time, there is unique continuous (in time) function w5 (t,-) such that the identity holds for all 2 and for P-almost every (C1, Cz). The same argument yields the construction of of and wig.
27

---

Published as a conference paper at ICLR 2021
Step 3: Measurability of constituent quantities. We show that H2 (X, C2: W (t)) is S
measurable. Recall that
H2 C2; C2; W (f)) Ecc [w2 (t, C2, C2) p1 ((0, (t. ()),)))
By the existence of wf and W2 for each 2 0, there exists Borel function fr such that almost surely
82 (1, C1, C2) pI (( (1, C1,, = ft (x, wf (C1), w2 (C1, C2), ug (C2))
We recall that pt and p² are the laws of wi (Ci) and wn (C1, C2). We analyze the following:
H2 (X, (2; - f fo (X, u1, U2 wy (C2)) p¹ (duz) p² (duz)
E
= E H2 (X, C2: W (t)) p]+ +B // f (X, u2 wy (C2)) ot (duz) o² (duz)
- -2E H2 H2, Cz; W (t)) / (X, ug wa (C2)) p¹ (du1) p² (duz)
(x,
Let us evaluate the first term:
E
H2 (x,C2; C2; W (t)
=E [Ec, [5e (x, of w
w2 (C,, C2), wg (C=
@ E. [[e (X, u (C1) 1 (G, C2) wg (C2)) St (X, of (C)), w0 (2), wg (C2))
(X, (01) (11) P2 (0), 02) (A2). Pa (02) (A2))
x (X, (01) (X1) P2 (8, 02) (22). Ps (02) (Aa))
(
E
Je (X,pi (0.) (A1). pa (01, 02) (A2). pa (02) (A2))
x (X, p1 (0)) (X1) P2 (8% 0z) (A2) Ps (02) (22)) I (01 * 87)
(d)
Ez If fe (X, u1, u2 us) fe (X, uh. ulg us) pt (du1) ) pt (du() p² (duz) x² (dub) pa (dus)
where in step (a), we define C₁ to be an independent copy of C1; in step (b), we recall C₁ 011; in step (c), we recall 01, 0z ~ Unif ([0, 1J) and since Ci is independent of C₁, we have 01 * 0f almost surely; step (d) is owing to the independence property of the construction of the functions Pl, P2 and ps. We calculate the second term:
/ fe (X, u1, U2 wg (C2)) ot (du1) p² (duz)
E
u1, (X, = -f E [J (X,
u2, w (C2)) Se (X, u/j uby wg (C2))] p¹ (dui) o² (duz) p¹ (dul) o² (du)) = -f Ez Ue (X, un, u2, wa) Sa (X, u4, uhy uz)] p² (dwi) p² (duz) p² (du') r² (dub) pa (dus);
as well as the last term:
Hz (x, C2; (1) / Se (X, un, u2, w (Cz)) ot (dui) 2² (dus))
E
= Se (X, IX, wf (Cz), wg (C1, C2), wg (C2)) / Se (x, un u2 wg (C2)) p² (dui) 0² (duz)
= =f Ez [fe (X, U2, U2 uz) fe (X, 2/, u/g uz)l p² (dui) x² (duz) pt (dul) p² (dub) P² (dus).
28

---

Published as a conference paper at ICLR 2021
It is then easy to that
27
Hz (t) / fo un u2, wg (C2)) p¹ (dul) p² (duz)
E
=
That is, we have almost surely
= Note that the right-hand side is S2 -/eeauuur,,,,, and hence so is H2 (X, C2; W (t)).
H2 (2, ( (1) = / ft (X, un u2, wg (C2)) o¹ (dui) p² (duz).
Next we consider A (Z, C2; W (t)). Recall that
A H (z, G2; W (1)) 0g6 (v, ú (a; W (1)) & (Ha (z; W (1)) wy (1, (1, 4/2 (H2 (I, c2; W (I)).
Then together with the existence of wg. we have A/ (Z, C2; W (1) is Sf- meaauraabee..
Now we consider Ec2 A (z, C2; W (t)) U2 (t, C₁ C2) With the existence of wz. there exists Borel function 9r such that
A/ (z, C2; W (1) U2 (1, Ci, C2) gi (Z, w (Ci) ; w (Ci, C2), wa (C2)). Then with the same argument as the treatment of H2 (X, Cz; W (t)), one can show that
Bc, [ (2, Cz; W (1) 22 (1, C1, C2)] = / gr (Z, w (C1) u2, wa) p² (dwz) p² (dua),
which is S7 meaauuraaee..
Using these facts together with the existence of wj. W2 and wg, we have As (C2: measurable, A2 (C1, C2; W (t)) is maa-muauree and A (Ci; W (f) is -ymmesurrall.
W (+)) is S3-
Step 4: Closeness between the MF dynamics and the reduced dynamics. We shall use [W - -WI,, with the same meaning as the distance between two sets of MF parameters. Recall by Lemma 11 that ||[]|| S Kr since I|lW I K. By the same argument, by Lemma 13, we have for any L 5 T.
|| Il I Kr. Then
eus-supsup As (C2; (5)) - As (C2; (+)) < Kr || - W|, WWW- 25t
ess-sup sup A (C1, C2; (8)) - A2 (C1, (s)) < Kr || -W W|,
15
sus-supsup A (C;; - (C:WW()) s Kr ||W - Wl.-
-.W.
We have:
ess-sup
A (C.. -
ess-sup E [2 (C1, (O))
(C
S ess-sup
E A2 (2: (t) A2 (G,,
+ ess-sup
E Ag (C1,C;: (t)) A2 (g;
(a)
ess-sup
E A2
- A2 C2; W (t) S2]
ess-sup
A2 (C1, W (1) - A2 (C,, (2; W (I))
where step (a) is because A2 (C1, W (1)) is Sja-measurabb from Step 3 and S13 C S2. As such,
ess-sup Ae (t, C1, C2) - A. (C1, C2: W (t)) < 2Kr || - W
almost surely for all rational t < T. By continuity in t of both sides, the same holds for all t < T. Hence by Assumption 1,
0t U2 (1, C1, C2) - Øt
U2 (1, C1 C2)
& (t. A (2
29

---

Published as a conference paper at ICLR 2021
$ 2Kr || Wll+.
for all t S T almost surely, which leads to 1i2 (1, Ci C2) w2 C1, C2)1 S 2Kr / I - WIl, ds
almost surely, One can obtain similar results for ūj versus w1 and by versus Wy Therefore,
I| w 11, s Kr / || -wll. ds.
Since W (0) = W (0), by Gronwall's inequality, ||W - -Wl,, = 0 for all t s T. In other words, since T is arbitrary,
Wi C1) wi (t, Ci), W2 Ct, C2) W2 (t, C1 C2), Tig (t, C2) W3 (t. C2),
for all t 2 0 almost surely.
The first claim of the lemma is proven by the conclusion of Step 4 and by choosing wi पj. w T and w Wg. well as the measurability facts from Step 3. To prove the second claim, since A (Z, C2; W (t)) is S measurable and W Wil. = for all
Step 5: Concluding.
t t 0, there exists a Borel function A H+ such that
A (Z, C2; W (t)) D (Z, C2; W (0) = A (t, Z, w/ (C2))
for all 2 0 almost surely, by the same argument in Step 2. These facts, together with the dynamics of w and w2, imply that almost surely, for all t 2 0,
at w2 (t, C1, C2)
(t) Ez |A (t, Z, w/ (C2)) pi (wwi (t, (CI))
0
Ot wi wf (C1)
- (1) Bz [Bc, A (t, 7, wg (C2) u2 (1, C1, C2) s (fwf (1. w/ ( x]:
with initialization wi (0, w| (C1) = f/ (C1). Substituting the first equation into the second one, we get:
on wi w1 (C1)) = -1 (1) Ez [Ec, A
th (1, Z, ug (C2)) w/ (Ci, C2)] & ((wi (4, (Ci)) xy) x
+ + (1) 1' & (s) Ez,7 Bos [ (1, Z, wg (Cz) A (s, Z', w/ (C2))]
x x (wwi (s, wp (Ci)) x²) s (wi (1, u/ (C1)) x)) x ds.
Note that by an argument similar to Step 3,
Ec, A H+ (t, Z, wg (C2)) da (C1, C2) = / / A (t, ua) uzp² (duz) p² (du3).
which holds for all / 0 almost surely by the same argument in Step We thus obtain:
01 w (t, uz) = -1 (t) / Ez A H+ (t, Z, uz) uzyi (uw) (t, un) xx) X] p² (duz) d² (dus)
/' Eo (s) E2) [f 1 (t, Z, us) A # (s. Z', us) pa (dus)
+ + (t)
x (ww) (8, un), X')) أ (iw) (1, u1). x)) x ds.
with initialization wi (1, un) un for all un E supp (pt) and t > 0.
30

---

Published as a conference paper at ICLR 2021
An important ingredient of the proof is that the distribution of w1 (t. C1) has full support at all time t 2 O, even though we only need to assume this property at initialization = O. This key property is proven by topology argument, supported by the measurability result of Lemma 16. We remark that a similar property for two-layer networks is established in Chizat & Bach (2018) using a different topology argument.
Lemma 17. Consider the same setting as Theorem 8. For all finite time t 2 0, the support of Law (w) (1, ()) TR...
Proof. By Lemma 16, one can choose a neural embedding such that there exists Borel functions wi and A Hx for which almost surely, for all / 2 0,
Wi (t, C1) = " (t, w/ (Ci)),
A (z, Cz; W (1)) 2 (1, 2, wg (C2)),
where W (t) is the MF dynamics formed under the coupling procedure with this neuronal embedding as described in Section 3. Furthermore,
Ot wi (t, u1) = - / Ez A (t, Z, us) uppi ((wi (t, u1) x)) x] p² (dwz) p² (dus)
L/ E2,z' {
A (t, Z. us) A (s. Z', us) p² (dus)
+
x 41 (107 u1) ,x') s (mw1 (t, u1), x)) ds,
with initialization wi (0, un) = #n for all in E supp (p') and t 2 0, where Z' is an independent copy of Z. We recall from Lemma 12 that
ess-supA (1, Z, wg (C2)) = ess-sup (Z, C2; W (1) S K.,
where Kr denotes a generic constant that depends on t and is finite with finite t. Therefore, by Assumption 1, for t < T' and u1 u' E supp (p').
a wi (t, un)) 0 wi (t, u1) - ot
s Ke lun (t, u1) - wi (t, u())= + Ke L/" lui (s, 21) - wi us) ds
< Ky sup lui (1) - wt (s,u())
S5t
w (t, un
s Kr.
Applying Gronwall's lemma to the first bound:
sup Iwj (t, u1) - wi (t, ut)) s Kr lwi (0, u1) - wi (0, x1)
esT
ee lui - uill
Furthermore the second bound implies
sup S'''
|w] (1) wi (1,, S KT |t 10--11 f'..
Therefore (t, (1))) + wi (t, un) is continuous mapping on [0, T] x Rdd for an arbitrary T 2 0.
Given this continuity, we show the thesis by a topology argument. Consider the sphere gd which is a compactification of Rd We can extend wi to function M : T x Sd + Sd fixing the point at infinity, which remains a continuous map since |M (1, u1) - wn M (t, un) - M (0, u1)l S Krl.
Let M, I - Rd be defined by M, (u1) M (t, W1). We claim that M, is surjective for all finite Indeed, if M, fails to be surjective for some t, then for some p E St. M, : Sd + 8d {p} + Sd is homotopic to the constant map, but M then gives homotopy from the identity map Mo on the sphere to a constant map, which is contradiction as the sphere gd is not contractible. Hence wi (t, is surjective for all finite t. Recall that wa (t, C1) wf (t, w (Ci)) almost surely and wf (C1) has full support. Now let us assume that wa (1, C1) does not have full support some
31

---

Published as a conference paper at ICLR 2021
time t, which implies there is an open ball B in R for which (w1 (t) B) = 0. Then P (wij (1, w/ (C1)) E B) = O. Since wi (1,-) has full support, there is an open set U such that wi ut) E B for all in U. Then P (wf (C1) E U) O, contradicting the assumption that wf (Ci) has full support. Therefore wi (1, Ci) must have full support at all 2
With this, we are ready to prove Theorem 8.
Proof of Theorem 8. Recall, by Theorem 1, the solution to the MF ODEs exists uniquely, and by Lemma 17, the support of Law (wi (t, C1)) is IRd all t. By the convergence assumption, we have that for any c > O, there exists T (c) such that for all t 2 T (c) and P-almost every c1:
Ec2 E2 [ (Z, C2; W (4)) fi ((wr (4, x>]] S.
Since Law (wi (t, C1) has full support, we obtain that for in a dense subset of R'
Ec, [ A (Z, C2: W (1) p1 ((u, x)]]] S.
By continuity of + pd (1 x)), extend the above to all Rdd. Since yo is bounded,
Ec,
E [ C2; W (t)) A (Z, C2; w1 u2 wz)) s [)]]]
KE |
H (Z, C2; W (0) - A (Z, C2; Wi, W2 ]3
KE
(1 + |0g 12) [ (t, C2) u3 (C2) |+ + [0,( (02 C2) 12 C2)
+ Ty (C2) 12 (C1, C2)1 w1 (t, C1) - w1 (C1)
where the last step is by Assumption 1. Recall that the right-hand side converges to 0 as / - co. We thus obtain that for all u Rd.
Bc. B2
A H (Z, C2: w1. U2 W3) X = x), p1 ((u, (ppl
=Ec,
Ez A (Z, C2 W1 U2 wz) 41 (nu, x))]
=0,
which yields that for all u E TRd and P-almost every C
(E [ (7, 62: wi, 22, W3) | = 2] 41 (lu, ((u,)))) = =0.
Here note that by Assumption 1,
Ez [ H (2, c2; W1 T2 (3) | = =]] S K 10g (c2)1,
Ez A (Z, 02 W1, 12, 133 =]] is in I.2 (Px) for P-almost { (fu, :) u Rd] has dense span in L (Px). we have [ (Z, c2 11 W2 Wz) X = x] = 0 for Px-almost every x and P-almost every c2. and hence
and
so
ery
C2.
Since
Ëz
[0.[ (Y, 9 (X; W1, W2 xx)) X = x] vp's HH (x; W1 W2 Wy)) Wy (c2) & (H2 (z, c2: 11 (2)) = 0. Ez
(0,02: We note that our assumptions guarantee that P (wa (C2) + 0) is positive. Indeed:
In the case wg (C2) + 0 with positive probability and & (+) = 0, the conclusion is obvious.
In the case L (wf, wy, w!) < Ez [C (Y, (0))], we recall the following standard prop- erty of gradient flows:
L (w) (t, wa (t, -)), W3 (t, S L (w) (t',-) -) we (t',,:), wa (t',)),,
for 1 2 t'. In particular, setting L' = 0 and taking + 0o, it is easy to see that
L (wi, W2, w3) < L (wf. un w.) < Ez [C (Y, sp (0))
If P (wg (C2) 0) then L (W1, W2 Wa) Ez [C (Y, Pa (0)) contradiction.
32

---

Published as a conference paper at ICLR 2021
Then since Ap/2 and sls strictly non-zero, have Ez [0,C (Y, (X; W1, W2 wz)) 2| for Px--allme Px-almost every I.
In Case 1, since L convex in the second variable, for any measurable function 5(zz).
C (v. (z)) L [y, D (z; 11 U2 Wz)) 2 0gC (y, 1 (z; W1 t2 wz)) (ū (x) (z: W1, W2 Wz)) Taking expectation, get E [C (C, (X)) 2 L (w1 U2 Wz). i.. (W1, W2 wz) is global mini- mizer of L.
In Case 2, since y is function of z, obtain 0z6 (v. (x; W1, U2 03)) and hence C (y. ÿ (z; W1, W2 W3)) = 0 for Px Px-almost every x.
Finally have from Assumptions 3:
1. (W (t)) L (w1, U2 wz3) = |Ez CC (Y, 9 (XX)
U2 Tx))]l
2 K 9 (X; W (t) - w1; 02 17)]l
< KE Iw's (t. C2) W3 (C2) + 100 (C2) w2 (t, C₁ C2) - W3 (C,, C2)
+ |iñ3 (C2) ا7ا (C1, C2) [m (1, (1) - 01 (Ci)
which tends to 0 as + co. This completes the proof.
E CONVERSE FOR GLOBAL CONVERGENCE: REMARK 9
We prove converse statement for global convergence in relation with the essential supremum dition (6).
Proposition 18. Consider a neuronal embedding (2, P. (W]) (u?] of (p',p', ,p,,ppp' )'')-iiid ini-
tialization. Consider the MF limit corresponding the network (1), such that they are coupled
together by the coupling procedure in Section 3.1, under Assumptions 1, 2, &1 (-) = &2 (-) 1. Assume that Cly, 9) + oo as |0l oo for each y. Further assume that there exists Wy such that as t + -0,
Ecg [10a(1, C2) - t.a (Ca(C2) + 0.
Then the following hold:
Case I (convex loss): If L. is convex in the second variable and
lim L (W (1) = inf L (V), 1000
then it must be that
sup Ec, ea Elh
uz
c1, Ca)
+0 as + 00.
Case 2 (generic non-negutive loss): Suppose that 2z (y, 1) 0 implies L (u. 1) 0, and y = y(x) is a function of æ. If L (W (t)) + as + sc, then the same conclusion also holds.
Proof. We recall
u2 (1, = ( (Y, (f)) wy (t)
s/s (H3 (X; W (t)) &e (H2 (X, W (+)) o ((t1(t.)),
for c E Q1 02 E 22 By Assumption 1.
0
U2 (t, (1,01,00)
S KEZ (1,22 (0)) wy (2))
0t
33

---

Published as a conference paper at ICLR 2021
Note that the right-hand side is independent of o1. Since Bcs [,yy(t, C2) Wy Wiy (C2) - 0 oo, we have for some finite to I K,
Ec, [1y (C2)l & Ec, N1os(to, C2)1 + K S K,
where the last step is by Lemma 11 and Assumption 2. As such, for all sufficiently large, we have:
sup Ec, er Esh
T 22 (1, c1, (2)
< KEz [| (Y, (X; (0))) Ilws (4, C2)1
00
Ky [1024 (X; (f))1 (K + Bcs [1ss (C2)l))
KE [10gC (()))l
The proof concludes once we show that E2 [022 (X; W (+))) - as
For a fixed (z, y). let us write L (t, z) [y, 0(;; (())) and 0z 02C(y, (a;; ((t))
for brevity. Consider Case 1. We claim that if there is an increasing sequence of time t, so that lim;-soa [C),, - infg flv. 5)] O, then lim; [0,((11;=) = 0. Indeed, it suffices
to show that for any subsequence li, of 4, there exists further subsequence 4/3 such that
lime-soc | | =) O. In any subsequence li of t. using that [(ti,: 2) is convergent and the fact C(y, 0) + oo as lúl + co, we have j(æ; ((4y)) is bounded. Hence, we obtain a subsequence ti for which (æ;; ((t,g)) converges to some limit g" By continuity, we have Cly. 8*) = lim [(te,, 2) inf, Cly.9). Thus, since convex in the second variable, we have 02f(y. 1/*) = O. Thus, lim 32 =) 102[(y, 8*)l = O, as claimed. Similarly, we obtain in Case 2 that there is an increasing sequence of time t, so that lim;- [C(t,, =)] = 0, then lim, (ay =)1 O.
To show that Ez |0 (t, 2) + 0 as + oo, suffices to show that for any increasing sequence of In Case 1, we have lim L (W (t+)) = infy L (V), so lim, Ez C (t, 2) - infy C(Y, =
times f; tending infinity, there exists a subsequence t., tf such that Ez [ (t.,, []
O. Since L (t,, 2) inf, L( is nonnegative, this implies that L (t,, Z) inf ((Y converges to 0 in probability. Thus, there is a further subsequence tiy for which L (ti, Z) infy CY, y) converges to P-almost surely. By the previous claim, |026 (li,, 2) con-
verges to 0 P-almost surely. Since |22C (44, 2) is bounded P-almost surely, we obtain that Ez [ (ly: Z)] O from the bounded convergence theorem. The result in Case 2 can be established similarly.
F USEFUL TOOLS
We first present a useful concentration result. In fact, the tail bound can be improved using the argument in Feldman & Vondrak (2018), but the following simpler version sufficient for our purposes.
Lemma 19. Consider an integer 2 1 and let n. Ct. ... Cn be mutually independent random variables. Let E, and Er denote the expectations w.r. only and {cr (G))en)) only, respectively. Consider a collection of mappings (fi (filietu) which map to separable Hilbert space F. Let fa (x) = E, [J (ci,z)). Assume that for some R > 0, ISe (ci, Ss (x:) S R almost surely, then for any 8> 500, 0,
8R
no82
8R exp
no2
F
-
28
5
cxp
S
nnd
8R2
8
8RR
Proof. For brevity, let us define
Zn (x) =
Σ (fi -1
34

---

Published as a conference paper at ICLR 2021
By Theorem 21,
P P[ (a)) 2 nó(x) - 2xxp (-n0// (4R²))
and therefore,
IP 2I (4) 2 nd) S 2 exp (-no²/ (4f²))
since the right-hand side is uniform in T. Next note that, w.r.t the randomness of æ only,
Ex 2n (=)]] = Ex 2| (x) 1 2n (x) 2 nó n2/ + E, Z, (I)) (a)) < n6/2)]
E. Zz (2)) ( (x) 2 n6/2)] + nó/2.
As such, by Markov's s inequality and Sehwar'S inequality,
(E I| (I) 2 no) s P (E,, 1 (x) (4) > n6/2)] 2 n6/2)
2
E
21, (c) (x) n6/2]]
nd
2
[Zz, (c))
1/2 (IZn (+) 2 no n5
E
no
4
E (x)1))
1/2 exp (- no² (-mo²
nd
8R2
Notice that since C1 Cn are independent and fi (x) = Eg [fi (ct,
Zn (x))
>
E
- (x)) s 2mRR.
We thus get:
8R
nd²
(Ez 2, (z)] 2 no)
S
exp
v no
8R22
This proves the claim.
We state martingale concentration result, which is special case of (Pinelis, 1994, Theorem 3.5) which applies to a more general Banach space.
Theorem 20 (Concentration of martingales in Hilbert spaces.). Consider a martingale Zn Z separable Hilbert space such that IZ,, Zn-1l < R and Zo = O. Then for any t 0,
max E
Z
N t
< 2 inf exp A20
-At + ess-sup
Σ
E
[ Zu Tu 1l - 112 - Zx-1l Fo-1]
km
In particular, for any ö > 0.
n252
P max Zyl
no
< exx
RSn
2RR2
The following concentration result for i.i.d. random variables in Hilbert spaces is a corollary.
Theorem (Concentration of i..dd. sum in Hilbert spaces.). Consider n i.i.d. random variables X ..... Xn in separable Hilbert space. Suppose that there exists constant R > 0 such that XX; E [X,]] S R almost surely. Then for any 0.
no²
P
Σ
X
EX
2
I 2 exp
2R
35


Reviews:
[{"review_id": "hvymXpDkiKU", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This article is concerned with convergence guarantees of online stochastic gradient descent for a rather generic class of three layers neural networks (instead of similar analyses that treated two layers). The main results state that in a proper limit of infinite width + vanishing learning rate, the dynamics of online SGD is proven to be tracked thanks a mean-field description in the form of coupled ordinary differential equations. Once this mean-field description at disposal, the main result is obtained: in the infinite width + vanishing learning rate + infinite time (= number of training samples), the generalization error tends to it minimal value for a broad class of models and losses (not necessarily convex, which is a novelty of the work) as well as generic data distribution.\n\nOverall this paper is very well written, enjoyable to read despite the technicality of the results, and understandable even for non-specialists of this line of works (like myself). I did not check the appendices and proofs. In the main part there are no typos, and I have no main concerns to bring about. Yet: I would find interesting to know more details about the differences with the refs Nguyen (2019); Araújo et al. (2019); Sirignano & Spiliopoulos (2019); that is not clear.  Also I would find useful to have some hints about the meaning of the (trained third layer) hypothesis in Theorem 8. Finally I find a bit surprising that there are nor restrictions whatsoever on the data distribution (or I missed that). The authors may comment on that in the final version.\n\nI recommend publication. Even if I'm not a specialist, it is obvious that the authors made a big effort of redaction, that the results are very solid, the proof technique seems original and requires less assumptions than previous works (I liked very much the \"idea of proof\" part). I have very few doubts about the quality of the paper despite I did not read the proof details, and the fact that I'm not aware of the literature in this specific field.", "extra_scores": {}, "extra_text": {}}, {"review_id": "T6WezFEzdFy", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the behavior of a 3-layer fully connected network when the width of the network is large. The authors define a mean field regime and prove that the behavior of the network under stochastic gradient descent converges to this mean field regime (for any finite time horizon).  This result complements nicely previous works like (Nguyen 2019) that contain an informal derivation of the mean field regime. This transient regime is complemented by a long-term analysis under quite restrictive assumptions (which imply essentially than the mean field regime always converge to the minimizer of the loss function).\n\nI did not check all details of the proof but the approach seems mathematically sound. Once the model is defined, the proof for the finite regime is relatively classical: it relies on Martingale concentration plus Gronwall's lemma. Yet, as always, the devil being in the details and defining the right model and using the right notations is a difficult task. The result of the stationnary regime seems also reasonable but I must admit that the proof of the infinite horizon is not really clear to me. I would have appreciated more pedagogical effort from the authors.\n\nTo summarize, the paper seems a nice theoretical contribution. Yet, to me one thing that this paper is missing is an explanation or illustration of how useful is their result to understand the behavior of deep neural networks.\n\nThat being said, one major concern that I have about the paper is the link with https://arxiv.org/pdf/2001.11443.pdf The arXiv paper considers a very similar model (but more general as it considers L layers instead of 3). It uses almost the exact same notations and the same structure (overall paper and proofs). I think that the authors should clarify the link between the two papers. Also, if I can admit that the present paper is a resubmission of the arXiv paper, I do not understand why does the current paper focus on 3 layers and not the more general model of the arXiv paper.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "FcX6yNDNQCl", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies some theoretical properties of three-layer neural networks (NNs) under the mean-field (MF) regime. The authors proposed neuronal embedding in order to study large-width neural networks. Then, the quantitative relation between finite-width NN and the MF limit was clarified. The global convergence of the continuous limit of the stochastic gradient descent (SGD) was proved without assuming the convexity of the loss function. The global convergence of the MF limit is used to establish the optimization efficiency of the neural network with SGD. \n\nThe problem considered in this paper is important. The definitions and the statement of theorems are clearly described. In order to prove the global convergence, the authors assumed the uniform approximation property rather than the convexity of the loss function. This approach is interesting. Though I'm not very familiar with the MF regime, the high-level idea to prove the theorem is well-written. I read some proofs in the appendix, and I found that the description is accessible to a wide range of audiences. Overall, this paper is thought to provide a promising idea to analyze not only three-layer NNs but deep NNs.\n\nSome comments are shown below:\n- Neuronal ensemble is introduced to analyze the dynamics of the MF limit. Apart from the MF regime, is there any similar idea of the neuronal ensemble? Showing some references would be beneficial to readers.\n- The constant K appears in the upper bound in Theorem 3. Showing a more concrete expression of K is informative. What is the typical K in this case? \n- In Theorem 8 and Corollary 10, the global convergence was proved. If the convexity was also assumed in addition to the assumptions in the theorems, is it possible to derive the convergence rate? \n- In the paper, the global convergence of three-layer NNs was analyzed. What is the main obstacle to investigating the global convergence of multi-layer NNs or deep NNs according to the idea proposed in this paper?\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "8O73d2rOD6N", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**: \nAnalysis of neural networks in the mean field regime gains more and more attention as it helps to study the dynamics in the wide regime. The paper extends the recent studies and provides global convergence guarantees for an unregularized feedforward three-layer NN. This is the first time global convergence is established for neural networks of more than two layers in the mean-field regime. \n\nI find the writing a bit chaotic and overdosed with notations. But, overall, I think the results are significant and will help to extend the line of research in the mean-field regime applied to neural networks. \n\n**Questions**: \n- In the paper, it is said several times that the convergence result “does not rely critically on convexity”. What do you mean by “critically”? You still assume the convexity. Do you mean that it can be easily relaxed in future works? I think it should be better stated throughout the paper.  \n\n**Minor suggestions**:\n- Mean field -> mean-field \n- Section 2.1 “k” is introduced at the very end. Maybe it would be better to add in the beginning, e.g. “the following network at time k”? \n- I would say “W(k) consist of the weights …” instead of “W(k) is the weight with ….” \n- It perturbs me a bit that the difference between NN notations and MF is its boldness. Boldness usually means a vector while in authors’ notations $\\mathbf{w}_2$ is an element. But I guess with superscripts or symbols like $\\hat w$ it would be heavier in notations... \n- Section 2.2. I would like to see the description of \\Omega, F and P  in the beginning. Like “Given a NN \\Omega_i would be a space  of…”. Just to connect it from the start and to easify the following read. \n- Definition 2. “The following hold” -> holds \n- Section 4.2 “where V a set…” -> is \n- “Helps avoiding” -> helps to avoid\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper provides a global convergence guarantee for feedforward three-layer networks trained with SGD in the MF regime. By introducing the novel concept of neuronal embedding of a random initialization procedure, SGD trajectories of large-width networks  are shown to be well approximated by the MF limit, a continuous-time infinite-width limit (Theorem 3). Furthermore, under some additional assumptions the MF limit is shown to converge to the global optimum when the loss is convex (Theorem 8, case 1) and for a generic loss when $y=y(x)$ is a deterministic function of input $x$ (Theorem 8, case 2). The global convergence guarantee presented in this paper is based on less restrictive assumptions compared with existing studies. All the reviewers rated this paper quite positively, with less confidence however, seemingly because of mathematical thickness of the proofs. Although the reviewers did not manage to check every detail of the proofs, they agreed that the reasoning seems mathematically sound as far as they can tell. The authors response adequately addressed minor concerns raised by the reviewers. I am thus glad to recommend acceptance of this paper.

Pros:
- Introduces the idea of a neuronal embedding, which allows establishing relation between SGD on large-width three-layer networks and its MF limit in a quantitative way with a less restrictive setting.
- Provides a global convergence guarantee under the iid initialization, in the sense that if the MF limit converges it attains the global optimum.
- Shows that the global convergence guarantee does not require convexity of the loss when a deterministic function is to be learned.

In particular, the uniform approximation property, rather than the convexity of the loss, plays a crucial role in proving the  global convergence guarantee (it allows translation of the vanishing gradient in expectation at convergence into the almost-sure vanishing gradient), which is a quite original contribution of this paper.
Decision: Accept (Oral)