TAD-Graph: Enhancing Whole Slide Image Analysis via Task-Aware Subgraph Disentanglement

Learning contextual features such as interactions among various biological entities is vital for whole slide images (WSI)-based cancer diagnosis and prognosis. Graph-based methods have surpassed traditional multi-instance learning in WSI analysis by robustly integrating local pathological and contextual interaction features. However, the high resolution of WSIs often leads to large, noisy graphs. This can result in shortcut learning and overfitting due to the disproportionate graph size relative to WSI datasets. To overcome these issues, we propose a novel Task-Aware Disentanglement Graph approach (TAD-Graph) for more efficient WSI analysis. TAD-Graph operates on WSI graph representations, effectively identifying and disentangling informative subgraphs to enhance contextual feature extraction. Specifically, we inject stochasticity into the edge connections of the WSI graph and separate the WSI graph into task-relevant and task-irrelevant subgraphs. The disentanglement procedure is optimized using a graph information bottleneck-based objective, with added constraints on the task-irrelevant subgraph to reduce spurious correlations from task-relevant subgraphs to labels. TAD-Graph outperforms existing methods in three WSI analysis tasks across six benchmark datasets. Furthermore, our analysis using pathological concept-based metrics demonstrates TAD-Graph's ability to not only improve predictive accuracy but also provide interpretive insights and aid in potential biomarker identification. Our code is publicly available at https://github.com/fuying-wang/TAD-Graph.