Learning Generalizable 3D Medical Image Representations from Mask-Guided Self-Supervision

TL;DR

MASS introduces a mask-guided self-supervised learning approach for 3D medical images, enabling the development of general-purpose representations that perform well on various downstream tasks without requiring annotations.

Contribution

The paper proposes MASS, a novel self-supervised method using class-agnostic masks to learn semantically rich 3D medical image representations, improving transferability and reducing annotation needs.

Findings

Effective in few-shot segmentation of new structures

Achieves comparable results to fully supervised methods with limited labeled data

Outperforms existing self-supervised baselines by over 20 in Dice score in low-data regimes

Abstract

Foundation models have transformed vision and language by learning general-purpose representations from large-scale unlabeled data, yet 3D medical imaging lacks analogous approaches. Existing self-supervised methods rely on low-level reconstruction or contrastive objectives that fail to capture the anatomical semantics critical for medical image analysis, limiting transfer to downstream tasks. We present MASS (MAsk-guided Self-Supervised learning), which treats in-context segmentation as the pretext task for learning general-purpose medical imaging representations. MASS's key insight is that automatically generated class-agnostic masks provide sufficient structural supervision for learning semantically rich representations. By training on thousands of diverse mask proposals spanning anatomical structures and pathological findings, MASS learns what semantically defines medical structures: the holistic combination of appearance, shape, spatial context, and anatomical relationships. We demonstrate effectiveness across data regimes: from small-scale pretraining on individual datasets (20-200 scans) to large-scale multi-modal pretraining on 5K CT, MRI, and PET volumes, all without annotations. MASS demonstrates: (i) few-shot segmentation on novel structures, (ii) matching full supervision with only 20-40\% labeled data while outperforming self-supervised baselines by over 20 in Dice score in low-data regimes, and (iii) frozen-encoder classification on unseen pathologies that matches full supervised training with thousands of samples. Mask-guided self-supervised pretraining captures broadly generalizable knowledge, opening a path toward 3D medical imaging foundation models without expert annotations. Code is available: https://github.com/Stanford-AIMI/MASS.