本文主要是介绍论文阅读:MuRCL: Multi-Instance Reinforcement Contrastive Learning for Whole Slide Image Classification,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
@[TOC]MuRCL: Multi-Instance Reinforcement Contrastive Learning for Whole Slide Image Classification
论文介绍
这篇文章是发表于TMI的关于WSI的多实例分类的文章,作者是厦门大学的学生。
论文地址:https://ieeexplore.ieee.org/abstract/document/9975198
代码地址:https://github.com/wwu98934/MuRCL
摘要
多实例学习 (MIL) 被广泛用于自动整张数字切片图像 (WSI) 分析,它通常包括两个阶段,即实例特征提取和特征聚合。然而,由于WSI级标签的“弱监督”,特征聚合阶段在训练有效的 MIL 模型时会受到严重的过度拟合的影响。在这种情况下,从有限的WSI级数据中挖掘更多信息对于 WSI 分析至关重要。与之前关于改进实例特征提取的工作不同,本文研究了如何利用不同实例(补丁)之间的潜在关系来对抗 MIL 中的过度拟合以进行更通用的 WSI 分类。特别是,我们提出了一种新颖的多实例强化学习框架(MuRCL),以深入挖掘不同补丁的固有语义关系,以促进 WSI 分类。具体来说,所提出的框架首先以自监督的方式进行训练,然后使用 WSI 级别的标签进行微调。我们将第一阶段制定为对比学习 (CL) 过程,其中正/负判别特征集由 WSI 的相同补丁级特征包构建。为了便于 CL 训练,我们设计了一种新的基于强化学习的智能体(agent),根据WSI级特征聚合的在线奖励逐步更新判别特征集的选择。然后,我们进一步使用标记的 WSI 数据更新模型,以规范学习到的特征以进行最终的 WSI 分类。在三个公共WSI分类数据集(Camelyon16、TCGA-Lung和TCGA-Kidney)上的实验结果表明,所提出的MuRCL优于最先进的MIL模型。此外,MuRCL 在 TCGA-Esca 数据集上可以达到与其他最先进的 MIL 模型相当的性能。
引言
WSI是癌症诊断的黄金标准[1],[2]。WSI分析对疾病研究的开始和进展以及靶向治疗的发展也至关重要[3]。虽然由于WSI的高分辨率(例如,典型大小为40,000 × 40,000),自动化WSI分析是医学图像分析中长期存在的具有挑战性的问题。最近在计算机视觉和医学图像分析领域取得了一些成功,用于WSI分析[4],[5],[6]。其中,当只有WSI级标签可用[7]、[8]时,多实例学习(MIL)被广泛采用,其中WSI被认为是一个包含许多图像补丁实例的包。在典型的基于 MIL 的 WSI 分析管道中,提取补丁级特征、聚合和分类以生成WSI级别的预测。
最近基于 MIL 的方法从深度神经网络中受益匪浅,用于使用WSI级监督 [8]、[9] 进行特征提取和聚合。然而,在开发用于准确 WSI 分析的稳健模型时仍然存在一些挑战。首先,如上所述,大多数现有的 MIL 方法依赖于弱监督的WSI级标签来训练整个框架。由于整个WSI图像数量有限,这些模型容易遭受过拟合,由于监督信号较弱[10],[11],无法学习丰富的表示。其次,由于端到端 MIL 训练过程对于 WSI 等大型特征包来说非常昂贵,因此当前模型仅对几个高分实例特征 [4]、[10]、[11]、[12] 进行采样以进行WSI级预测以降低计算成本。然而,这些模型会导致WSI分类的次优解决方案,因为这种学习方案没有考虑补丁之间的语义关系来进行判别特征选择。如何选择判别补丁特征来训练 MIL 特征聚合器对于准确的 WSI 分类至关重要,尤其是在有限的训练数据场景下。
为了应对上述挑战,我们旨在研究一种有效的学习范式来探索WSI中不同实例的有价值的语义关系,以实现可泛化的WSI分析,其中也可以自动选择判别实例特征来提高框架的性能。如图1所示,一个普通的MIL过程(图1(a)))可以被视为三分量范式:实例提取、实例选择和实例聚合。WSI 特征包 x \mathcal{x} x 是从给定的 WSI 中提取的,然后从 x \mathcal{x} x 中选择一个子集实例特征 x ~ \tilde{x} x~ 进行特征聚合以生成WSI级预测 p p p。为了挖掘WSI实例固有的语义关系,我们提出了一种新的多实例对比学习框架(图1(b)),其中来自同一WSI的两个判别特征集,然后训练该框架以最大化这两个集合在自监督阶段之间的一致性,然后进行微调过程进行最终的WSI预测。请注意,我们的方法与之前的实例级自监督方法[12]不同,该方法专门用于以自监督的方式训练特征提取网络,而我们专注于如何自监督地利用不同实例之间的关系,在提取的实例特征之上推进WSI级表示。
在本文中,我们提出了一种新颖的多实例强化学习 (MuRCL) 框架来推进 WSI 分类。我们提出的框架首先以自监督的方式进行训练,以利用不同实例的固有语义关系,然后对其进行微调以进行最终预测。我们将自监督阶段制定为 CL 过程,其中正/负判别特征集由 WSI 的实例级特征包构建。为了从WSI中选择信息丰富的实例特征,我们进一步设计了一种新的强化学习(RL)智能体来指导区分特征集的构建。特别是,给定一个 WSI 实例级特征包,我们维护一个由不同初始状态触发的智能体,分别从它构建两个判别特征集以形成 CL 训练中的正对,其中智能体将根据在线奖励逐步更新训练迭代中判别特征集的选择。我们使用输入正特征集之间的余弦相似度作为奖励,旨在引导智能体检索不同的特征,从而使所选特征集对 CL 更具信息性。我们进一步使用 WSI 标签微调框架以进行最终的 WSI 分类。值得注意的是,我们不仅将不同的现有模块组合成一个框架。相反,我们提出了一种新颖而全面的解决方案,通过将 RL 和 CL 中的优化策略无缝地适应 WSI 分析中的一个重要问题,即自监督 MIL 聚合。此外,如何设计同一样本的不同信息视图来实现基于集合的CL并非易事,我们首次通过训练RL模块提出不同的视图来最大化设计的奖励,即CL目标。
我们的主要贡献可以总结如下:
- 我们提出了一种新的多实例对比学习框架,通过挖掘不同实例的固有语义关系来推进WSI分类。
- 我们提出了一种新的基于集合的正对构造解决方案,以实现自监督多实例对比学习。
- 我们提出了一种新的基于 RL 的策略来从 WSI 中选择判别集。设计了一个智能体来动态细化自监督 MIL 对比学习的集合选择,并通过强化学习进行训练。
- 我们已经在四个基准 WSI 数据集上验证了我们的框架。我们的方法在 Camelyon16、TCGA-Lung 和 TCGA-Kidney 上优于以前的 SOTA 方法,并在 TCGA-Esca 上实现了与 SOTA 方法相当的性能。代码可在 https://github.com/wu98934/MuRCL 获得。
相关工作
多实例WSI分类
MIL广泛应用于可以观察到多个实例但只给出一个通用类别标签的各种应用中,包括 WSI 分析。通常将WSI划分为小块,这些小块可以被认为是具有单个标签[13]的包,可以使用MIL方法实现。最近,MIL已成功应用于WSI分析[3]、[13]、[14]、[15]。MIL方法[16]有两个主要流:实例级方法[4]和嵌入级方法[17]。在第一类中,通过基于包级标签为每个实例分配一个伪标签并选择前 k 个实例进行聚合来训练模型。对于第二类,首先将WSI的实例映射到特征空间,然后由算子聚合。当前的大多数 MIL 方法都基于第二种方法,并在 WSI 分析中取得了良好的性能。然而,当训练数据有限时,这些方法很容易受到过度拟合的影响。因此,我们提出了一种基于自监督 MIL 的方法来促进 WSI 分析。
自监督学习
自监督学习(SSL)在自然语言处理和医学图像分析任务中取得了相当大的成功[18],[19],[20]。现有的 SSL 方法构建不同的前置任务来利用输入数据中的监督信息。它们可以分为三类[18],[21]:1)基于上下文的,2)基于生成的,3)基于对比的方法。对于第一类,前置任务的设计通常基于特定领域的知识,如分割任务[22]、[23]、[24]、[25]和分类任务[19]、[26]。对于第二类,通过最小化像素空间中的重建损失[27],[28]来正则化模型。第三类是指一类新提出的方法,它学习在正/负对[29],[30]之间在潜在空间中强制相似性。已经有值得注意的技术被提出,这些技术将对比学习扩展到医学图像分析[11]、[31],但据我们所知,目前还没有关于多实例对比学习的工作来学习如何进行特征聚合。
视觉任务中的强化学习
深度学习在 RL 领域得到了广泛的研究。通过引入深度神经网络来构建策略函数,RL 智能体具有强大的能力来学习具有交互的环境的动态 [32]。最近,一些基于 RL 的视觉图像处理 [33]、[34] 方法被提出,取得了显着的性能。此外,已经提出了基于 RL 的方法来自动为分类任务 [35] 找到有效的增强或 RL [36]、[37]、[38] 中的泛化。然而,很少有作品使用强化学习来促进对比学习。据我们所知,这是第一次采用强化学习来促进基于 CL 的表示学习。值得注意的是,与之前将 CL 与 RL 相结合的工作(即 CURL [38])不同,我们的框架利用 RL 动态构建正/负对以促进基于集合的 CL 学习,而 CURL 将 CL 作为自监督辅助任务来促进特征表示以提高 RL 任务的性能。
方法
图 2(a) 概述了所提出的 MuRCL 的框架,我们的目标是通过智能体从 WSI 的输入特征包中选择两个独立的判别集来构建正/负对以进行对比学习。然后训练 MIL 聚合器 M ( ⋅ ) M(\cdot) M(⋅) 和预测头 f ( ⋅ ) f (\cdot) f(⋅) 以使用对比损失最大化正判别集的一致性。MuRCL 的每个分支称为 RL-MIL。在图 2(b) 中,我们说明了 RL-MIL 的顺序决策过程。使用输入特征包,RL-MIL 迭代地生成一系列判别集,并使用 MIL 聚合器和预测头输出特征向量序列。具体来说,在每一步,一组特征索引由奖励引导智能体(即判别集提议网络)确定。然后,通过将输入特征图包每个集群的索引特征复合来构建下一步的判别集,如图 2© 所示。最后,自监督预训练模型将使用 WSI 标签进行微调,以规范学习的表示以进行最终预测。
多实例对比学习
我们的多实例对比学习框架以WSI特征包(WSI-Fbag)作为输入,其中特征包由在ImageNet上预训练的ResNet18[39]提取的实例级嵌入组成。CL的一个关键组成部分是构建用于训练的逻辑正/负对(即语义相似/不相似的实例)。与以往基于图像增强的策略不同,我们提出从每个WSI-Fbag中抽取不同的WSI判别集(简称WSI-Fset),构建基于集的正/负对进行CL训练。
每个 WSI-Fset 是来自 WSI-Fbag 的多个特征集群的子集的组合。具体来说,给定一个 WSI-Fbag x x x。我们首先使用聚类算法(例如 K-means [40])将 WSI-Fbag x x x 拆分为 K 个特征簇 C k C_k Ck ( k ∈ [ 1 , 2 , … , K ] k\in[1,2,…,K] k∈[1,2,…,K])。通过对第k个簇的子簇进行采样,可以将复合WSI-Fset x ~ \tilde{x} x~ 形成为这些子簇的拼接。我们对每个簇使用相同的采样率,以便构建具有恒定数量的实例嵌入的 WSI-Fset。在得到WSI-Fset x ~ \tilde{x} x~后,我们按照分类头 f ( ⋅ ) f(\cdot) f(⋅)将其输入到MIL聚合器 M ( ⋅ ) M(\cdot) M(⋅)中,生成WSI级别的特征嵌入 p p p。值得注意的是,对于不同的采样策略,我们可以从相同的 WSI-Fbag 中获得不同的 WSI-Fset 和对应的 WSI 级嵌入。这些不同的WSI-Fset可以被视为同一WSI的不同视图,因此可用于在CL中形成正对。
假设 { x n ~ } n = 1 N \{\tilde{x_n}\}_{n=1}^N {xn~}n=1N是一组有N个WSI-Fset的集合,其中 x i ~ \tilde{x_i} xi~和 x j ~ \tilde{x_j} xj~是从同一个WSI-Fbag中采样的,其他的则是来自不同WSI-Fbag。那么CL损失可以计算为: L i , j = − l o g exp ( s i m ( p i , p j ) / τ ) ∑ n = 1 N I ( n ≠ i ) ⋅ exp ( s i m ( p i , p n ) / τ ) , ( 1 ) L_{i,j}=-\mathrm{log}\frac{\exp(sim(p_i,p_j)/\tau)}{\sum_{n=1}^N\mathbb{I}(n\neq i)\cdot\exp(sim(p_i,p_n)/\tau)},\quad\mathrm{(1)} Li,j=−log∑n=1NI(n=i)⋅exp(sim(pi,pn)/τ)exp(sim(pi,pj)/τ),(1)
其中 τ \tau τ表示温度参数, s i m ( ⋅ , ⋅ ) sim(·,·) sim(⋅,⋅)表示两个向量的余弦相似度。而 I ( n ≠ i ) ∈ { 0 , 1 } \mathbb{I}(n\neq i)\in\{0,1\} I(n=i)∈{0,1}是 n ≠ i n\neq i n=i时取值为1的标志函数。在这项工作中,我们采用NT-Xent损失[29]作为目标函数,以最大化正对之间的相似性和最小化负对之间的相似性。在这种情况下,将鼓励MIL聚合器学习聚合知识以进行准确的预测。
RL驱动的判别集构造
如前所述,如何在 MuRCL 中构建 WSI 判别集 (WSI-Fset) 很重要。因此,我们提出了一种新的强化学习驱动策略,称为 RL-MIL,用于从 WSI-Fbag 构建判别 WSI-Fset,其中 WSI-Fset 提出智能体 R(例如,循环神经网络)通过强化学习进行训练。如图2(b)所示,WSI-Fset构造可以表述为顺序决策过程。在每一步,MIL 聚合器 M(·) 和投影头 f (·) 将 WSI-Fset 作为输入并产生相应的语义预测 p。同时,智能体 R 还将通过将特征向量 v \mathcal{v} v 作为输入来为下一步生成另一个 WSI-Fset 提议 s \mathcal{s} s。
具体来说,给定输入 WSI-Fbag x x x,RL-MIL 模块迭代地处理一系列动态选择的 WSI-Fset { x 0 ~ , … , x t ~ , … } \{ \tilde{x_0} ,\dots, \tilde{x_t} ,\dots\} {x0~,…,xt~,…}。在过程的第 t 步中,MIL 聚合器 M(·) 接收当前 WSI-Fset x t ~ \tilde{x_t} xt~ 并输出特征向量 v t v_t vt ,并输入到分类头 f (·) 得到slide级的嵌入 p t p_t pt。然后发送slide级嵌入来计算等式(1)中的对比损失。同时,通过将获取的特征向量 v t v_t vt作为输入,还可以由WSI-Fset提出的智能体R确定选择下一个WSI-Fse x ~ t + 1 \tilde{x}_{t+1} x~t+1的特征索引的动作 s t + 1 s_{t+1} st+1。此外,WSI-Fset x ~ t + 1 \tilde{x}_{t +1} x~t+1将会通过如图 2 © 所描述的生成特征索引从WSI-Fbag x中选择。之后,它将作为下一步的输入。我们以 ABMIL [8] 和 CLAM [41] 作为 MIL 聚合器 M(·) 来聚合 WSI-Fset 中的输入特征,而分类头 f (·) 和 WSI-Fset 提议智能体 R是两个循环网络,如 [33],这样它们可以分别通过维护隐藏状态 h t − 1 R h_{t -1}^R ht−1R 和 h t f h_t^f htf来利用所有先前输入的信息。需要注意的是,我们的框架不是 RL 框架,而我们使用 RL 中的优化策略来训练我们的框架。在提出的 MuRCL 中,强化学习被用作使用多实例学习训练的辅助任务,用于判别集构建。在集合构建过程中,智能体多次扫描WSI来定位判别特征,并在每一步计算奖励以更新智能体以进行下一次决策。因此,我们在我们的框架中借用了强化学习中优化方法的思想来构建集合。
-
RL-Driven Selection:在 RL 驱动的判别集构建的每一步中,WSI-Fset 是从 WSI-Fbag 中选择的,特征索引如下过程。为了便于智能体生成空间一致的WSI-Fset,我们根据输入WSI-Fbag x x x中的簇标签对特征进行重新排序,即具有相同簇标签的特征被分配相邻的索引。之后,对于每个簇 C k C_k Ck 的特征,我们将它们沿相应补丁的坐标重新排列。重新排列的 WSI-Fbag 称为 Affinity 图,如图 2 © 所示。然后,我们可以根据智能体 R 预测的动作 s 从 Affinity图中复合 WSI-Fset,其中我们将动作 s 表示为重新排列的集群的一组特征索引。具体来说,在第 t 步,会产生由第 (t-1) 的智能体R 导出的特征索引向量 s t ∈ R K ∗ 1 s_t \in \mathbb{R}^{K ∗1} st∈RK∗1,其中元素 s t k s_t^k stk 表示第 k 个重新排列集群的特征索引。因此,对于集群 C k C_k Ck ,我们在第 ( s t k s_t^k stk )个特征的开头采样一个特征序列,其中序列的长度等于采样率乘以聚类特征的数量。然后,WSI-Fset x t x_t xt可以通过将不同簇的序列连接在一起来构建。
-
Reward: WSI-Fset 提议的智能体 R 使用策略梯度方法 [42] 进行训练。在训练过程中,如何设计奖励函数很重要,因为它控制智能体R优化的方向[32]。在我们的框架中,我们建议利用 WSI-Fset 的正对之间的相似性作为奖励来指导智能体定位信息特征。具体来说,在第 t 步,训练智能体的奖励函数设计为 Equ(2)。 r i , j ; t = s i m ( p i ; t − 1 , p j ; t − 1 ) − s i m ( p i ; t , p j ; t ) . ( 2 ) r_{i,j;t}=sim(p_{i;t-1},p_{j;t-1})-sim(p_{i;t},p_{j;t}).\quad(2) ri,j;t=sim(pi;t−1,pj;t−1)−sim(pi;t,pj;t).(2)其中 p i p_i pi和 p j p_j pj分别由相应的输入正WSI-Fset x i ~ \tilde{x_i} xi~和 x j ~ \tilde{x_j} xj~导出。通过扩大这些正WSI-Fset的余弦距离,引导智能体 R挑选出发散度较大的特征集,从而通过最小化CL损失来促进MIL模型专注于潜在聚集知识。
-
Discriminative Set Mixup: 为了在MIL聚合器训练中引入更多的扰动,我们还采用了一种有效的特征增强策略,称为set-mixup,以增加WSI-Fset的多样性。对于一个训练批次中的WSI-Fset,通过将一个WSI-Fset的 x l ~ \tilde{x_l} xl~与另一个的 x q ~ \tilde{x_q} xq~混合,生成的 x q ˉ \bar{x_q} xqˉ是 x q ~ \tilde{x_q} xq~的增广。这可以用等式(3)来说明,其中λ是从分布中采样的系数,例如 λ ∼ U ( α , 1.0 ) \lambda\sim U(\alpha,1.0) λ∼U(α,1.0),其中α在我们的实验中设置为0.9。 x ˉ q = λ x ~ q + ( 1 − λ ) x ~ l . ( 3 ) \bar{x}_q=\lambda\tilde{x}_q+(1-\lambda)\tilde{x}_l.\quad\quad\quad\quad\quad\quad(3) xˉq=λx~q+(1−λ)x~l.(3)通过混合包含多种wsi级表示的WSI-Fset,set-mixup 操作有利于增强MIL聚合器的语义概念学习。
RL-MIL的训练方案
如图2(a)所示,我们的对比学习框架有两个分支,其中MIL 聚类器 M ( ⋅ ) M(\cdot) M(⋅)、分类头 f ( ⋅ ) f(\cdot) f(⋅)和智能体R在两个分支中的参数是共享的。在训练批中,我们首先使用从WSI-Fbag中随机选择的两个WSI-Fset作为正对的初始化。然后,使用强化学习策略训练的智能体将构建正对。为了更清楚地说明,我们使用下标 ( ⋅ ) i (\cdot)_i (⋅)i和 ( ⋅ ) j (\cdot)_j (⋅)j分别表示两个分支的变量, { ( ⋅ ) t } t = 0 T \{(\cdot)_t\}_{t =0}^T {(⋅)t}t=0T表示每个分支产生的时间序列。T设为5,表示rnn对一个训练批执行5次迭代(步骤)。在开始的时候,将同一WSI-Fbag中随机选取的两个WSI-Fset x ~ ( i ; 0 ) \tilde{x}_{(i;0)} x~(i;0)和 x ~ ( j ; 0 ) \tilde{x}_{(j;0)} x~(j;0)作为初始化正对,然后分别由 M ( ⋅ ) M(\cdot) M(⋅)推断出两个不同的特征向量 v i ; 0 v_{i;0} vi;0和 v j ; 0 v_{j;0} vj;0。同时,计算wsi级特征嵌入 p i ; 0 p_{i;0} pi;0和 p j ; 0 p_{j;0} pj;0,并通过前向传播去计算训练批次中不同wsi-fbag的wsi-fset视为负对的CL损失 L 0 = L i , j ; 0 ( p i ; 0 , p j ; 0 ) L_0=L_{i,j;0}(p_{i;0}, p_{j;0}) L0=Li,j;0(pi;0,pj;0)。接下来,对R进行第一次迭代,分别取 v i ; 0 v_{i;0} vi;0和 v j ; 0 v_{j;0} vj;0作为R的初始状态,提出新的正对。通过这个新的正对,对第二个推论以同第一个一样的方式进行迭代,形成第二个R。经过5次迭代, f ( ⋅ ) f(\cdot) f(⋅)与R进行同步处理,从两个分支 { p i ; t , p j ; t } t = 0 5 \{p_{i;t}, p_{j;t}\}_{t = 0}^5 {pi;t,pj;t}t=05的 f ( ⋅ ) f(\cdot) f(⋅) 的输出将被用来通过等式1计算CL损失: L t = ∑ t = 0 5 L i , j ( p i ; t , p j ; t ) L_t=\sum_{t=0}^5L_{i,j}(p_{i;t},p_{j;t}) Lt=∑t=05Li,j(pi;t,pj;t)。此外,通过最大化折扣奖励 r i , j = ∑ t = 1 5 γ t − 1 r i , j ; t ( p i ; t , p j ; t ) r_{i,j}=\sum_{t=1}^5\gamma^{t-1}r_{i,j;t}(p_{i;t},p_{j;t}) ri,j=∑t=15γt−1ri,j;t(pi;t,pj;t)来同时训练智能体R,其中 γ \gamma γ设为0.1。此外,我们从两个不同采样的N个WSI-Fbag中获得了2N个WSI-Fset。对于每个正对,我们同[29]一样,将其他2(N−1)个WSI-Fset视为负例。综上所述,推理过程可以明确为算法1。
训练过程包括三个阶段。在第一阶段,WSI-Fset提出的智能体R不包括在RL-MIL中。相反,我们在每一步中,从输入的WSI-Fbag x中随机抽取一个WSI-Fset,然后训练MIL聚合器M(·)和分类头f(·),以最小化等式1中的对比损失。执行此步骤是为了确保模型能够适应任意输入序列。在第二阶段,固定MIL聚合器和分类头。我们随机初始化智能体R,然后按照上面描述的训练过程训练它。在第三阶段,利用第二步得到的固定R对MIL聚合器和分类头进行进一步微调,利用学习到的特征选择策略增强RL-MIL的性能。
框架微调和推理
利用上述多实例对比学习范式,该框架可以深入挖掘不同patch之间的语义关系,用于WSI级的WSI表示。对于最终的WSI级预测,我们进一步使用标记的wsi对框架进行微调。在这一阶段,我们将f(·)的最终输出维数从128改为类数,以满足我们的下游分类任务。如预训练中所述,微调训练也包括三个阶段。不同的是,我们通过对f(·)的输出进行softmax操作来产生置信度分数。然后,将置信度分数的增加,即: r ^ t = p ^ t − p ^ t − 1 \hat{r}_{t}=\hat{p}_{t}-\hat{p}_{t-1} r^t=p^t−p^t−1( ( ⋅ ^ ) (\hat{\cdot}) (⋅^)表示微调中对应的变量)作为WSI-Fset提议智能体R遵循[33]设置的奖励,其中 p ^ \hat p p^为带有真实标签的softmax预测概率。
测试阶段的MuRCL过程与微调阶段相同。在测试开始时,使用一个输入测试WSI-Fbag,由M(·)提供一个对智能体R随机采样WSI-Fset的初始状态,然后由智能体R确定策略选择的WSI-Fset,并由MIL聚合器M(·)和分类头f(·)在微调阶段进行处理。在这个过程中,agent用初始状态迭代处理状态向量,因为每次迭代的状态都会传递到下一次迭代。在训练和微调阶段,迭代设置为5次。我们将智能体R在最后一次迭代的输出作为WSI-Fset的提议,然后将f(·)的输出作为分类预测。对比损失可以促进模型扩大不同类别特征之间的空间距离,拉近同一类别特征之间的距离。经过微调后,该模型可以对空间距离小的特征分配相同的类别标签,对空间距离大的特征分配不同的类别标签,从而实现准确的分类预测。从而区分正常和癌变区域。此外,RL 智能体可以对从不同位置采样的特征进行迭代分析,从粗到细定位判别特征,类似于医生的重复诊断过程。因此,可以期望这些组件改进模型。
实验
数据集
我们报告了四个基准临床WSI数据集的结果,包括Camelyon16、TCGA-Lung、TCGA-Kidney[46]和TCGA-Esca,涵盖了平衡/不平衡和单/多类MIL问题。(1) Camelyon16是一个广泛使用的用于乳腺癌转移检测的公共数据集,包括270个训练wsi和129个测试wsi,已被广泛用于癌/非癌分类和定位任务[12]。预处理后,共获得约270万个实例,平均每个包约6881个实例。(2) TCGA-Lung包括两个子目标,即肺鳞癌(Lung Squamous Cell Carcinoma, TCGA-LUSC)和肺腺癌(Lung adenocarmicoma, TCGA-LUAD),共有1041例诊断WSI,其中LUAD有529例WSI,LUSC有512例,作为WSI子类分类和生存分析的基准。预处理后,每张WSI平均提取的实例数为11540个。(3) TCGA-Kidneyis分为肾憎色性肾细胞癌(TCGA-KICH)、肾透明细胞癌(TCGA-KIRC)和肾乳头状细胞癌(TCGA-KIRP) 3个子集合,共734个诊断wsi,其中KICH92张,KIRC411张,KIRP231张,适合多类分类任务。(4) TCGA-Esca分为两组子类,共156例诊断性wsi,分别为90例鳞状细胞癌和66例腺癌。
所有实验均采用20倍放大的wsi。将wsi裁剪成大小为256 × 256的patch,丢弃组织区域小于35%的patch。对于Camelyon16,我们将官方训练wsi随机分成80%和20%进行训练和验证,并在官方测试wsi上进行测试,遵循之前的工作[12]。对于TCGA数据集,我们以3:1:1的比例随机分割数据,用于训练、验证和测试,遵循[12]。
实现细节和评估指标
与[41]和[17]类似,每个WSI patch被嵌入到一个512维(d=512)的特征向量中,并使用预训练的编码器。我们对Camelyon16和TCGA-Lung使用了由[12]训练的SimCLR [29] ResNet18[39]编码器,而TCGA-Kidney和TCGA-Esca则由ImageNet上预训练的ResNet18进行编码。WSI特征包首先使用k -means[40]聚类成10个类别,每个簇C k的采样比在WSI- fset构造中设置为1024/u,其中u表示WSI包中的特征数量。此外,在我们的实验中,批量大小N设置为128,温度 τ \tau τ设置为1。在训练过程的第一阶段,我们使用一个初始学习率为1e-4和1e-5的Adam优化器分别对MIL聚合器M(·)和分类头f(·)进行优化。权重衰减都设为1e-5。第二阶段,我们使用初始学习率为1e-5的Adam优化器来训练agent r。第三阶段,由Adam优化器联合优化MIL聚合器M(·)和分类头f(·),其中初始学习率分别设置为5e-5和1e-5,权值衰减同样为1e-5。这三个阶段分别用100次、30次和100次训练。微调阶段的优化器设置与预训练阶段相同。所有实验均在两张RTX 3090 GPU卡上进行。我们用准确率(accuracy, ACC)、曲线下面积(area under the curve, AUC)分数和F1分数来评价分类性能,其中ACC在所有实验中都以0.5的阈值来计算。对于AUC,我们报告了多类分类情况下的平均one-versus-rest AUC(macro-averaged)。
与最好的方法比较
-
实验设置:我们给出了二类和多类分类的结果。二分类包括Camelyon16癌/非癌分类、TCGA-Lung cancer 分类和TCGA-Esca亚型分类。多分类包括TCGA-Kidney分类。我们将我们的框架与一组强基线进行比较,包括不同的高级MIL模型[4],[8],[41],[44]。由于我们的MuRCL对MIL聚合器是模型不可知的,我们用ABMIL[8]和CLAM_SB[41]作为MIL聚合器主干(分别表示为Ours(ABMIL)和Ours(CLAM_SB))来评估我们的MuRCL,以显示我们框架的泛化和鲁棒性。我们首先使用提出的MuRCL训练我们的框架,然后使用标记的wsi对其进行微调,同时我们还使用相同的标记wsi训练其他方法。我们用不同的seeds进行了5次实验,并以标准差报告了平均性能。
-
对比结果:我们首先比较了Camelyon16的性能,Camelyon16是最流行的WSI基准分类数据集。如表1和表2所示,我们使用CLAM_SB和ABMIL聚合器的框架优于相应的基线框架(即在ACC中Camelyon16上CLAM_SB和ABMIL分别为5.5%和3.9%),这表明我们的MuRCL预训练策略有助于获得聚合知识。总的来说,我们的框架比其他MIL方法实现了最好的性能,证明了我们的MuRCL框架的优越性。
我们可以在TCGA-Lung数据集上观察到类似的趋势。此外,我们的方法在所有三个指标上都优于所有其他竞争方法。观察到,由于提取的patch特征的表示能力有限,ABMIL在TCGA-Lung上表现不佳。然而,在我们提出的强化对比学习框架下,ABMIL的性能得到了很大的提高。我们认为这是因为我们的MuRCL可以指导骨干MIL网络基于从原始数据中学习到的通用知识集中于更有信息的特征。此外,我们的方法与不同的MIL主干兼容,得出与Camelyon16数据集相似的结论。
与之前的数据集不同,TCGA-Kidney由三类wsi组成。此外,TCGA-Kidney在癌症亚型中的分布不平衡,给准确分类增加了难度。在这种情况下,我们的MuRCL也适用于具有不平衡数据集的多类问题,并在所有指标中持续改进相应的基线对应项。在该数据集上的实验结果证明了我们提出的框架对各种站点的泛化和鲁棒性。
为了进一步验证所提出方法的鲁棒性,我们还在TCGA-Esca上进行了分析实验,该数据集(156张)少于其他三个数据集,即Camelyon16(399张),TCGA-Lung(1041张)和TCGA-Kidney(734张)。如表2所示,我们的方法实现了与其他最先进方法相当的性能。结果表明,在有限的训练数据场景下,强化对比学习框架仍然可以学习到分类任务的代表性特征。
因为肿瘤/正常分类比亚型分类更具挑战性。与TCGA数据集相比,我们的方法在Camelyon16上比其他mil更具优势,表明我们的框架可以提取和学习判别特征,以处理更具挑战性的WSI分析任务。考虑到肿瘤/正常分类的公开数据集很少,为了证明我们方法的鲁棒性,我们在三个TCGA数据集上进行了额外的实验来验证我们方法的有效性。此外,我们对不同的方法进行了统计检验,证明了所提出的方法显著优于其他方法。与our (CLAM_SB)和Ours(ABMIL)相比,DSMIL、CLAM_MB和CLAM_SB的AUC p值如表六所示。大多数p值小于0.05,该方法显著优于其他方法。
我们框架的分析
1)消融研究:我们进行了一项研究,以证明我们提出的多实例对比学习机制和强化学习模块在选择判别集方面的有效性。我们首先将MuRCL学习到的特征与几个自监督学习基线进行比较。表4给出了不同模型经过监督微调后的消融实验结果。报告的结果也是使用不同种子的五次运行的平均性能。
模型A ( M A M_A MA)是具有随机抽样WSI-Fset的普通MIL聚合器,模型B ( M B M_B MB)是使用RL智能体生成WSI-Fset的监督MIL聚合器。Model min( M m i n M_{min} Mmin)是典型的MIL聚合器,其WSI-Fset以提取的特征的最小余弦相似度采样。具体来说,我们从包中随机抽取10个特征集,并选择余弦相似度最小的一对特征集作为正对。类似地,Model max ( M m a x M_max Mmax)是典型的MIL聚合器,其WSI-Fset采样具有提取特征的最大余弦相似度。具体来说,我们从包中随机抽取10个特征集,并选择余弦相似度最大的一对特征集作为正对。我们观察到, M B M_B MB比 M A M_A MA得到了更大的改进,表明了基于强化学习的判别特征选择策略的有效性。在模型C( M C M_C MC)中,我们使用普通对比学习策略来训练模型,其中正/负对是用随机选择的特征构建的。 M C M_C MC的效果优于 M A M_A MA,说明对比学习预训练的有效性。对于模型D( M D M_D MD),我们使用提出的MuRCL框架来训练框架,但在微调期间不使用RL-agent来重建WSI-Fset(即在微调期间使用随机选择的WSI-Fset)。 M D M_D MD在TCGA-Kidney数据集和Camelyon16上的性能优于 M C M_C MC,在TCGA-Lung数据集上的性能与 M C M_C MC相当,表明基于RL的WSI-Fset构建可以提高正对的质量,从而更有效地进行CL训练。在表4中,我们的最终模型( M E M_E ME)在大多数情况下比所有消融模型实现了最好的性能,证明了我们提出的整个框架的能力。我们还对提出的方法进行了无设置mixup(我们的无设置mixup)的实验,结果见表7。结果表明,设置mixup可以有效地提高模型的性能。
2)RL-MIL设计分析:从表4的对比结果可以看出,基于RL的WSI-Fset构建不仅适用于自监督MIL训练( M D M_D MDvs. M C M_C MC),而且对全监督MIL训练( M B M_B MB vs. M A M_A MA)也有很强的促进作用。我们认为这是因为强化模块可以有效地从之前的动作中收集信息并做出更好的决策,这也证明了从WSI构造一个判别特征集对多实例学习是非常有益的。如表5所示,我们还研究了训练过程中运行RL的迭代次数。在表中可以看到,迭代次数越多,性能越好。我们认为这是因为智能体可以通过更充分的比较,从更多的迭代中选择最具判别性的特征。但是,每次迭代都需要存储中间变量参数,因此会消耗大量内存。由于内存消耗的限制,我们在实验中将迭代次数设置为5次。在未来,更多迭代的影响可以用更多的GPU内存来探索。
3)CL设计分析:如表4所示,通过强化对比学习预训练的 M D M_D MD比 M A M_A MA取得了更好的性能,这说明对比学习预训练在MIL中发挥了重要作用。此外,我们使用不同的编码网络,包括在ImageNet上预训练的ResNet18 (ImageNet ResNet18)和在[12]上使用SimCLR预训练的ResNet18 (SimCLR ResNet18)来提取Camelyon16和TCGA-Lung数据集的特征。由于[12]没有在TCGAKidney上提供预训练的SimCLR ResNet18,我们使用在ImageNet (ImageNet ResNet50)和ImageNet ResNet18上预训练的ResNet50从TCGA-Kidney中提取特征。如表IX和表X所示,在不同的特征表示下,所提出的强化对比学习仍然可以对分类性能的提高产生积极的影响(对比CLAM_SB与 M B M_B MB和 M E M_E ME)。另外,对比ImageNet ResNet18和SimCLR ResNet18的结果,我们可以发现用训练数据预训练的编码器(SimCLR ResNet18)可以帮助提升所有方法的性能。
4)WSI判别集可视化:我们将预训练的MuRCL模型所选择的WSI- Fset的位置可视化,以分析所提出的MuRCL的有效性。图3显示了两个选择的判别集的位置,图4显示了其特征在Camelyon16的WSI示例上的t-sne[47]可视化。由于缺少像素级标注,RL 智能体需要扫描WSI区域来确认是否存在癌斑。然而,由于模型存在不确定性误差,仅通过一个预测为转移瘤的实例来确定WSI的标签是不合理的。该模型需要从所有补丁中收集足够的特征以获得更精确的预测。因此,模型需要按照RL 智能体确定的特定顺序查看WSI区域,根据特征分布的变化发现判别特征。此外,我们在图5中给出了以预训练和微调方式训练的模型生成的预测热图,其中概率得分在0到1(蓝色到红色)的范围内。此外,我们还将热图与肿瘤区域的ground truth进行了比较,如图5中的黄色曲线所示。将这些地图与医生注释的癌症区域进行比较,很容易发现模型和医生往往具有相同的感兴趣区域。这一现象可以解释为什么该模型可以做出准确的诊断。一方面,我们在Camelyon16的所有WSI上测量模型构建的不同判别集的重叠度,得到肿瘤WSI和正常WSI的重叠度分别为10.24±9.17%和19.75±23.66%。在图3中,不同的颜色代表不同聚类的选定特征。如第三节的B1所示,我们将每个聚类的特征按照这些特征在WSI中的坐标进行排序,然后从每个聚类中选择特征构建判别集。由于同一聚类的特征在WSI中具有连续坐标,因此同一聚类中所选的特征将在可视化中显示为扫描线模式。如图3所示,所选择的两种视图相互区分,但与WSI包含相同的语义信息,这有助于模型在对比学习中学习潜在的类表示。此外,我们的MuRCL可以在没有像素级注释的情况下帮助定位关键的肿瘤区域(参见放大的区域),有助于提高WSI分类性能。
讨论
提出了一种新的多实例强化对比学习方法用于WSI分类。在具有代表性的基准WSI数据集上,我们的方法比以前的MIL方法有显著的提升。关键的技术创新是一种新的基于集合的自监督MIL训练范式,源自基于实例的对比学习。如前所述,我们有动机通过使用强化对比学习挖掘不同patch之间的更多关系来缓解WSI MIL中的过拟合问题。
此外,所提出的MuRCL框架在监督训练中也可以被视为一种强正则化。由于我们的MuRCL是一个两阶段的过程,为了评估其正则化效果,我们在Camelyon16数据集上进行了一个一阶段的多任务学习实验,其中有10%的标记训练数据。在本实验中,我们还在MuRCL框架中添加了一个监督任务(见图2 (a))。具体来说,我们在 p i p_i pi后面添加一个分类器并计算分类损失。结果如表x所示。虽然多任务模型的性能不如我们微调后的MuRCL,但它的性能优于MIL基线,显示了我们提出的强化对比学习框架的正则化效果。
为了证明我们提出的多实例对比学习范式的有效性,我们还使用线性评估协议[30],[48],[49]来评估我们的MuRCL的学习表征,其中线性分类器在固定编码器网络上进行训练,测试结果用作表征质量的代理。对于线性评估,我们使用学习率为1e-4的Adam优化器,训练过程在40个epoch中进行。结果见表八。观察到,在线性评价下, M C M_C MC的性能也优于 M A M_A MA和 M B M_B MB。而且,在大多数情况下, M B M_B MB比 M A M_A MA得到更好的结果,进一步证明了基于RL的WSI-Fset构建在CL训练中的有效性。
我们提出在自监督MIL对比学习中利用强化学习来构建正/负对的判别特征集。在强化学习中,如何设计动作和奖励是很重要的,我们利用提出的亲和图和对比学习损失来设计这些关键组件。我们将研究不同的策略来进一步提高所提出的RL-agent的有效性,例如优化奖励函数和agent的提议机制。例如,由于WSI包含数百万个实例,因此将WSI的特征与聚合器/分类器一起学习是一个耗时且消耗内存的过程。因此,使用内存有限的GPU对特征提取器和聚合器/分类器进行微调可能是不切实际的。我们可以使用基于集的多实例学习来构建判别实例集来训练模型,以减少GPU的消耗。然而,这种方法可能会显著增加训练时间消耗。在未来的工作中,我们将更有效地学习特征提取与聚合器/分类器的结合。
在实用性方面,我们相信我们的工作具有临床应用的潜力,在注释有限的情况下,为基于集合的自监督研究提供了新的见解。此外,由于图像的规模较大,WSI分析是一个繁琐、费力和耗时的过程。我们的方法可以自动帮助医生更快更好的进行诊断,减轻医生的工作量。此外,我们的自监督MIL训练概念可以推广到更高维度的数据(例如3D数据或视频)和更多MIL设置,这可能是我们未来的工作。
总结
在本文中,我们提出了一种新的高效的MIL框架用于WSI分类,这对其他MIL问题也有很大的潜力。为了解决由于训练数据不足导致的MIL任务的过拟合问题,我们将学习过程描述为一个自训练过程,以学习不同实例之间的潜在关系,从而使整个框架规范化。在此基础上,将自我训练过程设计为基于集合的学习过程,并结合强化学习来帮助选择判别正对。最后,使用滑动级标签对预训练模型进行微调,以进行最终预测。定量和定性评价都证明了该方法在三个独立数据集上的优越性。
论文评论
这篇文章主要是通过强化学习和对比学习来优化对实例的选取的,但由于我对强化学习和对比学习了解的不多,因此现阶段先找不出可能存在的问题(这部分在之后了解清楚之后可能会进行重新评论),且在讨论部分作者有提出一些问题,可以通过作者提出的问题去改进。
更多内容,见微信公众号:
这篇关于论文阅读:MuRCL: Multi-Instance Reinforcement Contrastive Learning for Whole Slide Image Classification的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!