KM-UNet KAN Mamba UNet for medical image segmentation

📄 arXiv: 2501.02559v1 📥 PDF

作者: Yibo Zhang

分类: eess.IV, cs.AI, cs.CV

发布日期: 2025-01-05

🔗 代码/项目: GITHUB


💡 一句话要点

提出KM-UNet,结合KAN和Mamba用于提升医学图像分割性能。

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

关键词: 医学图像分割 KAN Mamba 状态空间模型 UNet 深度学习

📋 核心要点

  1. 传统CNN方法在医学图像分割中难以捕捉长程依赖,Transformer方法计算复杂度高。
  2. KM-UNet结合KANs的高效特征表示和SSMs的可扩展长程建模,实现精度与效率的平衡。
  3. 在五个医学图像数据集上,KM-UNet取得了与SOTA方法相媲美的分割性能。

📝 摘要(中文)

医学图像分割是医学影像分析中的关键任务。传统的基于CNN的方法难以建模长程依赖关系,而基于Transformer的模型虽然成功,但计算复杂度呈二次方增长。为了解决这些限制,我们提出了KM-UNet,一种新颖的U型网络架构,它结合了Kolmogorov-Arnold Networks (KANs) 和状态空间模型 (SSMs) 的优势。KM-UNet利用Kolmogorov-Arnold表示定理进行高效的特征表示,并利用SSMs进行可扩展的长程建模,从而在准确性和计算效率之间取得平衡。我们在五个基准数据集上评估了KM-UNet:ISIC17、ISIC18、CVC、BUSI和GLAS。实验结果表明,KM-UNet在医学图像分割任务中实现了与最先进方法相比具有竞争力的性能。据我们所知,KM-UNet是第一个集成KANs和SSMs的医学图像分割框架。这项工作为开发更高效和可解释的医学图像分割系统提供了有价值的基线和新的见解。

🔬 方法详解

问题定义:医学图像分割旨在精确识别和分割医学图像中的特定区域,例如肿瘤、器官等。现有基于CNN的方法难以有效建模长程依赖关系,限制了分割精度。Transformer虽然能捕捉长程依赖,但计算复杂度高,难以应用于高分辨率医学图像。

核心思路:KM-UNet的核心思路是结合Kolmogorov-Arnold Networks (KANs) 和状态空间模型 (SSMs) 的优势。KANs能够通过学习非线性激活函数来高效地表示特征,降低计算复杂度。SSMs,特别是Mamba架构,擅长建模序列数据中的长程依赖关系,且计算复杂度为线性。

技术框架:KM-UNet采用U-Net的编码器-解码器结构。编码器部分使用KANs提取图像特征,并逐步降低分辨率。在编码器的底部,使用Mamba模块来建模长程依赖关系。解码器部分使用KANs进行特征上采样和融合,最终输出分割结果。跳跃连接用于将编码器特征传递到解码器,以保留细节信息。

关键创新:KM-UNet的关键创新在于将KANs和SSMs集成到U-Net架构中,用于医学图像分割。这是首次将KANs应用于医学图像分割任务,并结合Mamba架构进行长程依赖建模。这种结合既能保证特征表示的效率,又能有效捕捉图像中的全局信息。

关键设计:KM-UNet的具体实现细节包括:KANs的参数设置(例如,每层神经元的数量、激活函数的类型),Mamba模块的配置(例如,状态维度、选择机制的参数),以及损失函数的选择(例如,Dice Loss、交叉熵损失)。此外,跳跃连接的实现方式(例如,直接连接、卷积连接)也会影响最终的分割性能。

📊 实验亮点

KM-UNet在ISIC17、ISIC18、CVC、BUSI和GLAS五个医学图像分割数据集上进行了评估,实验结果表明,KM-UNet取得了与现有SOTA方法相媲美的性能。这验证了KANs和SSMs在医学图像分割中的有效性,并为未来的研究提供了新的方向。

🎯 应用场景

KM-UNet在医学图像分割领域具有广泛的应用前景,例如肿瘤检测、器官分割、病灶识别等。它可以辅助医生进行诊断和治疗计划,提高诊断效率和准确性。未来,KM-UNet可以扩展到其他医学影像模态(如CT、MRI),并应用于三维图像分割。

📄 摘要(原文)

Medical image segmentation is a critical task in medical imaging analysis. Traditional CNN-based methods struggle with modeling long-range dependencies, while Transformer-based models, despite their success, suffer from quadratic computational complexity. To address these limitations, we propose KM-UNet, a novel U-shaped network architecture that combines the strengths of Kolmogorov-Arnold Networks (KANs) and state-space models (SSMs). KM-UNet leverages the Kolmogorov-Arnold representation theorem for efficient feature representation and SSMs for scalable long-range modeling, achieving a balance between accuracy and computational efficiency. We evaluate KM-UNet on five benchmark datasets: ISIC17, ISIC18, CVC, BUSI, and GLAS. Experimental results demonstrate that KM-UNet achieves competitive performance compared to state-of-the-art methods in medical image segmentation tasks. To the best of our knowledge, KM-UNet is the first medical image segmentation framework integrating KANs and SSMs. This work provides a valuable baseline and new insights for the development of more efficient and interpretable medical image segmentation systems. The code is open source at https://github.com/2760613195/KM_UNet Keywords:KAN,Manba, state-space models,UNet, Medical image segmentation, Deep learning