Go to ICLR 2024 Conference homepageImproving Diffusion Models for Inverse Problems Using Optimal Posterior Covariance
Recent diffusion models provide a promising alternative zero-shot solution to noisy linear inverse problems without retraining for specific inverse problems. In this paper, we propose the first unified framework for diffusion-based zero-shot methods from the view of approximating conditional posterior mean for the reverse process. We reveal that recent diffusion-based zero-shot methods are equivalent to making isotropic Gaussian approximation to intractable posterior distributions over clean images given diffused noisy images, with only difference in handcrafted design of isotropic posterior covariances. Inspired by this finding, we develop the optimal posterior covariance of the posterior distribution via maximum likelihood estimation. We provide a general solution based on three approaches specifically designed for posterior covariance optimization, by training from scratch and using pre-trained models with and without reverse covariances. Remarkably, the proposed framework can be achieved in a plug-and-play fashion based on pre-trained unconditional diffusion models by converting reverse covariances or via Monte Carlo estimation without reverse covariances. Experimental results demonstrate that the proposed framework significantly outperforms existing zero-shot methods and enhances the robustness to hyper-parameters.