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

1. Problem Statement

Deep learning models, particularly Vision Transformers (ViTs) like MedFormer, have achieved state-of-the-art performance in medical image analysis. However, their deployment in clinical settings is hindered by a lack of trustworthiness.

• Overconfidence: Standard models produce poorly calibrated probabilities, often exhibiting extreme confidence on ambiguous or out-of-distribution inputs.
• Black Box Nature: They lack transparency regarding why a decision was made.
• Data Challenges: Clinical data is inherently noisy, imbalanced, and ambiguous (e.g., low-contrast mammograms, overlapping tissues), leading to high aleatoric (data) and epistemic (model) uncertainty.
• Limitations of Existing Solutions: Bayesian methods (e.g., Monte Carlo Dropout) are computationally expensive, and post-hoc calibration techniques (e.g., Temperature Scaling) do not address overconfidence during the training phase.

The paper proposes a framework that integrates uncertainty quantification directly into the learning process to ensure safe, reliable, and interpretable clinical integration.

2. Methodology: MedFormer-UR

The authors propose MedFormer-UR, an extension of the hierarchical MedFormer architecture. The framework introduces three core mechanisms without altering the core Dual Sparse Selection Attention (DSSA) backbone:

A. Per-Token Evidential Uncertainty Estimation

Instead of standard softmax outputs, the model employs an evidential learning framework using a Dirichlet distribution.

• Mechanism: A lightweight "evidential head" predicts non-negative evidence values ($e_i$) for each spatial token. These are transformed into Dirichlet parameters ($\alpha_i = e_i + 1$).
• Uncertainty Calculation: Uncertainty ($\sigma_i$) is derived from the Dirichlet strength ($S_i = \sum \alpha_{i,c}$). Low evidence results in high uncertainty.
• Benefit: This allows the model to quantify and localize ambiguity at the token level (spatial regions) in real-time, distinguishing between aleatoric and epistemic uncertainty.

B. Uncertainty-Guided Routing and Refinement (UGTR)

Uncertainty is not just an output; it acts as an active control signal during the forward pass.

• Dual Path Architecture: Transformer blocks feature a main attention path (capturing global context) and a refinement path (capturing local details).
• Routing Mask: A routing network predicts a soft mask to determine where refinement is needed.
• Uncertainty Gating: The refinement update is modulated by the token-level uncertainty.
  • High Uncertainty: The refinement is suppressed (gated down) to prevent the model from amplifying unreliable features. The model relies on the stable baseline representation.
  • Low Uncertainty: Refinement is applied to sharpen diagnostically relevant features (e.g., lesion boundaries).
• Global Gate: A global uncertainty coefficient ($\beta$) controls the overall magnitude of refinement based on the aggregated uncertainty of the image.

C. Prototype-Based Representation Learning

To improve interpretability, the classification head is replaced or augmented with a prototype-based approach.

• Concept: Instead of a linear classifier, the model learns class-specific prototypes (reference patterns).
• Decision Logic: Predictions are made based on the similarity (cosine similarity) between input features and these learned prototypes.
• Regularization: A cluster loss ensures prototypes match tokens, and a diversity loss prevents prototype collapse.
• Benefit: This grounds decisions in visual similarity to known patterns, making the reasoning process transparent and clinically explainable.

D. Training Objective

The model is trained using a composite loss function:
$$L = L_{CE} + \lambda_{route}L_{routing} + L_{prototype}$$
Where $L_{CE}$ is cross-entropy, $L_{routing}$ supervises the routing mask (using available tissue masks), and $L_{prototype}$ enforces prototype structure.

3. Key Contributions

1. Novel Architecture: A unified extension of MedFormer incorporating evidential heads for per-token uncertainty decomposition and uncertainty-guided feature routing.
2. Active Uncertainty Control: A mechanism where uncertainty dynamically suppresses or refines feature updates during training, rather than just being a post-hoc metric.
3. Interpretable Embedding Space: Integration of learnable class prototypes with contrastive loss, enabling similarity-based reasoning.
4. Comprehensive Evaluation: Extensive testing across four diverse modalities (Mammography, Ultrasound, MRI, Histopathology) demonstrating improvements in both accuracy and reliability.

4. Experimental Results

The model was evaluated on four datasets: CBIS-DDSM (Mammography), BUSI (Ultrasound), IDC (Histopathology), and Brain MRI.

• Classification Performance:

  • Ultrasound (BUSI): Significant improvement in Accuracy (0.800 $\to$ 0.893) and AUROC (0.919 $\to$ 0.956).
  • Histopathology (BreaKHis): Moderate but consistent gains (Accuracy 0.9206 $\to$ 0.9303).
  • Brain MRI: High baseline performance was further improved (Accuracy 0.9733 $\to$ 0.9809).
  • Mammography (CBIS-DDSM): Despite inherent data difficulty (low contrast, noise), the model improved Accuracy (0.629 $\to$ 0.673) and AUROC (0.727 $\to$ 0.755).

• Calibration and Reliability (Crucial Finding):

  • The most significant impact was on model calibration. The proposed method reduced the Expected Calibration Error (ECE) by up to 35% (from 0.1836 to 0.0688 at optimal $\beta$).
  • Selective Prediction: The model demonstrated superior ability to reject uncertain predictions. At optimal routing strength ($\beta \approx 0.8$), it achieved the lowest Area Under the Risk-Coverage Curve (AURC), indicating that when the model is confident, it is highly likely to be correct.
  • Non-Monotonic Behavior: The study found that increasing uncertainty routing strength ($\beta$) does not linearly improve performance. An intermediate value (e.g., $\beta=0.8$) offered the best balance between discrimination and reliability, while extreme values degraded calibration.

5. Significance and Conclusion

MedFormer-UR addresses the critical gap between high accuracy and clinical trustworthiness.

• Clinical Relevance: By reducing overconfidence and providing uncertainty maps that highlight ambiguous regions (e.g., dense tissue or unclear lesion boundaries), the model supports safer decision-making. It allows clinicians to trust high-confidence predictions and seek human review for low-confidence cases.
• Interpretability: The combination of uncertainty maps (showing where the model is unsure) and prototype similarity (showing what visual patterns drive the decision) provides a transparent rationale for predictions.
• Generalizability: The framework proves effective across diverse imaging modalities, suggesting that uncertainty-aware routing and prototype learning are robust strategies for medical AI.

In summary, the paper demonstrates that integrating evidential uncertainty and prototype learning into hierarchical Transformers creates a model that is not only accurate but also calibrated, reliable, and interpretable, which are essential prerequisites for the safe deployment of AI in healthcare.