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

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

AI in cancer detection被引用 501

一句话总结

TransMIL 引入基于 Transformer 的相关性 MIL 框架用于整张切片分类，利用形态和空间信息，在多个数据集上超越最先进的 MIL 方法，收敛更快并具备更好的可解释性。

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. Implementation is available at: https://github.com/szc19990412/TransMIL.

研究动机与目标

• 通过超越独立同分布实例假设，建模实例之间的相关性来解决对弱监督 WSI 分类的问题。
• 提出基于 Transformer 的 MIL（TransMIL）来同时捕捉补丁的形态和空间关系。
• 提供一个在病理学中收敛、可解释且高效的三步相关 MIL 算法。
• 展示在多个公开组学病理数据集上的优越性能和更快的收敛。

提出的方法

• 提出一个带有收敛性证明的通用三步相关 MIL 框架。
• 使用基于 Transformer 的 Temporal Patch Transformer (TPT) 模块开发 TransMIL，以建模补丁嵌入之间的长程相关性。
• 引入 Pyramid Position Encoding Generator (PPEG) 用于条件的多尺度空间位置编码。
• 在 TPT 中使用 Nyström-based 近似自注意力以高效处理长序列。
• 将形态嵌入（f）、空间嵌入（h）及其和（fh）结合起来，随后是一个学习的池化矩阵 P 和最终映射 g 以得到袋标签。
• 提供一个端到端的训练设置，使用交叉熵损失和 Lookahead 优化器，补丁嵌入基于 ResNet50（1024 维，降至 512 维）。

实验结果

研究问题

• RQ1 在 WSI 内建模实例相关性，是否能相较于基于独立同分布的 MIL 方法提升 MIL 基于分类的性能？
• RQ2 如何在保持空间上下文的前提下，将 Transformer 为基础的架构适配于大规模、可变长度的 WSI 补丁序列？
• RQ3 条件的多尺度位置编码（PPEG）是否提升病理学 MIL 的诊断性能和可解释性？
• RQ4 TransMIL 是否在二分类和多分类 WSI 分类任务以及不平衡数据集上均有效？
• RQ5 相关 MIL 方法是否比现有 MIL 方法收敛更快？

主要发现

Dataset	Metric	Mean-pooling	Max-pooling	ABMIL	PT-MTA	MIL-RNN	DSMIL	CLAM-SB	CLAM-MB	TransMIL
CAMELYON16	Accuracy	0.6389	0.8062	0.8682	0.8217	0.8450	0.7985	0.8760	0.8372	0.8837
CAMELYON16	AUC	0.4647	0.8569	0.8760	0.8454	0.8880	0.8179	0.8809	0.8679	0.9309
TCGA-NSCLC	Accuracy	0.7282	0.8593	0.7719	0.7379	0.8619	0.8058	0.8180	0.8422	0.8835
TCGA-NSCLC	AUC	0.8401	0.9463	0.8656	0.8299	0.9107	0.8925	0.8818	0.9377	0.9603
TCGA-RCC	Accuracy	0.9054	0.9378	0.8934	0.9059	-	0.9294	0.8816	0.8966	0.9466
TCGA-RCC	AUC	0.9786	0.9879	0.9702	0.9700	-	0.9841	0.9723	0.9799	0.9882

• TransMIL 在三个公开数据集 CAMELYON16、TCGA-NSCLC 和 TCGA-RCC 上达到最先进的性能。
• 在 CAMELYON16 上，TransMILACCURACY 0.8837 且 AUC 0.9309，优于基于 i.i.d. 的方法和一些非本地方法。
• 在 TCGA-NSCLC 上，TransMIL 达到 ACC 0.8835 和 AUC 0.9603，优于竞争方法。
• 在 TCGA-RCC 上，TransMIL 达到 ACC 0.9466 和 AUC 0.9882，在所报道的方法中最好。
• 消融实验表明 PPEG 与条件位置编码显著提升性能，相较于正弦位置编码或无位置编码。
• TransMIL 展现出更快的收敛，比竞争 MIL 方法少 2-3 倍训练轮次。
• 注意力图的可视化显示与癌变区域高度重叠，支持可解释性。