Mirror Mean-Field Langevin Dynamics

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.

Recent works on the optimization of two-layer neural networks have resulted in the study of the mean-field Langevin dynamics (MFLD), an extension of the classical Langevin dynamics underlying the wide success of Markov chain Monte Carlo (MCMC) methods. However, the analyses of MFLD typically assume full support of the measure. Motivated by the study of the mirror Langevin dynamics, we are interested in studying the mirror MFLD, where the measure is constrained to a convex subset of Euclidean space. Our experiments over the simplex, spectraplex, and bounded domains demonstrate the effectiveness of our method over the projected MFLD.