[논문 리뷰] Pretrained ViTs Yield Versatile Representations For Medical Images

C. Matsoukas, Johan Fredin Haslum외 3명|arXiv (Cornell University)|2023. 03. 13.

AI in cancer detection인용 수 7

한 줄 요약

이 논문은 ImageNet에서 사전학습된 비전 트랜스포머(ViTs)가 여러 2D 의료 영상 분류 벤치마크에서 컨volution 신경망(CNNs)과 동등한 성능을 달성함을 보여준다. 의료 데이터셋에서 미세조정을 거친 후, 오프더샤프트(off-the-shelf) ViTs는 CNNs와 유사하거나 약간 뛰어나며, 특히 자기지도 학습(self-supervised pretraining)을 통해 성능이 향상된다. 또한 주의 맵(attention maps)을 통해 해석 가능성(interpretability)이 향상되고, 인덕티브 바이어스(inductive biases)에 대한 의존도가 감소한다.

ABSTRACT

Convolutional Neural Networks (CNNs) have reigned for a decade as the de facto approach to automated medical image diagnosis, pushing the state-of-the-art in classification, detection and segmentation tasks. Over the last years, vision transformers (ViTs) have appeared as a competitive alternative to CNNs, yielding impressive levels of performance in the natural image domain, while possessing several interesting properties that could prove beneficial for medical imaging tasks. In this work, we explore the benefits and drawbacks of transformer-based models for medical image classification. We conduct a series of experiments on several standard 2D medical image benchmark datasets and tasks. Our findings show that, while CNNs perform better if trained from scratch, off-the-shelf vision transformers can perform on par with CNNs when pretrained on ImageNet, both in a supervised and self-supervised setting, rendering them as a viable alternative to CNNs.

연구 동기 및 목표

• 비전 트랜스포머(ViTs)가 의료 영상 분류에서 CNNs의 타당한 대안이 될 수 있는지 평가하기 위해.
• ImageNet 사전학습과 자기지도 학습이 다양한 의료 영상 작업에서 ViT 성능에 미치는 영향을 평가하기 위해.
• 의료 환경에서 CNNs와 비교해 ViTs가 해석 가능성과 특징 학습 측면에서 어떤 이점을 제공하는지 조사하기 위해.
• 의료 영상 환경에서 모델 용량과 데이터 크기 증가에 따라 ViTs가 얼마나 효과적으로 스케일링되는지 판단하기 위해.
• 특히 병변의 국소화에 있어 ViTs와 CNNs의 주의 패턴을 비교하기 위해.

제안 방법

• ISIC 2019, APTOS 2019, PatchCamelyon, CheXpert, CBIS-DDSM 등 6개의 공개 2D 의료 영상 벤치마크에서 DeiT-Small 및 DeiT-Base ViT 아키텍처를 미세조정하였다.
• 동일한 학습 프rotocol(감독 학습 및 자기지도 학습 사전학습 설정 포함)을 사용해 ViT 성능을 ResNet-50와 비교하였다.
• 의료 영상 데이터를 사용해 MoCo v2와 BYOL을 활용한 자기지도 학습 사전학습을 적용하여 도메인 특화 표현 학습 능력을 평가하였다.
• Grad-CAM과 클래스 토큰 주의 맵을 활용해 ViTs와 CNNs 간의 모델 해석 가능성 차이를 시각화하고 비교하였다.
• 패치 크기, 모델 용량, 주의 패턴에 대한 분석을 통해 아키텍처의 성능에 미치는 영향을 분석하기 위해 아블레이션 스터디를 수행하였다.
• 소수의 샘플(few-shot) 및 전체 감독 학습(full-supervision) 설정을 포함한 다양한 데이터 제약 조건에서 ViT 성능을 평가하여 데이터 효율성(data efficiency)을 분석하였다.

실험 결과

연구 질문

• RQ1 ImageNet에서 사전학습된 비전 트랜스포머(ViTs)가 2D 의료 영상 분류에서 CNNs와 동등한 성능을 낼 수 있는가?
• RQ2 의료 데이터에서 자기지도 학습을 통해 ImageNet 사전학습을 초월해 ViT 성능이 더욱 향상되는가?
• RQ3 의료 영상에서 병변을 국소화하는 데 있어 ViT 주의 맵은 CNN 활성화 맵과 어떻게 비교되는가?
• RQ4 ViTs가 전이 학습(transfer learning)에서 CNNs와 비교해 얼마나 유리한 성능 향상을 보이는가? 그리고 이는 초기 학습에서부터의 성능 향상과 어떻게 다를 수 있는가?
• RQ5 의료 영상 작업에서 ViTs가 모델 크기와 데이터 양 증가에 따라 CNNs보다 더 잘 스케일링되는가?

주요 결과

• ImageNet에서 사전학습된 ViTs는 ISIC 2019, APTOS 2019, PatchCamelyon, CheXpert, CBIS-DDSM 등 평가된 모든 의료 영상 분류 벤치마크에서 ResNet-50와 동등한 성능을 기록하였다.
• 초기 학습에서 ViTs는 CNNs보다 더 큰 성능 향상을 보였지만, ImageNet 사전학습을 통해 얻는 성능 향상은 CNNs보다 훨씬 더 크다.
• 의료 데이터에서 자기지도 학습을 통해 ViTs는 CNNs를 약간 앞서는 일관된 성능 향상을 보였으며, 특히 ImageNet 사전학습과 조합된 경우 더욱 두드러졌다.
• ViTs는 CNNs보다 더 국소화되고 해석 가능한 주의 맵을 생성했으며, 피부 병변, 출혈, 흑색종, 흉부 X-ray에서의 투명도 이상 등 병변에 명확히 집중하는 경향을 보였다.
• ViTs는 초기 레이어부터 전반적인 이미지 영역과 국소 영역을 모두 주시하며, 장거리 상관관계(long-range dependencies)를 조기에 포착하는 것으로 나타났다. 반면 CNNs는 전반적인 맥락을 얻기 위해 깊은 레이어에 의존한다.
• 모델 용량 증가와 작은 패치 크기에서 ViTs의 성능는 CNNs와 유사하게 스케일링되었으며, 인덕티브 바이어스를 줄였음에도 불구하고 강력한 성능 유지를 보였다.