Mean-field analysis for heavy ball methods: Dropout-stability, connectivity, and global convergence
Authors: Diyuan Wu
Vyacheslav Kungurtsev
Marco Mondelli
Venue: Accepted by TMLR ()

Abstract: The stochastic heavy ball method (SHB), also known as stochastic gradient descent (SGD) with Polyak's momentum, is widely used in training neural networks. However, despite the remarkable success of such algorithm in practice, its theoretical characterization remains limited. In this paper, we focus on neural networks with two and three layers and provide a rigorous understanding of the properties of the solutions found by SHB: \emph{(i)} stability after dropping out part of the neurons, \emph{(ii)} connectivity along a low-loss path, and \emph{(iii)} convergence to the global optimum.
To achieve this goal, we take a mean-field view and relate the SHB dynamics to a certain partial differential equation in the limit of large network widths. This mean-field perspective has inspired a recent line of work focusing on SGD while, in contrast, our paper considers an algorithm with momentum. More specifically, after proving existence and uniqueness of the limit differential equations, we show convergence to the global optimum and give a quantitative bound between the mean-field limit and the SHB dynamics of a finite-width network. Armed with this last bound, we are able to establish the dropout-stability and connectivity of SHB solutions. 

Published in Transactions on Machine Learning Research (02/2023)
Mean-field analysis for heavy ball methods: Dropout- stability, connectivity, and global convergence
Diyuan Wu
diyuan, diyuan.wn@ista.com www.sta.acccacccco
Institute of Science and Technology Austria (ISTA)
Vyacheslav Kungurtsev
kangunya&fee.comd.cc
Caech Technical University
Marco Mondelli
maroomoondelllhhttaacccc
Institute of Science and Technology Austria (ISTA)
Reviewed on OpenReview: https: // openreview. net/ forum? id= g2na3I10ft
Abstract
The stochastic heavy ball method (SHB), also known as stochastic gradient descent. (SGD) with Polyak's momentum, is widely used in training neural networks. However, despite the remarkable success of such algorithm in practice, its theoretical characterization remains limited. In this paper, we focus on neural networks with two and three layers and provide rigorous understanding of the properties of the solutions found by SHB: (i) stability after dropping out part of the neurons, (ti) connectivity along a low-loss path, and (ii) convergence to the global optimum. To achiove this goal, we take a mcan-field view and relate the SHB dynamics to a certain partial differential equation in the limit of large network widths. This mcan-field perspective has inspired a recent line of work focusing on SGD while, in contrast, our paper considers an algorithm with momentum. More specifically, after proving existence and uniqueness of the limit differential equations, show convergence to the global optimum and give quantitative bound between the mean-field limit and the SHB dynamics of finite-width network. Armed with this last bound, we are able to establish the dropout-stabiilty and connectivity of SHB solutions.
1 Introduction
Neural networks are one of the most popular modeling tools in machine learning tasks. In practice, they can be effectively trained by gradient-based methods, and are usually overparamcterized. However, despite the empirical success of various training algorithms, it is still not well understood why such algorithms have good convergence properties, given that the optimization landscape is known to be highly non-convex and to contain spurious local minima (Auer et al., 1996; Safran & Shamir, 2018; Yun et al., 2019). A popular line of work starting from (Mei et al., 2018; Chizat Bach, 2018; Rotskoff & Spiliopoulos, 2020a) has proposed a new regime to analyze the behavior of stochastic gradient descent (SGD), namely, the mean-field rogime. The idca that, as the number of neurons of the network grows, the SGD training dynamics converges to the solution of certain Wasserstein gradient flow. This perspective has facilitated the study of architectures with multiple layers (Araújo et al., 2019; Lu et al., 2020; Nguyen & Pham, 2020; Fang et al., 2021), and it has provided path to rigorously understand a number of properties, including convergence towards a global optimum (Mei et al., 2018; Chizat & Bach, 2018; Javanmard et al., 2020; Pham & Nguyen, 2021a), dropout-stability and connectivity (Shevchenko & Mondelli, 2020), and implicit bias (Williams et al., 2019; Chizat & Bach, 2020; Shevchenko et al., 2022).
Vanden-Eijjden. 2018; Sirignano
Optimization with momentum, e.g., the heavy ball method (Polyak, 1964) or Adam (Kingma Ba, 2015), is widely used in practice (Sutskever et al., 2013). However, all the aforementioned works consider the vanilla

---

Published in Transactions on Machine Learning Research (02/2023)
SGD algorithm and, in general, the theoretical understanding of algorithms with momentum has lagged behind. To address this gap, the recent paper by Krichene et al. (2020) defines a mean-field limit for the stochastic heavy ball (SHB) method also known as SGD with Polyak's momentum in a two-layer setup. In particular, the convergence to the mean-ficld limit is proved, as well as that the solution of the mean-field field equation approaches a global oplimum in the large time limit. However, Krichene et al. (2020) leave as an open problem finding a quantitative bound between the infinite-width limit and the finite-width neural network, and the analysis is restricted to two-layer networks.
In this paper, define mean-field limit for the heavy ball method applied to two-layer and three-layer networks. We show global convergence in the threc-layer setting, and give quantitative bounds for networks with finite widths. This last result opens the way to providing a rigorous understanding of effects commonly observed in practice, such as the connectivity of solutions via low-loss paths (Goodfellow & Vinyals, 2015; Garipov et al., 2018; Draxler et al., 2018; Entezari et al., 2022). More specifically, our main contributions can be summarized as follows:
1. We show existence and uniqueness of the mean-field differential equations capturing the SHB training for two-layer and three-layer networks (Theorem 4.1).
2. We give non-asymplolic convergence results of the SHB dynamics of a finite-idtth neural network to the corresponding mean-field limit (Theorem 5.1 for two layers, and Theorem 5.2 for three layers). Our bounds are dimension-free in the sense that the layer widths are not required to scale with the input dimension.
3.
We discuss how SHB solutions can be connected via a simple piece-wise linear path, along which the increase in loss vanishes as the width of the network grows (Section 6). This is a consequence of the stability against dropout displayed by the parameters found by SHB, which in turn follows from Theorems 1-5.2.
4. Finally, by exploiting a universal approximation property enjoyed by the activation function, we prove a global convergence result for three-layer networks, under certain assumptions on the mode of convergence of the dynamics (Theorem 7.2).
Organization of the paper. After discussing the related work in Section 2, the details concerning the network architecture and SHB training are presented in Section 3. In Section 4, we define the mean-field limit for both two-layer and three-layer networks, and we show that such limits exist and are unique. Next, in Section 5, we prove our quantitative convergence bounds between the discrete SHB dynamics and the mean- field limit. As an application of these bounds, we discuss the dropout-stability and connectivity displayed by the SHB solutions in Section 6. In Section 7, we prove global convergence for the three-layer mean-field differential equation. In Section 8, we show the results of some numerical experiments that illustrate the dropout stability of the SHB solution. Finally, we discuss future dircctions in Section 9.
2 Related work
Mean-field analysis for two-layer networks. A quantitative convergence result of the SGD dynamics towards a mean-field limit was presented in (Mei et al., 2018). This bound was refined to be independent of the input dimension in (Mei al., 2019), which also considers a setting where both layers are trained. Our approach extends this line of work to the heavy-ball method. Because of the presence of momentum, in this case the mean-field limit is described by a second-order differential equation (instead of the first-order one capturing the SGD dynamics). The mean-field limit for heavy-ball methods was first considered in (Krichene et al., 2020), which deals with a setting regularized by noise and does not provide quantitative bounds. The recent work by Schuh (2022) gives non-asymplotic guarantees, but the argument crucially relies on the presence of additive Gaussian noise. Global optimality of SGD was proved by Chizat & Bach (2018) in the noiscless setting, using the homogeneity of the activation and under an additional convergence assumption. This assumption is not needed in the noisy case, as the mean-field dynamics Wasserstein gradient flow for strongly convex free-energy functional (Mei et al., 2019). Convergence rates have been
2

---

Published in Transactions on Machine Learning Research (02/2023)
obtained by exploiting displacement convexity (Javanmard et al., 2020) or via the log-Sobolev constant of the stationary measure (Chizat, 2022; Nitanda et al., 2022). For heavy ball methods, although the limiting dynamics does not yield a gradient flow structure, the rate of convergence has been studied in the noisy case (Hu et al., 2019; Kazeykina et al., 2020; Schuh, 2022). For a noiscless dynamics, the explicit characterization of the convergence rale an open problem.
Mean-field analysis for multi-layers networks. The case of neural networks with more than two layers presents additional challenges in regards to defining and characterizing mean-field limit. Here, we follow the "neuronal embedding" framework (Nguyen & Pham, 2020; Pham & Nguyen, 2021a) to define mean- field limit for the heavy ball method in a threc-layer sctting. The key idea is to introduce a certain fized product probability space at initialization, and regard the weights as deterministic functions which evolve during training over this fixed probability space. Other approaches have been proposed as well: Sirignano & Spiliopoulos (2020a) give asymptotic results for three-layer networks, which require taking the largh-width limits in a specific sequential order; Fang et al. (2021) focus on the dynamics of the features, rather than on that of the weights; Araújo et al. (2019) consider neural networks with more than four layers, with un- trained first and last layer, and their result depends on the degeneracy phenomenon of middle layers under i.i.d. initialization. This degeneracy consists in the fact that all the weights in the middle layers (excopt the second and second-to-last remain i..d.., hence the corresponding neurons compute the same function. Pham & Nguyen (2021a) show global optimality of the mean-field dynamics in the noiseless case in the large time limit, under a convergence assumption in the same spirit of Chizat & Bach (2018). While the global convergence in Chizat & Bach (2018) crucially relies on the homogencity of the activation, Pham & Nguyen (2021a) exploit universal approximation property, similarly to Tu et al. (2020). For deeper networks, key technical hurdle is due to the degeneracy phenomenon mentioned above, which prevents universal approximation. To address this issue, Nguyen & Pham (2020) employ different initialization, and Lu et al. (2020); Fang et al. (2021) use skip connections.
For multi-layer networks, it is also possible to define a mcan-field limit beyond the parametcrization con- sidered the aforementioned works, see Hajjar et al. (2021); Chen et al. (2022). In particular, Hajjar et al. (2021) consider the regime of so called integrable parameterizations, which is a modification of the mazimum-update parameterization proposed by Yang & Hu (2021). Furthermore, Chen et al. (2022) provide convergence result for three-layer networks, where the first layer is random and fixed.
Additional related work on optimization with momentum. A recent line of work has considered training neural networks in the neural tangent kernel (NTK) regime (or lazy regime) (Jacot et al., 2018; Chizat et al., 2019). Here, the weights of the neural networks stay close to their initialization, so that the network is well approximated by its linearization and the convergence related to the smallest eigenvalue of the NTK, see e.g. (Du al., 2019; Allen-Zhu et al., 2019; Montanari & Zhong, 2020; Bombari et al., 2022) and references therein. This type of analysis has been adapted to the heavy ball method by Wang et al. (2021) and to other adaptive methods by Wu et al. (2019). However, we remark that, unlike in the mean-field regime, neural networks are unable to perform feature learning in the NTK regime (Yang & Hu, 2021). Beyond the training of neural networks, momentum-bassd stochastic gradient descent algorithms and their contimuuss variants have been widely studied in optimization: the continuous limit of these methods is studied in (Su et al., 2014; Wibisono et al., 2016; Shi et al., 2021), and such dynamics are known to be closely related to sampling methods such as MCMC (Ma al., 2021).
3 Problem setup
3.1 Network architecture
We consider neural networks with two and three layers. In the two-layer case, the network has n neurons and input æ E RD;
(o, =) = ww)
j [n),
n f(x; W)
(1)
Σ
(2)
n 1=1
3

---

Published in Transactions on Machine Learning Research (02/2023)
Here, we use the short-hand [n] := {1, (I,...,, n and, for j E [n], the parameters of the j-th neuron are denoted by {(j) (wn (f), wz())), with wi( E RD and w²(i) E R. In the three-layer case, the network has ny and neurons in the first and second hidden layer, respectively:
H1(1.1: f1: W) (i(i))ax h [m1l.
1
H2; W) = Σ
u2(it.b)o, W)), ji E [ma) j2 [m2].
n1
(2)
n=
1
f(x; W)
wy ( (2)
112
Here, æ RD, wi (ji) RD, (2( j2) E R and wa (ja) E R. In both cases, we use W to denote the collection of all parameters: in two-layer case W e [=((+)), and in three-layer case W R tnnguttn,
3.2 Training algorithm
Our training data z = (x, y) is generated i.idd. from a distribution D. The neural network is trained to minimize the population risk function R(W)- = Ey[R(,, W))] via the following one-pass stochastic heavy ball (SHB) method:
(k--11 W(k 1) (A)) (k) W(k - 1)) n ((y()) W())))
(3)
where we use V R(y(k), (k)(; W)) denote the scaled gradient, and the scaling factors for each pa- rameter are specified below. This is one-pass method in the sense that, at each step, sample new data point zkk) independent from the previous ones. The requirements tained in (A1) (scc Assumption 3.1) and (B1) (sce Assumption 3.3) for networks with two and three layers, respectively. In particular, we require R(y 9) to be different iable with respect to its second argument and to have a bounded derivative. We also remark that, for the convergence to the mean field limit (Theorem 5.1 and 5.2), the convexity the loss is not required. In contrast, to obtain the global convergence result of Theorem 7.2, we need additionally assume that Rly, g) is convex in the second argument, as mentioned in the statement of the theorem.
the loss function R(y, 9) are con-
Let be a constant which does not depend on the network width or the step size of gradient descent. Then, in order to define continuous- ODE for the heavy ball method, we pick 3 = (1 - -E) and n = =². so the one-pass SHB method can be equivalently written as follows:
W(k + 1) = (A) r(k), r rkk = (1- 71 W - 1)) 5 ((y(),,, W (+))).
(4)
A similar formulation is common in the literature, see for example (Shi et al., 2021, Eq. 1.2) The corresponding continuous ODE, also studied in (Krichene et al., 2020, equation 6), is given by
a (t) = r(t),
a,r(t)= = -yr(t - Vw R(W)))
(5)
where we recall that (w(()))) denotes the scaled gradient. We remark that there are different ways to derive a contimuons dynamics from (3), and (4) is obtained by applying the Euler scheme based on the second-order Taylor expansion. The corresponding ODE (5) is denoted as the low-resolution ODE by Shi et al. (2021). It an interesting and challenging task to analyze other types of ODEs associated to the SHB method, for example the high-resolution ODE proposed by Shi al. (2021). We leave this to future works. We also remark that a similar formulation of the contimuous counterpart of heavy ball methods with fixed momentum is considered in (Kovachki & Stuart, 2021; Kunin el al., 2021).
Note that in (Shi et al. 2021, Eq. 1.2), B 1-ne while here we let 8 = 1 7e. The two choices are basically the same Ioye'
when is small.
4

---

Published in Transactions on Machine Learning Research (02/2023)
We conclude this part by discussing the scaling factors for the gradient in (4). In the two-layer case, we have
V R(y W)) = ((A1" (ær JEE))) (4 (æ, g))e))))
(6)
where
A W (x, 1: W) nd By,, ((; W)) (W) W)),
(7)
A (x, W) = nV By. f(x; W)) - (W) ())
In the three-layer case, we have
Vw R(v. I(x W)) = ((AI" (i; W)) (42 Ji, s (A (x, J2; W) Walsenn)), (8)
where
A (z. W) n20 Rly. 02R(y, f(;
A Jz; W) == n20 w/ R(v. (a;; W)) 02 R. f(; ()) (H W)),
A4 (2,51,32 W) mmmn2 R[ f (a; W)) 4 (z (H.
(9)
(z, W) == m,0H (z.) Rly, f(æ; W))
A (x jz;
W)).
12=1
A (æ, W) = n1 V wiG R(y f(: W)) = A (x, jii W).
One can interpret this as indicating that in the two-layer case, the scaling factor is n. the thyee-layer case, the scaling factor is n2 for the third layer, n1 n2 for the second layer, and nı for the first layer. This choice of the scaling factors ensures that each component of the gradients is of order 1 (i.e., independent of the layer widths 71, 122. In the following sections, will use wi (t,j) or wn (k, to represent the weights at time L or time step k. The same notation also applies to w2 Wg.
3.3 Assumptions
Our assumptions are rather standard in the mean-field literature and appear e.g. in (Mei et al., 2018; 2019; Nguyen & Pham, 2020; Pham & Nguyen, 2021a). We write such assumptions separately for networks with two and three layers.
Assumption 3.1. We make the following assumptions for the training of two-layer network:
(A1) (Boundedness) There exists a universal constant K > 0 such that ||a IIo'l ||o" S K. The data distribution such that, almost surely, lyl. ||a/||2 S K. Furthermore, [0,R(y, fax; W))I is K --Lipschitz continuous in f(x; W) and K -boundedd for any W.
(A2) (Initialization of weights) At initialization, wi (0, w2, iid. Po. where is such that wi (0, is S.-sub-Gaussiay and [w2(0,,)l S K almost surely.
(A3) (Initialization of momentum) We assume that r(0) 0.
In order to state the assumption for thyee-layer neural network, we first define the notion of i.id. initial- ization.
Definition 3.2. We say that the three-layer neural network (2) has i.i.d. inilialization from pà x p x pă, if:
wi (0.51) kid. ph. w2(0..1.J2) iidd po. (2) i.d. pa. ju E [nu]. j2 E [nz]
(10)
In short, i.i.d. initialization means both cross-layer and in-layer independence. Assumption 3.3. We make the following assumptions for the training of three-layer network:
5

---

Published in Transactions on Machine Learning Research (02/2023)
(B1) (Boundedness) There crists universal constont K 0 such that ||aa]læe, ||o1llso || |02|| ||4 |oo 10% S K. The data distribulion 'D is such that lyl, ||æ||2 S K almost surely. Furthermore, d2(x) & 0 for all x, [0gR(y, f(; W)) is K -ippshhitz continuous with respect to the second argument and K --bounded for any W.
(B2) (Initialization of weights) wi (0. f2), (2) have an ii.dd initialization from pb ph x P.. Furthermore, w1(0, (1) is -ub Gaus-inns and (1, (0,(0,,2)) f2)) S K almost surely.
(B3) (Initialization of momentum) We assume that r(0) = O.
In the initialization of two-layer networks, Po is not required to be product measure, that is, wi (O, j2) are not required to be independent from each other. In other words, we do not assume cross-layer independence, but only in-layer independence. However, in the throc-layer case, we need to as- sume both cross-layer independence and in-layer independence, and it turns out later that the cross-layer independence is critical for proving our global convergence result.
The requirements above hold e.g. for tanh or sigmoid activation function, and logistic or Huber loss. The assumption that ayR(y, f(x; W)) is K-bounded does not hold for the square loss. However, expect the same results proved in this paper to hold also for the square loss, provided that the assumptions are modified as in Mei et al. (2019) (for the two-layer case) and in Nguyen & Pham (2020) (for the three-layer case). In lact, it suffices that 0g Ryy, (a; W)) is bounded with high probability, and then the arguments are similar.
Finally, we remark that the boundedness of the initialization w²(0,,) and w2 J1. ws(O,, (2) is purely to simplify the proof, which can be generalized to -a-sub-Ganssian Gaussian initialization. In particular, one can bound the sub-Gaussian norm of the weights, and then the absolute value of the weights will also be bounded with high probability.
4 Derivation of the mean-field limit
Two-layer networks. The idea is that the output of the network can be viewed an expectation over the empirical distribution of the weights, that is:
f(x:W)=
= n Σ w (2))) = -Ea=so ox;;0 (x;
where W (0() j E [n]} and 0 (wi (2)))w2())
w2()). Furthermore, we define o²(x; 0()))
(2())o(w1())" x) and po 1 - 8og)- Thus, the evolution of the parameters Ot) according to (5)
can be viewed as the evolution of (t) according to a certain distributional dynamics induced by (5).
Since we assume i.i.d. initialization, as the number of neurons n + oo, we expect that (0) - Po- In this limit, the distributional dynamics induced by (5), can be described by a certain ODE, with initial condition Pa. Let
f(z; p) = -Enwpo (Z;
R(z: p) =
R(p) E.R(y.
(Z2 0) 5R(z,P)
((: = 5R(p)
(11)
(0),
6p (0) Eg ((, p)..
8p
Here, f(:p): : IR + R; ((::p) : R D+1 + R, R()) : (2 + IF, where (2 denotes the space of
probability measures on D+1 with finite second moment; (:::))) RD1 x XRD+1 R; (:::p) DD1 - IR.
Then, we define the mean-field ODE associated to the heavy ball method as
d0(t) = r(tdtt
dr(t) (- V ( (1)) dt.
(12)
Three-layer networks. Among the various approaches to define a mean-field limit for multi- networks, follow the "neuronal embedding" framework proposed in (Nguyen & Pham, 2020; Pham & Nguyen, 2021a)
6

---

Published in Transactions on Machine Learning Research (02/2023)
to capture the dynamics of SGD training. We briefly summarize the key ideas, and then discuss how to obtain the mean-field limit for the stochastic heavy ball method.
Consider a three-layer neural network of the form (2) with weights wi ji), , w2 (0, j1. (0,)), obtained via an i.idd initialization from P1 X P2 x P3, according to Definition 3.2. Then, in (Pham & Nguyen, 2021a, Proposition 7), is proved that there exists a product probability space (Sl x f22 F1 x F2, PI x P2) and functions wi (0 -) Q1 + R2, w2(0,,,) -.) : 22
f22 R. w3(0,-) S22 R such that, for any ai n2,
fwi(0, Ci(i)) (2 (ib), C2(12)). (0 C2(i2)), for f1 E [maldde E [n2]]
(wi (0, Ji), w2 J1, 12), wz (0, for ji E [mil,j2 E [na]]:
where d denotes equality in distribution, C1 didd Pi, C2(22) i.d. Pg, for ji E [m1]. [n2] Here, the tuple {(S, Fi F2, x P2), wn (0, w00-), (0, wa wa(0; (0.-)) is called neuronal embedding. With slight abuse of notation, we use W1, W2, W3 to denote also the functions wy (0,-). w2(0,...), Wy(O,-). In later sections, we refer to the functions when we write wi (0,c))) w2(0, cI, 02), w3(0, o2), while we refer to the weights of the neural network when we write w1 (0. wy f2. wy (0.
At this point, we are ready to define the mean-field ODE for the heavy ball method. This ODE tracks the functions w wi(0,)) -), w2( w,,00,))))))))))) (0,-).
dwa(t, c2) = rt(t, dt, dwz(t, c1, 02) = rt,,, c dt, (t. 01) = (t, dt,
dr) = - - E,4, (z, c2; (())) dt, dr2(t,01,02 (2, E W (z, 02; 02; (())) dri(t,c1) (->r((tt, - - (z, ci (t))) dt,
(13)
where C E E0,00 C2 E S2 are dummy variables and W(L) refers to the collection of weights (w)((), (2)()), wl)..). The output of the neural network under the mean-field limit (13) is described via the following forward pass:
Hi e1 (()) w (t, x, ci Es,
Hz 02; (t)) = BO (t, (Z) f(x; W (f) = WCc-,,,a(tt,, C2) (H2( Cz; (())).
C.
E Ω2
(14)
Furthermore, the quantities AY A AW appearing in (13) are described via the backward pass:
A (z, c2: (t)) = 02 (H))) W(1)),
A (z, oz: (t)) := 02:(y: f(x; (()))
A 02 W()) := A (z, cz (H))
(15)
A o: (tt) == Ec, (z. C2 ((t)) c1 C2) W())),
(z,0; W (t)) := - (()))::(tt)z.
For convenience, we will also use the lighter notation Hi (I, x, ci), Ha(t, 2, 02), A (t, (1,,,))) AW (t, z, cg) (t 2,),), AW (t, e1) to denote the quantities H1 ct (()), H2(2, of (t)), A (z,02; W()), (2 (()), AW (z, 02; W()), ((),,,;;(())) AW (z,01; W()), respectively.
We note that the neuronal embedding framework can recover the distributional dynamics for two-layer networks as a special case (see (Nguyen & Pham, 2020, Corollary 22) for more details). It is also possible to define the mean-field limit for three-layer neural networks directly as a distributional dynamics (Araújo et al., 2019; Sirignano Spiliopoulos, 2020a), although this approach may require additional assumptions (namely, first and last layer not trained, see Araújo al. (2019)).
At this point, we are ready to prove the existence and uniqueness of the mean-field differential equations. Thcorem 4.1. For any < 5o, we have: (i) Under Assumption 3.1, there crists a unique solution of the mean-field PDE (12).
(ii) Under Assumption 3.3, there exists unique solution of the mean-field ODE (13).
7

---

Published in Transactions on Machine Learning Research (02/2023)
The proof of Theorem 4.1 follows from the analysis of Picard type of iteration (Sznitman, 1991), and it is deferred to Appendix B. We note that Krichene et al. (2020) consider a mean-field limit for two-layer networks in which an additional Brownian motion is applied to the momentum term. This extra noise term together with the additional assumption (A5) (sce p. 8 of Krichene ct al. (2020)) allows them to prove the existence and uniqueness of the mean-field limit. for any E [0, x)]. This includes the existence and uniqueness of the limiting point (for = oo). In contrast, Theorem 4.1 proves the existence and uniqueness of the solution for any finite and we do not have guarantees on the limiting point.
5
Convergence to the mean-field limit
5.1 Two-layer networks
We recall that the mean-field ODE is defined in (12), and the SHB dynamics can be expressed as
SHBB (k + 1,5) = SSHB () + (1 ye) (OSBB(( @SHB 1.j)) - (o SSHB (k. f); (k))
(16)
where (k,J) denotes the parameter associated to the j-th neuron at step k. T is defined in (11), and (k) = Σラ= 5gsum(k.)) demotes the empirical distribution of the parameters (0SHB (k.j)sen We couple the mean-field ODE (12) and the SHB dynamics (16), in the sense that they share the same initializatin:: ((0) (0) and (0H)( i.d.d. (0). Let us define the following distance metric that measures the difference between the mean-field dynamics and the SHB dynamics:
Dj (0,0SHHB max sup JEin sE[0.T] |OSHB(LL/E] - 0(1)l22
(17)
Theorem 5.1. Let Assumption 3.1 hold. Consider the mean-field ODE (12), the SHB dynamics (16) and the distance metric (17). Then, with probability at least 1 1- cxp((822),
(w
[VVogy +0) Vn + V log +0)
27(00SHHB K(TTT)
(18)
where K(),T) is constant depending only on T.
We remark that the RHS of (18) is also an upper bound on suptejo,r] Ws(p (Smn(())))))))) which follows directly from the definition of the Wasscrstein W2 distance.
Theorem 5.1 gives a quantitative characterization of the approximation error between the mean-field limit and the stochastic heavy ball dynamics. In particular, it shows that this approximation error scales roughly vlog n/n + Ve(D logn), i.e., it vanishes as the number of neurons grows large and the step size = o[LvVD + n). We remark that the bound in (18) is dimension-free in the sense that does not need to scale with the input dimension D. This is aligned with the bound provided for SGD by Mei et al. (2019) which is also dimensionn-fee.NNee deviation theory. This implies that it is not possible to improve the corresponding dropout stability and connectivity guarantees (sce Scction 6), in terms of the number of neurons mouron//ipput dimension. This means that the behavior of SGD and SHB is similar, as both algorithms are optimal in this regard. The constant K(y T) scales rather poorly in T, i.c., ((),T)= O(e*)) ). This is common shortcoming for "propagation of chaos arguments: for example, in Mei et al. (2018; 2019), the scaling of the bound in T' is O(e"). An interesting open problem is to improve such dependence, e.g., by using ideas from Schuh (2022).
that the order of the upper bound O ( is tight due to large
To the best of our knowledge, Thcorem 5.1 provides the first consistency guarantee of the mean-field limit (12). In the noisy case, a similar (although non-quantitative guarantee proved by Krichene et al. (2020). The injection of noise in the dynamics often simplifies the analysis and it allows to prove stronger results
8

---

Published in Transactions on Machine Learning Research (02/2023)
(e.g., existence and uniqueness of the limit). However, the noiseless dynamics is particularly interesting, since Brownian noise is typically not added in practice while training.
At the technical level, in order to deal with the second order dynamics arising from the heavy ball method, we exploit a second order Gronwall's lemma (ef. Lemma F.3) and use the Euler scheme to discretize the continuous dynamics. The detailed proof is provided Appendix C.
5.2
Three-layer networks
We recall that the mean-field ODE is defined in (13), and the SHB dynamics can be expressed as
SHB + 1,52 = w3 SHB (0, (1 70) SHB kk. j2) w3 SHB 1, 12)) <2 (z(k), (B)BBB()))
3
SHB (k + 1,51, = w2 SHB (0, (2) + 7e) 202 (k k, - w2 SHB (k 1,51,
SHB
2
- = (z(K)). W (k)),
w
SHB (l: + 1,51) = w SHB SHB (0, + (1 - re) s
SHB (k -
w 1,,1)) e²D'"" (z(+)) W SHB (k)),
(19)
where (HB =
Elm.l: (w) SHB (w (k, (t,),,,2)) Eim)) data point sampled at time step k and AW AW A are defined in (9).
SHB (w (k. z(k) is the
Before stating our result, let us discuss how to couple the SHB dynamics (19) and the mean-field ODE (13).
First, sample finite neural network w.r.t. the neuronal embedding, i.c., (i()) Lid. P1, C2(22) Ltd P2: for fi E [m1l.f2 E [ma] and wi (0, (0,)) wy( Given w1(0,,), -), ;) ws( let the mean-field ODE (13) evolve, thus obtaining wn (1, w2 (4,,,)) ;) wa((, Next, initialize the weights corresponding to the SHB evolution according to the initialization of the mean-field ODE, i.e., wf = wi (0, Ci
w W (0. = w2( (1))), C2(j2) and wg WSHB((.. (0. f2) = wa C2(12)). and let them evolve according to SHB dynamics (19), thus obtaining w/ (K, (1) w S (1. 52), wg SH
(K, j2). Finally, we define the following distance metric that measures the difference between the mean-field and the SHB dynamics:
(r WHB)= max sup max[]]w; (t,)) wSSB (Lt/E),5i)ll.. (ej
[12(1, C1(1), C2(j2) w2 w SHB (14/E).
(20)
[w3(t, Ca(je)) - w SHB (Lt)=),
Theorem 5.2. Let Assumption 3.3 hold. Consider the coupled the SHB dynamics (19) and mean-field ODE (13), and the distance metric (20). Then, with probability at least 1 - exp(-5"),
DW WSB < K(r,T) (Vlgg Rmmx + +6) + Ve(V D log mig + +01
(21)
Pimia
where Tmaax = max(n1, n2), Pmin min(n1: n2), and K(y,T) is a constant depending only on 7.T.
For three layers, Theorem 5.2 gives that the approximation error scales roughly flogmmax/ nmin Ve(D+ + log n172), i.e., it vanishes as long as n1, f2 both grow large, with Tmmax = o(e"min) and e o(1/VD + log (17)). As in the two-layer case, the bound dimension-free (in the sense that n1, my do not need to scale with D), and K(7,T) = O(e=". We remark that the scaling of the bound in T is also a shortcoming of the existing analysis for SGD in (Pham & Nguyen, 2021a). The detailed proof is provided in Appendix D.
6 Consequences of the fien-field analysis: Dropout stability and connectivity
The mean-field perspective put forward this paper leads to a precise characterization of the SHB train- ing dynamics. In particular, Theorems 5.1-5.2 offer a provable justification to two remarkable properties exhibited by solutions obtained via gradient-based methods, namely, dropout-stahillly and connectivity. The
9

---

Published in Transactions on Machine Learning Research (02/2023)
fact that solutions (often resulting from algorithms that use momentum) can be connected via simple paths with low loss was empirically observed in (Garipov et. al., 2018; Draxler et al., 2018), and this property related to dropout-stability by Kuditipudi et al. (2019). In (Shevchenko & Mondelli, 2020), it was shown that SGD solutions enjoy dropout-stability and connectivity and, by combining this analysis with Theorems 5.1-5.2, similarly strong guarantees can be obtained for heavy ball methods. We will keep the discussion at an informal level, as the details are similar those (Shevchenko & Mondelli, 2020).
Let's start with the two-layer case. Given non-empty set A C [n], say that W is en-dropout stable if
R(W) Rarop (W; A)] ED.
(22)
where Rarop (W; A) is obtained by replacing the two-layer network f(x; W) defined in (1) with the dropout network
((: WA) =
Σ JEA
w2 (20
(23)
A
w2(i))ees and |A| denotes the cardinality of the set A. Similarly, in the three-layer where WA (wr(()) case, given two non-empty sets Ai C [mi], C [m2], the dropout network is given by
H2 j2; WAL.A2)
Σ
w2( (i))
A
ЛEAm
(24)
f(a; WA).Aa) = A2
Σ
my o W WAL.)))).
J2EA2
and eD-dropout stability is defined analogously. Furthermore, say that two solutions W and W are cc-connected if there exists contimuous path in parameter space that starts at W, ends at W' and along which the population risk is upper bounded by max( R(W)) tcc.
At this point, the quantitative convergence result to the mean-field limit provided by Theorem 5.1 and 5.2 leads to quantitative bound on the dropout stability and connectivity of the solutions found by the stochastic heavy ball method. We remark that proving only the consistency of the mean-field limit, as done in (Krichene et al., 2020, Theorem 1), does not suffice to obtain such guarantees the structure of the SHB solution. In particular, for two-layer networks, after [T/e] steps of the iteration (16), the resulting parameters are cd-dropout stable and cc-connected, where
(vlog A +6) + Ve(VD+ + + +6)
ED K(x,T)
VIAI
(25)
(Vlgg + Vn
< T)
+ Velv D log n +5)
with probability at least 1 1- Here K (x, is universal constant depending only on 7, as before. Similarly, for thyee-layer networks, after k S [T/e] steps of the iteration (19), the resulting parameters are e er-rooouut stable and ec-connected, where
+ velVD + log 01722 + 8) ep K(?,T) (Vlog Amax +0)
VAnin
(26)
(Vlgg
log Amax + 8)
cc T)
+ V log n1 + 8)
Vmmin
with probability at least exp(-82). Here, Amax = max | Amin min |2/}, mmax max(ni, n2) and nmin = min(n1, n2). We remark that the path connecting the two solutions can be explicitly constructed as in (Kuditipudi et al., 2019; Shevchenko & Mondelli, 2020). More specifically, this
10

---

Published in Transactions on Machine Learning Research (02/2023)
path is pieco-wise linear, and the mumber of linear segments is fixed constant. We also note that the bounds for multi-layer networks in (Shevchenko & Mondelli, 2020) exhibit linear dependence on the input dimension D. In contrast, our bounds (25)-(26) are dimension-free.
Finally, let us highlight that our mean-field viewpoint can shed light on the (Entezari et al., 2022), where is empirically observed that, after a suitable permutation, the solutions of the optimization algorithm enjoy lineor connectivity. In fact, Theorems 5.1 and 5.2 show that, by running the corresponding SHB training algorithm multiple times, all the resulting solutions satisfy (18) and (21), respectively. This implies that, after permutation of the neurons, the distance between such solutions can also be upper bounded by the RHS of (18) and (21), hence the linear connoctivity is an immediate consequence of this closeness among the solutions.
conjecture in
7 Global convergence of the mean-field ODE for three-layer networks
In order to show the global convergence result, we first need to make some extra assumptions. Assumption 7.1. We make the following additional assumptions for the training of a three-layer neural network:
(C1) (Universal approzimation properly of the aclivalion) 0s exhibils a universal upprozimation property, that is: {o1((w.-)) w E RD has dense spon in (Dxx,, where Dx denotes the re-marginal of the data distribution. D.
(C2) (Full support at initialization) pb has full support.
(C3) (Mode of convergence) The mean-field ODE (13) converges to the limit (20, (00, Wa (0000)02)) Formally, we have that
Ex,,s, + Wo((,, C2)l) W/ (00, C2)1 [02(00, C1 C2) || (t,, C1) 100 Gi)lla 1 0, Ec,, [11 + wyy w (00, C2)l) (00, C2)1 (t, C1, C2) w2(00, C1, C2)1] 1 0,
Ec, + (00, C2)l) |wy C2) wy (oo, C2)l ( O,
csssupEc, E. 2t, C2)1 ( 0,
G,
Pr (x,, # 0.
The universal approximation property is the key assumption to obtain global convergence result. This requirement is mild, since most activation functions used in practice are universal approximators. The assumption on full support is also mild, since widely used initialized schemes (e.g., He's or LeCun's initial- ization) employ a Gaussian distribution, which indeed has full support. The assumption on the mode of convergence is purely technical, and is an open question whether can be relaxed. specifically, part (C3) is needed because of the lack of entropic and moment regularization, which makes difficult charac- terize the limiting points of the noiscless mean-field dynamics. We note that the uniform convergence of the gradient (the fourth assumplion in (C3)) could be replaced by Morse-Sard lype of regularity assumplions (see (Nguyen & Pham, 2020, Section 8)). We remark that these requirements also appear in Pham & Nguyen (2021a), with the exception of Pr [wa(co, C2) + 0] > O, which is needed to handle the heavy ball dynamics. Theorem 7.2. Let Assumptions 3.3 and 7.1 hold, and assume further that R(y, (x; W)) is conver in [(:; W). Let (I) be the solution of the mean-field ODE (13). Then, we have that
lim
inf E. E.R(y.((x)).
(27)
The detailed proof is deferred Appendix E, and we provide here a sketch. First, we show a degencracy prop- erty for the mean-field ODE, i.e., there exist deterministic functions w(.)) :) R220 RD RDD u ()))))))))
11

---

Published in Transactions on Machine Learning Research (02/2023)
time 2000
time 7500
-1.0
log of dropout error
time 12500
log of dropout error
time 15030
-1.5%
-4
-2.5
2500
üime 3600
-3.0
6
335
ume 12505
400
1600
3200
100
200
400
800
1600
3200
6400
# of neurons
# of neurons
(a) Two-layer networks
(b) Three-layer networks
Figure 1: Dropout error plotted as a function of the network width. The log-log plot is close to be lincar, matching the behavior of our theoretical predictions of Section
R x R X x :22 R - R such that
wilt (1) w w (0, C1)),
w2(1, C,,C w , w (0, Ci), w2(0, Ci C2) C2), C2)), wa(1, w3 (1, w0 C2)).
Next, we show that, for any finite 4, wi( is continuous both arguments and wi (t, Ci) has full support. Finally, the convergence to the global minimum is obtained by combining the argument that wi (t. C1 full support for all finite f with the assumption on the mode of convergence.
Theorem 7.2 is rather different from the global convergence result for the heavy ball method presented in (Krichene et al., 2020). Tn fact, Krichene et al. (2020) consider noisy dynamics, and show the convergence of the mean-field ODE to the global minimum of a certain free energy, which represents an entropic regular- ization of the loss function. In this setup, the convergence is guaranteed by the noise term in the dynamics and by the regularization term the free energy functional. In contrast, we consider noiseless dynamics and do not prove its convergence. Instead, we show that, when the mcan-field ODE converges, it must do so towards the global minimum of the un-regularized loss function. At the technical level, our proof strategy is an adaptation to the heavy ball case of the argument for SGD in (Pham & Nguyen, 2021a), which also crucially relies on the universal approximation property of the activation function. A similar idea was first proposed in (Lu et al., 2020), and it also appears e.g. in (Fang et al., 2021). However, our contribution is the first to lackle the case of optimization with momentum.
8 Numerical results
Experimental setup. We train a two-layer and a three-layer fully connected neural network in the mean- field regime on the MNIST datasct. The training algorithm is stochastic gradient descent with momentum, and we evaluate the dropout stability of the learnt models. For the two-layer network, we take the width 100 2 where & E (1...., 7}; for the three-layer network, we take n, n2 n and use the same grid for n. We use the PyTorch default initialization, pick the learning rate to be 0.05 and the momentum to be 0.9 (which implics that 2). We rescale the learning rate, so that the scaling of the gradient does not depend on n, as required by our theory. The batch size is 100 and we train for 25 epochs, which means that the neural network is trained for 15000 steps (each epoch contain 600 steps and there are 25 epochs). For each model, we perform 10 i.idd. experiments and report their average. We compute the population loss for both the original network and the dropout network, obtained by randomly dropping out half of the neurons. For each experiment, take 10 random dropout networks and report the average population loss.
12

---

Published in Transactions on Machine Learning Research (02/2023)
Experimental results. In Figure 1, we plot the log of the dropout error defined in (22) as function of the number of neurons in each layer. Different curves correspond to different numbers of trained epochs. Two remarks are in order. First, the dependence of the dropout error on the time of the dynamics is rather mild (and, hence, our bounds on the constant ((),T) appear to be pessimistic). Second, the dropout error scales as an inverse polynomial in the width, in agreement with (25)-(26).
9 Discussion and future direction
In this paper, we consider neural networks with two and three layers, and analyze the dynamics of stochastic gradient descent with momentum also known as the stochastic heavy ball method (SHB) from mean- field viewpoint. After showing the existence and uniqueness of the mean-field limit, we provide quantitative convergence result of the discrete SHB dynamics of a finile-width neural network to the corresponding limit differential equation, thus solving problem raised by Krichene et al. (2020). Then, we exploit the power of our mean-field perspective by (i) proving that the solutions found by SHB enjoy desirable properties, such as dropout stability and connectivity, and (ii) showing a global convergence result for three-layer networks.
At the technical level, our proof strategies build on the work by Mei et al. (2019) and Pham & Nguyen (2021a) for neural networks with two and three layers, respectively. However, these papers focus on vanilla SGD, and dealing with SHB requires a number of delicate technical results. In particular, (i) we establish several boundedness and smoothness properties of the mean-field dynamics, (ii) we track varions new quantities and exploit a second-order Gronwall lemma (Pachpatte's inequality) to bound them, (ii) we perform a discretization of the particle dynamics which is different. from the SGD case, and (iv) we prove a key measurability property for the second-order dynamics, which characterizes the dependency between layers during training. Our strategy is tailored to the stochastic heavy ball method, but similar idcas could potentially be applied also to Nesterov's 's accelerated method, due to the similarity in the continuous limit. The study of other popular training algorithms (e.g., Adam) most likely requires entirely different technical analysis, whose investigation is left
interesting open problem.
Let us conclude by discussing some extensions and future directions.
Beyond three layers and fully connected networks. While it should be possible to extend our analysis of SGD with momentum to feod-forward networks with more than three layers, remains unclear how to obtain global convergence guarantees in a more general setup. In fact, this is an open problem even for the case of optimization without momentum, see (Nguyen & Pham, 2020, Section 5). We also note that, although we only consider fully connected networks, our analysis could be generalized to other architectures, such as convolutional neural networks (CNNs) or ResNets, by slightly modifying the assumption. In particular, for CNNs, if we replace inner products by convolutions and let n or Rı, n2 be the number of filters in each layer, our analysis can be directly applicd. For ResNets, the presence of the skip connection modifies the backward path, and the mean-field limit would need to reflect this modification.
"Central limit pecoremmttype of results. Theorems 5.1-5.2 belong to a "law of large numbers" of results, in the sense that they show how close is the SHB dynamics to the mcan-field limit. A parallel line of work has studied the distribution of the perturbation of the finite-width neural network around its mean-field limit, see (Chen et al., 2020; Sirignano & Spiliopoulos, 2020b) for two-layer networks, and (Pham & Nguyen, 2021b) for multi-layer ones. However, all the existing results concern the vanilla SGD algorithm, and understanding how momentum can affect such a distribution is an interesting future direction.
Convergence of noisy dynamics. In this work, we consider noiseless dynamics, and the global con- vergence result for three-layer networks follows from the universality of the activation function. In contrast, in the noisy case, the mean-field limit for a two-layer network is under-damped mean-field Langevin dynamics, whose convergence follows from the convexity of related free-energy functional (Krichene et al., 2020). In a multi-layer setup, the free energy is not convex, which makes it challenging to obtain global convergence.
Comparison between SGD and heavy ball methods. Motivatod by the observation that, practice adding momentum helps to generalize better (Sutskever et al., 2013), an exciting avenue for future research is to exploit the urructure of the mean-field limit to draw insights concerning the solution found by heavy
13

---

Published in Transactions on Machine Learning Research (02/2023)
ball methods. While the recent work by Jelassi & Li (2022) theoretically proved the improvements in generalization error provided by momentum under certain settings, we are not aware about whether such results hold in the mean-field training regime. A key difficulty here is that, for the noiseless dynamics, global convergence not guaranteed in general, let alone an explicit characterization of the stationary solution. To circumventt this issue, one option may be to consider a suitably regularized noisy dynamics, which admits a stationary solution in Gibbs form, in the limit of vanishingly small noise and regularization. Finally, heavy ball methods are known to enjoy faster rates of convergence in the convex setting. Thus, establishing convergence rate for the mean-field dynamics is an exciting avenue for future research.
Acknowledgements
D. Wu and M. Mondelli are partially supported by the 2019 Lopez-Loreta Prize. V. Kungurtsev was sup- ported by the OP VVV project /0.00/00/0/0/0/66..... 019/0000765 'Rescarch Center for Informatics'.
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for doep learning via over- parameeterization. In International Conference on Machine Learning, pp. 242 252, 2019.
Luigi Ambrosio, Elia Brué, and Daniele Semola. Leclures on optimal transporl. Springer, 2021.
William F Ames and BG Pachpatte. Inequalities for differential and integral equalions, volume 197. Elsevier, 1997.
Dyego Araújo, Roberto Oliveira, and Daniel Yukimura. A mean-field limit for certain deep neural networks. arXiv preprint 19v: 00199, 2019.
P. Auer, M. Herbster, and M. Warmuth. Exponentially many local minima for single neurons. In Newral Information Processing Systems (NIPS), 1996.
Simone Bombari, Mohammad Hossein Amani, and Marco Mondelli. Memorization and optimization in deep neural networks with minimum wver-prraectriization. Advances in Neural Information Processing Systems, 2022.
Zhengdao Chen, Grant RotskofT, Joan Bruna, and Eric Vanden-Eijnden. A dynamical central limit theorem for shallow neural networks. In Advances in Newral Information Processing Systems, volume 33, pp. 22217-22230, 2020.
Zhengdao Chen, Eric Vanden-Eijnden, and Joan Bruna. A functional-space mean-field theory of partially- trained three-layer neural networks. arXiv preprint arXiv:2210.1622 2022.
Lénaic Chizat. Mcan-field langevin dynamics Exponential convergence and annealing. Transactions on Machine Learning Research, 2022.
Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterrzed models using optimal transport. In Advances in Newral Information Processing Systems, volume 31, 2018.
Lénaïc Chizat and Francis Bach. Implicit bias of gradient descent. for wide two-layer neural networks trained with the logistic loss. In Conference Learning Theory, pp. 1305 1338, 2020.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems, volume 32, 2019.
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprocht. Essentially no barriers in neural network energy landscape. In International Conference on Machine Learning, pp. 1308-1317, 2018.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-
parameterized neural networks. In International Conference Learning Representalions, 2019.
14

---

Published in Transactions on Machine Learning Research (02/2023)
Rahim Entezari, Hanic Sodghi, Olga Saukh, and Behnam Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. In International Conference on Learning Represenlations, 2022.
Cong Fang, Jason Lee, Pengkun Yang, and Tong Zhang- Modeling from features: mean-field framework for over- par-pearaeteri deep neural networks. In Conference on learning theory, pp. 1887-1936. PMLR, 2021.
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P. Vetrov, and Andrew G. Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. In Advances in Neural Information Processing Systems, pp. 8789-8798, 2018.
lan J. Goodfellow and Oriol Vinyals. Qualitatively characterizing neural network optimization problems. In International Conference on Learning Representations, 2015.
Karl Hajjar, Lénaïc Chizat, and Christophe Giraud. Training integrable parameterizations of deep neural networks in the infinite-width limit. arXiv preprint arXiv:2110 110.11001100 15596, 2021.
Kaitong Hu, Zhenjie Ren, David Siska, and Lukasz Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint XrXi::195.. 090,:1990.779 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Newral Information. Processing Systems, volume 31, 2018.
Adel Javanmard, Marco Mondelli, and Andrea Montanari. Analysis of a two-layer neural network via displacement convexity. The Annals of Slatistics, 48(6):3619 3644-26212,,,, 2020.
Samy Jelassi and Yuanzhi Li. Towards understanding how momentum improves generalization in deep learning. In International Conference on Machine Learning, pp. 9965 10040,00000000 2022.
Anna Kazeykina, Zhenjic Ren, Xiaolu Tan, and Junjian Yang. Ergodicity of the underdamped mean-field langevin dynamics. arXiv preprint vrXic:2007.16660 2020.
Diederik P. Kingma and Jimmny Ba. Adam: A method for stochastic optimization. In International Confer- Ence on Learning Representations, 2015.
Nikola B Kovachki and Andrew M Stuart. Continuous time analysis of momentum methods. Journal of Machine Learning Research, 22(17):1-40, 2021.
Walid Krichene, Kenneth F Caluya, and Abhishek Halder. Global convergence of second-order dynamics in two-layer neural networks. arXiv preprint arXim:2007.0688 2020.
Rohith Kuditipudi, Xiang Wang, Holden Lee, Yi Zhang, Zhiyuan Li, Wei Hu, Rong Ge, and Sanjeev Arora. Explaining landscape connectivily of low-cost solutions for multilayer nets. In Advances in neural infor- mation processing systems, volume 32, 2019.
Daniel Kunin, Javier Sagastuy-Brena, Surya. Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural me- chanics: Symmetry and broken conscrvation laws in deep learning dynamics. In International Conference on Learning Represenlations, 2021.
Yiping Lu, Chao Ma, Yulong Lu, Jianfeng Lu, and Lexing Ying. mean field analysis of deep resnet and beyond: Towards provably optimization via overparameterization from depth. In International Conference on Machine Learning, pp. 1426-6436, 2020.
Yi-An Ma, Niladri S Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L Bartlett, and Michael I Jordan. Is there an analog of Nesterov acceleration for gradient- based MCMC? Bernoulli, 27(3):1942-1992, 2021.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33):F7665 E7661,, 2018.
15

---

Published in Transactions on Machine Learning Research (02/2023)
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit. In Conference on Learning Theory, pp. 2388 2464, 2019.
Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. arXiv preprint orXiv:2007.18896 2020.
Phan-Minh Nguyen and Huy Tuan Pham. A rigorous framework for the mean field limit of multilayer neural networks. ar preprint XrX::2001.11444. 2020.
Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. arXiv preprint ar 220: 10469, 2022.
Huy Tuan Pham and Phan-Minh Nguyen. Global convergence of three-layer neural networks in the mean field regime. In International Conference on Learning Representations, 2021a.
Huy Tuan Pham and Phan-Minh Nguyen. Limiting fluctuation and trajectorial stability of multilayer neural networks with mean field training. In Advances in Neurol Information Processing Systems, volume 34, pp. 4843-4855, 2021b.
Boris T Polyak. Some methods of speeding up the convergence of iteration methods. Ussr computational mathemalics and mathemalical physics, 4(5):1 17, 1964.
Grant M. Rotskoff and Eric Vanden-Eijnden. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. In Advances in Neural Information Processing Systems, volume 32, 2018.
Itay Safran and Ohad Shamir. Spurious local minima common in two-layer ReLU neural networks. In International Conference on Machine Learning, 2018.
Katharina Schuh. Global contractivity for langevin dynamics with distribution-depeedeedt forces and uniform in time propagation of chaos. arXiv preprint arXiv:2206.0300 2022.
Aleksandr Shevchenko, Vyacheslav Kungurtsev, and Marco Mondelli. Mean-field analysis of piecewise linear solutions for wide relu networks. Journal of Machine Learning Research, 23(00)), 2022.
Alexander Shevchenko and Marco Mondelli. Landscape connectivity and dropout stability of SGD solutions for over- parammrrmmtee neural networks. In International Conference on Machine Learning, pp. 8773 8784, 2020.
Bin Shi, Simon S Du, Michael Jordan, and Weijie J Su. Understanding the acceleration phenomenon via high-resolution differential equations. Mathematical Programming, pp. 1-70, 2021.
Justin Sirignano and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A law of large numbers. STAM Journal on Applied Mathemalics, 80(2):725-775 2020a.
Justin Sirignann and Konstantinos Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 130(3):1820-135) 1852,, 2020b.
Weijie Su, Stephen Boyd, and Emmanuel Candes. A differential equation for modeling Nesterov's accelerated gradient method: theory and insights. In Advances in neural information processing systems, volume 27, 2014.
Tlya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning, pp. 1139 1147, 2013.
Alain-Sol Sznitman. Topics in propagation of chaos. In Paul-Louis Hennequin (ed.), Ecole d'Eté de Proba- bélités de Saint-Flour XIX, pp. 165-251. Springer Berlin Hcidelberg, 1991.
Jun-Kun Wang, Chi-Heng Lin, and Jacob D. Abernethy. A modular analysis of provable acceleration via Polyak's momentum: Training a wide ReLU network and deep linear network. In International Confer- ence on Machine Learning, pp. 10816-10822 2021.
16

---

Published in Transactions on Machine Learning Research (02/2023)
Andre Wibisono, Ashia C Wilson, and Michael I Jordan. A variational perspective accelerated methods in optimization. Proceedings of the National Academy of Sciences, 113(47):F7351--773) 2016.
Francis Williams, Matthew Trager, Claudio Silva, Daniele Panozzo, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate ReLU networks. In Advances in Newral Information Processing Systems, volume 32, 2019.
Xiaoxia Wu, Simon Du, and Rachel Ward. Global convergence of adaptive gradient methods for over-parametterieel neural network. Xiw preprint
2019.
Greg Yang and Edward Hu. Tensor programs TV: Feature learning in infinitewidthh neural networks. International Conference on Machine Learning, pp.
2021.
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small nonlinearities in activation functions create bad local minima in neural networks. In International Conference on Learning Representations (ICLR), 2019.
Organization of the appendix. In Appendix A, we provide some a-priori estimates that will be useful in the following arguments. In Appendix B, we prove Theorem 4.1, namely, the existence and uniqueness of the mean-field limit for two-layer and three-layer networks, respectively. In Appendix C and D, we prove Theorems 5.1 and 5.2, which show the convergence of the SHB dynamics to the corresponding mean-field limit for two-layer and threc-layer networks. Finally in Appendix E, we prove Theorem 7.2, which is the global convergence result the three-layer setup.
A A-priori estimates
A.1 Two-layer networks
Lemma A.1. Assume that (A1) (A3) hold, and let f(x; P)) (OV(0:p) be defined in (11). Then, for any ftxed T, there erist universal constants K, K2(n,T), T), where the latter depends only on 7, T, such that the following results hold.
1. (Boundedness) We have that, for any 0, P.
f(x: p) S KEE /22),
(0: p)] - K K|w2),
(28)
(oV(O:)) p) ll2 < K + [w2]
(Boundedness for mean-field ODE) We have that, for any t S T, (z(()) as governed by (12) satisfies
[w>(t)) s K2(y.
(29)
(Lipschila continuity)):
-00 - (00: < K(1 + /wal) (lwz - w//1 + || w/ 12212210 + Wo(p, p')),
(30)
Vo & K(1 + (w2) (lw2 - w61 + ||w
w ||2 + W2(p, p')).
(31)
Proof.
1. By the definition and assumption (A1), we have that
[f(:; /) [ [w] x) < KE, (22),
VO p) < E= [0. f(:; (z: x) |wz| < K|w2/,
V (0:p)) = EE> [0, f( (( | K,
W. v(0:0)ll || [a R(y: / (x;; () w22 (w) a)))]] || E K [w2]-
17

---

Published in Transactions on Machine Learning Research (02/2023)
2. By writing down the integral form of the ODE, we have
[wz(t)] - [22(0)] + + L/ ([wa(s)] + (22( ds + 1/' Lo Vu, plu) du
< (K + KT KT +y /-" |w2(s)]
$ {K + KT + KT2)
By setting K2(1,T) = (K + KT + KT?) the proof of (29) is complete.
For the Lipschitz continuity argument, we have
T ((: Ez [O2R(y. fa: ())
Es [8y f(:: p) uuz (wF
VoV(O:p) =
Ez [2PR(y: (xx (
Thus,
(v: v(0')))) xK(0:ll)ll wi| K w2 uz21
K k E Ex; ((; (32)
We define the Bounded Lipschitz (BL) divergence as follows:
dBL(P.p') = sup( Eo-p' f(0)) if
[f 1}.
We have the following relationship between the BL-divergence and the Wasserstein distance (see for example (Chizat & Bach, 2018, Appendix A) for more details):
dBL(P, S Wz(e, p').
Hence,
[E; 0) (2; s Kd s kWa(p,P), which implies that the RHS of (32) is upper bounded by
K [w2[([[ ([] willa + Wa(p. P')) + K|w>> w/l1 K(1 (w2l) (lw2 w/ll + l|u0 w/il2 + Wa(p,
This concludes the proof of (30). The Lipschitz continuity of Vo Vo p) follows from the argument.
A.2 Three-layer networks
Lemma A.2. Assume that (B1)-(B2) hold, and let H2.f. AW AA, be defined in (14) and (15). Then, for any fixed T, and given neuronal embedding
(() x 92 Fi F2, x P2), wi (1(00) :), w2 -,,,,,,0 )), ws(0,,)]
there exists a universal constant K and universal constants K3,((,, T), Ks,3l,, only depending on T such that the following results hold.
1. (Boundedness) We have that, for any W, 2, for any e [0, and for any eı 02,02 22
f(z; W (()) S K esssupc, [ta(t, C2)
H2(2,00 oz (()) < K esssuupc, [t, C1, Cz)
( oz W(I) < K
18

---

Published in Transactions on Machine Learning Research (02/2023)
A c2: W (t)) & K esssupc, (os(t, C2)
02 ((t) & K (esssupc,,ce (wz(t, C1, C2)l esssupc, wys(t, C2)
(" (t)) - K css sup Wta(t, C2)l ess w2 C1, C2)1 C₁
C2
(x c1: ( < K ess cu (y, C2) ess sussu C2 |ay (t, C1, C2)
2. (Boundedness for mean-field ODE) We have that, for any ST, ess sus wg C2) & Ka, ess (u2(t, C1, C2)1 S K3,((1T)) G.Os
(33)
C2
(Lipschitz continusity) We have that, for any TT
C1 (()) - Hi c1;
cs W (()) S -(,))
H2(m, ((t) - H2(ar 02 W())
- Kess (t, c2)l w (t. C1) - (t, (t, + l2 01, W2((, C1,
Gi
Isx; (()) (a; W())
< K ess sup( C2) |1, C1, C2)1 ((, ( (1,
+ +ws(1, -
C1,
(2, + 1ws(1,
(2)
A (2,02; (t)) - 45 c2 (t))
s K (1Hg(z,
H2(2, e2; W()) + (f(=; - f(:: (())))
A W (t)) (s)) cz. W(t))
K|lwa(t, 02) [ (()))))));((())) H2 cz W + ((, uz ca)l
+ KK 0z) If(x; W(t)) - f(x; (())
(2,C,, W (1)) (z) (z W())
I K (x. 02; W()) (z, ((()))
+ A (z, c2 (t)))-/[((ttc))
| (z ai W(1)
()))l:
S K EE. [ (z. C2; (t))aaottt Ca)] w/ - - (t,cel)
+ K Ec, [
(z, Cz: (()) (t,c1, C2) A Cz; (1)) C2) I.
Proof-
1. By the definition and assumption (B1), we have
[a= (()) |Ec, wa (2),, C2))
S K|EC,, (t, C2) S ess (ty(t, /ssssut
C2
|22 W())) Ee, u2 (HC
I ess sup (1, [u([
C1
< K esssup J2 (t, C1)
Cn.ss
A (x. c2; W()) = JarR(y; fx; (H (())) S K
A (x, C2 W (I)) = R(w: f(;; (1) oz (()))
S sus wtl (2)o < K css sup (Wy C2)
C2
C2
Oz. W(t)) (2,
A (z. cz W (H)) W(t)))
< K ess A (ar, C2 W (()) < K ess sup |W C2)
C2
C,
19

---

Published in Transactions on Machine Learning Research (02/2023)
A cz (()) = |Bc, (a, C2; Ca
I sus 1 (z, C2; (t))
ess [2 C1. C2) ess (H
C2
C..²2
C
S K ess sup (t, 02) ess sup 12 (t, C1 C2) C1,Cs
C2
|| (z. (8))l22 || (z. cj W (t))all2 I ess sus A a C1; (1))l
G.
S K ess sup (e, C2) ess (t, C1 C2) C₁,Cs
C2
2. We have that, for any t t
(ts(tt, S (ws(0, 02) + 7 % (1ws(0, c2)1 + |uss 02)l)
+ +%% [EAA (u, 62) du
< K + KTT + KT +7 L- [wa(s, c2)
$ (K + K-T + KT )e: = Ks,,(,,
which readily gives the first claim. Next, we write
(22(1,c1, S 1w2(0, c1,)l++ + S. c2)) +
+ *1 1 EEAA (a, e2; W du ds
SK+KTT+ ((1ma(ssa),22))d ds Ka,
which by Gronwall's lemma, implies that
ess sss (t, C2 C2)1 s (K + KYT+ Kas(1, = Kk,2(,, T).. Cr,Cz
3. For the Lipschitz continuity argument, we have
H1( (())) (1 c1; W()) = | (wi (t, w, ci)]
- kllwitt, - (t,
H2(m, 02 W(I)) H2 0z W())
=IEE, (t, C1, (w) 1)) x) Ec, (t Co (() (t,
S K (1w2 C1, 02) mm G)) wn (1, G1)l2 | (1, (1,
G.
Is(:; W()) f(; W())
.
< wa, (2) (H2 C2; W(t)) - Ec, uz (H), Cz; (()))
< K Kesssup(| C2) [w2(t, C2) wm (t. C1) (t,
C.C.
(2) + C2) o - wxt,
W (1)) A (z, cz; W())
J0,R(y, f(x (t)))
.22(H2(z oz (())) 0g f(x (1))
< K (HH2(1,02 e2; (()) H2 02 ( + (f
f(2))
20

---

Published in Transactions on Machine Learning Research (02/2023)
A (z, cz (t)) - (2
= -0, f(x, ()))))aa (()))
-O2R(y. (1))
< K|wasl 02)1 [H2 (()) H2(2, (4)) + Klua((, - (z((,
+ { 1f(z; (t)) f(; W())
(2,01, <2; (()) - (2,
= A 02
- K K H (x, c2: (t)) AZ W())
+ K / (z. c2 W(t))
iz (t,,)))))))
| c1; (()) - A (z, cii W (1))l22
S SKI A (z. C2: W (t) w2 (t)aa))
S K|Ec, (z, C2; [M
+ KB, C2; (C
B Existence and uniqueness of the mean-field limit
B.1 Two-layer networks
In this section, we prove the existence and uniqueness of the mean-field limit for two-layer networks. We recall the mcan-field ODE again here:
d@(t) rf)ddt, dr(() = (0) dt.
V
(34)
The proof follows from constructing Picard lype of iteration, similarly to (Sirignano & Spiliopoulos, 2020a, Section 4), (Javanmard et al., 2020, Theorem C.4). Below is an adaptation of the strategy in (Sznitman, 1991, Theorem 1.1). We first write the integral form of the mean-field ODE:
0(1) = 0(0) - 7 f (0(s) - 0(0)) ds - 1. ('Vo (w)) du du
(35)
r(1) = --)) - (0( - 0(0)) - -f. ())
(36)
where p(t) is the law of (0(t), r(t),, and we use 0 (1). (t) to denote the and r marginals, respectively. We define the space P2(RD x ( to be the space of probability measures on RD RDD equipped with Wasserstein metric W2, and we have p(t) E P2(R x R). We define the space C ([0,T), P2(RR R))) to be the space of continuous maps p(::T): : [0, + P2 D x RD).. We omit T when there's confusion. The space is equipped with the following metric: dr (p = (uP (2(P1 (t))(tt))2(t))
Note that the space (P2(RD RD), W2) is complete space (Ambrosio et al., 2021, Theorem 8.7). Thus for any fixed T < 0o, the space (C (10, PR( x R²)) dr) also complete.
Next, we define the operator HT(: 0(0) : C (0o, Ti, (R RD)) - (10, 2], P2 x IR))) as follows:
Hr(p::((0) : p, p(t) = [Law θ(t) F()))
0(t) 0 7 (0(s) (0)) ds 1/' ( du ds,
(37)
21

---

Published in Transactions on Machine Learning Research (02/2023)
where ((0) denotes the parameters of the mean-field ODE (35) at initialization, which means that the stochastic process we defined in (37) is coupled with the mean-field ODE.
Note that the (t) in (37) is no longer the law of (t). but the imput distribution. We use Hr(p(t)) to denote (p; (t)) that is the distribution of the solution (37) at time t. We omit ((0) when there is no confusion. The definition of Hr can be interpreted as follows: it maps E C ((o, T), P2 x RP)) as input to output the law of (0(1), which evolves according to the stochastic process induced by the probability measure p(t).
It casy to see that the fixed point of Hr is the solution of the non-linear dynamics (35). Thus, our goal is to show that there exist To such that HT, has unique fixed point, or equivalently that HT, strict contraction.
Proposition B.1. Under Assumptions (A1)-(A2), there erists a To only depending on K,y and C(T.) E (0, 1) such that, for all Pr, P2 C ([0, Tol, P2 BR)) with the some initialization (01(0), T2 (1)) (02(0),rz(0)) r20)), have:
dr, (Hr, (P1). Hr, (p2)) S C(To)dn (2).
Proof. We first fix any 0 T < 00, and the space C (2. R)).. C (10, (2 RRD), we define two dynamics as follows:
Given P1, P2
E
(t) (0) - 7 I (0)(s) (0)) ds L' L (V)(O(():)((())) du du ds,
02(1) = 00) - 7 / (0) 02())) ds S J (oV(02())) ds.
where (t) = wn (1) _ (1) and (2()) = w u
We want to upper bound the difference between these
two dynamics, which will give us an upper bound on
dr(ltT (pi), HT (p2))
For all E [0, T), we have
/ [[01(s) l2
[t0 02(1)ll2 $>
ds.
Now, by Lemma A.1, we have that
(1)) - (oV(O2(()>
<K(1 + (O1) (Iw5") (1) - == (1) + || - (0)22 ( +2666(0)20)))
<2K(1 + K2(7, T)) (10)((t) - -2(())l2 max Wa(o((0)
sElo.T]
where 0((1) (w) ( (0), m2 ( (4)) for i 1,2. Thus, have that:
0 (t) (t)l2 < 21 + K2ly, T))T2 max (s(pi(t), +7 [[((s)) - 02(s)]l2 ds LE|U,2
2K(1+ K2(1,T)) 1 [10((w) du
22

---

Published in Transactions on Machine Learning Research (02/2023)
Similarly for ||r((t) r2(())l2, we have that:
(|r(()) - r2(1))22 I 7 S.¹ [|r((s) - r2(5)]l2 ds + - ds
- 2K + K2(7, =E(0.7]
max
(s(e((t), +7 |mr(s) r2)
/ [10((s) - (s) + 2K + K2(,, T))
Putting these two results Logether we have:
[ ((i(e) (1((t)
(2(t) | S (10((()) 02(1) (mi()))))))))))))) (1)l2 r2(t))
< 4K(1 K2(7, T))
(1)
max 2E(o.T]
+7 [. (10)(()) - -22(s)l2++1r)) -02(())))0(())++
- (s) (8)
+4K(1 + K2(,,T) 1 I [10](u) 02(u)ll2 du ds
<<K(1- K2l7, T))T² max W2( (2()))))))(()))))))))))))))))))) + +2 L/' (1i13)
(02(s)
(2)(s) (
ds
2EJO.T]
4K(1 + K2(n, T)) 1/ / ((
(1)( - (2(u)) (2( 2
du.
By Corollary F.4, we have that:
(0((t) (1)())
(o2(t)
447 (
4K<(1 + K2(7, (12,, )
$ 4K(1+ + K2(y, T) 1 1xp exp
max p2)))
27
2E(O,T]
Thus, we have that
(()(3)
(o.(t) <T"K(T,T) (E(O.T]
max (2(o((0).
which implies that
(r(o(())
T) max (a(rf(0,,8()),, rE(O.T]
max k
where we set K(>,T) 4K Ka(,,) (1 (1+exp 43 KK(44
4%
Let C(T) T2K(;; T). Then, we could always find To such that C(Tb) < 1 since C(0) 0 and C(T) is continuous in T, which finishes our proof.
By Banach's fixed point theorem, there exist To 0 such that the mean-field ODE has unique solution in time interval [0, To]. Now, we show the existence and uniqueness of the solution of the mean-field ODE for any time period [0,7].
Theorem B.2. Under Assumptions (A1)-(A2), for any T > 0, there exists a unique solution for the mean- field ODE (12) in the interval [0, T..
Proof. The idea is to separate the time interval [0, T] into subintervals of length To, that is, we consider the intervals J0, Tol, [To, 2T,,..... .... 7L [To, Note that the contraction property we proved in Proposition B.1 only depends on the length of the time interval, the proof can be done recursively. That is:
23

---

Published in Transactions on Machine Learning Research (02/2023)
1. In the interval Tol, (12) with initialization (0()),r(0)) has unique solution {p(t))
2. In the interval [To: 27o), consider (12) with initial distribution p(To), and it has a unique solution (E(7,77))
Recursively do the above steps until the interval [17]] To: 7).
Thus we have that, for any T > O, there exists a unique solution for (12) in the interval [0, T).
B.2 Three-layer networks
In this section, we prove the existence and the unqqueness of the mean-field ODE (13). The integral form of the mean-field ODE is given by
wa(t, ww,,cc) -7 " (10((),02) ((x((,,)))=x(00.c))) ds 1." I-" E.Ag" (2) du
(38)
w2(t,c1,,2) (2.00,1..)) 02) - (wz(s (2,00, c1, 02)) ds / L' E.A. (u, 2, c1, du
(39)
wi(t,c1) = (0, ci) - (wi(s, (w](())))- (0, 1-' 1// E2Al (1, du
(40)
In order to prove the existence and the uniqueness, we follow the same Picard's itcration arguments for the two-layers case. We first define the following norm:
W Ir = max ess sup sup {1wa(t, C1, C2),, (us(t, C2)l).
(41)
C..ss tE[0,7]
Next, we define the following metric for two of mean-field parameters:
Dr(W, w') = max ess sup sup (t, Ci) wi (t,
Cr.22 rE(O.T]
(t, -z C1, 02) [t, C2) wz(t, C2)
(42)
Note that the metric we define above is not the metric induced by the norm, since in the definition of the norm we only require the boundedness of w2 and W3-
We define the following functional space of the mean-field parameters:
Wr(W(o)) = ( (t)
||]] W/ < co, W(0) W()),
(43)
which means that all the W E (W((()) have the same initialization W By Lomma A.2, we know that:
ess / C1. C2) < K3a(y, C,C,s
Wyy(t, ess / C2) < Kg,s(,,
(44)
C,
It easy to see that the space (r)W(O)) is complete w.r.t. the metric Dr (W, W'). Let us define the operator Hy: Wr (W) (W(()) as follows:
Input: (W)eee)
24

---

Published in Transactions on Machine Learning Research (02/2023)
Output: (eet)) such that:
02)(, wz(0, 02) ( - (wa((,)22) (f wz da 1/1 /- B.Ag (a. c2 ( du
(45)
02),, w2(0,,1, 7 (((0s)(,,,,)))) (2))
wy (y(00,1,, ds / I'¹ E.D (2, (()))
(46)
(((())==m000)))/((mmaa)))))) = = -7 - (wi ( (10))):))- (1(,ci)) - [. EE (2.01 (
(47)
We aim to show the following proposition.
Proposition B.3. Under Assumptions (B1)-(B3), there erists a To only depending on K, and C(To) e (0, 1), such that, for all W w2 E (W((())), we have:
Dr. (HtT (W"), HT. (W3)) s C(To)Dro (W W
(48)
Proof of Proposition B.3. For simplicity of notation, we denote the output of Hr((W to be W which is composed of wig Wiy, w|. The output of Hr(W2) is denoted similarly.
By the definition of the mean-field ODE, have that, for any t TT,
[ (w)(s,c)) wg(ss c2)
ct, c2) S/
/ E. E 02; c2; W (u)) A (x, o2: W du ds,
lub(s,c),, 02) cı, c2,)) ds
L.- 1 B (0,c1 W (u)) A (2. (
(t, (t, 3 4 L (s, w
f/ (el AM (a))))
By Lemma A.2, we have that:
max [Z AW (ar 02 (w()) A (0, W²(w))l,
E_ AW (r, (1,0,,02 W' A (2; (²())),
(2
Ex A (x, W (w)) A (z. (u))ll) z (W W2).
Thus, we have:
D.(W' W ST / D.(W', k( J. D. W
(7t ( W3),
which implies that
- (3T , Dw <
(49)
Since TT),,,) when T 0, and (YT + (T) continuous in T, we can pick To such that (o + Kg) (7,Th) Tb) <1, which finishes the proof.
K
25

---

Published in Transactions on Machine Learning Research (02/2023)
Since Wr(W(()). is complete, by Banach's fixed point theorem, there exists unique fixed point for the operator Hr. which implies that the mean-field ODE (13) has unique solution in [0, Tol. By following the same argument of the proof of Theorem B.2 (separate the interval [0, 7' into sub-intervals of length To and successively apply Proposition B.3 to each of them), we readily obtain our main result concerning the existence and uniqueness of (13) [).
Theorem B.4. Under Assumptions (B1)-(B3), for any T > O, there ezists a unique solution for the mean- field ODE (13) in the interval [0, T)
C Convergence to the mean-field limit - Two-layer networks
In this section, prove the convergence to the mean-field limit for two-layer neural networks (Theorem 5.1). Our proof's structure is inspired from Mei et al. (2019). Before going into the arguments, we first recall the definition of the mean-field ODE and the stochastic heavy ball method (SHB) for two-layer networks. Then, we define two auxiliary dynamics: the particle dynamics (PD) and the heavy ball dynamics (HB).
First, recall the mean-field ODE as follows:
(t r(t)dt,
dr(t) = (->r(t) - voN(O((;P dt,
(50)
and the corresponding integral form
O(1) 0 - I (0())) (0(s) (0(5)- 0(0)) ds S (O)( du
(51)
The SHB dynamics is as follows:
0 SHB + 1,5) SHBB (HHB((.))) + (1 (E (HHBBBB BSHB 1.5)) Vo ( (k
Vj E [n],
(52)
where (k) - E do(k.) is the empirical distribution.
In order to describe the convergence to mean-field limit, we define the following particle dynamics (PD):
do d0PP(L,)) =r (1, dr (t.,t) = (-yr (t))) dt, Vj [m].
(53)
where (t) 1 j= (1.) is the empirical distribution at time t. Furthermore, the heavy ball (HB)
dynamics is defined as
HB (k- + 1.j) Q (k.)) + (1 (0)BBB @B L.5)) No (K. Vj [n]
(54)
We remark that (52), (53) and (54) have the same initialization, that is
PP o (0.)) e @SHB((.)) [n].
Define the following distance metrics:
Dr(0, (P0) = max sup JE(n) =E(O,T]
(t, - (t.,
(55)
PD HB
max sup jE(n [0
HB ([t/e],)) 0 PP
(56)
(0 HB HB gSHB = max max
1 ( - SHB (HHB(()
(57)
je(n] [E[T/=]
26

---

Published in Transactions on Machine Learning Research (02/2023)
C.1 Bound between mean-field ODE and particle dynamics
In this section, we bound the difference between the mean-field ODE defined in (51) and the particle dynamics defined in (53), whose integral form is as follows:
0 (t,(t, = -PP()))) -1 S (PPP(s.))) -s"0 (PD(O,))) - [o (o du ds.
(u, (u,
Proposition C.1.
Under Assumplions (A1) (A3), we have that, with probability least exp(-
Dr (10 s k((,T) (8+ Vlogn
(58)
vn.
where K is constant depending only 7 T.
Before proving Proposition C.I, first prove the following lemma, which characterizes the Lipschitz conti- muity of the mean-field ODE.
Lemma C.2. Under Assumptions (A1) (A3), there erists universal constant K K(y, depending only on 7; T such that, for any I, > such that 1,1+ T <<x <T,
||0(1 + 0()]l2 - ()7,T)T, (t ( (t)) s I),
(59)
The same holds for the particle dynamics PD (t, E [n].
Proof of Lemma C.2. We only prove the results for the mean-field ODE, and the proof for the particle dynamics follows from the same arguments.
We first try bound the increments [@(t)= - ((0) By the definition of the mean-field dynamics, we have that:
[ [0(s) 0(0)ll2 ds S/' (Vox(
(t)(t) 0(0)1l2 S 7
S> 10(s)) 0(0)]l2 ds + K2(r,T),
where the last step we use that (O)); (u))ll2 S Ki T), which follows from Lemma A. By Gronwall's lemma, this implies that, for any t T
( - ( - K1 (7, T) exp(1T) == K2(1.T)
Next, by definition of the mean-field ODE, we have that:
1+ ||0(t + 7) 0(t)ll2 $> 1.- t+T [[0(s) 0(0) ds + 1. [ (VV) ( du ds.
Thus,
10( ( - 0())l22 - K((7,T)T,
where we use that fact that
[(()) 0(0)122 & K3(7, T),
Vo. p((( du - K3(7,T)
27

---

Published in Transactions on Machine Learning Research (02/2023)
For the Lipschitz contimuity of p², we just note that by definition of the W2 distance, we have:
[[0((+++] ( -
0(t)
Now we are ready to prove Proposition C.I.
Proof of Proposition C.I. In order to bound the difference, we first define iidd mean-field dynamics:
(tt. 0(0.)) ~ (((s.)))((0.j))) - ( (u. du ds,
where (t, is the law of (t,,)), and we coupled the iid dynamics with the particle dynamics at initialization, that is, let:
e(0.)) (P.(0.)). Vj E [n].
We also define the empirical distribution of ((, that is: (t) = 1/2 Σこ ) Since the n mean-field dynamics i.idd, we have that (1, (1), Vj [n], thus we use the notation of p (L) to denote the the law of jt for each j E [n].
By Lemma A.1 and Lemma C.2, we know that:
sup max [22 S K(y,T), T), <E[T] jE(n)
sup max ||0(t + r. (tt.j)llz & K(7,T)r.
*E(T] jEln]
We have that
10 (t.)(t))) 5(1 + (0.)0 + //" (s,) - -(ss,))ll2
* 1/ (u. -oV( (u)) du ds,
and our goal bound:
sup max (tn - -(t.)lll AE(O,T] JE(n)
Now we aim to bound the quantity
sup max jE
An application of the triangle inequality gives
(e -
(60)
+ " (u.))
(61)
Recall by definition that
Vw, V(O(u, (w) = Ez [a, (a;
x)))
V w. v(O(u,j): (w)) E, [agl(l,, f(x; p (u)) (o) (w (u, j) [))))
28

---

Published in Transactions on Machine Learning Research (02/2023)
Similarly,
V w, (): Es [0gR(y. f(x; ()) (wi(u,)) x))))
V w. j): (u) = E= [0. fa::
For the term in (60), we use concentration incqualities to give an upper bound. By the Lipschitz continuity of R2 Assumption 3.1, we have
Voo (v)
s K Ez (u)lu.))o((ux(") (), x)
|s<;; (u)) f(x; (w))
s (T(7,))][(a; (w) flx; 70 (u))..
For the term Is(z; fw:; 70 (u))), we have
(u)) f(x;; po (u))
(w ((co)(()))(ux)(())))))
Note that, by Lemina A.1, we know that
(w) Epta) (wtt))(()))(w s K2(7,
By Lemma F.1, we have that, with probability at least 1 - exp(-n((8)))
(ylu,))o(vv((u.))= (nnnps((uo( s (2(() (12+
By Lemma C.2, we know that 1. E.= (y(,. (x) Bp(w) (u)) (u) is K K(y, Lipschhtt
continuous in u. Thus, by taking union bound over j E [n] and e {0.7.... | [n) we have that, with probability at least oI exp(-n(8'))),
Σ wyy (w (ow)) s Ks(7)T) (1/ + 8' + +1)
max sup jEIn [E(o.T] n
82 når))
Take : 8'
Then, with probability at least exp(-
5 + vlog vn
E (ul) (u(,.))o(ws((,)))) s K4(y,
max sup JEIm tE(o,
n
which implies that, for term (60).
max sup Vo - (o))) s Ko(,, T)) 8 + Vlogn jEln [E(o.T]
with probability - exp(-6²).
For the term in (61), we use the Lipschitz continuity of V. By Lemma A.1, we have that, for each j E [n]]
(oN((((.)): Vo (o
P + W2 (t) (f))),
29

---

Published in Transactions on Machine Learning Research (02/2023)
Note that pbn)() discrete measures, thus we have:
1/2
1 (0(1.) -P
(2(7((0 v
S (,
n
Hence,
- Vo (t.)) (t0
Combining the above results, we have that, with probability - exp(-8²),
Di(0, 0D & () 8 Vlogn
+7 / D.(0,0"0 ds + Ks(n, / / Dul0, 0"D) du
Vn
An application of Corollary F.4 concludes the proof.
C.2 Bound between particle dynamics and heavy ball dynamics
In this section, we bound the difference between the particle dynamics defined in (53) and the heavy ball dynamics defined in (54). We recall that the distance we aim to bound is defined in (56). Note that the heavy ball dynamics is discretization of the particle dynamics. Thus we aim to bound the difference at time point ke.
Proposition C.3. Under Assumptions (A1) (A3), there exist universal constant K(7, T) depending only on T, such that
Dr 0HB) < T).
(22)
Proof of Proposition C.3. By the Taylor expansion, have the following approximation for the particle dynamics:
8 PP + 1)e,)) = (kE,))) + PD (ke, f) + 2 a.r PD 0
3,r
(63)
Also by Taylor expansion we have:
T PP(ke, -P 1)e,))
1, (ke
(64)
By plugging (61) into (63), we have that
8 8PD + 1)E,)) PP 1)e,5) 0, (KE, +
(ke,) (* - (0 1)8, (-yr(ke, - Ne. 8² O
= 0 (ke,)) + -1 PD (1 0 ( (ks
0(2))
Now get discrete iteration equation for the particle dynamics, with approximation error of at most OE:3). By accumulating the V. v (0 PD ((, (P)(()) term from 1 = l,.... k, we have
k-1
0 (P = (
Σ
q
(K)
o (h O(E)), (O
(65)
1=0
where g(t) 22 - (1 7 8 211-1--ee)"
78
30

---

Published in Transactions on Machine Learning Research (02/2023)
The heavy ball dynamics can be written in similar fashion:
0 =
Σ
(66)
1=0
Thus, we have that
( (ke,
(2)
(E. Vou (L.))
(67)
By Lemma A.1, we have that
(O VoV f): (t (inn(l)))
(K((7,T)(10>0((ee.)) - 0 0= (1e, + (ole
Since (pD((s),
(eels)),Hee(00 (Ie)) discrete distributions, we have that
1/2
Waldsp(le),0he0100 s
0 (Ie,)) - "d*((ee
S (D (
n j=1
which implics that
(L..) Vo (l. (e(,, (0PD,0HB)
As a result, we have
Die POPD OBB < K2(n, T) (Dir PD +
Finally, an application of the discrete Gronwall's lemma concludes the proof.
C.3 Bound between heavy ball dynamics and stochastic heavy ball dynamics
In this section, bound the difference between the heavy ball dynamics defined in (54) and the stochastic heavy ball dynamics defined in (52). We recall that the distance we aim to bound is defined in (57). The manipulations of the previous section imply that the heavy ball dynamics can be written as
0 (kB)) -
Σ (k) NoV (L)))
(68)
Similarly, the stochastic heavy ball dynamics can be written as
(HB( = 0 (0.))
( V (2()) 0 (SHB (. (s)n(())
(69)
1=0
Proposition C.4. Under Assumptions (A1) (A3), there erists universal constant K(r, T) depending only on 7, such that, with probability ex²(-
H8 SHB < K(1; T) (e(T))((V)))) D log + 8).
(70)
31

---

Published in Transactions on Machine Learning Research (02/2023)
Proof. By using (68) and (69), we have
[H. - (SHB((k. (k: 2 (I.) - SHB (1, (sHn(()))
By triangle incquality, we have that
Σ
(VoV (L)): VoW((z())
SHB (L. ((e(()))eeenn)))) (bHe(()))))))
(z(() -)) vo HB (1.3)
(71)
1 (Vo HB (1.)):PH(() (.
(72)
For the term in (71), by the Lipschitz continuity of Vo we obtain
(o SSB (HHB(1))): (sHB(()))()))) V (=(((=()) 0 0 (L:):
(1 (De.,. 8HB) (2(
Since PHB PSHB are discrete distributions, we have that W2 (0) < Dia (0 @SHB). Thus,
(=((= 0 SHB (1,i): (Z((=((), 0 nB (L.)): 3 K2(7, (0 OSHB).
For the term in (72), note that, since the z(1''s are sampled iidd at each step by definition, we have
Ex(1) (o 0
u (1, j): P):))))))))) Vov (L
Thus,
V (z(1), 0 HB (1.)) V VO (L.)
is martingalc difference. By Lemma A.1, we have that, for all le f1, [T/=]],
[[Vo((=(1), H (L.)::he)) (oV(O (1.)): = K2(7,
Thus, an application of Lemma F.2 gives that, with probability least 1- exp(-63),
max IE(1....IIP [e/ll (z((,, HB (L.))) VoV (LJ): & T)
By taking a union bounds on j E [n], we have that, with probability at least 1 - exp(-82)
(o((((()), 0 HB (L.)):Phn(() Vow (Li)) k,(r,,)) log 8).
max JEIn] (T/ee...TTTT))
By combining the above result, we conclude that
IT
HB SHBB < K((s,T) logn +5) 6)+ K2(7; T)
Die.e (B @SHB
(73)
Finally, application of the discrete Gronwall's lemma concludes the proof.
32

---

Published in Transactions on Machine Learning Research (02/2023)
C.4 Proof of Theorem 5.1
Proof of Theorem 5.1. The proof follows from combining Proposition C.I, C.3, C.4, and the fact that:
Dr(0,0 SHB) - 1r(0,000) - Dr HBB + (r SSH))
D Convergence to the mean-field limit Three-layer networks
In this section, we prove the convergence of the training dynarmics to the mean-field limit for a three-layer neural network. Our prool's structure is inspired from Pham & Nguyen (2021a).
Before going into the proofs, let's first recall the definition of the mean-field ODE and the SHB dynamics, and then define two auxiliary dynamics, namely the HB dynamics and the particle dynamics. For the convenience of further computation, random variable corresponding to the stochastic heavy ball dynamics, the heavy ball dynamics, the particle dynamics, and the mcan-field ODE as WSBB, HBB WPD, W respectively.
define these continuous dynamics in integral form. We define the
The mean-field ODE (13) in integral form is the following:
oz((, = c2) 7 S ( ds S." E (2,02; duds,
wr(t,cc,,g) = =.x(0.),.)))(00....))) ? ( (12())c).c))_(((0),).)))
( (z,c), W(u)) du
1/¹ / E. (. ((w)) du
wi(t,c = (0, (wi(),c1) - wl (w
(74)
The SHB dynamics as follows:
SHB (k + 12) = w SBB (2) + (1 w SHB (2)
w3 SHB 1, E"A," (z(1),, WSBB(()
W3
W2 SHRR (k: + 11,11 = u.2 (k, + +1 - 78) (w) SHB kk j1. -w2 w (k- 12)),
SHB
=² D (z(k).51..2 w. (l:)
wi SHB (k: + 1,11) = oi SHB (1) + - >e) SHB (t) wi
SHB (k: 1,j1)) e² 2 2 (z(k), W SSBB (x)),
(75)
where z(k) is the data point sampled at time step k. We define the particle dynamics as a contimuous dynamics without mean- field interaction:
w. =t (0))) -7 -, (wn (WFP(s,.))) -x) ds S.' E.AW (zi (P)) du ds,
wa - (00((.,.))))))) - 7 J/ (wF" (15"((,)11))))) w. (0, ds
% E.A (z.J1.32: w rD(( du ds,
wf (L(.))))) = (0,)),,))) -7 1 ('(10f (P "P (0(()))) ds S'' [ E.AW (z,j1; PPD du.s. ds. (76)
33

---

Published in Transactions on Machine Learning Research (02/2023)
We define the HB dynamics by replacing the stochastic gradient in the SHB dynamics by the true gradient. That is:
W (k 1,
(w) (k. f2) - w3 1 1, E (z, W (B))).
u12 (k + 1,11;j2) = (k. - (1 re) HB (k. j1: - w/2 11 1k-1,,1..1))
(77)
- E (z.J1. w (k)),
(k, B 1, f (k. (1) + (1--7)) (w B ji) w B 1, E (z, d1; W (k))...
Note that the HB dynamics can be viewed as the discrete version of the PD dynamics. In order to present our main theoretical results in this section, we first define a distance metric to quantify the level of correspondence of these dynamics. We define the following distance metries:
Dr(W, W PD) max sup { (P,,()))))) - wi (t.
1E[0.2] Iug (t,,i.))) - -=((t
(78)
(wg (t,,j)) (a(t, ; [mll, E [ma]]
Dr PD " max sup
xE[o.7] { [[][//e)))))=(((((
|w2 ([t/s],,h..e))
w2 PD (t,
(79)
HB (L1/8],
(1,,))))))) : in E [mn], E [ma]]
Dr. WHBB max max
HB
(k:,
w LuSHB(K (k, 11)112.
ke ke|T/=]
I22
HB (k. - w - SHB (k,
f1:22)1
(80)
Iw3
HB (k, - -u j2)l : E [m1]. E [ma]
We acknowledge the abuse in notation from reusing Dr for multiple metrics, which is done to avoid prolif- eration of notation. is clear that:
(W SHB) S Dr(W, PD PD W Dr. HB WSHB),
(81)
where Dr W is defined in (17).
In the following subscctions, we bound the terms (T WP), PD HBB and Dr. (W HB WSBB).
D.1 Bound between mean-field field ODE and particle dynamics
In this section, we bound the difference between the mean-field ODE defined in (74) and the particle dynamics defined in (76). We recall that the distance we aim bound is defined (78).
Proposition D.1. Under Assumptions (B1)-(B3), we have that, with probability at least exp(-82),
K(7,T) (Vlog Rmaxx + 6).
Dr PD)
(82)
Ammin
where (T) is universal constant depending only on 7, T, Temin min(n1, n23 and mmax max n2)..
In order prove Proposition D.1, we need the following auxiliary lemma, which characterizes the Lipschitz continuity of the mean-field ODE and of the particle dynamics.
34

---

Published in Transactions on Machine Learning Research (02/2023)
Lemma D.2. Under Assumptions (B1)-(B3), there erists universal constant (Ty, depending only 7, such that, for any T >> O with +, S S
esssup 22
(ay(t + T, ( (T,
ess sup
T, - w2 [ ((),)))
(83)
c1
esssup ott (Kn
T)..
e1
The same holds for the particle dynamics (L.((,,)); wh (1, J2), / (t..
Proof of Lemma D.2. We do the proof for wi (4, (1), and the same argument applies to w2(4, (2)),, ws (g(1)22) 62). First, we derive a bound on the increments up to time t, || a) w w1(0, By the definition of the mean-field ODE (74), we have that
= (w:(6.a))(((aa))) wi(0. 1/" ['EA' E.A," (z,c1; W(u)) du ds,,
wi (1) w11(0, (0,
which implies that
[wn ui (0, = / || (01,)alll (/ (E (w))ll2 du ds
w $ - dl) T2,,(T,
where in the last step we use that, for some constant Kı(r,T) depending only on 7.Y, | ess sup,, suPue(o, Al (z, ci; W(w) < T)( by Lemma A.2. Thus, by Gronwall's lemma, we have that
sup xE[0.7]
Im (1,c) un u. (0, TT T) K2(75T).
(84)
Now, by using again the definition of the mean-field ODE (74), we have that:
Jwn (1+ wi (1)22) -1- (wi(s,
++ [ EEAT (2, W(u)) du
w. w, c1 ds
S sup 1E(U,2]
|un - (0, Ki(9,
K2( Ki(,,
where in the second line that | ess sus, suPwejo.7] AW (2, W(u))ll $ Ki(y,T) by Lemma A.2, and in the last passage use (84). By setting K(n,T)- = K2(y,T) + T), we obtain the desired result and the proof is complete.
By the above Lemma D.2 and Lemma A.2, we immediately get the following corollary. Corollary D.3. Under Assumptions (B1)-(B3), there erists universal constant K K(y,T) depending only on T such that, for any t, > 0 with t, + T $ T, the following functions
f(x; W(),, H2 02; W()), Ec, (Z, Cz: (())w2(t, C2)
are K(y,T) -))-T)schippttt continuous in t. The same holds for the particle dynamics, i.e., the functions
n2
f(x; (()), H2 W()), ((()),
A (z. j2 (())
n2
K(y, -ipsshhit continuous in l.
35

---

Published in Transactions on Machine Learning Research (02/2023)
Proof of Corollary D.3. We do the proof for f(æ; (t)), and the same argument applics to the other cases. By Lemma A.2, we have that
If W(t + 7)) f(;; w(t)) & K (tssup((u((++t,cc)) l1z((t + T, 02)1 [t
/ - (1,tc))ll + w3 t, 02)1 /22 + T,, | w
qm2
By Lemma D.2, we have that max(I|w; (t + T, wi (t,ct))le Wa + 7, - w2(t, bay wg((t > T), Furthermore, by Lemma A.2, we have that:
ess sup u3 (t 7,02)) sup esssup wz((t, S k3,3(,,)), rE[o.T]
ess 61,62
10 ,,,,,,,) & sup ess wsspp/u/tt,,,ceeut & (,,,(n,T)) te az
Thus, we conclude
[(t:) (()) & K)T,T)T
which gives the desired result.
7
Now we are ready to prove Proposition D.1.
Proof of Proposition D.1. Let recall that the quantity AW is defined in (15). We start with computing the difference in the term AW:
B_48 (z, C2(J2); ((t) - E,AW (z, j2: W (())
<E, A (z, C2(12) (()) AW (z,12; W (P())
(85)
-E, 02R((, f(x; (H))) C2(12): W(t)) - 02R(y, f(x: W
j2: W (()))) <KE=/f(>; W(!) f(æ; (F) + H2H C2(12): (()) H2( j2; W W (VP((() (1))
where in the last inequality we use the boundedness and Lipschitz continuity of 0z and obtained from Lemma A.2.
Similarly, for AW, AW, we have that
ELAW (z, (in), C2(12): (()) E. J.. W PPP(I)
<E. A (z, (i))), W(t)) A (2,j1, W PP())
E. K_K[sstt, C2(12)) C()))
- H2( W PD(()))
(86)
+ (10a(1, C2(j2)) |ttp://((cclllooaaaaa/aaaaa-a/aa
E. [|
W (40
W PP)()))
and that
( E.AW W (z, (.Ii(i)) E.A (z.ji: (PP(
1 122
<K- E,
Ecz Cz: (())atCCa -
A (z.j2: W PP())) PD (t,,
122
j==1
K E.
(|w Ci())) (t.,
(87)
n2
<K Ez
Ec, [c,
(2. W PD (1...1)3)))) (..)))
11
+ ((T,T) - P
36

---

Published in Transactions on Machine Learning Research (02/2023)
Here, we remark that the expectation Bc, | (s, Cz: (2z(: C2)l for the mean-field ODE corresponds to the 1 average n2 2=z=1 A (u:2:;;)) for the particle dynamies.
Now, our goal is to upper bound the following quantities:
[f(; (()) - f(;; W (()),, (2 C2(12): (t) j2: W P())))
n2 A
Ec. (
(2, W (()) PD P (1,
122
12
To do so, we follow (Pham & Nguyen, 2021a, Appendix C.2, Proof Theorem 14, Claim 2). Then, we have that, for any 81, 82, 8z O,
[s:;;; W(t) - f(:; (P)())))
max
max
H2( C2(12): (()) j2; W (t))),
32
Rc, A (()) C2) 1 2g
Σ
A (2,
W 3 (1,
max fn
n2
f==1
() (Dr(W, w + 81 + 62 68)
with probability at least
n22 cxp
n220
n1
K K((T,T) })
n2
exp
exp
8,
K1
82
Ki
83
By Corollary D.3, we know that f(x; W H2(w, cz: W()), Bc, [ Cz ((1)) cı, C2)] are K(r, Lipschitt continuous, and the corresponding quantities for the particle dynamics are also K (7, T) Lipschitz continouu. Thus, by taking a union bound on t E {0, n..., [//n]], we have
[f(z; (()) f(æ; W (P)()).
max sup xE(O.T]
max
H2( C2(12); W(() (HH(ax - (z W PP()))
32
Ec, [ (z, C2; (()) C2)
A H jo: W (P))) (P
max
11
j==1
s K2(7, T) (W, WPD 81 82 53 + 1))
with probability at least
T f22 1
oz
)
1
{
22d
+ n1
n
)
1 8
exp
K2ly, T)
+
52
exp
K2(y, T)
83
exp
K2
T)
In particular, pick
K2(7,T)
(Viog Rmax + +6),
Ks(7, T)
(Vlog Umaa + +0).
77
51
82 = 83
V //max
vm
vnz
37

---

Published in Transactions on Machine Learning Research (02/2023)
Then,
[f(z; (()) f(a; (PP(()))
max sup (E[U,T]
max H2 C2(22) (() - H2; W PD()))
j2
Ec, [ (z, C2; W C1(1), - 112 j==1
A H (z,j2; W (t))) D (t. (t..)
max
21
I (a(T,T) Dr WPD) Ka(%,T)
(Vlog mmmas + 8))
v Amin
(88)
with probability at least 1 - exp(-60)
Next, we combine (88) with (85), (86) and (87) to provide high-probability bounds on A A AWW By rocalling the definition of the mean-field ODE (74) and the analogous definition of the particle dynamics (76), we finally obtain that, for all I. T, with probability at least 1 - exp(-6"),
D.(W, W S K(7:T) (Vlog +0) + +7 J'' ('D,(W, D.W,, w PP) ds S:' S. Du PP) du ds.
Vmii
An application of Corollary F.4 gives the desired result (82) and concludes the proof.
D.2 Bound between particle dynamics and heavy ball dynamics
In this part we bound the difference between the particle dynamics defined in (76) and the heavy ball dynamics defined in (77). We recall that the distance we aim to bound is defined in (79).
Proposition D.4. Under Assumptions (B1)-(B3), there erists a universal constant K(y, depending only on 7, T such that
Dr (W W.P,, S K(,,T)e
(89)
Proof of Proposition D.4. Note that the heavy ball dynamics is just a discretization of the particle dynamics, so we first bound the difference each time point ke. By a second order Taylor expansion, we have the following approximation for the particle dynamies. We do the computation for Wy as representative, and the proofs for w1, w2 the same.
We have that
PD W3 (Pk + 1)e,1)) = w jo) + P. (ke, + o (ke, + ((52)
(90)
Also by Taylor expansion, we have that.
P wf (Ke, (2) w. PP(( 1)E, + 1 PP (ks,
(91)
2
By plugging (91) into (90), we obtain
wgg PP 1)), w.g PD (ke, + w/y (ke, - w/ ( 1)e, j2) ) + S.ug PD Pkke f2)e² + O(e))
=w - + w PD (kE W/
PD (Dk( 1)e, j2) + (-70.w "0 - E. E.A W (z, jo; W WDD (ke))) +0 0(55)
=uP"(ke,j)) (1 (w)
PP((k PD (ke, j2) w3 ((K- 1)E,je)) - - (z,j2: W PD (kE))e +
where in the last step we again (91).
38

---

Published in Transactions on Machine Learning Research (02/2023)
By unrolling the recursion, can write the particle dynamics in the following form:
W/ (ke,
E.A (z, f2; W (Pee) + +Oe)),
where
k-1-1
1 - (1 g0)*=1
1 =
Σ
(1 nE)²
c²
e
7E
?
1=0
Similarly for w2, wi, we have:
2 w P0 (0,
Σ
(1) E (2: (Pe) 0(e),
w w ji), = w (1)
E, (z. (P + O(s).
t=0
We can write analogous expressions for the heavy ball dynamics:
w (k. w. B..
( E.AW (z, W (BB()))
102
HB (k, j1, = =w2 (0.5i())))))))
100 E ( W (I)),),
(92)
w HB (k. wiBB(00,11)
(00 E.AW (z, W (B))
Let us define the following notation, for & E {I,.... [ DD W k) max (1
HB (ki, ji) wi PD (0e (ke,j))ll,
| (1 (k,)1 w P (ke, 52:
103 HB (k, w (KE, j2)1 : :1 E [mil. e [n2]]}.
Recall that, by construction, wg (0, (0. j2). w. (0 j1. j2) w. 10 (0, j2) j2) and "0. 51)
(0, f1) for all f1 Thus, by computing the difference between w m uz and wHi w HB w have that DD W k) satisfies the following induction inequality:
k-1
Dr (W w HB. k))
<
Ki 1=U
PD W HB + O(s),
(93)
where we have used the Lipschitz contimuity of A A. and AW obtained via Lemma A.2. Thus, by the discrete Gronwall's lemma, we obtain that, for any {1,..., 11
(B: k) S K2(7.T)= Dr PD w MB,
Finally, an application of Lemma D.2 gives that wP w PD w PD are K2(7, Lipschhit continuous in time. Thus, for any < T, Iw (1,,j)) - ws F1 (L1/E),io)
S wy "P - W3 PD (11/e + ww PD (LL/eeee))))) ( ([4/e(()))
< [/w
PD (Lt/E/ w3 (1t/s (t/// jo) + (3(
Similar results hold also for |w? PD PP (1, w HB (1t/ f1)) and | PD we conclude that
PP (t.h1) w HB ([t/El)) As result,
( PD W1BB <
k) + K2l7, T)e S K(r,T)s, DD kg
max
ke (1)
1)
which gives the desired result.
39

---

Published in Transactions on Machine Learning Research (02/2023)
D.3 Bound between heavy ball dynamics and stochastic heavy ball dynamics
In this part we bound the difference between the heavy ball dynamics defined in (77) and the stochastic heavy ball dynamics defined in (75). We recall that the distance we aim to bound is defined in (80). Proposition D.5. Under Assumptions (B1)-(B3), we have that, with probability at least exp(-6"),
D W-BB) (7,,T)V + log(nin2) + 8),
(94)
where (7, is universal constant depending only on T.
Before proving Proposition D.5, we state and prove result concerning the boundedness of the SHB dynamics. Lemma D.6 (Boundedness of the SHB dynamics). Under Assumptions (B1)-(B3), we have that, for ony k e L
(k, je)l S K [+ ) I,
SHB W
1+ 2) (1 (1 7 w (k, f2)1 S K
where K is a universal constant.
Proof. By following passages analogous to those leading to (92), we have that the SHB dynamics can be written
w Sk j2) w SHB (0,
(k) A (=(1). (S)
(95)
SHB (k. (1,22) = uz (0)
A (2. SHB
w2
Recall that
AW J2; WSBB (1)) = 02 (, wSBB(()) (2(H2(2(()) W (I)))
which implies that
| f2: (0) (
18
SHB (I)) (2(H222((). W SHB (1)) s K.
Thus, WC have
--1
Ke KSK K 4 T.
BHB (k. S (w3" SHB (0. (2))2)
Σ
k K
7
1=0
where in the last step we use that ke < T.
For (wS (WSHB(( (k, J1, je)l. we recall that
AAW (=(1), (1. j2; WSBB(O) (0))
(Q ( SHB()))) (00)))/wHBBB (1,1)) o2((H(z()), (1))o() SHB (L.)))))))))))") x1))
< <KIws SHB SHB (1,j
Thus, we have
SHB (1, S 1 SHB (0,51:
SHB ( (1,j2))
102
10
T ke
< K + K
1+
1 =
1 12
S K
which gives the desired result.
40

---

Published in Transactions on Machine Learning Research (02/2023)
Proof of Proposition D.5. Throughout this argument, we fix and consider & L Recall that the HB and SHB dynamics can be written as in (92) and (95), respectively. Furthermore,
wi (HB wi
w (0.
Σ
(k) AW (=(1),1: W (HBB()) q
(96)
Recall that, by construction, w3 (0, (0, w (0, j2) w SHB (0, 12) and wi HB (0.
(0,
w for all J. Thus, by computing the difference between the expressions in (92) and (95)-
we have
HB (o.s (k. W' [w's
(uSHB( (He (ke j2)l
(1) (B.[A) (z, W 1R (D]] A3 (z(1). W ((SSH(()))
(0) (B.[A) (z. W HB (I)) - E (z.j2: SHB (I))
(97)
(E.IA) (z. SHB (1)) AW (z(0) W SHB (SSH(())
+
(98)
w (k. (4.J1.12) - SHB (k, f2.
(E,1AT (z, jo: W HB (B)) - AT (=(1). W (.HB((
(E.1A (2.d1,52 w (B)) - E. (z,J1,52; W SHB (()))
(99)
(E.IAW (z,J1. W SHB (0)) (1)) (z(1).11.2) SHBB(())
(100)
w HB (R:
w
wHHB
(H, (k, f1)l21
SHB
(1)) (E_ (z.01 W HB (0)) (= J: W (HSHB
[E (2, 3: W (1)) - A (z. in: W SHB (T)))
(101)
(B=|A) (z. SHB (1)) AW (z(1).n1 WSBB(() (T))
(102)
To bound (97), (99) and (101), we use the Lipschitz continuity of AW AW and AW, together with the fact that c(t) < E/y. In particular,
(z,je; W (D) (E,A)((z,jee
E,4," (z. j2: W (HB()))
1
Σ (E (z.ja; W HB HB)) E= (z.ja; W SHBB (HBB())
HB
(103)
$ K(y, T) 6 k-1
Σ
Die.e HB
HB SHBB
γ 1=0
41

---

Published in Transactions on Machine Learning Research (02/2023)
Similarly, we have
(0) (ELIAT (2, (D)] EZIAN (z,ji. W (1)1)
S K( T
Dis (W HB W.SH),
1=0
k-1
(K) (B_ (z, WHB(1) E.A, (z, W (I))
S <K(,,
Die,e D WSHB).
(104)
To bound (102).(100) and (98), we first define the filtration Fk(k) as the sigma-algebra generated by (mms(0,f))))xe1m)) z(0), (k)).. We define the filtration Fi the same way. Tet us recall that, in a one-pass algorithm, we take ii.d. samples at each step and, hence, we can write, for all I E {1, LE]).
Ez(() [A] (=((), (=(1),
W WHB (1) = E. (z, W WSHB(())
(1))) [ (=(1), J2: W (1) = J2: W'))(()),
(z( (z1 W SHB
(1) = =,,A (z,j1: SHB (I))
Clearly, we have that that
W
(H),, (1. [21 mutually independent, which implics
A (z(1). WSB E.A, j2; W SHB()))
is a martingale difference with respect to the filtration Fs((). Thus, EE a A (=(1), j2: (HB E_A! (z, J2: (I)]( E I }} is a martingalc (samc for AW and AW) Next, we show that the martingale differences are bounded, so that we can use martingale convergence results to bound these terms.
Combining Lemma D.6 with the same strategy of the a-priori estimations of Lemma A.2, we have the following upper bounds:
E_A, W (z. WSB DT (z(k), W (HB S K1,
Eg A (2, SHB((k) A Ji.k).,i..a SHB (k)] S Ki(r, T),
W E2 AW (2,51 - (=(((, W (H)) < ()
-
Thus, an application of Lemma F.2 gives
Pr
max ke/
(E.A) (z. W SHB SHB - (z(1), W SHB(())) 2 KVVTe[1+ + +5s)
<exp(-
Pr
max
(k) (E.AM (z.jn.32 W SHBB(() ( (z(),.1.)) S2: w (SBB())
2 + 62)
kE [T/T
S Sexp( (xx)
Pr
max ke|T/el
Σ (8)
(E_AW (z.J1: S-BH(() AT (=(1), f: W SHB (1))
2 2K(y + 51)
S <exp(-07)
42

---

Published in Transactions on Machine Learning Research (02/2023)
By taking union bound over J2, we have that, with probability at least 1 - exp(-8²),
k-1 (2, W SHB(()) A (z(1),52 W (HBB())
max max Jn. kE[T/E]
Σ
(E
1=0
(E247 (2, ())
W (())(()))
(EgA) (z.h: WS) AN (=()), W ( ||
K(y, (K(x,T)TT((v Diog(n122) + 6).
Combining the results, we conclude that, with probability at least exp(-x8
k-1
D DT HB W SHB) S K(o,, KKss:TVVEe(DDDollerrnm + + K(n, T)
E
Σ
Dia.e (W. (WHB,
HB WSSB).
1=0
An application of the discrcte Gronwall's lemma gives the desired result (94) and concludes the proof.
0
D.4 Proof of Theorem 5.2
Proof of Theorem 5.2. The proof follows from combining Proposition D.1, D.4, D.5 and the fact that: Dr WHBB < (W, WPD) DD W Dr. (W HB w-SH)).
E
Global convergence of the mean-field ODE
In this section, aim to prove the global convergence result through the recipe below:
1. We show the following degeneracy property for the mean-field ODE: there exist deterministic func-
tions wf(,) R220 x R + RD, w R200 x that
x
R
x
R
+ R. (.
R
x R +RR such
wi (t, C1) wf (t, un (0, Cı)), w2(t, C1, C22) w wT (0, C1), uz (0, C1. C2), w3 C2)),
(105)
Wy(t, C2) = w (t, (0, C2)).
We show that (i) (:.:) is continuous in both arguments for any finite t, and that (ti) if wn (0, C1) full support, then (l, C1) full support for any finite
Combining the argument that wn (1, Ci) is full support for all finite and the mode of convergence assumption, we show that the mean-field ODE must converge to the global minimum.
We first show the degeneracy property of the mean-field ODE in the following lemma: Lemma E.1. Under Assumptions (B1) (B3), there erist deterministic functions wi())) : R220 x R
R w)((,)))))) R220 RR x R R + R, wT (-..) R20 R + IR such that
wi(t, 01) = w (t. wi (0, C))),
wz Cr, C2) wj(t, wn (0, Ci), w2(0,, C1 C2), (0, C2)),
wa(1, C2) w (4, w/s w0(( C2))
43

---

Published in Transactions on Machine Learning Research (02/2023)
Proof of Lemma E.1. We follow the proof in (Pham & Nguyen, 2021a, Appendix D.2). To shorten the notations, we make the following definition: we define the sigma-algebras generated by wi (O, Ct)), (0, Ci), w2(0, C1, C2), w2,, lent to prove that wilt, C1), w2(t, C1, C2). wz(t, C2) are SI, S123, Sy-mcasurable, respectively.
wa(0, C2) as S123. S3 respectively. The lemma is equiva-
In order to prove the measurability result, we define a reduced dynamics as follows:
RD (t, = ( 02) -? (uufP(s, - wg (0, c2)) ds - [. / E (x, (2))1)) du ds, 33
RD (t) (0,,,,,,, 02) -7 / (w (2, w2 RD (0,
L: J E (z, C1 C2s (()))(S125) du ds,
w/ (t, w[(((0cc) -7 (s,))
(, (I du ds.
Note the reduced dynamics with w2 w who clearly S123, Sy-measurabl. Furthermore, the re- duced dynamics is not self-contained, in the sense that the gradient terms E [[ (z, Cz; ((t))sss). E | (z, C1, C2; (tt) and E | (z, Ci; (t))ssi are induced by the mean-field ODE W
In order to state the next result, define the following metric:
Dr W max { sup ess ( (t) (t,c))l2,
(E|0.T]
sup (ussuppu2(4,,.00 w/,t, c1,)) rE[o.T] e1.2 sup cessup/wa((,,0 w/(l, 02)1) (E|o,T] e2
Next, we aim to show that the reduced dynamics is equivalent to the mean-field ODE, i.e., for any T> 0,
Dr(W, wDD = 0.
The key step is to prove that
csssup sup
E [AN (z, C2: (())
E.A, (z, (z, C2: W()) (W,,,)),, WRD), WW.WD)
(106)
rE(O,T]
ess sup sup (E(O,T]
IE AAW
(z, C2 C2 W (()) -
K,AY (z,
C2, W (0)) S (W,, (W.WDD) WRD),
(107)
sup sup 1E[0,2]
I
[ (z, G1 W (O))) E.A (z. C3 W (1)ll S (W, (.))))
(108)
where K(,,T)
universal constant depending only on T,7. Here, IE [A (z, C2 (())
E, (z, C2; (t)) is random variable, and the ess sup in (106) is taken with respect to it. The same remark applies to the ess in (107) and in (108), which are intended to be taken with respect to the corresponding random variables.
We now prove that (106) holds. Note that
E | W (z. C2; W(t))ss) - E,4, (z, C2: w
E A (z, C2: - E [A"" (z, Cz; (t))s))
+ E A W (z, Cz W W (WWDD(()))))) (1)) (() - E.A, (z, W (()(
+ E.A (z, C2; W ()) (2;))))())))
(109)
Using the Lipschitz continuous property of AW, have that:
E A (z, C2; (t)) - [A (z, C2; W RD (1))(() | 1 K(y,T)D,,W, RD),
(110)
AA (z, W (1) - A (z. Cz: C2: & K(y, TyD, (W, (W, RRD.
44

---

Published in Transactions on Machine Learning Research (02/2023)
By following the argument in (Pham & Nguyen, 2021a, Appendix D.2) (which does not depend on the dynamics, but only on the structure of the gradient), we have that R W (z, C2; (R) is Sa-measurable, i.e,
E [ (z, C2: W (t))sss - E2A, (z, C2 W (D)( = 0.
(111)
By combining (109), (110) and (111), we obtain that (106) holds. The arguments giving (107) and (108) are analogous.
From this, we can compute the difference between the reduced dynamics and the mean-field ODE as
Dr W <> S.'' D.WW, WD) ds + k(7,1) J. D W RDD du ds,
which, after applying Corollary F.4, gives that Dr W 0. This implies that W W RD and, hence, wi (t. C1), (t, C1, C2), wttt, C2) are S1, S123, Sy Sy-maauraale,, respectively.
Next, we show the continuity of the function wi(:)) :) R2 x RDD + R in both arguments.
Lemma E.2. Under Assumptions (B1) - (B3), we have that, for all E [0, TJ and for all un ui E IR. | (t, un) - w (t', un 112 S ((y,T)) - -t', | (t. u1) wt (t, 21) S K(7, T) u' ll2-
(112)
(113)
Proof. In order to prove the lemma, we first need to derive the dynamics that characterize the evolution of the functions wi (t, u1), wā un, u2, uz), wū((, uz). This dynamics is induced by the mean-field ODE, whose form recall below:
5 (10((t.a)-w3(0c)))) oz( - S/¹ [E 02 W()) do ds,
wa(t, = w3(0,c²)
(114)
w2(t,c1,02) = w200,c1, c1, 02) (10)(( ())0))).))))) ds- f / E.A!" (z, e1, (()) do ds,
(115)
wi((,))) w.) (0, - ? S( (w)(), - wi (0, ds / [/ BLAY" (z, W(v) ds.
(116)
Recall also that wa((, 02) = w (4, c2)) Thus, in order to get the dynamics of w (4, u3), we replace wa (0, 02) by u3. w2(0,, c1, c2) by u2, and wi (0, c1) by un into (114). By doing the same replacements into (115) and (116) for w2(t, c1, c2) and wi (t, respectively, we obtain
uz)t, us -7 / (wā(s, u2) -2) ds 1-' I E.D. us) ds
(1,,, un, U2, uz) -2 7 / (1558, u1, u2,
u2, ua) do ds, 1'' ( E.AW (v, ui, u2,
u2, u3) u22 ds
(1) = u1 ? /''' (w) un) u1) ds 1' L E.AW (v. 2, u1) ds,
where we have the following modified forward and backward paths:
H1,t, u1) = (t, TT)) Hz(t, x. uz) = az w u1 u2 ua)an (H, (t, u))),
f(x; W W(t) Eug uy uz H2,, x, x, us).
A/ (1,,,13) =,)))a(((0 R, / (2; W (1))) æ uz)),
AW
(H2(t, x. va))o((H1()) x, 41)); AWr (1,2,11, W2, Ws, = 8 Rly, / (0; ((1) (H2()
AW W (t,,, = Euy 0 Rz f(: W (t))) us) (H2 x, ua)) un u2, ua)o u)))..
45

---

Published in Transactions on Machine Learning Research (02/2023)
Thus, we have that
(t, u1) - wi(t (t, u1ll22 S (1 mt)lee -uil2 7 %%/ |/00 u2) wi(ss 2i)ll ds
+ + Jo |R (,, -R,A (,, 12)
(117)
An application of Lemma A.2 gives that
|| (v, 21) E W (u, 2, ulll2 & K1(7, (1us u1, u2, uz) - wz(n; u1, 12, uz)l + [e; u1) - w (v, 22))22).
(118)
Similarly for wy, we have that [t, u1, u2, ua) - w un, u2, ua)l S ? L'' w5((, un, u2, us) wā un. u2 ua)
+ / //" |E2A u1, - E (u, z, u1, u2, ua)ll do ds,
(119)
and another application of Lemma A.2 gives that | (u, 2, U1, U2, us) E (v, 2, u1, uz, us)ll2 S K2(7, ()( un, u2, us) wzz (v. 2/, U2 us)l (120)
+1|wi (v. u1) wi u 21)l22).
By combining (117); (118), (119) and (120), we obtain
Iw (f, U2, uz) - (t, f, u/i, u2, us)l + |mll (t, u1) wi (t, 21)ll2
<(1 + +t)llll - - ll2 + S (1w5(s, u1, u2, us) - wr (s. u/, U2 uz)l + || (s, u1) - w u1)lla) ds
+ Kall,, 1) /'' 1'' (lwg(0,
un, un, u2, ua) wz un, u2, ua)l + (v, un) - wi(v u4)ll2) de ds.
Thus, by Corollary F.4, we have that:
[wā(t, u1, u2, uz) wh (t, u/, u2, uz)l + || (t, ui) - w u1) S (4(7,T)llu un which implies that ||wi(t: un) wt, ul)lla & K1(1; T) wil, and concludes the proof of (113). The Lipschitz continuity (112) of wi (t. un) is already proved in Lemma D.2.
At this point, we show that, if w1 (0,c1) c1) :Q + R. has full support, then w (t, u1) has full support. Lemma E.3. Under Assumplions (B1)-(B3) and Assumplion 7.1, we have that wii (1, u1) has full support for any t < co.
Proof. By the continuity argument in Lemma E.2, we have that
|| u1) - u ll2 = || (t, (i ui)lla S K)r,T)t, || un) w (t,, 2/2 S K({,T)
(121)
u1ll2
(122)
We want to show that, for any x E IRP,, there exist a v such that wi, v) = x. For any x E RRD, define a map go (t, v) = x - (wi(t, v) - v). It is easy to see that if a fixed point of gt, -), then w (t, v) = x as
(1. v) = x - -w (1, v) - v) = wi (1, v) = a.
By (121), we have that gn (1,.)) RDD + B(æ, K(7, T))), where B(x, K(T, T)) is the closed ball centered at x with radius K(y, Now, if we restrict (t, v) on B(x, K(y, T)t), we have that it is a map from B(x,, K(y T)t), which is compact set, to itsolf. Furthermore, g v) is continuous in v, since w (t, v) is continuous in by (122). Thus, by the Brouwer fixed point theorem, we have that there exist a fixed point B(æ,, K(y, T)t), which finishes the argument.
46

---

Published in Transactions on Machine Learning Research (02/2023)
Finally, we are ready to prove the main theorem. Our proof follows similar steps that of (Pham & Nguyen, 2021a, Proof of Theorem 8).
Proof of Theorem 7.2. By Assumption 7.1, we have that
lim Ec C1, ( = 0.
7oc
C1
By the definition of A C1, C2), have
lim sss Ec, |IE (W 70000
(()) 0.
Recall from Lemma E.3 that, for all finite t, wi (t, C1) has full support. Hence, we have that, for u1 in a dense subset of RD,
lim 700
Rc, E.A C2 (())oo =
Our aim is to conclude that, for almost all we have that E₂ 12, ((:; ((oo)) 0. By definition of the backward path, we have that
Ec, | (z, C2 W (t))or - E.A
Cz; (uf
<Bos ( (E_AA
(z, C2; (()) E2A Cz: W (u)))
KBc, Ex || (z, C2: ((t) - A." (z, C2; W (co))]]
< <KBcico [(1 + | (00, C2) (1103 C2) Wz C2) + (og C2)l C2, W2(1, C1. C2)
+ |u:s000 C2)1 122 (oo, C1, C2) um (oo, C1) (t. C1)l2)]
By Assumption 7.1, the RHS of (123) converges to 0 t - Hence, by taking the limit on both sides, we have that, for in a dense subset of D
Bc, (z, Cz: (( (u
lim Ec, |E (z, Cz: ((f)or(u()))) x) 0, 700
which implies that, for almost all C2,
E.A. Cz ((oo))o æ) =
By definition of (z, C2; W (2o)) we have that, for almost all
E- [8,R(y, f(; ((o))) (25))) (2), ((oo)))o
x)] =0.
(123)
Note that Assumption (B1) gives that o + 0, and Assumption 7.1 that wa wa(co, 02) with probability 0 (where the probability is intended over C2). Hence, we have that, with probability (over c2).
Eg B2 f(x; w (o0))))0 oz (oxs)))o 0.
(124)
Recall that 01 (uffx) is a function of æc but RgR(y, ((; W (xx))) depends on both and æ. Thus, we can re-write (124) as
Ex E. 102((y, f(x; W (z=)) cz; ((o))) ( = 0.
(125)
Now, we want to use the universal approximation property of of to conclude that, for almost every Ø;
By [0 ((,, ((x; W ((o))) (o W(o))) 0.
(126)
47

---

Published in Transactions on Machine Learning Research (02/2023)
The idea is that linear combinations of 01 x) approximate any function in C2(Da). Thus, if By [0z R(y, (a; ( (2,, ((o))) is in (2(D)), we have that there exist sequence of in- dex sets {1 MEN, such that:
lim Ex k
By 12 f(;; [=))) (22 (()o))
Σ an 4
=0 -0
To simplify the notation, define:
g(x) My 1yy(yR(y, f(x; (co))) df W(o)))),
(xx))) E ai ai (u. x).
From (125) and by linearity of expectation, we have that, for all k
Em (g(x))+ 0.
Thus we have
0 lim
En [g(x)) hx(
200
lim Em
[g(z))2 + (x) (2g())))2))))
7
= lim En [1(())) + [1)(2))
which implies that
Ex
|g(x)) = 0.
Hence, we have that
Ex B. [02R(y f(x: W(oo)) x (H2(H c2 ( =0,
which implies that (126) holds. Furthermore, to that E, [ [0, ((; (z W(x))) is indeed in C2 (D)),, it suffices to note that, by Assumption 3.3,
E, [02 f(x; ( oz (H2 ez ( - K².
By Assumption 3.3, we also have that (x) # 0 for all æ. Hence, (126) implies that, for almost every x,
Ey R(y, f(x: ((o)))) =0
(127)
Since the loss is convex in f(x; W (co)), we have
E- (((, R(y. (a: W (cc))) En Ey [2, f(x; ((xc))) f(x; W(x)))) = 0,
where the last passage follows from (127). Thus, we conclude that
E.R(y ((o))) inf E. [Ry, [())] f
(128)
Finally, we want to show that
lim Eg R(y, f(x; W())) = Eg R(y, f(x; ((o))) 200
(129)
48

---

Published in Transactions on Machine Learning Research (02/2023)
To this, write
B. Rv, (a; (()) Eg Rly, (a; fa; W(co))
<KB2 ((; W (t)) f(:; ((
SKBo, Co (0c, C2) wt C2) + w C2)l | (00, C1, wz(4, C1 C2) + lus, C2)1 (00, C1, C2)1 [m (00, C1) wi (t, Ci)lla],
and use again Assumption 7.1. By combining (128) and (129), we obtain the desired result.
F Technical lemmas
Lemma F.1 (Corollary of McDiarmid inoquality). (Mei et al., 2019, Lemma 30)
Let (Xi) E Rdd be a sequence. of i.id random variables, with ||x:|12 s K and [E[x,] 0, then we have:
T Pr
x 2 =((11/+=)) S eepp(n222 T1
2
Lemma F.2 (Azuma- Hoennain bound). (Mei et al., 2019, Lemma 31) Let (Xx)k20 a martingale taking values in D with respect to the filtration (F&), with Xo O. Assume that the martingale difference time l is Lg- subgaussssiin, which means the following holds almost surely for all X E RR,
{
[ 12111
E Texpf X exp
Then, we have
L (VD+6) +0) S expp(-82}
Pr
max kE(n)
Xxull2 2 2,
Note that, if La S L for all k, then
|x 2 nV (VD+ < exp(-82}
Pr
max
kE(n
Lemma F.3 (Pachpatte's inequality). (Ames & Pachpatte, 1997, Chapter 1, Theorem 1.7.1) Let u, f and be non-negative continuous functions defined on [0, [), for which the inequality
u(t)= < uo ( [(s))(()) ds [[5(s) [(s)) (% ((rr)(())))) dr) ds
holds, where uo is a non-negative constant. Then have:
u(t) S uo 1+ L'' e(s) (f + f(r)) dr)
ds
Corollary F.4 (Pachpatte's inequality for constants). Let u be a non-negative continuous function defined on [0, 2J, and y, K be positive real numbers. Assume the following inequality holds:
u(t) S no + S( u(()) ((u(s) ds + K S. u(() dr ds.
Then, we have
72
exp (22 K (***K))) S un 1 exe
(2+ (2+K1) 1)
u(() S un
1 72 + K
7
49


Reviews:
[{"review_id": "yKAKDsFCpY", "reviewer": "Reviewer_ESTq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The results on the heavy ball method presented in the paper are technical, interesting and seem to be new in the literature.  The earlier result on this method by Krichene et al. is only for noisy dynamics, while the present paper considers the noiseless dynamics with the only randomness from the stochastic gradient. Among all the results, the most interesting and novel one to me is the global-in-time convergence of the mean field ODE for three-layer networks.  \n\nWeaknesses: Despite the originality of the results, I do have some concerns and comments. My first concern is about the values and implications of the mean field limit of heavy ball method. The authors discussed the use of mean field results in dropout stability and connectivity. But I was not fully convinced by the discussions. I would imagine one can prove similar mean field results for many other training algorithms. Do they all enjoy dropout stability and connectivity?  Do these mean field results provide any insights on the pros or cons of heavy ball method compared to GD or SGD?   \n\nAt a technical level, the proof strategies used in the paper are standard and not original. The proofs of mean field limit follow from the standard coupling strategy with minor twists from dealing with the stochastic gradient and multiple layers. The proof of global convergence seems also to be an adaption of the paper by Pham & Nguyen. Part (C3) of Assumption 7.1 needs a proof rather than being made as an assumption. With this assumption, the result looks pretty trivial to me. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides a mean-field analysis of stochastic  heavy ball (SHB) method for training two-layer and three-layer neural networks. The authors proved that the SHB method converges to a PDE in the large width limit. Quantitative bounds between the dynamics of SHB and the mean-field dynamics were derived. These result extend the earlier mean field results of SGD, which is a first order dynamics, to the setting of second order dynamics. Dropout stability and connectivity were discussed as consequences of the mean field analysis.  The authors also obtained the global convergence of the mean-field ODE for three layer networks. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would strongly suggest the authors making a thorough discussion on the  values and implications of the mean field results by addressing the concerns raised above. ", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "5XIWv3IUFy", "reviewer": "Reviewer_LtDr", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The writing is precise and easy to follow. I checked the math, and it seemed correct.\n\nThe submission is certainly of interest to some readers of TMLR, since stochastic heavy ball methods are more commonly used in practice than vanilla SGD without momentum. Furthermore, analyses of the 2-layer SGD+momentum mean-field limit have been performed in the past, but without quantitative bounds on convergence -- so these quantitative bounds on convergence are new and will be appreciated by some members of the community.\n\nThis work seems like a good reference that other papers will be able to use whenever they want to analyze mean-field dynamics with momentum.\n\nWeaknesses: The results in this paper mainly seem like modifications of existing analyses of mean-field dynamics. I am unsure what is qualitatively new about the momentum term. I suppose that a take-away of this paper is that one can do the same analyses as before, but adding a momentum term, and not much changes.\n\nThe proofs of the claims in Section 6 are missing, although it is clear that they follow in a straightforward way from previous work + the new quantitative analyses of convergence. (So this is not so important.)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies training 2-layer and 3-layer neural networks with SGD with momentum. The neural networks are parametrized in the mean-field regime. The 2-layer parametrization of the network follows Mei et al'18, Chizat-Bach'18, Rotskoff-Vanden-Eijnden'18, and Sirignano-Spiriopoulous'20. The 3-layer parametrization follows Nguyen-Pham'20.\n\n(1) This paper proves a non-asymptotic error bound between the limiting mean-field dynamics and the dynamics of SGD+momentum as the number of neurons tends to infinity and step size tends to 0. The novelty relative to previous works is incorporating the additional momentum term in SGD.\n\n(2) It is claimed that the solutions of SGD+momentum have dropout-stability and connectivity. This extends Shevchenko-Mondelli'20 to dynamics with the additional momentum term. (The proof of these claims is omitted, but follows in a straightforward way from Shevchenko-Mondelli'20 and the new quantitative convergence bounds.)\n\n(3) It is shown that if the 3-layer network mean-field dynamics converge, then they converge to a global minimizer of the risk. The novelty relative to Pham-Nguyen is to incorporate momentum.\n\n\n\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The paper was well-written, so I do not think anything needs to be changed. However, it could be interesting to also mention in the literature review how the three-layer mean-field parametrization relates to the \"maximal-update\" parametrization of the Yang-Hu'21 paper which was also cited. Is the Yang-Hu parametrization incomparable to the one in this paper?\n\nTypo:\n-- In page 11, the second equation from the top should read $w_2*(t, w_1(0,C_1), w_2(0,C_1,C_2), w_3(0,C_2))$ instead? Where $w_2$ has been replaced by $w_3$.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "QkTLMoKnjW", "reviewer": "Reviewer_Vu68", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Let me first start by pointing out that I am unfortunately not a specialist of the mean-field regime setting as I have not personally worked on it. It is highly probable that some of my following comments are irrelevant for an \"mean-field expert\". However I believe that the paper should still be easily readable for someone (such as myself) which is familiar with the general NNs + momentum setting, but less with the mean-field setting.\n\nStrengths: \n- trying to understand the heavy ball method in the context of training NNs is very relevant to the community\n- the proposed results are non-trivial and seem to be technically involved \n- Theorem 5.1 and 5.2: obtaining non asymptotic results on finite-width networks will interest the community\n\nWeaknesses / questions:\nOverall when reading the paper I had a lot of difficulty to understand the setting, the notations and some equations as many objects are not properly introduced and there are many missing assumptions. To cite a few:\n- Beginning of section 3.2: the loss $R(y, \\hat{y})$ is never properly defined, and what are the assumptions on it ? is it convex as in (Krichene et al. 2020) ? is it differentiable ? \n- just before equation(4): parameter $\\gamma$ is never introduced, what is it ?\n- equation (5): $\\hat{\\nabla}R(W(t))$ does not seem to be defined, why is there a \"hat\" here ?\n- section 4: $\\mathbb{E}_{\\theta \\sim \\hat{\\rho}_\\theta}$ is an odd notation. The abuse of notation over $\\theta$ (vector vs random variable) is a bit misleading and should be explained or changed. \n- equation (11): all the definitions would be easier to understand if it was clearly stated what are the domain definitions of the introduced quantities.\n- before equation (13): what do you mean by \"(stochastic) ODE\" ?\n- equation (13): what does $W(t)$ correspond to once in the mean field regime ?\n- equation (16): what is $z$ ? should it be $z_k$ ?\n\nI have no doubt that nothing is wrong in all this, but it would be much easier for the reader to follow if all the objects are properly and rigorously defined. Also, the sloppy notations gives the impression that the tools are not properly handled. \n\nConcerning the results:\n- eq (12): are there any results showing the consistency of this ODE ( Theorem 1 in Krichene et al (2020) but in the noiseless case) ? if not then what is the validity of considering this regime ? \n-  Theorem 4.1 for the case of 2 layers: in Krichene et al (2020), to show the existence and uniqueness of solutions of their dynamics, they need an additional assumption on the initial distribution ((A5) in their paper), which you don't seem to need. Can you comment on this ?\n- Section 6: given the fact that $K(\\gamma, T)$ is exceedingly big, it is hard to deduce / believe from (25) a \"provable justification to two remarkable properties exhibited by solutions obtained via gradient-based methods, namely, dropout-stability and connectivity\" as stated.\n- some numerical simulations explaining / highlighting your results would have been much appreciated  \n\nMinor comments:\n- in the intro: \"This perspective has facilitated the study of architectures with multiple layers (Araújo et al., 2019; Lu et al., 2020; Nguyen & Pham, 2020; Fang et al., 2021), and it has given a rigorous justification to a number of properties displayed by SGD solutions, including convergence towards a global optimum\": I believe saying \"a rigorous justification\" is in my optinion an overclaim: there are still many things which are not at all understood concerning the global convergence of discrete GD and even more discrete of SGD !\n- top of page 4: should be N instead of D in $W \\in \\mathbb{R}^{n D + D}$\n- it is not clear to me why the mean-field ODE for 3 layers (equation 3) is written in a complete different way than in the case of 2-layers (eq 12)\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors analyse the stochastic heavy ball (SHB) method on 2-layer and 3-layer fully connected neural networks with non-linear activation. They consider the mean-field limit where the number of neurons grow to infinity. They show existence + uniqueness of the solutions of the considered mean-field ODEs and provide non-asymptotic bounds on the error between the discrete finite-width SHB dynamics and the continuous mean-field limit.  They then leverage this result to show dropout stability and connectivity of the iterates after k steps of  SHB.\n\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "As explained in the previous section, my major concerns are on the lack of rigor in definitions and notations.\n\nAlso I am skeptical about the conclusions which can be drawn from the bounds from theorem 5.1 and 5.2 given the fact the dependency on $T$ is prohibitive to drawing any meaningful conclusion on finite NNs. ", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper is well-written, after a round of revision three expert reviewers think that the paper has solid contributions. The authors have addressed a few minor concerns, added additional discussions, and several claims have been properly worded, which improved the quality of the paper even further. I agree with the reviewers and suggest accepting this paper.

Decision: Accept (poster)