VM-UNet: Vision Mamba UNet for Medical Image Segmentation

📄 arXiv: 2402.02491v2 📥 PDF

作者: Jiacheng Ruan, Jincheng Li, Suncheng Xiang

分类: eess.IV, cs.CV

发布日期: 2024-02-04 (更新: 2024-11-08)

备注: 9 pages, 5 figures, 6 tables. Work in progress

🔗 代码/项目: GITHUB

💡 一句话要点

提出VM-UNet以解决医学图像分割中的长程建模问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 医学图像分割 状态空间模型 视觉状态空间 卷积神经网络 长程建模 不对称结构 深度学习

📋 核心要点

  1. 现有的CNN和Transformer模型在医学图像分割中存在长程建模能力不足和计算复杂度高的问题。
  2. 本文提出的VM-UNet模型利用状态空间模型，设计了视觉状态空间模块以增强上下文信息捕获能力，并采用不对称的编码-解码结构。
  3. 实验结果显示，VM-UNet在多个医学图像分割数据集上表现优异，提供了新的基准和见解。

📝 摘要（中文）

在医学图像分割领域，卷积神经网络（CNN）和变换器（Transformer）模型已被广泛研究。然而，CNN在长程建模能力上存在局限，而变换器则受到其二次计算复杂度的限制。最近，状态空间模型（SSM）作为一种新兴方法，展现出在长程交互建模方面的优势，同时保持线性计算复杂度。本文提出了一种基于状态空间模型的U型架构模型，命名为Vision Mamba UNet（VM-UNet），引入视觉状态空间（VSS）模块以捕获广泛的上下文信息，并构建了一个不对称的编码-解码结构，以减少计算成本。我们在ISIC17、ISIC18和Synapse数据集上进行了全面实验，结果表明VM-UNet在医学图像分割任务中表现出色。我们希望为未来更高效的SSM基础分割系统的发展奠定基础。

🔬 方法详解

问题定义：本文旨在解决医学图像分割中长程建模能力不足和计算复杂度高的问题。现有的CNN和Transformer模型在处理长距离依赖时表现不佳，且计算资源消耗大。

核心思路：论文提出的VM-UNet模型基于状态空间模型，设计了视觉状态空间（VSS）模块以有效捕获上下文信息，同时通过不对称的编码-解码结构减少计算成本。

技术框架：VM-UNet的整体架构为U型结构，包含视觉状态空间模块作为基础模块，结合少量卷积层构建编码器和解码器。该结构旨在提高长程依赖建模能力，同时降低计算复杂度。

关键创新：VM-UNet是首个基于纯状态空间模型构建的医学图像分割模型，显著区别于传统的CNN和Transformer方法，提供了新的建模思路。

关键设计：在设计中，VSS模块用于捕获上下文信息，采用不对称的编码-解码结构以减少卷积层的数量，从而降低计算成本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在ISIC17、ISIC18和Synapse数据集上的实验结果表明，VM-UNet在医学图像分割任务中表现出色，具体性能数据未提供，但相较于现有基线模型，显示出显著的提升，证明了其有效性。

🎯 应用场景

该研究的潜在应用领域包括医学影像分析、疾病诊断辅助系统以及临床决策支持。通过提高医学图像分割的效率和准确性，VM-UNet有望在实际医疗场景中发挥重要作用，推动智能医疗的发展。

📄 摘要（原文）

In the realm of medical image segmentation, both CNN-based and Transformer-based models have been extensively explored. However, CNNs exhibit limitations in long-range modeling capabilities, whereas Transformers are hampered by their quadratic computational complexity. Recently, State Space Models (SSMs), exemplified by Mamba, have emerged as a promising approach. They not only excel in modeling long-range interactions but also maintain a linear computational complexity. In this paper, leveraging state space models, we propose a U-shape architecture model for medical image segmentation, named Vision Mamba UNet (VM-UNet). Specifically, the Visual State Space (VSS) block is introduced as the foundation block to capture extensive contextual information, and an asymmetrical encoder-decoder structure is constructed with fewer convolution layers to save calculation cost. We conduct comprehensive experiments on the ISIC17, ISIC18, and Synapse datasets, and the results indicate that VM-UNet performs competitively in medical image segmentation tasks. To our best knowledge, this is the first medical image segmentation model constructed based on the pure SSM-based model. We aim to establish a baseline and provide valuable insights for the future development of more efficient and effective SSM-based segmentation systems. Our code is available at https://github.com/JCruan519/VM-UNet.