论文阅读（1）：病理图像分类TransMIL: Transformer based Correlated Multiple Instance Learning

题目

Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification

概要

多实例学习(MIL)是解决基于全幻灯片图像弱监督分类(WSI)的病理诊断的有力工具。然而，目前的MIL方法往往基于独立相同的分布假设，忽视了不同实例之间的相关性。为了解决这个问题，我们提出了一个新的框架，称为相关MIL，并提供了收敛性的证明。在此基础上，我们设计了一种基于变形的MIL (TransMIL)，它同时探索了形态和空间信息。所提出的TransMIL能够有效地处理不平衡/平衡分类和二元/多元分类，具有良好的可视化和可解释性。我们对三个不同的计算病理学问题进行了各种实验，取得了更好的性能和更快的收敛速度，与最先进的方法相比。在CAMELYON16数据集上，二元肿瘤分类的检验AUC可达93.09%。TCGA-NSCLC数据集和TCGA-RCC数据集的AUC分别高达96.03%和98.82%。

Code

https://github.com/szc19990412/TransMIL

引入

病理全切片图像(whole slide image, WSI)扫描仪的出现为深度学习在数字病理领域的应用提供了很好的机会，它能将活检切片上的组织转化为完全保留原始组织结构的十亿像素图像。然而，基于深度学习的活检诊断在WSI中由于其巨大的尺寸和缺乏像素级的注释，而面临着巨大的挑战。为了解决这一问题，通常采用多实例学习(MIL)方法，将诊断分析视为弱监督学习问题。在基于深度学习的MIL中，一个简单的思路是对CNN提取的实例特征嵌入进行池化操作。Ilse等人提出了一种基于注意力的聚合算子，通过可训练的注意力权重为每个实例提供额外的贡献信息。此外，Li等将非局域注意力引入了MIL问题。通过计算得分最高的实例与其他实例的相似度，对每个实例给予不同的关注。

病理学家在作出诊断决定时，需要考虑单个区域的上下文信息和不同区域的信息，因而引入MIL机制和Transformer模型。

主体

1.相关工作

Application of MIL in WSI classification
Attention and Self-attention in Deep Learning

2.模型主要部分

TPT：聚合形态学信息
PPEG：编码空间位置信息

Overview：原始WSI图像裁剪成patches（丢掉背景信息），输入至ResNet50进行特征提取得到序列（包含class token、feature token、first M feature tokens），将得到的序列经过新的block（x-LN-MSA-A&N-PPEG-LN-MSA-A&N-）,输出序列再进入MLP head层，得到分类结果。

cls_tokenzai VIT的BERT中有介绍，feature tokens是网络的输入

TPT结构：输入一个feature embeddings包，输出包级别的预测标签Y，T代表Transformer层（LN+MSA+Add），P代表PPEG模块。

1. Squaring of sequence; √N ← d√ne, M ← N − n, HS ← Concat (hi,class, Hi, (hi,1, . . . , hi,M )), where hi,class ∈ R1×d represents class token, HS ∈ R(N+1)×d;
2. Correlation modelling of the sequence; HS ← MSA (HS), where denotes the layer index of the Transformer, H` S ∈ R(N+1)×d;
3. Conditional position encoding and local information fusion; HP S ← PPEG (H` S ), where HP S ∈ R(N+1)×d;
4. Deep feature aggregation; H+1 S ← MSA (HP S ), where H+1 S ∈ R(N+1)×d;
5. Mapping of T → Y; ˆ Yi ← MLP ( LN ( (H+1 S )(0))) , where (H+1 S )(0) ∈ R1×d represents class token.

Image embedding 和feature embedding都是降维的过程，把数据转换为固定大小的特征表示，以便处理和计算，其方法类比为一个没有激活函数的全连接层。

PPEG模块：输入一个模型建模后的feature embedding，输出位置编码后的feature embedding和本地融合信息。
其过程为：Class Token无变化（Identity层），Patch Tokens首先做维度恢复（reshape成为2Dimage space），然后进行群卷积（Identity和3×3，5×5，7×7的三个不同卷积核来做得到四个编码的空间信息），再把得到的空间信息进行融合，最后把得到的patch tokens拉平成为一维向量，把Class Token接上去得到最终输出给到下一层。

1. Split: H S is divided into patch tokens Hf and class token Hc; Hf , Hc ← Split (H S ), where Hf ∈ RN×d, Hc ∈ R1×d;
2. Spatial Restore: patch tokens Hf are reshaped to Hf S in the 2-D image space; Hf S ← Restore (Hf ), where Hf S ∈R √N ×√N ×d;
3. Group Convolution: using a set of group convolutions with kernel k and k−1 2 zero paddings(k = 3, 5, 7) to obtain Hf t , t = 1, 2, 3; Hf t ← Conv ( Hf S ) , where Hf t ∈R √N×√N×d, t = 1, 2, 3;
4. Fusion: Hf S and the Hf t , t = 1, 2, 3 obtained from the convolution block processing are added together to obtain HF S; HF S ← Hf S + Hf 1 + Hf 2 + Hf 3 , where HF S ∈R √N ×√N ×d;
5. Flatten: HF S are flattened into sequence Hse; Hse ← Flatten (HF S ), where Hse ∈ RN×d;
6. Concat: connect Hse and class token Hc to obtain HP S; HP S ← Concat (Hse, Hc), where HP S ∈ R(N+1)×d.

3.实验

本文使用公开数据集CAMELYON16和TCGA-NSCLC和TCGA-RCC：

CAMELYON16 is a public dataset for metastasis detection in breast cancer, including 270 training sets and 130 test sets. After pre-processing, a total of about 3.5 million patches at ×20 magnification, in average about 8,800 patches per bag were obtained.
Acc：0.8837 ACU：0.9309

TCGA-NSCLC includes two subtype projects, i.e., Lung Squamous Cell Carcinoma (TGCA-LUSC) and Lung Adenocarcinoma (TCGA-LUAD), for a total of 993 diagnostic WSIs, including 507 LUAD slides from 444 cases and 486 LUSC slides from 452 cases. After pre-processing, the mean number of patches extracted per slide at ×20 magnification is 15371.
Acc：0.8835 ACU：0.9603

TCGA-RCC includes three subtype projects, i.e., Kidney Chromophobe Renal Cell Carcinoma (TGCA-KICH), Kidney Renal Clear Cell Carcinoma (TCGA-KIRC) and Kidney Renal Papillary Cell Carcinoma (TCGA-KIRP), for a total of 884 diagnostic WSIs, including 111 KICH slides from 99 cases, 489 KIRC slides from 483 cases, and 284 KIRP slides from 264 cases. After pre-processing, the mean number of patches extracted per slide at ×20 magnification is 14627.
Acc：0.9466 ACU：0.9882

总结

本文基于MIL框架提出了一种基于Transformer的MIL，即TransMIL，用来探索在弱监督WSI分类任务中的形态和空间信息，设计了PPEG模块来进行位置编码，设计了TPT结构（两个Transformer层和一个PPEG位置编码层），实验结果描述上来看在三个公开数据集上的性能优于其他模型。