Gradient Step Denoiser for convergent Plug-and-Play

1 INTRODUCTION . Image restoration ( IR ) problems can be formulated as inverse problems of the form x∗ ∈ arg min x f ( x ) + λg ( x ) ( 1 ) where f is a term measuring the fidelity to a degraded observation y , and g is a regularization term weighted by a parameter λ ≥ 0 . Generally , the degradation of a clean image x̂ can be modeled by a linear operation y = Ax̂+ ξ , where A is a degradation matrix and ξ a white Gaussian noise . In this context , the maximum a posteriori ( MAP ) derivation relates the data-fidelity term to the likelihood f ( x ) = − log p ( y|x ) = 12σ2 ||Ax− y|| 2 , while the regularization term is related to the chosen prior . Regularization is crucial since it tackles the ill-posedness of the IR task by bringing a priori knowledge on the solution . A lot of research has been dedicated to designing accurate priors g. Among the most classical priors , one can single out total variation ( Rudin et al. , 1992 ) , wavelet sparsity ( Mallat , 2009 ) or patch-based Gaussian mixtures ( Zoran & Weiss , 2011 ) . Designing a relevant prior g is a difficult task and recent approaches rather apply deep learning techniques to directly learn a prior from a database of clean images ( Lunz et al. , 2018 ; Prost et al. , 2021 ; González et al. , 2021 ) . Generally , the problem ( 1 ) does not have a closed-form solution , and an optimization algorithm is required . First-order proximal splitting algorithms ( Combettes & Pesquet , 2011 ) operate individually on f and g via the proximity operator Proxf ( x ) = arg min z 1 2 ||x− z||2 + f ( z ) . ( 2 ) Among them , half-quadratic splitting ( HQS ) ( Geman & Yang , 1995 ) alternately applies the proximal operators of f and g. Proximal methods are particularly useful when either f or g is nonsmooth . Plug-and-Play ( PnP ) methods ( Venkatakrishnan et al. , 2013 ) build on proximal splitting algorithms by replacing the proximity operator of g with a generic denoiser , e.g . a pretrained deep network . ∗Corresponding author : samuel.hurault @ math.u-bordeaux.fr These methods achieve state-of-the-art results ( Buzzard et al. , 2018 ; Ahmad et al. , 2020 ; Yuan et al. , 2020 ; Zhang et al. , 2021 ) in various IR problems . However , since a generic denoiser can not generally be expressed as a proximal mapping ( Moreau , 1965 ) , convergence results , which stem from the properties of the proximal operator , are difficult to obtain . Moreover , the regularizer g is only made implicit via the denoising operation . Therefore , PnP algorithms do not seek the minimization of an explicit objective functional which strongly limits their interpretation and numerical control . In order to keep tractability of a minimization problem , Romano et al . ( 2017 ) proposed , with regularization by denoising ( RED ) , an explicit prior g that exploits a given generic denoiser D in the form g ( x ) = 12 〈x , x −D ( x ) 〉 . With strong assumptions on the denoiser ( in particular a symmetric Jacobian assumption ) , they show that it verifies ∇xg ( x ) = x−D ( x ) . ( 3 ) Such a denoiser is then plugged in gradient-based minimization schemes . Despite having shown very good results on various image restoration tasks , as later pointed out by Reehorst & Schniter ( 2018 ) or Saremi ( 2019 ) , existing deep denoisers lack Jacobian symmetry . Hence , RED does not minimize an explicit functional and is not guaranteed to converge . Contributions . In this work , we develop a PnP scheme with novel theoretical convergence guarantees and state-of-the-art IR performance . Departing from the PnP-HQS framework , we plug a denoiser that inherently satisfies equation ( 3 ) without sacrificing the denoising performance . The resulting fixed-point algorithm is guaranteed to converge to a stationary point of an explicit functional . This convergence guarantee does not require strong convexity of the data-fidelity term , thus encompassing ill-posed IR tasks like deblurring , super-resolution or inpainting . 2 RELATED WORKS . PnP methods have been successfully applied in the literature with various splitting schemes : HQS ( Zhang et al. , 2017b ; 2021 ) , ADMM ( Romano et al. , 2017 ; Ryu et al. , 2019 ) , Proximal Gradient Descent ( PGD ) ( Terris et al. , 2020 ) . First used with classical non deep denoisers such as BM3D ( Chan et al. , 2016 ) and pseudo-linear denoisers ( Nair et al. , 2021 ; Gavaskar et al. , 2021 ) , more recent PnP approaches ( Meinhardt et al. , 2017 ; Ryu et al. , 2019 ) rely on efficient off-theshelf deep denoisers such as DnCNN ( Zhang et al. , 2017a ) . State-of-the-art IR results are currently obtained with denoisers that are specifically designed to be integrated in PnP schemes , like IRCNN ( Zhang et al. , 2017b ) or DRUNET ( Zhang et al. , 2021 ) . Though providing excellent restorations , such schemes are not guaranteed to converge for all kinds of denoisers or IR tasks . Designing convergence proofs for PnP algorithms is an active research topic . Sreehari et al . ( 2016 ) used the proximal theorem of Moreau ( Moreau , 1965 ) to give sufficient conditions for the denoiser to be an explicit proximal map , which are applied to a pseudo-linear denoiser . The convergence with pseudo-linear denoisers have been extensively studied ( Gavaskar & Chaudhury , 2020 ; Nair et al. , 2021 ; Chan , 2019 ) . However , state-of-the-art PnP results are obtained with deep denoisers . Various assumptions have been made to ensure the convergence of the related PnP schemes . With a “ bounded denoiser ” assumption , Chan et al . ( 2016 ) ; Gavaskar & Chaudhury ( 2019 ) showed convergence of PnP-ADMM with stepsizes decreasing to 0 . RED ( Romano et al. , 2017 ) and RED-PRO ( Cohen et al. , 2021 ) respectively consider the classes of denoisers with symmetric Jacobian or demicontractive mappings , but these conditions are either too restrictive or hard to verify in practice . In Appendix A.3 , more details are given on RED-based methods . Many works focus on Lipschitz properties of PnP operators . Depending on the splitting algorithm in use , convergence can be obtained by assuming the denoiser averaged ( Sun et al. , 2019b ) , firmly nonexpansive ( Sun et al. , 2021 ; Terris et al. , 2020 ) or simply nonexpansive ( Reehorst & Schniter , 2018 ; Liu et al. , 2021 ) . These settings are unrealistic as deep denoisers do not generally satisfy such properties . Ryu et al . ( 2019 ) ; Terris et al . ( 2020 ) propose different ways to train deep denoisers with constrained Lipschitz constants , in order to fit the technical properties required for convergence . But imposing hard Lipschitz constraints on the network alters its denoising performance ( Bohra et al. , 2021 ; Hertrich et al. , 2021 ) . Yet , Ryu et al . ( 2019 ) manages to get a convergent PnP scheme without assuming the nonexpansiveness ofD . This comes at the cost of imposing strong convexity on the data-fidelity term f , which excludes many IR tasks like deblurring , super-resolution or inpainting . Hence , given the ill-posedness of IR problems , looking for a unique solution via contractive operators is a restrictive assumption . In this work , we do not impose contractiveness , but still obtain convergence results with realistic hypotheses . One can relate the ideal deep denoiser to the “ true ” natural image prior p via Tweedie ’ s Identity . In ( Efron , 2011 ) , it is indeed shown that the Minimum Mean Square Error ( MMSE ) denoiser D∗σ ( at noise level σ ) verifies Dσ ( x ) = x+ σ2∇x log pσ ( x ) where pσ is the convolution of p with the density ofN ( 0 , σ2 Id ) . In a recent line of research ( Bigdeli et al. , 2017 ; Xu et al. , 2020 ; Laumont et al. , 2021 ; Kadkhodaie & Simoncelli , 2020 ) , this relation is used to plug a denoiser in gradient-based dynamics . In practice , the MMSE denoiser can not be computed explicitly and Tweedie ’ s Identity does not hold for deep approximations of the MMSE . In order to be as exhaustive as possible , we detailed the addressed limitations of existing PnP methods in Appendix A.1 . 3 THE GRADIENT STEP PLUG-AND-PLAY . The proposed method is based on the PnP version of half-quadratic-splitting ( PnP-HQS ) that amounts to replacing the proximity operator of the prior g with an off-the-shelf denoiser Dσ . In order to define a convergent PnP scheme , we first set up in Section 3.1 a Gradient Step ( GS ) denoiser . We then introduce the Gradient Step PnP ( GS-PnP ) algorithm in Section 3.2 . 3.1 GRADIENT STEP DENOISER . We propose to plug a denoising operator Dσ that takes the form of a gradient descent step Dσ = Id−∇gσ , ( 4 ) with gσ : Rn → R. Contrary to Romano et al . ( 2017 ) , our denoiser exactly represents a conservative vector field . The choice of the parameterization of gσ is fundamental for the denoising performance . As already noticed in Salimans & Ho ( 2021 ) , we experimentally found that directly modeling gσ as a neural network ( e.g . a standard network used for classification ) leads to poor denoising performance . In order to keep the strength of state-of-the-art unconstrained denoisers , we rather use gσ ( x ) = 1 2 ||x−Nσ ( x ) ||2 , ( 5 ) which leads to Dσ ( x ) = x−∇gσ ( x ) = Nσ ( x ) + JNσ ( x ) T ( x−Nσ ( x ) ) , ( 6 ) where Nσ : Rn → Rn is parameterized by a neural network and JNσ ( x ) is the Jacobian of Nσ at point x . As discussed in Appendix A.2 , the formulation ( 5 ) for gσ has been proposed in ( Romano et al. , 2017 , Section 5.2 ) and ( Bigdeli & Zwicker , 2017 ) for a distinct but related purpose , and not exploited for convergence analysis . Thanks to our definition ( 6 ) for Dσ , we can parameterize Nσ with any differentiable neural network architecture Rn → Rn that has proven efficient for image denoising . Although the representation power of the denoiser is limited by the particular form ( 6 ) , we show ( see Section 5.1 ) that such parameterization still yields state-of-the-art denoising results . We train the denoiser Dσ for Gaussian noise by minimizing the MSE loss function L ( Dσ ) = Ex∼p , ξσ∼N ( 0 , σ2I ) [ ||Dσ ( x+ ξσ ) − x|| 2 ] , ( 7 ) or L ( gσ ) = Ex∼p , ξσ∼N ( 0 , σ2I ) [ ||∇gσ ( x+ ξσ ) − ξσ|| 2 ] , ( 8 ) when written in terms of gσ using equation ( 4 ) . Remark 1 . By definition , the optimal solution g∗σ ∈ arg minL is related to the MMSE denoiserD∗σ , that is , the best non-linear predictor of x given x + ξσ . Therefore , it satisfies Tweedie ’ s formula and ∇g∗σ = −σ2∇ log pσ ( Efron , 2011 ) i.e . g∗σ = −σ2 log pσ + C , for some C ∈ R. Hence approximating the MMSE denoiser with a denoiser parameterized as ( 4 ) is related to approximating the logarithm of the smoothed image prior of pσ with − 1σ2 gσ . This relation was used for image generation with “ Denoising Score Matching ” by Saremi & Hyvarinen ( 2019 ) ; Bigdeli et al . ( 2020 ) .

Summary: This paper presents a new form of the plug-and-play (PnP) half-quadratic splitting algorithm with provable convergence.  In contrast to directly formulating the denoiser as previous works, they propose to parameterize the denoiser with a learnable score-based function. This builds a bridge between the recently emerged score-based generative model and plug-and-play methods. And even more surprisingly, they show such a new formulation within the PnP framework leads to a very strong theoretical convergence guarantee under mild realistic assumptions. They also demonstrate good empirical results on three image restoration tasks i.e., deblurring, super-resolution and inpainting, verifying the empirical convergence rate is usually faster than $\mathcal{O}(\frac{1}{k})$ -- the worst-case rate established theoretically.