[論文レビュー] GCtx-UNet: Efficient Network for Medical Image Segmentation

Khaled Alrfou, Tian Zhao|arXiv (Cornell University)|Jun 9, 2024

Brain Tumor Detection and Classification被引用数 6

ひとこと要約

GCtx-UNetは、GC-ViTのグローバル/ローカル注意機構とCNNベースのダウンサンプリング/アップサンプリングを組み合わせた軽量なUNet様アーキテクチャで、複数の医用画像データセットにおいてモデルの複雑さを抑えつつ競合的なセマンティックセグメンテーション精度を達成します。

ABSTRACT

Medical image segmentation is crucial for disease diagnosis and monitoring. Though effective, the current segmentation networks such as UNet struggle with capturing long-range features. More accurate models such as TransUNet, Swin-UNet, and CS-UNet have higher computation complexity. To address this problem, we propose GCtx-UNet, a lightweight segmentation architecture that can capture global and local image features with accuracy better or comparable to the state-of-the-art approaches. GCtx-UNet uses vision transformer that leverages global context self-attention modules joined with local self-attention to model long and short range spatial dependencies. GCtx-UNet is evaluated on the Synapse multi-organ abdominal CT dataset, the ACDC cardiac MRI dataset, and several polyp segmentation datasets. In terms of Dice Similarity Coefficient (DSC) and Hausdorff Distance (HD) metrics, GCtx-UNet outperformed CNN-based and Transformer-based approaches, with notable gains in the segmentation of complex and small anatomical structures. Moreover, GCtx-UNet is much more efficient than the state-of-the-art approaches with smaller model size, lower computation workload, and faster training and inference speed, making it a practical choice for clinical applications.

研究の動機と目的

• 高い計算コストを伴わない正確な医用画像セグメンテーションのニーズに対応する。
• UNet様アーキテクチャでグローバル文脈モデリングと局所的な注意を統合する。
• MedNet事前学習を活用してImageNet事前学習よりドメイン内性能を向上させる。
• 複数データセット（Synapse、ACDC、Polyp）で一般化と効率を示す。

提案手法

• 局所自己注意とグローバル文脈クエリを組み合わせて長距離・短距離の依存性をモデル化するGC-ViTブロックを使用する。
• 誘導バイアスとチャネル間モデリングを注入するダウンサンプリング（Fused-MBConv）モジュールを組み込む。
• U字型アーキテクチャでスキップ接続を備えたGC-ViTベースのエンコーダー–ボトルネック–デコーダーを採用する。
• MedNet（医用画像）でGC-ViTを事前訓練し、ImageNet事前訓練と比較する。
• エンコーダ内でオーバーラップするパッチを作成するパッチ化レイヤーを用い、埋め込み投影を行う。

実験結果

リサーチクエスチョン

• RQ1 GC-ViTベースのブロックは、最先端のCNN/Transformerベースモデルと比較してパラメータ数を抑えつつ競合的なセグメンテーション性能を達成できるか。
• RQ2 ドメイン特化の医用データ（MedNet）での事前訓練は自然画像（ImageNet）事前訓練よりセグメンテーション精度を向上させるか。
• RQ3 GCtx-UNetはCT、MRI、ポリプ画像などの多様な医用画像タスクでDSCとHDの観点からどう性能を示すか。
• RQ4 アップサンプリング/ダウンサンプリング設計とハイパーパラメータがセグメンテーション性能に与える影響は何か。
• RQ5 提案アーキテクチャは、モデルサイズ、FLOPs、学習時間、推論速度の観点で同業他モデルと比較して効率的か。

主な発見

• GCtx-UNetは、最先端の結果に競合する性能を達成しつつ、調査対象の手法の中で最小のモデルサイズ（12.34Mパラメータ）および最も低いFLOPsを実現。
• MedNetで事前訓練したGCtx-UNetは、一般にImageNetで事前訓練した変種を上回る。
• Synapseでは、GCtx-UNet2（MedNet事前訓練）はDSC 82.39%、HD 15.94 mmを達成し、計算負荷が多くの同業より低いにもかかわらず報告上の優れた結果のひとつ。
• ACDCでは、GCtx-UNet2はDSCが91.23（RV）、89.88（心筋）、87.25（LV）を達成し、いくつかのTransformerベースおよびハイブリッドモデルを上回っている。
• ポリプデータセットでは、GCtx-UNet2は見落としデータセット（CVC-ColonDB、ETIS-LaribDB、CVC-300）への一般化が強く、しばしばトップまたはほぼトップのDSCスコアを示す。
• アブレーション研究により、最適な損失の組み合わせ（dice 0.3、クロスエントロピー 0.7）と学習率（0.0001）を特定し、SEブロックを含む転置畳み込みアップサンプリングが最良の性能を示すことを確認。