[Paper Review] Uncertainty Quantification in Computer-Aided Diagnosis: Make Your Model say "I don't know" for Ambiguous Cases

Max-Heinrich Laves, Sontje Ihler+1 more|arXiv (Cornell University)|Aug 2, 2019

Retinal Imaging and Analysis8 references4 citations

TL;DR

This paper proposes integrating uncertainty quantification into deep learning-based computer-aided diagnosis using Monte Carlo dropout (Bayesian ResNet) and variational inference (Variational ResNet) to enable models to express uncertainty. Both methods significantly increase uncertainty for incorrect predictions—mean uncertainty was 4.6 to 8.7 times higher for misclassified optical coherence tomography scans—demonstrating that uncertainty estimation enhances diagnostic reliability and patient safety.

ABSTRACT

We evaluate two different methods for the integration of prediction uncertainty into diagnostic image classifiers to increase patient safety in deep learning. In the first method, Monte Carlo sampling is applied with dropout at test time to get a posterior distribution of the class labels (Bayesian ResNet). The second method extends ResNet to a probabilistic approach by predicting the parameters of the posterior distribution and sampling the final result from it (Variational ResNet).The variance of the posterior is used as metric for uncertainty.Both methods are trained on a data set of optical coherence tomography scans showing four different retinal conditions. Our results shown that cases in which the classifier predicts incorrectly correlate with a higher uncertainty. Mean uncertainty of incorrectly diagnosed cases was between 4.6 and 8.1 times higher than mean uncertainty of correctly diagnosed cases. Modeling of the prediction uncertainty in computer-aided diagnosis with deep learning yields more reliable results and is anticipated to increase patient safety.

Motivation & Objective

• To improve patient safety in deep learning-based medical diagnosis by enabling models to express uncertainty.
• To address the lack of uncertainty estimation in standard deep learning models, which cannot say 'I don't know' for ambiguous cases.
• To evaluate and compare two uncertainty estimation methods—Monte Carlo dropout and variational inference—in a clinical diagnostic setting.
• To assess whether uncertainty correlates with prediction errors in retinal disease classification using optical coherence tomography.

Proposed method

• Applies Monte Carlo dropout at test time with dropout rate p=0.5, either before the final fully connected layer (bayesian1) or before every ResNet block (bayesian2), to approximate Bayesian inference.
• Uses variational inference to train a ResNet-18 to predict the mean μ and variance σ² of the output distribution, with reparameterization for differentiable sampling.
• Employs the reparameterization trick: ŷ = μ + σε, where ε ~ N(0, I), enabling end-to-end training with gradient-based optimization.
• Minimizes a loss function combining cross-entropy and Kullback-Leibler divergence (KLD) to regularize the posterior toward a standard normal distribution.
• Trains and evaluates three uncertainty-aware models (bayesian1, bayesian2, variational) and a standard ResNet-18 baseline on 84,484 OCT scans with four retinal conditions.
• Uses the variance of the posterior distribution as a measure of prediction uncertainty and compares uncertainty levels between correctly and incorrectly classified cases.

Experimental results

Research questions

• RQ1 Can Monte Carlo dropout with test-time dropout reliably estimate uncertainty in a diagnostic deep learning model for OCT scans?
• RQ2 Does variational inference with learnable posterior parameters improve uncertainty estimation compared to standard models?
• RQ3 Is there a significant correlation between high prediction uncertainty and incorrect classification in retinal disease diagnosis?
• RQ4 Do uncertainty estimates from Bayesian and variational approaches effectively flag ambiguous or misclassified cases?
• RQ5 Can uncertainty quantification enhance clinical safety by enabling models to 'say I don't know' for uncertain cases?

Key findings

• Mean uncertainty for incorrectly diagnosed cases was 8.7 times higher than for correctly diagnosed cases when using the bayesian1 method.
• Mean uncertainty for misclassified cases was 6.0 times higher with bayesian2 and 4.6 times higher with the variational approach.
• The bayesian1 method achieved the highest test set F1 score (0.96), while bayesian2 and variational methods showed slightly lower performance, likely due to increased noise from multiple dropout layers and reparameterization.
• All uncertainty estimation methods produced higher uncertainty for false predictions, confirming that uncertainty correlates strongly with prediction error.
• The integration of uncertainty quantification significantly improves model reliability and supports clinical decision-making by identifying ambiguous cases.
• The results demonstrate that uncertainty-aware models can enhance patient safety by enabling clinicians to review high-uncertainty cases.