0. TL; DR

这篇论文提出了TransMIL的框架，首次将Transformer架构引入多实例学习（MIL），以端到端的方式同时建模实例的形态和空间相关性。TransMIL将一个全切片图像中的所有实例（patches）视为一个序列，在2D空间使用多尺度卷积来进行条件位置编码，并采用了Nyström方法来近似自注意力计算。

在CAMELYON16（二分类）、TCGA-NSCLC（二分类）和TCGA-RCC（多分类）三个极具挑战性的公开WSI数据集上，TransMIL的性能全面超越了所有基于CNN和注意力机制方法。

1. 背景介绍

在利用多实例学习（MIL）处理全切片病理图像（WSI）时，一个普遍的做法是将WSI看作一个“包”，将其中的图块（patches）看作“实例”。然而，几乎所有现有的深度MIL方法都共享一个隐含的、但与现实严重不符的假设：包内的所有实例都是独立同分布的（i.i.d.）。

Transformer架构凭借其核心组件自注意力机制（Self-Attention）在自然语言处理和计算机视觉领域取得了成功。自注意力能够建模一个序列中任意两个元素之间的长距离依赖关系。既然WSI中的所有patches可以被看作一个序列，自注意力机制可以被用来计算任意两个patch之间的相关性。

然而，将Transformer直接应用于WSI也面临巨大挑战：

超长序列：一张WSI包含数千甚至上万个patches，远超传统Transformer的处理能力（通常小于1000），会导致$O(n²)$的计算和内存复杂度爆炸。
空间信息缺失: 标准的Transformer是为一维文本序列设计的，它本身不包含2D空间信息。如何将patches的空间坐标信息有效地融入模型，是一个亟待解决的问题。

2, TransMIL 架构

作者证明了任何一个连续的、对实例顺序不敏感的包级别评分函数，都可以被分解为一个三步过程：实例转换 -> 聚合 -> 包级别转换。进一步从信息熵的角度证明，考虑了实例间相关性的模型，其信息熵$H(Θ_1, …, Θ_n)$小于或等于i.i.d.假设下的信息熵$ΣH(Θ_t)$。

\[H(\Theta_1, \dots, \Theta_n) \le \sum_{t=1}^n H(\Theta_t)\]

这意味着引入相关性可以减少系统的不确定性，为模型提供更多有效信息。

TransMIL的核心是一个专门为WSI设计的TPT模块（Transformer-Position encoding-Transformer）。TPT模块是一个包含两个Transformer层和一个PPEG位置编码层的结构。

首先用一个预训练的ResNet50为WSI中的每个patch提取一个特征向量$h_i$。将所有实例的特征向量拼接成一个序列$H_i$。在序列的开头加入一个可学习的[class] token，它将最终用于代表整个WSI的全局表示。

将准备好的序列送入第一个Transformer层。这一层的多头自注意力（Multi-head Self-Attention, MSA）会计算序列中任意两个实例特征之间的相似性，捕捉它们在形态学上的相关性。为了解决WSI序列超长的问题，TransMIL在计算softmax$(QK^T)$时，采用了Nyström方法进行近似，将计算复杂度从$O(n²)$成功降至$O(n)$。

传统的绝对位置编码（如sin/cos编码）不适合实例数量可变、无固定方向的WSI，TransMIL使用PPEG层向模型注入空间信息。将一维的实例序列根据其原始的2D坐标，重新排列成一个二维的特征图。使用多个不同大小的卷积核（如3x3, 5x5, 7x7）对这个特征图进行卷积。将多尺度卷积得到的特征图与原始特征图相加，进行信息融合。将融合后的特征图重新展平为一维序列，并与之前分离的[class] token拼接回来。

将经过PPEG注入了空间信息的序列送入第二个Transformer层。这一层的作用是在已经融合了形态学和空间信息的基础上，进行更深层次的特征交互和聚合。

最后取出第二个Transformer层输出序列中的[class] token的表示，将这个表示（它已经聚合了整个WSI的所有信息）送入一个简单的MLP分类器，得到最终的WSI级别预测。

3. 实验分析

作者在三个大规模、公开的WSI数据集上，对TransMIL的性能进行了全面的评估。在极具挑战性的CAMELYON16（阳性实例极少）上，TransMIL的AUC比强基线ABMIL高出超过5个百分点。这表明当关键信息稀疏时，建模实例间的相关性（如肿瘤细胞聚集的趋势）能够提供至关重要的额外线索。在TCGA数据集上，TransMIL同样表现出色，证明了其方法的普适性。

作者进行了消融研究：

PPEG: 与不使用任何位置编码的模型或使用传统sin-cos绝对位置编码的模型相比，PPEG的性能更优。PPEG模块是TransMIL成功的关键，它所提供的多尺度、条件化的空间信息，远比简单的绝对位置编码更有效。
条件位置编码: 作者通过打乱输入序列的顺序来破坏空间连续性，发现模型的性能出现了下降。这证明了模型确实在利用PPEG提供的空间顺序信息，而不是简单地将位置编码当作一种无序的特征。

通过可视化TransMIL的注意力图谱，可以生成高质量的热力图。实验结果显示，模型高亮的区域与病理学家标注的肿瘤区域高度一致，证明了TransMIL在做出准确预测的同时，也具备良好的可解释性。

与ABMIL、DSMIL等方法相比，TransMIL的收敛速度要快2到3倍。在更少的训练轮数下，就能达到更高且更稳定的验证集AUC。更快的收敛速度可能得益于Transformer强大的全局信息整合能力。它能更高效地从整个WSI中学习判别性特征，而不是像传统注意力那样需要逐点、缓慢地发现关键区域。