仅 50KB！| 上海交大 MICCAI2023 最新提出超强 U-Net 变体，参数和计算量分别暴降 494 和 160 倍

导读

今天为大家介绍一篇上海交大发表在 MICCAI 2023 的最新研究工作，一个称为Efficient Group Enhanced UNet, EGE-UNet 的模型，基于 U-Net 进行魔改，用于解决医学图像（尤其是皮肤病变）分割中面临的问题。由于它是针对移动健康应用开发的，解决了当前许多模型所面临的高参数和计算负载问题。

简单来说，EGE-UNet融合了两个主要模块：

• Group multi-axis Hadamard Product Attention module (GHPA)
• Group Aggregation Bridge module (GAB)

其中，GHPA 利用哈达玛积注意力机制（HPA），通过将输入特征进行分组，对不同轴进行 HPA 操作，从多个视角提取病变信息。这种方法是受到多头自注意力机制(MHSA)的启发，而 HPA 则可以减少模型大小，因为它的复杂性被设计为线性的，而且不同于 MHSA 的二次复杂性。

另一方面，GAB 通过分组聚合将不同规模的高级语义特征和低级细节特征以及解码器生成的掩码进行融合，从而有效提取多尺度信息，这对医学图像的分割至关重要。

最终，通过融合上述两个模块，作者们提出了EGE-UNet模型，实现了在参数和计算复杂度极低的情况下，优秀的分割性能。模型不仅关注性能提升，更注重在现实环境中的可用性。

根据文中的实验报道，EGE-UNet 在 ISIC2017 和 ISIC2018 两个主流的皮肤病分割数据集上超越了现有的最先进方法，且相比于TransFuse，模型在保持优秀分割性能的同时，参数和计算成本分别降低了 494 倍和 160 倍。据作者们所知，这是第一种参数数量限制在 50KB 的模型，足以证明了它的高效性和实用性!

方法

Framework

如图所示，EGE-UNet的设计沿用了 U 形架构，包括对称的编码器-解码器部分。编码器由六个 stage 组成，各阶段的通道数量为{8, 16, 24, 32, 48, 64}。前三个阶段采用了普通卷积，而后三个阶段使用提出的 GHPA 来从多视角提取表征信息。

相较于 UNet 中简单的 Skip connection 连接，EGE-UNet 在编码器和解码器之间的每个阶段都集成了 GAB。此外，模型还利用深监督生成不同规模的掩膜预测，这些预测用于损失函数并作为 GAB 的输入之一。通过这些高级模块的集成，EGE-UNet 在比先前的方法提升了分割性能的同时，显著减少了参数和计算负载。

GHPA 模块

为了解决由 MHSA 带来的二次复杂性问题，作者提出了具有线性复杂性的HPA操作。具体地，对于输入 $x$ 和一个随机初始化的可学习张量 $p$，我们首先使用双线性插值将 $p$ 调整到与 $x$ 相匹配的维度。然后，在 $p$ 上使用深度可分卷积（DW），然后在 $x$ 和 $p$ 之间进行哈达玛积操作以获得输出。

需要注意的是，仅使用 HPA 是无法从多个角度提取信息，这意味着分割效果很一般。因此，作者根据 HPA 提出了 GHPA，如算法1所示：

我们将输入在通道维度上均等地划分为四个组，并分别在第一、二、三组的高-宽、通道-高、通道-宽轴上执行HPA。对于最后一组，我们只在特征映射上使用 DW。最后，通过沿着通道维度连接四个组，然后应用另一个 DW以整合来自不同角度的信息。

GAB 模块

众所周知，对于密集预测任务来说，获取多尺度信息至关重要。因此，本文引入了GAB，它接受三个输入：

1. 低级特征
2. 高级特征
3. 掩码

如上图所示，首先，使用深度可分卷积（DW）和双线性插值调整高级特征的大小，以匹配低级特征的大小。其次，我们将两个特征图沿着通道维度划分为四个组，并将低级特征的一组与高级特征的一组连接起来，得到四组融合特征。对于每组融合特征，掩码被连接。接下来，对不同的组应用带有3个核大小和不同扩张率的扩张卷积，以便提取不同尺度的信息。最后，沿着通道维度连接四个组，然后应用核大小为1的普通卷积以实现不同尺度特征之间的交互。

最后，由于不同的 GAB 需要不同尺度的掩码信息，因此这里采用深监督计算不同阶段的损失函数，以生成更准确的掩码信息。

实验

在实验部分，EGE-UNet 在两个公开皮肤病变分割数据集（ISIC2017和ISIC2018）上进行了验证，展现出了超过现有方法的表现。在ISIC2017数据集上，相比于更大的模型，比如TransFuse，EGE-UNet不仅有更优的性能，而且显著减少了参数和计算量，分别达到494倍和160倍。

此外，对于其他轻量级模型，EGE-UNet 超过了 UNeXt-S，增加了 1.55% 的 mIoU 和 0.97% 的DSC，同时参数和计算量减少了 17% 和 72%。此外，EGE-UNet 首次将参数减少到约 50KB，同时保持了优秀的分割性能。

在消融实验中，作者同样证明了提出的GHPA和GAB模块的有效性。它们不仅提升了性能，而且显著减少了参数和计算量。

总结

本文主要提出了两个新颖的模块，GHPA 和 GAB，这两个模块大大降低了模型的复杂性，并且提高了模型的性能。同时，本文基于这两个模块构建了EGE-UNet，用于皮肤病变分割任务，实验结果表明，该方法达到了最先进的性能，同时显著降低了资源需求。