Learning Diffusion Models with Flexible Representation Guidance | NeurIPS 2025

Chenyu Wang^1,* | Cai Zhou^1,* | Sharut Gupta¹ | Zongyu Lin² | Stefanie Jegelka^1,3 | Stephen Bates¹ | Tommi Jaakkola¹

¹ MIT ² UCLA ³ TU Munich

^* Equal Contribution

The repository contains the code for the REED method presented in the paper: Learning Diffusion Models with Flexible Representation Guidance (NeurIPS 2025). Check our project page here!

📢 News

• [2025/10/12] We release the project page.
• [2025/09/18] REED is accepted to NeurIPS 2025!
• [2025/07/12] Code is released!
• [2025/07/12] Paper is available on arXiv!

Overview

REED presents a comprehensive framework for representation-enhanced diffusion model training, combining theoretical analysis, multimodal representation alignment strategies, an effective training curriculum, and practical domain-specific instantiations (image, protein sequence, and molecule).

Image Generation

For the class-conditional ImageNet $256\times 256$ benchmark, REED achieves a $23.3 \times$ training speedup over the original SiT-XL, reaching FID=8.2 in only 300K training iterations (without classifier-free guidance); and a $4 \times$ speedup over REPA (Yu et.al, 2024), matching its classifier-free guidance performance at 800 epochs with only 200 epochs of training (FID=1.80). The detailed code and instructions are in image/.

Protein Sequence Design

For protein inverse folding, REED accelerates training by $3.6\times$ and yields significantly superior performance across metrics such as sequence recovery rate, RMSD and pLDDT. The detailed code and instructions are in protein/.

Molecule Generation

For molecule generation, REED improves metrics such as atom and molecule stability, validity, energy, and strain on the challenging Geom-Drug datasets. The detailed code and instructions are in molecule/.

Citation

If you find this work useful in your research, please cite:

@article{wang2025learning,
  title={Learning Diffusion Models with Flexible Representation Guidance},
  author={Chenyu Wang and Cai Zhou and Sharut Gupta and Zongyu Lin and Stefanie Jegelka and Stephen Bates and Tommi Jaakkola},
  journal={arXiv preprint arXiv:2507.08980},
  year={2025}
}