Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty (2006.06015v2)

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.

• The paper demonstrates a novel probabilistic framework that models joint label distributions to capture spatially correlated aleatoric uncertainty.
• It integrates a low-rank multivariate normal approximation into existing architectures like U-Nets, improving segmentation coherence and ambiguity resolution.
• Empirical results on medical datasets reveal enhanced Dice scores and superior uncertainty quantification compared to baseline models.

Analysis of Stochastic Segmentation Networks for Spatially Correlated Aleatoric Uncertainty

The paper introduces Stochastic Segmentation Networks (SSNs), a novel approach to modeling spatially correlated aleatoric uncertainty in image segmentation tasks. This method aims to address the pervasive issue of uncertainty in predictions when segmenting medical images, such as CT and MRI scans, which often involve ambiguous pixel-level labeling.

Core Contributions

Stochastic Segmentation Networks articulate a probabilistic framework for capturing spatial correlations in aleatoric uncertainty. This is achieved by moving beyond independent pixel-wise predictions typical of conventional Fully Convolutional Neural Networks (FCNNs). Instead, SSNs employ a multivariate normal distribution over label maps, bringing two key innovations to image segmentation tasks:

1. Modeling Joint Label Distributions: By leveraging a low-rank approximation of a multivariate normal distribution in the logit space, SSNs efficiently capture spatial dependencies and generate multiple plausible hypotheses for a given image.
2. Integration into Existing Architectures: This probabilistic framework is designed to be architecture-agnostic, allowing seamless integration with existing segmentation models like U-Nets, without requiring modifications to the network itself.

Methodological Insights

The SSN framework deviates from traditional settings where labels are assumed independent given the logits. By parameterizing the logits with a mean and a low-rank covariance factor, SSNs encapsulate the spatial relationship across various regions of an image, adequately addressing heteroscedasticity and coherence in predictions. The paper formulates its approach in a manner allowing for efficient Monte Carlo approximations to address the intractability associated with the integral over the logit distributions.

Results and Implications

Empirical validation on real-world medical datasets demonstrates that SSNs outperform existing baselines in both predictive performance and handling of label ambiguity:

• LIDC-IDRI Dataset: SSNs surpassed probabilistic U-Net and PHiSeg models in Dice Similarity Coefficient (DSC) and generalized energy distance metrics, highlighting superior ability to model correlated uncertainty.
• BraTS Dataset: In 3D MRI brain tumor segmentation, SSNs maintained performance parity with deterministic approaches while offering enhanced uncertainty quantification, essential for scenarios requiring multiple plausible outcomes.

These results underscore the practical value of SSNs in clinical settings, where decisions are sensitive to uncertainties inherent in medical images. The approach facilitates better-informed delineation between tissue types, crucial for treatment planning and diagnostic accuracy.

Future Directions

The introduction of SSNs opens multiple avenues for further exploration:

1. Extending to Broader Applications: While primarily validated within medical imaging, the SSN framework is applicable to any segmentation task where spatial correlations and uncertainty play a significant role, such as in autonomous driving or satellite imagery analysis.
2. Complexity and Scalability: Future work might focus on optimizing the computational efficiency and extending the method to handle even larger scale 3D data with complex spatial hierarchies.
3. Integration with Epistemic Uncertainty: Integrating SSNs with Bayesian approaches could offer a comprehensive view of uncertainty by simultaneously addressing both model and data-driven ambiguities.

In summary, the proposed stochastic segmentation networks mark a significant step toward robust, uncertainty-aware AI models for critical segmentation tasks, particularly within challenging domains like medical imaging where precision is paramount.