Technical Summary: Mirror Mean-Field Langevin Dynamics

Problem Statement

The paper addresses the optimization of an entropy-regularized functional over the space of probability measures constrained to a convex subset $X \subseteq \mathbb{R}^d$. The objective function is defined as:
$$\mathcal{L}(\mu) := F(\mu) + \lambda \text{Ent}(\mu)$$
where $\mu \in \mathcal{P}_2(X)$, $F: \mathcal{P}_2(X) \to \mathbb{R}$ is a linearly convex functional, and $\text{Ent}(\mu)$ is the entropy. This formulation arises naturally in statistical inference and machine learning, including infinite-width neural network limits, tensor decomposition, and density estimation.

While Mean-Field Langevin Dynamics (MFLD) effectively solves this problem when $X = \mathbb{R}^d$ by modeling the Wasserstein gradient flow, it fails for constrained domains. Standard Langevin dynamics rely on global diffusion terms that cause probability mass to escape the domain $X$. Existing constrained sampling methods, such as projecting particles back onto $X$ after each step, introduce undesirable behaviors, such as mass accumulation on the boundary $\partial X$, leading to suboptimal convergence and overfitting.

Methodology

To solve this constrained optimization problem, the authors propose Mirror Mean-Field Langevin Dynamics (MMFLD). This approach unifies the Mirror Langevin Dynamics (MLD) framework with the Mean-Field Langevin Dynamics (MFLD) framework.

Core Mechanism

The method employs a mirror map $\nabla \phi: X \to \mathbb{R}^d$, derived from a strictly convex barrier function $\phi$, to transform the constrained primal space $X$ into an unconstrained dual space $\mathbb{R}^d$. The dynamics are defined via a system of Stochastic Differential Equations (SDEs) in the dual space:
$$\begin{cases} X_t = \nabla \phi^*(Y_t) \\ dY_t = -\nabla \frac{\delta F(\mu_t)}{\delta \mu}(X_t) dt + \sqrt{2\lambda} \nabla^2 \phi(X_t) dB_t \end{cases}$$
where $\mu_t = \text{Law}(X_t)$, $\phi^*$ is the Legendre dual of $\phi$, and $B_t$ is standard Brownian motion. The diffusion term is adapted to the geometry of the domain via the Hessian of the mirror map, ensuring particles remain within $X$ without explicit projection.

Discretization and Propagation of Chaos

The paper analyzes a time- and particle-discretized version of MMFLD (Algorithm 1). Key technical components include:

• Forward Discretization: The authors adopt a forward discretization scheme for the drift term (inspired by Ahn and Chewi [1]) while simulating the diffusion step. This choice is motivated by its ability to achieve vanishing bias as the step size $\eta \to 0$, unlike standard Euler-Maruyama discretizations which may retain bias without strong assumptions.
• Uniform-in-Time Propagation of Chaos: The analysis extends recent advances in propagation of chaos to constrained domains. The authors establish that the error between the finite-particle system and the continuous mean-field limit converges to zero uniformly in time as the number of particles $N \to \infty$.
• Stochastic Gradients: The framework is extended to handle stochastic gradients, which is crucial for computational efficiency in large-scale applications like neural network training.

Key Contributions

1. Algorithm Proposal: The introduction of MMFLD, the first algorithm to minimize entropy-regularized functionals on general convex constrained domains with global convergence guarantees.
2. Convergence Guarantees:
   • Continuous Time: Proof of linear convergence for the continuous MMFLD flow under a uniform log-Sobolev inequality (LSI).
   • Discretized Setting: Non-asymptotic convergence rates for the time- and particle-discretized algorithm. The bounds are independent of the LSI constant in the particle approximation error, extending results from Nitanda [39] and Nitanda et al. [41] to the mirror setting.
3. Theoretical Extension: Demonstration that uniform-in-time propagation of chaos results, previously established for unconstrained MFLD, extend gracefully to constrained domains via the mirror map framework.
4. Stochastic Gradient Analysis: Convergence analysis for MMFLD utilizing stochastic gradients, providing error bounds that account for gradient noise.

Results

The paper provides theoretical convergence rates and validates the method through numerical experiments:

• Simplex Optimization: In optimizing a nonlinear mean-matching objective over the unit simplex, MMFLD significantly outperforms projected MFLD. Projected MFLD exhibits mass accumulation on the boundary, leading to higher loss, whereas MMFLD maintains a stable distribution and achieves lower loss.
• Spectraplex Optimization: For optimization over the spectraplex (positive semidefinite matrices with trace 1), projected MFLD makes negligible progress due to the geometry of the projection step canceling iteration gains. MMFLD successfully optimizes the objective.
• Neural Network Classification: In training norm-constrained two-layer ReLU networks for XOR classification, MMFLD converges faster and achieves lower test error than projected MFLD. The neuron distributions in MMFLD align tightly with the decision boundaries, whereas projected MFLD scatters neurons and saturates norm constraints.
• Scalability: Experiments in high dimensions ($d=100$) confirm that MMFLD scales effectively, converging to optimal solutions where projected methods stagnate.

Significance and Claims

The paper claims that MMFLD constitutes the first algorithm capable of minimizing entropy-regularized functionals on general convex constrained domains with global convergence guarantees.

The authors emphasize that their work bridges a critical gap in mean-field optimization: while MFLD is well-understood for unconstrained spaces, many practical applications (e.g., trajectory inference, Wasserstein barycenters, constrained neural networks) require constrained domains. By unifying mirror descent principles with mean-field dynamics, the authors provide a theoretically grounded and empirically effective solution that avoids the boundary accumulation issues inherent in projection-based methods. The work establishes that the geometric advantages of mirror maps (isometry between primal and dual spaces) can be successfully integrated into the mean-field setting to ensure stability and optimality in constrained optimization.