A Probabilistic U-Net for Segmentation of Ambiguous Images
Many real-world vision problems suffer from inherent ambiguities. In clinical applications for example, it might not be clear from a CT scan alone which particular region is cancer tissue. Therefore a group of graders typically produces a set of diverse but plausible segmentations. We consider the task of learning a distribution over segmentations given an input. To this end we propose a generative segmentation model based on a combination of a U-Net with a conditional variational autoencoder that is capable of efficiently producing an unlimited number of plausible hypotheses. We show on a lung abnormalities segmentation task and on a Cityscapes segmentation task that our model reproduces the possible segmentation variants as well as the frequencies with which they occur, doing so significantly better than published approaches. These models could have a high impact in real-world applications, such as being used as clinical decision-making algorithms accounting for multiple plausible semantic segmentation hypotheses to provide possible diagnoses and recommend further actions to resolve the present ambiguities.
Post rebuttal: Authors have responded well to the issues raised, and I champion publication of this work.    Main idea:   Use a conditional variational auto-encoder to produce well-calibrated segmentation hypotheses for a given input.   Strengths:   The application is well motivated and experiments are convincing and state of the art. Convincing baselines, good discussion and conclusion. Nice illustrations in the appendix.   Weaknesses:   No significant theoretical contribution.   Possibly in response, the manuscript is a little vague in its positioning relative to prior work. While relevant prior work is cited, the reader is left with some ambiguity and, if not familiar with this prior work, might be misled to think that there is methodological innovation beyond the specifics of architecture and application.   Based on its solid and nontrivial experimental contribution I advocate acceptance; but the manuscript would profit from a clearer enunciation of the fact that / what prior work is being built on.   Comments:   50: "treat pixels independently"; this is unclear (or faulty?), the quoted papers also use an encoder/decoder structure  As a consequence, contribution 1 (line 74) is dubious or needs clarification.   Contribution (2) is fine.   Contribution (3): The statement is true, but the contribution is unclear. If the claim refers to the fact that latent z is concatenated only at the fully connected layers, then this has been done before (e.g. in DISCO Nets by Bouchacourt et al., (NIPS 2016)).  Contribution (4): The claim is vague. If we take generative models in their generality, then it encompasses e.g. most of Bayesian Statistics and the claim of only qualitative evaluation is obviously wrong. If we only consider generative models in the deep learning world, the statement is correct insofar as many papers only contain qualitative evaluation of "My pictures are prettier than yours"; but there are nevertheless quantitative metrics, such as the inception score, or the metrics used in the Wasserstein AE.   Section 2:  It is not made very obvious to the reader which part of this VAE structure is novel and which parts are not. The paper does follow [4] and especially [5] closely (the latter should also be cited in line 82). So the only really novel part here is the U-net structure of P(y|z,x). Concatenating z after the U-net in (2) is new in this formulation, but not in general (e.g. as already mentioned by DISCO Nets (Bouchacourt at al., NIPS 2016)).   Finally, there is no justification for the appearance of $\beta$ in (4), but it is up to the parameter name identical to what Higgins et al., (ICLR 2017) do with their \beta-VAE. Especially since authors choose $\beta >= 1$, which follows the Higgins et al. disentangling argument, and not the usual $\beta_t <= 1$, in which case it would be a time dependent downscaling of the KL term to avoid too much regularization in the beginning of the training (but then again the references to earlier work would be missing).