[논문 리뷰] Rethinking Data Augmentation for Single-source Domain Generalization in Medical Image Segmentation

Zixian Su, Kai Yao외 4명|arXiv (Cornell University)|2022. 11. 27.

Domain Adaptation and Few-Shot Learning인용 수 8

한 줄 요약

이 논문은 복부 CT-MRI 및 심장 bSSFP-LGE 작업에서 ERM보다 4.3%와 2.1% 향상된 SOTA 성능을 달성하면서 일반화 위험을 이론적으로 경계짓는, 단일 소스 도메인 일반화를 위한 새로운 데이터 증강 프레임워크인 Saliency-balancing Location-scale Augmentation (SLAug)을 제안한다. Saliency-guided fusion와 함께 클래스 수준의 위치-스케일 증강을 조합함으로써, 증강된 데이터의 다양성과 정보성 모두 향상된다.

ABSTRACT

Single-source domain generalization (SDG) in medical image segmentation is a challenging yet essential task as domain shifts are quite common among clinical image datasets. Previous attempts most conduct global-only/random augmentation. Their augmented samples are usually insufficient in diversity and informativeness, thus failing to cover the possible target domain distribution. In this paper, we rethink the data augmentation strategy for SDG in medical image segmentation. Motivated by the class-level representation invariance and style mutability of medical images, we hypothesize that unseen target data can be sampled from a linear combination of $C$ (the class number) random variables, where each variable follows a location-scale distribution at the class level. Accordingly, data augmented can be readily made by sampling the random variables through a general form. On the empirical front, we implement such strategy with constrained B$\acute{ m e}$zier transformation on both global and local (i.e. class-level) regions, which can largely increase the augmentation diversity. A Saliency-balancing Fusion mechanism is further proposed to enrich the informativeness by engaging the gradient information, guiding augmentation with proper orientation and magnitude. As an important contribution, we prove theoretically that our proposed augmentation can lead to an upper bound of the generalization risk on the unseen target domain, thus confirming our hypothesis. Combining the two strategies, our Saliency-balancing Location-scale Augmentation (SLAug) exceeds the state-of-the-art works by a large margin in two challenging SDG tasks. Code is available at https://github.com/Kaiseem/SLAug .

연구 동기 및 목표

• 단일 소스 도메인만 이용 가능한 경우 의료 영상 분할에서 발생하는 도메인 이동 문제를 해결한다.
• 전반적인 또는 무작위 증강 전략의 한계를 극복하여 다양성과 정보성이 떨어지는 문제를 해결한다.
• 의료 영상에서의 클래스 수준 표현 불변성과 스타일 변동성을 활용하여 더 견고한 증강 전략을 설계한다.
• 제안된 증강 프레임워크 하에서 일반화 위험의 이론적 경계를 도출하여 그 효과성을 검증한다.
• 타겟 도메인 데이터가 필요 없이 모델의 견고성을 향상시킬 수 있는 플러그 앤 플레이 모듈을 개발한다.

제안 방법

• 각 클래스별로 위치-스케일 분포에서 샘플링된 증강 이미지를 선형 조합으로 조합하는 모델을 제안: $ x_{U}^{c} = \alpha^{c}x_{S}^{c} + \beta^{c} $, 여기서 $ \alpha^{c} $ 및 $ \beta^{c} $ 는 학습 가능한 스케일 및 위치 파라미터이다.
• 전체 이미지의 전반적 강도 이동을 시뮬레이션하기 위해 글로벌 위치-스케일 증강(GlA)을 구현한다.
• 지식 기반 마스크를 사용하여 클래스별로 특화된 변환을 적용함으로써 클래스 수준의 도메인 이동을 시뮬레이션하는 로컬 위치-스케일 증강(LLA)을 도입한다.
• 기울기 기반 샐런시 맵을 사용하여 적절한 방향성과 크기로 융합을 이끄는 Saliency-balancing Fusion(SBF) 메커니즘을 설계한다.
• GLA 및 LLA 모두에서 부드럽고 현실적인 이미지 변환을 보장하기 위해 제약 조건이 부여된 Bézier 변환을 적용한다.
• 이론적 분석을 통해 제안된 증강 전략이 SDG 설정에서 일반화 위험의 상한선을 유도함을 증명한다.

실험 결과

연구 질문

• RQ1 의료 영상에서의 클래스 수준 표현 불변성과 스타일 변동성을 활용하여 도메인 일반화를 위한 더 효과적인 데이터 증강 전략을 설계할 수 있는가?
• RQ2 위치-스케일 분포에서 추출한 클래스 수준의 랜덤 변수의 선형 조합이 전반적 또는 무작위 증강보다 더 나은 분포 이동을 포괄할 수 있는가?
• RQ3 증강 샘플의 샐런시 기반 융합이 고기울기 및 결정 경계 관련 영역에 집중함으로써 모델의 일반화 성능을 향상시킬 수 있는가?
• RQ4 제안된 증강 전략이 단일 소스 도메인 일반화에서 이론적으로 경계된 일반화 위험을 초래하는가?
• RQ5 각 개별 구성 요소(GlA, LLA, SBF)는 전체 성능에 어떻게 기여하며, 상호 보완적인가?

주요 결과

• 복부 CT-MRI 작업에서 SLAug는 ERM 대비 4.3% 향상된 성능을 기록하였으며, 복부의 경우 88.63의 Dice 점수, 심장의 경우 86.69의 Dice 점수를 기록하였다.
• 심장 bSSFP-LGE 작업에서 SLAug는 ERM 대비 2.1% 향상된 평균 87.66%의 Dice 점수를 달성하였다.
• 제거 실험 결과, GLA, LLA, SBF 모두 필수적임을 확인: 한 요소만 제거해도 성능이 유의미하게 저하됨.
• Variant 4(GlA + SBF)는 Variant 1(GlA only)보다 성능이 뛰어나, GlA 증강 이미지에서 유도된 샐런시 맵이 효과적인 분포 확장을 이끌어냄을 시사함.
• 무작위 융합은 SBF 대비 2.9% 낮은 성능을 보였으며, 샐런시 기반 융합이 더 정보적이고 설명 가능함을 확인함.
• SBF는 ERM 또는 무작위 융합보다 원본 도메인 지식을 더 잘 유지하며, 원본 테스트 성능을 각각 0.8 및 4.0 Dice 포인트 향상시켰다.