SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Convergence of Mean-field Langevin dynamics: Time-space discretization, stochastic gradient, and variance reduction
Taiji Suzuki $^{1,2}$ , Denny Wu $^{3,4}$ , Atsushi Nitanda $^{2,5}$
<sup>1</sup>University of Tokyo, <sup>2</sup>RIKEN AIP, <sup>3</sup>New York University,
4 Flatiron Institute, Kyushu Institute of Technology
taiji@mist.i.u-tokyo.ac.jp, dennywu@nyu.edu, nitanda@ai.kyutech.ac.jp
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.
In this work we consider the mean-field Langevin dynamics (MFLD) given by the following McKean-Vlasov stochastic differential equation:
$$
\mathrm {d} X _ {t} = - \nabla \frac {\delta F (\mu_ {t})}{\delta \mu} \left(X _ {t}\right) \mathrm {d} t + \sqrt {2 \lambda} \mathrm {d} W _ {t}, \tag {1}
$$
where  $\mu_t = \operatorname{Law}(X_t)$ ,  $F: \mathcal{P}_2(\mathbb{R}^d) \to \mathbb{R}$  is a convex functional,  $W_t$  is the  $d$ -dimensional standard Brownian motion, and  $\frac{\delta F}{\delta \mu}$  denotes the first-variation of  $F$ . Importantly, MFLD is the Wasserstein gradient flow that minimizes an entropy-regularized convex functional as follows:
$$
\min  _ {\mu \in \mathcal {P} _ {2}} \left\{F (\mu) + \lambda \operatorname {E n t} (\mu) \right\}. \tag {2}
$$
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  $\mu$  as a mixture of  $N$  particles  $(X^i)_{i=1}^N$  — 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
$$
X _ {k + 1} ^ {i} = X _ {k} ^ {i} - \eta_ {k} \tilde {\nabla} \frac {\delta F (\mu_ {k})}{\delta \mu} \left(X _ {k} ^ {i}\right) + \sqrt {2 \lambda \eta_ {k}} \xi_ {k} ^ {i}, \tag {3}
$$
where  $\eta_{k}$  is the step size at the  $k$ -th iteration,  $\xi_{k}^{i}$  is an i.i.d. standard Gaussian vector, and  $\tilde{\nabla}\frac{\delta F(\mu)}{\delta\mu}$  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.  $\mu$  is entirely represented by a finite set of particles:  $\mu_{k} = \frac{1}{N}\sum_{i = 1}^{N}\delta_{X_{k}^{i}}$ .
(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  $\mathcal{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?
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
Method (authors)	# of particles	Total complexity	Single loop	Mean-field
PDA* (Nitanda et al., 2021)	ε-2 log(n)	Gεε-1	×	✓
P-SDCA (Oko et al., 2022)	ε-1 log(n)	Gε(n+1/λ) log(n/ε)	×	✓
GLD (Vempala and Wibisono, 2019)	—	n/ε log(ε-1)/(λα)2	✓	×
SVRG-LD (Kinoshita and Suzuki, 2022)	—	(n+√n/ε) log(ε-1)/(λα)2	✓	×
F-MFLD (ours)	ε-1	nE* log(ε-1)/(λα)	✓	✓
SGD-MFLD* (ours)	ε-1	ε-1E* log(ε-1)/(λα)	✓	✓
SGD-MFLD* (ii) (ours)	ε-1	ε-1(1+√λE*) log(ε-1)/(λα)2	✓	✓
SVRG-MFLD (ours)	ε-1	√nE* log(ε-1)/(λα)+n	✓	✓
SVRG-MFLD (ii) (ours)	ε-1	(n1/3E*+√nλ1/4E*3/4) log(ε-1)/(λα)+n	✓	✓
Table 1: Comparison of computational complexity to optimize an entropy-regularized finite-sum objective up to excess objective value  $\epsilon$ , in terms of dataset size  $n$ , entropy regularization  $\lambda$ , and LSI constant  $\alpha$ . Label * indicates the online setting, and the unlabeled methods are tailored to the finite-sum setting. "Mean-field" indicates the presence of particle interactions. "Single loop" indicates whether the algorithm requires an inner-loop MCMC sampling sub-routine at every step. "(ii)" indicates convergence rate under additional smoothness condition (Assumption 4), where  $E_{*} = \frac{\bar{L}^{2}}{\alpha\epsilon} + \frac{\bar{L}}{\sqrt{\lambda\alpha\epsilon}}$ . For double-loop algorithms (PDA and P-SDCA),  $G^{*}$  is the number of gradient evaluations required for MCMC sampling; for example, for MALA (Metropolis-adjusted Langevin algorithm)  $G_{\epsilon} = O(n\alpha^{-5/2}\log(1/\epsilon)^{3/2})$ , and for LMC (Langevin Monte Carlo)  $G_{\epsilon} = O(n(\alpha\epsilon)^{-2}\log(\epsilon))$ .
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.
Consider the set of probability measure  $\mathcal{P}$  on  $\mathbb{R}^d$  where the Borel  $\sigma$ -algebra is equipped. Our goal is to find a probability measure  $\mu \in \mathcal{P}$  that approximately minimizes the objective given by
$$
F (\mu) = U (\mu) + \mathbb {E} _ {\mu} [ r ],
$$
where  $U:\mathcal{P}\to \mathbb{R}$  is a (convex) loss function, and  $r:\mathbb{R}^d\rightarrow \mathbb{R}$  is a regularization term. Let  $\mathcal{P}_2$  be the set of probability measures with the finite second moment. In the following, we consider the setting where  $F(\mu)\leq C(1 + \mathbb{E}_{\mu}[\| X\|^{2}])$ , and focus on  $\mathcal{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(\cdot)$  is a neuron with a parameter  $x \in \mathbb{R}^d$ , e.g.,  $h_x(z) = \sigma (w^\top z + b)$ , where  $w \in \mathbb{R}^{d - 1}$ ,  $b \in \mathbb{R}$ , and  $x = (w,v)$ . The mean-field neural network corresponding to a probability measure  $\mu \in \mathcal{P}$  can be written as  $f_{\mu}(\cdot) = \int h_{x}(\cdot)\mu (\mathrm{d}x)$ . Given training data  $(z_i,y_i)_{i = 1}^n \in \mathbb{R}^{d - 1} \times \mathbb{R}$ , we may define the empirical risk of  $f_{\mu}$  as  $U(\mu) = \frac{1}{n}\sum_{i = 1}^{n}\ell (f_{\mu}(z_i),y_i)$  for a loss function  $\ell :\mathbb{R}\times \mathbb{R}\to \mathbb{R}$ . Then, the objective  $F$  becomes
$$
F (\mu) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(f _ {\mu} \left(z _ {i}\right), y _ {i}\right) + \lambda_ {1} \int \| x \| ^ {2} \mathrm {d} \mu (x),
$$
where the regularization term is  $r(x) = \lambda_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: \mathcal{P}_2 \to \mathbb{R}$ . The first-variation  $\frac{\delta G}{\delta \mu}$  of a functional  $G: \mathcal{P}_2 \to \mathbb{R}$  at  $\mu \in \mathcal{P}_2$  is defined as a continuous functional  $\mathcal{P}_2 \times \mathbb{R}^d \to \mathbb{R}$  that satisfies  $\lim_{\epsilon \to 0} \frac{G(\epsilon \nu + (1 - \epsilon) \mu)}{\epsilon} = \int \frac{\delta G}{\delta \mu}(\mu)(x) \, \mathrm{d}(\nu - \mu)$  for any  $\nu \in \mathcal{P}_2$ . If there exists such a functional  $\frac{\delta G}{\delta \mu}$ , we say  $G$  admits a first-variation at  $\mu$ , or simply  $G$  is differentiable at  $\mu$ .
To avoid the ambiguity of  $\frac{\delta G}{\delta \mu}$  up to constant shift, we follow the convention of imposing  $\int \frac{\delta G}{\delta \mu}(\mu) \mathrm{d}\mu = 0$ . Using the first-variation of  $F$ , the MFLD is given by the following stochastic differential equation:
$$
\mathrm {d} X _ {t} = - \nabla \frac {\delta F (\mu_ {t})}{\delta \mu} (X _ {t}) \mathrm {d} t + \sqrt {2 \lambda} \mathrm {d} W _ {t}, \quad \mu_ {t} = \operatorname {L a w} (X _ {t}), \tag {4}
$$
where  $X_0 \sim \mu_0$ ,  $\operatorname{Law}(X)$  denotes the distribution of the random variable  $X$  and  $(W_t)_{t \geq 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  $\frac{\delta F(\mu_t)}{\delta \mu}(\cdot)$  depends on the distribution  $\mu_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):
$$
\mathcal {F} (\mu) = F (\mu) + \lambda \operatorname {E n t} (\mu), \tag {5}
$$
where  $\mathrm{Ent}(\mu) = -\int \log (\mathrm{d}\mu (z) / \mathrm{d}z)\mathrm{d}\mu (z)$  is the negative entropy of  $\mu$
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(\mu) = \int V(x)\mathrm{d}\mu (x)$ . In this case,  $\frac{\delta F}{\delta\mu} = V$  for any  $\mu$  and the MFLD Eq. (4) simplifies to
$$
\mathrm {d} X _ {t} = - \nabla V (X _ {t}) \mathrm {d} t + \sqrt {2 \lambda} \mathrm {d} W _ {t}.
$$
This is exactly the gradient Langevin dynamics for optimizing  $V$  or sampling from  $\mu \propto \exp (-V / \lambda)$ .
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 \in \mathbb{R}^d$ , e.g.,  $h_x(z) = \tanh(r\sigma(w^\top x))$ ,  $h_x(z) = \tanh(r)\sigma(w^\top x)$  for  $x = (r, w)$ , or simply  $h_x(z) = \sigma(x^\top z)$ . Then the learning of mean-field neural network  $f_\mu(\cdot) = \int h_x(\cdot)\mu(\mathrm{d}x)$  via minimizing the empirical risk  $U(\mu) = \frac{1}{n}\sum_{i=1}^{n}\ell(f_\mu(z_i), y_i)$  with a convex loss (e.g., the logistic loss  $\ell(f, y) = \log(1 + \exp(-yf))$ , or the squared loss  $\ell(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  $\mu$  and  $\nu$  is defined as  $\mathrm{MMD}^2(\mu, \nu) \coloneqq \int \int (k(x, x) - 2k(x, y) + k(y, y)) \, \mathrm{d}\mu(x) \, \mathrm{d}\nu(y)$ . We perform nonparametric density estimation by fitting a Gaussian mixture model  $f_{\mu}(z) = \int g_x(z) \, \mathrm{d}\mu(x)$ , where  $g_x$  is the Gaussian density with mean  $x$  and a given variance  $\sigma^2 > 0$ . The mixture model is learned by minimizing  $\mathrm{MMD}^2(f_{\mu}, p^*)$  where  $p^*$  is the target distribution. If we observe a set of training data  $(z_i)_{i=1}^n$  from  $p^*$ , then the empirical version of MMD is one suitable loss function  $U(\mu)$  given as
$$
\widehat {\mathrm {M M D}} ^ {2} (\mu) := \int \left[ \iint g _ {x} (z) g _ {x ^ {\prime}} (z ^ {\prime}) k (z, z ^ {\prime}) \mathrm {d} z \mathrm {d} z ^ {\prime} \right] \mathrm {d} (\mu \times \mu) (x, x ^ {\prime}) - 2 \int \left(\frac {1}{n} \sum_ {i = 1} ^ {n} \int g _ {x} (z) k (z, z _ {i}) \mathrm {d} z\right) \mathrm {d} \mu (x),
$$
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) = \nabla \log(\mu^{*})(z)$ , then for a positive definite kernel  $k$ , we define the Stein kernel as
$$
W _ {\mu^ {*}} (z, z ^ {\prime}) := s (z) ^ {\top} s (z ^ {\prime}) k (z, z ^ {\prime}) + s (z) ^ {\top} \nabla_ {z ^ {\prime}} k (z, z ^ {\prime}) + \nabla_ {z} ^ {\top} k (z, z ^ {\prime}) s (z ^ {\prime}) + \nabla_ {z ^ {\prime}} ^ {\top} \nabla_ {z} k (z, z ^ {\prime}).
$$
We take  $U(\mu)$  as the KSD between  $\mu$  and  $\mu^{*}$  defined as  $\mathrm{KSD}(\mu) = \int \int W_{\mu^{*}}(z,z^{\prime})\mathrm{d}\mu (z)\mathrm{d}\mu (z^{\prime})$ . 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  $\mathcal{X} = (X^i)_{i=1}^N \subset \mathbb{R}^d$ , we define its corresponding empirical distribution as  $\mu_{\mathcal{X}} = \frac{1}{N} \sum_{i=1}^{N} \delta_{X_i}$ . Let  $\mathcal{X}_k = (X_k^i)_{i=1}^N \subset \mathbb{R}^d$  be  $N$  particles at the  $k$ -th update, and define  $\mu_k = \mu_{\mathcal{X}_k}$  as a finite particle approximation of the population counterpart. Starting from  $X_0^i \sim \mu_0$ , we update  $\mathcal{X}_k$  as,
$$
X _ {k + 1} ^ {i} = X _ {k} ^ {i} - \eta_ {k} v _ {k} ^ {i} + \sqrt {2 \lambda \eta_ {k}} \xi_ {k} ^ {i}, \tag {6}
$$
where  $\eta_{k} > 0$  is the step size,  $\xi_k^i$  is an i.i.d. standard normal random variable  $\xi_k^i\sim N(0,I)$  and  $v_{k}^{i} = v_{k}^{i}(\mathcal{X}_{0:k},\omega_{k}^{i})$  is a stochastic approximation of  $\nabla \frac{\delta F(\mu_k)}{\delta\mu} (X_k^i)$  where  $\omega_{k} = (\omega_{k}^{i})_{i = 1}^{N}$  is a random variable generating the randomness of stochastic gradient.  $v_{k}^{i}$  can depend on the history  $\mathcal{X}_{0:k} = (\mathcal{X}_0,\dots ,\mathcal{X}_k)$ , and  $\mathbb{E}_{\omega_k^i}[v_k^i |\mathcal{X}_{0:k}] = \nabla \frac{\delta F(\mu_k)}{\delta\mu}(X_k^i)$ . We analyze three versions of  $v_{k}^{i}$ .
(1) Full gradient: F-MFLD. If we have access to the exact gradient, we may compute
$$
v _ {k} ^ {i} = \nabla \frac {\delta F (\mu_ {k})}{\delta \mu} (X _ {k} ^ {i}).
$$
(2) Stochastic gradient: SGD-MFLD. Suppose the loss function  $U$  is given by an expectation as  $U(\mu) = \mathbb{E}_{Z\sim P_Z}[\ell (\mu ,Z)]$ , where  $Z\in \mathcal{Z}$  is a random observation obeying a distribution  $P_Z$  and  $\ell :\mathcal{P}\times \mathcal{Z}\to \mathbb{R}$  is a loss function. In this setting, we construct the stochastic gradient as
$$
v _ {k} ^ {i} = \frac {1}{B} \sum_ {j = 1} ^ {B} \nabla \frac {\delta \ell (\mu_ {k} , z _ {k} ^ {(j)})}{\delta \mu} (X _ {k} ^ {i}) + \nabla r (X _ {k} ^ {i}),
$$
where  $(z_k^{(j)})_{j = 1}^B$  is a mini-batch of size  $B$  generated from  $P_Z$  in an i.i.d. manner.
(3) Stochastic variance reduced gradient: SVRG-MFLD. Suppose that the loss function  $U$  is given by a finite sum of loss functions  $\ell_i\colon U(\mu) = \frac{1}{n}\sum_{i = 1}^{n}\ell_i(\mu)$ , 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 \equiv 0 \mod m$ , where  $m$  is the update frequency, we set  $\dot{\mathcal{X}} = (\dot{X}^i)_{i=1}^n = (X_k^i)_{i=1}^n$  as the anchor point and refresh the stochastic gradient as  $v_k^i = \nabla \frac{\delta F(\mu_k)}{\delta \mu}(X_k^i)$ .
(ii) When  $k \neq 0 \mod m$ , we use the anchor point to construct a control variate and compute
$$
v _ {k} ^ {i} = \frac {1}{B} \sum_ {j \in I _ {k}} \nabla \left(\frac {\delta \ell_ {j} (\mu_ {\mathcal {X} _ {k}})}{\delta \mu} (X _ {k} ^ {i}) + r (X _ {k} ^ {i}) - \frac {\delta \ell_ {j} (\mu_ {\dot {\mathcal {X}}})}{\delta \mu} (\dot {X} ^ {i}) + \frac {\delta U (\mu_ {\dot {\mathcal {X}}})}{\delta \mu} (\dot {X} ^ {i})\right),
$$
where  $I_{k}$  is a uniformly drawn random subset of  $\{1, \dots, 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,
1.  $U:\mathcal{P}\to \mathbb{R}$  is a convex functional on  $\mathcal{P}$ , that is,  $U(\theta \mu +(1 - \theta)\nu)\leq \theta U(\mu) + (1 - \theta)U(\nu)$  for any  $\theta \in [0,1]$  and  $\mu ,\nu \in \mathcal{P}$ . Moreover,  $U$  admits a first-variation at any  $\mu \in \mathcal{P}_2$ .
2.  $r(\cdot)$  is twice differentiable and convex, and there exist constants  $\lambda_1, \lambda_2 > 0$  and  $c_r > 0$  such that  $\lambda_1 I \preceq \nabla \nabla^\top r(x) \preceq \lambda_2 I$ ,  $x^\top \nabla r(x) \geq \lambda_1 \|x\|^2$ , and  $0 \leq r(x) \leq \lambda_2 (c_r + \|x\|^2)$  for any  $x \in \mathbb{R}^d$ , and  $\nabla r(0) = 0$ .