Improving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance
Xinyu Peng, Ziyang Zheng, Wenrui Dai, Nuoqian Xiao, Chenglin Li, Junni Zou, Hongkai Xiong
TL;DR
This work reframes zero-shot diffusion-based solvers for inverse problems as using a Gaussian approximation to the denoising posterior with an isotropic covariance, and then optimizes the posterior covariance via maximum likelihood. It introduces plug-and-play strategies for both cases when reverse covariance is available and when it is not, including a scalable transform-domain variance model to capture pixel correlations. The approach unifies Type I (likelihood-score based) and Type II (proximal-based) guidance under a variational Gaussian posterior, derives fixed-point relations with DDPM reverse variances, and proposes Monte Carlo and transform-domain methods to estimate posteriors without retraining. Experiments on inpainting, deblurring, and super-resolution demonstrate significant improvements without hyperparameter tuning and without retraining, highlighting practical impact for diffusion-based inverse problem solvers. The work also discusses limitations of diagonal covariance and outlines directions for more expressive covariance designs and nonlinear transformation-based variance modeling.
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. We further propose a scalable method for learning posterior covariance prediction based on representation with orthonormal basis. Experimental results demonstrate that the proposed methods significantly enhance reconstruction performance without requiring hyperparameter tuning.