TNNLS|多样性,真实性,鲁棒性:面向可信学习的医学图像少标签分割算法的基础研究
-
转载: jbcmzdxgg (知乎)
研究背景介绍
医学图像上分割标注的高成本问题始终是大规模医学图像分割模型训练的巨大限制。由于常规医学图像的高维度,标注过程的高专业性和高耗时性等特点,获得可靠的医学图像分割标注是非常困难且昂贵的,通常难以快速获得大量高品质的分割标签。一种理想的情况是减少深度学习模型训练所需的标签数量,通过算法提高对标签的使用效率,从而获得更加高效的学习模型,如半监督学习,弱监督学习等。然而,当标签数量降低到非常少的情况时(少标签),如1个,这些已有的高效模型往往依旧会发生严重的过拟合问题,无法获得有效泛化的分割模型。
少标签学习的问题在于标签数量过少导致的数据分布和标签分布的不可信[5],我们这里称之为数据缺陷导致的不可信学习(以下简称“不可信学习”),致使分割模型拟合错误的从数据到分割结果的映射关系,最终导致不可信的预测结果。针对这个问题,我们对面向可信学习的医学图像少标签分割进行了基础性的探索,探讨了该问题下对于分割学习算法的多样性,真实性和鲁棒性需求,并从“学习配准以学习分割(Learning registration to learn segmentation, LRLS)”范式出发,提出了有效的解决方法,相关论文[1]已被TNNLS2022接受,开源代码可从https://github.com/YutingHe-list/BRBS上获取。
图 1 缺乏多样性、真实性和鲁棒性导致的不可信的少标签学习过程(在LRLS范式中的分布表现)
1.简介
1.1 少标签医学图像分割中不可信学习问题
由于少标签医学图像分割场景中数据的不完备性,导致了模型在学习这些有缺陷的数据时的不可信学习问题。具体表现为,**1)数据分布缺陷导致的模型泛化不可信。**在少标签设置下,可信的带标签图像数量也是非常小的,因此这导致了数据分布的缺陷。这些少量的可信数据稀疏的存在于样本空间中,难以为模型训练提供一个符合真实世界数据分布的输入分布,导致模型严重的过拟合问题,大大限制了模型在真实世界使用中的泛化能力。尽管一些方法利用数据生成算法为模型提供更加丰富和完整的输入分布,但是生成数据缺乏真实性的限制为模型的输入分布带来了额外的噪声,导致模型拟合到不符合真实情况输入分布,最终使得模型泛化不可信。**2)标签分布缺陷导致的映射关系不可信。**一些方法试图通过生成伪标签的方式为无标签数据提供监督信息,丰富输入分布,从而帮助模型拟合到更符合真实世界的数据分布上。然而,生成的伪标签通常是不精确的,存在大量的标签噪声,形成了从输入图像到学习目标之间的错误对应关系。这些错误的对应关系将使深度模型拟合到不可靠的映射函数上,最终导致预测结果的不可信。尤其是在少标签场景下,极少的标注加剧了这种标签噪声的产生风险。
图 2 多样性、真实性和鲁棒性构建少标签情况下的可信学习(在LRLS范式中的分布表现)
多样性,真实性和鲁棒性是应对少标签医学图像分割任务中不可信学习的研究关键。**1)针对数据分布缺陷,多样性试图解决缺乏数据导致的输入分布稀疏的问题,而真实性试图避免输入分布脱离真实数据分布范围导致的失真问题。已有一些工作关注到这方面的研究,从数据生成或使用大规模无标签数据等角度为模型提供更加丰富的输入分布,从而帮助模型学习到对更广泛真实数据分布的泛化能力。2)针对标签分布缺陷,鲁棒性试图构建拥有对抗标签噪声的鲁棒学习算法。一些工作通过设计一些数据筛选或正则化策略降低模型对标签噪声的拟合能力,使模型偏向符合真实标签分布的监督信息,从而降低输入图像到学习目标之间的错误对应关系对模型最终学到的映射关系的影响,使得学习过程更加可信,最终获得可靠的分割模型。
1.2 学习配准以学习分割范式
学习配准以学习分割范式(Learning registration to learn segmentation, LRLS)是一个高效的少标签医学图像分割范式。它利用配准弥合图像间的空间差异,将少量标签映射到大量的无标签数据上,从而生成大量伪标签数据,以训练具有良好泛化能力的分割模型。但是它受限于多样性,真实性和鲁棒性,限制了该范式下对数据潜力的挖掘。如图3所示,**1)真实性的限制:**图像配准的变换过程中会带来图像结构的形变,一些过度的形变往往会导致图像中器官对象真实的拓扑结构被破坏,使得生成图像失真。而配准失误又将带来图像间的不对齐问题,一些在LRLS中基于融合图像风格的方法[4,5]将导致生成数据语义区域信息的失真。**2)多样性限制:**在数据生成过程中,该范式直接配准标签到无标签数据上,并将最终将数据集增大到无标签图像的规模。一旦无标签数据量较少时,依旧极大地限制着训练数据的多样性。**3)鲁棒性的限制:**在分割学习的过程中,已有的LRLS工作缺乏抵抗生成数据中的标签噪声的能力,使得模型极容易拟合到不可信的从图像到分割结果的映射关系,得到不可靠的分割模型。因此一个科学问题在这项工作中被具体研究:如何促进该范式学习的多样性,真实性和鲁棒性,构建可信的数据生成和数据学习过程,获得可靠的少标签医学图像分割模型?
图 3 LRLS范式在医学图像少标签学习中的的基本架构图以及在各环节存在的缺乏多样性,真实性和鲁棒性的限制。
2. 方法
在这项工作中,我们提出了一个新的LRLS学习模型,我们称之为BRBS(Better registration for better segmentation),通过对模型的数据生成和学习过程中的真实性,多样性和鲁棒性进行设计,构建了一种可靠的少标签医学图像分割模型。更多方法细节请读者参考原论文[1]。
2.1 BRBS模型用于少标签医学图像分割
2.1.1 知识一致性约束策略提升真实性
图 4 知识一致性约束的算法架构图。通过同时约束配准模型学习拓扑保持和语义对齐的形变,提高数据生成的真实性。
知识一致性约束利用“对齐的图像其语义区域是对齐的“和”一个拓扑保持的形变是可逆和平滑的“这两个先验知识,构建了精准的配准学习框架,进而促进更加真实的数据生成。如图4所示,为了获得配准的可逆和平滑性,我们设计了一种双向一致性约束方法(BiC, LBiC),同时学习图像 F 到图像 M 的正向配准,和图像 M 到图像 F 的逆向配准(约束配准后形变场的平滑性 Lsmooth和形变后图像与目标图像的相似性 Lsim),并且学习双向配准形变场间互逆映射的一致性LI。为了获得配准图像对之间语义区域的对齐准确性,我们设计了一种语义一致性约束方法(SeC,LSeC),利用分割网络提取图像中的语义区域,并约束形变后图像间的语义区域对齐,从而获得语义区域级的精准对齐。
2.1.2 空间风格采样提升多样性
空间风格采样建模了图像间风格与空间关系的变换路径,利用这两个路径构建图像间风格和空间结构两个维度的连续变换过程,通过采样的方式进行图像生成,大大提升了生成图像的多样性。如图5所示,同时从数据集中获得两个无标签图像Uj,Uk 和一个标签图像Ai,利用配准模型将Ai配准到 Uj得到图像间的形变场 ϕi2j以描述这两个图像间的空间变换路径。再利用配准模型将 Uk配准到 Ai,得到空间上与 Ai对齐的图像 Uk2i, Uk2i与 Ai做差获得风格位移图 ψi2k以描述两个图像间的风格变换路径。接着对形变场 ϕi2j和风格位移图 ψi2k进行形变程度和风格变换程度的采样,以生成新的对 Ai的形变和风格变换类型。最后对 Ai进行风格和空间的变换产生新的伪标签图像。
图 5空间风格采样策略的算法结构图。通过同时采样图像间的风格和空间变换来提高生成图像的多样性。
2.1.3混合不对齐正则化提升鲁棒性
图 6混合不对齐正则化通过约束不对齐区域拟合线性分布,从而降低对错误信息的拟合程度,提高模型学习带噪声标签数据的鲁棒性。
混合不对齐正则化约束分割网络在生成图像的错误对齐区域学习更加线性的表征,从而降低对这些不准确信息的拟合程度,最终对符合真实标签分布的区域拥有更好的拟合能力。如图6所示,我们从无标签数据集中采样两个图像 XA 和XB,我们使用分割模型分别提取两个图像的分割结果 yA 和yB,并通过一个 [0,1] 的随机值 来加权混合两个原图XAB=γXA+(1−γ)XB,和两个分割结果 yAB=γyA+(1−γ)yB 。通过约束分割网络对混合的图像 XAB 的分割结果 y^AB与混合的分割结果 yAB一致。如图6所示,该过程约束两个数据点之间不对齐的区域拟合一个线性函数,从而降低模型对这块不对齐区域的拟合程度,实现对不对齐区域的有效正则化。
S(γXA+(1−γ)XB)⇔γS(XA)+(1−γ)S(XB)
2.2 BRBS模型表现背后的算法直觉
BRBS模型通过从真实性,多样性和鲁棒性的角度,克服了LRLS范式在少标签医学图像分割学习中的数据分布缺陷和标签分布缺陷,构建了可信的数据生成和数据学习过程,从而仅使用少量标签即获得了可靠的分割模型。如图7所示,**a1)**在少标签设置下,原始数据分布呈现少量的标签数据和一些无标签数据。**a2)**由于在数据生成过程中缺少真实性和多样性,尽管LRLS范式带来了更多有标签信息的训练数据,但是生成的数据依旧是稀疏的,并伴随着失真数据,输入分布依旧存在较大的缺陷。**a3)**当加入我们的风格空间风格采样(S3P),由于对变换过程的密集采样,提高了数据生成的多样性,因此获得了更加丰富的生成数据,从而覆盖更广范围的数据分布。然而,在提高多样性的同时,生成失真图像的数量也跟着被增加了,引入了更多的失真的数据分布。**a4)**再加入我们的知识一致性约束(KCC),进一步提高了生成数据的真实性,使得失真的数据数量减少,使模型生成过程中获得了更加贴合真实分布的数据分布。**b1)**在学习这些生成的数据时,由于生成数据的伪标签存在错误的监督信息,直接的学习这些错误的图像到标签的映射关系将拟合到不可信的分割预测。**b2)**我们的混合不对齐正则化(MMR)避免了模型对这些信息的过拟合,使其更趋向于拟合到符合真实的从数据到标签的分割模型,得到可靠的分割模型。
图 7 BRBS模型表现背后的算法直觉,通过增强数据生成和数据学习中的多样性,真实性和鲁棒性,从而获得可靠的分割模型。
3. 实验设计及结果
该文章使用公开数据集,在两个任务上进行了实验,包括少样本心脏结构分割和少样本脑结构分割。我们分别进行了只是用1个标签的实验和5个标签的实验。从表1的定量分析中,可以看出,我们的BRBS仅使用1个或者5个标签就可以在两个任务上获得可靠的分割结果,与其他LRLS的方法相比,BRBS取得了显著的性能提升。图8的定性分析可以发现,我们的BRBS同时对较大的结构和比较脆弱的小结构都拥有极好的分割品质。更多分析,请参考原文[1]。
表格 1对比分析实验,与其他LRLS范式的模型相比,我们的BRBS在心脏结构分割和脑组织分割任务的1-shot和5-shot设置中,都取得了显著的性能提升。
图 8定性分析结果图。我们的BRBS同时对较大的结构和比较脆弱的小结构都拥有极好的分割品质。
主要参考文献
[1] He, Y., Ge, R., Qi, X., Chen, Y., Wu, J., Coatrieux, J. L., … & Li, S. (2022). Learning Better Registration to Learn Better Few-Shot Medical Image Segmentation: Authenticity, Diversity, and Robustness. IEEE Transactions on Neural Networks and Learning Systems, 2022, doi: 10.1109/TNNLS.2022.3190452.
[2] He, Y., Li, T., Ge, R., Yang, J., Kong, Y., Zhu, J., … & Li, S. (2021). Few-shot learning for deformable medical image registration with perception-correspondence decoupling and reverse teaching. IEEE Journal of Biomedical and Health Informatics, 26(3), 1177-1187.
[3] He, Y., Li, T., Yang, G., Kong, Y., Chen, Y., Shu, H., … & Li, S. (2020, August). Deep complementary joint model for complex scene registration and few-shot segmentation on medical images. In European Conference on Computer Vision (pp. 770-786). Springer, Cham.
[4] Y. Ding, X. Yu, and Y. Yang, “Modeling the probabilistic distribution of unlabeled data for one-shot medical image segmentation,” in AAAI, 2021, pp. 1246–1254.
[5] Liang, W., Tadesse, G. A., Ho, D., Li, F. F., Zaharia, M., Zhang, C., & Zou, J. (2022). Advances, challenges and opportunities in creating data for trustworthy AI.Nature Machine Intelligence, 1-9.
[6] A. Zhao, G. Balakrishnan, F. Durand, J. V. Guttag, and A. V. Dalca, “Data augmentation using learned transformations for one-shot medical image segmentation,” in Proceedings of the IEEE conference on computer vision and pattern recognition, 2019, pp. 8543–8553.