Vision Mamba UNet：一种新型医学图像分割网络

人工智能咨询培训老师叶梓转载标明出处

在医学图像分割的研究中，卷积神经网络（CNN）与变换器（Transformer）模型各有其优势和局限。CNN由于其有限的局部感受野，在捕获长距离信息时能力较弱；而尽管Transformer在全局信息建模上具有显著优势，但其自注意力机制随着图像尺寸的增加，计算复杂度会以二次方的速度增长，这使得其在处理大规模图像数据时面临较大的计算压力。针对这些问题，上海交通大学的研究者提出了一种基于状态空间模型（State Space Models, SSMs）的新型U型网络架构——Vision Mamba UNet（VM-UNet）。该模型不仅在长距离依赖关系的建模上表现出色，而且保持了线性的计算复杂度。

方法

图 1(a) 展示了VM-UNet的整体架构。VM-UNet包括一个Patch Embedding层、一个编码器、一个解码器、一个Final Projection层和跳跃连接。与之前的方法不同，采用了非对称设计。

Patch Embedding层
该层将输入图像x划分为不重叠的4×4大小的块，然后将图像的维度映射到C维，其中C默认为96。这个过程产生了嵌入后的图像x'，尺寸为R H4 × W4 × C。在将x'送入编码器进行特征提取之前，使用Layer Normalization对其进行了归一化。

编码器
编码器由四个阶段组成，在前三个阶段的末尾应用了patch merging操作，以减少输入特征的高和宽，同时增加通道数。在四个阶段中分别使用了**[2, 2, 2, 2] VSS块**，每个阶段的通道数分别为[C, 2C, 4C, 8C]**。

解码器
解码器同样分为四个阶段。在最后三个阶段的开始，使用patch expanding操作来减少特征通道的数量并增加高度和宽度。在四个阶段中，使用了**[2, 2, 2, 1] VSS块**，每个阶段的通道数分别为[8C, 4C, 2C, C]。

Final Projection层
解码器之后是一个Final Projection层，用于将特征的大小恢复到与分割目标匹配。具体来说，通过patch expanding进行4倍上采样以恢复特征的高和宽，然后通过一个投影层来恢复通道数。

跳跃连接
对于跳跃连接，采用了简单的加法操作，没有引入任何额外的参数。

VSS块是VM-UNet的核心模块，源自VMamaba。在经过Layer Normalization后，输入被分成两个分支。第一个分支通过一个线性层后接一个激活函数；第二个分支则通过一个线性层、一个深度可分离卷积和一个激活函数处理后，送入2D-Selective-Scan (SS2D)模块进行进一步的特征提取。然后，使用Layer Normalization对特征进行归一化，接着执行元素间乘法操作，将两个分支的输出合并。最后，使用一个线性层混合特征，并通过与残差连接的组合形成VSS块的输出。默认使用SiLU作为激活函数。

SS2D由三个组件组成：scan expanding操作、S6块和scan merging操作。如图2(a)所示，scan expanding操作将输入图像沿四个方向（从左上到右下、从右下到左上、从右上到左下、从左下到右上）展开成序列。然后，这些序列由S6块处理以提取特征，确保从各个方向的信息都被彻底扫描，从而捕获多样化的特征。如图2(b)所示，scan merging操作将四个方向的序列求和并合并，恢复输出图像到与输入相同的尺寸。S6块在S4的基础上引入了选择性机制，通过根据输入调整SSM的参数，使模型能够区分并保留相关信息，同时过滤掉不相关的信息。算法 1展示了S6块的伪代码。

VM-UNet的引入旨在验证基于纯SSM模型在医学图像分割任务中的应用潜力。因此，专门使用了最基本的二元交叉熵和Dice损失（BceDice loss）以及交叉熵和Dice损失（CeDice loss）作为二元和多类分割任务的损失函数。公式5和6分别表示这两种损失函数：

其中，N表示样本总数，C代表类别总数。y_i和ŷ_i分别表示真实标签和预测。y_i,c是一个指示器，如果样本i属于类别c，则为1，否则为0。ŷ_i,c是模型预测样本i属于类别c的概率。|X|和|Y|分别代表真实值和预测值。λ_1和λ_2指的是损失函数的权重，这里默认都设置为1。

想要掌握如何将大模型的力量发挥到极致吗？叶老师带您深入了解 Llama Factory —— 一款革命性的大模型微调工具。实战专家1小时讲解让您轻松上手，学习如何使用 Llama Factory 微调模型。

评论留言“参加”或扫描微信备注“参加”，即可参加线上直播分享，叶老师亲自指导，互动沟通，全面掌握Llama Factory。关注享粉丝福利，限时免费录播讲解。

实验

实验部分在ISIC17、ISIC18和Synapse数据集上评估了VM-UNet在医学图像分割任务上的性能。

ISIC17和ISIC18数据集：国际皮肤成像协作2017年和2018年挑战数据集（ISIC17和ISIC18）包含2,150张和2,694张皮肤病变分割数据集，分别带有分割掩码标签。按照之前的工作，数据集按照7:3的比例划分为训练集和测试集。对于ISIC17数据集，训练集包括1,500张图像，测试集包括650张图像。对于ISIC18数据集，训练集包括1,886张图像，测试集包含808张图像。对于这两个数据集，提供了包括平均交并比（mIoU）、Dice相似系数（DSC）、准确率（Acc）、敏感性（Sen）和特异性（Spe）在内的多个指标的详细评估。

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

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

表1展示了ISIC17和ISIC18数据集上的比较实验结果。表2展示了Synapse数据集上的比较实验结果，其中也报告了每个类别的DSC。对于ISIC17和ISIC18数据集，VM-UNet在mIoU、DSC和Acc指标上优于其他模型。对于Synapse数据集，VM-UNet也取得了有竞争力的性能。例如，在DSC和HD95指标上，该模型超过了Swin-UNet（首个纯Transformer基模型）1.95%和2.34mm。结果证明了基于SSM模型在医学图像分割任务上的优越性。