Mirror Mean-Field Langevin Dynamics

TL;DR

The paper addresses constrained distributional optimization of an entropy-regularized functional $\mathcal{L}(\mu)=F(\mu)+\lambda\mathrm{Ent}(\mu)$ over $\mu\in\mathcal{P}_2(\mathcal{X})$ with convex $\mathcal{X}$. It introduces the mirror mean-field Langevin dynamics (MMFLD), combining mirror Langevin geometry with mean-field Langevin dynamics to respect domain constraints through a mirror map, enabling optimization on constrained domains. The authors prove linear convergence of the continuous-time MMFLD under a mirror log-Sobolev inequality (MLSI) and establish uniform-in-$N$ propagation of chaos and non-asymptotic bounds for the time- and particle-discretized MMFLD, supported by a discretization framework and a mean-field network setting. Numerical experiments on the unit simplex show that MMFLD yields stable, well-distributed mass away from the boundary and achieves lower objective values than projected MFLD, highlighting practical advantages for constrained problems. Overall, the work provides the first global convergence guarantees for constrained distributional optimization via mirror geometry and lays groundwork for extending to mirror Poincaré and broader geometric constraints.

Abstract

The mean-field Langevin dynamics (MFLD) minimizes an entropy-regularized nonlinear convex functional on the Wasserstein space over $\mathbb{R}^d$, and has gained attention recently as a model for the gradient descent dynamics of interacting particle systems such as infinite-width two-layer neural networks. However, many problems of interest have constrained domains, which are not solved by existing mean-field algorithms due to the global diffusion term. We study the optimization of probability measures constrained to a convex subset of $\mathbb{R}^d$ by proposing the \emph{mirror mean-field Langevin dynamics} (MMFLD), an extension of MFLD to the mirror Langevin framework. We obtain linear convergence guarantees for the continuous MMFLD via a uniform log-Sobolev inequality, and uniform-in-time propagation of chaos results for its time- and particle-discretized counterpart.

Figures (1)

• Figure 1: (a)-(c) MMFLD v.s. projected MFLD for optimizing a nonlinear mean-matching objective over the unit simplex; (d)-(f) with an additional boundary potential preventing mass accumulation.

Theorems & Definitions (23)

• Lemma 2.1: Properties of $\phi$
• Theorem 2.2
• Definition 1: First variation
• Definition 2
• Theorem 2.3
• Definition 3: Proximal Gibbs distribution
• Theorem 2.4
• Definition 4: Local and dual norms
• Theorem 3.1
• Theorem 3.2