论文题目:UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation
中文题目:UltraLight VM-UNet:并行视觉曼巴显著减少皮肤病变分割的参数
摘要
传统上,为了提高模型的分割性能,大多数方法倾向于添加更复杂的模块。这并不适合医疗领域,尤其是移动医疗设备,由于计算资源的限制,计算负载大的模型不适合真实的临床环境。最近,以曼巴(Mamba)为代表的状态空间模型(SSMs)已成为传统卷积神经网络(CNNs)和 Transformer 的有力竞争对手。在本文中,我们深入探索了曼巴中参数影响的关键因素,并在此基础上提出了 UltraLight 视觉曼巴 UNet(UltraLight VM-UNet)。具体而言,我们提出了一种在并行视觉曼巴中处理特征的方法,称为 PVM 层,它在保持整体处理通道数不变的同时,以最低的计算复杂度实现了优异的性能。我们在三个皮肤病变公共数据集上与几种最先进的轻量级模型进行了比较和消融实验,结果表明,UltraLight VM-UNet 仅用 0.049M 的参数和 0.060 的 GFLOPs 就展现出了同样强大的性能竞争力。此外,本研究深入探索了曼巴中参数影响的关键因素,这将为曼巴未来可能成为轻量级主流模块奠定理论基础。代码可从https://github.com/wurenkai/UltraLight-VM-UNet获取。
1 引言
随着计算机技术和硬件计算能力的不断发展,计算机辅助诊断在医疗领域得到了广泛应用,而医学图像分割是其中的重要组成部分。医学图像分割通常使用以卷积和 Transformer 为代表的深度学习网络来实现。卷积具有出色的局部特征提取能力,但在建立远程信息的相关性方面存在不足。在以往的工作中,研究人员提出利用大卷积核来缓解这个问题。至于基于 Transformer 的网络架构,近年来研究人员对其方法进行了深入研究。尽管自注意力机制可以通过连续补丁序列解决远程信息提取问题,但它也引入了更多的计算复杂性,因为自注意力机制的二次复杂性与图像大小密切相关。
此外,在提高计算机辅助诊断的准确性方面,算法模型的参数数量往往被丰富以提高模型的预测能力。然而,对于临床和真实的医疗环境,通常需要考虑实际的计算能力和内存限制。低参数和最小的计算内存占用是移动医疗任务中的基本考虑因素。因此,未来的移动医疗设备迫切需要一种计算复杂度低且性能良好的算法模型。
最近,状态空间模型(SSMs)在输入大小和内存占用方面表现出线性复杂度,这使其成为轻量级模型基础的关键。此外,SSMs 在捕捉远程依赖关系方面表现出色,这可以关键地解决卷积用于长距离信息提取的问题。在 Gu 等人的研究中,将时变参数引入 SSM 以获得曼巴(Mamba),并证明了曼巴能够以比 Transformer 更低的参数处理文本信息。在视觉方面,视觉曼巴(Vis5ion Mamba)的引入再次加深了人们对曼巴的理解,它在推理 1248×1248 大小的图像时无需注意力机制,节省了 86.8% 的内存。随着上述研究人员的杰出工作,5我们更有信心曼巴将在未来作为轻量级模型的基本构建块占据主要地位。
在本文中,我们正在构建一个基于视觉曼巴的轻量级模型。我们深入探索了曼巴的关键内存占用和性能权衡,并提出了 UltraLight 视觉曼巴 UNet(UltraLight VM-UNet)。据我们所知,提出的 UltraLight VM-UNet 是目前可用的最轻量级的视觉曼巴模型(参数为 0.049M,GFLOPs 为 0.060),并且在三个皮肤病变分割任务中表现出非常有竞争力的性能。具体而言,我们深入研究了影响曼巴计算复杂度的关键因素,并得出通道数是曼巴计算中内存占用爆炸的关键因素。基于我们的这一发现,我们提出了一种用于处理深度特征的并行视觉曼巴方法,称为 PVM 层,它同时保持整体处理通道数不变。提出的 PVM 层以惊人的低参数实现了优异的性能。此外,我们实现的 UltraLight VM-UNet 的深度特征提取仅使用包含曼巴的 PVM 层,如图 2 所示。在方法部分,我们将介绍提出的 UltraLight VM-UNet 的细节以及曼巴中参数影响的关键因素和性能平衡方法。
尽管在以前的研究中已经提到了类似的模块并行连接,但在使用 SSM 选择机制时,在曼巴中使用并行连接的影响仍然未知。这是因为并行连接导致每个 SSM 学习的特征通道数减少。在本文中,我们将详细给出答案。并行视觉曼巴或曼巴不仅在性能方面保持竞争力,而且在参数和计算复杂度方面得到了显著降低。
我们的贡献和发现可以总结如下: