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

TL;DR

MedFormer-UR introduces an uncertainty-guided transformer that improves medical image classification by quantifying ambiguity, enhancing calibration, and enabling reliable, interpretable predictions across multiple modalities.

Contribution

It integrates a Dirichlet-based uncertainty routing with prototype learning in a transformer, improving calibration and interpretability in medical imaging tasks.

Findings

Reduces expected calibration error (ECE) by up to 35%

Enhances model calibration and selective prediction across four modalities

Provides real-time uncertainty localization and filtering in predictions

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.