UltraLightUNet:重新思考带有多核轻量级卷积的U形网络用于医学图像分割

1、论文基本信息

1.1、论文名称:UltraLightUNet: Rethinking U-Shaped Network with Multi-Kernel Lightweight Convolutions for Medical Image Segmentation

1.2、论文下载链接:UltraLightUNet: Rethinking U-shaped Network with Multi-kernel Lightweight Convolutions for Medical Image Segmentation | OpenReview

1.3、论文代码链接:截止到目前源码尚未开源

1.4、论文其他信息:本文为ICRL 2025的投稿论文,处在匿名评审阶段,因此暂时无作者信息

2、摘要

        在本文中,作者介绍了UltraLightUNet(2D和3D),这是一种用于医学图像分割的超轻量级、多核U形网络。UltraLightUNet的核心是由一种新的多核反向残差(MKIR)块组成,该块可以通过多个内核高效处理图像,同时捕获复杂的空间关系。此外,作者的多核反向残差注意力(MKIRA)块通过新的复杂卷积多焦点注意力机制来提炼和强调图像中的重要特征。UltraLightUNet在编码器中策略性地使用MKIR块进行特征提取,在解码器中使用MKIRA块进行特征提炼,从而确保在每个阶段都有针对性地增强特征。凭借仅有0.316M个参数和0.314G FLOPs,UltraLightUNet提供了一种超轻量级但功能强大的分割解决方案,超越了11个医学影像基准测试中的最新方法(SOTA)。值得注意的是,UltraLightUNet在DICE得分上超越了TransUNet,同时使用的参数减少了333倍,FLOPs减少了123倍。与UNeXt相比,UltraLightUNet在DICE得分上提高了6.7%,参数减少了4.7倍。UltraLightUNet还在使用显著更少的参数和FLOPs的情况下,超越了最近的轻量级模型,如MedT、CMUNeXt、EGE-UNet、Rolling-UNet和UltraLight VM UNet。此外,作者的3D版本,UltraLightUNet3D-M(1.42M个参数和7.1G FLOPs),在MSD前列腺和FETA基准测试上超越了SwinUNETR(62.19M个参数,328.6G FLOPs)和nn-UNet(31.2M个参数,110.4G FLOPs)。这一卓越的性能,加上巨大的计算收益,使得UltraLightUNet成为资源受限环境中实时、高保真医学诊断的理想解决方案,例如在临床护理点服务中。作者将在论文接受后公开源代码。

3、主要贡献

  • 新型超轻量级UNet架构:提出了UltraLightUNet,一个在参数数量和计算复杂度上极为精简的网络,同时保持了优异的分割性能。

  • 多核反向残差块(MKIR):引入了MKIR,一种新型的多核卷积块,能够以较少的计算资源捕获多尺度特征。

  • 多核反向残差注意力(MKIRA):提出了MKIRA块,通过注意力机制增强了特征的提炼,提升了模型对关键特征的识别能力。

  • 分组注意力门(GAG):开发了GAG,通过分组卷积和注意力机制提高了特征融合的效率,增强了模型对医学图像中关键区域的分割精度。

4、研究背景

        医学图像分割对于诊断疾病至关重要,但现有模型(如 UNet、ResUNet、TransUNet 等)在资源受限环境中难以应用,因为它们计算量大、参数多。 轻量级模型(如 UNeXt、CMUNeXt 等)虽然计算量较小,但在分割复杂结构时性能不佳。3D 医学图像分割模型计算量更大,难以应用于临床环境。

5、网络设计

5.1、模型架构图

5.2、多核反向残差(MKIR)

        作者首先介绍多核反向残差(MKIR)块来生成和提炼特征图(见图2c)。通过利用不同的内核大小,MKIR允许更好地理解细节和更广泛的上下文,从而实现对输入的全面表示。如图2c所示,该过程首先通过点卷积PWC1、批量归一化BN和ReLU6激活来扩展#channels(即,扩展因子=2)。然后是多核深度卷积MKDC来捕获多尺度和多分辨率空间上下文。随后的点卷积PWC2和BN恢复原始的#channels。MKIR定义在方程1中,显著降低了计算开销,同时确保了丰富的特征表示:

其中MKDC对多个内核(K)定义在方程2和图2d中:

其中。这里是内核k × k的深度卷积。为了解决深度卷积中的通道独立性问题,使用通道混洗(CS)以确保通道间信息流动。

5.3、卷积多焦点注意力(CMFA)

        作者的CMFA(见图2f)利用统一的注意力机制,有效捕获通道和空间特征(Rahman & Marculescu,2023a),从而优化网络对图像中关键方面的关注能力,同时抑制无关细节。作者首先通过应用自适应最大池化(AMP)和平均池化(AAP)来增强相关通道,这允许网络在局部结构变化中保持鲁棒性。池化输出然后通过一系列点卷积(PWC)来降低(r = 16)维度并激活ReLU,随后是第二个PWC层进行扩展。Sigmoid激活生成注意力权重,然后与输入元素相乘(⊛),从而强调重要通道。这个注意力过程定义在方程3中:

随后,为了捕获空间依赖性并进一步提炼特征图,作者对通道进行池化操作以生成两个空间描述符:Channelmax(x)和Channelavg(x)。通过对连接的池化值应用大内核卷积(LKC),作者捕获更广泛空间上下文中的上下文关系,从而加强网络对图像重要区域的关注。提炼的特征图如方程4所示:

 本质上,结合平均和最大池化有助于平衡对高强度(max)区域和整体特征一致性(mean)的关注。同样,整合通道和空间注意力有助于在复杂场景中精确重建和分割,从而提高分割性能。

5.4、多核反向残差注意力(MKIRA)

        作者的新MKIRA块(见图2b)通过利用卷积多焦点注意力机制(CMFA)和多核反向残差(MKIR)有效地提炼特征图。CMFA的使用增强了网络关注关键通道和空间区域的能力,确保最重要的特征被增强,不相关信息被抑制。这种双重注意力机制有助于提高特征辨别力和表示力,特别是在重要结构可能显著变化的具有挑战性场景中。此外,MKIR块的整合进一步丰富了特征图,通过多个感受野捕获上下文关系。综合这些组件,网络在保持高精度的同时最小化计算开销。MKIRA定义在方程5中:

5.5、分组注意力门(GAG)

        作者设计了一个新的分组注意力门(GAG,见图2e),它混合特征图和注意力系数以增强相关特征并抑制不相关特征。利用来自更高分辨率特征的门控信号,GAG指导信息流,从而提高医学图像分割精度。与Attention UNet(Oktay等人,2018)不同,后者使用1×1卷积处理信号,作者的方法对门控(g)和输入(x)特征图分别应用3×3组卷积。卷积后,特征经过批量归一化(BN)并通过加法组合,随后是ReLU激活。随后,1×1卷积和批量归一化(BN)产生统一的特征图,经过Sigmoid激活(σ)后生成注意力系数。这些系数调整输入特征x,并创建注意力增强的输出。GAG定义在方程6中:

6、实验结果

6.1、对比实验

6.2、消融实验

7、总结

        在本文中,我们提出了UltraLightUNet,这是一种用于医学图像分割的新网络,通过超轻的设计实现了高精度。 UltraLightUNet在多个基准测试中优于最先进的模型,同时保持显着降低的计算足迹。 例如,UltraLightUNet在DICE分数上超过TransUNet,其Params减少了333x,FLOPs减少了123x。 同样,与UNeXt相比,UltraLightUNet将分割精度提高了6.7%,同时使用的Params减少了4.7倍。 我们的设计有效地捕捉了复杂的空间关系,并细化了显著特征,从而使其成为资源受限环境的理想选择,例如即时护理服务,在这些环境中,实时、高保真诊断是必不可少的。 在未来,我们计划探索该网络在其他密集预测任务中的适用性,如2D或3D图像重建、翻译、增强和去噪。