Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
The 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.
strength_and_weaknesses: **Strengths:**
- The method is novel in the context of segmentation uncertainty and well presented
- Good structure of the paper 
- comparison to many benchmark models 
- Results show improvement 

**Weaknesses:**
- Not really clear if the benchmark models are comparable in terms of number of parameters etc. Are the benchmark models own implementations? 
- 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?
- There is no discussion of the model beyond two limitations in the appendix. Why does it work?  
- 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.* 

**Questions:**

- 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. 
Can 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? 

- How would your model behave in the limit case of many experts K? Do the curves in Fig. 5 and Fig. 8 keep decreasing?  

- I would be interested into a row of Table 2 where you test  stochastic experts, uniform expert weights and the IoU-loss.

- Can you elaborate if there is a relation to ensemble models, especially if you set S=1? 

- 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?

- 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.  

- 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? 

- 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. 
clarity,_quality,_novelty_and_reproducibility: *Clarity:*
The paper is well structured and presented in a accessible fashion.

*Quality:*
I critic the missing comments on the comparability with the benchmark models, this needs to be enhanced. The authors conducted many experiments to back their claims and provide extensive supplementary material that clarifies many questions during reading.   

*Novelty:*
The proposed method is as far as I know novel in the context of segmentation uncertainty. While Wasserstein losses and the method of constraint relaxation have been used before in multi class segmentation I am not aware of their application in the context of aleatoric segmentation uncertainty. 

*Reproducibility:*
The authors publish their code upon publication.
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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
strength_and_weaknesses: Strengths:
- The paper is well written and easy to follow. It tackles an important problem of capturing uncertainty in segmentations.
- 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.
- The paper compares to relevant baselines on standard benchmarks and shows competitive results.

Weaknesses:
- It would be interesting to compare the proposed method to an ensemble of predictive methods that would be able to capture multi-modality.

Misc:
- Why are methods like the Prob. U-Net incapable of capturing multi-modality?
- 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?
- Do you have an intuition why "compact" performs better than the "standard sampling"?

[1] Osband, Ian, et al. "Deep exploration via bootstrapped DQN." Advances in neural information processing systems 29 (2016).
clarity,_quality,_novelty_and_reproducibility: The paper is well-written and easy to follow. Most details for reproducibility are included in the paper.
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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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).
strength_and_weaknesses: **Strengths:**

* 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.

* The experiments are performed on two commonly reported multi-rater datasets, with relevant baselines. The reported results show improved performance for the proposed method.   

**Weaknesses:**
* **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?

* **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. 

* **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. 

* **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?

* **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. 

* **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].
 

[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).

[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.

[3] Bø, Hans Kristian, et al. "Intra-rater variability in low-grade glioma segmentation." Journal of neuro-oncology 131.2 (2017): 393-402.

[4] Dilokthanakul, Nat, et al. "Deep unsupervised clustering with gaussian mixture variational autoencoders." arXiv preprint arXiv:1611.02648 (2016).

[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).

clarity,_quality,_novelty_and_reproducibility: The paper is largely clearly written. There are some missing details or justification of choices which I have pointed out in the main review. While the method in itself is original, the approach and the problem being addressed has issues which cannot be overlooked. Also, there are connections to other existing work which are not explored/discussed.
summary_of_the_review: **Edit (Scores Update)**: Updating from 3 to 6. 

Capturing 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.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: Summary:
This paper presents a new mixture of stochastic experts method for modeling aleatoric segmentation uncertainty.

Strengths:
- Well written paper with novel ideas on segmentation uncertainty
- Thorough experimental validation

Weaknesses:
- Limited discussion of the model beyond the appendix. Some discussion should be added to the paper itself.
- Missing references
- Some relevant baselines are omitted, also from discussion, e.g. VAE with Gaussian micture prior
- Validation with GED metric has limited ability to draw conclusions
justification_for_why_not_higher_score: The reviews do not support it
justification_for_why_not_lower_score: All reviewers suggest acceptance
Decision: Accept: poster