PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ge Wu, Shen Zhang, Ruijing Shi, Shanghua Gao, Zhenyuan Chen, Lei Wang, Zhaowei Chen, Hongcheng Gao, Yao Tang, Jian Yang, Ming-Ming Cheng, Xiang Li

arXiv:2507.01467v2cs.CV

TL;DR

Existing external alignment methods do not use discriminative representations throughout denoising, limiting how directly those representations guide generation. REG entangles image latents with a pretrained foundation model’s class token, and on ImageNet it substantially accelerates convergence with less than 0.5% additional FLOPs and latency.

  • Problem

    External alignment between diffusion features and pretrained visual representations is absent during denoising inference, limiting direct use of discriminative representations.

  • Method

    REG synchronously noise-injects and spatially concatenates image latents with one pretrained foundation model class token for joint denoising.

  • Results

    63× and 23× faster training convergence are achieved by SiT-XL/2 + REG than by SiT-XL/2 and SiT-XL/2 + REPA, respectively, on ImageNet 256×256.

  • Takeaways & Limitations

    REG jointly reconstructs image latents and global semantics, allowing semantic knowledge to guide generation while adding less than 0.5% FLOPs and latency.

Abstract

from arXiv · show

REPA and its variants effectively mitigate training challenges in diffusion models by incorporating external visual representations from pretrained models, through alignment between the noisy hidden projections of denoising networks and foundational clean image representations. We argue that the external alignment, which is absent during the entire denoising inference process, falls short of fully harnessing the potential of discriminative representations. In this work, we propose a straightforward method called Representation Entanglement for Generation (REG), which entangles low-level image latents with a single high-level class token from pretrained foundation models for denoising. REG acquires the capability to produce coherent image-class pairs directly from pure noise, substantially improving both generation quality and training efficiency. This is accomplished with negligible additional inference overhead, requiring only one single additional token for denoising (<0.5\% increase in FLOPs and latency). The inference process concurrently reconstructs both image latents and their corresponding global semantics, where the acquired semantic knowledge actively guides and enhances the image generation process. On ImageNet 256$\times$256, SiT-XL/2 + REG demonstrates remarkable convergence acceleration, achieving $\textbf{63}\times$ and $\textbf{23}\times$ faster training than SiT-XL/2 and SiT-XL/2 + REPA, respectively. More impressively, SiT-L/2 + REG trained for merely 400K iterations outperforms SiT-XL/2 + REPA trained for 4M iterations ($\textbf{10}\times$ longer). Code is available at: https://github.com/Martinser/REG.

1 Introduction

REG addresses the limitations of external representation alignment by entangling image latents with a pretrained model’s class token during denoising. It improves generation quality, convergence speed, and semantic learning with negligible overhead, including 63× and 23× faster convergence than SiT and REPA.

  • Motivation: External pretrained-vision alignment improves generative training, but its features remain stronger than generative-model features and are absent during inference.This motivates incorporating discriminative information directly into denoising rather than relying only on external alignment.
  • Contribution: REG entangles low-level image latents with a single high-level class token from pretrained foundation models for denoising.The method explicitly reflows discriminative information into the generation process.
  • Contribution: REG enhances generation quality, training convergence speed, and discriminative semantic learning while adding negligible computational overhead.The method adds only one token, with less than 0.5% additional FLOPs and latency.
  • Results: 63× and 23× faster training convergence are achieved by SiT-XL/2 + REG than by SiT and REPA on ImageNet generation benchmarks.The comparison is reported for class-conditional ImageNet generation at 256×256 resolution.

2 Related work

Prior work advances diffusion models through latent-space generation, transformer architectures, representation learning, auxiliary components, and pretrained visual features. REG differs by making a single discriminative class token part of the denoising input instead of using only external alignment or additional models.

  • Generative models for image generation: Diffusion generation has progressed from pixel-space DDPM and DDIM to latent-space LDM and transformer-based DiT and SiT architectures.SiT uses continuous-time stochastic interpolants for diffusion training.
  • Generative models as representation learners: Intermediate diffusion features encode semantic information useful for segmentation, depth estimation, controllable editing, and knowledge transfer.These results position generative models as representation learners across multiple vision tasks.
  • Generative models with external representations: RCG generates a class token with a secondary diffusion model, whereas REG uses one class token as part of the input without an additional model.REG uses the token to provide discriminative guidance while supporting conditional generation.
  • Generative models with external representations: REPA and REPA-E align diffusion features with pretrained vision representations, but do not use discriminative representations as denoising inputs.REG addresses this distinction by incorporating a class token directly into the denoising process.

3 Method

REG extends SiT by jointly denoising image latents and a foundation-model class token, while aligning hidden representations with visual features. Its design incorporates discriminative guidance into inference with minimal computational overhead.

  • REG framework: REG jointly denoises image latents and a single foundation-model class token from paired noisy inputs.Both inputs receive synchronized noise, are spatially concatenated, and enter the SiT blocks together.
  • Representation alignment: REG aligns projected hidden states with concatenated foundation-model class and image representations using cosine similarity at selected transformer layers.The alignment uses layer n = 4 for SiT-B/2 + REG and n = 8 for other variants.
  • Training objective: The training objective combines velocity prediction for image latents and class tokens with a weighted representation-alignment loss.β controls the relative image-token prediction weighting, while λ controls alignment relative to denoising.
  • REG framework: The class token provides discriminative guidance by enabling REG to reconstruct global semantics alongside image latents from random noise.The inference process requires no auxiliary network to generate the class token.
  • Efficiency: REG adds less than 0.5% FLOPs and latency at 256×256 resolution by introducing only one global class token.This overhead is reported as almost negligible compared with the base generation process.

4 Experiments

Experiments evaluate REG’s performance, computational cost, discriminative guidance, and ablations on ImageNet 256×256. REG consistently improves generation quality and semantic alignment while accelerating convergence with negligible overhead.

  • Setup: REG is evaluated on ImageNet 256×256 using SiT architectures, standardized preprocessing, and FID, sFID, IS, precision, and recall metrics.Models use Stable Diffusion VAE latents and 2 × 2 patch processing; evaluation uses 50K generated samples.
  • Model performance: 63× and 23× faster convergence are achieved by SiT-XL/2 + REG than SiT-XL/2 and REPA, respectively.REG reaches comparable performance in 110K and 170K steps, versus 7M and 4M steps.
  • Model performance: 17× faster training lets REG match SiT-XL’s quality in 80 epochs instead of 1400, while surpassing REPA’s 800-epoch performance at 480 epochs.These comparisons use classifier-free guidance and the same guidance interval as REPA.
  • Computational cost: REG adds only 0.30% parameters, 0.38% FLOPs, and 0.49% latency relative to REPA while reducing FID by 56.46% and increasing IS by 50.19%.The reported latency is 6.21s for REG versus 6.18s for REPA.
  • Ablation studies: DINOv2-B provides the best target representation, attaining FID 15.22 and IS 94.64 among evaluated pretrained vision encoders.All evaluated target representations outperform REPA in this comparison.
  • Ablation studies: REG improves over REPA across alignment depths, with FID reductions ranging from 4.19 to 7.16 points.The authors attribute this to direct class-token insertion, which supplies global guidance across layers.
  • Ablation studies: β = 0.03 achieves the best overall performance across evaluation metrics and is adopted as the default loss weight.β controls the contribution of the class-token alignment loss.
  • Ablation studies: DINOv2 class-token entanglement lowers FID by 9.18 and raises IS to 94.64, outperforming learnable-token and averaged-latent alternatives.The results identify high-level discriminative information and the entanglement method as important factors.

5 Conclusion

REG introduces image-class denoising by entangling low-level image latents with a high-level class token from a pretrained foundation model. The method uses synchronized noise injection and spatial concatenation so denoising reconstructs both components.

  • 5 Conclusion: REG introduces an image-class denoising paradigm by entangling low-level image latents with one high-level class token from a pretrained foundation model.This replaces a pure image-denoising formulation with joint image-class reconstruction.
  • 5 Conclusion: Synchronized noise injection and spatial concatenation connect the image latents and class token during training and denoising.The denoising process simultaneously reconstructs the image and its global semantics.

A Discriminative semantics in inference

REG preserves discriminative semantics during inference by jointly denoising image latents and a class token. Compared with OLT, REG’s class token retains substantially stronger similarity to pretrained reference representations.

  • Architecture: REG jointly denoises low-level latent features and one high-level noised class token, applying velocity prediction losses to both components.OLT instead concatenates a learnable token with noised latents and computes velocity loss only on dense features.
  • Inference evaluation: REG’s inference concatenates noised latents with a noise-initialized class token before multi-step denoising.The comparison processes 10,000 ImageNet validation images with identical noise injection and a SiT-B/2 backbone.
  • Inference evaluation: 0.458 cosine similarity and 0.514 CKNNA are achieved by REG’s class token, versus 0.007 and 0.014 for OLT’s learnable token.The reference representation is a DINOv2-g class token.

B Analysis of training overhead in REG

REG reaches the performance levels of much longer-trained SiT and REPA models with substantially fewer training steps and GPU hours. The reported overhead comparison emphasizes reduced training cost rather than added inference complexity.

  • Training overhead: 110K steps let REG reach SiT’s 7M-step performance level, reducing GPU hours by 98.36%.All overhead experiments use 8 NVIDIA A40 GPUs.
  • Training overhead: REG matches REPA’s 4M-step performance with only 170K iterations.The supplied passage reports the iteration comparison but truncates the corresponding GPU-hour reduction.
  • Training overhead: Table 7 compares training overhead while reporting comparable performance and significantly reduced training time for REG.The table summarizes the comparison across other models.

C 256×256 ImageNet

REG shows accelerated convergence and strong performance on ImageNet at 256×256 and 512×512 resolutions, including comparisons across training durations and model scales.

  • 256×256 ImageNet: 1.40 FID at 480 epochs demonstrates REG’s strong ImageNet 256×256 performance under classifier-free guidance.This performance is comparable to REPA at 800 epochs while using fewer than 40% of the training iterations.
  • 256×256 ImageNet: Fewer than 40% of REPA’s training iterations are required for REG to achieve comparable performance at 800 epochs.
  • 256×256 ImageNet: REG’s performance analysis across model scales further highlights its robustness and accelerated convergence without classifier-free guidance.
  • 512×512 ImageNet: At 80 epochs on ImageNet 512×512, REG surpasses REPA trained for 200 epochs and SiT-XL/2 trained for 600 epochs in FID.
  • Guidance analysis: Table 10 evaluates SiT-XL + REG at 2.4M training iterations across classifier-free guidance scales using the guidance interval method.

E Experimental setup

The experiments use DINOv2-based representation alignment and standardized optimization and inference settings across SiT model scales.

  • Hyperparameter setup: DINOv2-B is adopted as the optimal alignment model, with 250 denoising steps used for all inference processes.
  • Hyperparameter setup: Table 12 reports hyperparameter configurations for SiT + REG across different model scales.
  • CKNNA score: CKNNA measures local topological alignment by retaining pairwise relationships shared among the k-nearest neighbors of two representation spaces.
  • CKNNA score: The CKNNA computation uses truncated kernels, weighted covariance, and normalization to remove scale dependencies.
  • CKNNA score: CKNNA scores are computed from spatially averaged dense features, explicitly excluding the class token for fair comparison.

F More visualization results

Additional visualizations present SiT-XL/2 + REG outputs across diverse ImageNet classes using classifier-free guidance with w = 4.0.

  • Visualization protocol: Visualization results use classifier-free guidance with w = 4.0 across the reported image classes.
  • Animal classes: The visualizations cover animal classes including great white shark, bald eagle, great grey owl, macaw, koala, and timber wolf.
  • Animal classes: Additional animal visualizations include sulphur-crested cockatoo, American coot, lesser panda, border collie, polecat, and giant panda.
  • Non-animal classes: The remaining visualizations include castle, China cabinet, convertible, bubble, geyser, lakeside, and volcano.
Loading 2507.01467v2…