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

1. Problem Statement

Despite the success of foundation models in natural language and 2D computer vision (e.g., GPT, CLIP, DINO), 3D medical imaging lacks a comparable foundation model. Existing approaches face two primary limitations:

• Inadequate Pretext Tasks: Current self-supervised learning (SSL) methods rely on low-level reconstruction (e.g., Masked Autoencoders) or contrastive objectives. These fail to capture the anatomical semantics and spatial precision critical for medical analysis, where every voxel often carries clinical significance (e.g., distinguishing tumor texture from healthy tissue).
• Annotation Bottleneck: Supervised pretraining requires massive, expert-annotated datasets. However, annotations are confined to predefined taxonomies, limiting generalization to novel anatomical structures or pathologies. Furthermore, pixel-level delineation of 3D volumes is prohibitively expensive and time-consuming, preventing scalability.

The core challenge is to learn dense, spatially precise, and semantically rich representations across the full spectrum of anatomical and pathological concepts without expert annotations.

2. Methodology: MASS (MAsk-guided Self-Supervised learning)

The authors propose MASS, a novel framework that treats in-context segmentation as the pretext task, guided by automatically generated, class-agnostic masks.

A. Annotation-Free Mask Generation

Instead of relying on human labels, MASS uses SAM2 (Segment Anything Model 2) to generate structural supervision:

1. Preprocessing: Single-channel medical volumes (CT/MRI/PET) are converted to 3-channel inputs using complementary intensity windows (e.g., soft tissue, bone, contrast windows) to mimic RGB inputs for SAM2.
2. 2D Sampling: 2D slices are sampled along the highest-resolution axis.
3. Automatic Masking: SAM2 generates hundreds of class-agnostic region proposals per slice based on intensity boundaries and morphological coherence.
4. 3D Propagation: Using SAM2's video prediction capability, 2D masks are propagated bidirectionally through the volume to create 3D structural masks.

• Key Insight: While these masks lack semantic labels (they don't know "this is a liver") and contain noise, they capture anatomically and pathologically meaningful regions (organs, vessels, lesions) across diverse granularities.

B. In-Context Segmentation (ICS) Framework

MASS adopts an ICS paradigm similar to the Iris architecture, consisting of three components:

1. Image Encoder ($E_\theta$): Extracts features from the 3D volume.
2. Task Encoding Module ($T_\phi$): Takes a reference image-mask pair $(x_s, y_s)$ and generates a task embedding. This embedding encodes "what" structure to segment based on the visual appearance of the reference mask.
3. Mask Decoder ($D_\psi$): Takes a query image $(x_q)$ and the task embedding to predict the corresponding mask $(\hat{y}_q)$.

Training Objective:
The model is trained to predict the query mask given the query image and the reference task embedding. The loss function combines Dice loss and Binary Cross-Entropy:
$$\mathcal{L}_{Seg} = \mathcal{L}_{Dice}(\hat{y}_q, y_q) + \mathcal{L}_{BCE}(\hat{y}_q, y_q)$$

• Learning Mechanism: By forcing the model to map a reference mask to a query view under heavy data augmentation (spatial and appearance), the model learns invariance. It must identify the fundamental semantic identity of a structure (shape, texture, spatial context) that remains constant despite radical visual changes, effectively learning "what" defines a medical structure without ever seeing a semantic label.

3. Key Contributions

1. Novel Pretext Task: Introduces a segmentation-based self-supervised framework that aligns pretraining objectives with downstream clinical tasks (segmentation and classification), moving beyond low-level reconstruction.
2. Annotation-Free Scalability: Demonstrates that weak, class-agnostic masks generated by foundation models (SAM2) are sufficient to learn robust representations. This enables pretraining on 5,000+ diverse 3D volumes (CT, MRI, PET) without a single human annotation.
3. Strong Few-Shot Capabilities: MASS achieves state-of-the-art performance with minimal labeled data. It can perform zero-shot in-context inference (no fine-tuning) and matches full-supervision performance with only 20-40% of labeled data.
4. Broad Generalization: The learned representations transfer effectively to unseen pathologies and classification tasks, outperforming existing SSL baselines by significant margins (e.g., >20 Dice points in low-data regimes).

4. Experimental Results

The authors evaluated MASS across small-scale (single dataset) and large-scale (5K multi-modal volumes) pretraining scenarios.

• Segmentation Performance (Few-Shot):
  • In-Context Inference (No Fine-tuning): MASS-IC outperformed all baselines (Scratch, Model Genesis, MAE, SimCLR) in 1-shot settings on anatomical structures (e.g., BCV, AMOS MR).
  • Fine-tuning: MASS-FT achieved the best performance across all datasets. On the BCV dataset with only 1 shot, it outperformed the best baseline by 21.7 Dice points.
  • Data Efficiency: With only 10 shots (25-40% of training data), MASS matched the performance of full supervised training.
• Classification (Frozen Encoder):
  • Tested on unseen pathologies (RSNA Intracranial Hemorrhage, Trauma Detection).
  • MASS's frozen encoder matched or exceeded full supervised training from scratch using only 5-30% of the labeled data. For example, on RSNA ICH with 5% data, MASS achieved 75.4% AUC vs. 72.8% for training from scratch.
• Comparison with Supervised Pretraining:
  • While supervised models (SuPreM, Iris) excel when evaluation targets match their pretraining taxonomies, MASS demonstrated superior out-of-distribution generalization (e.g., on Pelvic and BraTS datasets) where supervised models struggled due to taxonomy mismatches.
• Ablation Studies:
  • Mask Quality: MASS learns effectively even when auto-generated masks have low overlap with ground truth (average Dice ~15%), proving that diverse, approximate supervision is sufficient.
  • Data Diversity: Performance gains were driven more by anatomical and modality diversity (CT, MRI, PET, whole-body) than raw data quantity.

5. Significance and Impact

• Democratizing Foundation Models: MASS provides a practical path to building 3D medical imaging foundation models without the prohibitive cost of expert annotation. It shifts the bottleneck from "labeling data" to "scaling unlabeled data."
• Semantic Understanding without Labels: The work demonstrates that semantic understanding (knowing what an organ or tumor is) can emerge implicitly from invariance constraints applied to weak structural supervision, rather than requiring explicit semantic labels.
• Clinical Applicability: The ability to adapt to novel structures and pathologies with minimal fine-tuning makes MASS highly suitable for real-world clinical environments where diseases and anatomical variations are unpredictable and expert annotations are scarce.
• Future Direction: The paper opens avenues for aligning these visual representations with textual modalities (radiology reports) to create true vision-language foundation models for medicine.

In summary, MASS successfully bridges the gap between self-supervised learning and clinical utility in 3D medical imaging by leveraging the structural priors of foundation models (SAM2) to learn generalizable, semantically rich representations without human supervision.