Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance
Xinyu Peng, Ziyang Zheng, Wenrui Dai, Nuoqian Xiao, Chenglin Li, Junni Zou, Hongkai Xiong
https://arxiv.org/abs/2402.02149v2

Abstract: Recent diffusion models provide a promising zero-shot solution to noisy linear inverse problems without retraining for specific inverse problems. In this paper, we reveal that recent methods can be uniformly interpreted as employing a Gaussian approximation with hand-crafted isotropic covariance for the intractable denoising posterior to approximate the conditional posterior mean. Inspired by this finding, we propose to improve recent methods by using more principled covariance determined by maximum likelihood estimation. To achieve posterior covariance optimization without retraining, we provide general plug-and-play solutions based on two approaches specifically designed for leveraging pre-trained models with and without reverse covariance. In addition, we propose a scalable method for learning posterior covariance prediction by leveraging widely-used orthonormal basis for image processing. Experimental results demonstrate that the proposed methods significantly enhance the overall performance and eliminate the need for hyperparameter tuning.

This code is based on:

• K diffusion: Provide the code structure.

• DPS: Provide the code for degradation opterators.

• DiffPIR: Provide tools for implementing closed-form solutions.

• GPyTorch: Provide tools for implementing differentiable Gaussian likelihoods, enabling auto-computed Type I guidance (Use --guidance autoI).

For creating the conda environment and installing dependencies run

conda env create -f environment.yml

Then activate the environment by

conda activate k-diffusion

From the link, download the FFHQ checkpoint ffhq_10m.pt, rename to diffusion_ffhq_10m.pt, and paste it to ../model_zoo.

To run guidance based on Analytic posterior covariance, download the precomputed Monte Carlo estimation from the link, and paste it to ./runs.

To run guidance based on DWT-Var posterior covariance, download the FFHQ checkpoint ffhq_dwt.ckpt from the link, and paste it to ../model_zoo.

From the link, download the validation data (the first 100 images from FFHQ and ImageNet datasets), unzip and paste it to ../data.

For reproducing results on FFHQ dataset in Table 2, run

bash quick_start/eval_guidance_I.sh ffhq # for Convert, Analytic, TMPD, DPS, PiGDM
bash quick_start/dwt_var/eval_guidance_I.sh 1 # for DWT-Var

For reproducing results on FFHQ dataset in Figure 3, run

bash quick_start/eval_guidance_diffpir.sh ffhq # for DiffPIR
bash quick_start/eval_guidance_II.sh ffhq # for PiGDM, Convert, Analytic
bash quick_start/dwt_var/eval_guidance_II.sh 1 # for DWT-Var

For reproducing results on FFHQ dataset in Table 3, run

bash quick_start/eval_complete_pgdm+mle.sh ffhq convert # for Convert
bash quick_start/eval_complete_pgdm+mle.sh ffhq analytic # for Analytic
bash quick_start/eval_complete_pgdm.sh ffhq # for PiGDM

For reproducing results on FFHQ dataset in Figure 4, run

bash quick_start/eval_complete_dps+mle.sh ffhq convert # for Convert
bash quick_start/eval_complete_dps+mle.sh ffhq analytic # for Analytic
bash quick_start/eval_complete_dps.sh ffhq # for DPS

If you find this repo helpful, please cite:

@inproceedings{
peng2024improving,
title={Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance},
author={Xinyu Peng and Ziyang Zheng and Wenrui Dai and Nuoqian Xiao and Chenglin Li and Junni Zou and Hongkai Xiong},
booktitle={Forty-first International Conference on Machine Learning},
year={2024},
url={https://openreview.net/forum?id=DrE7jVF4VW}
}