Efficient Mass Matrix Estimation with Gaussian Cooling

Abstract

We introduce an algorithm for efficiently preconditioning log-concave and log-smooth distributions that scales logarithmically with the condition number of the underlying distribution. Based on Gaussian cooling, our multistage method approximately samples from a sequence of well-conditioned distributions to construct a preconditioner for each subsequent stage of the cooling schedule. This method is motivated by existing practical approaches for mass matrix estimation in Markov chain Monte Carlo (MCMC), in which an effective preconditioner for the target distribution is estimated from samples produced during a warm-up period. However, our approach is conceptually distinct from existing approaches and moreover permits theoretical analysis. Our algorithmic framework is flexible, allowing essentially any log-concave sampling algorithm to act as the subroutine within each cooling stage. For instance, using first-order rejection sampling (FORS) as the sampler, the first-order query complexity of our method is $\tilde{O}(d^{4/3} \log κ)$ for a large class of log-concave and log-smooth distributions with condition number $κ$ and $\tilde{O}(d^{1/3} \log κ)$ for translation-invariant distributions. In the translation-invariant setting, we provide numerical experiments for the lattice $ϕ^4$ model, a simple lattice field theory.