[논문 리뷰] Rethinking U-net Skip Connections for Biomedical Image Segmentation

Frauke Wilm, Jonas Ammeling외 4명|arXiv (Cornell University)|2024. 02. 13.

Medical Image Segmentation Techniques인용 수 5

한 줄 요약

논문은 도메인 시프트가 U-net의 계층에 미치는 영향을 정량화하고 상단 스킵 연결(L1-pruned)을 가지치면 synthetic 및 임상 데이터셋에서 인도메인 및 교차도메인 분할 성능이 모두 향상된다는 것을 보여준다.

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에 도메인 시프트가 미치는 영향을 평가한다.
• 헬링거 거리(Hellinger distance)를 사용하여 계층별 고유 도메인 민감도를 정량화한다.
• 도메인 시프트 하에서 계층별 스킵 연결 가지치기(L1–L4)가 분할 성능에 미치는 영향을 평가한다.
• 합성 데이터와 임상 데이터셋 세 가지(조직병리학 두 가지, 심장 MRI 한 가지)에서 발견을 검증한다.
• 모델 개발 중 스킵 연결 사용에 대한 가이드라인을 제공하여 향상된 강건성을 달성한다.

제안 방법

• 밝기, 대비, 포화도 조정을 통해 합성 말라리아 데이터셋에서 도메인 시프트를 모델링한다.
• 특징 맵 전반에 걸친 헬링거 거리로 네트워크 각 계층별 도메인 시프트를 정량화한다.
• 상단에서 하향 가지치기로 L1–L4 U-net을 만들고 ResNet18/34 인코더로 학습한다.
• 합성 및 임상 데이터의 인도메인 및 교차도메인 테스트 세트에서 IoU로 분할을 평가한다.
• 세 가지 임상 데이터셋(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-가지치기 U-net은 인도메인 IoU 최대 10% 증가와 교차도메인 IoU 최대 13% 증가를 달성했다.
• 가장 우수한 전반적 성능은 모든 스킵 연결을 제거하는 것이 아니라 최상단 스킵 연결만 제거했을 때 달성되었다.
• 이런 결과는 도메인 강건성을 높이기 위해 모델 개발 중 스킵 연결을 재고하고 필요 시 가지치기를 고려할 것을 시사한다.