TransMIL：基于Transformer的多实例学习

MIL是弱监督分类问题的有力工具。然而，目前的MIL方法通常基于iid假设，忽略了不同实例之间的相关性。为了解决这个问题，作者提出了一个新的框架，称为相关性MIL，并提供了收敛性的证明。基于此框架，还设计了一个基于Transformer的MIL (TransMIL)。TransMIL可以有效地处理不平衡/平衡和二元/多分类，具有良好的可视化和可解释性。在CAMELYON16数据集上，二元肿瘤分类的测试AUC高达93.09%。在TCGANSCLC和TCGA-RCC数据集上，癌症亚型分类的AUC分别高达96.03%和98.82%。

来自：TransMIL: Transformer based Correlated Multiple Instance Learning for Whole Slide Image Classification

背景概述

WSI将活检切片上的组织转换成完全保留原始组织结构的十亿像素图像。然而，WSI中基于深度学习的活检诊断由于像素空间庞大导致缺乏像素级标注。为了解决这个问题，通常采用MIL将诊断分析作为一个弱监督学习问题。

在基于深度学习的MIL中，一个简单的想法是对CNN提取的instance特征嵌入进行pooling操作。Ilse等人提出了一种基于注意力的聚合算子，通过可训练的注意力为每个实例提供额外的贡献信息。此外，Li将非局部注意力引入了MIL问题。通过计算得分最高的实例与其他实例之间的相似度，赋予每个实例不同的注意力权重，从而得到可解释的注意力图。

然而，所有这些方法都基于这样的假设：每个bag中的所有实例都是独立且同分布的。虽然在许多任务中取得了一些改进，但在许多情况下，这种假设并不完全有效。实际上，在做出诊断决定时，病理学家通常会同时考虑单个区域周围的环境信息和不同区域之间的相关信息。因此，在MIL诊断中考虑不同instance之间的相关性是可取的。

目前，Transformer由于具有较强的描述序列中不同token之间的相关性以及对远距离信息建模的能力，被广泛应用于视觉任务中。如图1所示，Transformer采用自注意力机制，可以关注序列内每个token之间的两两相关性。然而，传统的Transformer受到其计算复杂性的限制，只能处理较短的序列(例如，小于1000或512)。因此，它不适合WSI等大尺寸图像。

• 图1：决策过程图示。MIL注意力机制：遵循iid假设。自注意机制：属于相关性MIL。

方法

以二元MIL为例，我们想要预测target value Y i ∈ { 0 , 1 } Y_{i}\in\left\{0,1\right\} Yi​∈{0,1}，给定一个bag $X_i$（instance为 { x i , 1 , x i , 2 , . . . , x i , n } \left\{x_{i,1},x_{i,2},...,x_{i,n}\right\} {xi,1​,xi,2​,...,xi,n​}），其中$i=1,..,b$，实例级标签是未知的： { y i , 1 , y i , 2 , . . . , y i , n } \left\{y_{i,1},y_{i,2},...,y_{i,n}\right\} {yi,1​,yi,2​,...,yi,n​}，bag标签是已知的，并且与实例标签有联系：

$b$是袋的总数，$n$是第$i$个袋里的实例数，$n$的个数可以根据不同的袋而变化。

关于相关性MIL的优势，文中给出了证明，但是在此略过。主要意思是：

• 考虑实例相关性可以具有更小的信息熵，从而减少不确定性，为MIL带来更多有效信息。TransMIL与过去方法的主要区别如图2

• 图2：不同pooling矩阵$P$的差异。假设从a中的WSI采样5个instance，$P\in {\mathbb{R}}^{5\times 5}$是对应的pooling矩阵，其中对角线内的值表示instance自身的注意力权重，其余值表示不同instance之间的相关性。b,c,d 都忽略了相关信息，因此$P$是对角矩阵。在b中，第一个实例是由Max-pooling算子选择的，所以在对角线位置只有一个非零值。在c中，由于Mean-pooling运算符，对角线内的所有值都是相同的。在d中，由于引入的是bypass注意力，对角线内的值可能会发生变化。但其余位置为0（独立同分布假设）。e服从相关性假设，因此在非对角线位置存在非零值，表示不同实例之间存在相关性。

对于MIL的pooling，这里有一个通用的三步法：

• 形态信息：morphological，空间信息：spatial

如何应用Transformer到相关性MIL

Transformer使用自注意力机制对序列中所有令牌之间的交互进行建模，位置信息的添加进一步增加了顺序信息。因此，将Transformer引入相关性MIL问题是合理的，其中函数$h$对实例之间的空间信息进行编码，pooling矩阵$P$使用自注意力进行信息聚合。为了说明这一点，进一步给出一个正式的定义。

给定一个bag集合 { X 1 , . . . , X b } \left\{X_{1},...,X_{b}\right\} {X1​,...,Xb​}，每个bag对应一个标签$Y_i$。目标是学习映射：$\mathbb{X}\to \mathbb{T}\to \mathbb{Y}$，即从bag空间，到Transformer空间，再到标签空间。

为了更好地描述$\mathbb{X}\to \mathbb{T}$的映射，作者设计了一个包含两个Transformer层和一个位置编码层的TPT模块，其中Transformer层用于聚合形态信息，PPEG (Pyramid position encoding Generator)用于编码空间信息。所提出的基于MIL (TransMIL)的Transformer的概述如图3所示。

• 图3：每个WSI被裁剪成patch(背景被丢弃)，并被ResNet50嵌入到特征向量中。然后用TPT模块对序列进行处理:1)序列的平方;2)序列相关性建模;3)条件位置编码与局部信息融合;4)深度特征聚合;5) $\mathbb{T}\to \mathbb{Y}$的映射。

序列来自每个WSI中的特征嵌入。TPT模块的处理步骤如算法2所示，其中MSA表示多头自注意力，MLP表示多层感知机，LN表示 Layer Norm。