总结
现有基于注意力的多实例学习(MIL)算法的问题:
- 现有的基于注意力的 MIL 算法在全幻灯片图像(WSI)分类任务中可能关注不相关的模式,如染色条件和组织形态,导致不可靠的补丁级预测和可解释性。
- 这些算法倾向于专注于显著实例,难以识别难以分类的实例(如肿瘤与正常细胞相似的区域),从而降低了在医学诊断中的准确性和可解释性。
提出新的方法——FocusMIL:
- 提出了一个名为 FocusMIL 的简单而有效的基于实例的 MIL 方法,旨在遵循标准的 MIL 假设并避免依赖注意力机制。
- 使用最大池化和前向摊销变分推理,使模型专注于肿瘤形态,而不是虚假相关性,从而提高模型的分类能力和可靠性。
理论分析和实验验证:
- 理论上分析了注意力机制在 WSI 分类中的缺陷,强调其不能保证基于肿瘤形态的准确预测。
- 实验评估表明,FocusMIL 在 Camelyon16 和 TCGA-NSCLC 基准测试中,在补丁级分类任务上显著优于现有基线方法,并且在幻灯片级别上具有竞争力。
- 提出并使用了一个标准 MIL 测试数据集来检验现有方法是否符合标准的 MIL 假设,结果显示基于注意力的算法在存在强虚假相关性时可能失败。
关键点和创新之处:
- 最大池化: 确保模型遵循MIL假设,即如果一个包(即一个幻灯片)中至少有一个实例(即一个补丁)是正的,那么整个包就被认为是正的。
- 前向摊销变分推理: 鼓励模型学习到一个具有良好生成能力的潜在表示,从而提高模型的泛化能力。
- 小批量梯度下降: 通过同时处理多个幻灯片,提高模型对不同类型肿瘤的识别能力。
Abstract
尽管基于注意力的多实例学习算法在幻灯片级全幻灯片图像 (whole slide image,WSI) 分类任务中取得了令人印象深刻的表现,但它们容易错误地关注不相关的模式,例如染色条件和组织形态,导致不正确的补丁级预测和不可靠的可解释性。此外,这些基于注意力的 MIL(Multi-Instance Learning) 算法往往专注于显着实例,难以识别难以分类的实例。在本文中,我们首先证明基于注意力的 WSI 分类方法不遵守标准的 MIL 假设。在标准 MIL 假设中,我们提出了一种基于最大池化和前向摊销变分推理的 WSI 分类 (FocusMIL) 的令人惊讶的简单而有效的基于实例的 MIL 方法。我们认为,使用变分推理合成标准 MIL 假设鼓励模型专注于肿瘤形态而不是虚假相关性。我们的实验评估表明,FocusMIL 在 Camelyon16 和 TCGA-NSCLC 基准上显着优于补丁级分类任务的基线。可视化结果表明,我们的方法还实现了更好的分类边界来识别硬实例,并减轻了包和标签之间的虚假相关性的影响。
1 Introduction
全幻灯片图像 (WSI) 分析已广泛应用于计算机辅助诊断和预后等临床应用(Lu et al. 2021; Yao et al. 2020; Chen et al. 2021; Li, Li, and Eliceiri 2021)。由于单个WSI通常包含数十亿个像素,获得补丁级注释是劳动密集型的(Lu et al. 2021),因此将WSI裁剪成更小的补丁并使用多实例学习(MIL)进行分类是一种标准做法。多实例学习 (MIL) 是一种弱监督学习范式,旨在从表示为实例(补丁)包(幻灯片)的样本中学习。由于计算机辅助诊断模型需要提供肿瘤组织的幻灯片级和补丁级预测,MIL特别适合。(Li、Li 和 Eliceiri 2021;Qu 等人 2022)。
大多数现有的基于 MIL 的 WSI 分类算法都专注于幻灯片级别的预测,并利用注意力 (Vaswani et al. 2017) 将补丁级表示聚合到幻灯片级别(Ilse、Tomczak 和 Welling 2018; Li、Li 和 Eliceiri 2021;Shao 等人。 2021; Qu et al. 2022a, 2023; Tang et al. 2023, 2024; Kum et al. 2024)。一般来说,这些算法利用注意力机制为包内的每个实例分配权重,并聚合它们的加权表示进行分类。然而,许多基于注意力的 MIL 方法偏离了标准的 MIL 假设,这可能会限制它们在医学诊断中的可解释性和可信度(Raff 和 Holt 2023)。
标准 MIL 假设指出,当且仅当其至少一个实例为正时,多实例包为正(Foulds 和 Frank 2010;Dietterich、Lathrop 和 Lozano-P ́erez 1997),这自然与 WSI 分类的目标一致:当且仅当幻灯片包含至少一个癌变补丁时,幻灯片应该被标记为癌变。不幸的是,基于注意力的 MIL 算法不会提出额外的约束来强制聚合注意力权重时的标准 MIL 假设。因此,这些模型有时可以将非因果或不相关的实例误解为有影响力的,从而导致医疗诊断中不太准确和可解释的预测。
图 1:(a) 基于注意力的 MIL 算法可能会产生不正确的补丁级注意力图,具有正确的幻灯片级预测。(b) 补丁生成过程及其标签的因果图。
WSI分类的另一个要求是,模型应该基于它们对与肿瘤细胞因果相关的因素的补丁级预测,而不是虚假相关性。图1(B)示出了patch x的生成过程、其真实标签y及其预测标签ˆy背后的因果结构。三组因素决定了补丁的特征:Zc,它因果影响细胞是癌变的,如细胞形态;Zs,它决定了风格元素,如斑块内细胞的定位、大小和压缩状态;ZB来源于整个切片,包括苏木精和伊红染色偏差和组织结构。尽管理想的模型应该基于 Zc 进行预测,但大多数当前模型不区分 Zc 和 ZB。
基于注意力的 MIL 方法不强制执行标准的多实例假设或patch级和slide级标签之间的因果关系。首先,如图 1(a) 所示,较高的注意力分数可以分配给与肿瘤无关的区域,显着降低了辅助临床诊断的可解释性(Lin et al. 2023)。其次,模型可能难以识别形态与正常细胞和恶性程度较低的肿瘤细胞(硬实例)(J̈ogi et al. 2012; Zhang et al. 2023; Tang et al. 2023),导致早期癌症患者的遗漏诊断。这是因为基于注意力的模型过度优先考虑显着实例(Tang et al. 2023; Zhang et al. 2023),并且倾向于在训练期间利用虚假相关性来拟合更具挑战性的幻灯片。
本文研究了弱监督WSI分类问题,而不依赖于注意机制,提出了一种简单而有效的基于实例的MIL方法:FocusMIL。我们的方法侧重于因果因素 Zc 来预测实例的标签,而忽略非因果因素,例如风格元素 Z 和包继承的信息 ZB。FocusMIL 包括一个前馈神经网络、重新参数化操作和最大池化。我们强调,尊重多实例假设固有地限制了模型根据肿瘤形态 Zc 进行分类。我们的方法与传统的基于最大池的 MIL 分类方法的主要区别在于,我们使用批量大小大于 1 的 mini-batch 梯度,而不是批量大小为 1 的 SGD。我们讨论了更大的批量大小如何使模型能够构建更好的分类边界来识别硬实例,并有助于 Zc 和 ZB 的更稳定和高效解开。
Camelyon16 和 TCGA-NSCLC 数据集的结果表明,我们的方法在补丁级分类任务中显着优于现有基线,并且在幻灯片级别上与基线相当。使用幻灯片级监督,我们的补丁级预测比现有方法更准确地预测显着和困难实例。此外,我们设计了一个半合成数据集来测试主流 WSI 分类算法是否尊重标准 MIL 假设,结果表明所有bag-level MIL 方法都失败了这个测试。代码可以从补充材料访问,这将在发布后公开。
图2:不同肿瘤切片和正常斑块的肿瘤斑块的比较。来自癌细胞小癌区域的载玻片的肿瘤细胞与正常细胞非常相似。相反,来自具有大癌区域的载玻片的肿瘤细胞表现出非常不同的肿瘤特征。图像来自 Camelyon16 数据集。
2 Related Work
2.1 Attention-based MIL for WSI classification
大多数基于 MIL 的 WSI 分类算法根据它们的注意力分数将补丁级特征聚合为袋级特征。早期的基于袋子的方法 ABMIL(Ilse、Tomczak 和 Welling 2018)利用可训练网络计算实例级分类的实例特征的注意力分数,并通过它们的加权和获得袋级表示。随后的方法继承了类似的想法,并广泛整合了各种注意力机制(Li、Li 和 Eliceiri 2021;Lu 等人 2021;Shao 等人 2021;Zhang 等人 2022a)。
基于注意力的 MIL 方法面临三个来源的过度拟合挑战(Zhang et al. 2022a; Lin et al. 2023; Tang et al. 2023; Zhang et al. 2023)。首先,许多癌变载玻片包含显着实例(salient instances),即具有高分化癌细胞的区域与正常细胞显着不同(Qu et al. 2024; J̈ogi et al. 2012)。其次,其他幻灯片可能有硬实例(hard instances),例如癌细胞与正常细胞非常相似的小癌区域(Bejnordi et al. 2017)。第三,数字扫描设备、载玻片制备和染色过程经常将染色偏差引入WSI数据集(Zhang et al. 2022c)。现有研究(Tang et al. 2023;Zhang et al. 2023;Kum et al. 2024)表明,基于注意力的 WSI 分类方法可能会过度关注显着实例,无法识别硬实例。对于仅包含硬实例的阳性幻灯片,这些方法通常依赖于染色偏差引起的虚假相关性来拟合训练集(Lin et al. 2023)。最近的方法,如WENO (Qu et al. 2022)、MHIM (Tang et al. 2023)、IBMIL (Lin et al. 2023)、ACMIL (Zhang et al. 2023)和Slot-MIL (Kum et al. 2024)试图缓解这些问题;然而,它们的泛化能力和补丁级分类性能仍然不理想。
图 3:我们的 FocusMIL 概述,其中 2 个幻灯片在一个小批量中进行说明。使用预训练的特征提取器从补丁中提取特征。实例特征通过变分编码器编码为潜在表示。分类器 P (Y |z) 使用推断的潜在表示 z 来获得补丁级预测分数。最大池化应用于每个 MIL bag 的实例预测,以得出幻灯片级别的预测。
2.2 Reproducibility in Multiple Instance Learning
最近的研究(Raff和Holt 2023)表明,许多广泛使用的MIL方法,如MI-Net (Wang et al. 2018)和TransMIL (Shao et al. 2021),不遵守标准的MIL假设,并可能利用归纳偏差,如没有某些实例,作为预测正包的信号。由于染色偏差在WSI分析中很普遍(Lin et al. 2023;Zhang et al. 2022c),不尊重标准MIL假设的模型可以学习使用这些虚假相关性(Raff和Holt 2023)。mi-Net (Zhou and Zhang 2002; Wang et al. 2018) 和最近提出的 CausalMIL (Zhang et al. 2022b) 可能是尊重 MIL 假设的唯一两种方法。CausalMIL 将变分自动编码器 (Kingma and Welling 2013) 应用于 MIL 问题,使用以包信息为条件的非分解先验分布来提供识别潜在因果表示的理论保证。然而,mi-Net 在 WSI 分类中表现不佳,CausalMIL 尚未针对 WSI 分类进行研究。
3 Method
3.1 Preliminaries
具有 n 个幻灯片的 WSI 数据集被视为 MIL bagB = {B1, · · · , Bi, · · · , Bn},其中每个幻灯片在训练期间都有一个幻灯片级别的标签 Yi ∈ {0, 1}。然后将每张幻灯片裁剪成与 MIL 实例 Bi ={pi1, · · · , pij , · · · , pini } 对应的补丁,其中 ni 是幻灯片中的补丁数量。对于每个补丁pij,存在一个对学习者未知的补丁级标签yij∈{0,1}。标准的多实例假设表明当且仅当其至少一个实例为正时,包为正:
这相当于对包内的实例使用最大池化:
Yi = max {yi,j }.
给定一个特征提取器f,每个实例被投影到d维特征向量xij = f (pij)∈Rd×1上。
3.2 FocusMIL
我们的主要目标是在保持标准多实例学习假设的同时,从幻灯片级监督训练补丁级分类器,并专注于跨幻灯片不变的因果因素Zc。形式上,我们假设
假设1确定细胞是否从斑块中的因果因素Zc癌变的机制在载玻片上是不变的
对应于假设 1,Zs 在不同的补丁之间变化,而 ZB 在不同的幻灯片之间变化。
我们提出的FocusMIL方法封装了弱监督变分推理中WSI分类的两个要求,如图3所示。首先,将原始输入补丁p使用预先训练的特征提取器处理为x。然后,编码器将 x 编码为低维潜在变量 z,用于预测补丁标签。预测的补丁标签与最大池化操作合并,以符合标准的 MIL 假设。值得注意的是,最大池化操作和损失函数的计算每一张幻灯片执行一次,并对编码器和分类器进行了优化,以小批量中多个包的总损失。这与大多数基于注意力的 WSI 算法根本不同,其中优化仅限于每个 epoch 的一张幻灯片。
标准MIL假设的重要性。关于 MIL 假设的模型仅基于正概念来预测袋子标签,例如肿瘤细胞的存在(Foulds 和 Frank 2010;Dietterich、Lathrop 和 Lozano-P ́erez 1997)。关于 MIL 假设 mi-Net 和 CausalMIL 的两个现有模型都使用与 MIL 问题设置对齐的最大池化。我们认为最大池化有效地将模型集中在因果因素 Zc 上以预测实例标签。由于最大池化的约束,模型必须仅根据包中最正的实例对包标签预测进行基础。如果分类器是使用 ZB 来确定实例标签,因为袋继承的信息(例如染色条件)在正载玻片上并不完全一致,并且在负载玻片上不同,模型将完全无法完全拟合训练集。因此,遵循 MIL 假设的模型必须学习肿瘤细胞 Zc 的形态结构等特征来拟合训练集。一个显着的优点是尊重在训练集上拟合良好的 MIL 假设的模型将获得更好的补丁级分类结果。
用于WSI分类的小批量梯度下降。对于WSI分类,由于患者的癌症恶性程度不同,肿瘤细胞的形态在载玻片上可能不同(Zhang et al. 2023;Qu et al. 2024;J̈ogi et al. 2012)。如图 2 所示,一些幻灯片包含大量肿瘤区域,与正常细胞的形态差异显着。相比之下,其他幻灯片具有较小的肿瘤区域,其中肿瘤细胞的分化更高,与正常细胞的差异不太明显(Qu et al. 2024),在一些文献中被称为“硬实例”(Tang et al. 2023)。当多个包在小批量中时,编码器和分类器将同时前向传播多个包。为了确保模型能够准确地预测批次中不同程度的癌变恶性肿瘤的幻灯片,优化过程必须考虑显着实例和硬实例的特征,从而学习更好的分类边界。考虑以下场景:如果模型仅基于显着实例建立分类边界,则无法准确分类较低恶性肿瘤的幻灯片,从而导致高损失。
Forward Armotized Variational Inference。仅最大池化和批量优化不足以进行 WSI 分类。从经典mi-Net的批处理版本的性能可以明显看出这一点(Wang et al. 2018)。mi-Net 经常正确分类袋子,但在训练期间无法对其补丁进行分类。此外,包和补丁分类都经常在测试集上失败。
我们认为这是因为 mi-Net 由于神经网络的灵活性,可能会学习过度拟合训练数据的极其复杂的决策表面(Raff 和 Holt 2023),并且不能有效地学习潜在的因果表示。
我们采用具有前向 KL 发散 z 的变分推理来解决这个问题。具体来说
其中k(X)是patches的重采样分布(Ambrogioni等人)。2019年)。由于第一项仅取决于变分分布,因此优化前向 KL 散度等效于优化
等式 4 中梯度的计算很简单,因为期望是针对 p 定义的,而梯度是针对 q 进行的。
为简单起见,我们将先验分布设置为 p(z) =N (0, I),并将后验限制为分解的高斯编码器。最后,优化目标可以得到如下:
其中 β 是平衡分类器和变分损失的超参数,z∗ 对应于使用最大池化获得的袋子相对于 p(Y |z∗) 的最正潜在表示,x∗ 对应于 z∗ 的原始特征。与 CausalMIL (Zhang et al. 2022b) 的一个值得注意的根本区别是,前向 KL 的使用完全消除了优化解码器 p(x|z) 的需要,这对于 WSI 分类是不必要的。
4 Experiments
4.1 Datasets, Metrics, and Baselines
Camelyon16数据集广泛用于乳腺癌的转移检测(Bejnordi et al. 2017)。该数据集由 270 个训练和 129 个测试 WSI 组成,在 10 倍放大(多分辨率金字塔中的第二级)下产生大约 270 万个补丁。像素级标签可用于肿瘤切片。每个WSI都被裁剪成256 × 256的补丁,没有重叠,背景补丁被丢弃。如果补丁包含 25% 或更多癌症区域,则该补丁被标记为阳性。由于Camelyon16训练集和测试集中正滑动的正斑块比分别约为8.8%和12.7%,肿瘤斑块与正常斑块的数量不平衡
TCGA-NSCLC 包括两种肺癌亚型:肺腺癌 (LUAD) 和肺鳞癌 (LUSC),共有 1,054 个诊断数字载玻片。该数据集只有幻灯片级别的标签可用。我们直接使用 Li 等人发布的补丁。 (Li、Li 和 Eliceiri 2021)。
为了测试现有的WSI分类方法是否符合标准MIL假设,我们提出了Camelyon16标准MIL测试数据集,灵感来自(Raff和Holt 2023)中的想法。具体来说,在训练集中,我们通过在正常玻片中随机选择20%的斑块并增加其绿色通道的强度来引入毒性。在测试集中,我们随机选择肿瘤载玻片中20%的斑块以同样的方式引入毒素。MIL模型不能合法地学习使用有毒信号,因为它只发生在正常的幻灯片中(Raff和Holt 2023)。如果一个模型的训练AUC > 0.5,但测试AUC < 0.5,则它依赖毒素进行预测,不尊重MIL假设。
由于补丁级别的高度类不平衡,我们报告精确度召回曲线下的面积(AUCPR)和F1分数进行评估。对于幻灯片级别的分类,我们报告AUC和准确率,因为肿瘤和正常类别是平衡的。
我们将我们的方法与最近发表的基线进行了比较,包括ABMIL (Ilse, Tomczak, and Welling 2018)、DSMIL (Li, Li, and Eliceiri 2021)、TransMIL (Shao等人2021)、DFTD-MIL (Zhang等人2022a)、IBMIL (Lin等人2023)和CausalMIL (Zhang等人2022b)。前五种方法是基于注意机制的,而CausalMIL是一种基于实例的方法,据我们所知,这种方法以前还没有应用于WSI分类。对于Camelyon16、TCGA-NSCLC和Camelyon16 Standard-MIL测试数据集,我们使用在ImageNet (Deng et al. 2009)上预训练的ResNet-18 (He et al. 2016)提取特征。此外,我们还利用强大的CTransPath (Wang et al. 2022)对Camelyon16数据集提取的特征,对各种MIL方法进行了更全面的比较。
4.2 Implementation details
FocusMIL的编码器由一个具有一个隐藏层的神经网络组成,以ReLU作为激活函数。我们使用初始学习率为1 ×10−4的AdamW优化器在训练过程中更新模型权值。潜在因子的维度设置为35。训练FocusMIL模型的小批量大小为3。对于滑动级预测,我们分别应用max-pooling和mean-pooling对Camelyon16和TCGA-NSCLC数据集的实例预测分数进行汇总。对于基于注意力的方法,我们根据其官方发布的代码构建了其他模型,并对关键超参数进行网格搜索。对于CausalMIL (Zhang et al. 2022b),其编码器和解码器都被设置为具有隐藏层神经元128的神经网络,另外两个完全连接的层用于根据包信息划分因式先验分布。所有实验均在Nvidia RTX4090 GPU上进行。更多的细节,请参考我们的代码在补充。
4.3 Results on Real-World Dataset
表1:使用Resnet-18和CTransPath特征的MIL方法的性能比较。下标是相应的95%置信区间。最优秀的表现用粗体突出显示,第二好的表现用下划线表示。
表2:在TCGA数据集上的性能比较
表3:Camelyon16 Standard-MIL测试数据集的性能。
表4:不同MIL方法的模型大小和每历元训练时间的比较。每个模型的输入维度设置相同,以确保比较的公平性。
4.7 Visualization
为了比较的公平性,我们没有刻意选择用于可视化的模型,训练批大小为3的FocusMIL和DTFD-MIL模型的测试幻灯片AUC都达到了约0.87。如图5所示,我们的模型准确地识别了大小肿瘤区域,没有任何假阳性预测。训练批大小为1的FocusMIL遗漏了一些小的肿瘤区域。对于DTFD-MIL,它错过了第一张幻灯片中的肿瘤区域和第二张幻灯片中的一些肿瘤斑块,并且存在许多假阳性预测。值得注意的是,在第四张幻灯片中,DTFD-MIL给许多与肿瘤无关的斑块打了高分。放大后,这些区域的组织明显模糊,这是由载玻片制备过程造成的。这证实了DTFD-MIL可以依靠袋遗传信息进行预测。我们在补充材料中提供了额外的可视化结果。
图5:FocusMIL和DTFD-MIL在Camelyon16数据集上的可视化结果。我们的(3)和我们的(1)分别是指批大小为3和1的FocusMIL模型。第四张幻灯片是负面的。对于FocusMIL,分类器置信度分数用于补丁强度。对于DTFD-MIL,使用patch概率派生,阈值为0.5001。
5 Conclusion
本文介绍了FocusMIL,这是一种尊重标准多实例学习假设并结合变分推理以确保模型专注于肿瘤形态的方法。我们进一步利用小批量梯度下降来获得更好的分类边界来识别硬实例。定量实验表明,我们的方法在补丁级分类中显著优于基于注意力的方法,并且在幻灯片级与当前最先进的方法相比,取得了具有竞争力的结果。可视化结果证明了该方法在肿瘤区域定位中的优越性。
本文提出的理论分析表明,基于注意力的方法不能保证基于肿瘤细胞形态(因果因素)的预测。我们的定量实验和可视化也表明,这些方法不能提供可靠的肿瘤区域预测。更糟糕的是,Camelyon16标准MIL测试数据集的结果表明,当袋子上下文和标签之间存在很强的虚假相关性时,基于注意力的方法可能会失败。因此,我们强调在WSI分类中应谨慎使用基于注意的MIL方法。
FocusMIL是一个更可靠的替代方案,它侧重于预测的因果因素。我们建议从业者从FocusMIL开始作为一个坚实的基础,以避免部署中的过度风险。此外,研究人员可以基于我们的FocusMIL设计更复杂的WSI分类方法,以满足MIL假设。