Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction
Authors: Taiji Suzuki
Denny Wu
Atsushi Nitanda
Venue: NeurIPS 2023 spotlight (2023)
Keywords: mean-field regime
interacting particle system
propagation of chaos
Neural network optimization
MMD minimization

Abstract: The mean-field Langevin dynamics (MFLD) is a nonlinear generalization of the Langevin dynamics that incorporates a distribution-dependent drift, and it naturally arises from the optimization of two-layer neural networks via (noisy) gradient descent. Recent works have shown that MFLD globally minimizes an entropy-regularized convex functional in the space of measures. However, all prior analyses assumed the infinite-particle or continuous-time limit, and cannot handle stochastic gradient updates. We provide a general framework to prove a uniform-in-time propagation of chaos for MFLD that takes into account the errors due to finite-particle approximation, time-discretization, and stochastic gradient. To demonstrate the wide applicability of our framework, we establish quantitative convergence rate guarantees to the regularized global optimal solution for $(i)$ a wide range of learning problems such as mean-field neural network and MMD minimization, and $(ii)$ different gradient estimators including SGD and SVRG. Despite the generality of our results, we achieve an improved convergence rate in both the SGD and SVRG settings when specialized to the standard Langevin dynamics.

Convergence of Mean-field Langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction
Taiji Suzuki?:?, Denny Wu Atsushi Nitanda
University of Tokyo,
"RIKEN AIP, "New York University,
FFlatioon Institute,
Kyyushu Institute of Technology
taiji@mist wtw.sit....ttoooo.a......o
dennyru@nyu.cm edu, kynndeaai.kyttehhi ac
Abstract
The mean-field Langevin dynamics (MFLD) is a nonlinear generalization of the Langevin dynamics that incorporates a distribution-dependen drift, and it naturally arises from the optimization of two-layer neural networks via (noisy) gradient descent. Recent works have shown that MFLD globally minimizes an entropy- regularized convex functional in the space of measures. However, all prior analyses assumed the infinite-particle or continuous-time limit, and cannot handle stochastic gradient updates. We provide general framework to prove uniform-in-time propagation of chaos for MFLD that takes into account the errors due to finite- particle approximation, cume-discretization and stochastic gradient. To demonstrate the wide applicability of our framework, we establish quantitative convergence rate guarantees to the regularized global optimal solution for (i) wide range of learning problems such as mean-fiell -ield neural network and MMD minimization, and (ii) different gradient estimators including SGD and SVRG. Despite the generality of our results, we achieve an improved convergence rate in both the SGD and SVRG settings when specialized to the standard Langevin dynamics.
Introduction
In this work we consider the mean-field Langevin dynamics (MFLD) given by the following McKean- Vlasov stochastic differential equation:
dX, - 5F(p)) (X,) 22ddW,
where f4 Law(X)) F P2 H is a convex functional, W, is the d-dimensional standard Brownian motion, and &F denotes the first-variation of F. Importantly, MFLD is the Wasserstein gradient flow that minimizes entropy regularrreul convex functional as follows:
min [F(u) + AEnt(p)).
pEP2
While the above objective can also be solved using other methods such as double-loop algorithms based on iterative linearization (Nitanda al., 2020, 2023; Oko al., 2022), MFLD remains attractive due to its simple structure and connection to neural network optimization. Specifically, the learning of two-layer neural networks can be lifted into an infinite-dimensional optimization problem in the space of measures (i.e., the mean-field limif), for which the convexity of loss function can be exploited show the global convergence gradient- 2018; Chizat and Bach, 2018; Rotskoff and Vanden Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020). Under this viewpoint, MFLD Eq. (1) corresponds the continuous time limit of the noisy gradient descent update on an infinite- width neural network, where the injected Gaussian noise encourages "exploration" and facilities global convergence (Mei al., 2018; Hu al., 2019).
optimization (Nitanda and Suzuki, 2017; Mei et al.,
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Quantitative analysis of MFLD. Most existing convergence results of neural networks in the mean-field regime are qualitative nature, that is, they do not characterize the rate of convergence nor the discretization error. A noticeable exception is the recent analysis of MFLD by Nitanda et al. (2022); Chizat (2022), where the authors proved exponential convergence to the optimal solution of Eq. (2) under a logarithmic Sobolev inequality (LSI) that can be verified in various settings including regularized empirical risk minimization using neural networks. This being said, there is still large gap between the ideal MFLD analyzed in prior works and a feasible algorithm. In practice, we parameterize μl as a mixture of N particles (X+) !=1 - this corresponds to a neural network with N neurons, and perform discrete-time update: at time step k, the update the i-th particle is given as
x = X - me V 8P(()) (xL)
v S
where 7e is the step size at the k-th iteration, E is an iid. standard Gaussian vector, and V (w) represents a potentially inexact (e.g., stochastic) gradient.
6y
Comparing Eq. (1) and Eq. (3), we observe the following discrepancies between the ideal MFLD and the implementable noisy partiele gradient descent algorithm.
(i) Particle approximation. i entirely represented by a finite set of particles: fa = N E 8 6xp-
Time discretization. We employ discrete gradient descent update opposed to gradient flow.
(iii) Stochastic gradient. In many practical settings, it is computationally prohibitive to obtain the exact gradient update, and hence it is preferable to adopt stochastic estimate of the gradient.
The control of finite-particle error (point (i)) is referred to as propagation of chaos (Sznitman, 1991) (see also Lacker (2021) and references therein). In the context of mean-field neural networks, discretization error bounds in prior works usually grow exponentially in time (Mei et al., 2018; Javanmard et al., 2019; De Bortoli et al., 2020), unless one introduces additional assumptions on the dynamics that are difficult to verify (Chen et al., 2020). Consequently, convergence guarantee in the continuous limit cannot be transferred to the snite-particle setting unless the time horizon is very short (c.g., Abbe et al. (2022)), which limits the applicability of the theory.
Very recently, Chen et al. (2022); Suzuki et al. (2023) established uniform-in--iee in-time propagation of chaos for MFLD, i.e., the "distance" between the N N-particle system and the infinite- carantte-prticleeeee limit is of order O(1/N) for all t > O. While this represents significant step towards an optimization theory for practical finite-width neural networks in the mean-field regime, these results assumed the continuous-time limit and access to exact gradient, thus cannot cover points (ii) and (iii).
In contrast, for the standard gradient Langevin dynamics (L.D) without the mean-field interactions (which is special case of MFLD Eq. (1) by setting F' to be linear functional), the time discretization is well-understood (see e.g. Dalalyan (2014); Vempala and Wibisono (2019); Chewi et al. (2021)), and its stochastic gradient variant (Welling and Tch, 2011; Ma et al., 2015), including ones that employ variance-reduced gradient estimators (Dubey et al., 2016; Zou et al., 2018; Kinoshita and Suzuki, 2022), have also been extensively studied.
The gap between convergence analyses of LD and MFLD motivates us to ask the following question.
Can we develop complete non-asymptotic convergence theory for MFLD that takes into account points (i) (iii), and provide further refinement over existing results when specialized to LD?
1.1 Our Contributions
We present a unifying framework to establish uniform-in-time convergence guarantees for the mean- field Langevin dynamics under time and space discretization, simultaneously addressing points (i)-(it)). The convergence rate is exponential up to an error that vanishes when the step size and stochastic gradient variance tend to 0, and the number of particles N tends infinity. Moreover, our proof is based on an LSI condition analogous to Nitanda et al. (2022); Chizat (2022), which is satisfied in wide range of regularized risk minimization problems. The advantages of our analysis summarized as follows.
Our framework provides a unified treatment of different gradient estimators. Concretely, we estab- lish convergence rate of MFLD with stochastic gradient and stochastie variance reduced gradient (Johnson and Zhang, 2013). While it is far from trivial to derive a tight bound on the stochastic
2

---

Method (authors)
of particles
Total complexity
Single loop Mean-field
PDA+
(Nitanda el al., 2021)
log(n)
G
x
P-SDCA
log(n)
G. log(()
V
(Oko et al., 2022)
GL.D
log(t (Aa)a
x
(Vempala and Wibisono, 2019)
SVRG-LD
+ n
logie-1) (Aa)=
v
x
(Kinoshita and Suzuki, 2022)
F-MFLD (ours)
6-1
n.E,
logg(e")
(Aa)
SGD-MFLD* (ours)
6-1
E,
log(t")
v
(Aa)
SGD-MFLD* (ii) (ours)
( VxE,) logie
/
[Aa]?
SVRG-MFLD (ours)
c-1
logie VTE
+n
(1x)
SVRG-MFLD (ii) (ours)
(n TI E lngie
(Aa)
Table 1: Comparison of computational complexity to optimize an entropy-regulatee
finite-sum objective up to excess objective value c, in terms of dataset size n, entropy regularization X, and I.SI constant a. Label indicates the online setting, and the unlabeled methods are tailored to the finite-sum setting. "Mean-field" indicates the presence of particle interactions. "Singlle loop" indicates whether the algorithm requires inner-loop MCMC sampling sub-routine every step. "(i)" indicates convergence rate under additional smoothness condition Vame For double-loop algorithms (PDA and P-SDCA), G" the number of gradient evaluations required for MCMC sampling: for example, for MALA (Metropolis-adjusted Langevin algorithm) G. = O 5/2 log(1 (1 l ))2) and for LMC (Langevin Monte Carlo) G. O(n(ae) log(e)).
(Assumption 4), where E. ac +
1.2
gradient approximation error because it requires evaluating the correlation between the randomness of each gradient and the updated parameter distribution, we are able to show that stochastic gradient effectively improves the computational complexity.
Noticeably, despite the fact that theorem simultaneously handles wider class of F, when specialized to standard Langevin dynamics (i.e., when F is linear), we recover state-of-the-art con- vergence rates for LD; moreover, by introducing an additional mild assumption on the smoothness of the objective, our analysis can significantly improve upon existing convergence guarantees.
Our analysis greatly extends the recent works of Chen et al. (2022); Suzuki et al. (2023), in that our propagation of chaos result covers the discrete time setting while the discretization error can still be controlled uniformly over time, i.e., the finite particle approximation error does not blow up as increases. Noticeably, we do not impose weak interaction large noise conditions that are common in the literature (c.g.. Delarue and Tse (2021); Lacker and Flem (2022)); instead, our theorem remains valid for any regularization strength.
2 Problem Setting
Consider the set of probability measure P on Ied where the Borel r-algebra is equipped. Our goal is to find probability measure y E P that approximately minimizes the objective given by
F(u) (p) + [r]
where U: P + IR is a (convex) loss function, and r Rd - IR is a regularization term. Let Pa be the set of probability measures with the finite second moment. In the following, we consider the setting where F(H)) < C(1 +
and focus on Pa so that F is well-defined.
As previously mentioned, important application of this minimization problem is the learning of two-layer neural network in the mean-field regime. Suppose that h2 (-) is a neuron with a parameter E Rdd e.g.. (z) aww" T + b). where E Rd-1.b R. and (w, v). The mean-field neural network corresponding to probability measure μ P can be written fn((-) = h2 (-)u(dx). Given training data (2,, 1.) E ddd-1 R. we may define the empirical risk of fn as Up)) =
n 1-1 (2), y.) for a loss function É: R x R + R. Then, the objective F becomes
m
F(a) ((..())))))) + +A1 [1ol)2du(())
3

---

where the regularization term is r(z) = |z|2. Note that the same objective can be defined for expected risk minimization. We defer additional examples the last of this section.
One effective way to solve the above objective is the mean-field Langevin dynamics (MFLD), which optimizes F via noisy gradient descent update. To define MFLD,, need to introduce the first-variatonn of the functional F.
Definition 1. Let P2 + TR. The first-variation of a functional G P2 - at jr E
by P2 defined as a continuous functional P2 x R² + R that satisfies lim,-0 Glev+(l-c)u) I 5C Su (p)(z))((v - ()) for any v E Pz. there exists such functional &G first- first-vrriatinn at μt, or simply G is differentiable at p.
we say G admits bp
To avoid the ambiguity of 5G up to constant shift, we follow the convention of imposing I 6G (p)dy oy
by
Using the first-variation of F, the MFLD is given by the following stochastic differential equation:
dx,- - - (F()) (X,)dt + 2Xd,,, 14 = Law (Xt),
6u
where Xo ~ 140, (X) denotes the distribution of the random variable X and (We) is the d-dimensional standard Brownian motion. Readers may refer to Huang al. (2021) for the existence and uniqueness of the solution. MFLD is an instance of distribution-depeedeen SDE because the drift term (F(m) (-) depends on the distribution μr of the current solution X, (Kahn and Harris, 1951; Kac, 1956; McKean, 1966). It is known that the MFLD is a Wasserstein gradient flow to minimize the following regulii-ee objective (Mei et al., 2018; Hu et al., 2019; Nitanda al., 2022; Chizat, 2022):
{(p) F()) AEnt(u), where (n)( - (z)dd)) is the negative entropy of p.
(5)
Reduction to standard Langevin dynamics. Note that the MFLD reduces to the standard gradient Langevin dynamics (LD) when F is linear functional, that is, there exists V such that F(u)
J' [z))ddp(z) In this case, 6y
&F V for any μ and the MFLD Eq. (4) simplifies to
dx, VV 2XdW.
This is exactly the gradient Langevin dynamics for optimizing V or sampling from μ x exp(-V/A).
2.1 Some Applications of MFLD
Here, we introduce few examples that can be approximately solved via MFLD.
Example 1 (Two-layer neural network in mean-field field regime.). Let h2) be neuron with parameter
& E Rd, e.g., (z) = tanh(ro T x)), (=) = tanh(r)o x) for x = (r, w.) or simply
(2) = ozx T 2). Then the learning of mean-fieed neural network fa(-) = f h (-)((zz) via minimizing the empirical risk U(u) = Σ (zf( y1) with a convex loss (e.g.. the logistic loss ((f.v) = log(1 + exp(-yf)), or the squared loss {(f.y) = (ʃ - 2) falls into our framework.
Example 2 (Density estimation via MMD minimization). For a positive definite kernel k. the Maximum Mean Discrepancy (MMD) (Gretton et al., 2006) between two probability measures and v is defined as MMD² (p,, v) := =J J [[(xx,,)) x) - 2k(r, u) + k(y, ())dd(())d(())) We perform nonparametric density estimation by fitting a Gaussian mixture model (z) = g (=)d((())) where Ge is the Gaussian density with mean æ and a given variance o² > O. The mixture model is learned by minimizing MMD2 (f
p") where p" is the target distribution. If we observe set of training data (= from p", then the empirical version of MMD is one suitable loss function U(4) given as
(M (z))))))))))))())))))))))))))))) 2 d')d (edpdddddd(((xpppee [(,xx))(x, -2fr)))))(())) " (2)
Note that we may also choose to directly fit the particles to the data (instead of a Gaussian mixture model), that is, we use a Dirac measure for each particle, as in Chizat (2022, Section 5.2) and Arbel et al. (2019). Here we state the Gaussian parameterization for the purpose of density estimation.

---

Example 3 (Kernel Stein discrepancy minimization). In settings where we have access to the target distribution through the score function (e.g., sampling from posterior distribution Bayesian infer- ence), we may employ the kernel Stein discrepancy (KSD) as the discrepancy measure (Chwialkowski al., 2016; Liu al., 2016). Suppose that we can compute s(2) v log(n=)(()) then for a positive definite kernel k, we define the Stein kernel as
Wi (2,2') == s()) a(=')k(2, =') + s(z) V (=,(2,z2)) + V., + V.
We take U(+) as the KSD between 1 and μ² defined as KSD(") = f [ =")dp("))pp(((""" By minimizing this objective via MFLD,, we attain kernel Stein variational inference with convergence guarantees. This can be seen as a "Langevin" version of KSD descent in Korba et al. (2021).
Remark 1. In the above examples, we introduce additional regularization terms in the objective Eq. (5) to establish the convergence rate of MFLD, in exchange for slight optimization bias. Note that these added regularizations often have statistical benefit due to the smoothing effect.
2.2 Practical implementations of MFLD
Although the convergence of MFLD (Eq. (4)) has been studied in prior works (Hu et al., 2019; Nitanda et al., 2022; Chizat, 2022), there a large gap between the ideal dynamics and a practical implementable algorithm. Specifically, need to consider (i) the finite maiite-article approximation, (ti) the time discretization, and (iti) stochastic gradient.
To this end, we consider the following space- and lime-discretized version of the MFLD with stochas- tic gradient update. For a finite set of particles 2 = (x)) C Rdd we define its corresponding empirical distribution as f.zz + > ōx,- Let 21 = (X C Rd be N particles at the k-th
update, and define μı = 11x; finite partiele approximation of the population counterpart. Starting from X ~ Po. we update 2i as,
X X nauk + 22
(6)
where T7& is the step size, & is an i..d. standard normal random variable 4 N0 I). and oj. = of{ (Zo::: w is a stochastic approximation of V y (Xi) where Wi = (w is a random variable generating the randomness of stochastic gradient. of can depend on the history Loo = (......, K2), and E [ot] Lo = v (F(we) (Xi). We analyze three versions of of.
8p
(1) Full gradient: F-MFLD. If we have access to the exact gradient, we may compute
(F(()) (X().
o/ V
8fz
Suppose the loss function U is given by an expectation as U (m) = Ez-P2 [e(H, 2),, where Z E Z is a random observation obeying a distribution Pz and e: P.P. x Z - IR is loss function. In this setting, we construct the stochastic gradient
(2) Stochastic gradient: SGD-MFLD.
56 5e(1u,, +3)) (x() vr(X)),
M
B
5p
where (= 1==1 is mini-batch of size generated from Pz in an i.i.d. manner.
(3) Stochastic variance reduced gradient: SVRG-MFLD. Suppose that the loss function U is given by finite sum of loss functions l; U(+) n Ei,, (e(w), which corresponds to the empirical risk in a usual machine learning setting. Then, the variance reduced stochastic gradient (SVRG) (Johnson and Zhang, 2013) is defined as follows:
ii) When & = 0 mod m, where m is the update frequency. we set 2 = (X:)) (Xi) the anchor point and refresh the stochastic gradient as v v
(F(m))) (XE).
iii) When & * 0 mod m, use the anchor point to construct control variate and compute
1
5l, (6z(12))
86 (Xxi) + (0())) (x))
c B
MV
(Xi) +(()))
Sjc
JETN
where 1 is a uniformly drawn random subset of (l....,n) nJ with size B without duplication.
5

---

3 Main Assumptions and Theoretical Tools
For our convergence analysis, make the following assumptions which inherited from prior works in the literature (Nitanda et al., 2020, 2022; Chizat, 2022; Oko et al., 2022; Chen et al., 2022).
Assumption 1. The loss function U and the regularization term are convex. Specifically,
U: P + RRir is a convex functional on P. that is, (0p + (1-00y) I (U(μ) + (1 - (V) any 0 [0. and u.v P. Moreover, U admits a first-variation any μ E P2.
r(.) is twice differentiable and convex, and there exist constants A1. X2 > and cr > such that A < VV r(x) 3 A21, x Vr(x)= 2 >11|=/1/? and 0 r(x) S X2(e, + [|=]](]]] for any Rd and Vr( = O.
V p (x) - V (x') S L(W2(H, + ||
Assumption 2. There exists L > such that
r'll) and (0U(m) (I,, 2²) 6n
< T(1 + cz. (11x=]1
for any ps n E P2 and x, E TRd. Also, there
exists R > 0 such that || (W(m) op
(z)]] S for any μ and x E 124
Verification of this assumption in the three examples (Examples 1, 2 and 3) is given in Appendix A in the supplementary material. We remark that the assumption on the second order variation is only required to derive the discretization error corresponding to the uniform-in in-time propagation of chaos (Lemma 8 in Appendix E.4). Under these assumptions, Proposition 2.5 of Hu et al. (2019) yields that I has unique minimizer μ² in P and μ² is absolutely continuous with respect the Lebesgue (F()") b4
measure. Moreover, J satisfies the self-consistent condition: μ² (X) exp (-x
(X))
3.1 Proximal Gibbs Measure & Logarithmic Sobolev inequality
The aforementioned self-consistent relation motivates to introduce the proximal Gibbs distribution (Nitanda et al., 2022; Chizat, 2022) whose density is given by
1 1 (x))
p& (X) x exp (-x
5μ
where 2 = W C Rd is a set of particles and E Rdd As we will see, the convergence of MFLD heavily depends logarithmic Sobolev inequality (LSI) on the proximal Gibbs measure. Definition 2 (Logarithmic Sobolev inequality). Let μ be a probability measure on (Rd, B(RR)) f satisfies the LSI with constant a > for any smooth function fonottfnnction & Rd - IR with E [ < 0o, we
have F.I log( - [0
S [[,
This is equivalent to the condition that the KI. divergence from μ is bounded by the Fisher divergence: log(de/ /djdde respect to μ.. Our analysis requires that the proximal Gibbs distribution satisfies the LSI as follows.
2/2 V (dg///d/))) dd for any E P which is absolutely continuous with
Assumption 3. n² and p.v. satisfy the LSI with
for any set of particles L (X C Rdd
Verification of LSI. The LSI of proximal Gibbs measure can be established via standard perturba- tion criteria. For U (1) with bounded first-variaton,, we may apply the classical Bakry-Emery and Dolley-Strrock arguments (Bakry and Émery, 1985a; Holley and Stroock, 1987) (see also Corollary 5.7.2 and 5.1.7 of Bakry et al. (2014)). Whereas for Lipschitz perturbations, we employ Miclo's 's trick (Bardet et al., 2018) or the more recent perturbation results in Cattiaux and Guillin (2022).
Theorem 1. Under Assumptions and 2, μe' and p.z. satisfy the log- hoo-oooleev inequality with
A
4R² YrK
{AA
AA JATx
n2 +
2
R
2x
exp
V
2d/x
V
A1
2+
g log
X1
A1
Furthermore,
(U()) 1l
S R satisfied for any i E Pz. then μ²' and pxr satisfy the LSI with
op
4ft
exp
See Lemma for the proof of the first assertion, and Section A. of Nitanda et al. (2022) or Proposition 5.1 of Chizat (2022) for the proof of the second assertion.

---

4 Main Result: Convergence Analysis
In this section, we present our the convergence rate analysis of the discretized dynamics Eq. (6). To derive the convergence rate, need to evaluate the errors induced by three approximations: (i) time discretization, (ii) particle approximation, and (iii) stochastic gradient.
4.1 General Recipe for Discretization Error Control
Note that for the finite-particle setting, the entropy term does not make sense because the negative entropy is not well-defined for a discrete empirical measure. Instead, we consider the distribution of particles, that is, let y(() P(N) be a distribution N particles 2 = (X where P(N)
the set of probability measures on (Rdx,, B(RdxN)). Similarly, we introduce the following objective on (M))
9e (1( = NE ~ [F(xx)] (Ent(u(()).
One can casily verify that if (x)) is a product measure of μ E P, then &N (n())) 2 (F) by the convexity of F. Propagation of chaos (Sznitman, 1991) refers to the phenomenon that, as the number of particles increases, the particles behave as they are independent; in other words, the joint distribution of the N particles becomes "close" to the product measure. Consequently, the minimum of I (n((N) (which can be obtained by the artiile-approxiimaee MFLD) is close to N N (μ²) Specifically, it has been shown in Chen et al. (2022) that
inf E(NDEPIN) N
1
I N -
F("))) Cx N
for some constant C3 > 0 (see Section for more details). Importantly, if we consider the Wasserstein gradient flow P(N), the convergence rate of which depends on the logarithmic Sobolev inequality (Assumption 3), we need to ensure that the LSI constant does not deteriorate as N increases. Fortunately. the propagation of chaos and the tensorization of LSI entail that the LSI constant with respect to the objective J. can be uniformly bounded over all choices of N (see Eq. (11) in the appendix).
To deal with the time discretization error, we build upon the one-step interpolation argument from Vempala and Wibisono (2019) which analyzed the vanilla gradient Langevin dynamics (see also Nitanda et al. (2022) for its application to the infinite-particle MFLD).
Bounding the enror induced by the stochastic gradient approximation is also challenging because the objective is defined on the space of probability measures, and thus techniques in finite-dimensional settings cannot be utilized in straightforward manner. Roughly speaking, this error is characterized In addition, to obtain refined evaluation, we also incorporate the following smoothness assumption.
by the variance of ui: o² == E [ V 6F(F)) [|a] I. öy
Assumption 4. (f w/) and D' F(2A)) v 8F(ux.) 6y and, for either G(22)) = ( wj) G( Di, (2k) function of Fh. it is satisfied that [G((2)) G( Σn VT
(Xt) are differentiable w.r.t. (X
x/ ((()) -X- x(DII Q Q 1X - +X X X XiII2
for
2N
some Q > ( and Vi. = (X and 24 (x 1ー=1- We also assume R.,, [ x (ooks wt)T T-
V xt Di ((2)) 2 [p] < Nº - and || (Li,, Lo 20 w/) - 0f ((27 L x 1): w 5
N2 22
2 Mzz)
+ +X II)...
Note that this assumption is satisfied if the gradient is twice differentiable and the second derivative is bounded. The factor N appears because the contribution of each particle is O(1/N) unless i In the following, we present both the basic convergence result under Assumption 2, and the improved rate under the additional Assumption 4.
After the acceptance of this manuscript, we noticed that by replacing o².u with square a conditional expectation Ex....4 [ 1oo Ro..eoalee [vv 4F(m2) ][["], we can obtain a refined bound as in Das et (2023) suggested by anonymous reviewer Y8on. We defer such refined analysis to future work.
7

---

Taking these factors into consideration, we can evaluate the decrease in the objective value after (N)
one-step update. Let Mis
E p(N) be the distribution of Li conditioned on the sequence wo:k
(ww =o Define
R = R EII
11
1
[(
+
1 x
(R² + X2cx) x2c)) Ad _d], o, CT Xn))
A1
422
where C1 = 8(R² >z(cr R2) + d) and I, = I. x²?
Theorem 2. Under Assumptions 1, and 3, mı. S 511/(422), we have
-F((e) where Yx 5m + (R+ X2 (1 v1/m))) (NRo,,Ook + Qn r + o with Assumption 4, and Yx =d o. T without this additional assumption; the expectation taken with respect to the randommess (wx) = (( of the stochastic gradient; and Cx =
E I -
Aame 2
$ | (ub - + nn (8m + G2 ++)++ TKk +
21Lo(11 + 201 R2) + 122
The proof can be found in Appendix B. We remark that to derive the bound for Y₁ = = o² is relatively straightforward, but to derive tighter bound with Assumption 4 is technically challenging, because we need to evaluate how the next-step distribution (N) is correlated with the stochastic gradient vig. Evaluating such correlations non-trivial because the randomness is induced not only by Wh but also by the Gaussian noise. Thanks to this general result, we only need to evaluate the variance k and u. for each method to obtain the specific convergence rate.
/i+1
Conversion to a Wasserstein distance bound. As a consequence of the bound on 1 I N (#4 (N))
N
F (f²)), we can control the Wasserstein distance between μu (,)) and ,º,, where N P(N) the (N-times) product measure of 4=" Let (2(p,,) 1/) be the 2- Wassersteen distance between and v, then
(N) p N) W
2 (5 Xce
( (x))))),
(Mz
under Assumptions 1 and (see Lemma 3 in the Appendix). Hence, if N I NN
(14 (1( F(#")
is small, the particles (X behaves like i.i.d. sample from μ², As example, for the mean-field neural network setting, if (z) ler L| x'll (Vx, x' E IR)) and Vi- (f (z))
(z)) for a fixed 2, then Lemma 4 in the appendix yields that
E B
(2) fa- (2))²
2L
W
P NNN + Vu-. (N).
2 N
N
which also gives 2 ([Huz (2) - (z))
112 X
(N-1
(p ) + 2 Vu,
This allows to monitor the convergence of the finite-width neural network to the optimal solution fr" in terms of the model output (up to 1/N error).
4.2 F-MFLD and SGD-MFLD
Here, we present the convergence rate for F-MELD and SGD-MFLD simultaneously. F-MFLD can be seen special case of SGD-MFLD where the variance of 0. We specialize the previous assumptions to the stochastic gradient setting as follows. Assumption 5.
(i) suPPER V
Af((.2) by
(z)]l <
R for all p E P Pand Z. (z)].op:
(it) sup, V V 6f(a.))
supa,a' V
(z, (z,2")]lop R.
b
Note that point (i) is required to bound the variance of (i.c,, oi RR/B), whereas point (ii) corresponds to the additional Assumption 4 required for the improved convergence rate. Let Ag == N ((
F (u²). We have the following evaluation the objective. Note that we can recover the evaluation for F-MFLD by formally setting B = xo so that o² = 0 and 02 = 0.
Thh constants not optimized (e.g.. ((++) ) in R can be improved) since prioritize for clean presentation.

---

Theorem 3. Suppose that The
(vk E No) and dam S 1/4 and S A1/(422). Under
Assumptions 2, and 5, it holds that
1 N (uk ) I F(x)) S exp (- Xonkk/2)) Ao
4 12C (>r) n/- + 4 r
40x
(8)
N
Aa
Aam
Na
where Y = Anp, + R + X2 R + (I. 22) /
+ ( X2 R(1 + v
こ n n²
under
B
Assumption 5-(ii), and T n², B 7) without this additional assumption.
The proof is given in Appendix C. This can be seen as a mean-field generalization of Vempala and Wibisono (2019) which provides convergence rate of discrete time vanilla GLD with respect to the KL divergence. Indeed, their derived rate D(exp(-Xanp) Ao + 2x is consistent to ours, since Eq. (26) in the Appendix implies that the objective } * (px
$ (n²)) upper bounds the KL divergence between ( and yr Moreover, our result also handles the stochastic approximation which can give better total computational complexity even in the vanilla GLD setting as shown below.
For given
0, if we take
Aac 4002G A EV 40C;
aC
7 S
A 1, k Xan
log log(2Ao/e),
with B 2 4 [(1+ R)(R + X2 + (L {2] R² vTTA)((acc) under Assumption 5-(ii) and B 2 4R2 / without such assumption, then the right hand side of (8) can be bounded as EB
N (ux (N)) (1 =+
4Ca, Hence we achieve +1((1/N) error with iteration complexity: XaN
k k= k - Vdac
)
L
- log(e λa
aE
If we neglect O(1/VXac) as a second order term, then the above can be simplified as O L2 logle
Aer
Noticeably, the mini-batch size B can be significantly reduced under the additional smoothness condition Assumption 5-(ii) (indeed, we have On + VTiA factor reduction). Recall that when the objective is an empirical risk, the gradient complexity per iteration required by F-MFLD is O(n).
Hence, the total complexity of F-MFLD is (nkk) where &o is given in Eq. (9). Comparing this with SGD-MFLD with mini-batch size B, we see that SGD-MFLD has better total computational
complexity (Bk) when n > B. In particular, if B (1(7 + vTJA)/ (Aacc) 2 2(x + (c)a)= which yields the objective value of order (c), SGD-MFLD achieves O B) = (m(AAV times smaller total complexity. For example, when À = 1//n,, ((Vn)((fad factor reduction of total complexity can be achieved by SGI-MFLLD, which is a significant improvement.
4.3 SVRG-MFLD
Now we present the convergence rate of SVRG-MFLD in the fixed step size setting where mhe (Vk). Instead of Assumption we introduce the following two assumptions. Assumption 6.
(i) For any i [n], holds that any p.pp' of E Pi and x, xf E Rd. and sup
V (6,(m) (x)
V dy
8t,(n')
(a')
S I(W2(H, n1)) 1| x' for
(l,
(x) R.for any / P.
hy
(ii) Additionally, we have sup V
86.lm) of
(x) lopp sup
V
(z, x' llop S R.
Here again, point (i) is required to bound of and point (ii) yields Assumption 4. We have the following computational complexity bound for SVRG-MFLD.
Theorem 4. Suppose that xan S 1/4 and n S /1/(2))). LerE= n-B B(n-1)- Then, under Assump-
tions 1, 2, 3 and 6, we have the same error bound Eq. (8) with different T:
T = 4njō, + (1+ CE..2m2 + nA)RE + (R+ Ro - >2) 221)
under Assumption 6-(ii), and T C (n + >) without Assumption 6-(ii).

---

The proof is given in Appendix D. Therefore, to achieve # for a given e O, it suffices to set
)
F(x) S O(e)
40
Aa.V
Vaae 11 40L3C1 40I.
A 2log(20o/e)
-0
L2 L aE
log(E-1)
Vλoc
(Aa)
V
int V
with B 2 be bounded
m
A n. In this setting, the total gradient complexity can
Vn((n)) (1+ (+v/)) 1 log(c-1)
nk: Bk
nd 1+ >²"
+o S
max
Aa + n,
m
1 v 1) and An). When small, the first term in the right hand side becomes the main term which is max/n1/3 vn(ae))1/4) log(e_?) There- fore, comparing with the full batch gradient method (F-MFLD), SVRG-MFLD achieves at least min { Vn(aa))- } times better total computational complexity when S A. Note that even without the additional Assumption 6-(ū), we still obtain /n-factor improvement (see Appendix D). This indicates that variance reduction is indeed effective improve the computational complexity, especially in a large sample size setting.
where m Q(n/B) (1 n (1+ / D vn(nA)+ (1
//A) vn(1 (n/A)-11 3/2 (V1X)-1/2) (V/n)))-)//>)--->>)))
B O
Aa=e
Finally, compare the convergence rate in Theorem 4 (setting F to be linear) against prior analysis of standard gradient Langevin dynamics (LD) under LSI. To our knowledge, the current best convergence rate of LD in terms of KL divergence
given in Kinoshita and Suzuki (2022),
which yields a O ((n+ y
log(e +1 iteration complexity to achieve error. This corresponds to (Aa)²
our analysis under Assumption 6-(i) (without (ii) (see Appendix D). Note that in this setting, our bound recovers their rate even though our analysis is generalized to nonlinear mean-field functionals (the O(A)-factor discrepancy due to the difference the objective their bound considers the KL divergence while our objective corresponds to À times the KL divergence). Furthermore, our analysis gives an even faster convergence rate under an additional mild assumption (Assumption 6-(i)).
Conclusion
We gave unified theoretical framework to bound the optimization error of the single-loop mean-field Langevin dynamics (MFLD) that applicable to the finite-patticle, discrete-time, and stochastic gradient algorithm. Our analysis is general enough to cover several important learning problems such the optimization of mean-field neural networks, density estimation via MMD minimization, and variational inference via KSD minimization. We considered three versions of the algorithms (F- MFLD, SGD-MFLD, SGLD-MFLD); and despite the fact that our analysis deals with a more general setting (mean-field interactions), are able to recover and even improve existing convergence guarantees when specialized to the standard gradient Langevin dynamics.
Acknowledgements
TS was partially supported by ISPS KAKENHI (20H00576) and JST CREST (JPMJCR2115). AN was partially supported by JSPS KAKENHI (22H03650). We thank the anonymous reviewers for their careful reading of our manuscript and their insightful comments.
References
E. Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-staircasse property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4787-4887. PMLR, 2022.
M. Arbel, A. Korba, SALIM, and A. Gretton. Maximum mean discrepancy gradi- ent flow. In H. Wallach, H. Larochelle, A. Beygelzimer, F, d'Alché-Buc, E. Fox, and
10

---

R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Cur- ran Associates, Inc., 2019. URL http:://prooceedng.. www.ss.eerris.cooo wwt..nuis.....paprr./eeee.esscee httpppp//.ccpppp////////... htt.aa/aeefaaccceeeooccbcbaa////
pdf.
D. Bakry and M. Emery. Diffusions hypereontractives. In Azéma and M. Yor, editors, Séminaire de Probabilités XIX 1983/84, pages 177-206, Berlin, Heidelberg, 1985a. Springer Berlin Heidelberg. ISBN 978-3-540-39997--.
D. Bakry and M. Emery. Diffusions hypercontractives in sem. probab. xix Inm 1123, 1985b.
D. Bakry, I. Gentil, M. Ledoux, et al. Analysis and geometry of Markov diffusion operators, volume 103. Springer, 2014.
J.-B. Bardet, N. Gozlan, F. Malricu, and P.-A. Zitt. Functional incqualities for Gaussian convolutions of compactly supported measures: Explicit bounds and dimension dependence. Bernoulli, 24(1): 333 353, 2018. doi: 10.3150116-BEJ879 URL https: org/ 10 3150/16 16 (---EJ89.
Cattiaux and A. Guillin. Semi Log-Concave Markov Diffusions, pages 231-292. Springer International Publishing, Cham, 2014.
P. Cattiaux and A. Guillin. Functional inequalities for perturbed measures with applications to log-concave measures and to some Bayesian problems. Bernoulli, 28(4):2294 2321, 2022. doi: 10.315021--EJJ411 URL httpt://///o/ .rgg 10. 3150/21-BBJJ11
F. Chen, 7. Ren, and S. Wang. Uniform-in-time propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:22/2.030500v, 2022.
Z. Chen, G. M. Rotskoff, J. Bruna, and Vanden-Eijnden. A dynamical central limit theorem for shallow neural networks. arXiv preprint arXiv:2008.06623 2020.
S. Chewi, M. A. Erdogdu, M. B. Li, R. Shen, and M. Zhang. Analysis of langevin monte carlo from poincar' 'e log-sobolev. arXiv preprint arXiv:21 2021.
L. Chizat. Mean-field langevin dynamics: Exponential convergence and annealing. Transactions on Machine Learning Research, 2022. URL httpt:://ppeerreiiieee www.:/enttfoouiiee/eeooeo BitpummmddBDDLLttt...c gEm.
L. Chizat and F. Bach. On the global convergence of gradient descent for over over-parameterized models using optimal transport. In Advances in Neural Information Processing Systems 31, pages 3040-3050, 2018.
K. Chwialkowski, H. Strathmann, and A. Gretton. kernel test of goodness fit. M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages York, New York, USA, 20-22 Jun 2016. PMLR. URL https: ptttps://pproceeinggg mlr, press/ #48/chwialkowst 16. html.
New
A. S. Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. arXiv preprint
2014.
A. Das, D. M. Nagaraj, and A. Raj. Utilising the clt structure in stochastic gradient based sampling Improved analysis and faster algorithms. In G. Neu and L. Rosasco, editors, Proceedings of Thirty Sixth Conference Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 4029. PMLR, 12-15 Jul 2023. URL https: https://prrceeddnggg mlr. www.prrevsiiie das23b. html.
V. De Bortoli, A. Durmus, X. Fontaine, and U. Simsekli. Quantitative propagation of chaos for sgd in wide neural networks. Advances Neural Information Processing Systems, 33:278-288, 2020.
Delarue and Tse. Uniform time weak propagation of chaos on the torus. arXiv preprint arXiv:2104.14973 14973, 2021.

---

K. A. Dubey, S. J. Reddi, S. A. Williamson, B. Poczos, A. J. Smola, and E. P. Xing. Variance reduction in stochastic gradient langevin dynamics. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Cur- ran Associates, Inc., 2016. URL https: httts:://prceeeeee http://ppoeeeeen..eeeee neurips.cc..cccooo wwtpp/pppee/////11/e. (at.aaaa110bb82222222222222222
Paper .df.
A. Gretton, K. Borgwardt, M. Rasch, B. Schölkopf, and A. Smola. kernel method for the two- sample-problem. In B. Schölkopf. J. Platt, and T. Hoffman, editors, Advances in Neural Information Processing Systems, volume 19. MIT Press, 2006. URL https: /pppceedingg neurips cc/ Paper, pdf.
R. Holley and Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics,
1987.
K. Hu, Ren, Siska, and L.. Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint arXiv: 2019.
X. Huang, Ren, and F.-Y. Wang. Distribution dependent stochastic differential equations. Frontiers of Mathematics in China, 16(2):257-301, 2021.
A. Javanmard, M. Mondelli, and Montanari. Analysis of a two-layer neural network via displace- ment convexity. arXiv preprint 191. 2019.
R. Johnson and Zhang. Accelerating stochastic gradient descent using predictive vari- ance reduction. In C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Wein- berger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. URL pttps://proceedingg wwwpeeiinnnnnnninn.omme httpanuipps..ppppee//////00 cc/paper/211//flee http::cccccceeaeeeeeeeeee6eeeeee
M. Kac. Foundations of kinetic theory. In Proceedings of The third Berkeley symposium on mathematical statistics and probability, volume 3, pages 171-197, 1956.
H. Kahn and T. E. Harris. Estimation of particle transmission by random sampling. National Bureau of Standards applied mathematics series, 12:27-30, 1951.
Y. Kinoshita and T. Suzuki. Improved convergence rate of stochastic gradient langevin dynamics with variance reduction and its application optimization. In Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https: // //opnrevieww..//////////a wtwpiiewweeeeoruneeeeee j=Sjzz_.... _2___iiw_--
A. Korba, P.-C. Aubin-Frankowski, S. Majewski, and P. Ablin. Kernel stein discrepancy descent, In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 5719-5730. PMLR, 18-24 Jul 2021. URL https: hattps://proceett http:///oocceenn..//// mlr. www.srressssieve https///////kkoobeaeaae wwt.koobaaaaaa.. html.
D. Lacker. Hierarchies, entropy, and quantitative propagation of chaos for mean field diffusions. arXiv preprint arXiv:2105.00000 2021.
D. Lacker and L. L. Flem. arXiv:2205. 12047, 2022.
Sharp uniform-in-time propagation of chaos. arXiv preprint
Q Liu, J. Lee, and M. Jordan. A kernelized stein discrepancy for goodness-of-fit tests. In M. F. Balean and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 276-284, New York, New York, USA, 20-22 Jun 2016. PMLR. URL htttps://proceeeeee v48/liub11 html.
mlr, .tress/
Y.-A. Ma, T. Chen, and F. Fox. complete recipe for stochastic gradient meme. Advances in neural information processing systems, 28, 2015.
H. McKean. class of markov processes associated with nonlinear parabolic equations. Proceed- ings of the National Academy of Sciences, 1676)1997-1 1966.
12

---

S. Mei, A. Montanari, and P.-M. Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33)E7665-E7777) 2018.
Mischler. An introduction to evolution PDEs, Chapter On the Gronwall
lemma, 2019. URL https://twatt cewemade........ wew...errmade.duuhiieee httr/mmischhere/Beeignnnna M2evo12018/ hat/e/001//ha.... .df.
A. Nitanda and T. Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint (7Xf::/712.05438,, 2017.
A. Nitanda, D. Wu, and T. Suzuki. Particle dual averaging: Optimization of mean field neural networks with global convergence rate analysis, 2020.
A. Nitanda, D. Wu, and T. Suzuki. Particle dual averaging: with global convergence rate analysis. M. Ranzato, A. Beygelzimer, Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 196211-1621. Curran Associates, Inc., 2021. URL https: /tppr//edings/sssseee paroceedings.neuriis.. newrips cc/ paper/2021/file /ttp://.........4.0.0000000000.. wwwp://aaaaa944aaaa---- Paper pdf.
Optimization of mean field neural network
A. Nitanda, D. Wu, and T. Suzuki. Convex analysis of the mean field langevin dynamics. In G. Camps- Valls, F. J. R. Ruiz, and I. Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statisties, volume 151 of Proceedings of Machine Learning Research, pages 7741-9757. PMLR, 28-30 Mar 2022.
A. Nitanda, K. Oko, D. Wu, N. Takenouchi, and T. Suzuki. Primal and dual analysis of entropic fictitious play for finite-sum problems. arXiv preprint arXiv:2303.02957, 2023.
K. Oko, T. Suzuki, A. Nitanda, and Wu. Particle stochastic dual coordinate ascent: Exponential convergent algorithm for mean field neural network optimization. In International Conference on Learning Representations, 2022. URL https: /ttp://oppeeeveeeeeeee seet/forummdd=PQpp/All
Otto and C. Villani. Generalization of an inequality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis, 173(2):361-44) 2000.
G. M. Rotskoff and E. Vanden-Fijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv: 2018
J. Sirignano and K. Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications, 1320312180--1822 2020.
Suzuki, A. Nitanda, and D. Wu. Uniform-in-tine propagation of chaos for the mean field gradient langevin dynamics. In Submitted to The Eleventh International Conference on Learning Representations, 2023. URL /ttttp:://opnrrevew/ wwt.//nrnuumnene. ?idr Jaaddd=JSSUKKTTBU.
A.-S. Sznitman. Topics in propagation of chaos. In Ecole d' ééété de probabilités de Saint-Flour XIX- XIX-1989, pages 165-251. Springer, 1991.
S. Vempala and A. Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems, pages 8094-8106, 2019.
M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681-688. 2011.
D. Zou, P. Xu, and Q Gu. Subsampled stochastic variance-reduced gradient langevin dynamics. In International Conference on Uncertainty Artificial Intelligence, 2018.
13

---

Table of Contents
1
Introduction
1.1
Our Contributions
2
Problem Setting
2.1
Some Applications MFLD
2.2
Practical implementations of MFLD
3
Main Assumptions and Theoretical Tools
3.1
Proximal Gibbs Measure & Logarithmic Sobolev inequality
4
Main Result: Convergence Analysis
4.1
General Recipe for Discretization Error Control
4.2
F-MELD and SGD-MFLD
4.3
SVRG-MFLD
5
Conclusion
10
A Verification of Assumptions
15
Proof of Theorem
15
B.1 Evaluation of Objective Decrease
15
B.2 Stochastic Gradient Error (Term C)
16
B.2.1 Analysis under Assumption
17
B.3 Putting Things Together...
22
C Proof of Theorem 3: F-MFLD and SGD-MFLD
23
D Proof of Theorem 4: SVRG-MFLD
24
Auxiliary lemmas
26
E.1 Properties of the MFLD Iterates
26
E,2 Wasserstein Distance Bound
27
E.3 Logarithmic Sobolev Inequality
29
E.4 Uniform Log-Sobolev Inequality
30
14

---

Appendix
A Verification of Assumptions
Assumption 2 can be satisfied in the following settings for the three examples presented in Section 2.1.
(i mean-field neural network. The neurons h-(()) and their gradients are bounded (e.g., tanh
activation), and the first derivative of the loss is Lipschitz continuous (e.g., squared loss, logistic
loss). That is, there exists > 0 such that sup, sup, [ S C. sup, sup, [z)l I C and || Vh(x)- Vh(x')l - Cl|x-='' x'll for all x. x' E IR., and sup, y sep ( v) S C,
[0y((f((z) (2). y) (2) & (1) - far (z) and (ffe(f((2)))) S C for all (y.2) and p4, P.
(ii) MMD minimization.
The kernel smooth and has light tail, c.g., the Gaussian RBF kernel. The kernel &o has a light tail such that sups, max{ (z,z'), 2²),, V. (2, S O: for example, k(2,z') exp
(iii
KSD minimization.
V W n- (z, =211)
|=:|=
20
1s'l"
-e=="'l²,
202 ) and max V log(u* log(u" | (og((((()))l) || log( (=))llop} S
201
C(1 + [2]]).
B Proof of Theorem 2
This section gives the proof of Theorem For notation simplicity, we write r to indicate me- Recall that for 2 = (X+) the proximal Gibbs distribution u. is defined by
1 )F(()))) _= (x))
p (X) x exp
(-X
5F
We also define another version of the proximal Gibbs distribution corresponding to FN as
N (-xr(x)))
p (2)) exp
B.1 Evaluation of Objective Decrease
By the same argument as Chen al. (2022), we have that
(F(ux) (xi) = log(p (X:)) = (N - - v.F(x))
N
where V, the partial derivative with respect X,. Therefore, we have
8. N (H((
(2) NV,F(uz) + XV, log(
8p
8F(pz)
(X;) + XV, (N) (2))
6h
= -VV (N) (M) + AV, log 1 (N) (M)
(10)
Remembering that 24 (X is updated by X X クou + 22 the solutions X and X can be interpolated by the following continuous time dynamics:
To = Lk, dx; vi + 22dW;
for 0 S S n). Then, L, obeys the same distribution as Lx+1. Let ūi (NN) be the law of V, The Fokker-Planck equation of the dynamics yields that
di(m) di 12 wox) Σ V.
(16
(2 (2 wor))b To.
+
Aip (N) (2 wo.)).
15

---

Hence, by taking expectation with respect to Lo.k, it holds that
djp (N)
(2)
dt
m V,
Σ AV.
Pi
(2)
(N)
plat (x)
6F z
Pi (2)
(22)
E
u (ooa wi))
Ht w/
xx)
20 al
8fc
where we omitted the influence of was to fi. (N), which should be written as j4 (N) precise manner. Combining this and Eq. (10) yields the following decomposition.
(X
wwb:)
more
d N (f4
(N)
dt
3X²
E
2 mp (N)
Vi log
plaj (2))
=:A
N
5F(16
8F 5F(
V
(Xx)
V.
(X;)
L L0
=.
(N) V2
F
A E
plat) (2t)
-V
(Xr)
Evaluation of term A: By the leave-one-out argument in the proof Theorem 2.1 of Chen et al. (2022), the first term of the right hand side can be upper bounded by
312
N
(N) p(N) (22
M
4N
E
~p
(N)
V, 1o
2=1
N
j (D) plaNy (2))
4N
E
2 (2
(N)
V.
1=1
λoe
( (N))
(1²")
+
(11)
2
N
with a constant Cx. The proof is given in Lemma for completeness, where we see that the inequality holds with Cx = 2XLa(1 2c, R2) 22
Evaluation of term B: By Lemma 2, we have
6F (N)
SF(T!
2
E 35.5%
(X6)
(X:)
S 8
54c
8j
for 6n C1L + Ama) = O(L2 ( + mx))).
B.2
Stochastic Gradient Error (Term C)
Now we evaluate the final term C. First, we derive bound without Assumption 4. By the Cauchy- Schwarz inequality, we have that
(NN) ppNN (2i)
5F g
(Xx) }
AE
V. log
(f
To
16

---

x2
(M) Vg
6F(11 go (Xx)
E
.00 %o
pDNN) (tt)
+ LE 54 2o
u uk (Ton
811
x
(i( p(()) (2t1)
|1
127
4
E
%
Vi
+ o n.k-
B.2.1 Analysis under Assumption 4
Next, we derive a tighter result under the additional Assumption 4. We decompose term C as follows:
(N)
6F(4
Vi log
ii plN) (
V
(Xx)
8p
6F g6
E L0
Vi log
(4tz (2t) (
(N)
(Xx)
8p
E Fo ( V, (p()) (2)x V
(F
}
(Xx)
Since V, log(7( (1) = v.ti( /x(a))))))))))))1)))( this is equivalent to
(1(156) (Xo)
[44 ( (2.) -1 i V. ( (2,) V, log(p .u (Zorn
K
8F(H2e (xxs) )
8F
EE Tol F
V sti
( (Xi 2 ) (c. (x 2/ (Zor - V
(N)
d
C-1
( Vo
(Too: -V 5F(1
(8x)
(12)
20 27
8p
c-m
where Loo is an independent copy of 2o.k Evaluation of C-(I):
(I) The first term (C-(I) of the right hand side of Eq. (12) can be evaluated as
f
E Dal
(N)
( ) (V) (d. (Zos,
8F(pz))
(xx)
d.
f ST
(s ti ( (%
oF zi)
(xi) )
Lo (1) (2 1o
d
8p
+ /
E. E // Vint (N)
(
(.. 26 (1
((Ttt (f) (Xt) 11 ((TT, Loo- wf) + V
8F zr.
(X;)
det..
81
5p
(13)
where (Te, 1o.-1)-= (Lo, Vi.... -i-1, L.) which is obtained by replacing Le of Losk to First, we evaluate the first term in the right hand side of Eq. (13). Let
T..
((2) 5F(ux) (X'),
6p
for 2' = ( and let D F F(2) =
Let
Au' = ( wi)) D',, F(2))
17

---

and let Au = (Av"), We also define
Z (V2Xt)
- -X
wf]]
Then Zi ~ ((0 1) conditioned by Fo. and is independent of w/. We also define
o wi) == E Voss 11 ,era [c] w))
where the expectation is taken conditioned by wirk but we omit w from the left hand side for the
simplicity of notation. Since the density of the conditional distribution ji! ( (2) Toi is proportional
to exp (-ER - -to (2 miDIl" it holds that, under second order differentiability of oj
2(2
and Dm
(X tu
(
O.R. whl)
((N)
2AL
n (f) - V (F( (x:)
d.t.
81
K-
Zi V2At" r2 - -1Dm 1DmFF(9 ((T)) + 22NLZ - LAV, w/)-
= 2
D F(Tx - -IDFF(%)) + V2NIZ - 14wi)]
Zi
= 22, V2NL'
+ (2% - tD (21) + v2MZ,wL - (2) V2NZ)
( - tuu Lo wk) + V2 w/) - (2)
1m -t + v22ZZw))
F(2 tual
+ V2xtZ) + F F - tD F (2)) + v2222)) (14)
By Assumption 4, we can evaluate
Eui
oe (2 - tox(2 WK) + v2XtZ, - -1 -1 (DmF((1)) + V2XtZ, w/)
F
L1R ( Wx) + V2XIZ) + F - -D F (()) + v2X1Z)
<Egge 1ink | tu (Z),, + 22177 2111 1)
- 11(( - +tDFF(2)) + V2XtZ, w)/
F(2 - tex( (Fir wx) + V2xtZ) + D', F(RK - tD -tD ((9)) + v2Xt2)||
N
<B g: - wo0
Σ
VJ
+ V2xtz, 2), wk)
VV] P( - 1D. ((2)) + V2N2, 20hk ,1)
||Acا]
op
Qt
Σ
|Auر 1|2 || 1|2
N
where Jensen's 's inequality is used in the first inequality. By using Assumption 4 again, the first term in the right hand side can be evaluated as
N
E or U.k
aj 2
1f - (DmF((2)) + V2MZ, 20'k-
1 A1 12
-V. m
((%% FDm + 2212 26 1))
2a,
18

---

S(N
2N v.k + _a_i
a-i
foo
+
ai
v.k 2
1 2a
t2 o.
for aj > 0 (j = I,.... where aj a-s (j + i) for some a-; > 0. By taking a-i k, Nlor,k/on, & and a; to oo the right hand side can be bounded as
(15)
Then, we return to the right hand side of Eq. (14). First, we note that
Z V2
- + v2Xtz)
E
(
= 0,
for fixed Z and 2x7 Hence, by taking this and Eq. (15) into account, the expectation of the right hand side Eq. (14) with respect w/ can be further bounded as
t x / 2
(16)
(2) Next, evaluate the second term of the right hand side of Eq. (13). We have
E8 Tal
(
2o.x) (Pxik).
5F((2r)
8F(
)
o/ (Kk
(Xx)
(o), V
(x:)
d
8jf
8μ
( V.
(T.)/ (
(2i)),
E Ti
8F( (Xi) }
5F((zz))
(xx V
(Xx) ( +
84
Vo (2))
E Zi
5F((
8F(
(2),
(X))
of (T,
(xi)
(17)
dyf
The first term in the right hand side of this inequality (Eq. (17)) can be upper bounded by
E
Vz
log
(2x)
(2t) (2i)/z (N Zi
1
/
E 27. 25
(1 V
(F(1) 8p
(X)) (fT,
V
8F(P
(X:)
2
8p
V. log (M (2 |1]
E
of Loa- 2om-1)) v
5F 5F(12 8f
8F
B
(xx) -o w/)z))))+ + +V
(X:)
84
where we used Jensen's inequality. By Lemma (the same bound applies of by the same proof). we see that the second term in the right hand side is bounded by
2
8
(18)
Finally, we evaluate the second term in the right hand side of Eq. (17). Note that
E To. R. [(V, log (2)))
19

---

5F((2z))
(Xx) (2T + 8F(15)
(Xx)
8u
27 K Vv.00g (D(N) -V, (008(/((111 (p)) + V log (DiN) ((o)
E 27
5F(()z))
(Xx) (/ w/) + of (1 5; (Xi)
20
+
E To T.
(V,
(20)
X
(F( (xi)
8F(
/
(T,,
/ (Z. + V
(x:)
84
8j
where we used Lemma 2 with the same argument as Eq. (18) and Young's 's inequality. Taking the expectation with respect to wi. it holds that
(o. (p() . (2(22, V
(F(H22)
))} O.
Euj
E 55555 57
(xx)
81
Hence, it suffices to evaluate the term
K
(F
(x:) )}
lo (pbN) (Fa)) o (Fr,w))
E 2t
Here, let 24 (X;))= be the following stochastic process:
Lo = 2x.
ax:- = v (F(()) (xt)dtt + 22XWW;
for 0 S L 17, where (Wi), is the same Brownian motion as that drives X Then, the term we are interested in can be evaluated as
{ 1o (p(() , (T.
8F( (x:)
E Ti
(vilog V. (o)N (T. v 8F(M sT.)
= E To Ki
(Xi)
ōu
8F sf (x:) + (fr
2Z (x:)
( wh)
v 8F( x.)
E % Zi
V, Vi log (p(N)( (T.
(X:)
sF (x:) }
-f v
8f
<E T. V (D(M
( (Zer - V 8F (
5F 8F sf
Ewi
(xx) - (2,, - V
(X:)
(19)
84
With the same reasoning as in Eq. (16), it holds that
8F(H)) (x:) (f(,))))))) V
- 8F(U
Eus
utf (Li wft)
(x:)
5h
84
< no, + 2Qt22 o2
20

---

Moreover, Assumption 2 and Lemm yield that
V log (ptM)
(U(m))
(a))
+ EIl
(+
sup
MEP
1
R + B Vr r(0) r(00)))
X
1
$
R + 12 x
X
1
R + >2 Bl1 |11
$
X
1 S X RR + R2
where the second and third inequalities are due to Assumption 2 and the last inequality is by Lemma Then, the right hand side of Eq. (19) can be bounded by
1 1 (R 12l1) (2tōv, + -QPPo..)
where we used Assumption 2 and Lemm 1. Hence, the second term of the right hand side of Eq. (13) (which is same as Eq. (17)) can be bounded by
8m + 1 +
X2 (21ō,, + 2Qi
(3) By summing up the results in (1) and (2), we have that
C-(1) s A V log (2.)
+ 46n + (R+ (1+ (1v/²) (2fōuxous + Q12
X ou
Evaluation of C-(Il):
Next, we evaluate the term C-(II). Here, let Lr = (8;) be the following stochastic process:
L0 Ln.
ax: = 8F(()) (xi) 22XW:;
for 0 < S where (W)), is the same Brownian motion as that drives X3. By definition, we notice that
d(x; x;) (4 (F(())) (xt)
v dt,
8y
which yields that
X
8 (v 8F((z))
(Xt) uk
(20)
By subtracting and adding the derivative Li from V, (N) we have
V, log (2))
1 V
8F (Xi)
81
8F T.)
8F(
5F si
(X))
V
(X:)
(8:).
8p
By Assumptions and 2, the first two terms in the right hand side can be bounded as
8F g.
5F(1 (X::)
(X:)
8p
21

---

8U (
SU(H) (x:) vr(()))
V
Vr((;)
8u
<L 5. 2 + |x - -;11- + 1211x; - Xiil
+1(
o
(F(R))
(x%)
12
+ +1102
8F(m))
(xx)
(21)
N
81
where the first inequality follows from
1./ VV r(Ox + (1 (1- ()))))(](= )do
[[r(()) vr(2')]l
I A2 / ||0- x'11d0 >2|1 - x'1l
(22)
by Assumption and we used Eq. (20) in the last inequality. Therefore, by noticing Lo Lk. we can obtain the following evaluation:
C-III)
N
( (ii (p(N)(Et))
8F 3
>
F Yock
(Xx)
6f
i=1
2(T
)²)2
.V
6F g
17
B (e
2 (2oz))
(X:i)
X
i=1
SIN 2(L + x²)²
O.kk-
X
Combining the bounds on C-(I) and C-(II):
N
N
Σ
AC
Σ
(C-()++-(II))
2m1
E
V, log
(1 1 1+ v// (2tx (210r,k0u,k 2Q12222 + +12((T + 22)2
N 48 (R+
22)
B.3 Putting Things Together
By Gronwall lemma (see Mischler (2019), for example), we arrive at
N Eu [ (14 (81) - F(1²)
1
exp( (spp->amp//)) N Wunn- 1 (1b (u())))] 3 (1m²)
(
mh (OTe
Ca
+
+
N
In addition, if Assumption 4 is satisfied, we have N Buoit | (" - -I"))
<exp(-Xomp
1
1
B (F F(1e)) m 6pm +
Ca
N
N
+ 4mp.m + + 2R+
[++/
(nioo,tov + Qmio? + u, (L + Ax) 2 o7
22

---

C
Proof of Theorem 3: F-MFLD and SGD-MFLD
Applying the Gronwall lemma (Mischler, 2019) with Theorem 2, we have that
Eupa
1
[N))] - (("))
N
h-1
R-1
k-1 1//2
)
[Ti+
66n, + N [
exp
Aa
n//2
Du
+
exp
-Aa
Σ
/
(23)
j=0
2m0
jm
Here, remember that Assumption 5 yields that
R² S B
R²
o
0 S
B
Under Assumption 5, we can easily check that Assumption 4 holds with Q R. When The = 17 for all k 2 O, we have a uniform upper bound of Yx as
4njr + [R + 22 + (L + )2)
1
+
1 17
n² B
T& <
+ (R
22
n? n
with Assumption 5-(ii), without Assumption 5-(ii).
B
Ra B
We denote the right hand side as T. Then we have
1
[ (4k (M)) (14))
E wo --
N
k-1
T T+ f (0, )
5 exp (->ankk222 Do Σ exp
- 1 -
- exp( \aknn/2)
T+1
86 N )
(a))
(upp(-ank//))a Do +
1 - exp( -aan/22
C2
]
S exp (-xank//2) Do +
T + 17 +n +
(: dan/2 S 1/2)
Aam
N
4 i (>n + 1 +
4 r
exp (-xmnta////) +
+
λa
dan
AaN'
Then, by taking
2
2 Aan
log( Ao/c),
then the right hand side can be bounded as
R N (14 (() - 8 (1x²)
N
4Cx
<et L'c, (An + 17)) +
T +
Xa
Aan
XaN
(I) Hence, without Assumption 5-(i), if take
" I Aaee
1 2Aoe
812 (C1A)-1 A
B 4R B24 ((acc)
8L V C₁
then the right hand side can be bounded as
EI E| N (1
(N) 5 (n*) < 3E +
40x \aN
We can easily check that the number of iteration & satisfies
k: = o ( I2
(Tee))
I Vice
log(e -11;
Aa
23

---

in this setting.
(2) On the other hand, under Assumption 5-(ii), if we take
7 Xae 407? (CIA) 1
Aae
T.V
40C;
A1,
B 2 4 [11 + R) + R2 + (L + A2) (n+ + (nA)e
then the right hand side can be bounded
[N)) - (2*) - 3E + 4C2
E B
(#4
\aN
We can again check that it suffices to take the number of iteration k as
I.2
L Viac/
k: ==
x
log(E "),
Aa
to achieve the targe accuracy.
D Proof of Theorem 4: SVRG-MFLD
The standard argument of variance yields
11 B 1
o2 S max ISEN
B1IA 1.5
[13)
B(n 1)
;=1
where A V 86
(xi)
V
80 (N)
(x))
V
(x:)
V
BU(H! 8
(X:). Here,
op
Assumption 2 and 6 give that
N
| 12 < ((T (M) .n. - S 4T.
Σ
|x
- x + 11 - -xiIP)
N
y'=1
Taking the expectation, we have
BIA 4T.2 1 k-1
2 |13]
27/d)
1=8
SC,L² (k ( +nA),
where we used that [[[[]] [22] < 2(R2 A2(c + R2)) Ci by Eq. (24) with Lemma 1. Hence, we have that
0º2 < C.EL' + nd):
where =- B B(n-1) 1)*
On the other hand, we have
2 < B
R² = ER?
B(n 1)
by the same argument with SGD-MFLD. We have uniform upper bound of Yg as
4n0 + (R+ R2 (1 / + 1² V3VC'LL +nAy +Aj RCE Ts
+ [2+ + (L + 222 (1+ + (t) C + n2),
with Assumption 6-(ii,, without Assumption 6-(ii).
0,LL3 + m))):
We again let the right hand side be Y. Then, we have that
E I (uL (N) $ (p²)
N
24

---

NEIS (2) -
+8+1 (
6n (0,
C2
S exx((-Xan/2)
Then, by the Gronwall's lemma yields that
N E
]] F(n))
Sexpl -2ankk/2
1
(16 (16² -F(1()
+
Σ (app--an((--aa 1)/2)
4
T + (5n((n))
6n Ci
)
N
N
1=1
1 Bl (MD [1) (
exp(- (>aanpkk)/2
T T+ on + C2
N }
< Eexp(-Aarphe
+
N
1 exp( (poo--amm2)
I exp(- (2ank))))/
1 Bl
(HD [] $ (1"))
+
4 + +
4
bn +
C2
N
Aam
ad
N
where we used Aan/2 I 1/2 in the last inequality.
(1) Without Assumption 6-(ii,, we have
C1 EL mmn (1+
C1 L'mm(n A)
(n B)
r =
Aan
Aam
λa
B(n 1)
obtain that
(x4 ((x)5()))))) s
4C, >n)) 4Cx/N
4C1
L2mnn(n + A) (n- B) B(n - 1)
λa
4C, + xn)(1 +
(n-B) B(n-1) m)
4Cx
=
+
da
AaN
when
1 log(e )
k 2Z Xan
and Aan/22 < 1/2. In particular, if we set B 2 m and
Vdac
" 4L20 4LVCT
we have
4Cx aAN'
E[
(m
-
N
with the iteration complexity and the total gradient computation complexity:
L
L
1
k S
log(E-
at
Vaee
(Aa)
nk
12
L
Bk +
S Vnk
n
as
Vnl
+
log(e -2) + n,
m
a6
v
(Ao)
where m = = Vñ.
(2) With Assumption 6-(ii), we have
1 T 4 Sn + da
-o
(nVmL(7² + xn) + mmmni² + xn))
(1+ [+
(n - B) B(n 1)
Xan
Aa
Then, we obtain that
B N (u4 [()))]]] - - F1²)
200 (17² + An) + 4CO/N
Lo (10mm// + 2n) + mml(1² + xn)
1
(n - B) B(n 1)
<e Aae
1+ #+//
Xa
25

---

when
1
2
log(E )
Aom
(1+ /n///// V m
(0+ (nt /nAaa
and Aan/22 S 1/2. In particular, if we set B
>
Vm
A n and
aE
VADE
1 =
4012C1
LV40C1
we have
1 - - O 4Cx
1 (pk
N
axN'
with the iteration complexity and the total gradient computation complexity:
L
L
l:
log(E ")
a6
Vdac
(Axx)
nk
1/33 (1// 4/3 .VV
(1 (1+vn/A) vn/a) 3/2
(VTA)1/22 k
Bk
+
+ n
max
m
x max { n/133 (1+ v1/2// 4/3
.VTl 1 vn/2) 3/2
(VTA)
- 11 ()cx)
log(e + n,
where m = Q)n/B)nn = 2([m2/3(1 + vn/A)
4/3 A vT(l+ + VT/A) (0/2))3/((1/) (/2/////2)
V 1) and
B
2 n (1+ v Vn(n)) (1 vX) An.
E
Auxiliary lemmas
E.1 Properties of the MFLD Iterates
Under Assumption 2 with Assumption for SGD-MFLD or for SVRG-MFLD, we can easily verify
that
(pR) ōp
(x))
S E. < 2(R + A2(c, + (((
(24)
Lemma 1. Under Assumption 2 with Assumption 5 for SGD-MFILD or Assumption 6 for SVRG- MFLD, in S (422), we have the following uniform bound of the second moment of X
[
A1
1)
E ~ A
(122 + )22x) + Ad
8A₂
2X1
for any
Proof. By the update rule of X and Assumption 1, we have
x ( not + 221Ac() Ellino + V 21
X V (U(m)) (x + vT(xx) >)
=|| -
+ 2n
5p
21
|| + 2n/² (R² + + | 3) + 2n7Ad
< - (1-2n)1 21 2 RB
X + 2n + >22))) + Xd)
< (1 7ソ ) 2 R2 + ) + Ad + R² /X1)
(:: 2II < X 1| 1/2 2R² />1).
where we used the assumption v - \1/ (4)2) Then, by the Gronwall lemma, it holds that
|[| |x:1 [1 1 (1-1--nm )*
E ||| < (1 ny (
nIn(R? + A22rr) + Ad + R² R²/M /Xi]
1/11
< | + A1 (n( + >2c7) + Xd + R$ ((1))
This with the assumption AAx(4 yields the assertion.
26

---

Lemma 2. Let Rd == = 1 [ +/+ (R2 A22,)) + sd)
Under Assumptions 1 and
Ai
nf S (4X2)) and we define
5n 8|22 + >2 + R22) + d) (rp + +2n),
then it holds that
8F( () (Xt)
5F(() 8p
JE 2o
V
(Xx)
8n.
8p
Proof. The proof is basically a mean field generalization of Nitanda et al. (2022). By Assumptions 1 and 2, the first two terms the right hand side can be bounded
(F
5F(m
(N)) (x:)
(Xx)
SU
(Xi) + vr((x)) - (0 (1)
(x:) - vo(X3))
8p
<LW2 (N)
1
X + x2l1 8i1l
(25)
Therefore, the right hand side can be further bounded as
8 FF (N)
5F (N) 8n
(x:)
6p
S(L +
22 (N) 1 + 11 x(112
<2
| x:11 + 21 X:11
N
<2
- 2
+2L2 221
Then, by taking the expectation, holds that
8F(1 (N)
(F( N)
E T. Lo
(X))
V
(X::)
8μc
8fc
<2L 1
Σ
B|lluti|l 2tAd) + 2L
+ 2tAd)
N
1=1
<412 A2(e, R2)) + 2tdd]
(:: Lemma
(F >2(G + R2) tAd]
<8[R²² + >2(c + R2)) + d df2 + tA).
Then, by noticing S we obtain the assertion.
E.2 Wasserstein Distance Bound
Recall that W2(u, v) is the 2- 2-Wasserstein distance between a and v. We let D(u, f log ( de be the KL-divergence between μ and v.
Lemma 3. Under Assumptions and 3, it holds that
(2( N s (INN N..
2 Xa
N. - -NF())))
27

---

Proof. By Assumption 3, fa" satisfies the LSI condition with constant a > 0. Then, it is known that its tensor product N also satisfies the LSI condition with the same constant (see, for example. Proposition 5.2.7 of Bakry et al. (2014)). Then, Otto- Villani theorem Otto and Villani (2000) yields the Talagrand's inequality of the tensorized measure v=N::
2 ( N - D(u (N) ).
Moreover, the proof of Theorem 2.11 of Chen et al. (2022) yields that, under Assumption 1, it holds that
1 D(u ,N S (NN - ())))))
(26)
Lemma 4. Suppose that (1)(2)) (1 (2)) S T.||x= x'1l (Vx, E (R)) and let Vi = Var (fw=)
(2) hr dux")) then it holds that
E 2impph (lfur (2) fa- (=))
2L
W .1 MM ) 2
(N)
Vu.
N
Proof. Consider a coupling y of random variable obeying the law ry Then,
2e and n and let (Fr. = (X)),.(X))) be a
(Uaz( - fz) (x))
1
Σ
N
xx (z) - -
2
N
-=1
N
Σ hx;(2) - / mmatehar(())) 2
N
Σ
(bx: (2) - (x)x +
N
N
i=1
N
2
N
2
Σ (mxg(2) (2) - -xxx(=))
+
Σ (x) -
- / he (=))))((()))
N
N
1=1
1
N
X - -xill
Σ hx((2) -J (h((e)dp i=1
N
2
2
+2
N
N
N
N
2
S
2L
|X
x +2
1xx
(2) (2)- [hol))d ()))))/)(())d
Then, by taking the expectation of the both side with respect to (2i, 2n), we have that
E Dinny (N) [Srrr
(z)
Sn (2))
N
< 2
X - xtI
+ N Va--
N
Hence, taking the infimum of the coupling 7 yields that
2L E 22 22 (N) [[ (z) fa- {z)))
N w
(N) (1)
,NNN)
2 Vu-
Remark 2. By the self-consistent condition of p², we can see that jr" is nab--Gassiinn Therefore, Vu- 00 is always satisfied.
28

---

E.3
Logarithmic Sobolev Inequality
Lemma 5. Under Assumptions I and 2, Je, and p satisfy the LSI condition with a constant
2
R2
41
R
2X
R²
d
A2
R²
exp
-4
1d/T
+
2A13
+ +
log
+
2x
A1
A
A1
2
A1
Proof. In the following, we give two lower bounds of a. By taking the maximum of the two, we obtain the assertion.
(1) By Assumption 2, (0))) is Lipschitz continuous with the Lipschitz constant R, which implies
that (U(H) is R/ R/X-Lisshhitz continuous. Hence, μ and p is Lipshitz perturbation of v(()) oc ou
exp(-2 1-'--rn (x)) Since X VV r(x) = 2 by Assumption 1, Miclo's trick (Lemma 2.1 of Bardet al. (2018)) yields that μ² and p satisfy the LSI with constant
A1 2X
R
A
R ArY
2
exp
4 A1
x
2d/T
2X
exp
-4
2d/
(2) We can easily check that V(x) = r(a)
X and (x)
(x:), for appropriately chosen jr, 2, cr, and L
1 (0(
by
satisfies the conditions in Lemma 6 with Ci = X
R. Hence, n² and X
pxx satisfy the LSI with a constant a such that
2
41 A1
2x
n2
A2
R
a 2
+
e TAyx
+ +
log
4
A1
A1
A
Lemma 6 (Log Sobolev inequality with Lipschitz perturbation). Let v(x) x exp(-V(xx) where VVTV(x) = I, VV(x) 2 =||2=|/ and S V(x) S ca(ev |x((²) with eı, C2 Cv > 0, and let H. Rd It is Lipschitz continuous function with the Lipschitz constant L. Suppose that μ E P is given by μ(x) & (x)(=H(c)))u(z)) Then, satisfies the LSI with a constanat such that
L
2
12
C2
+
2 d
log
C1
C1
Proof. Since v is a strongly log-concave distribution, the Bakry-Emery argument (Bakry and Émery, 1985b) yields that it satisfies the LSI condition with a constant a' = C₁ by the assump- tion VV V()) Z T.
Next, we evaluate the second moment of v because is required in the following analysis. Since we know that is the stationary distribution of the SDE dx,- = VV(XX))++2ddV. its corresponding infinitesimal generator gives that
E [[I - 2d 22, | X
2d.
dt
Then, by the Gronwall lemma, we have
X, < Xo + /' exp(-22c1 dssd
d
+ (1 exp( (2,,t))) C1
Hence, by taking X, = 0 (a.s.), we see that lim sup, E X, 1 < d/c1 that yields E, [| < d/c1- distribution j satisfies the Poincaré inequality with constant if it holds that
EaI(f
ExIf) à [)
29

---

for all bounded f with bounded derivatives. By Example (3) in Section 7.1 of Cattiaux and Guillin (2014), j satisfies the Poincaré inequality with a constant ā such that
1
1 2
2T. C1
L2
Z00I
ā
C1
Here, let G(z)= (()) ((). We can see that μ(x) can be expressed μ(æ) (xp(-G(x))edC
Za
with normalizing constant Zg O. Note that, by the Lipschitz continuity assumption of H, have
IG(=)) < L ||]|-||
Hence, we have the following evaluation of the normalizing constant Zg:
/ cxp(-G(z))u(ddddr
ZG =
exp
(-L|lz/ll
cs(ev + |,/111)
da
(2 (1((((
302
L
1
2
exp
|
c20V)
dx
2
202
(2x((1/11)
d/2
I
exp
02CV)
302
202
Theorem 2.7 of Cattiaux and Guillin (2022) claims that μ satisfies the LSI with constant such that
2 (6+ (8(1))(11 1)(1+ + - 2
1
(2 E. log(Ze))) + I 2
(1 01+0)(1+ + 8)
g
+
8
āa'
48
2
for any 8 >>0 and A >
0. The right hand side can be bounded by
(8 (1 + 1-
2
ī.2
d
2
(1+ [1+0)] () + B)
B²
+
2 LE
[|
+
+
log(3c2/(1)
+ 1
+
B
4
202
2
ac1
48
2
(8 (8+1) 11 +
1
2
d
1
d
(1 + 01+ + 3)
B²
+
2 + L
+
log
log(302/01)
+ 12 âci
B
4
C1
2c1
2
43
2
(B + (8(1)((1 + ()(1+))()) 1
2
1
(
d
2
(1 0) + ()
8
2+
+
(log
(og(3e2//5)
(1) +1
+ _
B
C1
C
2
âc1
43
2
1B 1)(1 +
0-1)
2
1
2 + (log(e2/c1) +
d
1.2
1
(1 + 0)
+ B)
188) ).
82
B
a
2
Hence, if we set 3
= 1
and 0 = 1, we have
1
2
1
d
1
<
2
2
+
(log(02/c) +2
21+2
2
2
9
25
2
Ci
4
L
2
2
L2
d
i2
<
+
+
20g
2 +
d +
log
+
C1
C1
C1
2
C1
E.4 Uniform Log-Sobolev Inequality
Lemma 7 (Uniform log- lobolee inequality). Under the same condition Theorem 2, it holds that
N
E
V:
(N) logog
pl (2)
)r"
Aa
( ( 1 (
- -F(1)))))
CA
2N
ps
2
N
1=1
(27)
with constant C22 = -2La(1 + Rcy 22
30

---

Proof. The derivation is analogous to Chen et al. (2022), but we present the full proof for the sake of the completeness. For L
(x let 2 (X₃) (x)) Then, it holds that
N
i (N)
E
2 App
(N)
v log
D(N) (2)
N
E L m
Vi log
148 (1 (2:3)
8F (F
(Xi)
12
x
i=1
N
E
(N)
V. log
(ute (2))
v
8F (
81
(Xi) + X
-
8F 6F
(Xi)
FF(2))
27
(Xi))
8p
1 N
E ~ (3)
||
v. log (1( (2))
5F 2)
(X:)
2 1=1
N
v 8F (1 (Xx))
o OF (1))))
27
E
2 xp
(N)
X
V
8p
(xi)
1
N
11 (29)
8F (1 (Xi)
12
2 1=1
E
L ~p (M
V log
6fc
N
1,2
2 ~m
(N)
2
E x mit
(M)
E Exr
Vi
-
X
8F 8F 1)
51
(Xx) 12
N
+
+ i=1
x>pp(') LE
(N) [W2 (U (1x
where the second inequality is due to the Lipschitz continuity of V UU in terms of the Wasserstein distance (Assumption 2). Here, the term corresponding to the Wasserstein distance can be upper bounded as
E 2' (m)) (u - x x
1
Σ
CII - X:112
N(N - 1)
2 N(N - 1)
E1x,1 N 11x,11
x x
iti
4
nta
N
Denote by - the conditional law of X, conditioned by 2' and denote by P the marginal
law of 2' where the joint law of 2 is zil Then, it holds that
vxPx,12- -(--))))()))) Px.xx- (Xx,)
V, Vogg(
(2))= (M
Px,1x--11
1 A 65F(Mx (Xi)
Σ E - x-
Ex,~xx
V: (Px,12-
1
< 2a
[Dppx-t, Px.1x--)]
i=1
31

---

by the LSI condition of the proximal Gibbs measure (Assumption 3), Let v, be the distribution with the density V, exp( exp(-r(x)/A). Note that the proximal Gibbs measure is the minimizer of the linearized objective:
{ 80(
[ (X.)dl (X,)) B.(rr + XEnt(())
P arrginn
PCP
{
[UU()=))) (X:)du (X AD(v..10) }
argmin
EEP
Then, by the optimality of the proximal Gibbs measure, it holds that
D1 Px,12--))
(I:) - -px-:) + (D Px,1 - AD(r..po0-
8U( (x:) - -dx)(xx;) -=))d==--- + AD(,r, Px, - AD(,,,H")
8U(12 8p
(28)
The expectation of the first term of the right hand side can be further evaluated as
N
[/
SU (0(Hx
(2:) - -x²)(dr))
E E
N
[
BU (U( (c,)0x((rr) - /
5U (1xx)-1))
E 2 mp!
(xi) (dxx))
N
{
(SU((LE)
(r;)8x(dxx)) - 88((1))
2)
>
E 2t' mM!
6f
i=1
2L
+ ((1x111112 Ex |||X||+ |X,
Σ
||x,,11 112)
N
N
NNE x mMM)
[ SU((12)) (x) (dr) - 8U(Hx)
(I) (1))
2L (2+ 4c2 R?)
NEE I x1(
[U(Ux)) - UU(²) - 2L (2 + 4c,, [[ (U(x) -U(x") - -2L (2 + 4c R22)
NE J. -xx
NE x w (U(rz) U(x")] - 4T (1 + 2c,, R²*)
where the first inequality due Lemma 8 and the second inequality due to Lemma I (we also notice that Ex- [|] < R2). The second term in the right hand side of Eq. (28) can be evaluated as
N
N n0 (Dur, px] Σ Egpo (Ent( (Ent(Px.1se Px.1 - E BNNN
N
Σ log(v,(Xx))
1=1
2 (())) - B,,cn
Σ
= , (2))))
L==1
where we used Lemma 3.6 of Chen et al. (2022) in the first incquality. Combining all of them, we arrive at
N
(N)
E (N)
V, Vogog
p(NJ (2))
x
< NE
2 xmpp( (N) (N) (p + >D(WN' (1(N) ((U())) -DD(r.H')) 4L (1 20, R2) A
32

---

NE 2 xxp
+ AEnt(id (F(xx)))
N(F(p)) *Enn((²))) AL (1+ 202 +
x (N)
4Lx x (1+ R2) RR.
Then, have the assertion with
Cx 2ALa(1 + 2er, 1t2 22
Lemma 8. With the same setting as Lemma 7, it holds that
6U (U(pz)
(U
L
(x)
(x)
2 cD
|||0||+
Σ
|1x,112
811
j
Proof.
Let 12 = 0ux + (11 0)
Then, see that
(U(xx)
(x) (1(
(x)) (x))
{/²
((u 8)
1
8 0)
(x, Xi)
+
(xr,
xx)
d0
8f
8p2
N(N - 1) iti
1 + CL
L N
2||2|| + +1,,112
1
Σ ||x,
N
where we used Assumption 2 (this is the only place where the boundedness of the second order variation of U used). Hence, we obtain the assertion.
33


Reviews:
[{"review_id": "IuPJWuJPkK", "reviewer": "Reviewer_YwhW", "summary": "This paper provides a set of results to analyze convergence of mean-field Langevin dynamics with a set of related algorithms. These results are in discrete time and space. Using log-Sobolev inequality techniques from optimal transport, which creates the opportunity for extensions to other settings. The propagation of chaos results are also proved in discrete time. ", "questions": "Assumption 4 is rather complicated. Can the motivation for it be explained more clearly in the text?\n\n", "limitations": "LSI conditions are obviously hard to satisfy for some problems, but this is well-known and well-discussed in the paper.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "My feeling is that this paper provides a useful advance in its effort to systematically control discretization error in the propagation of chaos. I am not familiar with other results that have done this. The rather explicit results obtained for the mean-field dynamics settings illustrates favorable scaling for single loop algorithms, at least in sufficiently smooth problems. ", "weaknesses": "I think the differences from Nitanda and Chizat could be articulated more clearly. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your supportive comments. We address the technical comments below.\n\n### Weaknesses \n**Q:** I think the differences from Nitanda and Chizat could be articulated more clearly.  \n**A:** Thanks for your suggestion. Please notice that the differences from Nitanda et al. (2022) and Chizat (2022) are summarized in the Introduction section. That is, they did not derive (1) the finite particle approximation error of MFLD and (2) the stochastic gradient approximation error. On the other hand, we established a unifying framework to evaluate the time discretization error, the finite particle approximation, and the stochastic gradient approximation. One of our contributions is to derive a tight bound on the stochastic gradient approximation error in the propagation of chaos analysis, which is challenging because it requires the evaluation of correlation between the randomness of each gradient and the updated distribution. \n\n### Questions  \n**Q:** Assumption 4 is rather complicated. Can the motivation for it be explained more clearly in the text?   \n**A:** Assumption 4 is essentially a higher-order smoothness condition, which is required to obtain an improved bound on the stochastic gradient approximation, i.e., smoother loss yields better stochastic approximation -- we will discuss this in more details in the revision. \nAs an example, Assumption 5 (ii) and Assumption 6 (ii) are sufficient conditions for Assumption 4 in each corresponding situation (gradient estimator). We believe that these conditions are rather intuitive because they only impose boundedness of the second order derivatives.  \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period."}, {"review_id": "TX6m2aPRVh", "reviewer": "Reviewer_wZfS", "summary": "This paper studies the mean-field Langevin dynamics (MFLD) with stochastic gradient updates. In particular, the authors propose a general framework to prove a uniform-in-time propagation of chaos for MFLD. The authors establish the convergence rate guarantees to the regularized global optimal solution, simultaneously addressing the uses of particle approximation, time discretization and stochastic gradient. \n", "questions": "It is suggested that the authors should provide some simple numerical experiments to demonstrate the effectiveness of the proposed algorithms.", "limitations": "Yes. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work provides a quite general framework for analyzing the mean-field Langevin dynamics (MFLD) with stochastic gradient updates. The authors establish the convergence rate guarantees to the regularized global optimal solution, simultaneously addressing the uses of particle approximation, time discretization and stochastic gradient. The results are interesting and appear to be novel. Various examples and practical implementations of MFLD are also provided to demonstrate the effectiveness of the proposed framework.", "weaknesses": "Purely theoretical paper. No numerical experiments are provided. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive evaluation. We address the technical comments below.\n\n### Weaknesses:  \n**Q:** Purely theoretical paper. No numerical experiments are provided.   \n**A:** Since the theoretical part already occupies the whole part of the paper, we chose not to include the numerical experiments in the main text. Following your suggestion, we will add numerical experiments in the appendix. Indeed, our preliminary experiments showed good convergence with respect to the number of particles even with stochastic gradient approximation. \n\n### Questions:  \n**Q:** It is suggested that the authors should provide some simple numerical experiments to demonstrate the effectiveness of the proposed algorithms.   \n**A:** Thank you for your suggestion. We have already done a preliminary experiment. \nWe would like to include more thorough experiments in the appendix of the camera-ready version.  \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n\n\n"}, {"review_id": "WzSDI78T6h", "reviewer": "Reviewer_rhAi", "summary": "In the work authors study mean field Langevin dynamics under stochastic gradient updates and prove uniform in time propagation of chaos that takes into account discretisation, stochastic errors which allows to establish convergence rates of MFLD.", "questions": "-", "limitations": "Yes", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. Strong theory with explicit bounds\n2. Excellent outline of results\n3. Thorough comparison of convergence rates with other approaches\n4. Wide applicability of results to different sg estimators and neural networks in mean field regime.", "weaknesses": "Seeing some numerical evaluation could make results even stronger to support the theory.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive evaluation. \nPlease find our answer to your concerns. \n\n**Q:** Seeing some numerical evaluation could make results even stronger to support the theory.   \n**A:** Thank you very much for your suggestion. \nSince the theoretical part already occupies the whole part of the paper, we chose not to include the numerical experiments in the main text. Following your suggestion, we will add numerical experiments in the appendix. Indeed, our preliminary experiments shows good convergence with respect to the number of particles and the stochastic gradient also converges properly. We would like to add more thorough experiments in the final version.\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period."}, {"review_id": "SuUp2hfXd4", "reviewer": "Reviewer_Y8on", "summary": "This work considers the analysis of mean field langevin dynamics as implemented algorithmically. i.e, \na) particle approximation b) time discretization and c) stochastic gradients. Under the assumption of certain logarithmic sobolev inequalities, prior works were mostly restrict", "questions": "See the weaknesses section", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Extensive and explicit non-asymptotic convergence rates are obtained for the practically implemented algorithm with finite particles, time discretization and stochastic gradients, which is a great addition to the literature.", "weaknesses": "1. What are the main technical insights in this work? It seems like an extension of Wibisono and Vempala's LSI analysis of LMC while accounting for the mean-field and stochastic gradients.\n\n2. Assumption 4 seems non-standard and the existence of a.s. bounded large order derivatives of the stochastic gradient seems restrictive compared to the assumptions in SGD/ SGLD literature. Similarly Assumption 5 seems too restrictive. \n\n3. What is the point of SVRG type algorithm when straightforward stochastic approximation like SGLD itself outperforms this? This is true in Table 1 for MFLD. We also refer to the Theorem 6 in [A1] for results on regular Langevin dynamics. \n\n4. Can you compare your technique with the recent results which utilized the CLT structure in the batched noise to obtain sharp analysis of stochastic approximations such as SGLD? (see [A1]). In Theorem 3 in the current work, with a batchsize $B$, the error bounds will have a $\\frac{\\eta}{B}$ (similar to [A2]) Whereas [A1] gets a rate of $\\frac{\\eta}{B^2}$.  Since the analysis techniques are close to that of SGLD, and one of the main contributions is the addition of stochastic gradients, it would be good if the authors compare the results to that of [A0,A1,A2] on a technical level.\n\n[A0] Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis\n\n[A1] Utilising the CLT Structure in Stochastic Gradient based Sampling: Improved Analysis and Faster Algorithms\n\n[A2] Faster Convergence of Stochastic Gradient Langevin Dynamics for Non-Log-Concave Sampling\n\n\nMinor: Theorem 2, definition of $\\mathcal{Y}_k$ needs to have $\\eta_k$\nThe notation is a bit cluttered at many points, which makes the paper hard to read sometimes. \n\n\nI will give a borderline accept for now. Happy to improve my score once the authors provide a satisfactory response.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful comments.   \n\n**Q:** *What are the main technical insights in this work? It seems like an extension of Wibisono and Vempala's LSI analysis of LMC while accounting for the mean-field and stochastic gradients.*   \n**A:** First, we note that the analysis of the mean-field extension of Langevin dynamics is far from trivial, even in the continuous-time and infinite-width limit. As for the discretized algorithm, we need to establish a uniform-in-time propagation of chaos, which is technically challenging as we have to handle (1) the finite particle approximation, (2) the stochastic gradient approximation, and (3) the time discretization. Note that such propagation of chaos calculation is not required in Wibisono and Vempala's analysis due to the absence of mean-field interaction. \nOur main contribution is to derive these errors in one unified framework. In particular, (1) and (2) are technically demanding, and we obtained faster rate than existing work taking into account these approximation errors with an additional smoothness condition.  \nPlease refer to our Introduction section for a more detailed summary of our contribution. \n\n**Q:** *Assumption 4 seems non-standard and the existence of a.s. bounded large order derivatives of the stochastic gradient seems restrictive compared to the assumptions in SGD/SGLD literature.  Similarly Assumption 5 seems too restrictive.*  \n**A:** Please note that our theorem also holds without Assumption 4. What we intend to convey is that the convergence rate can be further improved if we assume Assumption 4 additionally. \nIndeed, if we looked at the statement of Theorem 2, there are two types of bounds: one is without Assumption 4 and the other is with Assumption 4. We will make this point more explicit in the revision.   \nAlso, Assumptions 5 and 6 are not new and isolated assumptions, but rather, sufficient conditions for Assumption 4 when applied to different gradient estimators; and similarly, we have a valid convergence rate without these assumptions. Hence, our analysis includes the standard SGLD setting as a special case. \n\n**Q:** *Can you compare your technique with [A1]?*   \n**A:** Thank you very much for letting us know a relevant paper. \nIndeed, this paper is quite interesting and should be taken into account. \nOur bound (Theorem 2) is so general that the stochasticity is induced by not only the minibatch selection but also any other sources of randomness. Hence, we cannot directly compare with their result.\nOn the other hand, for certain special settings (SGD-MFLD and SVRG-MFLD), we can adapt their result into our analysis. \nThat is, we may replace the term $\\sigma_{v,k}^2$ by their evaluation of the square of the conditional expectation of the stochastic gradient, which would give better bound than our naive bound. On the other hand, if we assume Assumption 4, then for the SVRG setting, \nthe dependency on $\\epsilon$ is $\\min\\\\{\\sqrt{n/\\epsilon},1/\\epsilon\\\\}$ while their bound only achieves $1/\\epsilon$. Hence, we have better dependency on $\\epsilon$ in a small $\\epsilon$ regime with a stronger assumption.   \n\nThat being said, we have the following concern on the proof reasoning in [A1].  \nIn the proof of Lemma 2 which is the key lemma to obtain a tighter bound, it is argued that $N_1,\\dots,N_B$ \nare independent and identical even after conditioning on $x_{k\\eta}$ and $x_t$. \nHowever, it seems that this is not true because permutation invariance does not necessarily result in independence. \nMore precisely, it should be mentioned that the marginal distribution of $N_i$ are the same and the expectation of their average can be replaced by expectation of one of them. (We think the assertion itself is correct.)\n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n"}, {"review_id": "4wgCupRM1M", "reviewer": "Reviewer_WF3Y", "summary": "This paper studies the convergence of various algorithms falling under the umbrella of *Mean-Field Langevin Dynamics* (MFLD). Those algorithms encompass diverse approximations of the usual Langevin dynamics\n$$ dX_t = -\\nabla\\frac{\\delta F(\\mu_t)}{\\delta \\mu}(X_t) + \\sqrt{2\\lambda}dW_t $$\non three main points:\n- when the measure $\\mu_t$ is approximated by a finite-population version, i.e. the empirical measure $\\mu_X$ of $(X_t^1, \\dots X_t^N)$ all following the same dynamics,\n- when the dynamics are discretized according to a given time-step $\\eta$,\n- and finally when we only have access to a noisy version $v_t^i$ of $\\nabla\\frac{\\delta F(\\mu_X)}{\\delta \\mu}(X^i)$.\n\nThe authors provide quantitative bounds for the convergence of $\\mu_X$ to the minimum of an appropriately defined measure, which minimizes an entropy-regularized objective. They first provide a one-step bound in the general case, which is then turned into more precise convergence bounds depending on the version of $v_t^i$ considered.", "questions": "- When you remove the entropic regularization term, the existence of a limit measure for the dynamics is not guaranteed anymore; however, some papers (e.g. Chizat and Bach '18) manage to recover results by assuming that there is still a convergence. Do you think this could also be the case for your discretization results ?\n- Can your methods be applied to directly show an approximation result similar to (Mei et al. '18), i.e. an approximation bound between $\\mu_t$ and $\\mu_{X_t}$ instead of a direct control on $\\mathcal F(\\mu)$ ?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "This paper considers a very important and widely studied problem in the study of wide neural networks.The results obtained are impressive: even if the finite population approximation is borrowed from Chen et al., it is complemented by very precise and technical arguments about discretization and stochastic approximation, which makes this paper a major improvement on the state of the art.\n\nThe paper flows decently well, and even the appendix proof are nicely explained and made as easy to follow as possible (which is not a simple task, given their technicality). The significance of the paper is made clear by expanding on several well-studied instances of the general problem.", "weaknesses": "On the significance side, the main drawback of the paper is the necessity for two regularization terms (one strong convexity term in the objective function $U$, and the entropic regularization term), and the heavy dependency of the bounds on those two parameters. I am aware, unfortunately, that this is for now a common restriction for any time-independent bounds on Langevin dynamics.\n\nThe main drawback of this paper is that it is simply too technical for a Neurips submission. The presentation is good, but still suffers heavily from the 9-page (and even the 10-page, if accepted) limit, which forces a lot of inline math even for key equations (e.g. the log-Sobolev definition). This results in a very dense and sometimes hard to follow article, which could really benefit from a longer exposition. Notably, some examples (especially the first one) could be expanded upon, both to understand exactly the role of first-variation functional and the comparison with existing results on SGD/mean-field for 2LNNs.\n\nAll in all, this paper is more suited for a (very good) journal than for Neurips itself; however, I cannot recommend rejection based on the significance of the results.\n\nSome parts of the appendix are also a bit rushed, especially Appendix A, which only lists the necessary conditions for the examples to fit Assumptions 1 and 2 without any proof.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful comments. We address the technical points below. \n\n**Q:** *Some parts of the appendix are also a bit rushed, especially Appendix A, which only lists the necessary conditions for the examples to fit Assumptions 1 and 2 without any proof.*   \n**A:** Thank you very much for pointing it out. \nWe will extend some parts of Appendix including justifications of Assumptions 1 and 2 in the revision.  \n\n**Q:** *some papers (e.g. Chizat and Bach '18) manage to recover results by assuming that there is still a convergence. Do you think this could also be the case for your discretization results?*   \n**A:** Thank you for raising an interesting question. Indeed, Chizat and Bach (2018) showed global convergence by assuming convergence in the first place. That is, convergence is guaranteed only when the activation function is homogeneous and the solution satisfies some specific conditions which cannot be ensured beforehand (more precisely, if the solution converges to some distribution in $W_2$ distance, then the limit distribution is the global optimal, however whether the solution converges or not is not guaranteed). \nIn contrast, in our setting global convergence is guaranteed due to the entropy regularization, similar to (Mei et al. 2018). If we completely remove the regularization, then the quantitative convergence rate (which is exponential) would be lost unless additional assumptions are imposed, and consequently, the uniform-in-time propagation of chaos estimate would be very challenging to establish. We intend to investigate this problem in the future. \n\n\n**Q:** *Can your methods be applied to directly show an approximation result similar to (Mei et al. '18), i.e. an approximation bound between $\\mu_t$ and $\\mu_{X_t}$ instead of a direct control on $\\mathcal{F}(\\mu)$?*   \n**A:** Yes, we have a convergence of $\\mu_{X_t}$ and $\\mu_t$ to $\\mu^*$ in terms of the Wasserstein distance, which also yields convergence of the Wasserstein distance between $\\mu_{X_t}$ and $\\mu_t$ by the triangular inequality. Once again, the main difference from Mei et al. (2018) is that their discretization error blows up exponentially over the time, whereas our bound remains stable for all $t>0$. \n\nWe would be happy to clarify any concerns or answer any questions that may come up during the discussion period.\n"}]

Author rebuttal:

Meta-review: This paper show non-asymptotic convergence of the mean-field Langevin dynamics (MFLD), handling discretization errors due to particle approximation, time discretization, and stochastic gradients. This is an important problem, especially given the recent interest in MFLD. The analysis is applicable to general settings, and I agree with the reviewers that that this paper is an important contribution to the literature.
Decision: Accept (spotlight)