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 O(N^{\frac{1}{2}}). 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.