[논문 리뷰] Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification

Junfei Xiao, Yutong Bai외 2명|arXiv (Cornell University)|2022. 10. 23.

COVID-19 diagnosis using AI인용 수 6

한 줄 요약

이 논문은 266,340개의 레이블이 없는 흉부 X선 영상에서 마스크된 자동에코더(MAE)를 사용하여 비전 트랜스포머(ViT)의 강력한 사전 훈련 및 미세조정 레시피를 제안한다. 이로 인해 세 가지 벤치마크 데이터셋에서 다중 레이블 흉부 질환 분류 작업에서 최신 기술 수준의 컨볼루션 네트워크(CNN)와 비교해도 성능이 유사하거나 뛰어나며, 각각 NIH ChestX-ray14, CheXpert, COVIDx에서 mAUC 점수 82.3%, 89.2%, 99.3%를 기록하였다.

ABSTRACT

Vision Transformer (ViT) has become one of the most popular neural architectures due to its great scalability, computational efficiency, and compelling performance in many vision tasks. However, ViT has shown inferior performance to Convolutional Neural Network (CNN) on medical tasks due to its data-hungry nature and the lack of annotated medical data. In this paper, we pre-train ViTs on 266,340 chest X-rays using Masked Autoencoders (MAE) which reconstruct missing pixels from a small part of each image. For comparison, CNNs are also pre-trained on the same 266,340 X-rays using advanced self-supervised methods (e.g., MoCo v2). The results show that our pre-trained ViT performs comparably (sometimes better) to the state-of-the-art CNN (DenseNet-121) for multi-label thorax disease classification. This performance is attributed to the strong recipes extracted from our empirical studies for pre-training and fine-tuning ViT. The pre-training recipe signifies that medical reconstruction requires a much smaller proportion of an image (10% vs. 25%) and a more moderate random resized crop range (0.5~1.0 vs. 0.2~1.0) compared with natural imaging. Furthermore, we remark that in-domain transfer learning is preferred whenever possible. The fine-tuning recipe discloses that layer-wise LR decay, RandAug magnitude, and DropPath rate are significant factors to consider. We hope that this study can direct future research on the application of Transformers to a larger variety of medical imaging tasks.

연구 동기 및 목표

• 비전 트랜스포머(ViT)의 임의의 형태가 의료 영상, 특히 다중 레이블 흉부 질환 분류에 얼마나 잠재력을 지녔는지 평가하기 위해.
• ViT의 데이터 소비 성향을 해결하기 위해 266,340개의 레이블이 없는 흉부 X선 영상에서 마스크된 자동에코더(MAE)를 사용해 대규모 사전 훈련을 수행하기 위해.
• 저대비 해부학적 구조와 작은 병변 특징 등 의료 영상의 고유한 특성에 맞게 맞춤형 사전 훈련 및 미세조정 레시피를 개발하기 위해.
• 동일한 데이터에서 사전 훈련된 최신 기술 수준의 CNN과의 공정한 비교를 제공하기 위해.
• 재현 가능성과 의료 비전 분야의 향후 연구를 지원하기 위해 코드와 사전 훈련된 모델(ViT-Small, ViT-Base)을 공개하기 위해.

제안 방법

• 90%의 영상 패치를 마스크하고, 나머지 가시 패치에서 누락된 픽셀을 복원하도록 모델을 훈련하는 마스크된 자동에코더(MAE)를 사용해 비전 트랜스포머를 사전 훈련한다.
• 자연 이미지 사전 훈련에서 일반적으로 사용하는 0.2–1.0과는 대비로 훨씬 좁은 0.5–1.0의 랜덤 리사이즈 크롭 범위를 사용한 수정된 데이터 증강 전략을 적용한다.
• 자연 이미지 MAE에서 사용하는 75%와는 대비로 의료 영상 복원에 더 적합한 높은 마스크 비율(90%)을 적용한다.
• 레이어별 학습률 감소, magnitude 15의 RandAug 증강, DropPath 비율 0.1을 사용해 사전 훈련된 ViT를 미세조정한다.
• 동일한 데이터와 동일한 평가 프로토콜을 사용해 사전 훈련된 최신 기술 수준의 CNN(예: DenseNet-121)과 ViT 성능을 비교한다.
• 해석 가능성과 국소화 성능를 분석하기 위해 Grad-CAM을 사용해 ViT와 CNN 간의 주의 맵을 비교한다.

실험 결과

연구 질문

• RQ1 대규모 의료 데이터셋에서 MAE를 사용해 사전 훈련한 순수 비전 트랜스포머(ViT)가 다중 레이블 흉부 질환 분류에서 최신 기술 수준의 CNN과 유사하거나 뛰어난 성능을 내는가?
• RQ2 의료 영상, 특히 흉부 X선에 적용했을 때 MAE의 최적 사전 훈련 초모수(마스크 비율, 크롭 범위, 데이터 스케일)는 무엇인가?
• RQ3 레이어별 학습률 감소, RandAug magnitude, DropPath 비율 등의 미세조정 초모수는 다중 레이블 분류 작업에서 ViT 성능에 어떤 영향을 미치는가?
• RQ4 Grad-CAM과 바운딩 박스 IoU로 측정했을 때, ViT의 작은 병변 국소화 능력은 CNN과 비교해 어떻게 다른가?
• RQ5 의료 데이터에서의 도메인 내 전이 학습(의료 데이터에서 사전 훈련)이 ImageNet 사전 훈련보다 의료 비전 작업에서 더 효과적인가?

주요 결과

• 사전 훈련된 ViT는 NIH ChestX-ray14에서 mAUC 82.3%를 기록하여 최신 기술 수준의 DenseNet-121와 동일한 성능을 보였다.
• 스탠포드 CheXpert 데이터셋에서 ViT는 mAUC 89.2%를 기록하여 문헌에 보고된 바 있는 최고의 CNN 결과를 초월했다.
• COVIDx 데이터셋에서 ViT는 mAUC 99.3%를 기록하여 도전적이고 소규모이며 다양한 특성을 지닌 데이터셋에서 뛰어난 성능을 보였다.
• ViT는 작은 병변(예: Nodule)에 대해 CNN보다 뛰어난 국소화 능력을 보였고, CNN은 더 큰 이상징(예: Cardiomegaly)을 더 잘 국소화했다.
• 266,340장의 흉부 X선 영상에서의 도메인 내 사전 훈련은 ImageNet 사전 훈련보다 ViT 성능을 크게 향상시켜 도메인 특화 데이터의 중요성을 입증했다.
• 레이어별 학습률 감소, RandAug magnitude 15, DropPath 비율 0.1은 다중 레이블 분류 작업에서 효과적인 ViT 미세조정을 위해 핵심적인 초모수로 규명되었다.