💻 computer science

Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI

The paper introduces COJEPA, a self-supervised framework that combines contrastive loss with joint-embedding predictive architecture to learn robust 3D brain MRI representations, achieving state-of-the-art performance in twin retrieval, age regression, and tumor segmentation without requiring labeled data.

Original authors: Fabian Mager, Lars Kai Hansen

Published 2026-07-15

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Fabian Mager, Lars Kai Hansen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, 3D puzzle of the human brain, but you've never seen the picture on the box. You don't have a teacher to tell you which piece goes where, and you certainly don't have a label saying "this is a 25-year-old" or "this is a tumor." This is the challenge doctors face with MRI scans: there are millions of them, but very few have the expensive, time-consuming labels needed to teach computers how to understand them.

Enter COJEPA, a new "self-taught" AI model that tries to learn the secrets of the brain just by looking at the raw puzzle pieces.

The Two-Headed Monster

Most AI models trying to learn from unlabeled images usually pick one of two strategies:

The "Fill-in-the-Blank" Artist (JEPA): This model looks at a chunk of the brain, covers up a part of it, and tries to guess what's underneath based on the surrounding context. It's great at understanding local details, like how the texture of the brain's folds connects to the next fold.
The "Look-Alike" Detective (Contrastive Learning): This model takes two slightly different photos of the same brain and tries to make them look identical in its memory, while making sure they look totally different from other people's brains. It's great at figuring out "Who is this?"

The paper argues that doing just one of these isn't enough for 3D brain scans. The "Fill-in-the-blank" artist is too vague to tell twins apart, and the "Look-Alike" detective is too focused on the big picture to understand the tiny, local details of brain structure.

The Main Finding: The authors combined both strategies into one model called COJEPA. They found that this hybrid approach is the sweet spot. It learns to predict local brain details and recognize the unique identity of the person.

How They Taught It

The team trained this model on 2,286 healthy brain scans from two groups: young adults (ages 22–37) and older adults (ages 36–90+). They didn't use any labels. Instead, they used a clever trick:

They masked out (covered up) random 3D blocks of the brain, but they were careful to only cover up the "foreground" (the actual brain tissue), ignoring the empty space around it.
They asked the model to predict the hidden brain parts.
Simultaneously, they showed the model two different views of the same brain and forced it to realize, "Hey, these are the same person!"

The Results: What Did It Learn?

The paper tested this new brain-learner on three different tasks to see how smart it really was.

1. The Twin Test (Zero-Shot Retrieval)
This was the big test. The model had to look at a brain it had never seen before and find its twin in a database of other brains.

The "Fill-in-the-blank" only model failed miserably at finding identical twins (Monozygotic twins), getting the right match only 26% of the time at the top spot. It couldn't tell the difference between people well enough.
The "Look-Alike" only model did much better, getting the match 78% of the time.
COJEPA (The Hybrid) was the champion, getting the right twin 84% of the time at the top spot.
The Catch: The model was surprisingly bad at finding non-identical (Dizygotic) twins, getting the match only 17% of the time. This suggests that while it learned to spot the unique "fingerprint" of a brain, it still struggles with the more subtle genetic similarities between non-identical twins.

2. The Tumor Hunt (Segmentation)
They asked the model to find brain tumors in a dataset it had never seen before (BraTS 2024).

When the model was just "frozen" (not allowed to learn new things, just use what it already knew), the hybrid model and the "Look-Alike" model performed similarly, finding about 52% of the tumor correctly with very little training data.
However, when they gave the model a little bit of fine-tuning (letting it learn from the specific tumor data), all three models got much better, reaching about 71% accuracy.
The Takeaway: The hybrid model didn't hurt the local details; it was just as good at finding tumors as the others, proving it didn't lose its "artist" skills while becoming a "detective."

3. The Age Guess (Regression)
They tried to guess the person's age just by looking at their brain scan.

The hybrid model was the most efficient. With only 5% of the labeled data, it guessed the age with an error of 4.22 years, which was better than the other models.
When fully fine-tuned, it achieved the best result of all: an error of just 2.55 years.
Interestingly, the authors noticed that the part of the brain the model focused on most for guessing age was the cerebellum (the back part of the brain). They suggest this is an "emergent property," meaning the model figured out on its own that this area changes with age, even though no one ever told it to look there.

What the Paper Rules Out

The authors are very clear about what doesn't work as well as their new method:

They explicitly rule out the idea that just predicting missing parts (JEPA) is enough to create a model that can distinguish between different people. The data showed it failed at twin retrieval.
They also suggest that just looking for similarities (Contrastive Learning) might miss some of the fine-grained, local structural details that are crucial for medical tasks like segmentation, though the hybrid model managed to keep those skills.

How Sure Are They?

The authors are confident in their numbers but cautious about the bigger picture.

They measured the twin retrieval and age prediction results directly on test sets, so those numbers (like the 84% twin match and 2.55 year error) are solid facts from their experiment.
However, they suggest that the model's success comes from balancing "predictivity" (guessing the future) and "specificity" (knowing the identity). They don't claim this is the final, perfect solution for all medical AI.
They admit their training set of 2,286 people is actually quite small by modern AI standards. They suspect that if they had a much larger, more diverse group of people to train on, the model would likely get even better.
They also noted a limitation: the model struggled more with scans from a specific older machine (1.5T scanners), suggesting it might not work perfectly on every type of hospital scanner yet.

In short, COJEPA is a promising new way to teach computers to understand the brain without needing a teacher for every single scan. It learned to be both a good guesser of local details and a sharp detective of identity, but it still has room to grow before it can handle every brain scan in the world.

Technical Summary: Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI

Problem Statement

Self-supervised learning (SSL) has revolutionized representation learning in natural images but has yet to achieve comparable impact in 3D structural magnetic resonance imaging (MRI). The biomedical domain faces unique constraints: a scarcity of large, diverse, and shareable datasets due to acquisition costs, privacy, and ethical restrictions; heterogeneity across scanners and protocols; and the computational burden of processing high-resolution volumetric data. Existing brain MRI foundation models typically rely on one of three paradigms: Masked Image Modeling (MIM), which captures local context but incurs cubic reconstruction costs; Contrastive Learning (CL), which offers label efficiency but requires prohibitive batch sizes and struggles with medical invariances; or Self-Distillation, which often sacrifices volumetric context by operating on 2D slices. No existing model effectively combines the local predictivity of MIM with the global discriminability of CL in a 3D volumetric setting.

Methodology

The authors propose COJEPA (Contrastive Joint-Embedding Prediction Architecture), a self-supervised framework designed to learn representations that are simultaneously locally predictive and globally discriminative. The model is trained on T1-weighted structural MRI from two cohorts (HCP-YA and AABC, N=2286) without labels.

Core Architecture

COJEPA unifies two objectives within a Vision Transformer (ViT) framework adapted for 3D volumes:

Joint Embedding Predictive Architecture (JEPA): The model predicts the latent representation of masked target regions from a context view. Unlike pixel-space reconstruction, this operates in representation space, avoiding unnecessary pixel decoding.
Contrastive Loss (CO): An auxiliary objective that enforces global alignment between augmented views of the same subject using an InfoNCE loss on a global token.

Key Technical Adaptations for 3D MRI

Foreground-Aware Block Masking: Extending the block masking strategy from 2D to 3D, the sampling process prioritizes foreground patches (defined by voxel intensity thresholds) to ensure target regions contain rich anatomical information.
Hierarchical Convolutional Patch Embedding: Instead of a single linear projection, the model uses a hierarchical stem of three sequential 3D convolutions. This allows the network to learn local spatial features before projecting into the token space, stabilizing training on volumetric data.
World-Space Sinusoidal Positional Encodings: Leveraging the standardized space of MRI, the model assigns patches global coordinates based on the image affine rather than grid indices. This implicitly encodes voxel resolution and resolves brain symmetries.
Global Token Integration: A global token ( $z_{glob}$ ) is attached to the patch embeddings before the transformer layers. This token interacts with spatial context tokens via multi-head attention and serves as the target for the contrastive loss, while the predictive loss operates on spatial tokens.

Training Setup

The target encoder is updated via an Exponential Moving Average (EMA) of the context encoder. The total loss is a weighted sum of the prediction loss (MSE on foreground tokens) and the contrastive loss (InfoNCE on the global token). The model was trained for 4000 epochs on 8 AMD MI250x GPUs.

Key Contributions

Novel Architecture: The first application of joint-embedding predictive pretraining to 3D brain MRI, extended with a contrastive term to address representation collapse and enhance subject-level discriminability.
3D-Specific Adaptations: Introduction of foreground-aware masking, hierarchical convolutional stems, and world-space positional encodings tailored for volumetric medical data.
Comprehensive Evaluation: A rigorous assessment across three distinct downstream tasks: zero-shot twin retrieval, brain tumor segmentation (BraTS 2024), and age regression (OpenBHB).

Results

The model was evaluated against pure Contrastive (CO) and pure Predictive (JEPA) baselines:

Zero-Shot Twin Retrieval: COJEPA achieved the best performance in retrieving monozygotic (MZ) twins at rank 1 (Recall@1 = 0.84), outperforming pure JEPA (0.26) and pure CO (0.78). This demonstrates that the combined objective successfully captures genetic similarity and subject identity.
Brain Tumor Segmentation: On the out-of-distribution BraTS 2024 dataset, COJEPA matched the performance of pure CO in the frozen probe regime (Dice $\approx$ 0.66–0.67) and achieved comparable results to JEPA after end-to-end finetuning (Dice $\approx$ 0.71). This indicates that the contrastive term does not degrade local structural predictivity.
Age Regression: On the OpenBHB dataset, COJEPA achieved the best finetuning Mean Absolute Error (MAE) of 2.55 years. Notably, it showed superior data efficiency in the frozen probe regime at low label fractions (5% labels) compared to CO.
Feature Space Analysis: Principal Component Analysis (PCA) of the embeddings revealed that COJEPA utilizes the latent space more efficiently (flatter cumulative explained variance) than CO or JEPA, with clear separation of anatomical structures.

Significance and Claims

The authors claim that COJEPA successfully bridges the gap between two complementary properties of useful representations: local predictivity (capturing statistical dependencies of local anatomical structure) and global specificity (subject-level discriminability).

The paper argues that while pure predictive objectives (JEPA) capture sufficient structure for coarse identity retrieval (e.g., MZ twins), they fail to produce sufficiently discriminative global features for finer distinctions (e.g., DZ twins). Conversely, while pure contrastive objectives (CO) excel at global alignment, they may lack the fine-grained local structure required for tasks like segmentation without extensive finetuning. By combining these objectives, COJEPA yields representations that are simultaneously discriminative and locally structured.

The authors note that the model achieves competitive downstream performance despite being trained on a relatively small dataset (N=2286) compared to typical SSL scales, suggesting the efficiency of the proposed framework. They conclude that the learned representations capture morphological variations relevant to brain aging and genetic similarity as emergent properties of the pretraining objective, without requiring explicit demographic or genetic supervision.