[22’ PMLR] ConVIRT: Contrastive Learning of Medical Visual Representations from Paired Images and Text
카테고리: Medical
🔍 Abstract
이 논문이 나오기 전까지는 medical domain에서 contrastive learning을 적용한 선구적인 모델이 없었다. 따라서 text embedding, image embedding이 굉장히 원초적인 방법으로 이루어지고 있었고, 이 논문 이후 수많은 medical representation learning 논문이 나온 만큼 높은 중요도를 가진다. 이 논문에서는 CLIP과 동일한 구조의 ConVIRT(Contrastive VIsual Representation Learning from Text)를 제안한다.
1. Introduction
ConVIRT는 거의 최초로 medical domain에서 contrastive learning을 적용한 선구적인 모델이다. 이때까지만 해도 medical image의 representation은 ImageNet pretraining weight transfer로 이루어지고 있었다. 그러나 위 그림과 같이 두 CXR(흉부 X-ray) 이미지는 natural image로 학습한 네트워크에게는 매우 유사해 보이고, 따라서 다른 진단임에도 불구하고 비슷한 representation을 갖게 된다. 따라서 natural image based representation을 그대로 사용하는 것은 suboptimal하다.
이러한 문제를 해결하기 위해 위와 같이 rule-based label extraction을 통해 medical image-label 쌍을 만들어 fine-tuning을 수행하는 방법도 제안되었지만, 이러한 rule-based label은 부정확할 뿐만 아니라 domain-specific하여 일반화가 어렵다는 문제가 있다.
따라서 저자들은 CLIP과 유사한 접근으로, natural language supervision을 통해 representation을 학습하는 ConVIRT(Contrastive VIsual Representation Learning from Text)를 제안한다. 이 모델은 특정 가공을 거치지 않고도 image-text pair를 사용하여 학습할 수 있도록 설계되었다. 이러한 방법으로, medical domain에서도 pre-trained CLIP model을 만들 수 있게 되었다.
2. Method
전체 모델의 구조는 위와 같다. Image Encoder는 ResNet50, Text Encoder는 ClinicalBERT를 사용하였다. 이때 $t$는 cropping, flipping, affine, blur 등의 augmentation function을 추상적으로 나타낸 것이고, $g$는 simple MLP이다. Loss function은 CLIP과 같이 InfoNCE loss를 사용하였다. Pre-training 데이터셋은 MIMIC-CXR(217K Pair)와 Bone X-ray(Local, 48K)를 사용하였다.
3. Evaluation
저자들은 ConVIRT의 representative power를 측정하기 위해 (1) Image Classification (Linear Probing, Fine-tuning), (2) Zero-shot Retrieval (Image to Image, Text to Image) 성능을 측정하였다. 이때 두 task 모두 비교적 in-domain task로, CXR 또는 Bone X-ray를 활용하는 task로 한정되어 있다.
먼저 Image Classification task에서는 ConVIRT가 다른 모델보다 우월했고, medical domain pre-training 덕분에 다른 모델들보다 적은 데이터 비율로도 높은 성능을 내어 data efficiency가 높았음을 확인할 수 있었다.
Zero-shot Retrieval에서도 마찬가지 결과를 얻을 수 있었고, fine-tuning을 통해 성능을 더 향상시킬 수 있었다.
마지막으로 t-SNE visualization을 통해 representation learning이 효과적으로 이루어졌음을 확인할 수 있었다.
정리하자면, ConVIRT는 medical domain에서 처음으로 contrastive learning을 적용한 모델로, 기존 모델에 비해 data efficiency가 높다는 것이 장점이다.
💡 Summary
ConVIRT는 medical domain에서 contrastive learning을 적용한 최초의 모델로, 비교적 data-efficient한 모델이다. MIMIC-CXR(217K), Bone X-ray(Local, 48K) 데이터셋을 사용하여 pre-training을 진행하였다. 이 모델은 CLIP과 유사한 방식으로 image-text pair에 InfoNCE loss를 사용해 학습하였다.
댓글 남기기