[논문 리뷰] Vision Transformer using Low-level Chest X-ray Feature Corpus for COVID-19 Diagnosis and Severity Quantification

Sang Joon Park, Gwanghyun Kim외 7명|arXiv (Cornell University)|2021. 04. 15.

COVID-19 diagnosis using AI참고 문헌 34인용 수 8

한 줄 요약

이 논문은 대규모 데이터셋에서 사전 훈련된 저수준 흉부 X선(CXR) 특징 임bedding을 활용하여 코로나19 진단 및 중증도 정량화를 향상시키는 새로운 비전 트랜스포머(ViT) 프레임워크를 제안한다. 220,000장 이상의 CXR 영상으로 훈련된 백본 네트워크를 통해 공통적인 이상 소견(예: 융합, 투명도 감소)을 추출함으로써, 다양한 외부 데이터셋에서 뛰어난 일반화 능력을 보이며 기존의 CNN 및 ViT 기반 모델을 능가하는 최신 기술 수준(SOTA)의 성능을 달성한다.

ABSTRACT

Developing a robust algorithm to diagnose and quantify the severity of COVID-19 using Chest X-ray (CXR) requires a large number of well-curated COVID-19 datasets, which is difficult to collect under the global COVID-19 pandemic. On the other hand, CXR data with other findings are abundant. This situation is ideally suited for the Vision Transformer (ViT) architecture, where a lot of unlabeled data can be used through structural modeling by the self-attention mechanism. However, the use of existing ViT is not optimal, since feature embedding through direct patch flattening or ResNet backbone in the standard ViT is not intended for CXR. To address this problem, here we propose a novel Vision Transformer that utilizes low-level CXR feature corpus obtained from a backbone network that extracts common CXR findings. Specifically, the backbone network is first trained with large public datasets to detect common abnormal findings such as consolidation, opacity, edema, etc. Then, the embedded features from the backbone network are used as corpora for a Transformer model for the diagnosis and the severity quantification of COVID-19. We evaluate our model on various external test datasets from totally different institutions to evaluate the generalization capability. The experimental results confirm that our model can achieve the state-of-the-art performance in both diagnosis and severity quantification tasks with superior generalization capability, which are sine qua non of widespread deployment.

연구 동기 및 목표

• 팬데믹 기간 동안 코로나19 CXR 데이터의 레이블이 부족한 문제를 해결하기 위해 공통적인 이상 소견을 보이는 풍부한 레이블이 없는 CXR 데이터를 활용한다.
• 질병 분류에서 특징 학습을 분리함으로써 코로나19 진단 및 중증도 정량화의 일반화 능력을 향상시킨다.
• 다양한 임상 환경에서 도메인 전이가 최소한이 되도록 강력하고 이식 가능한 모델을 개발한다.
• 사전 훈련된 모델에서 유도된 저수준 CXR 특징이 의료 영상 작업에서 ViT의 최적의 코퍼스가 될 수 있음을 보여준다.

제안 방법

• CheXpert와 같은 공개 CXR 데이터셋에서 대규모 공개 데이터셋을 사전 훈련하여, 융합, 투명도 감소, 부종과 같은 공통적인 저수준 이상 소견을 탐지할 수 있도록 백본 네트워크를 훈련시킨다.
• 사전 훈련된 백본에서 특징 임bedding을 추출하여 비전 트랜스포머의 구조화된 코퍼스로 활용한다.
• ViT 아키텍처를 사용하여 학습된 특징 임bedding을 주의(attention)하여 CXR 전반에 걸친 장거리 상관관계를 모델링하고 고수준의 진단을 수행한다.
• ROI 최대 풀링을 적용하여 각 패치의 중증도 맵을 전역 중증도 점수 배열로 변환함으로써 임상 레이블과의 일치도를 향상시킨다.
• 백본은 동결한 채로 제한된 레이블이 있는 코로나19 데이터로 ViT를 미세 조정하여 과적합을 방지한다.
• 중증도 정량화에는 평균 제곱오차(MSE)를, 진단에는 이진 교차 엔트로피를 사용하며, 데이터 증강 및 기관 간 데이터 다양성을 통한 강건성 향상을 위해 모델을 훈련시킨다.

실험 결과

연구 질문

• RQ1 사전 훈련된 모델에서 유도된 저수준 CXR 특징 임bedding이 코로나19 진단에 있어 비전 트랜스포머의 효과적인 코퍼스가 될 수 있는가?
• RQ2 사전 훈련된 특징 코퍼스를 사용할 경우, 표준 ViT 또는 CNN 기반 모델 대비 외부 및 미사용된 데이터셋에서 일반화 성능이 향상되는가?
• RQ3 ROI 풀링 방식(최대 풀링 대 평균 풀링)의 선택이 중증도 정량화 성능에 어떤 영향을 미치는가?
• RQ4 단지 CXR 영상만을 사용하여도 모델이 임상 전문가 수준의 중증도 예측 성능을 달성할 수 있는가?
• RQ5 훈련 시 다양한 기관의 데이터를 포함시킬 경우 중증도 정량화의 일반화 능력이 크게 향상되는가?

주요 결과

• 제안된 모델은 세 개의 외부 데이터셋에서 평균 중증도 예측 MSE가 1.655로, 직접 패치 임베딩을 사용하는 기준 ViT 및 ResNet 특징을 사용하는 하이브리드 ViT보다 뛰어난 성능을 보였다.
• CNUH, YNU, KNUH의 외부 테스트 세트에서 각각 MSE 값이 1.682, 1.677, 1.607로, 뛰어난 일반화 성능을 확보했다.
• ROI 최대 풀링을 사용함으로써 평균 풀링 대비 MSE가 0.191 감소하여 국소적 이상 소견을 더 잘 포착함을 입증했다.
• 훈련 시 기관 데이터를 제외한 경우 성능이 크게 악화되었으며(CNUH에서 MSE = 2.503), 다양한 데이터가 안정적인 중증도 예측을 위해 필수적임을 확인했다.
• 진단 및 중증도 정량화 모두에서 최신 기술 수준의 성능을 달성했으며, 중증도 점수 예측에서는 임상 전문가 수준의 성능을 보였다.
• 다양한 기관에서의 강건한 성능은 본 모델이 자원이 부족하거나 자원이 제한된 의료 환경에서의 실질적 도입 잠재력을 확인한다.