Diffusion Model Guided Sampling with Pixel-Wise Aleatoric Uncertainty Estimation

TL;DR

This paper introduces a method to estimate pixel-wise aleatoric uncertainty during diffusion model sampling, improving image quality assessment and generation by filtering low-quality samples and enhancing FID scores.

Contribution

It proposes a novel uncertainty estimation technique for diffusion models and demonstrates its effectiveness in quality filtering and improved sample generation.

Findings

Uncertainty estimates correlate with the second-order derivative of the noise distribution.

Filtering based on uncertainty improves sample quality on ImageNet and CIFAR-10.

Guided sampling results in better FID scores compared to baseline diffusion models.

Abstract

Despite the remarkable progress in generative modelling, current diffusion models lack a quantitative approach to assess image quality. To address this limitation, we propose to estimate the pixel-wise aleatoric uncertainty during the sampling phase of diffusion models and utilise the uncertainty to improve the sample generation quality. The uncertainty is computed as the variance of the denoising scores with a perturbation scheme that is specifically designed for diffusion models. We then show that the aleatoric uncertainty estimates are related to the second-order derivative of the diffusion noise distribution. We evaluate our uncertainty estimation algorithm and the uncertainty-guided sampling on the ImageNet and CIFAR-10 datasets. In our comparisons with the related work, we demonstrate promising results in filtering out low quality samples. Furthermore, we show that our guided approach leads to better sample generation in terms of FID scores.