Command Palette
Search for a command to run...
RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation
Jiang Juntao ; Zhang Jiangning ; Liu Weixuan ; Gao Muxuan ; Hu Xiaobin ; Yan Xiaoxiao ; Huang Feiyue ; Liu Yong

Abstract
In recent years, there have been significant advancements in deep learningfor medical image analysis, especially with convolutional neural networks(CNNs) and transformer models. However, CNNs face limitations in capturinglong-range dependencies while transformers suffer high computationalcomplexities. To address this, we propose RWKV-UNet, a novel model thatintegrates the RWKV (Receptance Weighted Key Value) structure into the U-Netarchitecture. This integration enhances the model's ability to capturelong-range dependencies and improve contextual understanding, which is crucialfor accurate medical image segmentation. We build a strong encoder withdeveloped inverted residual RWKV (IR-RWKV) blocks combining CNNs and RWKVs. Wealso propose a Cross-Channel Mix (CCM) module to improve skip connections withmulti-scale feature fusion, achieving global channel information integration.Experiments on benchmark datasets, including Synapse, ACDC, BUSI, CVC-ClinicDB,CVC-ColonDB, Kvasir-SEG, ISIC 2017 and GLAS show that RWKV-UNet achievesstate-of-the-art performance on various types of medical image segmentation.Additionally, smaller variants, RWKV-UNet-S and RWKV-UNet-T, balance accuracyand computational efficiency, making them suitable for broader clinicalapplications.
Code Repositories
Benchmarks
| Benchmark | Methodology | Metrics |
|---|---|---|
| medical-image-segmentation-on-acdc | RWKV-UNet | Dice Score: 0.9217 |
| medical-image-segmentation-on-miccai-2015-1 | RWKV-UNet | Avg DSC: 84.02 Avg HD: 15.7 |
Build AI with AI
From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.