Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Enhancing Diffusion Models for Inverse Problems with Covariance-Aware Posterior Sampling

Shayan Mohajer Hamidi, En-Hui Yang

TL;DR

This work addresses solving linear inverse problems with pre-trained unconditional DDPM priors by improving posterior sampling. It derives a closed-form expression for the conditional posterior covariance Cov(x0|xt) in DDPMs and introduces a practical, covariance-aware approximation via a diagonal Hessian estimated with finite differences, yielding Σ̃_t. By combining the mean x̃0 from Tweedie’s formula with the covariance Σ̃_t, CA-DPS approximates p_t(x0|xt) as a Gaussian and computes a refined likelihood p_t(y|xt), enabling more accurate posterior sampling without retraining or hyperparameter tuning. Empirical results on FFHQ and ImageNet show CA-DPS consistently outperforms existing methods across inpainting, deblurring, and super-resolution, and a toy dataset demonstrates better posterior fidelity to the true posterior. The approach is designed to be easily applicable to pretrained DDPMs and scalable via conjugate gradient solves, making covariance-aware diffusion a practical improvement for inverse problems.

Abstract

Inverse problems exist in many disciplines of science and engineering. In computer vision, for example, tasks such as inpainting, deblurring, and super resolution can be effectively modeled as inverse problems. Recently, denoising diffusion probabilistic models (DDPMs) are shown to provide a promising solution to noisy linear inverse problems without the need for additional task specific training. Specifically, with the prior provided by DDPMs, one can sample from the posterior by approximating the likelihood. In the literature, approximations of the likelihood are often based on the mean of conditional densities of the reverse process, which can be obtained using Tweedie formula. To obtain a better approximation to the likelihood, in this paper we first derive a closed form formula for the covariance of the reverse process. Then, we propose a method based on finite difference method to approximate this covariance such that it can be readily obtained from the existing pretrained DDPMs, thereby not increasing the complexity compared to existing approaches. Finally, based on the mean and approximated covariance of the reverse process, we present a new approximation to the likelihood. We refer to this method as covariance-aware diffusion posterior sampling (CA-DPS). Experimental results show that CA-DPS significantly improves reconstruction performance without requiring hyperparameter tuning. The code for the paper is put in the supplementary materials.

Enhancing Diffusion Models for Inverse Problems with Covariance-Aware Posterior Sampling

TL;DR

This work addresses solving linear inverse problems with pre-trained unconditional DDPM priors by improving posterior sampling. It derives a closed-form expression for the conditional posterior covariance Cov(x0|xt) in DDPMs and introduces a practical, covariance-aware approximation via a diagonal Hessian estimated with finite differences, yielding Σ̃_t. By combining the mean x̃0 from Tweedie’s formula with the covariance Σ̃_t, CA-DPS approximates p_t(x0|xt) as a Gaussian and computes a refined likelihood p_t(y|xt), enabling more accurate posterior sampling without retraining or hyperparameter tuning. Empirical results on FFHQ and ImageNet show CA-DPS consistently outperforms existing methods across inpainting, deblurring, and super-resolution, and a toy dataset demonstrates better posterior fidelity to the true posterior. The approach is designed to be easily applicable to pretrained DDPMs and scalable via conjugate gradient solves, making covariance-aware diffusion a practical improvement for inverse problems.

Abstract

Inverse problems exist in many disciplines of science and engineering. In computer vision, for example, tasks such as inpainting, deblurring, and super resolution can be effectively modeled as inverse problems. Recently, denoising diffusion probabilistic models (DDPMs) are shown to provide a promising solution to noisy linear inverse problems without the need for additional task specific training. Specifically, with the prior provided by DDPMs, one can sample from the posterior by approximating the likelihood. In the literature, approximations of the likelihood are often based on the mean of conditional densities of the reverse process, which can be obtained using Tweedie formula. To obtain a better approximation to the likelihood, in this paper we first derive a closed form formula for the covariance of the reverse process. Then, we propose a method based on finite difference method to approximate this covariance such that it can be readily obtained from the existing pretrained DDPMs, thereby not increasing the complexity compared to existing approaches. Finally, based on the mean and approximated covariance of the reverse process, we present a new approximation to the likelihood. We refer to this method as covariance-aware diffusion posterior sampling (CA-DPS). Experimental results show that CA-DPS significantly improves reconstruction performance without requiring hyperparameter tuning. The code for the paper is put in the supplementary materials.
Paper Structure (18 sections, 3 theorems, 37 equations, 4 figures, 4 tables, 1 algorithm)

This paper contains 18 sections, 3 theorems, 37 equations, 4 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Diffusion Models for Inverse Problems
  4. Higher Order Approximation of reverse process
  5. Background and Preliminaries
  6. Diffusion Models
  7. Diffusion Models for Solving Inverse Problems
  8. Covariance-Aware Diffusion Posterior Sampling
  9. Experiments
  10. Quantitative Results.
  11. Qualitative Results
  12. Toy Dataset
  13. Conclusion
  14. Anonymous Repo of the Project
  15. Proof of Theorem 1
  16. ...and 3 more sections

Key Result

Proposition 1

Let $p({\boldsymbol{y}}|{\boldsymbol{\eta}})$ belong to the exponential family distribution where ${\boldsymbol{\eta}}$ is the canonical vector of the family, $T({\boldsymbol{y}})$ is some function of $\boldsymbol{y}$, and $\varphi({\boldsymbol{\eta}})$ is the cumulant generation function which normalizes the density, and $p_0({\boldsymbol{y}})$ is the density up to the scale factor when $

Figures (4)

  • Figure 1: Qualitative results on FFHQ dataset.
  • Figure 2: Visualizing the first two dimensions of the estimated posterior distributions for the configuration ($d=80$, $m=1$, $\sigma = 10^{-1}$) for a randomly generated $\bold{A}$.
  • Figure 3: Qualitative results on FFHQ dataset.
  • Figure 4: Qualitative results on FFHQ dataset.

Theorems & Definitions (5)

  • Proposition 1: Tweedie's formula
  • Theorem 1
  • proof
  • Corollary
  • proof