Beyond Representation Learning: A Systematic Study of Joint-Embedding Predictive Generation for 3D Brain MRI

Meng Zhou; Wenhao You; Yuxing Chen; Yueying Tian

cs.CVAug 28, 2026

Beyond Representation Learning: A Systematic Study of Joint-Embedding Predictive Generation for 3D Brain MRI

Authors: Meng Zhou, Wenhao You, Yuxing Chen, Yueying Tian

Organizations: University of Toronto · University of Waterloo · University of Alberta · University of Sussex

Abstract

Joint-embedding predictive architectures (JEPAs) have primarily been developed for self-supervised representation learning. Denoising JEPA (D-JEPA) recently demonstrated strong generative capabilities on natural images, yet the applicability to 3D medical imaging remains unexplored. Building on the D-JEPA framework, we present Med-D-JEPA, a systematic adaptation and evaluation of joint-embedding predictive generation for 3D brain MRI. Med-D-JEPA operates on continuous latent tokens produced by a 3D KL-regularized adversarial variational autoencoder, and combines masked context prediction, representation-level alignment, per-token diffusion, and iterative next-set-of-token sampling. We evaluate unconditional and class-conditional generation quality on BraTS2019 and OASIS-1 datasets; downstream classification utility; and preliminary whole-tumor segmentation on BraTS2020. Across different generation settings, Med-D-JEPA achieves superior or competitive performance compared to several strong baselines on fidelity and diversity metrics. Compared to training with real samples, Med-D-JEPA-based synthetic pretraining improves classification AUC from 0.63 to 0.85 on BraTS2019 and from 0.78 to 0.87 on OASIS-1. In the segmentation study, pretraining on Med-D-JEPA samples improves Dice from 0.74 to 0.80 and reduces HD95 from 13.40 to 9.56 mm. These findings establish joint-embedding predictive generation as a promising direction for 3D medical image synthesis and encourage further research in this direction.

Figures & tables

Figure 1: Synthetic MRI volumes generated by Med-D-JEPA. Left to right: orthogonal views of a $128^{3}$ volume on BraTS2019, a $128^{3}$ volume on OASIS-1 a $256^{3}$ volume on BraTS2020.

Figure 2: Overview of Med-D-JEPA. A 3D KL-VAE-GAN first compresses MRI volumes into continuous latent tokens. Med-D-JEPA then jointly optimizes a JEPA prediction loss and a per-token diffusion loss, with optional class conditioning. During inference, it progressively samples sets of latent tokens and decodes the completed 3D latent grid into a synthetic MRI volume. “reps.” denotes representations, hatched cells indicate masked locations, and orange highlights an example target location.

Figure 3: Qualitative comparison of real samples and samples generated by different methods on the BraTS2019 dataset. For each 3D volume, we show the central axial, coronal, and sagittal slices. Best viewed zoomed in.

(a) Unconditional
Type	Method	MMD $\downarrow$	MS-SSIM (%)	FID-Avg. $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
Mixed	3D-HA-GAN [ 47 ]	2.08	95.7 $(+9.0)$	90.0	0.23 $\pm$ 0.03	0.20 $\pm$ 0.05
	Medical Diffusion [ 29 ]	1.79	90.9 $(+4.2)$	33.4	0.45 $\pm$ 0.02	0.42 $\pm$ 0.04
	3D-MedDiffusion [ 51 ]	1.76	89.3 $(+2.6)$	33.9	0.45 $\pm$ 0.04	0.50 $\pm$ 0.02
	3D-VQGAN [ 58 ]	1.78	88.1 $(+1.4)$	32.9	0.47 $\pm$ 0.06	0.53 $\pm$ 0.09
	Med-D-JEPA (Ours)	1.68	85.4 $(-1.3)$	26.5	0.58 $\pm$ 0.03	0.57 $\pm$ 0.07

Table 1: Quantitative comparison of unconditional and LGG/HGG-conditioned generation on BraTS2019. The best and second-best results within each evaluation setting are shown in blue and green , respectively. For MS-SSIM, values in () denote the difference from the real-data score: 86.7 for mixed, 85.3 for LGG, and 88.6 for HGG.

Figure 4: Qualitative comparison of real samples and samples generated by different methods on the OASIS-1 dataset. For each 3D volume, we show the central axial, coronal, and sagittal slices. Best viewed zoomed in.

(a) Unconditional
Cohort	Method	MMD $\downarrow$	MS-SSIM (%)	FID-Avg. $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
Mixed	3D-HA-GAN [ 47 ]	1.82	82.1 $(+4.9)$	72.0	0.41 $\pm$ 0.03	0.35 $\pm$ 0.07
	Medical Diffusion [ 29 ]	1.58	80.2 $(+3.0)$	27.0	0.65 $\pm$ 0.09	0.42 $\pm$ 0.03
	3D-MedDiffusion [ 51 ]	1.60	75.2 $(-2.0)$	29.8	0.61 $\pm$ 0.06	0.51 $\pm$ 0.05
	3D-VQGAN [ 58 ]	1.50	81.2 $(+4.0)$	26.7	0.72 $\pm$ 0.05	0.53 $\pm$ 0.04
	Med-D-JEPA (Ours)	1.46	74.4 $(-2.8)$	23.6	0.71 $\pm$ 0.07	0.53 $\pm$ 0.01

Table 2: Quantitative comparison of unconditional and CN/AD-conditioned generation on OASIS-1. The best and second-best results within each evaluation setting are shown in blue and green , respectively. For MS-SSIM, values in () denote the difference from the real-data score: 77.2 for the mixed, 80.8 for CN, and 76.3 for AD.

Method comparison
Method	$N_{\mathrm{syn}}$ /class	AUC $\uparrow$	F1-Score $\uparrow$	Accuracy $\uparrow$
Real only	0	$0.63\pm 0.12$	$0.51\pm 0.23$	$0.55\pm 0.08$
Real + traditional augmentation	0	$0.64{\pm}0.02$	$0.62{\pm}0.02$	$0.58{\pm}0.04$
3D-HA-GAN [ 47 ]	300	$0.62{\pm}0.05$	$0.66{\pm}0.07$	$0.61{\pm}0.06$
Medical Diffusion [ 29 ]	300	$0.77{\pm}0.04$	$0.67{\pm}0.02$	$0.69{\pm}0.05$
3D-MedDiffusion [ 51 ]	300	0.79 $\pm$ 0.03	0.68 $\pm$ 0.04	0.70 $\pm$ 0.04

Table 3: Classification performance on BraTS2019. The top section shows the comparison of the performance of different methods. The bottom section shows the performance of Med-D-JEPA on different numbers of pretraining samples. Results are reported as mean

\pm

standard deviation over three runs. Best and second-best results are shown in blue and green , respectively.

Method comparison
Method	$N_{\mathrm{syn}}$ /class	AUC $\uparrow$	F1-Score $\uparrow$	Accuracy $\uparrow$
Real only	0	0.78 $\pm$ 0.08	0.60 $\pm$ 0.06	0.69 $\pm$ 0.03
Real + traditional augmentation	0	0.79 $\pm$ 0.04	0.68 $\pm$ 0.07	0.73 $\pm$ 0.05
3D-HA-GAN [ 47 ]	300	0.77 $\pm$ 0.06	0.68 $\pm$ 0.02	0.74 $\pm$ 0.06
Medical Diffusion [ 29 ]	300	0.80 $\pm$ 0.04	0.73 $\pm$ 0.09	0.75 $\pm$ 0.03
3D-MedDiffusion [ 51 ]	300	0.81 $\pm$ 0.03	0.74 $\pm$ 0.04	0.73 $\pm$ 0.03

Table 4: Classification performance on OASIS-1. The top section shows the comparison of the performance of different methods. The bottom section shows the performance of Med-D-JEPA on different numbers of pretraining samples. Results are reported as mean

\pm

standard deviation over three runs. Best and second-best results are shown in blue and green , respectively.

Figure 5: Examples of the NV-Segment-CTMR segmentation results on the synthetic volume.

Training strategy	$N_{\mathrm{syn}}$	Dice $\uparrow$	HD95 (mm) $\downarrow$
Real only	0	0.74 $\pm$ 0.03	13.40 $\pm$ 5.61
Med-D-JEPA + Real	300	0.80 $\pm$ 0.06	9.56 $\pm$ 3.97

Table 5: Whole-tumor segmentation performance on BraTS2020. Results are reported as mean

\pm

standard deviation over three runs. Best results are shown in blue .

Dataset	Steps	FID-Avg. $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Time/Volume $\downarrow$
BraTS2019	16	26.2	0.51	0.52	5.90
	32	29.9	0.53	0.57	11.78
	64	24.6	0.57	0.58	20.60
	512	28.7	0.55	0.57	185.79
OASIS-1	16	23.3	0.65	0.51	5.88
	32	23.2	0.66	0.55	12.71

Table 6: Ablation on the number of autoregressive (AR) sampling steps. Results are based on 100 synthetic samples.

Variant	BraTS2019			OASIS-1
Variant	FID-Avg. $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$	FID-Avg. $\downarrow$	Prec. $\uparrow$	Rec. $\uparrow$
w/o $\mathcal{L}_{\mathrm{pred}}$	28.1	0.54	0.52	25.2	0.65	0.50
Med-D-JEPA	26.5	0.58	0.57	23.6	0.71	0.53

Table 7: Ablation of the JEPA prediction loss.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: More visualizations on the synthetic volumes from Med-D-JEPA across three datasets.

Figure 7: TSNE Visualization. Top: BraTS2019 HGG and LGG results. Bottom: OASIS-1 CN and AD results.

Figure 8: Reconstruction results of the trained 3D KL-VAE-GAN models on the held-out test sets of BraTS2020, OASIS-1, and BraTS2019. For each dataset, the columns show the central slice of the input volume, its reconstruction, and the absolute-error map, respectively.

Dataset	$N_{\mathrm{test}}$	PSNR	3D SSIM
BraTS2020	74	30.63 $\pm$ 0.89	0.91 $\pm$ 0.03
OASIS-1	60	30.70 $\pm$ 0.99	0.92 $\pm$ 0.05
BraTS2019	50	33.27 $\pm$ 0.10	0.95 $\pm$ 0.04

Appendix

Table 8: Reconstruction performance of the separately trained 3D KL-VAE-GAN models on their corresponding held-out test sets. Results are reported as mean

\pm

standard deviation across test volumes.

Tumor type	Method	MMD ( $\times 10^{4}$ ) $\downarrow$	MS-SSIM (%)	FID-A $\downarrow$	FID-C $\downarrow$	FID-S $\downarrow$	Precision $\uparrow$	Recall $\uparrow$
(a) Unconditional generation
Mixed	3D-HA-GAN [ 47 ]	2.08	95.7 $(+9.0)$	92.6	87.8	89.5	0.23 $\pm$ 0.03	0.20 $\pm$ 0.05
	Medical Diffusion [ 29 ]	1.79	90.9 $(+4.2)$	33.2	31.4	35.6	0.45 $\pm$ 0.02	0.42 $\pm$ 0.04
	3D-MedDiffusion [ 51 ]	1.76	89.3 $(+2.6)$	30.1	34.8	36.8	0.45 $\pm$ 0.04	0.50 $\pm$ 0.02
	3D-VQGAN [ 58 ]	1.78	88.1 $(+1.4)$	28.5	32.3	37.9	0.47 $\pm$ 0.06	0.53 $\pm$ 0.09
	Med-D-JEPA (Ours)	1.68	85.4 $(-1.3)$	24.8	25.7	28.9	0.58 $\pm$ 0.03	0.57 $\pm$ 0.07

Appendix

Table 9: Quantitative comparison of unconditional and class-conditional 3D brain tumor MRI generation on the BraTS2019 dataset . The best and second-best results within each evaluation setting are shown in blue and green , respectively. For MS-SSIM, values in parentheses denote the signed difference from the corresponding real-data score: 86.7 for mixed, 85.3 for LGG, and 88.6 for HGG.

Cohort type	Method	MMD ( $\times 10^{4}$ ) $\downarrow$	MS-SSIM (%)	FID-A $\downarrow$	FID-C $\downarrow$	FID-S $\downarrow$	Precision $\uparrow$	Recall $\uparrow$
(a) Unconditional generation
Mixed	3D-HA-GAN [ 47 ]	1.82	82.1 $(+4.9)$	69.7	72.4	73.8	0.41 $\pm$ 0.03	0.35 $\pm$ 0.07
	Medical Diffusion [ 29 ]	1.58	80.2 $(+3.0)$	26.3	27.8	26.9	0.65 $\pm$ 0.09	0.42 $\pm$ 0.03
	3D-MedDiffusion [ 51 ]	1.60	75.2 $(-2.0)$	31.7	30.4	27.3	0.61 $\pm$ 0.06	0.51 $\pm$ 0.05
	3D-VQGAN [ 58 ]	1.50	81.2 $(+4.0)$	27.8	25.9	26.5	0.72 $\pm$ 0.05	0.53 $\pm$ 0.04
	Med-D-JEPA (Ours)	1.46	74.4 $(-2.8)$	23.2	22.9	24.6	0.71 $\pm$ 0.07	0.53 $\pm$ 0.01

Appendix

Table 10: Quantitative comparison of unconditional and class-conditional 3D brain tumor MRI generation on the OASIS-1 dataset . The best and second-best results within each evaluation setting are shown in blue and green , respectively. For MS-SSIM, values in parentheses denote the signed difference from the corresponding real-data score: 77.2 for the mixed cohort, 80.8 for CN, and 76.3 for AD.

Dataset	Steps	FID-Avg. $\downarrow$	Precision $\uparrow$	Recall $\uparrow$	Time/Volume $\downarrow$
BraTS2019	8	27.4	0.49	0.47	3.11
	16	26.2	0.51	0.52	5.90
	32	29.9	0.53	0.57	11.78
	64	24.6	0.57	0.58	20.60
	256	29.8	0.56	0.54	92.47
	512	28.7	0.55	0.57	185.79

Appendix

Table 11: Ablation Study on the number of AR steps. using 100 generated volumes. Sampling time is reported as the average number of seconds per volume.

Dataset	CFG	FID-Avg. $\downarrow$	Precision $\uparrow$	Recall $\uparrow$
BraTS2019	2.0	25.3	0.57	0.55
	2.5	25.8	0.54	0.52
	3.0	25.6	0.53	0.53
	3.5	25.9	0.52	0.52
	4.0	25.5	0.51	0.52
OASIS-1	2.0	23.5	0.65	0.57

Appendix

Table 12: Ablation Study on different cfg values

Method	Resolution	Time/volume (s) $\downarrow$
3D-HA-GAN	$128^{3}$	1.25
Medical Diffusion	$128^{3}$	8.72
3D-VQGAN-cond	$128^{3}$	15.82
Med-D-JEPA	$128^{3}$	21.35
3D-MedDiffusion	$128^{3}$	27.41
Med-D-JEPA	$256^{3}$	42.54

Appendix

Table 13: Average sampling time per volume.

Explore similar work

Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI
Jul 12, 2026Fabian Mager, Lars Kai HansenContrastive Learning Brain Age Prediction
Department of Applied Mathematics and Computer Science Technical University of Denmark
MedForj: An open, large-scale foundational generative prior for high-resolution 3D brain MRI
Oct 29, 2025Samuel W. Remedios, Aaron Carass, Jerry L. Prince +1Flow Matching Healthcare
Image Analysis and Communications Laboratory, Johns Hopkins University · Department of Neurology, Johns Hopkins Hospital
BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation
Jun 17, 2026Max Van Puyvelde, Ibrahim Gulluk, Wim Van Criekinge +1Masked Autoencoders Latent Diffusion Models
Department of Biomedical Data Science, Stanford University School of Medicine · Department of Mathematical Modelling, Statistics & Bioinformatics, Ghent University · Department of Electrical Engineering, Stanford University

Jul 12, 2026J/K move · Enter open · S save

Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI

Fabian Mager, Lars Kai Hansen

Department of Applied Mathematics and Computer Science Technical University of Denmark

Contrastive Learning Brain Age Prediction Joint-Embedding Predictive Architecture Medical Image Representation Learning

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI

MedForj: An open, large-scale foundational generative prior for high-resolution 3D brain MRI

BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation