[논문 리뷰] SPCXR: Self-supervised Pretraining using Chest X-rays Towards a Domain Specific Foundation Model

Syed Muhammad Anwar, Abhijeet Parida외 5명|arXiv (Cornell University)|2022. 11. 23.

COVID-19 diagnosis using AI인용 수 4

한 줄 요약

이 논문은 비라벨링 흉부 X선 영상을 사용하여 영역 특화 표현을 학습하는 자기지도 학습 비전 트랜스포머 프리트레인 프레임워크인 SPCXR를 제안한다. 대규모 CXR 데이터에서 그룹 마스킹 자기지도 학습을 적용함으로써, 코로나19 및 폐렴 분류, 폐 세분화와 같은 후행 작업에서 최신 기술 수준(SOTA) 성능을 달성하였으며, 감염병이 드문 소아 환자 데이터에서는 지도 학습 기반 모델 대비 25% 높은 정확도를 기록하였다.

ABSTRACT

Chest X-rays (CXRs) are a widely used imaging modality for the diagnosis and prognosis of lung disease. The image analysis tasks vary. Examples include pathology detection and lung segmentation. There is a large body of work where machine learning algorithms are developed for specific tasks. A significant recent example is Coronavirus disease (covid-19) detection using CXR data. However, the traditional diagnostic tool design methods based on supervised learning are burdened by the need to provide training data annotation, which should be of good quality for better clinical outcomes. Here, we propose an alternative solution, a new self-supervised paradigm, where a general representation from CXRs is learned using a group-masked self-supervised framework. The pre-trained model is then fine-tuned for domain-specific tasks such as covid-19, pneumonia detection, and general health screening. We show that the same pre-training can be used for the lung segmentation task. Our proposed paradigm shows robust performance in multiple downstream tasks which demonstrates the success of the pre-training. Moreover, the performance of the pre-trained models on data with significant drift during test time proves the learning of a better generic representation. The methods are further validated by covid-19 detection in a unique small-scale pediatric data set. The performance gain in accuracy (~25%) is significant when compared to a supervised transformer-based method. This adds credence to the strength and reliability of our proposed framework and pre-training strategy.

연구 동기 및 목표

• 희귀 또는 소아 폐 질환에 대한 라벨이 부여된 의료 영상 데이터가 부족한 문제를 해결하기 위해, 대규모 비라벨링 CXR 데이터를 활용해 프리트레인하는 것.
• 다양한 임상 데이터 분포와 영상 장치에 걸쳐 일반화 가능한 흉부 레이저의 도메인 특화 기초 모델을 개발하는 것.
• 성인 CXR 데이터에서의 자기지도 프리트레인을 통해 자원이 제한된 후행 작업—특히 소아 인구에서의 성능 향상—을 도모하는 것.
• 통합된 프리트레인 모델을 사용해 분류 및 세분화 작업 모두에 대해 학습된 표현의 전이성을 입증하는 것.
• 실제 임상 환경에서 흔히 발생하는 데이터 드리프트 및 분포 이탈 상황에서 프리트레인 모델의 견고성을 검증하는 것.

제안 방법

• 다양한 출처에서 확보한 100,000장의 비라벨링 소아 및 성인 CXR 영상에서, 그룹 마스킹 자기지도 학습(GMML) 전략을 사용해 비전 트랜스포머(ViT-S)를 프리트레인한다.
• 입력 CXR의 패치를 마스킹하고 모델이 이를 재구성하도록 학습시켜, 전반적이고 국소적인 해부학적 특징을 학습하도록 유도한다.
• 심천 및 벨라루스 데이터셋을 사용해 UNETR 아키텍처 내에서 프리트레인된 ViT-S를 세분화를 위해 피지컬러닝한다.
• 추가적인 라벨링 없이도 RSNA, RICORD 및 소아 데이터셋에서 동일한 프리트레인 가중치를 후행 분류 작업(예: 코로나19, 폐렴, 건강)에 적용한다.
• 프리트레인 후 테스트 시 256×256 해상도의 이미지에서 400 에포크 동안 AdamW 옵timizer를 사용하고, 가중치 감소는 1e-5, 초기 학습률은 5e-4로 설정한다.
• 정확도, AUC-ROC, TPR, TNR, 평균 IoU, Dice 계수, Hausdorff 표면 거리 등의 표준 지표를 사용해 성능을 평가한다.

실험 결과

연구 질문

• RQ1 대규모 비라벨링 CXR에서의 자기지도 프리트레인은 소아 코로나19 탐지와 같은 자원이 제한된 후행 작업에서 성능 향상에 기여하는가?
• RQ2 성인 CXR에서 프리트레인한 모델이 라벨이 적고 분포 이탈이 심한 새로운 소아 데이터셋에 얼마나 잘 전이되는가?
• RQ3 제안된 GMML 기반 자기지도 프리트레인 방식은 다양한 CXR 데이터 분포에서 정확도와 견고성 측면에서 완전히 지도 학습 기반 미세조정 대비 어떻게 비교되는가?
• RQ4 단일 프리트레인 비전 트랜스포머 모델이 흉부 레이저에서 분류 및 세분화를 포함한 다수의 후행 작업에 효과적으로 활용될 수 있는가?
• RQ5 자기지도 표현 학습은 특히 라벨 데이터가 부족한 소규모 데이터셋에서 일반화 능력을 향상시키고 과적합을 줄이는 데 기여하는가?

주요 결과

• 자기지도 프리트레인된 SPCXR 모델은 소아 CXR 데이터셋에서 이진 분류(건강 대비 코로나19)에서 94.73%의 정확도를 기록하였으며, 지도 학습 기반 트랜스포머 기준 대비 25% 향상된 성능을 보였다.
• RSNA 및 RICORD 테스트 세트에서 성인 CXR 분류에 대해 98.25%의 정확도를 달성하여 표준 벤치마크에서 뛰어난 성능을 입증하였다.
• 폐 세분화 작업에서, 프리트레인된 모델은 완전히 지도 학습 기반 UNETR 대비 Hausdorff 표면 거리를 약 46% 감소시켜 경계 정확도 향상을 입증하였다.
• 모델는 흉부 방사선 전문의보다 코로나19를 더 높은 진짜 양성률과 낮은 가짜 양성률로 식별하여 임상적 선별의 신뢰도를 향상시켰다.
• 세균성 폐렴과 바이러스성 코로나19를 99.67%의 정확도로 구분할 수 있었으며, 이는 미세한 병리적 패턴을 효과적으로 학습할 수 있음을 시사한다.
• 주의 히트맵 분석 결과, 대부분의 경우 임상적으로 관련 있는 영역에 주의를 기울였지만, 일부 경우 폐 외부 영역에도 주의가 분포하는 현상이 관찰되어 해석 가능성 향상의 필요성이 제기되었다.