UltraLight-VM-UNet_ultralight vm-uney

标题：《UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation》
论文：https://arxiv.org/pdf/2403.20035.pdf
源码：https://github.com/wurenkai/UltraLight-VM-UNet

导读

在医学图像分割领域中，皮肤病变分割同眼底视网膜血管分割、息肉分割算是几个非常经典的研究课题。

随着移动医疗设备的兴起，对于能够在这些设备上高效运行的轻量级模型的需求日益增长。

然而，大多数深度学习模型往往因为其庞大的参数量和高计算需求，难以满足这一需求。

在此之前，CVHub 平台也介绍过许多轻量化的模型，比较有代表性的一个方法是此前上海交大团队发表在 MICCAI’23 上发表的 EGE-UNet，详情可点击下方链接：

随着 Mamba 热潮的兴起，作为一度被视为“魔改重灾区”的医学图像分割领域自然也被各家机构和高校给盯上，正如此前我们发布过的第一个基于纯 Mamba 架构的医学图像分割模型 VM-UNet:

今天，同样是来自上海这边的高校团队为我们带来了一个极具“性价比”的工作——UltraLight VM-UNet。

在这篇精彩的论文中，作者们巧妙地将最新的状态空间模型（SSMs）技术——Mamba，融入到经典的UNet架构中，创造出了一种名为 UltraLight VM-UNet 的新型轻量级模型。

这个模型不仅在参数数量上实现了前所未有的精简——仅有0.049M，而且在计算效率上也达到了新的高度，GFLOPs仅为0.060。更令人振奋的是，UltraLight VM-UNet在保持模型轻盈的同时，并没有牺牲分割性能，其在三个公开的皮肤病变数据集上的表现与那些参数量更大的模型相比毫不逊色。

这篇论文不仅展示了一种创新的模型架构，还深入探讨了影响Mamba参数数量的关键因素，为未来的轻量化模型设计提供了宝贵的理论支持。后续的阅读，我们将一同踏上这场探索轻量级医学图像分割模型的奇妙旅程，见证如何在资源受限的环境中实现高效、精准的图像分析。让我们跟随作者的步伐，深入了解UltraLight VM-UNet的设计理念、实现细节以及令人瞩目的实验成果，探索这一领域的未来可能。

背景

在过去的几年里，深度学习技术在医学图像分析领域取得了显著的进展，特别是在图像分割任务中。图像分割作为医学诊断过程中的一个关键步骤，其目的是将图像中的感兴趣区域（如病变、器官等）与背景区分开。准确的图像分割对于疾病的早期发现和治疗至关重要，尤其是在皮肤癌筛查等领域。

然而，尽管深度学习模型在图像分割任务中表现出色，但它们通常需要大量的计算资源和存储空间。这对于资源受限的移动医疗设备来说是一个巨大的挑战。例如，移动设备上的电池寿命有限，而且设备的处理能力通常远不及传统的台式计算机或服务器。因此，为了在这些设备上实现实时的医学图像分析，迫切需要开发出既高效又轻量级的模型。

此外，医疗图像分割模型在临床应用中还面临着其他挑战。例如，模型需要能够处理不同尺度和复杂度的病变，同时保持高精度的分割结果。此外，医疗专业人员通常需要模型具有较高的解释性，以便理解模型的决策过程。

在这样的背景下，轻量级模型的研究自然而然便成为了医学图像分析领域的一个热点。轻量级模型不仅能够在资源受限的环境中运行，还能够减少模型的训练和部署成本，使得先进的医疗图像分析技术更加普及和可及。

最近，状态空间模型（SSMs）作为一种新兴的深度学习架构，因其在处理序列数据方面的高效性和灵活性而受到关注。特别是Mamba模型，它通过引入时间变化的参数和选择性状态空间，展示了在处理文本和视觉信息时的优越性能。这为开发新型轻量级医学图像分割模型提供了新的可能性。

在这样的研究背景下，本文提出的UltraLight VM-UNet模型，正是为了解决上述挑战而设计。通过结合Mamba模型的优势和UNet架构的高效性，UltraLight VM-UNet旨在实现在保持高分割性能的同时，大幅度减少模型的参数量和计算需求。这不仅为移动医疗设备提供了强大的图像分析工具，也为轻量级模型的研究开辟了新的道路。

方法

如上所述，本文提出了一种名为UltraLight VM-UNet的创新模型，该模型专为减少参数数量和计算负载而设计，同时保持优秀的皮肤病变分割性能。

架构概述(Architecture Overview)

首先我们先来看下整体架构。UltraLight VM-UNet模型采用了U形的结构，这种结构通常包含一个编码器（用于提取图像特征）和一个解码器（用于根据提取的特征重建分割图）。模型的特点是其通道数设置为[8, 16, 24, 32, 48, 64]，这样的设计有助于在不同层次上捕捉图像的细节和上下文信息。

模型的前三层使用传统的卷积模块来提取浅层特征，而后三层则采用了作者提出的Parallel Vision Mamba Layer（PVM Layer）来提取更深层次的特征。在跳跃连接路径中，模型使用了通道注意力桥接（Channel Attention Bridge, CAB）和空间注意力桥接（Spatial Attention Bridge, SAB）模块，以实现多尺度信息的融合。

Mamba参数影响分析(Mamba Parameter Impact Analysis)

作者们对Mamba模型中的参数影响因素进行了深入分析。实验发现，输入通道的数量对模型参数数量有着显著的影响。特别是，Vision Mamba中的2D选择性扫描模块（SS2D）是影响参数数量的关键组件。SS2D模块包括扫描扩展操作、S6块特征提取和扫描合并操作。

通过对这些操作的分析，本文得出一个结论：输入通道数量的增加会导致参数数量呈指数级增长。

并行Vision Mamba层(Parallel Vision Mamba Layer)

为了解决参数数量过多的问题，作者提出了一种新的并行处理方法，即Parallel Vision Mamba Layer（PVM Layer）。

在PVM Layer中，输入特征首先通过LayerNorm层进行归一化，然后被分割成四个部分，每部分的通道数为原始通道数的四分之一。这四个部分分别通过VSS Block进行处理，然后通过残差连接和调整因子进行优化，最后再通过拼接（concatenation）操作合并回原始通道数的特征。这种方法在保持处理通道数不变的情况下，显著降低了模型的参数数量和计算负载。

跳跃连接路径(Skip-connection Path)

在UltraLight VM-UNet模型中，跳跃连接路径的设计同样至关重要。模型使用了CAB和SAB模块来进行多尺度和多级别的信息融合。SAB模块包含最大池化、平均池化和扩展卷积等操作，而CAB模块则包含全局平均池化、拼接操作、全连接层和Sigmoid激活函数。这些模块的有效结合，不仅提高了模型的收敛能力，还增强了模型对病变的敏感性。

实验

为了确保模型输入的一致性，相关实验所有图像都被标准化到256×256像素的大小。同时，在三个主流的 ISIC2017、ISIC2018和PH2 皮肤病变分割数据集上广泛用于评估皮肤病变分割模型的性能。

实现细节

• 实验环境：实验在Python 3.8和Pytorch 1.13.0环境下进行，使用NVIDIA V100 GPU进行模型训练和测试。
• 数据增强：为了提高模型的泛化能力，实验中使用了水平翻转、垂直翻转和随机旋转等数据增强技术。
• 损失函数和优化器：使用BceDice损失函数进行模型训练，优化器选择了AdamW，学习率调度策略采用了余弦退火。
• 评估指标：采用经典的Dice相似系数(DSC)、敏感性(SE)、特异性(SP)和准确性(ACC)作为评估模型性能的主要指标。

SOTA对比

UltraLight VM-UNet模型与多个现有的轻量级和经典模型进行了比较，包括U-Net、SCR-Net、ATTENTION SWIN U-NET、C2SDG、VM-UNet、MALUNet和LightM-UNet等。

从上述实验结果可以看出，尽管UltraLight VM-UNet的参数数量大幅减少，但其在三个数据集上的分割性能与参数数量更多的模型相当，甚至在某些指标上更优。

消融实验

• 并行处理的影响：通过消融实验，作者们评估了不同并行处理设置对模型性能的影响。实验结果表明，所提出的PVM Layer结构在减少参数数量的同时，能够有效地保持模型性能。
• 模型组件的影响：作者还进行了消融实验来评估PVM Layer在UltraLight VM-UNet中的作用。实验结果显示，移除PVM Layer会导致模型性能下降，证明了PVM Layer在模型中的重要性。

写在最后

答案是：Yes!

好了，如果你对本文有任何的观点或疑问，欢迎评论区下方留言讨论。同时也欢迎对前沿AI技术感兴趣的同学扫描屏幕下方二维码添加微信好友，备注“交流学习”即可。