Sampling from the Mean-Field Stationary Distribution

The minimization of energy functionals $\mathcal{E}$ over the Wasserstein space $\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})$ of probability measures has attracted substantial research activity in recent years, encompassing numerous application domains, including distributionally robust optimization [Kuh+19, YKW22], sampling [JKO98, Wib18, Che24], and variational inference [LW16, Lam+22, Dia+23, JCP23, Lac23a, YY23].

A canonical example of such a functional is $\mathcal{E}(\mu)=\int V\,\mathrm{d}\mu+\int\log\mu\,\mathrm{d}\mu$, where $V:\mathbb{R}^{d}\to\mathbb{R}$ is called the potential. Up to an additive constant, which is irrelevant for the optimization, this energy functional equals the KL divergence $\operatorname{\mathsf{KL}}(\mu\mathbin{\|}\pi)$ with respect to the density $\pi\propto\exp(-V)$, and the celebrated result of [JKO98] identifies the Wasserstein gradient flow of $\mathcal{E}$ with the Langevin diffusion. This link has inspired a well-developed theory for log-concave sampling, with applications to Bayesian inference and randomized algorithms; see [Che24] for an exposition.

The energy functional above contains two terms, corresponding to two of the fundamental examples of functionals considered in Villani’s well-known treatise on optimal transport [Vil03]. Namely, they are the “potential energy” and the entropy, the latter being a special case of the “internal energy.” However, Villani identifies a third fundamental functional—the “interaction energy”—with the pairwise form given by

More generally, in this work we consider minimizing the generic entropy-regularized energy

where $\mathcal{F}:\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})\to\mathbb{R}$ is a known functional. The minimization of the energy ($\mathsf{gE}$) has recently been of interest due to its role in analysing neural network training dynamics in the mean-field regime, including with [SNW22] and without [CB18, MMN18] entropic regularization, as well as with Fisher regularization [Cla+23].

For the sake of exposition, let us first focus on minimizing the pairwise energy ($\mathsf{pE}$). A priori, this question is more difficult than log-concave sampling; for instance, $\pi$ does not admit a closed form but rather is the solution to a non-linear equation

However, here too there is a well-developed mathematical theory which suggests a principled algorithmic approach. Just as the Wasserstein gradient flow of ($\mathsf{pE}$) in the case when $W=0$ can be identified with the Langevin diffusion, the Wasserstein gradient flow of ($\mathsf{pE}$) in the case when $W\neq 0$ corresponds to a (pairwise) McKean–Vlasov SDE, i.e. an SDE whose coefficients depend on the marginal law of the process, given below as

where $\pi_{t}=\operatorname{\mathsf{law}}(X_{t})$, $W$ is even, and $\{B_{t}\}_{t\geq 0}$ is a standard Brownian motion on $\mathbb{R}^{d}$. Since the McKean–Vlasov SDE is the so-called mean-field limit of interacting particle systems, we can approximately sample from the minimizer $\pi$ by numerically discretizing a system of SDEs, which describe the evolution of $N$ particles$\{X_{t}^{1:N}\}_{t\geq 0}\coloneqq\{(X^{1}_{t},\dotsc,X^{N}_{t})\}_{t\geq 0}$ as:

where $\{B^{i}:i\in[N]\}$ is a collection of independent Brownian motions. Moreover, the error from approximating the mean-field limit via this finite particle system has been studied in the literature on propagation of chaos [Szn91]. Similarly, the Wasserstein gradient flow for ($\mathsf{gE}$) corresponds to the mean-field Langevin dynamics and admits an analogous particle approximation.

The bounds for propagation of chaos have been refined over time, with [LL23] recently establishing a tight error dependence $\mathcal{O}(\nicefrac{{1}}{{N}})$ on the total number of particles $N$. These bounds, however, do not translate immediately into algorithmic guarantees. Existing sampling analyses study the propagation of chaos and discretization as a single entangled problem, which thus far have only been able to use weaker ${\mathcal{O}}(\sqrt{\nicefrac{{1}}{{N}}})$ rates for the former. Furthermore, there has been recent interest in using more sophisticated particle-based algorithms, e.g., “non-linear” Hamiltonian Monte Carlo [BS23] and the mean-field underdamped Langevin dynamics [FW23] to reduce the discretization error. Currently, this requires repeatedly carrying out the propagation of chaos and time discretization analyses from the ground up for each instance.

This motivates us to pose the following questions: (1) Can we incorporate improvements in the propagation of chaos literature, such as the $\mathcal{O}(\nicefrac{{1}}{{N}})$ error dependence shown in [LL23], to improve existing theoretical guarantees? (2) Can we leverage recent advances in the theory of log-concave sampling to design better algorithms?

Our main proposal in this work is to decouple the error into two terms, representing the propagation of chaos and discretization errors respectively. This simple and modular approach immediately allows us to answer both questions in the affirmative. Namely, we show how to combine established propagation of chaos bounds in various settings [[, including the sharp rate of]]lacker2023sharp with a large class of sophisticated off-the-shelf log-concave samplers, such as interacting versions of the randomized midpoint discretization of the underdamped Langevin dynamics [SL19, HBE20], Metropolis-adjusted algorithms [Che+21, WSC22, AC23], and the proximal sampler [LST21, Che+22, FYC23]. Our framework yields improvements upon prior state-of-the-art, such as [BS23, FW23], and provides a clear path for future ones.

Let $\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})$ be the set of probability measures on $\mathbb{R}^{d}$ that admit a density with respect to the Lebesgue measure and have finite second moment. We will also abuse notation and use the same symbol for a measure and its density when there is no confusion. We use superscripts for the particle index, and subscripts for the time variable. We will use $\mathcal{O},\widetilde{\mathcal{O}}$ to signify upper bounds up to numeric constants and polylogarithms respectively. We recall the definitions of convexity and smoothness:

For two probability measures $\mu,\nu\in\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})$, we define the $\mathsf{KL}$ divergence and the (relative) Fisher information by

with the convention $\mathsf{KL}(\mu\mathbin{\|}\nu)=\mathsf{FI}(\mu\mathbin{\|}\nu)=\infty$ whenever $\mu\not\ll\nu$.

We recall the definition of the log-Sobolev inequality, which is used both for propagation of chaos arguments as well as mixing time bounds.

When $\log(1/\pi)$ is $\alpha$-uniformly convex for $\alpha>0$, it follows from the Bakry–Émery condition that $\pi$ satisfies (LSI) with constant $C_{\mathsf{LSI}}\leq 1/\alpha$ [BGL14, Proposition 5.7.1].

We can also define the $p$-Wasserstein distance $\mathcal{W}_{p}(\mu,\pi)$, $p\geq 1$, between $\mu,\pi$ as

where $\Gamma(\mu,\pi)$ is the set of all joint probability measures on $\mathbb{R}^{d}\times\mathbb{R}^{d}$ with marginals $\mu,\pi$ respectively.

Lastly, we recall that the celebrated Otto calculus interprets the space $\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})$, equipped with the $\mathcal{W}_{2}$ metric, as a formal Riemannian manifold [Ott01]. In particular, the Wasserstein gradient of a functional $\mathcal{L}:\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})\to\mathbb{R}\cup\{\infty\}$ is given as $\nabla_{\mathcal{W}_{2}}\mathcal{L}=\nabla\delta\mathcal{L}$. Here, $\delta\mathcal{L}$ is the first variation defined as follows: for all $\nu_{0},\nu_{1}\in\mathcal{P}_{2,\rm ac}(\mathbb{R}^{d})$, $\delta\mathcal{L}(\nu_{0}):\mathbb{R}^{d}\to\mathbb{R}$ satisfies

The first variation is defined up to an additive constant, but the Wasserstein gradient is unambiguous. See [AGS08] for a rigorous development. As a shorthand, we will write $\delta\mathcal{L}(\nu_{0},x)\coloneqq\delta\mathcal{L}(\nu_{0})(x)$ and similarly $\nabla_{\mathcal{W}_{2}}\mathcal{L}(\nu_{0},x)\coloneqq\nabla_{\mathcal{W}_{2}}\mathcal{L}(\nu_{0})(x)$.

Our general approach for sampling from the stationary distribution $\pi$ in either (1.1) or (2.2) is to directly apply an off-the-shelf sampler for the finite-particle stationary distribution $\mu^{1:N}$. The theoretical guarantees for this procedure require two main ingredients: (1) control of the “bias”—i.e. the error incurred by approximating $\pi$ by the $1$-particle marginal of $\mu^{1:N}$—and (2) verification of isoperimetric properties which allow for fast sampling from the measure $\mu^{1:N}$.

In this section, we present results for sampling from $\pi$. As outlined in Algorithm 1, we use off-the-shelf log-concave samplers to sample from $\mu^{1:N}$, during which we access the first-order³³3For our results involving the proximal sampler, we also assume access to a proximal oracle for simplicity. oracle for $\mu^{1:N}$ (i.e. an oracle for evaluation of $\log\mu^{1:N}$ up to an additive constant, and for evaluation of $\nabla\log\mu^{1:N}$). For $N$ sufficiently large, the first particle given by Algorithm 1 is approximately distributed according to $\pi$: for $\hat{\mu}^{1:N}$ the law of the output of the log-concave sampler and its $1$-particle marginal distribution $\hat{\mu}^{1}$,

where the inequality follows from exchangeability (Lemma 27). A similar decomposition also holds for $\operatorname{\mathsf{KL}}$, although the argument is more technical. We defer its presentation to §E.

To bound the second term by $\varepsilon$, it suffices to choose $N$ according to the propagation of chaos results in §3.1. Our results are summarized in Table 1, in which we record the total number of oracle calls $M$ for $\mu^{1:N}$ made by the sampler $M$ and the number of particles $N$ needed to achieve $\varepsilon$ error in the desired metric, hiding polylogarithmic factors. Note that in the pairwise McKean–Vlasov setting, each oracle call to $\mu^{1:N}$ requires $N$ calls to an oracle for $V$, and $\binom{N}{2}$ calls to an oracle for $W$.

The algorithms in the table refer to: Langevin Monte Carlo (LMC); underdamped Langevin Monte Carlo (ULMC); discretizations of the underdamped Langevin diffusion via the randomized midpoint method [SL19] or the shifted ODE method [FLO21] (ULMC⁺); and implementation of the proximal sampler [LST21, Che+22] via the Metropolis-adjusted Langevin algorithm or via ULMC (MALA–PS and ULMC–PS respectively). Note that LMC applied to sample from $\mu^{1:N}$ is simply the Euler–Maruyama discretization of ($\mathsf{pMV}_{N}$), and likewise ULMC is the algorithm considered in [FW23]. We refer to §E for proofs and references.

To streamline the rates, we simplify the notation by defining $\beta=\beta_{V}+\beta_{W}$ if Assumption 1 holds, otherwise we use the value from Assumption 6. We let $\alpha=\alpha_{V}+{\alpha_{W}^{-}}$ under Assumption 4, $\alpha=\nicefrac{{\sigma^{2}}}{{2\max\{C_{\mathsf{LSI}}(\mu^{1:N}),C_{\mathsf{LSI}}(\pi)\}}}$ under Assumptions 2 and 9, and $\alpha=\nicefrac{{\sigma^{2}}}{{2\max\{C_{\mathsf{LSI}}(\mu^{1:N}),\overline{C}_{\mathsf{LSI}}\}}}$ in the general McKean–Vlasov setting.

Finally, we let $\kappa\coloneqq\beta/\alpha$ denote the condition number. We briefly justify this terminology. If the target and all proximal Gibbs measures are strongly convex with parameter $\alpha/\sigma^{2}$, then the Bakry–Émery condition implies that $\overline{C}_{\mathsf{LSI}}\leq\tfrac{\sigma^{2}}{\alpha}$. Hence, the scale-invariant ratio $\overline{C}_{\mathsf{LSI}}\beta/\sigma^{2}$ reduces to the classical condition number $\beta/\alpha$, the ratio of the largest to smallest eigenvalues of the Hessian matrices for $V$ and $W$. Therefore, $\overline{C}_{\mathsf{LSI}}\beta/\sigma^{2}$ is a generalization of the condition number to settings beyond uniform strong convexity which allows us to state more interpretable bounds. The additional assumption $\kappa\leq\sqrt{d}/\varepsilon$ will be used to simplify some of the rates.

In the following subsections, we discuss some of the results in greater detail.

In this work, we propose a framework for obtaining sampling guarantees for the minimizers of ($\mathsf{pE}$) and ($\mathsf{gE}$), based on decoupling the problem into (i) particle approximation via propagation of chaos, and (ii) time-discretization via log-concave sampling theory. Our approach leads to simpler proofs and improved guarantees compared to previous works, and our results readily benefit from any improvements in either (i) or (ii).

We conclude by listing some future directions of study. As discussed in §4.2, our uniform-in-$N$ LSI for the mean-field Langevin dynamics currently scales poorly in the problem parameters, and it is important to improve it. We also believe there is further room for improvement in the propagation of chaos results. For example, can the sharp rate in Theorem 3 be extended to stronger metrics such as Rényi divergences, as well as to situations when the weak interaction condition (Assumption 3) fails, e.g., in the strongly displacement convex case or in the setting of §3.1.2? For the sampling guarantees, the prior works [BS23, SWN23] considered different settings, such as potentials satisfying convexity at infinity or the use of stochastic gradients; these extensions are compatible with our approach and could possibly lead to improvements in these cases, as well as others. Finally, consider the case where $\int\nabla W(X_{t}-\cdot)\,\mathrm{d}\pi_{t}$ in ($\mathsf{pMV}$) is replaced with a generic function $\int b_{t}(X_{t},\cdot)\,\mathrm{d}\pi_{t}$, $b_{t}:\mathbb{R}^{n}\times\mathbb{R}^{n}\to\mathbb{R}^{n}$. It would be interesting to extend our analysis to this setting, as it arises in many applications [AD20].

We thank Zhenjie Ren for pointing out an error in a previous draft of this paper, and Daniel Lacker, Atsushi Nitanda, and Taiji Suzuki for important discussions and references. YK was supported in part by NSF awards CCF-2007443 and CCF-2134105. MSZ was supported by NSERC through the CGS-D program. SC was supported by the Eric and Wendy Schmidt Fund at the Institute for Advanced Study. MAE was supported by NSERC Grant [2019-06167] and CIFAR AI Chairs program at the Vector Institute. MBL was supported by NSF grant DMS-2133806.

pages8 rangepages-1 rangepages54 rangepages52 rangepages29 rangepages18 rangepages88 rangepages36 rangepages24 rangepages51 rangepages121 rangepages157 rangepages43 rangepages31 rangepages2 rangepages41 rangepages26 rangepages-1 rangepages32 rangepages-1 rangepages49 rangepages18 rangepages44 rangepages20 rangepages11 rangepages36 rangepages24 rangepages69 rangepages17 rangepages36 rangepages37 rangepages21 rangepages56 rangepages38 rangepages32 rangepages58 rangepages35 rangepages24 rangepages21 rangepages5 rangepages-1 rangepages54 rangepages17 rangepages17 rangepages74 rangepages37 rangepages40 rangepages49 rangepages47 rangepages28 rangepages87 rangepages49 rangepages235 rangepages-1 rangepages935 rangepages63 rangepages16 rangepages36