Provable Probabilistic Imaging using Score-Based Generative Priors

TL;DR

This work introduces plug-and-play Monte Carlo (PMC), a principled posterior-sampling framework that fuses score-based generative priors with plug-and-play priors to solve ill-posed imaging inverse problems while providing uncertainty quantification. By establishing a continuous-time gradient-flow interpretation of PnP/RED and discretizing Langevin dynamics, the authors formulate two PMC algorithms (PMC-PnP and PMC-RED) and two annealed variants (APMC-PnP and APMC-RED) that leverage weighted annealing to improve exploration. They prove Fisher-information-based stationary-distribution convergence with an $O(1/N)$ rate under broad conditions, accommodating non-log-concave likelihoods and imperfect score networks, and demonstrate convergence acceleration via annealing. Extensive experiments across compressed sensing, MRI, and nonlinear black-hole interferometry show that PMC variants deliver superior reconstruction quality and reliable uncertainty quantification, including multimodal posterior recovery in nonlinear settings where traditional score-based sampling can fail. Overall, PMC provides a model-agnostic, theoretically grounded pathway to accurate image reconstruction with uncertainty quantification using expressive score priors.

Abstract

Estimating high-quality images while also quantifying their uncertainty are two desired features in an image reconstruction algorithm for solving ill-posed inverse problems. In this paper, we propose plug-and-play Monte Carlo (PMC) as a principled framework for characterizing the space of possible solutions to a general inverse problem. PMC is able to incorporate expressive score-based generative priors for high-quality image reconstruction while also performing uncertainty quantification via posterior sampling. In particular, we develop two PMC algorithms that can be viewed as the sampling analogues of the traditional plug-and-play priors (PnP) and regularization by denoising (RED) algorithms. To improve the sampling efficiency, we introduce weighted annealing into these PMC algorithms, further developing two additional annealed PMC algorithms (APMC). We establish a theoretical analysis for characterizing the convergence behavior of PMC algorithms. Our analysis provides non-asymptotic stationarity guarantees in terms of the Fisher information, fully compatible with the joint presence of weighted annealing, potentially non-log-concave likelihoods, and imperfect score networks. We demonstrate the performance of the PMC algorithms on multiple representative inverse problems with both linear and nonlinear forward models. Experimental results show that PMC significantly improves reconstruction quality and enables high-fidelity uncertainty quantification.

Figures (23)

• Figure 1: Conceptual illustration of how weighted annealing improves the convergence of APMC algorithms by introducing the weighted posteriors $\{\pi^{({\alpha_k})}_{\sigma_k}\}$. The solid curves and shade respectively denote the mean and probability density of the distribution; the white area means $\nabla \mathsf{log\,} p({\bm{x}})=0$. In order to facilitate the vanilla PMC algorithm to escape from plateaus in $\nabla \mathsf{log\,} p({\bm{x}})$, weighted annealing progressively decreases 1) the smoothing strength (${\sigma_k}$) of the prior and 2) its relative weights (${\alpha_k}$) to the likelihood.
• Figure 2: Comparison of the sample statistics obtained by PMC-PnP and APMC-PnP versus the ground-truth posterior distribution. Each test algorithm is run to infer a batch of $1000$ samples, which are then classified into two modes according to their distance and angle with respect to the ground-truth modes. PMC-PnP, which does not use annealing, fails to identify the bimodal distribution, as can be seen by the two groups of classified images being nearly indistinguishable. Under the pre-trained score, APMC-PnP significantly improves in performance over PMC-PnP by distinguishing the female and male modes. These two modes look similar to the ground truth posterior modes, although some differences remain due to the use of an approximate learned score. In the ideal case of the analytical score, APMC-PnP recovers a distribution that closely resembles the ground-truth posterior.
• Figure 3: Visual illustration of BHI. The left two images together demonstrate the subsampling pattern in the Fourier spectrum. The Groundtruth image shows the ground-truth black hole simulation image used in this experiment. The Target image corresponds to the scenario where the low-frequency band is fully sampled, resembling a single-dish telescope the size of the Earth. This target image represents the intrinsic resolution of our telescope; an effort to recover sharper features would be classified as attempting superresolution.
• Figure 4: Visualization of the sampling results obtained by APMC-PnP. In total 100 samples were drawn. (a) The t-SNE plot (perplexity=20) shows the distribution of the samples. Note that this t-SNE plot shows there are two distinct image modes. (b) Pixel-wise statistics of each mode. (c) The distribution of the closure phase $\chi^2_\mathsf{cph}$ and log closure amplitude $\chi^2_\mathsf{camp}$ statistics for each mode. Note that APMC-PnP successfully recovers the two modes of the posterior distribution, with both modes resulting in $\chi^2$ statistics close to 1.
• Figure 5: Visualization of the sampling results obtained by scoreDPI. (a) shows the t-SNE plot (perpelxity=20), and (b) plots the sample statistics. Note that scoreDPI recovers a single-mode distribution rather than a bimodal distribution.

Theorems & Definitions (12)

• Theorem 1: PMC-RED
• proof
• Theorem 2: PMC-PnP
• proof
• Theorem 3: APMC-RED
• proof
• Theorem 4: APMC-PnP
• proof
• Lemma 1: Balasubramanian.etal2022
• proof