[論文レビュー] RWKV-UNet: Improving UNet with Long-Range Cooperation for Effective Medical Image Segmentation

Juntao Jiang, Jiangning Zhang+6 more|ArXiv.org|Jan 14, 2025

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）を提供。

提案手法

• IR-RWKVブロックとIRブロックで構成されたスタック可能なエンコーダーを持つRWKV-UNetアーキテクチャを導入。
• エンコーダにおける空間ミックスをVision RWKVと深さ方向畳み込みで取り入れ、局所/グローバル機能を統合。
• チャンネルごとのグローバルコンテキストを通じてマルチスケールエンコーダ機能を融合するCross-Channel Mix (CCM)モジュールを追加。
• 特徴を効率的にアップサンプリング・精錬する9x9深さ方向畳み込みを用いたCNNベースのデコーダを設計。
• エンコーダをImageNet-1Kで事前学習（300エポック、AdamW）して特徴抽出と収束を改善。
• 精度と計算負荷のバランスを取る小型バリアントRWKV-UNet-SおよびRWKV-UNet-Tを提供。

実験結果

リサーチクエスチョン

• RQ1 RWKVベースの長距離モデリングは、さまざまな医用画像タスクにおいて純粋なCNNや純粋なトランスフォーマーのU-Netsより分割性能を向上させるか？
• RQ2 IR-RWKVエンコーダへのRWKV統合は、グローバル-ローカル特徴融合を高めつつ計算コストを過度に増やさないか？
• RQ3 CCMモジュールはマルチスケールのスキップ接続融合と全体のDSCを改善するか？
• 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の特定構成を持つデコーダ設計は、分割性能と効率を向上させる。