[논문 리뷰] Semi-supervised Learning for Quantification of Pulmonary Edema in Chest X-Ray Images

Ruizhi Liao, Jonathan Rubin외 6명|arXiv (Cornell University)|2019. 02. 27.

COVID-19 diagnosis using AI참고 문헌 13인용 수 7

한 줄 요약

이 논문은 330,000장의 임상 데이터셋에서 레이블이 부여된 이미지 6,000장만을 사용하여 흉부 X선 영상에서 폐부종의 심각도를 정량화하는 준지도 학습 베이지안 모델을 제안한다. 레이블이 있는 데이터와 없는 데이터로부터 압축된 잠재 표현을 학습함으로써, 이 방법은 지도 학습, 엔트로피 최소화, 다중 VAE 접근 방식보다 유의하게 낮은 제곱근 평균 제곱 오차와 더 높은 진정도 값과의 상관관계를 달성한다.

ABSTRACT

We propose and demonstrate machine learning algorithms to assess the severity of pulmonary edema in chest x-ray images of congestive heart failure patients. Accurate assessment of pulmonary edema in heart failure is critical when making treatment and disposition decisions. Our work is grounded in a large-scale clinical dataset of over 300,000 x-ray images with associated radiology reports. While edema severity labels can be extracted unambiguously from a small fraction of the radiology reports, accurate annotation is challenging in most cases. To take advantage of the unlabeled images, we develop a Bayesian model that includes a variational auto-encoder for learning a latent representation from the entire image set trained jointly with a regressor that employs this representation for predicting pulmonary edema severity. Our experimental results suggest that modeling the distribution of images jointly with the limited labels improves the accuracy of pulmonary edema scoring compared to a strictly supervised approach. To the best of our knowledge, this is the first attempt to employ machine learning algorithms to automatically and quantitatively assess the severity of pulmonary edema in chest x-ray images.

연구 동기 및 목표

• 심부전에서 중요한 그러나 미세한 임상적 지표인 흉부 X선 영상에서의 폐부종 심각도를 정확하게 정량화하는 기계 학습 방법을 개발하는 것.
• 레이블이 부족한 의료 영상 문제를 해결하기 위해 대규모의 레이블이 없는 흉부 X선 영상을 활용하여 모델 성능을 향상시키는 것.
• 베이지안 프레임워크를 통해 영상 분포와 레이블 예측을 동시에 학습함으로써 순수 지도 학습 또는 자기지도 학습 방법에 비해 회귀 정확도를 향상시킬 수 있음을 보여주는 것.
• 심부전 및 기타 체액 상태 관련 질환 연구 및 임상 의사결정 지원을 향상시키기 위해 환자의 상태를 정량적으로 표현할 수 있도록 하는 것.

제안 방법

• 모든 330,000장의 흉부 X선 영상(레이블이 있는 데이터와 없는 데이터 포함)으로부터 공유 잠재 표현을 학습하는 단일 VAE를 사용하는 베이지안 생성 모델이 사용된다.
• 학습된 잠재 특징을 폐부종 심각도 점수(0–3)로 매핑하는 회귀기와 함께 VAE를 공동으로 훈련하며, 재구성 및 예측 목표를 동시에 최적화한다.
• 스토캐스틱 잠재 변수를 통한 역전파를 가능하게 하기 위해 재생성 기법을 사용하고, 사후 분포를 표준 정규 분포 사전 분포로 정규화하기 위해 KL 발산 손실을 사용한다.
• 영상 재구성 손실은 영상 크기의 역수로 가중되며, VAE의 잠재 공간은 척도 편향을 줄이기 위해 정규화된다.
• 학습 과정에서는 Adam 옵timizer를 사용하며 학습률은 0.001, 배치 크기는 4이며 검증 손실 기반 조기 정지 전략을 적용한다.
• 각 레이블 클래스별로 별도의 VAE를 훈련하는 대신 전체 영상 세트에 대해 단일 VAE를 훈련함으로써 연속적이고 미세한 심각도 평가에 더 적합하다고 주장한다.

실험 결과

연구 질문

• RQ1 대규모의 레이블이 없는 흉부 X선 영상을 활용하는 준지도 학습 접근 방식이 완전히 지도 학습 방법에 비해 폐부종 심각도 회귀 정확도를 향상시킬 수 있는가?
• RQ2 레이블이 있는 데이터와 없는 데이터를 동시에 사용하여 VAE와 회귀기를 공동으로 훈련할 경우, 각 레이블 카테고리별로 별도의 모델을 훈련하는 것보다 더 나은 특징 표현을 얻을 수 있는가?
• RQ3 제한된 레이블 데이터 조건에서 제안된 방법의 성능이 엔트로피 최소화 및 다중 VAE 기반 모델과 비교해 어떻게 되는가?
• RQ4 단일 공유 VAE 표현이 연속적인 폐부종 심각도 점수와 관련된 미세한 영상 특징을 효과적으로 포착할 수 있는가?
• RQ5 전체 영상 분포를 함께 모델링함으로써 심각도 예측의 일반화 능력과 강인성은 어느 정도 향상되는가?

주요 결과

• 제안된 VAE-회귀 모델은 테스트 세트에서 제곱근 평균 제곱 오차(RMSE) 0.57을 달성하여 지도 학습 기준선(RMSE: 0.63), EM 기반 자기지도 학습(RMSE: 0.65), 다중 VAE 접근 방식(RMSE: 0.68)보다 뛰어난 성능을 보였다.
• 모델은 진정도 심각도 점수와 피어슨 상관계수(CC) 0.78을 기록했으며, 이는 지도 학습 기준선(CC: 0.72), EM(CC: 0.70), 다중 VAE(CC: 0.69)보다 유의미하게 높았다.
• 잠재 표현과 레이블 예측의 공동 학습은 모든 심각도 카테고리에서 더 정확하고 일관된 예측을 가능하게 했으며, 이는 예측 점수 분포도에서 명확히 확인되었다.
• 모델는 레이블이 없는 데이터를 효과적으로 활용함으로써 일반화 능력을 향상시키고, 검증 및 테스트 세트에서 과적합을 줄이며 강인성을 향상시켰다.
• 결과적으로, 제한된 레이블과 함께 영상 분포를 함께 모델링하는 것이, 특히 폐부종 심각도와 같이 연속적이고 미세한 특징을 갖는 경우, 각 레이블 카테고리별 별도의 VAE를 훈련하는 것보다 더 효과적임을 시사한다.
• 본 연구는 흉부 X선 영상에서 자동으로 정량화된 폐부종 점수를 산출하는 데 있어 기계 학습의 첫 성공적인 적용 사례를 제시하며, 더 나은 임상 의사결정 지원 및 표현형 분석을 가능하게 한다.