Robustness via Uncertainty-aware Cycle Consistency

Unpaired image-to-image translation refers to learning inter-image-domain mapping without corresponding image pairs . Existing methods learn deterministic mappings without explicitly modelling the robustness to outliers or predictive uncertainty , leading to performance degradation when encountering unseen perturbations at test time . To address this , we propose a novel probabilistic method based on Uncertainty-aware Generalized Adaptive Cycle Consistency ( UGAC ) , which models the per-pixel residual by generalized Gaussian distribution , capable of modelling heavy-tailed distributions . We compare our model with a wide variety of state-of-the-art methods on various challenging tasks including unpaired image translation of natural images , using standard datasets , spanning autonomous driving , maps , facades , and also in medical imaging domain consisting of MRI . Experimental results demonstrate that our method exhibits stronger robustness towards unseen perturbations in test data . Code is released here : https : //github.com/ExplainableML/UncertaintyAwareCycleConsistency . 1 Introduction . Translating an image from a distribution , i.e . source domain , to an image in another distribution , i.e . target domain , with a distribution shift is an ill-posed problem as a unique deterministic one-to-one mapping may not exist between the two domains . Furthermore , since the correspondence between inter-domain samples may be missing , their joint-distribution needs to be inferred from a set of marginal distributions . However , as infinitely many joint distributions can be decomposed into a fixed set of marginal distributions [ 1 , 2 , 3 ] , the problem is ill-posed in the absence of additional constraints . Deep learning-based methods tackle the image-to-image translation task by learning inter-domain mappings in a paired or unpaired manner . Paired image translation methods [ 4 , 5 , 6 , 7 , 8 , 9 ] exploit the inter-domain correspondence by penalizing the per-pixel residual ( using l1 or l2 norm ) between the output and corresponding ground-truth sample . Unpaired image translation approaches [ 1 , 10 , 11 , 12 , 13 , 14 ] often use adversarial networks with an additional constraint on the image or feature space imposing structure on the underlying joint distribution of the images from the different domains . Both paired and unpaired image translation approaches often learn a deterministic mapping between the domains where every pixel in the input domain is mapped to a fixed pixel value in the output domain . However , such a deterministic formulation can lead to mode collapse while at the same time not being able to quantify the model predictive uncertainty important for critical applications , e.g. , medical image analysis . It is desirable to test the performance of the model on unseen perturbed input at test-time , to improve their applicability in the real world . While robustness to outliers is a focus in some domains [ 15 , 16 , 17 , 18 ] , it has not attracted as much attention in unpaired translation . To address these limitations , we propose an unpaired ( unsupervised ) probabilistic image-to-image translation method trained without inter-domain correspondence in an end-to-end manner . The probabilistic nature of this method provides uncertainty estimates for the predictions . Moreover , modelling the residuals between the predictions and the ground-truth with heavy-tailed distributions 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . makes our model robust to outliers and various unseen data . Accordingly , we compare various state-of-the-art models and our model in their capacity to handle samples from similar distribution as training-dataset as well as perturbed samples , in the context of unpaired translation . Our contributions are as follows . ( i ) We propose an unpaired probabilistic image-to-image translation framework based on Uncertainty-aware Generalized Adaptive Cycle Consistency ( UGAC ) . Our framework models the residuals between the predictions and the ground-truths with heavy-tailed distributions improving robustness to outliers . Probabilistic nature of UGAC also provides uncertainty estimates for the predictions . ( ii ) We evaluate UGAC on multiple challenging datasets : natural images consisting Cityscapes [ 19 ] , Google aerial maps and photos [ 4 ] , CMP Facade [ 20 ] and medical images consisting of MRI from IXI [ 21 ] . We compare our model to seven state-of-the-art image-to-image translation methods [ 12 , 22 , 1 , 11 , 10 , 23 ] . Our results demonstrate that while UGAC performs competitively when tested on unperturbed images , it improves state-of-the-art methods substantially when tested on unseen perturbations , establishing its robustness . ( iii ) We show that our estimated uncertainty scores correlate with the model predictive errors ( i.e. , residual between model prediction and the ground-truth ) suggesting that it acts as a good proxy for the model ’ s reliability at test time . 2 Related Work . Image-to-image translation . Image-to-image translation is often formulated as per-pixel deterministic regression between two image domains of [ 24 , 25 , 26 ] . In [ 4 ] , this is done in a paired manner using conditional adversarial networks , while in [ 10 , 1 , 12 , 22 , 11 ] this is done in an unpaired manner by enforcing additional constraints on the joint distribution of the images from separate domains . Both CycleGAN [ 10 ] and UNIT [ 1 ] learn bi-directional mappings , whereas other recent methods [ 12 , 22 , 11 ] learn uni-directional mappings . Quantification of uncertainty in the predictions made by the unpaired image-to-image translation models largely remains unexplored . Our proposed method operates at the intersection of uncertainty estimation and unsupervised translation . Critical applications such as medical image-to-image translation [ 27 , 28 , 29 , 30 , 31 , 32 ] is an excellent testbed for our model as confidence in the network ’ s predictions is desirable [ 33 , 34 ] especially under the influence of missing imaging modalities . Uncertainty estimation . Among two broad categories of uncertainties that can be associated with a model ’ s prediction , epistemic uncertainty in the model parameters is learned with finite data whereas aleatoric uncertainty captures the noise/uncertainty inherent in the data [ 35 , 36 ] . For image-to-image translation , various uncertainties can be estimated using Bayesian deep learning techniques [ 36 , 37 , 38 , 39 , 40 ] . In critical areas like medical imaging , the errors in the predictions deter the adoption of such frameworks in clinical contexts . Uncertainty estimates for the predictions would allow subsequent revision by clinicians [ 41 , 42 , 43 , 44 , 45 , 46 , 47 , 48 , 49 , 50 , 51 ] . Existing methods model the per-pixel heteroscedasticity as Gaussian distribution for regression tasks [ 36 ] . This is not optimal in the presence of outliers that often tend to follow heavy-tailed distributions [ 52 , 53 ] . Therefore , we enhance the above setup by modelling per-pixel heteroscedasticity as generalized Gaussian distribution , which can model a wide variety of distributions , including Gaussian , Laplace , and heavier-tailed distribution . 3 Uncertainty-aware Generalized Adaptive Cycle-consistency ( UGAC ) . We present the formulation of the unpaired image-to-image translation problem . We discuss the shortcomings of the existing solution involving the cycle consistency loss called CycleGAN [ 10 ] . Finally , we present our novel probabilistic framework ( UGAC ) that overcomes the described shortcomings . 3.1 Preliminaries . Formulation . Let there be two image domains A and B . Let the set of images from domain A and B be defined by ( i ) SA : = { a1 , a2 ... an } , where ai ∼ PA ∀i and ( ii ) SB : = { b1 , b2 ... bm } , where bi ∼ PB ∀i , respectively . The elements ai and bi represent the ith image from domain A and B respectively , and are drawn from an underlying unknown probability distribution PA and PB respectively . Let each image have K pixels , and uik represent the kth pixel of a particular image ui . We are interested in learning a mapping from domain A to B ( A→ B ) and B to A ( B → A ) in an unpaired manner so that the correspondence between the samples from PA and PB is not required at the learning stage . In other words , we want to learn the underlying joint distribution PAB from the given marginal distributions PA and PB . This work utilizes CycleGANs that leverage the cycle consistency to learn mappings from both directions ( A→ B and B → A ) , but often we are only interested in one direction and the second direction is the auxiliary mapping that aids in learning process . We define the mapping A→ B as primary and B → A as auxiliary . Cycle consistency . Learning a joint distribution from the marginal distributions is an ill-posed problem with infinitely many solutions [ 3 ] . CycleGAN [ 10 ] enforces an additional structure on the joint distribution using a set of primary networks ( forming a GAN ) and a set of auxiliary networks . The primary networks are represented by { GA ( · ; θGA ) , DA ( · ; θDA ) } , where GA represents a generator and DA represents a discriminator . The auxiliary networks are represented by { GB ( · ; θGB ) , DB ( · ; θDB ) } . While the primary networks learn the mapping A→ B , the auxiliary networks learn B → A ( see Figure 1 ) . Let the output of the generator GA translating samples from domain A ( say ai ) to domain B be called b̂i . Similarly , for the generator GB translating samples from domain B ( say bi ) to domain A be called âi , i.e. , b̂i = GA ( ai ; θGA ) and âi = GB ( bi ; θ G B ) . To simplify the notation , we will omit writing parameters of the networks in the equation . The cycle consistency constraint [ 10 ] re-translates the above predictions ( b̂i , âi ) to get back the reconstruction in the original domain ( āi , b̄i ) , where , āi = GB ( b̂i ) and b̄i = GA ( âi ) , and attempts to make reconstructed images ( āi , b̄i ) similar to original input ( ai , bi ) by penalizing the residuals with L1 norm between the reconstructions and the original input images , giving the cycle consistency Lcyc ( āi , b̄i , ai , bi ) = L1 ( āi , ai ) + L1 ( b̄i , bi ) . Limitations of cycle consistency . The underlying assumption when penalizing with the L1 norm is that the residual at every pixel between the reconstruction and the input follow zero-mean and fixed-variance Laplace distribution , i.e. , āij = aij + aij and b̄ij = bij + b ij with , aij , b ij ∼ Laplace ( ; 0 , σ√ 2 ) ≡ 1√ 2σ2 e− √ 2 | −0| σ , ( 1 ) where σ2 represents the fixed-variance of the distribution , aij represents the jth pixel in image ai , and aij represents the noise in the j th pixel for the estimated image āij . This assumption on the residuals between the reconstruction and the input enforces the likelihood ( i.e. , L ( Θ|X ) = P ( X|Θ ) , where Θ : = θGA ∪ θ G B ∪ θDA ∪ θDB and X : = SA ∪ SB ) to follow a factored Laplace distribution : L ( Θ|X ) ∝ ∏ ijpq e− √ 2|āij−aij | σ e− √ 2|b̄pq−bpq| σ , ( 2 ) where minimizing the negative-log-likelihood yields Lcyc with the following limitations . The residuals in the presence of outliers may not follow the Laplace distribution but instead a heavy-tailed distribution , whereas the i.i.d assumption leads to fixed variance distributions for the residuals that do not allow modelling of heteroscedasticity to aid in uncertainty estimation .

Summary: Authors propose a relaxation of the cycle consistency loss that takes into account per-pixel prediction uncertainty. More specifically, authors assume that cycle-reconstruction residuals are distributed as iid generalized Gaussian random variables, and maximize likelihood of these reconstruction residuals given predicted per-pixel parameters. Authors report two metrics of sensitivity to noise and three segmentation accuracy metrics across four real datasets with known GT and seven baseline methods, confirming that in all cases their method outperforms or performs on par with SotA. Authors also show that restricting parameters of their generalized Gaussian log-likelihood loss to constant values (effectively reducing it to standard CycleGAN) degrades the performance of their method according to these performance metrics. Moreover, authors showed that the overall prediction uncertainty strongly correlates with the magnitude of the true prediction error computed using GT.