PTCMIL: Multiple Instance Learning via Prompt Token Clustering for Whole Slide Image Analysis

TL;DR

PTCMIL introduces a novel ViT-based multiple instance learning method with prompt token clustering, enhancing whole slide image analysis by improving robustness, efficiency, and interpretability across various tasks.

Contribution

It proposes a unified end-to-end framework that integrates learnable prompt tokens with clustering for WSI analysis, addressing heterogeneity and computational challenges.

Findings

Outperforms state-of-the-art methods on eight datasets.

Demonstrates superior accuracy in classification and survival analysis.

Shows robustness and interpretability through ablation studies.

Abstract

Multiple Instance Learning (MIL) has advanced WSI analysis but struggles with the complexity and heterogeneity of WSIs. Existing MIL methods face challenges in aggregating diverse patch information into robust WSI representations. While ViTs and clustering-based approaches show promise, they are computationally intensive and fail to capture task-specific and slide-specific variability. To address these limitations, we propose PTCMIL, a novel Prompt Token Clustering-based ViT for MIL aggregation. By introducing learnable prompt tokens into the ViT backbone, PTCMIL unifies clustering and prediction tasks in an end-to-end manner. It dynamically aligns clustering with downstream tasks, using projection-based clustering tailored to each WSI, reducing complexity while preserving patch heterogeneity. Through token merging and prototype-based pooling, PTCMIL efficiently captures task-relevant patterns. Extensive experiments on eight datasets demonstrate its superior performance in classification and survival analysis tasks, outperforming state-of-the-art methods. Systematic ablation studies confirm its robustness and strong interpretability. The code is released at https://github.com/ubc-tea/PTCMIL.