Modal Uncertainty Estimation via Discrete Latent Representations
Authors: Di Qiu
Zhanghan Ke
Peng Su
Lok Ming Lui
Venue:  (2021)
Keywords: uncertainty estimation
one -to-many mapping
conditional generative model
discrete latent space
medical image segmentation

Abstract: Many important problems in the real world don't have unique solutions. It is thus important for machine learning models to be capable of proposing different plausible solutions with meaningful probability measures.
In this work we propose a novel deep learning based framework, named {\it modal uncertainty estimation}  (MUE), to learn the one-to-many mappings between the inputs and outputs, together with faithful uncertainty estimation.
Motivated by the multi-modal posterior collapse problem in current conditional generative models, MUE uses a set of discrete latent variables, each representing a latent mode hypothesis that explains one type of input-output relationship, to generate the one-to-many mappings. Benefit from the discrete nature of the latent representations, MUE can estimate any input the conditional probability distribution of the outputs effectively. Moreover, MUE is efficient during training since the discrete latent space and its uncertainty estimation are jointly learned.
We also develop the theoretical background of MUE and extensively validate it on both synthetic and realistic tasks. MUE demonstrates (1) significantly more accurate uncertainty estimation than the current state-of-the-art, and (2) its informativeness for practical use.



Under review as a conference paper at ICLR 2021
UNCERTAINTY ESTIMATION VIA DISCRETE LATENT REPRESENTATION
Anonymous authors Paper under double-blind review
ABSTRACT
Many important problems in the real world don't have unique solutions. It is thus important for machine learning models to be capable of proposing differ- ent plausible solutions with meaningful probability measures. In this work we propose a novel deep learning based framework, named modal uncertainty esti- mation (MUE), to learn the one-to-many mappings between the inputs and out- puts, together with faithful uncertainty estimation. Motivated by the multi-modal posterior collapse problem in current conditional generative models, MUE uses a set of discrete latent variables, each representing a latent mode hypothesis that explains one type of input-output relationship, to generate the one-to-many map- pings. Benefit from the discrete nature of the latent representations, MUE can estimate any input the conditional probability distribution of the outputs effec- tively. Moreover, MUE is efficient during training since the discrete latent space and its uncertainty estimation are jointly learned. We also develop the theoreti- cal background of MUE and extensively validate it on both synthetic and realistic tasks. MUE demonstrates (1) significantly more accurate uncertainty estimation than the current state-of- of-the-art, and (2) its informativeness for practical use.
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 unique and deterministic prediction. For example, when inspecting where single CT scan of patient contains lesion, without more information it is possible for radiologists to reach different conclusions, 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 one-to-many mapping problem and can be formulated as follows. Suppose the observed information is x E X' in the input space, we are asked to estimate the conditional distribution p(y|x) for y E 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 very high dimensional spaces with very complex structures. Secondly, only the unorganized pairs (x,y), not the one-to- many mappings x + {y.Ji. 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 assigned to x. And hopefully, variation of e 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 eVAE 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, collapse problem in the VAE literature Alemi et al. (2018); Razavi et al. (2018). This issue is particularly easy understand in setting, where we assume there are multiple y's for given
form of what is known the posterior

---

Under review as a conference paper at ICLR 2021
pustetior mixture peior
(a)
(b)
(c)
Figure 1: Comparison between Gaussian latent representations and discrete latent representations in muulti-modll situation. Gaussian latents are structurally limited in such setting. (a) The ideal situation when there is no posterior collapse as multiple modes appear, but the prior distribution poor approximation of the posterior. (b) Posterior collapse happens, and no multi-modal in- formation is conveyed from the learned prior. (c) Discrete latent representation can ameliorate the posterior collapse problem while the prior can approximate the posterior more accurately when both restricted to be discrete.
Let us recall that one key ingredient of the VAE framework is to minimize the KL-divergence be- tween the latent prior distribution p(c[x) and the latent variational approximation ( y) of the posterior. Here á denotes the model parameters of the "recognition model" in VAE. It does not mat- ter if the prior is fixed p(c) x) p(c) Kingma & Welling (2014) learned p(c|x) po(c)x) Sohn et al. (2015), as long as both prior and variational posterior are parameterized by Gaussians. Now suppose for particular x, there there are two modes yI, 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( (()) 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 Gaussian, i.e. te.posterior collapsed as in 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 delicate matter to train a conditional VAE. The situation is schemat- ically illustrated in Figure 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 cannot 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 modify the strength of KL-divergence or the vari- ational lower bound, while keeping the Gaussian parametrization, and has been explored in the literature extensively, as in Higgins et al. (2017); Alemi 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 pro- posed in Nalisnick et al. (2016); Tomezak & Welling (2018), but the number of Gaussians in the mixture is usually fixed apriori. Hence making it 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 partial information often leads to a sig- nificantly 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 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)
2

---

Under review as a conference paper at ICLR 2021
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 layout our general framework and model details. We conducted a series of experiments on both synthetie and real datasets described in Section 4. The paper 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 tech- niques, 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 various computer vision and graphics tasks such as image synthe- sis, style transfer, image in-painting, etc. Early works in this direction focused on learning the uni- modal mapping, in Isola al. (2017) and Zhu 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 specific category are translated to another category, while keeping the desired semantic content. These methods achieved the goal through meta supervision technique known as the adversarial loss as in the GAN frame- work, 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 cannot guarantee faithful distribution approximation, where issues known 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 the setting of one-to-many mappings by learning disen- tangled 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 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 su- perior 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 al. (2018), 8), drop-out of Kendall al. (2015) and Image21mage VAE of Zhu et al. (2017b). How- ever, as discussed in Section 1, Probabilistic U-Net cannot solve the multi-modal posterior collapse problem since uses Gaussian latent parameterization. Therefore, in case the conditional distribu- tion is varying for different input data, the performance is expected to degrade. Furthermore, the latent prior density learned has no interpretation, and thus cannot 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- VAF framework of van den Oord et al. (2017). With such 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 collection of codes arranged in the spatial order. As such, the joint distribution of the codes cannot be obtained directly, and has to be estimated or sampled using e.g.an auto-regressive model in the spatial dimension, such PixeICNN Van den Oord et al. (2016).
3

---

Under review as a conference paper at ICLR 2021
In contrast, 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 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 help training. Others have explored using continuous approximations the discrete distributions, know as Concrete Maddison et al. (2016) or Gumbel-Softmax Jang et al. (2016) dis- tributions. As is noted in van den Oord 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 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 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. c.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 relation- ship in the dataset. Interestingly. Kohl et al. (2018) has experimented using Dropout as comparison to the c- VAB 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.
3 METHOD
3.1 GENERAL FRAMEWORK
Let yx denote the data-label pair. We model the generation of y conditioned on x using the conditional VAE framework as in Sohn et al. (2015). First, a latent variable is generated from some prior distribution po (c(x)) parametrized by a neural network. Then the label y is generated from some conditional distribution p (y|c, x). We use 0 to represent the collection of model parameters at testing time. The major distinction of our approach is that we assume c takes value in a finite set C, thought of as a code book for the latent mode hypotheses of our multi-modal data distribution. Our goal is to learn the optimal parameters 8" and the code book C, so that possibly multiple of the latent modes corresponding to x can be identified. and label predictions can be faithfully generated from x. The latter means the marginal likelihood po x) should be maximized.
The variational inference approach in Kingma & Welling (2014) starts by introducing posterior encoding model (c(c)x, y) with parameters o, which is used only during training. Since the label information is given, will assume the posterior encoding model is deterministic, meaning there is no "modal uncertainty" for the posterior encoding model. So the posterior distribution will be delta distribution for each data-label pair (x, y). In any case, the marginal log-likelihood of y can now be written as
log 2o(c/x,)) y) pa po(c)x,))
log po(c, y x) go (c(c))
log Po (y|x) = (
+ B (clx,))
(1)
4

---

Under review as a conference paper at ICLR 2021
Since the first term in the RHS of (1) is non-negative, we have the variational lower bound
log Pa(c,y|x))
log pespely/x) 2 B (c(xx))))
(c(c)/,,))
log lo(cly,
= -
+ E (epx.) [log Pyyle,
po(c|x)
We further lower bound Equation (2) by observing that the entropy term - -Eg() (log (-)) )) is positive and is constant if y)) is deterministic. This yields a sum of negative cross entropy and a conditional likelihood
log Poypo(y)x) 2 Byg(c(x.y) po(lx) + (olx,)) Po x))
For our optimization problem, we will maximize the lower bound (3). Since c takes value in the finite code book C, the probability distribution xc)x can be estimated using multi-class classification, and the cross entropy term can be estimated efficiently using stochastic approximation.
It is important to note that since assume ( y) is deterministic, will not regularize it by pulling to the prior distribution, contrast to the previous conditional VAE frameworks. This means that the probability values of the posterior is not influenced by the probability values of the prior distribution. Instead, we will let the prior encoding model to be trained by the posterior encoding model, as classification task with ground-truth being the class index obtained from the posterior encoder and the code book C. The lacking of prior regularization is also featured in the vq- VAE approach in van den Oord et al. (2017) for unconditional generation, and in Razavi et al. (2018) it is argued that restricting the latent space C to be finite set is itself structural prior constraint for the VAE framework. Note that here the discrete latent space C should be considered just finite set of indices, without any other structure between these indices, Below we will discuss how to realize it in I so that the actual representation can be useful the decoder.
First of all, because C is finite set, the objective (3) is not fully differentiable. We tackle this problem using a simple gradient approximation method and an extra regularization loss, following the approach of van den Oord et al. (2017); Razavi al. (2019). In details, denote the prior encoding network by Egg. the posterior encoding network by Eg. and the decoder as Dg. Since we assume delta distribution for the posterior encoding model (c(c/x y), we can let the posterior encoder produce a deterministic output for the given input-output pair (x. y). In other words, no sampling performed by the posterior encoder. Suppose the output of the posterior encoding network is e (x), y). Its nearest neighbor e in C in le distance
c = arg min ||c' - ell
ECEC
will become the input to the decoder network. And we simply copy the gradient of c to that of so that the posterior encoder can obtain gradient information from the label prediction error. To make sure the gradient approximation is accurate, we need to encourage the posterior encoder's outputs to approximate values in C' as close as possible. To achieve this we use an e fe-penalization of the form Blle - sg [c]]]] with parameter 3 > and sg is the stop-gradient operation. The code is updated using exponential moving average of the corresponding posterior encoder's outputs. In the above notation, our loss function be minimized for a single input pair (x, y) is
C(O, の) CE(E, ide) Recon x).) + (x) sg/c]]?
(4)
where CE denotes the cross entropy loss, Fx(x) is the probability vector of length |C| and ide is corresponding code index for the input pair (x, Recon denotes the label reconstruction loss in lieu of the negative log- liceeiod...od.
During training, we learn the prior encoding model po(c))), the posterior encoding model [s(c) y). the decoding model po x). together with the code book C in an end-to-end fashion. The posterior encoder plus decoder will learn good representation of the latent code, and the prior encoder will learn faithful uncertainty estimation from the stochastic training. At inference time, use the learned prior encoding model to output a conditional distribution given x on C, where each of the code will correspond to decoded label prediction with the associated probability.
5

---

Under review as a conference paper at ICLR 2021
3.2 MODEL DESIGN AND TRAINING
Our proposed framework in Section 3.1 is general and can be applied to a wide range of network architecture designs. Nonetheless it is worth discussing how the latent code с should be learned and utilized, so that the posterior encoding model can really learn the mode hypothesis of the given data-label pair, not simply memorize the data which causes an overwhelming number of codes and over-fitting One important principle is that the latent code should not be learned or used in spatially dependent manner, especially in pixel-level recognition tasks such as segmentation. This can also ensure that the prior encoding network is learning to solve the majority of the recognition problem while the latent code supplies only the additional but necessary information to reconstruct distinct outputs from the same (or similar) input. For this purpose we have adopted simple approach: the code output by posterior encoding network is obtained by global average pooling of its last layer; for the incorporation of the code into the intermediate feature of the decoding network that has spatial dimension (h, w) we will simply make h x w copies of the code, reshape it to have spatial size (h, w) and concatenate it to the corresponding feature.
In the experiments in Section 4 we will consider applications in computer vision and thus will use convolutional neural networks (CNNs) for the prior and posterior encoder, as well for the decoder, which together is similar to the U-Net architecture. Specifically, each encoding network consists of a sequence of downsampling residual blocks, and the decoder a sequence of upsampling residual blocks, where the decoder also has skip connections to the prior encoder by receiving its feature at each resolution level. The latent code is incorporated at a single level L into the decoder, which depends the task.
Suppose the latent code is c-dimensional. We initialize the code book as a size (nc.c) i.i.d ran- dom normal matrix, where each column represent an individual code. The statistics of the normal distribution is computed from the output of the posterior encoding network at initialization on the training data. We have found it to be beneficial since it allows the entire model to be initialized at a lower loss position on the optimization landscape. We have also found that the number of code utilized during training follows an interesting pattern: at the very first stage only very few codes will be used, then the number gradually grows to maximum before it slowly declines and becomes stable when the reconstruction loss plateaus. We therefore allow the network warm up in the first 17 epochs by training without the cross-entropy loss, since during this stage the number of uti- lized codes is unstable. This will not impair the learning of the posterior encoder, since receive no gradient information from the prior. We have found ne 512 to be well sufficient for all of our tasks, and the actual number of codes utilized after training is usually a fraction of it. Because of these observations, we did not try to explicitly enforce different codes to have different outputs, since for one reason the final number of codes are usually compact, and for the other we would like to allow different codes to have similar outputs, which means the possible situation where different hypotheses lead to similar predictions. We expect there will be some connection with the informa- tion bottleneck theory Tishby& & Zaslavsky (2015) and leave this direction for future work. We will release our open source implementation to promote future research.
4 EXPERIMENTS
To rigorously access our method's ability to approximate the distribution of one-to-many mappings, in Section 4.1 we first conduct a synthetic experiment with known ground truth conditional distribu- tions. In Section 4.2 we then demonstrate our method's performance on the realistic and challenging task of lesion segmentation on possibly ambiguous lesion scan on the LIDC-IDRI benchmark. In both experiments we compare with the state-of-the-an method Probabilistic U-Net Kohl et al. (2018) of the same model complexity as ours. More details on the experimental setting and additional re- sults can be found in Appendix A.
4.1 QUANTITATIVE ANALYSIS ON SYNTHETIC TASKS
MNIST guess game To test the ability for multi-modal prediction quantitatively, we design simple guessing game using the MNIST dataset LeCun & Cortes (2010) as follows. We are shown collection of images and only one of them is held by the opponent. The image be- ing held is not fixed and follows certain probability distribution. We need to develop a genera- tive model to understand the mechanism of which image is being held based on the previously

---

Under review as a conference paper at ICLR 2021
Mode probability predictions (ours)
Mode probability predictions (Prohabilistic Li-Ner)
0.8
Ground trut
Ground uh frequency
Prosicted mean
Prad stad meen
Mose prediction frequency
Mode nuttmal Count
10)
0.5
08
0.3
0.1
000
44
10
1b1
181
(a) Ours
(b) Probabilistic U-Net
Figure 2: Quantitative comparison on the MNIST guessing task. The small dots represent the pre- dictions for 1000 testing samples. Our method in (a) successfully produces accurate uncertainty estimate for each mode. Probabilistic U-Net uses conventional Gaussian latent parametrization, thus the sample's density provides no useful information about the uncertainty level, as shown the left axis. We also count the frequencies for each category and plot it the right axis. However, the approximation is far less accurate than ours even calculated on the entire testing dataset.
seen examples. In details, the input x will be an image that consists of four random digits, and belongs to one of the four categories: (A) (1, 2, 4); (B) (3, 4,5, 6); (C) (5, 6, 7, 8); (D) (7, 9, 0). The number represents the label of the image sampled. The output will be an im- age of the same size but only one of the input digit present. Specifically, for (A) the proba- bility distribution is (0.25, 0.25, 0.25, 0.25): for (B) (0.1, 0.4, 0.1, 0.4); for (C) (0.3, 0.5, 0.1, 0.1); for (D) (0.1, 0.1, 0.1, 0.7). Note that the distribution of the output, conditioned on each cate- gory's input, consists of four modes, and is designed to be different for each category. We re- quire the model to be trained solely based on the observed random training pairs (x,y). and thus no other information like digit categories should be used. The model would therefore need to learn to discriminate each category and assign the correct outputs with corresponding probabilities.
Thus for instance, an input image of Category (A) will be the combination of four random samples from Digit to 4 in that order, and the output can be
/ 2 3 4 2
/
1
2
3
4
3
the same digit 1 in the input with probability 0.25, or it can be the same digit 2 with probability 0.25. and forth. The images in the first row in Fig.3 il- lustrate an input image, where we also annotate the ground truth probability on the upper-left corner.
/
4 4
3
2
4
We trained our model samples from the training dataset of MNIST and tested on samples from the testing dataset, during both stages the random com- bination is conducted on the fly. The model here for demonstrating the results used total of 11 codes af- ter training. Please refer to Appendix A. for train- ing specifics and more results.
(a) Ours
(b) Probabilistic U-Net
Figure 3: Results visualization on the MNIST guess game. Note that our model can predict calibrated uncertainty estimation (e.g. "2" appears with probability 0.25 ver- sus our prediction 0.2634, best to zoom on screen to see the probability annotation in (a)). While Probabilistic U-Net cannot pre- diet such estimates and can produce non- sensible output from random sampling.
From the second to fifth row in Fig.3 we show the results of different models. Ours in Fig.3(a) are the top-4 predictions with explicit probability estimates annotated the upper-left in each row. For exam- plc, the second row has probability 0.2634, which is very close to the ground truth 0.25. In contrast, Probabilistic U-Net cannot rank its outputs and hence four random samples are drawn. Consequently one has little control when generating the samples and it's likely to obtain non- sen-sell outputs as in the fourth row of Fig.3(b).
7

---

Under review as a conference paper at ICLR 2021
Sradert
Srader2
Enaferd
Brador2
AJ222
A.1516
3.5428
K.004
A1244
0.130
R.007
0,301 2.000
(a)
(b)
Figure 4: Visualization of our results on the highly ambiguous samples from LIDC-IDRI dataset. The first row shows the input samples and their segmentations, and the next two rows show the top-8 predictions from our method. The uncertainty estimation for each segmentation proposal is annotated the upper-left corner.
Our method also performs much better quantitatively, as shown in Fig.2 with the results on 1000 random testing samples. We classify both models' outputs into the ground truth modes and aggregate the corresponding probabilities. We can see in Fig.2(a) that our method successfully discovered the distributional properties of the one-to-many mappings, and provides accurate uncertainty estimate. In contrast, due to the Gaussian latent parametrization, neither the individual density of each input provide useful information, shown by the left axis of 2(b. By the
nor their averages right axis of Fig.2(b) we also count the mode frequencies for each category for Probabilistic U-Net. However, calculated the entire testing dataset, the distribution approximation is still far from accurate compared to ours. Note that our method can directly output the uncertainty estimate for each imput accurately. This clearly demonstrates method's superiority and practical value.
4.2 REAL APPLICATIONS
Lesion segmentation of possibly ambiguous lung CT scans We use the LIDC-IDRI dataset pro- vided by Armato III et al. (2015; 2011 1); Clark et al. (2013), which contains 1018 lung CT scans from 1010 patients. Each scan has lesion segmentations by four (out of totally twelve) expert graders. The identities of the graders for each scan are unknown from the dataset. Samples from the testing set can be found the first row of Fig.4. As can be seen, the graders are often in disagreement about whether the scan contains lesion tissue. We hypothesize that the disagreement is due to the different assumptions the experts have about the scan. For example, judging from the sean's 's appearance, of the graders might have believed that the suspicious tissue is in fact normal tissue based on his/her experience, and thus gave null segmentation. There are also other possible underlying assumptions for the graders to come up with different segmentation shapes.
Our task is to identify such ambiguous scenarios by proposing distinct segmentation results from the corresponding latent hypotheses with their associated probabilities, which will be helpful for clini- cians to easily identify possible mis-identifications and ask for further examinations of the patients.
Our network architecture for this task is scaled-up version of the same model used in the MNIST guessing task. At training time, we randomly sample an CT scan x from the training set, and we randomly sample one of its four segmentations as y. The model we used to report the results has total of 31 codes. The specifics of the training and more results can be found in Appendix A.3.
Some sample testing results predicted by our model have high uncertainty are illustrated in Fig.4. The first row is the input and its four segmentations, and the last two rows are our top-8 predictions, where the probability associated to each latent code is annotated on the upper-left corner. We can that method can capture the uncertainty that is contained in the segmentation labels with notable probability scores, as well as other type of segmentations that seem plausible without further information.
Since no ground truth distribution for LIDC-IDRI dataset is available, quantitative evaluation has to be conducted differently from the MNIST guessing task. We follow the practice of Kohl et al. (2018) to adopt the generalized energy distance metric D found in Bellemare et al. (2017); Salimans al. (2018); Székely & Rizzo (2013), which is statistical quantity that measures the
8

---

Under review as a conference paper at ICLR 2021
LIDC-IDRI
2.00
1.50
125
para
0.75
0.30
0.5
0.0
Our
Probabllate U-M
mocel
(a) Quantitative comparison on LIDC
(b) Segmentation-code visualization
Figure 5: Results visualization on LIDC segmentation task. (a) The small dots represent test data instances' DRED values and the triangles mark the mean values. Our performance is competitive with the fate-of-thheare (b) We show the segmentation results for some frequently used code, annotated the bottom. The segmentation from the codes of probability 10 are left blank.
incoherence between two subsets of metric space. Please refer to Appendix A.2 for calculation details of this metric for segmentations based on Intersection over Union (IoU). The lower the value of Deg the closer the two subsets are. We report the results on the entire testing dataset in Fig.5(a). For model, the mean D of all testing data is 0.3354, the standard deviation is 0.2947. Our performance is thus competitive with that of Probabilistic U-Net, whose mean is 0.3470 and the standard deviation is 0.3139. Moreover, our model can give quantitative uncertainty estimate directly for each input scan, unlike Probabilistic U-Net that needs to perform sampling and clustering using metric such as IoU obtain uncertainty estimate.
Finally, we visualize some segmentation results for some frequently used codes in Fbb. The code used annotated at the bottom. The segmentations from the codes of negligible probability le.g.less than 10-4) are left blank. For example, the fourth column for code#1 may correspond latent hypothesis that leads the conclusion of no lesion, and the scan in the third row is not compatible with that particular latent hypothesis. It would be interesting future work to explore semantics of the latent codes if more information about the patient and the sean given.
DISCUSSION AND CONCLUSION
We have proposed MUE, novel framework learning one-to-many mapping with calibrated un- certainty estimation. As effeetive solution of the multi-modal posterior collapse problem, the discrete latent representations are learned to explain the corresponding types of input-output rela- tionship in one-to-many mappings. It also allows effectively and efficiently perform uncertainty estimation of the model prediction. We have extensively validated our method's performance and usefulness both synthetic and realistic tasks, and demonstrate superior performance the state-of-the-arr methods.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy. Fixing broken elbo. International Conference on Machine Learning, pp- 159-168, 2018.
Samuel G Armato III, Geoffrey Mcl.ennan, Lue Bidaut, Michael F McNitt-Gray, Charles Meyer, Anthony P Reeves, Binsheng Zhao, Denise R Aberle, Claudia I Henschke, Erie A Hoffman, et al. The lung image database consortium (lide) and image database resource initiative (idri): a completed reference database of lung nodules on ct scans. Medical physics, 38(2)9155-931, 2011.
Samuel G Armato III, Geoffrey McLennan, Luc Bidaut, Michael F McNitt-Gray, Charles Meyer, Henschke, Eric A Hoffman,
Anthony P Reeves, Binsheng Zhao, Denise R Aberle, Claudia I et al. Data from lidc-idri, 2015. URL https: pettpp::/wikk. /twkw/wiiii////.......... http://ww...aaarrnnreee www.ernnnccaa.eee maanceriiagii ngaa rchive net/ x/ /re.
Mare G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The cramer distance as a solution to biased wasserstein gradi- ents. arXiv preprint arXiv:1705. 10743, 2017.
Kenneth Clark, Bruce Vendt, Kirk Smith, John Freymann, Justin Kirby, Paul Koppel, Stephen Moore, Stanley Phillips, David Maffitt, Michael Pringle, et The cancer imaging archive (tcia): maintaining and operating a public information repository. Journal of digital imaging, 26(6): 1045-1057, 2013.
Yarin Gal. Uncertainty in deep learning. University of Cambridge, 133,20116 2016.
Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059, 2016.
Ian Goodfellow, Jean Pouget-Abadie, Mchdi Mirza, Bing Xu, David Warde-Farley. Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural infor- mation processing systems, pp. 2672-2680, 2014.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Ålexander Lerchner. beta-vae: Learning basie visual concepts with constrained variational framework. In ICLR, 2017.
Xun Huang, Ming-Yo Yu Liu, Serge Belongie, and Jan Kautz. Multimodal unsupervised image-to- image translation. Proceedings of the European Conference on Computer Vision (ECCV), pp. 172-189, 2018.
Eddy lgg, Ozgun Cicek, Silvio Galesso, Aaron Klein, Osama Makansi, Frank Hutter, and Thomas Brox. Uncertainty estimates and multi-hypotheses networks for optical flow. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 652-667, 2018
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with
conditional adversarial networks. CVPR, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizztion with gumbel-softmax. arXiv preprint 1rXhv:1611.01144 2016.
Alex Kendall and Yarin Gal. What uncertainties do we need bayesian deep learning for computer vision? In Advances in neurul information processing systems, pp. 5574-5584, 2017.
Alex Kendall, Vijay Badrinarayanan, and Roberto Cipolla. Bayesian segnet: Model uncertainty in deep convolutional encoder-decoder architectures for scene understanding. arXiv preprint arXiv:1511.02680 2015.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. URL http: //arxiv.orge /112 6980.
10

---

Under review as a conference paper at ICLR 2021
Diederik P. Kingma and Max Welling. Auto- Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014. URL httttt////// wwwxiv.xivvv../...aa.a.......... org/ abs/ 1312. 14.
Simon Kohl, Bernardino Romera-Paredes. Clemens Meyer, Jeffrey De Fauw, Joseph R Ledsam, Klaus Maier- Hein, SM Ali Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A probabilis- tic u-net for segmentation of ambiguous images. In Advances in Neural Information Processing Systems, pp. 6965-6975, 2018.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predic- tive uncertainty estimation using deep ensembles. In Advances in neural information processing systems, pp. 6402-6413, 2017.
Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URLhttp://yanr. /mmist/.
lecun. cwm/
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv: 1.0::///11.00222 2016.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. arXiv preprint
2014.
Andriy Mnih and Danilo J Rezende. Variational inference for monte carlo objectives. arXiv preprint arXiv:1602.000 06725, 2016.
Eric Nalisnick, Lars Hertel, and Padhraic Smyth. Approximate inference for deep latent gaussian mixtures. In NIPS Workshop on Bayesian Deep Learning, volume 2, 2016.
Ali Razavi, Aaron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with daesa-vaes. In International Conference on Learning Representations, 2018
Ali Razavi, Aaron den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems, pp. 14837-14847, 2019.
Danilo Jimenez Rezende and Fabio Viola. Taming vaes. arXiv preprint arXIv:1810.00597, 2018.
Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE International Conference on Computer Vision, pp. 3391-36600, 2017.
Tim Salimans, Han Zhang, Alec Radford, and Dimitris Metaxas. Improving gans using optimal transport. In International Conference on Learning Representations, 2018.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in neural information processing systems, pp- 3483-3491 2015.
Akash Srivastava, Lazar Valkov, Chris Russell, Michael U Gutmann, and Charles Sutton. Veegan: Reducing mode collapse in gans using implicit variational learning. In Advances in Neural Infor- mation Processing Systems, pp. 3308-3318, 2017.
Gábor J Székely and Maria L Rizzo. Energy statistics: A class of statistics based on distances. Journal of statistical planning and inference, 1239)):12491127 2013.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop (ITW), pp. 1-5. IEEE, 2015.
Jakub Tomezak and Max Welling. Vac with a vampprior. In International Conference on Artificial Intelligence and Statistics, pp. 1214 1243, 2018.
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. Con- ditional image generation with pixelcnn decoders. In Advances in neural information processing systems, pp. 4790-4798, 2016.
11

---

Under review as a conference paper at ICLR 2021
Aaron den Oord, Oriol Vinyals, al. Neural discrete representation learning. In Advances in 2017.
Neural Information Processing Systems, pp.
Tiancheng Zhao, Kyusong Lee, and Maxine Eskenazi. Unsupervised discrete sentence representa- tion learning for interpretable neural dialog generation. arXiv preprint arXiv: 08069, 2018.
Chuanxia Zheng. Tat-Jen Cham, and Jianfei Cai. Pluralistic image completion. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 1438-1447, 2019.
Jun- Yan Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. Proceedings of the IEEE international conference computer vision, pp. 2223-22332 2017a.
Jun- Yan Zhu, Richard Zhang, Deepak Pathak, Trevor Darrell, Alexei A Efros, Oliver Wang. and Eli Shechtman. Toward multimodal image-to-image translation. In Advances in neural information processing systems, pp. 465-476, 2017b.
12

---

Under review as a conference paper at ICLR 2021
APPENDIX
A EXPERIMENTAL DETAILS
As described in the main text, each of our encoding network consists of sequence of downsampling residual blocks, and the decoder sequence of upsampling residual blocks. The decoder also has skip connections to the prior encoder by receiving its feature at each resolution level. A residual block consists of three convolution layers. As the shortcut connection, the input is added to the output if the input and the output have the same channel size, or otherwise a 1 convolution is applied before the addition. Bi-linear down sampling and up sampling is applied before the residual blocks if the spatial size is changed. We fix the le penalization weight B = 0.25, number of initial candidate codes ne 512, and use the Adam optimizer Kingma & Ba (2015) with its default setting for all of our experiments. Hyper-parameters specific to each experiment are detailed in the following subsections.
A.1 MNIST GUESSING GAME
We use 6 layers in the prior and the posterior encoding networks, with output channel dimension [16, 32, 64, 128] + [128, 128). This notation means that the first 4 levels are used as in the U-Net which feed the decoder, and the last 2 levels used for the latent code learning. The posterior encoder further uses a 1 x 1 and global average pooling to obtain the code. The code is of dimension 128. The prior encoder uses linear layer learn the distribution on Our decoder has output channel dimension [64, 32, 16, 1]. We incorporate the code at the bottom level, namely the 1-st layer of the decoder. For the Probabilistic U-Net, since the architecture different, we used structure of similar capacity, with the parameter num num_fitte= [16, 32, 64, 128) in its released version, and we find the suggested hyperparameters in Kohl et al. (2018) for LIDC task works well in this case. For both networks, use the binary cross entropy loss, batch size of 256, and use the learning rate schedule [1e 5e 1eeb, 5e at [0, 30k, 90k, 120k] iterations.
Some additional results from model and Probabilistic U-Net are shown in 6ig.6 and Fig.7.
A.2 GENERALIZED ENERGY DISTANCE METRIC FOR SEGMENTATIONS
Denote Vx. Sx C y to be the set of segmentation labels and the set of model predictions corre- sponding to the scan x, respectively. ) equipped with the metric
d(y,s) IoU(y, s),
where loU( is the intersection--ver--ue operator that is suitable for evaluating the similarity between segmentations. The DIED statistic in our case is defined to be
('x, =
Σ Σ Pspydly,s)
Σ Σ dypy y'
Σ
Σ PsPs d's,s'),
vEYx
S'SxN'ES
where ps is model's probability prediction for the output and Py is the ground truth probability. In case the ground truth is not available like LIDC-IDRI, we use Py |3xx1' where [>x denotes the cardinality of V/-- For our model, we choose Sx to be the top-N predictions. To be rigorous we normalize the sum of their probabilities to be (though which in fact has negligible effect since N are usually chosen so that probability always almost sum up to 1). the case of Probabilistic U-Net, we use N random output samples and Pa is replaced by [.
N.
A.3 LIDC-IDRI SEGMENTATION
We use 6 layers in the prior and the posterior encoding networks, with output channel dimension [32, 64, 128, 192] + [256, 512]. The code is of dimension 128. Our decoder has output channel dimension [128, 64, 32, 1]. We incorporate the code at the bottom level, namely the 1-st layer of the decoder. For the Probabilistic U-Net, since the architecture is different,
used structure of similar capacity, with the parameter mum filter= [32, 64, 128, 192] in its released version, and follows the suggested hyperparameters in Kohl al. (2018). For both networks,
the binary
13

---

Under review as a conference paper at ICLR 2021
5
6
8
7
8
9
0
7
8
9
O
2
3
4
5
6
7
6
0
O
2
6
5
9
9
3
5
8
8
7
7
7
7
8
4
8
8
0
O
4
8
6
8
8
2
6
3
555
234
/
2
3 4
5
6
8
n
3
4
4
2
2
6
6
3
3
5
6
3
3
4
4
8
4
4
4
8
4
4
2
2
60
A.
7
8
9
0
3
4
5
8
9
0
3
4
5 6
7
8
1890
0
41
0
4
0
0
6
O
6
8
8
3
7
3
9
9
5
8
5
.7
6
9
6
O
8
L
8
4
8
3
4
156 6
5
6
7
8
7
8
9
3
4
5
7
8
0
4
6
4
0
6
5
O
6
0
7
9
3
9
5
8
8
S
7
6
8
7
6
8
4 6
6
8
4
8
5
7
8
7
8
9
0
7
8
9
0
5
6
7
8
7
8
9
6
0
0
6
5
0
0
5
9
9
8
7
7
8
7
7
8
9
7
J
9
8
Figure 6: Results from model the MNIST guessing task. Top-6 results with the predicted uncertainties are shown.
14

---

Under review as a conference paper at ICLR 2021
2
3
4
7
8
2
3
4
23 3
4
8
9
O
3
8
6
3
2
4
4
4
8
>
8
4
2
4
7
4
2
4
6
4
788
9
O
7
8
9
o
3
4
S
6
1
2
3
4
3
456
9
4
2
4
7
4
3
3
8
4
4
5
O
6
4
4
8
O
7
4
3
3
7
9
3
4.
7
8
9
0
2
3
4
7
8
9
0
3
4
5
6
7
8
9
0
4
0
6
9
0
6
8
0
9
4
2
p
4
8
4
8
5
O
2
9
6
7
3
123
4
3
4
5
6
S
6
7
8
5
6
7
8
5
60
7
8
4
3
6
5
5
3
6
5
8
6
2
5
g
8
6
2
3
5
5
6
2
3
6
8
2
6
8
8
60
8
7 8
9
0
7
$
9
O
6
7
8
/
2
3
4
3
4
5
7
6
4
9
7
$
4
6
7
3
6
7
00
7
4
4
7
E
4
6
Figure 7: shown.
Results from Probabilistic U-Net the MNIST guessing task. 6 random samples
15

---

Under review as a conference paper at ICLR 2021
Grader
Ginder
Gracera
Gradert
Grader2
Graterr
6.600
A-160
3.0667
06611
S.132
8,1861
6-1529
6:126
646666
18887
2rader:
Srader
Rnader1
Grader
0.1076
BLETAA
0.645
1.0281
5.3314
6.1122
6.311
Gisder
orader
Surdee
Sunder
araterd
Hraaerrt
arater
uraters
3-4079
8.2470
0.0817
1.091
8:1551
A:115
A:189
21812
5.2418
6.1111
8,1720
6.1878
U.C001
1.0000
53828
8:2812
0.3220
3.3155
Figure 8: Results from our model on the LIDC-IDRI segmentation task. Top-8 results with the predicted uncertainties are shown.
cross entropy loss, a batch size of 256, and use the learning rate schedule [1e-1, 11e-1,5e-b 55 1e-5 5e 56 at [0, 30k, 90k, 120k] itcrations.
Some additional results from model and Probabilistic U-Net are shown in F,g.8,9 and Fig. 10, 11, respectively.
16

---

Under review as a conference paper at ICLR 2021
brader
Sander
braerr
brblert
tueeen
5.3173
6-817
Grater 8
Grader
Goadee
Griadee
Gracer
Kratert
Grader Am
Grader3
3.1266
A.108
6.1321
6.1410
0.2118
31100
0.1790
1.1200
5.8111
R1120
6.5488
5.348
5.0111
R.T11
3.5500
5.3338
8:117
K:3125
Grader
Grader
Gender
Guader
Gracer
Wratert
Brsder
Grater
Figure 9: Results from our model on the LIDC-IDRI segmentation task. Top-8 results with the predicted uncertainties shown.
17

---

Under review as a conference paper at ICLR 2021
braders
Sanerr
arser1
bralert
Grader 8
Grader
Grnderr
Gracera
Gextert
Grader2
Grater3
Grader
Grader
Giader
Bracerd
Wratert
Brsder
brater
Figure 10: Results from Probabilistic U-Net on the LIDC-IDRI segmentation task. random sam- ples are shown.
18

---

Under review as a conference paper at ICLR 2021
brader
Sander
brsern
brotert
Geater 8
Grader
Gradee
Grander
Gracer
Krxtert
Grader Ag
Grater3
Grader
Grader
Grnder
Giader
Gracerd
Wratert
Brsder
drater
Figure 11: Results from Probabilistic U-Net on the LIDC-IDRI segmentation task. random sam- ples are shown.
19


Reviews:
[{"review_id": "wr02ZBmizNL", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a conditional VAE like framework to learn the one-to-many mappings between input and output, leading to an application of uncertainty estimation. Technically, the novel part is to utilize a deterministic (delta) distribution for approximate posterior.\n\nSome flaws may need future attention:\n\n1. In the introduction paragraph staring with \"Let us recall that one key ingredient of the VAE framework\", the main idea is understandable: discrete latent code has definitely advantages in coping with multimodal distributions than a Gaussian distribution which is in nature single mode.  However, its explanation is confusing as following: \n\tIn VAE, let's say we minimize KL divergence KL(P(c|x,y)|P(c|x)), where P(c|x,y) has two modes, P(c|x) has single mode. When minimization is successful, p(c|x) will spread out like figure 1(a), rather than 1(b). This is related to the difference between forward KL and reverse KL. However, it seems this paragraph suggests 1(b) as posterior collapse.\n\n2. Using discrete latent code is not new in VAE community. There are previous works (e.g. https://arxiv.org/pdf/1804.08069.pdf) noting that naively learned discrete code c for p(y|x,c) can not be interpreted alone, but need interpreted together with input x. Such statement argues this paper's novelty and contribution. \n\n3. Some typos and minor flaws, such as unclear references of figure rows in figure 3.\n\nOverall, the reviewer thinks this paper need some revisions for it to be more shining.", "extra_scores": {}, "extra_text": {}}, {"review_id": "vMfjvaJQypg", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This manuscript proposes to measure the \"modal uncertainty\" in conditional generative models by forcing a discrete latent intermediate representation (here, C), between inputs X and outputs Y. By then manipulating the estimated categorical distribution likelihoods, an uncertainty estimate can be produced.\n\nThis is, in essence, the kludge of the VQ-VAE (van den Oord et al. 2017) categorical latent variable model into the Probabilistic U-Net (Kohl et al. 2018), which had previously used the Gaussian latent variable model. Empirically this difference is important, as the Gaussian VAE specifies a single mode latent distribution, which is incorrect in many cases.\n\nStrong points:  \n*Presents a clear argument for VQ-VAE style latents vs. Gaussians (...non-unimodal uncertainty representation).  \n*Constructs and executes examples from a subset of previous literature (LeCun & Cortez 2010, Kohl et al. 2018) that illustrate empirical effectiveness.  \nWeak points:  \n*Limited characterization of uncertainty (i.e. uncertainty of a label image is one-of-K, or combinations thereof, and not uncertainty w.r.t. model/training/etc).  \n*Possibly incorrect claims (prior-free, disentanglement, Prob. U-Net ordering, VQ-VAE sampling(?)).  \n*Possible limited scope from the chosen definition of \"uncertainty\" (no model uncertainty, no parameter uncertainty).  \n\nI think this paper is marginal for this venue. It improves on the Prob. U-Net by borrowing improvements to VAE from VQ-VAE. While it does seem to better characterize label uncertainty, it is better suited to a venue where such improvements have intrinsic importance (e.g. radiology), and moreover importance that can be measured (i.e. experiments asking whether radiologists will use Non-Top-1 segmentations?). However, it appears to be an improvement on the prior art in my opinion, and I believe it would have some interest to the community as a poster.\n\nOne overall question: are we learning uncertainty of segmentations/labels, or are we simply learning the biases of the varying raters/radiologists/human segmenters, or are these two concepts indistinguishable? Supposing the number of codes equals the number of raters, why should $p(c|x)$ differ from $p(r)$? In an imbalanced rater case, why would an ordering of $p(c|x)$ not simply reflect the likelihood of one rater or another.\n\nThis seems to speak less about uncertainty in the segmentation w.r.t. the image and model parameters, and instead speak about biases in the segmentation w.r.t. the raters. While this is clearly also important, does this capture notions of uncertainty beyond the uncertainty of who is rating a certain image? Or is that only and exactly the uncertainty that the authors are attempting to capture?\n\nDiscrete Representation and Associated Loss\n----\n\nAs expected, there is significant discussion and derivation devoted to the discrete representation. This essentially mirrors the short derivations in the probabilistic u-net paper, except using a discrete likelihood.\n\nFrom the loss in Eq. (4) we have 1) a distribution matching term between $p_\\theta(c|x)$ and $q_\\phi(c|x,y)$, 2) a reconstruction term $D_\\theta$ given the output of the E(x,y) encoder and $x$ (using $\\arg\\min || c' - E_\\phi (c|x,y)||$ for $c$), and then 3) the VQ-VAE \"stop gradient\" term.\n\nIn the derivation slightly higher on the page, we assumed that $q_\\phi(c|x,y)$ is deterministic (i.e. $q_\\phi(c|x,y)$ is $1$ for one value of $c$ and otherwise $0$). While this results in a nice cross entropy term, why is it necessary? We otherwise have a still tractable difference between $q_\\phi(c|x,y)$ and $p_\\theta(c|x)$, summed over $c$. Is there added benefit to the forced discretization here?\n\nPerhaps it is useful to sample $c \\sim q_\\phi(c|x,y)$ or use the mode $c = \\arg\\min || c' - E_\\phi (c|x,y)||$ (this inducing an implicit gaussian likelihood structure on the embedding space), but it's not clear from just the theory that this is true.\n\nWhy is it reasonable to always use the mode reconstruction when training a loss (again, the argmax from VQ-VAE's embedding scheme), and then look at the other categories? Why should the outputs from the non-mode categories be reasonable (up to their likelihood)? (this question is in contrast to the convergence of p(c|x) to a rater distribution p(r)).\n\nThe authors claim that as in van den Oord et al 2017 there is no shrinkage to a prior over $c$. Is the regularization term $\\beta || E(x,y) - sg[c]||$ not such a prior? Doesn't this enforce a structure on the co-domain of E(x,y)? Are the embedding vectors determining $c'$s learned? If not, isn't this a prior?\n\n[minor] The authors use the terms \"prior encoder\" to describe $p_{\\theta}(c|x)$, a learned network, and similarly \"posterior encoder\" to describe $q_\\phi(c|x,y)$, also a learned network. Reading the probabilistic U-Net paper, it appears that the authors of that paper also use this terminology. It is a matter of viewpoint, but I personally think these terms conflict with \"the usual terminology\", where $p(c)$ would be \"the prior\", to which we shrink $p(c|x)$ to. Further, this implies a causal order to $x$ and $y$; this is the case with the empirical examples (tumor labels on CT images, etc.), but may not always be the case.\n\nU-Net Comparison\n----\n\nThe authors claim that \"Probabilistic U-Net\" outputs cannot be ranked; is this actually the case? It is my understanding that the difference between the proposed method and the Prob. U-Net method in architecture is the replacement of their Gaussian latent C (in their paper, Z) with a discrete categorical C. This, alongside the different loss, contains the majority of the changes.\n\nCould the Prob. U-Net have outputs ranked by their point-likelihoods? The C output are conditionally Gaussian, so why not use the $p(c|x)$ likelihood? Understandably, generating diverse samples from this is not as simple as querying the different codes in the proposed method, but it seems incorrect to say that the samples from the Prob. U-Net are unordered.\n\nOther Questions\n----\n\nThere are several mentions of \"disentanglement\". In what way are these representations disentangled? They're certainly categorical by construction, but both the outputs and the actual labels are highly correlated between different categories.\n\nAt the bottom of page 3 there is a claim about sampling from a VQ-VAE being auto-regressive? Is the particular decoder architecture in VQ-VAE important w.r.t. the theory discussion directly adjacent to that statement? Surely VQ refers to the Vector-Quantization phase, which is agnostic of the decoder architecture, auto-regressive or otherwise? This seems incorrect. Similarly, there is a claim that VQ-VAE was introduced to avoid\n\n> noise sampling, which is a different cause than ours which usually results in blurriness  \n\nIs this the cause of blurriness (not, e.g. the L2 loss in the decoder?)? And is the introduction of VQ not the same issue the authors here are attempting to address (misspecification of a single-mode latent distribution), not the additive noise sampling?\n\nRecommendations for improvement\n----\n-Remove the disentanglement sentences.  \n-Resolve questionable claims.  \n-Include an experimental case where the number of raters is larger than the number of latent categories.  \n\nEdit after Rebuttal/Response Period:\n----\n\nFirst, an apology to the authors that a dialogue did not occur during the response period; the authors response was prompt, and my (R4) response was not, thus they were not given an opportunity to respond this response to their response to the review (...the number of recurrences may indicate why this was not possible, given limited reviewer time resources).\n\nI think the authors misunderstand my questions about the nature of the uncertainty they're capturing:  \nIs there intrinsic uncertainty in the observed phenomena (e.g. medical images of tissues), are we capturing mixture proportions of deterministic states which have been mixed due to quantization, OR is the uncertainty due to the raters, i.e. found in the labels ONLY due to differences in label generation?  OR, a third case, is this moot because it does not change the outcome?\n\nI understand that there is no explicit modelling of raters. However, my concern was that what we are capturing is intrinsically the uncertainty due to raters, even though no actual rater indicator variable was provided. This would be analogous to learning, unsupervised, the writers of the various MNISTs digits. While for MNIST this is surely difficult due to the number of writers (and their anonymity), for medical images we will likely have a limited number of raters. Having the posterior code collapse to a rater indicator appears problematic, not a desirable outcome, and likely if any one rater has correlated outputs across samples (which seems reasonable; some raters may be more or less conservative with their tissue labeling, boundaries, better/more careful at delineating curves etc). What prevents the capture of this signal, or is this the actual variation we intend to capture in the first place?\n\nR3 further included this interesting question in their initial review:\n> In the shown samples, for a given input, many of the different outputs seem very similar and could be considered from the same mode. This can potentially make interpreting the probabilities more difficult than claimed in the paper, especially since the model is trained with a large number of codes. A very plausible scenario could be that one of the most likely modes is split between multiple low probability outputs and thus doesn't show up on the top ouput [sic]. Can the authors comment on this potential issue?\nIf outputs are correlated i.e. overlapping in the original image domain, should the measured uncertainty be aggregated across codes?\n\nI disagree with the characterization of the Gaussian VAE calibration in the \"Ranking Probabilistic U-NET\" response section. Simply because it does not accurately fit the function (one mode vs. many) does not mean we can't evaluate the learned unimodal beliefs. Yes, it is misspecified. Does this mean the rankings are meaningless? Surely the discrete model is misspecified (there are more possible masks than codes), but the ranking is claimed to be meaningful. The discrete model may have a better fit, and make the argument that it is better specified, but this doesn't mean you _can't_ evaluate the Gaussian VAE.\n\nI stand by my initial rating and reasoning, though I note to the AC that, given space, this could make an acceptable poster. It is, in my opinion and in gross summation, an improvement on the Prob. U-Net by way of improving the Gaussian VAE sub-model of the Prob. U-Net to the VQ-VAE. This allows for sampling from a discrete set of codes which hopefully correspond with modes of the generating distribution in the data domain, instead of sampling from a parametric density, which, while continuous, has only one local maximum.", "extra_scores": {}, "extra_text": {}}, {"review_id": "d47KsYwBrLc", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\n\n################################################\n\nStrong points:\n\n- The paper is clear and easy to follow. It is well-motivated and does a good job at highlighting the multi-modal posterior collapse problem.\n\n- The model outperforms the prior state-of-the-art on both a synthetic and realistic task.\n\n\nWeaknesses:\n\n- Although the application is very different, the proposed model is very similar, albeit lighter, to a prior work (cited by the authors).\n\n- In the shown samples, for a given input, many of the different outputs seem very similar and could be considered from the same mode. This can potentially make interpreting the probabilities more difficult than claimed in the paper, especially since the model is trained with a large number of codes. A very plausible scenario could be that one of the most likely modes is split between multiple low probability outputs and thus doesn't show up on the top ouput. Can the authors comment on this potential issue?\n\n\n################################################\n\nScore motivation:\n\nWhile the method is not particularly novel, the authors apply it in a way that could be of interest to the community.\nBesides, the behavior of VAEs with discrete latent space is a relevant topic that is little explored in the literature.\n\n\n################################################\n\nOther question:\n\nI am curious about the properties of the latent space succinctly mentioned at the end of section 4. What other properties have the authors observed? Also, do similar outputs tend to be represented with codes that are close in the latent space even in the case of discrete representations?\n\n\nMinor typo:\n\n3rd paragraph of the introduction: two consecutive commas", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a conditional discrete VAE for uncertainty estimation on high-dimensional data. Reviewers found the paper borderline, and two of the three reviewers stated it doesn't meet the acceptance bar due to lack of clarity in several aspects and limited technical novelty.
Decision: Reject