[论文解读] VM-UNET-V2 Rethinking Vision Mamba UNet for Medical Image Segmentation

Mingya Zhang, Yue Yu+3 more|arXiv (Cornell University)|Mar 14, 2024

Medical Image Segmentation Techniques被引用 7

一句话总结

VM-UNetV2 将 Vision State Space Models (VSS) 与 SDI 集成在一个类 UNet 架构中，以高效建模医疗图像分割的长程依赖，在多个数据集上实现具有竞争力的结果，同时具备有利的 FLOPs、Params 和 FPS。

ABSTRACT

In the field of medical image segmentation, models based on both CNN and Transformer have been thoroughly investigated. However, CNNs have limited modeling capabilities for long-range dependencies, making it challenging to exploit the semantic information within images fully. On the other hand, the quadratic computational complexity poses a challenge for Transformers. Recently, State Space Models (SSMs), such as Mamba, have been recognized as a promising method. They not only demonstrate superior performance in modeling long-range interactions, but also preserve a linear computational complexity. Inspired by the Mamba architecture, We proposed Vison Mamba-UNetV2, the Visual State Space (VSS) Block is introduced to capture extensive contextual information, the Semantics and Detail Infusion (SDI) is introduced to augment the infusion of low-level and high-level features. We conduct comprehensive experiments on the ISIC17, ISIC18, CVC-300, CVC-ClinicDB, Kvasir, CVC-ColonDB and ETIS-LaribPolypDB public datasets. The results indicate that VM-UNetV2 exhibits competitive performance in medical image segmentation tasks. Our code is available at https://github.com/nobodyplayer1/VM-UNetV2.

研究动机与目标

• 为医疗图像提出一个结合长程建模与线性复杂度的分割模型。
• 提出 VM-UNetV2，使用 VSS 块和 SDI 来融合低层和高层特征。
• 在皮肤病学和胃肠道息肉数据集上展示有竞争力的性能。
• 分析模型复杂度（FLOPs、Params、FPS）并对编码器深度和深度监督进行消融研究。

提出的方法

• 采用三个模块的架构：编码器、语义与细节注入（SDI）以及解码器。
• 将 Vision Mamba（VSS）块作为编码器主干，以线性复杂度捕获长程上下文。
• 通过 SDI 模块在 CBAM 基础的注意力引导下融合多尺度特征。
• 训练阶段应用跨阶段深度监督策略。
• 使用交叉熵和 Dice 损失（L = L_BCE + L_Dice）进行两类分割的训练。
• 用在 ImageNet-1k 上预训练的 VMamba 初始化编码器权重。

实验结果

研究问题

• RQ1 Vision State Space Models (SSMs) 能否在线性复杂度下为医疗图像分割提供具有竞争力的长程上下文建模？
• RQ2 将 SDI 融入以实现语义与细节注入，是否能在利用高级语义的同时更好地保留细节？
• RQ3 编码器深度和深度监督对不同医疗数据集的分割性能有何影响？

主要发现

• 与强基线相比，VM-UNetV2 在 ISIC17/18 和若干息肉数据集上实现了具有竞争力的 mIoU、DSC 和准确率。
• 在 ISIC17 上，VM-UNetV2 达到 mIoU 82.34、DSC 90.31、Acc 96.70、Spe 97.67、Sen 91.89；在 ISIC18 上，mIoU 81.37、DSC 89.73、Acc 95.06、Spe 97.13、Sen 88.64。
• 在 Kvasir-SEG、ClinicDB、ColonDB、ETIS 和 CVC-300 上，VM-UNetV2 相较于 VM-UNet 在 mIoU 和 DSC 有所提升，并在基线的 UNetV2 上具竞争力。
• VM-UNetV2 在 FLOPs、Params 和 FPS 方面优于若干基线，显示出有利的效率（如表 3）。
• 消融研究表明编码器深度在 [2,2,9,2] 附近且深度监督通常能提升性能，尽管收益因数据集而异。