用新型架构提升医学图像分割精度与效率
RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation
- 将RWKV结构融入U-Net，增强长距离依赖捕捉能力
- 在11个数据集上达到顶尖性能，小模型兼顾效率
- 适合需要高精度与轻量化部署的临床场景
近年来，深度学习在医学图像分割领域取得显著进展，主要依赖卷积神经网络（CNN）和变换器（transformer）模型。然而，CNN难以捕捉长距离依赖，而transformer计算复杂度高。为此，我们提出RWKV-UNet，将受控记忆机制的RWKV结构引入U-Net架构，显著提升模型对长程上下文的理解能力。通过设计结合CNN与RWKV的全局-局部空间感知（GLSP）模块构建强编码器，并提出跨通道混合（CCM）模块，实现多尺度特征融合与全局通道信息整合，优化跳跃连接。在11个基准数据集上的实验表明，RWKV-UNet在多种医学图像分割任务中达到当前最优性能。此外，其小型变体RWKV-UNet-S与RWKV-UNet-T在保持高精度的同时显著降低计算开销，更适合临床实际应用。
原文摘要 · Abstract (English)
In recent years, significant advancements have been made in deep learning for medical image segmentation, particularly with convolutional neural networks (CNNs) and transformer models. However, CNNs face limitations in capturing long-range dependencies, while transformers suffer from high computational complexity. To address this, we propose RWKV-UNet, a novel model that integrates the RWKV (Receptance Weighted Key Value) structure into the U-Net architecture. This integration enhances the model's ability to capture long-range dependencies and to improve contextual understanding, which is crucial for accurate medical image segmentation. We build a strong encoder with developed Global-Local Spatial Perception (GLSP) blocks combining CNNs and RWKVs. We also propose a Cross-Channel Mix (CCM) module to improve skip connections with multi-scale feature fusion, achieving global channel information integration. Experiments on 11 benchmark datasets show that the RWKV-UNet achieves state-of-the-art performance on various types of medical image segmentation tasks. Additionally, smaller variants, RWKV-UNet-S and RWKV-UNet-T, balance accuracy and computational efficiency, making them suitable for broader clinical applications.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。