Fully Unsupervised Diversity Denoising with Convolutional Variational Autoencoders

1 INTRODUCTION . The goal of scientific image analysis is to analyze pixel-data and measure the properties of objects of interest in images . Pixel intensities are subject to undesired noise and other distortions , motivating an initial preprocessing step called image restoration . Image restoration is the task of removing unwanted noise and distortions , giving us clean images that are closer to the true but unknown signal . In the past years , Deep Learning ( DL ) has enabled tremendous progress in image restoration ( Mao et al. , 2016 ; Zhang et al. , 2017b ; Zhang et al. , 2017 ; Weigert et al. , 2018 ) . Supervised DL methods use corresponding pairs of clean and distorted images to learn a mapping between the two quality levels . The utility of this approach is especially pronounced for microscopy image data of biological samples ( Weigert et al. , 2017 ; 2018 ; Ouyang et al. , 2018 ; Wang et al. , 2019 ) , where quantitative downstream analysis is essential . More recently , unsupervised content-aware image restoration ˚Shared first authors . : Shared last authors . ( CARE ) methods ( Lehtinen et al. , 2018 ; Krull et al. , 2019 ; Batson & Royer , 2019 ; Buchholz et al. , 2019 ) have emerged . They can , enabled by sensible assumptions about the statistics of imaging noise , learn a mapping from noisy to clean images , without ever seeing clean data during training . Some of these methods additionally include a probabilistic model of the imaging noise ( Krull et al. , 2020 ; Laine et al. , 2019 ; Prakash et al. , 2020 ; Khademi et al. , 2020 ) to further improve their performance . Note that such denoisers can directly be trained on a given body of noisy images . All existing approaches have a common flaw : distortions degrade some of the information content in images , generally making it impossible to fully recover the desired clean signal with certainty . Even an ideal method can not know which of many possible clean images really has given rise to the degraded observation at hand . Hence , any restoration method has to make a compromise between possible solutions when predicting a restored image . Generative models , such as VAEs , are a canonical choice when a distribution over a set of variables needs to be learned . Still , so far VAEs have been overlooked as a method to solve unsupervised image denoising problems . This might also be due to the fact that vanilla VAEs ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ) show sub-par performance on denoising problems ( see Section 6 ) . Here we introduce DIVNOISING , a principled approach to incorporate explicit models of the imaging noise distribution in the decoder of a VAE . Such noise models can be either measured or derived ( bootstrapped ) from the noisy image data alone ( Krull et al. , 2020 ; Prakash et al. , 2020 ) . Additionally we propose a way to co-learn a suitable noise model during training , rendering DIVNOISING fully unsupervised . We show on 13 datasets that fully convolutional VAEs , trained with our proposed DIVNOISING framework , yield competitive results , in 8 cases actually becoming the new state-ofthe-art ( see Fig . 2 and Table 1 ) . Still , the key benefit of DIVNOISING is that the method does not need to commit to a single prediction , but is instead capable of generating diverse samples from an approximate posterior of possible true signals . ( Note that point estimates can still be inferred if desired , as shown in Fig . 4 . ) Other unsupervised denoising methods only provide a single solution ( point estimate ) of that posterior ( Krull et al. , 2019 ; Lehtinen et al. , 2018 ; Batson & Royer , 2019 ) or predict an independent posterior distribution of intensities per pixel ( Krull et al. , 2020 ; Laine et al. , 2019 ; Prakash et al. , 2020 ; Khademi et al. , 2020 ) . Hence , DIVNOISING is the first method that learns to approximate the posterior over meaningful structures in a given body of images . We believe that DIVNOISING will be hugely beneficial for computational biology applications in biomedical imaging , where noise is typically unavoidable and huge datasets need to be processed on a daily basis . Here , DIVNOISING enables unsupervised diverse SOTA denoising while requiring only comparatively little computational resources , rendering our approach particularly practical . Finally , we discuss the utility of diverse denoising results for OCR and showcase it for a ubiquitous analysis task in biology – the instance segmentation of cells in microscopy images ( see Fig . 5 ) . Hence , DIVNOISING has the potential to be useful for many real-world applications and will not only generate state-of-the-art ( SOTA ) restored images , but also enrich quantitative downstream processing . 2 RELATED WORK . Classical Denoising . The denoising problem has been addressed by a variety of filtering approaches . Arguably some of the most prominent ones are Non-Local Means ( Buades et al. , 2005 ) and BM3D ( Dabov et al. , 2007 ) , which implement a sophisticated non-local filtering scheme . A comprehensive survey and in-depth discussion of such methods can be found in ( Milanfar , 2012 ) . DL Based Denoising . Deep Learning methods which directly learn a mapping from a noisy image to its clean counterpart ( see e.g . ( Zhang et al. , 2017a ) and ( Weigert et al. , 2018 ) ) have outperformed classical denoising methods in recent years . Two well known contributions are the seminal works by Zhang et al . ( 2017a ) and later by Weigert et al . ( 2018 ) . More recently , a number of unsupervised variations have been proposed , and in Section 1 we have described their advantages and disadvantages in detail . One additional interesting contribution was made by Ulyanov et al . ( 2018 ) , introducing a quite different kind of unsupervised restoration approach . Their method , Deep Image Prior , trains a network separately for each noisy input image in the training set , making this approach computationally rather expensive . Furthermore , training has to be stopped after a suitable but a priori unknown number of training steps . Recently , Quan et al . ( 2020 ) proposed an interesting method called SELF2SELF which trains a U-NET like architecture requiring only single noisy images . The key idea of this approach is to use blind spot masking , similar to Krull et al . ( 2019 ) , together with dropout ( Srivastava et al. , 2014 ) , which avoids overfitting and allows sampling of diverse solutions . Similar to DIVNOISING , the single denoised result is obtained by averaging many diverse predictions . Diverse results obtained via dropout are generally considered to capture the so called epistemic or model uncertainty ( Gal & Ghahramani , 2016 ; Lakshminarayanan et al. , 2017 ) , i.e . the uncertainty arising from the fact that we have a limited amount of training data available . In contrast , DIVNOISING combines a VAE and a model of the imaging noise to capture what is known as aleatoric or data uncertainty ( Böhm et al. , 2019 ; Sensoy et al. , 2020 ) , i.e . the unavoidable uncertainty about the true signal resulting from noisy measurements . Like in Ulyanov et al . ( 2018 ) , also SELF2SELF trains separately on each image that has to be denoised . While this renders the method universally applicable , it is computationally prohibitive when applied to large datasets . The same is true for real time applications such as facial denoising . DIVNOISING , on the other hand , is trained only once on a given body of data . Afterwards , it can be efficiently applied to new images . A detailed comparison of SELF2SELF and DIVNOISING in terms of denoising performance , run time and GPU memory requirements can be found in Appendix A.14 and Appendix Table 2 . Denoising ( Variational ) Autoencoders . Despite the suggestive name , denoising variational autoencoders ( Im et al. , 2017 ) are not solving denoising problems . Instead , this method proposes to add noise to the input data in order to boost the quality of encoder distributions . This , in turn , can lead to stronger generative models . Other methods also follow a similar approach to improve overall performance of autoencoders ( Vincent et al. , 2008 ; 2010 ; Jiao et al. , 2020 ) . VAEs for Diverse Solution Sampling . Although not explored in the context of unsupervised denoising , VAEs are designed to sample diverse solutions from trained posteriors . The probabilistic U-NET ( Kohl et al. , 2018 ; 2019 ) uses conditional VAEs to learn a conditional distribution over segmentations . Baumgartner et al . ( 2019 ) improve the diversity of segmentation samples by introducing a hierarchy of latent variables to model segmentations at multiple resolutions . Unlike DIVNOISING , both methods rely on paired training data . Nazabal et al . ( 2020 ) employ VAEs to learn the distribution of incomplete and heterogeneous data in a fully unsupervised manner . Babaeizadeh et al . ( 2017 ) build upon a VAE style framework to predict multiple plausible future frames of videos conditioned on given context frames . A variational inference approach was used by Balakrishnan et al . ( 2019 ) to generate multiple deprojected samples for images and videos collapsed in either spatial or temporal dimensions . Unlike all these approaches , we address the uncertainty introduced by common imaging noise and show how denoised samples can improve downstream processing . 3 THE DENOISING TASK . Image restoration is the task of estimating a clean signal s “ ps1 , . . . , sN q from a corrupted observation x “ px1 , . . . , xN q , where si and xi , refer to the respective pixel intensities . The corrupted x is thought to be drawn from a probability distribution pNMpx|sq , which we call the observation likelihood or the noise model . In this work we focus on restoring images that suffer from insufficient illumination and detector/camera imperfections . Contrary to existing methods , DIVNOISING is designed to capture the inherent uncertainty of the denoising problem by learning a suitable posterior distribution . Formally , the posterior we are interested in is pps|xq9ppx|sqppsq and depends on two components : the prior distribution ppsq of the signal as well as the observation likelihood pNMpx|sq we introduced above . While the prior is a highly complex distribution , the likelihood ppx|sq of a given imaging system ( camera/microscope ) can be described analytically ( Krull et al. , 2020 ) . Models of Imaging Noise . The noise model is usually thought to factorize as a product of pixels , implying that the corruption , given the underlying signal , is occurring independently in each pixel as ppx|sq “ N ź i pNMpxi|siq . ( 1 ) This assumption is known to hold true for Poisson shot noise and camera readout noise ( Zhang et al. , 2019 ; Krull et al. , 2020 ; Prakash et al. , 2020 ) . We will refer to the probability pNMpxi|siq of observing a particular noisy value xi at a pixel i given clean signal si as the pixel noise model . Various types of pixel noise models have been proposed , ranging from physics based analytical models ( Zhang et al. , 2019 ; Luisier et al. , 2010 ; Foi et al. , 2008 ) to simple histograms ( Krull et al. , 2020 ) . In this work , we follow the Gaussian Mixture Model ( GMM ) based noise model description of ( Prakash et al. , 2020 ) . The parameters of a noise model can be estimated whenever pairs px1 , s1q of corresponding noisy and clean calibration images are available ( Krull et al. , 2020 ) . The signal s1 « 1M řM j “ 0 x 1j can then be computed by averaging these noisy observations ( Prakash et al. , 2020 ) . In a case where no calibration data can be acquired , s1 can be estimated by a bootstrapping approach ( Prakash et al. , 2020 ) . Later , we additionally show how a suitable noise model can be co-learned during training .

Summary: This paper devises a novel unsupervised denoising paradigm, DIVNOISING, that allows us, for the first time, to generate diverse and plausible denoising solutions, sampled from a learned posterior. This approach only requires noisy images and a suitable description of the imaging noise distribution, providing a new perspective for the image denoising field. It has demonstrated that the quality of denoised images is highly competitive, typically outperforming the unsupervised state-of-the-art, and at times even improving on supervised results. This paper is well-written and good-organized. However, the reviewer has the following concerns.