Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks

Guotai Wang; Wenqi Li; Michael Aertsen; Jan Deprest; Sebastien Ourselin; Tom Vercauteren

Source-linked AI summary

Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks

Guotai Wang, Wenqi Li, Michael Aertsen, Jan Deprest, Sebastien Ourselin, Tom Vercauteren

arXiv:1807.07356v3Jul 19, 2018cs.CV

TL;DR

Deep CNNs perform strongly on medical image segmentation but rarely quantify uncertainty in pixel-wise outputs. The paper models test-time augmentation through Monte Carlo sampling of image-acquisition parameters, combining transformation- and noise-related aleatoric uncertainty with model uncertainty. Across 2D and 3D MRI segmentation, the approach improves uncertainty estimation, reduces overconfident incorrect predictions, and outperforms single-prediction and dropout-based baselines.

Problem
Deep CNNs rarely provide model or image-based uncertainty estimates for medical image segmentation outputs.
Method
The paper formulates test-time augmentation as Monte Carlo inference over image-acquisition parameters involving spatial transformations and noise, then combines aleatoric and epistemic uncertainty.
Results
Across 2D and 3D segmentation, test-time augmentation improves uncertainty estimation, reduces overconfident incorrect predictions, and outperforms single-prediction and dropout-based multiple-prediction baselines.
Takeaways & Limitations
Test-time augmentation provides a general framework for accounting for image noise and spatial transformations when estimating uncertainty in medical image segmentation.
Takeaways & Limitations
The optimal Monte Carlo sample count varies by dataset and application, so it should be selected using a validation-set performance plateau.

Abstract

from arXiv · show

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.

1. Introduction

Medical image segmentation remains difficult and CNNs generally do not provide uncertainty estimates, especially for pixel-wise predictions. This work formulates test-time augmentation to estimate aleatoric uncertainty from image noise and spatial transformations alongside model uncertainty.

Motivation: Medical image segmentation is challenging because image quality, pathology, imaging protocols, and patient variability affect target appearances.Uncertainty can identify potentially mis-segmented boundary regions and guide interactive refinement.
Motivation: Deep CNNs achieve state-of-the-art segmentation performance but do not by default estimate uncertainty, while large pixel-wise annotated datasets are difficult to collect.
Research gap: Prior uncertainty studies focused mainly on classification or regression, whereas segmentation requires pixel-wise uncertainty for interpreting and refining outputs.
Research scope: The paper investigates both epistemic model uncertainty and aleatoric input uncertainty for CNN-based medical image segmentation.Epistemic uncertainty can diminish with sufficient training data, whereas aleatoric uncertainty depends on noise or randomness in the test image.
Approach: Test-time augmentation estimates prediction distributions under transformations such as rotation, scaling, and flipping to capture transformation-related aleatoric uncertainty.
Contributions: The work proposes a theoretical test-time augmentation formulation and validates general aleatoric uncertainty estimation across 2D and 3D segmentation tasks.

2. Related Works

Related work spans uncertainty estimation for traditional medical segmentation and deep CNNs, including Bayesian approximations and test-time data augmentation. These approaches motivate combining model uncertainty with input-transformation uncertainty for segmentation.

Medical segmentation uncertainty: Traditional medical image segmentation studies estimated uncertainty using shape and appearance priors, graph cuts, lumen models, and uncertainty-guided interactions.
Deep CNN uncertainty: Deep CNN research has examined both epistemic and aleatoric uncertainty, with test-time dropout providing a practical Bayesian approximation for model uncertainty.Exact Bayesian networks are mathematically grounded but difficult to implement and computationally expensive.
Data augmentation: Data augmentation transforms training images through operations such as flipping, cropping, rotation, scaling, and elastic deformation to enlarge datasets.

3. Method

The method models test-time augmentation as inference over hidden image-acquisition variables, using Monte Carlo prediction distributions to estimate aleatoric uncertainty and combine it with epistemic uncertainty.

Image Acquisition Model: Test-time augmentation models observed images through spatial transformations and noise, with priors over transformation parameters and noise.The framework focuses on reversible spatial transformations such as flipping, scaling, rotation, and translation, while allowing extensions to other augmentation types.
Inference with Hidden Variables: For segmentation, transformed label maps are equivariant, so predictions are mapped back before estimating the distribution over outputs.The paper distinguishes invariant labels for classification from spatially transformed label maps for segmentation.
Inference with Hidden Variables: Monte Carlo simulation samples transformation and noise parameters, generates predictions, and averages or majority-votes them for regression or segmentation outputs.For segmentation, the most frequent sampled prediction is used as the final output.
Aleatoric Uncertainty Estimation: Aleatoric uncertainty is estimated from prediction diversity using entropy, including pixel-wise entropy for segmentation.The entropy approximation uses frequencies of unique sampled predictions or pixel labels.
Epistemic Uncertainty Estimation: Epistemic uncertainty is estimated with test-time dropout by sampling network parameters and computing entropy across the resulting predictions.The method interprets test-time dropout as an ensemble-like approach and compares it with augmentation-based sampling.
Structure-wise Uncertainty Estimation: Structure-wise uncertainty extends Monte Carlo sampling to segmented volumes and uses volume variation coefficient alongside test-time dropout and augmentation.The volume variation coefficient is agnostic to the segmented structure’s size.

4. Experiments and Results

The method was validated on fetal-brain and brain-tumor segmentation using 2D and 3D MRI, comparing augmentation-based, dropout-based, and hybrid uncertainties.

Datasets and Tasks: The experiments evaluated 2D fetal-brain MRI-slice segmentation and 3D brain-tumor segmentation from multi-modal MRI volumes.These two tasks were used to validate the proposed testing and uncertainty-estimation method.
Uncertainty Comparisons: Three uncertainty strategies were compared: test-time augmentation, test-time dropout, and their hybrid combination.Each strategy obtained uncertainty from N predictions using the same entropy-based formulation.
Evaluation Metrics: Dice was used as a segmentation metric based on true positives, false positives, and false negatives.The reported definition is Dice = 2 × TP / (2 × TP + FN + FP).
Evaluation Metrics: ASSD was defined using segmentation and ground-truth surface points and shortest Euclidean distances between them.S denotes predicted segmentation surface points and G denotes ground-truth surface points.

4.1. 2D Fetal Brain Segmentation from MRI

The fetal-brain MRI experiments evaluate test-time dropout (TTD), test-time augmentation (TTA), and their combination for segmentation accuracy and uncertainty estimation. Across visual, quantitative, and error-correlation analyses, TTA better identifies problematic regions and improves segmentation more consistently than TTD.

Data and setup: 60 fetal MRI scans provided 2D axial, sagittal, and coronal stacks for training, validation, and testing.The scans were acquired with SSFSE at 1.5 Tesla, with 3–6 mm slice thickness.
Data and setup: TTA models aleatoric uncertainty from sampled flips, rotations, scaling, and intensity noise, using the same parameter priors during training and testing.The transformation parameters include Bernoulli flipping, uniform rotation and scaling, and Gaussian intensity noise.
Segmentation results with uncertainty: In over- and under-segmented cases, TTD assigned low epistemic uncertainty to mis-segmented regions, whereas TTA marked larger regions corresponding to baseline errors.The hybrid TTA + TTD uncertainty showed behavior similar to TTA in these examples.
Segmentation results with uncertainty: TTA corrected large mis-segmentations and improved the baseline more noticeably than TTD, while combining TTA and TTD produced results similar to TTA.The comparison used Monte Carlo simulation with N = 20 for TTD, TTA, and TTA + TTD.
Quantitative evaluation: With or without training augmentation, TTA improved segmentation accuracy more than TTD, while adding TTD did not significantly outperform TTA (p-value > 0.05).The evaluation compared a single-prediction baseline, TTD, TTA, and TTA + TTD using Dice score and ASSD.
Correlation between uncertainty and segmentation error: TTA produced higher-average Dice distributions than TTD and showed stronger uncertainty–error relationships at both pixel and structure levels.TTA had fewer overconfident incorrect predictions than TTD, and its VVC increased more strongly with 1−Dice.

4.2. 3D Brain Tumor Segmentation from Multi-Modal MRI

The study evaluates TTA-based aleatoric, TTD-based epistemic, and hybrid uncertainty for 3D brain tumor segmentation from multi-modal MRI. Across qualitative and quantitative analyses, TTA better identifies segmentation errors and improves accuracy over single prediction and TTD.

Dataset and models: The BraTS 2017 dataset comprised 285 volumetric studies with four coregistered MRI sequences per study: T1w, T1wce, T2w, and FLAIR.Images were skull-stripped and resampled to isotropic 1 mm3 resolution; 20 studies were used for validation and 50 for testing.
Dataset and models: Experiments used 3D U-Net, V-Net, and W-Net to compare single prediction, TTD, TTA, and TTA + TTD.W-Net is a 2.5D network; the testing methods estimate epistemic, aleatoric, and hybrid uncertainties, respectively.
Qualitative uncertainty analysis: TTA better highlighted over- and under-segmented tumor regions than TTD in qualitative examples, while hybrid uncertainty was often similar to TTA.For a high-grade glioma over-segmentation, TTA highlighted the whole erroneous region; for a low-grade glioma under-segmentation, TTA better indicated the error.
Quantitative evaluation: Multi-prediction testing outperformed single prediction, and TTA achieved higher Dice and lower ASSD than TTD across the evaluated networks and training settings.Performance plateaued when N exceeded 40; results in Table 2 used N = 40. Combining TTA and TTD produced only a slight improvement over TTA, reported as not significant (p-value < 0.05).
Quantitative evaluation: TTA-based structure-wise uncertainty better indicated segmentation error than TTD-based uncertainty because its VVC increased more steeply with 1−Dice.The TTA and TTA + TTD distributions were similar in the 3D analysis.

5. Discussion and Conclusion

The study formulates test-time augmentation (TTA) as aleatoric uncertainty within an image acquisition model and evaluates it alongside model uncertainty for medical image segmentation. TTA reduces overconfident errors and improves segmentation accuracy, while combined TTA and test-time dropout does not consistently outperform TTA alone.

5. Discussion and Conclusion: TTA explicitly models spatial transformations and image noise as sources of aleatoric uncertainty in segmentation.The formulation can be extended to elastic deformations and simulated MRI bias fields.
5. Discussion and Conclusion: Model uncertainty alone can produce overconfident incorrect predictions, whereas TTA helps reduce them.This pattern is illustrated in the paper’s visual examples and conclusion.
5. Discussion and Conclusion: TTA plus test-time dropout does not always outperform TTA, although its average Dice is slightly larger across the overall testing images.The authors report that combined performance is close to TTA overall and can be lower for some test images.
5. Discussion and Conclusion: TTA is presented as applicable beyond segmentation to image classification, object detection, and regression.For regression, variation in the output distribution may be more suitable than entropy for uncertainty estimation.
5. Discussion and Conclusion: The suitable Monte Carlo sample count was around 20 to 40 for the tested CNNs and applications, but the optimal value may vary by dataset.The authors recommend selecting N using the performance plateau on a validation set; fixing N = 40 can waste computation when fewer samples suffice.
5. Discussion and Conclusion: TTA achieves higher segmentation accuracy than a single-prediction baseline and multiple predictions using test-time dropout.The conclusion reports this across 2D and 3D medical image segmentation tasks.