VM-UNet

首次探索纯粹基于SSM的模型在医学图像分割中的潜在应用。

在医学图像分割领域，基于CNN和Transformer的模型都已经被广泛研究。然而，CNN在长距离建模能力上存在局限，而Transformer则受到其二次计算复杂度的制约。最近，以Mamba为代表的态空间模型（SSMs）成为了一种有前景的方法。它们不仅在建模长距离交互方面表现出色，同时保持了线性计算复杂度。

在本文中，利用态空间模型，作者提出了一个U形架构模型，用于医学图像分割，命名为Vision Mamba UNet（VM-UNet）。具体来说，引入了视觉态空间（VSS）块作为基础块以捕捉广泛的上下文信息，并构建了一个非对称的编码器-解码器结构。作者在ISIC17、ISIC18和Synapse数据集上进行了全面的实验，结果表明VM-UNet在医学图像分割任务中表现出竞争力。据作者所知，这是第一个基于纯SSM模型的医学图像分割模型。作者旨在建立 Baseline ，并为未来开发更高效、有效的SSM-based分割系统提供宝贵的洞见。

代码：https://github.com/JCruan519/VM-UNet

1 Introduction

自动化医学图像分割技术帮助医生更快地进行病理诊断，从而提高了患者护理的效率。近来，基于CNN和基于Transformer的模型在各种视觉任务中展示了卓越的性能，特别是在医学图像分割方面。UNet作为基于CNN模型的代表，以其结构简单和强大的可扩展性而闻名，后续许多改进都是基于这种U形架构。

TransUnet，作为基于Transformer模型的先驱，首次在编码阶段采用视觉Transformer（ViT）[13]进行特征提取，并在解码阶段使用CNN，展示了其在获取全局信息方面的显著能力。随后，TransFuse融合了ViT与CNN的并行架构，同时捕捉局部和全局特征。此外，Swin-UNet将Swin Transformer与U形架构结合，首次引入了一个纯粹的基于Transformer的U形模型。

然而，基于卷积神经网络（CNN）的模型和基于Transformer的模型都具有固有的局限性。基于CNN的模型受到局部感受野的限制，这大大阻碍了它们捕捉远程信息的能力。这通常会导致提取的特征不充分，从而影响分割效果不佳。尽管基于Transformer的模型在全局建模方面表现出色，但自注意力机制在图像尺寸方面需要二次复杂度，导致了较高的计算负担，特别是对于像医学图像分割这类需要密集预测的任务。这些模型当前的不足促使作者开发一种新的医学图像分割架构，这种架构能够捕捉强烈的远程信息，并保持线性的计算复杂度。

近期，状态空间模型（SSMs）在研究行人中引起了极大的关注。在经典SSM研究的基础上，现代SSM（例如，Mamba）不仅建立了长距离依赖性，而且对于输入大小展现出线性复杂性。此外，基于SSM的模型在许多领域都得到了大量研究，包括语言理解，通用视觉等。特别是，U-Mamba最近引入了一种新颖的SSM-CNN混合模型，首次将其应用于医学图像分割任务中。SegMamba在编码器部分整合了SSM，而在解码器部分仍然使用CNN，为3D脑肿瘤分割任务提出了一个SSM-CNN混合模型。尽管上述工作已经将SSM用于医学图像分割任务，但基于纯SSM模型的性能尚未被探索。

受到V Mamba [20]在图像分类任务中成功的启发，本文首次引入了视觉猛蛇UNet（VM-UNet），这是一种纯粹的基于SSM的模型，旨在展示其在医学图像分割任务中的潜力。

具体来说，VM-UNet由三个主要部分组成：编码器、解码器和跳跃连接。编码器使用来自V Mamba的VSS块进行特征提取，并伴有图像下采样的像素合并操作。相反，解码器则由VSS块和像素扩展操作组成，以恢复分割结果的大小。对于跳跃连接部分，为了强调最原始的基于纯SSM模型在分割性能上的表现，作者采用了最简单的加性操作形式。

在器官分割和皮肤病变分割任务上进行了综合实验，以展示基于纯SSM模型在医学图像分割中的潜力。具体来说，作者在Synapse，ISIC17，和ISIC18数据集上进行了大量实验，其结果指出VM-UNet可以达到具有竞争力的性能。此外，值得注意的是，VM-UNet代表了基于纯SSM分割模型的最基本形式，因为它不包括任何特别设计的模块。

本文的主要贡献可以概括如下：

1. 作者提出了VM-UNet，标志着首次探索纯粹基于SSM的模型在医学图像分割中的潜在应用。

2. 在三个数据集上进行了全面实验，结果表明VM-UNet具有相当大的竞争力。

3. 作者为医学图像分割任务中纯粹基于SSM的模型建立了一个 Baseline ，提供了有价值的见解，为更高效、更有效的SSM-based分割方法的发展铺平了道路。

2 Preliminaries

3 Methods

在本节中，作者初步介绍了VM-UNet的整体结构。随后，作者将详细阐述其核心组件，即VSS块。最后，作者将描述在训练过程中使用的损失函数。

Vision Mamba UNet (VM-UNet)

如图1（a）所示，VM-UNet的整体架构被呈现出来。具体来说，VM-UNet包括一个Patch嵌入层、一个编码器、一个解码器、一个最终投影层以及跳跃连接。与之前的方法[9]不同，作者没有采用对称结构，而是使用了一种非对称设计。

同样地，解码器被组织成四个阶段。在最后三个阶段的开始，使用了一个块扩展操作来减少特征通道的数量并增加高度和宽度。在四个阶段中，作者使用了[2, 2, 2, 1] VSS块，每个阶段的通道数分别为[8C, 4C, 2C, 4]。在解码器之后，采用了一个最终的投影层来将特征的大小恢复到与分割目标相匹配。具体来说，通过块扩展执行了4倍的上采样来恢复特征的高度和宽度，然后通过一个投影层来恢复通道数量。

对于跳跃连接，采用了简单直接的加法操作，没有花哨的装饰，因此不会引入任何额外的参数。

VSS block

VSS块源自VMamaba的结构，是VM-UNet的核心模块，如图1（b）所示。在经过层归一化之后，输入被分为两个分支。在第一个分支中，输入通过一个线性层，然后是一个激活函数。在第二个分支中，输入经过一个线性层、深度可分离卷积和一个激活函数的处理，然后被送入2D-选择性扫描（SS2D）模块进行进一步的特征提取。

随后，特征通过层归一化进行标准化处理，并与第一个分支的输出进行逐元素乘法，以合并两条路径。最后，使用线性层混合特征，并将此结果与残差连接相结合，形成VSS块的输出。在本文中，默认采用SiLU作为激活函数。

SS2D由三个部分组成：扫描扩展操作、S6模块和扫描合并操作。如图2(a)所示，扫描扩展操作沿四个不同方向（从左上到右下，从右下到左上，从右上到左下，以及从左下到右上）展开输入图像成序列。这些序列随后由S6模块进行处理以提取特征，确保来自各个方向的信息得到彻底扫描，从而捕捉到多样的特征。接着，如图2(b)所示，扫描合并操作对来自不同方向的序列进行求和合并。

四个方向，将输出图像恢复到与输入图像相同的大小。S6模块源自Mamba，在S4之上引入了一种选择机制，通过根据输入调整SSM的参数。这使得模型能够区分并保留相关信息，同时过滤掉不相关的内容。S6模块的伪代码在算法1中给出。

Loss function

VM-UNet的引入旨在验证基于纯SSM模型的在医学图像分割任务中的应用潜力。因此，作者分别针对二值和多类分割任务，仅使用最基本的二元交叉熵和Dice损失（BceDice损失）以及交叉熵和Dice损失（CeDice损失）作为损失函数，如方程式5和6所示。

4 Experiments

在本节中，作者对VM-UNet在皮肤病变和器官分割任务上进行了全面的实验。具体来说，作者在ISIC17、ISIC18和Synapse数据集上评估了VM-UNet在医学图像分割任务上的性能。

Datasets

ISIC17和ISIC18数据集： 国际皮肤成像协作2017和2018挑战数据集（ISIC17和ISIC18）是两个公开可用的皮肤病变分割数据集，分别包含2,150和2,694张带分割 Mask 标签的皮肤镜图像。按照之前的工作，作者将数据集以7:3的比例分为训练集和测试集。

具体来说，对于ISIC17数据集，训练集由1,500张图像组成，测试集由650张图像组成。对于ISIC18数据集，训练集包括1,886张图像，而测试集包含808张图像。对于这两个数据集，作者提供了在几个指标上的详细评估，包括平均交并比（mIoU），Dice相似性系数（DSC），准确度（Acc），敏感性（Sen）和特异性（Spe）。         whaosoft aiot http://143ai.com

Synapse多器官分割数据集（Synapse）：Synapse是一个公开的多器官分割数据集，包含30个腹部CT病例，共3,779张轴向腹部临床CT图像，包括8种腹部器官（主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏、胃）。遵循之前的工作的设置，其中18个病例用于训练，12个病例用于测试。对于这个数据集，作者报告了Dice相似性系数（DSC）和95% Hausdorff距离（HD95）作为评估指标。

Implementation details

遵循之前的工作，作者将ISIC17和ISIC18数据集中的图像大小调整为256×256，将Synapse数据集中的图像大小调整为224×224。为防止过拟合，采用了数据增强技术，包括随机翻转和随机旋转。对于ISIC17和ISIC18数据集，使用BceDice损失函数，而对于Synapse数据集，采用CeDice损失函数。作者将批量大小设置为32，并采用AdamW优化器，初始学习率为1e-3。使用CosineAnnealingLR作为调度器，最大迭代次数为50，最小学习率为1e-5。训练周期设置为300。对于VM-UNet，作者用VMamba-S的权重初始化编码器和解码器的权重，VMamba-S是在ImageNet-1k上预训练的。所有实验都在单个NVIDIA RTX A6000 GPU上完成。

Main results

作者对比了VM-UNet与一些最先进模型的性能，并在表1和表2中展示了实验结果。

对于ISIC17和ISIC18数据集，VM-UNet在mIoU、DSC和Acc指标上超过了其他模型。对于Synapse数据集，VM-UNet也取得了具有竞争力的表现。例如，VM-UNet在DSC和HD95指标上分别超过了Swin-UNet——这是第一种纯粹的基于Transformer的模型——1.95%和2.34mm。这些结果证明了基于SSM的模型在医学图像分割任务中的优越性。

Ablation studies

在本节中，作者使用ISIC17和ISIC18数据集对VM-UNet的初始化进行了消融实验。作者分别用来自VMamba-T和VMamba-S的预训练权重来初始化VM-UNet。如表3所示，实验结果表明，更强大的预训练权重显著提高了VM-UNet在下游任务中的性能，这表明VM-UNet在很大程度上受到预训练权重的影响。

5 Conclusions and Future works

结论：在本文中，作者首次引入了一种基于纯SSM的医疗图像分割模型，并提出VM-UNet作为 Baseline 。为了发挥基于SSM模型的潜能，作者使用VSS块构建了VM-UNet，并用预训练的VMamba-S初始化其权重。在皮肤病变和多器官分割数据集上进行的全面实验表明，基于纯SSM的模型在医疗图像分割任务中具有高度竞争力，并在未来值得深入探索。

未来工作：

1. 基于SSM机制的特性，设计更适合分割任务的模块。

2. VM-UNet的参数数量约为3000万，这为通过人工设计或其他压缩策略简化SSM提供了机会，从而加强它们在现实医疗场景中的应用性。

3. 考虑到SSM在捕捉长序列信息方面的优势，进一步研究在高分辨率下的分割性能将具有价值。

4. 探索SSM在其他医疗成像任务中的应用，如检测、配准和重建等。