No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves

TL;DR

This work shows that diffusion transformers can generate and leverage meaningful internal representations without external priors by introducing Self-Representation Alignment (SRA), a self-distillation mechanism using an EMA teacher to align early noisy latent representations with later less-noisy, more discriminative ones. SRA combines a simple projection-based alignment loss with the standard generative objective, yielding consistent gains for DiT and SiT, especially in larger models, and outperforming approaches that rely on external representation priors. The study provides extensive ablations demonstrating the importance of alignment depth, noise-timing, and projection heads, and establishes a strong link between representation quality (via linear probing) and generation performance. Overall, SRA offers a lightweight, plug-in method to enhance diffusion transformers’ representation learning and generation quality without extra representation modules.

Abstract

Recent studies have demonstrated that learning a meaningful internal representation can both accelerate generative training and enhance the generation quality of diffusion transformers. However, existing approaches necessitate to either introduce an external and complex representation training framework or rely on a large-scale, pre-trained representation foundation model to provide representation guidance during the original generative training process. In this study, we posit that the unique discriminative process inherent to diffusion transformers enables them to offer such guidance without requiring external representation components. We therefore propose Self-Representation Alignment (SRA), a simple yet straightforward method that obtains representation guidance through a self-distillation manner. Specifically, SRA aligns the output latent representation of the diffusion transformer in the earlier layer with higher noise to that in the later layer with lower noise to progressively enhance the overall representation learning during only the generative training process. Experimental results indicate that applying SRA to DiTs and SiTs yields consistent performance improvements. Moreover, SRA not only significantly outperforms approaches relying on auxiliary, complex representation training frameworks but also achieves performance comparable to methods that are heavily dependent on powerful external representation priors.

Figures (24)

• Figure 2: We empirically investigate the representations in diffusion transformers across different blocks and timesteps with the original SiT-XL/2 checkpoint trained for 7M iterations. Left: Using PCA pca to visualize the latent features in SiT, we observe that the features lead a process from coarse to fine when increasing block layers and decreasing noise level. Right: A similar trend can also be seen in the linear probing results on ImageNet. Investigation of DiT is provided in Appendix \ref{['appen:dit_rep']}, which leads to the similar results as SiT.
• Figure 3: Overall framework. SRA aligns the student's latent representation in the earlier layer conditioned on higher noise (green branch) to that of the teacher in the later layer conditioned on lower noise (blue branch) to achieve self-representation alignment. We use a stop-gradient (sg) operator on the teacher to let gradients flow only through the student, and update the teacher's parameters with an exponential moving average (ema) of the student's parameters.
• Figure 4: FID comparisons with vanilla DiTs and SiTs across different model sizes on ImageNet 256$\times$256 without classifier-free guidance (CFG).
• Figure 5: Selected samples on ImageNet 256$\times$256 from the SiT-XL + SRA. We use classifier-free guidance with $w$ = 4.0. More uncurated samples are provided in Appendix \ref{['appen:samples']}.
• Figure 6: We empirically investigate the effect of representations in SRA. Left and Middle: Linear probing result of vanilla SiT-XL trained for 1400 epochs and SiT-XL + SRA trained for 800 epochs. Right: Linear probing vs. FID plot of SiT-XL + SRA with different teacher's output layers for alignment (similar plot of DiT + SRA is provided in Appendix \ref{['appen:ab_dit']}).