SNIPS: Solving Noisy Inverse Problems Stochastically

1 Introduction . Many problems in the field of image processing can be cast as noisy linear inverse problems . This family of tasks includes denoising , inpainting , deblurring , super resolution , compressive sensing , and many other image recovery problems . A general linear inverse problem is posed as y = Hx + z , ( 1 ) where we aim to recover a signal x from its measurement y , given through a linear degradation operator H and a contaminating noise , being additive , white and Gaussian , z ∼ N ( 0 , σ20I ) . In this work we assume that both H and σ0 are known . Over the years , many strategies , algorithms and underlying statistical models were developed for handling image restoration problems . A key ingredient in many of the classic attempts is the prior that aims to regularize the inversion process and lead to visually pleasing results . Among the various options explored , we mention sparsity-inspired techniques [ 13 , 55 , 11 ] , local Gaussian-mixture modeling [ 57 , 63 ] , and methods relying on non-local self-similarity [ 6 , 9 , 36 , 51 ] . More recently , and with the emergence of deep learning techniques , a direct design of the recovery path from y to an estimate of x took the lead , yielding state-of-the-art results in various linear inverse problems , such as denoising [ 25 , 59 , 61 , 52 ] , deblurring [ 22 , 48 ] , super resolution [ 10 , 17 , 54 ] and other tasks [ 29 , 28 , 19 , 16 , 37 , 58 ] . Despite the evident success of the above techniques , many image restoration algorithms still have a critical shortcoming : In cases of severe degradation , most recovery algorithms tend to produce washed out reconstructions that lack details . Indeed , most image restoration techniques seek a reconstruction that minimizes the mean squared error between the restored image , x̂ , and the unknown original one , x . When the degradation is acute and information is irreversibly lost , image reconstruction becomes a highly ill-posed problem , implying that many possible clean images could explain the given 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . measurements . The MMSE solution averages all these candidate solutions , being the conditional mean of the posterior of x given y , leading to an image with loss of fine details in the majority of practical cases . A recent work reported in [ 5 ] has shown that reconstruction algorithms necessarily suffer from a perception-distortion tradeoff , i.e. , targeting a minimization of the error between x̂ and x ( in any metric ) is necessarily accompanied by a compromised perceptual quality . As a consequence , as long as we stick to the tendency to design recovery algorithms that aim for minimum MSE ( or other distances ) , only a limited perceptual improvement can be expected . When perceptual quality becomes our prime objective , the strategy for solving inverse problems must necessarily change . More specifically , the solution should concentrate on producing a sample ( or many of them ) from the posterior distribution p ( x|y ) instead of its conditional mean . Recently , two such approaches have been suggested – GAN-based and Langevin sampling . Generative Adversarial Networks ( GANs ) have shown impressive results in generating realistically looking images ( e.g. , [ 14 , 35 ] ) . GANs can be utilized for solving inverse problems while producing high-quality images ( see e.g . [ 2 , 31 , 34 ] ) . These solvers aim to produce a diverse set of output images that are consistent with the measurements , while also being aligned with the distribution of clean examples . A major disadvantage of GAN-based algorithms for inverse problems is their tendency ( as practiced in [ 2 , 31 , 34 ] ) to assume noiseless measurements , a condition seldom met in practice . An exception to this is the work reported in [ 33 ] , which adapts a conditional GAN to become a stochastic denoiser . The second approach for sampling from the posterior , and the one we shall be focusing on in this paper , is based on Langevin dynamics . This core iterative technique enables sampling from a given distribution by leveraging the availability of the score function – the gradient of the log of the probability density function [ 38 , 3 ] . The work reported in [ 44 , 20 , 46 ] utilizes the annealed Langevin dynamics method , both for image synthesis and for solving noiseless inverse problems.1 Their synthesis algorithm relies on an MMSE Gaussian denoiser ( given as a neural network ) for approximating a gradually blurred score function . In their treatment of inverse problems , the conditional score remains tractable and manageable due to the noiseless measurements assumption . The question addressed in this paper is the following : How can the above line of Langevin-based work be generalized for handling linear inverse problems , as in Equation 1 , in which the measurements are noisy ? A partial and limited answer to this question has already been given in [ 21 ] for the tasks of image denoising and inpainting . The present work generalizes these ( [ 44 , 20 , 46 , 21 ] ) results , and introduces a systematic way for sampling from the posterior distribution of any given noisy linear inverse problem . As we carefully show , this extension is far from being trivial , due to two prime reasons : ( i ) The involvement of the degradation operator H , which poses a difficulty for establishing a relationship between the reconstructed image and the noisy observation ; and ( ii ) The intricate connection between the measurements ’ and the synthetic annealed Langevin noise . Our proposed remedy is a decorrelation of the measurements equation via a singular value decomposition ( SVD ) of the operator H , which decouples the dependencies between the measurements , enabling each to be addressed by an adapted iterative process . In addition , we define the annealing noise to be built as portions of the measurement noise itself , in a manner that facilitates a constructive derivation of the conditional score function . Following earlier work [ 44 , 20 , 46 , 21 ] , our algorithm is initialized with a random noise image , gradually converging to the reconstructed result , while following the direction of the log-posterior gradient , estimated using an MMSE denoiser . Via a careful construction of the gradual annealing noise sequence , from very high values to low ones , the entries in the derived score switch mode . Those referring to non-zero singular values start by being purely dependent on the measurements , and then transition to incorporate prior information based on the denoiser . As for entries referring to zero singular values , their corresponding entries undergo a pure synthesis process based on the prior-only score function . Note that the denoiser blends values in the evolving sample , thus intermixing the influence of the gradient entries . Our derivations include an analytical expression for a positiondependent step size vector , drawing inspiration from Newton ’ s method in optimization . This stabilizes the algorithm and is shown to be essential for its success . We refer hereafter to our algorithm as SNIPS ( Solution of Noisy Inverse Problems Stochastically ) . Observe that as we target to sample from the posterior distribution p ( x|y ) , different runs of SNIPS on the same input necessarily yield different results , all of which valid solutions to the given inverse 1The work reported in [ 18 , 42 , 26 ] and [ 15 , 23 ] is also relevant to this discussion , but somewhat different . We shall specifically address these papers ’ content and its relation to our work in section 2. problem . This should not come as a surprise , as ill-posedness implies that there are multiple viable solutions for the same data , as has already been suggested in the context of super resolution [ 31 , 2 , 34 ] . We demonstrate SNIPS on image deblurring , single image super resolution , and compressive sensing , all of which contain non-negligible noise , and emphasize the high perceptual quality of the results , their diversity , and their relation to the MMSE estimate . To summarize , this paper ’ s contributions are threefold : • We present an intricate derivation of the blurred posterior score function for general noisy inverse problems , where both the measurement and the target image contain delicately inter-connected additive white Gaussian noise . • We introduce a novel stochastic algorithm – SNIPS – that can sample from the posterior distribution of these problems . The algorithm relies on the availability of an MMSE denoiser . • We demonstrate impressive results of SNIPS on image deblurring , single image super resolution , and compressive sensing , all of which are highly noisy and ill-posed . Before diving into the details of this work , we should mention that using Gaussian denoisers iteratively for handling general linear inverse problems has been already proposed in the context of the Plugand-Play-Prior ( PnP ) method [ 53 ] and RED [ 39 ] , and their many followup papers ( e.g. , [ 60 , 30 , 1 , 49 , 7 , 50 , 40 , 4 ] ) . However , both PnP and RED are quite different from our work , as they do not target sampling from the posterior , but rather focus on MAP or MMSE estimation . 2 Background . The Langevin dynamics algorithm [ 3 , 38 ] suggests sampling from a probability distribution p ( x ) using the iterative transition rule xt+1 = xt + α∇xt log p ( xt ) + √ 2αzt , ( 2 ) where zt ∼ N ( 0 , I ) and α is an appropriately chosen small constant . The added zt allows for stochastic sampling , avoiding a collapse to a maximum of the distribution . Initialized randomly , after a sufficiently large number of iterations , and under some mild conditions , this process converges to a sample from the desired distribution p ( x ) [ 38 ] . The work reported in [ 44 ] extends the aforementioned algorithm into annealed Langevin dynamics . The annealing proposed replaces the score function in Equation 2 with a blurred version of it , ∇x̃t log p ( x̃t ) , where x̃t = xt + n and n ∼ N ( 0 , σ2I ) is a synthetically injected noise . The core idea is to start with a very high noise level σ and gradually drop it to near-zero , all while using a step size α dependent on the noise level . These changes allow the algorithm to converge much faster and perform better , because it widens the basin of attraction of the sampling process . The work in [ 20 ] further develops this line of work by leveraging a brilliant relation attributed to Miyasawa [ 32 ] ( also known as Stein ’ s integration by parts trick [ 47 ] or Tweedie ’ s identity [ 12 ] ) . It is given as ∇x̃t log p ( x̃t ) = D ( x̃t , σ ) − x̃t σ2 , ( 3 ) where D ( x̃t , σ ) = E [ x|x̃t ] is the minimizer of the MSE measure E [ ‖x−D ( x̃t , σ ) ‖22 ] , which can be approximated using a denoising neural network . This facilitates the use of denoisers in Langevin dynamics as a replacement for the evasive score function . When turning to solve inverse problems , previous work suggests sampling from the posterior distribution p ( x|y ) using annealed Langevin dynamics [ 20 , 46 , 21 ] or similar methods [ 15 , 18 , 42 , 26 ] , by replacing the score function used in the generation algorithm with a conditional one . As it turns out , if limiting assumptions can be posed on the measurements formation , the conditional score is tractable , and thus generalization of the annealed Langevin process to these problems is within reach . Indeed , in [ 44 , 20 , 46 , 42 , 26 ] the core assumption is y = Hx for specific and simplified choices of H and with no noise in the measurements . The works in [ 15 , 23 ] avoid these difficulties altogether by returning to the original ( non-annealed ) Langevin method , with the unavoidable cost of becoming extremely slow . In addition , their algorithms are demonstrated on inverse problems in which the additive noise is restricted to be very weak . The work in [ 21 ] is broader , allowing for an arbitrary additive white Gaussian noise , but limits H to the problems of denoising or inpainting . While all these works demonstrate high quality results , there is currently no clear way for deriving the blurred score function of a general linear inverse problem as posed in Equation 1 . In the following , we present such a derivation .

Summary: This paper proposes a technique based on annealed Langevin dynamics to solve noisy linear inverse problems. The technique is built on top of [18, 19] and proposes to perform the Langevin dynamics in the spectral space of the linear degradation operator $H$. To improve the performance of the proposed method, the authors use coordinate wise step sizes obtained by approximating the hessian with its diagonal. Qualitative results from the proposed method are presented for compress sensing, inpainting, super-résolution and deblurring.