[논문 리뷰] LVM-Med: Learning Large-Scale Self-Supervised Vision Models for Medical Imaging via Second-order Graph Matching

Duy M. H. Nguyen, Hoang Nguyen외 10명|arXiv (Cornell University)|2023. 06. 20.

Radiomics and Machine Learning in Medical Imaging인용 수 18

한 줄 요약

LVM-Med는 ~1.3백만 장의 이미지와 55개 공공 데이터셋에서 학습된 대규모 자기지도 의학 영상 모델을 도입하고, 새로운 2차 그래프 매칭 목표를 사용하여 강건한 표현을 학습하며 15개의 다운스트림 태스크에서 여러 SSL 및 파운데이션 모델을 능가한다.

ABSTRACT

Obtaining large pre-trained models that can be fine-tuned to new tasks with limited annotated samples has remained an open challenge for medical imaging data. While pre-trained deep networks on ImageNet and vision-language foundation models trained on web-scale data are prevailing approaches, their effectiveness on medical tasks is limited due to the significant domain shift between natural and medical images. To bridge this gap, we introduce LVM-Med, the first family of deep networks trained on large-scale medical datasets. We have collected approximately 1.3 million medical images from 55 publicly available datasets, covering a large number of organs and modalities such as CT, MRI, X-ray, and Ultrasound. We benchmark several state-of-the-art self-supervised algorithms on this dataset and propose a novel self-supervised contrastive learning algorithm using a graph-matching formulation. The proposed approach makes three contributions: (i) it integrates prior pair-wise image similarity metrics based on local and global information; (ii) it captures the structural constraints of feature embeddings through a loss function constructed via a combinatorial graph-matching objective; and (iii) it can be trained efficiently end-to-end using modern gradient-estimation techniques for black-box solvers. We thoroughly evaluate the proposed LVM-Med on 15 downstream medical tasks ranging from segmentation and classification to object detection, and both for the in and out-of-distribution settings. LVM-Med empirically outperforms a number of state-of-the-art supervised, self-supervised, and foundation models. For challenging tasks such as Brain Tumor Classification or Diabetic Retinopathy Grading, LVM-Med improves previous vision-language models trained on 1 billion masks by 6-7% while using only a ResNet-50.

연구 동기 및 목표

• 자연 이미지로부터의 도메인 차이로 인해 의학 영상에서 대규모의 도메인 특화 자가감독 학습의 필요성을 제시한다.
• 두 번째 차수의 그래프 매칭을 활용하여 강건한 표현을 학습하는 새로운 SSL 프레임워크(LVM-Med)를 제안한다.
• 의학에서 SSL 방법을 벤치마킹하기 위해 대규모의 다양하고 균형된 의료 영상 데이터셋(~1.3M 이미지, 55개 공공 데이터셋)을 만든다.
• 세분화/분할(segmentation), 분류(classification), 탐지(detection) 등을 포함한 15개의 다운스트림 태스크에서, in- 및 out-of-distribution 설정에서 최첨단 성능을 입증한다.

제안 방법

• 이미지당 두 개의 왜곡된 뷰를 구성하고 공유 백본으로 인코딩하여 임베딩을 얻는다.
• 배치당 두 개의 그래프를 구축하여 노드가 왜곡된 뷰를 나타내고 엣지가 로컬/전역 친화도를 인코딩한다.
• 전역 코사인 유사도와 로컬의 영역 인지 비용을 사용하여 정점 친화도를 정의하고 이를 결합하여 통합 친화도 c^v를 얻는다.
• 일치된 쌍 간의 관계 구조를 포착하기 위해 엣지 친화도 c^e를 갖는 2차 그래프 매칭을 도입한다.
• 조합적 목적 함수를 사용하여 그래프 매칭 문제를 해결하고 IMLE 기반 기울기 추정을 통해 엔드-투-엔드 학습을 위한 그래디언트를 학습한다.
• 비용에 Gumbel 잡음을 섞어 이산 해를 가진 해를 역전파하고, 유한 차 IMLE 방식으로 그래디언트를 추정한다.

실험 결과

연구 질문

• RQ1 두 번째 차수 그래프 매칭 SSL 목표가 전통적인 쌍대 대비 losses에 비해 의학 영상의 표현 학습을 향상시킬 수 있는가?
• RQ2 전역 및 로컬 친화도 정보를 그래프 기반 SSL 프레임워크에 통합하면 다양한 의학 모달리티와 태스크에서 강건하고 전이 가능한 특징을 얻을 수 있는가?
• RQ3 LVM-Med는 15개의 다운스트림 태스크에서 supervised, SSL, 및 파운데이션 모델과 비교하여 in- 및 out-of-distribution 조건에서 어떤 성능을 보이는가?
• RQ4 블랙박스 솔버에 대한 그래디언트 추정을 사용하여 다모달 공개 데이터셋에서 대규모 의학 SSL 모델을 효율적으로 학습하는 것이 가능한가?

주요 결과

• LVM-Med가 15개의 의학 태스크에서 여러 최첨단 감독 모델, 자기지도 학습 모델, 파운데이션 모델을 일관되게 능가한다.
• Brain Tumor Classification 및 Diabetic Retinopathy Grading에서 LVM-Med는 1B 마스크로 학습된 이전 비전-언어 모델보다 6–7 포인트의 성능 향상을 보이며 단지 ResNet-50 백본만으로도 달성한다.
• 정점 및 엣지 친화도를 모두 포함한 2차 그래프 매칭 형식은 순수 선형(쌍대) 매칭 방법에 비해 강건한 개선을 보인다.
• ResNet-50 및 SAM의 ViT 백본을 사용하는 LVM-Med는 2D 및 3D 세분화 태스크에서 강력한 결과를 보이며 SAM 기반 프롬프트 설정을 종종 능가한다.
• 이 접근법은 큰 데이터셋으로 확장 가능하며 그래프 매칭의 조합적 특성에도 불구하고 IMLE 기반 기울기 추정을 사용해 엔드-투-엔드 학습이 가능하다。）