[Paper Review] LVM-Med: Learning Large-Scale Self-Supervised Vision Models for Medical Imaging via Second-order Graph Matching

Duy M. H. Nguyen, Hoang Nguyen+10 more|arXiv (Cornell University)|Jun 20, 2023

Radiomics and Machine Learning in Medical Imaging18 citations

TL;DR

LVM-Med introduces a large-scale self-supervised medical imaging model trained on ~1.3 million images from 55 public datasets, using a novel second-order graph-matching objective to learn robust representations that outperform several SSL and foundation models on 15 downstream tasks.

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.

Motivation & Objective

• Motivate the need for large-scale, domain-specific self-supervised learning in medical imaging due to domain shifts from natural images.
• Propose a new SSL framework (LVM-Med) that leverages second-order graph matching to learn robust representations.
• Create a large, diverse medical image dataset (~1.3M images from 55 public datasets) to benchmark SSL methods in medicine.
• Demonstrate state-of-the-art performance across 15 downstream tasks including segmentation, classification, and detection, in in- and out-of-distribution settings.

Proposed method

• Construct two distorted views per image and encode them with a shared backbone to obtain embeddings.
• Build two graphs per batch where nodes represent distorted views and edges encode local/global affinities.
• Define vertex affinities using global cosine similarity and local, region-aware costs, combined into a unified affinity c^v.
• Introduce second-order graph matching with edge affinities c^e to capture relational structure between matched pairs.
• Solve the graph matching problem with a combinatorial objective and learn gradients through IMLE-based gradient estimation for end-to-end training.
• Backpropagate through the discrete solver by perturbing costs with Gumbel noise and estimating gradients via a finite-difference IMLE scheme.

Experimental results

Research questions

• RQ1 Can a second-order graph-matching SSL objective improve representation learning for medical imaging compared to traditional pairwise contrastive losses?
• RQ2 Does integrating global and local affinity information into a graph-based SSL framework yield robust, transferable features across diverse medical modalities and tasks?
• RQ3 How does LVM-Med perform on 15 downstream tasks under in- and out-of-distribution conditions compared to supervised, SSL, and foundation models?
• RQ4 Is it feasible to train large-scale medical SSL models efficiently on multi-modal, publicly available datasets using gradient estimation for black-box solvers?

Key findings

• LVM-Med consistently outperforms several state-of-the-art supervised, self-supervised, and foundation models across 15 medical tasks.
• On Brain Tumor Classification and Diabetic Retinopathy Grading, LVM-Med improves previous vision-language models trained on 1B masks by 6–7 percentage points using only a ResNet-50 backbone.
• The second-order graph-matching formulation, with both vertex and edge affinities, yields robust improvements over pure linear (pairwise) matching methods.
• LVM-Med with ResNet-50 and with SAM’s ViT backbone achieve strong results on both 2D and 3D segmentation tasks, often surpassing SAM-based prompt settings.
• The approach scales to large datasets and can be trained end-to-end using IMLE-based gradient estimation, despite the combinatorial nature of graph matching.