Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty

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.

arXiv commentPublished at Neurips2020. 17 pages, 11 figures, 2 tables

References in corpus (14)

• Auto-Encoding Variational Bayes
• Efficient Multi-Scale 3D CNN with Fully Connected CRF for Accurate Brain Lesion Segmentation
• Efficient Inference in Fully Connected CRFs with Gaussian Edge Potentials
• What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?
• Conditional Random Fields as Recurrent Neural Networks
• On Calibration of Modern Neural Networks
• Weight Uncertainty in Neural Networks
• Toward Multimodal Image-to-Image Translation
• Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks
• Evidential Deep Learning to Quantify Classification Uncertainty
• Why M Heads are Better than One: Training a Diverse Ensemble of Deep Networks
• Beyond temperature scaling: Obtaining well-calibrated multiclass probabilities with Dirichlet calibration
• Learning in an Uncertain World: Representing Ambiguity Through Multiple Hypotheses
• A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities

Cited by in corpus (16)

• A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges
• QU-BraTS: MICCAI BraTS 2020 Challenge on Quantifying Uncertainty in Brain Tumor Segmentation - Analysis of Ranking Scores and Benchmarking Results
• Compressed Sensing MRI Reconstruction Regularized by VAEs with Structured Image Covariance
• Spatially Varying Label Smoothing: Capturing Uncertainty from Expert Annotations
• Effect of latent space distribution on the segmentation of images with multiple annotations
• Propagation and Attribution of Uncertainty in Medical Imaging Pipelines
• Investigating and Improving Latent Density Segmentation Models for Aleatoric Uncertainty Quantification in Medical Imaging
• Improving Uncertainty-Error Correspondence in Deep Bayesian Medical Image Segmentation
• Application-driven Validation of Posteriors in Inverse Problems
• Medical Matting: A New Perspective on Medical Segmentation with Uncertainty
• Uncertainty quantification for White Matter Hyperintensity segmentation detects silent failures and improves automated Fazekas quantification
• That Label's Got Style: Handling Label Style Bias for Uncertain Image Segmentation
• Uncertainty Propagation in Node Classification
• Distributional Gaussian Process Layers for Outlier Detection in Image Segmentation
• Deep Classifiers with Label Noise Modeling and Distance Awareness
• Confidence Contours: Uncertainty-Aware Annotation for Medical Semantic Segmentation