Mean-Field Langevin Dynamics: Exponential Convergence and Annealing

Introduction

where $\eta>0$ is the step-size and $\mathbf{Z},\mathbf{Z},\dots$ are i.i.d. standard Gaussian vectors (see Eq. (10) for an equivalent definition of NPGD directly in terms of $G$ and its first-variation).

and it is thus sufficient to choose $m=1$ in that case. In the general case of a convex and non-linear $G$, the particles will interact in non-trivial ways and $m$ should be taken large, so that a mean-field behavior emerges.

The dynamics obtained in the many-particle $m\to\infty$ and vanishing step-size $\eta\to 0$ limit was called the Mean-Field Langevin dynamics in Hu et al. and is our object of interest. In this limit, the distribution $\mu_{t}$ of particles at time $t=k\eta$ solves the following drift-diffusion partial differential equation (PDE) of McKean-Vlasov type:

There is a long line of work around mean-field dynamics [Dobrushin, 1979, Sznitman, 1991] (see Lacker for an introduction and references) which guarantee that NPGD (3) indeed converges to the Mean-Field Langevin dynamics, sometimes with fine quantitative bounds [Lacker, 2021, Mei et al., 2019]. As for the behavior of the Mean-Field Langevin dynamics (5) itself, it is shown in [Mei et al., 2018, Hu et al., 2021] that, under suitable coercivity assumptions, $(\mu_{t})$ weakly converges to the unique minimizer of $F_{\tau}$ as $t\to\infty$. Moreover, Hu et al. remarks that the results from Eberle et al. to obtain quantitative rates apply here, but this argument is restricted to the large noise regime and does not exploit the convexity of $G$. These works leave open the question of quantitative guarantees without a strong noise assumption.

1 Contributions and related work

We prove that, under a certain uniform log-Sobolev inequality assumption (which is in particular satisfied in the settings of Mei et al. , Hu et al. ), solutions to (5) converge at a global exponential rate to the minimizer of $F_{\tau}$ (Theorem 3.2). The known convergence rate of the Langevin dynamics under a log-Sobolev inequality is recovered as a particular case when $G$ is linear.

We study the annealed dynamics where the noise $\tau=\tau_{t}$ is time-dependent and decays as $\alpha/\log(t)$ and prove that for $\alpha>0$ large enough, $G(\mu_{t})$ converges towards the minimum of the unregularized functional $F_{0}=G$ (Theorem 4.1).

In Section 5, we show that our results apply to noisy gradient descent on infinitely wide two-layer neural networks and we provide numerical experiments for $G$ being a kernel Maximum Mean Discrepancy (MMD).

Let us mention that other algorithms to solve problems of the form (1) are possible. Nitanda et al. proposed a dual averaging scheme which involves a sequence of Langevin diffusions and enjoys a $O(1/t)$ convergence rate in the mean-field limit. For low-dimensional problems, one can resort to discretizing the measure on a fixed grid, which leads to a convex problem amenable to standard (Bregman) gradient descent algorithms [Tseng, 2010].

The long-time behavior of drift-diffusion PDEs of the form Eq. (5) has been studied in the mathematical physics literature. General convergence rates (under assumptions that imply a large noise in our context) are proved in Eberle et al. . For interacting particule systems with an interaction kernel $k$ (discussed in Section 5.2), the case where $k(x,y)=h(x-y)$ with $h$ convex can be dealt with using the notion of displacement convexity, see [Villani, 2021, Chap. 9.6]. In contrast, we rely on standard convexity, which corresponds to a positive semi-definite interaction kernel $k$.

Upon completion of this work, we became aware of the paper [Nitanda et al., 2022] which also proves the exponential convergence of the Mean-Field Langevin dynamics with the same proof technique. The main differences between these two works is that they perform a discrete time analysis while we study the annealed dynamics. These works were conducted independently and simultaneously.

2 Notations

Assumptions and preliminaries

The Mean-Field Langevin dynamics in Eq. (5) involves the first-variation $V$ of $G$, defined as follows.

If it exists, the first-variation $V[\mu]$ is unique up to an additive constant.

for some time-dependent velocity field $v\in L^{2}((a,b),L^{2}(\mu_{t}))$. Observe that Eq. (5) is an equation of this form with $v_{t}=-\nabla V[\mu_{t}]-\tau\nabla\log(\mu_{t})$.

Let $(\mu_{t})_{t\in(a,b)}$ be a weakly continuous solution to Eq. (8) such that $\nabla\log(\mu_{t})\in L^{2}((a,b),L^{2}(\mu_{t}))$. Then $G(\mu_{t})$ and $H(\mu_{t})$ are absolutely continuous functions of $t$ and it holds for a.e. $t\in(a,b)$,

Using the vocabulary of analysis in Wasserstein space, the function $H$ is displacement convex with subdifferential $\nabla\log\mu$ [Ambrosio and Savaré, 2007, Thm. 4.16]. Also we prove in Lemma A.2 that $G$ is $(-2L)$-displacement convex with subdifferential $\nabla V[\mu_{t}]$. Then the claim is a consequence of [Ambrosio and Savaré, 2007, Sec. 4.4.E]. ∎

2 Characterization of the minimizer

We recall the optimality conditions for $F_{\tau}$ which have been proved in several works (see e.g. [Mei et al., 2018, Lem. 10.4] or [Hu et al., 2021, Prop. 2.5]) and require the following assumptions.

Under Assumption 1 and 2, the minimizer $\mu^{*}_{\tau}$ of $F$ is unique and satisfies

The uniqueness comes from the strict convexity of $H$. For Eq. (9), one first derives the first order optimality condition, which require that $V[\mu_{\tau}^{*}]+\tau\log(\mu_{\tau}^{*})$ must be a constant $\mu_{\tau}^{*}$-almost everywhere. Then one shows that $\mu_{\tau}^{*}$ has positive density everywhere due to the entropy term, and concludes. We refer to [Mei et al., 2018, Lem. 10.4] for details.

3 Noisy Particle Gradient Descent (NPGD)

The NPGD algorithm has been defined in Section 1 via the function $G_{m}$. We now give an alternative definition of this algorithm involving the first-variation $V$ of $G$.

where $\eta>0$ is the step-size and $Z_{i,k}\sim\mathcal{N}(0,I)$ are iid standard Gaussian random variables.

Under Assumption (1), the two definitions of NPGD in Eq. (3) and in Eq. (10) are equivalent.

This proves that $\forall i\in[m]$, $m\nabla_{X_{i}}G_{m}(\mathbf{X})=\nabla V[\mu_{0}](X_{i})$ and thus the update equations in Eq. (3) and Eq. (10) are the same. ∎

4 Mean-Field Langevin dynamics

Given Eq. (10) standard results about mean-field systems tell us that as $m\to\infty$, the random measure $\hat{\mu}_{k}$ becomes deterministic, so that in the limit (and taking also the small-step size limit $\eta\to 0$) the particles trajectories are given by i.i.d. samples from the following stochastic differential equation (SDE)

where $(B_{t})_{t\geq 0}$ is a Brownian motion. As mentioned in the introduction, the law $(\mu_{t})$ of a solution to this SDE solves the following PDE which is our main object of study:

where $\nabla\cdot$ stands for the divergence operator. Standard results about this class of PDEs guarantee its well-posedness, i.e. the existence of a unique solution, under Assumption 1 (see e.g. [Huang et al., 2021, Thm. 3.3] or Ambrosio and Savaré for an approach based on the gradient flow structure which applies here thanks to Lemma A.2 which states that $G$ is $(-2L)$-displacement convex).

Let us now study the convergence of $(\mu_{t})_{t\geq 0}$ to the global minimum of $F_{\tau}$.

Exponential convergence of Mean-Field Langevin dynamics

In the general case, we make an analogous assumption that such an inequality holds uniformly for $e^{-V[\mu_{t}]/\tau}$ throughout the dynamics.

These two criteria are standard, but many finer criteria are known, such as integral conditions [Wang, 2001], Lyapunov conditions [Cattiaux et al., 2010] or criteria for mixture distributions [Chen et al., 2021] (see also Section 5.1). Our main result regarding the Mean-Field Langevin dynamics (11) is the following.

Note that although Lemma 2.2 requires some regularity estimates, they can be bypassed here thanks to general results about Wasserstein gradient flows [Ambrosio and Savaré, 2007, Thm. 5.3 (v)]. Combining this energy identity with the log-Sobolev inequality and Lemma 3.4, it follows

which is a $2$-Łojasiewicz gradient inequality for $F_{\tau}$. By integrating in time we get Eq. (14). ∎

Under the assumptions of Theorem 3.2, for $t\geq 0$ we have

The following lemma establishes inequalities which are key to handle the non-linear aspect of the dynamics (when $G$ is linear, they become trivial equalities).

Invoking this inequality twice with the role of $\mu$ and $\mu^{*}$ exchanged, it holds

Recalling $F_{\tau}(\mu)=G(\mu)+\tau H(\mu)$, it holds, on the one hand,

On the other hand, using the fact that $\mu^{*}_{\tau}=e^{-V[\mu^{*}]/\tau}$ (Proposition 2.3), it holds

Convergence of the annealed dynamics

We now turn our attention to the “annealed” Mean-Field Langevin dynamics

Here we show that a similar guarantee holds in our more general context.

The lower-bound assumed on $\rho_{\tau}$ is natural when one has in mind the Holley and Stroock criterion given in Section 3. In Section 5, we show a lower bound of this form on a concrete example related to two-layer neural networks.

The bounds of Theorem 4.1 exhibit a two time-scales phenomenon: the dynamics $(\mu_{t})$ converges at a polynomial rate to the regularization path $(\mu^{*}_{\tau_{t}})$ (in relative entropy or $W_{2}^{2}$ distance, thanks to the “entropy sandwich” Lemma 3.4 or the Talagrand inequality) but the regularization path only converges at a logarithmic rate to the optimal value $\inf G$, because of the slow decay of $\tau_{t}$.

Our proof is partly inspired by Miclo , as revisited by Tang and Zhou .

Step 1. Consider the function that returns the values of the regularization path

As an infimum of affine functions, $h$ is concave and since the minimizer $\mu^{*}_{\tau}$ is unique, $h$ is differentiable for $\tau>0$ and its derivative is $h^{\prime}(\tau)=H(\mu^{*}_{\tau})$. We focus on $t\geq t_{0}$ so that $\tau_{t}=\alpha/\log(t)$. By Lemma 2.2 applied to $v_{t}=-\nabla V[\mu_{t}]+\tau_{t}\nabla\log(\mu_{t})$ (here again, the regularity assumptions of Lemma 2.2 can be bypassed using the gradient flow-like structure, see [Ferreira and Valencia-Guevara, 2018, Thm. 6.9]), we have

where we introduced the probability measure $\nu_{t}\propto e^{-V[\mu_{t}]/\tau_{t}}$. On the one hand, we have by the Log-Sobolev inequality and the “entropy sandwich” Lemma 3.4,

On the other hand, by Lemma 4.2 below, it holds for some $C_{2},C_{3}>0$ independent from $\mu$ and $t$,

where in the last step, we used that $G$ is lower bounded and $h(\tau)=F_{\tau}(\mu^{*}_{\tau})$ is bounded for $\tau\in[0,\tau_{0}]$. Combining the previous estimates, we get that for any $\epsilon>0$, there exists $C_{1},C_{2},C_{3}>0$ such that

where we used $\tau_{t}=\alpha/\log(t)$, $\tau^{\prime}_{t}=-\alpha/(t(\log t)^{2})$ and $\rho_{\tau_{t}}\geq C_{0}t^{-\alpha^{*}/\alpha}$. In passing, the first inequality in the above display guarantees that $F_{\tau_{t}}(\mu_{t})-F_{\tau_{t}}(\mu_{\tau_{t}}^{*})$ remains finite at all time because $\log\tau_{t}\in\mathcal{C}^{1}$, which justifies the fact that we can consider only $t$ large enough in the rest of the proof.

It follows that for any $\epsilon>0$ such that $\epsilon<1-\alpha^{*}/\alpha$, for $t$ large enough and some $C,C^{\prime}>0$,

and thus $F_{\tau_{t}}(\mu_{t})-F_{\tau_{t}}(\mu_{\tau_{t}}^{*})\leq C^{\prime\prime}t^{-\kappa}$ because $\kappa\coloneqq 1-\frac{\alpha^{*}}{\alpha}-\epsilon>0$ and $Q(t^{*})$ is finite. This proves Eq. (18).

where we have used the bound $-H(\mu_{t})\leq C_{1}F_{\tau_{t}}(\mu_{t})+C_{2}$ from Lemma 4.2, which is uniformly bounded for $t\geq 0$ by some $C^{\prime}$ thanks to Step 1.

On the other hand, we have by Jensen’s inequality for the convex function $\varphi:s\mapsto s\log(s)$ and Fubini’s theorem:

which is the entropy of the Gaussian distribution $g_{\sigma}$. Thus we have

by choosing $\sigma^{2}=\tau/L$. Plugging the value of $\tau_{t}=\alpha/\log(t)$ we get, for some $C,C^{\prime}>0$,

In the proof of Theorem 4.1, we used a lower bound on the value of $H(\mu)$ in terms of the functional value that is provided in the following lemma.

Invoking Lemma 4.3 with $\sigma^{2}=\tau/C_{1}$ we have

Summing the two previous equations (with the same value of $C_{1}$), we get that for $\tau\leq\tau_{0}$,

Combined with the fact that $-H(\mu)\leq C_{1}^{\prime}M_{2}(\mu)+C_{2}^{\prime}$, we get $-H(\mu)\leq C_{1}^{\prime\prime}F_{\tau}(\mu)+C_{2}^{\prime\prime}$.∎

See e.g. [Mei et al., 2018, Lem. 10.1] for a proof of the following lemma.

Applications and experiments

We now show that our results apply to the training dynamics of certain wide $2$-layer neural networks trained with noisy gradient descent.

When $\mu=\frac{1}{m}\sum_{i=1}^{m}\delta_{x_{i}}$ is an empirical distribution with $m$ atoms/particles, the function $G_{m}$ derived from $G$ as in Eq. (2) is exactly the risk with weight decay regularization for a two-layer neural network of width $m$. Thus noisy gradient descent for two-layer neural networks is equivalent to NPGD with $G$ defined in Eq (21).

Let us give simple conditions under which our convergence theorems apply in this case.

For the computation of the first variation and Assumptions 1, we refer e.g. to Hu et al. . For Assumption 2, $G$ is convex as a composition of a linear operator and a convex function. To see that $F_{\tau}$ admits a minimizer, notice that thanks to the regularization term, the sublevel sets of $G$ are tight and thus weakly-precompact by Prokorov’s theorem. Moreover, the loss term in $G$ is weakly continuous, the regularization term is weakly lower-semicontinuous (lsc) and $H$ is weakly lsc [Ambrosio and Savaré, 2007, Sec. 3.2] so, overall, $F_{\tau}$ is lsc. Thus a minimizer $\mu^{*}_{\tau}$ exists for all $\tau\geq 0$ by the Direct Method in the calculus of variations.

which is not covered by our assumptions (in particular because it is not bounded). In the case of the ReLU non-linearity $\phi(s)=\max\{0,s\}$, there is in addition a lack of smoothness issue. An interesting direction for future research would be to adapt this algorithm and analysis in order to cover the case of ReLU non-linearities.

then a uniform LSI holds (i.e. Assumption 3 holds).

2 Numerical illustration: kernel Maximum Mean Discrepancy

We conclude this paper with numerical experiments exploring the behavior of NPGDLink to Julia code to reproduce the experiments: https://github.com/lchizat/2022-mean-field-langevin-rate. defined in (3). Let us stress that our theoretical guarantees only apply to the mean-field Langevin dynamics – recovered in the many-particle and continuous time limit – so there remains a gap between the theory and the NPGD algorithm.

The first variation of $G$ is given, for $\mu\in\mathcal{P}(\mathcal{X})$ and $x\in\mathcal{X}$ by

Moreover, Assumptions 1,2 and 3 are satisfied with $\rho_{\tau}\geq((1+2\pi)d)^{-1}e^{(\inf k-\sup k)/\tau}$.

In our experiments, we consider $d=2$ and the translation invariant kernel k(x,y)=\prod_{i=1}^{d}\big{(}1+2\sum_{k=1}^{n}(1+k)^{-1}\cos(k(x_{i}-y_{i}))\big{)} with $n=5$ frequency components. Because of the frequency cut-off, this kernel is not strictly positive definite, so $G$ admits minimizers other than $\nu$ [Sriperumbudur et al., 2010, Simon-Gabriel et al., 2020] although in practice we observed that NPGD was in general attracted towards $\nu$. We take $\nu$ as a random empirical distribution of $m^{*}=10$ samples from the uniform distribution on $\mathcal{X}$. We run NPGD with $m=50$ particles, a step-size $\eta=0.08$ and $\mu_{0}$ being the uniform distribution on $\mathcal{X}$.

Figure 1(a) shows an example of a large-time particle configuration, with the atoms of $\nu$ is red and the atoms of $\hat{\mu}_{t}$ in black (with $t$ large), with a noise temperature $\tau=0.1$. Here the measure $\hat{\mu}_{t}$ is a noisy version of $\nu^{*}$.

Figure 1(b) shows the evolution of the objective $F_{\tau}=G+\tau H$ (up to a constant, adjusted for ease of comparison) along the iterations, where the entropy $H$ is estimated using the $1$-nearest-neighbor estimator [Kozachenko and Leonenko, 1987, Singh et al., 2003]. We observe the exponential decay of $F_{\tau}$ towards a plateau which we expect to be the global minimum of $F_{\tau}$, up to discretization errors. For small values of $\tau$, it is not excluded that the plateau corresponds instead to a suboptimal metastable state.

Finally, Figure 1(c) shows the advantage of NPGD with simulated annealing vs. PGD to minimize the unregularized function $G$. We used a noise temperature that decays polynomially as $\tau_{t}=20(t+1)^{-1}$ where $t$ is the iteration count, which is a faster decay than what the theory suggests. At iteration $800$, we stopped the noise in order to observe the “quality” of the configuration of particles. We see that the NPGD with simulated annealing consistently outperforms PGD, which gets stuck in poorer local minima.

Conclusion

We have proved the convergence of the Mean-Field Langevin dynamics to the global minimizer at an exponential rate, under natural assumptions that include all settings where (non-quantitative) convergence was previously shown. We have also proved the convergence of the annealed dynamics for a suitable noise decay.

From a higher perspective, our analysis—in particular the simple “entropy sandwich” Lemma 3.4—suggests that often, the guarantees about Langevin dynamics obtained via log-Sobolev inequalities can be generalized to mean-field Langevin dynamics. In this paper, we focused on exponential convergence and on simulated annealing, but other aspects could be considered, such as a direct analysis of the discrete dynamics, which could lead to computational bounds, as done in e.g. [Vempala and Wibisono, 2019, Ma et al., 2019] for the Langevin algorithm.

Another interesting direction for future work is to develop and study more applications of Mean-Field Langevin dynamics, since many problems can be cast as optimization problems of the form Eq. (1). This includes sparse deconvolution problems, mixture models fitting [Boyd et al., 2017] or problems involving optimal transport [Peyré and Cuturi, 2019, Chap. 9].

I would like to thank Loucas Pillaud-Vivien for orienting me through the literature of simulated annealing and Mo Zhou for noticing a gap in a previous version of the proof of Theorem 4.1. I would also like to thank the anonymous reviewers for their insightful comments and for suggesting the discussion on LSI in Section 5.1.

References

Appendix A Additional proofs

Let us start with a relation between $G$ and its first-variation $V$ that is more convenient for proofs.

where $\mu_{t}=(1-t)\mu_{0}+t\mu_{1}$ for $t\in$.

Since the same holds true for $-G$, we could say that $G$ is $2L$-smooth in the Wasserstein geometry, in the sense that it is both $2L$-semiconvex and $2L$-semiconvex.

For $\epsilon>0$ and $s\in$, let $\mu_{s}=(1-s)\mu+s\mu_{t+\epsilon}$. It holds by Lemma A.1,

where we used successively the Lipschitz continuity of $x\mapsto\nabla V[\mu](x)$ and of $\mu\mapsto\nabla V[\mu](x)$ in the last two lines. The first claim follows by taking the limit $\epsilon\to 0$. This also shows that $V[\mu]$ is the unique (strong) $W_{2}$-differential of $W_{2}$ at $\mu$, in the sense of [Ambrosio and Savaré, 2007, Def. 4.1].

For the semi-convexity claim, let $h(t)\coloneqq G(\mu_{t})$. For $s,t\in$, it holds by Cauchy-Schwarz

Since $W_{2}(\mu_{s},\mu_{t})\leq|t-s|\|v\|_{L^{2}(\mu)}$, it follows

which proves that $G$ is $(-2L)$-convex in the sense of [Ambrosio and Savaré, 2007, Remark 3.2]. ∎