Thinned Mean Field Langevin Dynamics

Abstract: Several important learning tasks can be formulated as minimizing an entropy-regularized objective over an appropriate space of probability distributions. Mean-field Langevin dynamics (MFLD) facilitate computation in this general context, casting the minimizer as the invariant distribution of a McKean--Vlasov process, which can be numerically discretized using $N$ particles and thus simulated. However, simulating this interacting particle system has computational complexity of order $N^2$. Motivated by recent research into \emph{kernel thinning}, we propose \texttt{KT-MFLD}, in which each particle interacts only with a thinned particle coreset of size $\mathcal{O}(N^{\frac{1}{2}})$. \texttt{KT-MFLD} thus reduces the computational complexity to order $N^{\frac{3}{2}}$ while, under mild regularity conditions, achieving the same convergence guarantees (up to logarithmic factors) as MFLD. Our theoretical analysis is empirically confirmed on tasks including the training of student-teacher neural networks, quantization with maximum mean discrepancy, and computation of predictively-oriented posteriors in a post-Bayesian framework.

• The paper introduces a kernel-thinned mean field Langevin dynamics method that reduces pairwise computational cost from O(N^2) to O(N^(3/2)).
• The paper establishes theoretical error bounds of O(N^(-1)(log N)^2) and validates these claims empirically across neural networks, MMD quantization, and mean-field games.
• The paper demonstrates that KT-MFLD can scale high-dimensional, multimodal target distributions with controlled bias, offering practical benefits for probabilistic learning.

Thinned Mean Field Langevin Dynamics: Accelerating Interacting Particle Systems via Kernel Thinning

Entropy-Regularized Objectives and Mean Field Langevin Dynamics

Entropy regularization is fundamental in probabilistic machine learning, underpinning a wide range of tasks that can be cast as minimizing functionals of the form $$F(\mu) = F_0(\mu) + \frac{\zeta}{2} E_\mu[\|x\|^2] - \sigma \mathrm{Ent}(\mu)$$ over measures $\mu$. When $F$ is nonlinear in $\mu$, the minimizer $\pi$ typically lacks a tractable explicit form, impeding traditional sampling schemes such as MCMC. Mean Field Langevin Dynamics (MFLD) provide a principled numerical solution by simulating an interacting particle system that approximates the McKean–Vlasov process whose invariant distribution coincides with $\pi$. MFLD facilitates computation in complex objectives encountered in mean-field neural networks, kernel-based quantization, and post-Bayesian inference frameworks.

However, the standard MFLD algorithm incurs an $\Omega(N^2)$ computational cost per iteration due to exhaustive pairwise interactions among $N$ particles. This complexity severely restricts scalability, often forcing practitioners to constrain $N$ to modest values, which limits the ability to approximate high-dimensional or multimodal target distributions.

Figure 1: KT-MFLD replaces the $\Omega(N^2)$ pairwise interactions in MFLD with interactions over a thinned coreset of $\mu$0 representative particles, reducing per-iteration cost to $\mu$1.

Kernel Thinning for Computational Acceleration

Motivated by advances in kernel thinning, the paper introduces Kernel-Thinned Mean Field Langevin Dynamics (KT-MFLD), which substitutes the full set of $\mu$2 particles with a thinned coreset of $\mu$3 particles at each iteration. The selection process uses fast kernel thinning algorithms (notably kt-split and KT-Compress) to ensure that the coreset retains integration fidelity relative to the original ensemble, controlled in the appropriate RKHS norm.

Theoretical guarantees are established under mild regularity assumptions: KT-MFLD achieves finite-particle approximation error scaling as $\mu$4, closely mirroring the original MFLD up to polylogarithmic factors. Notably, this is sharper than random subsampling, which delivers only $\mu$5 error. The reduction in complexity from $\mu$6 to $\mu$7 opens the door to larger particle populations with practical resource requirements.

Algorithmic Details and Theoretical Analysis

KT-MFLD's iteration replaces the empirical measure over all particles in the drift term of the SDE with an empirical measure over the thinned coreset. This modification can be applied to various objectives of interest, including those featuring pairwise or quadratic interactions:

$\mu$8

where $\mu$9 is the thinned subset.

Under regularity assumptions including boundedness, Lipschitzness, convexity, and appropriate kernel/RKHS properties, KT-MFLD retains non-asymptotic convergence guarantees for entropy-regularized objectives, with error terms characterized explicitly. A key result demonstrates that for sufficiently large $F$0 and $F$1, and suitably small $F$2, the joint distribution of KT-MFLD particles at termination closely approximates the $F$3-fold product of the target distribution, up to an extra $F$4 term arising from thinning.

Empirical Evaluation Across Applications

Comprehensive experiments validate KT-MFLD's theoretical predictions across several applications:

• Mean-field neural networks: In a student-teacher setup, KT-MFLD (using both kt-split and KT-Compress) consistently achieves lower test loss at fixed computational cost compared to both Random-MFLD and RBM-MFLD.
• MMD quantization: KT-MFLD variants outperform subsampling-based accelerations in minimizing MMD against target distributions, with KT-Compress yielding superior results due to its strong coreset guarantees.

Figure 2: KT-MFLD variants show improved generalization for student-teacher networks under fixed computational budgets compared to subsampling baselines.

Figure 3: KT-MFLD achieves lower MMD quantization error; KT-Compress outperforms both kt-split and full MFLD under appropriate parameter regimes.

• Predictively-Oriented (PrO) posteriors: KT-MFLD enables efficient approximation of PrO posteriors for misspecified Lotka–Volterra models, showing reduced kernel gradient discrepancy (KGD) relative to full MFLD and subsampled methods at equal wall-clock time.

  Figure 4: KT-MFLD consistently attains lower KGD in PrO posterior inference for Lotka–Volterra models, confirming practical reduction in bias.

• Mean-field games: KT-MFLD is further demonstrated to substantially reduce computation time in simulating mean-field games, achieving lower equilibrium risk versus subsampled variants, extending thinning applicability beyond standard MFLD.

  Figure 5: KT-MFLD accelerates mean-field game simulations, attaining lower risk functionals at parity of computational budget.

Practical and Theoretical Implications

KT-MFLD enables scalable simulation of interacting particle systems for entropy-regularized variational objectives. The reduction in computational cost, combined with provably controlled bias, allows researchers and practitioners to leverage larger populations—improving approximation quality in high-dimensional and multimodal regimes.

Theoretical implications include clarity on the trade-offs between coreset construction fidelity and computational efficiency. However, KT-MFLD's error term exhibits adverse dependence on dimensionality through the log-Sobolev constant, which may challenge its efficacy in extremely high-dimensional contexts.

Furthermore, kernel thinning is applicable to a broader class of mean-field models, general interacting particle systems (e.g., McKean–Vlasov processes), and potentially to deterministic variants (e.g., Stein variational gradient descent) pending advances in RKHS approximation theory for more general drift terms.

Outlook and Future Directions

The paper identifies several promising avenues:

• Combining kernel thinning with momentum-based or Nesterov-type accelerations in measure spaces.
• Extending kernel thinning approaches to interacting particle systems in ensemble Kalman filtering, consensus-based optimization, and sequential Monte Carlo.
• Addressing curse of dimensionality in coreset thinning error via adapted kernel architectures or domain-specific approximations.
• Scaling KT-MFLD to larger experimental regimes, including deep neural architectures and large-scale Bayesian inference.

Conclusion

KT-MFLD delivers a robust framework for accelerating mean-field Langevin dynamics and related interacting particle systems with controlled bias, exploiting kernel thinning algorithms to achieve near-optimal approximation rates at substantially reduced computational cost. The approach has broad implications for scalable probabilistic learning and numerical simulation, and further integration with advanced computational strategies promises enhanced practical utility across machine learning, statistical inference, and dynamical systems modeling.