[논문 리뷰] GAN-based Data Augmentation for Chest X-ray Classification

Shobhita Sundaram, Neha Hulkund|arXiv (Cornell University)|2021. 07. 07.

COVID-19 diagnosis using AI참고 문헌 22인용 수 5

한 줄 요약

이 논문은 CheXpert 데이터셋을 사용한 흉부 X-ray 분류에서 클래스 불균형 문제를 해결하기 위해 GAN 기반 데이터 증강을 제안한다. 부족한 병변 유형에 대해 합성 이미지를 생성함으로써 모델 성능을 향상시켜, 특히 데이터가 적은 환경에서 과적합을 줄이고 기존 증강 기법이나 증강 없음 대비 더 높은 AUC 향상을 달성한다.

ABSTRACT

A common problem in computer vision -- particularly in medical applications -- is a lack of sufficiently diverse, large sets of training data. These datasets often suffer from severe class imbalance. As a result, networks often overfit and are unable to generalize to novel examples. Generative Adversarial Networks (GANs) offer a novel method of synthetic data augmentation. In this work, we evaluate the use of GAN- based data augmentation to artificially expand the CheXpert dataset of chest radiographs. We compare performance to traditional augmentation and find that GAN-based augmentation leads to higher downstream performance for underrepresented classes. Furthermore, we see that this result is pronounced in low data regimens. This suggests that GAN-based augmentation a promising area of research to improve network performance when data collection is prohibitively expensive.

연구 동기 및 목표

• 작은 다양성 있는 훈련 세트에서 특히 발생하는 의료 영상 데이터셋의 클래스 불균형 문제를 해결한다.
• GAN으로 생성한 합성 데이터가 흉부 X-ray 분류에서 부족한 병변 유형에 대해 딥러닝 모델 성능을 향상시킬 수 있는지 조사한다.
• 다양한 데이터 규모에서 GAN 기반 증강 기법이 기존 증강 기법과 증강 없음 대비 얼마나 효과적인지 비교한다.
• 과적합과 일반화 부족이 가장 심각한 저데이터 환경에서 GAN 증강이 가장 유익한지 평가한다.
• 합성 데이터가 임상 적용에서 더 강건하고 일반화 능력이 뛰어난 모델을 만들어낼 수 있음을 입증한다.

제안 방법

• 부족한 병변 유형(폐 병변, 흉막 기타, 골절)을 갖는 합성 흉부 X-ray 이미지를 생성하기 위해 사전 훈련된 GAN을 미세조정한다.
• 생성된 GAN 이미지를 훈련 세트에 통합하여 클래스 분포를 균형 잡히게 하며, 특히 유병률이 5% 미만인 병변에 중점을 둔다.
• CheXpert 데이터셋의 부분집합(1%, 10%, 50%, 100% 데이터 크기)을 대상으로 DenseNet 모델을 GAN 증강, 표준 증강, 증강 없음 세 가지 조건에서 훈련한다.
• 이전 연구에서 성능 향상을 입증한 바와 같이, 불확실한 레이블(-1)을 1로 매핑하여 레이블 스무딩을 적용한다.
• 비교를 위한 기준선으로 표준 데이터 증강 기법(예: 무작위 수평 반전)을 사용한다.
• AUC 점수를 사용해 모델 성능을 평가하고, 훈련/검증 손실 곡선과 클래스 활성화 맵(CAMs)을 분석하여 해석 가능성과 과적합 여부를 평가한다.

실험 결과

연구 질문

• RQ1 GAN 기반 데이터 증강이 CheXpert 데이터셋에서 표준 증강 및 증강 없음 대비 부족한 병변 유형의 분류 성능을 향상시키는가?
• RQ2 특히 저데이터 설정에서 GAN 증강의 효과가 다양한 데이터 규모에 따라 어떻게 달라지는가?
• RQ3 작은 불균형 데이터 세트에서 GAN 증강이 과적합을 얼마나 줄이는가?
• RQ4 클래스 활성화 맵(CAMs)을 통해 GAN 증강 모델이 학습한 특징이 실제 데이터와 유사하고 일관된가?
• RQ5 표준 데이터 증강 기법이 대규모 데이터 세트에서 성능 저하를 초래하는가? 만약 그렇다면 그 이유는 무엇인가?

주요 결과

• 1% 데이터 규모에서 GAN 기반 증강은 폐 병변에 대해 0.08 AUC 향상, 흉막 기타에 0.06, 골절에 0.07의 AUC 향상을 기록했으며, 증강 없음 대비 유의미한 개선을 보였다.
• 10% 데이터 규모에서 GAN 증강은 모든 세 가지 부족 병변에 대해 AUC 향상을 보였으며, 표준 증강 대비 더 큰 향상을 기록했다.
• 1% 규모에서 과적합이 크게 감소했으며, GAN 증강 모델의 훈련 손실과 검증 손실 간 차이는 0.03으로, 증강 없음 모델의 0.06보다 유의미하게 낮았다.
• 100% 데이터 규모에서 표준 증강은 모든 병변에 대해 증강 없음 대비 성능이 열 劣했으며, 이는 반전에 의한 구조적 불변성 문제 때문일 수 있다.
• 클래스 활성화 맵(CAMs)은 실제 이미지와 GAN으로 생성된 이미지 간 유사한 활성화 패턴을 보였으며, 이는 모델이 합성 데이터에 과적합하지 않고 일관된 국소적 특징을 학습했다는 것을 시사한다.
• 대규모 데이터 환경에서도 GAN 증강은 가장 불균형한 클래스인 흉막 기타에 대해 0.006 AUC 향상을 기록했으며, 이는 극도로 불균형한 분포에 대해 지속적인 유용성을 보여준다.