Vision Mamba UNet for Medical Image Segmentation

Abstract

在医学图像分割领域，基于cnn和基于transformer的模型都得到了广泛的探索。然而，cnn在远程建模能力方面表现出局限性，而transformer则受到二次计算复杂性的阻碍。最近，以Mamba为例的状态空间模型(SSMs)作为一种很有前途的方法出现了。它们不仅在远程相互作用建模方面表现优异，而且保持了线性计算复杂度。本文利用状态空间模型，提出了一种用于医学图像分割的Ushape架构模型，命名为视觉曼巴UNet (VM-UNet)。具体来说，引入了视觉状态空间(VSS)块作为基础块来捕获广泛的上下文信息，并构造了一个非对称的编码器-解码器结构。我们在ISIC17, ISIC18和Synapse数据集上进行了全面的实验，结果表明VM-UNet在医学图像分割任务中具有竞争力。据我们所知，这是第一个基于纯ssm模型构建的医学图像分割模型。我们的目标是建立一个基线，并为未来更高效和有效的基于ssm的细分系统的发展提供有价值的见解。我们的代码可在https://github.com/JCruan519/VM-UNet上获得。

Keywords: Medical image segmentation State Space Models

1. Introduction

自动医学图像分割技术帮助医生更快地进行病理诊断，从而提高患者护理的效率。近年来，基于cnn的模型和基于transformer的模型在各种视觉任务中表现出了显著的性能，特别是在医学图像分割方面。UNet[27]作为基于cnn的模型的代表，以其结构简单、可扩展性强而闻名，后续的许多改进都是基于这种u型架构[11,37,28,29,30]。TransUnet[10]是基于Transformer的模型中的先驱，它在编码阶段首先使用Vision Transformer (ViT)[13]进行特征提取，在解码阶段使用CNN，显示出重要的全局信息获取能力。随后，TransFuse[36]采用了ViT和CNN的并行架构，同时捕获局部和全局特征。此外，swan-unet[9]将Swin Transformer[21]与u型架构相结合，首次引入了纯基于Transformer的u型模型。

然而，基于cnn的模型和基于transformer的模型都有固有的局限性。基于cnn的模型受到其局部接受域的限制，极大地阻碍了它们捕捉远程信息的能力。这通常会导致提取不充分的特征，从而导致次优分割结果。尽管基于transformer的模型在全局建模方面表现出优异的性能，但自注意机制在图像大小方面要求二次复杂度，导致计算负担很高[31,13]，特别是对于需要密集预测的任务，如医学图像分割。目前这些模型的缺点迫使我们开发一种新的医学图像分割架构，能够捕获强远程信息并保持线性计算复杂度。

近年来，状态空间模型(ssm)引起了研究人员的极大兴趣。在经典SSM[18]研究的基础上，现代SSM(如Mamba[16])不仅建立了长距离依赖关系，而且在输入大小方面表现出线性复杂性。此外，基于ssm的模型在许多领域都得到了大量的研究，包括语言理解[17,16]、通用视觉[38,20]等。特别是，U-Mamba[24]最近引入了一种新的SSM-CNN混合模型，这标志着它在医学图像分割任务中的首次应用。segamba[35]在编码器部分采用SSM，而在解码器部分仍然使用CNN，提出了一种SSM-CNN混合模型用于三维脑肿瘤分割任务。虽然上述工作已经将SSM用于医学图像分割任务，但纯粹基于SSM的模型的性能还有待探索。

受VMamba[20]在图像分类任务中取得成功的影响，本文首次引入了Vision Mamba UNet (VM-UNet)，这是一种纯粹基于ssm的模型，旨在展示其在医学图像分割任务中的潜力。具体来说，VM-UNet由三个主要部分组成:编码器、解码器和跳过连接。编码器由VMamba的VSS块组成，用于特征提取，以及用于下采样的补丁合并操作。相反，解码器包括VSS块和补丁扩展操作，以恢复分割结果的大小。对于跳跃连接组件，为了突出最原始的纯ssm模型的分割性能，我们采用了最简单的加法运算形式。

在器官分割和皮肤病变分割任务上进行了全面的实验，以证明纯ssm模型在医学图像分割中的潜力。具体来说，我们在Synapse[19]、ISIC17[8]和ISIC18[12]数据集上进行了大量的实验，结果表明VM-UNet可以达到有竞争力的性能。此外，重要的是要注意VM-UNet代表了纯基于ssm的分段模型的最基本形式，因为它不包括任何专门设计的模块。

本文的主要贡献如下:1)提出了VM-UNet，首次探索了纯粹基于ssm的模型在医学图像分割中的潜在应用。2)在三个数据集上进行了综合实验，结果表明VM-UNet具有相当的竞争力。3)我们为纯ssm模型在医学图像分割任务中建立了基线，为开发更高效、更有效的基于ssm的分割方法提供了有价值的见解。

2. Preliminaries

在现代基于ssm的模型中，即结构化状态空间序列模型(S4)和Mamba都依赖于一个经典的连续系统，该系统通过中间隐式状态h(t)∈RN将一维输入函数或序列映射为x(t)∈R到输出y(t)∈R。上述过程可以表示为线性常微分方程(ODE):

其中为状态矩阵，表示映射参数。

S4和Mamba将这个连续系统离散化，使其更适合深度学习场景。具体来说，他们引入一个时间尺度参数∆，并使用固定的离散化规则将A和B转换为离散参数和。通常采用零阶保持器(ZOH)作为离散化规则，其定义如下:

离散化后，基于ssm的模型可以通过线性递归或全局卷积两种方式进行计算，分别定义为式3和式4:

其中表示结构化卷积核，L表示输入序列x的长度。

3.Mehthods

在本节中，我们首先介绍VM-UNet的总体结构。随后，我们详细阐述了核心组件VSS模块。最后，我们描述了在训练过程中使用的损失函数。

3.1 Vision Mamba UNet (VM-UNet)

如图1 (a)所示，展示了VM-UNet的总体架构。具体来说，VM-UNet包括补丁嵌入层、编码器、解码器、最终投影层和跳过连接。与以往的方法[9]不同，我们没有采用对称结构，而是采用了不对称设计。

Patch Embedding层将输入图像划分为大小为4 × 4的不重叠的Patch，随后将图像的尺寸映射到C, C默认为96，该操作得到嵌入图像。最后，我们使用Layer Normalization对x '进行归一化[7]，然后将其输入编码器进行特征提取。编码器由四个阶段组成，在前三个阶段的末尾应用补丁合并操作，以降低输入特征的高度和宽度，同时增加通道数量。我们在四个阶段使用[2,2,2,2]个VSS区块，每个阶段的通道计数为[C, 2C, 4C, 8C]。

同样，解码器被组织成四个阶段。在最后三个阶段的开始，采用斑块扩展操作来减少特征通道的数量，增加特征通道的高度和宽度。在四个阶段中，我们使用[2,2,2,1]VSS块，每个阶段的通道计数为[8C, 4C, 2C, C]。在解码器之后，使用Final Projection层来恢复特征的大小以匹配分割目标。具体来说，通过patch expansion进行4次上采样来恢复特征的高度和宽度，然后通过投影层来恢复通道的数量。

对于skip连接，采用直接的加法操作，不需要附加参数，因此不会引入任何额外参数。

3.2 VSS Block

vamaba[20]衍生的VSS块是VM-UNet的核心模块，如图1 (b)所示。经过Layer Normalization后，输入被分成两个分支。在第一个分支中，输入经过一个线性层，然后是一个激活函数。在第二个分支中，输入通过线性层、深度可分离卷积和激活函数进行处理，然后输入到2d选择性扫描(SS2D)模块中进行进一步的特征提取。随后，使用Layer Normalization对特征进行规范化，然后使用第一个分支的输出执行逐元素的生产，以合并两条路径。最后，使用线性层混合特征，并将此结果与残差连接相结合，形成VSS块的输出。本文默认采用SiLU[14]作为激活函数。

SS2D由三个部分组成:扫描扩展操作、S6块操作和扫描合并操作。如图2(a)所示，扫描展开操作沿着四个不同的方向(左上到右下、右下到左上、右上到左下、左下到右上)将输入图像展开成序列。然后通过S6块对这些序列进行特征提取，确保各个方向的信息被彻底扫描，从而捕获不同的特征。随后，如图2(b)所示，扫描归并操作对来自四个方向的序列求和并归并，将输出图像恢复到与输入图像相同的大小。源自Mamba[16]的S6块在S4[17]之上引入了一种选择机制，通过根据输入调整SSM的参数。这使模型能够区分并保留相关信息，同时过滤掉不相关的信息。算法1给出了S6块的伪代码。

3.3 Loss Function

VM-UNet的引入旨在验证纯ssm模型在医学图像分割任务中的应用潜力。因此，我们专门利用最基本的二元交叉熵和骰子损失(BceDice loss)和交叉熵和骰子损失(CeDice loss)分别作为二元和多类分割任务的损失函数，如公式5和6所示。

式中，N为样本总数，C为类别总数。分别表示真实标签和预测。是一个指标，如果样本I属于c类，则等于1，否则等于0。为模型预测样本i属于类别c的概率。|X|和|Y|分别表示真实值和预测值。λ1， λ2为损失函数的权值，默认值均为1。

4.Experiments

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

4.1 Datasets

ISIC17和ISIC18数据集:国际皮肤成像协作2017年和2018年挑战数据集(ISIC17和ISIC18)[8,1,12,2]是两个公开可用的皮肤病变分割数据集，分别包含2,150和2,694张带有分割面具标签的皮肤镜图像。根据之前的工作[28]，我们将数据集以7:3的比例分割作为训练集和测试集。具体来说，对于ISIC17数据集，训练集由1500张图像组成，测试集由650张图像组成。对于ISIC18数据集，训练集包含1886张图像，而测试集包含808张图像。对于这两个数据集，我们提供了几个指标的详细评估，包括平均交联(mIoU)，骰子相似系数(DSC)，准确性(Acc)，灵敏度(Sen)和特异性(Spe)。

Synapse多器官分割数据集(Synapse): Synapse[19,3]是一个公开的多器官分割数据集，包含30例腹部CT病例3,779张轴向腹部临床CT图像，包括8种腹部器官(主动脉、胆囊、左肾、右肾、肝脏、胰腺、脾脏、胃)。按照之前作品[10,9]的设置，训练用18个案例，测试用12个案例。对于这个数据集，我们报告了骰子相似系数(DSC)和95%豪斯多夫距离(HD95)作为评估指标.

4.2 Implementation Details

根据之前的工作[28,9]，我们将ISIC17和ISIC18数据集中的图像调整为256×256，将Synapse数据集中的图像调整为224×224。为了防止过拟合，采用了随机翻转和随机旋转等数据增强技术。ISIC17和ISIC18数据集采用BceDice损失函数，Synapse数据集采用CeDice损失函数。我们将批大小设置为32，并使用AdamW[23]优化器，初始学习率为1e-3。使用CosineAnnealingLR[22]作为调度程序，最大迭代次数为50次，最小学习率为1e-5。训练周期设置为300。对于VM-UNet，我们使用在ImageNet-1k上预训练的vamba - s[20]的权重初始化编码器和解码器的权重。所有实验均在单个NVIDIA RTX A6000 GPU上进行。

4.3 Main Results

我们将VM-UNet与一些最先进的模型进行比较，给出了表1和表2中的实验结果。对于ISIC17和ISIC18数据集，我们的VM-UNet在mIoU, DSC和Acc指标方面优于其他模型。对于Synapse数据集，VM-UNet也取得了具有竞争力的性能。例如，我们的模型超过swwin - unet，这是第一个纯基于变压器的模型，在DSC和HD95指标上分别高出1.95%和2.34mm。结果证明了基于ssm的模型在医学图像分割任务中的优越性。

4.4 Ablation studies

在本节中，我们使用ISIC17和ISIC18数据集对VMUNet的初始化进行消融实验。我们分别用vamba - t和vamba - s预训练的权值初始化VM-UNet。1实验结果如表3所示，预训练的权值越强，VM-UNet的下游性能就越好，说明预训练的权值对VM-UNet的影响很大。

5. Conclusions and Future works

结论:在本文中，我们首次以VM-UNet为基线，引入了一个纯粹基于ssm的医学图像分割模型。为了利用基于ssm的模型的功能，我们使用VSS块构建VM-UNet，并使用预训练的vamba - s初始化其权重。在皮肤病变和多器官分割数据集上进行的综合实验表明，纯ssm模型在医学图像分割任务中具有很强的竞争力，值得未来深入探索。
未来的工作:1)基于ssm的机制，设计更适合分割任务的模块。2) VM-UNet的参数数约为30M，为通过手工设计或其他压缩策略精简ssm提供了机会，从而增强了ssm在现实医疗场景中的适用性。3)考虑到ssm在捕获长序列信息方面的优势，进一步研究ssm在更高分辨率下的分割性能是有价值的。4)探索ssm在其他医学成像任务中的应用，如检测、配准、重建等