Introduction
Abstract
Deep learning을 이용한 의료 영상 해석은 종종 전문가의 성능을 능가하지만, 대부분의 모델은 labeld(대규모의 전문가 주석) 데이터에 의존한다. 본 연구에서는 self-supervised learning 기법을 통해 별도의 수작업 주석 없이도 chest X-ray에 대한 zero-shot multi-label classification을 달성하는 CheXzero 방법을 제안하며, 방사선 전문의 수준의 진단 정확도를 입증한다.
Background
- 의료 영상 해석의 효율성
- 기존의 supervised 모델은 대량의 전문가 주석 데이터를 필요로 하며, 이는 비용과 시간이 많이 소요된다.
- Label 효율성 향상을 위한 시도
- pre-training이나 self-supervision 방식이 제안되었으나, 대부분의 방법은 결국 fine-tuning 단계에서 수동 레이블이 요구된다.
- 결과적으로, 모델은 학습 시 제공된 질병(Pathology)만 예측 가능하며, 새로운 병변에는 적용 불가하다.
Method: CheXzero
- 기본 아이디어
- zero-shot 설정에서, 학습 시 주석이 없는 이미지–텍스트 쌍을 이용해 representation을 학습하고, 테스트 시 미리 정의된 병변 이름만 제공하면 해당 병변 유무를 예측한다.
- Self-supervised Learning
- contrastive learning 기법으로 이미지와 임상 report 텍스트 간 연관성을 학습한다.
- 라벨이 아닌 자연어(Clinical report) 자체를 “natural-language supervision”으로 활용하여, fine-tuning 없이도 이미지 해석 능력을 갖추도록 한다.
- CheXzero 절차
- 방사선 보고서에 포함된 텍스트를 positive/negative 예제로 활용하여, 이미지와 텍스트 임베딩 공간을 맞춘다.
- 테스트 단계에서는 “Pneumothorax”, “Edema” 등 병변 이름을 텍스트로 입력하면, 각 병변에 대한 예측 점수를 즉시 계산한다.
Experiments
- 외부 검증 데이터셋(두 국가의 독립적 컬렉션)에서, 8개 병변 중 3개에서는 fully supervised 모델을 능가하는 성능을 달성하였다.
- 학습에 전혀 사용되지 않은 병변에 대해서도 높은 일반화 성능을 보였다.
- 기존의 label-efficient 접근법 대비 CheXzero가 우수한 정확도를 보이며, 별도의 수동 레이블 없이도 다중 병변 분류가 가능함을 확인하였다.
Contributions
- 대규모 수동 주석 작업의 필요성 감소: 방사선 보고서만 있으면, 별도 주석 없이도 고성능 모델을 구축할 수 있다.
- Zero-shot 적용 범위 확대: 학습 시 주어지지 않은 새로운 병변도 즉시 예측 가능하다.
- 의료 AI 워크플로우 효율화: 대규모 레이블링 비용과 시간을 절감하며, 다양한 의료기관 데이터에 적용할 수 있는 범용적 프레임워크를 제공한다.
Methods
1. Datasets 및 Training Data
- MIMIC-CXR 데이터셋: 공개 chest radiograph 이미지와 radiology report로 구성.
- 377,110개 이미지 (227,835개 radiographic study)
- 한 연구(study)에 이미지가 여러 장일 경우 AP/AP_VIEW 이미지를 선택
- 각 보고서는 examination, indication, impression, findings, technique, comparison 섹션으로 구성
- CheXpert 데이터셋: Stanford Hospital에서 수집된 224,316개 chest X-ray (65,240명 환자)
- 14개 condition(atelectasis, cardiomegaly, … support devices)에 대해 5명의 전문의 합의로 레이블링
- 각 이미지에 연관된 free-text report 포함
- Validation set: 5개 CheXpert competition 병변에 대한 threshold 튜닝용
- PadChest 데이터셋: 160,868개 chest X-ray
- 174개 radiographic finding, 19개 differential diagnosis 라벨
- 27%는 전문의 라벨, 나머지는 RNN 기반 자동 라벨러 사용
- 39,053개 전문의 검증용 예제만 평가에 활용
2. Pre-processing
- Image 처리
- 모든 X-ray 이미지를 224 × 224 크기로 리사이즈 후 zero padding
- 데이터셋 전체의 mean, standard deviation으로 정규화
- Text 처리
- Radiology report를 Byte Pair Encoding(vocabulary size 49,408)로 토크나이즈
- 최대 토큰 길이를 초과할 경우 “context length – 2” 토큰만 사용하고, 앞뒤에 [SOS], [EOS] 토큰 삽입
3. Architecture
- Image Encoder: Vision Transformer ViT-B/32 (pre-trained on CLIP), 입력 해상도 224×224
- Text Encoder: Transformer (63M 파라미터, 12 layers, width 512, 8 attention heads), 최대 토큰 길이 77
- 초기화: CLIP 논문과 동일한 weight initialization 사용
4. Training process
- Joint Training
- 이미지–텍스트 contrastive learning: 같은 report의 이미지–텍스트 쌍은 cosine similarity 최대화, 잘못된(pairing) 쌍은 최소화
- 텍스트 입력으로는 impressions 섹션만 사용 (ConVIRT의 random sentence 대신)
- Hyperparamter Tunning
- CheXpert validation set으로 batch size, learning rate 탐색
- 매 1,000배치마다 validation mean AUC 계산, 최고 성능 체크포인트 저장
- 최적 설정: SGD(learning rate 1e-4, momentum 0.9), batch size 64, 총 4 epochs
- Softmax Evaluation 기법
- 각 병변에 대해 positive prompt({label})와 negative prompt(no {label})로 logits 계산
- 두 logits 간 softmax 후 positive 확률을 해당 병변 존재 확률로 사용
- Ensembling
- validation에서 mean AUC 상위 10개 체크포인트를 평균 앙상블
- softmax 확률을 각 모델에서 계산하여 평균 후, threshold 적용
5. Auxiliary Task용 Knowledge Distillation
- full report(≦512 tokens)를 사용하기 위해
- Teacher: impressions만 학습된 text encoder
- Student: context length 512 text encoder에 teacher logits를 MSE로 distill
- 이후 teacher image encoder로 student 모델 학습
6. Prompt Engineering
- 테스트 시 사용되는 prompt는 레이블 그대로({label})와 부정형(no {label})
7. Metrics 및 통계 분석
- AUROC: CheXpert test set(500개 samples) 및 PadChest(39,053개 samples)
- MCC, F1: CheXpert test set의 14개 condition에 softmax 확률과 validation에서 최적화된 threshold 적용 후 계산
- Confidence Interval: non-parametric bootstrap(1,000 re-samples, α=0.05)으로 산출
본 방법은 contrastive learning 기반의 self-supervised 학습 절차를 통해 MIMIC-CXR의 이미지–텍스트 쌍만으로 zero-shot multi-label classification 모델을 학습하며, CheXpert 및 PadChest 외부 데이터셋에서 robust한 성능을 보인다.
Results
1. Data & Training
a. Training pipeline
- Image–text pair: MIMIC-CXR의 흉부 X선 이미지와 해당 radiology report(주로 impressions 섹션) 쌍을 입력으로 사용
- Contrastive learning: Vision Transformer(ViT-B/32)와 텍스트 Transformer를 joint training
- Positive pair: 올바른 이미지–텍스트 쌍의 cosine similarity 최대화
- Negative pair: 잘못된 쌍의 cosine similarity 최소화
- Natural-language supervision: 수동 라벨 없이, 보고서 텍스트를 감독 신호로 활용하여 표현 학습
b. Zero-shot 예측
- 각 병변(pathology)에 대해 positive prompt(예: “consolidation”)와 negative prompt(“no consolidation”) 생성
- 두 prompt에 대한 logits를 softmax로 정규화하여 positive 확률 계산
- 계산된 확률로 병변 존재 여부를 multi-label 분류
→ 학습에 사용된 데이터: MIMIC-CXR 데이터셋에서 추출한 377,110개의 chest X-ray 이미지와 해당 raw radiology report 쌍
→ 학습 방법: 이미지–텍스트 쌍을 이용한 zero-shot learning. 모델은 어떤 X-ray 이미지가 어떤 보고서에 대응하는지 예측하도록 학습됨
2. 방사선과 전문의와의 비교
a. F1 score 비교
- CheXpert test의 5개 병변에 대해, 모델(파란색) vs 전문의(초록색) F1
- Pleural effusion: 모델 유의미하게 우수
- Atelectasis: 모델 유의미하게 열세
- 나머지 병변(cardiomegaly, consolidation, oedema): 통계적 차이 없음
b. MCC 비교
- 동일 5개 병변에 대해 MCC 측정
- 모든 병변에서 모델과 전문의 간 통계적 차이 없음
c. ROC curve
- 모델 ROC 곡선이 전문의 operating point 위에 있을수록 모델이 더 우수함을 의미
- 전체 병변 평균 ROC에서 모델이 라인 상방에 위치
평균 MCC, F1에서 모델과 전문의 간 통계적 차이 없음
개별 병변: cardiomegaly(F1)에서 모델 우수, atelectasis(F1)에서 모델 열세
3. 기존 Label-efficient 방법과의 비교
CheXzero는 0% 레이블로 0.889를 달성, 최고 supervised 모델(DAM) 대비 −0.042 포인트 차이
ConVIRT, MedAug, MoCo-CXR 등 기존 self-supervised 방법들을 모두 능가
특히 pleural effusion 분류에서는 AUC 0.932로 MoCo-CXR (0.1%–10% 레이블 사용) 및 MedAug(1%)를 모두 뛰어넘음.
4. Auxiliary Task 수행
zero-shot 특성을 활용해 추가 과제에도 적용 가능하게 하였다.
PadChest Data Generalization 성능:
- 39,053장 전문의 주석 chest X-ray(57개 소견, n > 50)
- AUC ≥ 0.900: 6개 소견
- AUC ≥ 0.700: 38개 소견
- 학습 시 전혀 본 적 없는 다양한 radiographic finding에 대해 우수한 zero-shot 성능
PadChest에서 CheXNet vs CheXzero
- Atelectasis, Consolidation, Oedema 3개 병변에서 CheXzero가 CheXNet(100% 레이블)보다 우수
→ CheXzero는 학습 과정에서 단 한 장의 레이블도 사용하지 않고, CheXpert 및 PadChest 등의 다양한 과제에서 전문가 및 기존 레이블 효율적 방법들과 동등하거나 우수한 성능을 보였다.
5. Ensemble 효과
상위 10개 체크포인트 앙상블로 single model 대비 AUROC +0.011 향상
F1 역시 single model(0.563)에서 앙상블(0.606)로 개선
Discussion
Objective
본 연구의 목적은 어떠한 explicit labels 없이도 medical imaging에 적용 가능한 zero-shot classification 방법을 개발하고 그 성능을 입증하는 것이다.
Notable findings
- Zero-shot 방법은 fully supervised 모델과 expert radiologist의 성능과 통계적으로 유의미한 차이가 없으며, 학습에 사용되지 않은 병변에서도 높은 성능을 보인다.
- CheXpert Competition 기준으로, self-supervised 모델의 평균 AUC는 0.889로, 최고 성능의 fully supervised 모델 대비 −0.042 포인트 차이에 불과하다.
- 평균 MCC와 F1에서도 radiologist 수준의 성능을 달성하였다.
기존 Label-efficient 방법과의 비교
- 1% 레이블 사용 시 평균 AUC 0.870을 기록한 이전 연구들과 달리, 본 방법은 0% 레이블로 평균 AUC 0.889를 달성한다.
- 수작업 주석 없이 다양한 병변을 예측할 수 있어, 레이블 수집 비용과 시간을 크게 절감한다.
Generalization 능력
- MIMIC-CXR로 학습한 모델을 CheXpert 및 PadChest의 독립적 데이터셋에서 검증한 결과, site-specific bias 없이 높은 정확도를 유지하였다.
- PadChest 데이터셋의 57개 소견 중 6개에서 AUC ≥ 0.900, 38개에서 AUC ≥ 0.700을 기록하며, differential diagnosis 등 auxiliary task에도 적용 가능함을 보여주었다.
- 전처리 없이 즉시 사용할 수 있는 prompt engineering만으로 다양한 과제에 대응 가능하다.
Limitations
- 하이퍼파라미터 및 임계값 탐색: 레이블된 validation set을 이용해 반복적으로 성능을 확인해야 한다.
- 데이터 타입 제한: 현재는 2D chest X-ray에만 적용 가능하며, 3D 영상(CT, MRI)이나 time series, EHR 등 다른 modality로의 확장이 필요하다.
- 텍스트 표현 편향: 보고서에 특정 용어가 누락되거나 다양한 표현 방식이 존재할 경우 예측 성능이 저하될 수 있다.
- 해상도 한계: 작은 병변 탐지를 위해서는 더 큰 이미지 크기 지원이 요구된다.
Conclusion & Future works
본 연구는 chest X-ray와 corresponding radiology report의 contrastive learning을 기반으로 한 self-supervised 방법이 multiple pathology를 zero-shot으로 분류할 수 있음을 보여주었다.
- 라벨 없는 대규모 unstructured text 데이터를 활용하여, 의료 영상 해석에서 레이블 의존성을 크게 낮추고 clinical-workflow 효율을 제고할 수 있다.
- 향후에는 3D 영상, time series, non-imaging data와의 결합 연구 및 고해상도 확장, 다양한 의료 영역으로의 적용을 통해 방법의 범용성을 더욱 강화할 예정이다.