[논문 리뷰] Generalizable Medical Image Segmentation via Random Amplitude Mixup and Domain-Specific Image Restoration

Ziqi Zhou, Lei Qi외 1명|arXiv (Cornell University)|2022. 08. 08.

Medical Image Segmentation Techniques인용 수 7

한 줄 요약

이 논문은 다양한 영상 도메인 간에서 모델의 강건성을 향상시키기 위해 무작위 진폭 혼합(RAM)과 도메인 특화 이미지 복원(DSIR)을 사용하는 새로운 도메인 일반화 방법을 제안한다. 푸리에 기반 진폭 혼합을 통해 스타일화된 이미지를 합성하고, 자기지도 기반 복원 헤드를 훈련시킴으로써 특징 일반화를 향상시키며, 망막 및 전립선 영상 분할 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

For medical image analysis, segmentation models trained on one or several domains lack generalization ability to unseen domains due to discrepancies between different data acquisition policies. We argue that the degeneration in segmentation performance is mainly attributed to overfitting to source domains and domain shift. To this end, we present a novel generalizable medical image segmentation method. To be specific, we design our approach as a multi-task paradigm by combining the segmentation model with a self-supervision domain-specific image restoration (DSIR) module for model regularization. We also design a random amplitude mixup (RAM) module, which incorporates low-level frequency information of different domain images to synthesize new images. To guide our model be resistant to domain shift, we introduce a semantic consistency loss. We demonstrate the performance of our method on two public generalizable segmentation benchmarks in medical images, which validates our method could achieve the state-of-the-art performance.

연구 동기 및 목표

• 제한된 소스 도메인에서 훈련된 의료 영상 분할 모델의 일반화 능력 부족 문제를 해결한다.
• 더 강력한 모델 정규화를 도입하여 소스 도메인에 대한 과적합과 도메인 이동 문제를 완화한다.
• 대상 도메인 데이터에 접근할 필요 없이도 미리 보지 않은 대상 도메인에 대한 강건성을 향상시킨다.
• 이미지 복원과 일致성 손실을 통해 도메인 불변 특징 학습을 향상시키는 자기지도 기반 훈련 프레임워크를 개발한다.
• 망막 및 전립선 영상 분할을 위한 공개 도메인 일반화 벤치마크에서 최신 기술 수준의 성능을 입증한다.

제안 방법

• 두 개의 다른 도메인 영상에서 진폭 맵을 추출하기 위해 푸리에 변환을 적용하는 무작위 진폭 혼합(RAM) 모듈을 도입하여, 혼합된 도메인 특성을 가진 새로운 스타일화된 이미지를 생성한다.
• RAM에서 생성된 이미지를 데이터 증강 기법으로 사용하여 분할 모델을 훈련시킴으로써 훈련 분포의 다양성을 증가시킨다.
• RAM 증강 이미지에서 원본 소스 도메인 이미지를 재구성하는 자기지도 기반 도메인 특화 이미지 복원(DFIR) 모듈을 구현하여 특징의 일致성과 강건성을 강제한다.
• 표준 분할 손실, 실제 이미지와 증강 이미지 예측 간의 의미 일치 손실, 그리고 DSIR 모듈에서 유도된 이미지 복원 손실을 조합한 다중 작업 손실로 분할 모델을 훈련시킨다.
• 실제 소스 이미지와 그 RAM 증강 복제본 간의 예측 불일치를 최소화하기 위해 의미 일치 손실을 적용하여 도메인 불변 표현 학습을 촉진한다.
• 분할, 일치성, 복원 손실을 조합하여 전체 네트워크를 엔드 투 엔드로 최적화함으로써, 예측되지 않은 도메인으로의 일반화 능력을 향상시킨다.

실험 결과

연구 질문

• RQ1 무작위 진폭 혼합이 대상 도메인 데이터가 없더라도 다양한 도메인 특성을 혼합한 이미지를 효과적으로 합성하여 분할 일반화 능력을 향상시킬 수 있는가?
• RQ2 자기지도 기반 도메인 특화 이미지 복원 모듈을 통합함으로써 특징의 강건성과 모델의 일반화 능력이 향상되는가?
• RQ3 의미 일치 손실이 도메인 이동된 입력에 대해 예측 변동성을 줄이는 데 얼마나 효과적인가?
• RQ4 RAM 및 DSIR 구성 요소가 의료 영상 분할에서 예측되지 않은 대상 도메인에서 성능 향상에 개별적으로 및 공동으로 얼마나 기여하는가?
• RQ5 제안된 방법이 훈련 목표 함수에서 다양한 일치성 손실 함수 유형에 대해 강건한가?

주요 결과

• 제안된 RAM-DSIR 방법은 망막 및 전립선 영상 분할을 위한 두 개의 공개 도메인 일반화 벤치마크에서 최신 기술 수준의 성능을 달성한다.
• 제거 실험 결과, RAM 모듈을 추가함으로써 기준 모델 대비 망막 데이터셋에서 2.12% 향상되고 전립선 데이터셋에서 3.23% 향상됨을 확인하였다.
• DSIR 모듈은 망막 데이터셋에서 추가로 1.58% 향상되고 전립선 데이터셋에서 1.03% 향상됨을 기록하였다.
• 의미 일치 손실은 일반화 능력을 크게 향상시키며, 이를 제거할 경우 두 데이터셋에서 일관되게 성능 저하가 발생함을 확인하였다.
• 다양한 일치성 손실 유형(MSE, JS-Div, KL-Div)에 대해 방법이 강건하며, 비교적 유사한 성능를 기록하였고, KL-Div가 망막 데이터셋에서 최고의 전체 Dice 점수 88.94%를 기록하였다.
• 모든 구성 요소를 포함한 전체 모델이 최고의 성능를 기록하여 RAM, DSIR, 의미 일치 손실이 도메인 일반화 능력을 향상시키는 데 상호 보완적인 역할을 한다는 것을 확인하였다.