[논문 리뷰] RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation

Juntao Jiang, Jiangning Zhang외 6명|ArXiv.org|2025. 01. 14.

Radiomics and Machine Learning in Medical Imaging인용 수 3

한 줄 요약

RWKV-UNet은 Receptance Weighted Key Value (RWKV)를 U-Net과 결합하여 장거리 의존성을 포착하고, IR-RWKV 인코더와 Cross-Channel Mix 모듈을 사용해 다수의 의학 영상 분할 벤치마크에서 최첨단 결과를 달성하며, 효율성을 위한 더 작은 변형이 존재합니다.

ABSTRACT

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.

연구 동기 및 목표

• 의료 영상 분할에서 전통적인 CNN과 트랜스포머를 넘어 장거리 의존성 모델링을 개선할 것을 동기부여한다.
• 공간적 국소성과 RWKV 글로벌 컨텍스트를 혼합하는 하이브리드 인코더(IR-RWKV)를 제안한다.
• Cross-Channel Mix (CCM) 모듈을 도입하여 스킵 연결에서 다중 스케일 특징 융합을 강화한다.
• 다양한 의료 영상 데이터셋에서 최첨단 분할 성능을 입증한다.
• 정확성과 계산 효율성의 균형을 맞춘 더 작은 모델 변형(RWKV-UNet-S/T)을 제공한다.

제안 방법

• 적층된 Inverted Residual RWKV (IR-RWKV) 블록과 IR 블록으로 구성된 인코더를 가진 RWKV-UNet 아키텍처를 도입한다.
• 인코더에서 Vision RWKV의 공간 혼합(spatial-mix)과 깊이별 합성곱(depthwise convolutions)을 활용하여 로컬/글로벌 특징을 통합한다.
• 채널별 글로벌 컨텍스트를 통해 다중 스케일 인코더 특징을 융합하는 Cross-Channel Mix (CCM) 모듈을 추가한다.
• 효율적으로 특징을 업샘플링하고 정제하기 위해 9x9 깊이별 합성곱을 갖춘 CNN 기반 디코더를 설계한다.
• 특징 추출과 수렴을 개선하기 위해 ImageNet-1K에서 인코더를 사전 학습(300 에포크, AdamW)시킨다.
• 정확도와 계산 부하의 균형을 맞추기 위한 더 작은 변형 RWKV-UNet-S 및 RWKV-UNet-T를 제공한다.

실험 결과

연구 질문

• RQ1 RWKV 기반의 장거리 모델링이 다양한 의료 영상 태스크에서 순수 CNN 또는 순수 트랜스포머 U-Net 대비 분할 성능을 향상시킬 수 있는가?
• RQ2 IR-RWKV 인코더에 RWKV를 통합하는 것이 글로벌-로컬 특징 융합을 개선하되 계산 비용을 지나치게 증가시키지 않는가?
• RQ3 CCM 모듈이 다중 스케일 스킵 연결 융합과 전체 분할 정확도를 개선하는가?
• RQ4 ImageNet-1K에서 인코더를 사전 학습하는 것이 RWKV-UNet의 데이터셋 간 성능에 어떤 영향을 미치는가?
• RQ5 더 작은 RWKV-UNet 변형은 벤치마크 전반에서 전체 모델과 비교하여 정확도와 효율성에서 어떻게 다른가?

주요 결과

• RWKV-UNet은 Synapse, ACDC, BUSI, CVC-ClinicDB, CVC-ColonDB, Kvasir-SEG, ISIC 2017, GLAS 등의 주요 의학 영상 데이터셋에서 최첨단 분할 성능을 달성한다.
• 인코더 변형은 RWKV-UNet-T, RWKV-UNet-S, RWKV-UNet이 높은 정확도를 유지하면서 파라미터 수와 FLOPs를 줄여 강력한 성능을 보여준다.
• RWKV-UNet 인코더를 ImageNet-1K에서 사전 학습하는 것은 분할 지표를 크게 향상시키며 이 아키텍처의 전이 학습 이점을 확인한다.
• CCM 모듈은 다중 스케일 특징 융합과 전체 DSC를 개선하지만 약간의 계산 증가가 발생한다.
• 큰 커널(예: 9x9) 및 Conv1x1->Conv1x1->DW-Conv 구성을 갖춘 디코더 설계가 분할 성능과 효율성을 향상시킨다.