Variational Bayesian Imaging with an Efficient Surrogate Score-based Prior

TL;DR

This paper introduces a surrogate function for score-based priors in Bayesian imaging, significantly improving computational efficiency and accuracy in posterior estimation for ill-posed inverse problems.

Contribution

It proposes a novel surrogate prior based on the evidence lower bound, enabling faster and more accurate variational inference with score-based diffusion models.

Findings

Accelerates variational inference by over 100 times compared to exact priors.

Achieves more accurate posterior estimates than non-variational diffusion methods.

Provides a practical approach for using score-based models as general image priors.

Abstract

We propose a surrogate function for efficient yet principled use of score-based priors in Bayesian imaging. We consider ill-posed inverse imaging problems in which one aims for a clean image posterior given incomplete or noisy measurements. Since the measurements do not uniquely determine a true image, a prior is needed to constrain the solution space. 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.