点击下方卡片，关注「AI视界引擎」公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
本文介绍了MK-UNet，这是一种面向超轻量级、多核U形CNN的范式转变，专为医学图像分割而定制。MK-UNet的核心是作者设计的多核深度卷积块(MKDC)，它能够通过多个核熟练地处理图像，同时捕获复杂的多分辨率空间关系。MK-UNet还通过复杂的注意力机制强调图像的显著特征，包括通道注意力、空间注意力和分组门控注意力。
作者的MK-UNet网络仅拥有0.316M参数和0.314G FLOPs的适度计算开销，不仅代表了极其轻量级的解决方案，还显著提高了分割性能，在六个二元医学成像基准测试中提供了比最先进(SOTA)方法更高的准确性。
具体而言，MK-UNet在DICE分数上优于TransUNet，同时参数和FLOPs分别减少了近
和
。同样，与UNeXt相比，MK-UNet表现出卓越的分割性能，将DICE分数提高了高达
的幅度，同时使用的参数量减少了
。
作者的MK-UNet也以更低的计算资源优于其他最近的轻量级网络，如MedT、CMUNeXt、EGEUNet和Rolling-UNet。这种性能的飞跃，加上计算效率的显著提升，使MK-UNet成为资源受限环境（如即时诊断设备）中实时、高保真医学诊断的无与伦比的解决方案。
作者的实现可在https:// github.com/SLDGroup/MK-UNet获取。
1 引言
医学图像分割领域通过U形卷积神经网络(CNN)架构[12, 26, 31, 44]的发展经历了变革性增长，例如UNet [31]、
[44]、AttnUNet [26]、PraNet [12]、UACANet [18]和DeepLabv
[8]。这些模型在医学图像分割方面表现出色，能够精确分割关键肿瘤、病变或息肉。集成到这些架构中的注意力机制[12, 26, 41]有助于优化特征图，从而增强像素级分类。然而，这些模型（包括那些带有注意力机制的模型）的大量计算需求限制了它们在资源受限环境（如即时诊断）中的适用性。
vision transformers [4, 6, 27–30, 39]的引入，包括TransUNet [6]、SwinUNet [4]和MedT [39]，标志着一种转变，即利用self-attention来捕捉图像内的长距离依赖关系，以获得全面的global view。然而，transformers往往忽视像素之间关键的局部空间关系，而这些关系对于精确分割是必不可少的。此外，transformers在计算和融合attention与卷积机制时通常具有较高的内存和计算需求，这限制了它们的实际部署。
近年来，许多轻量级架构如UNeXt [38]、CMUNeXt [37]、MALUNet [32]、EGE-UNet [33]和Rolling-UNet [21]，通过结合CNN和多层感知器（MLP）的优势来弥合这一差距。然而，这些架构大多是为不太复杂或易于分割的应用而设计的，例如皮肤病变、超声中的乳腺癌以及微观细胞核/结构分割。因此，由于息肉在形状、大小和纹理上的高度可变性，这些架构在息肉分割等具有挑战性的应用中表现不佳。
为了解决这些计算和精度挑战，作者提出了MK-UNet，这是图像分割领域的一项重大突破，它利用多核视角（即对于
，有
或
）和深度卷积的优势，来解决现有基于CNN和transformer的模型中固有的计算复杂性和挑战。深度卷积大幅减少了计算负载，使网络更加高效，同时不牺牲捕获图像内细节特征的能力。此外，作者的多核特性使模型能够有效处理相同或不同感受野的特征表示，从而在各种应用中对复杂图像进行更加稳健和全面的分析。而且，通过仅在作者的解码器中融入复杂的卷积多焦点（通道和空间）注意力机制，进一步通过捕获图像的显著特征来优化特征图。作者注意到，作者的网络在两种场景下对分割都是有效的，无论感兴趣区域在大小和形状上差异显著还是保持相对均匀。通过整合这些新理念，MK-UNet在计算效率和分割精度之间取得了良好的平衡，从而提供了一个超轻量级模型，不仅在DICE分数上超越了重量级同类模型的性能，而且使用的#Params和#FLOPs显著更少。作者的贡献如下：
- 新型轻量级多核UNet：作者提出了一种新的端到端U型网络MK-UNet，用于医学图像分割，该网络使用轻量级多核卷积对图像进行编码，以捕获多分辨率空间表示。MK-UNet还使用多核卷积注意力逐步细化多分辨率空间表示。值得注意的是，作者的MK-UNet-T（tiny版本）分别仅有0.027M参数量和0.062G浮点运算次数，却提供了SOTA性能。此外，MK-UNet（标准版）仅有0.316M参数量和0.314G浮点运算次数。极低的模型大小（参数量）和计算量（浮点运算次数）使作者的MK-UNet易于在即时诊断或资源受限环境（如移动或边缘设备）中部署。
- 轻量级多核反向残差：作者引入了MKIR，一种新的多核反向残差块，它使用多个核执行深度卷积（即
或
，其中
。作者的编码器使用MKIR块提取特征；这一选择源于有效处理和编码医学图像中多样和复杂结构的需求，从而以最小的计算成本提供丰富的表示。 3. 轻量级多核倒置残差注意力：作者提出了多核倒置残差注意力（MKIRA），这是一个通过抑制无关区域来细化和增强多尺度显著特征的新模块。在作者的解码器中，MKIRA通过关注关键特征通道并突出图像中的重要空间区域来增强特征区分度。这确保了解码器能够仅通过关注编码特征中最关键的方面来重建精确准确的分割图。
- 在各种任务和基准上改进的性能：作者通过实验证明，在属于四个不同任务的六个二值医学图像分割基准上，MK-UNet与SOTA方法相比显著提高了医学图像分割的性能，同时计算成本显著降低（如图1所示）。
- 相关工作
2.1. 卷积神经网络 (CNNs)
CNN的出现标志着医学图像分割领域的重大转变[8, 12, 18, 26, 31, 44]。先驱性工作如全卷积网络(FCNs)[23]为端到端分割模型奠定了基础。FCNs用卷积层替代全连接层，从而实现了像素级预测和图像空间层次结构的高效学习。U-Net[31]因其具有 Shortcut 的编码器-解码器架构而成为医学图像分割的基石。这种设计有效地将来自编码器的高分辨率特征与来自解码器的上下文信息相结合，即使在有限的训练数据下也能实现精确分割。用于像素级分割任务的U型架构的精密设计促使作者在所提出的网络中选择U型设计。
U-Net的成功启发了众多变体和改进。受ResNet [13]中残差学习的启发，ResUNet [43]使用残差块来促进梯度 Stream 并提高收敛性，解决了深度网络中的梯度消失问题。Zhou等人[44]提出了
，它使用密集嵌套的 Shortcut 来进一步增强特征传播并提高分割精度。AttnUNet [26]融入了注意力机制，以关注特征图中的相关区域，通过抑制不相关的背景噪声来提高分割性能。Fan等人[12]提出了用于精确息肉分割的PraNet，该网络采用并行反向注意力和边缘引导来细化分割边界。UACANet [18]使用不确定性感知机制来提高分割结果的可靠性和鲁棒性。DeepLabv
[8]整合了空洞卷积和空间金字塔池化来捕获多尺度上下文信息。ACC-UNet [16]采用自适应上下文捕获机制，根据输入图像动态调整感受野。
2.2. 视觉Transformer
Vision Transformers (ViTs) [11, 22] 已经成为CNN的强大替代方案，通过利用自注意力机制[4, 6, 27–30, 39]为医学图像分析任务提供了新的范式。通过结合CNN在局部特征提取方面的优势和Transformer在捕获长距离依赖关系方面的能力，TransUNet [6] 在医学图像分割中取得了卓越的性能。SwinUNet [4] 基于Swin Transformer [22]架构引入，该架构利用移位窗口实现分层特征表示，从而实现高效计算。MedT [39] 是一种专为医学图像分割设计的轻量级Transformer模型，采用门控轴向注意力机制来关注相关区域并降低计算复杂度。Rahman等人提出了CASCADE [27]，这是一种使用标准卷积的级联注意力解码网络。
最近，EMCAD [30] 引入了一种基于深度卷积的多尺度解码器。尽管CASCADE和EMCAD在医学图像分割中表现良好，但它们的分割精度和计算复杂度完全取决于它们所使用的现有预训练transformer编码器的强度和复杂性，因此使它们不太适合资源受限的环境。相比之下，作者 Proposal 设计一种极其高效（超轻量级）的端到端（包括编码器和解码器）架构，使用多核技巧（其中，
或
对于
结合深度卷积。
2.3. 轻量级CNNs
最近的努力集中在使CNN在实时和资源受限的环境中更加高效。MobileNets [14] 和 EfficientNets [36] 分别引入了深度可分离卷积和复合缩放，以创建具有竞争性能的轻量级模型。此外，已经开发了几种新颖的轻量级架构，以进一步提高医学图像分割的效率 [21, 33, 37, 38]。UNeXt [38] 使用混合卷积和transformer块来高效捕获局部和全局特征，在保持计算效率的同时提高分割精度。CMUNeXt [37] 结合了卷积和多尺度特征来增强分割性能。
EGE-UNet [33] 集成了边缘引导机制来优化分割边界。Rolling-UNet [21] 融入了滚动卷积块，以增强模型捕获长距离依赖关系的能力。
- 方法
接下来，作者描述作者的核心构建模块。然后，作者通过将这些模块集成到 Baseline UNeXt [38]（图2a中的绿色框）中来介绍作者的MK-UNet架构。
3.1. 多核倒置残差 (MKIR)
作者首先引入多核倒置残差(MKIR)块来生成和优化特征图(图2c)。通过利用多个(相同或不同)的核大小，MKIR能够更好地理解细粒度细节和更广泛的上下文，从而实现对输入的全面表示。如图2c所示，该过程通过逐点卷积
、批归一化
[17]和
激活[19]来扩展通道数(即扩展因子
)。接下来是多核深度卷积
，用于捕获特定于应用的复杂空间上下文。随后的逐点卷积
和
恢复原始通道数。MKIR(公式1)在确保丰富特征表示的同时显著降低了计算成本：
其中多核
的
在公式2和图2d中定义：
其中
。这里，
是一个使用
核的深度可分离卷积。为了解决深度可分离卷积中的通道独立性问题，采用了通道混洗
来确保通道间的信息 Stream 。作者的
与
（在 EMCAD [30] 中）在核心理论概念上有所不同。作者的多核技巧对于
同时支持
（相同大小的核）和
（不同大小的核），而传统的多尺度（仅
）设计 [20, 30, 35] 则不然，因此允许自适应的上下文提取。这种概念上的区别使 MK-UNet 能够基于特定应用的需求调整核大小（例如，大核用于大物体，小核用于小物体，或混合用于两种物体的分割）。
3.2. 多核倒置残差Attention (MKIRA)
作者提出了一种轻量级的多核倒置残差注意力机制MKIRA，用于优化特征图
。MKIRA使用通道注意力
[15]来关注相关通道，使用空间注意力
[7]来捕获局部上下文，并使用多核倒置残差
来丰富特征图同时捕获上下文关系。MKIRA（图2b）在公式3中给出：
通道注意力
：作者使用
[15]通过应用自适应最大池化
和平均池化
来压缩空间信息 [27]，从而增强相关特征；随后通过带有
激活 [25]的点卷积
进行通道缩减（
），再通过
进行扩展。Sigmoid
激活生成注意力权重，然后通过Hadamard积
将其应用于输入，同时关注关键特征图。
在公式4中定义：
空间注意力（SA）：作者使用
[7]来关注特定的图像区域以突出关键特征，这对于精确分割至关重要。
聚合最大值
和平均值
通道值以突出局部细节，然后它采用大核
卷积
来加强上下文连接。Sigmoid
激活函数导出注意力权重，通过Hadamard积
应用于输入
，确保有针对性的细化。
由公式5导出：
3.3. 分组注意力门控 (GAG)
作者使用分组注意力门（
，图2e），它将特征图与注意力系数混合，以增强相关特征并抑制不相关特征。通过利用来自高分辨率特征的门控信号，
引导信息 Stream ，从而提高医学图像分割的准确性。与使用
卷积处理信号的Attention UNet [26]不同，yinqing-MK-UNet_2509分别对门控
和输入
特征图应用
组卷积
[30]。卷积后，特征经过批量归一化
并通过加法组合，然后进行ReLU
激活。随后，
卷积和批量归一化
产生一个统一的特征图，该特征图经过Sigmoid激活
后生成注意力系数。这些系数调整输入特征
，并创建一个注意力增强的输出。
在公式6中定义：
3.4. 多核UNet (MK-UNet)
作者的MK-UNet在五个编码和解码阶段采用多核方法来生成高分辨率分割图，如图2a所示。每个编码阶段使用多核反向残差(MKIR)块来生成
特征图，然后通过最大池化进行下采样，同时保留关键信息。最终编码阶段的输出通过解码器初始阶段的多核反向残差注意力(MKIRA)块，显著细化特征图。然后使用双线性插值对这些特征图进行上采样，用于后续的解码阶段。解码器阶段使用分组注意力门(GAG)将 Shortcut 与细化特征整合，然后进行加法聚合。生成的特征图通过MKIRA块进行细化，并上采样（仅使用双线性
，不使用卷积）以与后续阶段对齐。
最后四个阶段的分割头(SHs)输出分割图
、p3、
和
。作者将特征图
视为最终预测，并通过使用Sigmoid函数进行二值分割或使用Softmax激活函数进行多类分割来获得最终的分割输出。对于所有二值分割任务，作者仅优化最终预测
的损失。然而，作者建议在多类分割中使用深度监督。
- 实验与结果
4.1. 数据集
作者在涵盖四种分割任务的六个数据集上评估了MK-UNet的效能，包括乳腺癌（BUSI [1]，647张图像：437张良性图像和210张恶性图像）、息肉（ClinicDB [2]包含612张图像，ColonDB [40]包含379张图像）、皮肤病变（ISIC18 [9]，2,594张图像）以及细胞核/结构分割（DSB18 [3]包含670张图像，EM [5]包含30张图像）。这些数据集收集自各种成像中心，在图像特性上提供了广泛的多样性，确保了全面的评估。所有数据集均采用80:10:10的训练-验证-测试分割比例，并报告了测试集的DICE分数。
4.2. 实现细节
作者的网络使用Pytorch 1.11.0进行开发和评估，在配备48GB RAM的单个NVIDIA RTX A6000 GPU上运行。基于消融研究，作者在MKDC中利用多尺度核[1, 3, 5]。该架构在MK-UNet网络中采用了一系列并行深度卷积，在所有实验中统一使用[16, 32, 64, 96, 160]的通道配置，除非另有说明。模型优化通过AdamW [24]优化器实现，学习率和权重衰减均设置为
。训练持续200个轮次，批次大小为16，在此期间作者保存达到最高DICE分数的模型。
对于BUSI [1]、ISIC18 [10]、EM [5]和DSB18 [3]数据集，图像尺寸设置为
像素，而对于ClinicDB [2]和ColonDB [40]数据集，分辨率调整为
像素。作者采用多尺度训练方法，尺度为
，并将梯度裁剪强制设置为0.5。作者不应用任何形式的数据增强，并使用混合损失函数，该函数将(1:1)加权的BinaryCrossEntropy (BCE)与加权的Intersection over Union (IoU)损失相结合。
4.3. 结果
4.3.1. 乳腺癌分割
作者专注于使用BUSI数据集进行超声图像中的乳腺癌分割。作者的MK-UNet模型明显优于现有方法，在计算需求极低的情况下实现了最高的DICE分数
。这一成就突显了yinqing-MK-UNet_2509的效率和有效性，特别是与更复杂的模型相比时。例如，TransUNet尽管拥有更多的参数（105.32M）和FLOPs（38.52G），却仅实现了较低的DICE分数
。类似地，ACC-UNet拥有
参数和38.0G FLOPs，以
的DICE分数紧随其后。此外，与计算量相当的模型UNeXt相比，作者的MK-UNet表现出显著改进，以
更少的#Params超越了它
的差距。
4.3.2. 皮肤病变分割
在皮肤病变分割的ISIC18数据集上，作者的MK-UNet网络表现出值得称赞的性能，DICE分数为
。这一结果不仅展示了yinqing-MK-UNet_2509的鲁棒性，还展示了其处理皮肤病变图像中复杂变化的能力。作者注意到TransUNet以略高的分数
领先，但其成本显著更高（#Params和#FLOPs分别大了
和
）。作者的MK-UNet在保持最小计算需求的同时实现了接近顶级的性能——仅需0.316M #Params和0.314G #FLOPs。
4.3.3. 息肉分割
在Clinic和Colon数据集上的息肉分割任务中，作者的MK-UNet模型分别以
和
的领先得分表现出色，并且具有显著较低的计算开销。这一性能超越了DeepLabv
和ACC-UNet等显著的竞争对手，尽管这些模型消耗更多资源，但它们的精度仍无法与MK-UNet相媲美。具体而言，DeepLabv3
在Clinic和Colon数据集上分别达到
和
，但其参数和FLOPs分别大了
和
，而ACC-UNet虽然紧随其后但仍有所不及。
4.3.4. 微观细胞核/结构分割
在DSB18和EM数据集上进行细胞核/结构分割时，作者的MK-UNet网络展现了卓越的准确性，分别达到了
和
的DICE分数。相比之下，尽管TransUNet和UNeXt等其他模型采用了重型设计且计算需求更高，但它们的DICE分数并未超过MK-UNet。例如，TransUNet在DSB18上获得了
的较低分数，在EM上获得了
的分数，而UNeXt则落后MK-UNet达
。
4.4. 定性结果
在图3中，作者报告了代表性测试图像的乳腺肿瘤、皮肤病变、息肉和细胞分割的分割图。在乳腺肿瘤分割中，UNet、
和 UNeXt 显示出较大的错误分割，而 TransUNet 和作者的 MK-UNet 则产生了近乎完美的分割图。类似地，在皮肤病变分割中，UNet、ResUNet、
、AttnUNet、DeepLab
、PraNet、SwinuNet 和 UNeXt 遗漏了部分病变（红色矩形框内）。然而，UACANet、TransUNet、ACC-UNet 和作者的 MK-UNet 能够很好地分割那个具有挑战性的区域。作者的 MK-UNet 也能正确分割息肉，而所有其他方法错误地分割了一个
其他区域作为息肉。总的来说，作者的MK-UNet在所有四项任务中产生了最佳的重叠分割图。作者的MKUNet在计算预算非常低的情况下展现出这种全面性能的原因是使用了多核深度卷积以及门控和局部注意力机制。
- 消融研究
本节描述了三个关键的ablation studies。更多的ablation studies结果在附录中给出。
5.4. 编码器中MKIR与MKIRA的有效性
- 结论
在本文中，作者提出了MK-UNet，这是医学图像分割领域的一项重大进展。MK-UNet解决了长期存在的挑战，即在保持高分割准确性的同时实现计算效率的平衡。
通过利用深度卷积和多核处理，MK-UNet以极低的计算开销超越了TransUNet和UNeXt等模型（与TransUNet相比，参数量减少近
，复杂度降低
；与UNeXt相比，参数量减少
）。
计算资源的显著减少，加上性能的提升，使MK-UNet成为在计算资源有限的环境中部署的最佳选择，例如在即时医疗设备中。
参考
[1]. MK-UNet: Multi-kernel Lightweight CNN for Medical Image Segmentation
点击上方卡片，关注「AI视界引擎」公众号