PSEUDOINVERSE-GUIDED DIFFUSION MODELS FOR INVERSE PROBLEMS

Abstract

Diffusion models have become competitive candidates for solving various inverse problems. Models trained for specific inverse problems work well but are limited to their particular use cases, whereas methods that use problem-agnostic models are general but often perform worse empirically. To address this dilemma, we introduce Pseudoinverse-guided Diffusion Models (ΠGDM), an approach that uses problem-agnostic models to close the gap in performance. ΠGDM directly estimates conditional scores from the measurement model of the inverse problem without additional training. It can address inverse problems with noisy, non-linear, or even non-differentiable measurements, in contrast to many existing approaches that are limited to noiseless linear ones. We illustrate the empirical effectiveness of ΠGDM on several image restoration tasks, including super-resolution, inpainting and JPEG restoration. On ImageNet, ΠGDM is competitive with state-of-the-art diffusion models trained on specific tasks, and is the first to achieve this with problem-agnostic diffusion models. ΠGDM can also solve a wider set of inverse problems where the measurement processes are composed of several simpler ones.



(Top) Problem-agnostic diffusion models perform an iterative denoising operation to produce random samples. (Bottom) ΠGDM utilizes problem-agnostic diffusion models to solve inverse problems, a key component of which is pseudoinverse guidance (ΠG). ΠG converts the problem-agnostic score function into a problem-specific one, using information about the measurements y and measurement model, denoted as h here (h is JPEG compression + masking in this figure, best viewed zoomed in). The additional guidance term is a vector-Jacobian product (VJP) that encourages consistency between the denoising result and the measurements, after a pseudoinverse transformation h † . ΠGDM applies the denoising process from ΠG in an iterative fashion to generate valid solutions to the inverse problem. Most methods that solve inverse problems with diffusion models fall into one of the two paradigms. In the first paradigm, one trains a problem-specific, conditional diffusion model that is limited to specific inverse problems, such as super-resolution (Saharia et al., 2021; Whang et al., 2021; Saharia et al., 2022a) . In the second paradigm, one uses problem-agnostic diffusion models that are trained for generative modeling but not train on any specific inverse problem; solutions are obtained via a "plug-and-play" approach that combines the diffusion model and the measurement process, e.g., via Bayes' rule (Venkatakrishnan et al., 2013a; Bardsley, 2012; Laumont et al., 2022; Choi et al., 2021; Song et al., 2021b; Jalal et al., 2021; Chung et al., 2021; Kawar et al., 2021; 2022a; Chung et al., 2022b; Daras et al., 2022a) . These methods can easily adapt to different tasks without re-training the diffusion model but tend to perform worse than problem-specific diffusion models. To achieve the best of both worlds, we introduce pseudoinverse guidance (ΠG), which uses problemagnostic diffusion models to reach the empirical performance of problem-specific ones. Conditioned on the measurements and an explicit measurement model, ΠG estimates the problem-specific score function via Bayes' rule and uses these scores to draw samples. However, unlike classifier/classifierfree guidance (Dhariwal & Nichol, 2021; Ho & Salimans, 2022) , ΠG obtains the problem-specific score directly via the known measurement model, without training additional models. Intuitively, ΠG guides the diffusion process by matching the one-step denoising solution and the ground-truth measurements, after transforming both via a "pseudoinverse" of the measurement model (see Fig. 1 ). This perspective allows ΠG to be the first guidance-based approach for inverse problem solving that handles measurements with Gaussian noise, as well as some non-linear, non-differentiable measurement models, such as JPEG compression (Kawar et al., 2022b) . We evaluate our method, termed Pseudoinverse-Guided Diffusion Models (ΠGDM), on various inverse problems, such as super-resolution, inpainting, and JPEG restoration over ImageNet validation images, and show that it achieves similar performance when compared against state-of-the-art taskspecific diffusion models (Saharia et al., 2021; Dhariwal & Nichol, 2021; Saharia et al., 2022a) . To the best of our knowledge, ΠGDM is the first approach based on problem-agnostic models to achieve this quality on ImageNet. We further apply ΠGDM to a wider range of inverse problems, where the measurement process is composed of different types of measurements. This allows us to easily solve a much wider set of problems, including ones have never been solved with diffusion models (see Fig. 2 ), such as low-resolution + JPEG compression + masking.



Figure1: High-level illustration of ΠGDM. (Top) Problem-agnostic diffusion models perform an iterative denoising operation to produce random samples. (Bottom) ΠGDM utilizes problem-agnostic diffusion models to solve inverse problems, a key component of which is pseudoinverse guidance (ΠG). ΠG converts the problem-agnostic score function into a problem-specific one, using information about the measurements y and measurement model, denoted as h here (h is JPEG compression + masking in this figure, best viewed zoomed in). The additional guidance term is a vector-Jacobian product (VJP) that encourages consistency between the denoising result and the measurements, after a pseudoinverse transformation h † . ΠGDM applies the denoising process from ΠG in an iterative fashion to generate valid solutions to the inverse problem.

Figure 2: ΠGDM applies a single problem-agnostic diffusion model for various inverse problems, avoiding the cost of training multiple problem-specific ones. Best viewed zoomed in.

