点击下方卡片，关注 「AI视界引擎」 公众号
传统上，为了提高模型分割性能，大多数方法倾向于添加更复杂的模块。但这并不适用于医疗领域，尤其是在移动医疗设备上，由于计算资源限制，计算负担重的模型不适合真实的临床环境。最近，以Mamba为代表的状态空间模型（SSMs）成为传统卷积神经网络（CNNs）和Transformers的有力竞争者。在本文中，作者深入探讨了Mamba中参数影响的关键要素，并基于此提出了超轻量级视觉Mamba UNet（UltraLight VM-UNet）。
具体来说，作者提出了一种在并行视觉Mamba中处理特征的方法，称为PVM层，在保持处理通道总数不变的同时，实现了最低的计算负载和卓越的性能。作者在三个皮肤病变公共数据集上进行了与几个最先进的轻量级模型的比较和消融实验，并展示了仅拥有0.049M参数和0.060 GFLOPs的UltraLight VM-UNet具有同样强大的性能竞争力。此外，本研究深入探讨了Mamba中参数影响的关键要素，这将为Mamba可能在未来成为轻量化领域的新主流模块奠定理论基础。
1 Introduction
随着计算机技术和硬件计算能力的不断发展，计算机辅助诊断已经在医疗领域得到广泛应用，而医学图像分割是其重要的组成部分。医学图像分割通常是通过使用以卷积和 Transformer 为代表的深度学习网络来实现的。卷积在局部特征提取方面具有出色的能力，但在建立远程信息的关联方面存在不足。在之前的工作中，研究行人提出利用大卷积核来缓解这一问题。至于基于 Transformer 的网络架构，近年来，研究行人已经深入探讨了其方法。尽管自注意力机制可以通过连续的图像块序列解决远程信息提取的问题，但它也引入了更多的计算负载。这是因为自注意力机制的二次复杂度与图像大小密切相关。
此外，在提高计算机辅助诊断的准确性方面，算法模型的参数数量通常会被增加以增强模型的预测能力。然而，对于临床和实际医疗环境来说，现实的计算能力和内存限制往往需要被考虑。在mHealth任务中，低参数和最小的计算内存占用是至关重要的考虑因素。因此，对于未来移动医疗设备来说，迫切需要一种具有低计算负担和良好性能的算法模型。
图1：对ISIC2017数据集的比较结果可视化。X轴对应于参数和GFLOPs，越少越好。Y轴对应于分割性能（DSC），越高越好。
近期，状态空间模型（SSMs）在输入尺寸和内存占用方面显示出线性复杂性，这使得它们成为轻量级模型基础的关键。此外，SSMs擅长捕捉远程依赖关系，这可以关键性地解决在提取远距离信息时卷积的问题。在Gu等人[7]中，将时变参数引入到SSM中以获得Mamba，并且证明了Mamba能够使用比Transformers更低的参数处理文本信息。在视觉方面，Vision Mamba的引入再次加深了人们对Mamba的理解，它在不需要注意力机制的情况下，在推理1248×1248大小的图像时节省了86.8的内存。以上提到的研究行人所做的重要工作，使作者更有信心相信Mamba将在未来作为轻量级模型的基本构建块占据重要位置。
在本文中，作者基于Vision Mamba构建了一个轻量级模型。作者深入探讨了Mamba的关键内存占用及其性能权衡，并提出了一个超轻量级Vision Mamba UNet（UltraLight VM-UNet）。据作者所知，所提出的UltraLight VM-UNet是现有最轻量的Vision Mamba模型（具有0.049M的参数和0.060的GFLOPs），并且在三个皮肤病变分割任务中表现出极具竞争力的性能。具体来说，作者深入研究了对Mamba计算负载影响的关键因素，并得出结论：通道数是Mamba计算中内存占用激增的关键因素。
基于这一发现，作者提出了一种名为PVM层的并行Vision Mamba方法来处理深层特征，同时保持了总体处理通道数恒定。所提出的PVM层以惊人的低参数数量实现了卓越的性能。此外，作者仅使用包含Mamba的PVM层来实现所提出的UltraLight VM-UNet的深层特征提取，如图2所示。在方法部分，作者将介绍所提出的UltraLight VM-UNet的细节，以及Mamba中参数影响的关键因素和性能平衡方法。
作者的贡献和发现可以总结如下：
- 提出了一种超轻量级视觉曼巴UNet（UltraLight VM-UNet）用于皮肤病变分割。据作者所知，UltraLight VM-UNet是现有最轻的视觉曼巴模型（参数为0.049M，GFLOPs为0.060）。
- 提出了一种名为PVM层的并行视觉Mamba方法，用于处理深度特征，该方法在保持处理通道总数不变的同时，实现了最低的计算负载下的优异性能。
- 作者深入分析了影响Mamba参数的关键因素，并为Mamba在未来成为轻量级建模的主流模块提供了理论依据。所提出的UltraLight VM-UNet参数比传统的纯Vision Mamba UNet模型（VM-UNet）低99.82%，比当前最轻的Vision Mamba UNet模型（LightM-UNet）低87.84%。此外，UltraLight VM-UNet在所有三个公开可用的皮肤病变分割数据集上都保持了强大的性能竞争力。
2 Related Work
随着计算机计算能力的显著提升，计算机视觉已成为当今计算机技术中一个重要的领域。作为计算机视觉的一个分支，图像分割已经被相关研究者研究了数十年。传统的图像分割技术通过简单的阈值处理或数学方法来区分不同物体，这在不同像素区域之间造成了明显的差异。然而，传统方法难以应用于精细区域的分割。随着深度学习的不断发展，全卷积模型（FCN）首次在图像分割方法上展现了卓越的性能。FCN的出现也导致了深度学习图像分割方法的快速发展。在FCN出现后不久，另一个全卷积模型（U-Net）的出现再次引发了关注。U-Net中的跳跃连接操作能够很好地融合高层特征和低层特征，这对于图像分割尤其重要，特别是在需要细粒度分割的医疗图像分割中。
医学图像分割作为图像分割的重要分支之一，也是许多研究行人投入大量努力的研究方向。其中，多尺度变化问题和特征细化学习是医学图像分割中的关键问题。而皮肤病变分割具有丰富多变的特征信息以及由于其恶性黑色素瘤导致的高死亡率，这促使许多研究行人围绕皮肤病变分割进行一系列研究。
以皮肤病变为代表的医学图像分割算法在U-Net出现后得到了迅速发展。在Aghdam等人[1]的研究中，针对基于Swin U-Net[2]的皮肤病变分割，提出了级联操作中注意力机制的抑制操作。MHorUNet模型提出了一种用于皮肤病变分割的高阶空间交互UNet模型。在Wu等人[30]的研究中，提出了一种针对高阶交互的自适应选择UNet模型用于皮肤病变分割。
此外，基于U-Net改进的用于皮肤病变分割的算法非常多。然而，研究者们通常会在模型中添加更丰富的模块以提高识别的准确性，但这也会显著增加模型的参数数量和计算复杂性。在Vision Mamba出现后，LightM-UNet基于Mamba被提出以减少模型中的参数数量。LightM-UNet通过使用残差Vision Mamba进一步提取深层语义和远程关系，并在参数数量较少的情况下实现更优的性能。另外，U-Mamba首次将Vision Mamba引入到U型框架中，但其大量的参数（173.53M）限制了在真实临床环境中的应用。
在本文中，为了解决当前大型模型参数的问题，并揭示影响Mamba参数的关键因素，作者提出了一个基于Mamba的仅具有0.049M参数的超轻量级视觉Mamba UNet（UltraLight VM-UNet）。在三个公开的皮肤病变数据集上，UltraLight VM-UNet被证实仍然保持强大的竞争力。在下一节中，将详细描述作者的方法。
3 Method
Architecture Overview
提出的超轻量级视觉曼巴UNet（UltraLight VM-UNet）如图2所示。UltraLight VM-UNet总共有6层结构，由U形结构（编码器、解码器和跳跃连接路径）组成。6层结构中的通道数设置为。前3层浅层特征的提取由卷积模块（Conv Block）组成，其中每一层包括一个具有3x3卷积核的标准卷积和一个最大池化操作。从第4层到第6层的深层特征是作者的核心部分，每一层由作者提出的并行视觉曼巴层（PVM Layer）组成。解码器部分保持与编码器相同的设置。跳跃连接路径利用通道注意力桥（CAB）模块和空间注意力桥（SAB）模块进行多 Level 和多尺度的信息融合。
Mamba Parameter Impact Analysis
在视觉任务中，Mamba通常被嵌入到视觉状态空间（Visual State Space, VSS）块中以便使用，如图4(b)所示。VSS块主要由两个主要分支组成，第一个分支主要由线性层和SiLU激活函数[6]构成。第二个分支主要由线性层、深度卷积、SiLU激活函数、2D选择性扫描模块（SS2D）和层归一化（LayerNorm）组成。最后，两个分支通过逐元素的乘法进行合并以输出结果。
在VSS块中对参数影响最大的模块是SS2D。SS2D的组成部分如图3所示。它们包括扫描扩展操作、S6块特征提取以及扫描合并操作。首先通过扫描扩展操作，序列在四个方向上从左上到右下、从右下到左上、从右上到左下以及从左下到右上进行扩展。然后输入到S6块[7]进行特征提取。最后，通过扫描合并操作将其恢复到原始初始图像的大小。
其中，在SS2D中，输入通道的数量、S4D层状态维度的大小、内部卷积核的大小、投影扩张倍数以及投影矩阵的秩都影响着参数的数量。在这之中，输入通道数量的影响是巨大的，其影响主要来自四个方面：
- 首先，SS2D内部扩展投影通道的由投影扩展乘数和输入通道数量的乘积确定。这可以通过以下方程具体表示：
其中 是内部扩展投影通道， 是投影扩展倍数（默认固定为2），而 是输入通道的数量。作者可以看到，随着模型中每层通道数的急剧增加， 将指数级上升。
- 其次，在SS2D中，输入投影层和输出投影层的参数将直接与输入通道的数量相关。输入投影层和输出投影层的工作方式如下：
其中，输入投影（）层的参数为 ，输出投影层（）的参数为 。作者可以看到，输入通道数 是控制参数的关键要素，而内部扩展投影通道 也是受 控制的。
- 此外，SS2D网络中S6块内的四个线性投影层也是参数效果的关键。每个线性投影层具体指定如下：
其中 是投影矩阵的秩（）， 是S4D层状态维度的大小（固定为16），每个线性投影层的参数为
。
然而，总共有4个线性投影层，所以总参数量为
。
因此，从上述内容作者可以知道，所有参数仍然主要由输入通道数 控制。
同样，在SS2D模块中，控制S4D层不同状态注意力权重的参数矩阵的是一个重要的影响因素。 是一个形状为 的参数矩阵，而 是一个通常固定为4的超参数。因此，参数 可以推导为 。 总之，假设原始的输入通道数为1024，保持其他参数不变，将通道数减少到原来的四分之一（输入通道数变为256），根据上述参数公式计算，原始的总参数可以从7669760减少到525312。参数的大幅减少表明通道数减少了93.1，这进一步确认了输入通道数对SS2D参数具有非常关键的影响。这也进一步证实了输入通道数对SS2D参数具有极其重要的影响。
基于对影响Mamba参数的关键要素的深入研究发现，作者提出了一种在并行视觉Mamba中处理特征的方法，名为PVM层。在保持处理通道总数不变的情况下，实现了最低的计算负载和卓越的性能。具体细节将在下一节中详细说明。
Parallel Vision Mamba Layer
正如在前一小节中分析的那样，输入通道的数量对Mamba的参数有爆发式的影响。如图4(a)所示，作者提出了并行视觉Mamba层（PVMLayer）来处理深层特征。
具体来说，具有通道数的特征首先通过一个LayerNorm层，然后被分为具有每个通道数的特征、、和。之后，每个特征被输入到一个VSS块中，然后输出经过残差拼接和调整因子，以优化远程空间信息获取能力。最后，通过concat操作将四个特征合并为具有通道数的特征，然后分别通过LayerNorm和投影操作输出。具体操作可以用以下方程表示：
在这里，代表层归一化，是分裂操作，是VSS块操作，是对残差连接的调整因子，是拼接操作，而是投影操作。从方程3.6中，作者使用了并行的Vision Mamba处理特性，同时确保处理的总通道数保持不变，从而在最大化参数减少的同时保持了高准确性。
如图4(a)中方法A和B所示，再次假设通道计数大小为1024，方法A中的每个VSS块将参数减少了93.1。它包含4个此类操作，因此合计起来，比较方法B的参数总体减少了72.4。通过作者提出的并行Vision Mamba操作，在保持强大的性能竞争力的同时，实现了最大化的参数减少。
Skip-connection Path
Short-Cut 路径使用了由Ruan等人提出的空间注意力桥（SAB）模块和通道注意力桥（CAB）模块，如图2(b)所示。SAB和CAB的结合使用使得能够融合UltraLight VM-UNet不同尺度上的多阶段特征。
SAB模块包括最大池化、平均池化和共享权重的扩展卷积。CAB模块包括全局平均池化、拼接操作、全连接层和sigmoid激活函数。在之前的工作中已经显示，SAB和CAB均能有效提高模型的收敛能力并增强对病变的敏感性。
4 Experiment
Datasets
为了验证所提出的UltraLight VM-UNet在参数量为0.049M时也能达到具有竞争力的性能，作者在三个公开可获得皮肤病变数据集上进行了实验。ISIC2017 和 ISIC2018 数据集是由国际皮肤成像合作组织（ISIC）分别发布两个大型数据集。PH 数据集是一个小型公开的皮肤病变数据集，因此作者使用PH 作为外部验证，使用ISIC2017数据集来训练权重。
对于ISIC2017数据集，作者获得了2000张图像以及带有分割 Mask 标签的皮肤镜图像。其中，数据集是随机划分的，1250张用于模型训练，150张用于模型验证，600张用于模型测试。图像的初始大小为576767像素，作者在输入模型时将大小标准化为256256像素。
对于ISIC2018数据集，作者获得了2594幅图像以及带有分割 Mask 标签的皮肤镜图像。其中，数据集被随机划分，1815幅用于模型训练，259幅用于模型验证，520幅用于模型测试。图像的初始大小为20163024像素，作者在输入模型时将图像大小标准化为256256像素。
对于PH数据集，作者获得了200张图片以及带有分割 Mask 标签的皮肤镜图像。所有200张图片都被用于外部验证。图片的初始大小为768560像素，作者将大小标准化为256256像素以输入到模型中。
Implementation details
实验全部基于Python 3.8和Pytorch 1.13.0实现。所有实验均使用了一块拥有32GB内存的NVIDIA V100 GPU。为了更公平地确定模型的性能，所有实验采用了相同的数据增强操作，包括水平翻转、垂直翻转和随机旋转操作。采用了BceDice损失函数，优化器为AdamW，训练轮次为200，批大小为8，使用了余弦退火学习率调度器，初始学习率为0.001，最小学习率设置为0.00001。
Evaluation metrics
Dice相似系数（DSC）、灵敏度（SE）、特异性（SP）和准确度（ACC）是医学图像分割中最常用的评估指标。DSC用于衡量真实值与预测分割图之间的相似程度。SE主要用于衡量真正例在真正例和假负例中的百分比。SP主要用于衡量真负例在真负例和假正例中的百分比。ACC主要用于衡量正确分类的百分比。
其中 TP 表示真正例，TN 表示真负例，FP 表示假正例，FN 表示假负例。
Comparison results
为了验证在0.049M参数下，所提出的UltraLight VM-UNet的竞争力表现，作者与几种先进的轻量级和经典医学图像分割模型进行了对比实验。具体包括U-Net，SCR-Net，ATTENTION SWIN U-NET，CSDG，VM-UNet，MALUNet和LightM-UNet。
表1、表2和表3分别在ISIC2017、ISIC2018和PH数据集上展示了实验结果。如表所示，作者的模型的参数比传统的纯Vision Mamba UNet模型（VM-UNet）低99.82%，比当前最轻的Vision Mamba UNet模型（LightM-UNet）低87.84%。此外，作者模型的GFLOPs比VM-UNet低98.54%，比LightM-UNet低84.65%。在参数和GFLOPs大幅减少的情况下，作者模型的性能仍然保持优秀且极具竞争力。另外，MALUNet是基于卷积提出的轻量级模型，尽管它比VM-UNet和LightM-UNet具有更低的参数和GFLOPs，但作者的模型的参数和GFLOPs仍分别比它们低72.0%和27.71%。
特别是，基于卷积提出的轻量级模型MALUNet的性能远低于基于Mamba的模型，这反映出基于卷积的轻量级模型在平衡性能与计算负载之间的关系上存在困难。
图5展示了可视化的分割图，从可视化结果可以更直观地得出结论，所提出的UltraLight VM-UNet具有更平滑和更准确的轮廓和边界。
Ablation experiments
为了验证提出的具有不同并行度的Vision Mamba方法的有效性，作者进行了一系列的消融实验。
如图6所示，作者执行了3种不同的设置。设置1是VSS块的传统连接方式，设置2是使用并行连接两个具有一半通道数的VSS块，而设置3是使用并行连接四个每个具有通道数的VSS块。
通过分析本研究第3.2节中Mamba的参数，假设传统VSS块连接方法设置1的参数为，则设置2的参数可以计算为，设置3的参数为。表4展示了这次消融实验的结果，需要注意的是，这里的参数指的是整个模型（包含Conv块和跳跃连接部分）的参数。设置2和设置3的参数分别是传统VSS块连接方法设置1参数的51.47和36.03，而整体的GFIOPs变化不大。
在性能方面，设置3最低的参数仍然保持了更好的分割性能，因此，在本文中，作者采用设置3作为所提出并行Vision Mamba层（PVM层）的关键结构。
此外，为了验证在UltraLight VM-UNet中提出的并行视觉Mamba层（PVM层）的影响，作者进行了一系列的消融实验。如表5所示，作者分别用卷积核为3的标准卷积替换编码器和解码器中的PVM层。此外，作者还同时用标准卷积替换编码器和解码器中的PVM层。从表中可以得出，在分别替换编码器和解码器的PVM层后，参数增加了63.26，GFLOPs在两者中均增加，而性能在两者中均下降。特别是，在同时替换编码器和解码器的PVM层后，参数增加了151，GFLOPs增加了25。总之，这表明在替换PVM层后，所有性能方面都出现了下降，而参数和GFLOPs均有所增加。这再次证明了PVM层的重要作用。
参考
[1].UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation.
点击上方卡片，关注 「AI视界引擎」 公众号