PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, Saining Xie

arXiv:2410.06940v4cs.CVcs.LG

TL;DR

Diffusion transformers struggle to learn high-quality representations, limiting efficient generative training. REPA aligns their noisy hidden states with external self-supervised visual representations, enabling faster convergence and improved generation quality, including FID=1.42 with guidance interval.

  • Problem

    Diffusion models’ learned representations lag behind self-supervised visual representations, making representation learning a central challenge in large-scale generative training.

  • Method

    REPA regularizes diffusion transformers by aligning projections of noisy hidden states with clean representations from pretrained self-supervised visual encoders.

  • Results

    REPA consistently improves diffusion-transformer generation and convergence, achieving FID=7.9 at 400K SiT iterations and FID=1.42 with guidance interval.

  • Takeaways & Limitations

    External visual representations can substantially accelerate diffusion-transformer training while improving generation performance across DiT and SiT models.

  • Takeaways & Limitations

    The theoretical relationship between representations learned through instance discrimination and denoising remains an open direction for future work.

Abstract

from arXiv · show

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

The paper identifies high-quality internal representation learning as a central bottleneck in diffusion-transformer training and introduces REPA to align noisy hidden states with external self-supervised visual representations. REPA reduces the representation gap and substantially improves training efficiency and generation quality for DiTs and SiTs.

  • Contributions: The paper hypothesizes that learning high-quality representations is essential for improving diffusion-transformer generation performance.Its main contributions are the representation-learning hypothesis, REPA regularization, and improved ImageNet generation efficiency and FID scores.
  • Motivation: Diffusion transformers learn meaningful discriminative representations, but these are significantly inferior to those produced by DINOv2.This empirical gap motivates using external self-supervised visual representations to support generative-model training.
  • Method: REPA distills clean-image self-supervised representations into diffusion-transformer representations of noisy inputs.The regularizer addresses the mismatch between noisy diffusion inputs and clean-image encoder inputs while reducing the semantic gap.
  • Results: REPA improves generation performance and representation alignment when applied to DiTs and SiTs.The method is evaluated through a system-level comparison across these two recent diffusion-transformer architectures.
  • Results: 400K training iterations yields FID=7.9 on class-conditional ImageNet generation for SiT without classifier-free guidance.This is reported as >17.5× faster than vanilla SiTs.

2 PRELIMINARIES

The preliminaries unify flow- and diffusion-based generation through stochastic interpolants, which define noisy states, probability-flow ODEs, and equivalent SDE sampling procedures. The paper primarily uses a linear interpolant but states that the approach extends to related variants such as DDPM.

  • Stochastic interpolants: Stochastic interpolants define a continuous process x_t = α_t x* + σ_t ϵ from data x* ∼ p(x) and Gaussian noise ϵ ∼ N(0, I).The boundary conditions are α_0 = σ_T = 1 and α_T = σ_0 = 0.
  • Probability-flow ODEs: The interpolant admits a probability-flow ODE whose distribution at time t equals the marginal p_t(x), enabling sampling from Gaussian noise with ODE solvers.The paper cites the Euler sampler as an example.
  • Velocity modeling: The ODE velocity is v(x, t) = E[ẋ_t | x_t = x] = ẋα_t E[x* | x_t = x] + ẋσ_t E[ϵ | x_t = x].A model v_θ(x_t, t) approximates this velocity through a training objective.
  • Reverse SDEs: An equivalent reverse SDE shares the PF ODE’s marginal distributions and provides an alternative route for data generation using SDE solvers.Its formulation involves a diffusion coefficient w_t and a score s(x_t, t).
  • Interpolant choice: The main experiments use T = 1 with the linear interpolant α_t = 1 − t and σ_t = t, while the approach also applies to variants such as DDPM.DDPM uses a discretized process and different α_t, σ_t, with N(0, I) as an equilibrium distribution.

3 REPA: REGULARIZATION FOR REPRESENTATION ALIGNMENT

REPA addresses diffusion models’ representation-learning bottleneck by distilling pretrained self-supervised visual representations into diffusion transformers. It aligns projected hidden states from noisy inputs with clean, patch-wise visual representations and adds this alignment to standard diffusion objectives.

  • Method: REPA aligns projected diffusion-transformer hidden states with clean-image representations from pretrained self-supervised visual encoders.The method guides hidden states from noisy inputs toward noise-invariant, semantically useful representations for subsequent reconstruction.
  • Motivation: Diffusion transformers show a substantial semantic gap from state-of-the-art visual encoders, despite achieving a reasonably high linear-probing peak.The gap is measured against DINOv2 on ImageNet using a pretrained SiT model.
  • Motivation: Representation alignment improves with larger models and longer training, but remains below that of self-supervised visual encoders even after 7M iterations.This pattern is reported using CKNNA across model sizes and training iterations, including comparisons with MoCov3 and DINOv2.
  • Method: The alignment operates patch-wise by maximizing similarities between an MLP projection of the diffusion encoder output and the pretrained encoder output.For clean image x*, the pretrained representation y*=f(x*) has N patches and embedding dimension D; the projected hidden state has matching dimensions.
  • Method: REPA adds the alignment term to the original diffusion objective, with λ > 0 controlling the tradeoff between denoising and representation alignment.The paper primarily studies Improved DDPM in DiT and linear stochastic interpolants in SiT.

4 EXPERIMENTS

Experiments show that REPA substantially improves diffusion-transformer generation, representation quality, and convergence across model sizes and visual encoders. The gains persist across alignment choices and culminate in state-of-the-art generation quality.

  • Component Analysis: REPA consistently improves generation performance across diffusion-transformer design choices, yielding much better FID than vanilla models.
  • Target Representations: Stronger pretrained target representations improve both generation quality and linear-probing performance, with encoder quality strongly correlated with aligned-model performance.
  • Alignment Depth: Regularizing only the first few layers is sufficient, while attaching REPA at layer 6 or 8 can further enhance generation by leaving later layers for high-frequency details.
  • Alignment Objective: NT-Xent performs better early at 50-100K iterations, but its advantage diminishes over time, motivating cosine similarity for later experiments.
  • Scalability: As diffusion-transformer size increases, REPA provides increasingly significant convergence speed-ups, while larger models reach better performance more quickly with a fixed target encoder.

5 RELATED WORK

Related work connects diffusion models with representation learning through hybrid generation-and-understanding models and through analyses or uses of diffusion representations. Other studies leverage pretrained visual encoders to improve diffusion-model efficiency and performance, including two-stage text-to-image frameworks.

  • Bridging diffusion models and representation learning: Hybrid approaches train models for both diffusion-based generation and an understanding task such as classification or segmentation.Yang et al. (2022) and Deja et al. (2023) combine classification with generation, while Tian et al. (2024) combines segmentation with generation.
  • Bridging diffusion models and representation learning: Recent works also analyze and exploit representations learned from diffusion models.The passage places these studies within broader efforts to exploit or improve diffusion-model representations.
  • Diffusion models with external representations: Pretrained visual encoders have been used to enhance diffusion-model efficiency and performance.The cited studies include Pernias et al. (2024) and Li et al. (2024).
  • Diffusion models with external representations: Würstchen uses a two-stage text-to-image framework in which a text-conditioned model generates a semantic map before image synthesis.A second diffusion model synthesizes images based on the generated semantic map.

6 CONCLUSION … C ANALYSIS DETAILS

The paper concludes that REPA aligns diffusion-transformer representations with self-supervised representations, improving generation performance and convergence speed. Its supporting sections describe diffusion-based generative processes and the DiT/SiT architecture used.

  • 6 CONCLUSION: REPA aligns diffusion-transformer representations with recent self-supervised representations, improving generation performance while accelerating convergence.The authors present REPA as a simple regularization and identify faster convergence alongside improved generation.
  • A DESCRIPTIONS FOR DIFFUSION-BASED MODELS: The paper describes generative models that learn target distributions by training variants of denoising autoencoders.The overview covers denoising diffusion probabilistic models and stochastic interpolants.
  • A.1 DENOISING DIFFUSION PROBABILISTIC MODELS: Diffusion models learn a reverse denoising process from Gaussian noise N(0, I) toward the target distribution p(x), reversing a forward noise-adding process.Generation begins from random noise and iteratively applies the reverse process to reach a data sample.
  • A.1 DENOISING DIFFUSION PROBABILISTIC MODELS: DDPM trains its denoising network with a simple denoising autoencoder objective over the predefined forward and reverse processes.The forward process gradually adds Gaussian noise, while the reverse process removes it step by step.
  • A.2 STOCHASTIC INTERPOLANTS: Stochastic interpolants define a continuous-time path between data x∗∼p(x) and Gaussian noise ϵ∼N(0, I), enabling flow-based modeling.The path uses time-dependent functions α_t and σ_t, with a probability-flow ordinary differential equation matching the marginal distribution.
  • A.2 STOCHASTIC INTERPOLANTS: The interpolant velocity is represented through conditional expectations and approximated by a model trained with a corresponding objective.The score can likewise be modeled, but it can also be computed directly from the velocity for t > 0.
  • A.2 STOCHASTIC INTERPOLANTS: Stochastic interpolants require positive differentiable schedules with α_1=σ_0=0 and α_0=σ_1=1, allowing unbiased interpolation between data and noise.Linear and variance-preserving schedules are examples, and the diffusion coefficient can be chosen after training for reverse-SDE sampling.
  • B DIFFUSION TRANSFORMER ARCHITECTURE: The model follows DiT and SiT by patchifying an input sequence and operating on a downsampled latent image z=E(x) from a stable-diffusion VAE.The architecture is described as closely resembling a vision transformer.

C.1 EVALUATION DETAILS · C.2 DIT ANALYSIS · C.3 DESCRIPTION OF PRETRAINED VISUAL ENCODERS •

The appendices define CKNNA evaluation, analyze DiT representations against DINOv2, and describe pretrained visual encoders and their dataset effects. Together, they show that DiT representations are meaningful but weakly aligned and substantially behind DINOv2.

  • C.1 EVALUATION DETAILS: CKNNA relaxes CKA by measuring alignment through k-nearest-neighbor embeddings rather than all dataset pairs.CKA measures global model similarity using kernel matrices computed from network representations.
  • C.1 EVALUATION DETAILS: 10,000 ImageNet validation images are sampled randomly, and CKNNA is reported with k = 10 because smaller k showed better alignment.The evaluation follows an observation from Huh et al. (2024).
  • C.1 EVALUATION DETAILS: Linear probing trains a linear layer for 90 epochs with batch size 16,384, parameter-free batch normalization, Adam, cosine decay, and initial learning rate 0.001.This setup follows DAE (Chen et al., 2024c).
  • C.2 DIT ANALYSIS: DiT learns meaningful discriminative representations, but linear probing shows a large gap compared with DINOv2.The analysis uses the pretrained DiT-XL/2 model trained with 7M iterations.
  • C.2 DIT ANALYSIS: DiT is weakly aligned with DINOv2, and its absolute CKNNA value remains small.Figure 10 measures feature alignment between DINOv2-g and DiT-XL/2 using CKNNA.
  • C.3 DESCRIPTION OF PRETRAINED VISUAL ENCODERS •: The pretrained visual encoders include MAE, DINO, MoCov3, CLIP, DINOv2, I-JEPA, and SigLIP, spanning reconstruction, self-distillation, contrastive, joint-embedding, and image-text objectives.MAE reconstructs masked patches; DINO uses a momentum teacher; MoCov3 scales MoCo to vision transformers; CLIP uses image-text contrastive learning; DINOv2 combines pixel- and patch-level discrimination; I-JEPA predicts missing image parts; SigLIP uses pairwise sigmoid loss.
  • C.3 DESCRIPTION OF PRETRAINED VISUAL ENCODERS •: Better visual representations trained on massive image datasets provide greater improvement, regardless of whether their training data include ImageNet.Encoders such as SigLIP and CLIP are not fine-tuned on ImageNet when trained on other datasets, separating dataset-leakage effects.

D HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS · E EVALUATION DETAILS · F BASELINES

The paper specifies implementation choices for training, pretrained encoders, sampling, and computational resources, then defines evaluation procedures and summarizes baseline methods. Evaluation follows ADM’s reference setup and uses established generation metrics and diffusion-model baselines.

  • D HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS: The model follows the original SiT implementation and DiT/SiT structure, using AdamW with a constant 1e-4 learning rate, β1=0.9, β2=0.999, no weight decay, fp16, gradient clipping, and precomputed compressed latents.These choices are used throughout the experiments to accelerate training.
  • D HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS: Pretrained encoder checkpoints come from RCG’s MoCov3-B/-L implementation or official releases, with positional embeddings interpolated when patch counts differ.The interpolation adapts pretrained encoders to diffusion transformers with different numbers of patches.
  • D HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS: The sampler uses Euler-Maruyama with the Eq. (5) SDE, diffusion coefficient w_t=σ_t, and a final SDE step of 0.04, which significantly improves performance.The reported improvement is similar to that in the original SiT paper.
  • D HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS: 8 NVIDIA H100 80GB GPUs support training at about 5.4 step/s with batch size 256, with further speedups possible through engineering such as pretrained-feature precomputation.The stated training speed can be increased with additional optimization.
  • E EVALUATION DETAILS: Evaluation strictly follows ADM’s setup and reference batches, using H100 80GB or 4090Ti GPUs with tf32 precision, whose generation difference from fp32 is negligible.The reference batches and official implementation are retained for evaluation consistency.
  • E EVALUATION DETAILS: FID measures Inception-v3 feature distance between real and generated-image distributions, while sFID uses intermediate spatial features to capture spatial distribution.FID assumes multivariate Gaussian feature distributions; the supplied passage also introduces IS as an evaluation metric using Inception-v3.
  • F BASELINES: The evaluation compares against ADM, VDM++, Simple diffusion, and CDM, representing approaches involving classifier-guided sampling, adaptive noise schedules, simplified diffusion designs, and other diffusion baselines.The supplied baseline passage identifies ADM, VDM++, Simple diffusion, and CDM as baseline methods, though its CDM description is truncated.

G DETAILED QUANTITATIVE RESULTS · H MORE QUALITATIVE RESULTS · I MORE DISCUSSION ON RELATED WORK

The detailed results show that REPA improves SiT training speed and performance across settings, while qualitative figures illustrate uncurated class-conditional generations. Related work situates REPA among pretrained-encoder methods, denoising transformers, auxiliary self-supervised objectives, and denoising-based representation learning.

  • G DETAILED QUANTITATIVE RESULTS: SiT models are aligned with DINOv2-B representations using λ = 0.5 and negative cosine similarity, with layer selection varying by model size.The base model uses 4th-layer hidden states, whereas large and xlarge models use 8th-layer hidden states.
  • G DETAILED QUANTITATIVE RESULTS: SiT-XL/2+REPA is evaluated at 4M iterations with different classifier-free guidance scales, including settings using the guidance interval.The detailed evaluations cover both varying guidance scale w and guidance-interval application.
  • G DETAILED QUANTITATIVE RESULTS: REPA provides great speedup and performance improvement over the vanilla SiT-XL/2 model when FID is measured with classifier-guidance scale w = 1.35.Figure 11 reports training iteration versus FID, with values measured using classifier-free guidance.
  • H MORE QUALITATIVE RESULTS: Additional figures show uncurated SiT-XL/2+REPA generations with classifier-free guidance w = 4.0 across animal, object, food, and landscape classes.Examples include loggerhead sea turtle, macaw, golden retriever, acoustic guitar, cheeseburger, coral reef, and volcano.
  • I MORE DISCUSSION ON RELATED WORK: Pretrained visual encoders have been used in GAN training and distillation to accelerate convergence by serving as discriminators through intermediate features.The cited approaches include Sauer et al., Kumari et al., and Kang et al.
  • I MORE DISCUSSION ON RELATED WORK: Denoising transformers such as U-ViT, MDT, DiffiT, and DiT establish transformer backbones as effective architectures for diffusion or flow-based model training.The passage notes that skip connections can help, while DiT demonstrates they are not necessary for a pure transformer architecture.
  • I MORE DISCUSSION ON RELATED WORK: Other generative-model approaches accelerate training by adding auxiliary self-supervised tasks, including mask reconstruction, discriminative losses, and masked image modeling.Examples include MaskDiT, SD-DiT, and MAGE.
  • I MORE DISCUSSION ON RELATED WORK: Prior work also treats denoising as a self-supervised learning task for improving representation learning or producing compact latent representations.Examples include extensions of diffusion objectives, diffusion-guided encoders, and analyses of denoising-based representation learning.

J IMAGENET 512×512 EXPERIMENT

The ImageNet 512×512 experiment tests REPA’s scalability by adapting the SiT input to compressed 64 × 64 × 4 latents and resizing DINOv2 inputs to 448 × 448. With more than three times fewer training iterations, SiT-XL/2 with REPA outperforms vanilla SiT-XL/2 on FID, sFID, IS, and Prec.

  • Experimental setup: The 512×512 setup uses SiT inputs formed as 64 × 64 × 4 compressed latent images from 512 × 512 × 3 pixels.The latents are produced using the stable diffusion VAE.
  • Experimental setup: 448 × 448 resized images are used as inputs to DINOv2 in the ImageNet 512×512 experiment.The experiment otherwise follows the ImageNet 256×256 setup.
  • Quantitative results: >3× fewer training iterations still allow SiT-XL/2 with REPA to outperform vanilla SiT-XL/2 on FID, sFID, IS, and Prec.The comparison is reported in Table 11.
  • Qualitative results: Samples from SiT-XL/2+REPA on ImageNet 512×512 are shown using classifier-free guidance with w = 4.0.The qualitative results appear in Figure 36.
  • Quantitative results: The system-level comparison on ImageNet 512×512 uses classifier-free guidance with w = 1.35.These quantitative results are provided in Table 11.

K TEXT-TO-IMAGE GENERATION EXPERIMENT · L FEATURE MAP VISUALIZATION

REPA improves text-to-image generation in MMDiT models trained on MS-COCO, outperforming the vanilla model qualitatively and quantitatively. Feature visualizations further show that REPA produces coarse-to-fine representations, whereas vanilla features become noisy at large t.

  • K TEXT-TO-IMAGE GENERATION EXPERIMENT: The text-to-image experiments train MMDiT from scratch on MS-COCO for 150K iterations, following the U-ViT setup unless otherwise specified.MMDiT jointly computes attention with image patches and text embeddings, using separate train and validation splits.
  • K TEXT-TO-IMAGE GENERATION EXPERIMENT: REPA consistently outperforms the vanilla model in qualitative text-to-image comparisons on MS-COCO.The comparison uses classifier-free guidance with w = 4.0.
  • K TEXT-TO-IMAGE GENERATION EXPERIMENT: REPA yields considerable improvements in quantitative text-to-image generation, underscoring visual-representation alignment even when text representations are present.The quantitative comparison follows the MS-COCO setup with classifier-free guidance w = 2.0.
  • K TEXT-TO-IMAGE GENERATION EXPERIMENT: The authors identify training large-scale text-to-image models with large-scale data and REPA as a promising future direction.This statement is presented as the authors’ forward-looking view rather than a reported experiment.
  • L FEATURE MAP VISUALIZATION: PCA visualizations show that REPA produces coarse-to-fine feature maps across layers.The visualizations are provided for SiT-XL/2 and SiT-XL/2+REPA, following a presentation similar to DINOv2.
  • L FEATURE MAP VISUALIZATION: The vanilla model tends to produce noisy feature maps, particularly at large t.This contrasts with the coarse-to-fine structure observed under REPA.

M LIMITATIONS AND FUTURE WORK

The paper identifies future work on understanding and improving REPA’s alignment depth, extending it across data types and text-to-image settings, developing theory, and making its weighting time-varying.

  • Alignment depth: Applying REPA at layer 8 was more beneficial than applying it to later transformer-layer embeddings, motivating deeper analysis of alignment depth.The authors identify this analysis as a direction for further improving REPA.
  • Different input data types: REPA should be explored beyond latent image diffusion, including pixel-level diffusion, video domains, and large-scale text-to-image diffusion models.The text-to-image direction is motivated by the reported MS-COCO generation results.
  • Theoretical analysis: Future theoretical work should explain why REPA works and examine the relationship between representations learned through instance discrimination and denoising objectives.These objectives are proposed as a concrete comparison for theoretical analysis.
  • Time-varying REPA: A time-varying REPA scheme could use a weight function based on the diffusion process’s noise schedule, which the authors did not explore in this work.The authors leave this design as a future direction while focusing on other analyses, including target representations, alignment depth, and scalability.
Loading 2410.06940v4…