[논문 리뷰] Improving Zero-Shot Detection of Low Prevalence Chest Pathologies using Domain Pre-trained Language Models

Aakash Mishra, Rajat Mittal외 3명|arXiv (Cornell University)|2023. 06. 13.

COVID-19 diagnosis using AI인용 수 5

한 줄 요약

이 논문은 영상-보고서 정렬 데이터셋에서 노출된 바가 적은 희귀 병변에 대해 성능이 열악한 제로샷 흉부 X-ray 병변 탐지 문제를 해결하기 위해, CLIP 기반 텍스트 타워를 도메인 사전학습된 언어 모델(예: CXR-BERT, ClinicalBERT)으로 대체함으로써 성능 향상을 제안한다. 일반 병변에서는 성능 저하가 발생하지만, 도메인 최적화된 텍스트 인코더는 폐 종양 및 경추 골절과 같은 희귀 병변에서 최대 0.15까지 AUC를 향상시켜 자원이 제한된 상황에서의 우수성을 입증한다.

ABSTRACT

Recent advances in zero-shot learning have enabled the use of paired image-text data to replace structured labels, replacing the need for expert annotated datasets. Models such as CLIP-based CheXzero utilize these advancements in the domain of chest X-ray interpretation. We hypothesize that domain pre-trained models such as CXR-BERT, BlueBERT, and ClinicalBERT offer the potential to improve the performance of CLIP-like models with specific domain knowledge by replacing BERT weights at the cost of breaking the original model's alignment. We evaluate the performance of zero-shot classification models with domain-specific pre-training for detecting low-prevalence pathologies. Even though replacing the weights of the original CLIP-BERT degrades model performance on commonly found pathologies, we show that pre-trained text towers perform exceptionally better on low-prevalence diseases. This motivates future ensemble models with a combination of differently trained language models for maximal performance.

연구 동기 및 목표

• 이미지-보고서 정렬 데이터셋에서 노출된 바가 적은 흉부 병변에 대해 제로샷 성능이 열악한 문제를 해결하기 위해.
• 텍스트 타워의 도메인 특화 사전학습이 희귀 병변에 대한 제로샷 탐지 성능 향상에 기여하는지 조사하기 위해.
• CLIP의 일반 목적 텍스트 인코더를 의료 도메인 전용 모델로 교체할 경우, 일반 병변과 희귀 병변 간 성능의 상충 관계를 평가하기 위해.
• 도메인 사전학습된 텍스트 인코더가 자원이 제한된 제로샷 의료 영상 시나리오에서 표준 CLIP 기반 모델을 능가할 수 있는지 판단하기 위해.

제안 방법

• CheXzero의 CLIP 텍스트 타워를 도메인 사전학습된 언어 모델(CXR-BERT, BlueBERT, ClinicalBERT)으로 교체하며, 이는 MIMIC-CXR에서 파생된 파인튜닝을 거친다.
• 대비 학습을 통해 이미지 및 텍스트 임베딩을 대비 손실 함수를 통해 정렬하며, 배치 크기가 32이고 기본 학습률이 5e-6인 5 에포크 동안 훈련한다.
• 모든 텍스트 인코더의 최종 임베딩 차원을 128로 표준화하기 위해 학습 가능한 프로젝션 헤드를 사용하여 시각 타워와의 호환성을 확보한다.
• VinDr-CXR 데이터셋에서 양성 및 부정성 프롬프트를 사용해 제로샷 추론을 수행하고, 병변 확률 예측을 위한 로짓을 계산한다.
• AUC를 주요 평가 지표로 사용해, CLIP 가중치로 초기화된 CheXzero 베이스라인과 성능을 비교하며 12개의 병변에 대해 평가한다.
• 성능 차이의 통계적 유의성을 평가하기 위해 부트스트랩 기반 신뢰구간을 사용한다.

실험 결과

연구 질문

• RQ1 CLIP 텍스트 타워를 도메인 사전학습된 언어 모델로 교체하면, 희귀 흉부 X-ray 병변에 대한 제로샷 탐지 성능 향상이 가능할까?
• RQ2 텍스트 인코더에서 도메인 특화 사전학습을 적용하면, CLIP 기반 베이스라인 대비 고빈도 병변에서 성능 저하가 발생할까?
• RQ3 MIMIC-CXR 훈련 세트에서 100건 이하로 언급된 병변에서 CXR-BERT나 ClinicalBERT와 같은 도메인 사전학습된 텍스트 인코더가 CLIP 기반 베이스라인을 얼마나 뛰어넘을 수 있을까?
• RQ4 정렬 데이터에 희귀 병변이 거의 언급되지 않더라도, 도메인 사전학습된 텍스트 인코더가 희귀 병변에 대해 효과적으로 일반화될 수 있을까?
• RQ5 일반 목적 텍스트 인코더와 도메인 사전학습된 인코더를 조합하는 전략이 일반 병변과 희귀 병변 모두에서 성능을 극대화하는 데 유용한가?

주요 결과

• 폐 종양은 MIMIC-CXR 훈련 세트에서 단 한 번만 언급되었지만, CXR-BERT는 CheXzero 베이스라인 대비 AUC가 0.15 향상되었다.
• ClinicalBERT는 모든 희귀 병변에서 평균 AUC 향상 폭이 0.08에 달했다.
• 경추 골절(69건 언급)의 경우, CXR-BERT는 AUC 0.853을 기록해 베이스라인의 0.664를 크게 앞섰다.
• 폐렴 및 흉막 출혈과 같은 고빈도 병변에서는 CheXzero 베이스라인이 CXR-BERT 및 ClinicalBERT보다 평균 0.05~0.12 AUC 높은 성능을 보였다.
• BlueBERT는 모든 병변에서 성능이 열등했으며, 이는 모든 도메인 사전학습 모델이 성능 향상을 가져오진 않으며, 모델 선택이 매우 중요함을 시사한다.
• 결과는 도메인 사전학습된 텍스트 인코더가 PubMed 및 MIMIC-III와 같은 의료 문헌 코퍼스에서 광범위한 노출를 경험함으로써 희귀 병변에 특히 효과적이라는 점을 시사한다.