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.

    Source: §1 · pp. 1–2

  • 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.

    Source: §1 · p. 2, §3.3 · p. 4

  • 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.

    Source: §1 · p. 3, §6 · p. 8, §5.4 · p. 8

  • 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.

    Source: §6 · p. 8

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.Source: §1 · p. 1
  • Synthetic MRI must preserve 3D anatomical coherence, pathological characteristics, population diversity, and measurable downstream-task utility, not merely visual realism.Source: §1 · p. 1
  • During inference, Med-D-JEPA predicts and samples multiple latent tokens per autoregressive step, then decodes the completed latent grid into a synthetic volume.Source: Figure 2 · p. 2, §1 · p. 2
  • Med-D-JEPA extends D-JEPA from natural-image representation learning to systematic generation of 3D brain MRI.Source: §1 · p. 2
  • 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.Source: §1 · p. 2
  • 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.Source: §1 · pp. 2–3

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.Source: §2 · p. 3
  • Generated lesion images are often used to augment training data and address data-scarcity challenges in downstream medical-imaging tasks.Source: §2 · p. 3

Joint-Embedding Predictive Generation

  • JEPAs learn representations by predicting target features from visible context rather than reconstructing pixels.Source: §2 · p. 3
  • 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.Source: §2 · p. 3

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.Source: §3.1 · p. 3, §3.2 · p. 3
  • 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.Source: §3.2 · pp. 3–4
  • 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.Source: §3.2 · p. 4

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.Source: §3.3 · p. 4
  • 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.Source: §3.3 · p. 4

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.Source: §3.4 · p. 4

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.Source: §4.1 · p. 5, §4.2 · p. 5
  • 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.Source: §4.2 · p. 5

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.Source: Figure 3 · p. 5

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.Source: §4.2 · p. 5

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.Source: §5.1 · p. 6

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.Source: §5.2 · p. 6

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.Source: §5.3 · p. 7

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.Source: §5.4 · p. 8

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.Source: §5.5 · p. 8
  • 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.Source: §5.5 · p. 8
  • 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.Source: §5.5 · p. 8

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.Source: §6 · p. 8
  • Future work will evaluate Med-D-JEPA across additional organs and datasets and extend it to text-conditioned medical image generation.Source: §6 · p. 8

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.Source: A. Datasets and Implementation Details · p. 11
  • 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.Source: A. Datasets and Implementation Details · p. 11
  • 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.Source: A. Datasets and Implementation Details · p. 11

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.Source: A. Datasets and Implementation Details · p. 12

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.Source: A. Datasets and Implementation Details · p. 12

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.

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.Source: §C.1 · p. 13, Table 8 · p. 13
  • 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.Source: §C.1 · p. 13
  • BraTS2019 achieves the highest reconstruction scores, with 33.27 dB PSNR and 0.95 3D-SSIM.Source: §C.1 · p. 13
  • Reconstructions are evaluated using the central input slice, reconstruction, and absolute-error map for BraTS2020, OASIS-1, and BraTS2019.Source: Figure 8 · p. 13

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.Source: §C.2 · p. 13
  • 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.Source: §C.2 · p. 13
  • 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.Source: D. More on Segmentation Study · p. 13

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.Source: E. More on Ablation Study · p. 14
  • K = 64 achieves the lowest FID-Avg. and highest or near-highest Precision and Recall on both datasets.Source: E. More on Ablation Study · p. 14
  • 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.Source: E. More on Ablation Study · p. 14

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.Source: E. More on Ablation Study · pp. 14–15
  • Increasing guidance beyond scfg = 2.0 generally reduces Precision and Recall without improving FID, suggesting a more restricted generated distribution.Source: E. More on Ablation Study · p. 15

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.Source: §E.1 · p. 15
  • 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.Source: §E.1 · p. 15
  • Average sampling time per volume is reported in Table 13.Source: Table 13 · p. 15

Related papers

Loading 2608.28787v1…