Research map / Image generation / Fast diffusion sampling
Denoising with a Joint-Embedding Predictive Architecture
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes D-JEPA, which combines a joint-embedding predictive architecture (with an EMA target encoder and smooth-L1 feature prediction) with per-token diffusion or flow-matching loss to perform generalized next-set-of-token generation in continuous latent space for ImageNet class-conditional image synthesis.
Problem
Joint-embedding predictive architectures are underexplored for generative modeling, while autoregressive models that generate continuous tokens lack representation-learning benefits; the paper asks whether JEPA-style representation learning can improve continuous-token generative modeling.
Approach
Treat JEPA as masked image modeling reinterpreted as generalized next-token prediction, and add a per-token diffusion (or flow matching) loss on the predicted embeddings so the JEPA objective avoids collapse and the model generates continuous tokens autoregressively in sets.
Key results
- D-JEPA-L (687M) reaches FID 2.32 and IS 233.5 without CFG on ImageNet 256x256, outperforming MAR-H (943M, FID 2.35) with fewer epochs (480 vs 800).
- With CFG, D-JEPA-H reaches FID 1.54 (IS 324.2 at cfg=3.9) and D-JEPA-L reaches FID 1.58, close to the VAE upper bound of 1.199.
- D-JEPA-H achieves FID about 4.0 in roughly 43-47 ms per image on a single H800 with batch size 256.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add JEPA loss to the MAR-B setup (ablation) | MAR-B / D-JEPA-B · ImageNet 256×256 conditional generation, without CFG | FID: 3.48 -> 3.40 | Helped |
| add JEPA loss to the MAR-B setup (ablation) | MAR-B / D-JEPA-B · ImageNet 256×256 conditional generation, without CFG | IS: 192.4 -> 197.1 | Helped |
| increase autoregressive sampling steps from 32 to 64 (ablation) | D-JEPA-B, without CFG · ImageNet 256×256 generation | FID: 4.027 -> 3.452 | Helped |
| increase autoregressive sampling steps from 32 to 64 (ablation) | D-JEPA-L, with CFG (cfg=3.0) · ImageNet 256×256 generation | FID: 2.261 -> 1.633 | Helped |
| increase autoregressive sampling steps from 32 to 64 (ablation) | D-JEPA-H, with CFG (cfg=3.9) · ImageNet 256×256 generation | FID: 1.859 -> 1.551 | Helped |
| replace diffusion loss with flow matching loss (ablation) | D-JEPA · Image generation and continuous-data modeling | visual generation quality | Helped |
Surprising findings
- Best sampling FID does not require 256 autoregressive steps; for D-JEPA-H the best result is obtained with about 64 steps, and the optimal step count decreases as model size grows.
- Training in raw pixel space is not required; the authors report latent-space modeling is chosen for generation efficiency, although they claim raw-space training also works well.
Limitations
- Scaling beyond the huge model is stated as a significant open challenge.
- Results without CFG depend on tuned temperature and CFG grid search, making comparisons sensitive to sampling settings.
Experimental setup
Datasets: ImageNet-1K
Benchmarks: ImageNet 256x256 class-conditional generation
Metrics: FID Inception Score Precision Recall
Baselines: MAR-B MAR-L MAR-H MaskGIT MAGE GIVT MAGVIT-v2 DiT-XL SiT-XL MDTv2-XL VAR VDM++
Builds on: JEPA (I-JEPA) MAR DDPM flow matching MaskGIT ADM
Closest papers by meaning
- Connecting Joint-Embedding Predictive Architecture with Contrastive Self-supervised Learning (NeurIPS 2024 · similarity 0.89)
- How JEPA Avoids Noisy Features: The Implicit Bias of Deep Linear Self Distillation Networks (NeurIPS 2024 · similarity 0.89)
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.88)
- Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures (NeurIPS 2026 · similarity 0.87)
- LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels (NeurIPS 2026 · similarity 0.86)
- Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think (NeurIPS 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.