Multi-Cohort Framework with Cohort-Aware Attention and Adversarial Mutual-Information Minimization for Whole Slide Image Classification

TL;DR

A Cohort-Aware Attention module is introduced, enabling the capture of both shared and tumor-specific pathological patterns, enhancing cross-tumor generalization and developing a hierarchical sample balancing strategy to mitigate cohort imbalances and promote unbiased learning.

Abstract

Whole Slide Images (WSIs) are critical for various clinical applications, including histopathological analysis. However, current deep learning approaches in this field predominantly focus on individual tumor types, limiting model generalization and scalability. This relatively narrow focus ultimately stems from the inherent heterogeneity in histopathology and the diverse morphological and molecular characteristics of different tumors. To this end, we propose a novel approach for multi-cohort WSI analysis, designed to leverage the diversity of different tumor types. We introduce a Cohort-Aware Attention module, enabling the capture of both shared and tumor-specific pathological patterns, enhancing cross-tumor generalization. Furthermore, we construct an adversarial cohort regularization mechanism to minimize cohort-specific biases through mutual information minimization. Additionally, we develop a hierarchical sample balancing strategy to mitigate cohort imbalances and promote unbiased learning. Together, these form a cohesive framework for unbiased multi-cohort WSI analysis. Extensive experiments on a uniquely constructed multi-cancer dataset demonstrate significant improvements in generalization, providing a scalable solution for WSI classification across diverse cancer types. Our code for the experiments is publicly available at <link>.

Figures (5)

• Figure 1: Multi-cohort analysis paradigms. Left: Standard joint training, lacks mitigation of cohort biases. Right: Our framework, learning shared underlying mechanisms across cohorts.
• Figure 2: Cohort-aware attention module integrated into a VisionTransformer dosovitskiy2020image block. The dataset-wide query $Q_d$ and cohort-specific query $Q_c$ are processed through a token Query-Attention (QA) component to form $Q_{ca}$, which is then fed into the scaled dot-product attention.
• Figure 3: Hierarchical data structure of multi-cohort WSI analysis during pretraining and MIL phases. The pretraining phase balances cohorts, slides, and tiles, ensuring proportional representation at each level. The MIL phase further balances cohort-class combinations and uniformly distributes weights among slides. This strategy mitigates imbalances, promoting effective and unbiased model learning.
• Figure 4: AUC Performance with Varying $\lambda$ Values. A: MSS/MSI Task. B: GS/CIN Task.
• Figure 5: AUC Performance of different sample weighting techniques. A: MSS/MSI Task. B: GS/CIN Task.