Injecting Measurement Information Yields a Fast and Noise-Robust Diffusion-Based Inverse Problem Solver

Abstract

Diffusion models have been firmly established as principled zero-shot solvers for linear and nonlinear inverse problems, owing to their powerful image prior and iterative sampling algorithm. These approaches often rely on Tweedie’s formula, which relates the diffusion variate 𝐱_t to the posterior mean 𝔼 [𝐱_0 | 𝐱_t], in order to guide the diffusion trajectory with an estimate of the final denoised sample 𝐱_0. However, this does not consider information from the measurement 𝐲, which must then be integrated downstream. In this work, we propose to estimate the conditional posterior mean 𝔼 [𝐱_0 | 𝐱_t, 𝐲], which can be formulated as the solution to a lightweight, single-parameter maximum likelihood estimation problem. The resulting prediction can be integrated into any standard sampler, resulting in a fast and memory-efficient inverse solver. Our optimizer is amenable to a noise-aware likelihood-based stopping criteria that is robust to measurement noise in 𝐲. We demonstrate comparable or improved performance against a wide selection of contemporary inverse solvers across multiple datasets and tasks.