Title: U-REPA: Aligning Diffusion U-Nets to ViTs
Venue: NeurIPS 2025
Accepted: yes
Keywords: Representation Alignment

U-Net

Transformer

Diffusion
Abstract: Representation Alignment (REPA) that aligns Diffusion Transformer (DiT) hidden-states with ViT visual encoders has proven highly effective in DiT training, demonstrating superior convergence properties, but it has not been validated on the canonical diffusion U-Net architecture that shows faster convergence compared to DiTs. However, adapting REPA to U-Net architectures presents unique challenges: (1) different block functionalities necessitate revised alignment strategies; (2) spatial-dimension inconsistencies emerge from U-Net's spatial downsampling operations; (3) space gaps between U-Net and ViT hinder the effectiveness of tokenwise alignment. To encounter these challenges, we propose U-REPA, a representation alignment paradigm that bridges U-Net hidden states and ViT features as follows: Firstly, we propose via observation that due to skip connection, the middle stage of U-Net is the best alignment option. Secondly, we propose upsampling of U-Net features after passing them through MLPs. Thirdly, we observe difficulty when performing tokenwise similarity alignment, and further introduces a manifold loss that regularizes the relative similarity between samples. Experiments indicate that the resulting U-REPA could achieve excellent generation quality and greatly accelerates the convergence speed. With CFG guidance interval, U-REPA could reach FID<1.5 in 200 epochs or 1M iterations on ImageNet 256 $\times$ 256, and needs only half the total epochs to perform better than REPA under \textit{sd-vae-ft-ema}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes U-REPA, which is an adaptation of Representation Alignment$^{[1]}$ for U-Net-based diffusion models. In this paper, the authors mainly focus on how to deal with the difference between DiT and U-Net, facing challenges when directly applying REPA. In my opinion, although this paper finds a better way to align representations on U-Net-based diffusion models, the novelty and contribution are still limited. So I lean to reject of this paper.\n\n[1] Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThe paper explores a novel way to apply REPA on U-Net-based diffusion models, like SiTs and DiTs. When directly applying REPA on U-Net, the authors face several challenges and solve them properly:\n\n1. Divergent block functionalities due to skip connections: aligning intermediate U-Net layers.\n\n2. Spatial-dimension mismatches from U-Net downsampling: upscaling U-Net features and use post-MLP to solve spatial inconsistency.\n\n3. Feature space gaps hindering tokenwise alignment: replacing cosine similarity loss with the manifold loss.\n\n\nWeakness:\n\nI hold the belief that the authors have properly solved the issues when they apply REPA on U-Net. My main concern is that contribution and novelty of this paper. Comparing with REPA, U-REPA seems to change the position of alignment, and change the alignment loss.  \n\nBesides, a interesting point I find is that the authors claim that directly apply REPA cannot work well, so they propose a new manifold loss. However, in Eqn.(5), I still find the whole loss consists of the original REPA loss. The authors do not make any ablation on it. \n\nBesides, the authors do not provide any visualizations of the generated samples by using REPA and U-REPA. \n\nThe most important thing is that, the comparisons between REPA and U-REPA are unfair. In most of the tables, the U-Net used in U-REPA is better than the one used in REPA, this can be very strange. Why the authors do not put them under the same experiment settings?\n\nQUESTIONS:\n1. Make ablations on the loss term in Eqn.(5)\n\n2. Provide the visualizations of the generated samples of REPA and U-REPA.\n\n3. Can this method be applied on other U-Net-based diffusion models?\n\n4. Compare U-REPA and REPA under the same experiment setting, especially the trained U-Net.\n\nLIMITATIONS:\nAlthough the authors propose U-REPA, an adaptation of REPA, which can boost the performance of REPA on U-Net-based diffusion models, but the novelty and contributions are limited. Others please see the weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces U-REPA, an adaptation of the REPA alignment methods specifically for U-Net style architectures. The main contributions include:\n\n1. This paper suggests three key challenges in aligning features between diffusion U-Nets and ViT encoders: different block functionalities, spatial dimension inconsistencies, and larger feature space gaps.\n\n2. Through experiments, the authors verify that the fast convergence advantage of U-Net primarily stems from multi-scale hierarchical modeling rather than skip connection. And a series of SiT$\\downarrow$ models with modifications to U-Net style.\n\n3. The proposed U-REPA, a new framework developed to address the challenges through intermediate layer alignment, feature size adjustment, and manifold loss. \n\n4. Experiments demonstrate that U-REPA can achieve generation quality with FID < 1.5 on the ImageNet 256×256 dataset and outperforms REPA with only half the number of epochs under the same training settings.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**:\n\nThe paper proposes U-REPA, an effective representation alignment method for diffusion U-Nets that improves generation quality and training efficiency through intermediate layer alignment, feature upscaling, and manifold loss. The proposed architecture and alignment method can greatly improve the convergence and performance of image generation, supported by the extensive experiments.\n\n**Weaknesses**:\n\n1. Though the authors propose U-REPA to align features of the U-Net to ViT's, the alignment is conducted between a modified architecture SiT$\\downarrow$ rather than the original convolutional U-Net. Experiments concerning the convergence and performance of the convolutional U-Net are necessary.\n\n2. Despite the components of the macro components, skip connections and hierarchical design, the fast convergence of the original U-Net may come from the convolution modules. Discussion of convergence speed between convolutional nets and DiT should be included in this paper.\n\n3. The writing of the paper should be checked and improved. In L208 of the main body, there exists a meaningless sentence ``The similarity gap between SiT and SiT$\\downarrow$''. And the table, such as Table 1, is not cited in the main text. The symbol ${\\textbf{h}^{\\[n\\]}_{t}}$ is not specified in the text.\n\n4. It seems that the experiments are conducted only on global seed=0. Though this setting ensures the reproducibility, it still requires evaluating the quantitative experiment analyses and proposed methods across at least 3 random seeds, in order to avoid the fluctuation and cherry picking, and demonstrate robustness.\n\n5. The Time-aware MLP may not be compatible with the main story of the paper, and the experiment shows that it has a minor contribution to the whole performance. Thus, in my opinion, it should be treated as a trick and should be removed from the main paper.\n\n6. The proposed manifold loss is similar to the Marginal Distance Matrix loss in VA-VAE. It would be better to provide more discussion between the two losses in Section 3.3 and also refer to VA-VAE in that section, not just in the related work section.\n\nQUESTIONS:\nMy questions are listed in the weakness part above. If the authors address my concerns after the rebuttal process, I would like to raise the final score.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-30"}

comments: SUMMARY:
This manuscript aims to develop representation alignment method for UNet-based diffusion models. The authors perform experiments and point out that the challenges of this lies in the layer selection, dimension difference,  and feature compatibility. To this end, the authors take corresponding designs and propose to apply manifold loss for better alignment effect. Results show that the proposed U-REPA achieves better performance than vanilla REPA for UNet-based models.

STRENGTHS AND WEAKNESSES:
- Strengths:
1. The proposed method is well-motivated and the analysis is sufficient.
2. Extensive experiments demonstrate that the proposed U-REPA achieves great performance.
3. The analysis and ablation experiments are sufficient and reasonable.

- Weaknesses:
1. This work mainly involves some engineering analysis and engineering practice. The methods used, such as manifold loss, have also been proposed. 
2. This work is just a extension of REPA for UNet-based framework with some  engineering conclusions and do not give new findings for promoting AIGC fields.

QUESTIONS:
The author is recommended to optimize the layout to make it better.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the problem of adopting representation alignment (REPA) from the DiT to the diffusion U-Net. Despite U-Net-based diffusion models being known to converge faster than DiT, the authors found that adopting REPA to the U-Net architecture is not trivial due to (i) different functionalities of U-Net blocks, (ii) spatial-dimension inconsistencies from U-Net, and (iii) space gaps between U-Net and ViT. To address this issue, the authors propose U-REPA that aligns ViT features to the middle stage of U-Net with an upsampling strategy and incorporates a manifold loss that regularizes the relative similarity between samples. Experimental results show that U-REPA consistently and significantly outperforms REPA.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n1. The paper is overall easy to follow.\n2. Toy experiments demonstrate the potential of U-Net-based diffusion models, further motivating the necessity of applying REPA for diffusion U-Net. \n3. Comprehensive analysis of adopting REPA on diffusion U-Net.\n\n**Weaknesses**\n1. Important but missing comparison: In this paper, there is no comparison between SiT$\\downarrow$ and SiT$\\downarrow$ + U-REPA. While the proposed alignment scheme for the diffusion U-Net is intriguing, I am concerned whether the alignment loss is indeed beneficial for the diffusion U-Net architecture due to the lack of a comparison between SiT$\\downarrow$ and SiT$\\downarrow$ + U-REPA. For instance, DiT-XL/2 already outperforms DiT-XL/2+REPA at 400K iterations (DiT-XL/2 gets 11.02 FID scores vs DiT-XL/2+REPA gets 12.3 FID score.).\n2. Could the authors provide more analysis of incorporating the manifold alignment loss? I wonder whether the manifold alignment loss affects token-wise similarities. While SiT achieves larger token-wise similarity than SiT$\\downarrow$, I believe SiT$\\downarrow$ gets better generation quality than SiT. Based solely on these results, I think the direct alignment using L2 loss between the feature vectors is not beneficial for the U-Net.\n3. Because the proposed method includes an upsampling layer for alignment, I think a higher resolution experiment (e.g., ImageNet 512x512) is necessary.\n4. The font size in Figure 2 is too small to understand.\n\nQUESTIONS:\n1. Please answer the weaknesses.\n2. 1. Could the authors provide the SiT$\\downarrow$-XL/2 results in Table 4?\n\nLIMITATIONS:\nThe authors have not included the limitations and potential negative societal impact of the work, e.g., the societal impact of image generation models (deepfake, image copyright).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}