Thinned Mean Field Langevin Dynamics

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 kernel thinning, we propose KT-MFLD, in which each particle interacts only with a thinned particle coreset of size \(\mathcal{O}(N^{\frac{1}{2}})\). KT-MFLDthus 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.