U-RWKV: Lightweight medical image segmentation with direction-adaptive RWKV

TL;DR

U-RWKV introduces a lightweight, efficient medical image segmentation framework that captures long-range dependencies using novel direction-adaptive modules, outperforming existing methods especially in resource-limited settings.

Contribution

The paper presents U-RWKV, a new architecture with direction-adaptive and stage-adaptive modules, enabling effective long-range modeling with low computational cost for medical image segmentation.

Findings

Achieves state-of-the-art segmentation accuracy.

Maintains high efficiency suitable for resource-limited environments.

Demonstrates robustness across different medical imaging datasets.

Abstract

Achieving equity in healthcare accessibility requires lightweight yet high-performance solutions for medical image segmentation, particularly in resource-limited settings. Existing methods like U-Net and its variants often suffer from limited global Effective Receptive Fields (ERFs), hindering their ability to capture long-range dependencies. To address this, we propose U-RWKV, a novel framework leveraging the Recurrent Weighted Key-Value(RWKV) architecture, which achieves efficient long-range modeling at O(N) computational cost. The framework introduces two key innovations: the Direction-Adaptive RWKV Module(DARM) and the Stage-Adaptive Squeeze-and-Excitation Module(SASE). DARM employs Dual-RWKV and QuadScan mechanisms to aggregate contextual cues across images, mitigating directional bias while preserving global context and maintaining high computational efficiency. SASE dynamically adapts its architecture to different feature extraction stages, balancing high-resolution detail preservation and semantic relationship capture. Experiments demonstrate that U-RWKV achieves state-of-the-art segmentation performance with high computational efficiency, offering a practical solution for democratizing advanced medical imaging technologies in resource-constrained environments. The code is available at https://github.com/hbyecoding/U-RWKV.