[논문 리뷰] ProSFDA: Prompt Learning based Source-free Domain Adaptation for Medical Image Segmentation

Shishuai Hu, Zehui Liao외 1명|arXiv (Cornell University)|2022. 11. 21.

Domain Adaptation and Few-Shot Learning인용 수 11

한 줄 요약

이 논문은 소스 데이터에 접근할 수 없는 조건에서도 도메인 불일치를 줄이기 위해 도메인 인식 시각 프롬프트를 학습하여 타겟 도메인 이미지를 개선하는 프롬프트 학습 기반의 소스 프리 도메인 적응(ProSFDA) 방법을 제안한다. 배치 정규화 통계를 정렬하고 스타일 증강 특성 정렬을 통해 특성의 밀도를 강화함으로써 ProSFDA는 다중 도메인 벤치마크에서 비지도 도메인 적응 방법과 비교할 만한 성능을 달성한다.

ABSTRACT

The domain discrepancy existed between medical images acquired in different situations renders a major hurdle in deploying pre-trained medical image segmentation models for clinical use. Since it is less possible to distribute training data with the pre-trained model due to the huge data size and privacy concern, source-free unsupervised domain adaptation (SFDA) has recently been increasingly studied based on either pseudo labels or prior knowledge. However, the image features and probability maps used by pseudo label-based SFDA and the consistent prior assumption and the prior prediction network used by prior-guided SFDA may become less reliable when the domain discrepancy is large. In this paper, we propose a extbf{Pro}mpt learning based extbf{SFDA} ( extbf{ProSFDA}) method for medical image segmentation, which aims to improve the quality of domain adaption by minimizing explicitly the domain discrepancy. Specifically, in the prompt learning stage, we estimate source-domain images via adding a domain-aware prompt to target-domain images, then optimize the prompt via minimizing the statistic alignment loss, and thereby prompt the source model to generate reliable predictions on (altered) target-domain images. In the feature alignment stage, we also align the features of target-domain images and their styles-augmented counterparts to optimize the source model, and hence push the model to extract compact features. We evaluate our ProSFDA on two multi-domain medical image segmentation benchmarks. Our results indicate that the proposed ProSFDA outperforms substantially other SFDA methods and is even comparable to UDA methods. Code will be available at \url{https://github.com/ShishuaiHu/ProSFDA}.

연구 동기 및 목표

• 소스 데이터에 접근할 수 없는 새로운 도메인에 사전 훈련된 모델을 도입할 때 의료 영상 분할에서 도메인 불일치 문제를 해결하기 위해.
• 큰 도메인 이동 상황에서 신뢰할 수 없는 예측을 유발하는 가짜 레이블 기반 및 사전 지침 기반 SFDA 방법의 한계를 극복하기 위해.
• 사전 훈련된 소스 모델과 레이블이 없는 타겟 데이터만을 활용하여 적응 성능을 향상시키는 방법을 개발하기 위해.
• 가능한 도메인 인식 프롬프트와 스타일 증강 특성 정렬을 도입하여 특성의 밀도를 향상시키고 도메인 정렬을 강화하기 위해.
• 자료 및 개인정보 보호 제약 조건에도 불구하고 비지도 도메인 적응(UDA) 방법과 비교할 만한 성능을 달성하기 위해.

제안 방법

• 소스 도메인의 특성을 시뮬레이션하기 위해 타겟 도메인 이미지에 추가되는 학습 가능한 이미지 임베딩으로서 도메인 인식 시각 프롬프트를 도입한다.
• 소스 모델과 프롬프트 증강 타겟 이미지로 미세조정된 모델 간의 배치 정규화(BN) 통계 차이를 기반으로 계산된 통계 정렬 손실을 최소화함으로써 프롬프트를 최적화한다.
• 최적화된 프롬프트를 사용하여 소스 도메인 분포에 더 가까운 통계적 특성을 가진 개선된 타겟 이미지를 생성한다.
• 개선된 타겟 이미지와 그들의 스타일 증강 복제본 간의 특성 정렬을 적용하여 특성의 밀도를 향상시키고 개체 수준의 변동성을 줄인다.
• 특성 정렬 단계에서 분할 손실과 특성 정렬 손실을 함께 최적화하여 모델 일반화와 도메인 불변성을 동시에 향상시킨다.
• 재훈련 없이도 소스 데이터에 접근하지 않고도 지식 정착을 가능하게 하기 위해 고정된 소스 모델을 사용하여 프롬프트 증강 타겟 이미지에 대한 예측을 생성한다.

실험 결과

연구 질문

• RQ1 프롬프트 학습이 의료 영상 분할의 소스 프리 도메인 적응에서 도메인 불일치를 효과적으로 추정하고 완화하는 데에 유용한가?
• RQ2 큰 도메인 이동 상황에서 가령 레이블링 또는 사전 지침 기반 방법과 비교할 때 학습 가능한 도메인 인식 프롬프트의 사용은 어떤가?
• RQ3 스타일 증강 이미지와의 특성 정렬은 특성의 밀도와 분할 성능을 어느 정도 향상시키는가?
• RQ4 통계 정렬 손실에서 BN 레이어와 하이퍼파rameter(예: α)의 최적 설정은 무엇인가? 효과적인 도메인 불일치 추정을 위해.
• RQ5 소스 데이터가 없음에도 불구하고 제안된 방법이 비지도 도메인 적응(UDA) 방법과 비교할 만한 성능을 달성할 수 있는가?

주요 결과

• ProSFDA는 두 개의 다중 도메인 의료 영상 분할 벤치마크에서 기존 SFDA 방법들을 능가하며, 소스 프리 설정에서 최신 기술 수준의 성능을 달성한다.
• ‘BASE2’ 벤치마크에서 OD 분할 작업에서 평균 Dice 스코어 94.71 ± 0.01, OC 분할 작업에서 85.33 ± 0.08을 기록하여 기준 SFDA 방법들을 능가한다.
• 절단 실험 결과, 특성 정렬 단계에서 분할 손실과 특성 정렬 손실을 함께 사용할 경우 가장 우수한 성능을 기록하며, 단일 분할 손실 사용 대비 1.36% 향상된 성능을 보였다.
• t-SNE 시각화 결과, 프롬프트 증강 타겟 이미지에서 유래한 특성들이 소스 도메인 특성과 더 유사해졌음을 확인하여 도메인 불일치 추정의 효과성을 검증했다.
• 통계 정렬 손실에서 최적의 하이퍼파rameter α는 0.01으로 확인되었으며, 손실 계산에 최대 5개의 BN 블록을 사용할 경우 성능이 정점에 도달함을 관찰했다.
• 이 방법은 배치 정규화 레이어를 포함한 네트워크에서만 효과적이며, 도메인 불일치 추정에 BN 통계에 의존하기 때문에 현재 아키텍처 제약이 있음을 시사한다.