Dynamically Unfolding Recurrent Restorer: A Moving Endpoint Control Method for Image Restoration

1 INTRODUCTION . Image restoration , including image denoising , deblurring , inpainting , etc. , is one of the most important areas in imaging science . Its major purpose is to obtain high quality reconstructions of images corrupted in various ways during imaging , acquisiting , and storing , and enable us to see crucial but subtle objects that reside in the images . Image restoration has been an active research area . Numerous models and algorithms have been developed for the past few decades . Before the uprise of deep learning methods , there were two classes of image restoration approaches that were widely adopted in the field : transformation based approach and PDE approach . The transformation based approach includes wavelet and wavelet frame based methods ( Elad et al. , 2005 ; Starck et al. , 2005 ; Daubechies et al. , 2007 ; Cai et al. , 2009 ) , dictionary learning based methods ( Aharon et al. , 2006 ) , similarity based methods ( Buades et al. , 2005 ; Dabov et al. , 2007 ) , low-rank models ( Ji et al. , 2010 ; Gu et al. , 2014 ) , etc . The PDE approach includes variational models ( Mumford & Shah , 1989 ; Rudin et al. , 1992 ; Bredies et al. , 2010 ) , nonlinear diffusions ( Perona & Malik , 1990 ; Catté et al. , 1992 ; Weickert , 1998 ) , nonlinear hyperbolic equations ( Osher & Rudin , 1990 ) , etc . More recently , deep connections ∗Equal contribution . 1Supplementary materials and full code can be found at the project page . between wavelet frame based methods and PDE approach were established ( Cai et al. , 2012 ; 2016 ; Dong et al. , 2017 ) . One of the greatest challenge for image restoration is to properly handle image degradations of different levels . In the existing transformation based or PDE based methods , there is always at least one tuning parameter ( e.g . the regularization parameter for variational models and terminal time for nonlinear diffusions ) that needs to be manually selected . The choice of the parameter heavily relies on the degradation level . Recent years , deep learning models for image restoration tasks have significantly advanced the state-of-the-art of the field . Jain & Seung ( 2009 ) proposed a convolutional neural network ( CNN ) for image denoising which has better expressive power than the MRF models by Lan et al . ( 2006 ) . Inspired by nonlinear diffusions , Chen & Pock ( 2017 ) designed a deep neural network for image denoising and Zhang et al . ( 2017a ) improves the capacity by introducing a deeper neural network with residual connections . Chen et al . ( 2017 ) use the CNN to simulate a wide variety of image processing operators , achieving high efficiencies with little accuracy drop . However , these models can not gracefully handle images with varied degradation levels . Although one may train different models for images with different levels , this may limit the application of these models in practice due to lack of flexibility . Taking blind image denoising for example . Zhang et al . ( 2017a ) designed a 20-layer neural network for the task , called DnCNN-B , which had a huge number of parameters . To reduce number of parameters , Lefkimmiatis ( 2017 ) proposed the UNLNet5 , by unrolling a projection gradient algorithm for a constrained optimization model . However , Lefkimmiatis ( 2017 ) also observed a drop in PSNR comparing to DnCNN . Therefore , the design of a light-weighted and yet effective model for blind image denoising remains a challenge . Moreover , deep learning based models trained on simulated gaussian noise images usually fail to handle real world noise , as will be illustrated in later sections . Another example is JPEG image deblocking . JPEG is the most commonly used lossy image compression method . However , this method tend to introduce undesired artifacts as the compression rate increases . JPEG image deblocking aims to eliminate the artifacts and improve the image quality . Recently , deep learning based methods were proposed for JPEG deblocking ( Dong et al. , 2015 ; Zhang et al. , 2017a ; 2018 ) . However , most of their models are trained and evaluated on a given quality factor . Thus it would be hard for these methods to apply to Internet images , where the quality factors are usually unknown . In this paper , we propose a single image restoration model that can robustly restore images with varied degradation levels even when the degradation level is well outside of that of the training set . Our proposed model for image restoration is inspired by the recent development on the relation between deep learning and optimal control . The relation between supervised deep learning methods and optimal control has been discovered and exploited by Weinan ( 2017 ) ; Lu et al . ( 2018 ) ; Chang et al . ( 2017 ) ; Fang et al . ( 2017 ) . The key idea is to consider the residual block xn+1 = xn + f ( xn ) as an approximation to the continuous dynamics Ẋ = f ( X ) . In particular , Lu et al . ( 2018 ) ; Fang et al . ( 2017 ) demonstrated that the training process of a class of deep models ( e.g . ResNet by He et al . ( 2016 ) , PolyNet by Zhang et al . ( 2017b ) , etc . ) can be understood as solving the following control problem : min w ( L ( X ( T ) , y ) + ∫ τ 0 R ( w ( t ) , t ) dt ) s.t . Ẋ = f ( X ( t ) , w ( t ) ) , t ∈ ( 0 , τ ) ( 1 ) X ( 0 ) = x0 . Here x0 is the input , y is the regression target or label , Ẋ = f ( X , w ) is the deep neural network with parameter w ( t ) , R is the regularization term and L can be any loss function to measure the difference between the reconstructed images and the ground truths . In the context of image restoration , the control dynamic Ẋ = f ( X ( t ) , ω ( t ) ) , t ∈ ( 0 , τ ) can be , for example , a diffusion process learned using a deep neural network . The terminal time τ of the diffusion corresponds to the depth of the neural network . Previous works simply fixed the depth of the network , i.e . the terminal time , as a fixed hyper-parameter . However Mrázek & Navara ( 2003 ) showed that the optimal terminal time of diffusion differs from image to image . Furthermore , when an image is corrupted by higher noise levels , the optimal terminal time for a typical noise removal diffusion should be greater than when a less noisy image is being processed . This is the main reason why current deep models are not robust enough to handle images with varied noise levels . In this paper , we no longer treat the terminal time as a hyper-parameter . Instead , we design a new architecture ( see Fig . 3 ) that contains both a deep diffusion-like network and another network that determines the optimal terminal time for each input image . We propose a novel moving endpoint control model to train the aforementioned architecture . We call the proposed architecture the dynamically unfolding recurrent restorer ( DURR ) . We first cast the model in the continuum setting . Let x0 be an observed degraded image and y be its corresponding damage-free counterpart . We want to learn a time-independent dynamic system Ẋ = f ( X ( t ) , w ) with parameters w so that X ( 0 ) = x and X ( τ ) ≈ y for some τ > 0 . See Fig . 2 for an illustration of our idea . The reason that we do not require X ( τ ) = y is to avoid over-fitting . For varied degradation levels and different images , the optimal terminal time τ of the dynamics may vary . Therefore , we need to include the variable τ in the learning process as well . The learning of the dynamic system and the terminal time can be gracefully casted as the following moving endpoint control problem : min w , τ ( x ) L ( X ( τ ) , y ) + ∫ τ ( x ) 0 R ( w ( t ) , t ) dt s.t . Ẋ = f ( X ( t ) , w ( t ) ) , t ∈ ( 0 , τ ( x ) ) ( 2 ) X ( 0 ) = x . Different from the previous control problem , in our model the terminal time τ is also a parameter to be optimized and it depends on the data x . The dynamic system Ẋ = f ( X ( t ) , w ) is modeled by a recurrent neural network ( RNN ) with a residual connection , which can be understood as a residual network with shared weights ( Liao & Poggio , 2016 ) . We shall refer to this RNN as the restoration unit . In order to learn the terminal time of the dynamics , we adopt a policy network to adaptively determine an optimal stopping time . Our learning framework is demonstrated in Fig . 3 . We note that the above moving endpoint control problem can be regarded as the penalized version of the well-known fixed endpoint control problem in optimal control ( Evans , 2005 ) , where instead of penalizing the difference between X ( τ ) and y , the constraint X ( τ ) = y is strictly enforced . In short , we summarize our contribution as following : • We are the first to use convolutional RNN for image restoration with unknown degradation levels , where the unfolding time of the RNN is determined dynamically at run-time by a policy unit ( could be either handcrafted or RL-based ) . • The proposed model achieves state-of-the-art performances with significantly less parameters and better running efficiencies than some of the state-of-the-art models . • We reveal the relationship between the generalization power and unfolding time of the RNN by extensive experiments . The proposed model , DURR , has strong generalization to images with varied degradation levels and even to the degradation level that is unseen by the model during training ( Fig . 1 ) . • The DURR is able to well handle real image denoising without further modification . Quali- tative results have shown that our processed images have better visual quality , especially sharper details compared to others . 2 METHOD . The proposed architecture , i.e . DURR , contains an RNN ( called the restoration unit ) imitating a nonlinear diffusion for image restoration , and a deep policy network ( policy unit ) to determine the terminal time of the RNN . In this section , we discuss the training of the two components based on our moving endpoint control formulation . As will be elaborated , we first train the restoration unit to determine ω , and then train the policy unit to estimate τ ( x ) . 2.1 TRAINING THE RESTORATION UNIT . If the terminal time τ for every input xi is given ( i.e . given a certain policy ) , the restoration unit can be optimized accordingly . We would like to show in this section that the policy used during training greatly influences the performance and the generalization ability of the restoration unit . More specifically , a restoration unit can be better trained by a good policy . The simplest policy is to fix the loop time τ as a constant for every input . We name such policy as “ naive policy ” . A more reasonable policy is to manually assign an unfolding time for each degradation level during training . We shall call this policy the “ refined policy ” . Since we have not trained the policy unit yet , to evaluate the performance of the trained restoration units , we manually pick the output image with the highest PSNR ( i.e . the peak PSNR ) . We take denoising as an example here . The peak PSNRs of the restoration unit trained with different policies are listed in Table . 1 . Fig . 4 illustrates the average loop times when the peak PSNRs appear . The training is done on both single noise level ( σ = 40 ) and multiple noise levels ( σ = 35 , 45 ) . For the refined policy , the noise levels and the associated loop times are ( 35 , 6 ) , ( 45 , 9 ) . For the naive policy , we always fix the loop times to 8 . As we can see , the refined policy brings the best performance on all the noise levels including 40 . The restoration unit trained for specific noise level ( i.e . σ = 40 ) is only comparable to the one with refined policy on noise level 40 . The restoration unit trained on multiple noise levels with naive policy has the worst performance . These results indicate that the restoration unit has the potential to generalize on unseen degradation levels when trained with good policies . According to Fig . 4 , the generalization reflects on the loop times of the restoration unit . It can be observed that the model with steeper slopes have stronger ability to generalize as well as better performances . According to these results , the restoration unit we used in DURR is trained using the refined policy . More specifically , for image denoising , the noise level and the associated loop times are set to ( 25 , 4 ) , ( 35 , 6 ) , ( 45 , 9 ) , and ( 55 , 12 ) . For JPEG image deblocking , the quality factor ( QF ) and the associated loop times are set to ( 20 , 6 ) and ( 30 , 4 ) .

Summary: The authors proposes a new image restoration method that becomes particularly useful for blind restoration setting, e.g., the unknown noise variance setting for denoising. They utilized the moving endpoint control methodology, which essentially is applying reinforcement learning to the image restoration, and devised a method that adaptively decides the unfolding steps for given noisy image. The experimental results show encouraging results.