Learning Diffusion Models with Flexible Representation Guidance

We would like to thank the reviewer for their insightful review and their thoughtful feedback. In response to their review, we believe we have addressed their concerns as concretely as possible.

While REPA was motivated by improving training efficiency, REED uses Qwen2-VL 7B for image generation—a much heavier model than DINO—raising doubts about actual training speed improvements. A comparison of GFLOPS or latency is needed to confirm whether REED's 1M iterations in Table 1 are more efficient than REPA’s 4M, and whether 200 epochs in Table 2 are indeed more efficient than REPA’s 800. Figure 1 should also use GFLOPS or latency on the x-axis instead of training epochs.

We pre-calculate and store all VLM representations before diffusion model training. As a result, the additional computational overhead during REED training compared to REPA is minimal, limited only to a lightweight MLP projection head applied to the cached VLM representations. This overhead is negligible in practice, thus the actual training time for both REPA and REED is essentially the same. For example, training SiT-XL/2 for 1M iterations (i.e. 200 epochs) on 8×A100 GPUs takes approximately 180 hours, and 4M iterations (i.e. 800 epochs) takes about 720 hours for both methods.

In contrast, the one-time computational cost for generating captions and VLM representations is minor relative to the overall diffusion training cost: on 8×A100 GPUs, caption generation requires about 7.5 hours, and VLM representation extraction about 1.3 hours. Therefore, 1M iterations (200 epochs) of REED (+captioning+representation generation) takes 188.8 hours in total, much smaller than 4M iterations (800 epochs) of REPA (720 hours).

Furthermore, the pre-calculated representations can be reused across all experimental settings and training runs and the time cost is fixed, irrelevant to the number of training iterations. Therefore, the computational overhead introduced by REED is negligible compared to the total training cost. While we are unable to include an updated figure due to response restrictions, we will update Figure 1 accordingly and revise the manuscript with these clarifications and discussions.

Section 2's theoretical characterization is not novel, and seems unnecessary in this paper. The connection between ELBO and diffusion models is well-known [A], and the effectiveness of timestep weighting [B] has been established in prior work. I feel Sections 2 and 3 are loosely connected. [A] Kingma et al., Variational Diffusion Models, NeurIPS 2021. [B] Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation, NeurIPS 2023.

Our theoretical characterization is fundamentally different from [A] and [B]. Both [A] and [B] focus on diffusion models without external pretrained representations, whereas our work derives, to our knowledge, the first variational bound for diffusion models that incorporate external representation as additional guidance. This represents a substantial extension of the existing theory. Moreover, the timestep weighting discussed in our theoretical section specifically addresses how to balance the strength of representation guidance across different timesteps, which is conceptually distinct from the approaches explored in [B].

We believe this new theoretical framework is valuable for several reasons:

• It offers the first principled explanation of how and why representation-enhanced diffusion models–including prior works such as REPA and RCG–achieve their performance gains.
• It directly motivates key aspects of our new methods: for example, our use of multimodal representations (including synthetic modalities) is grounded in the analysis of multi-latent representation structures in Section 2.2. Similarly, our curriculum schedule for the alignment and diffusion losses is motivated by the decoupled structure of the variational bound in Section 2.1.

We will clarify these points in the revised manuscript and better highlight the novel connections between our theoretical analysis and the design of our method.

What would be the effect of using a smaller multimodal model, such as CLIP or SigLIP, instead of Qwen2-VL for image generation? Would REED retain its benefits while becoming more efficient?

As discussed in our response to the first point regarding computational efficiency, we pre-compute and store all captions and corresponding representations before diffusion model training. As a result, the computational cost of generating VLM representations is a one-time, fixed expense and does not scale with the number of training iterations. During diffusion training, the additional overhead from using these representations is negligible. Therefore, the size of the multimodal model used for representation extraction does not significantly impact the overall training efficiency of REED.

In the table below, we present the performance of REED using different multimodal representations with SiT-B/2 trained for 200K iterations and samples generated without classifier-free guidance (CFG). Both Qwen2-VL 7B and Qwen2-VL 2B representations result in substantial improvements over the setting without multimodal representations, while using OpenCLIP text representations yields slightly worse performance. This difference can be attributed to the fact that CLIP is trained with a contrastive objective, which is more suitable for classification tasks and does not provide the fine-grained semantic alignment needed for high-quality image generation. In comparison, VLMs are trained on a diverse set of vision-language tasks and have strong ability of vision-language understanding and generation. They are also better at bridging two modalities, and consequently, have better aligned visual and textual representations compared to CLIP. These satisfactory features enable them to generate richer and more relevant semantic representations for guiding image generation.

It is also worth noting that, even from a computational perspective, OpenCLIP actually needs more computational cost than Qwen2-VL 2B, since neither CLIP nor OpenCLIP can generate captions - hence the captioning step using Qwen2-VL 2B is necessary anyways and we can save the Qwen2-VL 2B embeddings along the way to avoid redundant computation.

The paper needs a figure illustrating the proposed framework, like the Figure 1 of the REPA paper.

We thank the reviewer for the suggestion. While we are unable to include an updated figure due to response restrictions, we will update the manuscript with a better figure illustration of the proposed framework in the revised version.

Thank you again for your constructive review and strong support of our work! We greatly appreciate your valuable suggestions and will incorporate the discussions into the final version.

We thank the reviewer for their time and feedback that helped us improve the work. In response to their review, we believe we have addressed their concerns as concretely as possible. Below, we share our thoughts on the questions asked and will update the manuscript to reflect these clarifications.

Could you provide more details on how the synthetic data is generated by VLM? Such as instruction prompts and caption examples.

We use Qwen2-VL 2B for image captioning, using the prompt “Describe this image.” and the raw image as the input. We adopt the official conversation template of Qwen2-VL-2B to arrange the text and image information. We set the number of maximum new tokens to be 200, which leads to concise and informative captions. While we are unable to include the corresponding images due to response restrictions, below are a few examples of the generated captions and the corresponding image class labels:

Class label: flamingo
Caption:
The image depicts a flamingo with its head turned slightly to the side. The flamingo has a long, curved neck and a pinkish-orange plumage. Its beak is black with a white tip, and it appears to be holding something in its beak, possibly a small object or food. The background is a dark green, possibly indicating a water body or a forested area. The flamingo's posture suggests it might be in the midst of a feeding or grooming behavior.

Class label: scabbard
Caption:
The image shows a knife with a brown leather sheath. The knife has a blade that appears to be made of a polished metal, possibly stainless steel, and it has a handle that is also made of a similar material, likely wood. The sheath is made of brown leather and has a decorative design on it. The knife and sheath are placed on a flat surface, possibly a table or a desk.

Class label: goldfish
Caption:
The image shows two goldfish in an aquarium. The goldfish on the left is facing the camera, with its body oriented towards the viewer. The goldfish on the right is facing away from the camera, with its body oriented towards the other goldfish. Both goldfish have a shiny, reflective surface, indicating they are likely well-maintained. The aquarium has a gravel bottom with some small rocks and a few large green plants. The lighting in the aquarium creates a warm, inviting atmosphere.

Experiments use Qwen2-VL 2B to generate synthetic captions and use representations from Qwen2-VL 7B? What’s the purpose to use different VLMs?

We use the smaller Qwen2-VL 2B model to generate the synthetic caption since it is already capable of generating captions with high quality and contains sufficiently useful information of the given image. Yet it is much faster than the 7B model in the generation speed. Since generating captions requires sequential inference and is more time-consuming compared with obtaining the representations in one forward pass for a given image-caption pair, we opt to generate captions with the 2B model and obtain representations with both models.

However, as shown in the table below, while the embeddings of the 2B model help, the 7B embeddings significantly improve the generation quality of REED. The table contains results for the SiT-B/2 models trained for 200K iterations and samples generated without classifier-free guidance. For the given image and caption pairs, the 7B model has stronger capability to understand the content and thus is able to generate more powerful multimodal representations that accelerate diffusion model training.

In summary, we use Qwen2-VL 2B since it is capable of generating high-quality captions efficiently. We leverage both Qwen2-VL 2B and 7B models for alignment, and the results show that the 7B model leads to superior performance.

In Section 3.1, the paper mentions using features from both a primary modality x and a secondary modality y. However, in the experiments, the authors chose to use features from VLMs. While VLMs do incorporate both text and image information, this appears more like a fused feature representation rather than distinct modalities as suggested by the theory. This seems to deviate significantly from the theoretical framework, particularly regarding how hyperparameters control the alignment between two separate features. Additionally, if VLMs are used as the aligned features, the key difference from REPA (which uses DINO features) seems to be primarily the choice of features (VLM vs. DINO). Could you elaborate more on the main experimental distinctions between these two works?

We actually use both the DINOv2 embeddings and the VLM embeddings in diffusion training following $\mathcal{L}_{\text{repgen}}$ in Equation 8 (Section 3.1). Specifically, for the image generation experiments, we set the coefficient of DINO embedding $\lambda_x$ as 1.0 and the coefficient of VLM embedding $\lambda_y$ as 0.5.

Therefore, compared with REPA, the main experimental distinctions are:

• We are the first to propose and demonstrate the effectiveness of using other auxiliary modalities to improve generation in a target modality. We believe the idea of leveraging (synthetic) multimodality to train single-modal generative models is novel and opens new directions for future research.
• We introduce a novel curriculum schedule between the representation alignment loss and diffusion loss, which phases in the diffusion loss while maintaining the representation alignment loss from the start.
• To the best of our knowledge, we are the first to apply representation alignment to diffusion models in the context of protein inverse folding and molecule generation, demonstrating the versatility and effectiveness of our approach across a diverse set of domains and tasks.

Each component of our method designs are well motivated by our theoretical framework and demonstrate significant improvement in the ablation study: as shown in the table with SiT-B/2 model trained for 200K iterations and samples generated without CFG, removing each component of REED leads to inferior performance. We will update the manuscript with more detailed explanations of the method and the experimental settings and these additional results in the revised version.

In Appendix C.1.2, from equation (21) to (22), it seems impossible to derive $p(z|x_0)p(x_0|x_1)=p(z|x_1)p(x_0|x_1,z)$

$p(z|x_1)p(x_0|x_1,z) = p(z, x_0|x_1) = p(x_0|x_1)p(z|x_0,x_1) = p(x_0|x_1)p(z|x_0)$ The last equation is due to the Markov chain $x_1-x_0-z$, which indicates $z$ and $x_1$ are conditionally independent given $x_0$.

Thank you for your time and response. We are glad that we have addressed your concerns. We sincerely appreciate your insightful comments and will incorporate the discussions into the final version. We would appreciate if you could update your final justification and rating.

We would like to thank the reviewer for their diligent and insightful review that helped us improve our work and their positive feedback regarding our theoretical framework, proposed method and experiments. In response to their review, we believe we have addressed their concerns as concretely as possible.

While the equivalence to the RCG method is clearly established through the variational bound and decomposed probabilistic framework, the connection to REPA is less explicit and lacks rigorous justification.

We establish the connection between our framework and REPA under the linear cumulative weight setting in Section 2.3 (lines 173–197) and Appendix C.1.6. Specifically, we show that in this setting, $\\mathcal{L}_{\\text{VB}}^z$ (Equation 4) corresponds directly to the REPA loss function. The connection is well-motivated and justified both theoretically and empirically:

• As detailed in Proposition 3 and Appendix C.1.6, the hybrid distribution $\\tilde{p}_{\theta}^{\\text{linear}}(x\_{t-1}|x\_t,z;A\_t)$ in the first term of $\\mathcal{L}\_{\\text{VB}}^z$ can be approximated by $p\_{\\theta}(x\_{t-1}|x\_t,\\hat{\\mu}\_t^z)$, where $\\hat{\\mu}\_t^z$ is an estimation of $z$ inferred from $x\_t$. This approximation renders that the first term is essentially equivalent to the diffusion loss in REPA, wherein the model extracts information about the pretrained representation $z$ from $x\_t$ in early layers of the transformer and then conducts score estimation based on $x\_t$ and the extracted features (i.e., $\\hat{\\mu}\_t^z$) in later layers.
• The third term of $\\mathcal{L}\_{\\text{VB}}^z$ aligns the predicted $z$ from $x\_t$ with the ground truth $z$, matching the cosine similarity alignment loss in REPA under the standard von Mises-Fisher distribution assumption for representations.
• We omit the log-normalization term due to its equivalent effect to the third term, and further justify this via experiments in Appendix F.1 (line 1195-1217) and Table 9.

Taken together, these points establish a rigorous connection between our theoretical framework and REPA, and provide clear motivation for the novel design choices in REED.

The paper would benefit from more comprehensive ablation experiments exploring how different choices of the curriculum schedules $\\alpha(n)$ and $\\beta(n)$ (which control the relative weighting of diffusion and representation alignment losses) affect model performance. Currently, the ablation on these schedules is limited; it remains unclear how sensitive the proposed method is to the selection or annealing strategy of these weights, and whether adaptive or learned schedules might further improve results.

This is indeed an essential ablation. We provide additional results on the image generation experiment as below with the SiT-B/2 architecture trained for 200K iterations and samples generated without classifier-free guidance. We will include these results in the revised manuscript.

Our findings clearly show that employing curriculum schedules leads to significant improvements: using a constant schedule (no curriculum), the FID is 29.4, whereas introducing a curriculum schedule for $\\alpha(n)$ reduces the FID to 24.6. This provides strong empirical support for the effectiveness of curriculum schedules, in line with our theoretical motivations.

We further examine different configurations for the schedules, including varying the phase-in period for $\alpha(n)$ and applying a cosine decay to $\beta(n)$. As shown in the table below, all curriculum variants of $\alpha(n)$ outperform the constant schedule, with 50K iterations chosen as the default for all image generation experiments. Notably, the choice of $\beta(n)$ schedule does not substantially affect generation performance.

We agree that adaptively or automatically learning these schedules could offer further benefits, for instance by adjusting the weighting based on the difficulty of individual samples or the training stage. For example, the training of samples with easy-to-extract representations could transition more quickly to a higher diffusion loss weight, while more challenging samples could undergo a longer phase-in period. However, designing and integrating such adaptive schedules presents additional challenges and is beyond the current scope of this work. We leave this as a promising direction for future research.

Thank you again for your constructive review and strong support of our work! We are glad that we have fully addressed your concerns. We greatly appreciate your valuable suggestions and will incorporate the discussions into the final version.

We would like to thank the reviewer for their diligent and insightful review that helped us improve our work and their positive feedback regarding our theoretical framework, presentation, experimental performance and broad applications. In response to their review, we believe we have addressed their concerns as concretely as possible.

Rigorous ablations to rule out simpler explanations for the performance gains.

This is indeed an essential ablation. We experiment on the class conditional image generation experiment on the SiT-B/2 architecture and report the results for models trained for 200K iterations and not use classifier-free guidance for generation. As shown in the table, removing each component leads to inferior performance than the full REED model, with FID increasing from 24.6 to 29.4 when removing the curriculum schedule and to 27.7 when removing multimodal representations. This demonstrates the effectiveness of each component in our proposed REED.

Further empirical analysis supporting why the linear phase-in for the diffusion loss is better than simpler schedules (e.g., constant weighting). Would adaptive or learned schedules yield further gains?

We provide additional empirical results with different curriculum schedules $\\alpha(n)$ and $\\beta(n)$ on the image generation experiment as below with the SiT-B/2 architectures trained for 200K iterations and samples generated without classifier-free guidance. We will include these results in the revised manuscript.

Compared to the REED without the curriculum schedules (i.e. constant schedule), the curriculum schedule leads to significant improvement, reducing the FID from 29.4 to 24.6. This further verifies the effectiveness of the curriculum schedules in accordance with our theoretical insights.

We further examine different configurations for the schedules, including varying the phase-in period for $\\alpha(n)$ and applying a cosine decay to $\\beta(n)$. As shown in the table below, all curriculum variants of $\\alpha(n)$ outperform the constant schedule, with 50K iterations chosen as the default for all image generation experiments. Notably, the choice of $\\beta(n)$ schedule does not substantially affect generation performance.

We agree that adaptively or automatically learning these schedules could offer further benefits, for instance by adjusting the weighting based on the difficulty of individual samples or the training stage. For example, the training of samples with easy-to-extract representations could transition more quickly to a higher diffusion loss weight, while more challenging samples could undergo a longer phase-in period. However, designing and integrating such adaptive schedules presents additional challenges and is beyond the current scope of this work. We leave this as a promising direction for future research.

The method's success is heavily dependent on extremely large, resource-intensive pretrained models like DINOv2, potentially limiting the application in resource-constrained settings.

Since it is a core motivation behind representation-enhanced generative methods to leverage powerful external pretrained representations to accelerate diffusion training, this is a common issue for this type of methods including REPA – which also utilizes DINOv2 representations.

However, it is worth noting that the specific DINOv2 variant we use in our experiments, DINOv2-B/14, has only 86M parameters. This is significantly smaller than the diffusion backbones used in our experiments, from SiT-B/2 (130M params) to SiT-XL/2 (675M params).

In terms of computational cost, for SiT-XL/2, the additional overhead of using DINOv2 representations per training step is minimal compared to the overall training cost–it takes about only 1.05x training time for each step–and is far outweighed by the 23.3x acceleration in training steps enabled by our method.

Moreover, some pretrained representations (e.g., the VLM representation from Qwen2-VL) can be pre-computed and reused, further reducing the computational cost. Thus, the overall FLOPS and memory consumption is almost the same as training a vanilla diffusion model.

The work is largely an extension of REPA. While the new components are innovative, the core contribution can be seen as an incremental improvement that combines existing ideas.

Our work has several fundamental and unique contributions compared to REPA and other representation-enhanced generative models:

• We provide the first theoretical framework for representation-enhanced diffusion model training, which not only explains the success of prior works such as REPA and RCG, but also well motivates our method design and offers guidance for future research in this area.
• We systematically explore the design space for representation alignment in diffusion models based on our theoretical insights, both theoretically grounded and empirically validated, which lead to significant performance gains.
• In particular, we are the first to propose and demonstrate the effectiveness of using other auxiliary modalities to improve generation in a target modality. We believe the idea of leveraging (synthetic) multimodality to train single-modal generative models is novel and opens new directions for future research.
• To the best of our knowledge, we are the first to apply representation alignment to diffusion models in the context of protein inverse folding and molecule generation, demonstrating the versatility and effectiveness of our approach across a diverse set of domains and tasks.

Overall, we believe that our work brings essential new insights into this field and can inspire a series of future research.

What failure modes might occur when the synthetic auxiliary modality is too dissimilar from the primary data? Could this lead to misaligned representations that degrade generation quality?

Our method is motivated by the relevance of the synthetic auxiliary modality to the generation task of the primary data. Accordingly, we are deliberate in constructing the auxiliary modality, carefully selecting modalities that are highly relevant and leveraging state-of-the-art models for both synthetic modality generation and multimodal representation extraction (see Section 3.3). For example, we use text modalities and vision-language models (VLMs) for image generation, 3D structures and AlphaFold3 for protein inverse folding, and 2D-3D joint molecule graphs with the UniMol model for molecule generation. This ensures that the auxiliary modality provides meaningful complementary information, maximizing its benefit during diffusion training while minimizing the risk of introducing irrelevant noise.

However, if the synthetic modality is too dissimilar from the primary data, there is a risk of misaligned representations. In such cases, the detrimental effects of irrelevant or distracting information could outweigh the benefits of any useful signal, ultimately leading to degraded generation quality. Therefore, careful selection and validation of the auxiliary modality remain essential to avoid these potential failure modes.

How does REED perform when using lightweight or less accurate pretrained representations?

We conduct additional experiments on the image generation experiment as below with the SiT-B/2 architectures trained for 200K iterations and samples generated without classifier-free guidance. We compare REED (using Qwen2-VL 7B representations) to versions using either a lightweight model (i.e., Qwen2-VL 2B) or intentionally noisier features (i.e., Qwen2-VL 7B with added Gaussian noise). All other settings were kept consistent. The results are presented below:

As shown in this table, using lightweight or less accurate pretrained representations still consistently leads to improved performance compared to the model without aligning with VLM representations. This demonstrates the robustness of our method to less accurate or noisy representations and suggests that REED retains its advantages even when ideal pretrained features are unavailable. Nevertheless, higher-quality representations (i.e., Qwen2-VL 7B) yield the strongest results, indicating that the quality of pretrained representations remains an important factor for maximizing performance.

How sensitive is REED to the differences between the pretrained representation space and the diffusion model's latent space?

During training, we align the diffusion model’s latent space with both the DINOv2 embedding space and the VLM representation space—each of which constitutes a distinct and independently pretrained representation space. Despite these substantial differences, our empirical results show that REED consistently achieves strong performance. This suggests that REED is rather robust to the differences among embedding spaces.

Thank you again for your insightful review and strong support of our work! We greatly appreciate your valuable suggestions and will incorporate the discussions and additional experimental results into the final version.