Uncertainty-aware consistency learning for semi-supervised medical image segmentation
半监督医学影像分割面临两个挑战性问题:(1)对潜在结构的探索不足,导致难以全面捕捉医学影像中的复杂特征和结构;(2)对噪声的敏感性,未标记的数据缺乏准确的标签信息,使得模型在学习过程中更容易受到噪声干扰。本文提出了一种不确定性感知一致性学习(UAC)方法,以改善半监督医学图像分割中因信息挖掘不足和对噪声敏感而导致的泛化能力差和性能不理想的问题。首先,通过在输入和输出层面采用多重扰动策略,特别是通过数据级和尺度级扰动,该模型能更好地捕捉器官内部的结构信息和影响分割性能的基本特征。其次,扰动不确定性利用扰动预测差异来衡量不确定性,有助于模型生成可靠的预测,避免过度关注预测中不可靠的区域。
Introduction:
虽然现有的方法已经取得了很大的优势,但要实现医疗图像分割仍面临挑战,原因在于:(1)对潜在信息的挖掘不足。在 SLL 中,由于未标记数据中缺乏明确的监督信号,因此模型无法直接依赖标签来指导特征提取。导致对潜在信息的利用不足。未标记数据中蕴含的宝贵信息往往被忽视,这可能导致模型学习的特征表征不够全面或准确,无法处理复杂区域的分割任务(图 1(a))。(2) 对噪声的敏感性。主观注释和噪声异常值会导致潜在的错误分割,并严重影响分割模型的性能。如图 1(a)所示,由于器官粘连区域和一些小分支的复杂性和错综复杂性,该模型容易在这些区域产生高度不确定和错误的分割。如果不对这些低置信度的预测加以限制,模型在训练过程中可能会过度关注这些噪声区域和不可靠的预测,从而导致学习到不正确的知识。
不确定性校准是实现半监督图像分割的一种潜在方法。通过在分割过程中加入不确定性校准,模型可以学习到更可靠的在无监督情况下,不确定性估算可通过设置阈值过滤掉不可靠的预测信息,从而减少噪声干扰的影响。一般来说,不确定性估计采用集合方法或蒙特卡罗抽样来计算模型预测置信度,并通过设置阈值过滤掉不可靠的预测,以避免模型过度关注噪声区域和不确定性较高的区域。然而,这些方法大多基于熵值越高对应不确定性越高的假设,当模型过于自信而数据分布不均匀时,就很难捕捉到模型的真实不确定性。此外,不可靠的预测往往发生在结构复杂、边界模糊的区域。直接过滤掉这些不可靠的预测可能会导致这些具有挑战性区域的信息丢失。
本文提出了一种新颖的基于一致性的训练策略--不确定性感知一致性学习(UAC,图 2),以实现少标签图像的半监督医学图像分割。如图 1(b)所示,在扰动前后容易识别的区域,模型的预测表现出鲁棒性,而在具有挑战性的区域,由于数据的复杂性和模糊性,模型的预测变得不稳定。在这些困难区域,扰动前后的预测概率存在显著差异。利用这一特点,不确定性可以在一次前向传递中得到估计,从而实现多层次一致性的动态调整。拟议的 UAC 由三个部分组成:(1) 设计了多扰动策略,通过同时在数据层面和尺度层面对数据进行扰动,增强模型的泛化能力;(2) 扰动不确定性的估计包括在特征空间中引入了体素对比学习,提供明确的监督,进一步提高了特征的可分离性,从而实现了对输出一致性的指导。修正后的网络更加关注可靠的预测信息,减少了噪声的影响;(3) 在特征空间中引入了体素对比学习,提供明确的监督,进一步提高了特征的可分离性,使模型能够更好地区分不同类别的特征。
Methodology
拟议的 UAC 框架(图 3)通过利用一致性学习,利用有限的标记数据进行半监督医学图像分割。UAC 通过引入多级扰动来学习一致性预测,从而提高分割泛化能力。具体来说,所提出的 UAC 包括三个部分:(1) 多扰动策略通过引入各种扰动来探索扰动空间,发掘未标注数据中的潜在信息,帮助模型学习更稳健的特征表征;(2) 扰动不确定性涉及测量扰动预测差异,以估计单次前向传递中的不确定性,然后用于约束一致性学习,防止模型过度关注噪声区域;(3) 体素对比学习是指在特征空间中将相似类别的特征拉近,将不同类别的特征拉远,从而进一步提高类别可分性,学习更具区分性的特征。
Multi-perturbation strategy for latent structure
多重扰动策略旨在帮助一致性学习模型更好地理解器官内部的微妙结构。多重扰动策略旨在引入不同类型的扰动,在训练过程中更全面地探索数据的特征空间,发现影响医学影像分割性能的重要信息,帮助模型更好地泛化到新样本中,提高模型的性能和鲁棒性。
在使用扰动时,存在扰动强度的平衡问题。如果扰动太弱,模型可能会过度依赖初始预测或局部细节,无法充分利用未标记数据中的潜在信息。反之,如果扰动强度过大,可能会破坏图像的结构和语义信息,使模型无法准确学习有效的特征和模式。要充分利用数据中的潜在信息,在避免过度扰动造成不稳定和混淆的同时,找到正确的平衡点至关重要。Fixmatch 通过采用简单的强弱一致性框架,实现了与最先进方法相当的图像分类性能。受此启发,强弱一致性学习框架被引入到医学图像分割任务中。
如图 3 所示,在输入空间中,对输入图像 x 应用随机亮度变换和随机噪声等弱增强以及随机裁剪和色彩抖动等强增强,以获得增强视图 xS 和 xW。考虑到医学图像中的常见结构往往存在于多个尺度上,解码器中引入了金字塔结构来生成多尺度预测,表示为 pc。这有助于模型捕捉医学图像中不同层次的结构和细节。在此分割任务中,c = {0, 1, 2, 3},c 值越大,输出结果的分辨率越高。对多尺度预测结果进行上采样,以获得相同大小的输出结果,用于对齐和比较。
Perturbation uncertainty for consistency learning
采用扰动不确定性来衡量预测置信度,以防止模型在训练过程中过度关注噪声区域。通过计算扰动图像的多尺度输出预测值与原始预测值之间的差值来表示预测不确定性图。预测不确定性图既可用于确保模型生成可靠的预测,也可用于动态调整一致性学习,以防止模型过度关注预测中不可靠的区域。过度关注预测中的不可靠区域可能会导致两个潜在问题:一个是模型有可能学习到错误的器官结构特征,从而影响分割性能;另一个是模型有可能忽略或混淆真正的器官结构特征,从而导致预测中关键信息的丢失。UAC 框架包括原始数据流和扰动数据流。对原始输入图像的概率图 p 进行锐化,以获得p ̂:
这里,T 是控制锐化程度的超参数。值得注意的是,随着 T 的减小,预测结果的熵约束会降低。但是,如果 T 设置过低,可能会导致预测过于自信,忽略模型的不确定性。选择合适的 T 可以帮助我们获得更好的分割结果。直接进行多尺度一致性学习可能会引入不必要的噪声或误差。因此,利用扰动前后预测概率的差值来估算一次前向传递的不确定性图:
这里,Ur 表示扰动输出与原始输出之间的差值图。差值越大,表明模型预测的不确定性越高,学习到错误知识的可能性越大。受文献[35]的启发,为了防止模型过度关注噪声区域和错误信息,估计的不确定性图被用来指导学习过程:
在公式 中,第一项表示即使在对不确定性进行动态调整后仍能获得的校正一致性损失,而第二项则计算总体不确定性,并将其用作正则化项。通过引入额外的惩罚项,可以防止模型持续产生具有高不确定性的预测结果。
Voxel contrastive learning for feature discrimination
进一步学习结构化特征空间,帮助模型捕捉更精细的特征表征,以处理复杂的结构和模糊的边界。本文引入了体素级对比度损失,通过在特征空间中拉近相似样本的距离,拉远不同类别样本的距离,为特征提供明确的监督,从而提高分割的准确性和鲁棒性。具体来说,在解码器的第二个上采样块引入投影头,以生成特征表示。两个增强视图被视为正样本,而数据集中的其他样本则被视为负样本。在特征空间中,通过最大化正样本对之间的相似性和最小化负样本对之间的相似性,鼓励相似样本具有更接近的特征表示,而不相似样本具有更分散的特征表示。特征对对比损失函数定义如下:
其中,rS、rW 和 r- 表示强增强视图和弱增强视图以及负样本的特征表示。此外,sim (⋅) 表示两个特征图之间的相似度,Λ- 表示负样本图像集。本文使用余弦距离作为度量两个特征表征之间相似性的指标。此外,未标记的图像被用作标记图像的负样本,缓冲区用于存储动态更新的负样本,以计算对比度损失。
在像素级对比学习任务中,正样本是强增强视图和弱增强视图中的相应像素。然而,按照图像级对比学习的方法选择负样本可能会导致资源紧张。此外,半监督分割任务对噪声很敏感,错误分类负样本会导致无效学习,甚至误导模型的学习方向。因此,采用了自信负采样策略。首先,为了防止在负样本图像中选取与正样本属于同一类别的像素,利用分割结果构建一个差异矩阵来选取负像素。随后,从负像元中抽取置信度最高的 K 个像素。