Guiding a Diffusion Transformer with the Internal Dynamics of Itself

TL;DR

This work tackles the challenge of diffusion models failing to cover low-probability regions by proposing Internal Guidance (IG), a plug-and-play strategy that adds auxiliary supervision to an intermediate layer and uses extrapolation between intermediate and final outputs during sampling to steer generation. IG improves both training efficiency and sample quality for Diffusion Transformers, and can be combined with classifier-free guidance (CFG) and guidance interval techniques to achieve state-of-the-art results (e.g., $\text{FID}=1.19$ on ImageNet-256 with LightningDiT-XL/1+IG+CFG). The authors demonstrate significant gains across multiple architectures (SiT, DiT, LightningDiT), show that IG reduces gradient-vanishing issues, and provide analyses of its compatibility with CFG, optimal guidance intervals, and training acceleration benefits. The approach offers a practical, scalable enhancement for large diffusion models with minimal computational overhead, broadening the applicability of high-fidelity image synthesis.

Abstract

The diffusion model presents a powerful ability to capture the entire (conditional) data distribution. However, due to the lack of sufficient training and data to learn to cover low-probability areas, the model will be penalized for failing to generate high-quality images corresponding to these areas. To achieve better generation quality, guidance strategies such as classifier free guidance (CFG) can guide the samples to the high-probability areas during the sampling stage. However, the standard CFG often leads to over-simplified or distorted samples. On the other hand, the alternative line of guiding diffusion model with its bad version is limited by carefully designed degradation strategies, extra training and additional sampling steps. In this paper, we proposed a simple yet effective strategy Internal Guidance (IG), which introduces an auxiliary supervision on the intermediate layer during training process and extrapolates the intermediate and deep layer's outputs to obtain generative results during sampling process. This simple strategy yields significant improvements in both training efficiency and generation quality on various baselines. On ImageNet 256x256, SiT-XL/2+IG achieves FID=5.31 and FID=1.75 at 80 and 800 epochs. More impressively, LightningDiT-XL/1+IG achieves FID=1.34 which achieves a large margin between all of these methods. Combined with CFG, LightningDiT-XL/1+IG achieves the current state-of-the-art FID of 1.19.

Figures (11)

• Figure 1: Visualization Results. We visualize our latent diffusion system with proposed IG together with LightningDiT-XL trained on ImageNet 256 × 256 resolution. With an IG scale of 1.4 and a CFG scale of 1.45, and further combining the guidance interval, we ultimately achieved the state-of-the-art FID = 1.19. More uncurated samples are provided in Supplementary Material.
• Figure 2: The overall framework of our proposed Internal Guidance. We introduce an additional auxiliary supervision loss during training, and utilize the intermediate layer outputs during sampling process to guide the final outputs.
• Figure 3: A fractal-like 2D distribution with two classes indicated with gray and orange regions. Approximately 99% of the probability mass is inside the shown contours. (a) Conditional sampling using a small denoising diffusion model generates outliers due to its limited fitting capability. (b) Classifier-free guidance (w = 2.5) eliminates the vast majority of outliers but reduces diversity by over-emphasizing the class. (c) Internal guidance (w = 2) can maintain diversity as Autoguidance karras2024guiding while allowing some outliers at the ends of the branches. (d) The combination of IG and CFG can significantly reduce outliers without reducing diversity.
• Figure 4: Internal guidance inspires new training acceleration methods. (a) Conditional sampling using a not well-trained denoising diffusion model generates a large number of outliers. (b) Outliers can gradually be eliminated with sufficient training. (c) Internal guidance can also eliminate outliers with not well-trained denoising diffusion model. (d) The loss function we proposed can accelerate convergence. In SiT-B/2 experiments, our proposed loss function demonstrates superior accelerated convergence performance compared to REPA yu2024representation.
• Figure 5: The combination of IG and CFG can yield a higher FID value, which is superior to simply apply CFG. In particular, the FID value generated by the combination of a lower IG coefficient and CFG is higher than that of a higher coefficient.