Source-linked AI summary
PHiSeg: Capturing Uncertainty in Medical Image Segmentation
Christian F. Baumgartner, Kerem C. Tezcan, Krishna Chaitanya, Andreas M. Hötker, Urs J. Muehlematter, Khoschy Schawkat, Anton S. Becker, Olivio Donati, Ender Konukoglu
TL;DR
Medical segmentation can admit multiple plausible masks because of ambiguous image evidence and differing annotator styles, while many methods learn only one image-to-mask mapping. PHiSeg models the conditional segmentation distribution with resolution-specific latent variables in a hierarchical probabilistic architecture. It produces more realistic and diverse samples than recent related work across single- and multi-annotator training, while its resolution-specific latent assignment is enforced by implementation.
Problem
Medical segmentation is ambiguous because image quality and annotator styles can support multiple plausible masks, but most methods learn a single image-to-segmentation mapping.
Method
PHiSeg uses a hierarchical conditional variational autoencoder with separate latent variables governing segmentation variation at different resolutions.
Results
PHiSeg produced significantly better distributional scores and wider sample variability than related methods, with similar results under single- and multi-annotator training.
Takeaways & Limitations
The authors conclude that modelling the segmentation probability distribution supports uncertainty estimation relevant to clinical acceptance and downstream applications.
Takeaways & Limitations
Each latent variable’s resolution-specific responsibility is enforced by the implementation rather than implied by the graphical model.
Abstract
from arXiv · showhide
Segmentation of anatomical structures and pathologies is inherently ambiguous. For instance, structure borders may not be clearly visible or different experts may have different styles of annotating. The majority of current state-of-the-art methods do not account for such ambiguities but rather learn a single mapping from image to segmentation. In this work, we propose a novel method to model the conditional probability distribution of the segmentations given an input image. We derive a hierarchical probabilistic model, in which separate latent variables are responsible for modelling the segmentation at different resolutions. Inference in this model can be efficiently performed using the variational autoencoder framework. We show that our proposed method can be used to generate significantly more realistic and diverse segmentation samples compared to recent related work, both, when trained with annotations from a single or multiple annotators.
1 Introduction
Medical image segmentation is ambiguous because image quality and expert annotation styles can support multiple plausible masks. PHiSeg addresses this by modelling a distribution of image-conditional segmentations rather than a single output.
- Motivation: Ambiguous image evidence and differing expert annotation styles can make multiple segmentation hypotheses plausible.Poor contrast and acquisition restrictions are cited as sources of inherent uncertainty.
- Motivation: Predicting only the most likely segmentation may lead to misdiagnosis and negatively affect downstream tasks.
- Related work: Bayesian weight inference, ensembles, and multi-headed networks address uncertainty incompletely by producing pixelwise variation or a fixed number of hypotheses.The cited approaches are contrasted with methods that can model richer segmentation distributions.
- Related work: Conditional variational autoencoders can generate unlimited segmentation samples conditioned on an input image, while probabilistic U-NET was reported to produce realistic samples from multiple expert annotations.
- Contribution: PHiSeg introduces a hierarchical probabilistic model that generates low-resolution segmentations and progressively refines them using separate latent variables at each resolution.The method is motivated by limited diversity in probabilistic U-NET samples.
2 Methods
PHiSeg represents image-conditional segmentations with a hierarchy of resolution-specific latent variables and neural-network parameterized prior, posterior, and likelihood distributions. Training maximizes a variational lower bound with cross-entropy supervision and heuristic KL weighting, while prediction samples latents from the prior and decodes them.
- Generative model: PHiSeg models p(s|x) using L hierarchical latent variables, with each latent variable assigned to a segmentation resolution.The conditional distribution factors through p(s|z1,...,zL) and successive conditional latent priors.
- Variational inference: The variational formulation approximates p(z|s,x) with q(z|s,x) and optimizes an evidence lower bound containing reconstruction and KL-divergence terms.
- Variational inference: Prior and posterior distributions are axis-aligned normal distributions parameterized by neural networks, while p(s|z) is a categorical distribution predicted by another neural network.This forms a hierarchical conditional variational autoencoder.
- Architecture: Structured latent variables form U-NET-like skip connections, refine representations across resolutions, and generate residual mask changes except at the bottom level.The latent dimensionality scales with spatial resolution, and sampling is required along each skip connection.
- Training: Training maximizes the lower bound using analytic KL terms, cross-entropy likelihood terms, single-sample expectations, KL weights αℓ=2ℓ−1, and deep supervision at each resolution.Deep supervision uses CE(ups(ŝℓ), sgt) for ℓ>1 and is intended to prevent latent variables from ignoring their resolution-specific roles.
- Prediction: At prediction time, segmentation samples are generated by sampling latent variables with the prior network and decoding them with the likelihood network.Testing uses only the prior and likelihood networks.
3 Experiments and Results
Experiments evaluated PHiSeg against probabilistic U-NET and related variants on LIDC-IDRI and prostate MR data, measuring distributional similarity, uncertainty localization, and segmentation accuracy. PHiSeg with multiple latent levels generally produced more diverse samples and stronger uncertainty-related results without reducing Dice performance.
- Datasets and setup: Experiments used LIDC-IDRI thoracic CT patches and an in-house prostate MR dataset with annotations from multiple radiologists and non-radiologists.The LIDC-IDRI dataset contains 1018 images, while the prostate dataset contains 68 patients.
- Datasets and setup: PHiSeg with L = 5 latent levels and seven total resolution levels was compared with probabilistic U-NET and a single-level PHiSeg variant.The comparison was designed to assess whether multi-resolution modeling, rather than network capacity, explains performance differences.
- Evaluation: The evaluation measured generalised energy distance, SNCC, and conventional Dice score using samples from the learned distributions.GED compared learned and ground-truth distributions, while SNCC assessed whether γ-maps predicted regions with large uncertainty; metrics used 100 generated samples.
- Results: With all annotators for training, PHiSeg (L = 5) produced significantly better GED and SNCC scores than the other methods.Its γ-maps were highly predictive of where average prediction errors occurred.
- Results: PHiSeg (L = 5) generated wider sample variability than probabilistic U-NET, while single-annotator training showed similar results and probabilistic U-NET sometimes became almost deterministic.PHiSeg (L = 1) performed worse, indicating that multiple resolution levels were important for the method.
- Results: No significant Dice differences were found among methods except PHiSeg’s (L = 1), indicating that stochastic modeling did not reduce segmentation performance for PHiSeg (L = 5) or probabilistic U-NET.The comparison included the deterministic U-NET.
4 Discussion and Conclusion
The paper concludes that PHiSeg models the conditional distribution of segmentation masks with a hierarchical probabilistic method. It reports improved metric performance and better prediction of the method’s own errors, while motivating uncertainty modeling for clinical acceptance and downstream applications.
- Conclusion: The paper introduced a hierarchical probabilistic method for modeling the conditional distribution of segmentation masks given an input image.This is the paper’s central methodological contribution.
- Conclusion: PHiSeg substantially outperformed the state of the art on a number of metrics and predicted its own errors better than previous work.The conclusion describes these as demonstrated outcomes of the experiments.
- Implications: The authors state that proper uncertainty modeling is important for clinical acceptance and that segmentation probability distributions may support downstream tasks.These statements are presented as the paper’s broader motivation and implications.
A Derivation of the Evidence Lower Bound (ELBO)
The appendix derives the evidence lower bound for the hierarchical conditional model by decomposing latent-variable terms and exploiting the graphical-model factorization. The resulting objective lower-bounds log p(s|x), with equality when the variational posterior equals the true posterior.
- Starting point: The derivation starts from the variational lower-bound decomposition for the conditional distribution p(s|x), with z defined as the collection of latent variables.The latent variables are denoted z1 through zL.
- Final ELBO: Because KL(q(z|s, x)||p(z|s, x)) is nonnegative, L(s|x) is a lower bound on log p(s|x), attaining equality when q(z|s, x) = p(z|s, x).This states the variational bound’s equality condition.
- Factorization: The graphical-model factorization separates zL and factorizes the variational distribution accordingly, allowing the KL-divergence terms to be simplified.The derivation uses independence of zL from the remaining latent variables where specified.
- Factorization: The KL terms are further decomposed across latent levels by regrouping integrals and using the factorized hidden-variable structure.The support of each latent variable is denoted Zℓ during the simplification.
- Final ELBO: Substituting the simplified expressions into the preceding decomposition yields the evidence lower bound used in the main article.The appendix explicitly identifies this substitution as the route to the final ELBO expression.
B Additional Samples
Appendix B presents additional segmentation samples for prostate and LIDC cases, covering both relatively small and substantial inter-expert disagreement. The LIDC examples distinguish disagreement about lesion shape from disagreement about lesion presence or absence.
- Additional samples are shown for the prostate and LIDC datasets using all investigated methods.
- Fig. B.1 presents a prostate MR slice with comparatively small inter-expert disagreement.
- Fig. B.2 presents a thoracic CT lesion where experts disagree about the lesion’s shape.
- Fig. B.3 presents a thoracic CT lesion where experts disagree over whether the lesion is present or absent.
- Each figure shows ground-truth annotations, method-generated samples, average error maps, and γ-maps, with the first and last columns expected to match for a perfect model.