Variational Bayesian Imaging with an Efficient Surrogate Score-based Prior
Abstract: We propose a surrogate function for efficient yet principled use of score-based priors in Bayesian imaging. Recent work turned score-based diffusion models into principled priors for solving ill-posed imaging problems by appealing to an ODE-based log-probability function. However, evaluating the ODE is computationally inefficient and inhibits posterior estimation of high-dimensional images. Our proposed surrogate prior is based on the evidence lower bound of a score-based diffusion model. We demonstrate the surrogate prior on variational inference for efficient approximate posterior sampling of large images. Compared to the exact prior in previous work, our surrogate accelerates optimization of the variational image distribution by at least two orders of magnitude. We also find that our principled approach gives more-accurate posterior estimation than non-variational diffusion-based approaches that involve hyperparameter-tuning at inference. Our work establishes a practical path forward for using score-based diffusion models as general-purpose image priors.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: **Adding results from a traditional prior (qKCT, jZTo).** We have added results from TV regularization to Figure 6. We find that our method and the diffusion-based baselines far outperform TV regularization in terms of image-reconstruction quality.
**Discussion of when one has a training dataset for the prior and would not want to use a paired training dataset (qKCT).** We have added such a discussion to the end of Section 2.1. In particular, we discuss how using a paired training dataset makes it difficult to analyze different measurement settings with the same prior since the network needs to be re-trained with each new measurement setting. It also makes it difficult to analyze the effects of different priors since an end-to-end network has an implicit prior that might change whenever it is re-trained. Training only the prior also allows one to better enforce the data likelihood while the network focuses on learning the prior, which is the unknown part of the problem. As for when the situation arises that one has a training dataset of clean images, there are many large-scale datasets for different applications, including the fastMRI dataset for MRI. It is also usually possible to create a clean dataset through manual acquisition of clean images or simulation.
**Other efforts to speed up diffusion models (jZTo).** We have added a discussion of work in this area at the end of Section 2.2. Many of the cited efforts aim to speed up sampling, not probability computation, and may even not allow for tractable probabilities. Our work aims for an efficient surrogate for probabilities using the same diffusion model. In fact the surrogate function may complement other efforts to speed up sampling/probability computation under the diffusion model, as any speedup in the diffusion process itself could also further improve the efficiency of evaluating the ELBO.
**Other writing suggestions (7FR5).** We have revised the abstract and introduction to clearly state that we consider ill-posed inverse imaging problems and better define “ill-posed image reconstruction.” Additionally, we have made the suggested typo fixes, additional references, and rephrasing in Eqs. 2 and 3 and Section 4.
**Broader impacts (jZTo).** We have added a discussion of broader impacts to the conclusion (Section 6).
Assigned Action Editor: ~Valentin_De_Bortoli1
Submission Number: 2151
Loading