Uncertainty-driven training for three-dimensional calibrated lung nodule classification

Giuseppe Tripodi; Alessandro De Rosis; Saleh Rezaeiravesh

eess.IVSep 17, 2026

Uncertainty-driven training for three-dimensional calibrated lung nodule classification

Authors: Giuseppe Tripodi, Alessandro De Rosis, Saleh Rezaeiravesh

Organizations: School of Engineering, The University of Manchester, Manchester (UK)

Abstract

In this work, we present an uncertainty-driven training framework for three-dimensional computed tomography (CT) lung nodule classification, where validation-based uncertainty estimates guide loss reweighting to enhance predictive performance and probability calibration. Two Uncertainty Quantification (UQ) methods are considered: Monte Carlo Dropout (MCD) and Evidential Deep Learning (EDL). Both provide per-class uncertainty estimates that modulate the loss and encourage focus on hard or unreliable classes. The framework is evaluated with ResNet, DenseNet, EfficientNet, Vision Transformer (ViT), and Swin Transformer backbones on two datasets: the clinical LIDC-IDRI cohort and the NoduleMNIST3D benchmark. Uncertainty-driven training achieves classification performance similar to conventional training while substantially improving calibration, with an expected calibration error (ECE) reduced by up to 65% on LIDC-IDRI. EDL attains competitive performance on shallower architectures with single-pass inference, whereas MCD is more robust on deeper networks. Analysis across architectural families reveals that uncertainty-driven training benefits convolutional backbones more consistently than transformer-based architectures: EDL in particular degrades on ViT, suggesting that the Dirichlet evidence parameterisation may interact unfavourably with attention-based architectures at lower input resolutions. A posteriori temperature scaling proves highly effective across all configurations, indicating that a simple scalar calibration can be competitive even without explicit uncertainty-aware training. Our results indicate that integrating UQ into the training loop can significantly improve probabilistic calibration and support more trustworthy deployment of three-dimensional medical imaging models.

Figures & tables

Figure 1: Overview of the proposed uncertainty-driven training framework for lung nodule classification. Starting from preprocessed 3D CT nodule patches, multiple 3D backbones are trained using the MCD and EDL methods. Epistemic and aleatoric uncertainties estimated on the validation set drive class reweighting and a weighted loss, while a posteriori temperature scaling applied to validation logits yields calibrated test-time probabilities for malignancy prediction.

	DenseNet121			DenseNet169			ResNet18			ResNet50			EfficientNet3D			SwinViT3D			ViT3D
Method	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS
FL	0.06	0.32	0.15	0.06	0.69	0.21	0.09	0.42	0.22	0.08	0.49	0.25	0.06	0.29	0.21	0.04	0.16	0.32	0.04	0.10	0.38
EDL	0.06	0.47	0.23	0.03	0.10	0.18	0.08	0.75	0.24	0.07	0.67	0.24	0.20	0.26	0.35	0.14	0.24	0.41	0.08	0.50	0.38
MCD	0.05	0.16	0.20	0.09	0.27	0.22	0.04	0.47	0.19	0.16	0.43	0.35	0.03	0.07	0.33	0.07	0.15	0.32	0.10	0.10	0.50
CA	0.11	0.53	0.21	0.08	0.15	0.17	0.07	0.41	0.19	0.14	0.85	0.27	0.19	0.88	0.37	0.08	0.39	0.34	0.08	0.15	0.38
	DenseNet121			DenseNet169			ResNet18			ResNet50			EfficientNet3D			SwinViT3D			ViT3D

Table 1: Calibration and classification performance on the LIDC-IDRI dataset after a posteriori temperature scaling. Architectures are columns, and training regimes are rows. Bold: best per architecture. For ECE, MCE, and BS, lower is better (

\downarrow

); for AUC, F1, and Accuracy, higher is better (

\uparrow

Figure 2: (a) Post-calibration classification accuracy vs. uncertainty quality (AUROC-ED) for each training mode across the two datasets. Each point represents the mean over all evaluated architectures. The dashed line marks AUROC-ED = 0.70, a reference threshold for practically useful error detection; points in the upper-right quadrant achieve both high accuracy and reliable uncertainty estimates. (b) Risk–coverage curves for selective prediction across the two datasets. Risk (error rate) is measured on the fraction of predictions retained after abstaining on the most uncertain samples (coverage). Solid lines show the mean across architectures; shaded bands show

\pm 1

standard deviation. Lower curves indicate better uncertainty-guided abstention, where the model concentrates errors in the samples it chooses to defer.

	DenseNet121			DenseNet169			ResNet18			ResNet50			EfficientNet3D			SwinViT3D			ViT3D
Method	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$	ECE ${}_{\text{b}}$	ECE ${}_{\text{a}}$	$\Delta$
Baseline	0.087	0.061	+29	0.118	0.056	+53	0.128	0.087	+32	0.160	0.083	+48	0.099	0.058	+42	0.072	0.039	+46	0.061	0.044	+28
EDL	0.089	0.060	+32	0.046	0.027	+40	0.130	0.075	+42	0.065	0.065	0	0.256	0.201	+22	0.112	0.143	-28	0.096	0.080	+16
MCD	0.119	0.052	+56	0.107	0.093	+14	0.106	0.043	+59	0.203	0.163	+20	0.039	0.029	+25	0.128	0.072	+44	0.098	0.098	0
CA	0.109	0.109	-1	0.091	0.081	+11	0.119	0.074	+37	0.155	0.140	+10	0.184	0.186	-1	0.070	0.085	-22	0.103	0.076	+26

Table 2: ECE before and after temperature scaling on LIDC-IDRI.

\Delta

denotes relative improvement (%) after scaling.

	DenseNet121			DenseNet169			ResNet18			ResNet50			EfficientNet3D			SwinViT3D			ViT3D
Method	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS	ECE	MCE	BS
Baseline	0.07	0.16	0.23	0.06	0.14	0.25	0.05	0.08	0.25	0.05	0.11	0.23	0.05	0.30	0.21	0.04	0.12	0.22	0.05	0.15	0.25
EDL	0.10	0.31	0.27	0.04	0.24	0.25	0.06	0.15	0.21	0.10	0.54	0.24	0.07	0.27	0.25	0.06	0.17	0.22	0.08	0.08	0.34
MCD	0.08	0.15	0.24	0.06	0.17	0.24	0.03	0.16	0.22	0.09	0.22	0.26	0.03	0.08	0.22	0.05	0.18	0.24	0.04	0.08	0.25
	DenseNet121			DenseNet169			ResNet18			ResNet50			EfficientNet3D			SwinViT3D			ViT3D
Method	AUC	F1	Acc	AUC	F1	Acc	AUC	F1	Acc	AUC	F1	Acc	AUC	F1	Acc	AUC	F1	Acc	AUC	F1	Acc

Table 3: Calibration and classification performance on the NoduleMNIST3D dataset after a posteiori temperature scaling. Architectures are columns and training regimes are rows. Bold: best per architecture. For ECE, MCE, and BS, lower is better (

\downarrow

); for AUC, F1, and Accuracy, higher is better (

\uparrow

Explore similar work

Beyond Boundary Noise: Aggregated Aleatoric Uncertainty Fails to Capture Presence Ambiguity in 3D Lung Nodule Segmentation
Aug 14, 2026Simon Baur, Arne Schernich, Ekin Böke +2Segmentation Uncertainty Estimation Medical Image Segmentation
Fraunhofer Heinrich-Hertz-Institut, 10587 Berlin, Germany · Technische Universität Berlin, 10623 Berlin, Germany · The Berlin Institute for the Foundations of Learning and Data (BIFOLD), 10587 Berlin, Germany
Confidence-Uncertainty Boundary Calibration for Bayesian Deep Learning in Medical Image Analysis
Feb 12, 2026Hua Xu, Julián D. Arias-Londoño, Juan I. Godino-LlorenteBayesian Neural Networks Post-Hoc Calibration
Escuela Técnica Superior de Ingenieros de Telecomunicación Universidad Politécnica de Madrid 28040, Madrid, Spain
MedFormer-UR: Uncertainty-Routed Transformer for Medical Image Classification
Apr 10, 2026Mohammed Maaz Sibhai, Abedalrhman Alkhateeb, Saad B. AhmedMedical Image Classification Vision Transformer
[1] Department of Computer Science, Lakehead University

Aug 14, 2026J/K move · Enter open · S save

Beyond Boundary Noise: Aggregated Aleatoric Uncertainty Fails to Capture Presence Ambiguity in 3D Lung Nodule Segmentation

Simon Baur, Arne Schernich, Ekin Böke, Wojciech Samek+1

Fraunhofer Heinrich-Hertz-Institut, 10587 Berlin, Germany · Technische Universität Berlin, 10623 Berlin, Germany · The Berlin Institute for the Foundations of Learning and Data (BIFOLD), 10587 Berlin, Germany

Segmentation Uncertainty Estimation Medical Image Segmentation 3D Medical Image Segmentation

Open paper AskPDF arXiv

Uncertainty estimation is critical for the safe clinical deployment of deep learning in medical image segmentation, with aleatoric uncertainty theoretically designed to capture irreducible data ambiguity. However, whether entropy-based measures reflect clinically meaningful ambiguity, i.e. case-level disagreement about whether a pathology is present at all, remains poorly understood. Contrary to most prior work, which focused on pixel-wise boundary disagreement, we systematically evaluate how well aleatoric uncertainty captures presence ambiguity. Our evaluation spans 3D lung nodule segmentation across four architectures with Monte Carlo dropout and deep ensembles, on LIDC-IDRI and an external validation cohort (LNDb). We find that entropy-based uncertainty maps align with boundary noise and minor drawing variation but carry insufficient discriminative signal for presence ambiguity. In contrast, a lightweight supervised ambiguity head trained on frozen segmentation features substantially outperforms all entropy-aggregation-based baselines across architectures, metrics, and both cohorts, and matches or exceeds methods that explicitly model ambiguity under disagreement supervision (Probabilistic U-Net, Annotator-Confusion 3D-UNet). A qualitative feature-space analysis shows that presence ambiguity is already encoded in the frozen encoder features of pixel-wise-trained networks, only to be discarded by the segmentation output and its entropy aggregation. Our findings expose a fundamental mismatch between the theoretical promise of aleatoric uncertainty and its practical behavior, and suggest that practitioners should not rely on entropy-based uncertainty as a proxy for clinical ambiguity in safety-critical applications.

Abstract

Explore similar work

Beyond Boundary Noise: Aggregated Aleatoric Uncertainty Fails to Capture Presence Ambiguity in 3D Lung Nodule Segmentation

Confidence-Uncertainty Boundary Calibration for Bayesian Deep Learning in Medical Image Analysis

MedFormer-UR: Uncertainty-Routed Transformer for Medical Image Classification