Improving score-based sampling via affine post-processing

Abstract

Sampling from a probability distribution with access only to its log-density is a foundational problem in statistics, machine learning, and the physical sciences. This problem is studied in the data-free score-based diffusion literature, which samples the target density via estimating scores of densities produced by convolution with Gaussian random variables and implementing a reverse SDE. In this framework, score estimators are constructed via Markov chain Monte Carlo, importance sampling, or rejection sampling. The estimators often operate far beyond the regime where mixing time and error control can be established and have large bias and variance when estimating the underlying score. In our method we propose post-processing: splitting a fixed computation budget across multiple evaluations, querying the estimator at different points, and fitting a single pooled estimate via linear combination to minimize bias and variance at a target location. The final value is an affine map of a vector of multiple score estimates. We show this reduces score estimation error when compared to the computation-matched base estimator queried at a single point. We demonstrate improved sample quality over other score estimation methods matched for computation cost on difficult target densities such as log-concave densities with large condition number, multimodal densities, and high-dimensional densities.