Learning Diffusion Models with Flexible Representation Guidance

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Wang, Chenyu, Zhou, Cai, Gupta, Sharut, Lin, Zongyu, Jegelka, Stefanie, Bates, Stephen, Jaakkola, Tommi
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866918159042215936
author Wang, Chenyu
Zhou, Cai
Gupta, Sharut
Lin, Zongyu
Jegelka, Stefanie
Bates, Stephen
Jaakkola, Tommi
author_facet Wang, Chenyu
Zhou, Cai
Gupta, Sharut
Lin, Zongyu
Jegelka, Stefanie
Bates, Stephen
Jaakkola, Tommi
contents Diffusion models can be improved with additional guidance towards more effective representations of input. Indeed, prior empirical work has already shown that aligning internal representations of the diffusion model with those of pre-trained models improves generation quality. In this paper, we present a systematic framework for incorporating representation guidance into diffusion models. We provide alternative decompositions of denoising models along with their associated training criteria, where the decompositions determine when and how the auxiliary representations are incorporated. Guided by our theoretical insights, we introduce two new strategies for enhancing representation alignment in diffusion models. First, we pair examples with target representations either derived from themselves or arisen from different synthetic modalities, and subsequently learn a joint model over the multimodal pairs. Second, we design an optimal training curriculum that balances representation learning and data generation. Our experiments across image, protein sequence, and molecule generation tasks demonstrate superior performance as well as accelerated training. In particular, on the class-conditional ImageNet $256\times 256$ benchmark, our guidance results in $23.3$ times faster training than the original SiT-XL as well as four times speedup over the state-of-the-art method REPA. The code is available at https://github.com/ChenyuWang-Monica/REED.
format Preprint
id arxiv_https___arxiv_org_abs_2507_08980
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Learning Diffusion Models with Flexible Representation Guidance
Wang, Chenyu
Zhou, Cai
Gupta, Sharut
Lin, Zongyu
Jegelka, Stefanie
Bates, Stephen
Jaakkola, Tommi
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
Diffusion models can be improved with additional guidance towards more effective representations of input. Indeed, prior empirical work has already shown that aligning internal representations of the diffusion model with those of pre-trained models improves generation quality. In this paper, we present a systematic framework for incorporating representation guidance into diffusion models. We provide alternative decompositions of denoising models along with their associated training criteria, where the decompositions determine when and how the auxiliary representations are incorporated. Guided by our theoretical insights, we introduce two new strategies for enhancing representation alignment in diffusion models. First, we pair examples with target representations either derived from themselves or arisen from different synthetic modalities, and subsequently learn a joint model over the multimodal pairs. Second, we design an optimal training curriculum that balances representation learning and data generation. Our experiments across image, protein sequence, and molecule generation tasks demonstrate superior performance as well as accelerated training. In particular, on the class-conditional ImageNet $256\times 256$ benchmark, our guidance results in $23.3$ times faster training than the original SiT-XL as well as four times speedup over the state-of-the-art method REPA. The code is available at https://github.com/ChenyuWang-Monica/REED.
title Learning Diffusion Models with Flexible Representation Guidance
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2507.08980