Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation

Mamba-UNet：用于医学图像分割的类UNet纯视觉Mamba网络

原文链接：Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation
本文代码链接：Mamba-UNet

摘要：在医学图像分析领域的最新进展中，卷积神经网络（CNN）和视觉转换器（ViT）已经树立了重要的基准。前者通过卷积操作在捕捉局部特征方面表现出色，而后者则通过利用自注意力机制实现了对全局上下文的显著理解。然而，这两种架构在有效建模医学图像中的长距离依赖关系方面都存在局限性，这对于精确分割至关重要。受以处理长序列和全局上下文信息并增强计算效率为特点的Mamba架构（作为状态空间模型SSM）的启发，我们提出了Mamba-UNet，这是一种将医学图像分割中的U-Net与Mamba能力相结合的新型架构。Mamba-UNet采用基于纯视觉Mamba（VMamba）的编码器-解码器结构，并融入跳跃连接，以在网络的不同尺度上保留空间信息。 这种设计促进了全面的特征学习过程，能够捕捉医学图像中的复杂细节和更广泛的语义上下文。我们在VMamba块内引入了一种新颖的融合机制，以确保编码器和解码器路径之间的无缝连接和信息流，从而提高了分割性能。 我们在公开可用的ACDC MRI心脏分割数据集和Synapse CT腹部分割数据集上进行了实验。结果表明，在相同的超参数设置下，Mamba-UNet在医学图像分割方面优于多种类型的UNet。

Keywords: Medical Image Segmentation，Convolution，Transformer，Mamba，State Space Models

1.介绍

  医学图像分割对于诊断和治疗至关重要，而基于深度学习的网络在这一领域表现出了卓越的性能。U-Net是其中最重要的架构之一，以其对称的编码器-解码器风格架构和跳跃连接而闻名，其中各种编码器和解码器在不同层次上提取特征信息，而跳跃连接则实现了特征信息的有效转换。大多数研究进一步探索了结合高级网络块技术的U-Net，如密集连接、残差块、注意力机制、深度卷积和空洞卷积，从而在CT、MRI、超声等医学图像分割中产生了各种改进的UNet。

  受自然语言处理中自注意力机制成功的启发，ViT首次将纯多头自注意力机制应用于图像识别任务，并取得了最先进的性能。这展示了其在建模长距离依赖方面的潜力。 像移动窗口这样的技术进一步定制了ViT，产生了Swin Transformer，增强了其在计算机视觉中密集预测任务（如图像分割和检测）的适用性。在医学图像分割中，受传统CNN设计的启发，ViT与U-Net架构的结合也导致了各种混合和纯ViT基U-Net的出现。例如，TransUNet是首个在U-Net编码器中利用ViT特征学习能力的工作。UNETR将ViT与U-Net结合用于3D分割，而Swin-UNet和DCSUnet则进一步探索了纯Swin Vision Transformer网络块与基于U-Net结构的结合。

图1. 简要介绍了在医学图像分割中，结合Transformer和状态空间模型（SSM）的UNet的近期发展演变。

  虽然Transformer在捕获长距离依赖方面表现出色，但由于自注意力机制的计算成本随输入大小呈二次方增长，这带来了高昂的计算成本，特别是对于高分辨率生物医学图像而言，构成了一个挑战。最近，状态空间模型（SSM）的发展，特别是结构化SSM（S4），以其处理长序列的高效性能提供了一个有前景的解决方案。Mamba模型通过选择性机制和硬件优化增强了S4，在密集数据域中表现出卓越的性能。在视觉状态空间模型（VMamba）中引入跨扫描模块（CSM）进一步增强了Mamba在计算机视觉任务中的适用性，它能够实现空间域的遍历，并将非因果视觉图像转换为有序的块序列。 受这些能力的启发，我们提出在U-Net架构中利用视觉Mamba块（VSS）来改进医学图像分析中的长距离依赖建模，从而产生了Mamba-UNet。 图1简要说明了U-Net随着各种网络块的演变以及我们提出的Mamba-UNet的定位。

图2. Mamba-UNet的架构，由编码器、瓶颈层、解码器和跳跃连接组成。编码器、瓶颈层和解码器都是基于视觉Mamba块构建的。

2.方法

2.1架构概览

  所提出的Mamba-UNet的架构如图2所示，该架构受UNet和Swin-UNet的启发。首先，将尺寸为H × W × 1的输入二维灰度图像分割成与ViT和VMamba类似的块，然后转换为维度为H/4 × W/4 × 16的一维序列。一个初始的线性嵌入层将特征维度调整为任意大小，表示为C。然后，这些块标记通过多个VSS块和块合并层进行处理，以创建层次化特征。块合并层负责下采样和维度增加，而VSS块则专注于学习特征表示。编码器每个阶段的输出分辨率分别为H/4 × W/4 × C、H/8 × W/8 × 2C、H/16 × W/16 × 4C和H/32 × W/32 × 8C。

  解码器包含VSS块和块扩展层，遵循编码器的风格，以输出与编码器相同大小的特征，从而通过跳跃连接增强在下采样过程中丢失的空间细节。在编码器和解码器的每个部分中，都使用了2个VSS块，并且在编码器中加载了预训练的VMamba-Tiny，这与Swin-UNet加载预训练的SwinViT-Tiny的过程相同。 VSS块、编码器中的块合并和解码器中的块扩展的详细信息将在以下小节中讨论。

2.2 VSS Block

  VSS网络块如图3所示，它主要基于Visual Mamba。具体来说，传统的SSM（状态空间模型）作为线性时不变系统函数，通过隐藏状态h(t) ∈ R^N将x(t) ∈ R映射到y(t) ∈ R，其中A ∈ C^(N×N) 是演化参数，B, C ∈ C^N是状态大小为N的投影参数，D ∈ C¹是跳跃连接。该模型可以表示为下面方程中的线性常微分方程（ODEs）。

  该线性模型的离散版本可以通过给定时间尺度参数∆ ∈ R^D的零阶保持器进行变换。

  其中B, C ∈ R^D×N（B, C属于实数集）。使用一阶泰勒级数对B进行细化后的近似值为

  Visual Mamba进一步引入了跨扫描模块（CSM），并将卷积操作集成到块中。在VSS块中，输入特征首先经过一个线性嵌入层，然后分叉成两条路径。一条路径经过深度卷积和SiLU激活函数，然后进入SS2D模块，并通过层归一化后，与另一条路径经过SiLU激活函数后的流合并。与典型的视觉Transformer不同，VSS块避开了位置嵌入，选择了没有MLP阶段的简化结构，从而在相同的深度预算内能够更密集地堆叠块。

图3. 视觉状态空间（VSS）块的详细结构

2.3 Encoder

  在编码器中，具有降低分辨率的C维标记化输入会经过两个连续的VSS块进行特征学习，同时保持维度和分辨率。在Mamba-UNet的编码器中，块合并作为下采样过程被使用了三次，每次通过将输入分割成四个象限（即1/4大小），然后将它们连接起来，并通过层归一化来规范维度，从而将标记数量减少一半，同时将特征维度加倍。

2.4 Decoder

  与编码器相对应，解码器使用两个连续的VSS块进行特征重建，并采用块扩展层而不是合并层来上采样深层特征。这些层通过增强分辨率（2倍上采样）并将特征维度减半（即1/2），从而实现特征维度的调整。以初始层为例，它首先加倍特征维度，然后重新组织和减少这些维度以增强分辨率。

2.5 Bottleneck & Skip Connetions

  在Mamba-UNet的瓶颈部分，使用了两个VSS块。编码器和解码器的每一层都采用了跳跃连接，以将多尺度特征与上采样输出进行融合，通过合并浅层和深层来增强空间细节。随后，一个线性层保持了这个综合特征集的维度，以确保与上采样分辨率的一致性。

3实验和结果

3.1数据集

  自动心脏诊断挑战赛

  我们的实验使用了MICCAI 2017挑战赛上公开可用的ACDC MRI心脏分割数据集。该数据集包含来自100名患者的MRI扫描图像，这些图像被标注了多种心脏结构，如右心室、左心室的心内膜和心外膜壁。它涵盖了多种病理状况，分为五个亚组：正常、心肌梗死、扩张型心肌病、肥厚型心肌病和异常右心室，确保了特征信息的广泛分布。在ACDC中验证了4类感兴趣区域（ROI）。

  Synapse多器官分割挑战赛

  我们还使用了MICCAI 2015多图谱腹部标记挑战赛中的30个腹部CT扫描图像，总共包含3779张轴向对比增强的腹部临床CT图像。在Synapse中验证了9类感兴趣区域（ROI）。为了满足Swin-Transformer和Visual-Mamba预训练网络的输入要求，所有图像都被调整为224×224像素。数据集被划分为测试集占20%，其余用于训练（包括验证）。

3.2实现细节

  实现是在Ubuntu 20.04系统上进行的，使用了Python 3.8.8、PyTorch 1.10和CUDA 11.3。硬件设置包括Nvidia GeForce RTX 3090 GPU和Intel Core i9-10900K CPU。对于ACDC数据集，平均运行时间约为5小时，而Synapse数据集则为12小时，这包括了数据传输、模型训练和推理过程。数据集特别针对2D图像分割进行了处理。Mamba-UNet模型进行了10,000次迭代训练，批处理大小为24。采用了随机梯度下降（SGD）优化器，学习率设置为0.01，动量为0.9，权重衰减设置为0.0001。每200次迭代在验证集上评估网络性能，并且仅在验证集上取得新的最佳性能时才保存模型权重。

3.3基线方法

  为了比较的目的，所有基线方法也都在相同的超参数配置下进行了训练。将Mamba-UNet与其他基线方法，包括UNet、Attention UNet、TransUNet和Swin-UNet进行了直接比较。

3.4评估指标

  与基线方法的比较中，对Mamba-UNet的评估采用了广泛的评估指标。相似性度量指标（其值越高越好）包括：Dice系数、交并比（IoU）、准确率、精确率、灵敏度和特异性， 这些指标用向上箭头（↑）表示，意味着更高的值代表更好的性能。相反，差异度量指标如95%的Hausdorff距离（HD）和平均表面距离（ASD），用向下箭头（↓）表示，其值越低越好，表明预测分割与真实分割之间的相似性更高。

  其中，TP代表真正例的数量，TN代表真反例的数量，FP代表假正例的数量，FN代表假反例的数量。

  其中，a和b分别代表预测表面和真实表面上点的集合。d(a, b)表示两点之间的欧几里得距离。95%的Hausdorff距离是Hausdorff距离的一个修改版本，它关注距离的第95百分位数，以减少离群值的影响。

3.5定性结果

  图4和图5分别展示了在ACDC和Synapse数据集上随机选取的三个样本原始图像，以及它们对应的真实标签（ground truth）和包括Mamba-UNet在内的所有基线方法的推理结果。

图4. 在ACDC MRI心脏分割数据集上，Mamba-UNet与其他分割方法相对于真实标签的分割结果的可视化比较。（a）原始MRI图像，（b）真实标签，（c）UNet，（d）Attention UNet，（e）TransUNet，（f）Swin-UNet，（g）Mamba-UNet。

3.6定量结果

  表1和表2报告了Mamba-UNet与其他分割网络在两个数据集上的直接比较，包括相似性度量和差异度量。最佳性能以粗体显示，Mamba-UNet的第二佳性能以下划线表示。量化结果表明，Mamba-UNet更有可能预测出精确的分割掩码。为了进一步在测试集上验证Mamba-UNet，我们还以逐图的方式进行了验证，并根据Dice系数绘制了分割预测的分布图（如图6所示），其中X轴为Dice系数，Y轴为预测数量。该直方图进一步证明了Mamba-UNet更有可能提供具有高Dice系数性能的预测。

4.结论

  在本文中，我们介绍了Mamba-UNet，这是一种基于纯视觉Mamba块的UNet风格网络，用于医学图像分割。实验结果表明，与经典的类似网络（如UNet和Swin-UNet）相比，Mamba-UNet具有优越的性能。在未来，我们计划对不同模态和目标的更多医学图像分割任务进行更深入的研究，并与更多的分割骨干网络进行比较。此外，我们计划将Mamba-UNet扩展到3D医学图像和半监督/弱监督学习，以进一步推动医学成像领域的发展。

图5. 在Synapse CT腹部分割数据集上，Mamba-UNet与其他分割方法相对于真实标签的分割结果的可视化比较。（a）原始CT图像（注意这里是CT图像，而不是MRI图像，可能是原文中的一个小错误），（b）真实标签，（c）UNet，（d）Attention UNet，（e）TransUNet，（f）Swin-UNet，（g）Mamba-UNet。

表1. ACDC MRI心脏测试集上分割网络性能的直接比较


图6. Mamba-UNet与其他分割方法相对于真实标签的Dice系数分布的直方图。
（a）在ACDC MRI心脏数据集上的测试。
（b）在Synapse CT腹部分割数据集上的测试。

表2. Synapse CT腹部测试集上分割网络性能的直接比较