Vision Transformers with Autoencoders and Explainable AI for Cancer Patient Risk Stratification Using Whole Slide Imaging

2025-04-07

AI Summary

This study addresses the challenge of balancing discriminative performance and clinical interpretability in whole-slide image (WSI)-based cancer risk stratification. We propose PATH-X, a novel framework that integrates a vision transformer (ViT) with an autoencoder to learn robust, compressed tissue representations. Crucially, PATH-X is the first to unify unsupervised clustering in the ViT feature space with SHAP-based spatial attribution analysis within a WSI survival prediction pipeline. Evaluated on breast cancer and glioma datasets, PATH-X achieves statistically significant risk stratification (log-rank *p* < 0.001), and SHAP visualizations accurately localize prognostically relevant tissue regions. Performance is attenuated in lung cancer due to limited data availability. Our core contribution is the establishment of the first interpretable, WSI-specific survival modeling paradigm—simultaneously enhancing predictive accuracy and clinical trustworthiness through transparent, biologically grounded explanations.

Technology Category

• Computer Vision: Interpretability, Explainability, and Transparency
• Machine Learning: Transparent, Interpretable, Explainable ML
• Humans and AI: Explainable AI (XAI) for Human Understanding

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Security and Privacy: Data transparency and provenance

Abstract

Cancer remains one of the leading causes of mortality worldwide, necessitating accurate diagnosis and prognosis. Whole Slide Imaging (WSI) has become an integral part of clinical workflows with advancements in digital pathology. While various studies have utilized WSIs, their extracted features may not fully capture the most relevant pathological information, and their lack of interpretability limits clinical adoption. In this paper, we propose PATH-X, a framework that integrates Vision Transformers (ViT) and Autoencoders with SHAP (Shapley Additive Explanations) to enhance model explainability for patient stratification and risk prediction using WSIs from The Cancer Genome Atlas (TCGA). A representative image slice is selected from each WSI, and numerical feature embeddings are extracted using Google's pre-trained ViT. These features are then compressed via an autoencoder and used for unsupervised clustering and classification tasks. Kaplan-Meier survival analysis is applied to evaluate stratification into two and three risk groups. SHAP is used to identify key contributing features, which are mapped onto histopathological slices to provide spatial context. PATH-X demonstrates strong performance in breast and glioma cancers, where a sufficient number of WSIs enabled robust stratification. However, performance in lung cancer was limited due to data availability, emphasizing the need for larger datasets to enhance model reliability and clinical applicability.

Problem

• Improving cancer risk stratification using Vision Transformers and Autoencoders
• Enhancing interpretability of WSI analysis with SHAP for clinical adoption
• Addressing data limitations in lung cancer for reliable risk prediction

Innovation

• Vision Transformers extract WSI feature embeddings
• Autoencoders compress features for risk stratification
• SHAP provides explainable AI for clinical insights