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

Lei Wang (Nankai University) · Xiang Li (Nankai University) · Ming-Ming Cheng (Nankai University) · Hongcheng Gao (University of the Chinese Academy of Sciences) · Shen Zhang (JIIOV Technology) · Zhaowei Chen · Ge Wu (Nankai University) · Yao Tang (Jiiov Technology) · Ruijing Shi (Nankai University) · Shanghua Gao (Harvard University) · Zhenyuan Chen (Nankai University) · jian Yang (Nankai University)

convergence accelerationdenoising networksdiffusion models discriminative representations flops global semanticshigh-level class tokenimage latentsimage-class pairsimagenet inference overheadreparepresentation entanglement for generationsemantic knowledge training efficiency visual representations

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 \textit{\textbf{R}epresentation \textbf{E}ntanglement for \textbf{G}eneration} (\textbf{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.