[논문 리뷰] Improve Model Generalization and Robustness to Dataset Bias with Bias-regularized Learning and Domain-guided Augmentation

Yundong Zhang, Hang Wu외 3명|arXiv (Cornell University)|2019. 10. 12.

COVID-19 diagnosis using AI참고 문헌 26인용 수 7

한 줄 요약

이 논문은 생물의학 영상에서 데이터셋 편향에 대한 딥러닝 모델의 강건성을 향상시키기 위해 편향 정규화 학습과 도메인 유도 데이터 증강을 결합한 도메인 일반화 프레임워크를 제안한다. 이 방법은 새로운 손실 함수와 다층 교차 기울기 훈련을 통해 도메인 불변 특징을 학습하도록 모델을 훈련시켜, 미세조정 없이도 새로운 흉부 X-ray 데이터셋에서 최신 기술 수준의 일반화 성능을 달성한다.

ABSTRACT

Deep Learning has thrived on the emergence of biomedical big data. However, medical datasets acquired at different institutions have inherent bias caused by various confounding factors such as operation policies, machine protocols, treatment preference and etc. As the result, models trained on one dataset, regardless of volume, cannot be confidently utilized for the others. In this study, we investigated model robustness to dataset bias using three large-scale Chest X-ray datasets: first, we assessed the dataset bias using vanilla training baseline; second, we proposed a novel multi-source domain generalization model by (a) designing a new bias-regularized loss function; and (b) synthesizing new data for domain augmentation. We showed that our model significantly outperformed the baseline and other approaches on data from unseen domain in terms of accuracy and various bias measures, without retraining or finetuning. Our method is generally applicable to other biomedical data, providing new algorithms for training models robust to bias for big data analysis and applications. Demo training code is publicly available.

연구 동기 및 목표

• 다양한 기관에서 수집된 생물의학 영상 데이터셋 내재적 편향으로 인한 모델 일반화 능력 저하 문제를 해결하기 위해.
• 훈련 중 외부 테스트 도메인에 대한 액세스 없이도 모델 강건성을 향상시키는 도메인 일반화 프레임워크를 개발하기 위해.
• 의료 영상 데이터셋에서 병원별 프로토콜, 영상 장비, 레이블링 편향 등의 혼란 요인의 영향을 완화하기 위해.
• 흉부 X-ray를 초과하여 EHR 및 기타 모odalities를 포함한 다양한 생물의학 데이터에 적용 가능한 일반화 가능한 방법을 만들기 위해.
• 효과적인 정규화와 데이터 증강을 통해 오직 소스 도메인 데이터만을 사용하여도 새로운 도메인에서의 모델 성능을 크게 향상시킬 수 있음을 입증하기 위해.

제안 방법

• 모델이 도메인 불변 표현을 학습하도록 유도하기 위해 도메인 특화 특징 활성화를 제재하는 편향 정규화 손실 함수를 도입한다.
• 다중 소스 도메인에서의 기울기를 활용해 다양한 도메인에 강건한 훈련 샘플을 합성하는 다층 교차 기울기 훈련(MCT)을 통한 도메인 유도 데이터 증강을 시행한다.
• 도메인 이동을 시뮬레이션하고 일반화 능력을 향상시키기 위해 여러 데이터셋(예: NIH, CheXpert, MIMIC)을 사용하는 다중 소스 훈련 설정을 적용한다.
• 주의 메커니즘을 분석하고 모델이 허위 아티팩트가 아닌 질병 관련 영역에 집중하고 있는지 확인하기 위해 기울기 유도 활성화 시각화(Grad-CAM)를 적용한다.
• 모델이 강건하고 일반화 가능한 표현을 학습하도록 보장하기 위해 도메인 특화 특징를 탐지하고 정규화하기 위해 별도의 편향 분류기를 훈련시킨다.
• 두 단계 훈련 프로세스를 사용한다: 첫 번째로, 편향 제거된 시각 세계 분류기를 훈련하고, 두 번째로, 이 분류기의 교차 기울기 신호를 사용하여 주 모델을 훈련시켜 불변성을 촉진한다.

실험 결과

연구 질문

• RQ1 여러 편향이 있는 생물의학 영상 데이터셋에서 훈련된 딥러닝 모델이 미세조정 없이도 새로운 도메인으로 일반화할 수 있는가?
• RQ2 편향 정규화 학습은 흉부 X-ray 영상에서 휴대용 스캔 태그나 삽입된 호흡관과 같은 허위 아티팩트에 대한 의존도를 얼마나 효과적으로 줄이는가?
• RQ3 도메인 유도 데이터 증강은 이질적인 의료 영상 도메인 간의 특징 표현 학습을 얼마나 향상시키는가?
• RQ4 제안된 방법은 표준 도메인 일반화 기준 대비 새로운 데이터셋에서 뛰어난 성능을 달성할 수 있는가?
• RQ5 제안된 방법을 사용할 경우, Grad-CAM으로 시각화된 모델의 주의 메커니즘이 방사선의사가 애너테이션한 병변 영역과 더 잘 일치하는가?

주요 결과

• Office-Caltech 데이터셋에서 제안된 방법은 테스트 정확도 47.8%를 기록하여 베이스라인 LeNet(41.6%)과 DANN(43.8%), CrossGrad(44.3%)와 같은 최신 기술 수준의 방법들을 능가했다.
• Rotated-MNIST 데이터셋에서 방법은 99.6%의 정확도를 달성하여 CCSA 기준(98.0%)과 다른 최신 기술 수준 방법들을 크게 앞섰다.
• Google Fonts 데이터셋에서 방법은 73.1%의 정확도를 기록하여 기준(68.5%)과 다른 방법들을 능가했으며, 비의료 도메인 일반화 설정에서도 효과적임을 입증했다.
• Grad-CAM 시각화에서 모델은 방사선의사가 애너테이션한 병변 영역과 강한 상관관계(Spearman’s rho)를 보였으며, 질병 관련 기능에 더 잘 집중함을 시사했다.
• t-SNE 시각화 결과, 제안된 모델에서 CheXpert, MIMIC 등 서로 다른 데이터셋의 특징 임베딩이 더 혼합되고 분리되지 않은 것으로 확인되어 도메인 특화 편향이 감소함을 확인했다.
• 단순한 CNN은 흉부 X-ray 소스 데이터셋을 거의 완벽한 정확도(약 100%)로 분류할 수 있었으며, 이는 기관 간에 강력하고 식별 가능한 데이터셋 편향이 존재함을 확인한다.