Nearest Neighbour Score Estimators for Diffusion Generative Models
This addresses the bottleneck of score estimation for researchers and practitioners in generative modeling, offering a novel method with potential broad impact, though it appears incremental as it builds on existing diffusion frameworks.
The paper tackles the problem of high variance in score function estimation for diffusion generative models by introducing a nearest neighbor estimator that uses multiple training samples, resulting in significantly increased convergence speed and sample quality for consistency models and enabling network-free probability-flow ODE integration in diffusion models.
Score function estimation is the cornerstone of both training and sampling from diffusion generative models. Despite this fact, the most commonly used estimators are either biased neural network approximations or high variance Monte Carlo estimators based on the conditional score. We introduce a novel nearest neighbour score function estimator which utilizes multiple samples from the training set to dramatically decrease estimator variance. We leverage our low variance estimator in two compelling applications. Training consistency models with our estimator, we report a significant increase in both convergence speed and sample quality. In diffusion models, we show that our estimator can replace a learned network for probability-flow ODE integration, opening promising new avenues of future research.