RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation

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.