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

We deeply appreciate your insightful comments and efforts in reviewing our manuscript. We respond to each of your comments one-by-one in what follows. In the revised draft, we mark our major revisions as “blue”.

[W1-1] Higher resolution experiments (e.g., 512x512)

Thanks for the suggestion! We agree that high-resolution results will make the paper more convincing. To address this concern, we additionally train SiT-XL/2 on ImageNet-512x512 with REPA and report the quantitative evaluation in the table below. Notably, as shown in this table, the model (with REPA) already outperforms the vanilla SiT-XL/2 in terms of four metrics (FID, sFID, IS, and Rec) using >3x fewer training iterations. We added this result in Appendix J in the revised manuscript with qualitative results. We will also add the final results in the final draft with longer training (please understand that it takes more than 2 weeks, so it will not be possible to update them until the end of the discussion period).

\begin{array}{lcccccc} \hline \phantom{0}\phantom{0} \text{Model} & \text{Epochs} & \text{FID $\downarrow$} & \text{sFID $\downarrow$} & \text{IS $\uparrow$} & \text{Pre. $\uparrow$} & \text{Rec. $\uparrow$} \newline \hline \text{Pixel diffusion} \newline \phantom{0}\phantom{0} \text{VDM++} & - & 2.65 & - & 278.1 & - & - \newline \phantom{0}\phantom{0} \text{ADM-G, ADM-U} & 400 & 2.85 & 5.86 & 221.7 & 0.84 & 0.53 \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-Net)} & 800 & 4.28 & - & 171.0 & - & - \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-ViT, L)} & 800 & 4.53 & - & 205.3 & - & - \newline \hline \text{Latent diffusion, Transformer} \newline \phantom{0}\phantom{0} \text{MaskDiT} & 800 & 2.50 & 5.10 & 256.3 & 0.83 & 0.56 \newline \phantom{0}\phantom{0} \text{DiT-XL/2} & 600 & 3.04 & 5.02 & 240.8 & 0.84 & 0.54 \newline \phantom{0}\phantom{0} \text{SiT-XL/2} & 600 & 2.62 & 4.18 & 252.2 & 0.84 & 0.57 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & \phantom{0}80 & {2.44} & 4.21 & 247.3 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 100 & 2.32 & \bf{4.16} & 255.7 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 200 & \bf{2.08} & 4.19 & \bf{274.6} & 0.83 & \bf{0.58} \newline \hline \end{array}

[W1-2] Text-to-image experiments.

To address your concern, we have additionally conducted a text-to-image (T2I) generation experiment on MS-COCO [1], following the setup in previous literature [2]. We tested REPA on MMDiT [3], a variant of DiT designed for T2I generation, given that the standard DiT/SiT architectures are not suited for this task. As shown in the table below, REPA also shows clear improvements in T2I generation, highlighting the importance of alignment of visual representations even under the presence of text representations. In this respect, we believe that scaling up REPA training to large-scale T2I models will be a promising direction for future research. We added these results with qualitative comparison in Appendix K in the revised manuscript.

\begin{array}{lcc} \hline \text{Method} & \text{Type} & \text{FID} \newline \hline \text{AttnGAN} & \text{GAN} & 35.49 \newline \text{DM-GAN} & \text{GAN} & 32.64 \newline \text{VQ-Diffusion} & \text{Discrete Diffusion} & 19.75 \newline \text{DF-GAN} & \text{GAN} & 19.32 \newline \text{XMC-GAN} & \text{GAN} & \phantom{0}9.33 \newline \text{Frido} & \text{Diffusion} & \phantom{0}8.97 \newline \text{LAFITE} & \text{GAN} & \phantom{0}8.12 \newline \hline \text{U-Net} & \text{Diffusion} & \phantom{0}7.32 \newline \text{U-ViT-S/2} & \text{Diffusion} & \phantom{0}5.95 \newline \text{U-ViT-S/2 (Deep)} & \text{Diffusion} & \phantom{0}5.48 \newline \hline \text{MMDiT (ODE; NFE=50, 150K iter)} & \text{Diffusion} & \phantom{0}6.05 \newline MMDiT+REPA (ODE; NFE=50, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.73} \newline \hline \text{MMDiT (SDE; NFE=250, 150K iter)} & \text{Diffusion} & \phantom{0}5.30 \newline MMDiT+REPA (SDE; NFE=250, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.14} \newline \hline \end{array}

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

[W2] Comprehensive analysis: different metric trends for ImageNet generation.

We mainly use FID (and IS) metrics as criteria for determining design choices, as these two metrics are the most popular, widely-used metrics, and they also showed quite clear trends across different design choices. Moreover, following your suggestion, we provide additional quantitative evaluation results for NT-Xent and cos. sim at the early training stage in the table below (due to the limited time we have, we used smaller NFE=50 here, different from NFE=250 in the manuscript). The results show that NT-Xent is more beneficial at early iterations, but the gap diminishes with longer training.

\begin{array}{lcccc} \hline \text{Method} & \text{50K} & \text{100K} & \text{200K} & \text{400K} \newline \hline \text{NT-Xent} & 26.29 & 18.11 & 15.61 & 14.38 \newline \text{Cos. sim.} & 29.21 & 18.42 & 15.43 & 14.17 \newline \hline \end{array}

[Q1] Can REPA achieve state-of-the-art results (outperforming FID=1.79 of MDT)?

First, in Appendix G, we already have shown that REPA can achieve the same FID of MDT (1.79) by using a slightly smaller CFG scale $w=1.325$. Moreover, note that MDT uses CFG coefficient scheduling that can dramatically improve the performance [4, 5], where we used a constant coefficient and did not apply this technique to achieve the same FID.

We also conducted an additional evaluation by applying CFG scheduling similar to MDT. If we apply a simple scheduling technique proposed in recent work [5], REPA achieves significantly improved FID=1.42. We added this result in Table 4 and Appendix G in the revised manuscript.

[4] Masked Diffusion Transformer is a Strong Image Synthesizer, ICCV 2023.
[5] Applying Guidance in a Limited Interval Improves Sample and Distribution Quality in Diffusion Models, to appear at NeurIPS 2024

[Q2] Why different CFG scale is used for quantitative and qualitative comparison?

This is a quite common practice in most diffusion model literature [6, 7, 8], and we simply followed their setup. Specifically, generating samples with a large CFG coefficient can provide high-quality individual samples, but they lack diversity [9], potentially leading to quantitative evaluation degradation.

[6] Scalable Diffusion Models with Transformers, ICCV 2023
[7] SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers, ECCV 2024
[8] High-Resolution Image Synthesis with Latent Diffusion Models, CVPR 2022
[9] Classifier-free Diffusion Guidance, arXiv 2022

[Q3] Why do minimum semantic differences appear at later layers without REPA?

We hypothesize that the denoising task used for diffusion model training, which is a sort of reconstruction, may not be a suitable task for learning good representations. Specifically, as mentioned in the Introduction of our original manuscript, it is not capable of eliminating unnecessary details in the image for representation learning [10, 11]. Thus, representation might be learned inefficiently and appear after unnecessarily many layers, even though the first few layers of the model have enough capacity to learn meaningful representations (note that we showed it through applying REPA at the 8th layer). Exploring such different behaviors will be an exciting future direction, and we leave it for future work.

[10] A Path towards Autonomous Machine Intelligence Version, OpenReview 2022
[11] Self-supervised learning from images with a joint-embedding predictive architecture, CVPR 2023

[Q4] Notation.

Thanks for the suggestion! We reflected them (notation and fixing typos) in the revised manuscript.

We deeply appreciate your insightful comments and efforts in reviewing our manuscript. We respond to each of your comments one-by-one in what follows. In the revised draft, we mark our major revisions as “blue”.

[W1] Higher resolution experiments (e.g., 512x512).

Thanks for the suggestion! We agree that high-resolution results will make the paper more convincing. To address this concern, we additionally train SiT-XL/2 on ImageNet-512x512 with REPA and report the quantitative evaluation in the table below. Notably, as shown in this table, the model (with REPA) already outperforms the vanilla SiT-XL/2 in terms of four metrics (FID, sFID, IS, and Rec) using >3x fewer training iterations. We added this result in Appendix J in the revised manuscript with qualitative results. We will also add the final results in the final draft with longer training (please understand that it takes more than 2 weeks, so it will not be possible to update them until the end of the discussion period).

\begin{array}{lcccccc} \hline \phantom{0}\phantom{0} \text{Model} & \text{Epochs} & \text{FID $\downarrow$} & \text{sFID $\downarrow$} & \text{IS $\uparrow$} & \text{Pre. $\uparrow$} & \text{Rec. $\uparrow$} \newline \hline \text{Pixel diffusion} \newline \phantom{0}\phantom{0} \text{VDM++} & - & 2.65 & - & 278.1 & - & - \newline \phantom{0}\phantom{0} \text{ADM-G, ADM-U} & 400 & 2.85 & 5.86 & 221.7 & 0.84 & 0.53 \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-Net)} & 800 & 4.28 & - & 171.0 & - & - \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-ViT, L)} & 800 & 4.53 & - & 205.3 & - & - \newline \hline \text{Latent diffusion, Transformer} \newline \phantom{0}\phantom{0} \text{MaskDiT} & 800 & 2.50 & 5.10 & 256.3 & 0.83 & 0.56 \newline \phantom{0}\phantom{0} \text{DiT-XL/2} & 600 & 3.04 & 5.02 & 240.8 & 0.84 & 0.54 \newline \phantom{0}\phantom{0} \text{SiT-XL/2} & 600 & 2.62 & 4.18 & 252.2 & 0.84 & 0.57 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & \phantom{0}80 & {2.44} & 4.21 & 247.3 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 100 & 2.32 & \bf{4.16} & 255.7 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 200 & \bf{2.08} & 4.19 & \bf{274.6} & 0.83 & \bf{0.58} \newline \hline \end{array}

[W2] Missing dataset analysis on the pretrained visual encoders.

Thanks for pointing this out. In the table below, we additionally provide the dataset analysis for visual encoders (we also added I-JEPA [1] and SigLIP [2] in Table 2 in the revision for a more extensive analysis). As shown in this table, better visual representations learned from massive amounts of image data provide more improvement, regardless of whether the dataset does not include ImageNet. We also added this discussion in Appendix C in the revised manuscript.

\begin{array}{llccc} \hline \text{Method} & \text{Dataset} & \text{w/ ImageNet-1K} & \text{Text sup.} & \text{FID$\downarrow$} \newline \hline \text{MAE} & \text{ImageNet-1K} & \text{O} & \text{X} & 12.5 \newline \text{DINO} & \text{ImageNet-1K} & \text{O} & \text{X} & 11.9 \newline \text{MoCov3} & \text{ImageNet-1K} & \text{O} & \text{X} & 11.9 \newline \text{I-JEPA} & \text{ImageNet-1K} & \text{O} & \text{X} & 11.6 \newline \hline \text{CLIP} & \text{WIT-400M} & \text{X} & \text{O} & 11.0 \newline \text{SigLIP} & \text{WebLi-4B} & \text{X} & \text{O} & 10.2 \newline \text{DINOv2} & \text{LVD-142M} & \text{O} & \text{X} & 10.0 \newline \hline \end{array}

[Q1] The pretrained vision encoder usually also finetunes on the ImageNet dataset; how do you separate the dataset leakage effect when measuring generative models trained on the ImageNet?

First, we did not fine-tune encoders (e.g., SigLIP and CLIP) with the ImageNet dataset, particularly if they were trained with another dataset, thereby separating the dataset leakage effect if we use these encoders. As shown in [W2], REPA also achieves significant improvements with these encoders, which validates that the improvement does not simply come from data leakages.

To further address your concern, we conduct an additional text-to-image (T2I) generation experiment on MS-COCO by using DINOv2 as the target representation. Since the training dataset of DINOv2 does not include MS-COCO, it can validate whether the gain comes from data leakage or from aligning diffusion representation with a good visual representation. As shown in the table below, REPA with DINOv2 also shows improvement in this setup, which further highlights our claim.

\begin{array}{lcc} \hline \text{Method} & \text{Type} & \text{FID} \newline \hline \text{AttnGAN} & \text{GAN} & 35.49 \newline \text{DM-GAN} & \text{GAN} & 32.64 \newline \text{VQ-Diffusion} & \text{Discrete Diffusion} & 19.75 \newline \text{DF-GAN} & \text{GAN} & 19.32 \newline \text{XMC-GAN} & \text{GAN} & \phantom{0}9.33 \newline \text{Frido} & \text{Diffusion} & \phantom{0}8.97 \newline \text{LAFITE} & \text{GAN} & \phantom{0}8.12 \newline \hline \text{U-Net} & \text{Diffusion} & \phantom{0}7.32 \newline \text{U-ViT-S/2} & \text{Diffusion} & \phantom{0}5.95 \newline \text{U-ViT-S/2 (Deep)} & \text{Diffusion} & \phantom{0}5.48 \newline \hline \text{MMDiT (ODE; NFE=50, 150K iter)} & \text{Diffusion} & \phantom{0}6.05 \newline MMDiT+REPA (ODE; NFE=50, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.73} \newline \hline \text{MMDiT (SDE; NFE=250, 150K iter)} & \text{Diffusion} & \phantom{0}5.30 \newline MMDiT+REPA (SDE; NFE=250, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.14} \newline \hline \end{array}

[Q2] Apply REPA to video generation model?

Similar to the success in the image domain, we strongly believe that REPA can be applied to video diffusion models by exploiting strong video self-supervised representations, such as V-JEPA [1]. Exploring this should be an interesting future direction.

[1] Revisiting Feature Prediction for Learning Visual Representations from Video, arXiv 2024.

Dear Reviewer i9uk,

Thank you again for your time and efforts in reviewing our paper.

As the discussion period draws close, we kindly remind you that two days remain for further comments or questions. We would appreciate the opportunity to address any additional concerns you may have before the discussion phase ends.

Thank you very much!

Many thanks,
Authors

We deeply appreciate your insightful comments and efforts in reviewing our manuscript. We respond to each of your comments one-by-one in what follows. In the revised draft, we mark our major revisions as “blue”.

[W1] Higher resolution experiments (e.g., 512x512).

Thanks for the suggestion! We agree that high-resolution results will make the paper more convincing. To address this concern, we additionally train SiT-XL/2 on ImageNet-512x512 with REPA and report the quantitative evaluation in the table below. Notably, as shown in this table, the model (with REPA) already outperforms the vanilla SiT-XL/2 in terms of four metrics (FID, sFID, IS, and Rec) using >3x fewer training iterations. We added this result in Appendix J in the revised manuscript with qualitative results. We will also add the final results in the final draft with longer training (please understand that it takes more than 2 weeks, so it will not be possible to update them until the end of the discussion period).

\begin{array}{lcccccc} \hline \phantom{0}\phantom{0} \text{Model} & \text{Epochs} & \text{FID $\downarrow$} & \text{sFID $\downarrow$} & \text{IS $\uparrow$} & \text{Pre. $\uparrow$} & \text{Rec. $\uparrow$} \newline \hline \text{Pixel diffusion} \newline \phantom{0}\phantom{0} \text{VDM++} & - & 2.65 & - & 278.1 & - & - \newline \phantom{0}\phantom{0} \text{ADM-G, ADM-U} & 400 & 2.85 & 5.86 & 221.7 & 0.84 & 0.53 \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-Net)} & 800 & 4.28 & - & 171.0 & - & - \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-ViT, L)} & 800 & 4.53 & - & 205.3 & - & - \newline \hline \text{Latent diffusion, Transformer} \newline \phantom{0}\phantom{0} \text{MaskDiT} & 800 & 2.50 & 5.10 & 256.3 & 0.83 & 0.56 \newline \phantom{0}\phantom{0} \text{DiT-XL/2} & 600 & 3.04 & 5.02 & 240.8 & 0.84 & 0.54 \newline \phantom{0}\phantom{0} \text{SiT-XL/2} & 600 & 2.62 & 4.18 & 252.2 & 0.84 & 0.57 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & \phantom{0}80 & {2.44} & 4.21 & 247.3 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 100 & 2.32 & \bf{4.16} & 255.7 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 200 & \bf{2.08} & 4.19 & \bf{274.6} & 0.83 & \bf{0.58} \newline \hline \end{array}

[Q1] Why are the experiments focused on diffusion transformers?

We mainly focused on the diffusion transformer (DiT) architecture for several reasons. First, as you mentioned, recent state-of-the-art visual encoders (e.g., DINOv2) use vision transformer architecture, so we thought enforcing the alignment of the features with DiTs makes more sense. Next, many papers have demonstrated that DiT architecture has great scalability, thereby being used in recent state-of-the-art diffusion models like Flux, Stable Diffusion 3.5, and Sora. Finally, DiT does not contain complicated components (e.g., up/downsampling blocks, skip connections) compared with U-Nets; thus, it is easy for us to perform extensive analysis of REPA under a simplified design space. In these respects, our experiments mainly focused on diffusion transformers. Applying REPA with other diffusion model architectures like U-Net would be an interesting direction in the future.

Dear Reviewer vfsY,

We are happy to hear that our rebuttal addressed your concerns well. Also, we appreciate your support for our work. If you have any further questions or suggestions, please do not hesitate to let us know. Moreover, thanks for letting us know about the issues about the ables that we added in the response. We will take a look and update to Markdown formatting if the problem persists.

Best regards,
Authors

We deeply appreciate your insightful comments and efforts in reviewing our manuscript. We respond to each of your comments one-by-one in what follows. In the revised draft, we mark our major revisions as “blue”.

[W1] Benefit of REPA with the classifier-free guidance: something like Figure 1?

First, we have already shown FID improvement with CFG: In Table 2 of our manuscript, REPA outperforms the vanilla SiT-XL/2 with CFG using 7x fewer iterations. We have also provided experimental results showing that the result can be further improved with longer training (1.96 at 1M iteration to 1.80 at 4M iteration). This has been highlighted in L482-L483 in our main manuscript.

Nevertheless, following your suggestion, we additionally provide a similar plot to Figure 1 with classifier-free guidance in Appendix G in the revised manuscript. The plot also shows a similar trend to Figure 1 in our manuscript, validating the effectiveness of REPA with CFG.

[W2] How does REPA work in text-to-image models?

To address your concern, we have additionally conducted a text-to-image (T2I) generation experiment on MS-COCO [1], following the setup in previous literature [2]. We tested REPA on MMDiT [3], a variant of DiT designed for T2I generation, given that the standard DiT/SiT architectures are not suited for this task. As shown in the table below, REPA also shows clear improvements in T2I generation, highlighting the importance of alignment of visual representations even under the presence of text representations. In this respect, we believe that scaling up REPA training to large-scale T2I models will be a promising direction for future research. We added these results with qualitative comparison in Appendix K in the revised manuscript.

\begin{array}{lcc} \hline \text{Method} & \text{Type} & \text{FID} \newline \hline \text{AttnGAN} & \text{GAN} & 35.49 \newline \text{DM-GAN} & \text{GAN} & 32.64 \newline \text{VQ-Diffusion} & \text{Discrete Diffusion} & 19.75 \newline \text{DF-GAN} & \text{GAN} & 19.32 \newline \text{XMC-GAN} & \text{GAN} & \phantom{0}9.33 \newline \text{Frido} & \text{Diffusion} & \phantom{0}8.97 \newline \text{LAFITE} & \text{GAN} & \phantom{0}8.12 \newline \hline \text{U-Net} & \text{Diffusion} & \phantom{0}7.32 \newline \text{U-ViT-S/2} & \text{Diffusion} & \phantom{0}5.95 \newline \text{U-ViT-S/2 (Deep)} & \text{Diffusion} & \phantom{0}5.48 \newline \hline \text{MMDiT (ODE; NFE=50, 150K iter)} & \text{Diffusion} & \phantom{0}6.05 \newline MMDiT+REPA (ODE; NFE=50, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.73} \newline \hline \text{MMDiT (SDE; NFE=250, 150K iter)} & \text{Diffusion} & \phantom{0}5.30 \newline MMDiT+REPA (SDE; NFE=250, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.14} \newline \hline \end{array}

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

[W3, Q1] In Figure 3, was the analysis performed by training different networks with RePA applied at different layers, or was the training process for a single layer and evaluation performed at different layers?

Thank you for pointing this out. The analysis was performed by using a single network trained with REPA at layer 8 and evaluated at different layers. This is why the model with REPA has the minimum representational gap at layer 8. This allows more layers after alignment to focus on capturing high-frequency details (based on a good representation) and achieves more improvement in generation performance, as explained in the “Alignment depth” paragraph in Section 4.2. As you suggested, we added clarification of Figure 3 in the caption in the revised manuscript.

[W4] No limitation or future works section.

Thanks for the suggestion. In this paper, we mainly focused on latent diffusion in the image domain. Exploring REPA with pixel-level diffusion or on other data domains like videos would be interesting future directions. Moreover, training large-scale text-to-image diffusion models with REPA will also be an interesting direction. Exploring theoretical insights into why REPA works well will also be an exciting future direction. We also added such discussions in Appendix M in the revised manuscript.

[Q2] Representation alignment with different levels of features from the foundation model?

We think it is a really interesting future direction to explore. Given that recent diffusion model literature has shown that coarse-to-fine generation can make diffusion model training more efficient, we think that exploring the alignment of representations from foundation models and diffusion models in a hierarchical manner will be an exciting direction. We think there are many other possible directions, and we leave it for future work.

[Q3] Time-varying REPA: might lead to an even better boost in performance?

Thanks for the insightful suggestion. Indeed, as diffusion models deal with noisy inputs over different noise scales, time-varying REPA could lead to further performance improvements. One possible direction can be designing a weight function based on the noise schedule used in the diffusion process. We have not explored this in this work as our main focus is more on performing extensive analysis on other perspectives, such as target representations used for alignment, alignment depth, scalability of the method, etc. We leave this as an exciting direction for future work.

[Q4] Performance of REPA with the state-of-the-art ViT classifier?

Following your suggestion, we conducted an additional experiment using the VIT-L classifier fine-tuned from SWAG [4], which has 88.1% accuracy on ImageNet and is better than the visual encoder that we used in our experiment. As shown in the table below, it shows less improvement than other powerful self-supervised ViTs like DINOv2. In this respect, we hypothesize that aligning with a generic good visual representation learned from massive amounts of data is more effective in diffusion transformer training compared to discriminative representation on the target dataset.

\begin{array}{lccccc} \hline \text{Method} & \text{FID $\downarrow$} & \text{sFID $\downarrow$} & \text{IS $\uparrow$} & \text{Pre. $\uparrow$} & \text{Rec. $\uparrow$} \newline \hline \text{DINOv2-L} & \phantom{0}9.9 & 5.34 & 111.9 & 0.68 & 0.65 \newline \text{ViT-L Classifier} & 11.4 & 5.24 & 100.3 & 0.68 & 0.64 \newline \hline \end{array}

[4] Revisiting Weakly Supervised Pre-Training of Visual Perception Models, CVPR 2022

I thank the authors for the detailed response. I have some doubts regarding some of the responses.

[W2] Was the MS-COCO finetuning performed by finetuning an existing model or was the training performed from scratch. In case it was trained with an existing models. From my understanding one of the main benefits RePA provides is faster convergence. Would this mean that finetuning with RePA may also lead to better performance?

[W2] Was the representation alignment performed using DinoV2 embeddings or CLIP embeddings. Would there be an added advantage in performing RePA with CLIP embeddings for T2I models

[W3] I thank the authors for this clarification. Seeing this response and after seeing the PCA visualization in Figure 38. Would performing RePA on later transformer layer embeddings (layers 16-24) than layer 8 be more suited. If Yes, I would appreciate if the authors may include this aspect in the future works section in the manuscript.

[Q4] Dear authors, I thank for the experiment with the ViT classifier fine-tuned on SWAG[4]. I think the authors may have misread my question in the questions section. My concern was whether utilizing a classifier trained just for ViT classification with labels may perform better when compared to a network trained in a self-supervised fashion. From my understanding, In the experiment provided by the authors, another self supervised network was utilized.

We deeply appreciate your insightful comments and efforts in reviewing our manuscript. We respond to each of your comments one-by-one in what follows. In the revised draft, we mark our major revisions as “blue”.

[W1-1] Focus is class-conditioned diffusion models rather than text-conditioned diffusion models.

To address your concern, we have additionally conducted a text-to-image (T2I) generation experiment on MS-COCO [1], following the setup in previous literature [2]. We tested REPA on MMDiT [3], a variant of DiT designed for T2I generation, given that the standard DiT/SiT architectures are not suited for this task. As shown in the table below, REPA also shows clear improvements in T2I generation, highlighting the importance of alignment of visual representations even under the presence of text representations. In this respect, we believe that scaling up REPA training to large-scale T2I models will be a promising direction for future research. We added these results with qualitative comparison in Appendix K in the revised manuscript.

\begin{array}{lcc} \hline \text{Method} & \text{Type} & \text{FID} \newline \hline \text{AttnGAN} & \text{GAN} & 35.49 \newline \text{DM-GAN} & \text{GAN} & 32.64 \newline \text{VQ-Diffusion} & \text{Discrete Diffusion} & 19.75 \newline \text{DF-GAN} & \text{GAN} & 19.32 \newline \text{XMC-GAN} & \text{GAN} & \phantom{0}9.33 \newline \text{Frido} & \text{Diffusion} & \phantom{0}8.97 \newline \text{LAFITE} & \text{GAN} & \phantom{0}8.12 \newline \hline \text{U-Net} & \text{Diffusion} & \phantom{0}7.32 \newline \text{U-ViT-S/2} & \text{Diffusion} & \phantom{0}5.95 \newline \text{U-ViT-S/2 (Deep)} & \text{Diffusion} & \phantom{0}5.48 \newline \hline \text{MMDiT (ODE; NFE=50, 150K iter)} & \text{Diffusion} & \phantom{0}6.05 \newline MMDiT+REPA (ODE; NFE=50, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.73} \newline \hline \text{MMDiT (SDE; NFE=250, 150K iter)} & \text{Diffusion} & \phantom{0}5.30 \newline MMDiT+REPA (SDE; NFE=250, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.14} \newline \hline \end{array}

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

[W1-2] it seems FID improvements with classifier-free guidance are not significant.

We politely disagree with your concern. We strongly believe that we have already shown that FID improvement is also quite significant with CFG; for instance, in Table 2 of our manuscript, REPA outperforms the vanilla SiT-XL/2 with CFG using 7x fewer iterations. We have also provided experimental results showing that the result can be further improved with longer training (1.96 at 1M iteration to 1.80 at 4M iteration). This has also been highlighted in L482-L483 in our main manuscript. Moreover, the above text-to-image generation results in [W1-1] are achieved with CFG, where REPA shows considerably better FID scores (5.30$\to$4.14). To further address your concern, we also conducted an additional experiment on ImageNet 512x512 and compared FID values with CFG: SiT-XL/2 trained with REPA outperforms the vanilla SiT-XL/2 FID by improving values from 2.62 to 2.08 using 3x fewer training iterations.

Finally, we additionally provide a similar plot to Figure 1 with classifier-free guidance in Appendix G in the revised manuscript. The plot also shows similar trends to those in Figure 1 in our manuscript, validating the effectiveness of REPA even with CFG.

[W2] Linear probing comparison is unfair: Performance for the diffusion model would be smaller on tasks like semantic segmentation.

We note that linear probing has been widely used in previous diffusion model representation analysis [4, 5, 6, 7]. Moreover, we could not perform the semantic segmentation that you mentioned because most open-sourced pretrained diffusion transformers (particularly DiT and SiT) are latent diffusion models. This makes the spatial resolution of DiTs much smaller (8x) in contrast to pretrained visual encoders, and thus it is difficult to perform semantic segmentation in this latent space. Nevertheless, we agree that exploring representational gaps using other tasks like dense prediction (e.g. [8]) should be an interesting direction in the future by training an additional pixel-level diffusion model with DiT architectures.

[4] Deconstructing Denoising Diffusion Models for Self-Supervised Learning, arXiv 2024
[5] Denoising Diffusion Autoencoders are Unified Self-supervised Learners, ICCV 2023
[6] Diffusion Models Beat GANs on Image Classification. arXiv 2023
[7] Do Text-free Diffusion Models Learn Discriminative Visual Representations?, ECCV 2024
[8] Diffusion Model as Representation Learner, ICCV 2023

[Q1] Add literature review: [a] and [b]

Thanks for introducing relevant works. [a] proposes a method to learn disentangled representations using a variety of frozen generative models, such as pretrained GANs. [b] also tries to solve a similar problem but focuses on disentangling diffusion model representations. Our goal is to improve diffusion model training using frozen visual representations. We also did several analyses using pretrained models. We also added this discussion and citations in the revised manuscript.

Dear Reviewer U4Y4,

Thank you again for your time and efforts in reviewing our paper.

As the discussion period draws close, we kindly remind you that two days remain for further comments or questions. We would appreciate the opportunity to address any additional concerns you may have before the discussion phase ends.

Thank you very much!

Many thanks,
Authors

We deeply appreciate your insightful comments and efforts in reviewing our manuscript. We respond to each of your comments one-by-one in what follows. In the revised draft, we mark our major revisions as “blue”.

[W1] Higher-resolution results (e.g., 512x512).

Thanks for the suggestion! We agree that high-resolution results will make the paper more convincing. To address this concern, we additionally train SiT-XL/2 on ImageNet-512x512 with REPA and report the quantitative evaluation in the table below. Notably, as shown in this table, the model (with REPA) already outperforms the vanilla SiT-XL/2 in terms of four metrics (FID, sFID, IS, and Rec) using >3x fewer training iterations. We added this result in Appendix J in the revised manuscript, along with qualitative results. We will also add the final results in the final draft with longer training (please understand that it takes more than 2 weeks, so it will not be possible to update them until the end of the discussion period).

\begin{array}{lcccccc} \hline \phantom{0}\phantom{0} \text{Model} & \text{Epochs} & \text{FID $\downarrow$} & \text{sFID $\downarrow$} & \text{IS $\uparrow$} & \text{Pre. $\uparrow$} & \text{Rec. $\uparrow$} \newline \hline \text{Pixel diffusion} \newline \phantom{0}\phantom{0} \text{VDM++} & - & 2.65 & - & 278.1 & - & - \newline \phantom{0}\phantom{0} \text{ADM-G, ADM-U} & 400 & 2.85 & 5.86 & 221.7 & 0.84 & 0.53 \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-Net)} & 800 & 4.28 & - & 171.0 & - & - \newline \phantom{0}\phantom{0} \text{Simple diffusion (U-ViT, L)} & 800 & 4.53 & - & 205.3 & - & - \newline \hline \text{Latent diffusion, Transformer} \newline \phantom{0}\phantom{0} \text{MaskDiT} & 800 & 2.50 & 5.10 & 256.3 & 0.83 & 0.56 \newline \phantom{0}\phantom{0} \text{DiT-XL/2} & 600 & 3.04 & 5.02 & 240.8 & 0.84 & 0.54 \newline \phantom{0}\phantom{0} \text{SiT-XL/2} & 600 & 2.62 & 4.18 & 252.2 & 0.84 & 0.57 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & \phantom{0}80 & {2.44} & 4.21 & 247.3 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 100 & 2.32 & \bf{4.16} & 255.7 & 0.84 & 0.56 \newline \phantom{0}\phantom{0} \text{+REPA (ours)} & 200 & \bf{2.08} & 4.19 & \bf{274.6} & 0.83 & \bf{0.58} \newline \hline \end{array}

[W2] DINO-FD on ImageNet-256x256.

This is a good point. Following your suggestion, we additionally measure the DINOv2-FD [1] of SIT-XL/2+REPA, as reported in the table below. It shows that our method achieves state-of-the-art DINO-FD on 256x256 resolution with a value of 63.90. Note that this improvement does not simply come from using SiT-XL/2 as a backbone; we observe that the vanilla SiT-XL/2 achieves an even worse DINOv2-FD of 90.93 than DiT-XL/2. We will also add the DINOv2-FD values on ImageNet 512x512 in the final draft after we complete training the 512x512 model with a large number of training iterations.

\begin{array}{lc} \hline \text{Method} & \text{DINOv2-FD}\downarrow \newline \hline \text{BigGAN} & 401.22 \newline \text{RQ-Transformer} & 304.05 \newline \text{GigaGAN} & 228.37 \newline \text{MaskGIT} & 214.45 \newline \text{LDM} & 112.40 \newline \text{ADMG-ADMU} & 111.24 \newline \text{DiT-XL/2} & \phantom{0}79.36 \newline \hline \text{SiT-XL/2} & \phantom{0}90.93 \newline \text{+REPA (ours)} & \phantom{0}\bf{63.90} \newline \hline \end{array}

Note that we follow the evaluation setup in [1] for a fair comparison with baseline results. Specifically, we do not tune the CFG coefficient for DINOv2-FD and use the same coefficient used for FID computation. To validate the improvement solely from REPA, We also report the DINOv2-FD of the vanilla SiT-XL/2 using the official implementation of SiT.

[1] Exposing Flaws of Generative Model Evaluation Metrics and Their Unfair Treatment of Diffusion Models, NeurIPS 2023

[W3] Feature visualization before vs. after REPA.

Following your suggestion, we added a feature map visualization of pretrained SiT models before and after alignment in Appendix L in the revised manuscript. As shown in the figure, vanilla SiT-XL/2 exhibits a noisy feature map in the early layers at large timestep $t$. In contrast, with alignment, the model shows cleaner feature maps with clear coarse-to-fine patterns across model depths.

[W4] Effect of different diffusion objectives.

We remark that our analyses on SiT (in Figure 2) and those on DiT (in Figure 9) already correspond to two different diffusion objectives, i.e., the v-prediction and $\epsilon$-prediction models, respectively. The results in Table 2 also validate the effectiveness of REPA for these models; we have clarified this point in the revised manuscript.

To further address your concern, we additionally report the performance of DiT-L/2 + REPA using the x0-prediction objective in the table below. Similar to our experimental results with v-prediction or $\epsilon$-prediction objectives in Table 2, REPA is also effective in this case.

\begin{array}{lc} \hline \text{Method} & \text{FID}\downarrow \newline \hline \text{DiT-L/2 ($\epsilon$-pred.)} & 23.3 \newline \text{DiT-L/2+REPA ($\epsilon$-pred.)} & 15.6 \newline \text{DiT-L/2+REPA ($\mathbf{x}_0$-pred.)} & 16.5 \newline \hline \end{array}

[W5] REPA for text-to-image or text-to-video generation

To address your concern, we additionally conducted a text-to-image (T2I) generation experiment on MS-COCO [2], following the setup in previous literature [3]. We tested REPA on MMDiT [4], a variant of DiT designed for T2I generation, given that the standard DiT/SiT architectures are not suited for this task. As shown in the table below, REPA also shows clear improvements in T2I generation, highlighting the importance of alignment of visual representations even under the presence of text representations. In this respect, we believe that scaling up REPA training to large-scale T2I models will be a promising direction for future research. We added these results with qualitative comparison in Appendix K in the revised manuscript.

\begin{array}{lcc} \hline \text{Method} & \text{Type} & \text{FID} \newline \hline \text{AttnGAN} & \text{GAN} & 35.49 \newline \text{DM-GAN} & \text{GAN} & 32.64 \newline \text{VQ-Diffusion} & \text{Discrete Diffusion} & 19.75 \newline \text{DF-GAN} & \text{GAN} & 19.32 \newline \text{XMC-GAN} & \text{GAN} & \phantom{0}9.33 \newline \text{Frido} & \text{Diffusion} & \phantom{0}8.97 \newline \text{LAFITE} & \text{GAN} & \phantom{0}8.12 \newline \hline \text{U-Net} & \text{Diffusion} & \phantom{0}7.32 \newline \text{U-ViT-S/2} & \text{Diffusion} & \phantom{0}5.95 \newline \text{U-ViT-S/2 (Deep)} & \text{Diffusion} & \phantom{0}5.48 \newline \hline \text{MMDiT (ODE; NFE=50, 150K iter)} & \text{Diffusion} & \phantom{0}6.05 \newline MMDiT+REPA (ODE; NFE=50, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.73} \newline \hline \text{MMDiT (SDE; NFE=250, 150K iter)} & \text{Diffusion} & \phantom{0}5.30 \newline MMDiT+REPA (SDE; NFE=250, 150K iter) & \text{Diffusion} & \phantom{0}\bf{4.14} \newline \hline \end{array}

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

[Q1] Code release.

Thank you! We plan to make all the training and eval code, model checkpoints, and supporting scripts from the paper open source, along with the necessary tools to replicate the results presented in the paper.

[Q2-1] Why direct DINOv2 fine-tuning fails?

We hypothesize that this is mainly due to the input mismatch problem: DINOv2 has never been exposed to noisy inputs, and they are trained with image pixels rather than compressed latent images computed from the Stable Diffusion VAE. Thus, DINOv2 initialization might not be that beneficial for training latent diffusion transformers.

[Q2-2] Diffusion model from scratch with partial DINOv2-based loss.

In our work, we did not consider such a training scheme for several reasons. First, since our method is inspired by analysis of the “internal” representation of diffusion transformers, we focus on a more explicit form of regularization that aligns between representations of diffusion transformers and powerful pretrained visual encoders. Moreover, applying an LPIPS-like loss using DINOv2 is both memory- and computation-intensive compared to REPA, particularly for latent diffusion models. This is because inputs of pretrained visual encoders (e.g., DINOv2) are image pixels, but outputs of diffusion transformers are compressed latent images. Thus, one should decode the output from diffusion transformers to image pixels through pretrained VAE at every training iteration before feeding them to pretrained visual encoders like DINOv2.

[Q3] How to handle if the spatial dimensions differ between DINOv2 and DiT?

When spatial dimensions differ, we align them by interpolating the positional embeddings of pretrained visual encoders and resizing input images accordingly (e.g., as done in [5]). For instance, the original MoCov3 has a spatial dimension of 14x14 after patch embedding. However, the spatial dimensions of DiT/SiT models that we used in our experiments are 16x16; thus, we interpolated positional embedding of MoCov3 and resized input image resolution (from 224x224 to 256x256) to make the spatial dimension to be 16x16. This strategy also worked well with higher-resolution experiments (512x512 resolution experiments in [W1]) with DINOv2.

[5] ResFormer: Scaling ViTs with Multi-Resolution Training, CVPR 2023.

1. experiments of high-resolution 512x512 are provided, remaining good improvement
2. experiments of text-to-image are provided, showing good improvement even with text condition
3. the results of DINO-Fid are provided, demonstrating robust performance improvement

Therefore, the rebuttal address most of my concerns, and i will raise my score to 10

Dear Reviewer XnZq,

We are happy to hear that our rebuttal addressed your concerns well! Also, we appreciate your support for our work. If you have any further questions or suggestions, please do not hesitate to let us know.

Best regards,
Authors