💻 computer science

Guiding a Diffusion Transformer with the Internal Dynamics of Itself

This paper proposes Internal Guidance (IG), a simple yet effective strategy that leverages the internal dynamics of Diffusion Transformers by introducing auxiliary supervision on intermediate layers during training and extrapolating layer outputs during sampling, achieving state-of-the-art generation quality and training efficiency on ImageNet.

Original authors: Xingyu Zhou, Qifan Li, Xiaobin Hu, Hai Chen, Shuhang Gu

Published 2026-04-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xingyu Zhou, Qifan Li, Xiaobin Hu, Hai Chen, Shuhang Gu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a brilliant but slightly chaotic artist (the Diffusion Model) how to paint perfect pictures of cats, dogs, and landscapes.

The Problem: The "Over-Confident" Artist

Right now, these AI artists are amazing, but they have a specific flaw. They are trained to cover every possibility. Sometimes, they get confused by the weird, rare, or "low-probability" parts of the data. When you ask them to draw a specific thing, they might wander off into a "low-probability" zone and produce a blurry mess or a distorted monster.

To fix this, humans usually use a technique called Classifier-Free Guidance (CFG). Think of this as a strict art teacher standing next to the artist, shouting, "No! That's not a cat! Make it more like a cat!"

The Catch: If the teacher shouts too loudly (high guidance), the artist gets scared and stops being creative. The result is a cat that looks like a plastic toy—perfectly shaped but soulless and distorted.
The Old Alternative: Some researchers tried to teach the artist by showing them their own "bad drafts" (a "bad version" of themselves) to say, "Don't paint like that." But this is like hiring a second, worse artist just to critique the first one. It's expensive, slow, and requires a lot of extra work.

The Solution: "Internal Guidance" (IG)

This paper introduces a clever trick called Internal Guidance (IG). Instead of hiring a second artist or shouting from the sidelines, the paper suggests listening to the artist's own inner voice during the painting process.

Here is how it works, using a simple analogy:

1. The Training Phase: The "Sketch vs. Final" Lesson

Imagine the artist is painting a picture. Usually, they only get graded on the final masterpiece at the very end.

The New Idea: The paper says, "Let's also grade the rough sketch you made halfway through!"
How it works: During training, the AI is forced to look at its own "intermediate layer" (the sketch) and try to make that sketch look like the real object too.
The Result: The artist learns to keep their "inner sketch" aligned with reality. This fixes a common problem where deep neural networks get confused and lose their way (vanishing gradients) as they get deeper into the painting process.

2. The Painting Phase: The "Self-Correction" Trick

Now, when the artist is actually painting a picture for you (the sampling stage), here is the magic:

The AI generates two things simultaneously:
1. The Deep Output: The final, polished prediction.
2. The Intermediate Output: The "weaker" prediction from the middle of the network (the sketch).
The Magic Move: The AI compares the two. It realizes, "Hey, my final prediction is drifting a bit too far from my solid sketch."
The Correction: It uses the "sketch" to pull the "final prediction" back toward the center of the good data. It's like the artist looking at their own rough draft and saying, "Okay, I need to steer my final brushstrokes back toward this solid foundation."

Why is this a Big Deal?

Think of it like driving a car:

Standard CFG is like a passenger screaming, "Turn left! Turn left!" If they scream too much, you might crash into a tree (distorted images).
Old "Bad Version" methods are like having a second, worse car driving next to you, and you try to avoid hitting it. It's complicated and slow.
Internal Guidance (IG) is like having a GPS built into your own dashboard. You don't need a second car or a screaming passenger. You just look at your own internal map (the intermediate layer) and gently steer yourself back on track.

The Results: Supercharged Art

The paper tested this on some of the most advanced AI models (like SiT and LightningDiT).

Speed: They reached top-tier quality in 800 steps that usually took other models 1,400 steps. It's like finishing a marathon in half the time.
Quality: They achieved a score called FID 1.19. In the world of AI art, a lower FID is better. A score of 1.19 is so low it's almost indistinguishable from real photos. It's currently the State-of-the-Art (the best in the world).
Efficiency: It doesn't require extra computers or extra time. It's a "plug-and-play" upgrade. You just turn it on, and the model becomes smarter and more stable.

Summary

Internal Guidance is a simple but powerful hack. It teaches the AI to trust its own "intermediate thoughts" to guide its final decisions. By listening to its own inner voice, the AI avoids getting lost, creates more diverse and realistic images, and does it all faster than before. It's the difference between an artist who needs a screaming coach and one who has finally learned to trust their own gut.

1. Problem Statement

Diffusion models, particularly those based on the Transformer architecture (DiT/SiT), have achieved state-of-the-art results in image generation. However, they face two primary challenges:

Low-Probability Area Coverage: Training objectives penalize the model for failing to cover low-probability areas of the data distribution. Without sufficient training data, the model struggles to generate high-quality images corresponding to these rare modes.
Limitations of Existing Guidance:
- Classifier-Free Guidance (CFG): While effective at aligning prompts and results, high CFG coefficients often push the sampling trajectory outside the conditional distribution, leading to over-simplified or distorted images and reduced diversity.
- Autoguidance (Bad Version Guidance): This approach guides a well-trained model using a "bad" (degraded) version of itself. While it maintains diversity better than CFG, it requires carefully designed degradation strategies, extra training for the degraded model, or additional sampling steps, making it computationally expensive and difficult to scale.

2. Methodology: Internal Guidance (IG)

The authors propose Internal Guidance (IG), a simple yet effective plug-and-play strategy that leverages the internal dynamics of a single Diffusion Transformer to improve generation quality without extra training or sampling steps.

A. Training Phase: Auxiliary Supervision

Instead of training a separate "bad" model, IG introduces an auxiliary supervision loss at an intermediate layer of the network during training.

Architecture: An additional output head is attached to an intermediate layer ( $D_i$ ) of the Diffusion Transformer, in addition to the final output head ( $D_f$ ).
Loss Function: The total training loss combines the standard denoising loss for the final layer and an auxiliary loss for the intermediate layer:
$L = L_{final} + \lambda L_{inter}$
Where $L_{inter} = \|D_i(x_t, t) - x_0\|^2$ and $L_{final} = \|D_f(x_t, t) - x_0\|^2$ .
Effect: This forces the intermediate layer to learn a "weaker" generative representation. The authors observe that this simple supervision alleviates gradient vanishing in deep networks and achieves convergence comparable to complex self-supervised regularization methods.

B. Sampling Phase: Extrapolation

During the sampling stage, the model utilizes the relationship between the intermediate output ( $D_i$ ) and the final output ( $D_f$ ) to guide the generation.

Mechanism: The intermediate output acts as the "bad version" of the model. The final output is guided away from the intermediate output's distribution toward the high-probability manifold.
Formula: The guided output $D_w$ is calculated via extrapolation:
$D_w(x; c) = D_i(x; c) + w(D_f(x; c) - D_i(x; c))$
Here, $w$ is the guidance scale. This effectively pushes the sample toward the direction of the stronger (final) model while maintaining the diversity inherent in the intermediate representation.
Advantage: Unlike Autoguidance, IG requires no extra training for a degraded model and no additional sampling steps. It is a "plug-and-play" modification.

C. Compatibility and Enhancements

Combination with CFG: IG and CFG are complementary. CFG provides class-conditional directionality, while IG maintains diversity and pushes samples toward the manifold. Combining them yields superior results.
Guidance Interval: The authors redefine the guidance interval for IG. Unlike CFG (which works best in low-noise ranges), IG is most effective when applied in high-noise and middle-noise ranges.

3. Key Contributions

Internal Guidance (IG) Strategy: A novel sampling guidance method that uses a model's own intermediate layers to guide the final output, eliminating the need for external "bad" models or complex degradation strategies.
Training Efficiency: The introduction of intermediate supervision acts as a regularizer that accelerates convergence and mitigates gradient vanishing, achieving performance comparable to self-supervised representation alignment methods (like REPA) but with simpler implementation.
State-of-the-Art Performance: The method significantly improves FID scores across various model scales (SiT, DiT, LightningDiT) with minimal computational overhead.
New Training Acceleration Insight: The authors demonstrate that the guidance signal (the difference between final and intermediate outputs) can be integrated directly into the training loss function to further accelerate convergence, offering a new direction for training acceleration.

4. Experimental Results

The paper evaluates IG on ImageNet 256×256 and 512×512 benchmarks.

SiT-XL/2 (Class-Conditional):
- Achieved FID = 5.31 in just 80 epochs (without CFG).
- Achieved FID = 1.75 at 800 epochs (without CFG), surpassing vanilla SiT-XL at 1400 epochs and REPA at 800 epochs.
- With CFG and guidance interval, achieved FID = 1.46 at 800 epochs.
LightningDiT-XL/1 (Self-Supervised Tokenizer):
- Achieved FID = 2.42 at 60 epochs.
- Achieved FID = 1.34 at 680 epochs (without CFG), outperforming all prior methods by a large margin.
- State-of-the-Art: Combined with CFG and guidance interval, achieved FID = 1.19 at 680 epochs.
Scalability: The performance gain of IG becomes more significant as the model size increases (from B to XL scales).
Computational Cost: IG introduces negligible overhead (approx. 0.44% increase in parameters and FLOPs) and maintains identical latency compared to baseline models.

5. Significance

Efficiency: IG offers a highly efficient way to boost generation quality without the heavy computational cost of training auxiliary models or increasing sampling steps.
Generalizability: The method is architecture-agnostic within the Diffusion Transformer family and can be seamlessly integrated into existing pipelines.
Theoretical Insight: It bridges the gap between intermediate supervision and sampling guidance, suggesting that the "internal dynamics" of a deep network contain sufficient information to correct sampling trajectories, potentially inspiring new training acceleration techniques based on loss function design.
Practical Impact: By achieving SOTA FID scores (1.19) with fewer training epochs and standard hardware, IG makes high-fidelity image generation more accessible and scalable.