[論文レビュー] Rethinking U-net Skip Connections for Biomedical Image Segmentation

Frauke Wilm, Jonas Ammeling+4 more|arXiv (Cornell University)|Feb 13, 2024

Medical Image Segmentation Techniques被引用数 5

ひとこと要約

本論文はドメインシフトがU-netの層に与える影響を定量化し、最上部のスキップ接続を剪定した（L1-pruned）場合が、合成データと臨床データセットの双方で同域内および異域間のセグメンテーションを改善することを示している。

ABSTRACT

The U-net architecture has significantly impacted deep learning-based segmentation of medical images. Through the integration of long-range skip connections, it facilitated the preservation of high-resolution features. Out-of-distribution data can, however, substantially impede the performance of neural networks. Previous works showed that the trained network layers differ in their susceptibility to this domain shift, e.g., shallow layers are more affected than deeper layers. In this work, we investigate the implications of this observation of layer sensitivity to domain shifts of U-net-style segmentation networks. By copying features of shallow layers to corresponding decoder blocks, these bear the risk of re-introducing domain-specific information. We used a synthetic dataset to model different levels of data distribution shifts and evaluated the impact on downstream segmentation performance. We quantified the inherent domain susceptibility of each network layer, using the Hellinger distance. These experiments confirmed the higher domain susceptibility of earlier network layers. When gradually removing skip connections, a decrease in domain susceptibility of deeper layers could be observed. For downstream segmentation performance, the original U-net outperformed the variant without any skip connections. The best performance, however, was achieved when removing the uppermost skip connection - not only in the presence of domain shifts but also for in-domain test data. We validated our results on three clinical datasets - two histopathology datasets and one magnetic resonance dataset - with performance increases of up to 10% in-domain and 13% cross-domain when removing the uppermost skip connection.

研究の動機と目的

• 生物医療分割におけるスキップ接続を有するエンコーダ–デコーダCNNがドメインシフトの影響を受ける程度を評価する。
• ヘリンガー距離を用いて層ごとの固有のドメイン感受性を定量化する。
• ドメインシフト下のセグメンテーション性能に対する層ごとのスキップ接続剪定（L1–L4）の評価。
• 合成データと3つの臨床データセット（組織病理2つ、心臓MRI1つ）で知見を検証する。
• 頑健性向上のためのモデル開発時のスキップ接続の使用に関する指針を提供する。

提案手法

• 明るさ・コントラスト・彩度を変更して、合成マラリアデータセット上でドメインシフトをモデル化する。
• 特徴マップ全体のヘリンガー距離を用いて、ネットワーク各層ごとのドメインシフトを定量化する。
• 上位から順にスキップ接続を剪定してL1–L4 U-netを作成し、ResNet18/34のエンコーダで学習する。
• 合成データと臨床データセットからの同域内および異域間のテストセットでIoUによるセグメンテーションを評価する。
• 臨床データセット3つ（MS-MF、MS-CCT、MV-MR）において剪定バリアントの性能をベースラインU-netと比較する。
• 事前学習済みImageNetエンコーダを使用し、50エポックの学習でクロスエントロピー＋Dice損失で最適化する。

実験結果

リサーチクエスチョン

• RQ1 スキップ接続を持つU-netアーキテクチャにおけるエンコーダ層全体で、ドメインシフトはどのように現れるか？
• RQ2 スキップ接続を削除することは、同域内の性能を犠牲にすることなく、ドメインシフトに対する頑健性を高めるか？
• RQ3 どのスキップ接続剪定レベル（L1–L4）が同域内および異域間のセグメンテーション精度の最適なトレードオフをもたらすか？
• RQ4 合成データと多様な臨床画像モダリティ（組織病理とMRI）において、所見は一貫しているか？

主な発見

• ヘリンガー距離で測定した結果、初期のネットワーク層ほど深い層よりもドメインシフトに対する感受性が高いことを示している。
• 最上位のスキップ接続を削除する（L1-pruned）は、データセットを跨ぐ同域内・異域間のテストで一貫してIoUを改善する。
• より深いスキップ接続（L2–L4）の剪定は全体の性能を一般に低下させるが、場合によっては非常に強い増強に対する感度を低減することがある。
• MS-MF、MS-CCT、MV-MRデータでは、L1-pruned U-nets が同域内IoUで最大10%、異域間IoUで最大13%の向上を達成した。
• 全てのスキップ接続を削除するのではなく、最上位のスキップ接続のみを削除することで最良の全体性能を達成した。
• この所見は、スキップ接続を再検討し、ドメイン堅牢性を高めるためにモデル開発時に剪定する可能性を示唆している。