[论文解读] Weak-Mamba-UNet: Visual Mamba Makes CNN and ViT Work Better for Scribble-based Medical Image Segmentation

Ziyang Wang, Chao Ma|arXiv (Cornell University)|Feb 16, 2024

AI in cancer detection被引用 14

一句话总结

Weak-Mamba-UNet 将 CNN、Vision Transformer 与 Visual Mamba 网络在基于涂抹的弱监督框架下结合，通过跨监督学习和伪标签来提升医疗图像分割。

ABSTRACT

Medical image segmentation is increasingly reliant on deep learning techniques, yet the promising performance often come with high annotation costs. This paper introduces Weak-Mamba-UNet, an innovative weakly-supervised learning (WSL) framework that leverages the capabilities of Convolutional Neural Network (CNN), Vision Transformer (ViT), and the cutting-edge Visual Mamba (VMamba) architecture for medical image segmentation, especially when dealing with scribble-based annotations. The proposed WSL strategy incorporates three distinct architecture but same symmetrical encoder-decoder networks: a CNN-based UNet for detailed local feature extraction, a Swin Transformer-based SwinUNet for comprehensive global context understanding, and a VMamba-based Mamba-UNet for efficient long-range dependency modeling. The key concept of this framework is a collaborative and cross-supervisory mechanism that employs pseudo labels to facilitate iterative learning and refinement across the networks. The effectiveness of Weak-Mamba-UNet is validated on a publicly available MRI cardiac segmentation dataset with processed scribble annotations, where it surpasses the performance of a similar WSL framework utilizing only UNet or SwinUNet. This highlights its potential in scenarios with sparse or imprecise annotations. The source code is made publicly accessible.

研究动机与目标

• 激励在医疗图像分割中使用涂抹式注释以降低标注成本。
• 提出一个多视角的跨监督框架，在弱监督下同时训练 CNN、ViT 与 VMamba 架构。
• 引入伪标签，将稀疏的涂抹监督转化为对所有网络的密集信号引导。
• 证明在 MRI 心脏数据上，整合 CNN、ViT 和 VMamba 相较于单骨干弱监督方法可获得更优的分割性能。

提出的方法

• 使用三种编码器-解码器骨干网络：UNet（CNN）、SwinUNet（ViT）和 MambaUNet（Visual Mamba）。
• 对涂抹标注的像素应用部分交叉熵损失（忽略未标记区域）。
• 通过三个网络预测的加权组合(Ypseudo = α fcnn + β ftjn + γ fmamba，迭代中随机取 α、β、γ，且 α+β+γ=1) 来创建密集伪标签。
• 使用 Dice 系数损失，结合预测的 argmax 与 Ypseudo，为所有网络提供密集监督。
• 损失函数将每个网络的涂抹损失和 Dice 损失结合起来：Ltotal = sum_i (Lpce^i + Ldice^i) 对 i ∈ {CNN, ViT, Mamba}。
• 网络分别独立初始化（UNet、SwinUNet、MambaUNet），以促进多样化视角，从而实现跨监督。

Figure 1: The Example Images of MRI Cardiac Scans, with the Corresponding Ground Truth, and Scribble-based Annotations.

实验结果

研究问题

• RQ1 在涂抹式监督下，基于 CNN、ViT 与 VMamba 的跨监督多视角框架，与单骨干的 WSL 方法相比，其分割性能如何？
• RQ2 来自不同骨干的伪标签是否能在标注有限的情况下提升密集分割？
• RQ3 在此背景下，Visual Mamba 架构对长程依赖建模的贡献是什么？

主要发现

网络	Dice ↑	Acc ↑	Pre ↑	Sen ↑	Spe ↑	HD ↓	ASD ↓
pCE + UNet	0.7620	0.9807	0.6799	0.9174	0.9823	151.0593	54.6531
USTM + UNet	0.8592	0.9917	0.8128	0.9257	0.9888	99.8293	26.0185
Mumford + UNet	0.8993	0.9950	0.8844	0.9200	0.9874	28.0604	7.3907
Gated CRF + UNet	0.9046	0.9955	0.8890	0.9304	0.9922	7.4340	2.0753
pCE + SwinUNet	0.8935	0.9950	0.8808	0.9129	0.9884	24.4750	6.9108
USTM + SwinUNet	0.9044	0.9957	0.8952	0.9187	0.9898	6.5172	2.2319
Mumford + SwinUNet	0.9051	0.9958	0.8996	0.9157	0.9889	6.0653	1.6482
Gated CRF + SwinUNet	0.8995	0.9955	0.8920	0.9175	0.9904	6.6559	1.6222
Weak-Mamba-UNet	0.9171	0.9963	0.9095	0.9309	0.9920	3.9597	0.8810

• Weak-Mamba-UNet 在 MRI 心脏数据集上实现了最佳的 Dice、准确率、精确度、灵敏度和特异性，优于所测试的其他 WSL 方法。
• 在基线方法中，它还获得了最低的 Hausdorff 距离（HD）和平均表面距离（ASD）。
• 消融实验表明 UNet+SwinUNet+MambaUNet 配置优于任一单骨干或两骨干组合。
• SwinUNet 基于的 WSL 方法通常优于 UNet 基于的方法，但多视图整合进一步带来收益。
• 该方法表明，在涂抹监督下的跨视图协作可在降低标注成本的同时实现高分割性能。

Figure 2: Semi-Mamba-UNet: The Framework of Contrastive Cross-Supervised Visual Mamba-based UNet for Semi-Supervised Medical Image Segmentation.