论文阅读 (79)：TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image

1 概述

1.1 题目

2021：用于WSI分类的Transformer相关多示例 (TransMIL: Transformer based correlated multiple instance learning for whole slide image classification)

1.2 动机

WSI–MIL方法通常基于独立同分布假设，这忽略了不同实例之间的相关性。为了处理这个问题，提出了一个称为相关多示例的新框架。基于该框架，部署了一个基于Transformer的MIL (TransMIL)，其能够同时探索形态和空间信息。

TransMIL可视化效果好、可解释性强，能够高效处理不平衡/平滑和二/多分类问题。实验验证了其性能及展示了收敛速度。

图1：决策过程：(a) 独立同分布假设下的注意力机制；(b) 相关多示例框架下的自注意力机制

1.3 代码

Torch：https://github.com/szc19990412/TransMIL

1.4 附件

https://proceedings.neurips.cc/paper/2021/file/10c272d06794d3e5785d5e7c5356e9ff-Supplemental.pdf

1.5 引用

@article{Shao:2021:21362147,
author		={Zhu Chen Shao and Hao Bian and Yang Chen and Yi Feng Wang and Jian Zhang and Xiang Yang Ji and Yong Bing Zhang},
title		={{TransMIL}: {T}ransformer based correlated multiple instance learning for whole slide image classification},
journal		={Advances in Neural Information Processing Systems},
volume		={34},
pages		={2136--2147},
year		={2021}
}
1
2
3
4
5
6
7
8

2 方法

2.1 相关多示例

问题定义：以二分类MIL为例，给定包 X i = { x i , 1 , x i , 2 , … , x i , n } \mathbf{X}_i=\{\boldsymbol{x}_{i,1},\boldsymbol{x}_{i,2},\dots,\boldsymbol{x}_{i,n}\} Xi​={xi,1​,xi,2​,…,xi,n​}，对于$i=1,\dots ,b$，这样的定义表明包类实例相互依赖且有一定顺序。实例标签 { y i , 1 , y i , 2 , … , y i , n } \{y_{i,1},y_{i,2},\dots,y_{i,n}\} {yi,1​,yi,2​,…,yi,n​}是未知的，包标签 Y i ∈ { 0 , 1 } Y_i\in\{0,1\} Yi​∈{0,1}是已知的。一个MIL二分类器可以被定义为：
Y i = { 0 , iff ∑ y i , j = 0 y i , j ∈ { 0 , 1 } , j = 1 … n 1 , otherwise (1) \tag{1} Y_i= \left\{

$$\begin{array}{ll} 0,&\text{iff}\sum y_{i,j}=0\quad y_{i,j}\in\{0,1\},j=1\dots n\\ 1,&\text{otherwise} \end{array}$$ \right. Yi​={0,1,​iff∑yi,j​=0yi,j​∈{0,1},j=1…notherwise​(1)$${\hat{Y}}_i=S({\mathbf{X}}_i)\text{\hspace{1em}(2)}$$其中$S$是评分函数、${\hat{Y}}_i$表示预测、$b$是包的总数、$n$是包中的实例数，其对不同的包是可变的。

与Attention-net相比，进一步引入实例之间的相关性。定理1和推理给出了$S(\mathbf{X})$的任意形式，定理2说明了相关多示例的一些优势。

定理1 假设$S:\mathcal{X}\to \mathbb{R}$是一个关于Hausdirff距离$d_H(\cdot ,\cdot )$的连续集合函数，对于任意的可逆图$P:\mathcal{X}\to {\mathbb{R}}^n$，存在函数$\sigma$和$g$，使得对于任意的$\mathbf{X}\in \mathcal{X}$有：
∣ S ( X ) − g ( P X ∈ X { σ ( x ∈ X ) } ) ∣ < ϵ (3) \tag{3} |S(\mathbf{X})-g(P_{\mathbf{X}\in\mathcal{X}}\{\sigma(\boldsymbol{x}\in\mathbf{X})\})|<\epsilon ∣S(X)−g(PX∈X​{σ(x∈X)})∣<ϵ(3)即一个Hausdorff连续函数$S(\mathbf{X})$能够被$g(\cdot )$中的一个函数任意近似。

推理 基于定理1，对于任意的$\mathbf{X}$有：
∣ S ( X ) − g ( P X ∈ X { f ( x ) + h ( x ) : x ∈ X } ) ∣ < ϵ (7) \tag{7} |S(\mathbf{X})-g(P_{\mathbf{X}\in\mathcal{X}}\{f(\boldsymbol{x})+h(\boldsymbol{x}):\boldsymbol{x}\in\mathbf{X}\})|<\epsilon ∣S(X)−g(PX∈X​{f(x)+h(x):x∈X})∣<ϵ(7)

定理2：包中的实例可以通过随机变量${\Theta }_1,{\Theta }_2,\dots ,{\Theta }_n$表示，在相关假设下包的信息熵可以被表示为$H({\Theta }_1,{\Theta }_2,\dots ,{\Theta }_n)$，包在独立同分布 (i.i.d.) 假设下的信息熵可以被表示为${\sum }_{t=1}^{n}H({\Theta }_t)$，则有：
$$H({\Theta }_1,{\Theta }_2,\dots ,{\Theta }_n)=\sum _{t=2}^{n}H({\Theta }_t|,{\Theta }_1,\dots ,{\Theta }_{t-1})+H({\Theta }_1)\le \sum _{t=1}^{n}H({\Theta }_t)\text{\hspace{1em}(8)}$$定理2证明了相关假设下有更小的信息熵，其可以减少不确定性和引入更多有用的信息。基于此，算法1展示了相关算法。图2展示了TransMIL与已有方法的主要区别。


  图2：池化矩阵$\mathbf{P}$的差异：(a) 假设一个WSI中有5个实例，$\mathbf{P}\in {\mathbb{R}}^{5\times 5}$是相应的池化矩阵，对角线表示和自己的注意力权重，其余的为与其他实例的；(b-d) 均忽略了相关信息，因此$\mathbf{P}$是对角矩阵；(b) 第一个实例通过最大池化选择，因此只有一个非零值；© 平均池化下对角线的值相等；(d) 注意力的引入使得对角线上的值出现变化；(e) 得益于相关假设，非对角线上的值表明了实例之间的相关性

2.2 Transformer应用到相关MIL

Transformer使用自注意力机制来建模一个序列中的所有token的相关性，并添加位置信息来增加序列顺序信息的有用性。因此，使用函数$h$来编码所有实例的空间信息，以及$\mathbf{P}$使用自注意力来汇聚信息的Transformer是很有用的。

Transformer MIL 给定一个包的集合 { X 1 , X 2 , … , X b } \{\mathbf{X}_1,\mathbf{X}_2,\dots,\mathbf{X}_b\} {X1​,X2​,…,Xb​}及相应标签$Y_i$，目的是习得一种映射$\mathbb{X}\to \mathbb{T}\to \mathcal{Y}$，其中$\mathbb{X}$是包空间、$\mathbb{Y}$是Transformer空间。以及$\mathcal{Y}$是标签空间。

$\mathbb{X}\to \mathbb{T}$被定义为：
$${\mathbf{X}}_i^0=[{\mathbf{x}}_{i,class};f({\mathbf{x}}_{i,1});f({\mathbf{x}}_{i,2});\dots ;f({\mathbf{x}}_{i,n})]+{\mathbf{E}}_{pos},{\mathbf{X}}_i^0,{\mathbf{E}}_{pos}\in {\mathbb{R}}^{(n+1)\times d}\text{\hspace{1em}(15)}$$$${\mathbf{Q}}^{\ell }={\mathbf{X}}_i^{\ell -1}{\mathbf{W}}_Q,{\mathbf{K}}^{\ell }={\mathbf{X}}_i^{\ell -1}{\mathbf{W}}_K,{\mathbf{V}}^{\ell }={\mathbf{X}}_i^{\ell -1}{\mathbf{W}}_V,\ell =1\dots L\text{\hspace{1em}(16)}$$$$\mathbf{head}=\text{SA}({\mathbf{Q}}^{\ell },{\mathbf{K}}^{\ell },{\mathbf{V}}^{\ell })=\text{softmax}\left(\frac{{\mathbf{Q}}^{\ell }({\mathbf{K}}^{\ell }{)}^T}{\sqrt{d_q}}\right){\mathbf{V}}^{\ell },\ell =1\dots L\text{\hspace{1em}(17)}$$$$\text{MSA}({\mathbf{Q}}^{\ell },{\mathbf{K}}^{\ell },{\mathbf{V}}^{\ell })=\text{Concat}({\mathbf{head}}_1,{\mathbf{head}}_2,\dots ,{\mathbf{head}}_h){\mathbf{W}}^O,\ell =1\dots L\text{\hspace{1em}(18)}$$$${\mathbf{X}}_i^{\ell }\text{MSA(LN}({\mathbf{X}}_i^{\ell -1}))+{\mathbf{X}}_i^{\ell -1},\ell =1\dots L\text{\hspace{1em}(19)}$$其中SA表示自注意力、$L$是MSA的数量、$h$是每个MSK中头的数量，以及$LN$是标准化层。

$\mathbb{T}\to \mathcal{Y}$被定义为：
$$Y_i=\text{MLP(LN}(({\mathbf{X}}_i^L{)}^{(0)}))\text{\hspace{1em}(20)}$$其中$({\mathbf{X}}_i^L{)}^{(0)}$表示类别token。$\mathbb{T}\to \mathcal{Y}$可以通过类别token或者全局池化完成。然而，目前直接在WSI中使用Transformer相对困难，因此包中的实例数量很多且变化巨大。因此接下来注重如何高效地部署Transformer。

2.3 TransMIL用于弱监督WSI分类

为了更好地描述$\mathbb{X}\to \mathbb{T}$，设计了包含两个Transformer层的TPT模块和一个位置编码层，其中Transformer层用于汇聚形态信息，金字塔位置编码生成器 (Pyramid position encoding generator, PPEG) 用于编码空间信息。TransMIL的总体架构如图3。

  图3：TransMIL架构。每个WSI被裁剪为多个区块 (背景被抛弃)，并通过ResNet50嵌入为特征向量，然后传递给TPT处理：1) 序列平方；2) 序列的相关性建模；3) 条件位置编码和信息融合；4）深度特征汇聚；5) $\mathbb{T}\to \mathcal{Y}$映射

2.3.1 使用TPT对长实例序列建模

序列来自于WSI的特征向量。TPT的的处理过程如算法2。

大多数情况下，用于视觉任务的Transformer中的softmax是按行处理的。而标准的自注意力机制需要计算每一对toekn之间的相似性得分，太慢太耗内存。为了处理WSI中的长实例序列问题，TPT中的softmax使用Nystrom方法。近似自注意力机制$\hat{\mathbf{S}}$被定义为：
$$\hat{S}=\text{softmax}\left(\frac{\mathbf{Q}{\tilde{\mathbf{K}}}^T}{\sqrt{d_q}}\right){\left(\text{softmax}\left(\frac{\tilde{\mathbf{Q}}{\tilde{\mathbf{K}}}^T}{\sqrt{d_q}}\right)\right)}^{+}\text{softmax}\left(\frac{\tilde{\mathbf{Q}}{\mathbf{K}}^T}{\sqrt{d_q}}\right)\text{\hspace{1em}(9)}$$其中$\tilde{\mathbf{Q}}$和$\tilde{\mathbf{K}}$是从$\mathbf{Q}$和$\mathbf{K}$中的$n$维序列中选择的$m$个landmark，以及${\mathbf{A}}^{+}$是$\mathbf{A}$的Moore-Penrose伪逆。最终的时间复杂度将从$O(n^2)$降为$O(n)$。由此，TPT可以满足包中实例很多的情况。

2.3.2 PPEG位置编码

在 WSI 中，由于载玻片和组织的可变大小，相应序列中的标记数量通常会有所不同。有研究表明，添加零填充可以为卷积提供绝对位置信息。受此启发设计了PPEG模块，相应的伪代码如算法3。

PPEG模块有以下优势：

1. 同一层使用了不同大小的卷积核，可以编码不同粒度的位置信息，以扩展PPEG的能力；
2. 可以获取序列中token的全局信息和上下文信息，这能够丰富每个token的特征。

3 实验及结果

3.1 数据集

• CAMELYON16：用于乳腺癌转移检测的公开数据集，包含270个训练集和130个测试集。预处理后有大约350万个区块，量级为$\times 20$，每个包平均有8800个。
• TCGA-NSCLC：包含两个子类，TGCA-LUSC和TCGA-LUAD，共993个诊断WSI，包含444种情形的507个LUAD和452种情况的486个LUSC。预处理后，每个包量级在$\times 20$的区块平均为15371。
• TCGA-RCC：包含三个子类，THCA-KICH、TCGA-KIRC，以及TCGA-KIRP，共884个WSI，三个子类的情形数分别为99、483，以及264，幻灯片数则111、489，以及284.预处理后平均为14627。

3.2 实验设置和度量指标

1. 每个WSI的裁剪为$256\times 256$的无交叠区块，饱和度$<15$的背景将被抛弃；
2. CAMELYON16的训练集划分10%作为验证集；
3. TCGA划分时，首先确保训练和测试集中不存在来自一名患者的不同幻灯片，然后训练:验证:测试=$60:15:25$；
4. 准确率和AUC作为评估指标，其中准确率附加计算阈值$0.5$；
5. CAMELYON16使用测试AUC；
6. TCGA-NSCLC使用平均AUC；
7. TCGA-RCC使用macro-averaged AUC；
8. TCGA使用4折交叉验证。

3.3 实现细节

1. 交叉熵损失；
2. Lookahead优化器，学习率$=2e-4$，权重衰减$=1e-5$；
3. 批次大小$=1$；
4. 每个区块通过ResNet50嵌入为1024维向量，在训练时通过全连接层降维到$512$；
5. 包的最终嵌入为${\mathbf{H}}_i\in {\mathbb{R}}^{n\times 512}$；
6. softmax用于标准化每个类别的预测得分。

3.4 基准线

1. 注意力网络ABMIL和PT-MTA；
2. 非定位注意力DSMIL；
3. 单注意力CLAM-SB；
4. 多头CLAM-MB；
5. 循环神经网络MIL-RNN。

3.5 结果

分类：


消融实验：

可视化：

收敛性：