[论文解读] TPMIL: Trainable Prototype Enhanced Multiple Instance Learning for Whole Slide Image Classification

Litao Yang, Deval Mehta+4 more|arXiv (Cornell University)|May 1, 2023

AI in cancer detection被引用 4

一句话总结

TPMIL 提出了一种可训练的原型增强型多实例学习框架，用于弱监督的全切片图像（WSI）分类，通过软聚类所有图像块嵌入来优化特征空间，并利用原型距离提升可解释性。该方法在两个 WSI 数据集上实现了最先进性能，AUC 和准确率均优于 CLAM 和 DGMIL 等先前方法。

ABSTRACT

Digital pathology based on whole slide images (WSIs) plays a key role in cancer diagnosis and clinical practice. Due to the high resolution of the WSI and the unavailability of patch-level annotations, WSI classification is usually formulated as a weakly supervised problem, which relies on multiple instance learning (MIL) based on patches of a WSI. In this paper, we aim to learn an optimal patch-level feature space by integrating prototype learning with MIL. To this end, we develop a Trainable Prototype enhanced deep MIL (TPMIL) framework for weakly supervised WSI classification. In contrast to the conventional methods which rely on a certain number of selected patches for feature space refinement, we softly cluster all the instances by allocating them to their corresponding prototypes. Additionally, our method is able to reveal the correlations between different tumor subtypes through distances between corresponding trained prototypes. More importantly, TPMIL also enables to provide a more accurate interpretability based on the distance of the instances from the trained prototypes which serves as an alternative to the conventional attention score-based interpretability. We test our method on two WSI datasets and it achieves a new SOTA. GitHub repository: https://github.com/LitaoYang-Jet/TPMIL

研究动机与目标

• 为解决因全切片图像（WSI）分辨率高且缺乏图像块级别标注而导致的弱监督 WSI 分类挑战。
• 通过利用所有实例嵌入而非依赖少数极端图像块来改进特征空间表征。
• 通过引入基于原型距离的可视化，超越注意力分数，提升可解释性。
• 通过学习到的原型距离建模肿瘤亚型之间的相关性。
• 开发一种可训练的端到端框架，联合优化特征学习与原型优化。

提出的方法

• 提出一种可训练的原型增强型 MIL 框架（TPMIL），通过将所有图像块嵌入分配给对应原型，实现对所有图像块嵌入的软聚类。
• 采用可微分的原型学习机制，基于所有实例特征的分布更新原型中心，实现全局特征空间优化。
• 采用基于原型的注意力机制，其中图像块的可解释性由其与预测类别最近原型的 L2 距离决定。
• 引入负向原型以捕捉非诊断区域，如正常组织、血管和伪影。
• 使用切片级别标签进行端到端训练，联合优化特征提取器、原型中心和分类器。
• 用原型距离热图替代传统注意力分数，以提升临床可视化中的鲁棒性与可解释性。

Figure 1: Overview of TPMIL. The key novelty of our proposed framework is the integration of Prototype Learning Module for learning an optimal distribution of instances.

实验结果

研究问题

• RQ1 与依赖极端图像块的方法相比，是否一种考虑所有图像块嵌入的原型方法能改善弱监督 WSI 分类中的特征空间表征？
• RQ2 学习到的原型距离与 WSI 中的生物肿瘤亚型及形态学差异之间是否存在相关性？
• RQ3 基于原型距离的热图是否比基于注意力分数的热图提供更鲁棒且更符合临床需求的可解释性？
• RQ4 所提出方法是否在多类别分类的基准 WSI 数据集上实现了最先进性能？
• RQ5 该框架是否能有效利用原型间距离区分肿瘤亚型？

主要发现

• 在 TCGA 脑肿瘤数据集上，TPMIL 实现了新的最先进（SOTA）性能，AUC 提高 0.24%，准确率提高 1.54%，优于 CLAM。
• 在 TCGA 肺癌数据集上，TPMIL 的 AUC 比 CLAM 高 0.11%，准确率高 1.41%，表现出一致的 SOTA 性能。
• 原型间距离矩阵显示，ACM（星形细胞瘤）与 ODG（少突胶质瘤）之间的距离比与 GBM（胶质母细胞瘤）更近，反映了它们同属低级别胶质瘤的生物学特征。
• 负向原型与 ACM 和 ODG 的距离比与 GBM 更近，表明 GBM 的形态多样性及更高的病理复杂性使其与正常样区域更为不同。
• TPMIL 中的距离热图相比注意力分数热图对伪影和背景噪声更具鲁棒性，尤其在突出胶质母细胞瘤中的坏死区域方面表现更优。
• 专家临床医生确认，原型距离热图比基于注意力的热图提供更可靠、更具临床可解释性的可视化效果，尤其在区分诊断特征与伪影方面。

Figure 2: The distance matrix of TCGA Brain Tumor prototypes after training