极市导读
所提出的方法可以即插即用的方式集成到现有的UNet架构中,并且不增加计算成本。实验结果表明,所提方法在4种医学图像分割数据集上持续提高了标准UNets的性能。 >>加入极市CV技术交流群,走在计算机视觉的最前沿
文章《SelfReg-UNet: Self-Regularized UNet for Medical Image Segmentation》由Wenhui Zhu、Xiwen Chen、Peijie Qiu等人撰写,发表于顶级学术会议MICCAI。通过Unet中自信息提升性能, 可以即插即用在大部分主流的Unets模型中.
Paper link: https://arxiv.org/pdf/2406.14896
Open-source code: https://github.com/ChongQingNoSubway/SelfReg-UNet
这篇文章提出了一种自正则化的即插即用的技术,旨在改善标准UNet的性能,特别是针对医学图像分割任务。UNet自推出以来,在医学图像分割任务中表现优异。然而,尽管有许多研究致力于提升UNet的性能,却鲜有深入分析UNet在医学图像分割中的底层模式的研究。在本文的研究中,我们深入分析了UNet在医学图像分割任务中的表现,发现了两个影响其性能的重要因素:非相关特征学习和特征图中的冗余信息。为了解决这些问题,我们提出了在编码器和解码器之间平衡监督,并减少UNet中的冗余信息的方法。
非相关特征学习
在UNet中,编码器负责将输入图像的高维像素数据转换为低维的特征表示,而解码器则将这些低维特征重新转换为高维的分割掩码。在这个过程中,编码器和解码器的监督信号存在不对称性,即解码器接受了更多的直接监督信号,而编码器则较少。具体表现为:
1.解码器的强监督信号:解码器的每一层在训练过程中都能获得来自最终分割结果的直接反馈,这使得解码器能够逐步调整其参数以更好地匹配真实标签。
2.编码器的弱监督信号:相比之下,编码器只能通过间接的方式获得监督信号,即通过解码器的反馈进行调整。这种不对称性导致编码器容易学习到一些与分割任务无关的特征,影响最终的分割效果。
这种不对称性监督导致了编码器倾向于捕获一些非相关的特征,无法有效地聚焦于分割任务所需的语义信息。这些非相关特征不仅增加了模型的复杂度,还可能对解码器的性能产生负面影响。
特征图中的冗余信息
在UNet的训练过程中,我们还发现了特征图中的冗余信息问题。具体来说:
1.深层特征图的冗余:在UNet的深层特征图中,不同通道之间的特征往往具有较高的相似性。这种高相似性的特征意味着模型在不同通道上学习到了类似的信息,导致计算资源的浪费和模型的过拟合。
2.浅层特征图的多样性:与深层特征图相比,浅层特征图展示了更高的多样性,不同通道之间的特征相对独立,包含更多有用的信息。
特征图中的冗余信息不仅增加了模型的计算成本,还可能导致特征表示的效率降低,从而影响分割性能。
平衡监督和减少冗余的方法
为了应对上述两个问题,本文提出了两种创新的方法:语义一致性正则化(Semantic Consistency Regularization,SCR)和内部特征蒸馏(Internal Feature Distillation,IFD)。
1.语义一致性正则化(SCR):我们通过在UNet的编码器和解码器之间引入额外的监督信号,平衡它们之间的监督强度。具体方法是利用解码器最后一层的特征图(包含了最丰富的语义信息)对其他层提供额外的监督信号。这种方法通过平均池化和随机通道选择操作,确保不同层之间的特征对齐,从而减少非相关特征的学习。
2.内部特征蒸馏(IFD):为了减少特征冗余,我们从浅层特征向深层特征进行信息蒸馏。具体做法是将特征通道分为上半部分和下半部分,利用上半部分(浅层)特征指导下半部分(深层)特征的学习。这种方法通过引入L2范数作为距离度量,确保深层特征能够有效地学习到有用的上下文信息,减少冗余。
通过这两种方法,我们不仅平衡了编码器和解码器之间的监督信号,还显著减少了特征图中的冗余信息,从而提高了UNet在医学图像分割任务中的性能。这些改进不仅适用于标准的CNN-UNet结构,还可以推广到基于ViT(Vision Transformer)的UNet结构,实现广泛的应用。
SelfReg-UNet的有效性通过在多个公认的医学图像数据集上进行的广泛测试得到了证实。这些数据集包括Synapse多器官CT数据集和ACDC心脏MRI数据集,其中SelfReg-UNet在所有主要性能指标上均显著超越了现有的先进方法。
-
Synapse多器官CT数据集: 在这个包含多种腹部器官的数据集上,SelfReg-UNet通过其创新技术显著提高了模型的Dice相似系数(DSC),平均提高了3.49%。这一改进显示了SelfReg-UNet在处理具有复杂结构和细微纹理的器官图像时的优势。
-
ACDC心脏MRI数据集: 在心脏MRI图像分割任务中,SelfReg-UNet同样展示了出色的性能,尤其是在分割心室和心肌区域时,模型能够更精确地识别并界定各个部分,从而为临床诊断提供了更为可靠的支持。

公众号后台回复“数据集”获取100+深度学习各方向资源整理
极市干货

点击阅读原文进入CV社区
收获更多技术干货

