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

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.

• The paper introduces a unified transformer model that integrates Dirichlet-based uncertainty estimation, uncertainty-gated feature refinement, and prototype-driven classification to improve calibration and interpretability.
• It demonstrates significant performance gains in accuracy and AUROC across diverse medical imaging modalities, with notable improvements in uncertainty management and selective prediction.
• The methodology enables risk-aware diagnostic support by coupling uncertainty routing with similarity-based decisions, paving the way for robust clinical AI deployments.

MedFormer-UR: A Unified Transformer Framework for Reliable and Interpretable Medical Image Classification

Introduction

MedFormer-UR introduces a formalized extension of the MedFormer hierarchical vision transformer for medical image classification by integrating uncertainty quantification and prototype-based reasoning into its operational pipeline. The motivating clinical requirement—robust deployment in high-stakes diagnostic settings—necessitates models that are not merely accurate but are also reliably calibrated, uncertainty-aware, and interpretable. This work evolves the MedFormer backbone by (1) embedding Dirichlet-based evidential uncertainty modeling at the token level, (2) dynamically gating feature propagation and refinement based on local and global uncertainty, and (3) structuring the decision process via class-specific prototypes that organize the embedding space for similarity-based classification.

Figure 1: Baseline MedFormer architecture with DSSA attention.

Methodological Contributions

Uncertainty Estimation via Dirichlet-Evidential Modeling

The architecture implements per-token evidential heads parameterized with Dirichlet distributions, enabling fine-grained decomposition of aleatoric and epistemic uncertainty. Each spatial token’s feature vector is mapped through a lightweight multilayer perceptron to non-negative evidence vectors, which serve as concentration parameters for a Dirichlet distribution, following established evidential deep learning protocols [sensoy2018evidential]. Token-level uncertainty is computed as the reciprocal of total Dirichlet strength, providing spatially localized uncertainty maps that correlate with ambiguous structures in the input.

Uncertainty estimation is not limited to post-hoc analysis but is directly integrated into the computational graph, influencing both forward and backward passes.

Uncertainty-Gated Routing and Feature Refinement

The core operational advancement is the uncertainty-gated token refinement module, implemented from intermediate hierarchy stages onward. At each eligible transformer block, dual pathways are maintained: the DSSA attention branch for global context and a local refinement branch for focused enhancement of spatial details.

A learnable routing network predicts a spatial mask, which is modulated by tissue supervision (when available) and the Dirichlet-derived uncertainty. Corrective updates along the refinement pathway are then selectively applied to tokens with low uncertainty, while high-uncertainty tokens are suppressed, mitigating over-amplification of unreliable features. The global uncertainty, pooled over all tokens, further gates the magnitude of refinements according to stage-dependent parameters, enforcing increasingly conservative updates at higher levels of abstraction.

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

Prototype-Based Classification Head

Departing from linear classification heads, MedFormer-UR can optionally classify via learned class-specific prototypes. For each class, multiple prototypes are maintained; final token embeddings from MedFormer are compared to these prototypes using cosine similarity, and the aggregated evidence across tokens/prototypes determines the output logits. Auxiliary losses enforce both cluster formation around prototypes and diversity between prototypes of a class, structuring the space for transparent, similarity-based decisions.

This protocol is compatible with case-based interpretability, akin to “this-looks-like-that” reasoning frameworks [chen2019protopnet], and augments the model’s capability to justify classifications via reference to specific learned patterns.

Figure 3: Prototype-Based Classification Head.

Experimental Results

Datasets and Evaluation Paradigm

Evaluation spans four publicly available and clinically heterogeneous datasets: CBIS-DDSM (mammography), BUSI (ultrasound), BreaKHis (histopathology), and Brain MRI. Tasks encompass both binary and multi-class scenarios, targeting difficult and noisy modalities to stress-test calibration and uncertainty management.

Classification and Calibration Results

Across all modalities, MedFormer-UR outperforms the baseline in accuracy and AUROC, with the largest gains on BUSI (accuracy +9%, AUROC +0.037) and consistent, if modest, improvements even in challenging mammography (CBIS-DDSM, accuracy +4.4%, AUROC +0.028).

Strikingly, calibration behavior is substantially improved: Expected Calibration Error (ECE) on CBIS-DDSM is reduced by up to 35%, reaching a minimum of 0.0688 (at optimal routing hyperparameter $\beta$). Negative Log Likelihood (NLL) and Brier scores are also consistently reduced, indicating better probabilistic reliability, notably under moderate uncertainty routing. Excessive uncertainty gating degrades calibration—a robust optimum is observed for $\beta$ in [0.7, 0.8].

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.

Selective Prediction and Uncertainty Analysis

Selective prediction performance (measured via AURC and accuracy at fixed coverage) is optimized under intermediate uncertainty routing, with clear stratification between confident and ambiguous instances. Uncertainty maps produced by the Dirichlet head localize regions of visual and annotation ambiguity, typically at mass margins or dense tissue transitions in mammograms, aligning with known diagnostic challenges.

Figure 9: Uncertainty maps of two mammogram.

Prototype similarity visualizations demonstrate that confident predictions align with strong prototype matches, while ambiguous inputs show diluted or competing prototype evidence, supporting interactive or human-in-the-loop triage scenarios.

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.

Implications and Future Directions

The introduction of uncertainty as an active routing mechanism distinguishes MedFormer-UR from classical deterministic ViTs and post-hoc calibration approaches. Rather than applying rejection or probability correction after learning, uncertainty is deeply coupled to both architectural pathway selection and loss landscape navigation. This yields models that are not just well-calibrated but fundamentally hesitant in the absence of strong evidence, providing a critical safeguard against overconfident errors on out-of-distribution or poorly labeled instances.

Prototype-based classification bolsters interpretability and offers a pathway for coupling with case retrieval and audit systems, while also regularizing the embedding space against spurious decision boundaries. The observed improvements in calibration, especially in the presence of low SNR and class imbalance, have direct applicability to clinical triage and second-reader workflows.

Practical extensions include: (1) integration with multi-modal imaging and non-image clinical data, critical for reducing error rates in real-world diagnostics [pfob2022multimodal]; (2) leveraging soft-prototype or prototype-critic systems to represent atypical or rare variants [kim2016criticism]; and (3) formal evaluation of uncertainty-guided abstention in prospective clinical settings for risk-constrained diagnostics.

On the theoretical axis, results reinforce that accuracy and calibration are neither antagonistic nor redundant objectives—a robust clinical AI must explicitly optimize both, and MedFormer-UR demonstrates a viable blueprint for fusing ViT expressivity with reliability constraints.

Conclusion

MedFormer-UR formalizes a principled approach to trustworthy deep medical image classification, uniting Dirichlet-based evidential uncertainty modeling, uncertainty-gated feature refinement, and prototype-based structured decision-making within a single transformer framework. Quantitative and qualitative results across four clinical modalities evidence improved calibration (ECE reduction up to 35%), enhanced selective prediction, and interpretable case support, particularly valuable in ambiguous and high-risk tasks such as mammography. The explicit coupling of uncertainty estimation to both representation and routing constitutes a robust advance towards credible clinical integration of medical vision transformers.