Title: Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion Model Acceleration; Representation Entanglement; Diffusion Transformers
Abstract: REPA and its variants effectively mitigate training challenges in diffusion models by incorporating external visual representations from pretrained models, through alignment between the noisy hidden projections of denoising networks and foundational clean image representations. We argue that the external alignment, which is absent during the entire denoising inference process, falls short of fully harnessing the potential of discriminative representations. In this work, we propose a straightforward method called $\textit{$\textbf{R}$epresentation $\textbf{E}$ntanglement for $\textbf{G}$eneration}$ ($\textbf{REG}$), which entangles low-level image latents with a single high-level class token from pretrained foundation models for denoising. 
REG acquires the capability to produce coherent image-class pairs directly from pure noise, substantially improving both generation quality and training efficiency.
This is accomplished with negligible additional inference overhead, requiring only one single additional token for denoising (<0.5\% increase in FLOPs and latency).
The inference process concurrently reconstructs both image latents and their corresponding global semantics, where the acquired semantic knowledge actively guides and enhances the image generation process.
On ImageNet 256$\times$256, SiT-XL/2 + REG demonstrates remarkable convergence acceleration, achieving $\textbf{63}\times$ and $\textbf{23}\times$ faster training than SiT-XL/2 and SiT-XL/2 + REPA, respectively. 
More impressively, SiT-L/2 + REG trained for merely 400K iterations outperforms SiT-XL/2 + REPA trained for 4M iterations ($\textbf{10}\times$ longer). Code is available at: https://github.com/Martinser/REG.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an improvement to REPA, a regularization designed to accelerate the training of diffusion transformers. While the original REPA aligns the spatial features of vision encoders like DINO with the features of the diffusion transformer, the proposed method, REG, further aligns the class token from DINO as well. In addition, the diffusion transformer is trained to predict not only the latent tokens but also an additional class token. This approach is simple, introduces minimal computational overhead, and significantly improves training speed, as demonstrated on ImageNet generation experiments.\n\nSTRENGTHS AND WEAKNESSES:\n[Strengths]\n\n- The proposed method is simple and effective.\n- The paper is clearly written and easy to understand.\n- Experimental validation is comprehensive, including both unguided and guided generation comparisons, ablation studies, and CKNNA analysis.\n\n[Weaknesses]\n\n- The paper only presents results for shorter training iterations compared to SiT and REPA. Although REG demonstrates higher training efficiency, it is unclear whether the performance remains strong or even improves with extended training. For instance, would training SiT-L/2 to 700K iterations or SiT-XL/2 to 4M iterations yield better results (Table 1)? Similarly, would REG benefit from training to 800 epochs (Table 2)?\n\n- REG incorporates four loss terms: velocity prediction and representation alignment for both latent and class tokens. However, ablation studies only remove the class token’s representation alignment loss (Table 6). There are no experiments that remove the class token’s velocity prediction loss or employ alignment using only class tokens. More detailed ablation is necessary to identify which loss components contribute most to REG’s performance.\n\n- Both REPA and this paper on REG omit a simple but strong baseline: allocating more training computation to timesteps that contribute more to performance. For example, the EDM [A, B] or Stable Diffusion 3 [C] paper samples timesteps from a logit-normal distribution instead of uniform, and emphasizes higher noise levels through timestep shifting for high-resolution images. Given that the core motivation of REG is to improve training speed, excluding this baseline weakens the claim.\n\n- Table 2 lacks baseline CFG performance without guidance interval, making comparisons incomplete.\n\n[A] Karras et al., Elucidating the Design Space of Diffusion-Based Generative Models, NeurIPS 2022.\n\n[B] Karras et al., Analyzing and Improving the Training Dynamics of Diffusion Models, CVPR 2024.\n\n[C] Esser et al., Scaling Rectified Flow Transformers for High-Resolution Image Synthesis, ICML 2024.\n\nQUESTIONS:\n- In the appendix, ImageNet 512×512 results are presented, but there is no EDM2 [B] baseline, and REG underperforms compared to EDM2. While outperforming EDM2 is not necessarily the goal, it would be valuable to include a discussion on REG’s applicability to EDM2 or on why EDM2 performs better.\n\n- In Table 3, why does DINO-B outperform DINO-L? An explanation or analysis would help clarify this counterintuitive result.\n\n[B] Karras et al., Analyzing and Improving the Training Dynamics of Diffusion Models, CVPR 2024.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose REG (Representation Entanglement for Generation) which improves/accelerates diffusion model training by integrating a high-level class token from pretrained vision models (e.g., DINO v2) into low-level image latents. Unlike REPA, which implicitly aligns the discriminative clean features with the image latents only during training, REG extends the alignments to be more explicit, by entangling the low-level image latents with the high-level class token from the pretrained vision models. This enables generation of coherent image-class pairs from noise, enhancing quality and training speed with minimal computational cost.\n\nSTRENGTHS AND WEAKNESSES:\n- Strengths:\n    - The paper is written clearly.\n    - The proposed method is simple but effective.\n    - Good visualizations of the proposed method.\n    - Experiment results are impressive.\n- Weaknesses\n    - Performance in T2I setting is questionable.\n    - Applicability of the proposed method to the widely-used pretrained diffusion models in T2I setup.\n    - Some details are not clear.\n    - Missing baseline; REPA-E.\n\nOverall, this paper lacks a baseline comparison with REPA-E, and performance report in T2I setting, which is more widely used than class-conditional setting. It would be much more impactful if the proposed method can be applicable in T2I setting, and even to the pretrained large models. However, the proposed method has a reasonable motivation with a suitable method, and shows good performance in class-conditional setting compared to the recent baseline. \n\nI think the proposed method is good enough to give accept. \nHowever, I am not sure if it could be impactful as the performance in T2I setting (which is more widely used in practice) is questionable.\nIf the authors can clarify some of my concerns, especially performance in T2I setting, I would raise my score to strong accept.\n\nQUESTIONS:\n- Can this method be still effective in T2I task? I guess the text embedding includes rich information as much as dino cls embeddings, and therefore the proposed method might not be as effective as in the simple class conditional setting.  \n- Where do you get the parameters for the linear layer mapping $cls$ to $cls’$? Is it trained together?\n- $F_0 = [f_0; cls_{0}]$ in $(N+1)\\times D_z$? How? $f_0 \\in \\mathbb{R}^{N \\times D_{vf}}$ and $cls_t \\in \\mathbb{R}^{1 \\times D_{vf}}$.\n- Dimension of $F_0$ and the output of $h_\\theta$. Are they in $\\mathbb{R}^{(N+1) \\times D_{vf}}$? I guess it should be, as $F_0$ is kinda proxy ground truth. Now, how do you maximize cos similarity? Inner product might yield $(N+1) \\times (N+1)$. Do you only maximize the diagonal terms (perhaps 1 if normalized)? \n- Related works involving the conditioning signal to the diffusion process, as this paper proposes, i.e., the condition is noised and denoised together. Can this be only used for training from scratch? Can it be used for finetuning pretrained large models?\n- How is FID measured?\n    - Are the same models used for Tab 1 and Tab 2?\n    - For Tab 1, I guess it is unconditional generation? 50k for synthetic and 50k for real?\n    - For Tab 2, it is confusing. CFG requires the model to take a conditional signal. Is it given through cross attention layer? Is one-hot class conditioned? How many samples are used to measure?\n- “Applying the loss in earlier layers yields superior results.” Why is this?\n\nLIMITATIONS:\n- Performance in T2I setting is not reported.\n- Applicability to pretrained T2I large models, e.g., SD1, SD3, FLUX.\n- Comparison with REPA-E", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nREPA allows training diffusion models faster through representation alignment with a pre-trained vision foundation models such as DINO. However, during sampling, models trained with REPA are not conditioned on any information from the vision backbone used for alignment. REG concatenates a noised class token from foundation models (e.g., DINOv2) with noised image latents as input to the diffusion transformer, enabling joint reconstruction of both image and semantic content. Experiments on ImageNet 256×256 show significant training acceleration and improved generation quality.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n**Simplicity**: REG is simple, general and easy to implement. It does not seem to require intensive hyperparameter tuning\n\n**Clarity**: The background section and explanation of the method are easy to follow, and can be understood on the first read. Figure 2 is easy to understand, and allows the reader to quickly understand their approach.\n\n## Weaknesses\n**Training overhead**: while the author discuss the sampling overhead, starting at line 219, they do not discuss the *training* overhead compared to training a regular SiT. In particular, they do not report the total cost of training with REG compared to a vanilla SiT without alignment.\n\n**Unfair claims**: occasionally, the authors make claims that I would judge to be unfair and that could mislead the reader. I would like the authors to revise their manuscript to clarify the following parts. Namely:\n- Lines 99-100\n> However, these methods rely on external alignment mechanisms that remain inactive during denoising inference, fundamentally limiting their capacity to fully exploit discriminative representations.\n\nREG does not use the external alignment mechanisms either during sampling. The above suggests that it does, hence it is misleading to the reader.\n- Lines 136-138\n>  To address the problems of REPA, the external alignment is absent during the entire denoising inference process, falling short of fully harnessing the potential of discriminative information. \n\nSimilarly, this is misleading, since REG does not rely on the external alignment module either.\n- Lines 165-166:\n> Full utilization of discriminative information.\n\n\"Full\" suggests that REG cannot be improved further. This cannot be proved empirically, hence it is not a fair claim.\n- Lines 169-189\n> This addresses a critical limitation of REPA, which relies on the external alignment mechanism that is absent during denoising inference.\n\nThis is not a fair claim: REG does not rely on the external alignment mechanism either during inference, only on the class token.\n\nQUESTIONS:\n**Noise schedule**: have the authors considered using a different noise scheduler for the diffusion process on the class embedding? Intuitively, I would expect that using a noise scheduler that is quicker to denoise the class token could be beneficial. Indeed, it means that the other tokens (that estimate the score for image tokens) could rely on the class embedding more, potentially leading to increased performance. I would appreciate if you could include such an experiment in your submission.\n\n**Additional parameters**: Why does REG requires 2M mode parameters than REPA (Table 1)? I did not see an explanation for this fact.\n\n**Failure Cases**: Under what conditions does REG perform poorly? Are there specific classes or image types where the method struggles? Including a more fine-grained evaluation (eg performance per class) would be interesting, for example compared to REPA of SiT.\n\nLIMITATIONS:\nThe only limitation discussed by the authors is that models might be under-trained. No other limitations are discussed. One could argue that while their approach is effective, just as REPA, it requires using an external model, hence the approach is not end-to-end.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-18"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new method for training diffusion models for image synthesis. \n\nA one-dimensional vector, typically obtained from self-supervised vision foundation models, that represents the \"semantic contents\" of the image is chosen and is noised similarly to the latents (from a VAE) corresponding to the image. \n\nThe DiT is then task to denoise (predict the flow field of) both the noisy latents and the noisy representational feature. \n\nThrough this auxiliary token injection into the image sequence and modeling of the task, the authors show that the representation quality and generation quality of the DiT thus obtained is greatly improved. Furthermore, they also show that training convergence is also dramatically accelerated.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\nThorough experimental study.\nGood ablations and analysis.\nWell-written and easy to understand. \n\nWeakness:\nEffectiveness of entanglement alone: While there are results showcasing the effectiveness of entanglement in terms of the generation quality, it's not clear how the representation quality benefits due to the addition of extra token and denoising task. \n\nREPA shows that the representation quality in DIT and performance improves with better vision encoder. This study seems to contradict that with the performance decreasing with bigger encoder (DINOv2-L). This might because of the hyperparam chosen / having the same timestep sampling schedule for the representation part as well as the latents. But there's some missing stuff here. \n\nWhile comparison with respect training iterations is helpful, it's also important to compare against SiT with respect to training FLOPs. REG, like REPA, has an alignment loss which requires additional flops compared to the standard SiT formulation. \n\nEvaluation of the generated representation: Experiments indicating the ability of the DiT to generate an accurate representation is missing.\n\nQUESTIONS:\nHow is it different than https://arxiv.org/abs/2504.16064v1?\n\nWhat happens to the representation quality in the standalone case (when you don't have the REPA loss)?\n\nI feel like a concrete explanation/investigation into how exactly the extra token and task helps image quality is lacking. It is mentioned that the key limitation of REPA is that it doesn't have some alignment option during inference, but it's not obvious as to why that is nedded or is a limitation in general?\n\nHow do you think the Vision foundational model chosen for the representation affects learning in the REG case compared ot REPA?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}