Modeling Multimodal Aleatoric Uncertainty in Segmentation with Mixture of Stochastic Experts
Authors: Zhitong Gao
Yucong Chen
Chuyu Zhang
Xuming He
Venue: ICLR 2023 poster (2023)
Keywords: Semantic Segmentation
Aleatoric Uncertainty
Stochastic Segmentation
Multiple Annotations

Abstract: Equipping predicted segmentation with calibrated uncertainty is essential for safety-critical applications. In this work, we focus on capturing the data-inherent uncertainty (aka aleatoric uncertainty) in segmentation, typically when ambiguities exist in input images. Due to the high-dimensional output space and potential multiple modes in segmenting ambiguous images, it remains challenging to predict well-calibrated uncertainty for segmentation. To tackle this problem, we propose a novel mixture of stochastic experts (MoSE) model, where each expert network estimates a distinct mode of the aleatoric uncertainty and a gating network predicts the probabilities of an input image being segmented in those modes. This yields an efficient two-level uncertainty representation. To learn the model, we develop a Wasserstein-like loss that directly minimizes the distribution distance between the MoSE and ground truth annotations. The loss can easily integrate traditional segmentation quality measures and be efficiently optimized via constraint relaxation. We validate our method on the LIDC-IDRI dataset and a modified multimodal Cityscapes dataset. Results demonstrate that our method achieves the state-of-the-art or competitive performance on all metrics.

Published as a conference paper at ICLR 2023
MODELING MULTIMODAL ALEATORIC UNCERTAINTY IN SEGMENTATION WITH MIXTURE OF STOCHASTIC EXPERTS
Zhitong Gao', Yucong Chen', Chuyu Zhang Xuming Het.3
ShanghaiTech University, Shanghai, China Lingang Laboratory, Shanghai, China Shanghai Engineering Research Cente of Intelligent Vision and Imaging, Shanghai, China (gaozht, chenys, zhangchy2, hexm www.neeeiannhhhihheeecmoo edu.cn
ABSTRACT
Equipping predicted segmentation with calibrated uncertainty is essential for safety-critichs applications. this work, we focus on capturing the data-inherent uncertainty (aka aleatoric uncertainty) in segmentation, typically when ambigu- ities exist in input images. Due to the high-dimensional output space and po- tential multiple modes in segmenting ambiguous images, it remains challenging to predict well-calibrated uncertainty for segmentation. To tackle this problem, we propose novel mixture of stochastic experts (MoSE) model, where each ex- pert network estimates distinct mode of the aleatoric uncertainty and a gating network predicts the probabilities of an input image being segmented those modes. This yields an efficient two-level uncertainty representation. To learn the model, we develop a Wasserstein-like loss that directly minimizes the distribution distance between the MoSE and ground truth annotations. The loss can easily integrate traditional segmentation quality measures and be efficiently optimized via constraint relaxation. We validate method on the LIDC-IDRI dataset and a modified multimodal Cityscapes dataset. Results demonstrate that our method achieves the Hate-of-the-art or competitive performance all metrics.
INTRODUCTION
Semantic segmentation, a core task in computer vision, has made sigmificant progress thanks to the powerful representations learned by deep neural networks. The majority of existing work focus on generating single (or sometimes fixed number of) segmentation output(s) to achieve high pixel-wise accuracy (Minaee al.) (02T). Such a prediction strategy, while useful in many scenar- ios, typically disregards the predictive uncertainty in the segmentation outputs, even when the input image may contain ambiguous regions. Equipping predicted segmentation with calibrated uncer- tainty, however, is essential for many safety-critical applications such medical diagnostics and autonomous driving to prevent problematic confidence decisions (Amodei et al. 2016).
An important problem of modeling predictive uncertainty in semantic segmentation is capture aleatoric uncertainty, which aims to predict multiple possible segmentation outcomes with cali- brated probabilities when there exist ambiguities in input images (Monteiro et al., 2020). In lung nodule segmentation, for example, an ambiguous image can either be annotated with a large nod- ule mask non-nodule with different probabilities (Armato III al. 2011). Such problem can be naturally formulated as label distribution learning (Geng & Ji, 2013), of which the goal is to estimate the conditional distribution of segmentation given input image. Nonetheless, due to the high- nigh-dimensional output space, typical multimodal characteristic of the distributions and limited annotations, it remains challenging predict well-calibrated uncertainty for segmentation.
Most previous works (Kohl et al. 2018 2019 Baumgartner et al. 2019 Hu et al. 2019) adopt the conditional variational autoencoder (cVAE) framework (Sohn al. (015) to learn the predictive
distribution of segmentation outputs, which has a limited capability to capture the multimodal dis- tribution due to the over-regularization from the Gaussian prior or posterior collapse (Razavi et al.)
Code available at http://iiiaiiiieooeee//e////////

---

Published as a conference paper at ICLR 2023
2018 Qiu & Lui 2021). Recently, Monteiro et al. (2020) propose multivariate Gaussian model in the logit space of pixel-wise classifiers. However, it has to use a low-rank covariance matrix for computational efficiency, which imposes a restriction problem, Kassapis et (202T) employ the adversarial learning strategy to learn an implicit density model, which requires additional loss terms to improve training stability (Luc et al.) 2016, [Samson et al. 2019). Moreover, all those methods have to sample large number of segmentation outputs to represent the predictive distribution, which can be inefficient in practice.
its modeling capacity. To alleviate this
In this work, we aim to tackle the aforementioned limitations by explicitly modeling the multimodal characteristic of the segmentation distribution. To this end, we propose novel mixture of stochastic experts (MoSE) (Masoudnia & Ebrahimpour 2014) framework, in which each expert network en- codes distinct mode of aleatoric uncertainty the dataset and its weight represents the probability of an input image being annotated by a segmentation sampled from that mode. This enables us to decompose the output distribution into two granularity levels and hence provides an efficient and more interpretable representation for the uncertainty. Moreover, we formulate the model learning as an Optimal Transport (OT) problem (Villani 2009), and design a Wasserstein loss that directly minimizes the distribution distance between the MoSE and ground truth annotations.
Specifically, our MoSE model comprises a set of stochastic networks and gating network. Given an image, each expert computes semantic representation via a shared segmentation network, which is fused with a latent Gaussian variable to generate segmentation samples of an individual mode. The gating network takes semantic feature embedding of the image and predicts the mode probabilities. As such, the output distribution can be efficiently represented by a set of samples from the experts and their corresponding probabilities. To learn the model, we relax the original OT formulation and develop an efficient bi-level optimization procedure to minimize the loss,
We validate our method on the LIDC-IDRI dataset (Armato III et al. 2011) and a modified mul- timodal Cityscapes dataset (Cordts et al. 2016; Kohl et al., (0118). Results demonstrate that our method achieves the state-of-the-art or competitive performance on all metrics.
To summarize, our main contribution is three-folds: (i) We propose a novel mixture of stochastic experts model to capture aleatoric uncertainty in segmentation; (ii) We develop Wasserstein like loss and constraint relaxation for efficient learning of uncertainty; (iii) Our method achieves the state-of-the-are or competitive performance on two challenging segmentation benchmarks.
2 RELATED WORK
Aleatoric uncertainty in semantic segmentation Aleatoric uncertainty refers to the uncertainty inherent in the observations and can not be explained away with more data (Hällermeier & Waege- man 2021). ) Early works mainly focus on the classification task, modeling such uncertainty by or post-training with calibrated loss (Sen- soy et al. 2018 Guo et al., 2017] Kull et al. 2019). In semantic segmentation, due to the high- dimensional structured output space, aleatoric uncertainty estimation typically employs two kinds of strategies. The first kind (Kendall& & Gal 2017] Wang et al. 2019 Ji et al. 2021) ignore the structure information and may sufter from inconsistent estimation. the second line of research, Kohl et al. (2018) first propose to learn a conditional generative model via the cVAR framework (Sohn et al.] 2015 Kingma & Welling 2014), which can produce joint distribution of the structured outputs. then improved by Baumgartner et al. (2019); Kohl et al. (2019) with hierarchical latent code injec- tion. Other progresses based on the eVEE framework attempt to introduce an additional calibration
exploiting Dirichlet prior (Malinin & Gales 2018 2019)
loss the sample diversity (Hu et al. 2019), change the latent code discrete (Qiu & Lui| 2021). or
augment the posterior density with Normalizing Flows (Valiuddin et al.) 2021). ). Recently, Monteiro
let al. (2200) directly model the joint distribution among pixels in the logit space with multivariate Gaussian distribution with low-rank parameterization, while Kassapis et al. (202)) adopt a cascaded framework and adversarial loss for model learning. Different from those methods, adopt an
explicit multimodal framework training with a Wasserstein- loss, which enables us to capture the multimodal characteristic of the label space effectively.
Optimal Transport and Wasserstein distance The optimal transport problem studies how to transfer the mass from one distribution to another in the most efficient way (Villani 2009). The corresponding minimum distance is known as Wasserstein distance Earth Mover's distance in
2

---

Published as a conference paper at ICLR 2023
Latent Priors
Samples
Annotations
3
(1,
2
Decoder
2
palRis
2. P2
Encoder
6.3
Image
3
:
Decoder
10x
3
PM
01 OT-based loss,
Gating Network
Figure/ 1: Given an ambiguous image input, our model consists of several expert networks and a gating network and produces segmentation samples with corresponding weights. It represents the aleatoric uncertainty at two granularity levels, (e.g. the input image be annotated either "2" loss
or "3" and also has variations near boundaries.) and is trained by an Optimal-Tranpport function which minimizes the distance between the MoSE outputs and ground truth annotations.
some special cases (Rüschendorf) 1985,, which has attracted much attention in the machine learn- ing community (Torres et al., (021). In multi-class segmentation, Fidon et al. (2017: 2020) adopt the Wasserstein distance at the pixel level to inject prior knowledge of the inter-class relationships. We apply the Optimal Transport theory for calibrating the aleatoric uncertainty in segmentation, where a Wasserstein-like loss between the predicted distribution and ground truth label distribution minimized. Unlike most work either dealing with fixed marginal probabilities (Daoodaran et al.) 2020) or fixed locations (Frogner et al., 2015) in the Optimal Transport problem,
parameter- ized marginal probabilities and jointly optimize marginal probabilities and segmentation predictions. Furthermore, we design novel constraint relaxation efficiently compute and optimize the loss.
Mixtures in deep generative models There has been a series of work exploring mixture models in VAE (Dilokthanakul et al., [016)] Diang al., 2017) or GAN (Yu & Zhou 2018: Ben- Yosef & Weinshall 2018) framework. However, VAEs require learning posterior network while GANs need to train a discriminator and often suffer from training instability. By contrast, our method does not require any additional network for training and utilize OT-based loss that can be easily integrated with traditional segmentation loss and be efficiently optimized.
Other multiple output problems There have been many works in the literature predicting multiple structured outputs (Guzman-Rivee Rivera et al. 2012, Rupprecht et al. 2017 Ilg et al. 2018). Unlike ours, those models are not necessarily calibrated and can only produce limited number of predictions. Another similar task is image-ta magee-to-image translation (Isola et al. (017). which considers transforming images from one style to another and is addressed by learning conditional generative model. However, these methods (Isola et al., 2017 Zhu et al.) 2017 Esser et al. 2018) seek more diverse images and place less emphasis on calibration.
contrast, we aim to model a well-calibrated conditional distribution segmentation that can generate varying segments during inference.
3 METHOD
In this section, we introduce our method for capturing the aleatoric uncertainty in semantic segmen- tation. Specifically, we first present the problem formulation in Sec. Sec.B.1] then describe the design of our MoSE framework in Sec. 3.2 and our OT formulation of the model training in Sec. 3.3
3.1 PROBLEM SETUP AND FORMULATION
We aim to build a flexible model for capturing the multimodal aleatoric uncertainty in semantic seg- mentation, where each input image could be annotated in multiple ways due its ambiguity. To this end, we formulate it as the problem of estimating the probability distribution of segmentation outputs given input image. Formally, assume annotated image dataset D {Ir, yo un Jn=l.--.
(8) (6) i=1.-
3

---

Published as a conference paper at ICLR 2023
is available for training, where In, E RR" IRxxxx3 x3 is image and has M (2 1) distinct segmentation
annotations yn yYn (AA with corresponding relative frequencies vn , (nhh)) E [0, 1], where E, th = 1. Each annotation yi (8) E ) == (1, C)HxW, where C is the number of classes and H, W are the height and width of the image, respectively. Here {ym on ( represent random samples from the underlying label distribution of 2n.-
To represent the alcatoric uncertainty, we introduce a conditional probabilistic model on segmenta- tion, denoted poly x) E Py, where Py is the space of segmentation distributions. The model learning problem corresponds to minimizing a loss function & Py x Py - IR between our model output (-(r,)) and the empirical annotation distribution of the image In, denoted Un, over the dataset DD. Namely, our learning task can be written as,
1 D
M 2
min = D
E(po(V) un) where =
8
In this work, we focus on developing an effective and easy-to-interpret uncertainty representation for semantic segmentation. In particular, we propose novel multimodal representation in our model μ y z) and Wasserstein-lite distribution distance as our loss E, which admits an efficient minimization strategy via partial constraint relaxation. Below we will present our model design and learning strategy in detail.
3.2 THE MIXTURE OF STOCHASTIC EXPERTS (MOSE) FRAMEWORK
In order to capture the multimodal structure of annotations, we develop a mixture of stochastic experts (MoSE) framework, which decomposes the output uncertainty into two granularity levels. At the fine-level, we employ an expert network to estimate a distinct mode of the aleatoric uncertainty in the dataset, which typically corresponds to certain local variations in segmentation. At the coarse- level, we use gating network to predict the probabilities of an input image being segmented by those stochastic experts, which depicts the major variations in the segmentation outputs.
Specifically, we assume there exist K stochastic experts {& each of which defines condi- tional distribution of the segmentation y given an input image x. denoted as (x(y)=)) Each expert Eg is associated with weight T& (x) predicted by a gating network, which represents the probability of an imput being annotated in the mode corresponding to &g. We instantiate (z(z)) as an implicit distribution that maps an expert-specific prior distribution Pe(*z) (where is latent variable) and the input x to the output distribution via deep network Fax, 2x). As such, given the input x. we are able to generate segmentation sample E ) as follows: we first sample expert id k from the categorical distribution (TT())) where т (x) = [TT (x), , k (x)]. Then we sample a latent code 2д from (2t) and compute the prediction by the network F(.):
F(x, 2k) where k G(T(()),, 2x (=()))
An overview of our MoSE framework is illustrated in Figure I Below we introduce our network design in Sec. 3.2.1 and present an efficient output representation Sec. 3.2.2
3.2.1 NETWORK DESIGN
We adopt an encoder-decoder architecture for our MoSE model, which consists of an encoder net- work that extracts an image representation, K expert-specific prior models and a decoder network that generates the segmentation outputs, as well as a gating module for predicting the expert proba- bilities. We detail the design of the experts and the gating module below.
Expert Network The stochastic expert network & comprises an encoder fe(.). a decoder fa(-) and latent prior models (Pa(zu)). We expect each expert to generate segmentation outputs for different images with a consistent uncertainty mode. To achieve this, we introduce a set of input- agnostic Gaussian distributions {P*(2*) priors, where m2 E R and is an L x L identity matrix. Given an input image x, we use a Conv- Net the encoder (a) to compute image representation My, followed by Deconv-Net the
Igl), omm,,og/) 1,... K)} the expert-specific
4

---

Published as a conference paper at ICLR 2023
decoder fa(Ug: 2x) which takes the feature Ug and sampled noise 2 to generate segmentation output for the k-tth expert. Concretely, assume the global representation uy R has Fi channels and spatial size of H' x W The decoder first generates a dense feature map ud E RFaxHxW and then fuses with the noise 2% to produce the output. Here we adopt a similar strategy as in Kassapis et al. (2021). which projects zn to two vectors z Wi 2 E R and 2k 222k E R and generates an expert-specific dense feature map as = = (mali, i + # 2x (i,j) E [1, H] [1, W.. Finally, we use three x 1 convolution layers generate pixel-wise predictions. For notation clarity, we denote all the parameters in the expert networks as 0g.
Gating Module The gating network shares the encoder fe(+) and uses the global representation My to predict the probabilities of experts. Specifically, we first perform the average pooling on Bg and pass the resulting vector U = AvgPool(ug) E RR through three-layer MLP to generate the probabilities (z)) We denote the parameters in the gating module 0.-.
We note that our MoSE framework is flexible, which can be integrated into most modern segmenta- tion networks with a small amount of additional parameters. Moreover, it can be easily extended to other model variants where the gating module takes the input features from other convolution layers the encoder or the latent noise is fused with other deconvolutional layers in the decoder (cf. C.5
3.2.2 UNCERTAINTY REPRESENTATION
Due to the implicit distributions used in our MoSE framework. expert networks. In this work, we consider two types of nonparametric representations for the pre- dictive uncertainty poly x)- The first representation utilizes the standard sampling strategy to generate set of N samples {s(') for input image (cf. Eq and approximates the uncer- tainty by the resulting empirical distribution N ΣN N ((s(t)). The second method utilizes the MoSE framework and approximates the uncertainty distribution by set of weighted samples. Specifically,
only generate samples from
we first generate S samples from each expert, denoting ()()) Ji=1 for the k-th expert, and then approximate the predictive distribution 1s0 with the following jup:
Po | x))
K
T
(x(2))
S
8(8
(
(3)
S
F(xx =()).. and 2k N(mR, OKI). It is worth noting that we sample from multiple experts simultaneously as above, which leads to an efficient and compact representation of the segmentation uncertainty. For elarity, we denote the nonparametric representation of both forms as ((s()) in the following, where ((9) is the sample weight and N the sample number.
where k
3.3 OPTIMAL TRANSPORT-BASED CALIBRATION LOss
The traditional segmentation loss functions typically encourage the model prediction to align with unique ground truth (e.g. Dice) or require an explicit likelihood function (c.g. Cross-Entropy To fully capture the label distribution with our implicit model, we formulate the model learning as an Optimal Transport (OT) problem, which minimizes a Wasserstein-like loss between the model distribution N (y) v..) and the observed annotations Vn. of x,, on the training set D.
Specifically, our loss measures the best way to transport the mass of model distribution no wn) to match that of the ground truth distribution vn,, with a given cost function y x ) + R. Based on this loss, our learning problem can be written as follows,
/ c(s, v)d v)-
min Σ We (Ho() vn), where We(m, = inf yell(u.)) yxy
(4)
Here [(,,, v/ denotes the set of joint distributions on y x y with i and v as marginals. Substituting in our discrete empirical distributions (o x)) = EN (()( and v = Σ (()) we derive the following objective function:
M
min 0,00-
ΣΣ P
P(
(2) (0.), v(1)) where p arg min
10
PEU
(5)
U = PE NXM
PI = un (0,) vn
5

---

Published as a conference paper at ICLR 2023
where C E RNXM is the ground cost matrix with each element Ci5 c(((s) (()) being the pairwise cost. In our problem, we can use a commonly-adopted segmentation quality measure to compute the cost matrix. P is the coupling matrix that describes the transportation path, and U the set for all possible values of which satisfy the marginal constraints. This formulation can be regarded as bi-level optimization problem, where the inner problem finds an optimal coupling matrix P and the outer problem optimizes the model output distribution under this coupling.
Constraint Relaxation We note that the marginal probabilities, u(x; 0z), occurs in the constraint of the inner problem in Eq [ As such, we can directly use backpropagation to optimize the gating network parameters B2 (Peyré et al. 2019). To tackle the challenge, we propose a novel imbalanced constraint relaxation. Our main idea is have fewer constraints when optimizing the inner problem, and use the optimally solved matrix P as supervision for the gating network in the outer problem. Formally, we introduce the following relaxation:
min ΣΣ PEPte
Pi (0.) BKLL Σ P un (0+))
0,,07
(6)
where P arg min
ΣPICC U (PERN PIM S S. 1 vn)
PEU 4
where 3 is used to balance the loss terms, and 7 E , 1] is designed for preventing sub-optimal solutions in the early training stage. Empirically, we start with a value 7o < 1 and annealing it to (meaning no constraint). is worth noting that in the upper-level problem, the first term can be regarded as normal segmentation loss for matched pairs of the predictions and ground truth labels, while the second term can be considered as optimizing the Cross- Fross-Entropy loss between the predicted probabilities and pseudo label generated from the coupling matrix.
Our constraint relaxation has several interesting properties for model learning. Firstly, by moving the constraint of the predictive marginal distribution from the inner problem the upper problem, we find it helps deal with the gradient bias (Bellemare et al. 2017 Fatras et al. (020)) when the ground truth annotations are limited for each image, which common real scenarios. In addition, our relaxed problem can be regarded as minimizing an approximation of the original problem, and achieves the same value when the KL term being O. Finally, our relaxation leads to a more efficient alternating learning procedure. In particular, when = 1, we have only one constraint in the inner problem and this can be solved linear time (Ignizio Cavalier] 1994). We show the training procedure, details solving the inner problem and theoretical analysis of Eq 6 in Appendix A B
EXPERIMENT
We evaluate our method on two public benchmarks, including LIDC-IDRI dataset (Armato III al, 2011) and multimodal Cityscapes dataset (Cordts et al., 2016). Below we first introduce dataset information Sec. 4.] and experiment setup in Sec. [.2]. Then we present our experimental results on the LIDC dataset in Sec. 4.3 and the Cityscapes dataset in Sec. 4.4
4.1 DATASETS
LIDC-IDRI Dataset The LIDC dataset (Armato III et al. 2011) consists of 1018 3D thorax CT scans, with each annotated by four anonymized radiologists For fair comparison, we use pre- processed 2D dataset provided by Kohl et al. (2018) with 15096 slices each cropped to 128 128 patches and adopt the 60-20-20 dataset split manner same as Baumgartner et al. (2019); MMnneiro et al. (2020). We conduct experiments in two settings, which train the model using either full labels or one randomly-sampled label per image as in the literature.
Multimodal Cityscapes Dataset The Cityscapes dataset (Cordts et al., 2016 contains images of street scenes with 2975 training images and 500 validation images, labeled with 19 different
A total of twelve radiologists participated in the image annotation and can revise their own annotations by reviewing the anonymized annotations of other radiologists. Therefore, the final disagreement among radiolo- gists indicates true ambiguities and typically represents the aleatoric uncertainty (Armato III al. 2011).

---

Published as a conference paper at ICLR 2023
Table 1: Quantitative results on the LIDC dataset. The number of sampled outputs is shown in the parentheses aside the metrics. Our method achieves SOTA with smaller or comparable model size.
Method
lahel GED J (16)
GED (50)
GED 4 (100)
M-ToU 1 (16)
ECE. J. (%) (16)
param..
Kohl al.
2018
0.320 0.030
0.239 = NA
0.500 : 0.030
76.15M
Kohi
2019
0,270
0.0100
87.51M
1u cl.
2019
0.267 1 0.012
20.30M
Baumgartier et al
2019
0.224 = NAA
74.82M
All
Monteiro et al. 2020
0,225 = 0.002
41.28M
Kassapis
0,264
0.248 0.004
0.243 0.004
0.005
1
175.36M
Ours light
0,219 0.002
0.195
0.002
0.190 0.003
0.620 + 0.003
0.070 0,013
9.37M
Ours
0.218 1: 0.003
0.195 1 0.002
0.189 = 0.002
0.624 1 0.004
0.064 = 0.015
41.60M
Ours compact
0.195 0.005
0.187
0.003
0.186 = 0.002
0.635
0.003
0.054
0.015
41.60M
Kohl et al. 2018
0,455 N/A
76.
Baumgarin
2019
0.323 N/A
74.82M
Monteiro al. 2020
One
0.365 0.005
41 41.28M
Onrs
0.252 0.004
0.229 + 0.005
0.223 0.005
0.596 + 0.003
0.105
= 0.009
41 61.60M
compact
0.004
0.004
0.220 0.005
0.090
0.011
41.60M
This score is taken from Baumgartner et al. (211))
* This score is generated by running the pre-trained model provided by Kassapis et al. (222
semantic classes. To inject ambiguity, we follow Kohl et al. (2018) to randomly flip five original semantic classes into five new classes with certain probabilities. Specifically, the original classes of 'sidewalk', 'person', 'car', 'vegetation', 'road' are randomly transformed to sidewalk2', 'person2". 'car2', 'vegetation2", 'road2" with corresponding probabilities of 8/17, 7/17, 6/17, 5/17, 4/17. This brings total of 24 semantic classes in the dataset, with 2 32 modes of ground truth segmentation labels for each image (See Figure3). For fair comparison, we follow the setting in (Kohl et al. 2018 Kassapis et al. 2021) and report results the validation
4.2 EXPERIMENT SETUP
Performance Measures To measure how well the model captures the aleatoric uncertainty in
segmentation, we follow Kohl et al. ( b0; Baumgartner (2019): Monteiro
Montero et al. (2020);
Kassapis et al. (202I) to adopt two main metrics for evaluation: the intersection-over--ue
(loU)
based Generalized Energy Distance (GED) (zzkelyy & Rizzo 2013) and the Hungarian- matchedd
IoU (Kuhn 1955). We extend the latter to Matched loU (M-ToU) by adopting network simplex algorithm (Cunningham] 1976) for the setting that the marginal distribution of predictions or ground truth is not uniform. Furthermore, we use the expected calibration error (ECE) (Guo et al.) (017)) with 10 bins as secondary metric to quantify the pixel-wise uncertainty calibration. For details on these metrics please refer Appendix A.2 Below, conduct all experiments in triplicate and report the results with mean and standard derivation.
Implementation Details We employ the same incoder-deoode architecture as in U-net (Ron- neberger et al.] 2015) and follow the model design as BBumgartner et al. (2019). On the LIDC experts each with S = 4 samples. In the loss function, we use the IoU (Mil- letari et al. 2016 as the pair-wise cost function and set the hyperparameters Yo 1/2, 3 I for full-annotation case and = 10 for one- mecannotation case. On the Cityscapes dataset, we use a slightly larger number K 35 of experts each with 5 = 2 samples. We use the CE the pair-wise cost function and set the hyperparameters "To = 1/32, B = for the loss. To stabilize the training, we adopt gradient-smoothing trick in some cases and refer the reader Appendix A.I for details.
dataset, we use K
4.3 EXPERIMENTS ON THE LIDC-IDRI DATASET
We compare our method with the previous state of the arts and summarize the results under sample number 16, 50, 100 in the Table I and Appendix Table 5 We validate our method with two kinds of output representation (cf. Sec. 3.2.2): the handard-ssmplingg as 'Ours' and the weighted form as 'Ours- compact'. For the latter, we report the results by taking S 4, 12, and 25 samples from each expert in order to generate the comparable number of total samples (i.e., 16, 48, and 100) as the sandard-sampling form. We conduct experiments with a lighter backbone denoted as 'Ours light' (cf. Appendix A) to make fair comparison to Hu al. (2019).
representation denoted
7

---

Published as a conference paper at ICLR 2023
Image
Annotations
Image
0.40
0.64
05
0.5
04
0.4
0.04
03
0.3
02
02
0.32
0.05
01
0.1
Samples
Entropy
0.0
Samples
Entropy
0.0
Figure 2: Qualitative results on the LIDC dataset. Each row of predictions represents a sampling sequence from expert with weight denoted on the left head. The last line shows the entropy map.
Table 2: Ablation study on our model components, evaluated with sample number 16.
Expert type
Expert weights
loss
GED +
M-IoU +
ECE (%6
stochastic stochastic
leamable uniform
loss
0.533 + 0.001
0.533 0.001
0,277 + 0.017
uniform lcarnable
OT-based
0.282 + 0.002
0.545 0.007
0.215 + 0.006
deterministic
OT-baaee
0.246 + 0.006
0.591 0.001
0.142 + 0.003
stochastic
Ieamable
OT-based
0.218 + 0.003
0.624 0.004
0.064 + 0.015
As shown in Table I our method is consistently superior others with the standard saandadd-samping form in all evaluation settings. Moreover, the compact weighted representation improves the performance further, especially in the small sample (16) case, and is more stable across different sample sizes. For the challenging one label setting, method is less affected and also outperforms the others by a large margin, demonstrating its robustness to the number annotations per image. We show the reliability plots in Appendix uncertainty, we further conduct experiments with varying training data sizes in Appendix CC3
C.2 To validate that MoSE model does capture the aleatoric
Figure visualizes compact uncertainty representation two cases. We can that each expert predicts a distinct mode of aleatoric uncertainty with a shared style across images. The image- specific weighting of experts shows coarse-level uncertainty, while the variations in samples from each expert demonstrate a fine-level uncertainty.
4.3.1 ABLATION STUDY
We conduct an ablation study of the model components on the LIDC dataset under the training setting with all (four) labels. We show the results with the sample number S 16 under the
-aappang-samping form in Table and leave the other settings to Appendix Table 6 In Row #1, we first analyze OT based loss by replacing it with the IoU loss (Milletari et al. 2016) calculated on every possible pair of predictions and labels. This leads to performance drop of more than 0.3 in the GED score and the resulting model tends to predict modify our MoSE design by adopting a uniform weighting of experts in Row #2, and replacing the stochastic experts with deterministic ones in Row #3. For both cases, compared with our full model in the last row, we see a clear performance drop. Those comparison results validate the efficacy of our design. We also conduct more analysis experiments by varying different aspects of the model architecture. Please refer to Appendix C.4 and C.5 for details.
averaged label mask. Moreover,
4.4
EXPERIMENTS ON THE MULTIMODAL CITYSCAPES DATASET
On this dataset, we compare our method with two previous works (Kohl et al, 2018 Kassapis et al,
2021 For fair comparison, conduct two sets of experiments: the first follows Kassapis et al.
(2021) to use predictions (Zoo 2020) from a pretrained model as an additional input and use a lighter backbone, and the second additional input as in Kohl et al. (018). We evaluate
8

---

Published as a conference paper at ICLR 2023
Table Quantitative results on the multimodal Cityscapes dataset.
Method
Additional input
Sample number
GED
M-ToU
ECE J.
param.
Kohl et al
2018
0.206 + NIA
0,512
3,730
71.82M
Ours
0.203 0.002
0,580 1 0.002
3.032 2: 0.018
41.63M
Kassapis al.
2021
0.164 = 0.01
0.633
3.274
17.31M
Ours
0.176 +: 0.002
0.638
0.002
3,030
0.021
41.3M
light
0.177 0.002
Kassapis al.
2021
0.139
0.672
2.196
17.31M
Ours
35
0.155 + 0.002
0.690
0,002
1.861 0.015
41.63M
Ours compact
0.142 + 0.001
0.761 + 0.002
0.352 0.025
41.63M
0.076
0:018
9: 9:053
0:009
060
[0]
buiding ng
0.076
0,009
058
161
Higetall
6.023
0
0. 0.045
0.632
00:01
tudk
(0)
bua
0.023
O 021
3. B.0
8.033
O 01111
edounie ovr
161
0.0
Figure 3: Qualitative results on the multimodal Cityscapes dataset. Within each case, (a) shows the image and part of ground truth labels with corresponding probabilities, and (b) shows the entropy map, predictive samples and predicted probabilities.
model with the standard- sampling representation under 16 and 35 samples, as well as the compact representation under 35 samples (I sample for each expert). As the previous works only provide their results on GED with 16 samples, we generate other results from the official pre-trained models and summarize the results in 3
Overall, our method achieves the -tate-of-the-art M-IoU and ECE under both sample number set- tings, as well as comparable GED scores with Kassapis et al. (2021). demonstrating the superiority of our approach. The discrepancy in the GED performance is likely caused by the evaluation bias of the GED on sample diversity (Kohl et al., 2019). as Kassapis et al. (202)) typically produces more diverse outputs due to its adversarial loss. Moreover, observe that our compact representation achieves better performance than the standard-sampling form under the same sample number (35), which demonstrates the effectiveness of such compact representation in multimodal uncertainty esti- mation. We show some visualization results in Figure 3 For clarity, we pair each ground truth label with the closest match in model predictions. We can see that our model can predict multimodal out- puts with calibrated mode probabilities well. Furthermore, we include additional reliability plots, experimental analysis on the calibration of mode ratio, detailed ablation on model architecture, and more visualization results in the Appendix D
5 CONCLUSION
In this work, we tackle the problem of modeling aleatoric uncertainty in semantic segmentation. We propose a novel mixture of stochastic experts model to capture the multi-modal segmentation distribution, which contains several expert networks and gating network. This yields efficient two-level uncertainty representation. To learn the model, we develop Wasserstein like loss that computes optimal coupling matrix between the model outputs and the ground truth label distri- bution. Also, it can easily integrate traditional segmentation metrics and be efficiently optimized via constraint relaxation. We validate our method on the multi-grader LIDC-IDRI dataset and the muti-modal Cityscapes dataset. Results demonstrate that our method achieves the state-of-the-an or competitive performance on all metries. We provide more detailed discussion on our method, including its limitations, in Appendix

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
This work was supported by Shanghai Science and Technology Program 21010502700 and Shanghai Frontiers Science Center of Human-centered Artificial Intelligence.
REPRODUCIBILITY STATEMENT
The complete source code and trained models are publicly released bhttps://githhut com/ tong:iotgggMM.... AUSeg] We provide the pseudo-code of our training procedure and explain the implementation details Appendix A.I We use the same segmentation backbone BBaumgart ner et al. (2019); Monteiro et al. (2020) and provide details on the network architectures in Appendix We use two public datasets: LIDC and multimodal Cityscapes and follow the data processed procedure previous work, as shown in Sec 4.1 and Appendix
A.1 A.1
REFERENCES
Dario Amodei, Chris Olah, Jacob Steinhardr, Paul Christiano, John Schulman, and Dan Mané. Con- crete problems in ai safety. arXiv preprint 1XXiv:6606.06665 2016.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference machine learning, 2017.
Samuel G Armato III, Geoffrey McLennan, Luc Bidaut, Michael F McNitt-Gray, Charles Meyer, Anthony P Reeves, Binsheng Zhao, Denise R Aberle, Claudia I Henschke, Eric A Hoffman, et al. The lung image database consortium (lidc) and image database resource initiative (idri): completed reference database of lung nodules on ct scans. Medical physics, 2011.
Christian F Baumgartner, Kerem C Tezcan, Krishna Chaitanya, Andreas M Hötker, Urs J Muehlematter, Khoschy Schawkat, Anton Becker, Olivio Donati, and Ender Konukoglu. Phiseg: Capturing uncertainty in medical image segmentation. In International Conference Medical Image Computing and Computer-Assisted Intervention, 2019.
Marc 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: 107433 2017.
Matan Ben-Yosef and Daphna Weinshall. Gaussian mixture generative adversarial networks for diverse datasets, and the unsupervised clustering of images. arXiv preprint arXiv: 10356, 2018.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint urXiv:1308.328 2013.
Liang-Chiel Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam. Encoder- decoder with atrous separable convolution for semantic image segmentation. In Proceedings of the European Conference on Computer Vision (ECCV), 2018
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urban scene understanding. In Proc. of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), 2016.
William H Cunningham. A network simplex method. Mathematical Programming, 1976.
Bharath Bhushan Damodaran, Rémi Flamary, Vivien Seguy, and Nicolas Courty. An entropic opti- mal transport loss for learning deep neural networks under label noise in remote sensing images. Computer Vision and Image Understanding, 2020.
Nat Dilokthanakul, Pedro AM Mediano, Marta Gamelo, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture varia- tional autoencoders. arXiv preprint / 1.06448, 2016.
10

---

Published as a conference paper at ICLR 2023
Patrick Esser, Ekaterina Sutter, and Bjöm Ommer. A variational u-net for conditional appearance and shape generation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018.
Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Learning with minibatch wasserstein: asymptotic and gradient properties. In AISTATS, 2020.
Jean Feydy, Thibault Séjourné, François- Saaniir--avvier
Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, 2019.
Lucas Fidon, Wenqi Li, Luis C Garcia-Peraza- Perrerraa---rreaa, Jinendra Ekanayake, Neil Kitchen, Sébastien Ourselin, and Tom Vercauteren. Generalised wasserstein dice score for imbalanced multi-class segmentation using holistic convolutional networks. In International MICCAI brainlesion work- shop, 2017.
Lucas Fidon, Séhastien Ourselin, and Tom Vercauteren. Generalized wasserstein dice score, distri- butionally robust deep learning, and ranger for brain tumor segmentation: Brats 2020 challenge. In International MICCAI Brainlesion Workshop, 2020.
Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya-Polo, and Tomaso Poggio. Learning with wasserstein loss. In Proceedings of the 28th International Conference on Neural Information Processing Systems- Volume 2, 2015.
Xin Geng and Rongzi Ji. Label distribution learning. In Proceedings of the 2013 IEEE 13th Inter- national Conference on Data Mining Workshops, 2013.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Weinberger. On calibration of modern neural networks. In International Conference on Machine Learning, 2017.
Abner Guzman-Rivens Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: learning to produce multiple structured outputs. In Proceedings of the 25th International Conference on Neural Information Processing Systems- Volume 2, 2012.
Shi Hu, Daniel Worrall, Stefan Knegt, Bas Veeling, Henkjan Huisman, and Max Welling. Su- pervised uncertainty quantification for segmentation with multiple annotations. In International Conference on Medical Image Computing and Computer-Assisted Intervention, 2019.
Eyke Hüllermeier and Willem Waegeman. Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and methods. Machine Learning, 2021.
James P Ignizio and Tom M Cavalier. Linear programming. Prentice-Hall, Inc., 1994.
Eddy Ilg, Özgün Çişek, Silvio Galesso, Aaron Klein, Osama Makansi, Frank Hutter, and Thomas Brox. Uncertainty estimates for optical flow with multi-hypotheses networks. arXiv preprint arXiv:1802.07095, 2018.
Phillip Isola, Jun- Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference computer vision and pattern recognition, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzaion with gumbel-softmax, arXiv preprint arXiv: 144, 2016.
Wei Ji. Shuang Yu, Junde Wu, Kai Ma, Cheng Bian, Qi Bi, Jingjing Li, Hanruo Liu, Li Cheng, and Yefeng Zheng. Learning calibrated medical image segmentation via multi-rater agreement mod- eling. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, 2021.
Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou. Variational deep embedding: an unsupervised and generative approach to clustering. In Proceedings of the 26th International Joint Conference Artificial Intelligence, 2017.

---

Published as a conference paper at ICLR 2023
Alain Jungo, Fabian Balsiger, and Mauricio Reyes. Analyzing the quality and challenges of uncer- tainty estimations for brain tumor segmentation. Frontiers in neuroscience, 2020.
Elias Kassapis, Georgi Dikov, Deepak K Gupta, and Cedric Nugteren. Calibrated adversarial finement for stochastic semantic segmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021.
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? Proceedings of the 31st International Conference on Neural Information Processing Systems, 2017.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Simon Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R. Ledsam, Klaus Maier-Hein, S. M. Ali Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A prob- abilistic u-net for segmentation of ambiguous images. In Advances in Neural Information Pro- cessing Systems, 2018
Simon AA Kohl, Bernardino Romera-Paredes, Klaus H Maier-Hein, Danilo Jimenez Rezende, SM Eslami, Pushmeet Kohli, Andrew Zisserman, and Olaf Ronneberger. hierarchical proba- bilistic u-net for modeling multi-scake ambiguities. In Medical Imaging meets Neurips Workshop, 2019.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics quarterly, 1955.
Meelis Kull, Miquel Perello Nieto, Markus Kängsepp, Telmo Silva Filho, Hao Song, and Peter Flach. Beyond temperature scaling: Obtaining well-calibrated multi-class probabilities with dirichlet calibration. Proceedings of the 33rd International Conference on Neural Informa- tion Processing Systems, 2019.
Dong Bok Lee, Dongchan Min, Seanie Lee, and Sung Hwang. Meta-gmvae: Mixture of gaussian vae for unsupervised meta-learning In International Conference Learning Representations, 2020.
Jesse Levinson, Jake Askeland, Jan Becker, Jennifer Dolson, David Held, Soeren Kammel, J Zico Kolter, Dirk Langer, Oliver Pink, Vaughan Pratt, et al. Towards fully autonomous driving: Sys- tems and algorithms. In 2011 IEEE intelligent vehicles symposium (IV), 2011.
Pauline Luc, Camille Couprie, Soumith Chintala, and Jakob Verbeek. Semantic segmentation using adversarial networks. NIPS Workshop on Adversarial Training, 2016.
A Malinin and M Gales. Predictive uncertainty estimation via prior networks. In NIPS' 18: Pro- ceedings of the 32nd International Conference on Neural Information Processing Systems, 2018
Andrey Malinin and Mark Gales. Reverse kl-divergence training of prior networks: improved uncer- tainty and adversarial robustness. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, 2019.
Saeed Masoudnia and Reza Ebrahimpour. Mixture of experts: literature survey. Artificial Intelli- gence Review, 2014.
Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi. V-net: Fully convolutional neural net- works for volumetric medical image segmentation. In 2016 fourth international conference on 3D vision (3DV), 2016.
S Minaee, YY Boykov, F Porikli. AJ Plaza, N Kehtarnavaz, and D Terzopoulos. Image segmentation using deep learning: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
12

---

Published as a conference paper at ICLR 2023
Miguel Monteiro, Loic Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Marques, Konstantinos Kamnitsas, Mark van der Wilk, and Ben Glocker. Stochastic segmentation net- works: Modelling spatially correlated alcatoric uncertainty. In Advances in Neural Information Processing Systems, 2020.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ Machine Learning, 2019.
Di Qiu and Lok Ming Lui. Modal uncertainty estimation via discrete latent representation. In Inter- national Conference on Medical Image Computing and Computer-Assister Intervention UNSURE Workshop, 2021.
Ali Razavi, Aaron van den Oord, Ben Poole, and Oriol Vinyals. Preventing posterior collapse with delta-vaes. In International Conference on Learning Representations, 2018.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference Medical image computing and computer- assisted intervention, 2015.
Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE International Conference Computer Vision, 2017.
Ludger Rüschendorf. The wasserstein distance and approximation theorems. Probability Theory and Related Fields, 1985.
Laurens Samson, Nanne van Noord, Olaf Booij, Michael Hofmann, Efstratios Gavves, and Mohsen Ghafoorian. I bet you are wrong: Gambling adversarial networks for structured semantic segmen- tation. In Proceedings of the IEEE/CVF International Conference on Computer Vision Workshops, 2019.
Murat Sensoy, Lance Kaplan, and Melih Kandemir. Evidential deep learning quantify classifi- cation uncertainty. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, 2018
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in Neural Information Processing Systems, 20155
Gábor Szêkely and Maria L Rizzo. Energy statistics: A class of statisties based on distances.
Journal of statistical planning and inference, 2013.
Luis Caicedo Torres, Luiz Manella Pereira, and M Hadi Amini. A survey optimal transport for machine learning: Theory and applications. arXiv preprint (XX.x:210600196. 2021.
MM Valiuddin, Christiaan GA Viviers, Ruud JG van Sloun, Fons van der Sommen, et al. Improving medical image segmentation with normal- izing flows. In Uncertainty for Safe Utilization of Machine Learning in Medical Imaging, and Perinatal Imaging, Placental and Preterm Image Analysis. 2021.
aleatoric uncertainty quantification
Cédric Villani. Optimal transport: old and new. Springer, 2009.
Guotai Wang, Wenqi Li, Michael Aertsen, Jan Deprest, Sébastien Ourselin, and Tom Vercauteren. Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks. Neurocomputing, 2019.
Yang Yu and Wen-Ji Zhou. Mixture of gans for clustering, In IJCAI, 2018.
Jun-Yan Zhu, Richard Zhang, Deepak Pathak. Trevor Darrell, Alexei A Efros, Oliver Wang, and Eli Shechtman. Toward multimodal image-to-image translation. Advances in neural information processing systems, 2017.
TensorFlow DeepLab Model Zoo. Pretrained model: xception71 dpe cityscapes trainfine. 2020.
13

---

Published as a conference paper at ICLR 2023
A
IMPLEMENTATION DETAILS
A.1
TRAINING PROCEDURE
The overall training procedure is shown in algorithm [ We use the weighted version of our non- parametric representation during training. displayed in lines 3 9. Another option is hundard-sampling representation (see Sec. 3.2 and utilize the Gumbel-Softme trick (Jag et al,
2016 to enable backpropagation when drawing expert id from the categorical distribution (TT).
Algorithm Training Procedure
Require:
Training data D with M ground truth labels per image, number of experts K, number of samples per expert S, batch size B, leaning rate loss hyperparameters 3,7-
Require: initial expert network parameters O,, initial gating network parameters O,..
1: while Oz, or not converged do
2: Sample { {zn; (5) (5) n=
Un 1j=1.-.,M M batch from the training dataset D.
3:
for ks 1.... K do
4:
for t = 1, ... do
Sample {zn (1:t) pr ~ PA a batch of expert prior codes.
5:
8n (k,tt) < F(x, 2n (1t)) (1t)), 0.) with n 1, ,,.., B.
6:
un (k.s) + S [x;1;;))) with n = 1, ... B.
7:
8:
end for end for
9:
10:
Rearrange all outputs as {sul Un M= .N. where - S.
Compute pairwise cost matrix Cn. with C, j] c(s.) y ) 1, ... B.
11:
12:
Solve the optimal coupling matrix P LP SOLVER(C,s Un Vn, 7), n = 1, B.
13:
L -n (PP: n. (9.)) BKL(P u, (0x)))
14:
0.. G, +- Adam ( n)
15:
end while
Details on solving the coupling matrix In each iteration of training, we need solve the best coupling matrix P This corresponds to the LP_SOLVER function module in algorithm | and is formally defined as:
P arg min P U' (PER PLM S 7,P"1N =v
PEU'
Especially, when = 1, the corresponding inequality constraint term vanished and U' = (P e R. N 1''' IN = vn We can then solve the problem by moving the mass of each ground truth label its nearest prediction. The time complexity is O(N). When y < 1, the problem can be solved by a standard linear programming solver like simplex or interior-point (Ignizio & Cavalier 1994). We adopt greedy strategy here, which empirically provides similar performance but with lower time complexity. Specifically, we first sort the ground truth labels according to their minimum cost to the predicted sample in ascending order. For each ground truth label, we also sort the predicted samples in terms of the corresponding cost in ascending order. Then, go through the sorted ground truth labels and transfer their mass to the predicted samples in the sorted order given the constraint P1M S y is satisfied. This algorithm has time complexity of (loMlo)NN)
Details computing Soft Gradient We observe that optimizing the loss (cf. Eq. 6) via direct gradient descent could sometimes suffer from slow convergence and low performance due to large variance in the target signal Pi in the second term. This is partially caused by the random samples generated by the stochastic experts and unstable P inferred during the learning process. To address this, we adopt a simple smoothing technique to produce a stable gradient. Specifically, group those experts with similar outputs during training, and average the gradients of the experts within each group. This prevents the gradients from changing dramatically in consecutive steps and resembles the SrriightTThroug method (Jang al. 2016: Bengio al. 2013
14

---

Published as a conference paper at ICLR 2023
We show the details for computing soft gradients in Algorithm Groups 91. .9 sample indices, and can be computed by a clustering method. We here take a simple strategy and merge experts with similar outputs. Specifically,
sets of
initialize each group Go. where s == s "x skit is the mean of samples for expert
expert, and merge
two experts k k if IoU(s*, ( s* k E [K]. Go is a hyper-parameter where we initialize as and linearly decreased during training. Notice that mainly apply this soft gradient operation on the Cityscapes dataset, where multiple experts generate similar predictions.
Algorithm Compute Soft Gradient
Input: loss function f, model outputs {s(0), J coupling matrix P", hyperparameter Go. Output: soft gradients V
1: procedure SOFTGRADIENT
2: 91.".. .9G Clustering(s, Go)
3:
t 1
4:
for m 1... do
5:
((m) + Iomm
Zieg.. (()
6: d(m) f l9ml
Σ fm (i)
7: Vaoft ((u(t), ,01)) x V Valm) El e( olm) (i e gm)
alt)
8: end for
9: return V
10:
end procedure
Details Network Architechtures For the backbone, use the same encoder-decoder archi- tecture as in U-net (Ronneberger et al. 2015) and follow the model design as (Baumgartner et al.) 2019). Specifically, there 6 down- and upsampling blocks, with each block consisting of three convolutional layers, each followed by batch normalization layer and a ReLU-activvtionn) The global feature map Ug E R H' xw" has feature dimension F1 = 192. The dense feature map Ud E R HxW has feature dimension Fi2 = 32. The gating network contains three-layer MLP, with the first two layers having 192 neurons and the last layer having neurons is the expert number). The three subsequent x convolution blocks after the decoder have 32 channels for the first two blocks and channels for the last block ( is the number of classes). In addition, order to keep fair comparison to the methods using a smaller backbone (Hu et al. (019), we conduct experiments with lighter U-net backbone same as (Hu et al. 2019) (decrease from the original
down-sample layer to 3 down-sample layer) 3 and keep other model design the same. We compare the model size of methods using the same segmentation backbone in Table
Table 4: Compare the model size of methods using the same segmentation backbone. Our model only adds small number of parameters to the original Unet backbone. (* denotes the score is calculated according to the re-implementation by Baumgartner et al. (2019).
Method
#Backbone param.
#Additional param. (relative ratio)
#All param.
Kohl et al. 2018
34.88M (84.5%)
76. 15M
Baumgartner al. 2019
33.55M (81.3%)
74.82M
41.27M
Monteiro et al. 2020)
0.01 (0.3%)
41.28M
Ours
0.33M (0.8%)
41.60M
Hu et al. 2019
11.32M (125,1%)
20.37M
9.05M
Ours light
0.32M (3.5%)
9.37M
'Other details on backbone architecture can be referred to Baumgartner et al. (019 and their code link https://github cww/ wwwngaahh//iieeeggaa..aaaass
"Other details on the lighter backbone architecture can be referred to their code link https:
//github. com/ shibux/ supervised_us and also the PyTorch implementation for the Probabilistic U-Net Kohl et al.] (2018) model from https://g" thub. com atte haaabb/c////ttefaeeeettt
Probabilise Unnticc-ntttttrrrrr
15

---

Published as a conference paper at ICLR 2023
Other Training Details On the LIDC dataset, we apply random flip and random rotation aug- mentations same as Baumgartner et al. 2019);; Monteiro et al. (2020). We use the latent code with dimension 1, which is empirically found to perform well and can be further refined. The mod- els are trained for 200 epochs with hatch size 16. On the Cityscapes dataset, we apply random horizontal flips, cropping, and scaling size 128x128 same as Kassapis al. (202T). We use batch size of 12, latent code with dimension 5, and train the models for 1000 epochs. The ground truth marginal probabilities v (cf. (6) are initialized in uniform and annealing the original values after some warm-up epochs. The additional prediction inputs are from the pretrained model xcep- tion71 sonn1_dde_ciin wwtysscapss.traiiiii trainfiie provided by DeepLab Model Zoo f
As commonly practiced (Kass- apis et al. [021; Kohl et al.) 2018). we ignore the unlabeled pixels when calculating loss. On both datasets, we use weight-decay with weight le- and Adam optimizer with an initial learning rate of 1e The learning rates and loss slack factor are linearly updated with scheduled updates. We do the model selection and tune hyperparameters according to the GED score on the validation set. We use two NVIDIA TITAN RTX GPUs the LIDC dataset and four NVIDIA TITAN RTX GPUs on the Cityscapes dataset.
A.2 EVALUATION METRICS
We aim measure how well the model captures the multimodal aleatoric uncertainty in segmenta- tion. The aleatoric uncertainty refers to the irreducible randomness in the data- Sata-generating process and can be represented by the ground truth label distribution (Monteiro et al.] 2020; Hüllermeier & Waegeman, 2021). Following the literature (Kohl et al.] 2018 2019; Baumgartner et al.) 2019; Monteiro et al. 2020, Kassapis et al. (02T), we use two distributional metries: generalized energy distance (GED) and mutched-loU (M-IoU) to evaluate how the model fits the ground truth label distribution in the segmentation output space, which also demonstrates the quality of alcatoric un- certainty quantification. Due to the complexity of the high- -igh-dimensional segmentation distribution, both metrics use the sampling strategy to approximate the distributional distance via two sets of sam- ples. In complementary to the sample-wise metrics, we also adopt the expected calibration error (ECH) (Guo et al.) 2017) as an additional metric to measure pixel-wise uncertainty calibration. We note that while the ECE measures the quality of more general predictive uncertainty, our problem setting is dominated by the aleatoric uncertainty (cf. C.3 and D.3). Besides, as it is mainly used in the calibration of classification problems, the ECE can not measure the important multimodal and structural characteristics of segmentation uncertainty, and thus is adopted as secondary metric. Be- low, we summarize the details of these three metrics. For clarity, we denote ý as the hard prediction (after 'argmax" operation) to distinguish the soft prediction (after 'softmax' operation) Other notations are kept consistent with previous content.
Generalized energy distance (GED) (Székely & Rizzo 2013) measures the distance between two distributions by calculating the pair-wise distance dl(-) = IoU of samples from predictive distri- bution and samples from ground truth distribution. The expected value is then subtracted with the sample diversity of each distribution to get the final score:
DGE v) == FF (d( v)] Fly.y' (d v')] Eggg [d ())]-
(8)
This metric is widely adopted by the previous literature (Kohl et al. 2018 2019 Baumgartner et al., 2019 Monteiro et al., 2020) Kassapis et al., (021). While effective in most situations, this metric potentially rewards sample diversity and sometimes gives a biased estimation of distance (Kohl et al.) 2019)
Matched-IoU (M-IoU) (Kuhn 1955
Cunningham 1976) measures the distributional distance by between samples from the predictive distribution and ground
solving an optimal matching matrix truth distribution:
WioU (14, v) == max Σ
Pr.,loU(D...)) s.t. PM = u, P 1N =
(9)
P
1.
It is previously adopted by Kohl et al. (011 Kassapis et al. 22221) as compensation for GED measure since this metric does not favor sample diversity. We here use an extended version to en- able measure when one of the marginal distributions is not uniform. Specifically, on the LIDC
Other details on this additional input can be referred to in Kassapis et al. 2021) and their public code: https://gi htwp:tt.bb..cccccooooooo.ee.eeea K1 aaska ssppsssssspa 8/ CARSSS
16

---

Published as a conference paper at ICLR 2023
dataset and evaluated with standard sampling representation, we adopt the Hungarian matching al- gorithm (Kuhn 1955) to solve matching matrix P same as Kohl et al. 2 Kassapis et al. (202I). When evaluated with compact representation, ningham 1976) to solve P. On the Cityscapes dataset where ground truth labels are aligned with non-uniform probabilities, we also adopt the network simplex algorithm.
adopt the network simplex algorithm (Cun-)
Expected calibration error (ECE) (Guo et al. 2017) measures the difference between probability prediction and the real accuracy. It is defined as:
ECE E P(Y YYP pl
(10)
Here, Y is the random variable for the label, Y is the class prediction, and P is the associated prob- ability. In our case, we get the pixel-wise label distribution and predictive distribution by marginal- ization and treat each pixel i.i.d. For the binary case in the LIDC dataset, we measure how the foreground class is calibrated. For the multi-class case in the Cityscapes dataset, we measure how the most-likely class is calibrated (aka confidence calibration) (Guo al. 2017; Kull et al.) 2019).
On the Cityscapes dataset, we follow the convention Kohl et al. 2019 Kassapis et al. calculate metrics of the ten switchable classes.
202 to only
B THEORETICAL ANALYSIS OF THE CONSTRAINT RELAXATION
We discuss the theoretical connection of our constraint relaxation to the original formulation in the following. For clarity, we consider bi-level optimization formulation in which we adopt alter- nating optimization process between 8g (outer loop) and {P, 0 (inner loop). We first focus on the inner loop when the parameters for segmentation predictions 0, is fixed, and analyze the optimiza- tion behavior of {P, 0x) } in the original form (Eq. 5) and relaxed form (Eq. 6 For simplicity, we rewrite them as following two functions:
1 P) = (P,C)
PlM = u, p" = =]. P
(11)
p IN =
(12)
where we use C short for C(s(O.() v). and short for u(0²) for notation elarity. We use A to denote all possible value of u that the gating network can generate, which also satisfies u E R Σ, 4; = 1. The inequality constraint in Eq. i not included here in la since it will eventually disappear due to annealing. We show in Proposition D that when 3 is sufficiently large, l2 achieves the same optimal solutions / when the KL term is zero. After that, when we alternates to optimizing 0., the two problems are exactly the same as they use the same P from the inner loop. Therefore the overall results of the original problem and the relaxed problem are the same.
Proposition 1 When is sufficiently large, la achieves the same optimal solutions as I
Proof: We denote the domains for variables as 22, P2 in l and l2 correspondingly. It's obvious that 01 S 92 and we denote the complementary set Q. = 222 Ω2. We denote the optimal solutions as Pa uu" = arg min p p,uss,
4, P. = arg min pwen, 122 In the following, we consider two cases:
When P", u² E Ωr, P. E Ω4, then PIM á is satisfied, and the KL term in l is O. Then we have (P",, u²) (P, û) = (((,,) a) < 12(PP,," ,u)) = (P(,,, u²)). Therefore, the equality should hold and 4 ()) = (1",, solution for 4 is also the optimal solution for 4
arg min P,aen (P,, u), which means the optimal
When P*, n² E Ωq, P,, û Pe. then KL(PIu|llu) > 0 must be true. Recall A is all possible values of a that the gating network can generate, we consider two cases:
(a) If PIM A, have 12(P", "u) 12(P,u) u') l2(P, û), which contradicts with the optimal of P, (b) If PL & A. with a sufficient large B. we have (i, C) + + (BKL(1MMMl (P*, BKKT(PP Iw mrllu)), which also contradicts with the optimal of P ú.
always find solution P, u' E Qq, by setting u' P1 Then
(P*, C) +0
17

---

Published as a conference paper at ICLR 2023
Therefore, when 3 is sufficient large, the optimal solution for l2 can only exist in Q2, which is also the optimal solution for 41.
In our optimization algorithm, take a further step, alternating between P and Specifically, first optimize P =arg m C²) and then optimize ū= arg min K (L(P1M as (u) (cf. Eq. 6 Compared to the joint optimization, this relaxed optimization strategy accelerating the optimization speed but have no guarantee to find the optimal solution as in the original form since the KL terms may not achieve zero. Specifically, we demonstrate in Proposition 2 that when the KL term can be optimized to 0, P, u is still the optimal solution for l2 and 4 Empirically, we find that our model achieves KL divergence around le-5, which is much smaller than the scale of the segmentation loss in le-1 on the constructed multi--oda Cityscapes dataset.
Proposition 2 our algorithm,, when the KL term can optimized to 0, P =arg min (e i= arg min, KL P1 M
[[u) is the optimal solution for 22 and 11.
Proof: K M u) = then P, ū E S2 Recall = EP R+ PlM = u, P 1y and
C E p IN v}. since B C C, we have arg millpee (P,C) S arg mimmpey
Therefore, (PPx, =²) < 12(P, u) 1 (P, ū) < 4PP"" .u)). Since K (MP""M]]"")) = 0, we have 12(P", u²) 1 (Px, u²),,, and hence the equality should hold, and P, a is also an optimal solution for 11,2
C ADDITIONAL EXPERIMENTS ON LIDC
C.I COMPLEMENTARY RESULTS
In complementary to the results shown in the main text, we show the performance of our method with M-IoU and ECE under sample-number of 50, 100 in Table 5 and the complementary ablation study under the sample number of 50, 100 in Table 6 For short, we name each experiment by its replaced term and refer to the comprehensive explanation in Sec. 4.3.1
Table 5: Other Quantitative results on LIDC-IDRI dataset. The ECE is in %.
Method
# labels
M-loU t(50)
M-IoU f(100)
ECE((50)
ECE4(100)
Ours
0.631 0.002
0.636 0.002
0.053 + 0.018
0.052 0.017
All
Ours compact
0.637 0.003
0.638 0.002
0.049 = 0.014
0.049 0.014
Ours
0.604 + 0.003
0,606 + 0.003
0.089 = 0.010
0.085 = 0.012
One
Ours compact
0.607 + 0.003
0.607 + 0.003
0.084 + 0.011
0.083 + 0.012
Table 6: Complementary results the ablation study with sample number 50, 100. The ECE is in
%.
Method
GED (50)
GED (100)
M-loU 1 (50)
M-IoU (100)
ECE (50)
BCE L (100)
IoU loss
0.533 +: 0.001
0.533 +: 0.001
0.533 = 0.001
0.533 +: 0.001
0.277
0.016
0.277 0.016
Uniform expert weights
0.256 + 0.003
0.249 + 0.003
0.554 = 0.006
0.556 + 0.007
0.209 = 0.008
0.209 + 0.009
Deterministic Expert
0.221 1: 0.004
0,216 + 0.004
0.601 = 0.008
0.603 1: 0.009
0.127 = 0.005
0.121 +: 0.003
Ours
0.195 1: 0.002
189 = 0.002
0.631 = 0.002
0.636 1: 0.002
0.053 = 0.018
0.052 :: 0.017
C.2 RELIABILITY PLOTS
In Figure 4 we show the reliability plot of our model (left) and compare with the plot of Kassapis
et al. 2022 (right). We generate the result by running the pre-trained model provided by Kassapis et al. 22021) For both methods, 16 samples are drawn for evaluation and we use the standard- sampling representation. This corresponds to the results the last column of Table T We that the confidence our model calibrates well with the accuracy, resulting in a small gap in the curve and achieving better pixel-wise calibration result than Kassapis et al. (2021)
18

---

Published as a conference paper at ICLR 2023
Ours
Kassapis et al. (2021)
Accuracy
Accuracy
Gap
Gap
ECE: 0.041%
ECE: 0.214%
Accuracy
Accuracy
26
22-
42
60
Confidence
Confidence
Figure 4: We show the reliability plot of our model (left) and compare with the plot of Kassapis al. (02T) (right) on the LIDC dataset. For both methods, 16 samples are drawn for evaluation and we use the standard-sand -amplling representation. We achieves 0.041 % ECE (one of the three random runs in Table
C.3 ANALYSIS THE TYPE UNCERTAINTY
While specifically designed for modeling the data-inherent uncertainty (aleatoric), the outputs of our model may also contain the model uncertainty (epistemic). We note that our major evaluation metries GED and M-IoU only measure how the model learns the aleatoric uncertainty. Therefore, the epistemic component in the output could lead to additional error in these metrics. To explore how much the epistemic uncertainty accounts for the total predictive uncertainty in our model and how it affects the final results, we conduct the following experiments.
Specifically, we follow Kendall Gal (207)) to utilize the property that epistemic uncertainty varies under different training data sizes while alcatoric part does not. By observing how the predictive uncertainty changes with different sizes of training data, we can infer how much the epistemic uncertainty accounts for in the predictive uncertainty. Besides, also show how the metrics GED, M-loU and ECE change, in order to analyze how the variation of epistemic uncertainty affects the final metries. Experiments done using all the labels and evaluated in 16 samples with standard- sampling representation.
As shown in Table 7 we change the size of the original LIDC training rize the uncertainty score by using pixel-wise entropy (same as Kendall & Gal (017)) and sample diversity. The sample diversity is calculated the sample-wise IoU distance, which the same as the diversity term in calculating GED (cf. Eq S) We find this score provides a better measure of the output variation under the class-imbalance situation. We see that both uncertainty scores increase slightly with the size of the training set decreasing. However, the magnitude of this change is in le-3, which is negligible when compared to the standard derivation of the repeated experiments. Therefore, we conclude that the predictive uncertainty contains only a small amount of epistemic uncertainty, which is nearly negligible. Besides, the changes in GED, M-IoU and ECE metrics are all in the same magnitude the standard derivation of the repeated experiments and are at least one order magnitude smaller than the difference in comparison results among methods in Table
1/2, 1/4 and summa-
C.4 ANALYSIS ON THE NUMBER OF EXPERTS
We analyze how expert number influences the model performance. To do this, we conduct exper- iments the LIDC full label dataset under expert number 1...6....0.15.20.... where 4 is our default expert number that is used other experiments. We show the GED score evaluated with 16 samples. As shown in Figure when the expert number is smaller than 15, with the increasing
19

---

Published as a conference paper at ICLR 2023
Table 7: We conduct experiments on the LIDC dataset with different training data sizes. Here, entropy and sample diversity are used to measure the uncertainty level of output in both pixel- wise and sample-wise. Experiments are done using all the labels and evaluated in 16 samples with standard-sampling representation.
Train dataset
Entropy (%)
Sample Diversity
GED 4
M-IOU 1
ECE (%)
LIDC/4
0.561 + 0.006
0.509 = 0.010
0.226
0.002
0.615 = 0.006
0.097 = 0.016
LIDC/2
0.559 1 0.007
0.505 i 0.011
0.225 + 0.002
0.619 1 0.004
0.076 1 0.010
LIDC
0.554 i 0.007
0.501 0.007
0.218 0.003
0.624 0.004
0.064 1 0.015
number of experts, the model achieves better GED score in general. However, when the expert number is much larger (eg. 20, 30 experts), the model performance decreases, corresponding an increase of GED score. Intuitively, more experts correspond to a more powerful model capacity. Though with too many experts, there could be some duplicate ones predicting roughly similar pre- dictions, which makes it difficult to learn a stable gating network. The results also demonstrate that the effectiveness of method does not owe large number of experts, which is different from the working mechanism of ensembling model. By using 4 experts, we keep balance between the model capacity and the expert compactness.
0,26
0.25
GED (16)
0.24
0.23
0.22
0.21
10
20
30
Number experts
Figure The GED score changes with the number experts, evaluated with 16 samples in coundrd-ssmpllnn representation. Experiments are conducted on the LIDC dataset with the full label set.
C.5 ANALYSIS OF THE MODEL ARCHITECTURE
Following, we explore the flexibility of our model design by conducting analysis experiments on three kinds of model variations: 1. different segmentation backbone; 2. the gating network takes input from different layers of the encoder; 3. the noise injects into the different layers the decoder (cf. Sec. 3.2). Experiments are conducted on the LIDC dataset with full labels and evaluated with 16 samples in the standard saandadd--ampling representation. Results are shown in Table 3 and we summarize that with all variations explored, the model shows satisfactory results and consistently performs the
20

---

Published as a conference paper at ICLR 2023
best compared with previous works (ef. Table I. which demonstrates the flexibility of our MoSE design. Following, we analyze the experimental results for each variation in detail.
On row # 1, we change the currently used Unet architecture to DeepLabv3+ (Chen al. 2018) with ResNet101 The gating module takes feature from the ASPP output and noise injection are added in the last layer of the encoder. Other architecture designs and hyperparameters the kept same as default model. We see that, without further tuning, the experiment using DeepLabv3++ as the backbone has achieved satisfactory results on all metries. This validates the feasibility of using other segmentation network as the backbone. On row 2 and #3, we make the gating network take the input features from the 2nd and 4th layers from the overall 6-layer encoder. On row # and #5, we conduct experiments with the latent noise fused with the 2nd and 4th layers from the overall 6-layer decoder. In comparison, we list the results of our default setting in the last row, which takes features from the last layer of the encoder and fuse latent noises in the last layer of the decoder. All other experimental settings are kept the same in these experiments. We observe that all these model variations show good performance on the final metrics, which validates the efficacy of using different levels of image features from the encoder network for the gating network input and injecting expert-specific noises in the different layers of decoder. For a new dataset, we suggest designing the gating network, its connection with the segmentation backbone, and noise injection way via prior knowledge by architecture tuning.
Table 8: We analyze different model variations by replacing the backbone and changing the noise injection layer and the gating input layer. Experiments are done on the full-label LIDC dataset and evaluated with 16 samples in the Haadard-sampling representation. Below denote I. as the full layer of encoder and Q as the full layer of decoder. In our default Unet architecture, L= Q = 6.
Backbone
Noise Injection Layer
Gating Input Layer
GED
M-IOU
ECE (%)
Deeplab v3+
0.237 + 0.005
0.609 0.005
0.104 + 0.018
Unet
1/3 L
Q
0.222 + 0.005
0.620 + 0.004
0.080 + 0.037
Unet
2/3 L
Q
0.218 + 0.001
0.626 0.002
0.082 +:2 0.009
Unet
L
1/3 Q
0.224 +:2 0,003
0.625 + 0.001
0.086 + 0.006
Unet L
2/Q
0.220 + 0.001
0.627 0.002
0.082 0.011
Unet L
Q
0.218 + 0.003
0.624 0.004
0.064 = 0.015
C.6 ANALYSIS ON TRAINING STRATEGY
Baumgarttnrr al.) One difference between our method and previous works (Kohl al. 2018 2019
2019 Monteiro et al. 2020 is that we use all ground truth annotations per image in one gradient step during training. To analyze this impact, we conduct an ablation study where we random sample one ground truth label in each gradient step the LIDC dataset and compare the results in Table We observe that the performance our method remains almost the same for such setting, which indicates that our advantage does not come from using all groundtruth annotations in each gradient step.
Table 9: Analysis of our model when using only one ground truth label per gradient step. Experi- ments are conducted on the LIDC dataset with full annotations and evaluated with 16 samples.
#gt per gradient step
GED +
M-IOU f
ECE J. (%)
One All (four)
0.218
0.003
0.627 +: 0.005
0.071 + 0.012
0.218
0.003
0.624 + 0.004
0.064 + 0.015
C.7 COMPARE WITH GMM- VAE MODELS
We also make empirical comparison between method and VAEs with GMM prior (Dilokthanakul
et al. 2016 Jing et al. 2017 Lee et al., 2020 To do this, we replace the original gaussian prior dis- tribution in Probablistic Unet Kohl al. 2018) by GMM with K components, and approximates
Details be found in httpa//tth/////roooooo//aaaaaaaa xeeppaab-xottto
21

---

Published as a conference paper at ICLR 2023
the ELBO via Monte Carlo sampling in Lee et al. (2020). Besides, also evaluate model variant by using input-agnostic Gaussian distributions and input-dependent categorical distribution model. We conduct experiments the LIDC dataset, use K = gaussian distirbutions (same as our expert number) and keep other settings the same as the Prohabilistic Unet. 'he results are summarized in Table TO We observe that the GMM prior do improves original Probabilistic Unet. However, our model still achieves the best performance, outperforming other models with a large margin (eg. 3% GED). This demonstrates the efficacy of our design, especially OT-based loss (for being one of the major difference).
Table An empirical comparison of our method and VAEs with GMM prior. We use the reimple- mented version of ProbUnet provided by Baumgartner al. 2019) to keep same Unet backbone.
Method
Gaussian Prior(s)
GED
M-IOU1 f
ECE (%)
param.
ProhUnet
Input dependent
0.298 +0.010
0.527 #0.007
0.118 #0.012
76.15
ProbUnet with GMM prior (1)
Input dependent
0,254 +0.008
0.597 20.005
0.111 20.013
76.18
ProhUnet with GMM prior
Input agnostic
0.251 +0.004
0.597 #0.007
0.101 #0.010
76.14
Ours
Input agmostic
0.218 20.003
0.624 +0.044
0.064 20.015
41.60
D ADDITIONAL EXPERIMENTS ON CITYSCAPES
D.I RELIABILITY PLOTS
On the left of Figure 6 we show the reliability plot of our model with the sample number 35 under the compact representation. For comparison, we generate the reliability plot of the previous state of the art (Kassapis et al. (202T) by running the pre-trained model provided by Kassapis et al.) (222I) with 35 samples and show the result the right. We that model achieves better calibration result
Ours
Kassapis al. (2021)
Arcuracy
Accuracy
Gap
Gap
ECE: 0.321%
ECE: 2.196%
Accuracy
Accuracy
02
3.
1.
00
Confidence
Confidence
Figure 6: We show the reliability plots of our model (left) and Kassapis al. 2021) (right) on the Cityscapes dataset evaluated with 35 samples. Our result is taken from one of the three random runs Table 3 last column, last row.
D.2 CALIBRATION IN MODE RATIO AND PIXEL FLIP RATIO
To have a further analysis of the predictive distribution and see how well it learns the characteristic of the multimodal ground truth distribution, we follow Kohl al. 2018 Kassapis al. 202 to
Notice that our model has all confidence scores less than 0.8. This is reasonable since the maximum occurrence probability of the evaluated ground truth class is $ s 0.76 (ef. See. 4.1).
22

---

Published as a conference paper at ICLR 2023
Ours
Estimated Frequencies
Ground truth
0.10
0.08
0.06
0.04
0.02
sidewalk
O
G
0
car
vegetation
Discrete Modes
Kassapis et al. (2021)
Estimated Frequendies
0.14
Ground truth
0.12
0.10
0.08
0.08
0.04
0.02
Plsii wise Priqurny Cranation
sidewalk
person
O
O
vagetation
Discrete Modes
Figure 7: We compare the calibration of our method and Kassapis et al. (2021) in terms of mode- wise probabilities and pixel-wise class flip ratios. the table below, we use '0" demonstrate the class is not flipped, and use '1' to demonstrate the class is flipped.
23

---

Published as a conference paper at ICLR 2023
calculate the mode ratio and pixel flip ratio of model outputs. To calculate the mode ratio, we match each output sample to the nearest ground truth label to determine its mode, where the distance is computed as I-IoU. For the class flip ratio, we count the relative proportions between the five flipped classes and the corresponding original classes for all pixels across all samples. In Figure 7 show the mode ratio and pixel-wise class flip ratio of method, and compare with [assapis iet al.j (202I). Specifically, we use the weighted output representation with 35 samples (last row in Table 3) and generate the output of Kassapis al. (202)) by running the pre-trained model provided, and evaluated with 35 samples (row #4 in Table 3). We see that our method has better calibration with the ground truth probabilittes both mode-level and pixel-level.
D.3 ANALYSIS ON THE TYPE OF UNCERTAINTY
We conduct experiments to analyze the type of predictive uncertainty on the Cityscapes dataset similar to the experiments conducted the LIDC dataset. We refer to Sec. C.3 for the detail explanation on experimental settings and do not repeat here. Results are evaluated in 35 samples with standard-sampling sampliin representation. As shown Table we get similar observations as in the LIDC where both the uncertainty scores and the evaluation metrics do not change much with the varies in training data size. Therefore, the epistemic uncertainty has less account the predictive uncertainty of our model.
We conduct experiments on the Cityscapes dataset with different training data sizes. Re- sults are evaluated in 35 samples with Sandard-samplingg representation.
Table 11:
Train datasct
Entropy
Sample Diversity
GED +
M-IOU f
ECE | (%)
Cityscapes /
0.823 i 0.006
0.598 + 0.002
0.157 1 0.002
0.686 1 0.004
1.878 1 0.009
Cityscapes/ 2
0.814 0.004
0.597 + 0.002
0.155 0.001
0.687 + 0.005
1.858 + 0.013
Cityscapes
0.817 = 0.010
0.595 = 0.004
0.155 + 0.002
0.690 + 0.002
1.861 = 0.015
D.4 ANALYSIS ON THE NUMBER OF EXPERTS
For the main experiments on the Cityscapes dataset, we assume the ground truth mode number (32) is known and hence choose slightly larger expert number K 35. This allows to validate how the model deals with the redundant experts. If we do not have any prior information on the mode number, we can also use validation set to draw the curve on how performance changes with the increase of expert numbers and select the K number at the inflection point of the curve.
Below, we show the experimental analysis of the expert number K for model training on the Cityscapes dataset. Specifically, we evaluate the model performance with expert-number K = 25, 30, 32, 35, 40, 50, 80 where 35 our default expert number. As shown in Figure g we see that when the expert number is smaller than 40, with the increasing number of experts, the model achieves a better GED score in general. However, when the expert number is much larger (eg. 50, 80 experts), the model performance decreases, corresponding to increase of GED score. This phenomenon similar to experiments on the LIDC dataset where we have discussed in C.4 and do not repeat here. Besides, K = 32 is the inflection point and can be chosen when prior information is known. Results are evaluated with 35 samples in standard -aapdddn-am representation.
D.5 COMPLEMENTARY QUALITATIVE ANALYSIS
In Figure 9 we show the predictions from all 35 experts. We matched 32 of them with 32 ground truth modes for easy comparison. We see that the expert networks learn the variety of modes well and the gating network gives well-calibrated mode probabilities. Also, we notice that our model learns relatively compact experts, with the extra experts shown at the bottom the predictions having relatively small weights.
E DISCUSSION
24

---

Published as a conference paper at ICLR 2023
0.185
GED (35)
0.180
0.175
0.165
30
50
60
80
Number experts
Figure The GED score changes with number experts, evaluated with 35 samples in standard saapding-aamping representation. Experiments are conducted the Cityscapes dataset.
E.1 WHY IT WORKS BETTER:
We consider the following two factors as the main reasons for the superior performance our method:
A flexible multi-modal representation: We explore explicit multi-modal uncertainty repre- sentation in a form mixture of stochastic experts, which differs from the previous work with more restricted structures (Kohl et al.) 2018 2019 Baumgartner et al.) 2019 Monteiro et al.) 2020). Our representation provides a divide-and dyvide-and-conuuee strategy for capturing complex multi-modal distribu- tions: each expert learns a clustering of segmentation outputs, which simpler model, while the gating network focuses on the overall frequency of those learned modes. This is validated in our empirical results, where our method captures the probability of each mode in Cityscapes much better than previous SOTA (Kassapis et al., 2021).
Wasserstein-like loss: We directly minimize the Wasserstein distance between segmentation mask distributions, while previous work either adopt the KL divergence (Kohl al. 2018 2019 Baum-
gartner et al.) 2019 Monteiro et al., 2020) or JS divergence (Kassapis et al. 2021 as the loss.
Wasserstain distance is known to have better characteristies measuring the geometry of distribu-
tions (Feydy et al. 2019) and mitigate the mode collapse problem (Arjovsky et al.) 2017) and show promising effect in the closely-related generative modeling tasks (Arjovsky al."
2017). We have
also conducted additional experiments by replacing the current loss and training pipelines to VAE style, and demonstrates the superiority of our Wasserstein-like loss (see Appendix C.7
E.2 LIMITATIONS:
We discuss some the limitations our method in the following: Firstly, when the expert number larger than the real mode number in the dataset, it is possible to have two expert networks giving similar predictions. As discussed Sec. this makes the matching process unstable and we adopt the soft gradient trick to handle this scenario. In the future, efforts be done by either stabilizing
25

---

Published as a conference paper at ICLR 2023
the matching process automatically pruning/merging those duplicated experts. Secondly, the matching is based on the cost matrix, however, some existed pairwise loss functions may not provide good measure for comparison. For example, the LIDC dataset where the class imbalance scenario exists, using the cross-entropy loss makes the measure dominated by the background pixels, which leads to less satisfactory matching results and further disturbs the training.
E.3 COMPARE TO ENSEMBLE MODELS
While our MoSE architecture shares certain similarity with the ensemble-hased models (i.e., com- bination of simpler models), there are three key differences between our framework and the ensem- ble methods:
Different goals: Most mixture-basse ensemble methods aim to learn a predictor with reduced model variance by aggregating multiple variants of the base models. By contrast, use the mixture architecture to explicitly capture the multi-modal distribution caused by the aleatoric uncertainty, in which each component corresponds different mode. We intend learn those modes from data without explicit supervision on the mode identity.
Learning strategy: The ensemble methods typically learn each component model based on original task loss (mostly with unique groundtruth for each input), and it is usually unclear which part of the data distribution it captures. In contrast, we design a Wasserstein- the distribution distance between the model outputs and the groundtruth. Thanks to the OT-based formulation, each expert of our model learns to capture a specific mode in the dataset.
loss to minimize
Representation: We develop a MoSE structure specifically for our uncertainty modeling task. Our stochastic expert models share most of their parameters and only vary in the latent distributions, which facilitates the learning of compact and meaningful modes. This differs from the conventional design of the component models with totally different parameters in the ensemble methods. Empir- ically, we have conducted experiments with our OT-based loss replaced by the IoU-loss (cf. Sec. 4.3.1 which is largely equivalent to learning standard ensemble model. As shown in Table #1, this leads to a performance drop of more than 0.3 in GED, demonstrating that naive ensem- bling technique may not suffice for this task. We note that Kohl et al. (2018) have also tested the performance of a Unet ensemble method and reached a similar conclusion.
26

---

Published as a conference paper at ICLR 2023
0:023
0.007
0.005
0:023
0.005
0:032
0.008
0.010
0.905
0.914
(2)
0:009
0:030
0:012
0.021
0:028
0.032
0.945
0:109
0:129
0:012
0:021
0:128
0:133
0:044
(3)
0:113 3
0.01
0:016
0:117
0:159
0.025
0:042
0:017
0:058
0.025
0:941
(4)
0:053
0:068
0:037
0:097
0:022
0:076
03115
0:052
0:070
0:036
0.093
0:021
0:076
0:316
(5)
0:2099
0:016
0:005
0:040
(1)
0.023
0.005
0.033
0.008
0.010
0.005
0.014
0.023
0.005
0.032
0.008
0.010
0.005
0.014
(2)
0.009
0.030
0.012
0.021
0.028
0.032
0.045
0.009
0.029
0.012
0.021
0.028
0.033
0.044
33)
0.013
0.0:1
0.016
0.017
0.059
0.025
0.042
0.013
0.0:1
0.017
0.017
0.058
0.025
0.041
(4)
0.053
0.068
0.037
0.097
0.022
0.076
0.015
0.052
0.070
0.038
0.093
0.021
0.076
0.016
(5)
0.109
0.0.8
0.00/
0.040
0.104
0.018
0.007
0.002
0.003
0.003
Figure 9: The extra visualization results on the Cityscapes. On the first line of each case, show the input image masked with ground truth label. Rest lines show samples from 35 experts and 32 ground truth modes. For clarity, we show matched pairs for ground truth masks on the top and predicted masks the bottom and divide them into 5 subgroups labeled (1) (5). The predictions are generated by sampling once from each expert, we denote the corresponding predicted probabilities on the left head of each prediction.
27


Reviews:
[{"review_id": "imZjQvu5F__", "reviewer": "Reviewer_bd5h", "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": "This work presents a two-step process to model the aleatoric (data) uncertainty in image segmentation using a mixture of experts type of an approach. The main assumption in this work is that each expert could correspond to a mode in the true, underlying uncertainty distributions that are multi-modal. Assuming the number of modes/experts to be known, a stochastic neural network predicts segmentation masks. A second gating network trained to mimic the expert behaviour by predicting the probabilities of experts. Experiments on two common multi-annotator datasets are reported and compared to relevant baseline methods using measures such as generalized energy distance (GED).", "summary_of_the_review": "**Edit (Scores Update)**: Updating from 3 to 6. \n\nCapturing segmentation uncertainty that is well calibrated can be useful in many applications, such as in medical image analysis. This work presents a two-step approach of modelling expert level behaviour and then within a mode capturing finer pixel-level uncertainties. The use of existing metrics such as GED is problematic. Also the connections to VAEs with Gaussian mixture priors is entirely missing.", "strength_and_weaknesses": "**Strengths:**\n\n* The primary motivation of capturing segmentation uncertainty as a measure of aleatoric uncertainty has been studied widely. This work presents a two-step approach to capture mode level differences (experts) and then finer, pixel level uncertainties per mode. The analogy of treating modes as different experts is reasonable but has issues which are desrcibed further below.\n\n* The experiments are performed on two commonly reported multi-rater datasets, with relevant baselines. The reported results show improved performance for the proposed method.   \n\n**Weaknesses:**\n* **Expert specific modes**: The assumption that each annotator can be modelled as a mode is not well motivated/justified. This implies that the uncertainty in segmentation is only between raters. In medical image segmentation tasks this is seldom the case. Variations in the data that cause ambiguity in segmentations could be due to other factors such as differences in acquistion, pathologies and intra-rater variability. This work explicitly models and trains each mode to correspond to a different rater, and in doing so it is highly restricted. How would this mixture of experts approach extend to differences arising due to medical data acquired from different centers, for example?\n\n* **Input agnostic, expert specific latent prior**: The latent distributions for $z_k \\sim \\mathcal{N}(m_k,\\sigma_k I)$ is said to be input agnostic with capabilities to model different experts. Specifically, how are the parameters of the latent priors obtained? In Sec 3.2.1 the expert specific prior parameters $m_k$ is mentioned but not how these values are trained. Also, the parameters of the categorical prior over the mixture components is unclear. \n\n* **VAE with Gaussian mixture prior**: The presented model with its categorical prior over the mixture components and then mode specific latents resembles a VAE with Gaussian mixture prior, which has been studied extensively such as in [4]. Is there a specific reason the authors deviate from this established notion of modelling Gaussian mixture type latents within generative model settings?  To me this model reduces to the probabilistic U-net with a Gaussian mixture prior [5]; making this connection to [4] and [5] also could help with contextualising the contribution. \n\n* **Gating module**: The key point of departure from [4,5] is the gating network; the role of which is unclear to me. Is the gating module injecting the input image features to choose which expert should be chosen? If yes, is the rater specific probabilities based on the input image?\n\n* **What happens when K=1**: In instances when there are no multiple raters to train from i.e., when K=1, does this model still capture aleatoric uncertainty? For instance, in [2,5] experiments when trained with only a single rater also show the model is capable of optimising for metrics such as GED. \n\n* **GED as the main metric**: GED is a distribution matching score which is trained to match the diversity in annotation variability and the predicted segmentations. This work primarily focuses on matching the annotator diversity as a means to quantify the uncertainty. While this is also what most existing literature is doing, there are concerns on the usefulness of this approach. For instance, this work in [1] takes up the usefulness of segmentation measures and clearly shows that GED is not a very useful measure. Using sample diversity to be somehow a measure of meaningful uncertainty is problematic, in my opinion. And I think as a community we need to move away from this notion and think about more useful ways of quantifying segmentation uncertainty, based on some of the discussions in [1].\n \n\n[1] Mehta, Raghav, Angelos Filos, Ujjwal Baid, Chiharu Sako, Richard McKinley, Michael Rebsamen, Katrin Dätwyler et al. \"QU-BraTS: MICCAI BraTS 2020 Challenge on Quantifying Uncertainty in Brain Tumor Segmentation--Analysis of Ranking Metrics and Benchmarking Results.\" arXiv preprint arXiv:2112.10074 (2021).\n\n[2] Christian F Baumgartner, Kerem C Tezcan, Krishna Chaitanya, Andreas M Hötker, Urs J Muehlematter, Khoschy Schawkat, Anton S Becker, Olivio Donati, and Ender Konukoglu. Phiseg: Capturing uncertainty in medical image segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention, 2019.\n\n[3] Bø, Hans Kristian, et al. \"Intra-rater variability in low-grade glioma segmentation.\" Journal of neuro-oncology 131.2 (2017): 393-402.\n\n[4] Dilokthanakul, Nat, et al. \"Deep unsupervised clustering with gaussian mixture variational autoencoders.\" arXiv preprint arXiv:1611.02648 (2016).\n\n[5] Kohl, Simon, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R. Ledsam, Klaus Maier-Hein, S. M. Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. \"A probabilistic u-net for segmentation of ambiguous images.\" Advances in neural information processing systems 31 (2018).\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "bChh7bLXq6q", "reviewer": "Reviewer_LSxs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "The paper proposes a novel method for modelling predictive uncertainties for segmentation tasks. The model consists of an encoder-decoder architecture with multiple decoder heads that are weighted by a separate gating network that predicts the weighting based on the encoded image code. The method is compared on the LIDC and cityscapes benchmarks and exhibits competitive performance.", "summary_of_the_review": "The paper proposes a simple yet novel and effective method at capturing the uncertainty segmentation predictions. The experiments support the claims of the paper.", "strength_and_weaknesses": "Strengths:\n- The paper is well written and easy to follow. It tackles an important problem of capturing uncertainty in segmentations.\n- The paper proposes a simple yet novel method of combining multi-head predictions (which reminds me of [1]) together with a gating network to capture the distribution of segmentation predictions.\n- The paper compares to relevant baselines on standard benchmarks and shows competitive results.\n\nWeaknesses:\n- It would be interesting to compare the proposed method to an ensemble of predictive methods that would be able to capture multi-modality.\n\nMisc:\n- Why are methods like the Prob. U-Net incapable of capturing multi-modality?\n- What's the overhead of this method compared to a single prediction head? Is the comparison fair given that the method might have more parameters etc?\n- Do you have an intuition why \"compact\" performs better than the \"standard sampling\"?\n\n[1] Osband, Ian, et al. \"Deep exploration via bootstrapped DQN.\" Advances in neural information processing systems 29 (2016).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "a6ugs3TiaM", "reviewer": "Reviewer_VqFH", "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": "This paper presents a new architecture for modelling the quantification of aleatoric segmentation uncertainty, i.e. predicting the distribution of segmentations in a given task. The model is tested on two established  datasets. Specifically, the author presents a mixture of stochastic experts model whose parameters are optimised by minimising a loss function based on the optimal transport problem. The coupling is calculated between the relative frequency of the ground truth annotations and the predicted probability for an expert given by the gating network. For the cost of transportation the IoU is used. To enable backpropagation, the original formulation of the optimal transport problem is substituted by an objective derived by constraint relaxation.\nThe author finds that their model outperforms all benchmark models in terms of common performance measures used in this domain of research on the LIDC dataset, as well as on some for the Cityscape dataset.", "summary_of_the_review": "The presented method is to the best of my knowledge novel and the authors present it well. However, the paper would benefit from a more detailed discussion and a provided intuition of why the proposed method is superior to the benchmark models. I encourage the authors to address the questions posed above. ", "strength_and_weaknesses": "**Strengths:**\n- The method is novel in the context of segmentation uncertainty and well presented\n- Good structure of the paper \n- comparison to many benchmark models \n- Results show improvement \n\n**Weaknesses:**\n- Not really clear if the benchmark models are comparable in terms of number of parameters etc. Are the benchmark models own implementations? \n- Not really clear whether the constraint relaxation leads theoretically to the same result as the original problem formulation. Can you clarify why those choices are justified beyond the fact, that they make the problem calculation tractable?\n- There is no discussion of the model beyond two limitations in the appendix. Why does it work?  \n- No reference to previous use of optimal transport based losses in segmentation, e.g.  *Fidon, Lucas, Sébastien Ourselin, and Tom Vercauteren. \"Generalized wasserstein dice score, distributionally robust deep learning, and ranger for brain tumor segmentation: BraTS 2020 challenge.\" International MICCAI Brainlesion Workshop. Springer, Cham, 2020.* \n\n**Questions:**\n\n- Sec. 3.1. In practice the relative frequency of each annotation is always 1/n in your model since it is quite unlikely that you get the exact same annotation from different experts, so v_n is always a uniform distribution, assigning same probability to all ground truth annotations. From an information theoretic point of view v_n contains no information. Therefore, all information must be inferred from the ground truth segmentation masks over the cost function. \nCan you elaborate on how the gating network still learns to predict probabilities that correspond to the relative mode frequencies? What do those probabilities express then? \n\n- How would your model behave in the limit case of many experts K? Do the curves in Fig. 5 and Fig. 8 keep decreasing?  \n\n- I would be interested into a row of Table 2 where you test  stochastic experts, uniform expert weights and the IoU-loss.\n\n- Can you elaborate if there is a relation to ensemble models, especially if you set S=1? \n\n- Sec. 4.2 In your experiments for LIDC you set the experts to 4 (equal to the number of annotations available per image) while you set the number of experts to 35 for the Cityscape dataset. This corresponds to the number of ground truth annotations for LIDC and approximately to the number of classes in Cityscapes. How do you justify the assumption that each ground truth annotation (or class in Cityscapes) contains exactly one mode?\n\n- Where do you draw the line between variation around one mode and different mode? For example in LIDC one could argue that one hypothesis is that there is no cancer and the other is that there is cancer and that these should be the modes and that all variation resulting from the experts uncertainty in drawing the segmentation mask should come from the variation around the second mode.  \n\n- How do you treat borderline cases of IoU in the cost function matrix? For example empty segmentation mask compared to segmentation mask that contains a mask? \n\n- What is the intuition why the model works better than the benchmarks? One major factor I see is that during training you use all ground truth annotations per image in one gradient step, which is a main difference to the compared models. The paper would benefit from a discussion of the method a lot. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster