Uncertainty Estimation in Pathology Foundation Models via Deep Mutual Learning

Technical Summary: Uncertainty Estimation in Pathology Foundation Models via Deep Mutual Learning

Problem Statement
Pathology Foundation Models (PFMs) have demonstrated strong generalization capabilities in extracting representations from whole-slide images (WSIs) for tasks such as survival analysis, biomarker discovery, and tumor subtyping. However, their translation to routine clinical practice is hindered by two critical limitations:

1. Lack of Confidence Estimates: Existing PFM pipelines typically output predictions without explicit measures of confidence or uncertainty. In clinical settings, models must not only be accurate but also capable of indicating when their predictions are uncertain to support safe decision-making.
2. Task-Dependent Performance: No single PFM is universally superior across all downstream tasks, datasets, and clinical contexts. Selecting the optimal model for deployment is challenging, and relying on a single model risks suboptimal performance or overconfidence in failure-prone cases.

Current fusion methods often improve performance but fail to provide explicit, task-agnostic uncertainty estimates. Conversely, uncertainty quantification methods like deep ensembles or Monte Carlo dropout are either not specifically designed for PFM ensembles or are task-specific (e.g., limited to glioblastoma diagnosis).

Methodology: DICE Framework
The authors propose DICE (Disagreement-Informed Coordination of Experts), a plug-and-play framework that ensembles $K$ frozen PFMs to estimate uncertainty and improve downstream performance. The core intuition is that the disagreement among heterogeneous PFMs serves as an informative signal for uncertainty, provided this disagreement is structured rather than dominated by optimization noise.

The framework operates as follows:

1. Architecture: The system utilizes $K$ frozen PFMs (e.g., Virchow2, UNI2-h, H-optimus-1, CONCHv1.5, Hibou-L). Each PFM acts as an "expert" $E^{(k)}$, paired with an independent Multiple Instance Learning (MIL) aggregator (specifically ABMIL) and a classification head.

2. Training Objectives: To ensure the ensemble's disagreement reflects meaningful uncertainty rather than noise, the experts are coupled during training via three objectives:

   • Supervised Loss ($L_{sup}$): Standard cross-entropy loss to ensure accurate slide-level classification.
   • Deep Mutual Learning ($L_{DML}$): A prediction-space alignment objective where each expert minimizes the Kullback–Leibler (KL) divergence between its predictive distribution and those of its $K-1$ peers. This encourages consensus while maintaining alignment with ground truth.
   • Gramian Regularization ($L_{Gram}$): A representation-space alignment objective based on the Gramian measure. It encourages the slide-level embeddings of the experts to span a low-volume subspace, preventing geometric misalignment that could masquerade as uncertainty.
   • Total Loss: The training objective combines these terms: $L = L_{sup} + L_{DML} + \lambda_{Gram}(t) L_{Gram}$, where $\lambda_{Gram}$ follows a warm-up schedule.

3. Inference and Uncertainty Estimation:

   • Uncertainty Proxy: At inference, the ensemble prediction is the average of the $K$ experts' distributions. The total predictive uncertainty is decomposed into aleatoric (data) and epistemic (model) components. The authors theoretically prove that the epistemic uncertainty (residual disagreement) is upper-bounded by the aggregate DML loss. Thus, minimizing the DML loss during training contracts the upper bound on model uncertainty.
   • Localization: Patch-level lesion localization is achieved unsupervised by computing the consensus of attention weights across all experts.

Key Contributions

1. First Task-Agnostic Framework: DICE is the first framework to leverage PFM ensembles specifically for task-agnostic uncertainty estimation in computational pathology.
2. Theoretical Foundation: The paper provides a theoretical proof that minimizing the training-time DML alignment objective contracts an upper bound on the ensemble's epistemic uncertainty, validating the use of residual disagreement as a reliable uncertainty proxy.
3. Empirical Performance: Extensive experiments demonstrate that DICE not only accurately flags failure-prone samples (both in- and out-of-distribution) but also matches or outperforms state-of-the-art (SOTA) single PFMs and fusion baselines in classification, calibration, and unsupervised lesion localization.

Experimental Results
The framework was evaluated on three weakly supervised WSI benchmarks: PANDA (prostate cancer grading), CAMELYON16, and CAMELYON17 (breast lymph-node metastasis detection).

• Classification & Calibration: DICE achieved the highest F1 scores and lowest Negative Log-Likelihood (NLL) across most datasets. On PANDA, it improved F1 by 3.8 percentage points over the best single PFM and 2.8 points over the best fusion baseline.
• Uncertainty Estimation: In selective prediction experiments, DICE effectively deferred high-uncertainty slides, resulting in the lowest retained error rates. It demonstrated superior separation between correct and incorrect predictions (measured by Cliff's $\Delta$) compared to Monte Carlo Dropout and other fusion methods. This capability held even under cohort shifts (e.g., training on CAMELYON16 and testing on CAMELYON17).
• Localization: By aggregating attention maps, DICE provided unsupervised patch-level localization that was competitive with or superior to single models, particularly in concentrating attention within annotated tumor regions rather than surrounding stromal tissue.
• Ablation Studies: The framework proved robust to the choice of MIL aggregator (DSMIL, TransMIL) and showed that increasing the number of experts yields diminishing returns once the feature space is sufficiently covered, highlighting the importance of expert complementarity over sheer ensemble size.

Significance and Claims
The paper positions DICE as a crucial step toward translating PFMs into uncertainty-aware decision-support systems. By establishing a theoretically grounded approach to fusing PFMs, DICE addresses the "black box" nature of current models, offering a mechanism to identify error-prone cases and safely defer to clinicians. The authors claim that their method makes PFMs more trustworthy in clinical settings by providing reliable uncertainty estimates without requiring explicit supervision for localization or confidence calibration.

The authors acknowledge limitations, including the computational cost of extracting features from $K$ frozen models (though training costs remain lower than end-to-end fine-tuning) and the need to carefully balance alignment and diversity to avoid suppressing the very disagreement required for uncertainty estimation. They conclude that while the number of experts is currently fixed, the framework paves the way for more reliable, uncertainty-aware pathology models integrated into clinical workflows.