Diff-ES: Stage-wise Structural Diffusion Pruning via Evolutionary Search

Diff-ES is een nieuw framework dat via evolutionaire zoektocht een geoptimaliseerd, geheugen-efficiënt structuurpruning-schema voor diffusion-modellen ontwikkelt, waardoor aanzienlijke versnelling wordt bereikt met minimale kwaliteitsverlies.

De kern

🎨 De Kunst van het Versnellen: Wat is Diff-ES?

Stel je voor dat een Diffusiemodel (zoals de AI die prachtige plaatjes maakt) een kunstenaar is die een schilderij maakt. Maar deze kunstenaar werkt heel langzaam: hij begint met een vaag, wazig potloodschetsje en moet honderden keer over het canvas gaan om het steeds scherper en gedetailleerder te maken.

Het probleem? Dit proces kost enorm veel tijd en rekenkracht. Het is alsof je een auto moet bouwen door elke bout met de hand vast te draaien, in plaats van met een boormachine.

Wetenschappers hebben al eerder geprobeerd dit sneller te maken door:

1. Minder stappen te nemen: De kunstenaar maakt het schilderij in 10 stappen in plaats van 100.
2. De kunstenaar te verkleinen: Ze halen onderdelen uit de "hersenen" van de AI weg (dit heet pruning), zodat hij lichter is.

Maar hier zit een addertje onder het gras: Niet elke stap is even belangrijk.

• De eerste stappen zijn cruciaal voor de grote lijnen (waar staat het huis? Wie zit erin?). Als je hier fouten maakt, is het hele schilderij verkeerd.
• De laatste stappen zijn voor de fijne details (de textuur van de huid, de glans in de ogen). Hier mag je wat meer "slordig" zijn zonder dat het er vreselijk uitziet.

🚧 Het oude probleem: De "MosaicDiff" methode

Vroeger probeerden mensen dit op te lossen met een methode genaamd MosaicDiff. Ze deelden het schilderijproces in drie grote blokken op en zeiden: "Oké, in het eerste blok doen we het heel precies, in het tweede blok iets minder, en in het laatste blok heel snel."

Het probleem? Ze deden dit op basis van gokjes en ervaring. Ze dachten: "Ik denk dat het midden het belangrijkst is."
Maar elke AI is anders! Wat voor de ene kunstenaar werkt, werkt niet voor de andere. Het was alsof je probeert een recept voor een taart te maken door te raden hoeveel suiker je nodig hebt, zonder te proeven. Soms werkt het, maar vaak is de taart te droog of te zoet.

✨ De oplossing: Diff-ES (De Slimme Zoeker)

De auteurs van dit paper hebben een nieuwe methode bedacht: Diff-ES.

Stel je voor dat je een evolutionaire zoektocht doet, net als in de natuur.

1. Het Idee: In plaats van te gokken, laten we de computer zelf ontdekken welke stappen belangrijk zijn.
2. De Populatie: We beginnen met 20 verschillende "plannen" (bijvoorbeeld: Plan A is streng in het begin, Plan B is streng in het midden).
3. De Test: We laten de AI met elk plan een paar plaatjes maken.
4. De Selectie: We kijken welke plannen de mooiste plaatjes opleveren. De slechte plannen worden weggegooid.
5. De Evolutie: De goede plannen worden "gemuteerd" (een beetje aangepast). Misschien is Plan A goed, maar als we in stap 3 iets minder streng zijn, wordt het nog beter?
6. Herhaling: Dit proces herhalen we honderden keren. Uiteindelijk houden we het perfecte plan over dat precies weet hoeveel energie er in elke stap moet worden gestoken.

🚂 De Magische Trein: "Weight Routing"

Er was nog een groot probleem: het testen van al deze plannen kostte te veel geheugen. Het was alsof je voor elk plan een hele nieuwe trein moest bouwen om te testen of hij snel rijdt.

Diff-ES lost dit op met een slimme truc genaamd Weight Routing (Gewicht-Routering).

• De Oude Manier: Je bouwt drie aparte treinen (voor de drie stappen) en rijdt ermee. Dat kost veel ruimte in je garage (het geheugen van je computer).
• De Diff-ES Manier: Je hebt één trein. Maar je hebt een slimme machinist die tijdens het rijden de motoren kan in- en uitschakelen of vervangen.
  • Als de trein in de "beginfase" zit, schakelt hij de krachtige motoren in.
  • Als hij in de "laatste fase" zit, schakelt hij over op zuinige motoren.
  • Je hoeft geen nieuwe trein te bouwen; je past alleen de onderdelen aan die op dat moment nodig zijn.

Dit bespaart enorm veel ruimte en maakt het proces veel sneller.

🏆 Wat is het resultaat?

Door deze slimme zoektocht en de efficiënte trein, kan Diff-ES:

• Veel sneller plaatjes maken (snelheidswinst).
• Beter kwaliteit behouden dan de oude methoden (de plaatjes zien er nog steeds scherp en mooi uit).
• Zich aanpassen aan elk type AI (of het nu een CNN is of een Transformer), in plaats van één starre regel te gebruiken.

Kortom: Diff-ES is als een slimme chef-kok die niet blindelings een recept volgt, maar eerst proeft en experimenteert om precies te weten hoeveel tijd en energie hij in elk onderdeel van het gerecht moet steken, zodat het gerecht perfect wordt zonder dat de keuken in brand vliegt.

Technische samenvatting

Probleemstelling

Diffusiemodellen hebben indrukwekkende resultaten geboekt op het gebied van hoogwaardige beeldgeneratie, maar zijn computationeel zeer veeleisend vanwege hun multi-stap denoising-proces en grote modelgroottes. Bestaande methoden voor efficiëntieverbetering richten zich vaak op het verminderen van het aantal sample-steps of het comprimeren van modelparameters.

Echter, bestaande benaderingen voor gestructureerde pruning (het verwijderen van kanalen of lagen) kampen met twee belangrijke beperkingen:

1. Niet-uniforme belangstelling: Het belang van diffusiestappen is sterk niet-uniform en afhankelijk van het specifieke model. Stappen aan het begin van het proces construeren vaak de globale semantische structuur, terwijl latere stappen fijne details verfijnen.
2. Suboptimale handmatige planning: Methoden zoals MosaicDiff gebruiken handmatig ingestelde, heuristische plannen voor de verdeling van de pruning over verschillende stadia. Ze verdelen het proces in vaste stadia (bijv. drie) en "naaien" tijdens de inferentie meerdere onafhankelijk geprepareerde modellen aan elkaar. Dit leidt tot:
   • Suboptimale prestaties omdat de heuristiek niet generaliseert naar verschillende architecturen (bijv. DiT vs. SDXL).
   • Hoge geheugenoverhead door het dupliceren van modelparameters voor elk stadium.

Methodologie: Diff-ES

De auteurs stellen Diff-ES voor, een raamwerk dat evolutionaire zoektocht combineert met gestructureerde pruning om een optimale, stadium-specifieke sparsiteitsplanning te vinden zonder het model opnieuw te trainen.

Kerncomponenten:

1. Probleemformulering:
   Het doel is om een vector van sparsiteitsniveaus $\mathbf{L} = \{L_1, L_2, ..., L_n\}$ te vinden voor $n$ stadia van het denoising-traject, zodat de beeldkwaliteit ($f(\mathbf{L})$) wordt gemaximaliseerd onder een vaste globale sparsiteitsbudget ($G(\mathbf{L}) \geq L_t$). Omdat de relatie tussen pruning en kwaliteit niet-differentieerbaar is, is gradient-based optimatie onmogelijk.

2. Evolutionaire Zoektocht (Level-Switch Mutation):

   • Initialisatie: Een populatie van kandidaat-planningen wordt gegenereerd (uniform, willekeurig of heuristisch).
   • Fitness Evaluatie: Elke kandidaat wordt beoordeeld door een kleine batch afbeeldingen te genereren en een lichtgewicht fitness-maatstaf (zoals CLIP-IQA, TOPIQ of SSIM) te berekenen.
   • Mutatie: Een "level-switch" mutatie wisselt sparsiteitsniveaus tussen twee willekeurige stadia ($L_i \leftarrow L_i + \Delta$, $L_j \leftarrow L_j - \Delta$) terwijl het totale budget constant blijft. Dit zorgt voor een efficiënte verkenning van de discrete zoekruimte.
   • Selectie: De beste individuen overleven naar de volgende generatie.

3. Stadium-specifieke Tweede-orde Structurale Pruning:
   Voor elk stadium wordt een specifieke pruning-methode toegepast (o.a. OBS, Wanda, LayerDrop).

   • SNR-bewuste Kalibratie: Er wordt een kalibratie-set gebruikt die is afgestemd op het Signaal-Ruisverhouding (SNR) van dat specifieke stadium, zodat de pruning gebaseerd is op de werkelijke ruiscondities tijdens de inferentie.
   • OBS (Optimal Brain Surgeon): Er wordt gebruikgemaakt van tweede-orde informatie (Hessian-matrix) om redundantie te verwijderen en de fout te compenseren, wat leidt tot hogere kwaliteit dan eerste-orde methoden.

4. Efficiënte Gewichtsroutering (Weight Routing):
   Een cruciale innovatie is het vermijden van het herberekenen van de dure tweede-orde pruning voor elke kandidaat in de evolutionaire zoektocht.

   • In plaats van meerdere modellen te dupliceren (zoals bij MosaicDiff), wordt één gedeeld model in het GPU-geheugen gehouden.
   • De volledige pruning-trajecten voor elk stadium worden eenmalig vooraf berekend en opgeslagen in een compacte database.
   • Tijdens de zoektocht worden de benodigde gewichten voor een specifiek stadium dynamisch "gerouteerd" (opgehaald) uit deze database. Dit elimineert de noodzaak tot modelstitching en verlaagt het geheugengebruik aanzienlijk.

Belangrijkste Bijdragen

1. Identificatie van een beperking: De auteurs tonen aan dat handmatige, heuristische planningen (zoals bij MosaicDiff) niet generaliseren over verschillende modelarchitecturen en vaak suboptimaal zijn.
2. Nieuw Evolutionair Raamwerk: Diff-ES automatiseert de toewijzing van computatiebudget over de diffusiestadia via evolutionaire zoektocht, wat leidt tot model-specifieke en optimale plannen.
3. Plug-and-Play Compatibiliteit: Het raamwerk werkt naadloos samen met bestaande pruning-methoden (Layer Dropping, Wanda, OBS) en is toepasbaar op zowel CNN-gebaseerde modellen (SDXL) als Transformer-gebaseerde modellen (DiT).
4. Efficiëntie zonder Kwaliteitsverlies: Door het gebruik van gewichtsroutering wordt de geheugenoverhead van modelstitching vermeden, terwijl de kwaliteit van de gegenereerde beelden behouden blijft of zelfs verbetert.

Resultaten

De methode is getest op DiT-XL/2 en SDXL-base-1.0 en vergeleken met state-of-the-art baselines zoals Diff-Pruning, DeepCache, OBS-Diff en MosaicDiff.

• Beeldkwaliteit: Diff-ES overtreft consistent alle baselines.
  • Op SDXL (30% sparsiteit): Diff-ES bereikt een FID van 25.87 (vs. 59.09 voor MosaicDiff met OBS) en behoudt een hoge structurele gelijkenis (SSIM).
  • Op DiT (50% sparsiteit): Diff-ES bereikt een FID van 4.25 (vs. 6.13 voor MosaicDiff).
• Snelheid: De methode levert aanzienlijke versnellingen op (tot 1.52x speedup op DiT) met minimale vertraging door de routering.
• Geheugen: Diff-ES gebruikt aanzienlijk minder GPU-geheugen dan MosaicDiff (bijv. ~42% minder geheugen op DiT) omdat het geen meerdere modellen hoeft te laden.
• Generalisatie: De gevonden plannen generaliseren goed naar prompts buiten de trainingsset en zijn robuust tegenover variaties in hyperparameters (zoals het aantal generaties of stadia).

Betekenis en Impact

Diff-ES vertegenwoordigt een significante doorbraak in de compressie van diffusiemodellen. Het bewijst dat de verdeling van rekenkracht over de tijdstappen van het denoising-proces niet uniform moet zijn en dat deze verdeling automatisch geoptimaliseerd kan worden in plaats van handmatig te worden ingesteld.

Door de combinatie van evolutionaire zoektocht en efficiënte gewichtsroutering, biedt Diff-ES een schaalbare oplossing die:

1. De kwaliteit van gegenereerde beelden behoudt bij hoge compressiepercentages.
2. De implementatiekosten verlaagt door het elimineren van dure hertraining en modelduplicatie.
3. Een nieuwe standaard zet voor gestructureerde pruning, waarbij de specifieke dynamiek van elke diffusiestap wordt benut voor maximale efficiëntie.

Het paper concludeert dat Diff-ES een robuuste en algemene oplossing biedt voor het versnellen van diffusiemodellen, met name voor real-time toepassingen waar zowel snelheid als beeldkwaliteit cruciaal zijn.