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

## Abstract

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.

## 1 Introduction

Generative models based on *denoising*, such as diffusion models (Ho et al., ; Song et al., ) and flow-based models (Albergo & Vanden-Eijnden, ; Lipman et al., ; Liu et al., ), have been a scalable approach in generating high-dimensional visual data. They achieve remarkably successful results in challenging tasks such as zero-shot text-to-image (Podell et al., ; Saharia et al., ; Esser et al., ) or text-to-video (Polyak et al., ; Brooks et al., ) generation.

Recent works have explored the use of diffusion models as representation learners (Li et al., ; Xiang et al., ; Chen et al., ; Mukhopadhyay et al., ) and have shown that they learn discriminative features in their hidden states, and better diffusion models learn better representations (Xiang et al., ). In fact, this observation is closely related to earlier approaches that employ *denoising score matching* (Vincent, ) as a self-supervised learning method (Bengio et al., ), which implicitly learns a representation ${\mathbf{h}}$ as a hidden state of a denoising autoencoder ${\mathbf{s}}_{\theta}(\tilde{{\mathbf{x}}})$ through a *reconstruction* of ${\mathbf{x}}$ from the corrupted data $\tilde{{\mathbf{x}}}$ (Yang & Wang, ). However, the reconstruction task may not be a suitable task for learning good representations, as it is not capable of eliminating unnecessary details in ${\mathbf{x}}$ for representation learning (LeCun, ; Assran et al., ).

Figure: (a) Semantic gap: Linear probing
Refer to caption: /html/2410.06940/assets/x3.png

Our approach.
In this paper, we identify that the main challenge in training diffusion models stems from the need to learn a high-quality internal representation ${\mathbf{h}}$. We demonstrate that the training process for generative diffusion models becomes significantly easier and more effective when supported by an external representation, ${\mathbf{y}}_{\ast}$. Specifically, we propose a simple regularization technique that leverages recent advances in self-supervised visual representations as ${\mathbf{y}}_{\ast}$, 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, ; Ma et al., ) and the state-of-the-art self-supervised vision model, DINOv2 (Oquab et al., ).
Similar to prior studies (Xiang et al., ), we first observe that pretrained diffusion models do indeed learn meaningful discriminative representations (as shown by the linear probing results in Figure [2(a)](#S1.F2.sf1)). 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 [2(b)](#S1.F2.sf2)) is still considered weak,(^1^11We describe this as “weak” because relatively, the alignments are much poorer than those seen with other self-supervised encoders (*e.g.*, MoCov3 (Chen et al., )), even after extensive training.) which we study by measuring their *representation alignment* (Huh et al., ). Finally, we observe this alignment between diffusion models and DINOv2 improves consistently with longer training and larger models (Figure [2(c)](#S1.F2.sf3)).

These insights inspire us to enhance generative models by incorporating external self-supervised representations. However, this approach is not straightforward when using off-the-shelf self-supervised 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 $\tilde{{\mathbf{x}}}$, whereas most self-supervised learning encoders are trained on clean images ${{\mathbf{x}}}$. This issue is even more pronounced in modern *latent diffusion* models, which take a compressed latent image ${\mathbf{z}}=E({\mathbf{x}})$ from a pretrained VAE encoder (Rombach et al., ) 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 a *regularization* technique that distills pretrained self-supervised representations into diffusion representations, offering a flexible way to integrate high-quality representations.

Figure: (a) Semantic gap: Linear probing
Refer to caption: /html/2410.06940/assets/x6.png

Specifically, we introduce *REPresentation Alignment* (REPA), a simple regularization technique built on recent diffusion transformer architectures (Peebles & Xie, ). In essence, REPA distills the pretrained self-supervised visual representation ${\mathbf{y}}_{\ast}$ of a clean image ${\mathbf{x}}$ into the diffusion transformer representation ${\mathbf{h}}$ of a noisy input $\tilde{{\mathbf{x}}}$. This regularization reduces the semantic gap in the representation ${\mathbf{h}}$ (Figure [3(a)](#S1.F3.sf1)) and better aligns it with the target self-supervised representations ${\mathbf{y}}_{\ast}$ (Figure [3(b)](#S1.F3.sf2)). Notably, this enhanced alignment significantly boosts the *generation* performance of diffusion transformers (Figure [3(c)](#S1.F3.sf3)). Interestingly, with REPA, we observe that sufficient representation 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 high-frequency details based on the aligned representations, further improving generation performance.

Based on our analysis, we conduct a system-level comparison to demonstrate the effectiveness of our scheme by applying it to two recent diffusion transformers: DiTs (Peebles & Xie, ) and SiTs (Ma et al., ). For SiT training, we show the model achieves FID$=$7.9 on class-conditional ImageNet (Deng et al., ) generation only using 400K training iteration (without classifier-free guidance; Ho & Salimans ) which is $>$17.5$\times$ 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-the-art results of FID$=$1.42 with guidance interval (Kynkäänniemi et al., ).

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 a 17.5$\times$ faster training for SiTs 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., ; Ma et al., ). Please refer to a more detailed explanation in Appendix [A](#A1).

We consider a continuous time-dependent process with a data ${\mathbf{x}}_{\ast}\sim p({\mathbf{x}})$ and a Gaussian noise ${\mathbf{\epsilon}}~{}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ on $t\in[0,T]$:

$$ ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{\ast}+\sigma_{t}{\mathbf{\epsilon}},\quad\alpha_{0}=\sigma_{T}=1,\,\,\alpha_{T}=\sigma_{0}=0,$ (1) $$

where $\alpha_{t}$ and $\sigma_{t}$ are a decreasing and increasing function of $t$, respectively. Given such a process, there exists a *probability flow ordinary differential equation* (PF ODE) with a velocity field

$$ $\dot{{\mathbf{x}}}_{t}={\mathbf{v}}({\mathbf{x}}_{t},t),$ (2) $$

where the distribution of this ODE at $t$ is equal to the marginal $p_{t}({\mathbf{x}})$. Thus, data can be sampled by solving this PF ODE in Eq. ([2](#S2.E2)) through existing ODE samplers (*e.g.*, Euler sampler) starting from a random Gaussian noise $\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ (Lipman et al., ; Ma et al., ).

This velocity ${\mathbf{v}}({\mathbf{x}},t)$ is represented as the following sum of two conditional expectations

$$ ${\mathbf{v}}({\mathbf{x}},t)=\mathbb{E}[\dot{{\mathbf{x}}}_{t}|{\mathbf{x}}_{t}={\mathbf{x}}]=\dot{\alpha}_{t}\mathbb{E}[{\mathbf{x}}_{\ast}|{\mathbf{x}}_{t}={\mathbf{x}}]+\dot{\sigma}_{t}\mathbb{E}[{\mathbf{\epsilon}}|{\mathbf{x}}_{t}={\mathbf{x}}],$ (3) $$

which can be approximated with model ${\mathbf{v}}_{\theta}({\mathbf{x}}_{t},t)$ by minimizing the following training objective:

$$ $\mathcal{L}_{\text{velocity}}(\theta)\coloneqq\mathbb{E}_{{\mathbf{x}}_{\ast},\bm{\epsilon},t}\big{[}||{\mathbf{v}}_{\theta}({\mathbf{x}}_{t},t)-\dot{\alpha}_{t}{\mathbf{x}}_{\ast}-\dot{\sigma}_{t}{\mathbf{\epsilon}}||^{2}\big{]}.$ (4) $$

Moreover, there exists a reverse *stochastic differential equation* (SDE) that the marginal $p_{t}({\mathbf{x}})$ coincides with the one of PF ODE in Eq. ([2](#S2.E2)) with a diffusion coefficient $w_{t}$ (Ma et al., ):

$$ $d{\mathbf{x}}_{t}={\mathbf{v}}({\mathbf{x}}_{t},t)dt-\frac{1}{2}w_{t}{\mathbf{s}}({\mathbf{x}}_{t},t)dt+\sqrt{w_{t}}d\bar{\mathbf{w}}_{t},$ (5) $$

where the score ${\mathbf{s}}({\mathbf{x}}_{t},t)$ is the following conditional expectation

$$ ${\mathbf{s}}({\mathbf{x}}_{t},t)=-{\sigma_{t}^{-1}}\mathbb{E}[{\mathbf{\epsilon}}|{\mathbf{x}}_{t}={\mathbf{x}}].$ (6) $$

and it can be directly computed using the velocity ${\mathbf{v}}({\mathbf{x}},t)$ for $t>0$ as

$$ ${\mathbf{s}}({\mathbf{x}},t)={{\sigma_{t}^{-1}}}\cdot\frac{\alpha_{t}{\mathbf{v}}({\mathbf{x}},t)-\dot{\alpha}_{t}{\mathbf{x}}}{\dot{\alpha}_{t}\sigma_{t}-{\alpha}_{t}\dot{\sigma}_{t}},$ (7) $$

implying that data can be alternatively generated through Eq. ([5](#S2.E5)) with SDE solvers.

Following Ma et al. (), we mainly consider a simple linear interpolant with restricting $T=1$: $\alpha_{t}=1-t$ and $\sigma_{t}=t$. However, our approach is applicable to any similar variants (*e.g.*, DDPM; Ho et al. ), which has a similar formulation but uses a discretized process and different $\alpha_{t},\sigma_{t}$ that $\mathcal{N}(\mathbf{0},\mathbf{I})$ becomes an equilibrium distribution (*i.e.*, ${\mathbf{x}}_{t}$ converges to $\mathcal{N}(\mathbf{0},\mathbf{I})$ only if $t\to\infty$).

## 3 REPA: Regularization for Representation Alignment

### 3.1 Overview

Let $p({\mathbf{x}})$ be an unknown target distribution for data ${\mathbf{x}}\in\mathcal{X}$. Our goal is to approximate $p({\mathbf{x}})$ through a model distribution using a dataset drawn from $p({\mathbf{x}})$. To lower computational costs, we adopt the recent prevalent *latent diffusion* (Rombach et al., ). This involves learning a latent distribution $p({\mathbf{z}})$, which is defined as the distribution of a compressed latent variable $\mathbf{z}=E({\mathbf{x}})$, where $E$ is an encoder from a pretrained autoencoder (*e.g.*, KL-VAE; Rombach et al. ), with ${\mathbf{x}}\sim p_{\text{data}}({\mathbf{x}})$.

We aim to learn this distribution by training a diffusion model ${\mathbf{v}}_{\theta}({\mathbf{z}}_{t},t)$ using objectives such as velocity prediction, as described in Section [2](#S2). Here, we revisit denoising score matching within the context of self-supervised representation learning (Bengio et al., ). From this perspective, one can think of the diffusion model ${\mathbf{v}}_{\theta}({\mathbf{z}}_{t},t)$ as a composition of two functions $g_{\theta}\circ f_{\theta}$ with an encoder $f_{\theta}:\mathcal{Z}\to\mathcal{H}$
with $f_{\theta}({\mathbf{z}}_{t})={\mathbf{h}}_{t}$ and a decoder $g_{\theta}:\mathcal{H}\to\mathcal{Z}$ with $g_{\theta}({\mathbf{h}}_{t})={\mathbf{v}}_{t}$, where the encoder $f_{\theta}$ implicitly learns a representation ${\mathbf{h}}_{t}$ that reconstructs the target ${\mathbf{v}}_{t}$.

However, learning a good representation through producing a prediction of the input space (*e.g.*, generating pixels) can be challenging, as the model is often not capable of eliminating unnecessary details, which is crucial for developing a strong representation (LeCun, ; Assran et al., ). We argue that a 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 Alignment* (REPA) using the recent diffusion transformer architectures (Peebles & Xie, ; Ma et al., ) (see Appendix [B](#A2) for an illustration). In a nutshell, our regularization distills pretrained self-supervised visual representations to diffusion transformers in a simple and effective way. This allows the diffusion model to leverage these semantically rich external representations for generation, leading to a substantial boost in performance.

Figure: Figure 4: REPA improves visual scaling. We compare the images generated by two SiT-XL/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-free guidance.
Refer to caption: /html/2410.06940/assets/x9.png

### 3.2 Observations

To take a deeper dive into this, we first investigate the layer-wise behavior of the pretrained SiT model (Ma et al., ) on ImageNet (Deng et al., ), which uses linear interpolants and velocity prediction for training. In particular, we focus on measuring the
*representation gap* between the diffusion transformer and the state-of-the-art self-supervised DINOv2 model (Oquab et al., ). 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. (), which involves linear probing on globally pooled hidden states of the diffusion transformer. Next, to measure *feature alignments*, we use CKNNA (Huh et al., ), a kernel alignment metric related to CKA (Kornblith et al., ), but based on mutual nearest neighbors. This allows for a quantitative assessment of alignment between different representations. We summarize the result in Figure [2](#S1.F2) and more details (*e.g.*, definition of CKNNA) in Appendix [C.1](#A3.SS1).

Diffusion transformers exhibit a significant semantic gap from state-of-the-art visual encoders. As shown in Figure [2(a)](#S1.F2.sf1), we observe that the hidden state representation of the pretrained diffusion transformer, in line with prior works (Xiang et al., ; Chen et al., ), 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-rich representations in order to generate images with high-frequency details.

Diffusion representations are already (weakly) aligned with other visual representations.
In Figure [2(b)](#S1.F2.sf2), we report representational alignments between SiT and DINOv2 using CKNNA. In particular, the SiT model representation already shows better alignment than MAE (He et al., ), which is also a 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., ) *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 [2(c)](#S1.F2.sf3), 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-supervised visual encoders (*e.g.*, MoCov3 and DINOv2), even after extensive training of 7M iterations.

These findings are not unique to the SiT model but are also observed in other denoising-based generative transformers.
For instance, in Figure [2](#S1.F2), we present a similar analysis using a DiT model (Peebles & Xie, ) pretrained on ImageNet with the DDPM objective (Ho et al., ; Nichol & Dhariwal, ). See Appendix [C.2](#A3.SS2) for more details.

### 3.3 Representation alignment with self-supervised representations

REPA aligns patch-wise projections of the model’s hidden states with pretrained self-supervised visual representations. Specifically, we use the *clean* image representation as the target and explore its impact. The goal of this regularization is 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 a pretrained encoder and consider a clean image $\mathbf{x}_{\ast}$. Let ${\mathbf{y}}_{\ast}=f({\mathbf{x}}_{\ast})\in\mathbb{R}^{N\times D}$ be an encoder output, where $N,D>0$ are the number of patches and the embedding dimension of $f$, respectively. REPA aligns $h_{\phi}({\mathbf{h}}_{t})\in\mathbb{R}^{N\times D}$ with ${\mathbf{y}}_{\ast}$, where $h_{\phi}({\mathbf{h}}_{t})$ is a projection of an diffusion transformer encoder output ${\mathbf{h}}_{t}=f_{\theta}({\mathbf{z}}_{t})$ through a trainable projection head $h_{\phi}$. In practice, we simply parameterize $h_{\phi}$ using a multilayer perceptron (MLP).

In particular, REPA achieves alignment through a maximization of patch-wise similarities between the pretrained representation ${\mathbf{y}}_{\ast}$ and the hidden state ${\mathbf{h}}_{t}$:

$$ $\mathcal{L}_{\text{REPA}}(\theta,\phi)\coloneqq-\mathbb{E}_{{\mathbf{x}}_{\ast},\bm{\epsilon},t}\Big{[}\frac{1}{N}\sum_{n=1}^{N}\mathrm{sim}({\mathbf{y}}_{\ast}^{[n]},h_{\phi}({\mathbf{h}}_{t}^{([n]}))\Big{]},$ (8) $$

where $n$ is a patch index and $\mathrm{sim}(\cdot,\cdot)$ is a pre-defined similarity function.

In practice, we add this term to the original diffusion-based objectives described in Section [2](#S2) and Appendix [A](#A1). For instance, for the training of a velocity model in Eq. ([4](#S2.E4)), the objective becomes:

$$ $\mathcal{L}\coloneqq\mathcal{L}_{\text{velocity}}+\lambda\mathcal{L}_{\text{REPA}}$ (9) $$

where $\lambda>0$ is a hyperparameter that controls the tradeoff between denoising and representation alignment. We primarily investigate the impact of this regularization on two popular objectives: Improved DDPM (Nichol & Dhariwal, ) used in DiT (Peebles & Xie, ) and linear stochastic interpolants used in SiT (Ma et al., ), though other objectives can also be considered.

## 4 Experiments

We validate the performance of REPA and the effect of the proposed components through extensive experiments. In particular, we investigate the following questions:

- •
Can REPA improve diffusion transformer training significantly? (Table [4.1](#S4.SS1), [4.2](#S4.SS2), [4.2](#S4.SS2), Figure [4](#S3.F4), [6](#S4.F6))
- •
Is REPA scalable in terms of model size and representation quality? (Table [4.1](#S4.SS1), Figure [5](#S4.F5))
- •
Can diffusion model representations be aligned with various visual representations? (Figure [8](#S4.F8))

### 4.1 Setup

Figure: Table 1: Model configuration details.

Implementation details.
We strictly follow the setup in DiT (Peebles & Xie, ) and SiT (Ma et al., ) unless otherwise specified. Specifically, we use ImageNet (Deng et al., ), where each image is preprocessed to the resolution of 256$\times$256 (denoted as ImageNet 256$\times$256), and follow ADM (Dhariwal & Nichol, ) for other data preprocessing protocols. Each image is then encoded into a compressed vector ${\mathbf{z}}\in\mathbb{R}^{32\times 32\times 4}$ using the Stable Diffusion VAE (Rombach et al., ). For model configurations, we use the B/2, L/2, and XL/2 architectures introduced in the DiT and SiT papers, which process inputs with a patch size of 2 (see Table [4.1](#S4.SS1) for details). To ensure a fair comparison with DiTs 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](#A4).

**Table 2: Component-wise analysis on ImageNet 256$\times$256. All models are SiT-L/2 trained for 400K iterations. All metrics except accuracy (Acc.) are measured with the SDE Euler-Maruyama sampler with NFE=250 and without classifier-free guidance. For Acc., we report linear probing results on the ImageNet validation set using the latent features aligned with the target representation. We fix $\lambda=0.5$ here. $\downarrow$ and $\uparrow$ indicate whether lower or higher values are better, respectively.**
| Iter. | Target Repr. | Depth | Objective | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Pre.$\uparrow$ | Rec.$\uparrow$ | Acc.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 400K | Vanilla SiT-L/2 (Ma et al., ) | 18.8 | 5.29 | 072.0 | 0.64 | 0.64 | N/A |  |  |
| 400K | MAE-L | 8 | NT-Xent | 12.5 | 4.89 | 090.7 | 0.68 | 0.63 | 57.3 |
| 400K | DINO-B | 8 | NT-Xent | 11.9 | 5.00 | 092.9 | 0.68 | 0.63 | 59.3 |
| 400K | MoCov3-L | 8 | NT-Xent | 11.9 | 5.06 | 092.2 | 0.68 | 0.64 | 63.0 |
| 400K | I-JEPA-H | 8 | NT-Xent | 11.6 | 5.21 | 098.0 | 0.68 | 0.64 | 62.1 |
| 400K | CLIP-L | 8 | NT-Xent | 11.0 | 5.25 | 100.4 | 0.67 | 0.66 | 67.2 |
| 400K | DINOv2-L | 8 | NT-Xent | 10.0 | 5.09 | 106.6 | 0.68 | 0.65 | 68.1 |
| 400K | DINOv2-B | 8 | NT-Xent | 09.7 | 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 | 09.8 | 5.22 | 108.9 | 0.69 | 0.64 | 65.7 |
| 400K | DINOv2-L | 6 | NT-Xent | 10.3 | 5.23 | 106.5 | 0.69 | 0.65 | 66.2 |
| 400K | DINOv2-L | 8 | NT-Xent | 10.0 | 5.09 | 106.6 | 0.68 | 0.65 | 68.1 |
| 400K | DINOv2-L | 10 | NT-Xent | 10.5 | 5.50 | 105.0 | 0.68 | 0.65 | 68.6 |
| 400K | DINOv2-L | 12 | NT-Xent | 11.2 | 5.14 | 100.2 | 0.68 | 0.64 | 69.4 |
| 400K | DINOv2-L | 14 | NT-Xent | 11.6 | 5.61 | 099.5 | 0.67 | 0.65 | 70.0 |
| 400K | DINOv2-L | 16 | NT-Xent | 12.1 | 5.34 | 096.1 | 0.67 | 0.64 | 71.1 |
| 400K | DINOv2-L | 8 | NT-Xent | 10.0 | 5.09 | 106.6 | 0.68 | 0.65 | 68.1 |
| 400K | DINOv2-L | 8 | Cos. sim. | 09.9 | 5.34 | 111.9 | 0.68 | 0.65 | 68.2 |

Evaluation.
We report Fréchet inception distance (FID; Heusel et al. ), sFID (Nash et al., ), inception score (IS; Salimans et al. ), precision (Pre.) and recall (Rec.)
(Kynkäänniemi et al., ) using 50,000 samples. We also include linear probing results (Acc.) and CKNNA (Huh et al., ) as discussed in Section [3.2](#S3.SS2). We provide more details of each metric in Appendix [E](#A5).

Sampler.
Following SiT (Ma et al., ), we always use the SDE Euler-Maruyama sampler (for SDE with $w_{t}=\sigma_{t}$) 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 diffusion*: ADM (Dhariwal & Nichol, ), VDM$++$ (Kingma & Gao, ), Simple diffusion (Hoogeboom et al., ), CDM (Ho et al., ), (b) *Latent diffusion with U-Net*: LDM (Rombach et al., ), (c) *Latent diffusion with transformer+U-Net hybrid models*: U-ViT-H/2 (Bao et al., ), DiffiT (Hatamizadeh et al., ), and MDTv2-XL/2 (Gao et al., ), and (d) *Latent diffusion with transformers*: MaskDiT (Zheng et al., ), SD-DiT (Zhu et al., ), DiT (Peebles & Xie, ), and SiT (Ma et al., ). 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](#A6).

### 4.2 Component-wise analysis

We answer the question of whether REPA leads to improved diffusion transformer training. As shown in Table [4.1](#S4.SS1), we discover that REPA consistently provides a substantially improved generation 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 transformers. 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.

Figure: (a) Different visual encoders
Refer to caption: /html/2410.06940/assets/x10.png

Figure: Figure 6: Selected samples on ImageNet 256$\times$256 from the SiT-XL/2 + REPA model. We use classifier-free guidance with $w=4.0$.
Refer to caption: /html/2410.06940/assets/assets/main_qual_12.png

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 are marginal, which we 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 [4.1](#S4.SS1). Interestingly, limiting regularization to the first few layers further enhances generation performance (*e.g.*, adding REPA to layer 6 or 8 yields best results). We hypothesize that this enables the remaining layers to concentrate on capturing high-frequency details, building on a 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-scaled Cross Entropy (NT-Xent; Chen et al. ) or negative cosine similarity (cos. sim.). Empirically, we find that NT-Xent offers advantages in the early stages (*e.g.*, 50-100K iterations), 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 Figure [5(a)](#S4.F5.sf1), 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 in size. We demonstrate this by plotting FID-50K of different SiT models with and without REPA in Figure [5(b)](#S4.F5.sf2): REPA achieves the same FID level more quickly with larger models. Lastly, Figure [5(c)](#S4.F5.sf3) 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 in both generation and linear evaluation) with longer training.

Figure: Table 3: FID comparisons with vanilla DiTs and SiTs on ImageNet 256$\times$256. We do not use classifier-free guidance (CFG). $\downarrow$ denotes lower values are better. Iter. indicates the training iteration.

### 4.3 System-level comparison

Based on the analysis, we perform a system-level comparison between recent state-of-the-art diffusion 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 [4.2](#S4.SS2), REPA shows consistent and significant improvement across all model variants. In particular, on SiT-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-XL/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 [4](#S3.F4), where we use the same initial noise across different models. The model trained with REPA exhibits better progression.

Finally, we provide a quantitative comparison between SiT-XL/2 with REPA and other recent diffusion model methods using classifier-free guidance (Ho & Salimans, ). Our method already outperforms the original SiT-XL/2 with 7$\times$ fewer epochs and it 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 state-of-the-art FID of 1.42 with a extra classifier-free guidance scheduling with guidance interval (Kynkäänniemi et al., ).
We also provide selected qualitative results of SiT-XL/2 with REPA in Figure [6](#S4.F6) and more examples in Appendix [H](#A8).

### 4.4 Ablation studies

Representation gap across different timesteps.
We begin by comparing the semantic gap (measured 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 [8](#S4.F8), REPA consistently reduces the representation gap across different noise levels, as indicated by better linear probing results and higher CKNNA values across all noise scales.

Alignment to different visual encoders.
In addition, we extend the analysis from Section [2](#S1.F2) to other visual encoders, not limited to the DINOv2 models. Specifically, we train SiT-L/2 models using REPA with MAE or MoCov3. As depicted in Figure [8](#S4.F8), 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 to DINOv2.

Figure: Table 5: Ablation study for $\lambda$.

Effect of $\lambda$. We also examine the effect of the regularization coefficient $\lambda$ by training SiT-XL/2 models for 400K with different coefficients 0.25 to 1.0 and comparing the performance. As shown in Table [4.4](#S4.SS4), the performance is robust to the values and it is quite saturated after $\lambda=0.5$.

Figure: (a) Linear probing
Refer to caption: /html/2410.06940/assets/x13.png

## 5 Related Work

We discuss with the most relevant literature here and provide a more discussion in Appendix [I](#A9).

Bridging diffusion models and representation learning.
Many recent works have attempted to exploit or improve representations learned from diffusion models (Fuest et al., ). First, there are hybrid model approaches: Yang et al. () and Deja et al. () try to train a single model capable of both classification and diffusion-based generation. Also, Tian et al. () introduces a hybrid model capable of segmentation and generation. Next, several works have analyzed and exploited representations in diffusion models: Xiang et al. () and Mukhopadhyay et al. () observe that the intermediate representations of diffusion models have discriminative properties. Moreover, Repfusion (Yang & Wang, ) and DreamTeacher (Li et al., ) propose knowledge distillation schemes from diffusion models to a student model to perform various downstream tasks such as classification and several dense prediction tasks. Finally, there have been some attempts to improve denoising in diffusion models as self-supervised learning: Abstreiter et al. () extends the diffusion objective for a better representation learning scheme, and Chen et al. () deconstructs diffusion models to improve denoising-based representation learning. Hudson et al. () introduces an encoder that learns a representation by guiding a diffusion with its output as a compact latent vector. Zaidi et al. () focuses on the molecular domain and proposes a pretraining scheme based on denoising. Our work also shares some similarities, where we focus on *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 et al., ; Li et al., ). Würstchen (Pernias et al., ) introduces a two-stage text-to-image diffusion model framework: a text-conditioned model that first generates a semantic map from a text prompt, followed by another diffusion model that synthesizes images based on the semantic map. RCG (Li et al., ) focuses on unconditional generation, where a compact 1D latent vector is produced by a diffusion model and subsequently used as a label for image generation by a second diffusion model. Our work also exploits pretrained representations for improving the diffusion model, but we improve the diffusion model without the need to train an additional diffusion model that learns the representation distribution. In particular, we train a single diffusion model by proposing a regularization that aligns diffusion representations with self-supervised visual representations.

## 6 Conclusion

In this paper, we have presented REPA, a simple regularization for improving diffusion transformers. 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 possible future research directions, including unifying discriminative and generative models and their representations or theoretical analysis in how this alignment improves the generation.

## Reproducibility Statement

We provide hyperparameter details in Section [4](#S4) and Appendix [D](#A4). We will also release the implementation and model checkpoints in the future to reproduce the results in the paper.

## Appendix 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 a denoising autoencoder. We first explain denoising diffusion probabilistic models (DDPM) in Section [A.1](#A1.SS1) and stochastic interpolants in Section [A.2](#A1.SS2). For detailed explanations and rigorous proofs, please refer to the original papers (Albergo et al., ; Ma et al., ) that provide excellent formulations and description.

### A.1 Denoising diffusion probabilistic models

*Diffusion models* (Sohl-Dickstein et al., ; Ho et al., ) model the target distribution $p({\mathbf{x}})$ via learning a gradual denoising process from Gaussian distribution $\mathcal{N}(\mathbf{0},\mathbf{I})$ to $p({\mathbf{x}})$. Formally, diffusion models learn a *reverse* process $p({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})$ of the pre-defined *forward* process $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$ that gradually adds the Gaussian noise starting from $p({\mathbf{x}})$ for $1\leq t\leq T$ with a fixed $T>0$.

For a given ${\mathbf{x}}_{0}\sim p({\mathbf{x}})$, $q({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1})$ can be formalized as $q({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1})\coloneqq\mathcal{N}({\mathbf{x}}_{t};\sqrt{1-\beta_{t}}{\mathbf{x}}_{0},\beta_{t}^{2}\mathbf{I})$, where $\beta_{t}\in(0,1)$ are pre-defined hyperparameters set to be small. In particular, DDPM (Ho et al., ) shows if one formalizes the reverse process $p({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})$ (with ${\alpha}_{t}=1-\beta_{t}$. $\bar{\alpha}_{t}\coloneqq\prod_{i=1}^{t}\alpha_{i}$ for $1\leq t\leq T$) as

$$ $\displaystyle p({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})\coloneqq\mathcal{N}\Big{(}{\mathbf{x}}_{t-1};\frac{1}{\sqrt{\alpha_{t}}}\big{(}{\mathbf{x}}_{t}-\frac{\sigma_{t}^{2}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}_{\bm{\theta}}({\mathbf{x}}_{t},t)\big{)},\mathbf{\Sigma}_{\theta}({\mathbf{x}}_{t},t)\Big{)},$ (10) $$

then $\bm{\epsilon}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ can be trained with a simple denoising autoencoder objective parameterized by $\bm{\theta}$:

$$ $\displaystyle\mathcal{L}_{\text{simple}}\coloneqq\mathbb{E}_{{\mathbf{x}}_{\ast},\bm{\epsilon},t}\Big{[}||\bm{\epsilon}-\bm{\epsilon}_{\bm{\theta}}({\mathbf{x}}_{t},t)||_{2}^{2}\Big{]}.$ (11) $$

For $\mathbf{\Sigma}_{\theta}({\mathbf{x}}_{t},t)$, (Ho et al., ) shows it is enough to simply define it as $\sigma_{t}^{2}\mathbf{I}$ with $\beta_{t}=\sigma_{t}^{2}$. After that, Nichol & Dhariwal () exhibits the performance can be improved if the model jointly learns $\mathbf{\Sigma}_{\theta}({\mathbf{x}}_{t},t)$ with $\bm{\epsilon}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ in dimension-wise manner through the following objective:

$$ $\mathcal{L}_{\text{vlb}}\coloneqq\exp(v\log\beta_{t}+(1-v)\log\tilde{\beta}_{t}),$ (12) $$

where $v$ denotes each component per dimension from the model output and $\tilde{\beta}_{t}=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\beta_{t}$.

With a sufficiently large $T$ and an appropriate scheduling of $\beta_{t}$, the distribution $p({\mathbf{x}}_{T})$ becomes almost an isotropic Gaussian distribution. Hence, one can generate a sample starting from a random noise and perform iterative reverse process $p({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})$ to reach the data sample ${\mathbf{x}}_{0}$ (Ho et al., ).

### A.2 Stochastic interpolants

Different from DDPM, *flow-based models* (Esser et al., ; Lipman et al., ; Liu et al., ) deal with the continuous time-dependent process with a data ${\mathbf{x}}_{\ast}\sim p({\mathbf{x}})$ and a Gaussian noise ${\mathbf{\epsilon}}~{}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ on $t\in[0,1]$:

$$ ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{0}+\sigma_{t}{\mathbf{\epsilon}},\quad\alpha_{0}=\sigma_{1}=1,\,\,\alpha_{1}=\sigma_{0}=0,$ (13) $$

where $\alpha_{t}$ and $\sigma_{t}$ are a decreasing and increasing function of $t$ (respectively). There exists a *probability flow ordinary differential equation* (PF ODE) with a velocity field

$$ $\dot{{\mathbf{x}}}_{t}={\mathbf{v}}({\mathbf{x}}_{t},t),$ (14) $$

where distribution of this ODE at $t$ is equal to the marginal $p_{t}({\mathbf{x}})$.

The velocity ${\mathbf{v}}({\mathbf{x}},t)$ is represented as the following sum of two conditional expectations

$$ ${\mathbf{v}}({\mathbf{x}},t)=\mathbb{E}[\dot{{\mathbf{x}}}_{t}|{\mathbf{x}}_{t}={\mathbf{x}}]=\dot{\alpha}_{t}\mathbb{E}[{\mathbf{x}}_{\ast}|{\mathbf{x}}_{t}={\mathbf{x}}]+\dot{\sigma}_{t}\mathbb{E}[{\mathbf{\epsilon}}|{\mathbf{x}}_{t}={\mathbf{x}}],$ (15) $$

which can be approximated with model ${\mathbf{v}}_{\theta}({\mathbf{x}}_{t},t)$ by minimizing the following training objective:

$$ $\mathcal{L}_{\text{velocity}}(\theta)\coloneqq\mathbb{E}_{{\mathbf{x}}_{\ast},\bm{\epsilon},t}\Big{[}||{\mathbf{v}}_{\theta}({\mathbf{x}}_{t},t)-\dot{\alpha}_{t}{\mathbf{x}}_{\ast}-\dot{\sigma}_{t}{\mathbf{\epsilon}}||^{2}\Big{]}.$ (16) $$

Note that this also corresponds to the following reverse *stochastic differential equation* (SDE):

$$ $d{\mathbf{x}}_{t}={\mathbf{v}}({\mathbf{x}}_{t},t)dt-\frac{1}{2}w_{t}{\mathbf{s}}({\mathbf{x}}_{t},t)dt+\sqrt{w_{t}}d\bar{\mathbf{w}}_{t},$ (17) $$

where the score ${\mathbf{s}}({\mathbf{x}}_{t},t)$ similarly becomes the conditional expectation

$$ ${\mathbf{s}}({\mathbf{x}}_{t},t)=-\frac{1}{\sigma_{t}}\mathbb{E}[{\mathbf{\epsilon}}|{\mathbf{x}}_{t}={\mathbf{x}}].$ (18) $$

Similar to ${\mathbf{v}}$, ${\mathbf{s}}$ can be approximated with a model ${\mathbf{s}}_{\theta}({\mathbf{x}},t)$ with the following objective:

$$ $\mathcal{L}_{\text{score}}(\theta)\coloneqq\mathbb{E}_{{\mathbf{x}}_{\ast},\bm{\epsilon},t}\Big{[}||\sigma_{t}{\mathbf{s}}_{\theta}({\mathbf{x}}_{t},t)+{\mathbf{\epsilon}}||^{2}\Big{]}.$ (19) $$

Here, since the score ${\mathbf{s}}({\mathbf{x}},t)$ can be directly computed using the velocity ${\mathbf{v}}({\mathbf{x}},t)$ for $t>0$ as

$$ ${\mathbf{s}}({\mathbf{x}},t)={\frac{1}{\sigma_{t}}}\cdot\frac{\alpha_{t}{\mathbf{v}}({\mathbf{x}},t)-\dot{\alpha}_{t}{\mathbf{x}}}{\dot{\alpha}_{t}\sigma_{t}-{\alpha}_{t}\dot{\sigma}_{t}},$ (20) $$

so it is enough to estimate only one of the two vectors.

*Stochastic interpolants* (Albergo et al., ) shows any $\alpha_{t}$ and $\sigma_{t}$ satisfy the three conditions

$$ $\displaystyle 1.~{}$ $\displaystyle\alpha_{t}^{2}+\sigma_{t}^{2}>0,\,\,\forall t\in[0,1]$ $\displaystyle 2.~{}$ $\displaystyle\alpha_{t}~{}\text{and}~{}\sigma_{t}~{}\text{are differentiable},\,\,\forall t\in[0,1]$ $\displaystyle 3.~{}$ $\displaystyle\alpha_{1}=\sigma_{0}=0,~{}\alpha_{0}=\sigma_{1}=1,$ $$

leads to a process that interpolates between ${\mathbf{x}}_{0}$ and ${\mathbf{\epsilon}}$ without bias. Thus, one can use a simple interpolant by defining them as a simple function during training and inference, such as linear interpolants with $\alpha_{t}=1-t$ and $\sigma_{t}=t$ or variance-preserving (VP) interpolants with $\alpha_{t}=\cos(\frac{\pi}{2}t)$ and $\sigma_{t}=\cos(\frac{\pi}{2}t)$ (Ma et al., ). One another advantage of stochastic interpolants is that the diffusion coefficient $w_{t}$ is independent in training any of a score or a velocity model. Thus, $w_{t}$ can be also explicitly chosen *after training* when sampling with the reverse SDE.

Note that existing score-based diffusion models, including DDPM (Ho et al., ), 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 $\mathcal{N}(\mathbf{0},\mathbf{I})$ at $t\to\infty$, where the training is done on $[0,T]$ with sufficiently large $T$ (*e.g.*, $T=1000$) that $p({\mathbf{x}}_{T})$ becomes almost isotropic Gaussian. Generation is done by solving the corresponding reverse SDE starting from a random Gaussian noise by assuming ${\mathbf{x}}_{T}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, where $\alpha_{t},\sigma_{t}$ and the diffusion coefficient $w_{t}$ is *implicitly* chosen from the forward diffusion process, which might lead to over-complicated design space of score-based diffusion models (Karras et al., ).

## Appendix B Diffusion Transformer Architecture

Figure: Figure 9: DiT block illustration.
Refer to caption: /html/2410.06940/assets/x17.png

We strictly follow the architecture used in DiT (Peebles & Xie, ) and SiT (Ma et al., ). The architecture is very similar to a vision transformer (ViTs; Dosovitskiy et al. ): an input is patchified, reshaped to a 1D sequence of patches with a length $N$, and then fed to the model. Similar to DiT and SiT, our architecture also uses a downsampled latent image ${\mathbf{z}}=E({\mathbf{x}})$ as an input, where ${\mathbf{x}}$ is a RGB image and $E$ is an encoder of the stable diffusion variational autoencoder (VAE) (Rombach et al., ). 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 a single multilayer perceptron (MLP) that projects a hidden state to the target representation space, which is only used in training. We provide an illustration of the DiT block in Figure [9](#A2.F9).

## Appendix C Analysis Details

### C.1 Evaluation details

CKNNA (Centered Kernel Nearest-Neighbor Alignment) is a *relaxed version* of the popular Centered Kernel Alignment (CKA; Kornblith et al. ) that mitigates the strict definition of alignment. We generally follow the notations in the original paper for an explanation (Huh et al., ).

First, CKA have measured *global* similarities of the models by considering all possible data pairs:

$$ $\displaystyle\mathrm{CKA}(\mathbf{K},\mathbf{L})=\frac{\mathrm{HSIC}(\mathbf{K},\mathbf{L})}{\sqrt{\mathrm{HSIC}(\mathbf{K},\mathbf{K})\mathrm{HSIC}(\mathbf{L},\mathbf{L})}},$ (21) $$

where $\mathbf{K}$ and $\mathbf{L}$ are two kernel matrices computed from the dataset using two different networks. Specifically, it is defined as $\mathbf{K}_{ij}=\kappa(\phi_{i},\phi_{j})$ and $\mathbf{L}_{ij}=\kappa(\psi_{i},\psi_{j})$ where $\phi_{i},\phi_{j}$ and $\psi_{i},\psi_{j}$ are representations computed from each network at the corresponding data ${\mathbf{x}}_{i},{\mathbf{x}}_{j}$ (respectively). By letting $\kappa$ as a inner product kernel, $\mathrm{HSIC}$ is defined as

$$ $\displaystyle\mathrm{HSIC}(\mathbf{K},\mathbf{L})=\frac{1}{(n-1)^{2}}\Big{(}\sum_{i}\sum_{j}\big{(}\langle\phi_{i},\phi_{j}\rangle-\mathbb{E}_{l}[\langle\phi_{i},\phi_{l}\rangle]\big{)}\big{(}\langle\psi_{i},\psi_{j}\rangle-\mathbb{E}_{l}[\langle\psi_{i},\psi_{l}\rangle]\big{)}\Big{)}.$ (22) $$

CKNNA considers a relaxed version of Eq. ([21](#A3.E21)) by replacing $\mathrm{HSIC}(\mathbf{K},\mathbf{L})$ into $\mathrm{Align}(\mathbf{K},\mathbf{L})$, where $\mathrm{Align}(\mathbf{K},\mathbf{L})$ computes Eq. ([22](#A3.E22)) only using a $k$-nearest neighborhood embedding in the datasets:

$$ $\displaystyle\mathrm{Align}(\mathbf{K},\mathbf{L})=\frac{1}{(n-1)^{2}}\Big{(}\sum_{i}\sum_{j}\alpha(i,j)\big{(}\langle\phi_{i},\phi_{j}\rangle-\mathbb{E}_{l}[\langle\phi_{i},\phi_{l}\rangle]\big{)}\big{(}\langle\psi_{i},\psi_{j}\rangle-\mathbb{E}_{l}[\langle\psi_{i},\psi_{l}\rangle]\big{)}\Big{)},$ (23) $$

where $\alpha(i,j)$ is defined as

$$ $\displaystyle\alpha(i,j;k)=\mathbbm{1}[i\neq j\,\,\text{and}\,\,\phi_{j}\in\mathrm{knn}(\phi_{i};k)\,\,\text{and}\,\,\psi_{j}\in\mathrm{knn}(\psi_{i};k)],$ (24) $$

so this term only considers $k$-nearest neighbors at each $i$. In this paper, we randomly sample 10,000 images in the validation set in ImageNet (Deng et al., ) and report CKNNA with $k=10$ based on observation in Huh et al. () that smaller $k$ shows better a better alignment.

Linear probing.
We follow the setup used in DAE (Chen et al., ). Specifically, we use parameter-free batch normalization layer and train a linear layer for 90 epochs with a batch size of 16,384. We use the Adam optimizer (Kingma, ) 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 [2(a)](#S1.F2.sf1) (linear probing) and [2(b)](#S1.F2.sf2) (CKA), and illustrate the result in Figure [10](#A3.F10). Overall is shows a 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.

Figure: (a) Semantic gap
Refer to caption: /html/2410.06940/assets/x18.png

### C.3 Description of pretrained visual encoders

- •
MAE (He et al., ) proposes a self-supervised representation learning objective for vision transformers, based on the reconstruction task of masked patches of input images.
- •
DINO (Caron et al., ) is a self-supervised learning method based on self-distillation through the mean of momentum teacher network.
- •
MoCov3 (Chen et al., ) studies empirical study to train MoCo (He et al., ; Chen et al., ) on vision transformer and how they can be scaled up.
- •
CLIP (Radford et al., ) proposes a contrastive learning scheme on large image-text pairs.
- •
DINOv2 (Oquab et al., ) proposes a self-supervised learning method that combines pixel-level and patch-level discriminative objectives by leveraging advanced self-supervised techniques and a large pre-training dataset.
- •
I-JEPA (Assran et al., ) predicts missing parts of an image by learning representations through joint-embedding, focusing on the context of the entire image without relying on pixel-level reconstruction.

## Appendix D Hyperparameter and More Implementation Details

**Table 6: Hyperparameter setup.**
|  | Figure 3 | Table 3 (SiT-B) | Table 3 (SiT-L) | Table 3 (SiT-XL) | Table 4 |
| --- | --- | --- | --- | --- | --- |
| Architecture |  |  |  |  |  |
| Input dim. | 32$\times$32$\times$4 | 32$\times$32$\times$4 | 32$\times$32$\times$4 | 32$\times$32$\times$4 | 32$\times$32$\times$4 |
| Num. layers | 28 | 12 | 24 | 28 | 24 |
| Hidden dim. | 1,152 | 768 | 1,024 | 1,152 | 1,024 |
| Num. heads | 16 | 12 | 16 | 16 | 16 |
| REPA |  |  |  |  |  |
| $\lambda$ | 0.5 | 0.5 | 0.5 | 0.5 | 0.5 |
| Alignment depth | 4 | 8 | 8 | 8 | 8 |
| $\mathrm{sim}(\cdot,\cdot)$ | cos. sim. | cos. sim. | NT-Xent | cos. sim. | NT-Xent |
| Encoder $f({\mathbf{x}})$ | DINOv2-B | DINOv2-B | DINOv2-L | DINOv2-B | DINOv2-B |
| Optimization |  |  |  |  |  |
| Training iteration | 1M | 400K | 700K | 4M | 4M |
| Batch size | 256 | 256 | 256 | 256 | 256 |
| Optimizer | AdamW | AdamW | AdamW | AdamW | AdamW |
| lr | 0.0001 | 0.0001 | 0.0001 | 0.0001 | 0.0001 |
| $(\beta_{1},\beta_{2})$ | (0.9, 0.999) | (0.9, 0.999) | (0.9, 0.999) | (0.9, 0.999) | (0.9, 0.999) |
| Interpolants |  |  |  |  |  |
| $\alpha_{t}$ | $1-t$ | $1-t$ | $1-t$ | $1-t$ | $1-t$ |
| $\sigma_{t}$ | $t$ | $t$ | $t$ | $t$ | $t$ |
| $w_{t}$ | $\sigma_{t}$ | $\sigma_{t}$ | $\sigma_{t}$ | $\sigma_{t}$ | $\sigma_{t}$ |
| Training objective | v-prediction | v-prediction | v-prediction | v-prediction | v-prediction |
| Sampler | Euler-Maruyama | Euler-Maruyama | Euler-Maruyama | Euler-Maruyama | Euler-Maruyama |
| Sampling steps | 250 | 250 | 250 | 250 | 250 |
| Guidance | - | - | - | - | 1.35 |

Further implementation details.
We implement our model based on the original SiT implementation (Ma et al., ). Throughout the experiments, we use the exact same structure as DiT (Peebles & Xie, ) and SiT (Ma et al., ). We use AdamW (Kingma, ; Loshchilov, ) with constant learning rate of 1e-4, $(\beta_{1},\beta_{2})=(0.9,0.999)$ without weight decay. To speed up training, we use mixed-precision (fp16) with a gradient clipping. We also pre-compute compressed latent vectors from raw pixels via stable diffusion VAE (Rombach et al., ) and use these latent vectors. Because of this, we do not apply any data augmentation, but we find this does not lead to a big difference, as similarly observed in EDM2 (Karras et al., ). We also use stabilityai/sd-vae-ft-ema decoder for decoding latent vectors to images. For MLP used for a projection, we use three-layer MLP with SiLU activations (Elfwing et al., ). We provide a detailed hyperparameter setup in Table [6](#A4.T6).

Pretrained encoders.
For MoCov3-B and -L models, we use the checkpoint in the implementation of RCG (Li et al., );(^2^22[https://github.com/LTH14/rcg](https://github.com/LTH14/rcg)) for other checkpoints, we use their official checkpoints released in their official implementations. To adjust a different number of patches between the diffusion transformer and the pretrained encoder, we interpolate positional embeddings of pretrained encoders.

Sampler.
For sampling, we use the Euler-Maruyama sampler with the SDE in Eq. ([5](#S2.E5)) with a diffusion coefficient $w_{t}=\sigma_{t}$. We use the last step of the SDE sampler as 0.04, and it gives a significant improvement, similar to the original SiT paper (Ma et al., ).

Computing resources.
We use 8 NVIDIA H100 80GB GPUs for experiments; our training speed is 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).

## Appendix E Evaluation Details

We strictly follow the setup and use the same reference batches of ADM (Dhariwal & Nichol, ) for evaluation, following their official implementation.(^3^33[https://github.com/openai/guided-diffusion/tree/main/evaluations](https://github.com/openai/guided-diffusion/tree/main/evaluations)) We use NVIDIA H100 80GB GPUs or 4090Ti GPUs for evaluation and enable tf32 precision for faster generation, and we find the performance difference is 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., ) measures the feature distance between the distributions of real and generated images. It uses the Inception-v3 network (Szegedy et al., ) and computes distance based on an assumption that both feature distributions are multivariate gaussian distributions.
- •
sFID (Nash et al., ) proposes to compute FID with intermediate spatial features of the Inception-v3 network to capture the generated images’ spatial distribution.
- •
IS (Salimans et al., ) also uses the Inception-v3 network but use logit for evaluation of the metric. Specifically, it measures a KL-divergence between the original label distribution and the distribution of logits after the softmax normalization.
- •
Precision and recall (Kynkäänniemi et al., ) are based on their classic definitions: the fraction of realistic images and the fraction of training data manifold covered by generated data.

## Appendix F Baselines

In what follows, we explain the main idea of baseline methods that we used for the evaluation.

- •
ADM (Dhariwal & Nichol, ) improves U-Net-based architectures for diffusion models and proposes classifier-guided sampling to balance the quality and diversity tradeoff.
- •
VDM++ (Kingma & Gao, ) proposes a simple adaptive noise schedule for diffusion models to improve training efficiency.
- •
Simple diffusion (Hoogeboom et al., ) proposes a diffusion model for high-resolution image generation by exploring various techniques to simplify a noise schedule and architectures.
- •
CDM (Ho et al., ) introduces cascaded diffusion models: similar to progressiveGAN (Karras et al., ), it trains multiple diffusion models starting from the lowest resolution and applying one or more super-resolution diffusion models for generating high-fidelity images.
- •
LDM (Rombach et al., ) proposes latent diffusion models by modeling image distribution in a compressed latent space to improve the training efficiency without sacrificing the generation performance.
- •
U-ViT (Bao et al., ) proposes a ViT-based latent diffusion model that incorporates U-Net-like long skip connections.
- •
DiffiT (Hatamizadeh et al., ) proposes a time-dependent multi-head self-attention mechanism for enhancing the efficiency of transformer-based image diffusion models.
- •
MDTv2 (Gao et al., ) proposes an asymmetric encoder-decoder scheme for efficient training of a diffusion-based transformer. They also apply U-Net-like long-shortcuts in the encoder and dense input-shortcuts in the decoder.
- •
MaskDiT (Zheng et al., ) proposes an asymmetric encoder-decoder scheme for efficient training of diffusion transformers, where they train the model with an auxiliary mask reconstruction task similar to MAE (He et al., ).
- •
SD-DiT (Zhu et al., ) extends MaskdiT architecture but incorporates self-supervised discrimination objective using a momentum encoder.
- •
DiT (Peebles & Xie, ) proposes a pure transformer backbone for training diffusion models based on proposing AdaIN-zero modules.
- •
SiT (Ma et al., ) extensively analyzes how DiT training can be efficient by moving from discrete diffusion to continuous flow-based modeling.

## Appendix G Detailed Quantitative Results

We provide evaluation results of different SiT models trained with REPA. All models are aligned with DINOv2-B representations with $\lambda=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 7: Detailed evaluation results with different model sizes. All results are reported without classifier-free guidance.**
| Model | #Params | Iter. | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Prec.$\uparrow$ | Rec.$\uparrow$ | Acc.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SiT-B/2 (Ma et al., ) | 130M | 400K | 33.0 | 06.46 | 043.7 | 0.53 | 0.63 | N/A |
| + REPA (ours) | 130M | 50K | 78.2 | 11.71 | 017.1 | 0.33 | 0.48 | 43.2 |
| + REPA (ours) | 130M | 100K | 49.5 | 07.00 | 027.5 | 0.46 | 0.59 | 50.9 |
| + REPA (ours) | 130M | 200K | 33.2 | 06.68 | 043.7 | 0.54 | 0.63 | 50.9 |
| + REPA (ours) | 130M | 400K | 24.4 | 06.40 | 059.9 | 0.59 | 0.65 | 61.2 |
| SiT-L/2 (Ma et al., ) | 458M | 400K | 18.8 | 05.29 | 072.0 | 0.64 | 0.64 | N/A |
| + REPA (ours) | 458M | 50K | 55.4 | 024.0 | 023.0 | 0.43 | 0.53 | 55.3 |
| + REPA (ours) | 458M | 100K | 24.1 | 06.25 | 055.7 | 0.62 | 0.60 | 61.8 |
| + REPA (ours) | 458M | 200K | 14.0 | 05.18 | 086.5 | 0.67 | 0.64 | 66.3 |
| + REPA (ours) | 458M | 400K | 10.0 | 05.20 | 109.2 | 0.69 | 0.65 | 69.4 |
| SiT-XL/2 (Ma et al., ) | 675M | 7M | 08.3 | 06.32 | 131.7 | 0.68 | 0.67 | N/A |
| + REPA (ours) | 675M | 50K | 52.3 | 31.24 | 024.3 | 0.45 | 0.53 | 56.1 |
| + REPA (ours) | 675M | 100K | 19.4 | 06.06 | 067.4 | 0.64 | 0.61 | 62.9 |
| + REPA (ours) | 675M | 200K | 11.1 | 05.05 | 100.4 | 0.69 | 0.64 | 67.3 |
| + REPA (ours) | 675M | 400K | 07.9 | 05.06 | 122.6 | 0.70 | 0.65 | 70.3 |
| + REPA (ours) | 675M | 4M | 05.9 | 05.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 8: Detailed evaluation results of SiT-XL+REPA at 4M iteration with different classifier-free guidance scale $w$.**
| Model | #Params | Iter. | $w$ | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Prec.$\uparrow$ | Rec.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SiT-XL/2 (Ma et al., ) | 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 |

Finally, we provide the results with the guidance interval (Kynkäänniemi et al., ).

**Table 9: Detailed evaluation results of SiT-XL+REPA at 4M iteration with different classifier-free guidance scale $w$. We apply the guidance interval (Kynkäänniemi et al., ).**
| Model | #Params | Iter. | Interval | $w$ | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Prec.$\uparrow$ | Rec.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SiT-XL/2 (Ma et al., ) | 675M | 7M | [0, 1] | 1.50 | 2.06 | 4.50 | 270.3 | 0.82 | 0.59 |
| + REPA (ours) | 675M | 4M | [0, 0.8] | 2.00 | 2.23 | 4.40 | 360.9 | 0.84 | 0.6 |
| + REPA (ours) | 675M | 4M | [0, 0.75] | 2.00 | 1.78 | 4.50 | 346.2 | 0.82 | 0.62 |
| + REPA (ours) | 675M | 4M | [0, 0.7] | 2.00 | 1.48 | 4.67 | 324.0 | 0.82 | 0.62 |
| + REPA (ours) | 675M | 4M | [0, 0.65] | 2.00 | 1.44 | 4.88 | 308.8 | 0.79 | 0.65 |
| + REPA (ours) | 675M | 4M | [0, 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 |

## Appendix H More Qualitative Results

Figure: Figure 11: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “loggerhead sea turtle” (33).
Refer to caption: /html/2410.06940/assets/x20.png

Figure: Figure 12: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “macaw” (88).
Refer to caption: /html/2410.06940/assets/x21.png

Figure: Figure 13: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “sulphur-crested cockatoo” (89).
Refer to caption: /html/2410.06940/assets/x22.png

Figure: Figure 14: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “golden retriever” (207).
Refer to caption: /html/2410.06940/assets/x23.png

Figure: Figure 15: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “husky” (250).
Refer to caption: /html/2410.06940/assets/x24.png

Figure: Figure 16: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “arctic wolf” (270).
Refer to caption: /html/2410.06940/assets/x25.png

Figure: Figure 17: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “arctic fox” (279).
Refer to caption: /html/2410.06940/assets/x26.png

Figure: Figure 18: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “lion” (291).
Refer to caption: /html/2410.06940/assets/x27.png

Figure: Figure 19: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “otter” (360).
Refer to caption: /html/2410.06940/assets/x28.png

Figure: Figure 20: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “red panda” (387).
Refer to caption: /html/2410.06940/assets/x29.png

Figure: Figure 21: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “panda” (388).
Refer to caption: /html/2410.06940/assets/x30.png

Figure: Figure 22: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “acoustic guitar” (402).
Refer to caption: /html/2410.06940/assets/x31.png

Figure: Figure 23: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “balloon” (417).
Refer to caption: /html/2410.06940/assets/x32.png

Figure: Figure 24: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “baseball” (429).
Refer to caption: /html/2410.06940/assets/x33.png

Figure: Figure 25: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “dog sled” (537).
Refer to caption: /html/2410.06940/assets/x34.png

Figure: Figure 26: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “fire truck” (555).
Refer to caption: /html/2410.06940/assets/x35.png

Figure: Figure 27: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “laptop” (620).
Refer to caption: /html/2410.06940/assets/x36.png

Figure: Figure 28: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “space shuttle” (812).
Refer to caption: /html/2410.06940/assets/x37.png

Figure: Figure 29: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “ice cream” (928).
Refer to caption: /html/2410.06940/assets/x38.png

Figure: Figure 30: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “cheeseburger” (933).
Refer to caption: /html/2410.06940/assets/x39.png

Figure: Figure 31: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “cliff drop-off” (972).
Refer to caption: /html/2410.06940/assets/x40.png

Figure: Figure 32: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “coral reef” (973).
Refer to caption: /html/2410.06940/assets/x41.png

Figure: Figure 33: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “lake shore” (975).
Refer to caption: /html/2410.06940/assets/x42.png

Figure: Figure 34: Uncurated generation results of SiT-XL/2 + REPA. We use classifier-free guidance with $w=4.0$. Class label = “volcano” (980).
Refer to caption: /html/2410.06940/assets/x43.png

## Appendix I 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. ) that try to accelerate training with better convergence using pretrained visual encoders (Sauer et al., ; Kumari et al., ; Sauer et al., ; ; Kang et al., ). 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 et al., ; ; Kang et al., ). Another line of work tries to exploit the pretrained visual encoders for improving diffusion model training from scratch (Pernias et al., ; Li et al., ), 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 our 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., ), MDT (Gao et al., ), and DiffiT (Hatamizadeh et al., ) show transformer-based backbones with *skip connections* can be an effective backbone for training diffusion models. Intriguingly, DiT (Peebles & Xie, ) show skip connections are not even necessary components, and a pure transformer architecture can be a scalable architecture for training diffusion-based models. Based on DiT, SiT (Ma et al., ) shows the model can be further improved with continuous stochastic interpolants (Albergo et al., ). Moreover, VDT (Lu et al., ) and Latte (Ma et al., ) show DiTs can be extended for video generation through a space-time attention (Arnab et al., ). Based on these improvements, Pixart-$\alpha$ (Chen et al., ), Pixart-$\Sigma$ (Chen et al., ), Stable diffusion 3 (Esser et al., ) show pure transformers can be scaled up for challenging text-to-image generation, and CMD (Yu et al., ), WALT (Gupta et al., ), and Sora (Brooks et al., ) demonstrates their success in text-to-video generation. Our work analyzes and improves the training of DiT (and SiT) architecture based on a simple feature matching regularization to the early layers.

Generative models with auxiliary self-supervised tasks. MaskDiT (Zheng et al., ) combines mask reconstruction in MAE (He et al., ) to diffusion model training for faster diffusion model training. Similarly, SD-DiT (Zhu et al., ) shows diffusion model training can be improved with an auxiliary discriminative self-supervised loss. MAGE (Li et al., ) bridge MAE training and masked image modeling (Chang et al., ) 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.

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
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.
2. 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.
3. 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.
soundness: 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.
2. 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.
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.
weaknesses: 1. High-resolution results (e.g., 512x512) are needed to support and further validate findings.
2. 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.
3. 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.
4. 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.
5. 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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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:

[a] Learning Disentangled Representation by Exploiting Pretrained Generative Models: A Contrastive Learning View. ICLR 2022.

[b] DisDiff: Unsupervised Disentanglement of Diffusion Probabilistic Models. NeurIPS 2023.
soundness: 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.

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

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

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

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
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. 
2. 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.
3. 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
4. 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 
Dinov2> CLIP>MoCov3>DINO>MAE
Hence 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
soundness: 4
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.
2. 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.
3. 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.
4. These insights are extremely valuable and novel, representing a significant step forward in the field.
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.
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
2. 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.
3. 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.
4. 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. 


   [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).

   [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).

    [3] Wah, C., Branson, S., Welinder, P., Perona, P. and Belongie, S., 2011. The caltech-ucsd birds-200-2011 dataset.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 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.

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

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

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
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?

-- 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?
soundness: 3
strengths: --the proposed method is simple and very easy to implement, also the idea is well-motivated by the semantic measurement metrics

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

-- 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?

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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.
2. 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?
3. 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?
4. 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).

[1] Gao, Shanghua, et al. "Masked diffusion transformer is a strong image synthesizer." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.
soundness: 4
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.
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.
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.
2. 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.

Meta-review: This paper introduces REPresentation Alignment (REPA), a simple yet effective regularization technique that significantly enhances the efficiency and performance of diffusion transformers by aligning their representations with pretrained self-supervised visual encoders. Reviewers praised the paper's strong empirical results, including substantial speedups in training (e.g., 17.5× faster for SiT) and improved generation quality, as well as the thorough experimental analyses across various settings. The clear motivation, innovative insights into representation learning, and demonstration of REPA's effectiveness in both class- and text-conditioned generative tasks make this work a strong contribution to the field.

Decision: Accept (Oral)