Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think

Thank you very much for your valuable comments and the time you devoted to reviewing our manuscript. We address each of your comments in detail below.

[W1]: REG improves the performance with extended training.

We further train REG, and as the training progressed, it consistently shows performance improvements across different task settings. For example, the two tables below present the ImageNet $256{\times}256$ results without and with CFG, respectively.

ImageNet $256{\times}256$ without CFG. Compared to other approaches, REG achieves the best performance, reaching an FID of 1.8 at 4M training steps, which is 4.1 lower than REPA.

ImageNet $256{\times}256$ with CFG. REG exhibits consistently improving performance with extended training time, achieving its best FID score of 1.36 at 800 epochs.

[W2]: The effectiveness of the an extra discriminative class token and the denoising task for it.

Ablation studies (Table 6) evaluate entanglement alone by removing alignment loss (REPA loss), assessing the extra noised class token and its denoising task (velocity prediction loss). Results show introducing a noised class token without a matching denoising objective causes a large performance drop (FID 98.67). This is because, without class token denoising, the model cannot autonomously generate a discriminative class token to guide inference. Moreover, there is an inconsistency between training and inference. This discrepancy disrupts the generation process. In inference, no explicit class token is provided; the diffusion model cannot rely on it as in training.

In contrast, introducing the noised class token with a dedicated training denoising task yields optimal performance (FID 26.67). This is because class token denoising loss enables the diffusion model to autonomously generate a semantically discriminative class token during inference. The generated class token can effectively guide the generation of other tokens. Moreover, this setup ensures consistency between training and inference.

Therefore, using the class token alone without corresponding denoising task (velocity prediction loss) yields no substantial benefit, while employing both jointly leads to clear improvements.

[W3, Q1]: REG achieve better performance in higher-resolution experiments on ImageNet $512{\times}512$.

Following your suggestion, we further train REG on ImageNet $512{\times}512$ and compare it with EDM2 [1]. Among all methods in the table below, REG achieves the best performance.

[1] Karras et al., Analyzing and Improving the Training Dynamics of Diffusion Models, CVPR 2024.

[W4]: REG perform better in CFG performance without the guidance interval.

Following your suggestion, we compare the missing CFG performance without guidance interval results in Table 2, and find that REG performs better than REPA.

[Q2]: The explanation and analysis about REG use DINOv2-B which perform better than DINOv2-L.

Thank you for your question. We provide the explanation and analysis for this issue below:

REG performs better with DINOv2-B, which is consistent with REPA. The observation that REG performs better with DINOv2-B than with DINOv2-L is consistent with findings reported in the REPA study. REPA does not exhibit a performance improvement when using larger encoders. Please refer to the content reproduced below from Table 2 of the original REPA paper, along with our experimental results.

Alignment difficulty and hyperparameter selection affect performance. Specifically, aligning shallow-layer representations of SiT (e.g., 4 layers) with deep-layer representations of DINOv2 (the final layer) may result in representational mismatches and increased training instability, thereby hindering effective alignment. For REG, final converged loss values reflect alignment and generation quality. Ablations on REG + SiT-B/2 (targeting DINOv2-B/L) use REPA’s loss formulation, reporting denoising (velocity prediction), alignment (negative cosine similarity), and total losses.

Specifically, Lower velocity prediction/total losses mean better generation; alignment loss ([−1,1]) closer to −1 indicates stronger alignment between SiT and DINOv2, reflecting more effective semantic understanding. Results show all three losses are lower when aligning with DINOv2-B than DINOv2-L. This highlights that SiT-B/2 aligns more easily and effectively with DINOv2-B, resulting in superior generative performance relative to DINOv2-L.

Furthermore, our hyperparameter ablation experiments for DINOv2-L show that, with appropriate tuning of β (0.03->0.05), which is the weight of the class token velocity prediction loss, the performance gap between DINOv2-B and DINOv2-L can be significantly reduced.

Overall, alignment difficulty and hyperparameter selection affect performance. To ensure a fair comparison, we follow the same experimental setup as REPA and adopt DINOv2-B as the alignment target across all models in our experiments.

Thank you very much for your valuable comments and the time you devoted to reviewing our manuscript. We address each of your comments in detail below.

[W1, Q1]: Performance in T2I setting.

Thank you for your question. We add text-to-image (T2I) results in the Appendix by adopting the same experimental configurations as REPA and UViT [1], utilizing MMDiT [2] as the backbone network. The results are presented below. REG demonstrates superior generation quality compared to alternative approaches on MS-COCO [3].

[1] All are Worth Words: A ViT Backbone for Diffusion Models, CVPR 2023
[2] Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, ICML 2024
[3] Microsoft COCO: Common Objects in Context, ECCV 2014

[W2, Q5]: Applicability to pretrained T2I large models.

REG is not limited to training from scratch; it can also be applied to fine-tuning pretrained large-scale models. However, due to time constraints and limited computational resources, we currently lack sufficient resources to fine-tune REG on pretrained large-scale text-to-image (T2I) models (e.g., SD1, SD3, FLUX). In future work, we plan to progressively extend REG to these models in order to evaluate its generalization capabilities and scalability in more complex generative settings.

[W3, Q2]: Explanation about the linear layer.

The parameter of the linear layer is randomly initialized and is trained with the diffusion model. The specific process is: The class token $c l s_t \in \mathbb{R}^{1 \times D_{v f}}$ of DINOv2-B is projected by a learnable linear layer (Linear: $\mathbb{R}^{1 \times D_{v f}}$ → $\mathbb{R}^{1 \times D_{z}}$) to get $c l s_t^{\prime} \in \mathbb{R}^{1 \times D_z}$. This linear layer is jointly trained with the entire diffusion model (SiT), to minimize the total loss function $\mathcal{L}_{\text {total }}$ (Equation 11).

[W3, Q3, Q4]: Details of calculating negative cost similarity.

Thank you for your question. In the final version of the manuscript, we will incorporate further revisions to improve clarity and facilitate better understanding. The following provides a detailed description:

• $F_0 \in \mathbb{R}^{(N+1) \times D_{vf}}$ is the reference representation of DINOv2 (pseudo ground truth). It is obtained by concatenating $f_0 \in \mathbb{R}^{N \times D_{v f}}$ and $cls_0 \in \mathbb{R}^{1 \times D_{v f}}$.
• $H_t^{[n]} \in \mathbb{R}^{(N+1) \times D_{z}}$ represents the output of the n-th SiT block, which is mapped by a trainable MLP projection $h_\phi$ to get $h_\phi\left(H_t^{[n]}\right) \in \mathbb{R}^{(N+1) \times D_{vf}}$
• $\mathcal{L}_ {REPA}$ (Equation 9) computes negative cosine similarity between $F_0$ and ${h_\phi\left(H_t^{[n]}\right)}$. It is minimized to encourage alignment.
• We compute negative cosine similarity only between corresponding token positions (i.e., the diagonal elements) to get the $(N+1)$ matrix rather than $(N+1)×(N+1)$.
• The detailed pseudocode is present below. F_0 denotes $F_0$, W denotes $h_\phi\left(H_t^{[n]}\right)$, D_vf represents the dimension $D_{vf}$.

# B: Batchsize   N+1: Length   D_vf: Dimension
# F_0: [B, N+1, D_vf]  - reference features (e.g., from SiT)
# W:   [B, N+1, D_vf]  - predicted features (e.g., from DINOv2)

# Step 1: Normalize both feature sets along the feature dimension
F_0 = F.normalize(F_0, dim=-1) # [B, N+1, D_vf]
W = F.normalize(W, dim=-1) # [B, N+1, D_vf]

# Step 2: Compute element-wise dot product between corresponding vectors
temp = F_0 * W  # [B, N+1, D_vf]

# Step 3: Sum across feature dimension to compute negative cosine similarity for each position
align_loss = -torch.sum(temp, dim=-1) # [B, N+1]

# Step 4: Average across all tokens and batches to get scalar alignment loss
align_loss = align_loss.mean() # Final scalar loss (e.g., -0.26)

[Q6]: How is FID measured?

We follow the evaluation protocol of REPA for computing the Fréchet Inception Distance (FID) [1], using the Inception-V3 model to extract 2048-dimensional features from both generated and real images. The FID is then calculated based on the mean and covariance of these feature distributions over 50,000 samples.

[1] Gans trained by a two time-scale update rule converge to a local nash equilibrium, NeurIPS 2017

[Q6.1]: Are the same models used for Table 1 and Table 2?

Yes, the same models are used for both Table 1 and Table 2. Table 1 reports results without using Classifier-Free Guidance (CFG), whereas Table 2 reports results with CFG.

[Q6.2]: Explanation of Table 1 and Table 2.

Table 1 and Table 2 all are trained for conditional generation using label inputs. In both tables, 50,000 images are generated and evaluated against real images using standard quantitative metrics.

[Q6.3]: Details of Classifier-Free Guidance (CFG) in REG.

Classifier-Free Guidance enhances the model's responsiveness to conditioning signals during inference by combining conditional and unconditional predictions. This improves the quality and diversity of generated images while avoiding the complexity of external classifier guidance. We sample 50K images to measure which follows SiT and REPA.

1. Form of conditioning input
   • In REG + SiT, the conditioning input is formed by the sum of the timestep embedding and the label embedding following DiT [1]. This combined vector serves as the conditioning signal.
2. Condition injection mechanism (AdaLN)
   • The conditioning vector is passed through an MLP and fed into the Adaptive Layer Normalization (AdaLN) [1] modules of each SiT block.
   • AdaLN maps the conditioning vector to layer-wise scale $\gamma$ and shift $\beta$ parameters to modulate the feature maps:
     $\text{AdaLN}(x) = \gamma(\text{cond}) \cdot \text{LN}(x) + \beta(\text{cond})$

[1] DiT: Scalable Diffusion Models with Transformers, ICCV 2023

[Q7]: Further explanation about applying the alignment loss in earlier layers yields superior results.

REG is built upon REPA and adopts the same configuration. We provide detailed ablation studies on alignment depth below. First, REG consistently outperforms REPA across different alignment depths, and applying the alignment loss in earlier layers yields superior results, which is consistent with REPA’s findings.

We believe this reflects an inherent trade-off: the alignment depth determines the diffusion model’s ability to understand and generate. Aligning at earlier layers allows the diffusion model to acquire DINOv2 semantics in its early stages, enabling subsequent layers to build on a strong representation and focus on capturing more high-frequency details for generation. However, aligning too early is not always better. For example, "depth = 2" performs worse than "depth = 4" indicating a trade-off between representation understanding and generation. Overall, earlier layers tend to yield superior results.

[Q8]: Comparison with REPA-E.

Thank you for your suggestion. Following the REPA-E paper, we compare the results of REG and REPA-E, and find that REG achieves better performance on ImageNet $256{\times}256$.

'*' indicates that VAE is updated during end-to-end training by REPA-E

Thank you very much for your valuable comments and the time you devoted to reviewing our manuscript. We address each of your comments in detail below.

[W1]: Analysis of training overhead in REG.

During training, REG requires alignment with DINOv2-B representations. We pre-process the DINOv2-B features offline and store them locally. As a result, during the training of the diffusion model, the DINOv2-B features can be directly loaded through the dataloader without almost no additional FLOPs.

We summarize the total training overhead in a below table, reporting the costs required to reach the same performance upper bounds claimed in the original SiT paper. All experiments are conducted on 8 NVIDIA A40 GPUs. Our results show that REG requires only 110K training steps to reach the performance level of SiT trained for 7M steps, reducing GPU hours by 98.36%. In addition, the single-step training speed of the two models shows no significant difference. These results highlight the training efficiency of REG, demonstrating faster convergence and significantly lower training overhead compared to prior methods.

[W2]: Modify some claims.

Thank you very much for pointing out the inappropriate phrasing in our manuscript. We will revise our manuscript accordingly in the final version to ensure clarity and accuracy.

• Lines 99-100 revision: However, these methods rely on external alignment mechanisms that do not take the discriminative representations as the input and denoising, which are unable to produce discriminative representations during inference to guide the generation process.

• Lines 136-138 revision: In REPA, the absence of the ability to autonomously generate discriminative representations to guide generation in inference may reduce the leverage of discriminative information effectively.

• Lines 165-166 revision: Improved utilization of discriminative information.

• Lines 169-189 revision: This design aims to address a limitation of REPA, which cannot autonomously generate discriminative representations to guide generation in inference. Because it relies on an external alignment mechanism during training to utilize discriminative features, rather than incorporating them as input and applying the corresponding denoising task.

[Q1]: Faster denoising noise scheduler for class token.

Thank you for your suggestion. Following your recommendation, we explore using different types of noise schedulers for different tokens in transforming velocity predictions to estimate the score. Specifically, we apply a cosine path type to the class token, aiming to achieve faster denoising. For the other tokens, we retain the linear path type, assuming they could benefit from stronger reliance on the class token during the generation process. We offer the detailed ablation results below.

However, this modification does not lead to performance improvements. We infer that the lack of gain stems from the training–inference inconsistency. During training, all tokens, including the class token, use the linear interpolant path for noise scheduling. Therefore, employing a different noise scheduler for the class token during inference creates a mismatch that negatively affects the model's performance.

Due to time constraints, we have not yet retrained the model under consistent training and inference configurations. In future work, we plan to retrain a model where the class token is trained and inferred under a faster denoising scheduler to further evaluate the effectiveness of this strategy.

[Q2]: Explanation of the additional parameters in REG.

Thank you for your comment. We explain the additional parameters (2M) introduced by REG below. Specifically, REG includes the following extra components: a Linear layer is used to map the class token from the DINOv2-B representation to the SiT space, a LayerNorm layer is applied to the mapped class token, and an additional Linear layer at the end of the network to decode the class token into the final output prediction. These components collectively enable the transformation and utilization of the class token within the REG framework. The detailed implementation and parameter calculations are provided below:

#DINOv2-B hidden_size
cls_token_dim = 768 

#SiT-XL hidden_size
hidden_size = 1152

# project the class token into the same embedding space via a linear layer to obtain 
self.cls_projectors2 = nn.Linear(in_features=cls_token_dim, out_features=hidden_size, bias=True)

# apply normalization to the mapped class token
self.norm = nn.LayerNorm(hidden_size, elementwise_affine=True, eps=1e-6)

# decode class token into the final output prediction
self.linear_cls = nn.Linear(hidden_size, cls_token_dim, bias=True)

Param = self.cls_projectors2 + self.norm + self.linear_cls
      = (768 × 1152 + 1152) + (1152 + 1152) + (1152 × 768 + 768) = 1.77M
      ≈ 2M

[Q3]: Discussion on failure cases.

Thank you for your suggestion. We conducted a large-scale comparison of the visualization results between REPA and REG. Overall, REG demonstrates a substantial improvement in image quality compared to REPA, However, for a small subset of challenging samples, such as tench (a kind of fish), both REG and REPA perform poorly. We attribute this to the extra noise in the ImageNet training data. For instance, these images depict a person holding a tench, but the provided label corresponds only to the tench. This discrepancy can introduce ambiguity during training and hinder the model’s ability to learn accurate class-conditional alignment. Due to time constraints, we will make a more fine-grained evaluation in the future.

[Q4]: Discussion on limitations.

Thank you for your comment. We agree that the reliance on an external model, such as DINOv2, limits the end-to-end nature of our approach. In addition, REG has not yet been extensively explored in terms of architecture and task diversity. There are several promising directions for future work:

• Extending to broader tasks and models, such as video generation and integration with autoregressive generative models, which may benefit from more complex discriminative guidance.
• Adopting stronger variants of VAE to enhance the model's ability to generate more accurate and diverse representations, potentially improving the overall performance in generative tasks.

We plan to investigate these directions, along with improved integration schemes and generalization to more diverse domains in future work.

Thank you very much for your valuable comments and the time you devoted to reviewing our manuscript. We address each of your comments in detail below.

[W1, Q2]: The effectiveness of the an extra noised class token and the denoising task for it.

Ablation studies (Table 6) evaluate entanglement alone by removing alignment loss (REPA loss), assessing the extra noised class token and its denoising task (velocity prediction loss). Results show introducing a noised class token without a matching denoising objective causes a large performance drop (FID 98.67). This is because, without class token denoising, the model cannot autonomously generate a discriminative class token to guide inference. Moreover, there is an inconsistency between training and inference. This discrepancy disrupts the generation process. In inference, no explicit class token is provided; the diffusion model cannot rely on it as in training.

In contrast, introducing the noised class token with a dedicated training denoising task yields optimal performance (FID 26.67). This is because class token denoising loss enables the diffusion model to autonomously generate a semantically discriminative class token during inference. The generated class token can effectively guide the generation of other tokens. Moreover, this setup ensures consistency between training and inference.

Therefore, using the class token alone without corresponding denoising task (velocity prediction loss) yields no substantial benefit, while employing both jointly leads to clear improvements.

[W2]: The explanation and analysis about REG use DINOv2-B which perform better than DINOv2-L.

Thank you for your question. We provide the explanation and analysis for this issue below:

REG performs better with DINOv2-B, which is consistent with REPA. The observation that REG performs better with DINOv2-B than with DINOv2-L is consistent with findings reported in the REPA study. REPA does not exhibit a performance improvement when using larger encoders. Please refer to the content reproduced below from Table 2 of the original REPA paper, along with our experimental results.

Alignment difficulty and hyperparameter selection affect performance. Specifically, aligning shallow-layer representations of SiT (e.g., 4 layers) with deep-layer representations of DINOv2 (the final layer) may result in representational mismatches and increased training instability, thereby hindering effective alignment. For REG, final converged loss values reflect alignment and generation quality. Ablations on REG + SiT-B/2 (targeting DINOv2-B/L) use REPA’s loss formulation, reporting denoising (velocity prediction), alignment (negative cosine similarity), and total losses.

Specifically, Lower velocity prediction/total losses mean better generation; alignment loss ([−1,1]) closer to −1 indicates stronger alignment between SiT and DINOv2, reflecting more effective semantic understanding. Results show all three losses are lower when aligning with DINOv2-B than DINOv2-L. This highlights that SiT-B/2 aligns more easily and effectively with DINOv2-B, resulting in superior generative performance relative to DINOv2-L.

Furthermore, our hyperparameter ablation experiments for DINOv2-L show that, with appropriate tuning of β (0.03->0.05), which is the weight of the class token velocity prediction loss, the performance gap between DINOv2-B and DINOv2-L can be significantly reduced.

Overall, alignment difficulty and hyperparameter selection affect performance. To ensure a fair comparison, we follow the same experimental setup as REPA and adopt DINOv2-B as the alignment target across all models in our experiments.

[W3]: Analysis of Training Overhead in REG.

During training, REG requires alignment with DINOv2-B representations. We pre-process the DINOv2-B features offline and store them locally. As a result, during the training of the diffusion model, the DINOv2-B features can be directly loaded through the dataloader without almost no additional FLOPs.

We summarize the total training overhead in a below table, reporting the costs required to reach the same performance upper bounds claimed in the original SiT paper. All experiments are conducted on 8 NVIDIA A40 GPUs. Our results show that REG requires only 110K training steps to reach the performance level of SiT trained for 7M steps, reducing GPU hours by 98.36%. In addition, the single-step training speed of the two models shows no significant difference. These results highlight the training efficiency of REG, demonstrating faster convergence and significantly lower training overhead compared to prior methods.

[W4]: The linear probing accuracy of the generated representation.

Following REPA’s setup, we compare REG and REPA’s linear probing accuracy on ImageNet. Results show REG’s superior performance, reflecting stronger representational understanding.

[Q1]: What the differences between REG and ReDi?

REG and ReDi [1] are proposed during the same period. REG is simpler, faster in both training and inference, and achieves superior performance. We provide a detailed comparison below:

1. Motivation：

   • REG: Overcomes REPA’s unaddressed issue: external alignment makes REPA unable to produce discriminative representations during inference to guide generation.
   • ReDi: Investigates integrating representation learning with generative modeling, enforcing the diffusion model to learn data-latent joint distributions.

2. Representation selection and usage methods:

   • REG (Representation Entanglement with Class Token): This method employs a single global class token which is concatenated with low-level latent tokens as the input.
   • ReDi (Representation Fusion with Dense Tokens): In contrast, this approach utilizes multiple semantic dense tokens, which are fused with low-level latent tokens via Principal Component Analysis (PCA) to construct the input.

3. Computational overhead of REG is lower:

   • Extra cost of REG is for handling one class token.
   • Extra cost of ReDi is PCA on all dense tokens.

4. REG performs significantly better than ReDi::

   Model	Iter	FID↓
   ReDi	1M	5.1
   ReDi	3M	3.7
   REG	400K	3.4
   REG	1M	2.7

[1] ReDi: https://arxiv.org/abs/2504.16064v1

[Q3]: Mechanism Difference between REPA and REG.

REPA aligns with discriminative representations during training, without incorporating them as input and applying the corresponding denoising loss. This design causes REPA to be unable to produce discriminative representations during inference to guide the generation process. In contrast, REG treats the discriminative class token as part of the input and applies a corresponding denoising loss, allowing it to guide training. During inference, REG can explicitly generate the discriminative class token and leverage it to guide the generation process.

[Q4]: The vision foundational model chosen compared with REPA.

The vision foundational model chosen strategy of REG is consistent with REPA, based on response to W2. Specifically, DINOv2-B performs better than DINOv2-L in both cases of REG and REPA, according to Table 2 of REPA and the ablation experiments of REG. Furthermore, to ensure a fair comparison, we adopt the same vision foundational model as used in REPA.