VM-UNet: Vision Mamba UNet for MedicalImage Segmentation

1、摘要

在医学图像分割领域，基于神经网络和基于变换器的模型都得到了广泛的研究。然而，神经网络在长距离建模能力方面表现出局限性，而变形金刚则受到其二次计算复杂度的阻碍。近年来，以Mamba为代表的状态空间模型（SSM）已成为一种有前途的方法。它们不仅在模拟长距离相互作用方面表现出色，而且保持了线性计算复杂度。本文利用状态空间模型，提出了一种用于医学图像分割的Ushape体系结构模型--Vision Mamba UNet（VM-UNet）.具体地说，引入视觉状态空间（VSS）块作为捕获大量上下文信息的基础块，并构造非对称的编解码器结构。我们在ISIC 17、ISIC 18和Synapse数据集上进行了全面的实验，结果表明VM-UNet在医学图像分割任务中具有较好的性能。据我们所知，这是第一个基于纯SSM模型构建的医学图像分割模型.我们的目标是建立一个基准，并为未来开发更高效、更有效的基于SSM的分割系统提供有价值的见解。我们的代码可在https://github.com/JCruan519/VM-UNet上获得。

2、引言

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

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

最近，状态空间模型（SSM）引起了研究人员的极大兴趣。建立在经典SSM研究的基础上[18]，现代SSM（例如，Mamba [16]）不仅建立了长距离依赖关系，而且还表现出与输入大小相关的线性复杂度。此外，基于SSM的模型在许多领域都得到了大量的研究，包括语言理解[17，16]，一般视觉[38，20]等。特别是，U-Mamba [24]最近引入了一种新的SSM-CNN混合模型，标志着其首次应用于医学图像分割任务。SegMamba [35]在编码器部分采用了SSM，同时在解码器部分仍然使用CNN，这表明SSM-CNN混合模型用于3D脑肿瘤分割任务。虽然上述工作已经利用SSM的医学图像分割任务，纯SSM为基础的模型的性能还有待探讨。

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

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

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

1）我们提出了VM-UNet，标志着第一次探索纯粹基于SSM的模型在医学图像分割中的潜在应用。

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

3)我们为医学图像分割任务中的纯SSM模型建立了一个基线，为开发更高效和更有效的基于SSM的分割方法铺平了道路。

3、s(4)和mamba的知识

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

其中A ∈ RN×N表示状态矩阵，而B ∈ RN×1和C ∈ RN×1表示投影参数。

S4和Mamba将这个连续系统离散化，使其更适合深度学习场景。具体地说，他们引入了一个时间尺度参数λ，并使用固定的离散化规则将A和B转换为离散参数A-和B-。通常，零阶保持（ZOH）被用作离散化规则，并且可以定义如下：

在离散化之后，基于SSM的模型可以以两种方式计算：线性递归或全局卷积，分别定义为等式3和4。

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

4、模型-Vision Mamba UNet (VM-UNet)

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

Patch Embedding层将输入图像x ∈ RH×W×3划分为大小为4 × 4的非重叠块，随后将图像的维度映射到C，C默认为96。这个过程的结果是嵌入图像x′ ∈ RH4 × W4 ×C。最后，我们使用Layer Normalization [7]对x '进行归一化，然后将其送入编码器进行特征提取。编码器由四个阶段组成，在前三个阶段的末尾应用patch merging操作，以减少输入特征的高度和宽度，同时增加通道的数量。我们在四个级上采用[2，2，2，2] VSS块，每个级的通道计数为[C，2C，4C，8C]。

类似地，解码器被组织成四个阶段。在最后三个阶段的开始，一个补丁扩展操作被用来减少特征通道的数量和增加的高度和宽度。在四个阶段中，我们使用[2，2，2，1] VSS块，每个阶段的通道计数为[8C，4C，2C，C]。在解码器之后，使用最终投影层来恢复特征的大小以匹配分割目标。具体而言，通过补丁扩展进行4倍上采样以恢复特征的高度和宽度，然后通过投影层恢复通道数量。对于跳过连接，采用简单的加法操作，没有花哨的东西，因此不会引入任何额外的参数。

5、VSS Block

从VMamaba [20]派生的VSS块是VM-UNet的核心模块，如图1（B）所示。在经历层规格化之后，输入被分成两个分支。在第一个分支中，输入通过线性层，之后是激活函数。在第二分支中，输入在被馈送到2D选择性扫描（SS 2D）模块中用于进一步特征提取之前经历通过线性层、深度可分离卷积和激活函数的处理。随后，使用层规范化对特征进行规范化，然后使用第一个分支的输出执行逐元素生成，以合并两个路径。最后，使用线性层混合这些特征，并将此结果与剩余连接组合以形成VSS块的输出。在本文中，默认使用SiLU 作为激活函数。

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

6、损失函数

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

表示样本总数，C表示类别总数。

yi, ˆyi 分别表示真实标签和预测。

yi,c 是指示符，如果样本i属于类别c，则该指示符等于1，否则等于0。

ˆyi,c 是模型预测样本i属于类别c的概率。

|X|和 |Y| 分别表示the ground truth 和 prediction。

λ1, λ2 是损失函数的权重，默认为1。

7、结论与展望

结论：在本文中，我们第一次介绍了一个纯粹的基于SSM的医学图像分割模型，提出VM-UNet作为基线。为了充分利用基于SSM的模型的功能，我们使用VSS块构造VM-UNet，并使用预训练的VMamba-S初始化其权重。在皮肤病变和多器官分割数据集上的综合实验表明，单纯的SSM模型在医学图像分割中具有很强的竞争力，值得进一步深入研究。

未来的工作：1）设计更适合分割任务的模块，基于SSM的机制。2)VM-UNet的参数计数约为30 M，可通过手动设计或其他压缩策略简化SSM，从而增强其在真实医疗场景中的适用性。3)考虑到SSM在捕获长序列信息方面的优势，进一步研究在更高分辨率下的分割性能将是有价值的。4)探索SSM在其他医学成像任务中的应用，如检测、配准和重建等。