U-REPA: Aligning Diffusion U-Nets to ViTs
Authors: Yuchuan Tian
Hanting Chen
Mengyu Zheng
Yuchen Liang
Chao Xu
Yunhe Wang
Venue: NeurIPS 2025 poster ()
Keywords: Representation Alignment
U-Net
Transformer
Diffusion
Primary area: deep_learning

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}.

U-REPA: Aligning Diffusion U-Nets to ViTs
Yuchuan Tian', Hanting Chen', Mengyu Zheng?, Yuchen Liang', Chao Xu', Yunhe Wang²
State Key Lab of General AI, School Intelligence Science and Technology, Peking University.
2 Huawei Noah's Ark Lab. The University of Sydney. School of Mathematical Sciences, Poking University.
cianyc@stu.pku pku. edu. en, {chenhanting, yunhe. wang) www.w/auaweii.ccommmmm wwwewhhhwww.com com xuchao@cis pku. edu.
Abstract
Representation Alignment (REPA) that aligns Diffusion Transformer (DiT) hidden- states with VIT visual encoders has proven highly effective in DiT training, demon- strating superior convergence properties, but it has not been validated on the canonical diffusion U-Net architecture that shows faster convergence compared DiT'. However, adapting REPA U-Net architectures presents unique chal- lenges: (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, repre- sentation alignment paradigm that bridges U-Net hidden states and ViT features as 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 aligmment, 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 1M iterations on ImageNet 256 x 256, and needs only half the total epochs to perform better than REPA under sd-vae-ft-ema. Codes: https: //tp::///tttaee ggtpgbtttubb..... com hwtp///uhhhaannnuuuaaaa
follows:
1 Introduction
Representation Aligument (REPA) 45 a methodology that aligns features from Diffusion Trans- formers (DiT) [30) modern visual encoders, has been demonstrated significantly accelerate DiT training. This approach holds particular significance given the growing prominence of DiTs, which have gained mainstream adoption in diffusion models and are extensively applied across image gener- ation 23 and video generation domains [88
However, emerging empirical evidence suggests that U-Net 33 architectures might present more advantageous alternative to DiTs in certain scenarios [8 39: 38): U-Net-based models exhibit substantially faster convergence while achieving generation quality comparable their transformer-based counterparts. This dichotomy mo- tivates our core research inquiry can modern Vision Transformer (ViT - visual encoders be effectively adapted guide diffusion U-Net training through alignment mechanisms similar to REPA, thereby potentially elevating the convergence speed ceiling of diffusion models?
However, establishing effecttve alignment between U-Net architectures and ViT-based encoders presents challenges. Unlike Diffusion Transformers (DiT's) that share structural similarities with Vision Transformers, U-Net architectures exhibit fundamentally different operational characteristics. Specifically, both DiT and VIT adopt isotropic architectures composed of uniformly stacked trans- former blocks, which inherently facilitates straightforward parameter alignment between the two
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Tokenwise A Algmment
Manifold A lignmient
Lineer
Transformer Block
Transformer
Transformer Black
Transformer
Transfoemer Block
Transfurmer Block
Trunsocmer Block
ViT-Based Encoder
Diffusion U-Net
Figure 1: The proposed U-REPA framework. We investigated and found that semantic-rich intermediate layers are the best for representation alignment, dimension and space gaps hinders alignment efficacy. To counter these challenges, we scale-up features and propose manifold alignment.
frameworks. In contrast, U-Net's skip connections create strong interdependencies between shallow and deep network layers by linking them together, resulting in different feature propagation dynamics. This architectural disparity renders conventional representation alignment strategies developed for DiT architectures inapplicable to U-Net frameworks. Furthermore, the progressive downsampling operations in U-Net generate feature maps with spatial dimension mismatches compared to the fixed-scale feature representations in ViT encoders, introducing additional complexity in establishing arehitecchtltttlall correspondence. In addition, features from high-stage U-Net and ViT have large space gaps, forming a barrier for cosine similarities as metrics. Forcibly using tokenwise similarity as loss is not necessarily the best option. This induces to rethink about the optimization objective.
In order to conquer these challenges, we propose U-REPA, a framework that aligns U-Net hidden states to features from ViT encoders. Firstly, Our analysis reveals that skip connections fundamentally alter the functional specialization transformer blocks in U-Net architectures. By establishing direct dependencies between early-stage and late-stage layers, these cross-connections induce a hierarchical redistribution of semantic information, with intermediate blocks exhibiting the highest semantic density. This pattern was empirically verified through controlled ablation studies on DiT augmented with skip connections, where progressive layer-wise evaluations demonstrated peak semantic richness median network depths.
The intermediate hither-stage layers, which contain semantically dense representations, require precise alignment with the ViT-based visual encoder. However, these critical layers undergo spatial downsampling in the U-Net architecture, necessitating explicit spatial dimension reconciliation between U-Net features and VIT features during representation alignment. Through empirical strategies, we identified an optimal solution: performing
exploration of various resolution-matchhe linear transformation via MLP U-Net features prior upsampling operations, which achieves superior alignment performance compared to alternative approaches.
Further analysis revealed fundamental incompatibility of measuring cosine similarity between the feature spaces of U-Net and ViT encoders. Enforcing strict token-wise similarity constraints proves excessively rigid due to inherent architectural discrepancies. To address this, introduce manifold loss that implements soft alignment through relational regularization. This loss operates on the relative geometric relationships between samples rather than imposing direct feature correspondence, thereby accommodating gross-architectural variations. Comprehensive experiments demonstrate that our proposed U-REPA framework effectively bridges the U-Net-ViT alignment gap while preserving the distinct advantages of both architectures.
Our contributions are as follows:

---

1. We identify U-Net's representation aligmment to be challenge due to different block functionalities, spatial-dimension inconsistencies. and larger feature space gaps.
2. We evaluate the contribution of downsampling and skips in U-Net and demonstrate U-Net's potential advantage over DiTs.
3. We propose U-REPA, a framework that evaluates layers, investigates the best scale-up policy, and introduces manifold-space loss in aid of alignment.
4. We conduct experiments and verify the effectiveness of our U-REPA framework in terms of fast convergence. Specifically, U- U-REPA reaches FID < 1.5 in just 200 epochs on ImageNet 256 x 256; and it reaches 1.41 FID, with only half the epochs of REPA under the same training setting.
2 Related Work
The development of diffusion architectures. The conventional diffusion works 17 leverages a U-Net [3 architecture, whose hasic block is a concatenation of convolution layers and self-attention. More recent architectural innovations in diffusion models have witnessed paradigm shift from conventional U-Net frameworks toward transformer-basee architectures. The emergence of Diffusion Transformers I 30] demonstrates their competitive performance despite abandoning the inductive biases inherent in U-Net designs. U-VIT (1] represents an intermediate architecture that preserves U-Net's hierarchical structure but replaces convolutional blocks with transformer layers, notably omitting the traditional downsampling operations. Subsequent developments have further streamlined the architecture: DiT 30 adopts a pure transformer backbone with isotropic scaling, while SiT 126) integrates the transformer architecture into the RectifiedFlow framework. Some other works either improve the micro-designs G 25),, or focuses on architectural efficiency
35 36
In contrary to these DiT works, some works still sticks to U-Net architectures and offer valuable rethinking on this conventional architectural preference: in pixel-space image generation, works including SimpleDiffusion [18) and HourglassDiT (7] still sticks to U-Net; in with its variants like U-DiT 39), Playground v3 (24), and DiC [8 extending its success latent-space diffusion through simple 3 designs. While these implementations empirically validate U-Net's accelerated con- vergence and stable training dynamics compared to transformer-batt based alternatives, current research predominantly focuses proposing architectural modifications rather than uncovering the reasons of U-Net's superior diffusion performance.
Techniques for better DiT performance. Building upon the success of self-supervised learning [6]. MDT 113 and MaskDiT (47 pioneer masked image modeling in diffusion frameworks by adaptively masking a good proportion of input patches during training. Other than the masking strategy, bunch of diffusion works refer higher-level semantic guidance from off-the-shelf pretrained models that significantly improves generation quality. REPA 451 establishes feature alignment between ViT-based encoder embeddings and diffusion latent spaces through contrastive learning. LightningDiT [4] innovates through an improved VAE distilled from MAE [16] and DINOv2 (28). Ma et al. (27] introduces CLIP [31] and DINO 22 to verify inference- time scaling of diffusion models. These methods demonstrate that a higher-level semantic-rich fcature-map map from pretrained vision encoders helpful diffusion-based generation.
Method
3.1 Preliminaries: REPA for DIT
Representation Alignment (REPA) 45 distills Diffusion Transformers with semantic features from off-the-shelf ViT-based vision encoders (e.g. DINOv2 28). CLIP BI. MAE [16), et cetera). Given a ViT-based vision encoder f and clean image X., let y. f(x.) ) N denote its patch embeddings, where N and / represent the number patches and embedding dimension, respectively. REPA establishes feature alignment between y. and the projected diffusion encoder ouiputs h (h+). where h, = fo (z.) the latent representation from the diffusion transformer at timestep i, and he is trainable multilayer perceptron (MLP).
3

---

SIT 5kip Connection
SIT (U-Net)
Thastrrma
Dept
Figure 2: Investigating alignment with respect to encoder depths on diffusion models with skip connections. Left: SiT with skip connections. Due to the change of block functionalities due to newly established skip dependencies, the most optimal encoder depth is shifted towards the middle of the model. Right: SiT the U-Net-based SiT model. Shadowed region represents higher U-Net stage. The plot infers that stage transitions (downsampling& upsampling in U-Net) bring large block functionality gaps. Alignment within higher U-Net stage is thus necessary for alignment performance.
The alignment is enforced by maximizing token-wise feature similarities, i.e. the similarity of a token from DiT hidden-state with its corresponding counterpart in the ViT encoder feature:
N
sim (y=", (boh)
CREPA Φ) == -,x,,,,
(1)
N
n=l
where sim(-:-) denotes similarity metric (e.g., cosine similarity). Typically, z, adopted as the output from early layers (the original work adopts layer index 8) in DiT for better alignment. This aligmment term is combined with the basic flow-based diffusion objective (i.e. SiT (26]) through tunable coefficient À >> O, and the final loss for diffusion model training is formulated as follows:
C = Lvelocity + ACREPA
3.2 Evaluating the Potential of U-Net
In diffusion models, U-Net and isotropic architectures (e.g., DiT) exhibit distinct design philosophies. While DiT achieves state-of-the-ar results through scalability and integration with advanced tech- niques, U-Net-based methods emphasize faster convergence (39). To dissect U-Net's efficacy, isolate its two core components: skip connections and downsampling.
1. Skip Connections:
Provide shorteuts between encoder and decoder layers, theoretically aiding gradient flow and feature reuse.
2. Downsampling:
Reduces spatial resolution (typically by scale factor of at each stage) to enable hierarchical, multi-scale feature learning. Critically, downsampling is always paired with skip connections to mitigate information loss.
Toy experiments on U-Net components. On top of DiT, we perform toy experiments that reveal the contribution of components mentioned above.
ImageNet 356x256, DiT 400K, cfg=1
Model
FLOPs (G)
FID.
ISt
DIT-XI/2
118.6
19.47
DiT-XL/2"
118.6
20.05
66.74
Skip Connections
114.1
19.86
67.29
Downsampling
108.8
13.78
88.93
DIT/-XL/2 (+Tricks)
108.8
11.02
100.35
Table Evaluating the contribution of U-Net components in terms of fast convergence. Experi- ments are conducted using hyperparameters from [30] for 400K iterations. Model depth is changed when a modification made such that the overall FLOPs is kept almost the same with DIT.
This suggests that U-Net's fast-convergence advantages primarily stem from multi-scale hierarchical modeling via downsampling. not skip connections. Downsampling compresses features into compact, semantically rich representations, accelerating learning while maintaining information flow through skip-augmented decoder layers. However, skip connection is not useless as compensates for the information loss due to downsampling.

---

Building on this insight, we propose DiT, (and S for the flow-based version, following the naming convention of (26) by adding tricks of RoPE [37 and SwiGLU following previous work 6 39
3.3
Aligning U-Net to ViT Encoders
Since U-Net has good potentials to achieve excellent generation, we are motivated to investigate whether REPA could also work U-Net. We are first focused on the block functionality pattern and investigates the most optimal position for alignment; then we are interested in feature size alignment problems; lastly, we are dedicated to merging space gaps between features from U-Net and ViT, respectively.
Position for alignment. Regarding the concern that block functionalities differ between U-Net and DiT, the comparison of prior studies [12 S demonstrates divergent hierarchical specialization: U-Net architectures typically employ mid-network layers for high-level semantic synthesis while reserving shallow layers for low-level image refinement, whereas DiTs exhibit a totally different pattern early layers primarily govern semantic-rich outline formation with deeper layers handling detailed image refinement.
These previous findings find empirical support in REPA's experimental findings, where representation alignment proves most effective when applied to initial transformer blocks. This phenomenon stems from DiT's early layers encoding semantic-rich representations that align well with the semantically dense outputs of ViT-based visual encoders, enabling meaningful guidance. Unlike the straightforward DiT architecture, the inherent skip connections in U-Net architectures induce fundamentally distinct block functionality compared to Diffusion Transformers. While all blocks in VIT or DiT maintain homogeneous computational roles, following continuous flow of transition from input to output, U-Net's cross-layer shortcuts establish direct dependencies between shallow and deep layers. fundamentally altering reature-map evolution patterns. As shown in Fig. (R), DiTs with skips indicates median layer is the best for representation alignment. The same pattern goes for U-Net (Fig. (L)) despite the downsampling stage.
Feature size alignment. The implementation of alignment between Diffusion U-Net's median stage and VIT encounters critical spatial resolution dilemma stemming from architectural disparities. While our analysis identifies mid-network U-Net features as optimal semantic carriers, their spatial dimensions drastically differ from VIT's full-resolution token sequence. This dimensional mismatch obstructs REPA's bosen-wise similarity computation, which requires strict cardinality matching between compared features.
In order to align two feature-maps (i.e. from U-Net and ViT encoder, respectively), from the macro level we advocate for upscaling the smaller-sized U-Net features rather than downscaling the larger visual encoder features. This design principle stems from the critical observation that compressing VIT's high-resolution features match U-Net's reduced dimensions inevitably discards fine-grainee visual information, thereby degrading alignment effectiveness. Preserving ViT's native resolution while expanding U-Net's bottleneck features proves essential for maintaining semantic fidelity.
At the implementation level, we empirically evaluated various upscaling strategies for U-Net features:
1. Upscale first and then MLP: feature upsampling is performed before passing the feature into the MLP.
2. Upscale within MIP: the MLP also acts a feature upsampler that receives a low- resolution input and outputs a high-resolution one via linear mapping and pixel un-shuffling.
3. MLP first and then upscale: the feature from higher-stage U-Net is first passed through MLP and then upsampled.
Among the three options, we found "MLP first and then upscale" is the best both in terms of performance and efficiency (minimum FLOPs cost), which will be discussed in the Ablation Study in Sec.
Manifold space alignment. Though we select the most suitable U-Net feature for alignment and keep dimensions between U-Net and VIT features aligned, challenges remain in feature space compatibility. First, compared to the structural congruence between DIT and VIT encoders, the architectural discrepancy of U-Net (with its skip connections and hierarchical downsampling) creates more pronounced feature distribution gap between U-Net hidden states and visual encoder outputs.
5

---

0.55
0.80
0.55
0.50
0.45
0.30
25K
50K
75K
1O0K
125K
150K
175K
2004
Traiiing Staps
Figure 3: The convergence of average tokenwise similarities. While SiT-L/2 could achieve better tokenwise similarities, SiTT converges at lower similarity value, indicating difficulties of feature alignment.
Second, the dimensional transformation required for alignment inevitably modifies U-Net's native feature space characteristics.
Some recent works on Diffusion U-Net 41 34 reveals that higher-stage U-Net features are low- frequency subspaces that discards higher frequency compoenents, including noises. Gaps are in- evitable when evaluating the cosine similarities of detail-rich, high-frequency--ice vectors and flat, vectors. In this sense, strict token-level alignment constraints like the original REPA loss prove suboptimal under these conditions, they assume implicit feature space homogeneity between aligned modalities.
www.frrqueney-dommieee
As is shown in Fig. 3 we conducted continuous measurements of token-wise cosine similarity against ViT features during training. The two models that we compare are SiT-L/2 of isotropic, standard transformer architecture and SiT/-L/2 of U-Net architecture. Our experiments revealed characteristic learning trajectory: while U-Net achieves slightly faster similarity improvement in early training phases thanks skip connection that helps convergence its progress stagnates beyond this point, ultimately plateauing at 0.60 notably inferior to DiT's s sustained growth reaching around 0.63 similarity. The similarity gap between SiT and SiT This phenomenon suggests that naively aligning U-Net with ViT encoders through angular similarity metrics alone encounters inherent limitations due to architectural incompatibilities.
Rather than strict token-wise regularization, we resort to looser objectives that does not require rigid augular alignment. Inspired by manifold knowledge distillation [14, hold that aligning similarities between samples from the same feature space could be promising solution. Hence, we define Manifold Loss CML
(0, Ex,is [d(y., he (h.))),
(3)
where
= ||sim (y!' ا(رر - sim (ns (b((s)) 112.
(4)
In the formula, cosine similarity adopted as the similarity metric, and F represents Frobenius Norm of matrices. By introducing affine hyperparameter w, the overall optimization target is then formulated as
C = Lvelocity + A (CREPA + wLM.).
(5)
3.4 Other Improvements
We also propose and evaluate some other improvements. Due to page limits, the proposed methods and corresponding ablations enclosed in the Appendix.

---

Experiments
4.1 Experiment Setup
Experiment settings. Our implementation completely adheres the training protocol established in REPA [45]. Following the architectural configuration of latent diffusion models [32), we employ the identical VAE variant (sd-vae-ft-ema) and adopt the AdamW optimizer. To ensure fair comparison, we maintain identical hyperparameter settings across all experiments: global batch size of 256, fixed learning rate of 1e - 4, and disabled weight decay (set 0). (31, 82) is set as (0.9, 0.999). All experiments are conducted on the ImageNet 2012 benchmark [8] under a controlled environment with fixed random seed (global seed=0). 8 NVIDIA A100 GPUs are used for main experiments.
For main experiments (Tab. 5 we apply guidance interval 21 [0, 0.7] and SDE sampling according to the convention of REPA 45 for fair comparison. We select smaller cfg of 1.65, because found it is better for our architecture, different from SiT. For all ablation experiments, we train models for 100K iterations, which sufficient show the trend of model performance; sampling is conducted with the default setting of the official REPA codebase, i.e. cf - 1.8 in ODE and guidance interval [0, 0.7].
Model settings. By aligning channel dimensions and FLOPs with standard Diffusion Transformers (DiTs SiTs), our PREPA-compatipe variants maintain architectural parity while introducing critical adaptations for U-Net principles. The base model (SiT)-B) employs a stage arrangement of [5,5,5], achieving 199.7G FLOPs. Scaling larger models, have the L variant (686.6M) and XL variant (954.4M params) that increases channel width (1024 vs. 1 1152 in base) through increased stage-wise block allocation [[9,14,9] vs. [10,16,10) 10))). Notably, when FLOPs are aligned, SiT models usually have more parameters than SiTs due to increased depth.
Model
Params (M)
FLOPs (G)
Patch Size
Channel
Heads
Blocks in Stages
SiT 199.7
24.1
768
12
[5,5,5]
-L
686.6
79.3
2
1024
16
[9,14,9]
S 954.4
109.3
1152
16
[10,16,10]
Table 2: Configurations of SiT architecture at different model sizes. The proposed SiT! in U-Net architectures aligned to DiT's in terms of FLOPPs and channel dimension.
4.2 The Advantage of U-REPA
Comparing S with SiT different scales. We evaluate our U-REPA alignment method on ImageNet 256 under generation setting with c/g (REPA framework without classifier-free guidance). As shown in Table 3. our approach consistently improves generation quality while significantly reducing computational costs across model scales. For the base-size SiT-B/2 variant, integrating U-REPA achieves a 39.3% improvement FID (from 24.4 to 15.3) with comparable FLOPs (24.1G vs. 23.0G) and identical training iterations (400K), demonstrating that feature alignment enhances parameter efficiency without additional training overhead. The acceleration effect becomes more pronounced in larger models: for SiT-L/2, U-REPA reduces required iterations by 42.9% (700K~40066 while simultaneously lowering FLOPs (79.3G vs. 80.8G) and achieving a 30.9% FID improvement (8.4+5.8). Most notably, the XL-scale variant with U-REPA (cf. Fig. for FIDs vs. Training iters) attains state-of-the-an FID (5.4) using 90 fewer iterations (400K vs. 4M) and fewer FLOPS (108.8G vs. 118.6G) compared to the baseline, proving our method's fast convergence.
We also demonstrate the advantage of the proposed U-REPA framework when measuring by parame- ters (rather than computation FLOPs), as shown in Fig. 4 Though U-Net brings extra parameters when FLOPs are aligned with DiT', the advantage of SiT +U-RAA is obvious as depicted in the Parameter versus FID plot.
Convergence performance. We also compare our method with previous State-of-the- Artttsss as shown Tab. 5 Our proposed SiT!+U-REPA achieves competitive FID of 1.48 with only 200 training epochs, significantly outperforming existing methods in training efficiency. Notably, while state-of- the-art masked diffusion transformers like MDTv2-XIL/2 require 1,080 epochs to reach 1. 1.58 FID, our method attains better performance (1.48) with 80% fewer iterations. Even compared to the SOTA
7

---

sir nots 1-ttRARA
ImageNet 256x256, w/o efg
Parameten
900
Model
FLOPs (G)
Iter.
FID
SiT-B/2+REPA
23.0
400K
24.4
SiT(-B/++-REEPA
24.1
400K
15.3
252
22.5
26.0
17.5
15.
12.5
10.0
5.0
SIT-L/2+REPA
80.8
700K
8.4
F10
SITU-L/2AURREPA
79.3
400K
5.8
SiT-XL/2+REPA
118.6
5.9
Comparing +UT REPA against SiT+REPA in terms of parameter scalability. While the U-Net architecture makes the diffusion model parameter-rich compared with same-FLOPs Dif- fusion Transformers. S models still outcompetes SiTs by large mar- gins in terms of parameters.
Figure 4:
SITU-XL/2U-REEPA
108.8
400K
5.4
Table 3: Comparing U-REPA against REPA across vari- model sizes without classifier-free guidance. U-Nets equipped with U-REPA show excellent capabilities. Notably, U-REPA achieves 10x faster convergence compared with REPA in terms of performance w/o CFG.
ImageNet 256x2566 w/ cfg
Model
Dep.
Feat. Dim.
FID,
1St
SiL/-XL/2
9.42
117.6
-IT/-XL/2 (REPA)
16
9.35
119.5
S-T/2
8
6.73
148.3
-
16
8
6.43
154.1
SiT
18
8
6,25
156.2
600K Training
SYT/-XL/2
20
8
6.36
155.8
SYT/-XL/2
24
8
6.77
150.5
SIT/-XL/2
28
16
8.40
130.5
Comparing SIT++U- REPA against SIT+REPA in terms of convergence speed. SiT XL/2 convergences much faster than SiT-XL/2 with the help of U- REPA.
Figure 5:
SiT/-XL/2
32
16
13.10
99.4
Table 4: Ablations on encoder depths for alignment in SIT Feat. Dim. stands for the spatial height& width at the certain layer. Compared with the default REPA setting, aligning at the centering layer (higher stage in U-Net) performs much better.
SiT-XL/2 REPA baseline (800 epochs for 1.42 FID), our approach uses only 1/2 of the training epochs (400) while achieving better generation quality (1.41 FID). The results demonstrate that the proposed U-REPA establishs new efficiency frontier for diffusion models.
4.3 Ablation Studies
Encoder depths. The ablation study encoder layer depths for feature alignment (Tab. coincides with the pattern of DiT with skip connections, as we analyzed in Sec. 4 despite progressive down- sampling operations that reduce spatial resolution, the centermost layers exhibit optimal alignment efficacy. For the SiT\-XL/2 model, aligning features at layer 18 (midway through the 36-layer architecture) achieves peak performance with 6.25 FID and 156.2 IS, outperforming both shallower and deeper alignment points. This phenomenon persists even as the spatial dimension (Feat. Dim.) halves from 16x 16 8x8 in the intermediate stage, indicating that semantic richness-nnt spatial oesolution-dominee alignment quality. Performance degradation occurs when alignment takes place at shallower deeper stages, even though the feature size kept the same with DINO in these stages.
Alignment dimension choices. The comparative results in Table о reveal that upsampling U-Net's higher-stage features (f2) to match DINOv2's native resolution achieves superior performance (5.72 FID, 161.6 IS), outperforming the alignment alternative in generation quality. This demonstrates that preserving ViT encoder's 's original feature granularity during alignment is beneficial for alignment.
Feature-map upscale choices. Among the three upscaling options mentioned in Sec. 4 we figure out that upsealing U-Net hidden states after getting passed through MI.P is the best option, achieving 5.72 FID and 161.6 IS. This option is also the most optimal one terms of computation cost analysis.

---

ImageNet 256x256, w/ cfg
Alignment Choices
FIDL
1St
U-Net DINOv2.2
5.99
158.8
ImageNet 256x256, w/ efg
U-Net DINOv2
5.72
161.6
Model
Epochs
FID
Table 6: Alignment dimension choices. Up- sampling higher-stage U-Net features in align-
Pixel diffusion
ment with ViT performs better due less infor- mation loss.
ADM-U 9
400
3,94
VDM+ 20
560
2.40
Simple diffusion
18
800
2.77
ImageNet 256 256x256, w/ cfg
Latent Diffiusion Transformer
Alignment
FID
IS1
U-ViT-H/2
240
2.29
Upscale before MLP
5.84
158.5
DiffiT [15]
1.73
Upscale in MLP
6.36
153.4
DiT-XL/2 [30
1400
2.27
Upscale after MLP
5.72
161.6
SiT-X1/2 266
1400
2.06
Table 7:
Feature-map upscale choices. Among the three options, Upscaling after pass-
Masked Diffusion Transformer
MaskDiT
47
1600
2.28
ing through MLP performs best; and it has lower cost as the small-sized feature map is passed through MLP.
2DTV2-XL/2
13
1080
1.58
Representation Alignment
SiT-XL/2 + REPA (45)
800
1.42
SiT U-REPA (Ours)
200
1.48
ImageNet 256 256, w/ cfg
SiT U-REPA (Ours)
400
1.41
Model
w
FID
IS1
SIT STT--XL/+U-RREPA
0
6.25
156.2
Table 5: Comparing U-REPA against State-of- the-Art baselines with classifier-free guidance. U-REPA could reach FID 1.5 in merely 200 epochs and FID = 1.41 in 400 epochs; The pro- posed method converge 2x faster while achieving lower FID.
-
2
5.81
160.8
SIT - SiT SIT--XL/2UU-REPA
3
5.72
161.6
5.79
160.6
Table 8: Adjusting weight in Eq. Mani- fold loss boosts U-Net's s alignment performance. The most optimal result is taken at 3.
Manifold loss weight w. The ablation study on alignment weight w in Eq. demonstrates clear performance peak 3 achieving the lowest FID (5.72) and highest IS (161.6) among tested configurations.
4.4 Higher Resolution Experiments
At the higher-resolution ImageNet 512x512 (w/ efg) setting, U-REPA remains clearly superior to the
REPA baseline (Tab. 99. Using SiTJ-XI/2, U-REPA reduces FID from 2.44 to 2.21 and raises IS from 247.3 to 274.7. These results indicate that U-REPA's benefits persist at 512 resolution, yielding better distributional fidelity and sample quality/diversityy and demonstrating strong scalability.
ImageNet 512x512, w/ cfg
Alignment Choices
FID
IS1
SiT-XL/2 + REPA
2.44
247.3
SiT U-REPA (Ours)
2.21
274.7
Table 9: Comparing U-REPA against REPA on ImageNet 512x512. On higher resolution, the
proposed U-REPA still maintain clear advantage.
4.5 The Energy Cost Advantage of U-REPA
also assess the energy-cost advantage of U-REPA over REPA. We train on eight NVIDIA A100 GPUs and record each GPU's power draw. Combining the measured power with the training duration.

---

ImageNet 256x256
Model
Avg. Pow. (W)
Training Hours
Est. Energy (J)
FID
SIT-XI./2RREPA (4M iter) SIT/-XL/2AU-REPA (2M iter)
373.2
302.7
x.25 10e9
1.42
295.3
230.7
1.96 x 10e9
1.41
Table 10: Energy cost comparison. We compare the energy cost of U-REPA (at 2M iters) and REPA (at 4M iters). U-REPA could significantly reduce the cost of training a State-of-the-AN diffusion model.
estimate the total energy consumed. The statistics for average power and estimated total energy used by all GPUs are summarized in Tab. I0 Results indicate that our U-REPA method is "greener", costing far less energy.
Reducing training energy directly curbs operational CO2 emissions. Methods that achieve comparable accuracy with lower energy, such as U-REPA vs. REPA in our study, advance both sustainability and the economic viability of large-scale AI.
Conclusion
In this paper, we propose U-REPA, an adapted version REPA on Diffusion U-Net. We identify key challenges in U-Net hidden state alignment and show that U-REPA effectively bridges the gap between U-Net-based diffusion models and ViT-based encoders. By aligning intermediate features, resolving spatial mismatches via post-MLP upsampling, and enforcing mawe regularization, U-REPA achieves faster convergence and an FID score of 1.41 on ImageNet-256x2222 at 2M iters.
Acknowledgement. This work is supported by the National Key R&D Program of China under grant No. 20222D0160300 and the National Natural Science Foundation of China under grant No. 62276007. This work is funded by Peking University-BHi Carbon and Climate Wei-Ming PhD Scholars Program (Program Name: Research on Low-Carbon and Energy-Efficient Large Model Architectures; Program Number: WM202505), We sincerely thank Sibo Fang for his generous help during this project.
References
[1] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: A vit backbone for diffusson models. In IEEE/CVF Conference Computer Vision and Pattern Recognition, CVPR 2023, Vancouver, BC, Canada, June 17-24, 2023, pages 22669 22679. IEEE, 2023.
[2] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In 2021 IEEF/CVF International Conference Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pages 9630-
9640. IEEF, 2021.
[3] Junsong Chen, Chongjian Ge. Enze Xie. Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-E: Weak-to-strong training of diffusion transformer for 4k text-to-image generation. CoRR, abs//403.04660 2024.
[4] Junsong Chen, Jincheng Yu, Chongjian Ge. Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, Zhenguo Li. Pixart-o: Fast training of diffusion transformer photorealistic text-to-image synthesis, 2023.
[5] Pengtao Chen, Mingzhu Shen, Peng Ye, Jianjian Cao, Chongjun Tu, Christos-Savvas Bouganis, Yiren Zhao, and Tao Chen. 5-dit: A training-free acceleration method tailored for diffusion transformers, 2024.
Xiangxiang Chu, Jianlin Su, Bo Zhang, and Chunhua Shen. Visionllama: A unified llama interface for vision tasks. CoRR, abs/2403.00522. 2024.
[7] Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Z. Kaplan, and Enrico Shippole. Scalable high resolutioooon pixel-space image synthesis with hourglass diffusion transformers. CoRR, abs/2401.11605, 2024.
10

---

[8] Jia Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 2009), 20-25 June 2009, Miami, Florida, USA, pages 248-255. IEEE Computer Society, 2009.
[9] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthesis. In Marc' Aurelio Ranzato, Alina Beygelzimer. Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan. editors, Advances in Neural Information Procesing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6- 14, 2021, virtual, pages $780-8794, 2021.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. CoRR, abs/2010.1 19929, 2020.
[11] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer. Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first Interna- tional Conference on Machine Learning, 2024.
[12] Gongfan Fang, Xinyin Ma, and Xinchao Wang. Structural pruning for diffusion models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[13] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Mdtv2: Masked diffusion transformer is strong image synthesizer, 2024.
[14] Zhiwei Hao, Jianyuan Guo, Ding Jía, Kai Han, Yehui Tang. Chao Zhang, Han Hu, and Yunhe Wang. Learning efficient vision transformers via fine- grainrd manifold distillation. In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave. K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
[15] Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. Diffit: Diffusion vision transformers for image generation. CoRR, abs/2312.021339 2023.
[16] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li. Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners, In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 15979-15988. IEEE, 2022.
[17] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. CoRR, abs/2006. 1239, 2020.
[18] Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, International Conference on Machine Learning, ICML 2023, 23-29 July 2023. Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 13213-13232. PMLR, 2023.
[19] Yang Jin, Zhicheng Sun, Ningyuan Li, Kun Xu, Hao Jiang, Nan Zhuang, Quzhe Huang, Yang Song, Yadong Mu. and Zhouchen Lin. Pyramidal flow matching for efficient video generative modeling. CoRR, abs/2410.05954. 2024.
[20] Diederik P. Kingma and Ruiqi Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Sacnko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[21] Tuomas Kynkäääniemii Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in limited interval improves sample and distribution quality in diffusion models. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang, editors, Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 - 15, 2024, 2024.
[22] PKU-Yua Yuan Lab and Tuzhan Al etc. Open-sora-phle April 2024.
11

---

[23] Black Forest Labs. Flux, August 2024.
[24] Bingchen Liu, Ehsan Akhgari, Alexander Visheratin, Aleks Kamko, Linmiao Xu, Shivam Shrirao, Chase Lambert, Joao Souza, Suhail Doshi, and Daiqing Li. Playground v3: Improving text-to-image alignment with deep-fusion large language models, 2024.
[25] Zeyu Lu, Zidong Wang. Di Huang, Chengyue Wu, Xihui Liu, Wanli Ouyang, and Lei Bai. Fit: Flexible vision transformer for diffusion model. CoRR, abs/2402. 12376, 2024.
[26] Nanye Ma, Mark Goldstein, Michael S. Alhergo, Nicholas M. Boffi, Frie Vanden-Fijnden and Saining Xie. Sit: Exploring flow and diffusion-based based generative models with scalable interpolant transformers. CoRR, aabs/241108744 2024.
[27] Nanye Ma, Shangyuan Tong, Haolin Jia, Hcxiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xuhui Jia, and Saining Xie. Inference-time scaling for diffusion models beyond scaling denoising steps, 2025.
[28] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaacldin El-Nouby, Mido Assran, Nicolas Ballas, Wojeiech Galuba, Russell Howes, Yao Huang, Shang- Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jégou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision. Truns. Mach. Learn. Res., 2024, 2024.
[29] Byeongjun Park. Sangmin Woo, Hyojun Go, Jin- LonYoung Kim, and Changick Kim. Denoising task routing for diffusion models, 2024.
[30] William Peebles and Saining Xie. Scalable diffusion models with transformers. In JEEE/CVF International Conference on Computer Vision, ICCV 2023, Paris, France, October 1-6, 2023, pages 4172-4182, IEBE, 2023.
[31] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabricl Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, PMLR, 2021.
Virtual Event, volume 139 of Proceedings of Muchine Learning Research, pages
[32] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High-resolution image synthesis with latent diffusion models. In IEEF/CVF Conference Computer Vision and Pattern Recognition. CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 10674-10685. IEEE, 2022.
[33] Olaf Ronneberger. Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells III, and Alejandro F. Frangi. editors, Intervention MICCAI 2015 18th International Conference Munich, Germany, October 5 - 9, 2015, Proceedings, Part III, volume 9351 of Lecture Notes in Computer Science, pages 234-241. Springer, 2015.
Medical Image Computing and Computer-
[34] Chenyang Si, Ziqi Huang, Yuming Jiang, and Ziwei Liu. Freeu: Free lunch in diffusion u-net. CoRR, abs/2309. 11409,..1997,.. 2023.
[35] Jiaming Song. Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. CoRR, abs/2010.02502, 2020.
[36] Yang Song, Jascha aahl-Dickttein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[37] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063 127063, 2024.
[38] Yuchuan Tian, Jing Han, Chengcheng Wang, Yuchen Liang, Chao Xu, and Hanting Chen. Die: Rethinking conv3x3 designs in diffusion models. CoRR, 0bs/2501.00603, 2025.
[39] Yuchuan Tian, Zhijun Tu, Hanting Chen, Jie Hu, Chao Xu, and Yunhe Wang. U-dits: Downsample tokens in u-shaped diffusion transformers, 2024.
[40] Jiahao Wang. Ning Kang. Lewei Yao, Mengzhao Chen, Chengyue Wu, Songyang Zhang, Shuchen Xue, Yong Liu, Taiqiang Wu, Xihui Liu, Kaipeng Zhang, Shifeng Zhang, Wengi Shao, Zhenguo Li, and Ping Luo. Lit: Delving into simplified linear diffusion transformer for image generation, 2025.
12

---

[41] Christopher Williams, Fabian Falck. George Deligiannidis, Chris C. Holmes, Arnadd Doucet, and Saifuddin Syed. A unified framework for u-net design and analysis. In Alice Oh, Tristan Naumann, Amir Globerson. Kate Saenko, Moritz Hardt, Sergey Levine, editors, Advunces in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
[42] Enze Xie, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li. Ligeng Zhu, Yao Lu, and Song Han. Sana: Efficient high-rcsolution image synthesis with linear diffusion transformer, 2024.
[43] Yixing Xu, Chao Li, Dong Li, Xiao Sheng, Fan Jiang, Lu Tian, and Ashish Sirasao. Fdvit: Improve the hierarchical architecture of vision transformer. In Proceedings the IEEE/CVF International Conference Computer Vision (ICCV), pages 5950-5960, October 2023.
[44] Jingfeng Yao and Xinggang Wang. Reconstruction generation: Taming optimization dilemma in latent diffusion models. CoRR, abs/2501.01 01420,1122222 2025.
[45] Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers casier than you think. CoRR, abs/2410.06940, 2024.
[46] Wangbo Zhao, Yizeng Han, Jiasheng Tang, Kai Wang, Yibing Song, Gao Huang, Fan Wang, and Yang You. Dynamic diffusion transformer. arXiv preprint arXiv:2410.03456, 2024.
[47] Hongkai Zheng, Weili Nic, Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. Trans. Mach, Learn. Res., 2024, 2024.
[48] Zangwei Zheng. Xiangyu Peng. Tianji Yang, Chenhui Shen, Shenggui Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all, March 2024.
13

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics. and societal impact. Do not remove the checklist: The papers not including the checklist precede the (optional) supplemental material, The checklist does NOT count towards the page limit.
desk rejected. The checklist should follow references and
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant
information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked "[Yes] is generally preferable to "[No] ". is perfectly acceptable to answer "[No] provided a proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset used"). In general, answering "[No] or "[NA] is not phrased in binary way. acknowledge that the answer write a justification elaborate.
the checklist as of the factors in their evaluation. While
grounds for rejection. While the questions often more nuanced, so please just use your best judgment evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
supporting
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract demonstrates motivation, the proposed ideas and brief summary of experiment results.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in 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 or 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 he expected to generalize other setrings.
It is fine to include aspirational goals motivation long it is clear that these goals attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The paper has discussed the limitations of the work in the appendix due to page limits.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors encouraged create separate "Limitations" section in their paper.
14

---

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 sell-sppeiiiiaatinn. asympiotic approximations only holding locally). The authors should reflect on 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 a few datasets or with a 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 or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because 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 be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't I acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Question:
Answer: [NA]
Justification: The paper does not inlcude theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross-referencedd
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper fully discloses all the information needed to reproduce the main experimental results of the paper.
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 important, regardless whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken 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 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 the model. In general. releasing code and data often one good way to 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 a large language model), releasing of model checkpoint, or other means that appropriate to the research performed.
15

---

While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new 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 a way to access this model for reproducing the results a way to reproduce the model (e.g.. with open-source datasct instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of booeddsouuce 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 rescarchers 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The paper will provide open access to the data and code during camera ready period.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines hattpa://nips.cc/pple cc/public/
for more details.
While we encourage the release of code and data, we understand that this might not be possihle, so "No" an acceptable answer. Papers cannot rejected simply for not including code, unless this is central to the contribution (e.g.. for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines hattpa://nipe http:///iipsccccppbp/e cc/public/ htteee//oodSutmiisssoooooooi for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental Setting/Detail
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer. etc.) necessary to understand the results?
Answer: [Yes]
Justification: This paper has specified all the training and test details.
Guidelincs:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of 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 or other appropriate informa- tion about the statistical significance of the experiments?
Answer: [NA]
Justification: This not relevant this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

The authors should answer "Yes" if the results are accompanied by enoor bars, confidence intervals, statistical significance tests, least for the experiments that support the main claims the paper.
factors variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions madc should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is report I-sigma error bars, one should it. The authors should preferably report a 2-sigma error than state they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric enror bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question: For each experiment, docs the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
The paper has indicated sufficient information on the computer resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide amount compute required each the individual experimental runs as well as 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 or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer:
[Yes]
Justification:
This research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer means that authors have not reviewed NeurIPS Code of Ethics.
the authors answer No, they should explain special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., if there is a special consideration duc to laws or regulations in 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]
Discussed in the appendix due page limits.
Guidelines:
The answer NA means that there is no socictal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance). fairness considerations (e.g., deploy- ment lechnologies that could make decisions that unfairly impact specific groups), privacy considcrations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 not needed point out that generic algorithm for optimizing neural networks could enable people 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible relcase of data or models that have a high risk for misuse (e.g. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
This paper has described safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by requiring users adhere to 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. Licenses for existing assets
Question: 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?
Answer: [Yes]
Justification: The utilization of code, data and models in this paper is in accordance with the license and the terms.
Guidelines:
The answer NA mcans that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version the asset is used and, possible, include a URL..
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.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 he provided. For popular datasets, papprsswi thcode. com/ datasetattetaaaaaaaaaa has curated licenses for some datasets. Their licensing guide can help determine the license a dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
this information not available online, the authors are encouraged 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]
18

---

Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/codc/model as part of their sub- missions via structured templates. This includes details about training, license, limitations,
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create anonymized URL or include anonymized zip file.
14. Crowdsourcing 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 as 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 contribution of the paper involves human subjects, then much detail possible should be included in the main paper.
According the NeurlPS Code of Ethics, workers involved in data collection, curation, other lahor should be paid at least minimum wage country data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does 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?
Question:
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 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 the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we 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 applica- ble), such as the institution conducting the review,
Declaration LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original. or non stanstandard methods in this research? Note that if the LLM is used only for writing. editing, formatting purposes and does impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
component of the
Answer: [Yes]
Justification: This paper centered around LLM architecture.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: http://eerrrsccccco hatprripp..c//oooeeeeeeeee//222 for what should or should not be described.
19

---

Appendix
A Other Improvements
In the appendix, we address the effect of proposed "Other Improvements" (Sec 4. including "Time-aware MLP" and "Weight schedule of loss".
Time-aware MLP. Some works in the diffusion task 29 46 sensitive and useful certain subset of time during sampling. As the time-variant feature be aligned to time-invariant vision encoder features, we hold that alignment could perform better when the MLP is time-aware and extract time-invarian information out of diffusion features for alignment.
reveals that channel dimension is particularly
Inspired by the conditioning of canonical Diffusion U-Net 17 36 9 and DiT B0), we add a module to predict a pair of channel-we -winn shift& scale vector (> (t), 3 (t)). The module is in parallel to MLP and follows the design of DiT's AdaLN, which is concatenation of SiLU and Linear layer. The shift& scale vectors are imposed the output MI.P follows:
K (bh])) = (+) (h("")) (t)
Weight schedule of loss. As is prompted in REPA 45 designing weight schedule is a future direction. We try various weight schedules (i.e. make À in Eq. 2 schedules bring very limited improvements on the proposed U-REPA. Hence, we stick to original constant weight strategy of REPA.
function A(t) with respect time) but found that these
These improvements bring slight increases on the generation metrics. Hence, we do not include them in the main experiments for the simplicity of the method. The effects of proposed measures are shown in Tab. [ and
Tab. 12
ImageNet 256x256 w/ efg
Alignment Choices
FID
1St
Ordinary MLP
5.72
161.6
Time-aware MLP
5.63
163.3
Table 11: The effect of time-aware MLPs.
ImageNet 256x256, w/ cfg
Alignment Choices
FID
1St
Constant
5.72
161.6
max (1, + 0.5)
5.85
161.3
max -t 1.5) min (1, max 1,,,
5.72
161.6
5.58
164.0
Table 12: The effect of different weight schedules of loss.
B Additional Experiments
Evaluating J-T/2 We also evaluated the proposed U-Net architecture the Scalable Interpolant Trans- formers (SiT) framework without the guidance of REPA. The results shown in Tab. 13
Notably, though the amount of FID improvement brought U-REPA SiT+REPA vs. SITL+U-REPA) we hold that this comparison invalid due to following reasons:
great REPA (i.e. SiT /
As generation performance gets stronger, is also becoming much harder to improve (especially for FID when it gets lower).
2. Aligning J and VIT is much harder than aligning SiT and VIT, because the backbone of SiT and ViT encoders very similar. Aligning SiT, VIT encoder is special case due great architecture difference.
However, we hold that comparing REPA and U-REPA on the same model of SiT/ is fair. The default REPA achieves FID 9.35 (as shown in Tab, 4 while our method achieves FID 5.72, both trained for 100K iterations with cfg and guidance interval adopted.
20

---

Figure 6: Samples generated by SITU+U-REPA at 1M iterations. The samples are generated following the setting of REPA, at efg = Best viewed on screen.
ImageNet 256x w/o cfg
Model
Iter.
FID
SiT-XL/2
400K
17.2
SiT
400K
9.2
SiT-XL/2+REPA
400K
7.9
SIT|-XI//+U-REEPA
400K
5.4
Table 13: Evaluating the performance of SiT SiTI-XII/2 performs much better than SiT-XI//2,
and U-REPA further reduces the FID of --T to 5,4 without classifier-free guidance.
Seed Sensitivity. In our paper, we take seed following the setting of REPA. We also tested other seeds (seed = 1, 2) in training to examine the seed sensitivity of our method, shown in Tab. 14 The experiments are for 600K iterations with guidance interval and cfg, following REPA.
ImageNet 256x256, w/ cfg
Model
seed
FID
-TTI-XL/2+U-REPA
1.618
SIT -XT--XL/+U--RPA
1.599
SIT STT--XL/+UU-REAA
1.588
SST/-XL/2UU-REPA
mean
1.602+0.012
Table 14: Examining seed sensitivity. We selected seed = 0, 1, and evaluate the performance with cfg. The performance fluctuation is limited to narrow interval (approximately 0.01).
Ablations on the REPA Loss. We also conduct ablations on the REPA loss (Crepa) while leaving the
proportion of manifold loss intact (keeping the multiplication Aw fixed). The results shown in Tab 15
21

---

ImageNet 256x256, w/ cfg
0.5
0.25
0 (No CREPA)
FID 5.72
6.42
10.91
Table 15: Adjusting the hyperparameter for REPA loss in Eq. 5 The REPA loss vital for the generation performance; removing CREPA would cause significant performance decay.
Figure 7: Comparing the visual quality of SIT+REPA (upper row) and SITU+U-REPA (lower row). The samples are generated following the sampling strategy that yields the State-of-the-Art FIDs in respective methods. Best viewed on screen.
One-on-one visualization comparison. Apart from quantitative comparisons, also provide qualitative comparisons in Fig. 7 by inserting the same madom noise into trained SiT+REPA (at 4M iterations, FID 1.42) and SiTL+U-REPA 2M iterations, 1.41). samples are not cherrypicked; we directly pick the first several samples at seed=0. Samples generated by SiT|+U-REPS has better visual quality.
C Limitations & Impact
Limitations and Future work. The U-Net architecture a simple one with only one intermediate stage. do not further refine the architecture as we want to show U-Net architectures as simple as SiT, could also achieve rapid convergence. Further improvements on the U-Net architecture includes efficient attention 42 39) non-integer down& up scaling factors [43). and could be applied downstream diffusion tasks that rely heavily on U-Nets (c.g. Low-Level Vision) remains investigated.
of convolutions [6 38 Besides, whether U-REPA
Broader Impact. As work centered around AIGC, is prohable that inappropriate contents may appear from the output. We should be aware of this negative societal impact.
22


Reviews:
[{"review_id": "bNZviY0cCt", "reviewer": "Reviewer_BL2s", "summary": "This 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.", "questions": "1. Please answer the weaknesses.\n2. 1. Could the authors provide the SiT$\\downarrow$-XL/2 results in Table 4?", "limitations": "The 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).", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "After reading the authors' rebuttal, my concerns are well addressed. Therefore, I would like to raise my final rating to \"Accept\".", "rebuttal": "We sincerely thank reviewer BL2s for their constructive comments. Our responses are as follows:\n\n## Q1: Important but Missing Comparison.\nWe hope to clarify that the value 11.02 FID is in fact for DiT$\\downarrow$-XL, the performance **our proposed U-Net architecture (instead of SiT-XL that is based on a conventional DiT architecture)**. Below we provide a table illustrating the comparison.\n\n| Model | FID @ 400K |\n| ------------------ | ---- |\n| DiT-XL             | 19.47 |\n| SiT-XL+REPA        | 7.9 |\n| DiT$\\downarrow$-XL | 11.02 |\n| SiT$\\downarrow$-XL+U-REPA (Ours) | 5.4 |\n\nFrom the table, it is obvious that both REPA and U-REPA could boost the generation performance.\n\n## Q2: More analysis of incorporating the manifold alignment loss.\n\nYes definitely. The manifold alignment loss is imposing very small amount of change on tokenwise cosine similarities (unfortunately, we cannot put a plot here to illustrate). Whether adding manifold loss or not, the tokenwise cosine similarity will always fluctuate around 0.6 (which we believe is the upper bottleneck for the alignment between U-Net and ViT encoders). In spite of this, manifold loss is adding regularization on the manifold aspect and thus making U-Net more performant.\n\n## Q3: Higher-Resolution (ImageNet 512x512) Experiments.\n\nThanks for your suggestion! We are happy to provide the performance of U-REPA on ImageNet 512x512 in the table below. U-REPA outperforms REPA by considerable margins, showing that the proposed U-REPA could also be generalized to larger images. Due to time and resource constraints (running ImageNet-512 is very resource intensive), we could reach 400K iterations at this moment.\n\n\n| FID$\\downarrow$  / IS $\\uparrow$ | SiT$\\downarrow$-XL/2+REPA | SiT$\\downarrow$-XL/2+U-REPA |\n| -------------------------------- | ------------------------- | --------------------------- |\n| 400K                             | 2.44 / 247.3              | 2.21 / 274.7                |\n\n\n## Q4: Font size is too small.\n\nWe are terribly sorry for the overly-small fontsize in the paper. We will enlarge it in the next revision for better readability."}, {"review_id": "pjTnsZBsdy", "reviewer": "Reviewer_i6qu", "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.", "questions": "The author is recommended to optimize the layout to make it better.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "- Strengths:\n1. The proposed method is well-motivated and the analysis is sufficient.\n2. Extensive experiments demonstrate that the proposed U-REPA achieves great performance.\n3. The analysis and ablation experiments are sufficient and reasonable.\n\n- Weaknesses:\n1. This work mainly involves some engineering analysis and engineering practice. The methods used, such as manifold loss, have also been proposed. \n2. 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.", "was_revised": true, "final_justification": "The rebuttal has addressed part of my concerns. I am inclined to accept this paper.", "rebuttal": "We sincerely thank reviewer i6qu for their constructive comments. Our responses are as follows:\n\n## Q1&2: Clarifying our core contributions.\n\nOur core contributions are as follows:\n\n1. While DiT is becoming the mainstream architecture and U-Net is gradually being phased out, prior work has shown that U-Net converges quickly, which is an exploitable advantage. However, its upper bound in terms of generation quality and convergence speed remains unexplored (previous U-Nets report FID>2, which is still a large gap from SOTA). We aim to investigate the full potential of U-Net in these aspects (Finally achieving 1.41 FID at merely 1M iterations).\n\n2. A plain use of REPA does not work well on U-Net due to architecture mismatch. We adapt REPA to U-Nets by making key modifications, including position, upsampling, and manifold loss.\n\nIn a nutshell, beyond proposing a method, this paper shows that U-Net is perhaps **a better choice of diffusion architecture** both in terms of convergence and generation quality, paving the way for future research.\n\n\n\n## Q3: Optimizing the layout is recommended.\n\nWe are sorry that our layout is too cramped due to space limits. We will move some less important part to the appendix part for better clarity and readability."}, {"review_id": "qHT9RYXE9n", "reviewer": "Reviewer_kqTd", "summary": "This 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.", "questions": "My 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.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I have read through all the reviews and author responses. I have discussed with the author during the phase and reached a final conclusion based on the rebuttal and discussion. The authors provide comprehensive explanations that have addressed all my main concerns. The new insight and supplement reinforce the quality and contribution of the paper. Thus, I raise the final rating from borderline reject to accept, which stands that the paper is ready for publication in the conference.", "rebuttal": "We sincerely thank reviewer kqTd for their constructive comments. Our responses are as follows:\n\n## Q1&2: Regarding the convolutional U-Net.\n\nUnfortunately, we hold that the original conventional U-Net is rather limited in capacity and are not suitable for scaling up. We conducted experiments on SongUNet and figure out that it is not as performant as SiT$\\downarrow$:\n\n\n|                        | 100K  | 200K  | 300K  | 400K  |\n| ---------------------- | ----- | ----- | ----- | ----- |\n| SiT+REPA\\* (DiT architecture) | 19.7 | 11.9 | 9.4  | 9.0  |\n| SiT$\\downarrow$+U-REPA | 12.8 |  6.6  |  6.2   | 5.4  |\n| SongUNet+U-REPA        | 29.0 | 18.8 | 14.9 | 12.9 |\n\n\\* Replicated according to the official codebase and setting. In the official paper, SiT+REPA @ 400K is reported as 7.9 FID.\n\nWe hold that the **capacity and property of architecture**, rather than **the use of convolution only**, is the key to fast convergence. SongUNet is a fused architecture of convolution and self-attention, but it fails to be converging faster than full transformer diffusion models.\n\n[1] Song, Yang, et al. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021.\n\n\n\n\n## Q3: Writing should be improved.\n\nWe apologize for the writing issues. The sentence in L208 should be \"The similarity gap between SiT and SiT is obvious\". Table 1 should be mentioned in the \"Toy experiments on U-Net components\" paragraph. $h_t^{[n]}$ refers to the n$^{th}$ token in an image. We will revise the paper in the next revision.\n\n## Q4: Multiple seed for robustness.\n\nThanks for your advice. To clarify, we didn't cherrypick seeds; we just used the identical training setting and codebase of REPA that adopted seed=0. We also add experiments when seed is 1 and 2 to verify the robustness of our method.\n\nWe extend the training to 600K iterations (due to limited time and resources) and reach the results as follows:\n\n| SiT$\\downarrow$-XL + U-REPA | seed=0 (Ours) | seed=1 | seed=2 | mean              |\n| --------------------------- | ------------- | ------ | ------ | ----------------- |\n| FID$\\downarrow$             | 1.618         | 1.599  | 1.588  | 1.602 $\\pm$ 0.012 |\n\nThe experiments indicate that fluctuation is 0.01 at merely 600K iterations, which is tiny compared to the values.\n\n\n\n## Q5: Time-aware MLP may not be compatible.\n\nThanks for your suggestions. We agree that these part is a trick and minor. We will remove this part in the next revision. We are not using these tricks for the main experiments (as commented with the ablations).\n\n## Q6: Discussion between VA-VAE loss and Manifold Loss\n\nThanks for your suggestions. The key difference between the Marginal Distance Matrix loss in VA-VAE and the manifold loss are as follows:\n\n1. Purpose Difference: in VA-VAE, loss is applied to VAE for better image compression; our manifold loss is applied to the model while the VAE is kept fixed based on the off-the-shelf sd-vae.\n\n2. Form Difference: VA-VAE is a l1-norm-based loss while our loss is based on l2-norm. We also tried l1 as an alternative but found it harmful our setting, slightly degrading from FID 6.24 to 6.28."}, {"review_id": "fNEKwh8McS", "reviewer": "Reviewer_U3kJ", "summary": "This 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", "questions": "1. 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.", "limitations": "Although 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.", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 1, "originality": 2, "strengths_and_weaknesses": "Strengths:\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?", "was_revised": true, "final_justification": "Although the author has addressed some doubts, given the novelty and contribution of this article, I will revise my score to Bordline Reject.", "rebuttal": "We sincerely thank reviewer U3KJ for their constructive comments. Our responses are as follows:\n\n## Q1: Contribution & Novelty.\n\nThe contribution of this paper goes beyond an engineering practice: it rethinks U-Net, which is a long-neglected architecture. Its most important contribution is that it demonstrates U-Nets perform no worse than the DiT-class.\n\n- To the best of our knowledge, we are the first to adapt REPA to U-Nets by making key modifications, including position, upsampling, and manifold loss.\n- There are some previous works that rethink U-Net (as cited in \"Related Work\"), but they fail to demonstrate U-Nets' best performance (in these work, usually >2 FID on ImageNet256). In contrast, we show U-Nets' **extreme best performance** is better than DiT at only half of the training cost. Our work reveals the great potential of U-Net, and hints a better architectural choice in diffusion.\n\n## Q2: Ablation on the REPA loss.\n\nThanks for your suggestions. To clarify, we don't reject the importance REPA loss (we will emphasize this point in the next revision); rather, we hold that REPA loss is reaching a bottleneck in the U-Net alignment setting and requires manifold guidance for assistance.\n\nIn our setting, we follow the hyperparameter setting from REPA, which sets $\\lambda=0.5$. We also enclose an ablation experiment table that reveals the importance of REPA loss while $\\lambda$ is reduced to 0 and manifold loss is kept intact. Results indicate that REPA loss is important for generation quality:\n\n| $\\lambda$        | 0.5  | 0.25 | 0 (No $\\mathcal{L}_{REPA}$) |\n| ---------------- | ---- | ---- | --------------------------- |\n| FID $\\downarrow$ | 5.72 | 6.42 | 10.91                       |\n\n## Q3: Lack of Visualization Comparison.\n\nWe have already put a visualization of our method in the supplementary materials due to space limits. Since we cannot attach a picture for comparison in the rebuttal, instead we conduct a human preference experiment that performs one-on-one comparison between REPA and U-REPA under the same sampling seed ($seed=0$) as follows.\n\nOur experiment consists of the first 100 images and no cherrypick is conducted. Due to content difference or similar quality, we find it hard to distinguish some image pairs. Results show the image quality of U-REPA is better than REPA, which coincides with the experiment results.\n\n| Which method is better? | REPA | U-REPA (Ours) | Can't Tell |\n| ----------------------- | ---- | ------------- | ---------- |\n| Counts                  | 17   | 38            | 45         |\n\n\n## Q4: Compare U-REPA and REPA under the same experiment setting.\n\nTo clarify, the whole process is started from scratch (according to the recipe of REPA);  we are not loading a trained U-Net for U-REPA (which might cause unfairness).\n\nWhile we have demonstrated that U-Net is more performant, this evaluation is only at initial phase (only @ 400K); we are not certain whether U-Net could achieve close SOTA performance. In our paper, we show that U-Net could also reach the comparable SOTA performance of DiT at a faster convergence speed under exactly the same, fair training recipe."}]

Author rebuttal:

Meta-review: Final rating: 5:Accept/ 4: Borderline Accept/ 5:Accept /3:Borderline Reject. This paper introduce U-REPA, extending REPA to diffusion U-Nets by addressing block heterogeneity, downsampling mismatches, and ViT–U-Net token gaps. It aligns at the U-Net mid-stage, upsamples MLP-projected features, and replaces token-wise similarity with a manifold loss. Experiments show faster convergence and FID < 1.5 on ImageNet 256×256 in 200 epochs (1M iters), outperforming REPA with ~½ the training under sd-vae-ft-ema.

Reviewers acknowledge the strong experiments and ablations but raise concerns about the missing comparison between SiT↓ and SiT↓+U-REPA, absent results with a convolutional U-Net, and other missing details and discussion. After the rebuttal, the concerns of reviewers BL2s, i6qu, and kqTd were largely addressed and they now lean toward acceptance; however, reviewer U3kJ still questions the fairness of key experiments demonstrating U-REPA’s contribution, notes missing experiments on SiT↓+REPA and SiT↓+U-REPA, points out the lack of direct visual comparisons between REPA and U-REPA, and requests additional runs on DDPM, iDDPM, or EDM.

The ACs conclude that fairness is ensured by comparing methods on the same U-Net architecture; they also note that the paper already contains related experiments for SiT↓+REPA and SiT↓+U-REPA, judge the absence of visual side-by-sides as non-critical given strong quantitative results (though encouraged for completeness), and consider that the DDPM/iDDPM/EDM request misaligned and  does not make sense given that those works are actually about diffusion sampling operation but reviewer U3kJ's concern is about vision backbone architectures.

Overall, the ACs find the rebuttal satisfactory, view the paper as a strong contribution to feature alignment in U-Nets—highlighting their potential and offering an alternative to the prevailing DiT paradigm—and recommend acceptance.
Decision: Accept (poster)