Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI

2026-07-12

AI Summary

This work addresses the challenge of scarce and costly annotations in medical imaging by proposing COJEPA, a novel framework that extends I-JEPA to 3D brain MRI for the first time, integrating contrastive learning with a joint-embedding predictive architecture. The method enhances self-supervised representation learning on unlabeled T1-weighted images by simultaneously optimizing local structural predictability and global representational discriminability through foreground-aware block masking, hierarchical convolutional block embeddings, and world-space sinusoidal positional encoding. Experimental results demonstrate strong performance: the model achieves a rank@1 retrieval accuracy of 0.84 in zero-shot monozygotic twin identification, a mean absolute error of 2.55 years in age regression on OpenBHB, and competitive whole-tumor segmentation Dice scores on BraTS, matching state-of-the-art supervised methods.

Abstract

Self-supervised learning offers a compelling approach for medical imaging, where labeled data are scarce and acquisition costs are high. We present COJEPA, a self-supervised framework for volumetric brain MRI that combines a joint-embedding predictive architecture (JEPA) with a contrastive loss (CO), targeting two complementary properties: local predictivity and global discriminability. The model is trained without labels on T1-weighted structural MRI from two cohorts (HCP-YA and AABC, $N{=}2286$, ages 22 to 90), extending I-JEPA to 3D with foreground-aware block masking, a hierarchical convolutional patch embedding, and world-space sinusoidal positional encodings. We evaluate all three objectives across zero-shot twin retrieval, brain tumor segmentation (BraTS 2024), and age regression (OpenBHB). COJEPA achieves the best monozygotic twin recall at rank@1 (0.84), the best finetuning age MAE (2.55 years on OpenBHB 3.0T), and matches CO on BraTS whole-tumor Dice, demonstrating that the combined objective yields representations that are simultaneously discriminative and locally structured.

Problem

• self-supervised learning
• structural MRI
• representation learning
• contrastive learning
• joint-embedding prediction

Innovation

• contrastive learning
• joint-embedding prediction
• self-supervised representation learning
• 3D structural MRI
• foreground-aware masking