Score-Based Diffusion Models as Principled Priors for Inverse Imaging

TL;DR

The paper introduces score-based priors as principled, hyperparameter-free Bayesian priors for inverse imaging, enabling exact log-probability computations via a probability-flow ODE. It then shows how to perform posterior sampling using a variational framework with a normalizing flow, demonstrated on denoising, deblurring, and interferometric imaging, including black-hole simulations. The results indicate richer, uncertainty-aware posteriors that automatically balance data fidelity with prior knowledge, and robustness to prior mismatches compared to existing baselines. This work bridges diffusion-based generative modeling with classic principled inference, offering a practical route for principled, data-driven imaging with uncertainty quantification in scientific applications.

Abstract

Priors are essential for reconstructing images from noisy and/or incomplete measurements. The choice of the prior determines both the quality and uncertainty of recovered images. We propose turning score-based diffusion models into principled image priors ("score-based priors") for analyzing a posterior of images given measurements. Previously, probabilistic priors were limited to handcrafted regularizers and simple distributions. In this work, we empirically validate the theoretically-proven probability function of a score-based diffusion model. We show how to sample from resulting posteriors by using this probability function for variational inference. Our results, including experiments on denoising, deblurring, and interferometric imaging, suggest that score-based priors enable principled inference with a sophisticated, data-driven image prior.

Figures (9)

• Figure 1: A score-based prior is a hyperparameter-free, probabilistic prior that is also expressive and data-driven. Paired with a set of measurements, the prior can be used for principled inference of a full posterior. In this example, a score-based prior was trained on face images ("Prior" shows samples from the learned prior). The inverse problem is interferometic imaging of a synthetic black hole. We simulated interferometric measurements from the actual telescope array used to capture the first black-hole image event2019first and sampled images from the posterior via variational inference. From the top to bottom row, the posterior stably moves away from the prior given more constraining measurements. With measurements from only three telescopes, the posterior shows strong influence from the prior and contains images resembling faces that are brighter on the left half. As more telescopes (measurements) are added, the posterior reveals the ring-like structure of the underlying image. Our framework finds the proper relative strengths of the prior and measurements automatically.
• Figure 2: Log-probabilities of the score-based prior vs. ground-truth. Black line indicates perfect agreement. In-Distribution. The log-probabilities of 128 samples from the Gaussian ground-truth distribution were evaluated (shown as scatter points). Score-based log-probabilities are strongly correlated with ground-truth log-probabilities ($R^2\approx0.98$). Out-of-Distribution. The log-probabilities of test images from CIFAR-10 (scaled to $16\times 16$) are shown. The score-based prior generalizes well out of distribution.
• Figure 3: ODE gradients vs. score-model outputs. Histogram shows density of cosine distance between the estimated gradient and true gradient for 128 samples from the ground-truth Gaussian. $\mathbf{s}_\theta(\mathbf{x},t=0)$ is the score-model-approximated gradient ($t=0.001$ is actually used for numerical stability). $\nabla_{\mathbf{x}}\log p_\theta(\mathbf{x})$ is computed numerically according to the probability flow ODE.
• Figure 4: Baseline methods song2022solvingjalal2021robustchung2023diffusion do not sample true posterior. Heatmap depicts $p(\mathbf{x}| y)$ approximated from samples; contour lines depict true posterior. All methods use the same (true) score function. For baselines, we did a grid search to find the optimal hyperparameter weight (ALD and DPS hyperparameters were distilled into a global hyperparameter). The KL divergence from the estimated posterior to the true posterior was approximated for each hyperparameter value. No matter the value, baselines do not get more accurate than our hyperparameter-free method. For instance, with "Best meas. weight", all baselines sample from both modes equally, even though the lower-left mode should have more density. A poor hyperparameter setting can even lead to unstable sampling (see "High meas. weight" of Score-ALD, where samples became NaNs due to slightly over-weighted measurements).
• Figure 5: True vs. estimated posterior. Measurements are the 6.25% lowest DFT spatial-frequencies of an image from the prior and have i.i.d. noise with $\sigma=\lvert 1 \rvert$. The true mean and variance were derived analytically since the prior is the Gaussian ground-truth distribution, and the likelihood distribution is also Gaussian. The mean and variance were estimated from 10240 samples.