Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think
Authors: Sihyun Yu
Sangkyung Kwak
Huiwon Jang
Jongheon Jeong
Jonathan Huang
Jinwoo Shin
Saining Xie
Venue: ICLR 2025 Oral (2025)
Keywords: Diffusion models
Representation learning
Primary area: generative models

Abstract: Recent studies have shown that the denoising process in (generative) diffusion models can induce meaningful (discriminative) representations inside the model, though the quality of these representations still lags behind those learned through recent self-supervised learning methods. We argue that one main bottleneck in training large-scale diffusion models for generation lies in effectively learning these representations. Moreover, training can be made easier by incorporating high-quality external visual representations, rather than relying solely on the diffusion models to learn them independently. We study this by introducing a straightforward regularization called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from external, pretrained visual encoders. The results are striking: our simple strategy yields significant improvements in both training efficiency and generation quality when applied to popular diffusion and flow-based transformers, such as DiTs and SiTs. For instance, our method can speed up SiT training by over 17.5$\times$, matching the performance (without classifier-free guidance) of a SiT-XL model trained for 7M steps in less than 400K steps. In terms of final generation quality, our approach achieves state-of-the-art results of FID=1.42 using classifier-free guidance with the guidance interval.

Published as a conference paper at ICLR 2025
REPRESENTATION ALIGNMENT FOR GENERATION: TRAINING DIFFUSION TRANSFORMERS
IS EASIER THAN YOU THINK
Sihyun Yuu Sangkyung kwak Huiwon Jang'
Jongheon Jeong² Jonathan Huang
Jinwoo Shin Saining Xie
KAIST Korea University SScaled Foundations "New York University
ABSTRACT
Recent studies have shown that the denoising process in (generative) diffusion models can induce meaningful (discriminative) representations inside the model, though the quality of these representations still lags behind those learned through recent self- supervieed learning methods. We argue that one main bottleneck in training large- scale diffusion models for generation lies in effectively learning these representations. Moreover, training be made easier by incorporating high-quality external visual representations, rather than relying solely on the diffu- sion models learn them independently. We study this by introducing straight- forward regularization called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from external, pretrained visual encoders. The results are striking: our simple strategy yields sigmificant improvements in both training ef- ficiency and generation quality when applied to popular diffusion and flow-based transformers, such as DiT's and SiTs. For instance, our method can speed up SiT training by over 17.5x, matching the performance (without classificr-free guid- ance) of a SiT-XL model trained for 7M steps in less than 400K steps. In terms final generation quality, our approach achieves of-te-of-the-ar results of FID=1142 using classifier-fres guidance with the guidance interval.
9000
Dancising Objective
DIT-XLIZ
DITISIT Block
500
SIT-XL2
DIT-XL/2 + REPA
Representation Aligrment
DITISIT
FID-50K
SIT-XL/2 REPA
DITISIT Block
25.0
DITISIT Block
Visual Enceders
Pretrained
DIT/SIT Block
DITISIT Block
100
50K
400K
1M
2M
4M
7M
Training Iteration
Figure 1: Representation alignment makes diffuusio transformer training significantly easier. Our framework, REPA, explicitly aligns the diffusion model representation with powerful pretrained visual representation through a simple regularization. Notably, model training becomes significantly more efficient and effective, and achieves 11.5xx faster convergence than the vanilla model.
INTRODUCTION
Generative models based on denoising, such as diffusion models (Ho et al., 2020; Song et al., 2021) and flow-based models (Albergo & Vanden-Eijnden, 2023; Lipman et al., 2022: Liu et al., 2023), have been a scalable approach in generating high- digensiinii visual data. They achieve remarkably successful results in challenging tasks such as zero-shot text-to-image (Podell al., 2023: Saharia et al., 2022; Esser et al., 2024) or text-to-video (Polyak et al., 2024; Brooks et al., 2024) generation.
"Equal advising.
Project page: httpas //st hyu..mee///eom........mmm.

---

Published as a conference paper at ICLR 2025
Mocov3-B
0.3
Mocov3-B
1-0.0
CKNNA (vs. DINOV2-9: *)
Validation Acc., (%)
Large
CKNNA (vs. DINOv2-9:
DIVO-E"
DINO-B
60-
0.2
0.2
40-
t:0.0
DINGv2-g
0.1
MAEL
0.1
MAE-L
t-0.2b
SIT-8/2
20-
SIT -IT-X/22
10
15
15
20
50K100K
400K
7M
Layer Index
Layer Index
Training Iteration
(a) Semantic gap: Linear probing
(b) Alignment to DINOv2-g
(c) Alignment progression
Figure 2: Alignment behavior for a pretrained SiT model. We empirically investigate the feature alignment between DINOv2-g and the original SiT-X1-/2 checkpoint trained for 7M iterations. (a) While SiT learns semantically meaningful representations, a significant gap remains compared to DINOv2. (b) Using CKNNA (Huh et al., 2024), we observe that SiT already shows some alignment with DINOv2, though its absolute value is lower compared to other vision encoders. (c) Alignment improves with larger model and longer training. but the progress remains slow and insufficient.
Recent works have explored the use of diffusion models as representation learners (Li et al., 2023a; Xiang al., 2023; Chen al., 2024c; Mukhopadhyay et al., 2021) and have shown that they learn discriminative features in their hidden states, and better diffusion models learn better representations (Xiang al., 2023). In fact, this observation closely related to earlier approaches that employ denoising score matching (Vincent, 201 1) as self- self-super learning method (Bengio et al., 2013), which implicitly learns a representation as a hidden state of denoising autoencoder Sy sx(x) through reconstruction of from the corrupted data & (Yang & Wang, 2023). However, the reconstruction task may not be suitable task for learning good representations, as it is not capable climinating unnecessary details in for representation learning (LeCun, 2022; Assran et al., 2023).
Our approach. In this paper, we identify that the main challenge in training diffusion models stems from the need to learn high-quality internal representation h. We demonstrate that the training process for generative diffusion models becomes significantly easier and more effective when sup- ported by an external representation, y.. Specifically, we propose simple regularization technique that leverages recent advances in self-supervised visual representations as y., leading to substantial improvements in both training efficiency and the generation quality of diffusion transformers.
We start by performing an empirical analysis with recent diffusion transformers (Peebles & Xie, 2023; Ma al., 2024a) and the state-of-the-art self-supervised vision model, DINOv2 (Oquab al., 2024). Similar to prior studies (Xiang et al., 2023), we first observe that pretrained diffusion models do indeed learn meaningful discriminative representations (as shown by the linear probing results in Figure 2a). However, these representations are significantly inferior to those produced by DINOv2. Next, we find that the alignment between the representations learned by the diffusion model and those of DINOv2 (Figure 2b) is still considered weak,' which study by measuring their repre- sentation alignment (Huh et al., 2024). Finally, we observe this alignment between diffusion models and DINOv2 improves consistently with longer training and larger models (Figure 2c).
These insights inspire us to enhance generative models by incorporating external self supervised rep- resentations. However, this approach is not straightforward when using off- the-shelf self-superviser visual encoders (e.g., by fine-tuning an encoder for generation tasks). The first challenge is an input mismatch: diffusion models work with noisy inputs X, whereas most self-supervised learning en- coders are trained on clean images x. This issue is even more pronounced modern latent diffusion models, which take compressed latent image E(x) from a pretrained VAE encoder (Rombach et al., 2022) as input. Additionally, these off-the-shelf vision encoders are not designed for tasks like reconstruction or generation. To overcome these technical hurdles, we guide the feature learning of diffusion models using regularization technique that distills pretrained self-supervised representa- tions into diffusion representations, offering flexible way to integrate high-quality representations.
We describe this as "weak" hecause relatively, the alignments are much poorer than those seen with other self-supervised encoders (e.g., MoCov3 (Chen et al., 2021)), even after extensive training.
2

---

Published as a conference paper at ICLR 2025
SIT-XL/2
Validation Acn. (1)
80
Smaiter yep
OKNNA (vs. DINOV2-91
0.6
MIE
SIT-XLy2 REPA
SIT SIT-
80+
60-
0.4
FID-50K
Better FID,
40
VaL Acc.
40
0.2
DINDVI-9
SIT-30//2
20+
20
SIT-XL/2 REPA
0.0
12
16
20
24
40
70
Layer Index
Layer Index
Valldation Acc. (%)
(a) Semantic gap: Linear probing
(b) Alignment to DINOv2-g
(c) Acc. and FID progression
Figure 3: Bridging the representation gap: (a) Our method, REPA significantly reduces the "se- mantic gap" between diffusion transformers and DINOv2, as demonstrated by the linear probing results ImageNet classification.
(b) With REPA, the alignment between diffusion transformers and DINOv2 improves substantially, even after just a few (e.g., 8) layers. (c) Notably, with improved alignment, can push the SIT model's ber of training iterations, it delivers both better generation quality and stronger linear probing results. We single network trained with REPA layer 8 and perform the evaluation at different layers.
rreressttttereentttttinnn envelope: within the same num-
Specifically, we introduce REPresentation Alignment (REPA), a simple regularization technique built on recent diffusion transformer architectures (Peebles & Xie, 2023). In essence, REPA dis- tills the pretrained self-supervised visual representation y. of a clean image into the diffusion transformer representation of noisy input This regularization reduces the semantic gap in the representation h (Figure 3a) and better aligns it with the target self-supervised representations (Figure 3b). Notably, this enhanced aligmment significantly boosts the generation performance of diffusion transformers (Figure 3c). Interestingly, with REPA, we observe that sufficient represen- tation alignment can be achieved by aligning only the first few transformer blocks. This, in turn, allows the later layers of the diffusion transformers to focus on capturing figh-rreuuecy details based on the aligned representations, further improving generation performance.
Based on our analysis, we conduct system-level comparison to demonstrate the effectiveness of our scheme by applying it to two recent diffusion transformers: DiTs (Peebles Xie, 2023) and SiTs (Ma et al., 2024a). For SiT training, show the model achieves FID=7.9 on class-conditions ImageNet (Deng et al., 2009) generation only using 400K training iteration (without classifier free guidance; Ho & Salimans 2022) which is 17.5x faster than the vanilla SiTs. Moreover, with classifier-free guidance, our scheme shows an improved FID at the final from 2.06 to 1.80 and achieves state-of-ther arr results of 1ID-1.42 with guidance interval (Kynkäänniemí et al., 2024).
We highlight the main contributions of this paper below:
We hypothesize that learning high-quality representations in diffusion transformers is essential for improving their generation performance.
We introduce REPA, a simple regularization for aligning diffusion transformer representations with strong self-supervised visual representations.
Our framework improves the generation performance of diffusion transformers, e.g., for SiTs, we achieve 17.5x faster training for SiT's and improved FID) scores on ImageNet generation.
2 PRELIMINARIES
We present a brief overview of flow and diffusion-based models through the unified perspective of stochastic interpolants (Albergo et al., 2023; Ma et al., 2024a); Appendix A for more details.
We consider a continuous time-dependent process with a data x, ~ p(x) and a Gaussian noise (0, 1 E 0, T):
X = = Oic.
ag =- 1, or 0 0.
where ar and or are decreasing and increasing function of I, respectively. Given such a process, there exists probability flow ordinary differential equation (PF ODE) with velocity field
*; =v(x;,0 1),
3

---

Published as a conference paper at ICLR 2025
where the distribution of this ODE at is equal to the marginal px()). Thus, data can be sampled by solving this PF ODE in Eq. (2) through existing ODE samplers (e.g., Euler sampler) starting from a random Gaussian noise e (0, (Lipman et al., 2022; Ma et al., 2024a).
This velocity v(x, t) is represented as the following sum of two conditional expectations
v(x. t) = |x x] = d, |x: x] + o = x),
(3)
which can be approximated with model (x( 1) by minimizing the following training objective:
Livelacity (0) == Exx.,c, [ix () â - o
(4)
Moreover, there exists a reverse stochastic differential equation (SDE) that the marginal pa (x) coin- cides with the one of PF ODE in Eq. (2) with diffusion coefficient w, (Ma et al., 2024a):
dx. v(x,, t)dt 2 w t)dt + vwidwe.
(5)
where the score s(x,, t) is the following conditional expectation
s(x,, () = of -o B = x].
(6)
and it can be directly computed using the velocity v(x, t) for t >> 0 as
s(x, 1) = or vx, d;
0+0; agtt
implying that data can be alternatively generated through Eq. (5) with SDE solvers.
Following Ma et al. (2024a), we mainly consider a simple linear interpolant with restricting T = 1: ar = 1 - and a, = l. However, our approach is applicable to any similar variants (e.g., DDPM; Ho et al. 2020), which has similar formulation but uses discretized process and different a, de that NN, I) becomes an equilibrium distribution (i.e., X converges to (0, I) only if I -> oo).
3 REPA: REGULARIZATION FOR REPRESENTATION ALIGNMENT
3.1 OVERVIEW
Let p(x) be unknown target distribution for data E X. Our goal is to approximate p(x) through a model distribution using a dataset drawn from p(x). To lower computational costs, we adopt the recent prevalent latent diffusion (Rombach et al., 2022). This involves learning a latent distribution p(z), which is defined as the distribution of a compressed latent variable z = E(x), where E is an encoder from pretrained autoencoder (e.g., KL-VAE, Rombach et al. 2022), with Pama (x).
We aim to learn this distribution by training a diffusion model Vo t) using objectives such as velocity prediction, described in Section 2. Here, we revisit denoising score matching within the context of self-supervised representation learning (Bengio et al., 2013). From this perspective, one can think of the diffusion model vo (o () composition of two functions go So with an encoder fe : Z + H. with fe(z,) h, and a decoder ge H. + Z with ge (h,) = Vz. where the encoder fa implicitly learns a representation h, that reconstructs the target Vt.
However, learning a good representation through producing a prediction of the input space (e.g., generating pixels) can be challenging, the model often not capable of eliminating unnecessary details, which is crucial for developing strong representation (LeCun, 2022; Assran et al., 2023). We argue that key bottleneck in the training of large-scale diffusion models for generation lies in representation learning, an area where current diffusion models fall short. We also hypothesize that the training process can be made easier by guiding the model with high-quality external visual representations, rather than relying solely on the diffusion model to learn them independently.
To address this challenge, we introduce a simple regularization method called REPresentation Align- ment (REPA) using the recent diffusion transformer architectures (Peebles & Xie, 2023; Ma et al., 2024a) (see Appendix B for an illustration). In a nutshell, our regularization distills pretrained self- supervised visual representations to diffusion transformers in simple and effective way. This al- lows the diffusion model to leverage these semantically rich external representations for generation, leading to substantial boost in performance.
4

---

Published as a conference paper at ICLR 2025
Training Iteration
100K
200K
400K
100K
200K
400K
100K
200K
400K
ST--L/2
S0--XU22 RR
SIT-L//2 REPA SIT-XU22
Figure 4: REPA improves visual scaling. We compare the images generated by two SiT-X1-/2 models during the first 400K iterations, with REPA applied to one of the models. Both models share the same noise, sampler, and number of sampling steps, and neither uses classifier guidance.
3.2 OBSERVATIONS
To take deeper dive into this, first investigate the layer-wise behavior of the pretrained SiT model (Ma et al., 2024a) on ImageNet (Deng et al., 2009), which uses linear interpolants and velocity prediction for training. In particular, we focus measuring the representation gap between the diffusion transformer and the state-of-the-art self-supervised DINOv2 model (Oquab et al., 2024). We examine this from three angles: semantic gap, feature alignment progression, and their final feature alignment. For the semantic gap, we compare linear probing results using DINOv2 features with those from SiT models trained for 7M iterations, following the same protocol as in Xiang et al. (2023), which involves linear probing on globally pooled hidden states of the diffusion transformer. Next, measure feature alignments, we use CKNNA (Huh al., 2024), a kernel aligmment metric related to CKA (Kornblith et al., 2019), but based mutual nearest neighbors. This allows for quantitative assessment of alignment between different representations. We summarize the result in Figure 2 and more details (e.g., definition of CKNNA) in Appendix C.I.
Diffusion transformers exhibit a significant semantic gap from state-of-the-art visual encoders. As shown in Figure 2a, we observe that the hidden state representation of the pretrained diffusion transformer, in line with prior works (Xiang et al., 2023; Chen et al., 2024c), achieves a reasonably high linear probing peak at layer 20. However, its performance remains well below that of DINOv2, indicating a substantial semantic gap between the two representations. Additionally, we find that after reaching this peak, linear probing performance quickly declines, suggesting that the diffusion transformer must shift away from focusing solely on learning semantically-rin representations in order generate images with high-frequency details.
Diffusion representations are already (weakly) aligned with other visual representations. In Figure 2b, report representational alignments between SiT and DINOv2 using CKNNA. In particular, the SiT model representation already shows better alignment than MAE (He et al., 2022), which is also self-supervised learning approach based on the reconstruction of masked patches. However, the absolute alignment score remains lower than that observed between other self-supervised learning methods (e.g., MoCov3 (Chen et al., 2021) vs. DINOv2). These results suggest that while diffusion transformer representations exhibit some alignment with self-supervised visual representations, the alignment remains weak.
Alignment improves with larger models and extended training. We also measure CKNNA values across different model sizes and training iterations. As depicted in Figure 2c, we observe improved alignment with larger models and extended training. However, the absolute alignment remains low and does not reach the levels observed between other self-super visual encoders (e.g., MoCov3 and DINOv2), even after extensive training of 7M iterations.
5

---

Published as a conference paper at ICLR 2025
These findings not unique to the SIT model but also observed in other denoising basen gener- ative transformers. For instance, in Figure 2, we present a similar analysis using a DiT model (Pee- bles & Xie, 2023) pretrained ImageNet with the DDPM objective (Ho et al., 2020; Nichol & Dhariwal, 2021). ). See Appendix C.2 for more details.
3.3
REPRESENTATION ALIGNMENT WITH SELF-SUPERVISED REPRESENTATIONS
REPA aligns patch-win wise projections of the model's hidden states with pretrained self--uuprvvied visual representations. Specifically, we use the clean image representation the target and ex- plore impact. The goal of this regularization for the diffusion transformer's hidden states to predict noise-invariant, clean visual representations from noisy inputs that contain useful semantic information. This provides meaningful guidance for the subsequent layers to reconstruct the target.
Formally, let f be pretrained encoder and consider clean image Let y. f(x.) E RNND be an encoder output, where N, > 0 are the number of patches and the embedding dimension of f. respectively. REPA aligns hg (h,) E RN with y.. where he (hr) is projection of diffusion transformer encoder output h, fz( that through a trainable projection head hg. In practice, we simply parameterize he using multilayer perceptron (MIP).
In particular, REPA achieves alignment through a maximization patch-wise similarities between the pretrained representation y. and the hidden state h,:
Σ sim(y!") hhay"""hhhhh (ba (b.")))
CREPA (0, Φ) = -Ex.,
where n is patch index and sim(..-) is pre-defined similarity function.
In practice, we add this term to the original diffusion- objectives described in Section 2 and Appendix A. For instance, for the training of velocity model in Eq. (4), the objective becomes:
L = Evelecity + ACREPA
(9)
where > is a hyperparameter that controls the tradeoff between denoising and representation alignment. We primarily investigate the impact of this regularization two popular objectives: Im- proved DDPM (Nichol & Dhariwal, 2021) used in DiT (Peebles Xie, 2023) and linear stochastic interpolants used in SiT (Ma et al., 2024a), though other objectives can also be considered.
EXPERIMENTS
We validate the performance of REPA and the effect of the proposed components through extensive experiments. particular, we investigate the following questions:
Can REPA improve diffusion transformer training significantl? (Table 2, 3, 4, Figure 4, 6)
REPA scalable terms of model size and representation quality? (Table 2, Figure 5)
Can diffusion model representations be aligned with various visual representations? (Figure 8)
4.1
SETUP
Implementation details. We strictly follow the setup in DiT (Peebles & Xie, 2023) and SiT (Ma et al., 2024a) unless otherwise specified. Specifically, we use Ima- geNet (Deng et al., 2009), where each image is prepro- cessed the resolution of 256 (denoted as Ima- geNet 256x256) and follow ADM (Dhariwal & Nichol, 2021) for other data preprocessing protocols. Each image then encoded into a compressed vector 32 using the Stable Diffusion VAE (Rombach et al., 2022). For model configurations, we use the B/2, L/2, and XL/2 architectures introduced the DiT and SiT papers, which process in- puts with patch size of 2 (see Table 1 for details). To ensure fair comparison with DiT's and SiTs, we consistently use a batch size of 256 during training. Additional experimental details, including hyperparameter settings and computing resources, are provided in Appendix D.
Table 1: Model configuration details.
Config
#Layers
Hidden dim
#Heuds
B/2
768
12
1024
XL/2
28
1152
16

---

Published as a conference paper at ICLR 2025
Table 2: Component-wise analysis on ImageNet 256x256. All models are SiT-L/2 trained for 400K iterations. All metrics except accuracy (Acc.) are measured with the SDE Euler- Maaler-ar sampler with NFE=250 and without classifier-free guidance. For Ace., we report linear probing results on the ImageNet validation set using the latent features aligned with the target representation. We fix = 0.5 here. & and 1 indicate whether lower or higher values are better, respectively.
Iter.
Target Repr.
Depth
Objective
FID
sFIDJ
1st
Pre.t
Rec.t
Acc.1
400K
Vanilla SiT-L/2 (Ma et al., 2024a)
18.8
5.29
72.0
0.64
0.64
N/A
400K
MAE-I
NT-Xent
12.5
4.89
90.7
0.68
0.63
57.3
400K
DINO-B
NT-Xent
11.9
5.00
92.9
0.68
59.3
400K
MoCov3-L
8
NT-Xent
11.9
5.06
92.2
0.68
0.64
63.0
400K
I-JEPA-H
8
NT-Xent
11.6
5.21
98.0
0,68
0.64
62.1
400K
CLIP-L
NT-Xent
11.0
5,25
100.4
0.67
0.66
67.2
400K
SigLIP-L
NT-Xent
10.2
5.15
107.0
0.69
0.64
68.8
400K
DINOv2-L
NT-Xent
10.0
5.09
106.6
0.68
0.65
68.1
400K
DINOv2-B
NT-Xent
5.13
107.5
0.69
0.64
65.7
400K
DINOv2-L
8
NT-Xent
10.0
5.09
106.6
0.68
0.65
68.1
400K
DINOv2-g
8
NT-Xent
9.8
5.22
108.9
0.69
0.64
65.7
400K
DINOv2-L.
NT-Xent
10.3
5.23
106.5
0.69
0.65
66.2
400K
DINOv2-L
8
NT-Xem Xent
10.0
5,09
106.6
0,68
0.65
68.1
400K
DINOv2-L
10
NT-Xem
10.5
5,50
105.0
0.68
0.65
68.6
400K
DINOv2-L
NT Ne-Xent
11.2
5.14
100.2
0.68
69.4
400K
DINOv2-I.
NT-Xent
11.6
5.61
99.5
0.67
70.0
400K
DINOv2-L.
16
NT-Xent
12.1
5.34
96.1
0.67
0.64
71.1
400K
DINOv2-L
NT-Xent
10.0
5.09
106.6
0.68
0.65
68.1
400K
DINOv2-L
8
Cos. sim.
9.9
5.34
111.9
0.68
0.65
68.2
Evaluation. We report Fréchet inception distance (FID; Heusel et al. 2017), sFID (Nash et al., 2021), 1), inception score (IS; Salimans et al. 2016), precision (Pre.) and recall (Rec.) (Kynkäänniemi et al., 2019) using 50,000 samples. We also include linear probing results (Acc.) and CKNNA (Huh ct al., 2024) as discussed Section 3.2. We provide more details cach metric in Appendix
Sampler. Following SiT (Ma et al., 2024a), we always the SDE Euler-Maruyama sampler (for SDE with we = σ,) and set the number of function evaluations (NFE) as 250 by default.
Baselines. We use several recent diffusion-based generation methods as baselines, each employing different inputs and network architectures. Specifically, we consider the following four types of approaches: (a) Pixel diffussion: ADM (Dhariwal & Nichol, 2021), VDM++ (Kingma & Gao, 2024), Simple diffusion (Hoogeboom et al., 2023), CDM (Ho et al., 2022). (b) Latent diffusion with U-Ner: LDM (Rombach et al., 2022), (c) Latent diffusion with transformer+ U-Net hybrid models: U-ViT-H/2 (Bao et al., 2023), DiffIT' (Hatamizadeh et al., 2024), and MDTv2-XL.2 (Gao et al., 2023), and (d) Latent diffusion with transformers: MaskDiT (Zheng al., 2024), SD-DiT (Zhu al., 2024), DiT (Peebles & Xie, 2023), and SiT (Ma et al., 2024a). Here, we refer to Transformer+U- Net hybrid models that contain skip connections, which are not originally used in pure transformer architecture. Detailed descriptions of each baseline method are provided in Appendix F.
4.2 COMPONENT- WISE ANALYSIS
We answer the question of whether REPA leads to improved diffusion transformer training. As shown in Table 2, we discover that REPA consistently provides substantially improved genera- tion performance across various design choices, achieving a much better FID score than the vanilla model. Below, we provide a detailed analysis of the impact of each component.
Target representation. We begin by analyzing the effect of using different pretrained self- supervised encoders as the target representation. Notably, there is a strong correlation between the quality of these encoders and the performance of the corresponding aligned diffusion transform- ers. When a diffusion transformer is aligned with a pretrained encoder that offers more semantically meaningful representations (i.e., better linear probing results), the model not only captures better semantics but also exhibits enhanced generation performance, as reflected by improved validation accuracy with linear probing and lower FID scores.

---

Published as a conference paper at ICLR 2025
90.0
70.0
MAEL
500
50.0
McCov3-L
SIT-B/2
FID-S0K
DING.
FID-50K
FID-50K
JEPA-H
25.0
11-
SIT-L/2
CEIP- S4gUP-
SIT-XL/2
DINDv2-1
10.0
B1T-8/2
10.0
SI--12 REPA
7.5
SIT-L/2 REPA
SIT SIT-XU2 REPA
ST-XU/2 REPA
50
DO
70
80
50KL0OK
200K Training Iteration
400K
60
60
70
Validation Acc. (%))
Validation Acc. (96
(a) Different visual encoders
(b) Relative convergence
(c) Validation
FID
Figure 5: Scalability of REPA. (a) Linear probing vs. FID plot of REPA with different target encoders (400K iterations). stronger encoder improves both discrimination and generation per- formance. (b) The relative improvement of REPA over the vanilla model becomes increasingly significant as the model size grows. (c) With a fixed target encoder, larger models reach better per- formance more quickly. In the line plot, results are marked at 50K, 100K, 200K, and 400K iters.
Figure 6: Selected samples on ImageNet 256x256 from the SIT-XL/2 + REPA model. We use classifier-free guidance with w 4.0.
Target encoder size. Next, we investigate the impact of different target representation encoder sizes by evaluating various DINOv2 models (i.e., DINOv2-B, -L., -g). We observe that the performance differences marginal, which hypothesize is due to all DINOv2 models being distilled from the DINOv2-g model and thus sharing similar representations.
Alignment depth. We also examine the effect of attaching the REPA loss to different layers. We find that regularizing only the first few layers (e.g., 8) in training is sufficient, as indicated by the linear probing results in Table 2. Interestingly, limiting regularization to the first few layers further enhances generation performance (e.g., adding REPA to layer or 8 yields best results). We hy- pothesize that this enables the remaining layers concentrate capturing high-frequency details, building on strong representation. In future experiments, we apply REPA to the first 8 layers.
Alignment objective. We compare two simple training objectives for alignment: Normalized Temperature-scaledd Cross Entropy (NT-Xent; Chen et al. 2020a) or negative cosine similarity (cos. sim.). Empirically, we find that NT-Xent offers advantages in the early stages (e.g. 50-100K itera- tions), but the gap diminishes over time. Thus, we opt for cos. sim. in future experiments.
Scalability, Lastly, we investigate the scalability of REPA by varying the model sizes of both the target representation encoders and the diffusion transformers. In general, as summarized in Fig- ure 5a, aligning with stronger representations improves both the generation results and the linear probing performance. Moreover, the convergence speed-up from REPA becomes more significant as the diffusion transformer model increases size. We demonstrate this by plotting FID-50K of different SiT models with and without REPA in Figure 5b: REPA achieves the same FID level more quickly with larger models. Lastly, Figure 5c highlights the relationship between linear probing results and FID scores as model size varies, while keeping the target representation encoder fixed as DINOv2-B. Larger models exhibit a steeper performance improvement (i.e., faster gains both generation and linear evaluation) with longer training.
8

---

Published as a conference paper at ICLR 2025
Table 3:
FID comparisons with vanilla DiTs and SiTs on ImageNet 256x256. We do not use classifier-free guidance (CFG). 4 denotes lower values are better. Iter. indicates the training iteration.
Table 4: System-level comparison ImageNet x56x 256 with CFG. + and t indicate whether lower or higher values better, respectively. Results that include additional CFG scheduling are marked with an asterisk (*), where the guidance interval from (Kynkäänniemi et al., 2024) is applied for REPA.
Model
#Params
Iter.
FID4
Model
Epochs
FIDL
sFIDL
1St
Pre.f
Rec."
DIT-L/2
458M
400K
23.3
Pixel diffusioo
+ REPA (ours)
458M
400K
15.6
ADM-U
400
3.94
6.14
186.7
0.82
0.52
DiT-XL/2
675M
400K
19.5
VDM++
560
2,40
225.3
+ REPA (ours)
675M
400K
12.3
Simple diffusion
800
2.77
211.8
CDM
2160
4,88
158.7
DiT-XL/2
675M
7M
9.6
+ REPA (ours)
675M
850K
9.6
Latent diffusion, U-Net
LDM-4
200
3.60
247.7
0.87
0.48
SiT-B/2
130M
400K
33.0
Latent diffusion, Transformer
U-Ner hybrid
+ REPA (ours)
130M
400K
24.4
U-VIT-H/2
240
2.29
5.68
263.9
0.82
0.57
SiT-L/2
458M
400K
18.8
DifiT
1.73
276.5
0.80
0.62
+ REPA (ours)
458M
400K
9.7
MDTV2-XL/2*
1080
1.58
4.52
314.7
0.79
0.65
+ REPA (ours)
458M
700K
8.4
Latent diffusion, Transformer
S/T-XL/2
675M
400K
17.2
MaskDiT
1600
2,28
5.67
276.6
0.80
0.61
SD-DiT
480
3.23
REPA (ours)
675M
150K
13.6
DiT-XL/2
1400
2.27
4.60
278.2
0.83
0.57
SiT-XL/2
675M
7M
8.3
SIT-XL/2
1400
2.06
4.50
270.3
0.82
0.59
+ REPA (ours)
675M
400K
7.9
+ REPA (ours)
200
1.96
4.49
264.0
0.82
0.60
REPA (ours)
675M
1M
6.4
REPA (ours)
1.80
284.0
0.81
+ REPA (ours)
675M
4M
5.9
REPA (ours)"
800
1.42
4.70
305.7
0.80
0.65
4.3
SYSTEM-LEVEL COMPARISON
Based on the analysis, we perform a system-level comparison between recent state-of-the-art dif- fusion model approaches and diffusion transformers with REPA. First, we compare the FID values between vanilla DiT or SiT models and the same models trained with REPA. As shown in Table 3, REPA shows consistent and significant improvement across all model variants. In particular, XL/2, aligning representation leads to FID=7.9 at 400K iteration, which already exceeds the FID of the vanilla SiT-XL at 7M iteration. Note that the performance continues to improve with longer training; for instance, with SiT-XI S/2- FID becomes 6.4 at 1M iteration and 5.9 at 4M iteration. We also qualitatively compare the progression of generation results in Figure where we use the same initial noise across different models. The model trained with REPA exhibits better progression.
SiT-
Finally, we provide quantitative comparison between S/T-X1/2 with REPA and other recent dif- fusion model methods using classifier-free guidance (Ho & Salimans, 2022). Our method already outperforms the original SiT-XL/2 with 7x fewer epochs and is further improved with longer training. At 800 epochs, SiT-XL/2 with REPA achieves FID of 1.80 with a classifier-free guidance scale of w 1.35, and achieves ofate-of-the-art FID of 1.42 with a extra classifier-free guidance scheduling with guidance interval (Kynkäänniemi et al., 2024). We provide selected qualitative re- sults SiT-XL/2 with REPA in Figure 6 and more examples Appendix H. Moreover, we provide experimental results on ImageNet 512x512 and text-to-image generation in Appendix J and K; we show that REPA also provides significant improvements in such setups.
4.4 ABLATION STUDIES
Representation gap across different timesteps. We begin by comparing the semantic gap (mea- sured through linear probing results) using outputs of the SiT models with different noise scale (i.e,, different timesteps), and maximum CKNNA values using clean DINOv2-g representations. As shown in Figure 7, REPA consistently reduces the representation gap across different noise levels, as indicated by better linear probing results and higher CKNNA values across all noise seales.
Alignment to different visual encoders. In addition, extend the analysis from Section 2 to other visual encoders, not limited to the DINOy2 models. Specifically, we train SiT-L/2 models using REPA with MAE MoCov3. As depicted in Figure 8, these models demonstrate higher CKNNA values across the corresponding target representations than the vanilla model. This indicates that REPA is effective in aligning various visual representations, not limited DINOy2.

---

Published as a conference paper at ICLR 2025
06
05-
0.5
Validatonn Ace (N)
CINNNA Iv.. D.ND-2-EE
02
CRNNAA 22-
D4-
2923
022
25 15
92
01-
Timestop
Index
(a) Linear probing
(b) Alignmeent
(a) McCov3-L
(b) MAE-L
Figure 7: Representation gap across different
Figure 8: Alignment to different target repre- sentations. CKNNA values of SiT-L/2 models using REPA, with (a) MoCov3-L and (b) MAE- L target representations. After that, mea- sure CKNNA using these encoders. REPA con- sistently improves the aligmment regardless of the target representations.
timesteps. We plot the linear probing results and maximum CKNNA values (using DINOv2- g) different timesteps, comparing the vanilla SiT-XL/2 model and the same model trained us- ing REPA. REPA consistently reduces the rep- resentation gap across different noise levels.
Effect of X. We also examine the effect of the regulariza- tion coefficient À by training SiT-X1122 models for 400K with different coefficients 0.25 1.0 and comparing the performance. As shown in Table 5, the performance is robust the values and it is quite saturated after À 0.5.
Table 5: Ablation study for A.
A
0.25
0.5
0.75
1.0
FIDU
8.6
7.9
7.8
7.8
1S1
118.6
122.6
124.4
124.8
5 RELATED WORK
We discuss with the most relevant literature here and provide more discussion in Appendix
Bridging diffusion models and representation learning. Many recent works have attempted to exploit or improve representations learned from diffusion models (Fuest et al., 2024). First, there are hybrid model approaches; Yang et al. (2022) and Deja et al. (2023) train single model capable of both classification and diffusion-based generation. Also, Tian al. (2024) introduces hybrid model capable of segmentation and generation. Next, several works have analyzed and exploited representations in diffusion models: Xiang et al. (2023) and Mukhopadhyay al. (2021) observe that the intermediate representations of diffusion models have discriminative properties. Moreover, Repfusion (Yang Wang, 2023) and DreamTeecher (Li et al., 2023b) propose knowledge distilla- tion schemes using diffusion models to perform various downstream tasks. Our work also shares some similarities, where we focus on designing a regularization for alignments between recent self- supervised and diffusion representations and how they affect generation.
Diffusion models with external representations. Several recent studies have explored leveraging pretrained visual encoders to enhance efficiency and performance of diffusion models (Pernias al., 2024; Li et al., 2024). Wüirstchen (Pernias et al., 2024) introduces a two-stage text-to-image diffu- sion model framework: text-conditioned model that first generates semantic map from text prompt, followed by another diffusion model that synthesizes images based on the semantic map. RCG (Li al., 2024) focuses on unconditional generation, where a compact 1D latent vector is produced by diffusion model and subsequently used as a label for image generation by second diffusion model. We also exploit pretrained representations for improving the diffusion model but without the need of training an additional model that learns the representation distribution.
CONCLUSION
In this paper, we have presented REPA, simple regularization for improving diffusion transform- ers. In particular, we investigated whether diffusion transformer representations can be aligned with recent self-supervised representations, and if it can improve the generation performance of diffusion transformers. We showed REPA can significantly improves generation performance of diffusion transformers with faster convergence speed. We hope our work would facilitate many pos- sible future research directions, including unifying discriminative and generative models and their representations or theoretical analysis; We provide discussion in Appendix M.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We provide hyperparameter details in Section 4 and Appendix D. We also release the implementation and model checkpoints to reproduce the results in the paper.
ACKNOWLEDGMENT
This work was supported by Institute for Information & communications Technology Promo- tion(IITP) grant funded by the Korea MSST) (No. RS-2019-1110075 Artificial In- telligence Graduate School Program (KAIST); No. RS-2024-00500279, Global AI Frontier Lab).
SY thanks the anonymous reviewers, Hankook Lee, Jihoon Tack, Sukmin Yun, and Yisol Choi for their insightful discussions; Myungkyu Koo and Daewon Choi for proofreading; and Kyungmin Lee for providing an excellent diffusion model implementation. SY also thanks Yewon Kim for the assistance in improving the visualizations in the paper.
acknowledges support from the Institute Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. RS-2019-11190079, Arti- ficial Intelligence Graduate School Program (Korea University)), and the ITP--TRCC (Information Technology Research Center) grant funded by the Korea government (MSIT) (No. IITP-2025-RS- 2024-00436857, 10%).
SX thanks discussions with Boyang Zheng and Willis Ma. SX also acknowledges support from the ITTP grant funded by the Korean Government (MSIT) (No. RS-2024-00457882, National Al Re- search Lab Project). the Open Path AI Foundation, Amazon Research Award, Google TRC program, and NSF Award IIS-2443404
REFERENCES
Korbinian Abstreiter, Sarthak Mittal, Stefan Bauer, Bernhard Schölkopf, and Arash Mehrjou. Diffusion-based representation learning, arXiv preprint wwXiv:2105.14257 2021.
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unify- ing framework for flows and diffusions. arXiv preprint arXiv:2303.0877 2023.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In International Conference on Learning Representations, 2023.
Anurag Amab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučié, and Cordelia Schmid. ViVIT: video vision transformer. In IEEE International Conference on Computer Vision, 2021.
Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with joint-embedding predictive architecture. In IEEE Conference on Computer Vision and Pattern Recognition, 2023.
Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: VIT backbone for diffusion models. In IEEE Conference on Computer Vision and Pattern Recognition, 2023.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Transactions Pattern Analysis and Machine Intelligence, 3588) 2013.
Tim Brooks, Bill Peebles. Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. OpenAl Blog, 2024.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In IEEE International Conference Computer Vision, 2021.
11

---

Published as a conference paper at ICLR 2025
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. MaskGiT: Masked gen- erative image transformer. In IEEE Conference on Computer Vision and Pattern Recognition, 2022.
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. arXiv preprint arXiv:2403.0460 044692 2024a.
Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Yue Wu, Zhongdao Wang. James Kwok, Ping Luo, Huchuan Lu, et al. PixArt-a: Fast training of diffusion transformer for photore- alistic text-to-image synthesis. In International Conference on Learning Representations, 2024b.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. International Conference on Machine Learning, 2020a.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.c wrXiv:2003.04297, 2020b.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training celf-superviee vision transformers. In IEEE International Conference on Computer Vision, 2021.
Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He. Deconstructing denoising diffusion models for self superviee learning. arXiv preprint 1XXi:22401.14400 2024c.
Kamil Deja, Tomasz Trzciński, and Jakub M Tomezak. Learning data representations with joint diffusion models. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In International Conference on Machine Learning, 2024.
Wan-Cyuan Fan, Yen-Chun Chen, DongDong Chen, Yu Cheng, Lu Yuan, and Yu-Chiang Frank Wang. Frido: Feature pyramid diffusion for complex scene image synthesis. In AAAI Conference on Artificial Intelligence, 2023.
Michael Fuest, Pingchuan Ma, Ming Gui, Johannes S Fischer, Vincent Tao Hu, and Bjomn Ommer. Diffusson models and representation learning: A survey. arXiv preprint arXiv:2407.0078, 2024.
Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. MDTv2: Masked diffusion transformer a strong image synthesizer. arXiv preprint arXiv:2303. 133.999.3.14389 2023.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Infor- mation Processing Systems, 2014.
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In IEEE Conference Computer Vision and Pattern Recognition, 2022.
12

---

Published as a conference paper at ICLR 2025
Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. European Conference on Computer Vision, 2024.
Ali Hatamizadeh, Jiaming Song, Guilin Liu, Jan Kautz, and Arash Vahdat. DiffIT: Diffusion vision transformers for image generation. In European Conference on Computer Vision, 2024.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In IEEE Conference on Computer Vision and Pattern Recognition, 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. IEEE Conference on Computer Vision and Pattern Recognition, 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Processing Systems, 2017.
Jonathan Ho and Tim Salimans. arXiv:2207. 12598, 2022.
Classifier-free
diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Sali- mans. Cascaded diffusion models for high fidelity image generation. Journal of Machine Learning Research, 23(47):1-35 2022.
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans. Simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, 2023.
Drew Hudson, Daniel Zoran, Mateusz Malinowski, Andrew K Lampinen, Andrew Jaegle, James McClelland, Loic Matthey, Felix Hill, and Alexander Lerchner. SODA: Bottleneck diffusson mod- els for representation learning. In IEEE Conference on Computer Vision and Pattern Recognition, 2024.
Minyoung Huh, Brian Cheung, Tongzhou Wang, and Phillip Isola. The platonic representation hypothesis. In International Conference on Machine Learning, 2024.
Minguk Kang, Jun- Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up GANs for text-to-image synthesis. In JEEE Conference on Computer Vision and Pattern Recognition, 2023.
Minguk Kang, Richard Zhang, Connelly Barnes, Sylvain Paris, Suha Kwak, Jaesik Park, Eli Shecht- man, Jun- Yan Zhu, and Taesung Park. Distilling diffusion models into conditional GANs. In European Conference on Computer Vision, 2024.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for im- proved quality, stability, and variation. In International Conference on Learning Representations, 2018.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Advances in Neural Information Processing Systems, 2022.
Tero Karras, Miik Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, and Samuli Laine. Analyz- ing and improving the training dynamics of diffusion models. In IEEE Conference on Computer Vision and Pattern Recognition, 2024.
Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. Advances Neural Information Processing Systems, 2024.
Diederik Kingma. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
13

---

Published as a conference paper at ICLR 2025
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International Conference on Machine Learning, 2019.
Nupur Kumari, Richard Zhang, Eli Shechtman, and Jun- Yan Zhu. Ensembling off-the-shelf models for GAN training. In IEEE Conference on Computer Vision and Puttern Recognition, 2022.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. In Advances in Neural Information Processing Systems, 2019.
Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in a limited interval improves sample and distribution quality in diffusion models. arXiv preprint arXiv:2404.07724 2024.
Yann LeCun. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62(1):1-62, 2022.
Alexander C Li, Mihir Prabhudesai, Shivam Duggal, Ellis Brown, and Deepak Pathak. Your dif- fusion model is secretly a zero-shot classifier. In IEEE International Conference on Computer Vision, 2023a.
Daiqing Li, Huan Ling, Amlan Kar, David Acuma, Seung Wook Kim, Karsten Kreis, Antonio Tor- ralba, and Sanja Fidler. DreamTeacher: Pretraining image backbones with deep generative mod- els. In IEEE International Conference on Computer Vision, 2023b.
T Li, D Katabi, and K. He. Return of unconditional generation: A self-supervised representation generation method. In Advances in Neural Information Processing Systems, 2024.
Tianhong Li, Huiwen Chang, Shlok Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan. MAGE: Masked generative encoder to unify representation learning and image synthesis. In IEEE Con- ference Computer Vision and Pattern Recognition, 2023c.
Tsung-Y Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference on Computer Vision, 2014.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations, 2023.
I Loshchilov. Decoupled weight decay regularization. In International Conference on Learning Representations, 2017.
Haoyu Lu, Guoxing Yang. Nanyi Fei, Yuqi Huo. Zhiwu Lu, Ping Luo, and Mingyu Ding. VDT: Gereral-purpose video diffusion transformers via mask modeling. In International Conference Learning Representations, 2024.
Nanye Ma, Mark Goldstein, Michael Albergo, Nicholas M Boffi, Eric Vanden-Eijnden and Sain- ing Xic. SiT: Exploring flow and iiffusion-basee generative models with scalable interpolant transformers. In European Conference on Computer Vision, 2024a.
Xin Ma, Yaohui Wang, Gengyun Jia, Xinyuan Chen, Ziwei Liu, Yuan-Fang Li, Cunjian Chen, and Yu Qiao. Latte: Latent diffusion transformer for video generation. arXiv preprint arXiv:2401.0300 2024b.
Soumik Mukhopadhyay, Matthew Gwilliam, Vatsal Agarwal, Namitha Padmanabhan, Archana Swaminathan, Srinidhi Hegde, Tianyi Zhou, and Abhinav Shrivastava. Diffusion models beat GANs on image classification. Advances in Neural Information Processing Systems, 2021.
Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W Battaglia. Generating images with sparse representations. In International Conference Machine Learning, 2021.
14

---

Published as a conference paper at ICLR 2025
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, 2021.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. DINOv2: Learning robust visual features without super- vision. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
William Peebles and Saining Xie. Scalable diffusion models with transformers. IEEE Interna- lional Conference on Computer Vision, 2023.
Pablo Pernias, Dominic Rampas, Mats Leon Richter, Christopher Pal, and Marc Aubreville. Würstchen: An efficient architecture for large-scale text-to-image diffusion models. Inter- national Conference Learning Representations, 2024.
Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhom, Jonas Müller, Joe Penna, and Robin Rombach. SDXL: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952. 2023.
Adam Polyak, Amit Zohar, Andrew Brown, Andros Tjandra, Animesh Sinha, Ann Lee, Apoorv Vyas, Bowen Shi, Chih- Yao Ma, Ching-Yac Yao Chuang. David Yan, Dhruv Choudhary, Dingkang Wang, Geet Sethi, Guan Pang. Haoyu Ma, Ishan Misra, Ji Hou, Jialiang Wang, Kiran Ja- gadeesh, Kunpeng Li, Luxin Zhang, Mannat Singh, Mary Williamson, Matt Le, Mitesh Ku- mar Singh, Peizhao Zhang, Peter Vajda, Quentin Duval, Rohit Girdhar, Roshan Sumbaly, Sai Saketh Rambhatla, Sam Tsai, Samaneh Azadi, Samyak Datta, Sanyuan Chen, Sean Bell, Sharadh Ramaswamy, Shelly Sheynin, Siddharth Bhattacharya, Tao Xu, Tingbo Hou, Wei- Ning Hsu, Xi Yin. Xiaoliang Dai, Yaniv Taigman, Yaqiao Luo, Yen-Cheng Liu, Yi-Chiao Wu, Yue Zhao, Yuval Kirstain, Zecheng He. and Zijian He. MovieGen: A cast me- dia foundation models. Meta Al Blog Post, 2024. URL https: metaa com/blog/ movie-gen-mediaa foundat- ioniioonmmooeee s-
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In IEEE Conference Computer Vision and Pattern Recognition, 2022.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et al. Photorealistic tat-to-immae diffusion models with deep language understanding. In Advances in Neural Infor- mation Processing Systems, 2022.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training GANs. In Advances in Neural Information Processing Systems, 2016.
Axel Sauer, Kashyap Chitta, Jens Müller, and Andreas Geiger. Projected GANs converge faster. Advances in Neural Information Processing Systems, 2021.
Axel Sauer, Katja Schwarz, and Andreas Geiger. StyleGAN-XI. Scaling StyleGAN to large diverse datasets. In ACM SIGGRAPH conference proceedings, 2022.
Axel Sauer, Tero Karras, Samuli Laine, Andreas Geiger, and Timo Aila. StyleGAN-T: Unlocking the power of GANs for fast large-scale text-to-image synthesis. In International Conference Machine Learning, 2023a.
Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion dis- tillation. arXiv preprint arXiv:2311.1704. 2023h.
15

---

Published as a conference paper at ICLR 2025
Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rom- bach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. arXiv preprint drXiv:2403.120155 2024.
Jascha bohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibriun thermodynamics. In International Conference on Machine Learn- ing, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- tional Conference on Learning Representations, 2021.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the Inception architecture for computer vision. In IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Ming Tao, Hao Tang, Fei Wu, Xiao Yuan Jing, Bing-Kun Bao, and Changsheng Xu. DF-GAN: A simple and effective baseline for text-to-image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and paltern recognition, pp. 16515-16525, 2022.
Changyao Tian, Chenxin Tao, Jifeng Dai, Hao Li, Ziheng Li, Lewei Lu, Xiaogang Wang. Hong- sheng Li, Gao Huang, and Xizhou Zhu ADDP: Learning general representations for image recognition and generation with alternating denoising diffusion process. In International Confer- ence on Learning Representations, 2024.
Pascal Vincent. connection between score matching and denoising autoencoders. Neural compu- tation, 23(7):1661-1674 2011.
Weilai Xiang. Hongyu Yang. Di Huang, and Yunhong Wang. Denoising diffusion autoencoders are unified self-supervised learners. In IEEE International Conference on Computer Vision, 2023.
Tao Xu, Pengchuan Zhang, Qiuyuan Huang, Han Zhang, Zhe Gan, Xiaolei Huang, and Xiaodong
He. AttnGAN: Fine-grained text to image generation with attentional generative adversarial net- works. In IEEE Conference on Computer Vision and Pattern Recognition, 2018.
Xingyi Yang and Xinchao Wang. Diffusion model as representation learner. In IEEE International Conference Computer Vision, 2023.
Xiulong Yang, Sheng-Min Shih, Yinlin Fu, Xiaoting Zhao, and Shihao Ji. Your VIT is secretly hybrid discriminative-geerraiit diffusion model. arXiv preprint arXiv:2208.0773 2022.
Sihyun Yu. Weili Nic, De-An Huang, Boyi Li, Jinwoo Shin, and Anima Anandkumar. Efficient video diffusion models via content-frame motion-latent decomposition. In International Conference on Learning Representations, 2024.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzallz, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via de- noising for molecular property prediction. International Conference on Learning Representa- tions, 2023.
Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training In IEEE Conference on Computer Vision and Pattern Recognition, 2023.
Han Zhang, Jing Yu Koh, Jason Baldridge, Honglak Lee, and Yinfei Yang. Cross-modal contrastive learning for text-to-image generation. In IEEE Conference on Computer Vision and Pattern Recognition, 2021.
Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar. Fast training of diffusion models with masked transformers. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Yufan Zhou, Ruiyi Zhang, Changyou Chen, Chunyuan Li, Chris Tensmeyer, Tong Yu, Jiuxiang Gu, Jinhui Xu, and Tong Sun. LAFITE: Towards language-fret free training for text-to- text-to-image generation. In IEEE Conference on Computer Vision and Pattern Recognition, 2021.
16

---

Published as a conference paper at ICLR 2025
Minfeng Zhu, Pingbo Pan, Wei Chen, and Yi Yang. DM-GAN: Dynamic memory generative adver- sarial networks for coxt-to-image synthesis. IEEE Conference on Computer Vision and Pattern Recognition, 2019.
Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao, Zhenglong Sun, Tao Mei, and Chang Wen Chen. SD- DiT: Unleashing the power of self-supervised discrimination in diffusion transformer. In /EEE Conference on Computer Vision and Pattern Recognition, 2024.
17

---

Published as a conference paper at ICLR 2025
A DESCRIPTIONS FOR DIFFUSION-BASED MODELS
We provide an overview of two types of generative models that we use in this paper, which learn the target distribution by training variants of denoising autoencoder. We first explain denoising diffusion probabilistic models (DDPM) in Section A. and stochastic interpolants in Section A.2. For detailed explanations and rigorous proofs, please refer the original papers (Albergo et al., 2023; Ma et al., 2024a) that provide excellent formulations and description.
A.1 DENOISING DIFFUSION PROBABILISTIC MODELS
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) model the target distribution p(x) via learning gradual denoising process from Gaussian distribution (N(0, p(x). Formally, diffusion models learn a reverse process p(x,-11 1 lxx) the pre-defined forward process q(x, (x)) that gradually adds the Gaussian noise starting from p(x) for
I T with fixed 7 0.
For a given Xg ~ p(x), ((x,/xx-1) can be formalized as q(x; xx-1) = (x,;; where Be (0, 1) pre-defined hyperparameters set to be small, In particular, DDPM (Ho al., 2020) shows if one formalizes the reverse process p(xx-1(x;) (with a 3c. ā₁ II 0g for 1 < < 7) as
1TT1,
of
p(x--1(x,) = (xxx111
Vau (x+
t)). Eo (xxxxt)))
(10)
then (x t) can be trained with simple denoising autoencoder objective parameterized by 0:
Cimple == Ex.,e.( [|le Ex(,,,))1
(11)
For (x t), (Ho et al., 2020) shows it is enough to simply define of1 with Be oA. After that, Nichol & Dhariwal (2021) exhibits the performance can be improved if the model jointly learns Eg (x2, f) with tx),,)) dimension-we wise manner through the following objective:
Cvb == exp(v log B, (1 v) log Bt),
(12)
where denotes each component per dimension from the model output and A 1-
Bt.
1-a,
With sufficiently large T and an appropriate scheduling 8,, the distribution p(xr) becomes almost an isotropic Gaussian distribution. Hence, one can generate sample starting from a random noise and perform iterative reverse process p(x, xx-- (x.) to reach the data sample Xo (Ho al., 2020).
18

---

Published as a conference paper at ICLR 2025
A.2 STOCHASTIC INTERPOLANTS
Different from DDPM, flo-bbaod models (Esser et al., 2024; Lipman et al., 2022; Liu et al., 2023) deal with the continuous time-dependent process with data x. p(x) and Gaussian noise e ~ (0, on E [0, l]:
X = 0, + OcE,
ag = 01 ai = 00 = 0,
(13)
where a and og are decreasing and increasing function of (respectively). There exists a proba- bility flow ordinary differential equation (PF ODE) with velocity field
v(x,, t),
(14)
where distribution of this ODE is equal to the marginal pe px()..
The velocity vxx, t) is represented as the following sum of two conditional expectations
v(x, 1) [x X x] = áx, |x: = x] o, - x],
(15)
which can be approximated with model vo(x,, t) by minimizing the following training objective:
Evelucity (0) == Ex..,t (xx,,, t) - á,x, - dtcl
(16)
Note that this also corresponds to the following reverse stochastic differreniaal equation (SDE);
dx, = v(xi, ()dl 2 (xx, l)dl + wundwe
(17)
where the score sx,, t) similarly becomes the conditional expectation
s(x,, t) = 10--EE[eh E|e|x; = =x].
(18)
Similar to v, s can be approximated with a model so t) with the following objective:
Lecr (f) = Ex..e. .e. (1or)s((xe,) t) + <ll]]
(19)
Here, since the score s(x, 1) can be directly computed using the velocity v(x, () for >> 0 as
a.v t) - â,x
s(x, t) = ae â,o; ajoz
(20)
so it is enough to estimate only of the two vectors.
Stochastic interpolants (Albergo et al., 2023) shows any a and or satisfy the three conditions
a + of 0, Vt [0,0011]
2. a and are differentiable, Vt E [0, 1]
01 do = 0, og = =0
leads process that interpolates between Xo and without bias. Thus, one can use a simple interpolant by defining them as a simple function during training and inference, such as linear in- terpolants with 0₁ = 1 and a; and cos( Şt) (Ma et al., 2024a). One another ad advantage of stochastic interpolants is that the diffusion coefficient w, independent training any of score or velocity model. Thus, w, can be also explicitly chosen after training when sampling with the reverse SDE.
pariance-preserving (VP) interpolants with o, cos( / ()
Note that existing score-based diffusion models, including DDPM (Ho al., 2020), can be similarly interpreted as an SDE formulation. In particular, their forward diffusion process can be interpreted as a pre-defined (discretized) forward SDEs that have an equilibrium distribution as N( I) at 1 oo, where the training is done on [0, T] with sufficiently large T (e.g., T 1000) that pl (())) becomes almost isotropic Gaussian. Generation is done by solving the corresponding reverse SDE starting from random Gaussian noise by assuming X7 (), where at, or and the diffusion coefficient w, is implicitly chosen from the forward diffusion process, which might lead to over- complicated design space of ccore-based diffusion models (Karras et al., 2022).
19

---

Published as a conference paper at ICLR 2025
B DIFFUSION TRANSFORMER ARCHITECTURE
Scale, Shift
Pointwise Feedforward
Scale,
Layer
Scalee
Multi-Head
Attention
Scale,
Layer
Tokens
Timestep
Figure 9: DiT block illustration.
We strictly follow the architecture used in DIT (Peebles & Xie, 2023) and SiT (Ma al., 2024a).
The architecture is very similar to a vision transformer (ViTs; Dosovitskiy et al. 2021): an input is patchified, reshaped to ID sequence of patches with a length N, and then fed the model. Similar to DiT and SiT, our architecture also uses a downsampled latent image z E(x) as an input, where x is RGB image and E is an encoder of the stable diffusion variational autoencoder (VAE) (Rombach et al., 2022). Different from the original ViT, our architecture also includes additional modulation layers at each attention block called AdaIN-zero layers. These layers scale and shift each hidden state with respect to the given timestep and additional conditions. We also consider single multilayer perceptron (MLP) that projects a hidden state to the target representation space, which is only used in training. We provide illustration of the DiT block in Figure 9.
C ANALYSIS DETAILS
C. EVALUATION DETAILS
CKNNA (Centered Kernel Nearest- Veirest-Neighhor Alignment) is relaxed version of the popular Cen- tered Kernel Alignment (CKA; Kornblith et al. 2019) that mitigates the strict definition of alignment. We generally follow the notations in the original paper for explanation (Huh al., 2024).
First, CKA have measured global similarities of the models by considering all possible data pairs:
HSIC(K, L)
CKA(K, V (HSIČ(K, K K)HSIC(L,L)
(21)
where K and L are two kernel matrices computed from the dataset using two different networks. Specifically, it is defined Kiý n(cbe, Φ5)) and Lij "e (15))))))) where di, dj and vh, th representations computed from each network at the corresponding data x,, (respectively). By letting inner product kernel, HSIC is defined
HSIC(K, = (n - 1)2
ΣΣ ((o,, 0z) - (o, oi))) ((v,, bb) - ((((,...1)))
(22)
CKNNA considers a relaxed version of Eq. (21) by replacing HSIC(K, L) into Align(K, L). where Align( K, L) computes Eq. (22) only using k-nearest neighborhood embedding in the datasets:
Align(K L) L) (n - 1)2 ΣΣ (().)) o5) - Erl(ón @())) ((1h, bby) - Ex((dh v1)1)
(23)
20

---

Published as a conference paper at ICLR 2025
where (ii,j) is defined as
alt.j:k)= = + and dy knn( R) and th knn(s/s; A)],
(24)
this term only considers k-nearest neighbors at each In this paper, we randomly sample 10,000 images in the validation set in ImageNet (Deng et al., 2009) and report CKNNA with & = 10 based on observation Huh et (2024) that smaller & shows better better alignment.
Linear probing. We follow the setup used in DAE (Chen et al., 2024c), Specifically, we use parameter-free batch normalization layer and train linear layer for 90 epochs with batch size of 16,384. We use the Adam optimizer (Kingma, 2015) with cosine decay learning rate scheduler, where the initial learning rate is set to 0.001.
C.2 DIT ANALYSIS
We also perform a similar analysis have done in Figure 2a (linear probing) and 2b (CKA), and illustrate the result in Figure 10. Overall is shows similar trend; the model includes discriminative representation but the gap is still large compared with DINOv2, as shown in the linear probing results, and also weakly aligned with DINOv2 representations.
0.3
Mocov3
[-]
80
CKNNA (vs, DINOv2-g:
t-101
Validation Acc. (%)
Large
t-201
0.2
80-
40
DENDV2
0.1
L-201
0.0
10
15
20
10
15
20
Layer Index
Layer Index
(a) Semantic gap
(b) Final alignment
Figure 10: Empirical study with the pretrained DiT model. Similar to Figure 2, we compare the semantic gap and measure the feature alignment between DINOv2-g and the DiT-XL/2 model trained with 7M iterations. (a) DiT learns meaningful (discriminative) representation but it still have large gap between DINOv2. (b) Measured with CKNNA (Huh et al., 2024), DiT already shows weak aligmment with DINOv2, but absolute value is still small.
C.3 DESCRIPTION OF PRETRAINED VISUAL ENCODERS
MAE (He et al., 2022) proposes self-supervised representation learning objective for vision transformers, based on the reconstruction task masked patches of input images.
DINO (Caron et al., 2021) is a self-supervised learning method based on self-distillation through the mean of momentum teacher network.
MoCov3 (Chen al., 2021) studies empirical study to train MoCo (He et al., 2020; Chen al., 2020b) on vision transformer and how they can be scaled up.
CLIP (Radford et al., 2021) proposes contrastive learning scheme on large image-text pairs.
DINOv2 (Oquab al., 2024) proposes self-supervised learning method that combines pixel- level and patch-level discriminative objectives by leveraging advanced self-supervised tech- niques and large pre-training dataset.
I-JEPA (Assran et al., 2023) predicts missing parts of an image by learning representations through joint-embedding. focusing on the context of the entire image without relying pixel- level reconstruction.
SigLIP (Zhai et al., 2023) replaces the traditional softmax loss with a pairwise sigmoid loss, enhancing performance and efficiency of image-text representation leaming.
21

---

Published as a conference paper at ICLR 2025
Moreover, in Table 6, also provide the datasets used for training of each of pretrained visual encoder. 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. Note that we do not fine-tune encoders (e.g., SigLIP and CLIP) with the ImageNet dataset, particularly they were trained with another dataset, thereby separating the dataset leakage effect if we these encoders. REPA also achieves significant improvements with these encoders, which validates that the improvement does not simply come from data leakages.
Table 6: Dataset analysis used for pretrained visual encoders.
Method
Dataset
w/ ImageNet- 1K
Text sup.
FID
MAE
ImageNet-Il IK
X
12.5
DINO
ImageNet-Il 1K
0
X
11.9
MoCov3
ImageNet-IK
11.9
I-JEPA
ImageNet-1K
11.6
CLIP
WIT-400M
11.0
SigLIP
WebLi-4B
X
10.2
DINOv2
LVD-142M
O
X
10.0
22

---

Published as a conference paper at ICLR 2025
D HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS
Table 7: Hyperparameter setup.
Figure 3
Table 3 (SiT-B)
Table 3 (SiT-L.)
Table 3 (SIT-XI.)
Table
Architecture
Input dim.
32x32x4
32x32x4
32x32x4
32x32x4
32x32x4
Num. layers
28
12
24
28
28
Hidden dim.
1.152
768
1,024
1,152
1.152
Num. heads
16
12
16
16
16
REPA X
0.5
0.5 4
0.5
0.5
0.5
Alignment depth sim(..)
8
8 NT-Xent DINOv2-I.
8 cos. sim. DINOv2-B
8
cos. sim. DINOv2-B
cos. sim. DINOv2-B
cos, sim.
Encoder f(x)
DINOv2-B
Optimization
Training iteration
1M
400K 256
700K
4M
4M
Batch size Optimizer
256 Adam
256
256 Adam
256 Adam
Adam
Adam
Ir
0.0001 (0.9, 0.999)
0.0001 (0.9. 0.999)
0.0001
0.0001 (0.9, 0.999)
0.0001 (0.9, 0.999)
32)
0.999)
Interpolants
Training objective
v-prediction Marle-Maruyama 250
v-prediction Mareer-aruyyaaa 250
v-prediction Muler-Maruyama 250
v-prediction Maler-Maruyamaa 250
v-prediction Maler-Maruyama
Sampler
Sampling steps
250
Guidance
1.35
Further implementation details. We implement our model based on the original SiT implementa- tion (Ma et al., 2024a). Throughout the experiments, we the exact structure DiT (Pee- bles & Xie, 2023) and SiT (Ma et al., 2024a). We use Adam (Kingma, 2015; Loshchilov, 2017) with constant learning rate of le-4, ,8,,82) (0.9, 0.999) without weight decay. To speed up training. we use mixed-prece precision (fp16) with gradient clipping. We also pre-compute compressed latent vectors from raw pixels via stable diffusion VAE (Rombach al., 2022) and use these la- tent vectors. Because of this, we do not apply any data augmentation. but we find this does lead a big difference, as similarly observed in EDM2 (Karras et al., 2024). We also use decoder for decoding latent vectors to images. For MLP used
stabilitya wwt.alisdd-vaeeeaaaaaaeeeeeeee //sddda-at-eeeee for a projection, we use three-layer MLP with SiLU activations (Elfwing et al., 2018). We provide detailed hyperparameter setup in Table 7.
Pretrained encoders. For MoCov3-B and -L models, we use checkpoint of RCG (Li et al., 2024);² for other checkpoints, official implementations. To adjust a different number of patches between the diffusion transformer and the pretrained encoder, we interpolate positional embeddings of pretrained encoders.
implementation their official checkpoints released in their
Sampler. For sampling, we use the Euler- Mariyyma sampler with the SDE in Eq. (5) with diffu- sion coefficient wy dt, We use the last step of the SDE sampler 0.04, and gives significant improvement, similar to the original SiT paper (Ma et al., 2024a).
Computing resources. We 8 NVIDIA H100 80GB GPUs for experiments; our training speed about 5.4 step/s with a batch size of 256. Note that this can be further boosted with additional engineering (e.g., pre-computation of pretrained encoder features).
"https://githube http://tttttt..oo.././oo/.t //To////11/
23

---

Published as a conference paper at ICLR 2025
E EVALUATION DETAILS
We strictly follow the setup and use the same reference batches of ADM (Dhariwal Nichol, 2021) for evaluation, following their official implementation." 3 We use NVIDIA H100 80GB GPUs or 4090Ti GPUs for evaluation and enable tf32 precision for faster generation, and we find the perfor- mance difference negligible to the original fp32 precision.
In what follows, we explain the main concept of metrics that we used for the evaluation.
FID (Heusel et al., 2017) measures the feature distance between the distributions of real and generated images. It uses the Inception-v3 network (Szegedy et al., 2016) and computes distance based on an assumption that both feature distributions are multivariate gaussian distributions.
sFID (Nash et al., 2021) proposes to compute FID with intermediate spatial features of the Inception-v3 network to capture the generated images' spatial distribution.
IS (Salimans et al., 2016) also uses the Inception-v3 network but use logit for evaluation of the metric. Specifically, measures KL-divergence between the original label distribution and the distribution of logits after the softmax normalization.
Precision and recall (Kynkäänniemi et al., 2019) are based on their classic definitions: the fraction of realistic images and the fraction of training data manifold covered by generated data.
F BASELINES
In what follows, we explain the main idea of baseline methods that used for the evaluation.
ADM (Dhariwal & Nichol, 2021) improves U-Net-based architectures for diffusion models and proposes classifier-guided sampling to balance the quality and diversity tradeoff.
VDM++ (Kingma & Gao, 2024) proposes simple adaptive noise schedule for diffusion models to improve training efficiency.
Simple diffusion (Hoogeboom et al., 2023) proposes a diffusion model for high-resolution im- age generation by exploring various techniques to simplify noise schedule and architectures.
CDM (Ho al., 2022) introduces cascaded diffusion models: similar to progressiveGAN (Kar- ras et al., 2018), it trains multiple diffusion models starting from the lowest resolution and ap- plying or more super-resolution diffusion models for generating high-fidelity images.
LDM (Rombach et al., 2022) proposes latent diffusion models by modeling image distribution in a compressed latent space improve the training efficiency without sacrificing the generation performance.
U-VIT (Bao et al., 2023) proposes ViT-based latent diffusion model that incorporates U-Net- like long skip connections.
DiffiT (Hatamizadeh al., 2024) proposes time-dependent multi-head head self-attention mecha- nism for enhancing the efficiency of transformer-base image diffusion models.
MDTv2 (Gao et al., 2023) proposes an asymmetric encoder-decoder scheme for efficient training of diffusion-based transformer. They also apply U-Net-like long-shortcuts in the encoder and dense input-shortcuts the decoder.
MaskDiT (Zheng et al., 2024) proposes an asymmetric encoder-decodee scheme for efficient training of diffusion transformers, where they train the model with an auxiliary mask reconstruc- tion task similar to MAE (He et al., 2022).
SD-DiT (Zhu al., 2024) extends MaskdiT architecture but incorporates self-supervised dis- crimination objective using momentum encoder.
DiT (Peebles & Xie, 2023) proposes a pure transformer backbone for training diffusion models based on proposing AdaIN-zero modules.
SIT (Ma et al., 2024a) extensively analyzes how DiT training can be efficient by moving from discrete diffusion continuous flow-based modeling.
hhttps://gt thub. www.bcco//opeaata ded--
ded-di fff.ffffuu wwwssioottrr///mmeeenm i onn
24

---

Published as a conference paper at ICLR 2025
G DETAILED QUANTITATIVE RESULTS
We provide evaluation results of different SiT models trained with REPA. All models are aligned with DINOv2-B representations with 0.5 and negative cosine similarity. We use the 4th layer hidden states for the base model and use the 8th layer hidden states for the large and xlarge model.
Table 8: Detailed evaluation results with different model sizes, classifier-free guidance.
All results are reported without
Model
#Params
Iter.
FID
sFID
ISt
Prec.t
Rec.1
Acc.1
SIT-B/2 (Ma et al., 2024a)
130M
400K
33.0
6.46
43.7
0.53
0.63
N/A
REPA (ours)
130M
50K
78.2
11.71
17.1
0.33
0.48
43.2
REPA (ours)
130M
100K
49.5
27.5
0.46
0.59
50,9
REPA (ours)
130M
200K
33.2
6.68
43.7
0.54
0.63
50,9
REPA (ours)
130M
400K
24.4
6.40
59.9
0.59
0.65
61.2
SiT-L/2 (Ma et al., 2024a)
458M
400K
18.8
5.29
72.0
0.64
0.64
N/A
REPA (ours)
458M
50K
55.4
24.0
23.0
0,43
0.53
55.3
REPA (ours)
458M
100K
24.1
6.25
55.7
0.62
0.60
61.8
REPA (ours)
458M
200K
14.0
5.18
86.5
0.67
0.64
66.3
REPA (ours)
458M
400K
10.0
5.20
109.2
0.69
0.65
69.4
SiT-XL/2 (Ma ct al., 2024a)
675M
7M
8.3
6.32
131.7
0.68
0.67
N/A
REPA (ours)
675M
50K
52.3
31.24
24.3
0.45
0.53
56.1
REPA (ours)
675M
100K
19.4
6.06
67.4
0,64
0.61
62.9
REPA (ours)
675M
200K
11.1
5.05
100.4
0.69
0.64
67.3
REPA (ours)
675M
400K
7.9
5.06
122.6
0.70
0.65
70.3
REPA (ours)
675M
4M
5.9
5.73
157.8
0.70
0.69
74.6
We also provide SiT-XL/2+REPA at 4M iteration with classifier-free guidance with different class- free guidance scales.
Table 9: Detailed evaluation results of SiT-XL+REPA at 4M iteration with different classifier free guidance scale
Model
#Params
Iter.
FID,
sFID
ISt
Prec.
Rec.
SYT-XL/2 (Ma et al., 2024a)
675M
7M
1.500
2.06
4.50
270.3
0.82
0.59
REPA (ours)
675M
4M
1.300
1.80
4.55
268.6
0.80
0.63
REPA (ours)
675M
4M
1.325
1.79
4.51
276.8
0.81
0.62
REPA (ours)
675M
4M
1.350
1.80
4.50
284.0
0.81
0.61
REPA (ours)
675M
4M
1.375
1.84
4.48
291.7
0.82
0.61
REPA (ours)
675M
4M
1.400
1.90
4.48
297.5
0.82
0.60
Moreover, we provide the results with the guidance interval (Kynkäääniemi et al., 2024).
Model
#Params
Iter.
Interval
FIDJ
sFIDJ
ISt
Prec.
Rec.
SIT-XL/2 (Ma al., 2024a)
675M
[0. 1]
1.50
2.06
4.50
270.3
0.82
+ REPA (ours)
675M
4M
00.0.8]
2.00
2.23
4.40
360.9
0.84
+ REPA (ours)
675M
[0. 0.75]
2.00
1.78
4.50
346.2
0.82
REPA (ours)
675M
0.71
2.00
1.48
4.67
324.0
0.82
+ REPA (ours)
675M
4M
[0, 0.651
2.00
1.44
4.88
308.8
0.79
0.65
+ REPA (ours)
675M
4M
0.6]
2.00
1.56
5.11
290.7
0.78
0.66
+ REPA (ours)
675M
4M
[0, 0.7]
1.90
1.45
4.68
317.6
0.80
0.64
+ REPA (ours)
675M
4M
[0, 0.7]
1.80
1.42
4.70
305.7
0.80
0.64
Table 10: Detailed evaluation results of SiT-LL+REPA at 4M iteration with different classifier-free guidance scale w. We apply the guidance interval (Kynkäääniemi et al., 2024).
25

---

Published as a conference paper at ICLR 2025
Finally, in Figure 11, provide plot similar Figure 1 using SiT-XL/2 but FIDs obtained with a classifier guidance scale w = 1.35. Similar to Figure 1, REPA provides great speedup and performance improvement compared with the vanilla model.
70.0
SiT XL/2
35.0
SIT-XL/2 REPA
FID-50K
15.0
10.0
6.29
2.43
1.8
400K
Training Iteration
Figure 11: Training iteration vs. FID plot. All values are measured using a classifier free guidance. REPA demonstrates notable speedup and enhanced performance.
26

---

Published as a conference paper at ICLR 2025
H MORE QUALITATIVE RESULTS
Figure 12: Uncurated generation results SiT-XT./2 REPA. We use classifier-free guidance with w 4.0. Class label "loggerhead sea turtle" (33).
Figure 13: Uncurated generation results SIT-XL/2 + REPA. We use classifier-fres free guidance
with w 4.0. Class label "macaw" (88).
27

---

Published as a conference paper at ICLR 2025
Figure 14: Uncurated generation results of SIT-XL/2 + REPA. We use classifier free guidance
with w 4.0. Class label "sulphur-crested cockatoo" (89).
Figure 15: Uncurated generation results SIT-XL/2 + REPA. We use classifier-free guidance
with w 4.0. Class label "golden retriever" (207).
28

---

Published as a conference paper at ICLR 2025
Figure 16: Uncurated generation results of SIT-XL/2 + REPA. We use classifier free guidance
with w 4.0. Class label "husky" (250).
Figure 17: Uncurated generation results SiT-XL/2 REPA. We use classifier-free guidance
with w 4.0. Class label "aretic wolf" (270).
29

---

Published as a conference paper at ICLR 2025
Figure 18: Uncurated generation results of SIT-XL/2 REPA. We use classifier-free guidance
with w 4.0. Class label "arctic fox" (279).
Figure 19: Uncurated generation results SIT-XL/2 REPA. We use classifier-free free guidance
with w 4.0. Class label "lion" (291).
30

---

Published as a conference paper at ICLR 2025
Figure 20: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free free guidance
with w 4.0. Class label "otter" (360).
Figure 21: Uncurated generation results SIT-XL/2 + REPA. We use classifier-free free guidance
with w 4.0. Class label "red panda" (387).
31

---

Published as a conference paper at ICLR 2025
Figure 22: Uncurated generation results of SIT-XL/2 + REPA. We use classifier free guidance
with w 4.0. Class label "panda" (388).
Mare
Figure 23: Uncurated generation results SIT-XL/2 + REPA. We use classifier-free guidance
with w 4.0. Class label "acoustic guitar" (402).
32

---

Published as a conference paper at ICLR 2025
flu
Figure 24: Uncurated generation results SIT-XL/2 + REPA. We use classifier-free guidance with w 4.0. Class label "balloon" (417).
HUP WANNE
Figure 25: Uncurated generation results SIT-XL/2 REPA. We use classifier-free guidance with w 4.0. Class label "baseball" (429).
33

---

Published as a conference paper at ICLR 2025
Figure 26: Uncurated generation results of SIT-XL/2 + REPA. We use classifier free guidance
with w 4.0. Class label "dog sled" (537).
Figure 27: Uncurated generation results SIT-XL/2 REPA. We use classifier-fret guidance
with w 4.0. Class label "fire truck" (555).
34

---

Published as a conference paper at ICLR 2025
Figure 28: Uncurated generation results of SIT-XL/2 + REPA. We use classifier free guidance
with w 4.0. Class label "laptop" (620).
Figure 29: Uncurated generation results SIT-XL/2 + REPA. We use classifier- free guidance
with w 4.0. Class label "space shuttle" (812).
35

---

Published as a conference paper at ICLR 2025
Figure 30: Uncurated generation results of SiT-XL/2 REPA. We use classifier-free guidance
with w 4.0. Class label "ice cream" (928).
Figure 31: Uncurated generation results SIT-XL/2 REPA. We use classifier-free guidance
with w 4.0. Class label "cheeseburger" (933).
36

---

Published as a conference paper at ICLR 2025
Figure 32: Uncurated generation results of SIT-XL/2 + REPA. We use classifier free guidance
with w 4.0. Class label "cliff drop-off" (972).
Figure 33: Uncurated generation results SIT-XL/2 REPA. We use classifier-free free guidance with w 4.0. Class label "coral reef" (973).
37

---

Published as a conference paper at ICLR 2025
Figure 34: Uncurated generation results of SIT-XL/2 + REPA. We use classifier-fre free guidance
with w 4.0. Class label "lake shore" (975).
Figure 35: Uncurated generation results SIT-XL/2 + REPA. We use classifier-free guidance
with w 4.0. Class label "volcano" (980).
38

---

Published as a conference paper at ICLR 2025
1 MORE DISCUSSION ON RELATED WORK
Pretrained visual encoders for generative models. First, there have been several approaches in generative adversarial network (GAN; Goodfellow et al. 2014) that try accelerate training with better convergence using pretrained visual encoders (Sauer et al., 2021; Kumari et al., 2022; Sauer 2022; 2023a; Kang et al., 2023). They usually use pretrained visual encoders as a discriminator by leveraging their intermediate features. This approach has also been applied to the distillation of diffusion models with adversarial objectives (Sauer al., 2023b; 2024; Kang et al., 2024). Another line of work tries to exploit the pretrained visual encoders for improving diffusion model training from scratch (Pernias et al., 2024; Li al., 2024), usually by training two diffusion models where one model generates the pretrained representations and the other model generates the target data conditioned on the generated representation. Our method also tries to improve diffusion model training through pretrained visual encoders, but motivation is in the alignment between the diffusion model representation and recent self-supervised visual representations.
Denoising transformers. Many recent works have tried to use transformer backbones for diffusion or flow-based model training. First, several works like U-ViT (Bao et al., 2023), MDT (Gao et al., 2023), and DiffIT (Hatamizadeh et al., 2024) show transformer- tions can be an effective backbone for training diffusion models. Intriguingly, DiT (Peebles & Xie, 2023) show skip connections are not even necessary components, and a pure transformer architec- ture can be a scalable architecture for training diffusion-based models. Based on DiT, SiT (Ma et al., 2024a) shows the model can be further improved with continuous stochastic interpolants (Albergo et al., 2023). Moreover, VDT (Lu et al., 2024) and Latte (Ma et al., 2024b) show DiTs can be extended for video generation through a space-time attention (Arnab et al., 2021). Based on these improvements, Pixart-o: (Chen et al., 2024b), Pixart-E (Chen et al., 2024a), Stable diffusion 3 (Esser et al., 2024) show pure transformers can be scaled up for challenging text-to-image generation, and CMD (Yu et al., 2024), WALT (Gupta et al., 2024), and Sora (Brooks al., 2024) demonstrates their success in text-to-video generation. Our work analyzes and improves the training of DiT (and SiT) architecture based simple feature matching regularization to the early layers.
backbones with skip connec-
Generative models with auxiliary self-supervised tasks. MaskDiT (Zheng et al., 2024) combines mask reconstruction in MAE (He et al., 2022) diffusion model training for faster diffusion model training. Similarly, SD-DiT (Zhu et al., 2024) shows diffusion model training can be improved with auxiliary discriminative self-supervised loss. MAGE (Li et al., 2023c) bridge MAE training and masked image modeling (Chang et al., 2022) by adjusting the masking ratio in training, which leads to a single model both capable of discrimination and generation tasks. Our method also has a similarity to these works, where our training scheme has an additional distillation loss to projection of diffusion transformer hidden states.
Denoising as self-supervised learning task. There have been some attempts to improve self- supervised learning with denoising: Abstreiter et al. (2021) extends the diffusion objective for a better representation learning scheme, and Chen et al. (2024c) deconstructs diffusion models to improve denoising-basee
representation learning. Hudson et al. (2024) introduces an encoder that compact latent vector. Zaidi et al.
learns representation by guiding diffussin with its output (2023) focuses on the molccular domain and proposes a pretraining scheme based on denoising. Our work also analyzes representational gap between popular self-supervised networks and those learned from pretrained diffusion models.
39

---

Published as a conference paper at ICLR 2025
J IMAGENET 512x 512 EXPERIMENT
To validate the scalability of REPA with respect to the input image resolution, we conduct an ad- ditional experiment ImageNet 512x 512. We strictly follow the setup used in ImageNet 256x256 experiment except an input dimension. Specifically, the input dimension for SiT becomes 64 64 4 compressed latent image from 512 512 image pixels using stable diffusion VAE (Rombach et al., 2022). Moreover, we use an image resized to 448 x 448 resolution for an input to the DINOv2 (Oquab et al., 2024) encoder with an interpolation of the positional embedding.
We provide quantitative results in Table 11 and qualitative result in Figure 36. 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 Prec) using >3x fewer training iterations.
Figure 36: Samples on ImageNet 512x512 from SiT-XL/2+REPA using CFG with w = 4.0.
Table 11: System-level comparison on ImageNet 512x512. We use CFG with = 1.35.
Model
Epochs
FID4
sFID
ISt
Pre.f
Rec.t
Pixel diffassion
VDM++
2.65
278.1
ADM-G. ADM-U
400
2.85
5,86
221.7
0.84
0.53
Simple diffusion (U-Net)
800
4.28
171.0
Simple diffusion (U-VIT, L)
800
4.53
205.3
Latent diffusiio, Transformer
MaskDiT
800
2.50
5.10
256.3
0.83
0.56
DIT-XL/2
600
3.04
5.02
240.8
0.84
0.54
SiT-XL/2
600
2.62
4.18
252.2
0.84
0.57
+ REPA (ours)
80
2.44
4.21
247.3
0.84
0.56
+ REPA (ours)
100
2.32
4.16
255.7
0.84
0,56
+ REPA (ours)
200
2.08
4.19
274.6
0.83
0.58
40

---

Published as a conference paper at ICLR 2025
K TEXT-TO-IMAGE GENERATION EXPERIMENT
We also validate REPA in exx-tto-image generation. We mostly follow the experimental setup used in U-ViT (Bao et al., 2023) unless otherwise specified: we train the model from scratch train split of the MS-COCO dataset (Lin et al., 2014) and use a validation split for evaluation. We use MMDiT (Esser al., 2024), a simple variant of DiT that design attention layers to be jointly computed with image patches and text embeddings. We train MMDiT models for 150K iterations with a batch size of 256. We set a hidden dimension as 768 and a model depth as 24, and we use the CLIP (Radford et al., 2021) text encoder to compute text prompts from captions.
We report the results in Table 12 and Figure 37. First, as shown in the qualitative comparison in Figure 37, REPA shows consistently better results than the vanilla model. Moreover, as shown in Table 12, REPA also shows considerable improvements in T2I generation, highlighting the impor- tance of alignment of visual representations even under the presence of text representations.
In this respect, we strongly believe that training large-scale T2I models with large-scale data using REPA will be promising direction in the future.
A small kitchen with
bus pulls over the curb close intersection.
driving city with traffie signs.
living with television table.
A coming tracks.
baseball player
MNCN
MMDIT (ours)
A zebra standing the pond.
of skiers
Some fond is cnoking small
A close a plate of broccsl
A lable satting of colorful food.
group paropia walking and hoiding
MMOIT
MMDIT-REPA (ours)
Figure 37: Qualitative comparison text-to-image generation (MS-COCO). We use classifier-
free guidance with w = 4.0.
41

---

Published as a conference paper at ICLR 2025
Method
Type
FID
AttnGAN (Xu et al., 2018)
GAN
35.49
DM-GAN (Zhu et al., 2019)
GAN
32.64
VQ-Diffusion (Gu et al., 2022)
Discrete Diffusion
19.75
DE-GAN (Tao et al., 2022)
GAN
19.32
XMC-GAN (Zhang ct al., 2021)
GAN
9.33
Frido (Fan al., 2023)
Diffusion
8.97
LAFITE (Zhou et al., 2021)
GAN
8.12
U-Net (Bao et al., 2023)
Diffusion
7.32
U-ViT-S/2 (Bao et al., 2023)
Diffusion
5.95
U-ViT-S/2 (Deep) (Bao al., 2023)
Diffusion
5.48
MMDiT (ODE; NFE=50)
Diffusion
6.05
MMDiT+REPA (ODE; NFE=50)
Diffusion
4.73
MMDIT (SDE; NFE-250)
Diffusion
5,30
MMDIT+REPA (SDE; NFE=250)
Diffusion
4.14
Table 12: Quantitative comparison rext-to-image generation (MS-COCO). We use classifier guidance with w 2.0 following the setup in (Bao et al., 2023).
L FEATURE MAP VISUALIZATION
We provide PCA visualizations of feature map, similar to those in DINOv2 (Oquab et al., 2024). As shown in Figure 38, REPA shows coarse-to-fine feature maps, while the vanilla model tends show noisy feature map particular for large I.
W: REPA SiT-XU/2
Layer
Layer
Layer
Layer 4
Layer
Layer
Layer
Layer 14
Layer 21
Layer 28
(a) 0.5
W/ RAPA SIT-XL/2
Layer
Layer
Layer
Layer 4
Layer
Layer
Layer
Layer
Layer 21
Layer 28
(b) 0.7
W/PEPA SIT-XU22
Layer
Layer
Layer
Layer
Layer
Layer
Layer 7
Layer Layer 21
Layer
(c) 0.9
Figure 38: PCA visualization of layer-wise features of SiT-XL/2 and SiT-XL/++REPA.
42

---

Published as a conference paper at ICLR 2025
M LIMITATIONS AND FUTURE WORK
In this section, we enumerate several possible future research directions in what follows.
Alignment depth. Recall that empirically showed that applying REPA to layer 8 is more bene- ficial than to later transformer layer embeddings (see Table 2); conducting extensive analysis of this result will be interesting direction further improve REPA.
Different input data types. 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 an interesting future work. Moreover, based on our text-to-image generation results on MS-COCO, training large-scale text-to-image diffusion models with REPA will also be interesting direction.
Theoretical analysis. Exploring theoretical insights into why REPA works well will also be an exciting future direction. For instance, will be interesting to explore the relationship between representations learned with an instance discrimination objective and denoising objective.
Time-varying REPA. We think one of the interesting possible directions can be designing a weight function based on a noise schedule used in the diffusion process. We have not explored this in this work 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 future work.
43


Reviews:
[{"review_id": "kIhFzftsQH", "reviewer": "Reviewer_XnZq", "summary": "The paper presents a novel argument that a primary bottleneck in training large-scale diffusion models for generation lies in learning effective representations. The authors suggest that training can be significantly simplified when the model is supported by strong external visual representations. To address this, they introduce a straightforward regularization technique called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from pretrained external visual encoders, such as DINOv2. Comprehensive experiments validate this perspective, reinforcing the paper’s innovative viewpoint.", "questions": "1. I recommend releasing the training and evaluation codes to support future research efforts, as feature representation plays a crucial role in diffusion models.\n2. You mention that directly fine-tuning a DINOv2 model within a diffusion framework fails due to the noisy input. This raises some questions, as fine-tuning is generally quite robust and should be capable of handling noise. Additionally, have you considered training a diffusion model from scratch while incorporating a partial DINOv2-based loss (similar to a LPIPS loss) alongside the denoising loss? This approach might aid in feature preservation during denoising.\n3. How do you handle cases where the spatial dimensions (height × width) differ between DINOv2 and DIT? Would using cross-attention mechanisms be a feasible method for aligning features in such cases? This might offer a flexible solution to manage discrepancies in spatial resolution between the models.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The paper’s key finding—that self-supervised representations, like those from DINOv2, can serve as a teacher to distill base features in a diffusion model—is both novel and insightful. By using these high-quality representations, the diffusion model can more quickly learn a robust semantic foundation, allowing deeper layers to focus on higher-frequency details. This approach is innovative and inspiring.\n2. The experimental setup is thorough, covering various aspects such as types of self-supervised encoders, encoder scalability, the number and location of layers used for regularization, and different alignment objectives. These comprehensive experiments add credibility to the proposed method.", "weaknesses": "1. High-resolution results (e.g., 512x512) are needed to support and further validate findings.\n2. FID may be unreliable here, as it relies on the Inception network, which tends to lose significant spatial information and focusing on class information. I recommend re-evaluating using DINO-FID to obtain a more accurate measure. I would consider raising the score to 10 if state-of-the-art fid is demonstrated at resolution 512x512 and DINO-FID at 256x256. For DINO-FID, it’s necessary to pass the ImageNet training set through DINOv2 to store features for dino-FID calculation, which might be challenging to implement. If resources are limited, DINO-FID can be computed at 256x256 resolution, while 512x512 results can use the Inception FID npz provided by ADM.\n3. It would be beneficial to include visualizations of the features before and after alignment. Although differences can already be observed in the generated images, since the core concept is feature alignment, visualizing the impact directly at the feature level would strengthen the argument and provide deeper insights.\n4. I recommend exploring the effects of different diffusion objectives. For example, the intermediate features of an epsilon-prediction diffusion model might differ from those of a v-prediction or x0-prediction model. An analysis of these variations could provide valuable insights into the interaction between diffusion objectives and feature alignment.\n5. The use of class conditions may not fully represent the entire generative domain, particularly in areas like text-to-image and text-to-video generation. This raises concerns about the generalization of the claim that feature alignment is universally important for generative tasks. It would be beneficial to clarify or qualify this claim in light of these broader application areas. In scenarios that already include rich semantic prompts, such as text-to-video generation, would the proposed method still be effective? DINO’s semantic extraction during training may differ from the semantics derived by an LLM when captioning or re-captioning. It would be helpful to include experiments that use detailed textual descriptions as conditions to evaluate the method’s performance under these circumstances.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ne7Uq2zn2q", "reviewer": "Reviewer_U4Y4", "summary": "This paper presents REPresentation Alignment (REPA), a regularization technique designed to enhance the efficiency and quality of training in diffusion models. By aligning the noisy input states within denoising networks with clean image representations taken from pretrained visual encoders, REPA helps diffusion models learn better internal representations. The approach significantly improves training efficiency—such as a 17.5× speed-up for SiT models—and boosts generation quality, achieving a FID score of 1.80 on ImageNet with guidance. The authors emphasize the importance of aligning diffusion model representations with powerful self-supervised representations to make training easier and more effective.", "questions": "I believe this is a solid piece of work and an important first step toward combining representation learning and diffusion model training. My main concern is whether this approach is generalizable to text-conditioned diffusion models. Additionally, I would suggest that the authors add [a] and [b] to the literature review:\n\n[a] Learning Disentangled Representation by Exploiting Pretrained Generative Models: A Contrastive Learning View. ICLR 2022.\n\n[b] DisDiff: Unsupervised Disentanglement of Diffusion Probabilistic Models. NeurIPS 2023.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- I like the idea of combining diffusion model training and representation learning. The authors highlight the importance of aligning diffusion model representations with powerful self-supervised representations to make training easier and more effective.\n\n- The writing in the paper is clean and easy to follow. The authors start with their observations from the empirical study on pretrained diffusion models, which provides the audience with enough background to understand the problem.\n\n- The experimental section of this paper is solid. They conduct extensive ablation studies on the proposed method, including aspects such as encoder type, encoder size, alignment depth, and more.", "weaknesses": "- My main concern with this work is that the authors focus on class-conditioned diffusion models rather than text-conditioned diffusion models. The success of diffusion heavily depends on text prompts, as better prompts lead to improved performance and faster training of the diffusion model. I am worried that representation learning may not be necessary for text-conditioned diffusion models, as the text itself already acts as a regularizer to guide the training. Additionally, the improvement in FID with classifier-free guidance (CFG) in this paper is not significant, which partially validates my concern that additional representation learning may not be needed when there is strong guidance.\n\n- The empirical study presented in Figure 2 is based on linear probing, which is unfair to the diffusion model as it is trained with reconstruction objective. In this context, I suspect that the performance gap for the diffusion model would be smaller on tasks like semantic segmentation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sUSKrkBpY2", "reviewer": "Reviewer_K7ig", "summary": "The paper introduces representation alignment for diffusion models, a new technique that aligns diffusion features with those of a pretrained foundation model during training. The method is motivated by the observation that the features of pretrained diffusion models inherently align with the features of pretrained, diffusion-based transformer models. Inspired by this, the authors conducted a thorough analysis of the correlation between features at different layers of a foundation model and a diffusion model, finding that certain layers exhibited higher feature correlation. To enhance this alignment, the authors introduced a representation alignment loss, which results in a performance boost. An extensive study was conducted with various foundation models and across different diffusion model layers to identify the most suitable model and layer for enforcing similarity. The authors achieved 17.5x faster convergence on DiT when trained with RePA.", "questions": "1. 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 was performed at different layer? This aspect is not clear from the paper. Could the authors please include this clarification in the figure caption or the related section. \n2. In cross attention layers, the main learning process happens through alignment of the image space and text space . Diffusion models leveraging cross attention functions better if the  cross attention is applied at different layers. Considering a contrast of REPA for diffusion alignment. Would doing representation alignment based training leveraging different levels of features from the foundation model and the diffusion model lead to an improve in performance? Could the authors comment on this.\n3. Diffusion features are generally noisy towards the start of the sampling process and becomes clearer towards the end. In Figure 2a, the authors showed that the semantic gap for cleaner features tends to be lower. However does this also suggest that a time varying repa might lead to a even better boost in performance\n4. In Figure 5, the authors have shown the benefit of different visual foundation models for represenation alignment. However, there is also a correlation that the linear probing performance also follows the same trend i.e \nDinov2> CLIP>MoCov3>DINO>MAE\nHence I'm curious if repa using a pre-trained SOTA VIT based classifier may further boost the performance. An experiment for a limited number of training iterations or a detailed explanation would suffice for this issue", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-motivated. The introduction section clearly states the logic and reasoning behind why and how the features of foundation models can be leveraged to improve the convergence of a diffusion model.\n2. The authors analyze different visual foundation models for representation learning and demonstrate that RePA is invariant to the choice of foundation model, with almost all visual foundation models enhancing generation performance and convergence.\n3. The authors also study the correlation between generation performance and discriminative classification performance using validation accuracy on ImageNet, showing that both exhibit a linear correlation. A model trained to provide better representation capabilities automatically achieves better generation quality as well.\n4. These insights are extremely valuable and novel, representing a significant step forward in the field.", "weaknesses": "1. The main claims of the paper is valid in the cases where classifier free guidance is not applied. Could the authors provide an analysis of the benefit of RePA in the presence of classifier free guidance. Specifically it would be great to see the results in Figure 1 with classifier free guidance included\n2. How well does RePA work when there is already an alignment with text like cross attention in text-to-image models. Will RePA cause a conflict in representation alignment between CLIP/T5 and Dinov2 representations and reduce performance? If time permits I would request the authors to add an analysis in this aspect. This could be a simple experiment like finetuning the existing model an small text to image datasets like CUB-200[3]. If not a detailed explanation of why/why not this may work would suffice.\n3. Diffusion features are generally noisy in the encoder layers and tends to become clearer in the decoder[1,2]. However in REPA we find that the best suited layer is indeed in the encoder which suggests cleaner representations in the encoder itself. Could the authors comment on why although the correlations seems better in the decoder features, the encoder layer 8 is best suited for representation alignment? An answer to Q1 in the questions section would suffice for this query.\n4. There are no limitations section or future works section in the manuscript. I would advise the authors to please include these as well in an updated version. The request for these sections is for enabling the research community to further tackle problems that the authors might not have time to tackle. \n\n\n   [1] Tumanyan, N., Geyer, M., Bagon, S. and Dekel, T., 2023. Plug-and-play diffusion features for text-driven image-to-image translation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 1921-1930).\n\n   [2] Cao, M., Wang, X., Qi, Z., Shan, Y., Qie, X. and Zheng, Y., 2023. Masactrl: Tuning-free mutual self-attention control for consistent image synthesis and editing. In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 22560-22570).\n\n    [3] Wah, C., Branson, S., Welinder, P., Perona, P. and Belongie, S., 2011. The caltech-ucsd birds-200-2011 dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BDOiQSAKDd", "reviewer": "Reviewer_vfsY", "summary": "This paper presents a technique called REPA, which speeds up the process of training diffusion transformers. The high-level idea of REPA is to align the diffusion model representation with the visual representation of pretrained models. Experimental results show that after using REPA, the training process of diffusion transformers can be much faster, and also achieve better performance when training for the same number of steps.", "questions": "-- I understand that this paper mainly focuses on speeding up diffusion transformers, as mentioned in the title. However, I am curious about why the application scenario is only limited in diffusion transformers instead of more general diffusion models. Is it because pretrained visual encoders like DINOv2 are mainly in vision transformer architecture, so enforcing the alignment of the features make more sense? Or the authors think diffusion transformers are the future trend for developing diffusion models, so studying the training dynamics of diffusion transformers is more meaningful?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "++ The idea flow of this paper is very clear. First, the paper presents the semantic gap between diffusion transformer features and discriminative features. Then, the paper shows that the quality of the feature representations from diffusion transformers are correlated with the generation quality. Finally, the proposed method REPA is an intuitive solution based the previous observations and arguments. Therefore, the proposed technique has clear and reasonable motivations.\n\n++ The benefit of applying REPA in training diffusion transformers seem to be significant to me. Experimental results like the quantitative ones from Tables 2-4 show that REPA helps the models reach the same performance with fewer iterations, and further achieve better generation results when trained for the same number of steps. Therefore, this technique can be useful in training diffusion models.\n\n++ The proposed solution is simple but effective, just to align the features from diffusion models with pretrained visual foundation models. Therefore, it is a neat yet efficient solution for speeding up diffusion models.", "weaknesses": "-- The quantitative evaluations of REPA on DiT/SiT are only conducted on $256\\times 256$ images, which is kind of low-resolution with respect to the capability of state-of-the-art diffusion models. I am wondering how REPA performs in higher-resolution cases, like $512\\times 512$ which is also be supported by DiT/SiT. Is it still useful in training higher-resolution diffusion transformers? I think this evaluation is crucial, otherwise the application scenario of this technique will be relatively restricted to lower-resolution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JDPzwyOJMj", "reviewer": "Reviewer_i9uk", "summary": "This paper proposes a simple method to improve the diffusion-based generative model’s training efficiency. It discovers the semantic difference in self-supervised vision encoders and diffusion generative models, and the motivation is to use pre-trained vision encoder representation to regularize the diffusion transformer’s feature representation. The paper is well-written, the motivation, metrics, and methods are easy to understand, and the experiments conducted on various pre-trained vision encoders under different settings demonstrate the effectiveness of the proposed method in both generation quality measured in FID and representation ability measured in linear-probing accuracy.", "questions": "--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?\n\n-- The same techniques can be applied to the video generative model, where training cost/efficiency is more important, do you think the proposed method can be transferred to the video domain?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "--the proposed method is simple and very easy to implement, also the idea is well-motivated by the semantic measurement metrics\n\n--the experiments are strong and complete, and examine various vision encoders under various setups (model size, depth) and with different DiT model sizes, all show promising improvement in terms of FID", "weaknesses": "-- the paper only shows results under resolution 256x256, which is the scale most of the vision encoders are pretrained on, it raises the concern if the pre-trained low-res representation can transfer to the generative models when train on higher resolution 512, 1024, 2048 etc\n\n-- missing dataset analysis on the pre-trained vision encoder, i.e. is the performance of the encoder pre-trained on ImageNet the same as on in-the-wild dataset?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZXxUUVp7T4", "reviewer": "Reviewer_MZYg", "summary": "This paper finds that learning a good representation in diffusion transformers will speed up the training process and achieve better performance. By simply aligning diffusion transformer representations with recent self-supervised representations (DINOv2), the DDPM-based models (DiT) and flow-based models (SiT) show significant FID metric drops in early training.", "questions": "1. The experimental results indicate that the speedup for SiT training is quite significant. However, can it achieve state-of-the-art performance in the final results? MDT (ICCV'23) [1] seems to have achieved an FID of 1.79.\n2. Let's discuss the gap between FID metrics and actual image performance, in Figure 6, your visualization are under classifier-free guidance with $w = 4.0$, but SiT-XL/2 with REPA achieves best FID of 1.80 under classifier-free guidance scale of $w = 1.35$. What do you think of the difference w between best visualization and best metrics, and whether this measure is unreasonable?\n3. I agree with \"the diffusion transformer model needs some capacity for capturing high-frequency details after learning a meaningful representation that contains good semantics.\", but minimum semantic gap are at layer 20 before alignment with REPA, which is different between layer 8 after alignment with REPA  (I note this because layer 8 yields the best metrics). Could you further explain this difference?\n4. To more clearly represent the process of alignment in the representation, I recommend including the definition of layers in formula (8). I suggest using $h^{[n,m]}$ to represent diffusion transformer encoder output, where $n$ is a patch index and $m$ is layer index. There is a small parenthesis error in formula (8).\n\n[1] Gao, Shanghua, et al. \"Masked diffusion transformer is a strong image synthesizer.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper astutely captures the feature representation gap between generative and discriminative models, providing a feasible approach to enhance model representation and accelerate the training process of DiT models. The extensive analysis and observations regarding representation clarify the underlying motivation. Additionally, the experiments are detailed and validate the effectiveness of the proposed method across different settings. The presentation of Table 2 and the experiment section is very clear.", "weaknesses": "1. All the experiments in this paper are conducted under the setting of ImageNet 256x256. Although this work shows good results in this setting, it is still far from being applied to the currently popular high-resolution transformer-based text-to-image generation models, such as Flux. I would be more excited to see results and observations at higher resolutions or in text-to-image generation. If the generation is limited to the DiT/SiT class condition, its applications will be severely restricted.\n2. Metrics for ImageNet generation often confuse me; the FID, sFID, IS, and other metrics frequently show differing trends. And FID does not effectively reflect actual quality differences at low values. How do you consider these metrics comprehensively to make a reasonable judgment about performance? For instance, in the last two rows of Table 2, NT-Xent outperforms in sFID, while Cos. sim. outperforms in FID and IS. And it would be beneficial to provide some experimental results to explain why you chose NT-Xent, rather than simply stating \"Empirically, ...\" in line 421. I would be glad to discuss this further with you in the Questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)