参数仅有0.049M！基于Mamba的医学图像分割新SOTA来了！

news2024/11/25 21:25:29

目前，基于CNN和Transformer的医学图像分割面临着许多挑战。比如CNN在长距离建模能力上存在不足，而Transformer则受到其二次计算复杂度的制约。

相比之下，Mamba的设计允许模型在保持线性计算复杂度的同时，仍然能够捕捉到长距离的依赖关系。因此基于Mamba的医学图像分割能够结合CNN的局部特征提取能力和Transformer的全局上下文理解能力，更有效地处理医学图像中复杂的结构和模式。

以上海交大提出的VM-UNet为例：

作为首个将Mamba结构融入UNet的模型，VM-UNet引入了视觉态空间（VSS）块作为基础块以捕捉广泛的上下文信息，并构建了一个非对称的编码器-解码器结构。在ISIC17、ISIC18和Synapse数据集上超越UNet++/UNet v2等SOTA。

受此启发，研究者们提出了更多Mamaba医学图像分割改进方案，我整理了其中10个值得学习的最新成果分享，论文以及开源代码也列上了，方便同学们复现。

论文原文以及开源代码需要的同学看文末

UltraLight VM-UNet

UltraLight VM-UNet: Parallel Vision Mamba Significantly Reduces Parameters for Skin Lesion Segmentation

方法：论文提出了一种用于处理深层特征的并行Vision Mamba层(PVM层)。PVM层使用四个并行的VSS块来处理特征，每个VSS块处理的通道数是初始通道数的四分之一。由于Mamba中输入通道数对参数数量有爆炸性影响，处理四分之一通道数的VSS块参数只是原始VSS块参数的6.9%，减少了93.1%。

基于PVM层，作者提出了参数仅为0.049M，GFLOPs仅为0.060的UltraLight Vision Mamba UNet。UltraLight VM-UNet的参数比传统的纯Vision Mamba UNet模型(VM-UNet)低99.82%，比可用的最轻量级Vision Mamba UNet模型(LightM-UNet)低87.84%

创新点：

提出了一种用于处理深度特征的并行视觉曼巴方法，名为PVM Layer，它在保持总体处理通道数不变的同时，以最低的计算负载实现了出色的性能。
对曼巴参数影响的关键因素进行了深入分析，并基于此提出了用于处理深度特征的并行视觉曼巴层（PVM Layer）。

VM-UnetV2

VM-UNET-V2: Rethinking Vision Mamba UNetfor Medical Image Segmentation

方法：论文提出了VM-UnetV2算法，是对医学图像分割中基于SSM的算法的改进探索。作者在七个数据集上进行了详尽的实验，结果表明VM-UNetV2具有显著的竞争力。作者是首次将基于SSM的算法与Unet变种相结合的探索者，推动了更高效、更有效的基于SSM的分割算法的发展。

创新点：

作者提出了VM-UnetV2，首次在医学图像分割中探索了更好的基于SSM的算法。
在七个数据集上进行了全面的实验，结果表明VM-UNetV2展现出显著的竞争力。
作者首次探索了将基于SSM的算法与Unet变体相结合，推动了更高效、更有效的基于SSM的分割算法的发展。

LMa-UNet

Large Window-based Mamba UNet for Medical Image Segmentation: Beyond Convolution and Self-attention

方法：论文引入了一种基于Mamba的UNet模型用于医学图像分割，实现了大窗口空间建模。此外，作者设计了一种分层和双向的SSM，进一步增强了Mamba在局部和全局特征建模方面的能力。

创新点：

提出了基于大窗口的Mamba U-Net（LMa-UNet）用于二维和三维医学图像分割。相比于基于小核的CNN和基于小窗口的Transformer，LMa-UNet利用大窗口在局部空间建模方面具有优势，在全局建模方面保持了超过二次复杂度的自注意力的高效性能。
设计了一种新颖的分层和双向的Mamba模块，进一步增强了Mamba在全局和局部空间建模能力。通过引入双向扫描，模型能够更好地关注图像中具有更多器官和病变的中心区域，并能够很好地对每个补丁的绝对位置信息和相对位置信息进行建模。

Mamba-UNet

Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation

方法：论文介绍了一种用于医学图像分割的纯视觉Mamba块基于UNet风格的网络，称为Mamba-UNet。研究结果表明，与UNet和Swin-UNet等经典相似网络相比，Mamba-UNet具有更好的性能。作者还计划将Mamba-UNet扩展到3D医学图像，并进行半/弱监督学习，以进一步推动医学成像领域的发展。

创新点：

Mamba-UNet：引入了一种纯视觉Mamba块为基础的UNet风格网络，用于医学图像分割。与经典的UNet和Swin-UNet等类似网络相比，Mamba-UNet表现出更出色的性能。
Visual Mamba块：在U-Net架构中引入了Visual Mamba块（VSS），以改善医学图像分析中的远距离依赖建模。这种创新提供了一种新的方法来处理长序列数据，并在医学图像分割领域中取得了优越的性能。