Uncertainty Estimation for Graph-Based Learning of WSIs

Whole slide images (WSIs) are high-resolution digital scans of tissue samples that contain spatial and contextual information vital for cancer diagnosis. Deep learning models used for classifying WSIs often struggle to provide trustworthy predictions when faced with data that differs from the training distribution, commonly referred to as out-of-distribution (OOD) data. The limitation becomes an existential challenge for clinical tasks, where trustworthy decision-making is crucial. To address this, our research enhances the uncertainty estimation of WSI classification using graph-based deep learning models. Unlike traditional models that tend to overconfidently rely on unseen data, our proposed method identifies and ranks WSIs based on the network's uncertainty for potential additional review. This capability is essential in clinical workflows, as uncertainty estimation can bring to light edge cases that need pathologist review, improving diagnostic safety and decision transparency.Our research provides a comprehensive study of two graph neural network (GNN) architectures-Graph Attention Networks (GAT) and GraphSAGE-for uncertainty-aware classification. Specifically, we propose multi-head GNN frameworks that incorporate multiple output branches. These branches introduce variability during inference, and we use the divergence among their predictions to estimate model uncertainty.Our experiments studied binary and multi-class breast cancer classification tasks both in in-distribution (ID) and OOD data. Our proposed multi-head GNN models are compared to baseline single-head GNNs, Monte Carlo dropout, and ensemble-based uncertainty estimation methods. Particularly, the three-head MH-GraphSAGE model consistently delivered the best balance of classification performance and uncertainty calibration across all datasets, including difficult OOD scenarios. Experiments demonstrate that our multi-head designs consistently achieve higher classification accuracy, along with the quality of uncertainty estimates.These findings confirm that multi-head GNNs offer a robust and computationally efficient solution for uncertainty estimation within digital pathology. By consistently annotating uncertain predictions, especially for distribution shifts, this approach enhances model trustworthiness in clinical settings and idetifies uncertain WSIs for additional review.

Learning Objectives: 

1. Graph-based Learning of WSIs
2. Uncertainty Estimation in Digital Pathology
3. WSI-level uncertainty ranking visualization