[论文解读] TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classication

Zhuchen Shao, Hao Bian+5 more|arXiv (Cornell University)|Jun 2, 2021

AI in cancer detection参考文献 38被引用 26

一句话总结

TransMIL 提出了一种基于 Transformer 的相关性多实例学习框架，通过建模全切片图像（WSIs）中补丁之间的空间和形态相关性，实现了最先进的性能，具有更快的收敛速度和更高的可解释性。在 CAMELYON16 数据集上，其二分类肿瘤检测 AUC 达到 93.09%，在 TCGA-RCC 数据集上，癌症亚型分类 AUC 高达 98.82%。

ABSTRACT

Multiple instance learning (MIL) is a powerful tool to solve the weakly supervised classification in whole slide image (WSI) based pathology diagnosis. However, the current MIL methods are usually based on independent and identical distribution hypothesis, thus neglect the correlation among different instances. To address this problem, we proposed a new framework, called correlated MIL, and provided a proof for convergence. Based on this framework, we devised a Transformer based MIL (TransMIL), which explored both morphological and spatial information. The proposed TransMIL can effectively deal with unbalanced/balanced and binary/multiple classification with great visualization and interpretability. We conducted various experiments for three different computational pathology problems and achieved better performance and faster convergence compared with state-of-the-art methods. The test AUC for the binary tumor classification can be up to 93.09% over CAMELYON16 dataset. And the AUC over the cancer subtypes classification can be up to 96.03% and 98.82% over TCGA-NSCLC dataset and TCGA-RCC dataset, respectively.

研究动机与目标

• 为解决现有 MIL 方法假设实例独立同分布（i.i.d.）的局限性，通过建模全切片图像中补丁之间的相关性来实现。
• 开发一种弱监督学习框架，以有效捕捉 WSIs 中的形态学和空间关系。
• 在类别不平衡和多分类的 WSI 数据集上提升分类性能，同时保持高水平的可解释性。
• 为所提出的关联性 MIL 框架提供理论收敛性证明。

提出的方法

• 提出一种关联性 MIL 框架，用结构化注意力机制替代 i.i.d. 假设，以建模补丁实例之间的依赖关系。
• 采用 Transformer 编码器，捕捉 WSI 中多个补丁实例之间的长距离空间和形态依赖关系。
• 应用可学习的类别标记和多头自注意力机制，将补丁级特征聚合为全局图像级表征。
• 使用全局平均池化层后接分类器头，实现弱监督标签下的端到端训练。
• 引入位置编码，以在补丁嵌入中保留空间上下文信息。
• 为所提出的关联性 MIL 框架提供理论收敛性证明。

实验结果

研究问题

• RQ1 在 WSI 中建模实例相关性是否能带来相较于标准 MIL 方法的性能提升？
• RQ2 基于 Transformer 的注意力机制在弱监督 WSI 分类中如何增强特征表示学习？
• RQ3 所提出方法在病理图像分析中对注意力图的可解释性和可视化改善程度如何？
• RQ4 关联性 MIL 框架是否具有理论保证的收敛性？与现有方法相比，其训练稳定性如何？

主要发现

• 在 CAMELYON16 数据集上，TransMIL 在二分类肿瘤检测任务中 AUC 达到 93.09%，优于当前最先进的方法。
• 在 TCGA-NSCLC 数据集上，TransMIL 在癌症亚型分类任务中 AUC 达到 96.03%，展现出在多分类问题上的强大性能。
• 在 TCGA-RCC 数据集上，TransMIL AUC 达到 98.82%，表明其在复杂肾癌亚型分类任务中表现极为出色。
• 与现有 MIL 方法相比，该模型在平衡和非平衡数据集上均表现出更快的收敛速度和更强的泛化能力。
• TransMIL 生成的注意力图具有高度可解释性，能清晰定位 WSI 中相关的肿瘤区域。