Scale-wise Distillation of Diffusion Models

Probleemstelling

Diffusiemodellen (DM's) voor het genereren van hoge-resolutie afbeeldingen en video's staan bekend om hun hoge kwaliteit, maar lijden onder een significant prestatieprobleem: de generatie vereist een sequentiële sampling van 20 tot 50 stappen. Hoewel bestaande distillatiemethoden het aantal stappen succesvol hebben gereduceerd tot ongeveer 4, wordt het verder verlagen van dit aantal (bijv. naar 1-2 stappen) steeds uitdagender.

De huidige aanpak richt zich voornamelijk op het verminderen van het aantal stappen terwijl de modelarchitectuur en de invoerresolutie constant blijven. Dit leidt tot inefficiëntie, omdat het model op elke stap, zelfs bij hoge ruisniveaus waar hoge frequenties nog niet zichtbaar zijn, berekeningen uitvoert op de volledige resolutie. Er is een onderzochte richting nodig om efficiëntie te winnen door de modelarchitectuur of het sampling-proces zelf te optimaliseren in plaats van alleen het aantal iteraties te verkleinen.

Methodologie: SwD Framework

De auteurs introduceren SwD (Scale-wise Distillation), een raamwerk dat een vooraf getraind diffusiemodel omzet in een few-step model dat progressief de resolutie verhoogt tijdens het generatieproces.

1. Spectrale Analyse van Latent Spaces:
De basis van SwD ligt in de observatie dat het omgekeerde diffusieproces vergelijkbaar is met spectrale autoregressie. De auteurs tonen aan dat in de latent spaces van VAE's (zoals bij SD3.5 en Wan2.1) hoge ruimtelijke en temporele frequenties in de vroege stadia van het proces (hoge ruis) grotendeels onderdrukt zijn. Dit betekent dat het modelleren van hoge frequenties op deze tijdstippen overbodig is en dat het model op lagere resoluties kan werken zonder informatie te verliezen.

2. Scale-Wise Sampling en Training:
In plaats van te starten met een vaste resolutie, start SwD met de laagste resolutie ($s_1$) en verhoogt de resolutie stap voor stap ($s_1 \to s_2 \to \dots \to s_N$) volgens een vooraf bepaald schema.

• Upsampling Strategie: Een cruciale technische innovatie is de manier waarop tussenliggende samples worden geupsampled. In plaats van het ruwe, ruisende latent ($x_t$) direct te upsamplen (wat artefacten introduceert), voorspelt het model eerst een schone sample ($\hat{x}_0$) op de lagere resolutie. Dit $\hat{x}_0$ wordt vervolgens geupsampled en opnieuw "geruisd" (renoised) naar het volgende ruisniveau. Dit behoudt de juiste ruisstatistieken en minimaliseert artefacten.
• Training: Het model wordt getraind op paren van aangrenzende schalen. Het leert om van een geupsampled en geruisd latent op schaal $s_i$ een schone sample op schaal $s_{i+1}$ te voorspellen.

3. Nieuwe Distillatie-Doelstelling (LMMD):
Naast het schaalraamwerk introduceren de auteurs een nieuwe distillatie-objektive gebaseerd op Maximum Mean Discrepancy (MMD) op patch-niveau.

• In plaats van een discriminator te trainen (zoals bij GAN-gebaseerde distillatie), wordt MMD berekend in de feature space van een vooraf getraind teacher-model.
• De loss functie ($L_{MMD}$) vergelijkt de gemiddelde features van de gegenereerde samples met die van de target samples op het niveau van ruimtelijke tokens (patches).
• Dit doel is computatie-efficiënt, vereist geen extra trainbare modellen en werkt zelfs als standalone loss zeer goed.

Belangrijkste Bijdragen

1. Spectrale Autoregressie in Latent Spaces: Het bewijs dat latent diffusion models effectief kunnen opereren op lagere resoluties tijdens de vroege, ruisrijke fasen van het generatieproces, zowel voor afbeeldingen als video.
2. SwD Framework: Een unificatie van multi-scale generatie in één enkele few-step model, wat leidt tot significante snelheidswinst zonder kwaliteitsverlies.
3. MMD-based Distillatie: Een eenvoudige maar krachtige loss-functie die de convergentie van distillatie verbetert en een competitieve baseline biedt zonder de complexiteit van adversarial training.
4. Efficiënte Upsampling: Een bewezen strategie om upscaling-artefacten te minimaliseren door upscaling toe te passen op de voorspelde schone sample ($\hat{x}_0$) in plaats van op het ruisende latent.

Resultaten

De auteurs hebben SwD toegepast op state-of-the-art modellen zoals SDXL, SD3.5 (Medium en Large), FLUX.1-dev en Wan2.1 (video).

• Snelheid: SwD bereikt een sampling-snelheid die vergelijkbaar is met 2 full-resolution stappen, maar met een veel lagere rekentijd.
  • Voor tekst-naar-afbeelding: ~2x versnelling ten opzichte van de snelste bestaande few-step modellen (bij gelijke stap-aantallen).
  • Voor tekst-naar-video: ~3x versnelling.
  • Totaal gezien is SwD >10x sneller dan de oorspronkelijke teacher-modellen.
• Kwaliteit:
  • Automatische Metrieken: SwD presteert op niveau met of beter dan de teacher-modellen en andere distillatiemethoden (zoals DMD2, Turbo, Hyper-SD) op metrieken zoals FID, HPSv3, ImageReward en GenEval.
  • Menselijke Evaluatie: In side-by-side vergelijkingen prefereerden beoordelaars SwD vaak boven concurrenten, vooral op het gebied van beeldcomplexiteit en aesthetiek, met vergelijkbare tekst-relevantie en minder defecten dan sommige alternatieven.
  • Video: Voor Wan2.1-SwD werd een 72x versnelling bereikt ten opzichte van de teacher, met betere of vergelijkbare resultaten op VBench-2.0.

Betekenis en Impact

Dit paper markeert een verschuiving in de optimalisatie van diffusiemodellen. In plaats van alleen te focussen op het reduceren van het aantal iteraties (stappen), toont SwD aan dat het dynamisch aanpassen van de resolutie tijdens het generatieproces een veelbelovende route is voor efficiëntie.

De combinatie van schaal-gebaseerde distillatie en de MMD-loss biedt een nieuw paradigma voor het trainen van snelle generatieve modellen. Het elimineert de noodzaak van zware adversarial training en maakt het mogelijk om bestaande grote modellen snel en efficiënt om te zetten in real-time generatoren, wat cruciaal is voor de toepassing van AI in interactieve en resource-beperkte omgevingen. De methode is breed toepasbaar op zowel afbeeldingen als video's.