Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks (1807.07356v3)

Abstract: Despite the state-of-the-art performance for medical image segmentation, deep convolutional neural networks (CNNs) have rarely provided uncertainty estimations regarding their segmentation outputs, e.g., model (epistemic) and image-based (aleatoric) uncertainties. In this work, we analyze these different types of uncertainties for CNN-based 2D and 3D medical image segmentation tasks. We additionally propose a test-time augmentation-based aleatoric uncertainty to analyze the effect of different transformations of the input image on the segmentation output. Test-time augmentation has been previously used to improve segmentation accuracy, yet not been formulated in a consistent mathematical framework. Hence, we also propose a theoretical formulation of test-time augmentation, where a distribution of the prediction is estimated by Monte Carlo simulation with prior distributions of parameters in an image acquisition model that involves image transformations and noise. We compare and combine our proposed aleatoric uncertainty with model uncertainty. Experiments with segmentation of fetal brains and brain tumors from 2D and 3D Magnetic Resonance Images (MRI) showed that 1) the test-time augmentation-based aleatoric uncertainty provides a better uncertainty estimation than calculating the test-time dropout-based model uncertainty alone and helps to reduce overconfident incorrect predictions, and 2) our test-time augmentation outperforms a single-prediction baseline and dropout-based multiple predictions.

• The paper introduces a novel framework that integrates test-time augmentation to quantify aleatoric uncertainty in medical image segmentation.
• It employs Monte Carlo sampling of input transformations to capture inherent image variability and mitigate overconfident predictions.
• Experiments on fetal brain and brain tumor MRI datasets demonstrate enhanced segmentation accuracy and reliability.

Uncertainty Estimation in Medical Image Segmentation with CNNs

The paper examines uncertainty estimation in convolutional neural network (CNN)-based medical image segmentation, focusing on aleatoric uncertainty enhanced by test-time augmentation. This paper presents a theoretical framework for estimating aleatoric uncertainty using input transformations, combined with efforts to contrast this with epistemic uncertainty derived from model parameters like test-time dropout.

Methodology and Theoretical Framework

The authors propose a novel framework wherein test-time augmentation is modeled mathematically as part of an image acquisition process. This process incorporates randomness through transformations such as scaling, rotation, and noise, which are probabilistically modeled using Monte Carlo simulations. By leveraging the variability in these augmented inputs, aleatoric uncertainty is quantified, potentially uncovering different aspects from traditional model-based uncertainty measurements.

Key considerations in their framework include:

• Image Acquisition Model: The authors depict the observed image as a transformation of a latent image, influenced by stochastic processes, providing a basis for modeling real-world variability.
• Probability Distribution and Monte Carlo Sampling: Sampling over the space of potential transformations allows the examination of prediction distributions, offering insights into uncertainties not purely tied to model weights.

Experimental Evaluation

Experiments were conducted on fetal brain and brain tumor MRI datasets, showcasing improvements in segmentation accuracy and uncertainty estimation when using the proposed test-time augmentation techniques. The aleatoric uncertainty helped highlight regions with higher segmentation errors, distinguishing itself as a critical tool for practical deployment scenarios. The findings suggest:

• Increased Accuracy: Incorporating test-time augmentation markedly improved the segmentation results compared to baseline methods and test-time dropout, providing empirical evidence for its efficacy.
• Reduced Overconfidence: Aleatoric uncertainty appeared to mitigate overconfident errors present in predictions, offering a robust mechanism for uncertainty quantification.

Implications and Future Directions

This paper firmly situates aleatoric uncertainty—particularly when expressed through test-time augmentation—at the forefront of improving segmentation reliability. For future developments, the integration of this methodology with larger datasets and more intricate transformations could further capitalize on its potential.

The framework could also be extended to incorporate more complex and biologically relevant transformations, enhancing the scope of its applicability. Additionally, combined approaches utilizing both epistemic and aleatoric uncertainty metrics may offer even richer predictive insights, particularly in clinical settings where precision is paramount.

Ultimately, this paper's contributions provide a robust mechanism for enhancing the reliability and interpretability of CNN-based medical image segmentation, advancing both theoretical understanding and practical implementation.