Source-linked AI summary
Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty
Miguel Monteiro, Loïc Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Marques, Konstantinos Kamnitsas, Mark van der Wilk, Ben Glocker
TL;DR
Image segmentation often has multiple plausible label maps because ambiguous boundaries create inherent aleatoric uncertainty, while many methods produce only pixel-wise estimates. SSNs model a joint label-map distribution with a low-rank Gaussian in logit space, and experiments show stronger correlated-uncertainty modelling than baselines while remaining simpler and efficient.
Problem
Segmentation models often lack the ability to represent multiple spatially coherent hypotheses for ambiguous images with correlated aleatoric uncertainty.
Method
SSNs model the logit map with a low-rank multivariate normal distribution that captures label correlations and works with existing segmentation architectures.
Results
SSNs outperform baselines for modelling correlated uncertainty while maintaining efficient inference and no loss in performance relative to deterministic models.
Takeaways & Limitations
The method enables multiple plausible segmentation hypotheses and post-inference manipulation for applications such as human-in-the-loop radiology.
Takeaways & Limitations
The toy problem exhibited covariance overflow and suboptimal optimization, requiring mean pre-training and early stopping.
Abstract
from arXiv · showhide
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.
1 Introduction
Image segmentation must represent spatially consistent label maps and multiple plausible outputs, but common pixel-wise models do not explicitly capture this structured uncertainty. SSNs address this gap by modelling spatially correlated aleatoric uncertainty efficiently across existing segmentation architectures.
- Semantic segmentation requires spatially consistent label maps, yet dense pixel-wise classification largely ignores the underlying spatial structure.
- Ambiguous boundaries can produce multiple plausible segmentations, so models should capture uncertainty over alternative outputs.
- Aleatoric uncertainty is spatially correlated and heteroscedastic, with different image regions exhibiting different uncertainty levels.
- Pixel-wise uncertainty estimates yield spatially incoherent samples and therefore cannot fully represent structured uncertainty in label maps.
- SSNs model the label-map distribution through a low-rank multivariate normal over logits, supporting efficient, architecture-independent sampling of coherent hypotheses.
2 Related work
Prior methods model joint segmentation uncertainty but often require complex architectures, tuned objectives, costly inference, or uncorrelated latent variables. SSNs offer a simpler alternative that models correlations directly in logit space and supports efficient sampling.
- Bayesian, ensemble, and multi-head methods primarily target epistemic uncertainty by modelling uncertainty over neural-network weights.
- Graphical-model approaches explicitly represent label dependencies but commonly restrict inference to MAP or fixed numbers of solutions with high computational cost.
- Conditional variational approaches generate multiple segmentations using latent stochasticity, but their uncorrelated latent variables rely on decoders to create spatial variation.
3 Background
The standard segmentation cross-entropy follows conditional-independence assumptions for logits and labels. Because neighboring segmentation labels are correlated, these assumptions omit dependencies that should be modelled.
- A classic CNN maps an image to per-pixel logits and obtains categorical label probabilities through a softmax layer.
- The standard cross-entropy formulation assumes logits are deterministic functions of the image and model parameters.
- It also assumes labels are conditionally independent given their respective logits, so neighboring true labels provide no additional information.
- These independence assumptions are inappropriate for segmentation because labels at different pixels are correlated.
4 Stochastic segmentation networks
SSNs model correlated label uncertainty by placing a low-rank multivariate normal distribution over logits, then efficiently sampling spatially coherent predictions through a neural network.
- SSNs use a multivariate normal distribution over logits whose neural-network-produced parameters model dependencies between pixels and classes.The distribution uses mean µ(x) and covariance Σ(x), rather than assuming logits are independent.
- The covariance uses a low-rank factorization with a diagonal term, making its mean, covariance factor, and diagonal efficient to compute.The rank R controls the covariance factor size, while the diagonal term contains S × C elements.
- Monte-Carlo integration approximates the otherwise intractable softmax-normal integral, producing a differentiable loss via the re-parameterisation trick.The few degrees of freedom give the Monte-Carlo integral low variance.
- At inference, one forward pass supports repeated sampling of logit maps that can become probability or label maps, while the mean gives the most likely logit sample.
- The method captures label dependencies with minimal overhead: three output maps and low-rank sampling add O(rank) cost, leaving the underlying network dominant.This contrasts with CRF inference, which explicitly models dependencies but has an expensive inference procedure.
5 Experiments and Results
Experiments show that low-rank stochastic segmentation models capture correlated uncertainty and generate diverse, spatially coherent predictions while retaining predictive performance. Across toy, lung-nodule, and brain-tumour tasks, the method supports repeated sampling and post-inference manipulation.
- 5.1 Toy problem: The low-rank model learned the toy problem’s noise structure, matched the desired distribution, and achieved log-likelihood -0.93 versus -4.87 for the diagonal model.The diagonal model learned marginal uncertainty but produced uncorrelated noise.
- 5.2 Lung nodule segmentation in 2D: On LIDC-IDRI, the proposed low-rank model outperformed baselines in predictive performance for both one- and four-annotation settings and had the lowest D2It was also the only method whose predictive performance improved with additional annotations.
- 5.3 Brain tumour segmentation in 3D: On BraTS 2017, stochastic models matched the deterministic model’s performance, while additional spatial context did not improve performance or calibration.The amount of spatial context was application dependent.
- 5.3 Brain tumour segmentation in 3D: Brain-tumour samples varied coherently: entire structures appeared or disappeared, and uncertain lesion borders consistently shrank or expanded between samples.Some errors in the deterministic model were corrected in at least one stochastic sample.
- 5.3 Brain tumour segmentation in 3D: 26.0% of brain-tumour samples were better than the mean prediction, while the average best-sample 95% quantile reached 70.3% versus 66.8% for the deterministic model.The improvement was larger for lower-performing cases.
- 5.3 Brain tumour segmentation in 3D: After inference, scaling class-specific covariance components enabled samples to increase or reduce a class, while temperature scaling traded sample diversity for quality.The paper demonstrates this manipulation for correcting mistakes or adjusting borders.
6 Discussion
The paper presents an efficient method for modelling spatially correlated aleatoric uncertainty in segmentation and generating multiple plausible hypotheses. Its simplicity supports deployment across architectures and applications, including 3D medical imaging and human-in-the-loop radiology.
- 6 Discussion: SSNs model spatially correlated aleatoric uncertainty efficiently and outperform baselines while remaining simpler than alternative approaches.The method improves predictive performance with added uncertainty and can generate samples better than deterministic predictions.
- 6 Discussion: The method can be implemented over existing neural network architectures, enabling application to 3D segmentation.
- 6 Discussion: Multiple plausible hypotheses can support semi-automatic segmentation in human-in-the-loop radiology workflows.
- 6 Discussion: Spatially correlated uncertainty is also relevant to autonomous systems because uncertainty about an object should extend across its pixels.
Broader Impact
Reliable uncertainty estimates can support clinical decision making and indicate when automatically derived measurements should be ignored. They can also be propagated to downstream tasks such as radiotherapy planning.
- Broader Impact: Uncertainty estimates could inform clinical decisions and signal when clinicians should disregard automatically derived measurements.
- Broader Impact: Uncertainty estimates could be propagated to downstream clinical tasks such as radiotherapy planning.
A.1 Training procedure for the LIDC and BraTS datasets.
The LIDC and BraTS experiments used separate training procedures, optimizers, schedules, sampling strategies, and data augmentation settings.
- A.1 Training procedure for the LIDC and BraTS datasets.: LIDC models were trained for 500000 Adam iterations with learning rate 0.001, batch size 12, and flipping, rotation, and scaling augmentation.
- A.1 Training procedure for the LIDC and BraTS datasets.: BraTS models were trained for 1200 epochs using RMSProp with momentum 0.6, lesion-or-background patch sampling, and scheduled learning-rate halving.
A.2 Evaluation details.
Evaluation distances use one minus multiclass IoU, while DSC is reported with a convention that differs from PHiSeg for empty slices.
- A.2 Evaluation details.: The distance between label maps is d = 1 − IoU, with multiclass IoU averaged across non-background classes and empty pairs assigned distance zero.
- A.2 Evaluation details.: The reported DSC is lower than PHiSeg’s because this work does not assign DSC 1.0 when predicted and ground-truth slices are both empty.
A.3 Rank ablation study
The rank ablation study examines how the multivariate normal rank affects sampling complexity and performance on BraTS. Visual analyses compare rank-dependent metrics, sample diversity, and stochastic samples against independent categorical sampling.
- A.3 Rank ablation study: Rank controls how many independent pixel clusters are jointly controlled, thereby limiting the maximum possible sample complexity.The study evaluates this effect on BraTS using models trained on 110 mm image patches.
- A.3 Rank ablation study: The rank ablation measures sample diversity, generalised energy distance, and average class DSC across rank settings.Figure A1 reports results as means and standard errors over five random seeds.
- A.3 Rank ablation study: Sampling from independent categorical distributions produces grainy label noise, whereas the proposed model generates spatially structured samples.Figure A3 compares five random samples using the T1ce slice, ground truth, marginal entropy, and sample outputs.
- A.3 Rank ablation study: Additional figures show stochastic-model samples alongside T1ce slices, ground truth, deterministic and stochastic predictions, marginal entropy, and seven random samples.Figures A4–A7 provide examples for multiple test cases.
- A.3 Rank ablation study: Figure A2 visualizes how ranks from 1 to 20 affect sampled outputs, with rows representing ranks and columns representing different samples.Rank increases from top to bottom across [1, 2, 5, 10, 15, 20].