[论文解读] UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation

Renkai Wu, Yinghao Liu+2 more|arXiv (Cornell University)|Mar 29, 2024

Cutaneous Melanoma Detection and Management被引用 28

一句话总结

提出 UltraLight VM-UNet，一种基于 Parallel Vision Mamba（PVM Layer）的高参数效率皮肤病变分割模型，在公开数据集上以 0.049M 参数实现具有竞争力的性能。

ABSTRACT

Traditionally for improving the segmentation performance of models, most approaches prefer to use adding more complex modules. And this is not suitable for the medical field, especially for mobile medical devices, where computationally loaded models are not suitable for real clinical environments due to computational resource constraints. Recently, state-space models (SSMs), represented by Mamba, have become a strong competitor to traditional CNNs and Transformers. In this paper, we deeply explore the key elements of parameter influence in Mamba and propose an UltraLight Vision Mamba UNet (UltraLight VM-UNet) based on this. Specifically, we propose a method for processing features in parallel Vision Mamba, named PVM Layer, which achieves excellent performance with the lowest computational load while keeping the overall number of processing channels constant. We conducted comparisons and ablation experiments with several state-of-the-art lightweight models on three skin lesion public datasets and demonstrated that the UltraLight VM-UNet exhibits the same strong performance competitiveness with parameters of only 0.049M and GFLOPs of 0.060. In addition, this study deeply explores the key elements of parameter influence in Mamba, which will lay a theoretical foundation for Mamba to possibly become a new mainstream module for lightweighting in the future. The code is available from https://github.com/wurenkai/UltraLight-VM-UNet .

研究动机与目标

• 在计算资源有限的移动/临床场景中，推动轻量级医疗影像分割。
• 研究 Mamba 中参数减少如何影响视觉任务的性能。
• 开发一种并行处理策略（PVM Layer）以在保持准确性的同时控制参数增长。
• 在 ISIC 2017、ISIC 2018 和 PH2 数据集上展示 UltraLight VM-UNet 的有效性。

提出的方法

• 提出 Parallel Vision Mamba Layer（PVM Layer），将输入通道分成四等份并由 VSS Blocks 并行处理。
• 使用基于 Vision Mamba 的核心（SS2D、S4D及相关投影），并通过对通道数的精心控制来最小化参数。
• 整合带跳过连接的 U-Net 风格编码器-解码器，使用 Channel and Spatial Attention Bridges 进行多尺度特征融合。
• 进行消融研究，分析通道数和并行 VSS Blocks 如何影响参数和性能。
• 在三个公开的皮肤病变数据集上采用标准增广和 BCE-Dice 损失进行评估，报告 DSC、SE、SP 和 ACC。

实验结果

研究问题

• RQ1 在 Mamba 组件中减少输入通道数量如何影响参数量和性能？
• RQ2 并行处理方案（PVM Layer）是否能在显著减少参数的同时保持或提升分割性能？
• RQ3 在 ISIC2017、ISIC2018 和 PH2 数据集上，参数减少与分割精度之间的权衡是什么？
• RQ4 在极轻量设置下，跳跃连接融合模块（CAB/SAB）是否对性能有显著贡献？

主要发现

• UltraLight VM-UNet 实现了 0.049M 参数和 0.060 GFLOPs，同时在三个皮肤病变数据集上保持竞争力。
• 所提出的 PVM Layer 通过在四个并行块中分配四分之一通道输入，在 VSS Block 路径中将参数数量降低多达 93.1%。
• UltraLight VM-UNet 在 ISIC2017、ISIC2018 和 PH2 上的 DSC 约为 0.909–0.926，且 ACC 高，SE/SP 指标稳健。
• 消融显示用标准卷积替代 PVM Layer 会增加参数并降低性能，证实了并行 Vision Mamba 设计的重要性。
• 与 VM-UNet 和 LightM-UNet 相比，UltraLight VM-UNet 在参数方面分别实现 99.82% 和 87.84% 的降低，同时保持竞争性能。