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

Mohammed Maaz Sibhai; Abedalrhman Alkhateeb; Saad B. Ahmed

eess.IVApr 10, 2026

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

Authors: Mohammed Maaz Sibhai, Abedalrhman Alkhateeb, Saad B. Ahmed

Organizations: [1] Department of Computer Science, Lakehead University

Abstract

To ensure safe clinical integration, deep learning models must provide more than just high accuracy; they require dependable uncertainty quantification. While current Medical Vision Transformers perform well, they frequently struggle with overconfident predictions and a lack of transparency, issues that are magnified by the noisy and imbalanced nature of clinical data. To address this, we enhanced the modified Medical Transformer (MedFormer) that incorporates prototype-based learning and uncertainty-guided routing, by utilizing a Dirichlet distribution for per-token evidential uncertainty, our framework can quantify and localize ambiguity in real-time. This uncertainty is not just an output but an active participant in the training process, filtering out unreliable feature updates. Furthermore, the use of class-specific prototypes ensures the embedding space remains structured, allowing for decisions based on visual similarity. Testing across four modalities (mammography, ultrasound, MRI, and histopathology) confirms that our approach significantly enhances model calibration, reducing expected calibration error (ECE) by up to 35%, and improves selective prediction, even when accuracy gains are modest.

Figures & tables

Dataset	Modality	Task Type	Classes
CBIS-DDSM	Mammography (X-ray)	Binary	2
BUSI	Ultrasound	Multi-class	3
Histopathology	Microscopy	Binary	2
Brain MRI	MRI	Multi-class	4

Table 1: List of datasets used for model performance evaluation

Figure 1: Baseline MedFormer architecture with DSSA attention.

Figure 2: MedFormer backbone with Uncertainty-Gated Token Refinement (UGTR).

Figure 3: Prototype-Based Classification Head.

Dataset	Model	Accuracy	Macro-F1	AUROC
BUSI	MedFormer	0.800	0.761	0.919
	Proposed	0.893	0.882	0.956
BreaKHis	MedFormer	0.9206	0.8650	0.9676
	Proposed	0.9303	0.8760	0.9755
Brain Tumor	MedFormer	0.9733	0.9721	0.9965
	Proposed	0.9809	0.9801	0.9982

Table 2: Classification performance across datasets(

\beta

= 0 for proposed model)

Figure 4: Training and validation loss curves for the BUSI ultrasound dataset.

Figure 5: Training and validation loss curves for the Brain Tumor MRI dataset.

Figure 6: Validation performance on the Brain Tumor MRI dataset across training epochs.

Figure 7: Training and validation loss curves for the CBIS-DDSM mammography dataset (single-view UG2RLPR).

Figure 8: Training and validation AUC curves for the BreakHis histopathology dataset.

Model setting	AUROC	Accuracy	F1-score
Regular MedFormer (epoch 20)	0.727	0.619	0.644
Proposed ( $\beta$ = 0.0)	0.727	0.688	0.642
Proposed ( $\beta$ = 0.3)	0.730	0.675	0.639
Proposed ( $\beta$ = 0.35)	0.729	0.636	0.642
Proposed ( $\beta$ = 0.4)	0.752	0.695	0.632
Proposed ( $\beta$ = 0.45)	0.710	0.648	0.642

Table 3: CBIS-DDSM classification performance versus

\beta

(MC Test, T = 20 stochastic forward passes)

$\beta$	ECE	Brier	NLL	MCE
Regular MedFormer	0.1836	0.2434	0.8115	0.3505
Proposed ( $\beta$ = 0.0)	0.1174	0.2156	0.6579	0.1976
Proposed ( $\beta$ = 0.3)	0.1562	0.2354	0.7177	0.2851
Proposed ( $\beta$ = 0.35)	0.1182	0.2183	0.6530	0.2456
Proposed ( $\beta$ = 0.4)	0.1510	0.2317	0.6749	0.3098
Proposed ( $\beta$ = 0.45)	0.2333	0.2707	0.8589	0.3098

Table 4: Calibration metrics on CBIS-DDSM (MC Test, T = 20 stochastic forward passes)

$\beta$	AURC	ACC @ 50%	ACC @ 70%	ACC @ 90%
Regular MedFormer	0.2134	0.793	0.722	0.680
Proposed ( $\beta$ = 0)	0.2179	0.801	0.744	0.703
Proposed ( $\beta$ = 0.3)	0.2059	0.767	0.724	0.694
Proposed ( $\beta$ = 0.35)	0.2038	0.767	0.732	0.688
Proposed ( $\beta$ = 0.4)	0.2251	0.773	0.714	0.659
Proposed ( $\beta$ = 0.45)	0.2517	0.739	0.712	0.669

Table 5: Selective prediction results on CBIS-DDSM (MC Test, T = 20 stochastic forward passes)

Panel of Figure 9 — Figure 9: Uncertainty maps of two mammogram

Panel of Figure 10 — Figure 10: Prototype similarity visualization of three different mammograms. The query mammogram is compared to learned prototypes, demonstrating a high cosine similarity with typical class patterns.

Explore similar work

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
Uncertainty of Vision Medical Foundation Models
Aug 31, 2026Haoxu Huang, Narges RazavianMedical Imaging Foundation Models Uncertainty Quantification
Center for Data Science, New York University · Department of Radiology, NYU Grossman School of Medicine · Department of Population Health, NYU Grossman School of Medicine
Delving Aleatoric Uncertainty in Medical Image Segmentation via Vision Foundation Models
Apr 13, 2026Ruiyang Li, Fang Liu, Licheng Jiao +8Medical Image Segmentation Segmentation Uncertainty Estimation
Key Laboratory of Intelligent Perception and Image Understanding of Ministry of Education, International Research Center for Intelligent Perception and Computation, Joint International Research Laboratory of Intelligent Perception and Computation, School of Artificial Intelligence, Xidian University, Xi’an 710071, China

Feb 12, 2026J/K move · Enter open · S save

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

Hua Xu, Julián D. Arias-Londoño, Juan I. Godino-Llorente

Escuela Técnica Superior de Ingenieros de Telecomunicación Universidad Politécnica de Madrid 28040, Madrid, Spain

Bayesian Neural Networks Post-Hoc Calibration Medical Image Analysis Medical Image Classification Uncertainty Calibration

Open paper AskPDF arXiv

In critical decision support systems based on medical imaging, the reliability of AI-assisted decision-making is as relevant as predictive accuracy. Although deep learning models have demonstrated significant accuracy, they frequently suffer from miscalibration, manifested as overconfidence in erroneous predictions. To facilitate clinical acceptance, it is imperative that models quantify uncertainty in a manner that correlates with prediction correctness, allowing clinicians to identify unreliable outputs for further review. To address this necessity, this paper proposes a probabilistic optimization framework grounded in Bayesian deep learning. Specifically, the Confidence-Uncertainty Boundary Curve (CUBC) is first explored as an intermediate operational target. Grounded in this target, a novel Confidence-Uncertainty Boundary Loss (CUB-Loss) is proposed to regularize the alignment between prediction confidence and uncertainty estimates during training, imposing penalties on high-certainty errors and low-certainty correct predictions. Upon completion of training optimization, a Boundary Curve Calibration Error (BCCE) metric is further introduced to measure the degree of boundary alignment in the calibrated model. Building on this measurement, a Dual Temperature Scaling (DTS) strategy is devised to perform post-hoc refinement, further adjusting the posterior predictive distribution across different confidence-uncertainty regions. The proposed framework is validated on three distinct medical imaging tasks: automatic screening of pneumonia, diabetic retinopathy detection, and identification of skin lesions. Empirical results demonstrate that the proposed approach improves uncertainty calibration across diverse modalities, maintains robust performance in data-scarce scenarios, and remains effective on severely imbalanced datasets, underscoring its potential for real clinical deployment.

Abstract

Explore similar work

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

Uncertainty of Vision Medical Foundation Models

Delving Aleatoric Uncertainty in Medical Image Segmentation via Vision Foundation Models