Title: Convergence of mean-field Langevin dynamics: time-space discretization, stochastic gradient, and variance reduction
Venue: NeurIPS 2023
Accepted: yes
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

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 entropyregularized 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 finiteparticle 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.

1 Introduction

In this work we consider the mean-field Langevin dynamics (MFLD) given by the following McKeanVlasov stochastic differential equation:

where µ t = Law( X t ) , F : P 2 ( R d ) → R is a convex functional, W t is the d -dimensional standard Brownian motion, and δF δµ denotes the first-variation of F . Importantly, MFLD is the Wasserstein gradient flow that minimizes an entropy-regularized convex functional as follows:

While the above objective can also be solved using other methods such as double-loop algorithms based on iterative linearization (Nitanda et al., 2020, 2023; Oko et 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 limit ), for which the convexity of loss function can be exploited to show the global convergence of gradient-based optimization (Nitanda and Suzuki, 2017; Mei et al., 2018; Chizat and Bach, 2018; Rotskoff and Vanden-Eijnden, 2018; Sirignano and Spiliopoulos, 2020). Under this viewpoint, MFLD Eq. (1) corresponds to 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 et al., 2018; Hu et al., 2019).

Quantitative analysis of MFLD. Most existing convergence results of neural networks in the mean-field regime are qualitative in 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 a large gap between the ideal MFLD analyzed in prior works and a feasible algorithm. In practice, we parameterize µ as a mixture of N particles ( X i ) N i =1 -this corresponds to a neural network with N neurons, and perform a discrete-time update: at time step k , the update to the i -th particle is given as

where η k is the step size at the k -th iteration, ξ i k is an i.i.d. standard Gaussian vector, and ˜ ∇ δF ( µ ) δµ represents a potentially inexact (e.g., stochastic) gradient.

Comparing Eq. (1) and Eq. (3), we observe the following discrepancies between the ideal MFLD and the implementable noisy particle gradient descent algorithm.

(i) Particle approximation. µ is entirely represented by a finite set of particles: µ k = 1 N ∑ N i =1 δ X i k .

(ii) Time discretization. We employ discrete gradient descent update as 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 a 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 finite-particle setting unless the time horizon is very short (e.g., Abbe et al. (2022)), which limits the applicability of the theory.

Very recently, Chen et al. (2022); Suzuki et al. (2023) established a uniform-in-time propagation of chaos for MFLD, i.e., the 'distance' between the N -particle system and the infinite-particle limit is of order O (1 /N ) for all t > 0 . While this represents a 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 (LD) without the mean-field interactions (which is a special case of MFLD Eq. (1) by setting F to be a 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 Teh, 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 a 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 meanfield Langevin dynamics under time and space discretization, simultaneously addressing points ( i ) -( iii ) . 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 to infinity. Moreover, our proof is based on an LSI condition analogous to Nitanda et al. (2022); Chizat (2022), which is satisfied in a wide range of regularized risk minimization problems. The advantages of our analysis is summarized as follows.

Our framework provides a unified treatment of different gradient estimators. Concretely, we establish convergence rate of MFLD with stochastic gradient and stochastic variance reduced gradient (Johnson and Zhang, 2013). While it is far from trivial to derive a tight bound on the stochastic

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 a stochastic gradient effectively improves the computational complexity.

Noticeably, despite the fact that our theorem simultaneously handles a wider class of F , when specialized to standard Langevin dynamics (i.e., when F is linear), we recover state-of-the-art convergence 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 t increases. Noticeably, we do not impose weak interaction / large noise conditions that are common in the literature (e.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 R d where the Borel σ -algebra is equipped. Our goal is to find a probability measure µ ∈ P that approximately minimizes the objective given by

where U : P → R is a (convex) loss function, and r : R d → R is a regularization term. Let P 2 be the set of probability measures with the finite second moment. In the following, we consider the setting where F ( µ ) ≤ C (1 + E µ [ ∥ X ∥ 2 ]) , and focus on P 2 so that F is well-defined.

As previously mentioned, an important application of this minimization problem is the learning of two-layer neural network in the mean-field regime . Suppose that h x ( · ) is a neuron with a parameter x ∈ R d , e.g., h x ( z ) = σ ( w ⊤ z + b ) , where w ∈ R d -1 , b ∈ R , and x = ( w,v ) . The mean-field neural network corresponding to a probability measure µ ∈ P can be written as f µ ( · ) = ∫ h x ( · ) µ (d x ) . Given training data ( z i , y i ) n i =1 ∈ R d -1 × R , we may define the empirical risk of f µ as U ( µ ) = 1 n ∑ n i =1 ℓ ( f µ ( z i ) , y i ) for a loss function ℓ : R × R → R . Then, the objective F becomes

where the regularization term is r ( x ) = λ 1 ∥ x ∥ 2 . Note that the same objective can be defined for expected risk minimization. We defer additional examples to the last of this section.

One effective way to solve the above objective is the mean-field Langevin dynamics (MFLD), which optimizes F via a noisy gradient descent update. To define MFLD, we need to introduce the first-variation of the functional F .

Definition 1. Let G : P 2 → R . The first-variation δG δµ of a functional G : P 2 → R at µ ∈ P 2 is defined as a continuous functional P 2 × R d → R that satisfies lim ϵ → 0 G ( ϵν +(1 -ϵ ) µ ) ϵ = ∫ δG δµ ( µ )( x )d( ν -µ ) for any ν ∈ P 2 . If there exists such a functional δG δµ , we say G admits a first-variation at µ , or simply G is differentiable at µ .

To avoid the ambiguity of δG δµ up to constant shift, we follow the convention of imposing ∫ δG δµ ( µ )d µ = 0 . Using the first-variation of F , the MFLD is given by the following stochastic differential equation:

where X 0 ∼ µ 0 , Law( X ) denotes the distribution of the random variable X and ( W t ) t ≥ 0 is the d -dimensional standard Brownian motion. Readers may refer to Huang et al. (2021) for the existence and uniqueness of the solution. MFLD is an instance of distribution-dependent SDE because the drift term δF ( µ t ) δµ ( · ) depends on the distribution µ t of the current solution X t (Kahn and Harris, 1951; Kac, 1956; McKean, 1966). It is known that the MFLD is a Wasserstein gradient flow to minimize the following entropy-regularized objective (Mei et al., 2018; Hu et al., 2019; Nitanda et al., 2022; Chizat, 2022):

where Ent( µ ) = -∫ log(d µ ( z ) / d z )d µ ( z ) is the negative entropy of µ .

Reduction to standard Langevin dynamics. Note that the MFLD reduces to the standard gradient Langevin dynamics (LD) when F is a linear functional, that is, there exists V such that F ( µ ) = ∫ V ( x )d µ ( x ) . In this case, δF δµ = V for any µ and the MFLD Eq. (4) simplifies to

This is exactly the gradient Langevin dynamics for optimizing V or sampling from µ ∝ exp( -V/λ ) .

2.1 Some Applications of MFLD

Here, we introduce a few examples that can be approximately solved via MFLD.

Example 1 (Two-layer neural network in mean-field regime.) . Let h x ( z ) be a neuron with a parameter x ∈ R d , e.g., h x ( z ) = tanh( rσ ( w ⊤ x )) , h x ( z ) = tanh( r ) σ ( w ⊤ x ) for x = ( r, w ) , or simply h x ( z ) = σ ( x ⊤ z ) . Then the learning of mean-field neural network f µ ( · ) = ∫ h x ( · ) µ (d x ) via minimizing the empirical risk U ( µ ) = 1 n ∑ n i =1 ℓ ( f µ ( z i ) , y i ) with a convex loss (e.g., the logistic loss ℓ ( f, y ) = log(1 + exp( -yf )) , or the squared loss ℓ ( f, y ) = ( 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 ν is defined as MMD 2 ( µ, ν ) := ∫ ∫ ( k ( x, x ) -2 k ( x, y ) + k ( y, y ))d µ ( x )d ν ( y ) . We perform nonparametric density estimation by fitting a Gaussian mixture model f µ ( z ) = ∫ g x ( z )d µ ( x ) , where g x is the Gaussian density with mean x and a given variance σ 2 > 0 . The mixture model is learned by minimizing MMD 2 ( f µ , p ∗ ) where p ∗ is the target distribution. If we observe a set of training data ( z i ) n i =1 from p ∗ , then the empirical version of MMD is one suitable loss function U ( µ ) given as

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 a posterior distribution in Bayesian inference), we may employ the kernel Stein discrepancy (KSD) as the discrepancy measure (Chwialkowski et al., 2016; Liu et al., 2016). Suppose that we can compute s ( z ) = ∇ log( µ ∗ )( z ) , then for a positive definite kernel k , we define the Stein kernel as

We take U ( µ ) as the KSD between µ and µ ∗ defined as KSD( µ ) = ∫ ∫ W µ ∗ ( z, z ′ )d µ ( z )d µ ( z ′ ) . 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 a slight optimization bias. Note that these added regularizations often have a 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 is a large gap between the ideal dynamics and a practical implementable algorithm. Specifically, we need to consider ( i ) the finite-particle approximation, ( ii ) the time discretization, and ( iii ) stochastic gradient.

To this end, we consider the following space- and time-discretized version of the MFLD with stochastic gradient update. For a finite set of particles X = ( X i ) N i =1 ⊂ R d , we define its corresponding empirical distribution as µ X = 1 N ∑ N i =1 δ X i . Let X k = ( X i k ) N i =1 ⊂ R d be N particles at the k -th update, and define µ k = µ X k as a finite particle approximation of the population counterpart. Starting from X i 0 ∼ µ 0 , we update X k as,

where η k > 0 is the step size, ξ i k is an i.i.d. standard normal random variable ξ i k ∼ N (0 , I ) , and v i k = v i k ( X 0: k , ω i k ) is a stochastic approximation of ∇ δF ( µ k ) δµ ( X i k ) where ω k = ( ω i k ) N i =1 is a random variable generating the randomness of stochastic gradient. v i k can depend on the history X 0: k = ( X 0 , . . . , X k ) , and E ω i k [ v i k | X 0: k ] = ∇ δF ( µ k ) δµ ( X i k ) . We analyze three versions of v i k .

Full gradient: F-MFLD. If we have access to the exact gradient, we may compute

Stochastic gradient: SGD-MFLD. Suppose the loss function U is given by an expectation as U ( µ ) = E Z ∼ P Z [ ℓ ( µ, Z )] , where Z ∈ Z is a random observation obeying a distribution P Z and ℓ : P × Z → R is a loss function. In this setting, we construct the stochastic gradient as

where ( z ( j ) k ) B j =1 is a mini-batch of size B generated from P Z in an i.i.d. manner.

Stochastic variance reduced gradient: SVRG-MFLD. Suppose that the loss function U is given by a finite sum of loss functions ℓ i : U ( µ ) = 1 n ∑ n i =1 ℓ i ( µ ) , 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:

(i) When k ≡ 0 mod m , where m is the update frequency, we set ˙ X = ( ˙ X i ) n i =1 = ( X i k ) n i =1 as the anchor point and refresh the stochastic gradient as v i k = ∇ δF ( µ k ) δµ ( X i k ) .

(ii) When k ̸≡ 0 mod m , we use the anchor point to construct a control variate and compute

where I k is a uniformly drawn random subset of { 1 , . . . , n } with size B without duplication.

3 Main Assumptions and Theoretical Tools

For our convergence analysis, we make the following assumptions which are 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 r are convex. Specifically,

U : P → R is a convex functional on P , that is, U ( θµ +(1 -θ ) ν ) ≤ θU ( µ ) + (1 -θ ) U ( ν ) for any θ ∈ [0 , 1] and µ, ν ∈ P . Moreover, U admits a first-variation at any µ ∈ P 2 .

r ( · ) is twice differentiable and convex, and there exist constants λ 1 , λ 2 > 0 and c r > 0 such that λ 1 I ⪯ ∇∇ ⊤ r ( x ) ⪯ λ 2 I , x ⊤ ∇ r ( x ) ≥ λ 1 ∥ x ∥ 2 , and 0 ≤ r ( x ) ≤ λ 2 ( c r + ∥ x ∥ 2 ) for any x ∈ R d , and ∇ r (0) = 0 .

Assumption 2. There exists L > 0 such that ∥ ∥ ∥ ∇ δU ( µ ) δµ ( x ) -∇ δU ( µ ′ ) δµ ( x ′ ) ∥ ∥ ∥ ≤ L ( W 2 ( µ, µ ′ ) + ∥ x -x ′ ∥ ) and ∣ ∣ ∣ δ 2 U ( µ ) δµ 2 ( x, x ′ ) ∣ ∣ ∣ ≤ L (1+ c L ( ∥ x ∥ 2 + ∥ x ′ ∥ 2 )) for any µ, µ ′ ∈ P 2 and x, x ′ ∈ R d . Also, there exists R > 0 such that ∥∇ δU ( µ ) δµ ( x ) ∥ ≤ R for any µ ∈ P and x ∈ R d .

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-time propagation of chaos (Lemma 8 in Appendix E.4). Under these assumptions, Proposition 2.5 of Hu et al. (2019) yields that F has a unique minimizer µ ∗ in P and µ ∗ is absolutely continuous with respect to the Lebesgue measure. Moreover, µ ∗ satisfies the self-consistent condition: µ ∗ ( X ) ∝ exp ( -1 λ δF ( µ ∗ ) δµ ( X ) ) .

3.1 Proximal Gibbs Measure & Logarithmic Sobolev inequality

The aforementioned self-consistent relation motivates us to introduce the proximal Gibbs distribution (Nitanda et al., 2022; Chizat, 2022) whose density is given by

where X = ( X i ) N i =1 ⊂ R d is a set of N particles and X ∈ R d . As we will see, the convergence of MFLD heavily depends on a logarithmic Sobolev inequality (LSI) on the proximal Gibbs measure. Definition 2 (Logarithmic Sobolev inequality) . Let µ be a probability measure on ( R d , B ( R d )) . µ satisfies the LSI with constant α > 0 if for any smooth function ϕ : R d → R with E µ [ ϕ 2 ] < ∞ , we have E µ [ ϕ 2 log( ϕ 2 )] -E µ [ ϕ 2 ] log( E µ [ ϕ 2 ]) ≤ 2 E µ [ ∥∇ ϕ ∥ 2 ] .

α 2

This is equivalent to the condition that the KL divergence from µ is bounded by the Fisher divergence: ∫ log(d ν/ d µ )d ν ≤ 1 2 α ∫ ∥∇ log(d ν/ d µ ) ∥ 2 d ν, for any ν ∈ P which is absolutely continuous with respect to µ . Our analysis requires that the proximal Gibbs distribution satisfies the LSI as follows. Assumption 3. µ ∗ and p X satisfy the LSI with α > 0 for any set of particles X = ( X i ) N i =1 ⊂ R d .

Verification of LSI. The LSI of proximal Gibbs measure can be established via standard perturbation criteria. For U ( µ ) with bounded first-variation, we may apply the classical Bakry-Emery and Holley-Stroock 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 trick (Bardet et al., 2018) or the more recent perturbation results in Cattiaux and Guillin (2022).

Theorem 1. Under Assumptions 1 and 2, µ ∗ and p X satisfy the log-Sobolev inequality with

Furthermore, if ∥ δU ( µ ) δµ ∥ ∞ ≤ R is satisfied for any µ ∈ P 2 , then µ ∗ and p X satisfy the LSI with α ≥ λ 1 λ exp ( -4 R λ ) .

See Lemma 5 for the proof of the first assertion, and Section A.1 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, we 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 N particles, that is, let µ ( N ) ∈ P ( N ) be a distribution of N particles X = ( X i ) N i =1 where P ( N ) is the set of probability measures on ( R d × N , B ( R d × N )) . Similarly, we introduce the following objective on P ( N ) :

One can easily verify that if µ ( N ) is a product measure of µ ∈ P , then F N ( µ ( N ) ) ≥ N F ( µ ) by the convexity of F . Propagation of chaos (Sznitman, 1991) refers to the phenomenon that, as the number of particles N increases, the particles behave as if they are independent; in other words, the joint distribution of the N particles becomes 'close' to the product measure. Consequently, the minimum of F N ( µ ( N ) ) (which can be obtained by the particle-approximated MFLD) is close to N F ( µ ∗ ) . Specifically, it has been shown in Chen et al. (2022) that

for some constant C λ > 0 (see Section E for more details). Importantly, if we consider the Wasserstein gradient flow on 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 F N 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 error 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 a straightforward manner. Roughly speaking, this error is characterized by the variance of v i k : σ 2 v,k := E ω k , X 0: k [ ∥ v i k - ∇ δF ( µ k ) δµ ∥ 2 ] 1 . In addition, to obtain a refined evaluation, we also incorporate the following smoothness assumption.

̸

Assumption 4. v i k ( X 0: k , ω i k ) and D i m F ( X k ) = ∇ δF ( µ X k ) δµ ( X i k ) are differentiable w.r.t. ( X j k ) N j =1 and, for either G ( X k ) = v i k ( X 0: k , ω i k ) or G ( X k ) = D i m F ( X k ) as a function of X k , it is satisfied that ∥ G ( X ′ k ) -G ( X k ) -∑ N j =1 ∇ ⊤ X j k G ( X k ) · ( X ′ j k -X j k ) ∥ ≤ Q ∑ j = i ∥ X ′ j k -X j k ∥ 2 + N ∥ X ′ i k -X i k ∥ 2 2 N for some Q > 0 and X k = ( X j k ) N j =1 and X ′ k = ( X ′ j k ) N j =1 . We also assume E ω k [∥ ∥ ∇ X j k v i k ( X 0: k , ω i k ) ⊤ -∇ X j k D i m F ( X k ) ⊤ ∥ ∥ 2 op ] ≤ 1+ δ i,j N 2 N 2 ˜ σ 2 v,k , and ∥ v i k (( X k , X 0: k -1 ) , ω i k ) -v i k (( X ′ k , X 0: k -1 ) , ω i k ) ∥ ≤ L ( W 2 ( µ X k , µ X ′ k ) + ∥ X i k -X ′ i k ∥ ) .

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 = j . In the following, we present both the basic convergence result under Assumption 2, and the improved rate under the additional Assumption 4.

Taking these factors into consideration, we can evaluate the decrease in the objective value after one-step update. Let µ ( N ) k ∈ P ( N ) be the distribution of X k conditioned on the sequence ω 0: k = ( ω k ′ ) k k ′ =0 . Define

where C 1 = 8[ R 2 + λ 2 ( c r + ¯ R 2 ) + d ] and ¯ L = L + λ 2 . 2

Theorem 2. Under Assumptions 1, 2 and 3, if η k ≤ λ 1 / (4 λ 2 ) , we have

where Υ k = 4 η k δ η k + ( R + λ 2 ¯ R ) (1 + √ λ/η k ) η 2 k ˜ σ v,k σ v,k + Qη 3 k σ 2 v,k + η 2 k ( L + λ 2 ) 2 σ 2 v,k with Assumption 4, and Υ k = σ 2 v,k η k without this additional assumption; the expectation is taken with respect to the randomness ( ω k ′ ) k k ′ =1 = (( ω i k ′ ) N i =1 ) k k ′ =1 of the stochastic gradient; and C λ = 2 λLα (1 + 2 c L ¯ R 2 ) + 2 λ 2 L 2 ¯ R 2 .

The proof can be found in Appendix B. We remark that to derive the bound for Υ k = σ 2 v,k η k is relatively straightforward, but to derive a tighter bound with Assumption 4 is technically challenging, because we need to evaluate how the next-step distribution µ ( N ) k +1 is correlated with the stochastic gradient v i k . Evaluating such correlations is non-trivial because the randomness is induced not only by ω k but also by the Gaussian noise. Thanks to this general result, we only need to evaluate the variance σ v,k and ˜ σ v,k for each method to obtain the specific convergence rate.

Conversion to a Wasserstein distance bound. As a consequence of the bound on 1 N F N ( µ ( N ) k ) -F ( µ ∗ ) , we can control the Wasserstein distance between µ ( N ) k and µ ∗ N , where µ ∗ N ∈ P ( N ) is the ( N -times) product measure of µ ∗ . Let W 2 ( µ, ν ) be the 2-Wasserstein distance between µ and ν , then

under Assumptions 1 and 3 (see Lemma 3 in the Appendix). Hence, if 1 N F N ( µ ( N ) k ) -F ( µ ∗ ) is small, the particles ( X i k ) N i =1 behaves like an i.i.d. sample from µ ∗ . As an example, for the mean-field neural network setting, if | h x ( z ) -h x ′ ( z ) | ≤ L ∥ x -x ′ ∥ ( ∀ x, x ′ ∈ R d ) and V µ ∗ = ∫ ( f µ ∗ ( z ) -h x ( z )) 2 d µ ∗ ( x ) < ∞ for a fixed z , then Lemma 4 in the appendix yields that

which also gives E X k ∼ µ ( N ) k [( f µ X k ( z ) -f µ ∗ ( z )) 2 ] ≤ 4 L 2 λα ( N -1 F N ( µ ( N ) k ) -F ( µ ∗ ) ) + 2 N V µ ∗ .

This allows us to monitor the convergence of the finite-width neural network to the optimal solution µ ∗ 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-MFLD and SGD-MFLD simultaneously. F-MFLD can be seen as a special case of SGD-MFLD where the variance σ 2 v,k = 0 . We specialize the previous assumptions to the stochastic gradient setting as follows.

Assumption 5.

Note that point ( i ) is required to bound the variance σ 2 v,k (i.e., σ 2 v,k ≤ R 2 /B ), whereas point ( ii ) corresponds to the additional Assumption 4 required for the improved convergence rate. Let ∆ 0 := 1 N E [ F N ( µ ( N ) 0 )] -F ( µ ∗ ) . We have the following evaluation of the objective. Note that we can recover the evaluation for F-MFLD by formally setting B = ∞ so that σ 2 v,k = 0 and ˜ σ 2 v,k = 0 .

Theorem 3. Suppose that η k = η ( ∀ k ∈ N 0 ) and λαη ≤ 1 / 4 and η ≤ λ 1 / (4 λ 2 ) . Under Assumptions 1, 2, 3 and 5, it holds that

Assumption 5-(ii), and ¯ Υ = R 2 B η 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 a convergence rate of discrete time vanilla GLD with respect to the KL divergence. Indeed, their derived rate O (exp( -λαηk )∆ 0 + η αλ ) is consistent to ours, since Eq. (26) in the Appendix implies that the objective 1 λ ( 1 N E [ F N ( µ ( N ) k )] -F ( µ ∗ )) upper bounds the KLdivergence between µ ( N ) k and µ ∗ N . 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 a given ϵ > 0 , if we take

with B ≥ 4 [ (1 + R )( R + λ 2 ¯ R ) + ( L + λ 2 ) 2 ] R 2 ( η + √ ηλ ) / ( λαϵ ) under Assumption 5-(ii) and B ≥ 4 R 2 / ( λαϵ ) without such assumption, then the right hand side of (8) can be bounded as 1 N E [ F N ( µ ( N ) k )] -F ( µ ∗ ) = ϵ + 4 C λ λαN . Hence we achieve ϵ + O (1 /N ) error with iteration complexity:

If we neglect O (1 / √ λαϵ ) as a second order term, then the above can be simplified as O ( ¯ L 2 λα 2 log( ϵ -1 ) ϵ ) .

Noticeably, the mini-batch size B can be significantly reduced under the additional smoothness condition Assumption 5-(ii) (indeed, we have O ( η + √ ηλ ) 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 O ( nk ) where k 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 + √ ( ϵλα ) -1 ) which yields the objective value of order O ( ϵ ) , SGD-MFLD achieves O ( n/B ) = O ( n ( λ ∧ √ ϵλα )) times smaller total complexity. For example, when λ = αϵ = 1 / √ n , a O ( √ n ) -factor reduction of total complexity can be achieved by SGD-MFLD, which is a significant improvement.

4.3 SVRG-MFLD

Nowwepresent the convergence rate of SVRG-MFLD in the fixed step size setting where η k = η ( ∀ k ) . Instead of Assumption 5, we introduce the following two assumptions.

Assumption 6.

Here again, point (i) is required to bound σ 2 v,k and point (ii) yields Assumption 4. We have the following computational complexity bound for SVRG-MFLD.

Theorem 4. Suppose that λαη ≤ 1 / 4 and η ≤ λ 1 / (4 λ 2 ) . Let Ξ = n -B B ( n -1) . Then, under Assumptions 1, 2, 3 and 6, we have the same error bound as Eq. (8) with different ¯ Υ :

under Assumption 6-(ii), and ¯ Υ = C 1 Ξ L 2 mη 2 ( η + λ ) without Assumption 6-(ii).

The proof is given in Appendix D. Therefore, to achieve 1 N E [ F N ( µ ( N ) k )] -F ( µ ∗ ) ≤ O ( ϵ ) + 4 C λ λαN for a given ϵ > 0 , it suffices to set with B ≥ [ √ m ( η + √ η/λ ) 2 ηλ ∨ m ( η + √ η/λ ) 3 ηλ ] ∧ n . In this setting, the total gradient complexity can be bounded as

where m = Ω( n/B ) = Ω([ n 2 / 3 (1 + √ η/λ ) -4 / 3 n (1 + √ η/λ ) -3 / 2 ( ηλ ) -1 / 2 )] ∨ 1) and B = O ([ n 1 3 ( 1 + √ η λ ) 4 3 ∨ √ n ( ηλ ) 1 4 ( 1 + √ η λ ) 3 2 ] ∧ n ) . When ϵ is small, the first term in the right hand side becomes the main term which is max { n 1 / 3 , √ n ( αϵλ ) 1 / 4 } log( ϵ -1 ) λα 2 ϵ . Therefore, comparing with the full batch gradient method (F-MFLD), SVRG-MFLD achieves at least min { n 2 3 , √ n ( αϵλ ) -1 4 } times better total computational complexity when η ≤ λ . Note that even without the additional Assumption 6-(ii), we still obtain a √ n -factor improvement (see Appendix D). This indicates that variance reduction is indeed effective to improve the computational complexity, especially in a large sample size setting.

Finally, we 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 was given in Kinoshita and Suzuki (2022), which yields a O (( n + √ n ϵ ) log( ϵ -1 ) ( λα ) 2 ) iteration complexity to achieve ϵ error. This corresponds to 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 ( λ ) -factor discrepancy is due to the difference in 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-(ii)).

5 Conclusion

We gave a unified theoretical framework to bound the optimization error of the single-loop mean-field Langevin dynamics (MFLD) that is applicable to the finite-particle, discrete-time, and stochastic gradient algorithm. Our analysis is general enough to cover several important learning problems such as 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 (FMFLD, SGD-MFLD, SGLD-MFLD); and despite the fact that our analysis deals with a more general setting (mean-field interactions), we are able to recover and even improve existing convergence guarantees when specialized to the standard gradient Langevin dynamics.

References

Abbe, E. B. Adsera, and T. Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory , pages 4782-4887. PMLR, 2022.

Arbel, A. Korba, A. SALIM, and A. Gretton. Maximum mean discrepancy gradient flow. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and

R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper/2019/file/ 944a5ae3483ed5c1e10bbccb7942a279-Paper.pdf .

Bakry and M. Émery. Diffusions hypercontractives. In J. 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-39397-9.

Bakry and M. Émery. Diffusions hypercontractives in sem. probab. xix lnm 1123, 1985b.

Bakry, I. Gentil, M. Ledoux, et al. Analysis and geometry of Markov diffusion operators , volume 103. Springer, 2014.

J.-B. Bardet, N. Gozlan, F. Malrieu, and P.-A. Zitt. Functional inequalities for Gaussian convolutions of compactly supported measures: Explicit bounds and dimension dependence. Bernoulli , 24(1): 333 - 353, 2018. doi: 10.3150/16-BEJ879. URL https://doi.org/10.3150/16-BEJ879 .

Cattiaux and A. Guillin. Semi Log-Concave Markov Diffusions , pages 231-292. Springer International Publishing, Cham, 2014.

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.3150/21-BEJ1419. URL https://doi.org/10.3150/21-BEJ1419 .

Chen, Z. Ren, and S. Wang. Uniform-in-time propagation of chaos for mean field langevin dynamics. arXiv preprint arXiv:2212.03050v1 , 2022.

Chen, G. M. Rotskoff, J. Bruna, and E. Vanden-Eijnden. A dynamical central limit theorem for shallow neural networks. arXiv preprint arXiv:2008.09623 , 2020.

Chewi, M. A. Erdogdu, M. B. Li, R. Shen, and M. Zhang. Analysis of langevin monte carlo from poincar \ 'e to log-sobolev. arXiv preprint arXiv:2112.12662 , 2021.

Chizat. Mean-field langevin dynamics : Exponential convergence and annealing. Transactions on Machine Learning Research , 2022. URL https://openreview.net/forum?id=BDqzLH1gEm .

Chizat and F. Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. In Advances in Neural Information Processing Systems 31 , pages 3040-3050, 2018.

Chwialkowski, H. Strathmann, and A. Gretton. A kernel test of goodness of fit. In M. F. Balcan and K. Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 2606-2615, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https://proceedings.mlr.press/ v48/chwialkowski16.html .

S. Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. arXiv preprint arXiv:1412.7392 , 2014.

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 on Learning Theory , volume 195 of Proceedings of Machine Learning Research , pages 4072-4129. PMLR, 12-15 Jul 2023. URL https://proceedings.mlr.press/v195/ das23b.html .

De Bortoli, A. Durmus, X. Fontaine, and U. Simsekli. Quantitative propagation of chaos for sgd in wide neural networks. Advances in Neural Information Processing Systems , 33:278-288, 2020.

Delarue and A. Tse. Uniform in time weak propagation of chaos on the torus. arXiv preprint arXiv:2104.14973 , 2021.

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, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016. URL https://proceedings.neurips.cc/paper/2016/file/ 9b698eb3105bd82528f23d0c92dedfc0-Paper.pdf .

Gretton, K. Borgwardt, M. Rasch, B. Schölkopf, and A. Smola. A kernel method for the twosample-problem. In B. Schölkopf, J. Platt, and T. Hoffman, editors, Advances in Neural Information Processing Systems , volume 19. MIT Press, 2006. URL https://proceedings.neurips.cc/ paper/2006/file/e9fb2eda3d9c55a0d89c98d6c54b5b3e-Paper.pdf .

Holley and D. Stroock. Logarithmic sobolev inequalities and stochastic ising models. Journal of statistical physics , 46(5-6):1159-1194, 1987.

Hu, Z. Ren, D. Siska, and L. Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. arXiv preprint arXiv:1905.07769 , 2019.

Huang, P. Ren, and F.-Y. Wang. Distribution dependent stochastic differential equations. Frontiers of Mathematics in China , 16(2):257-301, 2021.

Javanmard, M. Mondelli, and A. Montanari. Analysis of a two-layer neural network via displacement convexity. arXiv preprint arXiv:1901.01375 , 2019.

Johnson and T. Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Weinberger, editors, Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013. URL https://proceedings.neurips.cc/paper/2013/file/ ac1dd209cbcc5e5d1c6e28598e8cbbe8-Paper.pdf .

Kac. Foundations of kinetic theory. In Proceedings of The third Berkeley symposium on mathematical statistics and probability , volume 3, pages 171-197, 1956.

Kahn and T. E. Harris. Estimation of particle transmission by random sampling. National Bureau of Standards applied mathematics series , 12:27-30, 1951.

Kinoshita and T. Suzuki. Improved convergence rate of stochastic gradient langevin dynamics with variance reduction and its application to optimization. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https: //openreview.net/forum?id=Sj2z__i1wX-.

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://proceedings.mlr.press/v139/korba21a.html .

Lacker. Hierarchies, entropy, and quantitative propagation of chaos for mean field diffusions. arXiv preprint arXiv:2105.02983 , 2021.

Lacker and L. L. Flem. Sharp uniform-in-time propagation of chaos. arXiv preprint arXiv:2205.12047 , 2022.

Liu, J. Lee, and M. Jordan. A kernelized stein discrepancy for goodness-of-fit tests. In M. F. Balcan 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 https://proceedings.mlr.press/ v48/liub16.html .

Y.-A. Ma, T. Chen, and E. Fox. A complete recipe for stochastic gradient mcmc. Advances in neural information processing systems , 28, 2015.

P. McKean. A class of markov processes associated with nonlinear parabolic equations. Proceedings of the National Academy of Sciences , 56(6):1907-1911, 1966.

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-E7671, 2018.

Mischler. An introduction to evolution PDEs, Chapter 0: On the Gronwall lemma, 2019. URL https://www.ceremade.dauphine.fr/~mischler/Enseignements/ M2evol2018/chap0.pdf .

Nitanda and T. Suzuki. Stochastic particle gradient descent for infinite ensembles. arXiv preprint arXiv:1712.05438 , 2017.

Nitanda, D. Wu, and T. Suzuki. Particle dual averaging: Optimization of mean field neural networks with global convergence rate analysis, 2020.

Nitanda, D. Wu, and T. Suzuki. Particle dual averaging: Optimization of mean field neural network with global convergence rate analysis. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 19608-19621. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/ paper/2021/file/a34e1ddbb4d329167f50992ba59fe45a-Paper.pdf .

Nitanda, D. Wu, and T. Suzuki. Convex analysis of the mean field langevin dynamics. In G. CampsValls, F. J. R. Ruiz, and I. Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , volume 151 of Proceedings of Machine Learning Research , pages 9741-9757. PMLR, 28-30 Mar 2022.

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.

Oko, T. Suzuki, A. Nitanda, and D. Wu. Particle stochastic dual coordinate ascent: Exponential convergent algorithm for mean field neural network optimization. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=PQQp7AJwz3 .

Otto and C. Villani. Generalization of an inequality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis , 173(2):361-400, 2000.

M. Rotskoff and E. Vanden-Eijnden. Trainability and accuracy of neural networks: An interacting particle system approach. arXiv preprint arXiv:1805.00915 , 2018.

Sirignano and K. Spiliopoulos. Mean field analysis of neural networks: A central limit theorem. Stochastic Processes and their Applications , 130(3):1820-1852, 2020.

Suzuki, A. Nitanda, and D. Wu. Uniform-in-time propagation of chaos for the mean field gradient langevin dynamics. In Submitted to The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=_JScUk9TBUn .

A.-S. Sznitman. Topics in propagation of chaos. In Ecole d'été de probabilités de Saint-Flour XIX-1989 , pages 165-251. Springer, 1991.

Vempala and A. Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. In Advances in Neural Information Processing Systems , pages 8094-8106, 2019.

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.

Zou, P. Xu, and Q. Gu. Subsampled stochastic variance-reduced gradient langevin dynamics. In International Conference on Uncertainty in Artificial Intelligence , 2018.

Table of Contents

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 x ( · ) 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 C > 0 such that sup z sup x | h x ( z ) | ≤ C , sup z sup x ∥∇ x h x ( z ) ∥ ≤ C and ∥∇ h ( x ) -∇ h ( x ′ ) ∥ ≤ C ∥ x -x ′ ∥ for all x, x ′ ∈ R d , and sup y,z sup µ ∈P | ∂ f ℓ ( f µ ( z ) , y ) | ≤ C , | ∂ f ℓ ( f µ ( z ) , y ) -∂ f ℓ ( f µ ′ ( z ) , y ) | ≤ C | f µ ( z ) -f µ ′ ( z ) | and | ∂ 2 f ℓ ( f µ ( z ) , y ) | ≤ C for all ( y, z ) and µ, µ ′ ∈ P .

(ii) MMDminimization. The kernel k is smooth and has light tail, e.g., the Gaussian RBF kernel.

(iii) KSD minimization. The kernel k has a light tail such that sup z,z ′ max {| W µ ∗ ( z, z ′ ) | , ∥∇ z W µ ∗ ( z, z ′ ) ∥ , ∥∇ z ∇ ⊤ z W µ ∗ ( z, z ′ ) ∥ op } ≤ C ; for example, k ( z, z ′ ) = exp ( -∥ z ∥ 2 2 σ 2 1 -∥ z ′ ∥ 2 2 σ 2 1 -∥ z -z ′ ∥ 2 2 σ 2 2 ) , and max {∥∇ log( µ ∗ ( z )) ∥ , ∥∇ ⊗ 2 log( µ ∗ ( z )) ∥ op , ∥∇ ⊗ 3 log( µ ∗ ( z )) ∥ op } ≤ C (1 + ∥ z ∥ ) .

B Proof of Theorem 2

This section gives the proof of Theorem 2. For notation simplicity, we write η to indicate η k . Recall that for X = ( X i ) N i =1 , the proximal Gibbs distribution µ X is defined by

We also define another version of the proximal Gibbs distribution corresponding to F N as

B.1 Evaluation of Objective Decrease

By the same argument as Chen et al. (2022), we have that

where ∇ i is the partial derivative with respect to X i . Therefore, we have

Remembering that X k = ( X i k ) N i =1 is updated by X i k +1 = X i k -ηv i k + √ 2 ληξ i k , the solutions X i k and X i k +1 can be interpolated by the following continuous time dynamics:

for 0 ≤ t ≤ η . Then, ˜ X η obeys the same distribution as X k +1 . Let ˜ µ ( N ) t be the law of ˜ X t . The Fokker-Planck equation of the dynamics yields that

Hence, by taking expectation with respect to X 0: k , it holds that

where we omitted the influence of ω 0: k to ˜ µ ( N ) t , which should be written as ˜ µ ( N ) t ( X | ω 0: k ) in a more precise manner. Combining this and Eq. (10) yields the following decomposition,

Evaluation of term A : By the leave-one-out argument in the proof of Theorem 2.1 of Chen et al. (2022), the first term of the right hand side can be upper bounded by

with a constant C λ . The proof is given in Lemma 7 for completeness, where we see that the inequality holds with C λ = 2 λLα (1 + 2 c L ¯ R 2 ) + 2 λ 2 L 2 ¯ R 2 .

Evaluation of term B : By Lemma 2, we have

B.2 Stochastic Gradient Error (Term C )

Now we evaluate the final term C. First, we derive a bound without Assumption 4. By the CauchySchwarz inequality, we have that

B.2.1 Analysis under Assumption 4

Next, we derive a tighter result under the additional Assumption 4. We decompose term C as follows:

where X ′ 0: k is an independent copy of X 0: k .

Evaluation of C -(I):

(1) The first term ( C -(I)) of the right hand side of Eq. (12) can be evaluated as

where ( ˜ X t , X 0: k -1 ) = ( X 0 , X 1 , . . . , X k -1 , ˜ X t ) which is obtained by replacing X k of X 0: k to ˜ X t . First, we evaluate the first term in the right hand side of Eq. (13). Let

for X = ( X i ) N , and let D m F ( X ) = ( D F ( X )) N . Let

and let ∆ v = (∆ v i ) N i =1 . We also define

Then Z i ∼ N (0 , 1) conditioned by X ′ 0: k and is independent of ω i k . We also define

where the expectation is taken conditioned by ω 1: k but we omit ω 1: k -1 from the left hand side for the simplicity of notation. Since the density of the conditional distribution ˜ µ ( N ) t ( ˜ X t | X ′ 0: k ) is proportional to exp ( -∑ N i =1 ∥ ˜ X i t -( X ′ i k -tv i k ( X ′ 0: k ,ω i k )) ∥ 2 2(2 λt ) ) , it holds that, under second order differentiability of v i k and D m F ,

where Jensen'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

̸

for a j > 0 ( j = 1 , . . . , N ) where a j = a -i ( j = i ) for some a -i > 0 . By taking a -i = Ntσ v,k / ˜ σ v,k and a i = tσ v,k / ˜ σ v,k , the right hand side can be bounded as

Then, we return to the right hand side of Eq. (14). First, we note that

for fixed Z and X ′ k . Hence, by taking this and Eq. (15) into account, the expectation of the right hand side Eq. (14) with respect to ω k can be further bounded as

(2) Next, we evaluate the second term of the right hand side of Eq. (13). We have

The first term in the right hand side of this inequality (Eq. (17)) can be upper bounded by

where we used Jensen's inequality. By Lemma 2 (the same bound applies to v i k by the same proof), we see that the second term in the right hand side is bounded by

Finally, we evaluate the second term in the right hand side of Eq. (17). Note that

where we used Lemma 2 with the same argument as Eq. (18) and Young's inequality. Taking the expectation with respect to ω i k , it holds that

Hence, it suffices to evaluate the term

Here, let ˆ X t = ( ˆ X i t ) n i =1 be the following stochastic process:

for 0 ≤ t ≤ η , where ( W i t ) t is the same Brownian motion as that drives ˜ X i t . Then, the term we are interested in can be evaluated as

With the same reasoning as in Eq. (16), it holds that

Moreover, Assumption 2 and Lemm 1 yield that

where the second and third inequalities are due to Assumption 2 and the last inequality is by Lemma 1. Then, the right hand side of Eq. (19) can be bounded by

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

(3) By summing up the results in (1) and (2), we have that

Evaluation of C -(II):

Next, we evaluate the term C -(II). Here, let ˆ X t = ( ˆ X i t ) n i =1 be the following stochastic process:

for 0 ≤ t ≤ η , where ( W i t ) t is the same Brownian motion as that drives ˜ X i t . By definition, we notice that which yields that

By subtracting and adding the derivative at ˆ X t from ∇ i log( p ( N ) ( ˜ X t )) , we have

By Assumptions 1 and 2, the first two terms in the right hand side can be bounded as

where the first inequality follows from

by Assumption 1 and we used Eq. (20) in the last inequality. Therefore, by noticing ˜ X 0 = X k , we can obtain the following evaluation:

Combining the bounds on C -(I) and C -(II):

B.3 Putting Things Together

By Gronwall lemma (see Mischler (2019), for example), we arrive at

In addition, if Assumption 4 is satisfied, we have

C Proof of Theorem 3: F-MFLD and SGD-MFLD

Applying the Gronwall lemma (Mischler, 2019) with Theorem 2, we have that

Here, remember that Assumption 5 yields that

Under Assumption 5, we can easily check that Assumption 4 holds with Q = R . When η k = η for all k ≥ 0 , we have a uniform upper bound of Υ k as

We denote the right hand side as ¯ Υ . Then we have

Then, by taking

then the right hand side can be bounded as

(1) Hence, without Assumption 5-(ii), if we take

then the right hand side can be bounded as

We can easily check that the number of iteration k satisfies

in this setting.

(2) On the other hand, under Assumption 5-(ii), if we take

then the right hand side can be bounded as

We can again check that it suffices to take the number of iteration k as

to achieve the targe accuracy.

D Proof of Theorem 4: SVRG-MFLD

The standard argument of variance yields

where ∆ i,j k = ∇ δℓ j ( µ ( N ) k ) δµ ( X i k ) - ∇ δU ( µ ( N ) k ) δµ ( X i k ) - ∇ δℓ i ( µ ( N ) s ) δµ ( X i s ) + ∇ δU ( µ ( N ) s ) δµ ( X i s ) . Here, Assumption 2 and 6 give that

Taking the expectation, we have

where we used that E [ ∥ v i k ∥ 2 ] ≤ 2( R 2 + λ 2 ( c r + ¯ R 2 )) ≤ C 1 by Eq. (24) with Lemma 1. Hence, we have that σ 2 v,k ≤ C 1 Ξ L 2 m ( η 2 + ηλ ) , where Ξ = n -B B ( n -1) .

On the other hand, we have

by the same argument with SGD-MFLD. We have a uniform upper bound of Υ k as

We again let the right hand side be ¯ Υ . Then, we have that

we have

Then, by the Gronwall's lemma yields that

where we used λαη/ 2 ≤ 1 / 2 in the last inequality.

(1) Without Assumption 6-(ii), we have

we obtain that

when

and λαη/ 2 ≤ 1 / 2 . In particular, if we set B ≥ m and

with the iteration complexity and the total gradient computation complexity:

where m = B = √ n .

(2) With Assumption 6-(ii), we have

Then, we obtain that

when

we have

with the iteration complexity and the total gradient computation complexity:

E Auxiliary lemmas

E.1 Properties of the MFLD Iterates

Under Assumption 2 with Assumption 5 for SGD-MFLD or 6 for SVRG-MFLD, we can easily verify that

Lemma 1. Under Assumption 2 with Assumption 5 for SGD-MFLD or Assumption 6 for SVRGMFLD, if η ≤ λ 1 / (4 λ 2 ) , we have the following uniform bound of the second moment of X i k :

for any k ≥ 1 .

Proof. By the update rule of X i k and Assumption 1, we have

where we used the assumption η ≤ λ 1 / (4 λ 2 ) . Then, by the Gronwall lemma, it holds that

This with the assumption η ≤ λ 1 / (4 λ 2 ) yields the assertion.

Lemma 2. Let ¯ R 2 := E [ ∥ X i 0 ∥ 2 ] + 1 λ 1 [( λ 1 4 λ 2 + 1 λ 1 ) ( R 2 + λ 2 c r ) + λd ] . Under Assumptions 1 and 2, if η ≤ λ 1 / (4 λ 2 ) and we define

then it holds that

Proof. The proof is basically a mean field generalization of Nitanda et al. (2022). By Assumptions 1 and 2, the first two terms in the right hand side can be bounded as

Therefore, the right hand side can be further bounded as

Then, by taking the expectation, it holds that

Then, by noticing t ≤ η , we obtain the assertion.

E.2 Wasserstein Distance Bound

Recall that W 2 ( µ, ν ) is the 2-Wasserstein distance between µ and ν . We let D ( µ, ν ) = ∫ log ( ν µ ) d ν be the KL-divergence between µ and ν .

Lemma 3. Under Assumptions 1 and 3, it holds that

Proof. By Assumption 3, µ ∗ satisfies the LSI condition with a constant α > 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 µ ∗ N :

Moreover, the proof of Theorem 2.11 of Chen et al. (2022) yields that, under Assumption 1, it holds that

Lemma 4. Suppose that | h x ( z ) -h x ′ ( z ) | ≤ L ∥ x -x ′ ∥ ( ∀ x, x ′ ∈ R d ) and let V µ ∗ := Var µ ∗ ( f µ ∗ ) = ∫ ( f µ ∗ ( z ) -h x ( z )) 2 d µ ∗ ( x ) , then it holds that

Proof. Consider a coupling γ of µ ( N ) X k and µ ∗ N and let ( X k , X ∗ ) = (( X i k ) N i =1 , ( X i ∗ ) N i =1 ) be a random variable obeying the law γ . Then,

Then, by taking the expectation of the both side with respect to ( X k , X ∗ ) , we have that

Hence, taking the infimum of the coupling γ yields that

Remark 2. By the self-consistent condition of µ ∗ , we can see that µ ∗ is sub-Gaussian. Therefore, V µ ∗ < ∞ is always satisfied.

E.3 Logarithmic Sobolev Inequality

Lemma 5. Under Assumptions 1 and 2, µ ∗ and p X satisfy the LSI condition with a constant

Proof. In the following, we give two lower bounds of α . By taking the maximum of the two, we obtain the assertion.

(1) By Assumption 2, δU ( µ ) δµ is Lipschitz continuous with the Lipschitz constant R , which implies that 1 λ δU ( µ ) δµ is R/λ -Lipschitz continuous. Hence, µ ∗ and p X is a Lipshitz perturbation of ν ( x ) ∝ exp( -λ -1 r ( x )) . Since λ -1 ∇∇ ⊤ r ( x ) ⪰ λ 1 λ I by Assumption 1, Miclo's trick (Lemma 2.1 of Bardet et al. (2018)) yields that µ ∗ and p X satisfy the LSI with a constant

(2) We can easily check that V ( x ) = r ( x ) λ and H ( x ) = 1 λ δU (˜ µ ) δµ ( x ) , for appropriately chosen ˜ µ , satisfies the conditions in Lemma 6 with c 1 = λ 1 λ , c 2 = λ 2 λ , c V = c r , and ¯ L = R λ . Hence, µ ∗ and p X satisfy the LSI with a constant α such that

Lemma 6 (Log Sobolev inequality with Lipschitz perturbation) . Let ν ( x ) ∝ exp( -V ( x )) where ∇∇ ⊤ V ( x ) ⪰ c 1 I , x ⊤ ∇ V ( x ) ≥ c 1 ∥ x ∥ 2 and 0 ≤ V ( x ) ≤ c 2 ( c V + ∥ x ∥ 2 ) with c 1 , c 2 , c V > 0 , and let H : R d → R is a Lipschitz continuous function with the Lipschitz constant ¯ L . Suppose that µ ∈ P is given by µ ( x ) ∝ exp( -H ( x )) ν ( x ) . Then, µ satisfies the LSI with a constanat α such that

Proof. Since ν is a strongly log-concave distribution, the Bakry-Emery argument (Bakry and Émery, 1985b) yields that it satisfies the LSI condition with a constant α ′ = c 1 by the assumption ∇∇ ⊤ V ( x ) ⪰ c 1 I .

Next, we evaluate the second moment of ν because it is required in the following analysis. Since we know that ν is the stationary distribution of the SDE d X t = -∇ V ( X t )d+ √ 2d W t , its corresponding infinitesimal generator gives that

Then, by the Gronwall lemma, we have

Hence, by taking X 0 = 0 (a.s.), we see that lim sup t E [ ∥ X t ∥ 2 ] ≤ d/c 1 that yields E ν [ ∥ X 2 ∥ ] ≤ d/c 1 .

A distribution ˜ µ satisfies the Poincaré inequality with a constant ˜ α if it holds that

for all bounded f with bounded derivatives. By Example (3) in Section 7.1 of Cattiaux and Guillin (2014), µ satisfies the Poincaré inequality with a constant ˜ α such that

Here, let G ( x ) = H ( x ) -H (0) . We can see that µ ( x ) can be expressed as µ ( x ) = exp( -G ( x )) ν ( x ) Z G with a normalizing constant Z G > 0 . Note that, by the Lipschitz continuity assumption of H , we have

Hence, we have the following evaluation of the normalizing constant Z G :

Theorem 2.7 of Cattiaux and Guillin (2022) claims that µ satisfies the LSI with a constant α such that

for any β > 0 and θ > 0 . The right hand side can be bounded by

Hence, if we set β = 1 and θ = 1 , we have

E.4 Uniform Log-Sobolev Inequality

Lemma 7 (Uniform log-Sobolev inequality) . Under the same condition as Theorem 2, it holds that

with a constant C λ = 2 λLα (1 + 2 c L ¯ R 2 ) + 2 λ 2 L 2 ¯ R 2 .

̸

Proof. The derivation is analogous to Chen et al. (2022), but we present the full proof for the sake of the completeness. For X = ( X i ) N i =1 , let X -i := ( X j ) j = i . Then, it holds that

where the second inequality is due to the Lipschitz continuity of ∇ δU δµ in terms of the Wasserstein distance (Assumption 2). Here, the term corresponding to the Wasserstein distance can be upper bounded as

̸

̸

Denote by P X i | X -i the conditional law of X i conditioned by X -i and denote by P X -i the marginal law of X -i where the joint law of X is ˜ µ ( N ) t . Then, it holds that

by the LSI condition of the proximal Gibbs measure (Assumption 3). Let ν r be the distribution with the density ν r ∝ exp( -r ( x ) /λ ) . Note that the proximal Gibbs measure is the minimizer of the linearized objective:

Then, by the optimality of the proximal Gibbs measure, it holds that

The expectation of the first term of the right hand side can be further evaluated as

̸

where the first inequality is due to Lemma 8 and the second inequality is due to Lemma 1 (we also notice that E X ∼ µ ∗ [ ∥ X ∥ 2 ] ≤ ¯ R 2 ). The second term in the right hand side of Eq. (28) can be evaluated as

where we used Lemma 3.6 of Chen et al. (2022) in the first inequality. Combining all of them, we arrive at

Then, we have the assertion with

Lemma 8. With the same setting as Lemma 7, it holds that

Proof. Let µ X ,θ := θµ X +(1 -θ ) µ X -i . Then, we see that

̸

where we used Assumption 2 (this is the only place where the boundedness of the second order variation of U is used). Hence, we obtain the assertion.

̸

̸

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nOn 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.\n\nQUESTIONS:\n- 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)$ ?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\n\nSTRENGTHS:\nExtensive 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSee the weaknesses section\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

STRENGTHS:
1. Strong theory with explicit bounds
2. Excellent outline of results
3. Thorough comparison of convergence rates with other approaches
4. 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.

QUESTIONS:
-

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

comments: 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.

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.

QUESTIONS:
It is suggested that the authors should provide some simple numerical experiments to demonstrate the effectiveness of the proposed algorithms.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-16

comments: 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.

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.

QUESTIONS:
Assumption 4 is rather complicated. Can the motivation for it be explained more clearly in the text?

LIMITATIONS:
LSI conditions are obviously hard to satisfy for some problems, but this is well-known and well-discussed in the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-27