[Paper Review] Towards Foundation Models Learned from Anatomy in Medical Imaging via Self-Supervision

Mohammad Reza Hosseinzadeh Taher, Michael B. Gotway+1 more|arXiv (Cornell University)|Sep 27, 2023

AI in cancer detectionComputer Science3 citations

TL;DR

This paper proposes Adam, a self-supervised foundation model for medical imaging that learns anatomical hierarchies through a coarse-to-fine, hierarchical contrastive learning strategy. By encoding locality and compositionality of anatomical structures into dense embeddings (Eve), Adam achieves state-of-the-art few-shot segmentation performance, improving accuracy by 9–30% over prior SSL methods while enhancing annotation efficiency.

ABSTRACT

Human anatomy is the foundation of medical imaging and boasts one striking characteristic: its hierarchy in nature, exhibiting two intrinsic properties: (1) locality: each anatomical structure is morphologically distinct from the others; and (2) compositionality: each anatomical structure is an integrated part of a larger whole. We envision a foundation model for medical imaging that is consciously and purposefully developed upon this foundation to gain the capability of "understanding" human anatomy and to possess the fundamental properties of medical imaging. As our first step in realizing this vision towards foundation models in medical imaging, we devise a novel self-supervised learning (SSL) strategy that exploits the hierarchical nature of human anatomy. Our extensive experiments demonstrate that the SSL pretrained model, derived from our training strategy, not only outperforms state-of-the-art (SOTA) fully/self-supervised baselines but also enhances annotation efficiency, offering potential few-shot segmentation capabilities with performance improvements ranging from 9% to 30% for segmentation tasks compared to SSL baselines. This performance is attributed to the significance of anatomy comprehension via our learning strategy, which encapsulates the intrinsic attributes of anatomical structures-locality and compositionality-within the embedding space, yet overlooked in existing SSL methods. All code and pretrained models are available at https://github.com/JLiangLab/Eden.

Motivation & Objective

• To develop a foundation model for medical imaging that explicitly learns from the intrinsic anatomical hierarchy—locality and compositionality—rather than relying on generic self-supervision.
• To address the gap in self-supervised learning for medical imaging, where current methods fail to capture the structural and semantic foundations of human anatomy.
• To improve few-shot segmentation and annotation efficiency by training a model that generates semantically meaningful, anatomy-aware embeddings.
• To validate the model’s anatomical understanding through novel evaluation metrics measuring locality and compositionality in the embedding space.

Proposed method

• A hierarchical contrastive learning framework that progressively decomposes anatomical structures from coarse to fine granularity during pretraining.
• The use of an anatomy decomposer to generate multi-scale patches from medical images, enabling coarse-to-fine learning of anatomical parts.
• A purposive pruner that filters out noisy negative pairs in the contrastive memory bank, improving feature discriminability.
• Training with a contrastive objective that encourages identical anatomical structures across patients and scales to have near-identical embeddings.
• The model, named Adam, produces dense, semantics-rich embeddings (Eve) that preserve both locality and compositionality in the embedding space.
• The framework is evaluated across 9 downstream tasks and adapted to fundoscopy images, demonstrating generalizability across modalities.

Experimental results

Research questions

• RQ1 Can self-supervised learning be restructured to explicitly learn the hierarchical nature of human anatomy in medical imaging?
• RQ2 Does encoding anatomical locality and compositionality into learned embeddings improve downstream performance in medical image segmentation?
• RQ3 To what extent can a self-supervised model achieve few-shot segmentation performance comparable to fully supervised models without explicit anatomy supervision?
• RQ4 How does the coarse-to-fine learning strategy affect the model’s ability to capture anatomical structure representations?
• RQ5 Can the proposed method generalize across different medical imaging modalities, such as chest X-rays and retinal fundoscopy?

Key findings

• Adam outperformed state-of-the-art self-supervised and fully supervised baselines in few-shot segmentation, with performance gains ranging from 9% to 30% across multiple datasets.
• The embedding space of Adam demonstrated strong locality, with distinct anatomical structures forming well-separated clusters in t-SNE visualization, unlike baseline models.
• Adam’s embeddings preserved compositionality, showing a mean cosine similarity of nearly 1 between whole anatomical patches and the aggregate of their sub-patches, significantly outperforming baselines.
• Ablation studies confirmed that increasing data granularity in the coarse-to-fine learning strategy consistently improved downstream performance, validating the hierarchical approach.
• The purposive pruner significantly enhanced model performance by reducing noisy contrastive pairs, confirming its role in improving feature discriminability.
• The method generalized to fundoscopy images, achieving a 1.4% improvement in blood vessel segmentation over top-performing SSL baselines on the EyePACS dataset (p < 0.05).