[논문 리뷰] Learning to Unlearn: Building Immunity to Dataset Bias in Medical Imaging Studies

Ahmed Ashraf, Shehroz S. Khan외 3명|arXiv (Cornell University)|2018. 12. 03.

Domain Adaptation and Few-Shot Learning참고 문헌 8인용 수 19

한 줄 요약

이 논문은 의료 영상에서 데이터셋에 특화된 편향을 제거하기 위한 프레임워크를 제안한다. 데이터를 연구자별로 구분할 수 없는 잠재 공간으로 매핑함으로써 임상 변수(예: MCI 대 비질환자)의 분류 성능를 유지하면서도, 데이터셋에 특화된 편향을 제거한다. 이 방법은 단순한 한 데이터셋을 제외한 교차 검증 대비 평균 AUC가 47점 향상되었고, 내부 데이터셋 기반 기준 대비 성능 저하를 20% 이하로 줄였다.

ABSTRACT

Medical imaging machine learning algorithms are usually evaluated on a single dataset. Although training and testing are performed on different subsets of the dataset, models built on one study show limited capability to generalize to other studies. While database bias has been recognized as a serious problem in the computer vision community, it has remained largely unnoticed in medical imaging research. Transfer learning thus remains confined to the re-use of feature representations requiring re-training on the new dataset. As a result, machine learning models do not generalize even when trained on imaging datasets that were captured to study the same variable of interest. The ability to transfer knowledge gleaned from one study to another, without the need for re-training, if possible, would provide reassurance that the models are learning knowledge fundamental to the problem under study instead of latching onto the idiosyncracies of a dataset. In this paper, we situate the problem of dataset bias in the context of medical imaging studies. We show empirical evidence that such a problem exists in medical datasets. We then present a framework to unlearn study membership as a means to handle the problem of database bias. Our main idea is to take the data from the original feature space to an intermediate space where the data points are indistinguishable in terms of which study they come from, while maintaining the recognition capability with respect to the variable of interest. This will promote models which learn the more general properties of the etiology under study instead of aligning to dataset-specific peculiarities. Essentially, our proposed model learns to unlearn the dataset bias.

연구 동기 및 목표

• 의료 영상에서 모델이 연구자별로 고유한 특징에 과적합하는 대신 일반화 가능한 특징을 학습하도록 하는 데이터셋 편향 문제를 해결하기 위해.
• 목표 데이터셋에 대한 재학습 없이도 교차 데이터셋 일반화를 가능하게 하여, 기본적인 질병 관련 패턴을 학습하는 모델을 만드는 데 목적이 있다.
• 기존에 인종이나 성별과 같은 민감한 속성에 대해 사용되던 공정한 표현 학습 기법을 의료 영상에서 데이터셋 소속 정보를 흐리게 하기 위해 응용하는 것.
• 네 개의 뇌 MRI 데이터셋을 대상으로 MCI 대 비질환자 분류를 위한 한 데이터셋을 제외한 교차 검증 전략을 사용하여 방법을 평가하는 것.
• 데이터셋 소속 정보를 제거하는 것이 더 견고하고 일반화 능력이 뛰어난 모델을 얻는 데 기여하며, 다양한 데이터셋 간 성능 저하를 최소화할 수 있음을 입증하는 것.

제안 방법

• 딥 오토에인코더를 사용하여 입력 특징을 잠재 공간으로 매핑하고, 프로토타입 기반 정규화 성분을 도입하여 다양한 데이터셋 간 데이터 포인트의 구분이 불가능하도록 강제한다.
• 분류기의 소스 데이터셋 예측 능력을 최소화하면서도 임상 목표(예: MCI 대 비질환자)의 분류 정확도를 최대화함으로써 도메인 불변 표현을 학습한다.
• 재구성 손실, 분류 손실(소프트맥스), 도메인 식별 손실을 조합한 목적 함수를 사용하여 모델이 데이터셋 소속 정보를 제거하도록 유도한다.
• 멀티헤드 아키텍처를 사용한다: 하나의 헤드는 분류를 위해, 다른 하나는 도메인 식별을 위해 사용하며, 후자는 잠재 표현에 의해 속임을 당하도록 적대적으로 훈련된다.
• 교차 검증을 통해 튜닝된 초기화 파라미터들(예: 무게 조정 계수 αJ, αE, αL 및 L2 정규화)을 사용하여 ROC-AUC를 최적화한다.
• 모델은 한 데이터셋을 제외한 교차 검증 프로토콜을 사용하여 평가되며, 훈련 과정에서 타겟 데이터셋이 사용되지 않도록 보장하고, AUC와 내부 데이터셋 기반 기준 대비 상대적 성능 저하율을 성능 측정 지표로 사용한다.

실험 결과

연구 질문

• RQ1 입력 영상이 어떤 데이터셋에 속해 있는지 알 수 없도록 하면서도 임상 결과(예: MCI 대 비질환자)를 분류할 수 있는 딥 러닝 모델을 훈련시킬 수 있는가?
• RQ2 데이터셋 소속 정보를 제거하는 것이 기존의 한 데이터셋을 제외한 교차 검증 및 편향이 없는 SVM 기준 대비 더 나은 교차 데이터셋 일반화 능력을 제공하는가?
• RQ3 제안된 방법을 사용할 경우, 교차 데이터셋 모델의 성능 저하율이 내부 데이터셋 기반 기준 대비 어느 정도 수준이 되는가?
• RQ4 ADNI1, ADNI2, AIBL 등의 다양한 데이터셋은 각각 어떤 고유한 특징을 지니며, 이는 일반화 능력에 어떤 영향을 미치는가?
• RQ5 공정한 표현 학습 기법에서 유래한 프레임워크를 의료 영상에서 데이터셋 편향을 줄이기 위해 효과적으로 재사용할 수 있는가?

주요 결과

• 제안된 방법은 모든 한 데이터셋을 제외한 실험에서 평균 AUC 0.72를 달성했으며, 단순한 한 데이터셋을 제외한 교차 검증 기준 대비 유의미하게 높은 성능을 보였다. 이는 모든 데이터셋에서 우연의 성능 이하(AUC < 0.5)를 기록한 기준 대비 뚜렷한 향상이다.
• 단순한 한 데이터셋을 제외한 교차 검증 기준 대비 평균 47점의 AUC 향상을 기록했고, 편향이 없는 SVM 기준 대비 18점의 향상을 기록했다.
• 모든 데이터셋에서 내부 데이터셋 기반 기준 대비 성능 저하율이 20% 이하로 유지되어 강력한 일반화 능력을 보였다.
• ADNI1과 ADNI2는 '연구자 식별' 작업에서 더 쉽게 식별되었으며, 이는 더 높은 데이터셋 특화된 고유 특징을 지닌 것으로 나타났다. 반면 AIBL은 다른 데이터셋과 유사하여 더 잘 일반화되었다.
• 모든 데이터셋에서 두 기준 모델 대비 일관되게 뛰어난 성능를 보이며, 데이터셋 고유의 변동성에 대해 뛰어난 내성적 저항력을 입증했다.
• 결과는 데이터셋 소속 정보를 제거함으로써 모델이 더 일반화 가능한, 병태생리학적으로 관련된 특징을 학습하고, 데이터셋 고유의 잡음이나 특징을 학습하지 않음을 확인했다.