EGE-UNet: an Efficient Group Enhanced UNet for skin lesion segmentation

1. Abstract

目前的医学图像分割模型大多是 Transformer + Unet，这些模型的大量参数和计算负载使得它们不适合移动健康应用。

作者提出的EGE-UNet 模型轻量、高效。（与 TransFuse 相比，参数和计算成本分别降低了 494 倍和 160 倍，模型参数量只有50KB）

创新点：组多轴哈达玛产品注意力模块（GHPA）和组聚合桥模块（GAB）。

1.GHPA 对输入特征进行分组，并在不同轴上执行哈达玛产品注意力机制（HPA），以从不同角度提取病理信息。

2.GAB 通过对低级特征、高级特征以及解码器在每个阶段生成的掩码进行分组，有效地融合了多尺度信息。

2. Introduction

背景: 恶性黑色素瘤是世界上增长最快的癌症之一。据美国癌症协会估计，2020 年约有 100,350 例新发病例，超过 6,500 例死亡。因此，自动化皮肤病变分割系统势在必行，因为它可以帮助医疗专业人员快速识别病变区域并促进后续治疗过程。

相同方式可引入脑瘤、肺癌。

为了提高分割性能，最近的研究倾向于采用具有更大参数和计算复杂度的模块，例如结合视觉变换器（ViT）的自注意力机制[7]。例如，Swin-UNet [4]，基于Swin Transformer [11]，利用自注意力机制的特征提取能力来提高分割性能。 TransUNet [5] 开创了用于医学图像分割的 CNN 和 ViT 的串行融合。 TransFuse [26]采用双路径结构，利用 CNN 和 ViT 分别捕获局部和全局信息。UTNetV2[8]利用混合分层架构、高效的双向注意力和语义图来实现全局多尺度特征融合，结合了CNN和ViT的优点。 TransBTS [23] 将自注意力引入脑肿瘤分割任务中，并用它来聚合高级信息。

Abstract提到当前医学分割模型大部分是Transformer + Unet，这里做出具体阐述。

先前的工作通过引入复杂的模块来提高性能，但忽略了实际医疗环境中计算资源的限制。因此，迫切需要为移动医疗中的分割任务设计一种低参数、低计算负载的模型。最近，UNeXt [22] 结合了 UNet [18] 和 MLP [21] 开发了一种轻量级模型，该模型可以获得优异的性能，同时减少参数和计算量。此外，MALUNet [19]通过减少模型通道数并引入多个注意力模块来减小模型大小，从而比 UNeXt 具有更好的皮肤病变分割性能。然而，尽管MALUNet大大减少了参数数量和计算量，但其分割性能仍然低于一些大型模型，例如TransFuse。因此，在本研究中，我们提出了 EGE-UNet，这是一种轻量级皮肤病变分割模型，可实现最先进的效果，同时显着降低参数和计算成本。此外，据我们所知，这是第一个将参数减少到大约 50KB 的工作。

提出问题：医疗环境中计算资源的限制，复杂模块难以落地 —> 解决办法：轻量化模型

当前轻量化发展历程 —> 轻量化的模型分割效果不好 —> EGE-Unet 轻量+分割能力强

具体来说，EGE-UNet 利用两个关键模块：群组多轴 Hadamard 产品注意力模块（GHPA）和群组聚合桥模块（GAB）。

一方面，由于多头自注意力机制（MHSA），最近基于 ViT [7] 的模型已经显示出前景。 MHSA将输入划分为多个head，并在每个head中计算self-attention，这使得模型能够从不同的角度获取信息，整合不同的知识，提高性能。尽管如此，MHSA 的二次复杂度极大地增加了模型的大小。因此，我们提出了具有线性复杂度的哈达玛产品注意力机制（HPA）。HPA 采用可学习的权重，并使用输入执行哈达玛乘积运算以获得输出。随后，受到 MHSA 中多头模式的启发，我们提出了 GHPA，它将输入分为不同的组，并在每个组中执行 HPA。然而，值得注意的是，我们在不同组的不同轴上进行HPA，这有助于进一步从不同的角度获取信息。

另一方面，对于GAB，由于医学图像中分割目标的大小和形状不一致，因此获得多尺度信息至关重要[19]。因此，GAB基于组聚合融合不同大小的高层和低层特征，并额外引入掩模信息来辅助特征融合。通过将上述两个模块与UNet相结合，我们提出了EGE-UNet，它以极低的参数和计算量实现了出色的分割性能。与以前仅注重提高性能的方法不同，我们的模型还优先考虑现实环境中的可用性。图 1 显示了 EGEUNet 与其他网络的清晰比较。

具体介绍为什么引入两个创新模块（GHPA、GAB）、以及模块是基于什么论文。（模块背景+创新方法）

(1)提出了GHPA和GAB，前者有效地获取和集成多视角信息，后者接受不同尺度的特征，以及用于高效多尺度特征融合的辅助掩模。

(2)我们提出了EGEUNet，这是一种专为皮肤病变分割而设计的极其轻量级的模型。

(3) 我们进行了广泛的实验，证明了我们的方法在以显着降低的资源需求实现最先进性能方面的有效性。

主要贡献：（1）写模块作用 （2）写整体网络优势 （3）实验效果

3. Method

3.1EGE-Unet网络结构

EGE-UNet由对称编码器-解码器部分组成的 U 形架构之上。

编码器由六级组成，每级通道数为{8,16,24,32,48,64}。解码器同理

前三个阶段采用内核大小为 3 的普通卷积，后三个阶段利用提出的 GHPA 从不同的角度提取表示信息。

与 UNet 中的简单跳跃连接相比，EGE-UNet 在编码器和解码器之间的每个阶段都采用了 GAB。

利用深度监督生成不同规模的掩模预测，这些预测用于损失函数并作为 GAB 的输入之一。

通过集成这些高级模块，EGE-UNet 显着减少了参数和计算负载，同时与之前的方法相比增强了分割性能。

3.2 GHPA (Group multi-axis Hadamard Product Attention module)

为了克服 MHSA 带来的二次复杂度问题，我们提出了具有线性复杂度的 HPA。给定输入 x 和随机初始化的可学习张量 p，首先使用双线性插值来调整 p 的大小以匹配 x 的大小。然后，我们在 p 上采用深度可分离卷积（DW）[10][20]，然后在 x 和 p 之间进行哈达玛乘积运算以获得输出。然而，仅利用简单的HPA不足以从多个角度提取信息，导致结果不理想。受 MHSA 中多头模式的启发，我们引入了基于 HPA 的 GHPA，如算法 1 所示。我们将输入沿通道维度平均分为四组，并在高度-宽度、通道-高度和通道上执行 HPA - 分别为前三组的宽度轴。对于最后一组，我们只在特征图上使用DW。最后，我们沿着通道维度连接四组，并应用另一个数据仓库来整合不同角度的信息。请注意，DW 中使用的所有内核大小均为 3。

首先对输入的特征分为四组进行处理：高度-宽度、通道-高度、通道-宽度、深度可分离卷积

然后连接4组特征，进行可分离卷积融合特征。

具体过程：

第一步，按通道数将输入张量分为四组。（x1, x2, x3, x4）

设置初始化三个全一张量，分别为高度-宽度、通道-高度、通道-宽度（Pxy, Pzx, Pzy）。

第二步，将 x1, x2, x3 的对应切片分别使用双线插值法（bilinear）在Pxy, Pzx, Pzy中进行插值。

第三步，对插值后的Pxy, Pzx, Pzy，进行深度可分离卷积，然后分别和x1, x2, x3进行哈达玛乘积

第四步，连接4组特征信息，然后经过深度可分离卷积融合特征。

3.3 GAB (Group Aggregation Bridge module)

多尺度信息的获取被认为对于密集预测任务（例如医学图像分割）至关重要。因此，如图 3 所示，我们引入了 GAB，它接受三个输入：低级特征、高级特征和掩码。首先，采用深度可分离卷积（DW）和双线性插值来调整高层特征的大小，以匹配低层特征的大小。其次，我们沿通道维度将两个特征映射分为四组，并将一组低级特征与一组高级特征连接起来，以获得四组融合特征。对于每组融合特征，掩码被连接起来。接下来，将内核大小为3和不同扩张率{1,2,5,7}的扩张卷积[25]应用于不同的组，以提取不同尺度的信息。最后，将四组沿通道维度连接起来，然后应用内核大小为 1 的普通卷积，以实现不同尺度的特征之间的交互。

GAB模块作用： 将高级特征、低级特征、低级特征的预测掩码进行特征融合，作为新的输入特征进行解码。

具体过程： 高级特征、低级特征、低级特征的预测掩码 (xh、xl 、Mask)

首先，采用深度可分离卷积（DW）和双线性插值来调整高层特征 (xh) 的大小，以匹配低层特征 (xl) 的大小。

其次，沿通道维度将两个特征映射分为四组。（对应不同空洞卷积的扩张率：d1 = 1, d2 = 2, d3 = 5, d4 = 7）

并将每一组的低级特、高级特征和掩码连接起来，总共四组融合特征。

最后，将四组特征进行连接，并进行1x1卷积得到输出。

3.4 Loss Function

在本研究中，由于不同的GAB需要不同尺度的掩模信息，因此采用深度监督来计算不同阶段的损失函数，以生成更准确的掩模信息。我们的损失函数可以表示为方程（1）和（2）。其中 Bce 和 Dice 表示二元交叉熵和dice损失。 λi是不同阶段的权重。在本文中，我们默认将i=0到i=5之间的λi设置为1、0.5、0.4、0.3、0.2、0.1。

分为6个阶段，逐一计算每个阶段的损失。然后按照权重对损失进行求和。

4.Experiments

4.1 Datasets and Implementation details

为了评估我们模型的有效性，我们选择了两个公共皮肤病变分割数据集，即 ISIC2017 [1][3] 和 ISIC2018 [2][6]，分别包含 2150 个和 2694 个皮肤镜图像。与之前的研究[19]一致，我们以 7:3 的比例将数据集随机划分为训练集和测试集。

EGE-UNet是由Pytorch[17]框架开发的。所有实验均在单个 NVIDIA RTX A6000 GPU 上执行。图像被归一化并调整大小为 256×256。我们应用各种数据增强，包括水平翻转、垂直翻转和随机旋转。 AdamW [13] 用作优化器，以 0.001 的学习率初始化，CosineAnnealingLR [12] 用作调度器，最大迭代次数为 50，最小学习率为 1e-5。总共训练了 300 个 epoch，批量大小为 8。为了评估我们的方法，我们采用并集平均交集 (mIoU)、Dice 相似度得分 (DSC) 作为指标，并进行 5 次训练

​ 在公共皮肤病变分割数据集（ISIC2017 和 ISIC2018 ）进行对比实验，在ISIC2018进行消融实验

采用并集平均交集 (mIoU)、Dice 相似度得分 (DSC) 作为评估指标

4.2 Comparison Experiments

4.3 Ablation Experiments

4.4 Qualitative Comparisons

5. ConClusions

在本文中，我们提出了两个高级模块。我们的 GHPA 使用一种新颖的 HPA 机制将自注意力的二次复杂度简化为线性复杂度。它还利用分组来充分捕获来自不同角度的信息。我们的 GAB 融合了低级和高级特征，并引入了一个掩模来集成多尺度信息。基于这些模块，我们提出了用于皮肤病变分割任务的 EGE-UNet。实验结果证明了我们的方法在显着降低资源需求的情况下实现最先进的性能的有效性。我们希望我们的工作能够激发医学图像界对轻量级模型的进一步研究。

作者提出的EGE-UNet实现了轻量、准确的皮肤病变分割任务