Guiding a Diffusion Transformer with the Internal Dynamics of Itself

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

Auteurs: Xingyu Zhou, Qifan Li, Xiaobin Hu, Hai Chen, Shuhang Gu.

---

1. Het Probleem

Diffusiemodellen hebben bewezen krachtig te zijn in het genereren van hoge-kwaliteit beelden door de volledige (conditionele) data-distributie te leren. Echter, er zijn twee belangrijke uitdagingen:

• Ondervertegenwoordiging van lage-kansgebieden: Modellen worden zwaar gestraft als ze lage-kansgebieden niet goed leren, maar ze missen vaak voldoende training en data om daar hoogwaardige beelden voor te genereren.
• Beperkingen van bestaande geleidingsstrategieën:
  • Classifier-Free Guidance (CFG): De standaardmethode om samples naar hoge-kansgebieden te leiden. Een te hoge CFG-coëfficiënt leidt echter vaak tot oververeenvoudigde of vervormde samples en vermindert de diversiteit.
  • Autoguidance (gebruik van een "slechte" versie): Een alternatieve aanpak waarbij een goed getraind model wordt geleid door een slechtere versie van zichzelf. Hoewel dit veelbelovend is, vereist deze methode vaak zorgvuldig ontworpen degradatiestrategieën, extra trainingstijd of extra sampling-stappen, wat de schaalbaarheid beperkt.

Er is behoefte aan een eenvoudige, effectieve strategie die de generatiekwaliteit verbetert zonder extra sampling-kosten of ingewikkelde degradatiemethoden.

2. Methodologie: Internal Guidance (IG)

De auteurs stellen Internal Guidance (IG) voor, een plug-and-play strategie die gebruikmaakt van de interne dynamiek van de Diffusion Transformer (DiT) zelf. De methode bestaat uit twee fasen:

A. Training: Intermediäre Supervisie

Tijdens het trainingsproces wordt een hulpverlies (auxiliary supervision loss) geïntroduceerd op een tussenschicht (intermediate layer) van het netwerk.

• Het netwerk heeft zowel een eindlaag ($D_f$) als een tussenschicht ($D_i$).
• Er wordt een extra verliesfunctie gedefinieerd voor de output van de tussenschicht, naast het hoofdverlies voor de eindoutput.
• Het totale verlies is: $L = L_{final} + \lambda L_{inter}$, waarbij $\lambda$ een hyperparameter is.
• Doel: Dit helpt het gradient-vanishingsprobleem in diepe netwerken te verminderen en zorgt ervoor dat het model een "zwakkere" generatieve output kan produceren vanuit de tussenschicht-representaties.

B. Sampling: Extrapolatie van Interne Outputs

Tijdens het generatieproces (sampling) wordt de relatie tussen de output van de tussenschicht en de eindlaag benut om de sample te sturen.

• De output van de tussenschicht ($D_i$) fungeert als de "slechte versie" van het model, zonder dat er extra training of degradatie nodig is.
• De uiteindelijke output ($D_w$) wordt berekend door extrapolatie, vergelijkbaar met CFG maar met interne dynamiek:
  $$D_w(x; c) = D_i(x; c) + w(D_f(x; c) - D_i(x; c))$$
  Waarbij $w$ de geleidingsfactor is.
• Voordeel: Deze methode behoudt de diversiteit (in tegenstelling tot hoge CFG-waarden) en verbetert de kwaliteit door samples weg te duwen van de "slechte" verdeling van de tussenschicht, zonder extra sampling-stappen.

3. Belangrijkste Bijdragen

1. Internal Guidance (IG): Een eenvoudige maar effectieve methode die gebruikmaakt van tussenschicht-outputs binnen een Diffusion Transformer om de generatiekwaliteit te verbeteren.
2. Plug-and-Play Implementatie: IG vereist geen extra trainingstijd voor specifieke "slechte" modellen, geen extra sampling-stappen en geen complexe degradatiestrategieën. Het kan naadloos worden geïntegreerd in bestaande DiT-architecturen.
3. Combinatie met CFG: IG is compatibel met Classifier-Free Guidance. De combinatie van beide strategieën leidt tot nog betere resultaten, waarbij CFG de class-conditional richting biedt en IG de diversiteit behoudt en outliers reduceert.
4. Training Acceleratie: De auteurs tonen aan dat het inzicht uit IG kan worden gebruikt om nieuwe trainingsverliesfuncties te ontwerpen die de convergentie versnellen, vergelijkbaar met geavanceerde self-supervised regularisatiemethoden maar eenvoudiger.

4. Resultaten

De methode is uitgebreid getest op ImageNet 256×256 en 512×512 met verschillende baselines (SiT, DiT, LightningDiT).

• SiT-XL/2:

  • Bereikt een FID van 5.31 na slechts 80 epochs (zonder CFG).
  • Bereikt een FID van 1.75 na 800 epochs.
  • Met CFG en Guidance Interval bereikt het een FID van 1.46 na 800 epochs.
  • Dit presteert aanzienlijk beter dan het standaard SiT-XL na 1400 epochs en andere geavanceerde methoden zoals REPA.

• LightningDiT-XL/1 (met self-supervised tokenizer):

  • Bereikt een FID van 2.42 na slechts 60 epochs.
  • Bereikt een FID van 1.34 na 680 epochs (zonder CFG), wat een grote marge is ten opzichte van alle bestaande methoden.
  • State-of-the-Art: In combinatie met CFG en Guidance Interval bereikt het een FID van 1.19 na 680 epochs. Dit is momenteel het beste resultaat voor deze specifieke setup.

• Efficiëntie: IG introduceert slechts een marginale toename in parameters en FLOPs (ongeveer 0.44%) en heeft geen extra latentie tijdens het sampling-proces.

5. Betekenis en Impact

Dit paper biedt een doorbraak in de efficiëntie en kwaliteit van diffusiemodellen:

• Kostenefficiëntie: Het bereikt state-of-the-art resultaten met aanzienlijk minder trainings-epochs dan concurrenten, wat rekenkracht en tijd bespaart.
• Schaalbaarheid: De methode werkt effectief op verschillende modelgroottes en verbetert de prestaties des te meer naarmate het model groter wordt.
• Eenvoud: Het lost het dilemma op tussen hoge kwaliteit en diversiteit zonder de complexiteit van extra training of ingewikkelde sampling-procedures.
• Toekomstperspectief: Het inzicht dat interne dynamiek kan worden gebruikt voor geleiding, opent de deur voor nieuwe trainingssnelheidsstrategieën en maakt diffusiemodellen toegankelijker voor grootschalige toepassingen.

Kortom, Internal Guidance is een krachtige, eenvoudige toevoeging die de grenzen van wat mogelijk is met Diffusion Transformers aanzienlijk verlegt, met name in termen van trainingsduur en eindkwaliteit.