DIFFUSION POSTERIOR SAMPLING FOR GENERAL NOISY INVERSE PROBLEMS

Abstract

Diffusion models have been recently studied as powerful generative inverse problem solvers, owing to their high quality reconstructions and the ease of combining existing iterative solvers. However, most works focus on solving simple linear inverse problems in noiseless settings, which significantly under-represents the complexity of real-world problems. In this work, we extend diffusion solvers to efficiently handle general noisy (non)linear inverse problems via approximation of the posterior sampling. Interestingly, the resulting posterior sampling scheme is a blended version of diffusion sampling with the manifold constrained gradient without a strict measurement consistency projection step, yielding a more desirable generative path in noisy settings compared to the previous studies. Our method demonstrates that diffusion models can incorporate various measurement noise statistics such as Gaussian and Poisson, and also efficiently handle noisy nonlinear inverse problems such as Fourier phase retrieval and non-uniform deblurring. Code is available at https: //github.com/DPS2022/diffusion-posterior-sampling.

1. INTRODUCTION

Diffusion models learn the implicit prior of the underlying data distribution by matching the gradient of the log density (i.e. Stein score; ∇ x log p(x)) (Song et al., 2021b) . The prior can be leveraged when solving inverse problems, which aim to recover x from the measurement y, related through the forward measurement operator A and the detector noise n. When we know such forward models, one can incorporate the gradient of the log likelihood (i.e. ∇ x log p(y|x)) in order to sample from the posterior distribution p(x|y). While this looks straightforward, the likelihood term is in fact analytically intractable in terms of diffusion models, due to their dependence on time t. Due to its intractability, one often resorts to projections onto the measurement subspace (Song et al., 2021b; Chung et al., 2022b; Chung & Ye, 2022; Choi et al., 2021) . However, the projection-type approach fails dramatically when 1) there is noise in the measurement, since the noise is typically amplified during the generative process due to the ill-posedness of the inverse problems; and 2) the measurement process is nonlinear. One line of works that aim to solve noisy inverse problems run the diffusion in the spectral domain (Kawar et al., 2021; 2022) so that they can tie the noise in the measurement domain into the spectral domain via singular value decomposition (SVD). Nonetheless, the computation of SVD is costly and even prohibitive when the forward model gets more complex. For example, Kawar et al. (2022) only considered seperable Gaussian kernels for deblurring, since they were restricted to the family of inverse problems where they could effectively perform the SVD. Hence, the applicability of such methods is restricted, and it would be useful to devise a method to solve noisy inverse problems without the computation of SVD. Furthermore, while diffusion models were applied to various inverse problems including inpainting (Song et al., 2021b; Chung et al., 2022b; Kawar et al., 2022; Chung et al., 2022a ), super-resolution (Choi et al., 2021; Chung et al., 2022b; Kawar et al., 2022) , colorization (Song et al., 2021b; Kawar et al., 2022; Chung et al., 2022a) , compressed-sensing MRI (CS-MRI) (Song et al., 2022; Chung & Ye, 2022; Chung et al., 2022b) , computed tomography (CT) (Song et al., 2022; Chung et al., 2022a) , etc., to our best knowledge, all works so far considered linear inverse problems only, and have not explored nonlinear inverse problems. In this work, we devise a method to circumvent the intractability of posterior sampling by diffusion models via a novel approximation, which can be generally applied to noisy inverse problems. Specifically, we show that our method can efficiently handle both the Gaussian and the Poisson measurement noise. Also, our framework easily extends to any nonlinear inverse problems, when the gradients can be obtained through automatic differentiation. We further reveal that a recently proposed method of manifold constrained gradients (MCG) (Chung et al., 2022a ) is a special case of the proposed method when the measurement is noiseless. With a geometric interpretation, we further show that the proposed method is more likely to yield desirable sample paths in noisy setting than the previous approach (Chung et al., 2022a). In addition, the proposed method fully runs on the image domain rather than the spectral domain, thereby avoiding the computation of SVD for efficient implementation. With extensive experiments including various inverse problems-inpainting, super-resolution, (Gaussian/motion/non-uniform) deblurring, Fourier phase retrieval-we show that our method serves as a general framework for solving general noisy inverse problems with superior quality (Representative results shown in Fig. 1 ).

2. BACKGROUND 2.1 SCORE-BASED DIFFUSION MODELS

Diffusion models define the generative process as the reverse of the noising process. Specifically, Song et al. (2021b) defines the Itô stochastic differential equation (SDE) for the data noising process (i.e. forward SDE) x(t), t ∈ [0, T ], x(t) ∈ R d ∀t in the following form 1 dx = - β(t) 2 xdt + β(t)dw, where β(t) : R → R > 0 is the noise schedule of the process, typically taken to be monotonically increasing linear function of t (Ho et al., 2020), and w is the standard d-dimensional Wiener process. The data distribution is defined when t = 0, i.e. x(0) ∼ p data , and a simple, tractable distribution (e.g. isotropic Gaussian) is achieved when t = T , i.e. x(T ) ∼ N (0, I). Our aim is to recover the data generating distribution starting from the tractable distribution, which can be achieved by writing down the corresponding reverse SDE of (1) (Anderson, 1982) : dx = - β(t) 2 x -β(t)∇ xt log p t (x t ) dt + β(t)d w,



In this work, we consider the variance preserving (VP) form of the SDE(Song et al., 2021b) which is equivalent to Denoising Diffusion Probabilistic Models (DDPM)(Ho et al., 2020).



Figure 1: Solving noisy linear, and nonlinear inverse problems with diffusion models. Our reconstruction results (right) from the measurements (left) are shown.

