Research map / Image generation / Autoregressive and latent generation
Learning Diffusion Models with Flexible Representation Guidance
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 REED, a variational framework for representation-guided diffusion that generalizes REPA and RCG, adding multimodal synthetic-pair representations and a phase-in curriculum, validated on ImageNet, protein inverse folding, and molecule generation.
Problem
Lack of a systematic understanding of how pretrained representations should be incorporated into diffusion training, and how to choose when and how to inject them, for faster and better generative training.
Approach
Derives a variational bound with decomposed parameterization and a cumulative-weight hybrid conditional distribution that interpolates guided and unguided reverse steps; extends to multi-latent representations from synthetic paired modalities and uses a curriculum that ramps up the diffusion loss while representation loss starts from the outset.
Key results
- On ImageNet 256x256 without CFG, REED reaches FID 8.2 on SiT-XL/2 at 300K iterations, beating vanilla SiT-XL at 7M steps (8.3).
- With CFG, REED matches REPA's 800-epoch FID of 1.80 with only 200 epochs.
- On GEOM-Drug, SemlaFlow+REED improves molecule stability from 97.4 to 98.1 and lowers energy from 114.44 to 105.63 with the same 100 NFE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add REED instead of REPA | SiT-B/2 · ImageNet 256×256, no CFG, 400K iterations | FID: 24.4 -> 19.5 | Helped |
| add REED instead of REPA | SiT-L/2 · ImageNet 256×256, no CFG, 400K iterations | FID: 9.7 -> 8.9 | Helped |
| add REED instead of REPA | SiT-XL/2 · ImageNet 256×256, no CFG, 1M iterations | FID: 6.4 -> 4.7 | Helped |
| add REED rather than use REPA | SiT-XL/2 · ImageNet 256×256, CFG | FID: 1.80 at 800 epochs -> 1.80 at 200 epochs | Mixed |
| add REED | DiscreteDiff · PDB protein inverse folding, 100 epochs | Seq. Recovery (%): 40.74 (0.04) -> 41.91 (0.07) | Helped |
| add REED | DiscreteDiff · PDB protein inverse folding, 100 epochs | scRMSD: 1.789 (0.015) -> 1.780 (0.017) | Helped |
| add REED | SemlaFlow · GEOM-DRUG | Molecule stability: 97.4 -> 98.1 | Helped |
| add REED | SemlaFlow · GEOM-DRUG | Energy: 114.44 -> 105.63 | Helped |
| add multimodal representations (ablation) | REED on SiT-B/2 · ImageNet 256×256, 200K iterations, no CFG | FID: 27.7 -> 24.6 | Helped |
| add curriculum schedule (ablation) | REED on SiT-B/2 · ImageNet 256×256, 200K iterations, no CFG | FID: 29.4 -> 24.6 | Helped |
| set image alignment depth to 4 and VLM alignment depth to 8 (ablation) | REED on SiT-B/2 · ImageNet 256×256, 200K iterations | FID: 25.9 -> 24.6 | Helped |
| vary VLM alignment depth from 6 to 8 while image depth is 4 (ablation) | REED on SiT-B/2 · ImageNet 256×256, 200K iterations | FID: 24.8 -> 24.6 | Helped |
Surprising findings
- Removing the curriculum schedule and removing multimodal representations each degrade FID on SiT-B/2, with the curriculum contributing a larger drop (29.4 vs 27.7 without multimodal reps).
Limitations
- Linear representation weighting schedule is the only one explored in experiments; adaptive schedules left to future work.
- Depends on quality of pretrained multimodal and folding models, and on synthetic auxiliary data generation (captions, AF3 structures).
Experimental setup
Datasets: ImageNet PDB (Multiflow split) GEOM-Drug
Benchmarks: ImageNet 256x256 GEOM-Drug PDB protein inverse folding
Metrics: FID IS precision recall sequence recovery scRMSD pLDDT atom stability molecule stability validity
Baselines: REPA SiT-XL/2 ADM LDM-4 DiT-XL/2 MDTv2-XL/2 ProteinMPNN MiDi EQGAT-diff SemlaFlow
Builds on: REPA RCG DDPM stochastic interpolants SiT
Closest papers by meaning
- Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think (ICLR 2025 · similarity 0.93)
- Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-based Decoding (NeurIPS 2025 · similarity 0.91)
- Why DiT Models Underperform as Representation Learners without Long Skip Connections (NeurIPS 2026 · similarity 0.91)
- Elucidating the design space of classifier-guided diffusion generation (ICLR 2024 · similarity 0.90)
- On improved Conditioning Mechanisms and Pre-training Strategies for Diffusion Models (NeurIPS 2024 · similarity 0.90)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.90)
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.