Thinned Mean Field Langevin Dynamics
For researchers using particle-based methods for entropy-regularized optimization, this work offers a computationally cheaper alternative to mean-field Langevin dynamics with minimal loss in convergence rate.
The authors propose KT-MFLD, a method that reduces the computational complexity of mean-field Langevin dynamics from O(N^2) to O(N^{3/2}) by using kernel thinning to create a smaller particle coreset, while maintaining similar convergence guarantees. Empirical validation on neural network training, quantization, and Bayesian inference tasks confirms the theoretical results.
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.