Interpretable Unsupervised Diversity Denoising and Artefact Removal

1 INTRODUCTION . Deep Learning ( DL ) based methods are currently the state-of-the-art ( SOTA ) for image denoising and artefact removal with supervised methods typically showing best performance ( Zhang et al. , 2017 ; Weigert et al. , 2017 ; Chen et al. , 2018 ; Delbracio et al. , 2020 ) . However , in order to be trained , supervised methods need either paired high and low quality images ( Zhang et al. , 2017 ; Weigert et al. , 2017 ; Delbracio et al. , 2020 ) or pairs of low quality images ( Lehtinen et al. , 2018 ; Buchholz et al. , 2019 ) . Both requirements are often not , or at least not efficiently satisfiable , hindering application of these methods to many practical applications , mainly in microscopy and biomedical imaging . Hence , unsupervised methods that can be trained on noisy images alone Krull et al . ( 2019 ) ; Batson & Royer ( 2019 ) ; Xie et al . ( 2020 ) ; Quan et al . ( 2020 ) present an attractive alternative . Methods additionally using models of imaging noise have also been proposed ( Krull et al. , 2020 ; Laine et al. , 2019 ; Prakash et al. , 2019b ) and can further boost the denoising performance of unsupervised models . Still , unsupervised methods suffer from three drawbacks : piq they fail to account for the inherent uncertainty present in a corrupted image and produce only a single restored solution ( i.e . point estimation ) 1 , piiq they typically show weaker overall performance than their supervised counterparts , and piiiq they are , by design , limited to pixel-wise noise removal and can not handle structured noises or other image artefacts ( spatially correlated noise ) . Recently , the first of these drawbacks was addressed by DIVNOISING ( DN ) ( Prakash et al. , 2021 ) which proposed a convolutional Variational Autoencoder ( VAE ) architecture for unsupervised denoising and generates diverse denoised solutions , giving users access to samples from a distribution of sensible denoising results . But DN exhibits poor performance on harder ( visually more complex and varied ) datasets , e.g . diverse sets of natural images . Additionally , it does not improve on 1This is also true for supervised methods . the performance achieved with other unsupervised denoisers on existing microscopy benchmark datasets ( Prakash et al. , 2021 ) , where supervised methods , whenever applicable , lead to best results . Last but not least , DN does not address the problem of artefact removal ( structured noise removal ) . We hypothesize that the performance of methods like DN is limited by the used VAE architecture , which can neither capture longer range dependencies , nor the full structural complexity of many practically relevant datasets . Although more expressive hierarchical VAE ( HVAE ) architectures have been known for some time in the context of image generation ( Sønderby et al. , 2016 ; Maaløe et al. , 2019 ; Vahdat & Kautz , 2020 ; Child , 2021 ) , so far they have never been applied to image restoration or denoising . Besides , well performing HVAE architectures are computationally very expensive , with training times easily spanning many days or even weeks on multiple modern GPUs ( Vahdat & Kautz , 2020 ; Child , 2021 ) . Hence , it remains unexplored if more expressive HVAE architectures can indeed improve unsupervised image restoration tasks and if they can do so efficiently enough to justify their application in practical use-cases . Contributions . In this paper , we first introduce a new architecture for unsupervised diversity denoising called HIERARCHICAL DIVNOISING ( HDN ) , inspired by ladder VAEs ( Sønderby et al. , 2016 ) , but introducing a number of important task-specific modifications . We show that HDN is considerably more expressive than DN ( see Fig . 2 ) , leading to much improved denoising results ( see Fig . 1 ) , while still not being excessively computationally expensive ( on a common 6 GB Tesla P100 GPU training still only takes about 1 day ) . Most importantly though , HDN leads to SOTA results on natural images and biomedical image data , outperforming 8 popular unsupervised denoising baselines and considerably closing the gap to supervised results ( see Table 1 ) . Additionally , we investigate the application of diversity denoising to the removal of spatially correlated structured noise in microscopy images , an application of immense practical impact . More specifically , we show that HDN , owing to its hierarchical architecture , enforces an interpretable decomposition of learned representations and we devise an effective approach to remove structured noise in microscopy data by taking full advantage of this interpretable hierarchical representation . We showcase this on multiple real-world datasets from diverse microscopy modalities , demonstrating that our method is an important step towards interpretable image restoration and is immediately applicable to many real-world datasets as they are acquired in the life-sciences on a daily basis . Related Work . Apart from the DL based works already discussed earlier , classical methods such as Non-Local Means ( Buades et al. , 2005 ) and BM3D ( Dabov et al. , 2007 ) have also found widespread use for denoising applications . A detailed review of classical denoising methods can be found in Milanfar ( 2012 ) . An interesting unsupervised DL based image restoration method is Deep Image Prior ( DIP ) ( Ulyanov et al. , 2018 ) which showed that convolutional neural networks ( CNNs ) can be used to restore corrupted images without supervision when training is stopped at an apriori unknown time before convergence . Methods such as DIP ( or other methods that require training for each input image separately , e.g . SELF2SELF ( Quan et al. , 2020 ) ) are , however , computationally demanding when applied to entire datasets . Diversity pixel denoising based on VAEs was introduced by DN ( Prakash et al. , 2021 ) , and a similar approach using GANs by Ohayon et al . ( 2021 ) . For structured noise removal , Broaddus et al . ( 2020 ) extended Krull et al . ( 2019 ) to specifically remove structured line artefacts that can arise in microscopy data . Their method needs exact apriori knowledge about the size of artefacts , which can not span more than some connected pixels . Dorta et al . ( 2018 ) model structured noises as the uncertainty in the VAE decoder represented by a structured Gaussian distribution with non-diagonal covariance matrix which is learned separately by another network . Jancsary et al . ( 2012 ) learn a structured noise model in a regression tree field framework . Unlike these approaches , we take an interpretability-first perspective for artefact removal and remove artefacts even without having/learning any structured noise model . 2 THE IMAGE RESTORATION TASK . The task of image restoration involves the estimation of a clean and unobservable signal s “ ps1 , s2 , . . . , sN q from a corrupted reconstruction x “ px1 , x2 , . . . , xN q , where si and xi , refer to the respective pixel intensities in the image domain . In general , the reconstructed image comes from solving an inverse imaging problem giving by the forward model , y “ Apsq ` e , ( 1 ) where A is the forward operator ( tomography , blurring , sub-sampling , etc . ) , e is noise in the measurements typically assumed iid , and y is the noisy measurements . An image reconstruction algorithm is needed to recover an estimation x of s from the noisy measurements y . Typically , the image reconstruction is obtained through an optimization formulation , where we seek a solution x that fits the observed values and is compatible with some image prior R , x “ argmin s1 } Aps1q ´ y } 2 ` λRps1q , ( 2 ) where s1 is the auxiliary variable for the signal being optimized for and λ ě 0 is related to the level of confidence of the prior R. There exists an extensive amount of work defining image priors ( Rudin et al. , 1992 ; Golub et al. , 1999 ; Bruckstein et al. , 2009 ; Zoran & Weiss , 2011 ; Romano et al. , 2017 ; Ulyanov et al. , 2018 ) . Without loss of generality , we can decompose the reconstructed image as x “ s ` n , where n is the residual ( noise ) between the ideal image and the reconstructed one . Generally , the noise n on the reconstruction x is composed of pixel-noise ( such as Poisson or Gaussian noise ) and multi-pixel artefacts or structured noise that affect groups of pixels in correlated ways . Such artefacts arise through dependencies that are introduced by the adopted reconstruction technique and the domainspecific inverse problem ( e.g . tomography , microscopy , or ISP in an optical camera ) . For example , consider the case where the reconstruction is done using Tikhonov regularization Rpsq “ } s } 2 , in the particular case of linear operator . In this case , the solution to Eq . 2 is x “ pATA ` λIq´1AT y . Thus , x “ pATA ` λIq´1ATAs ` pATA ` λIq´1AT e “ s ` n , ( 3 ) where n “ ppATA ` λIq´1ATA´ Iqs ` pATA ` λIq´1ATe . ( 4 ) The reconstructed image is affected by colored noise ( term in Eq . 4 depending on e ) , and also by a structured perturbation introduced by the reconstruction method ( term in Eq . 4 depending on s ) . This structured perturbation appears even in the case when the measurements are noiseless . Accurately modeling the noise/artefacts on the reconstructed image is therefore challenging even in the simplest case where the reconstruction is linear , showing that structured noise removal is non-trivial . In Section 5 , we deal with image denoising where noise contribution ni to each pixel xi is assumed to be conditionally independent given the signal si , i.e . ppn|sq “ ś i ppni|siq ( Krull et al. , 2019 ) . In many practical applications , including the ones presented in Section 6 , this assumption does not hold true , and the noise n is referred to as structured noise/artefacts . 3 NON-HIERARCHICAL VAE BASED DIVERSITY DENOISERS . Recently , non-hierarchical VAE based architectures ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ; Prakash et al. , 2021 ) were explored in the context of image denoising allowing a key advantage of providing samples sk drawn from a posterior distribution of denoised images . More formally , the denoising operation of these models fψ can be expressed by fψpxq “ sk „ pps|xq , where ψ are the model parameters . In this section , we provide a brief overview of these models . Vanilla VAEs . VAEs are generative encoder-decoder models , capable of learning a latent representation of the data and capturing a distribution over inputs x ( Kingma & Welling , 2019 ; Rezende et al. , 2014 ) . The encoder maps input x to a conditional distribution qφpz|xq in latent space . The decoder , gθpzq , takes a sample from qφpz|xq and maps it to a distribution pθpx|zq in image space . Encoder and decoder are neural networks , jointly trained to minimize the loss Lφ , θpxq “ Eqφpz|xqr´ log pθpx|zqs ` KL pqφpz|xq||ppzqq “ L R φ , θpxq ` LKLφ pxq , ( 5 ) with the second term being the KL-divergence between the encoder distribution qφpz|xq and prior distribution ppzq ( usually a unit Normal distribution ) . The network parameters of the encoder and decoder are given by φ and θ , respectively . The decoder usually factorizes over pixels as pθpx|zq “ N ź i “ 1 pθpxi|zq , ( 6 ) with pθpxi|zq being a Normal distribution predicted by the decoder , xi representing the intensity of pixel i , and N being the total number of pixels in image x . The encoder distribution is modeled in a similar way , factorizing over the dimensions of the latent space z . Performance of vanilla VAEs on pixel-noise removal tasks is analyzed and compared to DN in Prakash et al . ( 2021 ) . DIVNOISING ( DN ) . DN ( Prakash et al. , 2021 ) is a VAE based method for unsupervised pixel noise removal that incorporates an explicit pixel wise noise model pNMpx|sq in the decoder . More formally , the generic Normal distribution over pixel intensities of Eq . 6 is replaced with a known noise model pNMpx|sq which factorizes as a product of pixels , i.e . pNMpx|sq “ śN i pNMpxi|siq , and the decoder learns a mapping from z directly to the space of restored images , i.e . gθpzq “ s. Therefore , pθpx|zq “ pNMpx|gθpzqq . ( 7 ) The loss of DN hence becomes Lφ , θpxq “ Eqφpz|xq « N ÿ i “ 1 ´ log ppxi|gθpzqq ff ` KL pqφpz|xq||ppzqq “ LRφ , θpxq ` LKLφ pxq . ( 8 ) Intuitively , the reconstruction loss in Eq . 8 measures how likely the generated image is given the noise model , while the KL loss incentivizes the encoded distribution to be unit Normal . DN is the current SOTA for many unsupervised denoising benchmarks ( particularly working well for microscopy images ) , but performs poorly for complex domains such as natural image denoising ( Prakash et al. , 2021 ) .

Summary: This work extends the work of DivNoising (Prakash et al., 2021) to build an interpretable approach for unsupervised image restoration. The proposed method, namely Hierarchical DivNosing (HDN), is based on the well-studied concept of hierarchical Variational Autoencoder. This work is the first to use hierarchical Variational Autoencoder for the task of unsupervised image denoising. In doing so, HDN achieves state-of-the-art unsupervised image denoising results on twelve benchmark image datasets. The method is also shown to remove structured artifacts on three real microscopy datasets.