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 proposes a new method of noise removal using convolutional VAE.  An observed image with noise is input to VAE, and after the expression $z$ in the latent space, the noise removed image is finally output.  After that, it is possible to generate a pseudo noisy observation image according to the noise model.  The noise model part is flexibly designed using the Gaussian mixture model.  In the training, VAE and noise model can be learned at the same time. Since VAE is a generative model, and a clean denoising image can be obtained by averaging a large number of candidates of clean images, $s$, sampled from the periphery of the latent space representation $z$.