Improving score-based sampling via affine post-processing

2026-10-03

AI Summary

This study addresses the limited sampling quality in score-based diffusion models caused by large bias and high variance in score estimation. To mitigate this, we propose an affine post-processing technique that introduces a computational budget allocation mechanism, transforming single evaluations into multi-point queries. Specifically, the method constructs base estimators using Markov chain Monte Carlo (MCMC) and importance sampling, and achieves joint optimization of bias and variance through affine combinations of multiple score estimates via linear regression. Experimental results demonstrate that, under equivalent computational budgets, the proposed approach significantly reduces estimation error and effectively improves sample generation quality for complex target distributions, outperforming conventional score estimation methods.

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.

Problem

• score-based sampling
• log-density
• score estimation
• bias and variance
• diffusion models

Innovation

• score-based diffusion
• affine post-processing
• score estimation
• variance reduction
• sampling