Bayesian Conditioned Diffusion Models for Inverse Problems

TL;DR

BCDM addresses conditioning of diffusion models for linear inverse imaging problems by deriving the conditional score $\nabla_{x_t}\log q(x_t;y)$ and estimating the conditional expectation $E[x_0; x_t,y]$ via a correlated denoising framework. The method unrolls a diffusion solver with a principled data-consistency step and trains a denoiser to recover $x_0$ from forward-model-corrupted inputs, using Tweedie-inspired identities to connect the scores. Empirically, it achieves state-of-the-art results across denoising, MRI reconstruction, inpainting, and super-resolution, and demonstrates robustness to mismatched SNRs and masks, at the cost of losing unconditional generation capability. The work provides a practical,Bayesianly-grounded conditioning mechanism for diffusion models in inverse problems with forward operators, offering improved performance in challenging, low-SNR settings.

Abstract

Diffusion models have recently been shown to excel in many image reconstruction tasks that involve inverse problems based on a forward measurement operator. A common framework uses task-agnostic unconditional models that are later post-conditioned for reconstruction, an approach that typically suffers from suboptimal task performance. While task-specific conditional models have also been proposed, current methods heuristically inject measured data as a naive input channel that elicits sampling inaccuracies. Here, we address the optimal conditioning of diffusion models for solving challenging inverse problems that arise during image reconstruction. Specifically, we propose a novel Bayesian conditioning technique for diffusion models, BCDM, based on score-functions associated with the conditional distribution of desired images given measured data. We rigorously derive the theory to express and train the conditional score-function. Finally, we show state-of-the-art performance in image dealiasing, deblurring, super-resolution, and inpainting with the proposed technique.