[论文解读] Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation

Ziyang Wang, Jian-Qing Zheng+3 more|arXiv (Cornell University)|Feb 7, 2024

Medical Image Segmentation Techniques被引用 79

一句话总结

Mamba-UNet 在类似 UNet 的编码器-解码器中使用纯 Visual Mamba 模块，以改进长程特征建模，在 MRI 心脏数据上获得的分割准确度高于 UNet 和 Swin-UNet。

ABSTRACT

In recent advancements in medical image analysis, Convolutional Neural Networks (CNN) and Vision Transformers (ViT) have set significant benchmarks. While the former excels in capturing local features through its convolution operations, the latter achieves remarkable global context understanding by leveraging self-attention mechanisms. However, both architectures exhibit limitations in efficiently modeling long-range dependencies within medical images, which is a critical aspect for precise segmentation. Inspired by the Mamba architecture, known for its proficiency in handling long sequences and global contextual information with enhanced computational efficiency as a State Space Model (SSM), we propose Mamba-UNet, a novel architecture that synergizes the U-Net in medical image segmentation with Mamba's capability. Mamba-UNet adopts a pure Visual Mamba (VMamba)-based encoder-decoder structure, infused with skip connections to preserve spatial information across different scales of the network. This design facilitates a comprehensive feature learning process, capturing intricate details and broader semantic contexts within medical images. We introduce a novel integration mechanism within the VMamba blocks to ensure seamless connectivity and information flow between the encoder and decoder paths, enhancing the segmentation performance. We conducted experiments on publicly available ACDC MRI Cardiac segmentation dataset, and Synapse CT Abdomen segmentation dataset. The results show that Mamba-UNet outperforms several types of UNet in medical image segmentation under the same hyper-parameter setting. The source code and baseline implementations are available.

研究动机与目标

• 推动在医学图像分割中对长程依赖建模的改进.
• 提出一种使用 Visual Mamba 模块（VSS）作为编码器、瓶颈和解码器的 UNet 向导架构。
• 通过跳跃连接与基于 VMamba 的框架内的 patch 合并/扩展来保留空间细节。
• 在公开的 MRI 心脏数据集上评估分割性能并与基线进行对比。

提出的方法

• 采用类似 UNet 的编码器-解码器结构，以纯 Visual Mamba 模块作为核心构建块。
• 将输入图像表示为 patch token，然后通过分层 VSS 块进行处理，并进行 patch 合并/扩展。
• 在编码器与解码器之间使用跳跃连接以融合多尺度特征。
• 在编码器中加载预训练的 VMamba-Tiny 以改进初始化。
• 在固定超参数下使用 SGD 进行训练，并用标准分割指标进行评估。

Figure 1: A brief introduction of the evolution of recent developments of UNet with incorporation of Transformer and State Space Models (SSM) for medical image segmentation.

实验结果

研究问题

• RQ1 与传统 UNet 和 ViT 基于方法相比，基于 VMamba 的模块是否能改善医学图像分割中的长程依赖建模？
• RQ2 纯 VMamba UNet 在 MRI 数据上是否在保持计算效率的同时实现更高的分割准确性？
• RQ3 在相同训练设定下，Mamba-UNet 与 UNet 与 Swin-UNet 的性能比较如何？

主要发现

• Mamba-UNet 在 MRI 心脏测试集上实现 Dice 0.9281、IoU 0.8698 和准确度 0.9972。
• 在相同超参数条件下，Mamba-UNet 的 Dice 和 IoU 均优于 UNet。
• Mamba-UNet 在基线方法下对比中显示出有竞争力的 Hausdorff Distance (HD 95%) 和 ASD 分数，HD 2.4645 和 ASD 0.7677。
• Swin-UNet 的 Dice 0.9188 和 IoU 0.8545，均低于 Mamba-UNet，而 UNet 的 Dice 为 0.9248。
• Mamba-UNet 在所报告的指标中表现出更高的边界精度。

Figure 2: The architecture of Mamba-UNet, which is composed of encoder, bottleneck, decoder and skip connections. The encoder, bottleneck and decoder are all constructed based on Visual Mamba block.