【医学图像分割 2024】VM-UNet

论文题目：VM-UNet Vision Mamba UNet for Medical Image Segmentation

中文题目：VM-UNet:医学图像分割的Mamba视觉UNet

论文链接：https://arxiv.org/abs/2402.02491

论文代码：JCruan519/VM-UNet: This is the official code repository for “VM-UNet: Vision Mamba UNet for Medical Image Segmentation”. (github.com)

论文团队：上海交通大学

发表时间：

DOI:

引用：

引用数：

摘要

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

1. 介绍

自动医学图像分割技术帮助医生更快地进行病理诊断，从而提高患者护理的效率。近年来，基于CNN的模型和基于Transformer的模型在各种视觉任务中表现出了显著的性能，特别是在医学图像分割方面。UNet作为基于CNN的模型的代表，以其结构简单、可扩展性强而闻名，后续的许多改进都是基于这种U型架构。TransUNet是基于Transformer的模型中的先驱，它在编码阶段首先使用Vision Transformer (ViT)[13]进行特征提取，在解码阶段使用CNN，显示出重要的全局信息获取能力。随后，TransFuse[36]采用了ViT的并行架构和CNN，同时捕捉本地和全球特征。此外，SwinUNet[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. 预备知识

在现代基于SSM的模型中，即结构化状态空间序列模型(S4)和Mamba都依赖于一个经典的连续系统，该系统通过中间隐式状态$h(t)\in {\mathcal{R}}^N$将一维输入函数或序列映射为$x(t)\in \mathcal{R},$到输出$y(t)\in \mathcal{R}$。上述过程可以表示为线性常微分方程(ODE):
$$\begin{array}{rl}{h}^{\prime }(t)& =\mathbf{A}h(t)+\mathbf{B}x(t)\\ y(t)& =\mathbf{C}h(t)\end{array}$$
其中，$\mathbf{A}\in {\mathcal{R}}^{N\times N}$表示状态矩阵，$\mathbf{B}\in {\mathcal{R}}^{N\times 1},\mathbf{C}\in {\mathcal{R}}^{N\times 1}$表示投影参数。

S4和Mamba将这个连续系统离散化，使其更适合深度学习场景。具体来说，他们引入一个时间尺度参数$\mathbf{\Delta }$，并使用固定的离散化规则将$A$和$B$转换为离散参数$\overline{\mathrm{A}}$和$\overline{\mathrm{B}}$。通常采用零阶保持器(ZOH)作为离散化规则，其定义如下:
$$\begin{array}{rl}\overline{\mathbf{A}}& =\exp (\mathbf{\Delta }\mathbf{A})\\ \overline{\mathbf{B}}& =(\mathbf{\Delta }\mathbf{A}{)}^{-1}(\exp (\mathbf{\Delta }\mathbf{A})-\mathbf{I})\cdot \mathbf{\Delta }\mathbf{B}\end{array}$$
离散化后，基于ssm的模型可以通过线性递归或全局卷积两种方式进行计算，分别定义为式3和式4。
$$\begin{array}{rl}{h}^{\prime }(t)& =\overline{\mathbf{A}}h(t)+\overline{\mathbf{B}}x(t)\\ y(t)& =\mathbf{C}h(t)\end{array}$$
公式4
$$\begin{array}{rl}& \overline{K}=(\mathbf{C}\overline{\mathbf{B}},\mathbf{C}\overline{\mathbf{AB}},\dots ,\mathbf{C}{\overline{\mathbf{A}}}^{L-1}\overline{\mathbf{B}})\\ & y=x\ast \overline{\mathbf{K}}\end{array}$$
式中，$\overline{\mathbf{K}}\in {\mathcal{R}}^L$表示一个结构化卷积核，L表示输入序列x的长度。

3. 方法

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

图1所示。(a)虚拟机- unet的总体结构。(b) VSS区块是VM-UNet的主要构建区块，SS2D是VSS区块中的核心操作。

3.1 VM-UNet

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

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

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

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

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块的伪代码。

$\overline{\text{Algorithm1Pseudo-codeforS6blockinSS2D}}$ $\overline{\text{Input:}x,\text{thefeaturewithshape[B,L,D](batch size, token length, dimension)}}$
Params: $\mathbf{A}$, the nn.Parameter; $\mathbf{D}$, the nn.Parameter
Operator: Linear(.), the linear projection layer
Output: $y$, the feature with shape $[\mathbb{B},\mathbb{L},\mathbb{D}]$
$1{: }\underline {\boldsymbol{\Delta}}, \mathbf{B} , \mathbf{C} = $Linear$( x) , $ Linear$( x) , $ Linear$( x) $

$2:\overline{\mathbf{A}}=\exp (\mathbf{\Delta }\mathbf{A})$

$3:\overline{\mathbf{B}}=\underline{(}\underline{\mathbf{\Delta }}\mathbf{A}{)}^{-1}(\underline{\exp }(\mathbf{\Delta }\mathbf{A})-\mathbf{I})\cdot \mathbf{\Delta }\mathbf{B}$
$4:h_t=\overline{\mathbf{A}}{h}_{t-1}+\overline{\mathbf{B}}{x}_t$
$5:y_t=\mathbf{C}{h}_t+\mathbf{D}{x}_t$ $6:y=[y_1,y_2,\cdots ,y_t,\cdots ,y_L]$

$7:\text{return}y$

3.3 损失函数

VM-UNet的引入旨在验证纯ssm模型在医学图像分割任务中的应用潜力。因此，我们专门利用最基本的二元交叉熵和骰子损失(BceDice loss)和交叉熵和骰子损失(CeDice loss)分别作为二元和多类分割任务的损失函数，如公式5和6所示。
$$\begin{array}{rl}{L}_{\mathrm{BceDice}}& ={\lambda }_1{L}_{\mathrm{Bce}}+{\lambda }_2{L}_{\mathrm{Dice}}\\ & \\ L_{\mathrm{CeDice}}& ={\lambda }_1{L}_{\mathrm{Ce}}+{\lambda }_2{L}_{\mathrm{Dice}}\end{array}$$

$$\{\begin{array}{l}{L}_{\mathrm{Bce}}=-\frac{1}{N}{\sum }_{i=1}^N\left[y_i\log ({\hat{y}}_i)+(1-y_i)\log (1-{\hat{y}}_i)\right]\\ \\ L_{\mathrm{Ce}}=-\frac{1}{N}{\sum }_{i=1}^N{\sum }_{c=1}^C{y}_{i,c}\log ({\hat{y}}_{i,c})\\ \\ L_{\mathrm{Dice}}=1-\frac{2|X\cap Y|}{|X|+|Y|}\end{array}$$
式中，N为样本总数，C为类别总数。$y_i,{\hat{y}}_i$分别表示真实标签和预测。$y_{i,c}$是一个指标，如果样本$i$属于c类，则等于1，否则等于0。${\hat{y}}_{i,c}$为模型预测样本i属于类别c的概率。|X|和|Y|分别代表真实值和预测值。${\lambda }_1$和${\lambda }_2$为损失函数的权值，默认值均为1。

4. 实验

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

4.1 数据集

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 实现细节

根据之前的工作[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 主要结果

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

4.4 消融实验

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

5. 结果

结论:在本文中，我们首次以VM-UNet为基线，引入了一个纯粹基于ssm的医学图像分割模型。为了利用基于ssm的模型的功能，我们使用VSS块构建VM-UNet，并使用预训练的vamba - s初始化其权重。在皮肤病变和多器官分割数据集上进行的综合实验表明，纯ssm模型在医学图像分割任务中具有很强的竞争力，值得未来深入探索。

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