[Paper Review] Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty

Miguel Monteiro, Loïc Le Folgoc+6 more|arXiv (Cornell University)|Jun 10, 2020

Explainable Artificial Intelligence (XAI)Computer Science48 references79 citations

TL;DR

Stochastic Segmentation Networks (SSNs) model the joint distribution of segmentation maps using a low-rank multivariate normal in logit space, enabling multiple spatially coherent hypotheses and efficient inference across 2D and 3D medical images.

ABSTRACT

In image segmentation, there is often more than one plausible solution for a given input. In medical imaging, for example, experts will often disagree about the exact location of object boundaries. Estimating this inherent uncertainty and predicting multiple plausible hypotheses is of great interest in many applications, yet this ability is lacking in most current deep learning methods. In this paper, we introduce stochastic segmentation networks (SSNs), an efficient probabilistic method for modelling aleatoric uncertainty with any image segmentation network architecture. In contrast to approaches that produce pixel-wise estimates, SSNs model joint distributions over entire label maps and thus can generate multiple spatially coherent hypotheses for a single image. By using a low-rank multivariate normal distribution over the logit space to model the probability of the label map given the image, we obtain a spatially consistent probability distribution that can be efficiently computed by a neural network without any changes to the underlying architecture. We tested our method on the segmentation of real-world medical data, including lung nodules in 2D CT and brain tumours in 3D multimodal MRI scans. SSNs outperform state-of-the-art for modelling correlated uncertainty in ambiguous images while being much simpler, more flexible, and more efficient.

Motivation & Objective

• Motivate the need to capture inherent uncertainty in image segmentation where multiple plausible boundaries exist (e.g., medical imaging).
• Propose a lightweight probabilistic framework that models joint label-map distributions without changing network architectures.
• Enable generation of multiple spatially coherent segmentation hypotheses from a single forward pass.
• Demonstrate applicability to 2D CT lung nodules and 3D multimodal MRI brain tumor segmentation.

Proposed method

• Model the logit map as a low-rank plus diagonal multivariate normal: η|x ~ N(μ(x), Σ(x)) with Σ = P P^T + D.
• Parameterize μ(x), P, and D as outputs of a neural network, enabling efficient computation and scalability.
• Approximate the intractable integral over η with Monte Carlo sampling in the logit space.
• Use the re-parameterization trick to backpropagate through Monte Carlo samples with a cross-entropy-like objective.
• Inference yields multiple samples of logit maps in a single forward pass; the mean logit provides the most likely segmentation.

Experimental results

Research questions

• RQ1 Can modeling spatially correlated aleatoric uncertainty via a low-rank logit-space distribution produce coherent, multiple segmentation hypotheses?
• RQ2 Does SSN improve predictive performance and uncertainty calibration relative to deterministic and existing probabilistic segmentation methods on medical imaging tasks?
• RQ3 How does the approach scale to high-dimensional 3D data and integrate with common segmentation backbones?

Key findings

• SSNs outperform state-of-the-art methods in modeling correlated uncertainty for ambiguous medical images.
• The low-rank covariance model yields higher predictive performance and better uncertainty calibration than a diagonal-covariance variant and deterministic baselines on lung nodule segmentation.
• SSNs enable sampling of multiple plausible segmentations in 2D and 3D with minimal architectural changes and modest computational overhead.
• In BraTS 2017 3D brain tumor segmentation, stochastic models achieve comparable Dice scores to deterministic models while providing a distribution over segmentations.
• The framework allows post-inference manipulation of samples (e.g., adjusting class presence) and temperature-based control of diversity.