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

TL;DR

This paper introduces Representation Entanglement for Generation (REG), a simple method that entangles image latents with high-level class tokens from pretrained models, significantly improving diffusion model training efficiency and generation quality with minimal inference overhead.

Contribution

The paper proposes REG, a novel technique that entangles low-level image latents with high-level class tokens, enabling faster training and better generation quality in diffusion models.

Findings

Achieves 63x and 23x faster training on ImageNet compared to baseline methods.

Produces coherent image-class pairs directly from noise with minimal additional inference cost.

Outperforms longer-trained models with significantly fewer training iterations.

Abstract

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.