Can Diffusion Models Provide Rigorous Uncertainty Quantification for Bayesian Inverse Problems?

2025-03-04

AI Summary

This work addresses Bayesian inverse problems by coupling pretrained diffusion priors with task-specific likelihoods to enable posterior sampling and rigorous uncertainty quantification. To this end, we propose BIPSDA—a unified framework jointly modeling diffusion priors, inverse problem likelihoods, and sampling mechanisms—supporting image inpainting, X-ray tomography, and phase retrieval. Our contributions include: (i) the first systematic evaluation of diffusion models’ uncertainty calibration on analytically tractable Gaussian mixture priors; (ii) a modular, composable posterior sampling paradigm integrating annealing, MCMC variants, and other strategies. Experiments demonstrate high-fidelity posterior sampling and reliable uncertainty estimation for inpainting and tomography. For phase retrieval, we uncover a fundamental theoretical limitation of existing diffusion samplers arising from inherent bimodality—leading to inconsistent posterior approximations. BIPSDA thus advances both the practical applicability and theoretical understanding of diffusion-based Bayesian inference in ill-posed inverse problems.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Programming
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs

Abstract

In recent years, the ascendance of diffusion modeling as a state-of-the-art generative modeling approach has spurred significant interest in their use as priors in Bayesian inverse problems. However, it is unclear how to optimally integrate a diffusion model trained on the prior distribution with a given likelihood function to obtain posterior samples. While algorithms that have been developed for this purpose can produce high-quality, diverse point estimates of the unknown parameters of interest, they are often tested on problems where the prior distribution is analytically unknown, making it difficult to assess their performance in providing rigorous uncertainty quantification. In this work, we introduce a new framework, Bayesian Inverse Problem Solvers through Diffusion Annealing (BIPSDA), for diffusion model based posterior sampling. The framework unifies several recently proposed diffusion model based posterior sampling algorithms and contains novel algorithms that can be realized through flexible combinations of design choices. Algorithms within our framework were tested on model problems with a Gaussian mixture prior and likelihood functions inspired by problems in image inpainting, x-ray tomography, and phase retrieval. In this setting, approximate ground-truth posterior samples can be obtained, enabling principled evaluation of the performance of the algorithms. The results demonstrate that BIPSDA algorithms can provide strong performance on the image inpainting and x-ray tomography based problems, while the challenging phase retrieval problem, which is difficult to sample from even when the posterior density is known, remains outside the reach of the diffusion model based samplers.

Problem

• Integrate diffusion models with likelihood for Bayesian inverse problems.
• Assess uncertainty quantification in diffusion-based posterior sampling algorithms.
• Evaluate BIPSDA framework on image inpainting and x-ray tomography.

Innovation

• Introduces BIPSDA for diffusion-based posterior sampling.
• Unifies diffusion model algorithms for Bayesian problems.
• Tests on image inpainting, x-ray tomography, phase retrieval.