[논문 리뷰] How Can We Tame the Long-Tail of Chest X-ray Datasets?

Arsh Verma|arXiv (Cornell University)|2023. 09. 08.

COVID-19 diagnosis using AIMedicine인용 수 3

한 줄 요약

이 논문은 목표 데이터셋에서 훈련하기 전에 흉부 X선(CXR) 데이터로 시각 모델을 프리트레인하는 것이 희귀하고 긴 尾(롱테일) 레이블에서 성능을 크게 향상시킨다—복잡한 손실 함수나 하이퍼파라미터 튜닝 없이도 가능하다. 이미지넷을 거쳐 CXR 전용 데이터셋으로 이중 단계 프리트레인 전략을 적용함으로써, 초희귀 레이블에서 최대 3%p 향상되고, 앙상블 모델을 통해 mAP가 6%p 절대적으로 향상됨을 입증하였다. 이는 의료 영상에서 장기적 분포 일반화에 있어 초기화 품질이 핵심 요소임을 시사한다.

ABSTRACT

Chest X-rays (CXRs) are a medical imaging modality that is used to infer a large number of abnormalities. While it is hard to define an exhaustive list of these abnormalities, which may co-occur on a chest X-ray, few of them are quite commonly observed and are abundantly represented in CXR datasets used to train deep learning models for automated inference. However, it is challenging for current models to learn independent discriminatory features for labels that are rare but may be of high significance. Prior works focus on the combination of multi-label and long tail problems by introducing novel loss functions or some mechanism of re-sampling or re-weighting the data. Instead, we propose that it is possible to achieve significant performance gains merely by choosing an initialization for a model that is closer to the domain of the target dataset. This method can complement the techniques proposed in existing literature, and can easily be scaled to new labels. Finally, we also examine the veracity of synthetically generated data to augment the tail labels and analyse its contribution to improving model performance.

연구 동기 및 목표

• 흉부 X선(CXR) 분류 작업에서 희귀하고 장기적 분포 레이블에서의 성능이 열악한 문제를 해결하기 위해.
• 손실 함수나 데이터 샘플링 전략을 수정하지 않고도 도메인 특화 프리트레인을 통해 모델 초기화를 개선하면 희귀 레이블로의 일반화가 향상되는지 조사하기 위해.
• 모델 앙상블과 합성 데이터 증강이 꼬리 레이블 성능 향상에 얼마나 효과적인지 평가하기 위해.
• 최소한의 재훈련으로 새로운 희귀 레이블에 대한 모델 성능을 향상시킬 수 있는 확장 가능한, 즉시 적용 가능한 방법을 제공하기 위해.

제안 방법

• 이 방법은 이중 단계 프리트레인 전략을 사용한다: 먼저 이미지넷 가중치로 초기화한 후, 세 개의 오픈소스 CXR 데이터셋(NIH CXR, PadChest, CheXpert)으로 파인튜닝하여 도메인 적응된 초기화를 생성한다.
• 모델은 MIMIC CXR 데이터셋에서 엔드 투 엔드로 파인튜닝되며, 이 데이터셋은 26개 레이블의 장기적 분포를 포함하고 있고, 프리트레인 데이터셋에 포함되지 않은 5개의 고유한 꼬리 레이블이 존재한다.
• 모델 앙상블은 서로 다른 백본(DenseNet161 및 ResNeXt101)을 사용한 모델의 예측을 앙상블하여 자원이 적은 레이블에서의 강건성을 향상시킨다.
• 합성 데이터 증강은 RoentGen을 사용해 생성한 소규모 합성 데이터셋(5,000개 샘플)과 MIMIC CXR 데이터셋을 함께 파인튜닝함으로써 꼬리 레이블 성능 향상을 탐색한다.
• 모든 모델는 시그모이드 활성화와 이진 크로스엔트로피 손실을 사용하며, 448×448 해상도에서 20 에포크 동안 훈련되어 미세한解부학적 세부 정보를 유지한다.
• 훈련 파이프라인은 환자 수준에서 적용되며, 각 환자에서 이용 가능한 모든 영상 뷰를 사용하고, 데이터 분할은 환자 수준에서 생성되어 데이터 泄露를 방지한다.

실험 결과

연구 질문

• RQ1 다양한 CXR 데이터셋으로 프리트레인하면, 손실 함수나 데이터 샘플링 전략을 수정하지 않고도 희귀한 장기적 분포 레이블에서 성능 향상이 가능한가?
• RQ2 이중 단계 프리트레인 전략(이미지넷 → CXR 전용 데이터)은 표준 이미지넷 초기화보다 꼬리 레이블에서 더 나은 일반화를 이끌어내는가?
• RQ3 모델 앙상블은 특히 목표 데이터셋에만 존재하는 낮은 빈도 레이블에서 성능 향상에 얼마나 기여하는가?
• RQ4 합성 데이터 증강은 꼬리 레이블의 데이터 부족 문제를 완화하는 데 도움이 되며, 성능 향상이 측정 가능한가?

주요 결과

• 연쇄적 프리트레인 전략은 DenseNet161 기준으로 고유 꼬리 레이블(Tail-U)에서 3%p 향상되었고, ResNeXt101의 경우 0.5%p에 그쳐 초희귀 클래스에서 뚜렷한 성과 향상을 보였다.
• 모델 앙상블은 Tail-U 레이블에서 평균 평균 정확도(mAP)를 6%p 절대적으로 향상시켰다(0.1806 → 0.2430)로, 모든 개별 모델보다 뛰어난 성능을 기록했다.
• 앙성 모델은 검증 세트에서 26개 레이블 중 19개에서 최고 성능을 기록했으며, 헤드, 미디엄, 테일 카테고리 전반에서 모두 선도했다.
• 합성 데이터 증강은 크기가 제한된(5,000개 샘플) 상태에서도 꼬리 레이블에서 측정 가능한 성능 향상을 보였으며, 특히 꼬리 카테고리에서 가장 뚜렷한 성과를 보여 향후 확장 가능성에 기대를 품게 했다.
• 이 방법은 매우 확장 가능하다: 새로운 레이블을 추가하려면 몇 번의 파인튜닝 에포크만으로도 가능하므로, 임상 환경에서의 실용적 구현에 매우 유용하다.
• 결과는 모델 초기화가 장기적 분포 일반화에 있어 매우 중요하고, 아직 충분히 인식되지 않은 요소임을 시사한다. 특히 희귀 병변이 임상적으로 중요한 의료 영상 분야에서 그러하다.