DINO Contrastive Learning in Medical Imaging

DINO Algorithm

 

DINO Algorithm (Multi-view cropper)

 

DINO Algorithm을 간단하게 얘기하면,

이미지에서 큰 부분 (Global view)와 작은부분 (Local view)으로 이미지를 떼어내어 준비한다.

 

DINO Algorithm Model Architecture

 

그리고 Global view는 teacher에게 보여주고,

Global + Local view를 student에게 보여줌으로써, 작은부분만 보고도

넓은 부분의 feature와 동일한 정보를 추출하도록 하는 것이다.

 

이렇게 계속 학습하다 보면, 이미지의 일부분만 보고도 전체이미지의 특징을 잡게 되므로

이미지의 주요한 부분에 집중하게 되고, 좋은 feature를 추출하게 된다.

 

바로 Label 없이 학습하는 self-supervised learning이다.

 

 

 

Stanford의 Artificial Intelligence in Medicine & Imaging (AIMI) 그룹에서도

이러한 방식으로 foundation model을 만드는 것에 대해서 깊히 고민하고 있었다.

 

 

DINO in Medical Imaging

 

이 DINO 방법론을 활용해서, 공개데이터베이스 및 교내 데이터를 활용해서 학습해보았다.

 

총 이미지는 638,302개 이다.

모든 이미지를 활용하여, 480x480 이미지를 default input으로,

Crop view 사이즈는 Gloval / Local에 대하여 (224, 224) / (96, 96),

총 4개의 NVIDIA RTX 4090 GPU를 활용하여, 72개의 batch_size로 학습해봤다.

 

1 epoch당 대략 1시간정도가 걸렸다.

 

 

Example in Brain Tumor Image

 

Malignant 병변이 명확하게 보이는 하나의 Brain tumor image를 기반으로,

Epoch에 따라서 Visition transformer의 Attention을 heatmap으로 나타내보았다.

 

 

epoch가 진행될 수록, 한 이미지에서 환자의 특징을 가장 잘 나타내는 Lesion에 집중을 하는 것을 볼 수 있다.

 

AI모델이 Local / Global view에서 동일한 feature를 뽑게 하는 것은,

결국 view에 따라서 마구잡이로 바뀌는 background noisy가 아닌

사진에서 가장 중요한 object를 집중하게 만들게 된다.

 

비록 아무런 Label 없이 이미지만 학습한 것이지만, 훌륭히 중요위치를 살펴본다.

 

Normal brain에 대한 heatmap!

 

 

 

많은 곳에서 foundation model의 역할을 하는 Back-bone model을 제작중에 있다.

 

self-supervised, 그리고 Contrastive learning는

Radiology, Pathology 등 모든 곳에서 엄청난 역할을 할 것으로 기대된다.