SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Assumption 2. There exists  $L > 0$  such that  $\left\| \nabla \frac{\delta U(\mu)}{\delta \mu}(x) - \nabla \frac{\delta U(\mu')}{\delta \mu}(x') \right\| \leq L(W_2(\mu, \mu') + \| x - x' \|)$  and  $\left| \frac{\delta^2 U(\mu)}{\delta \mu^2}(x, x') \right| \leq L(1 + c_L (\| x \|^2 + \| x' \|^2))$  for any  $\mu, \mu' \in \mathcal{P}_2$  and  $x, x' \in \mathbb{R}^d$ . Also, there exists  $R > 0$  such that  $\| \nabla \frac{\delta U(\mu)}{\delta \mu}(x) \| \leq R$  for any  $\mu \in \mathcal{P}$  and  $x \in \mathbb{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  $\mathcal{F}$  has a unique minimizer  $\mu^{*}$  in  $\mathcal{P}$  and  $\mu^{*}$  is absolutely continuous with respect to the Lebesgue measure. Moreover,  $\mu^{*}$  satisfies the self-consistent condition:  $\mu^{*}(X)\propto \exp \left(-\frac{1}{\lambda}\frac{\delta F(\mu^{*})}{\delta\mu} (X)\right)$ .
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
$$
p _ {\mathcal {X}} (X) \propto \exp \left(- \frac {1}{\lambda} \frac {\delta F (\mu_ {\mathcal {X}})}{\delta \mu} (X)\right),
$$
where  $\mathcal{X} = (X^i)_{i=1}^N \subset \mathbb{R}^d$  is a set of  $N$  particles and  $X \in \mathbb{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  $\mu$  be a probability measure on  $(\mathbb{R}^d,\mathcal{B}(\mathbb{R}^d))$ .  $\mu$  satisfies the LSI with constant  $\alpha >0$  if for any smooth function  $\phi :\mathbb{R}^d\to \mathbb{R}$  with  $\mathbb{E}_{\mu}[\phi^{2}] < \infty$ , we have  $\mathbb{E}_{\mu}[\phi^{2}\log (\phi^{2})] - \mathbb{E}_{\mu}[\phi^{2}]\log (\mathbb{E}_{\mu}[\phi^{2}])\leq \frac{2}{\alpha}\mathbb{E}_{\mu}[\| \nabla \phi \| _2^2 ]$ .
This is equivalent to the condition that the KL divergence from  $\mu$  is bounded by the Fisher divergence:  $\int \log (\mathrm{d}\nu /\mathrm{d}\mu)\mathrm{d}\nu \leq \frac{1}{2\alpha}\int \| \nabla \log (\mathrm{d}\nu /\mathrm{d}\mu)\| ^2\mathrm{d}\nu$ , for any  $\nu \in \mathcal{P}$  which is absolutely continuous with respect to  $\mu$ . Our analysis requires that the proximal Gibbs distribution satisfies the LSI as follows.
Assumption 3.  $\mu^{*}$  and  $p_{\mathcal{X}}$  satisfy the LSI with  $\alpha >0$  for any set of particles  $\mathcal{X} = (X^i)_{i = 1}^N\subset \mathbb{R}^d$
Verification of LSI. The LSI of proximal Gibbs measure can be established via standard perturbation criteria. For  $U(\mu)$  with bounded first-variation, we may apply the classical Bakry-Emery and Holley-Stroock arguments (Bakry and Emery, 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,  $\mu^{*}$  and  $p_{\mathcal{X}}$  satisfy the log-Sobolev inequality with
$$
\alpha \geq \frac {\lambda_ {1}}{2 \lambda} \exp \left(- \frac {4 R ^ {2}}{\lambda_ {1} \lambda} \sqrt {2 d / \pi}\right) \vee \left\{\frac {4 \lambda}{\lambda_ {1}} + e ^ {\frac {R ^ {2}}{2 \lambda_ {1} \lambda}} \left(\frac {R}{\lambda_ {1}} + \sqrt {\frac {2 \lambda}{\lambda_ {1}}}\right) ^ {2} \left[ 2 + d + \frac {d}{2} \log \left(\frac {\lambda_ {2}}{\lambda_ {1}}\right) + 4 \frac {R ^ {2}}{\lambda_ {1} \lambda} \right] \right\} ^ {- 1}.
$$
Furthermore, if  $\| \frac{\delta U(\mu)}{\delta\mu}\|_{\infty}\leq R$  is satisfied for any  $\mu \in \mathcal{P}_2$ , then  $\mu^{*}$  and  $p_{\mathcal{X}}$  satisfy the LSI with  $\alpha \geq \frac{\lambda_1}{\lambda}\exp \left(-\frac{4R}{\lambda}\right)$ .
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  $\mu^{(N)}\in \mathcal{P}^{(N)}$  be a distribution of  $N$  particles  $\mathcal{X} = (X^i)_{i = 1}^N$  where  $\mathcal{P}^{(N)}$  is the set of probability measures on  $(\mathbb{R}^{d\times N},\mathcal{B}(\mathbb{R}^{d\times N}))$ . Similarly, we introduce the following objective on  $\mathcal{P}^{(N)}$ :
$$
\mathcal {F} ^ {N} (\mu^ {(N)}) = N \mathbb {E} _ {\mathcal {X} \sim \mu^ {(N)}} [ F (\mu_ {\mathcal {X}}) ] + \lambda \mathrm {E n t} (\mu^ {(N)}).
$$
One can easily verify that if  $\mu^{(N)}$  is a product measure of  $\mu \in \mathcal{P}$ , then  $\mathcal{F}^N (\mu^{(N)})\geq N\mathcal{F}(\mu)$  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  $\mathcal{F}^N (\mu^{(N)})$  (which can be obtained by the particle-approximated MFLD) is close to  $N\mathcal{F}(\mu^{*})$ . Specifically, it has been shown in Chen et al. (2022) that
$$
0 \leq \inf  _ {\mu^ {(N)} \in \mathcal {P} ^ {(N)}} \frac {1}{N} \mathscr {F} ^ {N} \left(\mu^ {(N)}\right) - \mathscr {F} \left(\mu^ {*}\right) \leq \frac {C _ {\lambda}}{N}, \tag {7}
$$
for some constant  $C_{\lambda} > 0$  (see Section E for more details). Importantly, if we consider the Wasserstein gradient flow on  $\mathcal{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  $\mathcal{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_{k}^{i}$ :  $\sigma_{v,k}^{2} \coloneqq \mathbb{E}_{\omega_{k},\mathcal{X}_{0:k}}[\| v_{k}^{i} - \nabla \frac{\delta F(\mu_{k})}{\delta\mu}\|^{2}]^{1}$ . In addition, to obtain a refined evaluation, we also incorporate the following smoothness assumption.
Assumption 4.  $v_{k}^{i}(\mathcal{X}_{0:k},\omega_{k}^{i})$  and  $D_{m}^{i}F(\mathcal{X}_{k}) = \nabla \frac{\delta F(\mu_{\mathcal{X}_{k}})}{\delta\mu}(X_{k}^{i})$  are differentiable w.r.t.  $(X_{k}^{j})_{j=1}^{N}$  and, for either  $G(\mathcal{X}_{k}) = v_{k}^{i}(\mathcal{X}_{0:k},\omega_{k}^{i})$  or  $G(\mathcal{X}_{k}) = D_{m}^{i}F(\mathcal{X}_{k})$  as a function of  $\mathcal{X}_{k}$ , it is satisfied that  $\| G(\mathcal{X}_{k}') - G(\mathcal{X}_{k}) - \sum_{j=1}^{N} \nabla_{X_{k}^{j}}^{\top} G(\mathcal{X}_{k}) \cdot (X_{k}'j - X_{k}^{j}) \| \leq Q \frac{\sum_{j \neq i} \|X_{k}'j - X_{k}^{j}\|^{2} + N \|X_{k}'i - X_{k}^{i}\|^{2}}{2N}$  for some  $Q > 0$  and  $\mathcal{X}_{k} = (X_{k}^{j})_{j=1}^{N}$  and  $\mathcal{X}_{k}' = (X_{k}'j)_{j=1}^{N}$ . We also assume  $\mathbb{E}_{\omega_{k}}\left[\|\nabla_{X_{k}^{j}} v_{k}^{i}(\mathcal{X}_{0:k},\omega_{k}^{i})^{\top}-\nabla_{X_{k}^{j}} D_{m}^{i}F(\mathcal{X}_{k})^{\top}\|_{\mathrm{op}}^{2}\right] \leq \frac{1 + \delta_{i,j} N^{2}}{N^{2}} \tilde{\sigma}_{v,k}^{2}$ , and  $\|v_{k}^{i}((\mathcal{X}_{k},\mathcal{X}_{0:k-1}),\omega_{k}^{i}) - v_{k}^{i}((\mathcal{X}_{k}',\mathcal{X}_{0:k-1}),\omega_{k}^{i})\| \leq L(W_{2}(\mu_{\mathcal{X}_{k}},\mu_{\mathcal{X}_{k}'}) + \|X_{k}' - X_{k}^{\prime i}\|$ .
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.
<sup>1</sup>After the acceptance of this manuscript, we noticed that by replacing  $\sigma_{v,k}^2$  with the square of a conditional expectation  $\mathbb{E}_{\mathcal{X}_{0:k+1}}\left[\|\mathbb{E}_{\omega_k}\|_{\mathcal{X}_{0:k+1}}[v_k^i - \nabla \frac{\delta F(\mu_k)}{\delta \mu}]\|^2\right]$ , we can obtain a refined bound as in Das et al. (2023) suggested by anonymous reviewer Y8on. We defer such refined analysis to future work.
Taking these factors into consideration, we can evaluate the decrease in the objective value after one-step update. Let  $\mu_k^{(N)} \in \mathcal{P}^{(N)}$  be the distribution of  $\mathcal{X}_k$  conditioned on the sequence  $\omega_{0:k} = (\omega_{k'})_{k'=0}^k$ . Define
$$
\bar {R} ^ {2} := \mathbb {E} [ \| X _ {0} ^ {i} \| ^ {2} ] + \frac {1}{\lambda_ {1}} \left[ \left(\frac {\lambda_ {1}}{4 \lambda_ {2}} + \frac {1}{\lambda_ {1}}\right) (R ^ {2} + \lambda_ {2} c _ {r}) + \lambda d \right], \delta_ {\eta} := C _ {1} \bar {L} ^ {2} (\eta^ {2} + \lambda \eta),
$$
where  $C_1 = 8[R^2 + \lambda_2(c_r + \bar{R}^2) + d]$  and  $\bar{L} = L + \lambda_2$ .
Theorem 2. Under Assumptions 1, 2 and 3, if  $\eta_k \leq \lambda_1 / (4\lambda_2)$ , we have
$$
\frac {1}{N} \mathbb {E} \left[ \mathcal {F} ^ {N} (\mu_ {k + 1} ^ {(N)}) \right] - \mathcal {F} (\mu^ {*}) \leq \exp \left(- \frac {\lambda \alpha \eta_ {k}}{2}\right) \left(\frac {1}{N} \mathbb {E} \left[ \mathcal {F} ^ {N} (\mu_ {k} ^ {(N)}) \right] - \mathcal {F} (\mu^ {*})\right) + \eta_ {k} \left(\delta_ {\eta_ {k}} + \frac {C _ {\lambda}}{N}\right) + \Upsilon_ {k},
$$
where  $\Upsilon_{k} = 4\eta_{k}\delta_{\eta_{k}} + (R + \lambda_{2}\bar{R})(1 + \sqrt{\lambda / \eta_{k}})\left(\eta_{k}^{2}\tilde{\sigma}_{v,k}\sigma_{v,k} + Q\eta_{k}^{3}\sigma_{v,k}^{2}\right) + \eta_{k}^{2}(L + \lambda_{2})^{2}\sigma_{v,k}^{2}$  with Assumption 4, and  $\Upsilon_{k} = \sigma_{v,k}^{2}\eta_{k}$  without this additional assumption; the expectation is taken with respect to the randomness  $(\omega_{k^{\prime}})^{k}_{k^{\prime} = 1} = ((\omega_{k^{\prime}}^{i})_{i = 1}^{N})_{k^{\prime} = 1}^{k}$  of the stochastic gradient; and  $C_\lambda = 2\lambda L\alpha (1 + 2c_L\bar{R}^2) + 2\lambda^2 L^2\bar{R}^2$
The proof can be found in Appendix B. We remark that to derive the bound for  $\Upsilon_{k} = \sigma_{v,k}^{2}\eta_{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  $\mu_{k + 1}^{(N)}$  is correlated with the stochastic gradient  $v_{k}^{i}$ . Evaluating such correlations is non-trivial because the randomness is induced not only by  $\omega_{k}$  but also by the Gaussian noise. Thanks to this general result, we only need to evaluate the variance  $\sigma_{v,k}$  and  $\tilde{\sigma}_{v,k}$  for each method to obtain the specific convergence rate.
Conversion to a Wasserstein distance bound. As a consequence of the bound on  $\frac{1}{N}\mathcal{F}^N (\mu_k^{(N)}) - \mathcal{F}(\mu^*)$ , we can control the Wasserstein distance between  $\mu_k^{(N)}$  and  $\mu^{*N}$ , where  $\mu^{*N}\in \mathcal{P}^{(N)}$  is the  $(N$ -times) product measure of  $\mu^*$ . Let  $W_{2}(\mu ,\nu)$  be the 2-Wasserstein distance between  $\mu$  and  $\nu$ , then
$$
W _ {2} ^ {2} \left(\mu_ {k} ^ {(N)}, \mu^ {* N}\right) \leq \frac {2}{\lambda \alpha} \left(\mathscr {F} ^ {N} \left(\mu_ {k} ^ {(N)}\right) - N \mathscr {F} \left(\mu^ {*}\right)\right),
$$
under Assumptions 1 and 3 (see Lemma 3 in the Appendix). Hence, if  $\frac{1}{N}\mathcal{F}^N (\mu_k^{(N)}) - \mathcal{F}(\mu^*)$  is small, the particles  $(X_{k}^{i})_{i = 1}^{N}$  behaves like an i.i.d. sample from  $\mu^{*}$ . As an example, for the mean-field neural network setting, if  $|h_x(z) - h_{x'}(z)|\leq L\| x - x'\|$ $(\forall x,x^{\prime}\in \mathbb{R}^{d})$  and  $V_{\mu^{*}} = \int (f_{\mu^{*}}(z) - h_{x}(z))^{2}\mathrm{d}\mu^{*}(x) < \infty$  for a fixed  $z$ , then Lemma 4 in the appendix yields that
$$
\mathbb {E} _ {\mathcal {X} _ {k} \sim \mu_ {k} ^ {(N)}} [ (f _ {\mu_ {\mathcal {X} _ {k}}} (z) - f _ {\mu^ {*}} (z)) ^ {2} ] \leq \frac {2 L ^ {2}}{N} W _ {2} ^ {2} (\mu_ {k} ^ {(N)}, \mu^ {* N}) + \frac {2}{N} V _ {\mu^ {*}},
$$
which also gives  $\mathbb{E}_{\mathcal{X}_k\sim \mu_k^{(N)}}[(f_{\mu_{\mathcal{X}_k}}(z) - f_{\mu^*}(z))^2 ]\leq \frac{4L^2}{\lambda\alpha}\left(N^{-1}\mathcal{F}^N (\mu_k^{(N)}) - \mathcal{F}(\mu^*)\right) + \frac{2}{N} V_{\mu^*}$ . This allows us to monitor the convergence of the finite-width neural network to the optimal solution  $\mu^{*}$  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  $\sigma_{v,k}^2 = 0$ . We specialize the previous assumptions to the stochastic gradient setting as follows.
Assumption 5.
(i)  $\sup_{x\in \mathbb{R}^d}\| \nabla \frac{\delta\ell(\mu,z)}{\delta\mu} (x)\| \leq R$  for all  $\mu \in \mathcal{P}$  and  $z\in \mathcal{Z}$
(ii)  $\sup_x\| \nabla_x\nabla_x^\top \frac{\delta\ell(\mu,z)}{\delta\mu} (x)\|_{\mathrm{op}},\sup_{x,x'}\| \nabla_x\nabla_{x'}^\top \frac{\delta^2\ell(\mu,z)}{\delta^2\mu} (x,x')\|_{\mathrm{op}}\leq R.$
Note that point  $(i)$  is required to bound the variance  $\sigma_{v,k}^2$  (i.e.,  $\sigma_{v,k}^2 \leq R^2 / B$ ), whereas point  $(ii)$  corresponds to the additional Assumption 4 required for the improved convergence rate. Let  $\Delta_0 := \frac{1}{N}\mathbb{E}[\mathcal{F}^N (\mu_0^{(N)})] - \mathcal{F}(\mu^*)$ . We have the following evaluation of the objective. Note that we can recover the evaluation for F-MFLD by formally setting  $B = \infty$  so that  $\sigma_{v,k}^2 = 0$  and  $\tilde{\sigma}_{v,k}^2 = 0$ .
2The constants are not optimized (e.g.,  $O\left(\frac{1}{\lambda_1}\right)$  in  $\bar{R}$  can be improved) since we prioritize for clean presentation.
Theorem 3. Suppose that  $\eta_k = \eta$  ( $\forall k \in \mathbb{N}_0$ ) and  $\lambda \alpha \eta \leq 1/4$  and  $\eta \leq \lambda_1/(4\lambda_2)$ . Under Assumptions 1, 2, 3 and 5, it holds that
$$
\frac {1}{N} \mathbb {E} [ \mathcal {F} ^ {N} (\mu_ {k} ^ {(N)}) ] - \mathcal {F} (\mu^ {*}) \leq \exp (- \lambda \alpha \eta k / 2) \Delta_ {0} + \frac {4}{\lambda \alpha} \bar {L} ^ {2} C _ {1} (\lambda \eta + \eta^ {2}) + \frac {4}{\lambda \alpha \eta} \tilde {\Upsilon} + \frac {4 C _ {\lambda}}{\lambda \alpha N}, (8)
$$
where  $\bar{\Upsilon} = 4\eta \delta_{\eta} + \left[R + \lambda_2\bar{R} +(L + \lambda_2)^2\right](1 + \sqrt{\frac{\lambda}{\eta}})\eta^2\frac{R^2}{B} +\left(R + \lambda_2\bar{R}\right)R(1 + \sqrt{\frac{\lambda}{\eta}})\eta^3\frac{R^2}{B}$  under Assumption 5-(ii), and  $\bar{\Upsilon} = \frac{R^2}{B}\eta$  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 (-\lambda \alpha \eta k)\Delta_0 + \frac{\eta}{\alpha\lambda})$  is consistent to ours, since Eq. (26) in the Appendix implies that the objective  $\frac{1}{\lambda}\left(\frac{1}{N}\mathbb{E}[\mathcal{F}^N (\mu_k^{(N)})] - \mathcal{F}(\mu^*)\right)$  upper bounds the KL divergence between  $\mu_{k}^{(N)}$  and  $\mu^{*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  $\epsilon > 0$ , if we take
$$
\eta \leq \frac {\alpha \epsilon}{4 0 \bar {L} ^ {2} C _ {1}} \wedge \frac {1}{\bar {L}} \sqrt {\frac {\lambda \alpha \epsilon}{4 0 C _ {1}}} \wedge 1, k \geq \frac {2}{\lambda \alpha \eta} \log (2 \Delta_ {0} / \epsilon),
$$
with  $B \geq 4\left[(1 + R)(R + \lambda_2\bar{R}) + (L + \lambda_2)^2\right]R^2 (\eta +\sqrt{\eta\lambda}) / (\lambda \alpha \epsilon)$  under Assumption 5-(ii) and  $B \geq 4R^{2} / (\lambda \alpha \epsilon)$  without such assumption, then the right hand side of (8) can be bounded as  $\frac{1}{N}\mathbb{E}[\mathcal{F}^N (\mu_k^{(N)})] - \mathcal{F}(\mu^*) = \epsilon +\frac{4C_\lambda}{\lambda\alpha N}$ . Hence we achieve  $\epsilon +O(1 / N)$  error with iteration complexity:
$$
k = O \left(\frac {\bar {L} ^ {2}}{\alpha \epsilon} + \frac {\bar {L}}{\sqrt {\lambda \alpha \epsilon}}\right) \frac {1}{\lambda \alpha} \log \left(\epsilon^ {- 1}\right). \tag {9}
$$
If we neglect  $O(1 / \sqrt{\lambda\alpha\epsilon})$  as a second order term, then the above can be simplified as  $O\left(\frac{\bar{L}^2}{\lambda\alpha^2}\frac{\log(\epsilon^{-1})}{\epsilon}\right)$ .
Noticeably, the mini-batch size  $B$  can be significantly reduced under the additional smoothness condition Assumption 5-(ii) (indeed, we have  $O(\eta + \sqrt{\eta\lambda})$  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 = \Omega((\eta + \sqrt{\eta\lambda}) / (\lambda\alpha\epsilon)) \geq \Omega(\lambda^{-1} + \sqrt{(\epsilon\lambda\alpha)^{-1}})$  which yields the objective value of order  $O(\epsilon)$ , SGD-MFLD achieves  $O(n/B) = O(n(\lambda \wedge \sqrt{\epsilon\lambda\alpha}))$  times smaller total complexity. For example, when  $\lambda = \alpha\epsilon = 1/\sqrt{n}$ , a  $O(\sqrt{n})$ -factor reduction of total complexity can be achieved by SGD-MFLD, 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  $\eta_{k} = \eta$  ( $\forall k$ ). Instead of Assumption 5, we introduce the following two assumptions.
Assumption 6.
(i) For any  $i\in [n]$ , it holds that  $\left\| \nabla \frac{\delta\ell_i(\mu)}{\delta\mu} (x) - \nabla \frac{\delta\ell_i(\mu')}{\delta\mu} (x')\right\| \leq L(W_2(\mu ,\mu ') + \| x - x'\|)$  for any  $\mu ,\mu^{\prime}\in \mathcal{P}_2$  and  $x,x^{\prime}\in \mathbb{R}^{d}$ , and  $\sup_{x\in \mathbb{R}^d}\| \nabla \frac{\delta\ell_i(\mu)}{\delta\mu} (x)\| \leq R$  for any  $\mu \in \mathcal{P}$ .
(ii) Additionally, we have  $\sup_x\| \nabla_x\nabla_x^\top \frac{\delta\ell_i(\mu)}{\delta\mu} (x)\|_{\mathrm{op}},\sup_{x,x'}\| \nabla_x\nabla_{x'}^\top \frac{\delta^2\ell_i(\mu)}{\delta^2\mu} (x,x')\|_{\mathrm{op}}\leq R.$
Here again, point (i) is required to bound  $\sigma_{v,k}^2$  and point (ii) yields Assumption 4. We have the following computational complexity bound for SVRG-MFLD.
Theorem 4. Suppose that  $\lambda \alpha \eta \leq 1/4$  and  $\eta \leq \lambda_1/(4\lambda_2)$ . Let  $\Xi = \frac{n-B}{B(n-1)}$ . Then, under Assumptions 1, 2, 3 and 6, we have the same error bound as Eq. (8) with different  $\Upsilon$ :
$$
\bar {\Upsilon} = 4 \eta \delta_ {\eta} + \left(1 + \sqrt {\frac {\lambda}{\eta}}\right) \left\{\left(R + \lambda_ {2} \bar {R}\right) \eta^ {2} \sqrt {C _ {1} \Xi L ^ {2} m (\eta^ {2} + \eta \lambda) R ^ {2} \Xi} + \left[ \left(R + \lambda_ {2} \bar {R}\right) R \eta^ {3} + (L + \lambda_ {2}) ^ {2} \eta^ {2} \right] \right\},
$$
under Assumption 6-(ii), and  $\widetilde{\Upsilon} = C_1\Xi L^2 m\eta^2 (\eta +\lambda)$  without Assumption 6-(ii).
The proof is given in Appendix D. Therefore, to achieve  $\frac{1}{N}\mathbb{E}[\mathcal{F}^N (\mu_k^{(N)})] - \mathcal{F}(\mu^*)\leq O(\epsilon) + \frac{4C_\lambda}{\lambda\alpha N}$  for a given  $\epsilon >0$  , it suffices to set
$$
\eta = \frac {\alpha \epsilon}{4 0 \bar {L} ^ {2} C _ {1}} \wedge \frac {\sqrt {\lambda \alpha \epsilon}}{4 0 \bar {L} \sqrt {C _ {1}}}, k = \frac {2 \log (2 \Delta_ {0} / \epsilon)}{\lambda \alpha \eta} = O \left(\frac {\bar {L} ^ {2}}{\alpha \epsilon} + \frac {\bar {L}}{\sqrt {\lambda \alpha \epsilon}}\right) \frac {\log (\epsilon^ {- 1})}{(\lambda \alpha)},
$$
with  $B \geq \left[\sqrt{m} \frac{(\eta + \sqrt{\eta / \lambda})^2}{\eta \lambda} \vee m \frac{(\eta + \sqrt{\eta / \lambda})^3}{\eta \lambda}\right] \wedge n$ . In this setting, the total gradient complexity can be bounded as
$$
B k + \frac {n k}{m} + n \lesssim \max \left\{n ^ {\frac {1}{3}} \left(1 + \sqrt {\frac {\eta}{\lambda}}\right) ^ {\frac {4}{3}}, \sqrt {n} (\eta \lambda) ^ {\frac {1}{4}} \left(1 + \sqrt {\frac {\eta}{\lambda}}\right) ^ {\frac {3}{2}} \right\} \frac {1}{\eta} \frac {\log (\epsilon^ {- 1})}{\lambda \alpha} + n,
$$
where  $m = \Omega (n / B) = \Omega ([n^{2 / 3}(1 + \sqrt{\eta / \lambda})^{-4 / 3}\wedge \sqrt{n} (1 + \sqrt{\eta / \lambda})^{-3 / 2}(\sqrt{\eta\lambda})^{-1 / 2})]\vee 1)$  and  $B = O\left(\left[n^{\frac{1}{3}}\left(1 + \sqrt{\frac{\eta}{\lambda}}\right)^{\frac{4}{3}}\vee \sqrt{n} (\eta \lambda)^{\frac{1}{4}}\left(1 + \sqrt{\frac{\eta}{\lambda}}\right)^{\frac{3}{2}}\right]\wedge n\right)$ . When  $\epsilon$  is small, the first term in the right hand side becomes the main term which is  $\max \{n^{1 / 3},\sqrt{n} (\alpha \epsilon \lambda)^{1 / 4}\} \frac{\log(\epsilon^{-1})}{\lambda\alpha^2\epsilon}$ . Therefore, comparing with the full batch gradient method (F-MFLD), SVRG-MFLD achieves at least  $\min \left\{n^{\frac{2}{3}},\sqrt{n} (\alpha \epsilon \lambda)^{-\frac{1}{4}}\right\}$  times better total computational complexity when  $\eta \leq \lambda$ . Note that even without the additional Assumption 6-(ii), we still obtain a  $\sqrt{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\left(\left(n + \frac{\sqrt{n}}{\epsilon}\right)\frac{\log(\epsilon^{-1})}{(\lambda\alpha)^{2}}\right)$  iteration complexity to achieve  $\epsilon$  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(\lambda)$ -factor discrepancy is due to the difference in the objective - their bound considers the KL divergence while our objective corresponds to  $\lambda$  times the KL divergence). Furthermore, our analysis gives an even faster convergence rate under an additional mild assumption (Assumption 6-(ii)).