[논문 리뷰] Self-supervised Image-text Pre-training With Mixed Data In Chest X-rays

Xiaosong Wang, Ziyue Xu외 3명|arXiv (Cornell University)|2021. 03. 30.

COVID-19 diagnosis using AI참고 문헌 44인용 수 14

한 줄 요약

이 논문은 흉부 X-레이에서 쌍체, 비쌍체, 그리고 병원 간 데이터를 포함한 혼합 데이터 입력을 활용하여 공동 표현 학습을 위한 자기지도 학습 기반 트랜스포머 기반 이미지-텍스트 사전학습 프레임워크인 UWOX를 제안한다. 다중 척도 마스킹 비전 모델링과 교차상관관계 모듈을 도입함으로써 UWOX는 MIMIC-CXR, NIH14-CXR, OpenI-CXR 데이터셋에서 분류, 검색, 이미지 재구성 작업에서 최고 성능을 기록하며, 이질적인 데이터 원천에도 불구하고 강건성과 일반화 능력을 입증한다.

ABSTRACT

Pre-trained models, e.g., from ImageNet, have proven to be effective in boosting the performance of many downstream applications. It is too demanding to acquire large-scale annotations to build such models for medical imaging. Meanwhile, there are numerous clinical data (in the form of images and text reports) stored in the hospital information systems. The paired image-text data from the same patient study could be utilized for the pre-training task in a weakly supervised manner. However, the integrity, accessibility, and amount of such raw data vary across different institutes, e.g., paired vs. unpaired (image-only or text-only). In this work, we introduce an image-text pre-training framework that can learn from these raw data with mixed data inputs, i.e., paired image-text data, a mixture of paired and unpaired data. The unpaired data can be sourced from one or multiple institutes (e.g., images from one institute coupled with texts from another). Specifically, we propose a transformer-based training framework for jointly learning the representation of both the image and text data. In addition to the existing masked language modeling, multi-scale masked vision modeling is introduced as a self-supervised training task for image patch regeneration. We not only demonstrate the feasibility of pre-training across mixed data inputs but also illustrate the benefits of adopting such pre-trained models in 3 chest X-ray applications, i.e., classification, retrieval, and image regeneration. Superior results are reported in comparison to prior art using MIMIC-CXR, NIH14-CXR, and OpenI-CXR datasets.

연구 동기 및 목표

• 병원 데이터베이스에서의 대규모 애너테이션된 의료 영상 데이터 부족 문제를 해결하기 위해, 미사용된 원시 임상 이미지-텍스트 쌍을 활용하고자 한다.
• 쌍체, 비쌍체, 그리고 병원 간 이미지-텍스트 데이터를 포함한 혼합 입력 데이터를 처리할 수 있는 통합 사전학습 프레임워크를 개발하고자 한다.
• 자기지도 사전학습을 통해 흉부 X-레이 응용 분야인 분류, 검색, 이미지 재구성 작업의 최종 성능을 향상시키고자 한다.
• 충분한 데이터가 확보될 경우 트랜스포머 기반 아키텍처와 다중 척도 비전 모델링이 CNN 기반 모델보다 우수하거나 동등한 성능을 낼 수 있음을 입증하고자 한다.

제안 방법

• 자기지도 사전학습을 활용한 통합 트랜스포머 아키텍처를 제안하여 이미지 및 텍스트 표현 학습을 공동으로 수행한다.
• 이미지 패치 재구성이라는 자기지도 학습 과제로 다중 척도 마스킹 비전 모델링을 도입하여 공간적 특징 학습을 향상시킨다.
• 쌍체가 아니거나 다른 기관에서 유래한 데이터일 경우에도 이미지와 텍스트 간 관계를 모델링하기 위해 교차상관관계 모듈을 활용한다.
• 쌍체 및 비쌍체 설정 모두에서 이미지와 텍스트 표현 간 정렬을 향상시키기 위해 쌍 매칭 손실을 적용한다.
• 텍스트에 대해 마스킹 언어 모델링을 적응하고, 이를 마스킹된 이미지 패치 재구성과 결합하여 공동 사전학습을 수행한다.
• 다양한 기관에서 온 데이터(예: 한 곳의 이미지와 다른 곳의 보고서)를 포함한 혼합 입력 데이터를 활용하여 학습을 지원함으로써 데이터 활용도를 극대화한다.

실험 결과

연구 질문

• RQ1 자기지도 학습 기반 이미지-텍스트 사전학습 프레임워크는 다수 기관의 쌍체 및 비쌍체 이미지-텍스트 데이터를 포함한 혼합 입력 데이터로부터 효과적으로 학습할 수 있는가?
• RQ2 의료 영상에서 단일 척도 접근 방식에 비해 다중 척도 마스킹 비전 모델링은 이미지 표현 학습에 어떻게 향상되는가?
• RQ3 데이터가 비쌍체이거나 다른 기관에서 유래한 경우, 제안된 교차상관관계 모듈이 이미지 및 텍스트 특징 간 정렬을 얼마나 향상시키는가?
• RQ4 제안된 UWOX 프레임워크는 분류, 검색, 이미지 재구성과 같은 최종 흉부 X-레이 작업에서 기존 방법보다 뛰어난 성능을 보일 수 있는가?
• RQ5 충분한 훈련 데이터가 확보될 경우 순수 트랜스포머 기반 아키텍처가 의료 영상 분석에서 CNN 기반 모델과 경쟁력 있거나 슈퍼어리어한 성능을 낼 수 있는가?

주요 결과

• UWOX는 MIMIC-CXR 데이터셋에서 최고 성능을 기록하였으며, 이미지 기반 및 텍스트 기반 쿼리에서 각각 P@10이 0.621과 0.603으로 기존 방법을 초월하였다.
• 블록 크기 B=8인 다중 척도 이미지 인코더 및 디코더가 가장 높은 이미지 재구성 품질을 달성하였으며, 평균 SSIM은 0.872, PSNR는 356.5를 기록하였다.
• 쌍 매칭 손실 적용 시 이미지 전용 검색 성능이 크게 향상되었으며, AUC는 0.727에서 0.732로 상승하였다.
• 추가로 NIH14-CXR의 이미지를 활용한 UWOX-mixup는 가장 높은 검색 정밀도를 기록하였으며, 텍스트 쿼리의 P@1은 0.763, 이미지+텍스트 쿼리의 P@10은 0.657이었다.
• 모델은 강력한 일반화 능력을 보이며, 혼합 시나리오에서 NIH14-CXR에서 최고의 전체 성능(P100%+uP)을 기록하여 데이터 이질성에 대한 강건성을 입증하였다.
• 충분한 훈련 데이터가 확보될 경우 트랜스포머 기반 이미지 모델링 접근 방식이 분류 및 재구성 작업에서 CNN 기반 베이스라인과 동등하거나 슈퍼어리어한 성능을 낼 수 있었다.