Source-linked AI summary
Guiding a Diffusion Transformer with the Internal Dynamics of Itself
Xingyu Zhou, Qifan Li, Xiaobin Hu, Hai Chen, Shuhang Gu
TL;DR
Diffusion models may lack training and data to generate high-quality images in low-probability regions, while existing guidance approaches have trade-offs or added costs. The paper proposes Internal Guidance (IG), supervising an intermediate layer and extrapolating its output with the final output at sampling. On ImageNet 256×256, LightningDiT-XL/1+IG reaches FID 1.19 with CFG and guidance interval, and the paper also reports a training-side benefit from intermediate supervision.
Problem
Diffusion models may lack sufficient training and data to generate high-quality images in low-probability regions, while CFG can distort or oversimplify samples and other self-guidance methods may require extra design, training, or sampling.
Source: §1 · p. 1Method
IG adds auxiliary supervision at an intermediate layer during training, then extrapolates intermediate and deep outputs during sampling to guide generation.
Source: §1 · p. 1, §3.2 · p. 3Results
On ImageNet 256×256, LightningDiT-XL/1+IG achieves FID 1.19 at 680 epochs when combined with CFG and guidance interval.
Source: §1 · p. 2Takeaways & Limitations
IG improves generation quality while maintaining diversity without additional sampling steps, and intermediate supervision alone alleviates partial vanishing gradients.
Source: §1 · pp. 1–2A large IG coefficient applied before the model is fully trained may introduce new outliers.
Source: §4.1 · p. 4
Abstract
from arXiv · showhide
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.
1. Introduction
The paper introduces Internal Guidance (IG), which uses intermediate Diffusion Transformer outputs to guide sampling, as an alternative to guidance approaches with trade-offs or added costs. Experiments report improved generation results and training benefits across several model families.
1. Introduction
- CFG can improve prompt alignment, but excessive guidance may produce oversimplified or distorted samples; other self-guidance approaches can require degradation strategies, extra training, or additional sampling steps.The paper frames these costs against the challenge of learning high-quality outputs in low-probability regions with insufficient training and data.Source: §1 · p. 1
- IG adds intermediate-layer supervision during training, then extrapolates intermediate and deep outputs during sampling to improve generation quality without additional sampling steps.The authors also report that intermediate supervision alone alleviates partial vanishing gradients and converges comparably to, or better than, self-supervised intermediate-layer regularization.Source: §1 · pp. 1–2
- On ImageNet 256×256, SiT-XL/2+IG reaches FID 5.31 at 80 epochs and 1.75 at 800, while LightningDiT-XL/1+IG reaches FID 1.34 at 680 epochs and 1.19 with CFG and guidance interval.The reported SiT results without CFG exceed vanilla SiT-XL at 1400 epochs and REPA at 800 epochs.Source: §1 · p. 2
- The reported experiments find IG improves performance across DiTs, SiTs, and LightningDiT, and can be combined with CFG and guidance interval for further gains.The strongest reported combination reaches FID 1.19 on LightningDiT-XL/1 at 680 epochs.Source: §1 · p. 2
2. Related Work
Related work covers sampling-time guidance and training-time regularization for Diffusion Transformers. The paper positions IG as a simple intermediate-supervision approach that avoids some costs of prior guidance strategies.
Sampling Guidance in Diffusion Models
- CFG extrapolates from an unconditional score, while noise-dependent weights and weaker-model guidance are alternatives aimed at reducing CFG’s diversity downside.The paper notes that weaker-model approaches may depend on degradation strategies, extra training, or additional sampling steps.Source: §2 · p. 2
Intermediate Representation Regularization for Diffusion Transformers
- Prior Diffusion Transformer training methods include masked or routed-token strategies and self-supervised representation alignment or regularization.The paper groups self-supervised approaches into alignment with pretrained models and representation learning used as a training regularizer.Source: §2 · p. 3
- The authors use auxiliary intermediate-layer supervision to alleviate gradient vanishing, reporting convergence comparable to self-supervised regularization.They describe the design as simpler than approaches relying on self-supervised representations and note a conceptual similarity to Deep-Flow [64].Source: §2 · p. 3
3. Methodology
The method adds auxiliary supervision to an intermediate Diffusion Transformer output and uses that output during sampling to guide the final output. It builds on diffusion and flow-matching models that transform Gaussian noise into data samples.
3.1. Preliminaries
- Diffusion and flow-matching models transform Gaussian noise into data samples through stochastic processes, with their formulations differing in how time and the noise-to-data process are defined.Flow-matching typically interpolates noise and data over a finite interval, while diffusion defines a forward stochastic differential equation converging toward Gaussian noise.Source: §3.1 · p. 3
3.2. Internal Guidance: Guiding a Diffusion Model with Its Own Internal Dynamics
- The framework adds an auxiliary loss at an intermediate network output alongside the final-output loss, with λ controlling their trade-off.The intermediate and final outputs are denoted Di and Df, respectively.Source: §3.2 · p. 3
- At sampling time, IG extrapolates the intermediate and final denoised outputs so the intermediate output guides the final output away from its own poor distribution.The intermediate output is obtained through the auxiliary supervision introduced during training.Source: §3.2 · p. 4
- IG is presented as a plug-and-play approach that maintains diversity while improving image quality, without setup-specific bad-version training or extra sampling consumption.The framework figure summarizes the training supervision and sampling-time use of intermediate outputs.Source: Figure 2 · p. 3, §3.2 · p. 4
4. Discussion
The discussion examines how Internal Guidance (IG) interacts with classifier-free guidance (CFG) and how its effect depends on the noise range. It reports complementary guidance effects and different preferred noise intervals for IG and CFG.
4.1. The compatibility between IG and CFG
- In a 2D toy dataset, increasing CFG reduces outliers and improves inter-class separability, but suppresses branches near other classes and reduces diversity.The synthetic data has narrow, anisotropic support and local details that emerge hierarchically during denoising.Source: §4.1 · p. 4
- Autoguidance uses class-independent inward gradients, while IG performs similarly; increasing the IG coefficient narrows samples, and large coefficients before full training may introduce outliers.The narrower IG distribution did not appear to adversely affect images in practice.Source: §4.1 · p. 4
- Combining IG and CFG achieves superior generative quality across training stages because CFG’s class-conditional directionality complements IG when the stronger output poorly targets high-probability regions.The stronger output can be less effective before the model is fully trained.Source: §4.1 · p. 5
4.2. Guidance Interval in Internal Guidance
- CFG guidance is used in middle- and low-noise ranges, whereas IG guidance is applied in high- and middle-noise ranges, not low-noise ranges.The discussion notes that CFG at high noise can sharply reduce variation, producing a handful of template images per class.Source: §4.2 · p. 5
5. Experiments
Experiments evaluate IG on ImageNet 256×256, test auxiliary-supervision and guidance design choices, and compare SiT and LightningDiT variants with state-of-the-art methods.
5.1. Setup
- ImageNet-1K training uses 256×256 images; evaluation uses 50,000 samples, with samplers following each baseline family’s setup.DiT and SiT use 250-step SDE Euler-Maruyama sampling; LightningDiT uses 125-step ODE Heun sampling.Source: §5.1 · p. 5
5.2. Ablation Studies
- Early-layer supervision, such as the first four SiT-B/2 layers, is effective, whereas later-half or multiple intermediate losses do not speed convergence.The authors hypothesize that later or multiple losses interfere with training the deep-layer output.Source: §5.2 · p. 5, Table 1 · p. 5
- Combining IG with CFG lowers the best SiT-B/2 FID from 7.30 with CFG alone to 6.50, with little additional computational cost.The best combination uses a lower IG coefficient than the optimum for IG alone.Source: §5.2 · p. 6, Figure 4 · p. 6
- The best rough-search setting is IG coefficient 2.3 over interval [0.3, 1); without an interval, the optimal coefficient is around 1.9.Increasing the coefficient helps when IG is not applied in the low-noise range.Source: §5.2 · p. 6, Table 3 · p. 6
- IG reaches the same FID level more quickly as model size grows; experiments use coefficient 2.3 and interval [0.3, 1), applying IG at layer 4 for B/2 and layer 8 otherwise.The reported trend covers DiT and SiT models without CFG.Source: §5.2 · p. 6, Figure 5 · p. 6
- In the 400K-iteration SiT-B/2 coefficient sweep, performance stabilizes once λ ≤ 0.5.The sweep tests λ values from 0.25 to 1.0.Source: §5.2 · p. 6, Table 4 · p. 6
5.3. Comparsion with State-of-the-art Methods
- FID=5.31 at 80 epochs and FID=1.75 at 800 epochs for SiT-XL/2+IG; LightningDiT-XL/1+IG reaches FID=1.34 without CFG and FID=1.19 with CFG at 680 epochs.The comparisons use class-conditional ImageNet 256×256 generation.Source: Table 5 · p. 7, §5.3 · p. 7
6. Further Discussion on Training Acceleration
The authors use IG’s outlier-reduction insight to modify the training objective, reporting faster convergence than REPA in SiT-B/2 experiments.
6. Further Discussion on Training Acceleration
- The training objective adds a detached difference between final- and intermediate-layer outputs to the final objective, scaled by ω.The intermediate objective remains x0; the final objective is x0 + ω · sg(Df(xt) − Di(xt)).Source: §6 · p. 8
- In the toy example, IG reduces outliers near branches, motivating a training loss that incorporates the difference between intermediate and final outputs.The authors suggest continuing to use IG in applications because it is flexible and tunable during image generation.Source: §6 · p. 8, Figure 6 · p. 8
- After 400K iterations, the proposed loss significantly outperformed REPA and demonstrated performance comparable to IG in SiT-B/2 experiments.The loss also significantly reduced outliers in the toy example.Source: §6 · p. 8, Figure 6 · p. 8
7. Conclusion
The conclusion presents Internal Guidance as a strategy for improving generation quality and training efficiency, with CFG further improving image-generation results.
7. Conclusion
- Internal Guidance is presented as a simple strategy for enhancing both generation quality and training efficiency.It introduces auxiliary supervision in a deep Diffusion Transformer and uses intermediate-layer outputs to guide deep-layer outputs during sampling.Source: §7 · p. 8, Figure 7 · p. 8
- With CFG, LightningDiT-XL/1+IG achieves FID=1.19.The conclusion identifies this as the state-of-the-art FID result for image generation.Source: Figure 7 · p. 8
- The authors analyze IG’s compatibility with CFG and its guidance interval, and discuss training-acceleration methods inspired by IG.They suggest the work may facilitate research on more applicable sampling methods for visual generation.Source: Figure 7 · p. 8
A. Evaluation Metric
The appendix describes the evaluation metrics and notes that class-balanced sampling is used for a comprehensive comparison with recent methods.
A. Evaluation Metric
- FID measures feature-distribution distance between real and generated images using Inception-v3 features and a multivariate Gaussian assumption.The evaluation also reports sFID, inception score, precision, and recall.Source: A. Evaluation Metric · p. 12
- sFID uses intermediate spatial Inception-v3 features; inception score uses logits, while precision and recall measure realism and coverage of the training-data manifold..Source: A. Evaluation Metric · p. 12
- The comparison uses class-balanced sampling because it more closely approximates the training set’s label distribution than random sampling.The authors apply it to compare recent VQ- and diffusion-based methods.Source: A. Evaluation Metric · p. 12
B. Computational Cost Comparison
IG improves generation quality over REPA at the same model scale with only marginal additional computational cost and nearly identical latency.
B. Computational Cost Comparison
- IG improves generation quality over REPA at the same model scale, with marginal parameter and FLOPs increases and nearly identical latency.The comparison reports a relative FID reduction and an IS increase for IG.Source: B. Computational Cost Comparison · p. 13, Table 7 · p. 13
- The reported efficiency comparison uses SiT-XL/2+REPA and SiT-XL+IG at the same model scale.Source: B. Computational Cost Comparison · p. 13
- Table 7 compares REPA and IG on ImageNet 256×256 across model size, FLOPs, sampling steps, latency, and generation-quality metrics.Source: Table 7 · p. 13
C. Comparison of Guidance Methods
The guidance comparison reports that IG outperforms CFG-only and Autoguidance-only variants, while combining IG with CFG further improves results. At 512×512, IG surpasses REPA trained for 200 epochs and SiT-XL/2 trained for 600 epochs on FID after 60 epochs.
C. Comparison of Guidance Methods
- IG outperforms CFG-only and Autoguidance-only guidance, and combining IG with CFG achieves a significant improvement over other guidance methods.The comparison uses an initialized 12-layer LightningDiT trained for 10 epochs as Autoguidance's bad version, with coefficient 1.15; vanilla LightningDiT-XL/1 CFG is trained for 800 epochs.Source: C. Comparison of Guidance Methods · p. 13, Table 9 · p. 13
D. 512×512 ImageNet
- At 512×512, IG surpasses REPA trained for 200 epochs and SiT-XL/2 trained for 600 epochs on FID after 60 epochs.The experiment follows REPA's protocol [79] and uses the Muon optimizer.Source: D. 512×512 ImageNet · p. 13
E. Hyperparameter and More Implementation Details
The supplementary material specifies implementation and toy-example settings, then relates IG to prior work on transformer training, sampling guidance, and training efficiency.
E. Hyperparameter and More Implementation Details
- Implementation uses mixed precision, gradient clipping, FlashAttention, precomputed VAE latents, and no data augmentation; largest models use A6000 GPUs and smaller models use 4090 GPUs.The default hyperparameter setup is tabulated, and SiT-XL/2 + IG samples at 512×512 use CFG with w = 1.8.Source: E. Hyperparameter and More… · pp. 13, 15, Figure 8 · p. 14, Table 11 · p. 15
- Muon alleviates early instability when aligning a VAE encoder with pretrained self-supervised representations, but its early convergence advantage narrows and late instability can recur.Resuming training in full precision can solve the instability problem.Source: E. Hyperparameter and More… · pp. 14–15
F. Details of the 2D toy example
- The 2D toy dataset models each class as a thin tree-like mixture of Gaussians, with 1,016 components per class and 2,032 total.Recursive branch subdivision and randomized branch geometry define the construction; this permits analytic score and density calculations.Source: F. Details of the 2D toy example · pp. 15–16
- Toy denoisers are MLPs whose scalar output represents log unnormalized probability density; preconditioning connects the raw network Fθ to Gθ.The architecture uses one input layer, four hidden layers, SiLU activations, and an output after the first hidden layer for a weaker output.Source: F. Details of the 2D toy example · p. 16
- Toy-model training uses exact score matching with an additional intermediate-layer supervision loss, and the final loss combines deep and intermediate losses.The setup runs 4,096 iterations with batch size 4,096; denoising score matching showed no significant difference in model behavior or training dynamics.Source: F. Details of the 2D toy example · pp. 16–17
- Toy sampling uses the EDM sampler with 32 Heun steps, chosen N and σmax well above what the dataset requires to avoid discretization effects.The reported toy-example guidance weights differ by method, including IG + CFG with w1 = 1 and w2 = 1.5.Source: F. Details of the 2D toy example · p. 17
G. More Discussion on Related Work
- The discussion situates IG among transformer-backbone work and describes it as auxiliary supervision to early layers for analyzing and improving Diffusion Transformer training.Source: G. More Discussion on Related Work · p. 17
- Unlike approaches that use weaker-model guidance, IG requires no additional sampling steps and is described as plug-and-play for deep diffusion Transformers.The discussion also positions IG as addressing gradient vanishing while improving sampling results with less training consumption.Source: G. More Discussion on Related Work · p. 17
H. More Qualitative Results
The qualitative-results section presents uncurated ImageNet samples for SiT-XL and LightningDiT models using IG with classifier-free guidance.
H. More Qualitative Results
- The SiT-XL + IG (1.8) examples are ImageNet 256×256 samples using CFG with w = 2.5.Source: H. More Qualitative Results · p. 17
- The LightningDiT-XL/1 + IG (1.8) figures show uncurated results using CFG with w = 2.5.Source: Figure 9 · p. 18, Figure 10 · p. 19, Figure 11 · p. 20