摘要

在医学图像分析的最新进展中，卷积神经网络（CNN）和视觉变换器（ViT）已经设定了重要的基准。前者擅长通过卷积运算捕捉局部特征，而后者通过利用自注意机制实现了显著的全局上下文理解。然而，这两种架构在有效地建模医学图像中的长程依赖性方面都存在局限性，这是精确分割的关键方面。受Mamba架构的启发，我们提出了Mamba UNet，这是一种将医学图像分割中的U-Net与Mamba的能力协同的新架构。Mamba UNet采用了纯基于视觉Mamba（VMamba）的编码器-解码器结构，融入了跳跃连接，以保存网络不同规模的空间信息。这种设计有助于全面的特征学习过程，捕捉医学图像中复杂的细节和更广泛的语义上下文。我们在VMamba块中引入了一种新的集成机制，以确保编码器和解码器路径之间的无缝连接和信息流，从而提高分割性能。我们在公开的MRI心脏多结构分割数据集上进行了实验。结果表明，在相同的超参数设置1下，Mamba UNet在医学图像分割方面优于UNet、SwinUNet。源代码和基准实现可在https://github.com/ziyangwang007/Mamba-UNet.

关键词：医学图像分割·卷积·变换器·曼巴·状态空间模型。

设置1：超参数设置包括：损失函数、优化器、训练迭代、批量大小、学习率、相同数据分割等。

1.引言

医学图像分割对于诊断和治疗至关重要，基于深度学习的网络在该领域表现出了主导性能[20]。U-Net是以其对称编码器-解码器风格的架构和跳过连接[24]而闻名的最基本的架构之一，其中各种编码器和解码器提取不同级别的特征信息，并且跳过连接能够实现特征信息的有效转换。大多数研究使用高级网络块技术进一步探索U-Net，如密集连接[12]、残差块[10]、注意力机制[30]、深度卷积[11]和萎缩卷积[33，35]，从而在CT、MRI和超声医学图像分割中产生各种改进的UNet[23，13，15，29，34，36]。

受自然语言处理的自注意机制成功的启发[26]，ViT首次将纯多头自注意机制用于图像识别任务，具有最先进的性能[5]。这展示了其在建模长期依赖关系方面的强大能力。移位窗口等技术进一步定制了ViT，产生了SwinTransformer[18]，增强了其在计算机视觉中密集预测任务中的适用性，如图像分割和检测[19，31，17]。在医学图像分割中，受传统CNN设计的启发，ViT与U-Net架构的集成也产生了各种混合和纯基于ViT的U-Net。例如，TransUNet是第一个在UNet[4]的编码器中利用ViT的特征学习能力的工作。UNETR将ViT与UNet相结合用于3D分割[9]，而Swin UNet和DCSUnet进一步探索了具有基于U-Net结构的纯Swin Vision Transformer网络块[3，28]。

虽然Transformers擅长捕捉长程依赖关系，但由于自注意机制与输入大小的二次缩放，其高计算成本带来了挑战，尤其是对于高分辨率生物医学图像[32，21]。状态空间模型（SSM）[6，22，27]，特别是结构化SSM（S4）[8]的最新发展，以其在处理长序列方面的高效性能提供了一种有前途的解决方案。Mamba模型通过选择性机制和硬件优化增强了S4，在密集数据域中显示出卓越的性能[7]。在视觉状态空间模型（VMamba）中引入交叉扫描模块（CSM），通过实现空间域的遍历并将非因果视觉图像转换为有序的补丁序列，进一步增强了Mamba对计算机视觉任务的适用性[16]。受这些功能的启发，我们建议在U-Net架构中利用Visual Mamba块（VSS）来改进医学图像分析中的长期依赖性建模，从而产生Mamba UNet。U-Net与各种网络块的演变以及我们提出的Mamba UNet的定位如图1所示。

2.方法

2.1 体系结构概述

图2概述了拟议的Mamba UNet的架构，其动机是UNet[24]和Swin UNet[3]。输入的H×W×1的二维灰度图像首先被吐出类似于ViT和VMamba[5，16]的块，然后被吐出尺寸为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[16]，遵循与Swin UNet加载预训练SwinViT Tiny[3]相同的过程。以下小节将讨论VSS块、编码器的补丁合并和解码器的补丁扩展的细节。

2.2 VSS块

VSS网络块如图3所示，主要基于Visual Mamba[16]。在VSS块中，输入特征首先遇到线性嵌入层，然后分叉为双路径。一个分支经历深度卷积[11]和SiLU激活[25]，继续到SS2D模块，以及层后归一化，与SiLU激活后的备用流合并。与典型的视觉转换器不同，这种VSS块避免了位置嵌入，选择了无MLP阶段的流线型结构，从而在相同的深度预算内实现了更密集的块堆叠。

2.3 编码器

在编码器中，分辨率降低的C维标记化输入经过两个连续的VSS块进行特征学习，保持尺寸和分辨率。在Mamba UNet的编码器中，补丁合并作为下采样过程被使用了三次，通过将输入分割成4个象限，将它们连接起来，然后每次通过层形式对维度进行归一化，将令牌计数减少了1/2，同时将特征维度增加了2倍。

2.4 解码器

与编码器类似，解码器利用两个连续的VSS块进行特征重构，采用补丁扩展层而不是合并层来放大深度特征[3]。这些层提高了分辨率（2×放大），同时将特征尺寸减半12，例如，初始层在重组和缩小特征尺寸以提高分辨率之前将特征尺寸加倍。

2.5 瓶颈和跳过连接

两个VSS块用于Mamba UNet的瓶颈。每一级编码器和解码器都采用跳跃连接，将多尺度特征与放大输出相融合，通过合并浅层和深层来增强空间细节。随后的线性层保持了该集成特征集的维度，确保了与放大分辨率的一致性。

3.实验和结果

3.1 数据集

我们使用MICCAI 2017挑战赛[1]中公开的ACDC MRI心脏分割数据集进行了实验。该数据集包括来自100名患者的MRI扫描，注释了多个心脏结构，如右心室以及左心室的心内膜和心外膜。它包括各种各样的病理状况，分为五个亚组：正常、心肌梗死、扩张型心肌病、肥厚型心肌病和异常右心室，确保了特征信息的广泛分布。为了满足ViT分割骨干网络的输入要求，将所有图像的大小调整为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。平均运行时间约为5小时，包括数据传输、模型训练和推理过程。该数据集专门用于2D图像分割。Mamba UNet模型接受了10000次迭代的训练，批量大小为24。采用随机梯度下降（SGD）优化器[2]，学习率为0.01，动量为0.9，权重衰减设置为0.0001。每200次迭代在验证集上评估网络性能，只有在验证集获得新的最佳性能后，才能保存模型权重。

3.3 基线方法

为了进行比较，UNet和Swin-UNet也在相同的超参数配置下进行训练。直接比较Mamba UNet以及其他基线方法，包括UNet[24]和Swin UNet[3]。

3.4 评估指标

根据基线方法对曼巴UNet进行的评估采用了广泛的评估指标。相似性度量，最好更高，包括：骰子、并集交集（IoU）、准确性、精度、灵敏度和特异性，用向上箭头表示(↑) 以指示更高的值表示更好的性能。相反，差异度量，如Hausdorff距离（HD）95%和平均表面距离（ASD），用向下箭头标记(↓), 当较低时是可取的，这表明预测的分割和基本事实分割之间更接近相似性。

其中，TP表示真阳性数，TN表示真阴性数，FP表示假阳性数，FN表示假阴性数。

其中，a和b分别表示预测曲面和地面实况曲面上的点集。d（a，b）表示两点之间的欧几里得距离。95%是Hausdorff距离的修改版本，专注于距离的第95个百分位，以减少异常值的影响。

3.5 定性结果

图4显示了三个随机选择的样本原始图像，根据包括Mamba UNet在内的所有基线方法的公布地面实况进行相应推断，其中不同的颜色显示了地面实况的边界。

3.6 定量结果

表1报告了Mamba UNet与其他分割网络的直接比较，包括相似性度量和差异性度量。表现最好的是Bold，其次是Mamba UNet的Underline。定量结果表明，Mamba UNet更有可能预测精确的分割掩模。为了在测试集上进一步验证MambaUNet，我们还以逐图像的方式在图像上进行验证，根据骰子系数的分割预测分布如图5所示，其中X轴为骰子系数，Y轴为预测次数。该直方图进一步表明，Mamba UNet更有可能提供具有高骰子系数性能的预测。

4.结论

在本文中，我们介绍了Mamba UNet，它是一种纯粹基于视觉Mamba块的UNet风格的医学图像分割网络。性能表明，与UNet和Swin UNet等经典的类似网络相比，Mamba UNet具有优越的性能。未来，我们的目标是从不同的模式和目标对更多的医学图像分割任务进行更深入的探索，并与更多的分割骨干进行比较。此外，我们的目标是将Mamba UNet扩展到3D医学图像，以及半/弱监督学习[14]，以进一步促进医学成像的发展。

引用

1. Bernard, O., et al.: Deep learning techniques for automatic mri cardiac multistructures segmentation and diagnosis: is the problem solved? IEEE transactions on medical imaging 37(11), 2514–2525 (2018)
2. Bottou, L.: Stochastic gradient learning in neural networks. In: Proceedings of Neuro-Nîmes 91. EC2, Nimes, France (1991)
3. Cao, H., Wang, Y., Chen, J., Jiang, D., Zhang, X., Tian, Q., Wang, M.: Swin-unet: Unet-like pure transformer for medical image segmentation. In: European conference on computer vision. pp. 205–218. Springer (2022)
4. Chen, J., Lu, Y., Yu, Q., Luo, X., Adeli, E., Wang, Y., Lu, L., Yuille, A.L., Zhou,Y.: Transunet: Transformers make strong encoders for medical image segmentation.arXiv preprint arXiv:2102.04306 (2021)
5. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner,T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.: An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 (2020)
6. Gu, A.: Modeling Sequences with Structured State Spaces. Ph.D. thesis, Stanford University (2023)
7. Gu, A., Dao, T.: Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 (2023)
8. Gu, A., Goel, K., Ré, C.: Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 (2021)
9. Hatamizadeh, A., Tang, Y., Nath, V., Yang, D., Myronenko, A., Landman, B.,Roth, H.R., Xu, D.: Unetr: Transformers for 3d medical image segmentation. In:Proceedings of the IEEE/CVF winter conference on applications of computer vision. pp. 574–584 (2022)
10. He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In:Proceedings of the IEEE conference on computer vision and pattern recognition.pp. 770–778 (2016)
11. Howard, A.G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., Adam, H.: Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv:1704.04861 (2017)
12. Huang, G., Liu, Z., Van Der Maaten, L., Weinberger, K.Q.: Densely connected convolutional networks. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 4700–4708 (2017)
13. Ibtehaz, N., Rahman, M.S.: Multiresunet: Rethinking the u-net architecture for multimodal biomedical image segmentation. Neural networks 121, 74–87 (2020)
14. Jiao, R., Zhang, Y., Ding, L., Xue, B., Zhang, J., Cai, R., Jin, C.: Learning with limited annotations: A survey on deep semi-supervised learning for medical image segmentation. Computers in Biology and Medicine (2023)
15. Li, X., Chen, H., Qi, X., Dou, Q., Fu, C.W., Heng, P.A.: H-denseunet: hybrid densely connected unet for liver and tumor segmentation from ct volumes. IEEE transactions on medical imaging 37(12), 2663–2674 (2018)
16. Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang, Y., Ye, Q., Liu, Y.: Vmamba:Visual state space model. arXiv preprint arXiv:2401.10166 (2024)
17. Liu, Z., Hu, H., Lin, Y., Yao, Z., Xie, Z., Wei, Y., Ning, J., Cao, Y., Zhang, Z., Dong,L., et al.: Swin transformer v2: Scaling up capacity and resolution. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp.
    12009–12019 (2022)
18. Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin-transformer: Hierarchical vision transformer using shifted windows. In: Proceedings of the IEEE/CVF international conference on computer vision. pp. 10012–10022
    (2021)
19. Liu, Z., Ning, J., Cao, Y., Wei, Y., Zhang, Z., Lin, S., Hu, H.: Video swin transformer. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 3202–3211 (2022)
20. Long, J., Shelhamer, E., Darrell, T.: Fully convolutional networks for semantic segmentation. In: Proceedings of the IEEE conference on computer vision and pattern recognition. pp. 3431–3440 (2015)
21. Ma, J., Li, F., Wang, B.: U-mamba: Enhancing long-range dependency for biomedical image segmentation. arXiv preprint arXiv:2401.04722 (2024)
22. Mehta, H., Gupta, A., Cutkosky, A., Neyshabur, B.: Long range language modeling via gated state spaces. arXiv preprint arXiv:2206.13947 (2022)
23. Oktay, O., Schlemper, J., Folgoc, L.L., Lee, M., Heinrich, M., Misawa, K., Mori,K., McDonagh, S., Hammerla, N.Y., Kainz, B., et al.: Attention u-net: Learning where to look for the pancreas. arXiv preprint arXiv:1804.03999 (2018)
24. Ronneberger, O., et al: U-net: Convolutional networks for biomedical image segmentation. In: MICCAI (2015)
25. Shazeer, N.: Glu variants improve transformer. arXiv preprint arXiv:2002.05202(2020)
26. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser,Ł., Polosukhin, I.: Attention is all you need. Advances in neural information processing systems 30 (2017)
27. Wang, J., Zhu, W., Wang, P., Yu, X., Liu, L., Omar, M., Hamid, R.: Selective structured state-spaces for long-form video understanding. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 6387–
    6397 (2023)
28. Wang, Z., Su, M., Zheng, J.Q., Liu, Y.: Densely connected swin-unet for multiscale information aggregation in medical image segmentation. In: 2023 IEEE International Conference on Image Processing (ICIP). pp. 940–944. IEEE (2023)
29. Wang, Z., Zhang, Z., Voiculescu, I.: Rar-u-net: a residual encoder to attention decoder by residual connections framework for spine segmentation under noisy labels. In: 2021 IEEE International Conference on Image Processing (ICIP). pp.
    21–25. IEEE (2021)
30. Woo, S., Park, J., Lee, J.Y., Kweon, I.S.: Cbam: Convolutional block attention module. In: Proceedings of the European conference on computer vision (ECCV).pp. 3–19 (2018)
31. Xie, Z., Lin, Y., Yao, Z., Zhang, Z., Dai, Q., Cao, Y., Hu, H.: Self-supervised learning with swin transformers. arXiv preprint arXiv:2105.04553 (2021)
32. Xing, Z., Ye, T., Yang, Y., Liu, G., Zhu, L.: Segmamba: Long-range sequential modeling mamba for 3d medical image segmentation. arXiv preprint arXiv:2401.13560
    (2024)
33. Yu, F., Koltun, V.: Multi-scale context aggregation by dilated convolutions. arXiv preprint arXiv:1511.07122 (2015)
34. Zhang, Y., Yuan, L., Wang, Y., Zhang, J.: Sau-net: efficient 3d spine mri segmentation using inter-slice attention. In: Medical Imaging With Deep Learning. pp.903–913. PMLR (2020)
35. Zhou, X.Y., Zheng, J.Q., Li, P., Yang, G.Z.: Acnn: a full resolution dcnn for medical image segmentation. In: 2020 IEEE International Conference on Robotics and Automation (ICRA). pp. 8455–8461. IEEE (2020)
36. Zhou, Z., Siddiquee, M.M.R., Tajbakhsh, N., Liang, J.: Unet++: Redesigning skip connections to exploit multiscale features in image segmentation. IEEE transactions on medical imaging 39(6), 1856–1867 (2019)