A Probabilistic U-Net for Segmentation of Ambiguous Images

Simon A. A. Kohl; Bernardino Romera-Paredes; Clemens Meyer; Jeffrey De Fauw; Joseph R. Ledsam; Klaus H. Maier-Hein; S. M. Ali Eslami; Danilo Jimenez Rezende; Olaf Ronneberger

Source-linked AI summary

A Probabilistic U-Net for Segmentation of Ambiguous Images

Simon A. A. Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R. Ledsam, Klaus H. Maier-Hein, S. M. Ali Eslami, Danilo Jimenez Rezende, Olaf Ronneberger

arXiv:1806.05034v4Jun 13, 2018cs.CVcs.LGcs.NEstat.ML

TL;DR

Ambiguous images may admit multiple plausible segmentations, motivating a model that learns their conditional distribution rather than a single map. The proposed Probabilistic U-Net combines a U-Net with a conditional variational autoencoder and samples latent variants to generate hypotheses. Across lung-abnormality and Cityscapes tasks, it reproduced segmentation variants and their frequencies better than the considered baselines, including rare modes.

Problem
Image context can leave segmentation ambiguities unresolved, while decisions based on a single or pixel-wise-probabilistic map may be difficult or risky.
Method
The Probabilistic U-Net combines a U-Net with a conditional variational autoencoder and uses an input-conditioned latent distribution to generate multiple segmentation hypotheses.
Results
The model significantly outperformed considered baselines on matching coherent segmentation variants and their frequencies, including 32 Cityscapes modes with occurrence likelihoods from 0.5% to 10.9%.
Takeaways & Limitations
The model can support downstream diagnosis workflows by propagating multiple hypotheses, suggesting tests to resolve ambiguities, or enabling expert selection.
Takeaways & Limitations
Baseline comparisons used models sharing the same U-Net core and the same training schedule for fairness and comparability.

Abstract

from arXiv · show

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.

1 Introduction

Ambiguous images can support multiple plausible segmentation maps, yet pixel-wise uncertainty or deterministic predictions may fail to represent coherent alternatives. The paper introduces the Probabilistic U-Net to model these variants jointly, including their frequencies, for quantitative evaluation and downstream decision support.

Ambiguities in medical images can remain unresolved even with full image context, such as whether a visible lung lesion is cancerous.
Providing only the most likely segmentation may contribute to misdiagnoses and sub-optimal treatment when later decisions depend on the segmentation map.
Pixel-wise probabilities ignore pixel co-variances, whereas multiple consistent hypotheses can be propagated, tested, or selected in subsequent diagnosis pipelines.
The Probabilistic U-Net combines a CVAE with a U-Net to generate consistent segmentation maps, model joint pixel likelihoods, and reproduce low-probability variants at corresponding frequencies.
The framework provides joint mode likelihoods, can represent rare modes with calibrated probabilities, samples cheaply, and supports quantitative evaluation that penalizes missing modes.

2 Network Architecture and Training Procedure

The model combines a conditional variational autoencoder with a U-Net to learn a conditional distribution over segmentations. A latent sample selects a segmentation variant, while training aligns posterior and prior distributions and enforces agreement with ground-truth maps.

The architecture combines a conditional variational autoencoder with a U-Net to learn a conditional density model over segmentations given an image.
Sampling: A low-dimensional latent space represents segmentation variants, and an input-conditioned Gaussian prior supplies random samples used to generate multiple segmentations.
Sampling: Each latent sample is broadcast into a feature map and combined with U-Net activations by three 1 × 1 convolutions to produce its segmentation map.
Sampling: For repeated samples of one image, the prior output and U-Net activations can be reused, so only the combination function must be reevaluated.
Training: Training minimizes a conditional-VAE variational lower bound while learning an embedding of segmentation variants in the latent space.
Training: The posterior network maps an image and ground-truth segmentation to a latent distribution whose sampled representation must reconstruct that ground-truth segmentation.
Training: A cross-entropy term penalizes segmentation differences, while a weighted KL divergence penalizes differences between posterior Q and prior P.

3 Performance Measures and Baseline Methods

The paper evaluates segmentation distributions with generalized energy distance and compares the proposed framework against several U-Net-based stochastic or multi-output baselines. These baselines vary in how they generate uncertainty, diversity, and multimodal outputs.

3.1 Performance measures: Generalized energy distance compares predicted and ground-truth distributions using distances between segmentation observations.The paper uses d(x, y) = 1 − IoU(x, y), making the metric sensitive to agreement between segmentation samples.
3.2 Baseline methods: The comparison includes Dropout U-Net, U-Net Ensemble, M-Heads, and Image2Image VAE baselines built around U-Net components.The baselines are trained comparably, with the same U-Net core used where applicable.
3.2 Baseline methods: Dropout U-Net samples stochastic predictions by applying dropout to selected encoder and decoder activations during training and sampling.The dropout probability is p = 0.5 in the three innermost encoder and decoder blocks.
3.2 Baseline methods: U-Net Ensemble produces one output per independently trained member, with the number of members matched to the required number of samples.The deterministic U-Net is treated as the one-sample corner case of the ensemble.
3.2 Baseline methods: M-Heads branches into multiple output heads and uses an adjusted cross-entropy loss to promote diverse hypotheses while reducing idle heads.The best-performing head receives weight 1 − ϵ, while the remaining heads each receive ϵ/(M − 1).
3.2 Baseline methods: Image2Image VAE uses broadcast normal latents and cross-entropy training, but its prior and posterior are not conditioned on the input image.Unlike the proposed model, it omits adversarial supervision in the fully supervised segmentation setting.

4 Results

The experiments evaluate multimodal segmentation on expert-annotated lung CT images and artificially ambiguous Cityscapes images. The proposed model improves distributional agreement and reproduces the frequencies of many segmentation modes, including rare ones.

4 Results: The evaluation uses LIDC-IDRI with four expert annotations per input and Cityscapes with artificially created ambiguous semantic classes.The Cityscapes task introduces uncertainty by randomly flipping five classes into synonymous alternatives.
4.1 Lung abnormalities segmentation: The LIDC-IDRI test task focuses on cropped lesion-centered CT images, which isolates regions where annotators disagree but is not directly clinically relevant in isolation.Cropping removes large areas where all labelers agree.
4.1 Lung abnormalities segmentation: Squared generalized energy distance decreases for all models as more samples are drawn, indicating better matching of the ground-truth distribution and greater sample diversity.The result is reported on 1992 large-lung-abnormality test images.
4.1 Lung abnormalities segmentation: The Probabilistic U-Net outperforms all baselines when sampling 4, 8, and 16 times.At 16 samples, its performance is significantly higher than the baselines, with p-value ∼O(10^-13).
4.2 Cityscapes semantic segmentation: The Cityscapes construction yields 32 discrete segmentation modes with ground-truth probabilities ranging from 10.9% to 0.5%.These modes arise from independently flipping five semantic classes with different probabilities.
4.2 Cityscapes semantic segmentation: The proposed model closely captures frequencies across the combinatorial Cityscapes hypothesis space, including very rare modes, thereby supplying calibrated mode likelihoods.The Image2Image VAE captures all variants but has inaccurate frequencies, while the other baselines fail to represent the modes adequately.

5 Discussion and conclusions

The Probabilistic U-Net models coherent segmentation distributions, including rare variants, and outperforms the considered baselines. Its latent space also supports inspecting how segmentation variants are organized and may support downstream clinical decision-making.

5 Discussion and conclusions: The model produces coherent segmentation maps that match multi-modal ground-truth distributions in both lung abnormalities and Cityscapes tasks.The IoU-based energy distance evaluates coherence and expected variant frequencies, including penalties for missing variants.
5 Discussion and conclusions: The model significantly outperforms the considered baselines on matching the joint likelihood of segmentation variants.The comparison evaluates whether samples are both close to ground truth and produced with expected frequencies.
5 Discussion and conclusions: With 32 Cityscapes modes occurring from 0.5% to 10.9%, the model best matches expected frequencies, including very unlikely modes.It does so without prior assumptions about the number of modes or hypotheses.
5 Discussion and conclusions: The VAE-like latent space exposes the likelihood landscape underlying the modeled segmentation distribution.Unlike GAN-based approaches, the model explicitly parameterizes distributions, enabling latent-space inspection.
5 Discussion and conclusions: Separating the prior from the segmentation network can condition on an entire 3D scan while processing it tile by tile under limited GPU memory.This design is especially relevant to medical imaging with large volumetric scans.
5 Discussion and conclusions: The model’s hypotheses and likelihoods could inform diagnosis probabilities or guide actions that resolve ambiguities in medical imaging.The authors also suggest replacing deterministic U-Nets and applying latent-space inspection beyond explicitly multi-modal tasks.

A.1 Lung Abnormalities Segmentation

In the lung-abnormalities example, the learned latent space organizes lesion-size and shape variation while aligning grader masks with corresponding regions of the prior.

A.1 Lung Abnormalities Segmentation: The z0 component roughly encodes lesion size, including a transition to complete lesion absence.The particular example assigns relatively little probability mass to absence, consistent with one of four graders assessing the image as lesion-free.
A.1 Lung Abnormalities Segmentation: The z1 component appears to encode shape variations among lesion segmentations.The latent representation separates this variation from the lesion-size component in the described example.
A.1 Lung Abnormalities Segmentation: The KL-divergence ties posterior and prior distributions so grader masks can be projected into the same latent space.Three graders agreeing on lesion presence lie within the prior’s 1-sigma interval, while the absence grader lies near the region encoding absence.

A.2 Street Scene Segmentation

For Cityscapes, the model uses a higher-dimensional latent space to represent many discrete grader modes as coherent regions, while displayed slices show only part of that structure.

A.2 Street Scene Segmentation: The Cityscapes model uses more latent dimensions than the lung task to provide sufficient capacity for grader modes.The best-performing model used a 6D latent space, while the discussion visualizes a 3D version.
A.2 Street Scene Segmentation: The precisely defined grader modes form coherent, well-confined regions in the latent space.Because the space is 3D, not all 32 modes are fully visible in the displayed z2 slice.
A.2 Street Scene Segmentation: Transitions between latent regions can produce mixtures of discrete grader modes, so ideally they should be sharp relative to mode probability scales.The passage describes these transitions as an inherent consequence of representing discrete modes continuously.
A.2 Street Scene Segmentation: The 3D visualization maps 19 × 19 samples into the z0-z1 plane with z2 = 0, using white numbers for grader-mask positions.Yellow circles indicate deviations from the mean in sigma, and samples are shown as squeezed rectangles under a spherical unit-Gaussian rescaling.

Appendix B Metrics

The appendix evaluates distributions of segmentation samples using mode-wise and pixel-wise frequencies, with IoU-derived distances and task-specific ground-truth handling.

Appendix B Metrics: For LIDC, the evaluation uses m = 4 ground-truth samples and n model samples in the reported statistic.The statistic compares distributions rather than a single deterministic prediction and ground truth.
Appendix B Metrics: The lesion-mask distance is d(x, y) = 1 − IoU(x, y), with empty-mask pairs assigned distance 0.This definition rewards agreement when both prediction and ground truth indicate lesion absence.
Appendix B Metrics: For Cityscapes, the known ground-truth distribution is a mixture of M = 32 Dirac delta distributions and is used directly in the estimator.Sampling from the ground-truth distribution is therefore unnecessary.
Appendix B Metrics: Cityscapes mode weights ω_j represent delta distributions concentrated at their corresponding ground-truth modes.The distance depends on average IoU over the 10 switchable classes, and absent predicted classes receive a 0 score.
Appendix B Metrics: Mode-wise frequencies count samples matched to their closest ground-truth mode, while pixel-wise frequencies estimate p(predicted class|ground-truth class).Figures 8–10 compare these frequency summaries across baselines and the proposed model.

Appendix D Ablation analysis

The ablation study compares three architectural variations with the original Probabilistic U-Net and finds that the original design achieves better agreement with the ground-truth segmentation distribution. Injecting latent variables at the final U-Net stage outperforms injecting them at the beginning, while preserving sampling efficiency.

The ablation study evaluates three variations: a fixed Gaussian prior, a fixed prior with a context-independent posterior, and latent-feature injection at the U-Net input.
Lower energy distance indicates better agreement between predicted and ground-truth segmentation distributions.
The original approach performs better than the architectural variations in energy distance.
Latent injection at the U-Net end outperforms injection at the beginning, despite the original choice being motivated by sampling efficiency.
The authors hypothesize that final-stage latent injection makes it easier to represent different segmentations for the same input.

Appendix E Predicting ground truth ambiguity from models’ samples

The appendix tests whether models’ sampled lesion-presence histograms can distinguish ambiguous from unambiguous lesions. Using 16 samples per instance, the Probabilistic U-Net shows an advantage over competing models on this discrimination task.

An instance is ambiguous when at least one grader disagrees about the presence of abnormal tissue.
Each model produces 16 samples per instance, and the number of samples containing a lesion is binned into separate histograms for ambiguous and unambiguous instances.
The best threshold separating ambiguous and unambiguous cases is selected on the validation set and evaluated for accuracy on the test set.
The Probabilistic U-Net has an advantage over competing models in distinguishing ambiguous from unambiguous lesions.

Appendix F Sampling LIDC masks using different models

This appendix presents qualitative samples from the Probabilistic U-Net and baseline models on LIDC, plus Probabilistic U-Net samples on Cityscapes. The figures compare model-generated masks with expert or ground-truth segmentation variants.

LIDC results: The LIDC numerical results corresponding to the qualitative examples are reported in Table 2.
Cityscapes results: The Cityscapes numerical results corresponding to the examples are reported in Table 3.
LIDC qualitative examples: LIDC examples compare 16 Probabilistic U-Net samples with the four graders’ ground-truth masks for the same 15 test subjects.
LIDC qualitative examples: The Dropout U-Net, U-Net Ensemble, M-Heads, and Image2Image VAE are shown using the same qualitative-example layout as the Probabilistic U-Net.
Cityscapes qualitative examples: Cityscapes examples show input images, four of 32 ground-truth modes, and 16 random network samples, with masked pixels marked in black.

Appendix H Training details

The appendix describes the training and preprocessing settings for the LIDC and Cityscapes experiments. It specifies lesion filtering, data augmentation, image sizes, and shared U-Net architecture details.

The experiments use a common section describing architecture settings and training procedures for both datasets.
LIDC training details: LIDC preprocessing retains lesions specified by polygon outlines, corresponding to lesions larger than 3mm, and filters selected DICOM files by slice-location consistency.
LIDC training details: LIDC training randomly samples image-grader pairs and applies elastic deformation, rotation, shearing, scaling, and translated cropping from 180 × 180 to 128 × 128 pixels.
Cityscapes training details: Cityscapes images and label maps are downsampled to 256 × 512 and augmented with geometric transformations plus random color augmentation.