HMT-UNet: A hybird Mamba-Transformer Vision UNet for Medical Image Segmentation

📄 arXiv: 2408.11289v2 📥 PDF

作者: Mingya Zhang, Zhihao Chen, Yiyuan Ge, Xianping Tao

分类: eess.IV, cs.CV

发布日期: 2024-08-21 (更新: 2024-09-07)

备注: arXiv admin note: text overlap with arXiv:2403.09157; text overlap with arXiv:2407.08083 by other authors

🔗 代码/项目: GITHUB

💡 一句话要点

提出HMT-UNet，一种混合Mamba-Transformer的UNet，用于提升医学图像分割性能。

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

关键词: 医学图像分割 Mamba Transformer UNet 状态空间模型

📋 核心要点

  1. 现有CNN难以建模长程依赖，Transformer计算复杂度高，限制了医学图像分割的性能。
  2. 利用Mamba等状态空间模型建模长程依赖的优势，并结合Transformer的自注意力机制。
  3. 提出的HMT-UNet在多个医学图像分割数据集上取得了有竞争力的性能。

📝 摘要（中文）

本文针对医学图像分割领域，提出了Hybird Transformer vision Mamba UNet (HTM-UNet)。该模型基于状态空间模型(SSM)的混合机制，特别是Mamba，并结合Transformer的优势。传统CNN模型在建模长程依赖关系方面存在局限性，而Transformer的计算复杂度较高。Mamba等SSM模型在建模长程交互方面表现出色，并保持线性计算复杂度。通过精心设计，SSM和Transformer的混合机制能够增强模型对视觉特征的建模能力。实验结果表明，在Mamba架构层后集成自注意力机制能够显著提高捕获长程空间依赖性的能力。在ISIC17、ISIC18、CVC-300等多个公共数据集和ZD-LCI-GIM私有数据集上的实验表明，HTM-UNet在医学图像分割任务中表现出竞争性的性能。

🔬 方法详解

问题定义：医学图像分割任务需要精确地识别和分割图像中的目标区域，例如病灶、器官等。现有的基于CNN的方法在建模长程依赖关系方面存在局限性，难以充分利用图像中的语义信息。而Transformer虽然能够有效地建模长程依赖，但其计算复杂度较高，难以应用于高分辨率的医学图像。

核心思路：本文的核心思路是结合状态空间模型（SSM），特别是Mamba，和Transformer的优势，设计一种混合架构，既能有效地建模长程依赖，又能保持较低的计算复杂度。通过将Mamba的线性计算复杂度和Transformer的自注意力机制相结合，提升模型对视觉特征的建模能力。

技术框架：HTM-UNet采用U型架构，类似于传统的UNet，包含编码器和解码器两个主要部分。编码器负责提取图像的特征表示，解码器负责将特征表示恢复为分割结果。在编码器和解码器的每一层中，都使用了Mamba模块和Transformer模块的混合结构。Mamba模块用于建模长程依赖，Transformer模块用于增强特征表示。

关键创新：最重要的技术创新点是Mamba和Transformer的混合机制。通过将自注意力机制集成到Mamba架构的层后，可以显著提高模型捕获长程空间依赖性的能力。这种混合机制既能利用Mamba的线性计算复杂度，又能利用Transformer的全局建模能力。

关键设计：具体的网络结构细节包括Mamba模块和Transformer模块的堆叠方式、自注意力机制的参数设置、损失函数的选择等。论文中可能还涉及了数据增强、学习率调整等训练技巧，以进一步提升模型的性能。具体细节需要参考论文原文。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

HTM-UNet在ISIC17、ISIC18、CVC-300等多个公共数据集和ZD-LCI-GIM私有数据集上进行了实验，结果表明其在医学图像分割任务中表现出竞争性的性能。具体的性能数据和对比基线需要在论文原文中查找。

🎯 应用场景

该研究成果可应用于多种医学图像分割任务，例如皮肤病灶分割、息肉分割等。通过提高分割精度，可以辅助医生进行诊断和治疗，提高医疗效率和准确性。未来，该模型可以进一步扩展到其他医学图像分析任务，例如图像配准、图像重建等。

📄 摘要（原文）

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. 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. The hybrid mechanism of SSM (State Space Model) and Transformer, after meticulous design, can enhance its capability for efficient modeling of visual features. Extensive experiments have demonstrated that integrating the self-attention mechanism into the hybrid part behind the layers of Mamba's architecture can greatly improve the modeling capacity to capture long-range spatial dependencies. In this paper, leveraging the hybrid mechanism of SSM, we propose a U-shape architecture model for medical image segmentation, named Hybird Transformer vision Mamba UNet (HTM-UNet). We conduct comprehensive experiments on the ISIC17, ISIC18, CVC-300, CVC-ClinicDB, Kvasir, CVC-ColonDB, ETIS-Larib PolypDB public datasets and ZD-LCI-GIM private dataset. The results indicate that HTM-UNet exhibits competitive performance in medical image segmentation tasks. Our code is available at https://github.com/simzhangbest/HMT-Unet.