Learning Diffusion Models with Flexible Representation Guidance
Authors: Chenyu Wang
Cai Zhou
Sharut Gupta
Zongyu Lin
Stefanie Jegelka
Stephen Bates
Tommi Jaakkola
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion Models
Representation Learning
Image Generation
Biomolecule Generation
Theoretical Analysis of Generative Model
Primary area: deep_learning
TL;DR: We introduce a systematic approach to flexibly incorporating representation guidance into diffusion models, resulting in both accelerated training and better performance across image, protein, and molecule generation tasks.

Abstract: 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.

Learning Diffusion Models with Flexible Representation Guidance
Chenyu Wang Cai Zhou
Sharut Gupta'
Zongyu Lin?
Stefanie Jegelka
Stephen Bates!
Tommi Jaakkola
MIT UCLA STU Munich
Project Page
Abstract
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 frame- work 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 represen- tations 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 joint model over the multimodal pairs. Second, design 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 consiconniiional ImageNet 256 x 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-ase method REPA [[].
Introduction
ImageNet 256x256
Protein Inverse Folding
Diffusion models 20 50
SIT-XL2
have achieved significant empirical success across im- ages, videos, audio, and
SIT-XL/2 REPA
Sequence Recovery
41.9
3,6x faster
FID-50K
STT-XL/2+ REED
41.1
25.0
biomolecules [46 6 377
I [56]. While their pri- mary objective is model the data distribution, recent
10.0
23.3x faster
DiscreteDiff
4x faster
38.3
DiscreteDiff' REED
work 110 32 58) show that diffusion models implicitly learn some discriminative features in their latent rep- resentations. Nevertheless,
2M
Training Iteration
Training Epoch
Figure 1: REED achieves superior performance and accelerated train- ing on image generation and protein inverse folding.
the learned representations still fall behind those of pretrained representation learning models, which potentially limits the performance of generative modeling. Consequently, integrating high-quality representations obtained from other pretraining tasks such Self-supervised Learning (SSL) unlocks
additional possibilities to improve diffusion training. Representations can integrated through flexible designs, for example, via conditioning 33
35 or alignment [6]]. Specifically, RCG (34)
Equal Contribution: (wangchy, 2aizt22))0mit.eee
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

trains generative model to produce semantic representations from a self-supervised encoder and conditions image generation on these representations, while REPA [6]] aligns diffusion model hidden states with clean representations from pretrained encoders. Despite the empirical effectiveness of these methods, there is a lack of systematic understanding of the role of representation learning in training diffusion models.
In this paper, we present a theoretical framework to analyze how high-quality, pretrained represen- tations can enhance diffusion model training. Building on the DDPM [20 framework, we derive a variational bound on the negative log sigllikllooodd that incorporates the pretrained representations to guide the generation through probabilistic decomposition of the joint distribution and a weighted aggregation process. This leads to two key components: representation modeling. which promotes ex- tracting representations from noisy input, and data generation, which fuses representationncondrrinneed conditioned and unconditional generation into a hybrid model governed by a weight schedule (Section 2.T). The framework generalizes to multi-laten structures, enabling integration of diverse representations (Section 2.2). Our unified perspective subsumes existing approaches like REPA [6] and RCG [44, offering theoretical insight into representation--nnaanee emerges as special case with a linear weight aggregation schedule and a single latent, where the output of the hybrid model is approximated by conditioning solely representations extracted from noisy inputs to mitigate the training-inferenee inference gap (Section 2.3). We also provide a TV distance error bound for the sample distribution of representation-allge alligned diffusion models (Appendix C.2).
diffusion models. In particular, REPA
Built on the theoretical insights, our general framework allows for much broader design space for exploiting representation alignment [61] in diffusion models. In particular, we demonstrate the effectiveness of two novel strategies, designated as REED
prtresennationEEhhanee Elucidation of Diffusion). First, we argue that diffusion models can profit from diverse representation sources, hence we incorporate multi-latent representations from different modalities to provide complementary information. Our new multimodal alignment scheme uses synthetic paired data and multimodal pretrained models trained on datasets from different corpora using discriminative or self-supervised tasks to generate the target representations (Section [..]p Second, we analyze optimal training curricula for representation-aiie aligned diffusion models, demonstrating that representation learning is crucial in the early training stage to yield substantial advantages. Consequently, we introduce a novel curriculum that applies the representation learning loss from the outset, while progressively increasing the diffusion loss coefficient from Finally, we propose three practical instantiations of our framework that yield improvements in image, protein, and molecule generation (Section 3.3)
its maximum value a phase-in protocol (Section (3.2).
We extensively evaluate our proposed REED across image generation, protein inverse folding, and molecule generation tasks. For the class-conditional ImageNet 256 x 256 benchmark, aligning SiT-XL/2 with DINOv2 42) and Qwen2-VL (55) representations using our improved curriculum significantly improves the generation quality measured by FID scores. REED achieves 2 training speedup over the original SiT-XL, reaching FID=8.2 2 in only 300K training iterations (without free guidance [19]; and 4x speedup over REPA [61) matching its classifier free guidance performance at 800 epochs with only 200 epochs of training (FID=1.80). For protein inverse folding, aligning discrete diffusion models trained on the PDB dataset with AlphaFold3 [I] structure and sequence representations accelerates training by 3.6x and yields significantly superior performance across metries such as sequence recovery rate, RMSD and PLDDT. For molecule generation, aligning state-of-the-art geometric equivariant flow matching (SemlaFlow (25)) with pretrained Unimol [64] representations improves metrics such atom and molecule stability, validity, energy, and strain on the challenging Geom-Drug (3) datasets.
We summarize main contributions follows:
We conduct rigorous analysis for representation eepresentaioon-enaaneed diffusion model training, delivering theoretical insights into prior empirical methods and informing the design of new strategies.
We systematically investigate the design space for representation guidance in training diffusion models, proposing a novel multimodal representation alignment scheme via synthesized paired multimodal data, an effective training curriculum, and practical domain-specific instantiations.
Extensive experiments demonstrate that our method consistently accelerates training and improves performance across various diffusion model types and diverse tasks, including image, protein sequence, and molecule generation tasks.
2

---

2 Theoretical Characterizations
In this section, we introduce a general theoretical framework for representation-enaaree diffusion models. Building on the DDPM framework (20), we explore how representations can offer additional information and enhance the generation process. Notably, due to the known equivalence between DDPM, score matching [50], stochastic interpolants [2]. and flow matching [88, our insights are applicable to broad class of diffusion model variants. Proofs for this section are detailed in
Appendix C
2.1 Variational Bound with Decomposed Probabilistic Framework
Problem Formulation. We consider a generative modeling setting involving discrete-time latent variables Toer and an auxiliary representation 2, structured according to the Markov chain:
FT æT- -1 270
(I)
The data distribution follows the factorization q[<0o:; =) (3 (x1)))(111r)))) ( ( (2))) where (3rg) denotes the empirical data distribution, (: (m)) corresponds to the forward diffusion process, and (2) represents the mapping of the data semantic latent space through pretrained encoder.
Our goal learn a generative model pe (zo). defined implicitly through joint model pe (T 2). such that pa (zo) ;= pal (o,T. To 2)dx1. Following the variational framework used in diffusion- based generative models [20 48 (see Appendix [B for details), derive variational bound on the negative log-like ligelikld of (C)) that incorporates the latent representation guide the generation process. The bound LB is formalized in Equation (2):
log ((=|zo) IIT ((11-1 ((rT 20) 2 -gg(zo) log (x0o
LiB Egixo
(2)
(o(TT, 2)
Decompasition of Model Parameterization. The denominator of Lin. po (o 2). defines how the latent variable model parameterized for the reverse diffusion process. The way in which is integrated into the reverse diffusion steps- speciifiill the inference of z;-1 from noisier state rz-is crucial for representation -eneenatttoo-eeenneeed diffusion models. When z introduced earlier in the process (at larger values of t), the model benefits from more accurate estimation via po (2 (2 2) rather than po(s,-1 [at)) since conditioning on 2 leads to a better approximation of the posterior q( 20).
However, this comes at the cost of less accurate estimation with (=(2),,), as the dependency path between Fe and becomes longer based on the Markov structure, making harder to approximate the true posterior q( (0)) This is formally reflected in the decomposition of joint probability p(zO.T, z):
T II
p(zorT 2) = p(=|xx) II p(=
p(==-1lz;) p(ar) = p¹ (TOTT. 2). Vt= 0, ..., 7 #=++1
(3)
Remark 1. Each decomposition above is an exact equality for any t. The optimal components p(æ;-1læ,, 2) and ((z;-1(x;) matching the forward process q 2) are consistent across all ts. Thus, identical components can be used in every decomposition without loss of accuracy.
Remark 2. At any timestep I, the marginal distribution of x; remains invariant across all decomposi- tions. Consequently, = can be resampled at any step while preserving the same marginal distribution.
The above remarks indicate that representation-guide without incurring additional loss. To control this combination, we introduce a novel timestep weighting schedule [ax)]=m where a, 2 are hyperparameters with Σ as = This produces an equivalent formula of the joint distribution, i.c., pe (Iror, 2) = Σt cxpp (wo-r, 2). which allows us to adjust when and how the representation 2 is injected into the generative process. Such formula results in a hybrid conditional distribution ğe (2:--//a,, 2; A.) (Equation (5)) governed by the cumulative weights A; := Σ a. A, decreases with and serves as schedule for gradually introducing guidance from z by interpolating between guided and unguided reverse diffusion. By choosing dg and thus A,, we can flexibly control the influence of z throughout the diffusion process. The formal expression for the resulting Ln is given below.
decompositions can be arbitrarily combined
WWhen f = 0 or t T'. the corresponding products reduce to 1.
3

---

Proposition (Decomposition Structure of the Variational Bound), Let (as 2 03%-1 be set of weights summing to one, and define A. E [0,1] = Eiz as Then, the variational bound Lip in Equation (2) can be written as:
Live
E T.)
log g(24-1 Fe, vo)
Ed( log Ziyzz,rr.. 2; Ai]
Pa (0 12,, 2j As)
ar (( (ea(y()lxo)lpacaaaa)))
(4)
where po (me-1 10--1 2; Ae) and its normalization Ze(ar, 2; As) are defined as:
po Ac)
2) Po p (2 2)pb At (2
(5)
(z 23 /
pg" (10-1111, (x (
(6)
The first term in Equation corresponds to a KL. divergence between the data distribution (1 and the hybrid model estimation po ( 1 2; 2; At), serving as the primary data generation objective. The second term, involving the expectation of log-normalization constants, can be calculated as the discrepancy between conditional and unconditional model predictions, multiplied by the coefficient A,(1 - At) (see Appendix 4.1.4 for details). The third term is the KL divergence between the predicted (=( and the reference distribution (=(xx)) Optimizing this term drives po(2(2;)) to better approximate the true (but intractable) (=2(2;) complementing the second term. Together, these terms facilitate the representation learning procedure where the model uses 24 to recover meaningful semantic information about 2.
2.2 Multi-Latent Representation Structure
The expressiveness of latent representations can be significantly enhanced through multi-latent structure, as shown in prior work such as NVAE [52). Instead of modeling a single representation 2. extend formulation to series of representations (=)f=1+ each indicating different level of abstraction. This leads to a similar Markov chain: 2r -T-1 representing the collection of representations from pretrained models, zt (zo). Thus, the variational bound Cū extends naturally to C where all instances of are replaced with z². Notably, the generation process of zz in the third KL divergence term of Z can be realized with a sequential structure, where each representation level depends on both the noisy input xy and all previous latents %
z² with z' = { H
21
Dxa,(((z" ((oo) [ps lax)))
Σ Eg(eer mo)
(2oo(((( z<1))
This multi-latent formulation enables the model to leverage semantic features across different gran- ularity. Previous hierarchical VAE literature 49 52 found that deeper layers tend encode coarse-grained global structure while shallower layers capture local details. Our experimental results (Table 4 confirm this pattern, showing that different hierarchical layers contribute complementary information at varying resolutions, improving both expressiveness and controllability in generation.
2.3 Connecting Existing Approaches to pntrrseentatin-EEhaaeeee Generation
Our theoretical formulation provides general and flexible framework for representationneehamme generation by introducing two key components: customizable weighting schedule {ar Y and hierarchical latent structure [= This framework recovers several existing approaches as special cases, in particular RCG (34 and REPA [1..
When 1, and a 8r, i.c. the representations come in fully at the beginning, we have A, 1, vt as a constant. Thus, the hybrid distribution pronstamt (1 [r, z; At) = P z).
pa
"Equation assumes conditional independence must still account for their interactions when inferring from noisy zt. For dependent ž) 's given do. a similar equation applies by replacing g( (ra) with an appropriate factorization of 9(z (20) that reflects the specitic dependency structure.
given Zo. i.e.. q(z (ro) II gi (z) (zo). though ps

---

and the second log- normaalaliziinnnn term in Equation two-stage procedure: first, a latent representation z sampled from the generative model po approximating the true marginal (=)): second, the sequence -1,-,,-.; is generated conditionally on the sampled ~ (z)... This formulation corresponds to the two-stage design in RCG [34), where representation modeling and data generation are decoupled. However, such decoupling poses challenges for learning po(z), particularly when z is high-dimensional or encodes rich semantic information. Since must be inferred independently of Iy, the model is unable to leverage context from the current diffusion state during inference, limiting its expressiveness and sample efficiency.
4 equal to zero. In this case, generation follows
In contrast, when L = and 0 1/T. Y t, the cumulative weights decrease linearly over time, yielding A, = t/T. In this setting, the hybrid distribution p9 plineer (xc-1 2; At) x po(#:-11 [r: 2) (=(=(1112x) generation gradually increases as the noise level decreases (i.e., as becomes smaller). However, this setup leads to a training-inference mismatch. During training, the conditional model (o (rt receives the ground truth representation
1 reflecting that the influence of the representation on x-
(==(x))) expectations in Equation (4) are taken over the data distribution q(.), while during inference this representation is not available. Instead, the model must rely
estimate based the current available Iı, i.c., ji the mean of pe (=(2)))) Notably, this estimation becomes more accurate at smaller t, when 2y contains less noise, which naturally matches the increased reliance in the hybrid distribution plincrr as t. decreases.
To address this mismatch, we approximate the hybrid distribution pe ( (x1-1 2; A²) using po( (25) throughout both training and inference. This allows the model to progressively infer z from z, through Po (2) ensuring that the representation guidance becomes more prominent as the process denoises, in line with the linear schedule (see Appendix for details). The second log-normalization term in the objective, which empirically is much smaller than the main denoising loss as supported by our experiments (Appendix [.T]. J. can be omitted from training-it also suffers from training-ine rraiinn-inffrrnce inference discrepancies. Further, under the von Mises-Fisher distribution assumption, the KL divergence in the third term simplifies to the cosine similarity between mean vectors of the respective distributions (sce Appendix C for details). Altogether, this results in the REPA framework [61], which unifies representation extraction and data generation within a single coherent architecture via a representation alignment loss, and different parts of the network specialize different stages of the generation process. Our experiments (Appendix [II validate the menteneet approximations, showing that training with the exact Lia (Equation (3)) achieves similar performance to REPA, while the latter offers a concise objective, simpler architectures, and lower computational cost.
Our method in Section builds upon the REPA setting, adopting the linear weighting schedule and ensuring end-to-end coherence between training and inference. Importantly, we move beyond the original L =1 setting to the more general case of L > enabling the novel use of richer multimodal representations through synthetic data. Additionally, inspired by our theoretical framework, we propose a better training curriculum that dynamically balances the data generation and representation modeling objectives, further enhancing model performance and flexibility.
3 REED: Representation-EEnanned Elucidation of Diffusion
Building on the insights provided by our generic analytic framework, we present two novel strategies: integrating multimodal representations with synthetic data, and designing an improved training curriculum. We further showcase the versatility of our approach by introducing three practical instantiations across diverse domains, spanning image, protein, and molecule generation.
3.1 Multimodal Representations with Synthetic Data
As discussed in Section 2.2 the effectiveness of representation-enhaiee generation depends on both the quality of pretrained representations and the interplay between hierarchical levels {zr (= To effectively harness the information from different sources, we propose integrating representation z in secondary modality y alongside the representation z" of the primary modality æ. Representations n that are both helpful for score estimation and complementary to z" yield the greatest improvement.
This cross-modal approach enhances generation quality by leveraging distinct yet related information encoded different perceptual channels. also allows integration of different embedding spaces (c.g., VAE space for latent diffusion and DINO embedding space for images), leading to an optimal
5

---

Alignment
Sequence & Structure Embedding
mage
stallow The faningo pintish-eeoen beti disiccivee yelow head tuned image
Model
Diffusion
Diffusion Model
Model
Image
sightly the
Captioning
cover nagarine in hast ooer MANAGEMENTT
tile depiyed promincnll
Alignment
Alignment
VIM
VLM
2D-3D Embedding
MRL Model
Embedding
Figure 2: We synthetic auxiliary data modalities and multimodal representations enhance representation alignment in diffusion model training. Shown are examples from class-conditional image generation (left), protein inverse folding (lop ligh/), and molecule generation (bottom right).
shared solution space. Such advantage also aligns with the Platonic representation hypothesis 23 which suggests that representations learned on different data and modalities tend to converge toward a unified model of reality. Consequently, with finite available data, incorporating data from other modalities should enhance generation comprehension in the original modality of interest.
Utilizing the representations 2" (ar) and 2yy) of the data pair (a. (). the resulting representation generation loss, corresponding to the third term LVe. is:
Cregee (0, 2/2, viy) == E A (2)imm((()))))))))))))))"()( (x)), proje, (h/=)) + sy.sim(="(v) (v), proj (b4
(8)
where I1- and N1 represent the features extracted from the Lt and ty-th layers the diffusion transformer encoder h; == fa(x)) (x, denotes the noisy inpu to the diffusion model at time step t), proja, and projž, are trainable projection heads parameterized by tha and ehy, sim(-,) :) is predefined similarity function (we cosine similarity in practice), and hyperparameters A Ag +0 controls the relative importance of aligning each representation.
However, directly using multimodal data requires paired samples across modalities, which are often scarce and costly to obtain. the context of conditional generation, where the model generates outputs based on specific conditioning variable (c.g.. class label input sequence), it may also introduce shortcuts between the conditioning variable and the representations. For instance, aligning the representation of the input condition an additional modality may cause the model to simply replicate the condition, rather than extract meaningful features. To address this issue, we generate synthetic multimodal data y == fz( using an auxiliary multimodal model fa that bridges modalities x and y (see Figure2), This introduces stochasticity and diversity into the data pairs, enabling efficient data augmentation that fosters model generalization and mitigates shortcut risks.
In practice, we use multimodal models (between x and y) instead of single-modality models to obtain the representations z for other auxiliary modalities because their representations are inherently aligned across modalities x and y. For images, we use text as an auxiliary modality, generate captions with Vision Language Models (VLMs), and extract multimodal features from intermediate layers of the VLMs. For protein inverse folding, we generate auxiliary structures with folding models and use their single, pair, and structure embeddings representations. For 3D molecules, we treat 2D graph structure as an auxiliary modality, using external tookits like RDKit to extract atom and bond features. We precompute and store these representations before diffusion training, so REED incurs minimal additional computational overhead compared to REPA. Further analysis of computational cost provided in Appendix E.
3.2 Curriculum of Representation Learning and Diffusion Training
We analyze optimal curricula between the two components, representation modeling and data gen- eration, introduced in Section 2.1 We contend that in REPA [6] type of representation-alignee
NNiiy VAE features for images, directly noisy data for proteins and molecules.

---

diffusion model training. integrating pretrained representations is particularly useful in the early stage of diffusion training. This also supported by our experimental observation that decaying or dropping REPA loss in the late plateau of training phase yields comparable or even superior performance. Our theoretical analysis in Section 2 key insights into the phenomenon. Recall the training loss bound in Equation 4 where the first term is for the approximate conditional distribution based on estimated representations pi, and the last term forces the model to predict ground truth representa- tions. Intuitively, better representation estimations jig naturally benefit oinesentttiooncoonditineedd distribution modeling in the first term.
To leverage these insights, we introduce a joint training curriculum for Lis with special emphasis on representation generation in the initial training phasc-rather than adopting rigid two-stage approach which may suffer from disjoint optimization dynamics and misaligned representation spaces (see more discussions in Appendix D). Specifically, the resulting REED loss at epoch n is given by:
REED a()) Laiffusion + 3(m) Grepgen
Here, the diffusion loss weight ex(n) progressively increases from zero following phase-in protocol, while the representation alignment weight ((n) remains fixed or decays. This curriculum enables coherent training flow. Moreover, by exposing the model to imperfect j early on, we encourage it to learn under noisy representation conditions, which can improve its generalization and robustness. The reasonability of our arguments and the effectiveness of the proposed schedule are well supported by experimental evidence.
3.3 Instantiations of the Framework for Domain Representations
In this section, we consider several application areas and incorporate additional inductive biases into the design of representation learning techniques. Our designs demonstrate the effectiveness of representation learning methods for Euclidean, sequential, and graph-structured/33 point-cloud data. To the best of our knowledge, are the first apply representation alignment diffusion models in the context of protein inverse folding and molecule generation.
Images. As discussed in Section 3.1 text embeddings-to enhance single-modal diffusion model training. Image features are drawn from pretrained self- supervseee models offer low-level details, while VLMs generate synthetic captions and corresponding cross-modal embeddings for high-level semantic guidance. In particular, for better alignment between different representation spaces, we use pooled image and text embeddings from VLMs, rather than relying soiely on pure text embeddings in dingle-modal LLMs.
we adopt multimodal
ciweeenttatins-combiiiii image and
Our analysis in Section 2.2 suggests that shallow latent layers capture fine details, whereas deeper layers represent global semantics. Accordingly, we align shallow layers in the diffusion transformer with SSL image patch embeddings, and deeper layers with pooled VLM embeddings. This aligns with empirical observations: due to the discrepancy between the latent space of external models and the diffusion output (VAE) space, local regularization should be performed in shallow layers as observed in [61;; however, global alignment via pooled embeddings is relatively light and can be applied to deeper layers for improved semantic control.
Proteins. We focus on protein sequence design conditioned on the protein 3D backbone structure protein inverse folding which can be framed as a conditional generation task using discrete diffusion models. To expand structural variability and avoid potential shortcuts mentioned in Section 3.1 we employ folding model, AlphaFold3 (AF3) [ to generate auxiliary structures from target sequences. Since diffusion sampling with AF3 introduces detailed variability into synthetic structures, using the structure representation in the AF3 diffusion head can improve generalization and controllability of our inverse folding model. We obtain sequence representations from both single token embeddings and pair embeddings from AF3's Pairformer, where the latter models residue interactions.
Since protein inverse folding can be regarded as the reverse of the folding process, we align layers of our discrete diffusion model with those of folding models in the reverse order. Specifically, early encoder layers in our discrete diffusion model align with representations from folding decoders (guided by structure), while deeper decoder layers of our denoising network are trained to match single and pair sequence encoder representations from folding models.
Molecules. Generative models for molecules must respect key symmetries, such as permutation symmetry (S(N) for 2D graphs and Euclidean symmetries (E(3) or SE(3)) for 3D structures.
7

---

Table 1: FID comparisons with vanilla SiT's and REPA of class-conditional gen- eration ImageNet 256 256. No classifier-free guidance (CFG) is used. Iter. denotes training iteration.
Table 2: System-level comparison ImageNet 256 x 256 with CFG. Models with additional CFG schedul- ing 30] are marked with asterisk (*).
Model
Epochs
FID
1SI
Pre.
Rec.
Pixel diffusion
Model
#Params
Iter.
FID
ADM-U
400
3.94
186.7
0.82
0.52
VDM++
560
2.40
225.3
SiT-B/2
130M
400K
33.0
Simple diffusion
2.77
211.8
REPA
130M
400K
24.4
CDM
158.7
REED (ours)
130M
400K
19.5
Latent diffusion, U-Net
SiT-LJ2
458M
400K
18.8
LDM-4
200
3.60
247.7
0.87
0.48
REPA
458M
400K
REED (ours)
458M
400K
8.9
Latent diffusion, Transformer
U-Net hybrid
U-VIT-H/2
240
2.29
263.9
0.82
0.57
SiT-XL/2
675M
400K
17.2
DiffTT
1.73
276.5
0.80
0.62
REPA
675M
150K
13.6
MDTv2-XL/2
1080
1.58
314.7
0.79
0.65
REED (ours)
67SM
150K
11.6
Latent diffusion, Transformer
SIT-XL/2
675M
7M
8.3
MaskDiT
1600
2.28
276.6
0.80
0.61
REPA
675M
400K
7.9
SD-DIT
480
3.23
REPA
675M
IM
DIT-XL/2
1400
2.27
278.2
0.83
0.57
REPA
675M
4M
SiT-XL/2
1400
2.06
270.3
0.82
0.59
REED (ours)
675M
300K
REPA
200
1.96
264.0
0.82
0.60
REED (ours)
675M
400K
REPA
1.80
284.0
0.81
0.61
REED (ours)
675M
IM
4.7
RRED (ours)
200
1.80
267.5
0.81
0.61
Following 35 we use graph-level invariant representations from pretrained molecular representation learning (MRL) model improve generation quality with minimal computation overhead. We adopt diffusion and flow matching models with equivariant GNN backbones to maintain these symmetries, and incorporate lightweight graph-level invariant representations from pretrained S molecular encoders. Instead of the strictly two-stage representation-condiiooee generation in (35), we show that directly aligning diffusion models with pretrained global representations significantly accelerates training and enhances generation quality.
4 Experiments
We evaluate the effectiveness of our proposed REED through extensive experiments. The superior performance across image, protein and molecule domains validates the universal benefit of our framework for both continuous and discrete diffusion or flow-matching models.
4.1 Image Generation
Experimental Setup. For continuous diffusion on Euclidean data, we test REED on class-conditional image generation following the setup in REPA (61. We conduct experiments ImageNet [12] at 256 x 256 resolution, adopting ADM [3] data preprocessing and encoding each image into a latent vector with the Stable Diffusion VAE [46]. We utilize the B/2, L/2, and XL/2 architectures from SiT [40]. For sampling, consistent with SiT and REPA, we employ the SDE Euler-Maruyama sampler and fix the diffusion timesteps to 250 for all runs. We report Fréchet inception distance (FID (18)), inception score (IS
[47 precision (Pre.) and recall (Rec.) 29 using 50,000 samples for evaluation.
2B to generate a synthetic caption for each image and compute joint image- text representations with Qwen2-VL 7B, using both the image and caption as input. We extract the 15th-layer representations from Qwen2-VL 7B, averaging across all tokens for high-level joint embedding, and align this with the averaged SiT image latents using a cosine similarity loss and MLP projector. We choose the 15th layer for its strong semantic representation;
We use Qwen2- VI. 55
ablation of different
layers is provided in Appendix F.1 For training curriculum (Equation ( we keep 3(n) constant and linearly increase (n) over the first 50K training iterations, after which it remains fixed for the rest of training. Detailed experimental settings are provided in E
Baselines. We compare against recent diffusion- generation baselines with varying inputs and
architectures (details Appendix E. (a) pixel diffusion models (ADM [3], VDM++ (7). Simple diffusion 1221. CDM 21),, (b) latent diffusion with U-Net (LDM 46). (c) latent diffusion using

---

transformer+U-Net hybrids (U-ViT-H/2 4 DiffIT 16 MDTv2-XL/2 [5] and (d) latent diffusion
with transformers (MaskDiT 63) SD-DiT 651, DiT 43 SiT 40 and SiT with REPA 6T]).
Overall Results. We provide FII comparisons in Table I REED consistently outperforms vanilla SiT and REPA the same training iterations across all SiT model scales. In particular, REED achieves FID=8.2 on SiT-XL/2 at 300K steps, surpassing vanilla SiT-XL at 7M steps, and reaches FID=4.7 at IM steps, outperforming REPA 4M. Furthermore, we provide comparison between REED and other recent diffusion model for images using classifier-free guidance (CFG (19). As shown in Table REED matches REPA performance with 4x fewer epochs, achieving FID=1.80 at just 200 epochs. Further details and additional metrics are provided in Appendix F.I.
Ablation Study on Each Algorithm Component. To examine the effect of each component of REED, including leveraging multimodal representations and applying the curriculum training schedule, we experiment the SiT-B/2 architecture and report the results for models trained for 200K iterations and not use classífier-free guidance for generation. As shown in T removing each component leads inferior performance than the full REED model, with FID increasing from 24.6 29.4 .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 ablations on different types of curriculum schedule and multimodal representation are provided in Appendix F.
Table 3: Ablation study each component of REED. We report results SiT-B/2 training for 200K iterations. No classifier-free guidance is used.
Model (SIT-B/2. 200K Iter.)
FID
ISt
Prec.t
Rec.1
+ REPA
33.2
43.7
0.54
0.63
+ REED
24.6
62.2
0.58
0.64
+ REED w/o Multimodal Representations
27.7
55.5
0.56
0.64
+ REED w/o Curriculum Schedule
29.4
50.8
0.56
0.64
Alignment Depth for Multimodal Representations. We in- vestigate the optimal alignment depth for image (i.e., Dep.-I) and VLM (i.e., Dep.-VL) representations in Table 4 Aligning image features at a shallow layer (i.e., 4) and text at a deeper layer (i.e., 8) yields the best results, consistent with our theo- retical analysis (Section (.2) that shallow layers capture fine details while deeper layers encode high-level semantics. In
Table 4: FID comparison differ- ent alignment depth SiT-B/2.
Iter.
Dep.-I DTp
FID.
200K
25.9
200K
4
25.9
200K
4
6
24.8
200K
4
24.6
contrast, for VLM-only aligmment, the best performance comes
200K
4
10
26.4
from aligning at shallow layer 4, indicating the above effect is specific to interactions between multi-latent representations.
200K
45,5
200K
45.6
4.2 Protein Sequence Generation
Experimental Setup. For discrete diffusion sequence data, experiment on protein inverse
protein sequences conditioned on backbone 31) structure. Following [54], we train an inverse folding model using the Multiflow 7 objective and the ProteinMPNN []] architecture, with the PDB training set from [1]. We filter for single-chain proteins with sequences under 256 residues, yielding 13,753 training and 811 test proteins.
We use AlphaFold3 I to generate auxiliary structures from target sequences, extracting latents from its diffusion head structure representations, well both single and pair embeddings of Pairformer as amino acid representations. We align the ProteinMPNN encoder output with the structure representation, and the decoder's first-layer output with the single and pair token representations, reflecting the reverse nature between inverse folding and folding. Regarding training curriculum, (m) is kept constant, while a<n) linearly increases over the first. 50 epochs, then decays with a cosine schedule.
We assess inverse folding performance using sequence recovery rate (Seq. Recovery), self-consistency RMSD (scRMSD), and PLDDT from ESMFold (36) on generated sequences. For each test structure, we generate sequence and report the median scRMSD and pLDDT across the test set. We also compute the proportion of cases with scRMSD below 2 and PLDDT above 80, standard thresholds for successful inverse folding 7 41. Further details of the metrics are in Appendix E.2
Evaluations.

---

Table 5: Model performance on protein inverse folding. REED achieves high sequence recovery rate and pLDDT and low scRMSD across various epochs and significantly accelerates training. We report the mean across 3 random seeds, with standard deviations in parentheses.
Methods
Epochs
Seq. Recovery(S))
scRMSD]
S(scRMSD=2) f
PLDDT
*<pLDDT>80($>
ProteinMPNN
200
41.73 (0.08)
1.801 (0.035)
54.45 (0.68)
82.68 (0.25)
65.66 (1.39)
DistreteDiff
50
38.29 (0.03)
2,047 (0.018)
49.11 (0.29)
81.52 (0.13)
58.80 (0.21)
REED (ours)
50
(0.03)
1,788 (0.033)
54,37 (1.09)
82.74 (0.08)
67.41 (0.38)
DiscreteDiff
100
40.74 (0.04)
1.789 (0.015)
54.17 (0.61)
82.80 (0.08)
67.57 (0.62)
REED (ours)
100
41.91 (0.07)
1.780 (0.017)
56.23 (0.62)
83,09 (0.14)
68.41 (0.38)
Diff
150
41.18 (0.08)
1.828 (0.046)
54.87 (1.02)
82,92 (0.08)
67.25 (1.02)
REED (ours)
42.26 (0.12)
1.799 (0.016)
54.66 (0.53)
83.17 (0.12)
69.07 (0.37)
Results. We compare REED against the de facto inverse folding method ProteinMPNN and and discrete diffusion models without REED (DiscreteDiff), with results in Table REED consistently outperform other methods in all metrics, demonstrating the effectiveness of incorporating powerful representations in the protein domain. Mcanwhile, REED accelcrates training, reaching 41.5% sequence recovery rate just 70 epochs, whereas DiscreteDiff requires 250 epochs. Ablation studies on different combinations among single, pair and structure representations are in Appendix F2
4.3 Molecule Generation
Experimental Setup. For structured and geometric graph data, we choose the 3D molecule generation task with equivariant point-cloud diffusion models. Leveraging canmetry-preservinn representations and equivariant backbones, approach effectively models molecular graph distributions while
respecting their inherent symmetries. We adopt the challenging Geom-Drug 3 dataset, which
contains large drug-like molecules (average 44 atoms). Following SemlaFlow [25]. we filter
molecules with more than 72 atoms in the training set and use standard splits. Our base model is
SemlaFlow [25] with ((3)-qquivariant backbone and the ODE flow matching framework, and use
the Unimol (64] further finetuned on GEOM-DRUG as the pretrained encoder. All other training and
evaluation settings follow 25
Evaluations. We adopt standard benchmark evaluation metrics: atom stability, molecule stability, validity, well as additional metrics energy, strain, and sampling efficiency (NFE). Full metric details are in Appendix atoms and bonds, including MiDi [53), EQGAT-diff (31), and SemlaFlow 255 since those approaches only generating atoms fall short in producing high-quality molecules as is explained in [25]. Notably, train SemlaFlow+RERS for just 200 epochs, compared to 800 epochs for EQGAT-diff.
E3 We compare against state-of-the-art 3D generators that jointly generate
Table 6: Performance of REED with SemlaFlow as the base model on GEOM-DRUG dataset. Atom stability, molecule stability, and validity reported as percentages, while energy and strain energy are expressed in kcal-mol- - Results marked with were reproduced in our own experiments. All
results are calculated based on 5k randomly sampled molecules.
Methods
Atom Stab
Mol Stab
Valid
Energy
Strain +
NFE
MiDi
99.8
91.6
77.8
500
EQGAT-diff
99.8 8
93.4
94.6
148.8
140.2
500
SemlaFlow*
99.88
97.4
93.4
114.44
74.18
100
+ REED (ours)
99.93
98.1
94.5
105.63
65.33
100
Results. As shown in Table SemlaFlow trained with our REED significantly improves almost all metrics compared to the baseline and outperforms other state-of-the-art models. Remarkably, the superior performance is achieved within much less training computation and sampling NFE compared to EQGAT-diff, verifying the effectiveness of REED.
Conclusion
In this paper, we introduced general theoretical framework for incorporating representation guidance into diffusion model training. Our proposed design strategies consistently achieve superior perfor- mance and accelerated training. across range of domains. While we adopt linear representation weighting schedule in our experiments, exploring adaptive schedules remains an intriguing direction for future research. Broader impacts include data efficiency in machine learning and biology.
10

---

Acknowledgement
CW, CZ and TJ acknowledge support from the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the NSF Expeditions grant (award 1918839) Understanding the World Through Code, and the DSO Singapore grant on next generation techniques for protein ligand binding. CZ and SB were partly supported by the ARPA-H ADAPT program. SG is supported by the MathWorks Engineering Fellowship. SG and SJ acknowledge the support of the NSF Award CCF-21 I 12665 (TILOS AI Institute), an Alexander von Humboldt fellowship, and Office of Naval Research grant N00014-20- 20-1-2023------- (MURI ML-SCOPE).
References
[1] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-500 2024.
[2] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023,
[3] Simon Axelrod and Rafael tions for property prediction and molecular generation. Scientific Data, (1):1885, 2022.
Geom, energy-annotated molecular conforma-
[4] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao. Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 22679, 2023.
[5] Normand Beaudry and Renato Renner. An intuitive proof of the data processing inequality. arXiv preprini arXiv: 1 2011.
[6] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo. Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, et al. Video generation models as world simulators. OpenAl Blog, 1:8, 2024.
[7] Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gener- ative flows on discrete spate-spaces; Enabling multimodal flows with applications to protein co-design. arXiv preprint arXi::2402.0499) 2024.
[8] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang, Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference Learning Representations, 2023. URL https
[9] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang Sampling is easy learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference on Learning Representations, 2023. URL https
/ wet.://rree.tttooommo////yyyyyg
[10] Xinlei Chen, Zhuang Liu. Saining Xie, and Kaiming He. Deconstructing demoising diffusion models for self-supervised learning. arXiv preprint 7xXi:22401144000 2024.
[11] Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, al. Robust deep Haaming-basee protein sequence design using proteinmpnn. Science, 378(6615):49-5) 2022.
[12] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. 2009 IEEE conference on computer vision and pattern recognition, pages 248 255. Ieee, 2009.
[13] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances neural information processing systems, 34:8780-8788 2021.
[14] Stefan Elfwing, Fiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
11

---

[15] Shanghua Gao, Pan Zhou, Ming-Ming Ming Cheng, and Shuicheng Yan. Mdtv2: Masked diffusion transformer is a strong image synthesizer. arXiv preprint 1...:22333.14389 2023.
[16] Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. Diffit: Diffusion vision transformers for image generation. In European Conference on Computer Vision, pages 37-55. Springer, 2024.
[17] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000-16009, 2022.
[18] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems. 30, 2017.
[19] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207. /259 2022.
[20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, Advances in neural information processing systems, 33:6840 6851, 2020.
[21] Jonathan Ho. Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47):1-33, 2022.
[22] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pages 13213- 13232. PMLR, 2023.
[23] Minyoung Huh, Brian Cheung. Tongzhou Wang, and Phillip Isola. The platonic representation hypothesis. arXiv preprint arXiv:2405. 07987, 2024.
[24] Ross Irwin, Alessandro Tibo, Jon Paul Janet, and Simon Olsson. Efficient 3d molecular generation with flow matching and scale optimal transport. In ICML 2024 Al for Science Workshop, 2024.
[25] Ross Irwin, Alessandro Tibo, Jon-Paul Janet, and Simon Olsson. Efficient 3d molecular generation with flow matching and scale optimal transport. arXiv preprint arXiv:2406.07266. 2024.
[26] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv: 2017.
[27] Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36:65484-65516. 2023.
[28] Diederik P Kingma. Adam: method for stochastic optimization. arXiv preprint 14Xiv:1412.6990 2014.
[29] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
[30] Tuomas Kynkäänniemi, Miik Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance a limited interval improves sample and distribution quality in diffusion models. In Proc. NewrIPS, 2024.
[31] Tuan Lc, Julian Cremer, Frank Noé, Djork-Amé Clevert, and Kristof Schütt. Navigating the design space of equivariant diffusion- generative models for de novo 3d molecule generation. arXiv preprint arXiv:2309. 2023.
[32] Alexander Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak. Your diffusion model is secretly zero-shot classifier. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 2206-2217, 2023.
12

---

[33] Daiqing Li, Huan Ling, Amlan Kar, David Acuna, Seung Wook Kim, Karsten Kreis, Antonio Torralba, and Sanja Fidler. Dreamteacher. Pretraining image backbones with deep generative models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 16798-16708, 2023.
[34] Tianhong Li, Dina Katabi, and Kaiming He. Return of unconditional generation: A self- supervised representation generation method. Advances in Neural Information Processing Systems, 37:125441-125468 2024.
[35] Zian Li, Cai Zhou, Xiyuan Wang, Xingang Peng, and Muhan Zhang. Geometric representation condition improves equivariant molecule generation. arXiv preprint arXiv:2410.03655, 2024.
[36] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary-scale prediction of atomic-level protein structure with language model. Science, 379(6637):1123- 2023.
[37] Zongyu Lin, Wei Liu, Chen Chen, Jiasen Lu, Wenze Hu, Tsu-Jui Fu, Jesse Allardice, Zhengfeng Lai, Liangchen Song, Bowen Zhang, et al. Stiv: Scalable text and image conditioned video generation. arXiv preprini 24 2024.
[38] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.0274 2022.
[39] Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rXivv1711.05101, 2017.
[40] Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference Computer Vision, pages 23-40. Springer, 2024.
[41] Hunter Nisonoff, Junhao Xiong, Stephan Allenspach, and Jennifer Listgarten. Unlocking guidance for discrete state-space diffusion and flow models. arXiv preprint 1775 2024.
[42] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07195 2023.
[43] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195-4205, 2023.
[44] A Polyak, A Zohar, A Brown, A Tjandra, A Sinha, A Lee, A Vyas, B Shi, CY Ma, CY Chuang. al. Movie gen: A cast of media foundation models. Meta Al Blog Post, 2024, 273403698. 2024.
[45] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
[46] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[47] Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung. Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
[48] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-22666666 2265. pmlr, 2015.
[49] Casper Kaae Sønderby, Tapani Raiko, Lars Maalge, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Advances in neural information processing systems, 29, 2016.
13

---

[50] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1. 13456, 2020.
[51] Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 2818 28222, 2016.
[52] Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Advances in neural information processing systems, 33: 33:19667-19679 2020.
[53] Clement Vignac, Nagham Osman, Laura Toni, and Pascal Frossard. Midi: Mixed graph and
3d denoising diffusion for molecule generation. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 560-576. Springer, 2023.
[54] Chenyu Wang, Masatoshi Uehara, Yichun He, Amy Wang, Tommaso Biancalani, Avantika Lal, Tommi Jaakkola, Sergey Levine, Hanchen Wang, and Aviv Regev. Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design. arXiv preprint arXiv:2410. 13643, 2024.
[55] Peng Wang. Shuai Bai, Sinan Tan, Shijie Wang, Zhihao Fan, Jinze Bai, Keqin Chen, Xuejing Liu, Jialin Wang, Wenbin Ge, al. Qwen2-vl: Enhancing rision-language model's perception of the world at any resolution. arXiv preprint arXiv:2409. 12191, 2024.
[56] Joseph I. Watson, David Juergens, Nathaniel R Bennett, Brian I. Trippe, Jason Yim, Helen Eisenach, Woody Ahern, Andrew Borst, Robert J Ragotte, Lukas Milles, al. De novo design of protein structure and function with rfdiffusion. Nature, 60099976)10889-1100 2023,
[57] Talal Widatalla, Rafael Rafailov, and Brian Hie. Aligning protein generative models with experimental fitness via direct preference optimization. bioRxiv, pages 2024-05, 2024.
[58] Weilai Xiang, Hongyu Yang, Di Huang, and Yunhong Wang. Denoising diffusion autoencoders unified self-supervised learners. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1582215812, 2023.
[59] Xingyi Yang and Xinchao Wang. Diffusion model as representation learner. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 18938-18949, 2023.
[60] Xiulong Yang, Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your vit is secretly hybrid discriminative-geneene geeerrtivegeeetiieeeeee diffusion model. arXiv preprint arXiv:2208.07 077991 2022.
[61] Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong. Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers easier than you think. arXiv preprint 2rXiv22410.06944 2024.
[62] Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training In Proceedings of the IEEE/CVF international conference on computer vision, pages 1995--11886. 2023.
[63] Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. arXiv preprini (tXXix23366093000 2023.
[64] Gengmo Zhou, Zhifeng Gao. Qiankun Ding. Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: A universal 3d molecular representation learning framework. In The Eleventh International Conference on Learning Representations, 2023. URL https
//penreview. (igeeneeieer.ee//foumeieeee
[65] Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao, Zhenglong Sun, Tao Mei, and Chang Wen Chen. Sd-dit: Unleashing the power of self-supervised discrimination in diffusion transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8435 8445, 2024.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:
[Yes]
Justification: We support all claims through rigorous proofs and extensive experiments various domains.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Limitations of our work are discussed in Section 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification:
Proofs to all theoretical results are provided in Appendix C
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All experimental details to reproduce the main results are provided in Section and Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification:
We use public data in our experiments and have released our code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ Soaedddeee/CoSSuuiisssseeee
www.SSubmisiooooiiiooooooooo for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We mention all experimental details and training protocols in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We run experiments on protein inverse folding over three random seeds and report the mean and standard deviation. For other experiments, we run one seed due to computation resource constraint and the stability of model performance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
17

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources and requirements are mentioned Appendix E
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics /ttppp////aaurpp.cc//////////c/a
Answer: [Yes]
Justification:
We obey all aspects of the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This is a method oriented paper and is not related to societal impacts. However, we still discuss potential social impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
18

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
This paper does not involve such models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the authors for their models and datasets
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
We used LLMs solely for minor writing edits.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https://neure) httppp////////ooeeeeeee///////// for what should or should not be described.
21

---

Appendix
A
Related Work
Enhancing Diffusion Models with Pretrained Representations. Recent empirical studies have explored various strategies for incorporating pretrained representations into diffusion model training. CrtrrsentationCConiiiiinnn
Generation (RCG) [34] addresses unconditional image generation by first training a diffusion model generate semantic features based on a pretrained self-supervised encoder, and subsequently conditioning second image diffusion model on these features to achieve high-qualin outputs. Representation Alignment (REPA) 61 introduces regularization technique that aligns intermediate noisy states within the diffusion process to clean image representations extracted from external pretrained visual encoders, thereby improving both representation learning and generation performance. Despite their empirical effectiveness, there remains gap in the theoretical understanding of how pretrained representations can enhance diffusion model training. Additionally, they focus primarily on image generation, leveraging visual encoder representations, without considering a diverse array of modalities or a wide range of application domains.
Unitying Representation Learning and Diff'usion Models. Many recent works seek to bridge representation learning and diffusion- based generation. Several works leverage or refine represen- tations obtained from diffusion models. For example, [60] proposes a hybrid framework capable of performing both discriminative tasks and diffusion-bar based generation, while [58] uncovers the discriminative nature of intermediate representations within diffusion models. Subsequent studies, such as [59 33 further distill these learned representations for various downstream applications. Ad- ditionally, denoising objectives have been harnessed to enhance representation learning; for instance, [T0 deconstructs diffusion models to advance denoising -based representation learning. Diffusion models have also benefited from auxiliary self- superviedd learning signals- 165] demonstrates that incorporating an auxiliary discriminative self-supervised loss can strengthen diffusion model training. In contrast to these approaches, our focus is improving diffusion model training by utilizing high-quality pretrained representations from a variety of modalities.
Generative Models for Different Domains. Diffusion-based generative models have demonstrated remarkable performance across diverse range of data types and application areas. In image generation-where data is continuous and resides in Euclidean space-denoising transformers such as DiT (Diffusion Transformer) and SiT 40 (which incorporates continuous stochastic inter- polants to enhance DiT) have been widely adopted. For protein sequence generation, which involves discrete sequential data, discrete diffusion and flow models have been successfully applied, as seen in works such as 7 41 (5). In the domain of molecular generation, where data is represented as structured and geometric graphs, -equiiuuia generative approaches have been developed to model 3D molecular structures. Notable examples include the MiDi model
(53) for joint 2D and 3D denoising diffusion, EQGAT-diff (31] for 3D diffusion, and SemlaFlow [25] an ODE flow matching model for 3D molecule data. Despite their widespread use, training diffusion- based generative models remains computationally intensive, often requiring significant resources and time. This poses challenges for broader adoption, especially in domains with complex and high-dimensional data.
B Preliminaries
Our theoretical analysis in Appendix C.2 builds on the variational framework of denoising diffusion probabilistic models (DDPM) (20 48 Below, we briefly summarize the key formulas and concepts; for further details, refer readers to the original papers.
Diffusion models can be formulated as latent variable models, po (xro) == (o where latent variables F1.T are generated by a forward diffusion process g (2o) that forms Markov chain JT- TT-1- I1 - Fo. with a predefined variance schedule Bi...., Br:
T
q (20) = II (()) (( == -N(m v Bes 3.1)
(10)
A variational upper bound for -Ellog (zo)] is then given by
22

---

(zo) (og S E (( log ((11)100)
= LvB
(11)
(o(
where Lvs be decomposed
log ((27(20) IT ( bo,, 2o)
LvB =
By((or))
(12)
po(xT) IT po (0 (it)
T
[( (xl(@(erlzollp<((r)) + Dxl. - log po
=
(13)
The reverse process (2 typically parameterized a Gaussian / (1 (2, as in [20], enabling tractable computation of the KI. divergences.
extend this framework by introducing additional latent variable 2, repre- senting high-quality pretrained representation of zp- We show how incorporating such semantic representations can enhance diffusion model training.
In Appendix C.2
Note that there are other variants of diffusion models from a different theoretical perspective, such as score matching [50]. stochastic interpolants (2). and flow matching (38). Due to their equivalence with the denoising diffusion models framework, our theoretical insights in Section are broadly applicable to these approaches. In our experiments, for image generation, discrete flow models (as in Multiflow (7]) for protein sequence generation, and -equiivrrrr flow matching (as in SemlaFlow (5 for molecule generation.
adopt stochastic interpolants (as in SiT (40)
C Additional Theoretical Results
C.1 Additional Theoretical Characterizations and Proofs for Section 2
C.1.1 Derivation of Equation
Similar to the variational upper bound of DDPM, by incorporating additional latent variable following the Markov chain 2T - IT-1
21
2g z, we have
log pe (xo) s - log (o + DKL (g(2111 2 xo) ppo (x1.7, 2|2o))
(14)
log ((11.1 2/20) + log po (zoo)
- log (o(o) + Bo(ar (moo)
(15)
pe (Z 2)
g( 2|20) log po ( 2)
(16)
Therefore,
log g0
BBy(z (ogo)( S Eg(som)))
L
(17)
po (
Further decomposing the numerator of Lip leads to:
((2)00)((F)))(( log po (
LiB = Ex((zoir.))
(18)
log g(2(20) IT ((x--11 zo) g( (zo)
(19)
Po (
Note that when 1, g( 12,,
=
23

---

C.1.2 Derivation of Equation 3
The marginal distribution p(zo.T 2) can be decomposed as
() = =)
II p(a)-1 p(=(ser)
II p(
(20)
following the order TT, ...., Ft. : 2, .... To for any 0,
Based on the Markov chain 2r Ty [.:r = p(=:((111) 1ss 2). and p(z(zer) = p(2(2x). Thus,
11 To - the following conditional independencies
hold: ( = p(zi-1la;), p (
for any & = 0, 1,.., T, we can exactly decompose p(Zu.r: 2) follows:
p(ZOTT 2) p(z|20) p(zol=1) ((21(22)
(x2) ... p(5T-1 p(oT)
(21)
p(z|01) p(zol=1, 2) p(21(12)) ... ((r--122)) p(zT)
(22)
= -((s/z)) II p(a:-1123
II p(x,-1(z;; p((T)) ==++1
(23)
= =((2|or-1) ( (1 2) ..... p(or- ((0r-2/0r-1,3)) p(ar - (xT) p(rr)
(24)
p(2|0r) p(zoll1122) 2) p(zi (22, 2) p(øT-1 (T 2) p(or)
(25)
C.1.3 Proof for Proposition]
1
Proposition (Decomposition Structure of the Variational Bound). Let {ar 2 O)T= be set of weights summing to one, and define A, E [0, == El ag Then, the variational bound Liv in Equation can be written as:
T
log (( [12, (00) 20) po 111, A1)
Lip
B Bx( r
B ( log Ze 2; Ao)
1=1
Σ ar E ((Bor()))))
(26)
where po (o Ir, 2; As) and its normalization Ze As) are defined as:
(1( Z(ze, 2) 2 2)pb (x1-11
(27)
Ze Ze 2; Ar) = /
pp A: (xe-1l 2)pb 1-44 Ar (x) xr )d2111
(28)
Proof. The marginal distribution Po(rurr, 2). which appears in the denominator of Li (Equa- tion (2),, can be decomposed as shown in Equation (3) (3) for any t = 1,.... T. By aggregating these decompositions according to the weight schedule {a, [t=1- we obtain an equivalent form of Cin
'When f 0 or t T'. the corresponding products in Equation
20 reduce 1.
24

---

T
((2(20) IT 9 mo) g(0r)220))
vB =
at B (xa.))))
log
(29)
p(z(=r))
II p(a 1 2) II p(xz-1/2;;) p(xr)
t=1
log ( F0)
Oge (g)
log
(30)
Po (re)
po (3:4-1 [i 2) i=1
t=1
((3 1 xo)
+ log g mr) po (o(T)
log
(31)
po pe (x,)
i=#+1
T
[ 1mt
T
2o) log pol (111)122(12)
1-1
az log g 102, xo)
ar)
+
(32)
po(x1-112 126)
i=
log 2(21210) po(2 (are)
log ((JT((rO) (x(r))
Org (ofooorss)
E (o(
(33)
Ae log ((30- 2ie, x00)
g [ 20) + As) log
E (o(
(34)
Po p0 ( (re, z)
po 20( (re))
T
log ((=120)
log g (2o) po(oT)
Orp Be (.fogtoor)))
Ey(xurr)))
(35)
po (mra)
g( 1 2o)
log
(36)
P8 (x p) 1-A (x0-11201)
t=1
log g(=1210) po(2(10))
Ba( log (( (aro) (x(r))
Or ((
(37)
Since po (2T) a pre-specified fixed distribution (typically standard Gaussian) that is easy to sample from, and g( converges to this Gaussian as T is large, the last term is approximately zero Additionally, is not involved in the optimization model parameters 0, thus it can omitted from Cig Following the definitions of (a(xx-11 12,, Ar) and Ze 2; At). the variational bound Lyz can then written as:
q( [0,, F'o)
LVp = By By(zurris) log
(38)
Po (xx- 1 |11,
P) A-Ar (xx-1 (re)
1=1
log ((2(200)
Σ Bg((oorr
log Ze (Ite As) + Ey(xu o Bo(zl1o)
(39)
log q(I1~11 23 mo)
E ,x(20:T,=)
E Be(eeer,)) .2)
log (e z; Ar)
Po po(x:-1 [1t, z; As)
T
+
Ore Be (((
(40)
t=1
C.1.4 Bounds on the Log-Normalization Term
To clarify the role of the second log- ne-oommallzzaionn term in the expression for LVp (Equation (4). we present the following proposition, which provides both bound and a closed-form solution under the Gaussian assumption for the reverse process:
25

---

Proposition (Bounds on the Log-Normalization Term), The log- ongnnommalizaiinn term in Equation satisfies:
T
T
0 S
log (20, 2; An) & min 2))
Σ
B
2)),
T
Specifically, under the common Gaussian assumption for the reverse process [20]. i.e., po (0 (x2) Na (1g( (z,,t), t),t),,)) of) and Po(2z-11244 ~ (ug (2,, 2,t,, o?), the term has the following closed-form
expression:
T
Ar(1
Ad) 207
:)) log Ze 2; Ai)
Σ
(o 12 (x), t)l² (41)
Proof. First, we consider the lower bound of the log-normalizatio term. Utilizing Hölder's inequality with exponents 1/A, and 1/(1 At), we have
po Al (12-11372 =) A- ( /
(e( At)
(42)
$ (f po(oL poloseslon 2)00111
Ar
1-Ac
12-A,
(43)
Thus, this establishes the lower bound:
Σ =x( log Zel 2; Ao) 2 Σ E ())))))))) log
(44)
Second, investigate its upper bound. Note that Ze(ses 2; Ar) can be written as:
7, 2; Ar) = f po(
pa p [po(sr-iluts 1252, 2) As
[ dr Fpp(
(po(sr- 12,, 2x), 2)
Ac
po(
po ( 12xe)
(45)
Leveraging the concavity of the logarithm and Jensen's inequality, have
log (1 z; Ad) 2 Epo( log (o(( 4,
po( (
(46)
log
Po Po( 12),,)
-A
-A, A, (po( 1e (47)
po (I1-1 12x))
Therefore,
T
Σ B log Ze( 2; Ao)
< (a
(48)
Similarly, using the alternative expression of Ze (I z; Ar) follows:
Ze (1 z; Ar) = / po ( 1
Po (0 1-1,
dri-1
(49)
po ( ( 1,2,,
pe (p(x==11 (It)
1-A77
-Boo
(50)
1o(
26

---

we have
log po(x (3re) 1-A4
log (z4, 2; Ac) 2 (pe 12(e00)
(51)
2)
log po po(22-1 (it)
-11-Ar (o. Po
=- Dx) z)lpo (
po ( Dt, T1,
(52)
and subsequently
T B -
B (.))
log Zo (; An)
(53)
Putting both together, we derive the upper bound the log-normalization term as
min
E /a
(54)
T
E ( - At0.://10/...x.0/-o11oo//oo (a
(55)
Finally, we calculate the closed-form expression under the Gaussian assumption for the reverse process po (x and po (1 z) and denoting d as the dimensionality of 21, i.e.
14-1 14 (3re, 2, 2,11
pe x),2)
(2mo))a/2 exp
(56)
20
1
Tt-1 p4 (3 1)12
po (4))
(2mo?)d/2 exp
(57)
20
Multiplying them together yields:
So 10 At) 20²
pe (ctt-1/20 2)p - As (2)4) = (2Tof) exp
(58)
where fo( A.) A 1 (xt, 2, 2, term. For simplicity of notation, use Me to represent ( 1) and tou to represent (x f). The exponent term f 2, A₁) can be expanded as:
1+) 12 (1x, 1) 22 is the exponent
fo [0 2, A,) = A, A, Hol2 + (1 A,) - pul 2
(59)
= AA(z1-1 - Mc)² (xt-1 - mc) - (1 (1) Uu (I -M
(60)
A1(of _ 241 It-1 + 1º Me) (1 -A) 1 x
24 x P pa)
(61)
x = 2(AMMc + (1 Ar) )² 21-1 + Aep o Me +(1 AA) Mu
(62)
= [ - Aweighted 2 + A _ Me + (1 Ad) Hu. - T
Wweiihtte
(63)
where Mweighted
Apple (1 At)juu Denote
C := Aqp Me + (1 At) T Wweiighed//wwigheedd
(64)
Then,
Po (21-1/01, 2)p 1 (x,-
( exp
[2:-1
20
Mwwiihttd
exp
20
C
(65)
By integrating over _z-1 and using the fact that the integral of normalized Gaussian density is 1, we obtain the expression for Ze (xe, 2; Ar):
C
1
|weiihted
( )
C
Z, (z,, 2; Ar) exp
20
(2ma
exp
20
dx exp
20
(66)
27

---

Further, C' can be simplified as
C Astle Me + (1 - Ae) (Arfle + (11-A1) pm) (As + (1 Ar)
(67)
A,(1 - Ad) - Me + As(1 - Ad) 1 - -2Ar(1 Ac) Nu
(68)
A.(1 - A,) Mull
(69)
This gives us the final closed-form expression:
T
Ar(1 - At)
log Ze 2; Ai]
Ie -uu
(70)
207
1,11 Ar)
(1((2,,,))) (x) (71)
207
According to Equation (41). the log-normalization term quantifies the discrepancy between the conditional and unconditional model predictions, scaled by the variance and the product of A and (1 - At). Since the first KL divergence term already encourages extracting useful information from to aid estimation, we can interpret the log- oogmalaaaiiationn term stop-gred stop-gradient applied to the first conditional term (z), 2, t). As such, it acts as an inherent regularization mechanism that encourages alignment between the two. This implicitly promotes the model extract from 2; the aspects most informative about 2, thereby improving the unconditional estimation.
C.1.5 Derivation of Equation 7
In the multi-laters representation setting, the single latent variable in Proposition T and Equation (4) can be replaced by the collection of latent variables z". In particular, the third KL divergence term which corresponds to the generation of z z"-can be structured sequentially in the order 1, 2...... L.. Formally, this is expressed as:
q(z" (z0) log
(aL(a(z"1an [xo) (2)))))))) (x=)) E (2a))
(72)
pol (z (xt)
I ((21(200)
log
(73)
IT 1 pe (ze Ze1)
((21(20))
EEg(al (10)
log
(74)
Po (0 1,,, 221)
q(2i ((2)) log pe (z) 2ee 221)
Eg(zed (o)) Be(z((
(75)
[(( [
(76)
C.1.6 Further Justifications of Section 2.3
we address the training-inference discrepancy arising from the lack of ground-truth representation during inference by approximat- ing the hybrid distribution plincer(z
Justification of the Distribution Approximation. In Section 2.3
[rz 2; (dl)-associater with linear cumulative weights- with po(#s-1 (r,, (.) for both training and inference. Here, juf denotes the mean of the current best esti- mate Po Po(2|20 (2(2z). We contend that this approximation enhances representation guidance as the process denoises, which is consistent with the linear weight schedule. We provide addition justifications such approximation in this subsection.
We first examine the behavior of the hybrid model as described in Proposition
28

---

Proposition: (Score Function of the Hybrid Distribution). The score function of the marginal
distribution Po(a (1(2--122; At) implied by the hybrid reverse process Po As) defined in
Equation satisfies:
R2-1 log 1og
= A,V x. Po 112) +
+ (1 A,) a Po
(77)
where po(se-1lz) and po 11) are the marginal distributions implied by the conditional and unconditional reverse processes, p0 [Te, and pg(#e-112z.) respectively.
Proof. For any reverse process p( (xa) and its implied marginal p(æ;-1), the following relation holds under the fixed forward process (2 (2)-)) based the Bayes' rule:
p(=) = (z (==-)))(()))))))))))))))))))))))
(78)
By taking the logarithm and the gradient with respect to rt-1 and rearranging, we have:
1i logp(Es-1) V z (2r V Vx,-1 1og
(79)
Applying this to different reverse processes, including the unconditional (i.e., po (= ( (xx)). condi- tional (i.e., (1 [Fr, =)), and the hybrid (i.e.. (o | Ar)) ones, we have
1 log (
V - ( (x(-)) + V #1-1 log pe (2 [z+)
(80)
V #t-1 log po (x = - -.... log q (2 (x + V zx-1 log po (2
(81)
A&) = -Va log ( + V 2 log Po (+ [2t, As)
(82)
Further, the gradient of the log-hybrid reverse process can be written as
Re-1 logio(zz
V -.- log
2) Po Po (xt (21-1110 2)p A (
(83)
Zi(z,,
A V a
logpo(Ea-1 2) (1 - -A, m- logpo(r1-11z4)
(84)
Therefore,
Va.- 1og Ac(-Vz,-) logg(z/( v.._, 1ogo
(85)
+ (1 A,)( V-,,, log + V z,-1 log po (It-1 (2x)))
(86)
A V x, log V .n log
(87)
Proposition shows that the hybrid distribution's score function interpolates between those of the conditional and unconditional models according to the weight A,, which decreases over time. Intuitively, at earlier (less noisy) timesteps, the model places more emphasis the conditional path (leveraging 2), while at later timesteps, it increasingly relies on the unconditional model. In practice, by leveraging po( pt) ), approximate the hybrid score "x log (1 A:) using the tractable score V log (1t-11 where ji is predicted during inference.
Observe that, by Bayes' theorem, the marginal distributions admit the following factorizations:
P (2)-112)pb (z(-1) pe(x:-1 (=)pp=1)/pp (z) pa [ [(t]
(88)
(pt 1x-1)/ po (At)
(89)
Taking the gradients of their logarithms, we obtain the following decompositions:
log (0 A,) V æ,-, log + A V s,_, log (o(2 1t-1 log V ..- 1og + V Vz.., = log po (-)))
(90)
(91)
These expressions highlight that Vo,-, log po (iz 11-1) serves as a practical surrogate for the in- tractable term A, V (2 Note that the property of the approximation naturally varies
29

---

is closer to the true la- tent 2 ~ (z) since less noise has been injected into the input following the Markov structure 20 Iz. As a result, the information content fi is more aligned with 2 when é small, which matches the increasing influence of the conditional component via A,.
with t. With smaller t, the estimation pi = (=] R Bo(a(a;)
This time-dependent behavior facilitates smooth interpolation between conditional and unconditional estimation, closely resembling the mechanism classifier free guidance [19]. Instead of explicitly sampling from both modes, our approach achieves this interpolation implicitly, using the predicted latent representation jęç This built-in mechanism ensures coherence between the training objective and the inference procedure.
KL Divergence under the von Mises-Fisher (vMF) Distribution. The third term in Cin (Equa- tion 44) represents the KL divergence between the predicted distribution pe (z|2z) and the reference distribution (2( This measure of discrepancy can be expressed in tractable form by making suitable assumptions about the distributions. In particular, if both p=(2(2)) and g(2(2r)) follow von Mises-Fisher (vMF) distributions-icd po (2] ~ vMF( (ft, k) and ((=(20) (M Ri), where is concentration parameter controlling uncertainty-the the KL divergence between these distributions is:
= kAp(K) po(()))) where (p(K) == ( p=(()=))) order
(<p(K)) ig
(92)
and I, denotes the modified Bessel function of the first kind at
Thus, minimizing the KL divergence amounts to maximizing the cosine similarity between the model's predicted mean direction ji and the ground truth mean (2z, up to scaling factor determined by In practice this objective is modulated by the representation alignment weight (n) (Fquation (
which controls the relative emphasis placed on the representation generation component.
C.2 Provable Distribution Approximation with Representation Aligned Diffusion Models
In this subsection, we quantitatively analyze how representation alignment can help to improve the sampling quality of diffusion models. We provide the result for representation-aliee alligged DDPM in Theorem
Theorem Consider the random variable x E Rd with ground truth distribution x ~ (()) == g(20). Assume that the second moment of x is bounded m² == Ey(z)|1æ x| < CC, where I : Egt=r) jz). and that the score V q(zz) is L-Lipschitz for all t. For practically trained
diffusion model parameterized by 0, the score estimation error timestep t bounded by 6,,0, i.e.. (xx( [ ( L, 20,8)
log (=²)))]²²] S éτ.0. Denote the step size as h T/N, where T the total diffusion time and N is the number of discretization steps, and assume that 1 1/ Denote
k-dim isotropic Gaussian distribution as nth Then the following holds,
N GRh
TV 3 (T.(all1))))p expl -T) T) + (L.V + LInhyTT
(93)
convergence forwand process
discretization crrur
Score extimation ermor
Proof. Recall the notation that Pa == Po is the distribution predicted by the denoising network 0 starting from Gaussian noise yd Consider the two measures over the path space: (i) QT where (X1)te(o.7) follows the law of the reverse process; (ii) Per under which (Xe)xejo,7) has the law of the score matching generative model algorithm initialized gr instead of 7d, where gr is the distribution at timestep I' the forward process starting from go == Denote the end distribution of pJT Po have the following inequality:
TV(Po- 4) S TV(Po.P8") TV .46)
(94)
The convergence of the OU process in KL divergence
( bounds the first term,
TV pg") S TV(2 gr) < (xxel)"))xp((-1
(95)
30

---

The second term essentially consists of score estimation error and discretization enror. We next show the following holds,
T/h
TP q0)² (Lollpp = h(
tn + (I-2dh + L2m2"²yT
(96)
k=1
We start proving Equation (96) by proving
pkh
N ② k=1
BoF
1100 (kh))
(100 (rans 2.,) Vlngi
kh 1h (I.2hh 12
(k-1)h
(97)
For t [(k: 1)h,khi we decompose
Eq- [11$) (1)) (IRR, kh, 8x V In (xx)(11)
(98)
E |1|s0) [||s(/h)]] (Ians kh, Exn,0) Vinn )l + Eez Vinn( (ee(()))
(99)
+ Eoz (ee(nn) - V (x+))111
(100)
16hh,0 +Bg In
Dch (ITA) |1
+ +L 150 [l|æRa xell?]
(101)
qt
Utilizing Lemma 16 from I we bound the second term follows.
V In 9kh
(Fan)]l 3 dh + =2 (L² + 1>h3 (
(102)
where
Ving(
/ -aaaaaan--Vngggnaaan (x:)(
(103)
|| In In
| 21112
(104)
Combining all the terms, we obtain
Eqi [1|. (1th) (IuR) kh, 0) Vn
(105)
3ch.o + dh + L' 1| + L h |I| In (xx)))) L'E Fkn 24 2x
(106)
Sehn. + d2 1,2 m²) + L Lddh²² + 1.2 2012+000 dd)
(107)
S6hn, 1.2dh mm
(108)
Similarly to { we can apply the Girsanov theorem and complete stochastic integration using the
properties of Brownian motions and local martingales. Note that go is the end of the reverse SDE, by the lower semicontinuity of the KL divergence and the data-processing inequality [5]. we take the
limit and obtain
N
T kh,0
(L(golPP =
+ (L² L2h2
(109)
where we recall that N = T. We finally conclude with Pinsker's inequality (TV2 KL) and Cauchy-Schwarz inequality.
Theorem ii the first theoretical result that gives a strict TV distance bound between the distribution generated by representation-alin aliggned diffusion models and the ground truth. This is a tighter bound compared to g thanks the benefit of external representations in improving the score estimation. In particular, a common bound of the score estimation error 2 assumed at all timesteps in [9]. In our case, however, since the model latents 2,,9 is good estimation of the ground truth representation it provides nonzero information about the score, resulting in smaller error. To interpret this, note that the score log q(x;)) ) |,, =3 is not a function only of the individual sample 5, but function the entire distribution g(x;) that is related to other samples. External representations, taking advantage of their pretraining tasks and additional training corpus, are able to capture part of the distributional infor- mation. Therefore,
express the score function as V log g(x;) l (( (e)),, r(xx)) ll.==)
31

---

where (T) is the estimated representation of ă and r( denotes the remaining information not captured by either ă or (.). When training DDPM, the diffiusion loss only helps in predicting (due to the equivalence between t-parameterization and e-parameterization), while incorporating the representation alignment loss additionally facilitates predicting (a)), which is crucial in improving score estimation. Therefore, by choosing suitable representations, a appresentation-alignne diffusion model is capable of predicting more accurate scores with the help of estimated representations 2,, i.e. Eg Ill (1), t, 21,0) V log (<x;))l 1 č. S Score. Since predicting representations at smaller is easier, can arguably expect that the improvement is even larger at small timesteps I.
Another advantage of pretrained representations is their rich and casy-to-obtained information for downstream predictions, which is particularly beneficial in conditional generation. In class- conditioned image generation, for example, class labels can be easily inferred from the DINO representations (e.g., through linear probing), thus aligning diffusion models with these representa- tions provides strong guidance towards the conditions.
D Additional Method Details
D.1 Further Justifications for the Joint Training Curriculum
As discussed in Section 3.2 improved estimates of the representations juę naturally enhance the
modeling of the representation-comee cooddiiioond distribution in the first term of L One simple and natural way fulfill this goal is two-stage training paradigm, where we first initialize the model by aligning with pretrained representations, and then optimize through the original diffusion loss. This approach is equivalent to optimizing the first and the last term in Equation 44 independently, offering the model better initializations.
However, argue that two-staged training is suboptimal for the following reasons: (i) the sudden change in loss paradigm leads to gradient direction discontinuities in the optimization dynamics, causing the learned representations to collapse (as observed in our experiments): (ii) the representation spaces of the pretrained encoder and the diffusion model (particularly, the VAE space for latent diffusion) are disjoint, while for diffusion models both representation learning and distribution generation are coherent tasks learned by a single neural network, requiring coherent training flow, Therefore, only conducting representation learning may extract the non-informative features for downstream generation while neglecting the informative part. In other words, simply optimizing the third term in Equation (4) can result in a suboptimal starting point to optimize the first term, suggesting that diffusion and representation generation loss cannot be fully decoupled.
The above reasoning highlights the criticality of joint training all loss terms in Equation (4) with the special emphasis on representation learning in the early training period. Hence, as described in Equation (2), we propose single-stage training curriculum where the diffusion loss weight progressively increases from zero, with a fixed or decaying representation alignment loss.
D.2 Additional Information for Molecule Representations
Generative models over molecules should respect permutational and geometric symmetries, such as S(N) symmetry group for 2D graphs and (special) Euclidean group E(3) or SEE(3 for 3D point clouds. To preserve symmetries, the denoising network should be equivariant with respect to the group transformations. Equivariance ensures that the resulting diffusion process produces invariant distributions. Analogously, molecule representations can be constructed to be either equivariant or invariant. For example, node (atom) level and edge (bond) level features can be made S(N) x E(3)- or S(N) x (E -equivrrvui while graph or molccular level features remain invariant.
To ensure these properties, we follow B5 where graph-level invariant representations were shown to improve generation quality with minimal additional computational cost. In our work, we adopt diffusion and flow matching models with equivariant GNN backbones to maintain permutation and geometric equivariances. Additionally, we incorporate lightweight, semantic graph-level invariant representations derived from pretrained S E3) invariant molecular graph encoders. While equivariant node and edge level representations can also incorporated as guidance for diffusion model training, we leave this as future directions.
32

---

E
Experimental Details
E.1 Image Generation
For continuous diffusion on Euclidean data, we experiment on clads-condiiional image generation using the ImageNet [12] dataset resolution of 256 256, following the experimental setup in REPA [6]]
Evaluation Metries. We generate 50,000 samples for evaluation, using the SDE Buler-Maruyama sampler with the number of diffusion timesteps fixed at 250 for all experiments. Consistent with SiT [4 and REPA [6]], we set the last step of the SDE sampler to 0.04. We adopt the evaluation protocol and reference batches of ADM [13], utilizing their official implementation" For evaluation, we use NVIDIA A100 80GB A6000 50GB GPUs, and enable tf32 precision to accelerate generation as in REPA. We use the following metrics to measure the generation performance:
FID is a standard metric for evaluating image generation, Fréchet Inception Distance (FID) 18
measuring the similarity between the distributions of real and generated images. It does so by calculating the distance between their feature embeddings obtained from the Inception-v3 network 51 assuming that both sets of features follow multivariate Gaussian distributions.
Inception Score (IS) [47] Also leveraging the Inception-v3 network, IS evaluates the quality and diversity of generated samples by measuring the KL-divergence between the original label distribution and the distribution of logits after softmax normalization.
Precision and Recall [29]. These metrics follow classic definitions, with precision indicating the proportion of realistic images, and recall reflecting the fraction of training data manifold covered by generated data.
Baselines. We compare REED with recent diffusion- -based generation baselines that utilize different inputs and architectural designs:
Pixel diffusion models. ADM 13 enhances U-Net-hased diffusion models and introduces classifier- guided sampling to balance quality and diversity. VDM++ 27 presents an adaptive noise schedule improve training efficiency. Simple diffusion 22 focuses high-resolution image synthesis by simplifying both noise schedules and architectures. CDM (21] proposes a cascaded approach, training several diffusion models at increasing resolutions, similar to progressiveGAN (26). to generate high-fidelity images via successive super-resolution stages.
Latent diffusson with U-Net. LDM 46 proposes latent diffusion models that model image distribu- tions in compressed latent space, significantly improving training efficiency while maintaining generation quality.
Latent diffusion using transformer and U-Net hybrids. U- VIT-H/2 introduces a ViT-based latent diffusion model with U-Net-style skip connections. DiffTT [16] improves terssforerrbbasee diffusion with time-dependent multi-head self-attention mechanism. MDTv2-XL.22 [15] proposes asymmetric encoder-decoder framework, utilizing U-Net-like long shortcuts in the encoder and dense input shortcuts the decoder for efficient diffusion transformer training.
Latent diffusion with transformers. MaskDiT (63 employs asymmetric encoder-decoder for efficient transformer- based diffusion, training the model using an auxiliary mask reconstruction task similar to MAE [17]. SD-DiT (65] builds on MaskDiT by adding self-supervised discrimination introduces a pure transformer architecture for diffusion models, utilizing AdaIN-zero modules. SiT 440] systematically investigates training efficiency by shifting from diffusion to continuous flow-based modeling. SiT with REPA [61 further accelerates training by enforcing representation alignment between intermediate layers and DINOv2 42 patch representations.
objective with a momentum encoder. DiT 43
Image Captioning. For each ImageNet image, we generate synthetic caption using Qwen2- VL (55) 2B, using the prompt "Describe this image." and the raw image as the input. We adopt the official conversation template of Owwen-VVL-25 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. We
33

---

provide a few examples of the generated captions and the corresponding image and class labels in Figure
Class: Flamingo
The image depicts flamingo with its head numed slightly to the side. The flamingo has long, curved neck and pinkish-orange plumage. Its beak is black with a white tip, its beak, possibly a small object or food. The background is a dark green, possibly indicating water body or a forested area. The flamingo's posture suggests it might be in the midst of feeding grooming behavior.
Caption:
appears to holding something
Class: Scabbard
The image shows a knife with brown leather sheath. The knife has blade that appears to be made of polished metal, possibly stainless steel, and has handle that also made of similar material, likely wood. The sheath made of brown leather and has a decorative design it. The knife and sheath are placed on a flat surface, possibly table or desk.
Caption:
Class: Goldfish
The image shows two goldfish in 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 shiny, reflective surface, indicating they likely well-maintained. The aquarium has gravel bottom with some small rocks and a few large green plants. The lighting in the aquarium creates warm, inviting atmosphere.
Caption:
Figure 3: Examples of images with their ground-truth class labels and the generated captions.
VLM Embeddings. Both the generated caption and the image are then provided as input to the Qwen2-VL 7B model, which contains 28 joint vision-language transformer layers. We extract representations from the 15th layer of Qwen2 VL 7B and average across all image and text tokens to obtain unified vision-language embedding of dimension 3584. We then align the averaged SiT image patch latents from an intermediate SiT layer, after passing through an MLP projector, with this joint VLM embedding, using cosine similarity loss (Equation (8)). We choose VLM embeddings over unimodal pretrained text embeddings because VLMs offer a shared embedding space that is already partially aligned between modalities, facilitating better transfer and alignment across image and text representations.
It is noteworthy that we utilize VL.Ms with different sizes for captioning and representation generation. In particular, we use the smaller Qwen2-VI VL 2B model to generate the synthetic caption since 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. In contrast, during representation generation, the representations image-caption pair, which is much less cone-ccnsuming than the sequential inference the captioning step. Therefore, opt to use the embeddings of the 7B model, which has stronger capability to understand the input content and is thus able to generate more powerful multimodal representations that accelerate diffusion model training. This is also validated empirically in Table T2]
be obtained in one forward pass for a given
Model Architecture and Further Implementation Details. We adopt the same SiT architectures as REPA and SiT, specifically SiT-B/2, SiT-L/2, and SiT-X1//2 which have 12, 24, and 28 layers, hidden dimensions of 768, 1024, and 1 1152, and 12, 16, and 16 attention heads, respectively. Latent vectors are pre-computed from raw images using the stable diffusion VAE (46), resulting in 32 32
34

---

latent vectors that serve input to SiT. For decoding, we use the stablityal//sl--a--ttoom decoder to reconstruct images from latent vectors.
For image SSL. embeddinng we use DINOv2-B and set the alignment coefficient X in Equation (8) to 1.0. For VLM embeddings, we use the aforementioned VL.M representation with the coefficient Ay set to 0.5. Alignment depth is as follows: for SiT-B, the 4th layer is aligned with the image embedding and the 8th layer with the VLM embedding; for SiT-L and SiT-XL, the 8th layer is aligned with the image embedding and the 16th layer with the VLM embedding. The projection heads for alignment are implemented as 3-layer MLPs with SiLU activations [14]. Regarding the training curriculum (Equation (9). we keep (n) fixed at 0.5 and increase (n) linearly during the first 50K iterations, after which remains at 1.0 for the remainder of training.
For optimization, we use Adam W 39 with a constant learning rate of 1 x 10-1, parameters (Bi, (22) (0.9, 0.999), and no weight decay. The batch size is set to 256. To accelerate training, use precisson (fp16) and apply gradient clipping with a threshold of 1.0.
Computing Resources. Our experiments are conducted using NVIDIA A100 80GB GPUs.
Analysis on Computational Cost. We pre-caleulate and store all VLM representations before diffusion model training. As a result, the additional computational overhead during REED training compared REPA is minimal, limited only to lightweight MLP projection head applied to the cached VLM representations. This overhead negligible in practice, thus the actual training time for both REPA and REED is essentially the same. For example, training SiT-X1/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: A100 GPUs, captioning requires about 7.5 hours, and VLM representation extraction about 1.3 hours. Therefore, 1M iterations (200 epochs) of REED, togther with captioning and 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.
On the other hand, following REPA, the DINOv2 representations are calculated during diffusion train- ing and not precomputed, which introduces additional computational cost to both REED and REPA in comparison to the vanilla SiT training. However, the specific DINOv2 variant we use, DINOv2-B/14 has only 86M parameters, significantly smaller than the diffusion backbones used in our experiments, from SiT-B/2 (130M params) to SiT-XL/2 (675M params). Also, the representation calculation does not require gradient backpropagation. Therefore, the additional overhead of calculating DINOv2 representations per training step is minimal compared to the overall training cost. For SiT-XL/2, it takes about only 1.05 x training time per step compared with the vanilla SiT training, which is far outweighed by the 23.3x acceleration in training steps enabled by our method. Overall, the FLOPS and memory consumption of REED almost the same as training a vanilla diffusion model.
E.2 Protein Sequence Generation
For sequence data, we investigate discrete flow models in the context of protein inverse folding, where the goal is to generate protein sequences conditioned on given 3D backbone structure.
Dataset and Evaluation Metrics. We train the inverse folding models on the PDB training set from [Iل] and evaluate the corresponding PDB test set, using the same data splitting. We focus on single-chain proteins with fewer than 256 residues, resulting in 13,753 training proteins and 811 test proteins after filtering. For each test structure, each model generates one protein sequence. Both ProteinMPNN [T]] and the discrete flow model use a temperature of 0.1, while the discrete flow model uses 500 diffusion timesteps.
We assess inverse folding performance sequence recovery rate, self-consistency RMSD (scRMSD), and PLDDT from ESMFold (36). Details for each metric provided below:
35

---

Sequence Recovery Rate. This metric measures the proportion of residues in the generated sequence that match the ground truth. The average computed the token level, dividing the number of correctly predicted residues by the total number of residues across all test sequences.
scRMSD, Self-consistency RMSD evaluates how closely the generated sequence, when folded by ESMFoll [6]] matches the target backbone structure. We report the median scRMSD) across all test proteins in Table Since an scRMSD below 2 is typically considered indicative of successful
inverse folding
we also report the fraction of test cases with scRMSD less than 2Å.
PLDDT. To further assess sequence quality, we use pLDDT scores from ESMFold predictions on generated sequences. pLDDT reflects the confidence of structure predictions, providing insight into the plausibility of the generated protein sequences. We report the median pLDDT across all test proteins, as well as the percentage of predictions with pLDDT above 80, following standard practice [57]
Baselines. We benchmark REED against ProteinMPNN [T]. the de facto approach for inverse folding, as well as discrete flow models trained without REED. Following [41 [54, the discrete flow baselines use the Multiflow 7 objective and the ProteinMPNN [T]] architecture. All models are trained on our curated dataset as described above, employing a backbone noise of 0.1 and constructing the graph using the 30 nearest neighbors. All other hyperparameterss with ProteinMPNN (using 3 encoder and 3 decoder layers). In the original ProteinMPNN, edge representations are fixed and node embeddings are not learnable. To enable better representation alignment with REED, we modify the architecture to allow edge representations to be updated and node features to be learnable. For ablation, we also report the performance of discrete flow models with these architectural adaptations but without REED in Table 13
kept consistent
Protein Folding and AF3 Representations. We utilize AlphaFold3 [ to generate auxiliary structural information from target protein sequences. For faster inference, omit the use of Multiple Sequence Aligmment (MSA). Structural representations from the penultimate layer of AlphaFold3's diffusion head are extracted latents, while amino acid representations are obtained from the single and pair embeddings of the last Pairformer layer (corresponding to node and edge features, respectively). When sampling through the diffusion module, diffusion steps. Given that protein inverse folding is conceptually the reverse of protein folding, we align the layers of our discrete diffusion model with those of the folding model in reverse order. Specifically, the output of the ProteinMPNN encoder, which captures the protein backbone structure, aligned with the structure representation from AF3. Additionally, after passing through an MLP projector, the output of the decoder's 's first layer is aligned with the single and pair token embeddings from Pairformer: decoder node embeddings are aligned to single embeddings, and edge embeddings are aligned to pair embeddings.
10 rounds of recycles and 200
Further Implementation Details. For the training curriculum, we keep ((n) fixed at 0.2, while (nn) is linearly increased to 1.0 over the first 50 epochs and then decayed using a cosine schedule. The aligmment coefficients As are set to 0.5 for single representations, 2.0 for pair representations, and 1.0 for structure representations. The projection heads for alignment are implemented as 2-layer MLPs with SiLU activations. Optimization is performed using Adam 28 with a constant learning rate of x 10-3
Computing Resources. AlphaFold3 representation generation is performed on 8 NVIDIA A100 80GB GPUs, with the extracted embeddings stored for subsequent use. All other experiments are conducted on single NVIDIA A100 80GB GPU.
E.3 Molecule Generation
For structured and geometric graph data, we tackle the 3D molecule generation task using equivariant diffusion models, ensuring molecular data symmetries are preserved throughout the process.
Datasets and Evaluation Metries. QM9 45 and GEOM-DRUG 3 are commonly used benchmarks for unconditional molecule generation. Compared with QM9 where molecule contain approximately 18 atoms on average and up to heavy atoms, GEOM-DRUG a more challenging and useful dataset of larger, dikg-like molecules with an average size of around 44 atoms, which is used to assess
36

---

method. Following SemlaFlow [25], we use the same data splits and discard molecules with more than 72 atoms in the training set (corresponds about 1% of training data), while keeping validation and test sets unchanged.
We adopt standard benchmark evaluation metrics: atom stability, molecule stability, validity, well as additional metrics energy and strain. We also include the number of function evaluations (NFE) to sample one batch of molecules for fair comparison, taking sampling efficiency into account. thorough description of these metrics is provided as below:
Atom Stability. Fraction of atoms with correct valency.
Molecule Stability. Fraction of molecules in which all atoms are stable.
Validity. Fraction of molecules convertible to valid SMILES strings using RDKit.
Energy. The energy (xx) of a conformation x; computed with MMFF94 in RDKit, a commonly used molecular modeling framework [24], where lower values indicate greater physical plausibility.
Strain. Energy difference ((x) - (z) between the generated conformation r and the MMFF94- optimized conformation ē, with lower values indicating less distortion.
Baselines. Recent 3D molecular generators consist of two main types depending on the approaches they obtain bonds: generate atoms only and infer bonds based on coordinates with external rules, or directly generate bonds jointly with atoms. As explained in [25 the latter methods can produce higher-quality samples, including MiDi [53], EQGAT-diff [31]. SemlaFlow [25], which we consider
as the current -fate-of-theearrs We only take these advanced 2D&3D methods that directly learn to
generate bonds as baselines. In our experiments, we choose SemlaFlow, an F(3) -())-equivariant ODE
flow matching, as our base model.
For further context, modeling molecular graph distributions requires careful treatment of both permutational and geometrie symmetries. For 2D molecular graphs, the relevant symmetry group is S(N), capturing all permutations of atom indices. For 3D molecular structures, symmetries governed by the special Euclidean group SE(3) S3)3) R$, where SO(3) represents rotations and Rd represents translations. In some cases, the full Euclidean group E(3) O(3) x RR, which also includes reflections, may be considered, but SE(3) often more practical as preserves molecular handedness. To maintain these symmetries, the denoising network must be equivariant to the appropriate groups, and the diffusion process must model distributions invariant to these transformations. Accordingly, node-level (atom) and edge-level (bond) features should be permutation equivariant, while graph-level (molecule) features should be permutation invariant. For 3D coordinates, analogously, their node-level embeddings should be SE(3) equivariant, while graph-level features should be SE) invariant.
Pretrained Representation. For the pretrained encoder, we select the Unimol 64 with S3) transformers, and additionally add GEOM-DRUG to its pre-training dataset described in [55 We use graph-level invariant representations to align with SemlaFlow intermediate layers. For our 12-layer model, we empirically find that aligning the fourth layer yields the best performance.
Further Implementation Details. For the training curriculum, we keep ((n) fixed at 1.0, while a(n) is linearly increased to 1.0 over the first 30 epochs and then fixed as a constant for the remaining epochs. The alignment coefficient À set to 0.2. The projection heads for alignment are implemented as 3-layer MLPs with SiLU activations. We follow all other hyperparameters regarding model, training and evaluation settings in (25).
Computing Resources. We train SemlaFlow with REED for 200 epochs on a single NVIDIA A100 80GB GPU. Notably, it requires significantly fewer GPU hours than EQGAT-diff, which trained for 800 epochs on 4 NVIDIA A100 GPUs.
F Additional Experimental Results
F.1 Image Generation
Additional Evaluation Results for Table I. We present evaluation results across multiple metrics for our method trained with SiT models of various sizes, compared to both the vanilla SiT and REPA
37

---

Table 7: Comprehensive evaluation results for comparisons with SiT and REPA across various model
sizes. No classifier free guidance is used. Iter. denotes training iteration.
Model
#Params
Iter.
FID
IS
Pree.t
Rec/t
SIT-B/2
130M
400K
33.0
43.7
0.53
0.63
REPA
130M
50K
78.2
17.1
0.33
0,48
+ REED (ours)
130M
50K
66.0
22.9
0.37
0,51
+ REPA
130M
100K
49.5
27.5
0.46
0.59
+ REED (ours)
130M
100K
36.8
42.5
0.51
0.61
+ REPA
130M
200K
33.2
43.7
0.54
0.63
REED (ours)
130M
200K
24.6
62.2
0.58
0.64
+ REPA
130M
400K
24.4
59.9
0.59
0.65
+ REED (ours)
130M
400K
19.5
75.9
0.61
0.65
SiT-L/2
458M
400K
18.8
72.0
0.64
0.64
+ REPA
458M
S0K
55.4
23.0
0.43
0.53
+ REED (ours)
458M
50K
41.7
36.7
0.50
0.59
REPA
458M
100K
24.1
55.7
0.62
0.60
+ REED (ours)
458M
100K
19.1
75.3
0.63
0.63
+ REPA
458M
200K
14.0
86,5
0.67
0.64
+ REED (ours)
458M
200K
12.1
101.8
0.67
0,65
+ REPA
458M
400K
10.0
109.2
0.69
0,65
REED (ours)
458M
400K
8.9
122.4
0.69
0.66
SiT-XL/2
675M
7M
8.3
131.7
0.68
0.67
+ REPA
675M
50K
52.3
24.3
0.53
+ REED (ours)
675M
50K
38.0
40.4
0.52
0.59
+ REPA
675M
100K
19.4
67.4
0.64
0.61
+ REED (ours)
675M
100K
16.0
84.4
0.65
0.63
+ REPA
675M
200K
11.1
100.4
0.69
0.64
+ REED (ours)
675M
200K
9.8
114.7
0.65
+ REPA
675M
400K
7.9
122.6
0.70
0.65
+ REED (ours)
675M
400K
7.3
134.5
0.70
0,66
+ REPA
675M
4M
5.9
157.8
0.70
0,69
+ REED (ours)
675M
IM
4.7
166.1
0.72
0.66
Table 8: Detailed evaluation results for SiT-B/2 models aligned using VLM embeddings from different layers of the Qwen2- VL 7B model. Dep.-I and Dep.-VL denote the alignment depth for image and VLM representations respectively. No classifier-free guidance is used. Iter. denotes training iteration.
Iter.
Dep.-I
Dep.-VL
VLM Embedding
FID1
1St
Prec.t
Rec.
200K
4
8
Layer 0
29.9
50.6
0.56
0.64
200K
4
Layer
30.2
49.6
0.55
0.64
200K
Layer
29.4
50.8
0.56
0.64
200K
8
Layer 27
29.5
51.2
0,56
0.65
in T Hyperparameter settings and implementation details are as described in Appendix E.I REED consistently outperforms both REPA and vanilla SiT across different model sizes, training iterations, and evaluation metrics.
Results with Different VLM Embeddings. We evaluate the performance REED on SiT-B/2 using VLM embeddings extracted from various layers of the Qwen2-VL. VL 7B model, as shown in Table For all experiments, the image (DINOv2) embedding is aligned at the 4th layer and the VLM embedding at the 8th layer, using the same hyperparameters detailed in Appendix E. Notably, the best results are achieved when aligning with the 15th layer (out 28) of Qwen2-VL. 7B. Early layers do not provide well-fused vision and text token embeddings, while later layers are more similar to the final VL.M output (text) space. Therefore, aligning middle layer offers balanced and robust semantic representation.
38

---

Table 9: Detailed evaluation results at different alignment depth on SiT-B/2. Dep.-I and Dep.-VL denote the alignment depth for image and VLM representations respectively. No classifier-free guidance is used. Iter. denotes training iteration.
Iter.
Dep.-I
Dep.-VI.
FID.
IS
Pree.t
Rec.
200K
2
33.2
46.9
0.53
0.64
200K
27.7
55.5
0.56
0.64
200K
6
27.9
54.7
0.56
0.65
200K
8
30.1
51.8
0.55
0.64
200K
4
26.5
57.7
0.57
0.64
200K
4
25.9
59.5
0,57
0.64
200K
4
6
24.8
61.9
0.57
0.63
200K
4
8
24.6
62.2
0,58
0.64
200K
4
10
26.4
59.1
0.57
0.65
200K
45.5
33.6
0.47
0.60
200K
8
45.6
32.3
0.47
0.61
Additional Results for Different Alignment Depths in Table 4 We present further evaluation results across various alignment depths for both image and VLM representations in Table 9 The observed trends remain consistent across most metrics (except recall): aligning image representations shallower layer (i.e., layer 4) and VLM representations at deeper layer (i.e., layer 8) out of the total 12 layers yields the best performance. This observation is consistent with findings from the Hierarchical VAE literature, where shallower latents tend capture fine-grate ine-grained details, while deeper latents encapsulate higher-level semantic information. These results provide empirical support for theoretical analysis in Section 2.2
Empirical Verification of Approximations in Section 2.3
To assess the validity of the approxima- tions introduced in Section 2.3 we train the model using an alternative objective that directly follows Equation (4) without additional approximations. Specifically, we use the SiT-B/2 model and apply the cosine similarity representation alignment loss (corresponding to the third term in Equation ( the 4th layer. Additionally, we introduce an extra input condition via cross attention at the start of the 5th layer in SiT. With this architectural adjustment, po (o | 2) is computed using the ground truth representation as the input condition, while po (It-1 (are) obtained with a learnable dummy latent vector (representing no conditioning) as the input. The second term is derived according to Equation (1 with its relative coefficient set to 1.0.
The weighted product of two Gaussians with equal variance remains Gaussian, where the mean is the weighted average of the individual means and the variance is unchanged. As a result, under the Gaussian assumption of pe (e( 2) and (= the distribution Pe (xx,-1/zz.
A,) in Equation (5) can be expressed closed form an aggregation of the conditioned and unconditioned model outputs. This allows the first term in Equation 4 to be calculated as standard score matching loss but relative to the weighted average of the two model outputs. During inference, since the ground-truie truth representation is not available, we sample using the unconditional model po (1
We conduct experiments using DINOv2 representations, keeping all other settings consistent with those used for REPA. The results, presented in Ta]le show that the model trained directly with Equation 4 performs comparably
REPA, thereby empirically supporting the validity of the approximations we made in Section 2.3 Additionally, we observe that the second log-normalization term the objective remains below 0.04. which is considerably smaller than the denoising loss in the first term (approximately 0.7). This further justifies our decision to exclude it from training, as described in Section 2,
Ablation Study on Different Curriculum Schedules. To measure the effect of the curriculum schedule between representation learning and diffusion training, we conduct ablations on different configurations for the schedules a(n) and 3(n), including varying the phase-in period for a(n) and applying cosine decay to (n). As shown in Table 11 all curriculum variants of (n))) outperform the constant schedule, with 50K iterations chosen as the default for all image generation experiments. Notably, the choice of ((n) schedule does not substantially affect generation performance.
39

---

Table 10: Results of SiT-B/2 model trained with Equation 4 and with DINOv2 representations compared with REPA. No classifier free guidance is used. Iter. denotes training iteration.
without approximation (w/o Approx.)
Model (SiT-B/2)
Paarams
Iter.
FID
IS
Pree.t
Rec.
REPA + Equation
130M
200K
33.2
43.7
0.54
0.63
w/o Approx.
130M
200K
34.7
41.9
0,54
0.62
Table 11: Ablation study on different curriculum schedules. We report results on SiT-B/2 training for 200K iterations. No classifier-fres guidance is used.
Model
((1)
3(n)
FID
1St
Prec.
Rec.t
+ REPA
33.2
43.7
0.54
0.63
+ REED
Constant
Constant
29.4
50.8
0.56
0.64
+ REED
Linear increase in 50K iterations
Constant
24.6
62.2
0.58
0.64
REED
Linear increase in 25K iterations
Constant
25.2
60.7
0.58
0.64
REED
Linear increase in 75K iterations
Constant
24.7
62.2
0.58
0.64
REED
Linear increase in 50K iterations
Cosine decay in 200K iterations
24.7
61.4
0.59
0.63
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 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.
Ablation Study on Different Multimodal Representations. We conduct additional ablations on the multimodal representations obtained from different pretrained models. We compare REED (using Qwen2-VL. 7B representations) to versions using either a more lightweight model (i.e., Qwen2-Vi VL 2B, OpenCLIF and SigLIP 622 intentionally noisier features (i.e., Qwen2-VL 7B with Gaussian noise). All other settings were kept consistent.
The results are shown in Table 12 Using lightweight or less accurate VLM representations still consistently leads to improved performance compared to the model without aligning with multi- modal representations. This demonstrates the robustness of method to less accurate noisy representations and suggests that REED retains its advantages even when ideal pretrained features 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.
However, utilizing text representations from OpenCLIP or SigLIP yields little improvement or slightly worse performance than the setting without multimodal representations. This can be attributed to the fact that OpenCLIP and SigLIP are trained with a contrastive objective on short and concise captions and are primarily optimized for downstream classification tasks without capturing fine-grained semantic details needed for high-quality image generation. As a result, the information content in their text embeddings may be insufficient fully support image generation. In comparison, VLMs are trained on a diverse set of dision-language tasks and have strong ability of vision- llaguael understanding and generation. They also better at bridging two modalities, and consequently, have better aligned visual and textual representations compared to OpenCLIP and SigLIP. These enable VLMMs to generate richer and more relevant semantic representations for guiding image generation.
Moreover, from the computational perspective, OpenCLIP and SigLIP requires a higher computational cost than Qw 2B. Since neither of them can generate text captions, the captioning step using Qwen2-VL VL 2B is always necessary, and we can save the Qwen2-VL 2B embeddings along the way to avoid redundant computation in the representation generation stage.
The text encoder in http:///wwegggeeeeoo///////////a
base-1.0
40

---

Table 12: Ablation study using different multimodal representations for REED. We report results on SiT-B/2 training for 200K iterations. No classifier free guidance is used.
Multimodal Representation
FID
ISt
Pree.
Rec.
27.7
55.5
0,56
0.64
Qwen2- 2L22-VLL7 7B Qwen2- VL
24.6
62.2
0.58
0.64
26.3
59.9
0.57
0.65
Qwen2-VI + noise (scale=0.1)
25.0
61.8
0.58
0.64
OpenCIIP
29.0
55.8
0.55
0.64
SigL.IP
27.6
57.6
0.56
0.64
F.2 Protein Sequence Generation
Ablation Study on Architectural Modifications. As described in Appendix E.2 original ProteinMPNN architecture to enhance representation alignment with REED by enabling updates edge representations and making node features learnable. For the ablation study, evaluate discrete flow models incorporating these architectural changes but without applying REED, shown in Table 13 The results comparable those of the original DiscreteDiff model, suggesting that the observed performance improvements stem from the REED algorithm itself.
modify the
Table 13: Ablation study on the base discrete flow model with the same architectural modifications as DiscreteDiff-REED. The performance remains similar the original DiscreteDiff model, indicating that the improvements are attributable to REED itself.
Methods
Epochs
Seq. Recovery(So)
scRMSD
$(scMMSD<2)
PLDDT
$(PLDDT>80(())
DiscreteDiff 50
38.04
2.034
49.32
81.36
58.39
Architectural
100
40.60
1.794
53.42
82.82
68.32
Modifications
150
41.35
1.817
54.78
83.00
68.20
Ablation Study Aligned Representations. We present ablation results for various combina- tions of aligned representations in Table including using only single representations, only pair representations, both single and pair representations, and only structure representations. These are compared against the results obtained when utilizing all three representation types. For each ablation (except "all"), we fix (nz) at 0.2 and the alignment coefficients 1.0 for all representations. As shown in Table 14, each representation type independently contributes substantially to model performance, with all outperforming the baseline DiscreteDiff model each training epoch. Notably, the pair representation yields the largest performance gain, highlighting the importance of capturing amino acid interactions for improved protein structure understanding in the inverse folding task.
Table 14: Ablation studies on different combinations of aligned representations: single, pair, and structure. "all" denotes using all three types of representations.
Representation
Epochs
Seq. Recovery(%) 1
scRMSD
$(scRMSD<2)
PLDDT <(pLDDT>>0))
50
40.35
1.790
53.91
82.76
67.21
single
100
41.43
1.778
55.03
83.20
67.83
150
42.06
1.768
55,28
83.50
71.68
50
41.07
1.775
55.53
83.13
68.94
pair
100
41.92
1.781
55,03
82.89
68.70
150
42.21
1.778
55,53
83.02
69.32
50
40.62
1,806
55,28
83.12
68.32
single+pair
100
41.69
1,766
56,52
83.04
66.96
150
41.85
1.753
54.53
83.13
69.57
50
40.40
1,839
$4.04
82,50
64.97
structure
100
41.47
1.774
53.91
83.08
67.58
150
41.88
55.90
83.32
70.81
50
41.06
1.788
$4.37
82.74
67.6 .1
all
100
41.91
1.780
56.23
83.09
68.41
150
42.26
1.799
54.66
83.17
69.07
41

---

G
Qualitative Results
Figure 4: Selected samples on ImageNet 256 256 generated by the SIT-XL/2+REED model after IM training iterations. We use classifier-free guidance with w 1.275.
42

---

Figure 5: Selected samples on ImageNet 256 256 generated by the SIT-XI//+RRFFD model after IM training iterations. We use classifier-free guidance with w = 4.0. Each row corresponds to the same class label. From top to bottom, the class labels are: "macaw" (88), "flamingo" (130), "borzoi, Russian wolfhound" (169), "Samoyed" (258), "Egyptian cat" (285), "otter" (360), and "dogsled" (537), respectively.
43

---

Figure 6: Selected samples on ImageNet 256 256 generated by the SiT-XL/2+REED model after 1M training iterations. We use classifier-fres free guidance with 4.0. Each row corresponds to the same class label. From top to bottom, the class labels are: "jigsaw puzzle" (611), "laptop" (620), "school bus" (779), "cheeseburger" (933), "cliff" (972), "earthstar" (995), and "toilet tissue" (999), respectively.
44

---

pdb_ie __b 1eyw
pdb id: 2o09
scRMSD: 1.38
scRMSD: 0.91
Seq. Recovery: 39.0% pL.DDT 86.7
Seq. Recovery: 37.9% pL.DDT: 85.9
pdb id: Ihqz
pdb id: 2ooj
scRMSD: 1.27
scRMSD: 1.27
Scq. Recovery: 44.6% PLDDDT 87.4
Seq. Recovery: 38.9%
PLDDD: 87.3
pdb id: 1o8p
pdb id: 2pb9
scRMSD: 0.53
scMSSD: 0.99
Seq. Recovery: 53.4%
Seq. Recovery: 47.8%
PLDDT: 89.5
pLDDT: 84.5
pdb id: 1s8n
pdb 2xom
scRMSD: 1.27
ScRMSD: 0.42
Seq. Recovery: 55.2%
Seq. Recovery: 49.5%
pLDDT: 84,5
PLDDT: 92.5
pdb id: 2xqx
pdb id: 1tje
scRMSD: 1.69
scRMSD: 0.61
Seq. Recovery: 43.5%
Seq. Recovery: 40.0% PLDDT: 90.1
PLDDT: 78.5
Figure 7: Selected samples on protein inverse folding. Green color denotes the ground truth structure and blue color denotes the generated sequence folded by ESMFold [36].
45

---

pdb id: 3bjo
scRMSD: 1.05
pdb id: 4iau
Seq. Recovery: 45.6%
pL.DDT: 88.9
scRMSD: 1.44
Seq. Recovery: 47.8% PLDDT: 85.0
pdb id: 3hkl
pdb id: 5hj1
scRMSD: 0.88
Seq. Recovery: 41.8%
scRMSD: 1.66
PLDDT: 87.8
Seq. Recovery: 45.4% PL.DDT: 76.7
pdb id: 5oa9
pdb id: 3nb6
ScRMSD: 1.28
scRMSD:
Seq. Recovery: 44.4%
Seq. Recovery: 46.4%
PLDDT: 86.5
PLDDT: 83.6
pdb id: 3qp3
pdb id: 6eez
scRMSD: 1.13
scRMSD: 0.68
Seq. Recovery: 45.5%
Seq. Recovery: 44.1% PLDDT:
PLDDT: 90.3
Figure 8: (Continued.) Selected samples protein inverse folding. Green color denotes the ground truth structure and blue color denotes the generated sequence folded by ESMFold (36).
46

---

Figure 9: Selected samples molecule generation.
47


Reviews:
[{"review_id": "uS9sHfQ6Bo", "reviewer": "Reviewer_LKbx", "summary": "This paper presents a unified theoretical framework for integrating pretrained representations into diffusion model training, enabling flexible guidance via multimodal and hierarchical alignment. The proposed method, REED, demonstrates strong improvements in sample quality and training speed across diverse tasks including image generation, protein design, and molecule generation.", "questions": "NA", "limitations": "NA", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n\n1. The paper develops a rigorous and general variational framework for understanding and designing representation-guided diffusion models. This subsumes and explains prior works (e.g., REPA, RCG) as special cases, and provides new design insights (e.g., multi-latent structures, flexible scheduling).\n\n2. The proposed approach goes beyond simple alignment with a single pretrained representation. It supports multi-level, multimodal representations (e.g., combining image and text embeddings, or structure and sequence in proteins), and shows empirically that this leads to better controllability and generation quality.\n\n3. The paper motivates and validates a practical curriculum that introduces representation learning loss early and phases in diffusion loss, leading to faster convergence and improved stability.\n\n**Weakness:**\n\n1. 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.\n\n2. 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.", "was_revised": true, "final_justification": "The authors solved my concerns.", "rebuttal": "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.\n\n> 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.\n\nWe 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:\n- 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.\n- 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.\n- 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.\n\nTaken together, these points establish a rigorous connection between our theoretical framework and REPA, and provide clear motivation for the novel design choices in REED.\n\n> 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.\n\nThis 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. \n\nOur 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.\n\nWe 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.\n\n| Model | $\\alpha(n)$ | $\\beta(n)$ | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|-----|-----|\n| REPA | - | - | 33.2 | 43.7 | 0.54 | 0.63 |\n| REED | constant | constant | 29.4 | 50.8 | 0.56 | 0.64 |\n| REED | linear increase to 50K iterations | constant | **24.6** | **62.2** | 0.58 | **0.64** |\n| REED | linear increase to 25K iterations | constant | 25.2 | 60.7 | 0.58 | 0.64 |\n| REED | linear increase to 75K iterations | constant | 24.7 | 62.2 | 0.58 | 0.64 |\n| REED | linear increase to 50K iterations | cosine decay in 200K iterations | 24.7 | 61.4 | **0.59** | 0.63 |\n\nWe 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."}, {"review_id": "MoNxpEGcNl", "reviewer": "Reviewer_edpA", "summary": "Paper introduces aframework for enhancing diffusion models through representation guidance, i.e. aligning the internal representations of diffusion models with pretrained representations like DINO features. The authors derive a variational bound incorporating auxiliary representations and show how a weighting schedule can control when and how these representations guide the denoising process. \ntwo key strategies are proposed: (1) multimodal representation alignment using synthetic paired data (2) training curriculum that phases in the diffusion loss. REED is evaluated across three tasks: image generation, protein inverse folding, and molecule generation. Showing improved performance and faster training compared to prior methods like REPA.", "questions": "- Can you provide 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?\n- What failure modes might occur when the synthetic auxiliary modality is too dissimilar from the primary data? For example, could this lead to misaligned representations that degrade generation quality?\n- how does REED perform when using lightweight or less accurate pretrained representations? This is crucial for understanding the method's generalizability\n- how sensitive is REED to the differences between the pretrained representation space and the diffusion model's latent space?", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- The paper is built on a theoretical motivation, deriving a variational bound that incorporates pretrained representations. The experimental section demonstrates clear performance gains.\n\n- The paper is logically organized, moving from theory to methods and then to experiments. Key ideas, such as the use of synthetic multimodal data, are illustrated with figures.\n\n- The paper addresses the important challenge of improving the efficiency and quality of diffusion models. The application to protein and molecular domains shows potential for impact in scientific discovery.\n\n\n**Weakness**\n- The experiments could be stronger with more rigorous ablations to rule out simpler explanations for the performance gains. FOr example,  could the authors report results for a baseline where the full REED model is trained without the proposed curriculum? i.e. with constant weighting for both the diffusion loss and the representation loss? \n\n- The method's success is heavily dependent on extremely large, resource-intensive pretrained models like DINOv2. This reliance is a barrier to wider adoption, and application in resource-constrained settings, potentially limiting the practical impact of the work.\n\n- 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 rather than presenting a fundamentally new framework", "was_revised": true, "final_justification": "I rated this paper highly because it offers a well-motivated and technically solid extension of REPA. Despite some dependence on large pretrained models, the strong empirical results and broad applicability justify an acceptance.", "rebuttal": "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.\n\n> Rigorous ablations to rule out simpler explanations for the performance gains.\n\nThis 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.\n\n| Model | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|\n| REPA | 33.2 | 43.7 | 0.54 | 0.63 |\n| REED | **24.6** | **62.2** | **0.58** | **0.64** |\n| REED w/o curriculum schedule | 29.4 | 50.8 | 0.56 | 0.64 |\n| REED w/o multimodal representations | 27.7 | 55.5 | 0.56 | 0.64 |\n\n> 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?\n\nWe 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.\n\nCompared 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.\n\nWe 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.\n\n| Model | $\\alpha(n)$ | $\\beta(n)$ | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|-----|-----|\n| REED | constant | constant | 29.4 | 50.8 | 0.56 | 0.64 |\n| REED | linear increase to 50K iterations | constant | **24.6** | **62.2** | 0.58 | **0.64** |\n| REED | linear increase to 25K iterations | constant | 25.2 | 60.7 | 0.58 | 0.64 |\n| REED | linear increase to 75K iterations | constant | 24.7 | 62.2 | 0.58 | 0.64 |\n| REED | linear increase to 50K iterations | cosine decay in 200K iterations | 24.7 | 61.4 | **0.59** | 0.63 |\n\nWe 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.\n\n> The method's success is heavily dependent on extremely large, resource-intensive pretrained models like DINOv2, potentially limiting the application in resource-constrained settings.\n\nSince 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. \n\nHowever, 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).\n\nIn 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.\n\nMoreover, 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.\n\n> 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.\n\nOur work has several fundamental and unique contributions compared to REPA and other representation-enhanced generative models:\n- 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.\n- 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.\n- 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.\n- 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.\n\nOverall, we believe that our work brings essential new insights into this field and can inspire a series of future research.\n\n> 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?\n\nOur 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.\n\nHowever, 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.\n\n> How does REED perform when using lightweight or less accurate pretrained representations?\n\nWe 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:\n\n| Model | Multimodal representation | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|-----|\n| REED | - | 27.7 | 55.5 | 0.56 | 0.64 |\n| REED | Qwen2-VL 7B | **24.6** | **62.2** | **0.58** | 0.64 |\n| REED | Qwen2-VL 2B | 26.3 | 59.9 | 0.57 | **0.65** |\n| REED | Qwen2-VL 7B + noise (scale=0.1) | 25.0 | 61.8 | 0.58 | 0.64 |\n\nAs 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.\n\n> How sensitive is REED to the differences between the pretrained representation space and the diffusion model's latent space?\n\nDuring 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."}, {"review_id": "v4ci8G9uhu", "reviewer": "Reviewer_EPRe", "summary": "This paper presents a systematic framework REED (Representation-Enhanced Elucidation of Diffusion) for incorporating representation guidance into diffusion models. The authors provide theoretical insights into how high-quality pretrained representations can enhance diffusion model training, proposing two novel strategies: (1) integrating multimodal representations with synthetic data, and (2) designing an optimal training curriculum that balances representation learning and data generation. The framework is demonstrated across three domains: image generation, protein sequence generation, and molecule generation, showing good performance and accelerated training compared to existing methods.", "questions": "1. The paper mentions using synthetic paired data to address the scarcity of multimodal data. Could the authors provide more details on how the synthetic data is generated? Clarifying this could strengthen the reproducibility of the results.\n2. Can you explain why use Qwen2-VL 2B to generate synthetic captions and use representations from Qwen2-VL 7B?\n3. Could you provide more details between the theoretical analysis and experiment details?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strength\n\n1. The paper is technically sound with well-supported claims through both theoretical analysis and extensive experiments. \n2. The paper tests the algorithm on three tasks: image generation, protein generation, and molecule generation.\n3. REED seems to show better performance compared to baselines.\n\n### Weaknesses\n\n1. Could you provide more details on how the synthetic data is generated by VLM? Such as instruction prompts and caption examples.\n2. Experiments use Qwen2-VL 2B to generate synthetic captions and use representations from Qwen2-VL 7B? What’s the purpose to use different VLMs?\n3. 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.\n    \n    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?\n    \n4. 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)$.", "was_revised": true, "final_justification": "The author addressed all my concerns, so I raise the score.", "rebuttal": "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. \n\n> Could you provide more details on how the synthetic data is generated by VLM? Such as instruction prompts and caption examples.\n\nWe 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:\n\n```\nClass label: flamingo\nCaption:\nThe 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.\n```\n\n```\nClass label: scabbard\nCaption:\nThe 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.\n```\n\n```\nClass label: goldfish\nCaption:\nThe 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.\n```\n\n> Experiments use Qwen2-VL 2B to generate synthetic captions and use representations from Qwen2-VL 7B? What’s the purpose to use different VLMs?\n\nWe 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.\n\nHowever, 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 \nfor 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.\n\n| Model | Multimodal representation | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|-----|\n| REPA | - | 33.2 | 43.7 | 0.54 | 0.63 |\n| REED | - | 27.7 | 55.5 | 0.56 | 0.64 |\n| REED | Qwen2-VL 7B | **24.6** | **62.2** | **0.58** | 0.64 |\n| REED | Qwen2-VL 2B | 26.3 | 59.9 | 0.57 | **0.65** |\n\n**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.**\n\n> 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.\nAdditionally, 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?\n\nWe 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.\n\nTherefore, compared with REPA, the main experimental distinctions are:\n- 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.\n- 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.\n- 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.\n\nEach 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.\n\n| Model | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|\n| REPA | 33.2 | 43.7 | 0.54 | 0.63 |\n| REED | **24.6** | **62.2** | **0.58** | **0.64** |\n| REED w/o curriculum schedule | 29.4 | 50.8 | 0.56 | 0.64 |\n| REED w/o multimodal representations | 27.7 | 55.5 | 0.56 | 0.64 |\n\n\n> 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)$\n\n$$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)$$\nThe 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$."}, {"review_id": "FVn9hwIRgm", "reviewer": "Reviewer_WXGL", "summary": "This paper extends REPA, a method originally developed to accelerate diffusion transformer training for image generation, to various modalities. While REPA focuses on aligning intermediate features between DINO and the diffusion transformer, REED further aligns representations from multiple modalities—such as text—during image generation. Additionally, REED is the first to introduce REPA-style regularization terms into protein sequence and molecule generation, demonstrating its applicability beyond vision tasks such as protein sequence generation and molecule generation.", "questions": "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?", "limitations": "The section 5 should further discuss the computation overhead by using large external models.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "[Strengths]\n\n- Novel generalization of REPA to multimodal representation\n\n- First to apply REPA-like regularization to non-visual generative tasks\n\n- Demonstrates consistent performance improvements over baselines across all tasks\n\n[Weaknesses]\n\n- 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.\n\n- 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.\n\n[A] Kingma et al., Variational Diffusion Models, NeurIPS 2021.\n\n[B] Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation, NeurIPS 2023.", "was_revised": true, "final_justification": "The rebuttal addresses my concerns. Pre-computation of representations addresses my concern on computational overhead. The theoretical characterization gives valuable insights. Thus, I raised my score to 5.", "rebuttal": "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.\n\n> 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.\n\nWe **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. \n\nIn 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).\n\nFurthermore, 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.\n\n\n> 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.\n[A] Kingma et al., Variational Diffusion Models, NeurIPS 2021.\n[B] Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation, NeurIPS 2023.\n\n\nOur 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].\n\nWe believe this new theoretical framework is valuable for several reasons:\n- 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.\n- 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.\n\nWe will clarify these points in the revised manuscript and better highlight the novel connections between our theoretical analysis and the design of our method.\n\n> 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?\n\nAs 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.\n\nIn 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.\n\nIt 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.\n\n| Model | Multimodal representation | FID | IS | Pre. | Rec. |\n|-----|-----|-----|-----|-----|-----|\n| REPA | - | 33.2 | 43.7 | 0.54 | 0.63 |\n| REED | - | 27.7 | 55.5 | 0.56 | 0.64 |\n| REED | Qwen2-VL 7B | **24.6** | **62.2** | **0.58** | 0.64 |\n| REED | Qwen2-VL 2B | 26.3 | 59.9 | 0.57 | **0.65** |\n| REED | OpenCLIP | 29.0 | 55.8 | 0.55 | 0.64 |\n\n\n> The paper needs a figure illustrating the proposed framework, like the Figure 1 of the REPA paper.\n\nWe 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."}]

Author rebuttal:

Meta-review: This paper presents REED (Representation-Enhanced Elucidation of Diffusion), a systematic framework for incorporating representation guidance into diffusion models. The authors provide theoretical insights into how high-quality pretrained representations can improve diffusion model training, and propose two strategies: (1) integrating multimodal representations with synthetic data, and (2) designing an optimal curriculum that balances representation learning with data generation. The framework is validated across three domains—image generation, protein sequence generation, and molecule generation—showing strong performance and accelerated training compared to existing methods.

Most reviewer concerns were effectively addressed during the rebuttal, and all reviewers agree that this is a significant contribution. The authors are encouraged to incorporate the provided feedback into the final revision.
Decision: Accept (poster)