PaperlayerView exact revision on arXiv

Source-linked AI summary

Score-Based Diffusion Models as Principled Priors for Inverse Imaging

Berthy T. Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L. Bouman, William T. Freeman

arXiv:2304.11751v2cs.CV

TL;DR

Inverse imaging needs priors that capture complex image distributions without sacrificing principled posterior inference. This paper uses score-based diffusion models as probabilistic priors and applies their probability function in variational inference. Experiments across inverse-imaging tasks support accurate, modular posterior inference, while computational and variational-approximation costs remain limitations.

  • Problem

    Existing principled inverse-imaging methods use simple priors, while deep-learned priors are difficult to analyze or require hand-tuned measurement-prior weighting.

  • Method

    The paper trains a score-based diffusion prior once, computes its probability function, and uses it with a RealNVP variational distribution to approximate measurement-conditioned posteriors.

  • Results

    The framework supports principled posterior sampling across denoising, deblurring, and interferometric imaging, with deblurring results outperforming comparison methods on MSE, PSNR, and SSIM.

  • Takeaways & Limitations

    Score-based priors provide a modular interface that applies the same learned prior across inverse-imaging tasks without hand-tuning and supports uncertainty-focused scientific imaging.

  • Takeaways & Limitations

    Exact log-probabilities require potentially repeated ODE solves, and posterior quality depends on the expressiveness of the variational approximation.

Abstract

from arXiv · show

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.

1. Introduction

Inverse imaging is ill-posed, and existing approaches trade principled inference against expressive learned priors. The paper proposes score-based diffusion models as probabilistic priors that combine learned image statistics with Bayesian inference.

  • Noisy and incomplete measurements can admit many image solutions, so priors constrain reconstructions according to desired image statistics.
  • Principled Bayesian methods require tractable prior values or gradients, limiting them to simple distributions and weighted regularizers.
  • Deep-learned priors can produce convincing solutions but are difficult to probe, task-specific, or dependent on hand-tuned hyperparameters.
  • The paper proposes a score-based diffusion model as a Bayesian image prior with an empirically validated exact probability function.
  • A single trained score-based prior can pair with arbitrary likelihoods and inference algorithms, demonstrated through variational posterior sampling for denoising, deblurring, and interferometry.

2. Related Work

The paper frames inverse imaging as Bayesian recovery under incomplete measurements, contrasting principled but limited priors with expressive yet difficult-to-analyze deep generative models. Score-based diffusion models provide a route to learned distributions with probability-based inference.

  • 2.1. Inverse Problems in Imaging: Inverse imaging seeks a hidden image from measurements generated by a known forward model with random noise, and missing information makes recovery ill-posed.
  • 2.1. Inverse Problems in Imaging: Bayesian inference decomposes the posterior into a measurement likelihood and an image prior, enabling MAP estimation and posterior sampling for uncertainty quantification.
  • 2.2. Deep Learning for Inverse Problems: Deep neural networks learn complex image distributions, but implicit priors are difficult to analyze, require retraining for new tasks, or do not support uncertainty analysis.
  • 2.2. Generative Models: Classical and deep generative models generally simplify image distributions or lack exact tractable probabilities, whereas score-based diffusion models support a continuous-time probabilistic framework.
  • 2.2.1 Diffusion Models: Diffusion-based inverse-problem methods enforce measurement consistency through projections or likelihood gradients but require hyperparameter tuning and are not guaranteed to sample the true posterior.

3. Score-Based Priors

Score-based priors turn diffusion models into differentiable image distributions whose probabilities can be computed through a probability-flow ODE. Experiments validate their probabilities and gradients against a known Gaussian ground truth, supporting their use in principled inference.

  • 3. Score-Based Priors: A score-based prior is a differentiable image prior trained on a dataset and used through log pθ(x) to model posteriors for inverse-imaging tasks.
  • 3.1. Log-Probability Computation: The probability-flow ODE creates an invertible mapping between diffusion-time distributions without Brownian motion, enabling probability computation from a score model.
  • 3.1. Log-Probability Computation: The log-probability of an image combines its terminal Gaussian log-probability with the integrated divergence of the learned ODE vector field.
  • 3.2. Log-Probability Validation: For a 16 × 16 Gaussian image distribution, score-based log-probabilities were validated against ground truth on both in-distribution and out-of-distribution images.
  • 3.2. Log-Probability Validation: ODE-derived gradients are closer to true gradients than score-model outputs, while using score outputs directly can produce an incorrect posterior.
  • Practical use assumes a well-trained score model and incurs ODE computation costs, while variational posterior quality depends on the expressiveness of its parameterized distribution.

4. Posterior Sampling with Score-Based Priors

Score-based priors enable direct posterior modeling by treating a diffusion model as a fixed probabilistic prior and combining it with measurement likelihoods. Variational inference then produces posterior samples without hand-tuned measurement weights, with experiments showing close agreement to analytical posteriors.

  • Framework: Score-based priors allow posterior sampling with arbitrary image priors and measurements when sufficient training data and a known forward model are available.The approach plugs the score-based prior into an established variational-inference method, DPI.
  • Framework: Previous diffusion-based methods entangle measurements with diffusion and approximate the likelihood score, preventing principled posterior sampling across diffusion times.The true likelihood score is defined only at t = 0, while posterior reverse diffusion requires it for every t.
  • Framework: Hyperparameter weighting creates a tradeoff between prior and measurement information, and even ideally tuned baselines cannot capture the true posterior.Under-weighting biases samples toward the prior; over-weighting can collapse samples onto an implausible subspace.
  • Variational inference: The method minimizes KL divergence between a RealNVP variational family and the desired posterior using the likelihood, score-based-prior probability, and variational entropy.The expectation is Monte-Carlo approximated, while the RealNVP makes the entropy term tractable.
  • Variational inference: The objective requires no hand-tuned hyperparameters because the score-based prior is a true probability distribution rather than a weighted regularizer.This distinguishes the prior itself from the DPI inference procedure.
  • Empirical validation: In a Gaussian-prior deblurring experiment, the estimated posterior mean and covariance closely matched analytical values, while direct score-model output produced an incorrect posterior.The posterior was estimated from samples and compared against analytically derived statistics.

5. Posterior Sampling Results

The paper compares score-based priors and posterior-sampling methods across denoising, deblurring, and interferometric imaging. Score-based priors produce informative, data-driven posteriors and show stronger performance and robustness than traditional priors and diffusion-based baselines.

  • Baseline methods: Existing diffusion-based posterior samplers rely on likelihood-gradient approximations or projections, often requiring hand-tuned measurement-weight hyperparameters.SDE+Proj is restricted to compressed-sensing linear inverse problems, while Score-ALD and DPS use methods to adjust likelihood-score magnitude.
  • Denoising: Score-based priors achieve the highest average SSIM and PSNR for denoising posterior samples among the compared priors.The comparison includes TV regularization, PCA-G, RealNVP NF, and score-based priors, with metrics averaged across posterior samples.
  • Denoising: Score-based priors provide richer, data-driven uncertainty, assigning greater uncertainty to facial features than traditional priors.The empirical standard-deviation maps show uncertainty concentrated around structures such as eyes, ears, and hair.
  • Deblurring: 24.75 average PSNR was achieved by the proposed method on the CelebA deblurring example, compared with 20.37 for DPS.The proposed posterior samples also outperform other methods in MSE, PSNR, and SSIM in the reported experiments.
  • Deblurring: The proposed posterior-sampling method remains more robust under mismatched priors, whereas baselines hallucinate faces or introduce unnatural artifacts as measurement weight changes.This behavior is illustrated when a CelebA prior is applied to CIFAR-10 measurements.
  • Interferometric imaging: Score-based priors support scientific imaging by combining measurement consistency with image statistics from trained distributions.The paper motivates this use in black-hole interferometric imaging, where real black-hole images are unavailable for training and faithful posterior modeling is important.

6. Limitations

The score-based prior is an approximation whose correctness depends on the learned score model, and exact probability evaluation can be computationally expensive. Posterior sampling also depends on the expressiveness of the chosen variational family.

  • Limitations: The score-based prior is a learned approximation, so its correctness is tied to the correctness of the score model.The prior also involves a complexity tradeoff with compute time and memory.
  • Limitations: Exact log-probabilities require solving an ODE, often repeatedly during optimization, increasing computational cost.The paper identifies more generalizable score models and related improvements as possible ways to reduce this cost.
  • Limitations: The variational posterior approximates the true posterior, and its quality depends on the expressiveness of the parameterized variational distribution.

7. Conclusion

The paper turns score-based diffusion models into principled priors for posterior inference. Experiments support their modularity, expressiveness, and robustness across inverse-imaging tasks, including black-hole interferometric imaging.

  • Conclusion: The paper establishes score-based diffusion models as principled priors and demonstrates variational posterior sampling with them.
  • Conclusion: Experiments show modularity across denoising and deblurring, expressiveness relative to traditional priors, and robustness relative to diffusion-based inverse-problem methods.
  • Conclusion: Score-based priors are especially useful for scientific imaging such as interferometric black-hole imaging.

Supplemental

The supplemental section describes numerical methods for evaluating score-based-prior log-probabilities and gradients, including ODE solving, trace estimation, and continuous adjoints. Solver choice and multiple trace estimators are used to balance accuracy, efficiency, and variance.

  • Log-probability computation: Log-probabilities under a score-based prior are computed by solving an initial-value problem based on the probability-flow ODE.The computation starts from x0 = x and integrates the modified vector field defined by the pretrained score model.
  • Implementation considerations: The implementation uses Diffrax to swap ODE solvers and adaptively select time steps for log-probability estimation.Solver choice is one of the two implementation decisions identified as most affecting accuracy.
  • Implementation considerations: Hutchinson-Skilling trace estimation with multiple trace estimators reduces variance in estimated log-probabilities and gradients.The same trace estimators are applied to each image in a batch.
  • ODE solver: Bogacki–Shampine’s 3/2 and Dormand-Prince’s 5/4 methods provide a good efficiency–accuracy balance on the Gaussian ground-truth experiment.The authors caution that solver trends may differ for score-based priors trained on other datasets.
  • Adjoint ODE: The continuous adjoint method computes gradients through a secondary ODE while avoiding the memory cost of direct backpropagation through the ODE solve.The method is selected to balance memory, speed, and accuracy requirements.

A.2. Posterior Sampling Experiments

The posterior-sampling experiments use Gaussian and image-trained score models with RealNVP variational inference. Sampling is substantially faster than the compared diffusion-based baselines, while the framework is presented as reducing manual prior design effort.

  • Experimental setup: The Gaussian ground-truth distribution is fitted from CelebA images resized to 16 × 16 grayscale resolution and regularized by adding 0.01 to the covariance diagonal.Samples from this fitted Gaussian train the score model used in the ground-truth experiments.
  • Posterior sampling: Posterior sampling uses a RealNVP with 64 affine-coupling layers optimized by SGD with Adam at learning rate 0.0002.The optimization uses batch size 64 and clips gradients to norm 1.
  • Posterior sampling: The optimized RealNVP is sampled to obtain approximate posterior samples, with samples containing pixel magnitudes above 2 discarded as outliers.This postprocessing is applied before computing statistics of DPI-estimated posteriors.
  • Sampling efficiency: 128 samples of 32 × 32 RGB images take about 2.76 seconds, compared with 20.8 seconds for SDE+Proj, 51.8 seconds for Score-ALD, and 34.1 seconds for DPS.The authors note that optimization with the score-based prior is comparatively time-consuming.
  • Practical implication: The framework is described as producing a reliable, rich posterior while reducing the human effort required to handcraft and validate regularizers or priors.This claim concerns the practical workflow enabled by the framework.

A.3. 2D Experiments

The 2D experiments compare the proposed posterior-sampling approach with diffusion-based baselines using a common true score model. Because the baselines provide samples rather than posterior probabilities, their densities are approximated with KDE.

  • 2D posterior comparison: The proposed approach and SDE+Proj, Score-ALD, and DPS are compared on a toy 2D posterior using the same true score model.The proposed samples use a RealNVP with 32 affine-coupling layers.
  • 2D posterior comparison: Baseline posterior densities are approximated with kernel density estimation from 10000 samples because those methods provide samples but not posterior probabilities.The KDE procedure is used to construct approximate PDFs for comparison.
  • Gradient estimation: Figure 2 evaluates the mean and variance of estimated score gradients using 10 versus 50 trace estimators on 32 × 32 grayscale CelebA images.The test image is drawn from the model distribution, so the learned score should equal the true gradient at the evaluation point.

B. Image-Restoration Metrics

The deblurring experiments compare posterior samples from the proposed approach with SDE+Proj, Score-ALD, and DPS using standard image-restoration metrics. The authors caution that these metrics do not establish posterior correctness, although they suggest the framework remains preferable for high-quality posterior samples.

  • Deblurring evaluation: Deblurring results are evaluated against SDE+Proj, Score-ALD, and DPS using MSE, SSIM, and PSNR.These are standard image-restoration metrics applied to the compared posterior-sampling approaches.
  • Metric limitation: The authors state that MSE, SSIM, and PSNR do not reflect whether the posterior is correct.They nevertheless report that the metrics suggest their framework is preferable when applications require high-quality posterior samples.

C. Score-Based Priors vs. Discrete-Flow Priors

Score-based diffusion priors are more expressive than discrete normalizing-flow priors and support more stable, visually convincing posterior inference in the reported experiments, while RealNVP performance depends on problem constraint and initialization.

  • Discrete normalizing flows are limited by invertible architectures and probability functions that generalize poorly outside training data.
  • Compared with a score-based prior trained on the same dataset, the RealNVP prior produced less visually convincing samples and unstable DPI posterior optimization.
  • The limited expressiveness of discrete normalizing flows may explain their poorer samples, while poor generalization may contribute to optimization instability.
  • The reported RealNVP limitation is qualified: clever initialization or more expressive normalizing-flow architectures could improve performance.
  • The deblurring comparison trained score-based diffusion and RealNVP priors on the same CelebA dataset and applied both probability functions in DPI.
  • RealNVP may be adequate when measurements sufficiently constrain the posterior but may fail to capture it when the inverse problem is extremely ill-posed.
Loading 2304.11751v2…