Guiding a Diffusion Transformer with the Internal Dynamics of Itself

Titel: Guiding a Diffusion Transformer with the Internal Dynamics of Itself

Autoren: Xingyu Zhou, Qifan Li, Xiaobin Hu, Hai Chen, Shuhang Gu (u.a. UESTC, NUS, SYSU)

---

1. Problemstellung

Diffusionsmodelle haben sich als äußerst leistungsfähig erwiesen, um die gesamte (bedingte) Datenverteilung zu erfassen. Dennoch bestehen zwei wesentliche Herausforderungen:

• Unterdosierung bei niedrigen Wahrscheinlichkeiten: Modelle werden stark bestraft, wenn sie Bereiche mit niedriger Wahrscheinlichkeit nicht abdecken, haben aber oft nicht genügend Trainingsdaten, um hochwertige Bilder für diese Bereiche zu generieren.
• Limitationen bestehender Führungsstrategien:
  • Classifier-Free Guidance (CFG): Die Standardmethode, um Proben in Bereiche hoher Wahrscheinlichkeit zu lenken, führt bei hohen Gewichtungsfaktoren oft zu übermäßig vereinfachten oder verzerrten Bildern und reduziert die Vielfalt (Diversity).
  • Autoguidance (mit einer „schlechten Version" des Modells): Ansätze, die ein gut trainiertes Modell mit einer absichtlich verschlechterten Version desselben Modells führen, sind vielversprechend, erfordern jedoch sorgfältig gestaltete Degradationsstrategien, zusätzlichen Trainingsaufwand oder extra Sampling-Schritte. Dies schränkt ihre Skalierbarkeit für große Modelle ein.

Das Ziel ist es, eine Methode zu finden, die die Generierungsqualität verbessert, ohne die Vielfalt zu opfern, ohne zusätzliche Sampling-Schritte zu benötigen und ohne komplexes Training für „schlechte Versionen" zu erfordern.

---

2. Methodik: Internal Guidance (IG)

Die Autoren schlagen eine einfache, aber effektive Strategie namens Internal Guidance (IG) vor. Diese nutzt die internen Dynamiken des Diffusions-Transformers (DiT/SiT), indem sie Zwischenschichten des Netzwerks während des Trainings und des Sampling-Prozesses ausnutzt.

A. Training: Auxiliary Supervision (Zusätzliche Überwachung)

• Ansatz: Während des Trainings wird eine zusätzliche Ausgabe-Schicht an einer Zwischenschicht des Netzwerks eingefügt.
• Verlustfunktion: Neben dem finalen Verlust ($L_{final}$) wird ein Hilfsverlust ($L_{inter}$) berechnet, der den Fehler zwischen der Ausgabe der Zwischenschicht ($D_i$) und dem Zielbild ($x_0$) misst.
  $$L = L_{final} + \lambda \cdot L_{inter}$$
• Effekt: Dies zwingt das Modell, auch aus den mittleren Repräsentationen eine schwächere, aber gültige Generierung zu produzieren. Dies hilft effektiv gegen das Problem des Vanishing Gradients in tiefen Netzwerken und beschleunigt die Konvergenz, ähnlich wie bei komplexeren selbstüberwachten Regularisierungsmethoden, aber mit deutlich geringerem Aufwand.

B. Sampling: Extrapolation interner Ausgaben

• Prinzip: Während des Sampling-Prozesses stehen nun zwei Ausgaben zur Verfügung: die finale Ausgabe ($D_f$) und die Ausgabe der Zwischenschicht ($D_i$).
• Führungsmechanismus: Die Zwischenschicht-Ausgabe fungiert als die „schlechte Version" des Modells. Durch Extrapolation zwischen $D_i$ und $D_f$ wird die finale Ausgabe so gelenkt, dass sie sich von der schlechteren Verteilung der Zwischenschicht wegbewegt.
• Formel:
  $$D_w(x; c) = D_i(x; c) + w \cdot (D_f(x; c) - D_i(x; c))$$
  Hierbei ist $w$ der Führungsparameter (Scale).
• Vorteile:
  • Plug-and-Play: Keine zusätzlichen Sampling-Schritte nötig.
  • Kein Extra-Training: Es wird keine separate „schlechte Version" des Modells trainiert; sie entsteht natürlich durch die Zwischenschicht.
  • Kompatibilität: IG kann nahtlos mit CFG und Guidance Intervals kombiniert werden.

C. Guidance Interval für IG

Im Gegensatz zu CFG, das in Bereichen mit hohem Rauschen (frühe Sampling-Schritte) oft schädlich ist, zeigt IG das umgekehrte Verhalten: Es wirkt am besten in den hohen und mittleren Rauschbereichen. Die Anwendung von IG nur in diesen Bereichen (z. B. $\sigma \in [0.3, 1.0)$) führt zu besseren Ergebnissen.

---

3. Wichtige Beiträge

1. Einführung von Internal Guidance (IG): Eine neue, effiziente Führungsstrategie, die die interne Struktur von Diffusion-Transformern nutzt, um die Generierungsqualität zu steigern, ohne die Vielfalt zu beeinträchtigen.
2. Trainingseffizienz: Die Einführung einer einfachen Hilfsüberwachung an Zwischenschichten verbessert die Konvergenzgeschwindigkeit und die Trainingsstabilität signifikant, vergleichbar mit komplexen selbstüberwachten Methoden, aber ohne deren Overhead.
3. Synergie mit CFG: Die Arbeit zeigt, dass IG und CFG komplementär wirken. IG erhält die Vielfalt und lenkt die Proben in die richtige Richtung, während CFG die Klassen-Treue erhöht. Die Kombination führt zu State-of-the-Art-Ergebnissen.
4. Skalierbarkeit: Die Methode skaliert effizient mit der Modellgröße; der relative Gewinn durch IG nimmt bei größeren Modellen (z. B. SiT-XL, LightningDiT-XL) zu.

---

4. Ergebnisse

Die Methode wurde auf ImageNet (256×256 und 512×512) mit verschiedenen Architekturen (SiT, DiT, LightningDiT) evaluiert.

• SiT-XL/2 (ohne CFG):
  • Nach nur 80 Epochen erreicht das Modell mit IG einen FID von 5.31.
  • Nach 800 Epochen sinkt der FID auf 1.75.
  • Vergleich: Dies übertrifft das Vanilla SiT-XL nach 1400 Epochen und REPA (ein State-of-the-Art-Modell mit vortrainierten Repräsentationen) nach 800 Epochen.
• LightningDiT-XL/1 (ohne CFG):
  • Nach 60 Epochen: FID = 2.42.
  • Nach 680 Epochen: FID = 1.34.
  • Dies stellt einen großen Vorsprung gegenüber allen zuvor genannten Methoden dar.
• State-of-the-Art (mit CFG + Guidance Interval):
  • Die Kombination von LightningDiT-XL/1 + IG + CFG erreicht bei 680 Epochen einen FID von 1.19. Dies ist der aktuell beste Wert (State-of-the-Art) für diese Aufgabe.
• Effizienz: IG fügt nur einen minimalen Overhead hinzu (ca. 0,44% mehr Parameter, vernachlässigbare Latenzsteigerung), liefert aber massive Qualitätsverbesserungen.

---

5. Bedeutung und Ausblick

• Paradigmenwechsel: Die Arbeit zeigt, dass man die „schlechte Version" eines Modells nicht extern trainieren muss, sondern sie intern durch die Dynamik der Zwischenschichten extrahieren kann.
• Praktische Anwendbarkeit: Da IG als Plug-and-Play-Modul in bestehende Diffusion-Transformer integriert werden kann und keine zusätzlichen Sampling-Schritte erfordert, ist es ideal für den Einsatz in großen, kommerziellen Generatoren.
• Training Acceleration: Die Autoren leiten aus den Erkenntnissen von IG neue Trainingsverlustfunktionen ab, die die Konvergenz beschleunigen können, indem sie den Führungsgradienten direkt in den Trainingsprozess integrieren.
• Zukunft: Die Methode ebnet den Weg für effizienteres Training und hochwertigere Generierung in zukünftigen Multimodal-Modellen (Bilder, Videos), insbesondere in Kombination mit vortrainierten Repräsentationsmodellen.

Zusammenfassend bietet „Internal Guidance" einen eleganten, rechen-effizienten Weg, um die Grenzen aktueller Diffusionsmodelle zu überwinden und gleichzeitig die Trainingszeit zu verkürzen.