Guiding a Diffusion Transformer with the Internal Dynamics of Itself

AI summary

Overview

Research area: Generative computer vision, specifically sampling guidance for latent diffusion transformers built on flow matching and diffusion.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching objectives, denoising networks, classifier-free guidance (CFG), and guidance intervals.

Scope (one sentence): The paper proposes Internal Guidance (IG), a training-time auxiliary supervision on an intermediate layer of a diffusion transformer that is reused at sampling time to extrapolate between intermediate and final denoiser outputs, improving FID at negligible extra cost.

What This Paper Is About

Diffusion models are trained to cover an entire data distribution, but they lack the training and data to generate high-quality images in low-probability regions, so samples drift toward poor outputs. Guidance strategies such as CFG push samples toward high-probability areas, yet a large CFG coefficient produces over-simplified or distorted results, and the alternative "guide with a bad version of yourself" line of work requires carefully designed degradation strategies, extra training, or additional sampling steps. The goal of this paper is a simple, plug-and-play alternative that uses the diffusion transformer's own intermediate-layer outputs as the weaker model to guide against, without any of those extra requirements.

Key Contributions

1. Internal Guidance (IG): a training scheme that adds an auxiliary supervision loss at an intermediate layer of a deep diffusion transformer, and a sampling scheme that extrapolates between the intermediate output D_i and the final output D_f using a guidance factor w (Equation 5), achieving an Autoguidance-like effect without extra sampling steps.
2. A plug-and-play method with negligible cost: reported at only +0.44% parameters, +0.01% FLOPs, and +0.16% latency relative to SiT-XL/2 in the computational comparison against REPA, while cutting FID by 70.34% and raising IS by 44.87%.
3. Analysis of two properties of IG: its compatibility with CFG (verified quantitively and on a 2D toy distribution with a fractal-like two-class structure), and a guidance interval for IG, where the useful interval is the opposite of CFG's, i.e., high-noise and middle-noise rather than low-noise ranges.
4. A training-acceleration direction derived from IG: moving the discrepancy between final and intermediate outputs into the training objective (Equation 7) to reduce outliers during training, tested on SiT-B/2 against REPA.

Main Findings

• State-of-the-art FID of 1.19: LightningDiT-XL/1+IG reaches FID = 1.34 at 680 epochs without CFG and FID = 1.19 with CFG (CFG scale w = 1.45) plus a guidance interval, which the paper calls the current state-of-the-art. Figure 1 describes the corresponding setting as an IG scale of 1.4 and a CFG scale of 1.45.
• Fast convergence on SiT-XL/2: with IG, FID = 5.31 at 80 epochs and FID = 1.75 at 800 epochs without CFG, and FID = 1.46 at 800 epochs with CFG (scale w = 1.35) plus a guidance interval. The 80-epoch result already exceeds vanilla SiT-XL at 1400 epochs (FID = 8.61) and REPA at 800 epochs (FID = 5.90), which uses a pre-trained representation model for alignment.
• Fast convergence on LightningDiT-XL/1: FID = 2.42 at 60 epochs with IG (1.7), already exceeding vanilla LightningDiT-XL/1 at 64 epochs (FID = 5.14); FID = 1.34 at 680 epochs.
• Balanced sampling results: under uniform balanced sampling on ImageNet 256×256, LightningDiT+IG (680 epochs, 678M parameters) reaches FID = 1.24 without CFG and FID = 1.07 with CFG, compared to RAE(DiT^DH) at 1.51/1.13, REPA-E at 1.70/1.15, and LightningDiT at 2.11/1.28.
• Intermediate supervision placement matters: on SiT-B/2 at 80 epochs, regularizing only the first few layers helps, with layer 4 giving FID = 30.60 without IG and 19.02 with IG (w = 1.5); placing the loss in the later half or at multiple intermediate layers (layers 2 and 6: FID = 31.55 with IG; layer 8: 34.05; layer 10: 35.58) does not improve convergence and is hypothesized to interfere with deep-layer training. Large-scale experiments apply the loss and guidance at the 8th layer.
• Competitive with self-supervised regularization: on SiT-B/2 at 80 epochs without IG or CFG, the auxiliary supervision loss gives FID = 30.45 / IS = 47.97, versus SRA at 29.10 / 50.20 and Disperse Loss at 31.45 / 47.05, and a baseline of 33.02 / 43.71.
• IG coefficient has an optimum: on SiT-B/2 without a guidance interval, FID improves from 30.60 (w = 1.0) through 19.02 (1.5), 17.61 (1.7), to 17.38 at w = 1.9, then degrades to 17.75 (2.1) and 18.53 (2.3).
• Guidance interval for IG differs from CFG: with an IG coefficient of 2.3, restricting guidance to [0.3, 1) gives FID = 16.19 / IS = 72.95, compared to 18.53 / 66.77 with no interval, 18.23 / 66.21 for [0.3, 0.7), and 20.69 / 60.14 for [0, 0.7). The authors state IG should not be applied in the low-noise range but in the high-noise and middle-noise ranges.
• IG combines well with CFG: on SiT-B/2, adding CFG after IG reduces the best FID from 7.30 (CFG alone) to 6.50, and a lower IG coefficient is optimal when combined with CFG than when IG is used alone.
• Loss weight is robust below a threshold: with SiT-B/2 trained for 400K using λ from 0.25 to 1.0, FID is 30.38 (0.25), 30.60 (0.5), 31.58 (0.75), and 31.24 (1.0); performance is reported as stable for λ <= 0.5.
• Scalability: the relative improvement of IG over vanilla models grows with model size, and IG is reported to reach a given FID level faster as models get larger.
• Training acceleration: folding the intermediate/final discrepancy into the training loss (ω = 0.5, following prior work), the SiT-B/2 loss significantly outperformed REPA after 400K iterations.
• Comparison with Autoguidance: on EDM2-S with ImageNet 64, FID is 1.58 for the base model, 1.01 with Autoguidance, and 0.99 with IG.

Methodology in Plain English

During training, the authors attach an extra output head to an intermediate layer of the diffusion transformer and supervise it with the same denoising objective as the final layer, giving two losses: L_inter and L_final, combined as L_final + λ · L_inter. This intermediate head is trained to produce a weaker denoised prediction from mid-level features.

At sampling time, the network produces two denoised estimates per step: the intermediate one and the final one. The guided prediction is an extrapolation, D_w = D_i + w(D_f - D_i), which moves the result away from the weaker intermediate distribution toward the stronger final distribution, analogous to how CFG extrapolates away from an unconditional prediction but without requiring an unconditional pass or a degraded model. The guidance factor can be replaced by a piecewise-constant function of the noise level so that guidance applies only inside a chosen noise interval.

Experiments train on ImageNet-1K at 256×256, encoding images to z ∈ R^(32×32×4) with the Stable Diffusion VAE or z ∈ R^(16×16×32) with the VA-VAE. DiT and SiT follow their original setups; LightningDiT is modified to use the Muon optimizer instead of AdamW and an EMA weight of 0.9995 instead of 0.9999 because the original recipe caused instability and slow EMA convergence in early epochs. Evaluation uses 50,000 generated samples over 1,000 class labels, FID/sFID/IS/precision/recall, an SDE Euler-Maruyama sampler with 250 steps for DiT and SiT, and an ODE Heun sampler with 125 steps for LightningDiT.

Why This Matters

Impact on research. The paper shows that a guidance mechanism can be obtained essentially for free by supervising and then reusing a model's own intermediate layers, rather than building a separate degraded model, running extra sampling steps, or adding self-supervised representation alignment. It also reports that a plain auxiliary supervision loss reaches convergence comparable to self-supervised regularization on intermediate layers, and it reframes guidance as a training-acceleration idea.

Real-world applications.

• Text-to-image and class-conditional image generation systems where sampling speed and FID both matter.
• Large-scale generative model training pipelines that want faster convergence without adding a pre-trained representation network.
• Video and 3D generation pipelines built on diffusion transformers, where the same plug-and-play layer supervision could be applied.
• Content creation and design tooling that benefits from higher-quality, more diverse samples with light tuning of a guidance scale.

Industry relevance. The overhead figures (roughly +0.44% parameters, +0.01% FLOPs, +0.16% latency in the reported comparison) and the ability to combine with existing CFG and guidance-interval schedules mean the method can be dropped into existing latent diffusion transformer stacks with minimal retuning, and reportedly brings the best results when combined with CFG rather than replacing it.

Future Directions

• Whether the training-acceleration objective (Equation 7) can be scaled beyond the reported SiT-B/2 comparison with REPA, and how it interacts with longer training runs.
• How the optimal intermediate layer, IG coefficient, λ, and guidance interval scale jointly with model size and training duration; the paper fixes an IG coefficient of 2.3 and interval [0.3, 1) for the scalability study and uses layer 4 for B/2 and layer 8 for other scales.
• Why IG's preferred guidance interval (high and middle noise) is the opposite of CFG's (middle and low noise), which the paper observes empirically but attributes to guidance behavior rather than deriving theoretically.
• Whether IG transfers to other modalities and resolutions, since the main text reports 256×256 results and defers 512×512 results and additional uncurated samples to the supplementary material.

Target Audience

Researchers and engineers working on diffusion transformers, flow matching, and sampling guidance who want a low-cost alternative or complement to CFG; practitioners deploying large latent diffusion models who care about FID, diversity, and inference overhead; and students with a background in denoising diffusion who want a clear example of how training-time design and sampling-time guidance can be tied together.

Authors’ 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.