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 introduce two PMC algorithms which can be viewed as the sampling analogues of the traditional plug-and-play priors (PnP) and regularization by denoising (RED) algorithms. We also establish a theoretical analysis for characterizing the convergence of the PMC algorithms. Our analysis provides non-asymptotic stationarity guarantees for both algorithms, even in the presence of 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.

Method

Plug-and-play Monte Carlo (PMC) is built on the fusion of PnP/RED and score-based generative model (SGM): it leverages powerful score-based generative priors in a plug-and-play fashion, similar to PnP/RED, while enabling provable posterior sampling by incorporating the Markov Chain Monte Carlo (MCMC) formulation used in the SGM. An additional weighted annealing is employed to accelerate the sampling speed. PMC equipped with weighted annealing is referred to as annealed PMC (APMC) here after.



Experimental Results

Linear inverse problems

Visual comparison of the reconstructions obtained by APMC algorithms and baseline algorithms for 10% CS (1st row) and 8x MRI (2nd row) tasks. The final images of the sampling algorithms are obtained by averaging 10 image samples. The visual difference is highlighted in the zoom-in images. Note how APMC algorithms restore the fine details.

Visualization of the pixel-wise statistics associated with the CS and MRI reconstructions shown in last figure. Figure (a) corresponds to CS, and figure (b) to MRI. In each figure, the left columns plot the absolute error and standard deviation (SD), and the right columns plot the 3-SD credible interval with the outlying pixels highlighted in red. Note that APMC algorithms lead to a better UQ performance than the baselines by recovering an accurate mean and thus avoiding the need for an arbitrarily large SD.

Nonlinear black-hole imaging (BHI)

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. Note that prior literature has discovered two modes in the posterior distribution for this particular BHI task.

Visualization of the sampling results obtained by APMC. 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 and log closure amplitude Chi-square statistics for each mode. Note that APMC successfully recovers the two modes of the posterior distribution, with both modes resulting in Chi-square statistics close to 1.

Acknowledgments

This work is sponsored by the Heritage Medical Research Fellowship, S2I Clinard Innovation Award, and NSF award 1935980. Z.W. is sponsored by the Kortschak Fellowship and Amazon AI4Science Partnership Discovery Grant. Y.C. acknowledges the support from the Courant instructorship. We thank Charles Gammie, Ben Prather, Abhishek Joshi, Vedant Dhruv, and Chi-kwan Chan for providing the black hole simulations. We thank Charles A. Bouman for his valuable feedback.

BibTeX


      @article{sun2023provable,
        title={Provable Probabilistic Imaging using Score-Based Generative Priors},
        author={Yu Sun and Zihui Wu and Yifan Chen and Berthy Feng and Katherine L. Bouman},
        journal={ArXiv},
        year={2023},
        volume={abs/2310.10835},
        url={https://api.semanticscholar.org/CorpusID:264172914}
      }