Gaussian is All You Need: A Unified Framework for Solving Inverse Problems via Diffusion Posterior Sampling

TL;DR

This paper introduces a unified likelihood approximation method for diffusion models that improves inverse problem solving by enhancing convergence and computational efficiency, without propagating gradients through the model.

Contribution

A novel likelihood approximation with covariance correction that enhances diffusion-based inverse problem solving and is computationally efficient.

Findings

Improved convergence towards true data posterior.

Enhanced performance on natural image datasets.

Efficient covariance matrix factorization for inverse problems.

Abstract

Diffusion models can generate a variety of high-quality images by modeling complex data distributions. Trained diffusion models can also be very effective image priors for solving inverse problems. Most of the existing diffusion-based methods integrate data consistency steps by approximating the likelihood function within the diffusion reverse sampling process. In this paper, we show that the existing approximations are either insufficient or computationally inefficient. To address these issues, we propose a unified likelihood approximation method that incorporates a covariance correction term to enhance the performance and avoids propagating gradients through the diffusion model. The correction term, when integrated into the reverse diffusion sampling process, achieves better convergence towards the true data posterior for selected distributions and improves performance on real-world natural image datasets. Furthermore, we present an efficient way to factorize and invert the covariance matrix of the likelihood function for several inverse problems. Our comprehensive experiments demonstrate the effectiveness of our method over several existing approaches. Code available at https://github.com/CSIPlab/CoDPS.