轻量级医学图像分割模型,高效捕捉长距离上下文关系
U-RWKV: Lightweight medical image segmentation with direction-adaptive RWKV
- 用方向自适应RWKV模块实现全局上下文建模
- 计算复杂度仅为O(N),在多个数据集上达顶尖性能
- 适合移动端或边缘设备部署,推动医疗影像普惠
实现医疗可及性的公平性需要在资源受限环境中部署轻量且高性能的医学图像分割方案。现有方法如U-Net及其变体通常存在全局有效感受野(ERF)有限的问题,难以捕捉长程依赖。为此,我们提出U-RWKV框架,利用循环加权键值(RWKV)架构,在O(N)计算成本下实现高效的长程建模。该框架引入两项关键创新:方向自适应RWKV模块(DARM)和阶段自适应挤压-激励模块(SASE)。DARM采用双RWKV与四扫描机制,跨图像聚合上下文信息,缓解方向偏差,同时保持全局上下文与高计算效率。SASE动态适配不同特征提取阶段,平衡高分辨率细节保留与语义关系捕获。实验表明,U-RWKV在保持高计算效率的同时达到最先进的分割性能,为资源受限环境中的先进医学影像技术普及提供了可行方案。代码已开源:https://github.com/hbyecoding/U-RWKV。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。