Modal Uncertainty Estimation via Discrete Latent Representations

1 INTRODUCTION . Making predictions in the real world has to face with various uncertainties . One of the arguably most common uncertainties is due to partial or corrupted observations , as such it is often insufficient for making a unique and deterministic prediction . For example , when inspecting where a single CT scan of a patient contains lesion , without more information it is possible for radiologists to reach different conclusions , as a result of the different hypotheses they have about the image . In such an ambiguous scenario , the question is thus , given the observable , which one ( s ) out of the many possibilities would be more reasonable than others ? Mathematically , this is a one-to-many mapping problem and can be formulated as follows . Suppose the observed information is x ∈ X in the input space , we are asked to estimate the conditional distribution p ( y|x ) for y ∈ Y in the prediction space , based on the training sample pairs ( x , y ) . There are immediate challenges that prevent p ( y|x ) being estimated directly in practical situations . First of all , both X and Y , e.g.as spaces of images , can be embedded in very high dimensional spaces with very complex structures . Secondly , only the unorganized pairs ( x , y ) , not the one-tomany mappings x 7→ { yi } i , are explicitly available . Fortunately , recent advances in conditional generative models based on Variational Auto-Encoder ( VAE ) framework from Kingma & Welling ( 2014 ) shed light on how to tackle our problem . By modelling through latent variables c = c ( x ) , one aims to explain the underlying mechanism of how y is assigned to x . And hopefully , variation of c will result in variation in the output ŷ ( x , c ) , which will approximate the true one-to-many mappings distributionally . Many current conditional generative models , including cVAE in Sohn et al . ( 2015 ) , BiCycleGAN in Zhu et al . ( 2017b ) , Probabilistic U-Net in Kohl et al . ( 2018 ) , etc. , are developed upon the VAE framework , with Gaussian distribution with diagonal covariance as the de facto parametrization of the latent variables . However , in the following we will show that such a parametrization put a dilemma between model training and actual inference , as a form of what is known as the posterior collapse problem in the VAE literature Alemi et al . ( 2018 ) ; Razavi et al . ( 2018 ) . This issue is particularly easy to understand in our setting , where we assume there are multiple y ’ s for a given x . Let us recall that one key ingredient of the VAE framework is to minimize the KL-divergence between the latent prior distribution p ( c|x ) and the latent variational approximation pφ ( c|x , y ) of the posterior . Here φ denotes the model parameters of the “ recognition model ” in VAE . It does not matter if the prior is fixed p ( c|x ) = p ( c ) Kingma & Welling ( 2014 ) or learned p ( c|x ) = pθ ( c|x ) Sohn et al . ( 2015 ) , as long as both prior and variational posterior are parameterized by Gaussians . Now suppose for a particular x , there there are two modes y1 , y2 for the corresponding predictions . Since the minimization is performed on the entire training set , p ( c|x ) is forced to approximate a posterior mixture p ( c|x , y ( · ) ) of two Gaussians from mode y1 and y2 . In the situation when the minimization is successful , meaning the KL divergence is small , the mixture of the variational posteriors must be close to a Gaussian , i.e.posterior collapsed as in Fig.1 ( b ) , and hence the multi-modal information is lost . Putting it in contrapositive , if multi-modal information is to be conveyed by the variational posterior , then the minimization will not be successful , meaning higher KL divergence . This may partly explain why it can be a delicate matter to train a conditional VAE . The situation is schematically illustrated in Figure 1 in one dimension . Note that the case in Figure 1 ( a ) is usually more preferable , however the density values of the prior used during testing can not reflect the uncertainty level of the outputs . We quantitative demonstrate this in Section 4 and Fig.2 . One direction to solve the above problem is to modify the strength of KL-divergence or the variational lower bound , while keeping the Gaussian parametrization , and has been explored in the literature extensively , as in Higgins et al . ( 2017 ) ; Alemi et al . ( 2018 ) ; Rezende & Viola ( 2018 ) . However , besides the need of extensive parameter tuning for these approaches , they are not tailored for the multi-modal posterior collapse problem we described above , thus do not solve the inaccurate uncertainty estimation problem . Mixture or compositions of Gaussian priors have also been proposed in Nalisnick et al . ( 2016 ) ; Tomczak & Welling ( 2018 ) , but the number of Gaussians in the mixture is usually fixed apriori . Hence making it a conditional generative model further complicates the matter , since the number in the mixture should depend on the input . We therefore adopt another direction , which is to use a latent distribution parameterization other than Gaussians , and one that can naturally exhibit multiple modes . The simplest choice would be to constrain the latent space to be a finite set , as proposed in van den Oord et al . ( 2017 ) , so that we can learn the conditional distribution as a categorical distribution . We argue that the approach of discrete latent space may be beneficial particularly in our setting . First , different from unconditional or weak conditional generative modelling tasks where diversity is the main consideration , making accurate predictions based on partial information often leads to a significantly restricted output space . Second , there is no longer noise injection during training , so that the decoder can utilize the information from the latent variable more effectively . This makes it less prone to ignore the latent variable completely , in contrast to many conditional generation methods using noise inputs . Third , the density value learned on the latent space is more interpretable , since the learned prior can approximate the variational posterior better . In our case , the latent variables can now represent latent mode hypotheses for making the corresponding most likely predictions . We call our approach modal uncertainty estimation ( MUE ) . The main contributions of this work are : ( 1 ) We solve the MUE problem by using c-VAE and justify the use of a discrete latent space from the perspective of multi-modal posterior collapse problem . ( 2 ) Our uncertainty estimation improves significantly over the existing state-of-art . ( 3 ) In contrast to models using noise inputs that require sampling at the testing stage , our model can directly produce results ordered by their latent mode hypothesis probabilities , and is thus more informative and convenient for practical use . The rest of paper is organized as follows . In Section 2 we sample some works that related to ours and stress the key differences between them . In Section 3 we layout our general framework and model details . We conducted a series of experiments on both synthetic and real datasets described in Section 4 . The paper is concluded in Section 5 . 2 RELATED WORK . Conditional generative models aim to capture the conditional distribution of the data and generate them according to some given information . Thanks to the recent advancement of deep learning techniques , especially the methods of generative adversarial networks ( GANs ) Goodfellow et al . ( 2014 ) and variational auto-encoders ( VAEs ) Kingma & Welling ( 2014 ) , conditional generative models have been effectively applied to various computer vision and graphics tasks such as image synthesis , style transfer , image in-painting , etc . Early works in this direction focused on learning the unimodal mapping , as in Isola et al . ( 2017 ) and Zhu et al . ( 2017a ) . They are called uni-modal because the mapping is between fixed categories , namely a one-to-one mapping . There are no latent codes to sample from , thus the generation is deterministic . In these works , images of a specific category are translated to another category , while keeping the desired semantic content . These methods achieved the goal through a meta supervision technique known as the adversarial loss as in the GAN framework , where one only needs to supply weak supervision for whether the generated image belongs to a certain category or not . Adversarial loss has been known for producing sharp visual look but it alone can not guarantee faithful distribution approximation , where issues known as mode collapse and mode dropping often occur for complicated data distribution Srivastava et al . ( 2017 ) . In Isola et al . ( 2017 ) it is noted that additional noise input in the conditional model in fact fails to increase variability in the output . How to ensure good approximation of output distribution for GANs is still an active area of research . Therefore , the above frameworks might not be suitable for approximating the distribution of one-to-many mappings . Many works have been proposed to extend to the setting of one-to-many mappings by learning disentangled representations , of e.g. “ content ” and “ style ” , and consequently some form of auto-encoding has to be used . Conditional generation can then be accomplished by corresponding latent code sampling and decoding . This includes the approaches of Zhu et al . ( 2017b ) ; Huang et al . ( 2018 ) for multi-modal image-to-image translation , Zheng et al . ( 2019 ) for image in-painting , and many others . Since the main objectives of these works are the visual quality and diversity of the outputs , they are usually not evaluated in terms of the approximation quality of the output distribution . One notable exception is Probabilistic U-Net proposed in Kohl et al . ( 2018 ) , which is based on the conditional VAE framework Sohn et al . ( 2015 ) and is close in spirit to ours . Probabilistic U-Net has shown superior performance over various other methods for calibrated uncertainty estimation , including the ensemble methods of Lakshminarayanan et al . ( 2017 ) , multi-heads of Rupprecht et al . ( 2017 ) ; Ilg et al . ( 2018 ) , drop-out of Kendall et al . ( 2015 ) and Image2Image VAE of Zhu et al . ( 2017b ) . However , as discussed in Section 1 , Probabilistic U-Net can not solve the multi-modal posterior collapse problem since it uses Gaussian latent parameterization . Therefore , in case the conditional distribution is varying for different input data , the performance is expected to degrade . Furthermore , the latent prior density learned has no interpretation , and thus can not rank its prediction . To perform uncertainty estimation for Probabilistic U-Net , one must perform extensive sampling and clustering . Our framework improves significantly upon Probabilistic U-Net by introducing discrete latent space . With this latent parameterization we can directly output the uncertainty estimation and we can rank our predictions easily . The discrete latent space has been proposed in the vq-VAE framework of van den Oord et al . ( 2017 ) . With such a latent space it can get rid of the noise sampling , which enables the latent variable to be more effectively utilized by the decoder and produce outputs with better visual quality . While our use of discrete latent space is motivated by the multi-modal posterior collapse problem . The major technical difference compared to our framework is that the image in vq-VAE framework is encoded by a collection of codes arranged in the spatial order . As such , the joint distribution of the codes can not be obtained directly , and has to be estimated or sampled using e.g.an auto-regressive model in the spatial dimension , such as PixelCNN Van den Oord et al . ( 2016 ) . In contrast , we learn disentangled representations and only the necessary information to produce different outputs goes into the discrete latent space . In particular , we model the each mode of y given x by a single latent code , thus our model enjoys much simpler sampling . Besides vq-VAE van den Oord et al . ( 2017 ) , the use of discrete latent variables in neural network has been explored in various previous works , including the early work of Mnih & Gregor ( 2014 ) and Mnih & Rezende ( 2016 ) that use single or multiple samples objectives with variance reduction techniques to help training . Others have explored using continuous approximations to the discrete distributions , know as Concrete Maddison et al . ( 2016 ) or Gumbel-Softmax Jang et al . ( 2016 ) distributions . As is noted in van den Oord et al . ( 2017 ) , in general the above approaches have fallen short of their continuous counterparts . Worth mentioning is a recently proposed neural dialogue generation method Zhao et al . ( 2018 ) that uses Gumbel-Softmax approximation , which treats the dialogue generation as a one-to-many mapping problem . Our method diverge from theirs by the assumption about the model . In Zhao et al . ( 2018 ) , they designed the learned discrete representation for an utterance to be “ context free ” . This is in contrast to our assumption that the latent hypothesis of an input should depend on the input itself . Taking the task of medical image segmentation for an example , if we encode the hypotheses from the segmentation alone as in Zhao et al . ( 2018 ) , likely there will either be two modes ( benign vs malignant ) or a huge number of modes if the shape of the segmentation is taken into account . Moreover , it will not contain any information about what kinds of actual biological tissue they might be , which on the other hand can be judged from the actual scan image . In our case , we have deliberately separated the recognition task learning , e.g . segmenting the image , and the hypothesis learning , so that together they can approximate the variation of the outputs given the input . Finally , we briefly summarize the differences between MUE and existing uncertainty estimation methodologies in deep learning . Many existing works Gal & Ghahramani ( 2016 ) ; Gal ( 2016 ) ; Kendall et al . ( 2015 ) ; Kendall & Gal ( 2017 ) focus on model uncertainty , which try to capture the calibrated level of confidence of the model prediction by using stochastic regularization techniques . Such uncertainty will be of major interest for model predictions on unseen data and long-tail rare cases , or when model is trained on limited data . While ours is more about learning from conflicting or ambiguous training data , and estimating the calibrated uncertainty of the input-output relationship in the dataset . Interestingly , Kohl et al . ( 2018 ) has experimented using Dropout as comparison to the c-VAE framework in the MUE setting , but found it only achieved inferior performance . In general , since MUE is independent from the model uncertainty , our framework can be used jointly with existing techniques for prediction confidence estimation .

Summary: This paper introduces a novel conditional generative model for high dimensional data with multimodal output distributions. The proposed method, called modal uncertainty estimation (MUE), is a conditional VAE but with discrete latent representations. This discrete latent space allows the model to better handle multimodal outputs and provide confidence scores for the different modes predicted by the model. These capabilities are applied to the task of segmenting lesions in medical scans.