PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Meng Zhou, Wenhao You, Yuxing Chen, Yueying Tian

arXiv:2608.28787v1cs.CV

TL;DR

Existing JEPA generative evidence centers on natural images, leaving its applicability to 3D medical imaging unresolved. Med-D-JEPA adapts the framework to brain MRI using continuous latent tokens, joint prediction and diffusion, and iterative multi-token sampling, achieving competitive or superior generation performance and improving downstream classification and segmentation with synthetic pretraining.

  • Problem

    The applicability of joint-embedding predictive generation to 3D medical imaging, including its effects on anatomical consistency, dataset-size robustness, and quality–efficiency trade-offs, remains unexplored.

  • Method

    Med-D-JEPA compresses MRI volumes into continuous latent tokens and combines masked context prediction, per-token diffusion, and progressive sampling of token sets before decoding synthetic volumes.

  • Results

    Med-D-JEPA achieves competitive or superior performance to strong baselines across 3D brain MRI generation evaluations, while synthetic pretraining improves downstream classification and preliminary segmentation.

  • Takeaways & Limitations

    Joint-embedding predictive generation is a promising alternative for 3D medical image synthesis and warrants evaluation across additional organs, datasets, and conditioning settings.

Abstract

from arXiv · show

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.

1. Introduction

Med-D-JEPA adapts joint-embedding predictive generation to 3D brain MRI, addressing the need for synthetic volumes that preserve anatomy, pathology, diversity, and downstream utility. It combines continuous latent-token modeling with masked prediction, per-token diffusion, and progressive multi-token sampling across diverse MRI settings.

  • 3D MRI learning is constrained by limited cohorts, costly annotations, heterogeneous acquisition protocols, and patient-data sharing restrictions.
  • Synthetic MRI must preserve 3D anatomical coherence, pathological characteristics, population diversity, and measurable downstream-task utility, not merely visual realism.
  • During inference, Med-D-JEPA predicts and samples multiple latent tokens per autoregressive step, then decodes the completed latent grid into a synthetic volume.
  • Med-D-JEPA extends D-JEPA from natural-image representation learning to systematic generation of 3D brain MRI.
  • A 3D KL-regularized adversarial variational autoencoder compresses volumes into continuous latent tokens, while Med-D-JEPA jointly optimizes JEPA prediction and per-token diffusion objectives.
  • The evaluation spans tumor-centered regions and whole-brain volumes, FLAIR and T1-weighted contrasts, and resolutions from 128^3 to 256^3 voxels across three MRI datasets.

2. Related Work

Prior 3D medical image synthesis has relied heavily on GANs, while continuous token generation remains relatively unexplored in data-limited volumetric imaging. JEPA methods predict target representations from visible context, and D-JEPA connects this principle to generation through diffusion over continuous masked tokens.

  • GAN-based Image Synthesis: GANs established realistic 2D and 3D medical image synthesis, but 3D adversarial training is susceptible to optimization instability and mode collapse.
  • GAN-based Image Synthesis: Generated lesion images are often used to augment training data and address data-scarcity challenges in downstream medical-imaging tasks.
  • Joint-Embedding Predictive Generation: JEPAs learn representations by predicting target features from visible context rather than reconstructing pixels.
  • Joint-Embedding Predictive Generation: D-JEPA connects joint-embedding prediction with generation by conditioning diffusion on continuous masked tokens, whereas prior studies focused mainly on large-scale 2D images.

3. Method

Med-D-JEPA compresses 3D MRI into continuous latent tokens, trains masked context prediction with per-token diffusion, and generates volumes through iterative masked-token sampling.

  • Continuous 3D Latent Compression: The method first compresses each MRI volume into continuous latent representations with a 3D KL-regularized adversarial variational autoencoder.The encoder downsamples the voxel-space volume, and the decoder reconstructs it from the latent features.
  • Continuous 3D Latent Compression: The compression network combines perceptual, adversarial, voxel-space reconstruction, gradient, and KL-regularization objectives.These objectives are used to constrain reconstruction quality and latent representations.
  • Continuous 3D Latent Compression: Frozen latent features are partitioned into non-overlapping 3D patches and linearly projected into continuous tokens for Med-D-JEPA training.Each token represents a spatial latent patch, and the resulting sequence is used as the model input.
  • Med-D-JEPA Training: Med-D-JEPA uses context and target encoders plus a feature predictor to predict masked-token representations from visible context.The target encoder is updated through an exponential moving average, while the context encoder and predictor are randomly initialized.
  • Med-D-JEPA Training: Per-token diffusion learns the conditional distribution of each masked latent token given its contextual representation.Gaussian noise is added at a sampled diffusion timestep, and a denoising MLP is trained for each masked token.
  • Iterative 3D Latent Generation: At inference, generation starts from a fully masked 3D latent grid and progressively samples token sets through reverse diffusion before decoding the completed grid.A cosine schedule controls how many tokens are generated per iteration; larger K provides richer context updates but fewer tokens per step.

4. Experiments

The experiments evaluate Med-D-JEPA across brain MRI datasets, generation settings, downstream classification, and multiple fidelity, diversity, and task metrics.

  • Datasets and Evaluation: Med-D-JEPA is evaluated on BraTS2019, OASIS-1, and BraTS2020 using different MRI contrasts, patient groups, spatial scopes, and downstream tasks.BraTS2019 provides FLAIR tumor ROIs for unconditional, HGG/LGG-conditioned generation, and classification; OASIS-1 provides T1-weighted MRI grouped by cognitive status.
  • Datasets and Evaluation: The comparison includes five state-of-the-art generation baselines rerun on the same datasets and splits.The baselines are 3D-HA-GAN, Medical Diffusion, 3D-MedDiffusion, 3D-VQGAN, and 3D-VQGAN-cond.
  • Qualitative Comparison: Figure 3 compares real and generated BraTS2019 volumes using central axial, coronal, and sagittal slices.The qualitative comparison covers samples from different generation methods and is intended for zoomed inspection.
  • Evaluation Metrics: Generation quality is assessed with FID-Avg., Precision, and Recall, while classification uses AUC, F1-score, and Accuracy.FID-Avg. averages FID across the three anatomical views, while Precision and Recall characterize fidelity and coverage separately.

5. Results

Med-D-JEPA performs competitively across tumor-ROI and whole-brain MRI generation, while synthetic pretraining improves downstream classification and segmentation. Ablations identify effective sampling and objective choices, with a modest inference-time overhead.

  • Tumor-ROI generation: Med-D-JEPA achieves competitive performance across unconditional and class-conditional tumor-ROI generation metrics.It reports competitive results on MMD, MS-SSIM, FID-Avg., precision, and recall.
  • Whole-brain generation: Med-D-JEPA attains the lowest MMD and FID-Avg. for unconditional whole-brain generation, while improving recall in CN- and AD-conditioned settings.For CN conditioning, precision is 0.76 versus 0.78 for 3D-VQGAN-cond; its lower MS-SSIM indicates greater inter-sample diversity.
  • Downstream classification: 22, 28, and 19 percentage points are the BraTS2019 AUC, F1-score, and accuracy gains from pretraining on 300 Med-D-JEPA samples per class versus real-only training.Med-D-JEPA also outperforms conventional augmentation and the strongest generative baseline, 3D-VQGAN, on these metrics.
  • Downstream segmentation: Dice increases from 0.74 to 0.80 and HD95 decreases from 13.40 to 9.56 mm after Med-D-JEPA synthetic pretraining for BraTS2020 segmentation.Both comparisons are statistically significant, although the evaluation is preliminary and uses pseudo-labels for synthetic volumes.
  • Ablation and efficiency: 64 autoregressive generation steps provide the best quality-efficiency trade-off, with the lowest FID and highest recall in approximately 21 seconds per volume.Sixteen steps take about 6 seconds but reduce distributional coverage, while 512 steps are 8.3–9.0× slower without quality improvement.
  • Ablation and efficiency: Adding the JEPA prediction loss reduces average FID by 1.6 and 1.8 on BraTS2019 and OASIS-1, respectively, while improving precision and recall.The gains are 4 and 5 percentage points in precision, and 6 and 3 percentage points in recall.
  • Ablation and efficiency: Med-D-JEPA generates a 128^3 volume in 21.35 seconds on average, compared with 15.82 seconds for 3D-VQGAN-cond.The reported overhead is characterized as modest relative to the generation-quality trade-off.

6. Conclusion

The paper presents Med-D-JEPA as a competitive joint-embedding predictive framework for 3D brain MRI generation. It identifies predictive generation as a promising alternative and proposes broader evaluation across organs, datasets, and conditioning modes.

  • Conclusion: Med-D-JEPA achieves competitive or superior performance to strong state-of-the-art methods for 3D brain MRI generation.The conclusion characterizes joint-embedding predictive generation as a promising alternative to existing generative paradigms.
  • Conclusion: Future work will evaluate Med-D-JEPA across additional organs and datasets and extend it to text-conditioned medical image generation.

A. Datasets and Implementation Details

The study evaluates three public brain MRI datasets spanning tumor-centered and whole-brain generation, multiple contrasts and resolutions, and classification and segmentation tasks. Implementation uses separate 3D KL-VAE-GAN and Med-D-JEPA models with 64 autoregressive iterations by default.

  • Datasets: The experiments use BraTS2019, BraTS2020, and OASIS-1 across tumor-centered and whole-brain generation, FLAIR and T1-weighted MRI, and 128^3 and 256^3 resolutions.The evaluation also includes downstream classification and segmentation tasks.
  • Datasets: BraTS2019 provides FLAIR tumor ROIs from 259 HGG and 76 LGG patients for generation and tumor-grade classification.The ROIs are resampled to 128 × 128 × 128 voxels.
  • Datasets: OASIS-1 provides T1-weighted MRI from 416 patients grouped as CN or AD using Clinical Dementia Rating values.CDR = 0 defines CN, while CDR ∈ {0.5, 1, 2} defines AD.
  • Implementation: The implementation trains separate 3D KL-VAE-GAN and Med-D-JEPA models for each dataset, using K = 64 autoregressive generation iterations by default.Class-conditional generation uses classifier-free guidance with a guidance scale of 2.0.
  • Evaluation: Generation quality is evaluated with plane-wise FID, whole-volume MMD, and volumetric-feature precision and recall.FID is computed across axial, coronal, and sagittal planes, while MMD uses complete 3D volumes.

B. More on Qualitative Results

Med-D-JEPA-generated volumes overlap substantially with real samples in learned feature space, although some cohorts show distributional shift. Across cohorts, generated samples occupy multiple regions rather than collapsing into a compact cluster.

  • BraTS2019 LGG and OASIS-1 AD samples show substantial intermixing between real and generated features, without clear separation by data source.
  • BraTS2019 HGG and OASIS-1 CN samples overlap with parts of the real feature space but also occupy broader, source-dominant regions.
  • Generated samples span multiple feature-space regions across cohorts rather than collapsing into a compact cluster.
  • The t-SNE visualization compares real and synthetic samples for BraTS2019 HGG/LGG and OASIS-1 CN/AD cohorts.

C.1. VAE Reconstruction Ability on Test set

The 3D KL-VAE-GAN latent compression stage preserves the principal volumetric structures of held-out MRI volumes. All evaluated models exceed 30 dB PSNR and 0.91 3D-SSIM, with BraTS2019 achieving the strongest reconstruction scores.

  • The reconstruction evaluation uses PSNR and volumetric 3D-SSIM on each model’s corresponding held-out test set.
  • All three 3D KL-VAE-GAN models achieve PSNR above 30 dB and 3D-SSIM above 0.91 on held-out test sets.These results indicate preservation of the principal volumetric structures during latent compression and decoding.
  • BraTS2019 achieves the highest reconstruction scores, with 33.27 dB PSNR and 0.95 3D-SSIM.
  • Reconstructions are evaluated using the central input slice, reconstruction, and absolute-error map for BraTS2020, OASIS-1, and BraTS2019.

C.2. Full Results on Quantitative Comparison

Med-D-JEPA achieves the lowest plane-wise FID across all unconditional comparisons and most class-conditional comparisons, while its synthetic volumes can support preliminary pseudo-mask generation.

  • Med-D-JEPA achieves the lowest FID in all three anatomical planes for unconditional generation on both BraTS2019 and OASIS-1.
  • For class-conditional generation, Med-D-JEPA achieves the lowest plane-wise FID in nine of twelve dataset-class combinations.The exceptions are BraTS2019 HGG FID-S, OASIS-1 CN FID-S, and OASIS-1 AD FID-A.
  • The frozen NV-Segment-CTMR model generates whole-tumor pseudo-masks for synthetic BraTS2020 volumes without task-specific fine-tuning.Its suitability is first evaluated on the real BraTS2020 training cohort using available ground-truth masks.

E. More on Ablation Study

Ablations identify K = 64 and scfg = 2.0 as practical defaults: they balance generation quality, distributional coverage, and sampling cost better than more extreme settings.

  • AR-step ablation: Increasing K from 8 to 64 generally improves generation quality and distributional coverage by providing progressively richer context for later token predictions.
  • AR-step ablation: K = 64 achieves the lowest FID-Avg. and highest or near-highest Precision and Recall on both datasets.
  • AR-step ablation: K = 512 is approximately 9.0× and 8.3× slower than K = 64 on BraTS2019 and OASIS-1, respectively.Further increasing K to 256 or 512 provides no consistent quality improvement.
  • Classifier-free guidance ablation: scfg = 2.0 provides the best overall balance between fidelity and distributional coverage across BraTS2019 and OASIS-1.It achieves the lowest FID-Avg. and strongest or tied-best Precision and Recall across the reported dataset settings.
  • Classifier-free guidance ablation: Increasing guidance beyond scfg = 2.0 generally reduces Precision and Recall without improving FID, suggesting a more restricted generated distribution.

E.1. Sampling Efficiency

Med-D-JEPA has higher sampling costs than some baselines because it uses iterative next-set-of-token sampling and per-token denoising, but remains faster than 3D-MedDiffusion at 128^3 while achieving stronger overall generation quality.

  • 21.35 seconds per volume at 128^3 makes Med-D-JEPA slower than Medical Diffusion and 3D-VQGAN-cond but approximately 22.1% faster than 3D-MedDiffusion.The comparison is based on average sampling time per volume.
  • 42.54 seconds per volume at 256^3 highlights the cost of iterative token generation and per-token denoising.The paper identifies reducing this cost without compromising generation quality as future work.
  • Average sampling time per volume is reported in Table 13.
Loading 2608.28787v1…