[论文解读] H-vmunet: High-order Vision Mamba UNet for Medical Image Segmentation

Renkai Wu, Yinghao Liu+2 more|arXiv (Cornell University)|Mar 20, 2024

Image Processing Techniques and Applications被引用 4

一句话总结

本文提出 H-vmunet，一种高阶视觉 Mamba UNet，通过引入高阶二维选择性扫描（H-SS2D）减少冗余信息，同时保留全局感受野，从而提升医学图像分割性能。通过将 H-SS2D 集成到 UNet 架构中的新型高阶视觉状态空间（H-VSS）模块，H-vmunet 在参数量上比 VM-UNet 减少 67.28%，并在 ISIC2017、Spleen 和 CVC-ClinicDB 数据集上优于多个基线模型。

ABSTRACT

In the field of medical image segmentation, variant models based on Convolutional Neural Networks (CNNs) and Visual Transformers (ViTs) as the base modules have been very widely developed and applied. However, CNNs are often limited in their ability to deal with long sequences of information, while the low sensitivity of ViTs to local feature information and the problem of secondary computational complexity limit their development. Recently, the emergence of state-space models (SSMs), especially 2D-selective-scan (SS2D), has had an impact on the longtime dominance of traditional CNNs and ViTs as the foundational modules of visual neural networks. In this paper, we extend the adaptability of SS2D by proposing a High-order Vision Mamba UNet (H-vmunet) for medical image segmentation. Among them, the proposed High-order 2D-selective-scan (H-SS2D) progressively reduces the introduction of redundant information during SS2D operations through higher-order interactions. In addition, the proposed Local-SS2D module improves the learning ability of local features of SS2D at each order of interaction. We conducted comparison and ablation experiments on three publicly available medical image datasets (ISIC2017, Spleen, and CVC-ClinicDB), and the results all demonstrate the strong competitiveness of H-vmunet in medical image segmentation tasks. The code is available from https://github.com/wurenkai/H-vmunet .

研究动机与目标

• 为解决 CNN 在捕捉长距离依赖方面的局限性，以及 ViT 在局部特征敏感性和计算成本方面的不足。
• 克服标准二维选择性扫描（SS2D）操作在状态空间模型（SSMs）中存在冗余和次优局部特征学习的问题。
• 通过集成基于 Mamba 的模块，提升 UNet 框架在医学图像分割中的性能，实现高阶空间交互。
• 开发一种参数高效、高性能的分割模型，在保持强全局上下文的同时最小化噪声与冗余。

提出的方法

• 提出高阶二维选择性扫描（H-SS2D），一种分阶段逐步减少特征处理过程中冗余信息的 SS2D 机制。
• 引入 Local-SS2D 模块，以在每一阶交互中提升局部特征学习能力，增强对细粒度细节的敏感度。
• 基于 H-SS2D 设计高阶视觉状态空间（H-VSS）模块，替代传统 VSS，实现分层、顺序感知的特征提取。
• 通过在编码器-解码器架构中集成 H-VSS 模块并引入跳跃连接，构建高阶视觉 Mamba UNet（H-vmunet），用于医学图像分割。
• 在编码器和解码器模块中采用 [2,3,4,5] 的混合阶次策略，以平衡性能与计算成本。
• 通过在三个公开数据集上的消融实验与对比分析，验证 H-SS2D 和 H-VSS 组件的有效性。

实验结果

研究问题

• RQ1 SS2D 操作中的高阶空间交互是否能在保持全局感受野的同时减少冗余信息？
• RQ2 将 H-SS2D 集成到视觉 Mamba 模块中，是否能提升局部特征学习能力与分割精度？
• RQ3 在医学图像分割基准上，所提出的 H-vmunet 与 VM-UNet 和 MHorUNet 相比，在性能与参数效率方面表现如何？
• RQ4 在分割精度与计算效率方面，H-SS2D 的最优阶次配置是什么？

主要发现

• H-vmunet 在 ISIC2017、Spleen 和 CVC-ClinicDB 三个医学图像数据集上均达到最先进性能，优于 MHorUNet 和 VM-UNet。
• 与基线模型 VM-UNet 相比，参数量减少 67.28%，显著提升参数效率。
• 消融实验表明，H-SS2D 和 Local-SS2D 模块至关重要，完整 H-SS2D 配置达到最高 Dice 评分。
• 混合阶次配置 [2,3,4,5] 的性能优于单一阶次或均匀阶次设置，表明分层交互具有优势。
• 可视化结果表明，H-vmunet 生成的病灶边界更清晰，对小病灶特征的敏感度优于 VM-UNet。
• H-SS2D 机制成功最小化冗余信息，同时保留全局上下文，验证了其设计原理的有效性。