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

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 xt to the posterior mean mathbbE [x0 | xt], in order to guide the diffusion trajectory with an estimate of the final denoised sample x0. However, this does not consider information from the measurement y, which must then be integrated downstream. In this work, we propose to estimate the conditional posterior mean mathbbE [x0 | xt, y], 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 y. We demonstrate comparable or improved performance against a wide selection of contemporary inverse solvers across multiple datasets and tasks.