分享自:

基于Transformer的关联多示例学习用于全切片图像分类

期刊:35th conference on neural information processing systems (NeurIPS 2021)

以下是一篇基于您所提供文本生成的综合性学术研究报告。

基于Transformer相关多示例学习的全切片图像分类研究

本文向您介绍一项在计算病理学领域取得显著进展的研究。该研究由Zhuchen Shao, Hao Bian, Yang Chen, Yifeng Wang, Jian Zhang, Xiangyang Ji以及通讯作者Yongbing Zhang共同完成。研究团队主要来自清华大学深圳国际研究生院、哈尔滨工业大学(深圳)、北京大学以及清华大学的自动化系。这项题为“TransMIL: Transformer based correlated multiple instance learning for whole slide image classification”的成果,发表于第35届神经信息处理系统大会(NeurIPS 2021)

一、研究背景与目的 该研究聚焦于计算病理学领域,特别是全切片图像(Whole Slide Image, WSI)弱监督分类问题。WSI是数字化病理活检的产物,能够将组织切片转化为千兆像素级别的超高分辨率图像,完整保留了原始的组织结构信息。这为深度学习在病理诊断中的应用带来了巨大机遇,但同时也带来了两个核心挑战:首先,WSI尺寸极大,无法直接输入常规的深度学习模型;其次,获取像素级或区域级的精确标注(即标注出图像中癌变的具体位置)非常耗时费力且成本高昂,这在临床实践中往往是不可行的。

为了解决这些问题,研究者们通常采用多示例学习(Multiple Instance Learning, MIL) 范式。在MIL框架下,一张WSI被视为一个“包”,其中包含从图像中裁剪出的成千上万个“补丁”(即“示例”)。训练时,我们只有包级别的标签(例如,整张WSI是否为癌变),而没有每个补丁的标签。一个经典的MIL假设是:如果包是阳性的,那么其中至少存在一个阳性示例;如果包是阴性的,那么所有示例都是阴性的。

然而,本研究指出,当前大多数基于深度学习的MIL方法都建立在示例之间独立同分布(i.i.d.) 的假设之上。这意味着模型在处理包内的补丁时,认为它们彼此之间是相互独立的。但病理学家在实际诊断时,不仅会观察单个区域的形态特征,还会综合考虑不同区域之间的空间上下文信息相关性(例如,肿瘤细胞的排列方式、与周围正常组织的对比等)。i.i.d.假设忽略了这种重要的相关性信息,可能导致模型无法学习到更全面、更具判别性的特征,从而限制了分类性能的进一步提升。因此,本研究的目标是打破传统MIL的i.i.d.假设,提出一个新的相关MIL(correlated MIL)框架,并在此基础上设计一个高效的、能够显式建模示例间相关性的深度学习模型,以实现更准确、更可解释的WSI分类。

二、研究方法与流程 本研究提出了一种名为TransMIL的新型框架,其核心思想是利用Transformer架构强大的序列建模和长距离依赖捕获能力,来学习WSI中不同补丁(示例)之间的相关性与空间结构信息。

  1. 理论基础与框架提出:研究首先从理论层面形式化地定义了“相关MIL”问题,并提供了定理证明来支持其框架的合理性。定理1及其推论表明,一个满足特定条件的连续集合函数(即评分函数)可以被一个特定形式的函数任意逼近,该形式包含了用于聚合示例间信息的“池化矩阵P”。关键创新在于,传统MIL方法(如最大池化、均值池化或注意力池化)中的P矩阵是一个对角矩阵,仅关注每个示例自身的权重,而忽略了示例之间的关系。而在相关MIL框架中,P矩阵允许非对角线元素非零,从而能够显式地编码任意两个示例之间的相关性。定理2进一步从信息论角度证明,考虑相关性的假设比i.i.d.假设具有更小的信息熵,这意味着前者可能包含更多有用信息,减少不确定性。基于此,研究者提出了一个通用的三步骤算法:第一步,通过函数f和h提取所有示例的形态信息和空间信息;第二步,通过一个能建模相关性的池化矩阵P聚合这些信息;第三步,通过函数g将聚合后的信息映射为包级别的预测标签。

  2. 模型架构设计:TransMIL是该框架的具体实现,其整体流程如下图所示(基于论文图3):

    • 数据预处理:首先,将每张WSI裁剪成256x256像素的非重叠图像块,并过滤掉背景区域(饱和度低于15)。然后,使用在ImageNet上预训练的ResNet-50模型,将每个图像块编码为一个1024维的特征向量。这样,一张WSI就被转化为一个长度为N(N为补丁数量,可变)的序列,每个元素是一个512维的特征嵌入(通过一个全连接层将维度从1024降至512)。
    • TPT模块:这是TransMIL的核心模块,负责处理长序列并建模相关性。其具体步骤包括:
      • 序列正方形化:为了便于后续的二维空间位置编码,将一维序列的长度N调整为一个近似正方形(边长为⌈√N⌉)的二维结构,不足部分用额外的补丁特征填充。
      • 序列相关性建模:这是第一个Transformer层。为了应对WSI序列通常极长(成千上万个补丁)带来的计算复杂度O(N²)问题,研究采用了Nyström方法来近似计算自注意力(Self-Attention),将复杂度降至O(N)。这使得Transformer能够高效地处理长序列,并捕获所有补丁对之间的全局相关性。
      • 条件位置编码与局部信息融合:这是金字塔位置编码生成器(PPEG) 模块。由于WSI的尺寸和形状各异,序列长度可变,传统的固定位置编码(如正弦编码)不适用。PPEG模块将经过相关性建模的补丁序列重塑回二维“图像”空间,然后使用一组不同尺寸的卷积核(如3x3, 5x5, 7x7) 对其进行卷积操作。这些卷积操作有两个关键作用:一是隐式地编码了补丁之间的相对空间位置信息(“条件位置编码”);二是像CNN一样,聚合了局部邻域内的上下文信息,丰富了每个补丁的特征表示。
      • 深度特征聚合:这是第二个Transformer层,它对经过PPEG增强后的特征序列进行进一步的特征整合与提炼。
    • 分类输出:最终,序列中的“类别令牌”([class] token)的输出特征被输入到一个多层感知机(MLP)中,得到整个WSI(包)的分类预测结果。
  3. 实验设计与评估

    • 数据集:研究在三个公开的病理图像数据集上进行了广泛的实验,以验证模型的通用性和鲁棒性:
      • Camelyon16:用于乳腺癌淋巴结转移检测的二进制分类任务(阳性/阴性)。包含270张训练WSI和130张测试WSI。阳性切片中肿瘤区域占比通常很小(<10%),属于极具挑战性的弱阳性场景。
      • TCGA-NSCLC:用于非小细胞肺癌亚型分类的二进制任务,区分肺腺癌(LUAD)和肺鳞癌(LUSC)。包含993张WSI。此数据集中肿瘤区域通常较大(>80%)。
      • TCGA-RCC:用于肾细胞癌亚型分类的多类别任务(三分类),包含三种亚型:肾嫌色细胞癌(KICH)、肾透明细胞癌(KIRC)和肾乳头状细胞癌(KIRP)。包含884张WSI,且数据分布不平衡。
    • 基线模型:研究比较了多种先进的MIL方法,包括传统池化方法(均值池化、最大池化)、基于注意力的方法(ABMIL, PT-MTA, CLAM-SB, CLAM-MB)、基于非局部注意力的方法(DSMIL)以及基于循环神经网络的方法(MIL-RNN)。
    • 评估指标与实现细节:主要评估指标为分类准确率(Accuracy)曲线下面积(AUC)。对于TCGA数据集,采用4折交叉验证。模型使用交叉熵损失和Lookahead优化器进行训练,学习率为2e-4。所有实验在一张RTX 3090 GPU上完成。

三、主要研究结果 1. WSI分类性能:TransMIL在三个数据集上的分类性能均超越了所有基线模型,具体数据如下: * Camelyon16:在肿瘤区域占比很小的困难任务中,TransMIL取得了93.09% 的AUC和88.37% 的准确率,显著优于其他方法。例如,其AUC比次优的注意力基线方法(如CLAM-SB)高出约5个百分点,证明了建模示例间相关性在识别稀疏阳性区域时的巨大优势。 * TCGA-NSCLC:TransMIL达到了96.03% 的AUC和88.35% 的准确率,AUC比第二名高出1.40%,准确率高出2.16%,再次验证了其有效性。 * TCGA-RCC:在不平衡的多分类任务中,TransMIL同样表现最佳,取得了98.82% 的宏观平均AUC和94.66% 的准确率,证明了其对不平衡数据和多分类问题的良好适应性。

  1. 消融实验分析:为了验证模型中关键组件的贡献,研究进行了详细的消融实验。

    • PPEG模块的作用:实验比较了无位置编码、使用正弦编码(Sinusoidal Encoding)以及使用不同大小卷积核的PPEG模块的效果。结果显示,任何形式的位置编码都能提升模型性能,而PPEG模块带来的提升最为显著。特别是,使用多尺寸卷积核(3x3, 5x5, 7x7)的PPEG性能最好,因为它能够进行多粒度的位置编码和上下文信息融合。
    • 条件位置编码的有效性:通过打乱输入补丁序列的顺序进行实验,发现当输入序列是随机无序的(w/o)时,模型性能会下降。而使用原始空间顺序(Order)时,PPEG能带来最显著的性能提升(在Camelyon16上AUC提升0.9%,在TCGA-NSCLC上提升0.61%),这直接证明了PPEG所编码的空间位置信息对于病理诊断至关重要
  2. 可解释性与可视化:TransMIL具有良好的可解释性。研究者将模型自注意力机制生成的注意力分数可视化为热力图,叠加在原始的WSI上。结果显示,模型所关注的高响应区域(热力图红色区域)与病理学家提供的精细标注的癌症区域高度重合。这表明TransMIL不仅能做出准确预测,还能自动定位出对诊断决策最重要的组织区域,其决策过程对临床医生而言是透明且可理解的。

  3. 快速收敛性:相较于ABMIL、DSMIL、CLAM等其他先进的MIL方法,TransMIL的训练收敛速度快了大约2到3倍。这意味着TransMIL能够用更少的训练周期达到相同甚至更好的性能,这得益于Transformer架构和相关性建模对信息的高效利用。

四、研究结论与价值 本研究提出并验证了一个新颖的相关多示例学习(correlated MIL)框架,并基于此框架开发了TransMIL模型。该模型通过引入Transformer的自注意力机制,成功建模了全切片图像中不同组织补丁之间的形态学相关性和空间上下文信息,有效克服了传统MIL方法中i.i.d.假设的局限性。

研究的科学价值在于:它从理论和算法两个层面推进了弱监督学习在计算病理学中的应用。定理证明为相关MIL提供了数学基础,而TransMIL的具体实现则为处理超长序列、整合空间信息提供了创新的技术方案(如Nyström近似自注意力、PPEG模块)。

应用价值非常显著:TransMIL在多个公开数据集上的分类性能达到了新的最先进水平,尤其在肿瘤区域稀疏的困难场景下优势明显。同时,其快速的收敛速度和强大的可视化能力,使得它更易于在实际的科研和临床辅助诊断环境中部署和应用。该模型能够处理不平衡数据和多分类问题,通用性强。

五、研究亮点 1. 理论创新:首次明确提出了“相关MIL”的框架,并提供了严格的收敛性证明和信息论优势分析,为后续研究奠定了理论基础。 2. 方法创新:首次将Transformer架构系统性地应用于WSI分类任务,并巧妙设计了TPT模块PPEG模块,解决了Transformer处理超长序列的计算难题以及WSI序列长度可变、需编码空间信息的技术挑战。 3. 性能卓越:在三个具有不同特性的公开数据集上,全面超越了现有最先进的MIL方法,尤其在具有挑战性的弱阳性检测任务上提升显著。 4. 实用性强:模型不仅精度高,而且收敛速度快,并具有优异的可解释性,能够生成与病理学家标注高度一致的注意力热力图,这对于建立临床信任至关重要。

六、其他有价值内容 文章最后展望了TransMIL的潜在应用方向,例如在需要更多相关性信息的生存分析癌细胞扩散检测等问题上可能具有更大潜力。同时,作者也指出了当前工作的一个局限:所有实验均在20倍放大倍率下进行,更高倍率的WSI会产生更长的序列,对计算和内存提出更大挑战,这将是后续工作的一个研究方向。研究代码已在GitHub上开源,促进了领域的可复现性和进一步发展。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com