OBS-Diff: Accurate Pruning For Diffusion Models in One-Shot

Dit paper introduceert OBS-Diff, een trainingsvrij one-shot pruning-framework dat door het aanpassen van de Optimal Brain Surgeon-methode met een tijdstap-bewuste Hessian en een efficiënte groepsstrategie, grote tekst-naar-beeld diffusiemodellen nauwkeurig comprimeert met minimale kwaliteitsverlies.

De kern

Stel je voor dat je een gigantische, hyperrealistische kunstenaar hebt die prachtige afbeeldingen kan maken op basis van een simpele tekstbeschrijving. Dit is wat moderne "diffusiemodellen" doen. Maar er is een probleem: deze kunstenaars zijn zo zwaar en complex dat ze een enorme hoeveelheid rekenkracht en geheugen nodig hebben. Het is alsof je een vrachtwagen vol met gereedschap gebruikt om een postzegel te schilderen.

Deze paper introduceert OBS-Diff, een slimme methode om deze kunstenaars te "versnellen" zonder hun talent te verliepen. Hier is hoe het werkt, uitgelegd met alledaagse vergelijkingen:

1. Het Probleem: De "Grote Snoeischaar" werkt niet

Bestaande methoden om deze modellen kleiner te maken (pruning) werken vaak alsof je zomaar willekeurige onderdelen uit een auto haalt om brandstof te besparen. Bij diffusiemodellen werkt dat niet goed, omdat ze werken in stappen. Ze beginnen met een wazig beeld en maken het steeds scherper, stap voor stap.

Als je een foutje maakt in de eerste stap (het wazige begin), vermenigvuldigt die fout zich door alle volgende stappen. Het is alsof je in het begin van een reis een verkeerd kompas neemt; hoe verder je komt, hoe verder je van je bestemming verwijderd bent. Bestaande methoden snijden vaak net die cruciale eerste stappen weg, waardoor het eindresultaat een puinhoop wordt.

2. De Oplossing: OBS-Diff (De Slimme Tuinman)

OBS-Diff is een nieuwe, "één-slag" methode (je hoeft het model niet opnieuw te trainen) die werkt als een zeer ervaren tuinman. In plaats van willekeurig takken weg te knippen, kijkt hij precies naar welke takken het belangrijkst zijn voor de groei van de plant.

Hier zijn de drie slimme trucs die OBS-Diff gebruikt:

A. De "Tijdsbewuste" Snoeiplan (Timestep-Aware)

Stel je voor dat je een film draait. De eerste scène is cruciaal; als die verkeerd is, is de hele film een ramp. OBS-Diff begrijpt dit. Het geeft extra aandacht aan de eerste stappen van het beeldgeneratieproces.

• De analogie: Het is alsof je een boek schrijft. Je besteedt meer tijd aan het perfectioneren van de eerste hoofdstukken (de basis), omdat die bepalen hoe de rest van het verhaal loopt. OBS-Diff "weegt" deze eerste stappen zwaarder en zorgt dat de belangrijkste onderdelen daar niet worden weggehaald.

B. De "Groepsaanpak" (Module Packages)

Het berekenen van welke onderdelen weg kunnen, is normaal gesproken extreem duur en tijdrovend (alsof je elke steen in een muur apart moet testen).

• De analogie: In plaats van elke steen één voor één te testen, groepeert OBS-Diff de stenen in blokken (pakketten). Het test een heel blok tegelijk. Dit bespaart enorm veel tijd en energie, net als het verpakken van verhuisdozen in plaats van losse borden te dragen. Hierdoor kan het proces snel en efficiënt plaatsvinden zonder dat de computer vastloopt.

C. Drie Manieren om te Snoeien

OBS-Diff is flexibel. Het kan op drie manieren werken, afhankelijk van wat je nodig hebt:

1. Ongeordend snoeien: Je haalt hier en daar kleine draden weg (zoals het verwijderen van stofdeeltjes).
2. Half-geordend snoeien: Je haalt groepjes van 4 draden weg, waarbij er altijd 2 overblijven (perfect voor speciale computerchips).
3. Geordend snoeien: Je haalt hele onderdelen weg, zoals een hele "oog" (attentie-head) of een hele "hersencel" (neuron) uit het model.

3. Het Resultaat: Sneller, Kleiner, Beter

De tests in de paper tonen aan dat OBS-Diff wonderen doet:

• Kwaliteit: Zelfs als je 50% of meer van het model weghaalt, blijven de gegenereerde afbeeldingen prachtig en duidelijk. Andere methoden leveren dan vaak alleen maar ruis of vervormde gezichten op.
• Snelheid: De modellen worden sneller en nemen minder geheugen in beslag, waardoor ze zelfs op krachtige laptops of servers kunnen draaien.
• Geen training nodig: Het grootste voordeel is dat je het model niet maandenlang opnieuw hoeft te "leren". Je snijdt het een keer, en het werkt direct.

Samenvatting

Kortom, OBS-Diff is als een meester-restaurator die een zwaar, oud schilderij (het grote AI-model) schoonmaakt en verlicht. Hij verwijdert alleen de stof en de overtollige verf die het beeld verstoren, maar doet dit met zoveel zorg dat de essentie van het kunstwerk intact blijft. Het maakt de zwaarste AI-modellen lichter en sneller, zonder dat de magie verdwijnt.

Technische samenvatting

Probleemstelling

Grootschalige tekst-naar-beeld diffusiemodellen (zoals Stable Diffusion 3 en Flux) hebben enorme rekenkracht en geheugenvereisten, wat hun toegankelijkheid beperkt. Bestaande methoden voor modelcompressie, zoals kwantisatie en pruning, stuiten op specifieke uitdagingen bij diffusiemodellen:

1. Iteratieve aard: Diffusiemodellen werken via een iteratief proces van ruisverwijdering over meerdere tijdstappen. Bestaande "one-shot" pruning-methoden (ontwikkeld voor LLM's) kunnen niet direct worden toegepast omdat ze aannemen dat de parameters slechts één keer worden doorlopen, terwijl diffusiemodellen parameters delen over een hele traject van denoising-stappen.
2. Gebrek aan generalisatie: Veel huidige pruning-methoden zijn specifiek ontworpen voor U-Net-architecturen en werken niet goed op moderne, complexe architecturen zoals Multimodal Diffusion Transformers (MMDiT).
3. Kosten: Veel methoden vereisen kostbare hertraining of fine-tuning na het pruning-proces, wat het voordeel van compressie tenietdoet.

De centrale vraag is: Kan een algemeen, trainingsvrij pruning-framework worden ontwikkeld dat diverse diffusiemodelformaten ondersteunt en meerdere granulariteiten van pruning (unstructured, semi-structured, structured) in één keer (one-shot) mogelijk maakt?

Methodologie: OBS-Diff

De auteurs introduceren OBS-Diff, een framework dat de klassieke Optimal Brain Surgeon (OBS) methode revitaliseert en aanpast aan de unieke dynamiek van diffusiemodellen. Het framework bestaat uit drie kerncomponenten:

1. Tijdstap-bewuste Hessian-constructie (Timestep-Aware Hessian Construction)

In tegenstelling tot LLM's, waar fouten lineair optreden, hebben fouten in de vroege stappen van het denoising-proces van een diffusiemodel een cumulatief effect op de uiteindelijke output.

• Aanpassing: De auteurs herschrijven het pruning-doel om een gewogen reconstructiefout te minimaliseren.
• Logaritmische weging: Ze introduceren een wegingsschema ($\alpha_t$) dat gebaseerd is op een logaritmisch afnemende functie. Dit geeft vroege tijdstappen (kleine $t$) een hogere weging dan latere stappen.
• Hessiaan: De Hessian-matrix (die de kromming van de foutfunctie beschrijft) wordt berekend als een gewogen som over alle tijdstappen. Hierdoor worden gewichten die kritiek zijn voor de vroege, vormende fasen van de generatie, beter beschermd.

2. Groepsweise Sequentiële Pruning Strategie ("Module Packages")

Het sequentieel calibreren van elke laag afzonderlijk is computatief onhaalbaar voor diffusiemodellen omdat elke calibratie een volledige multi-stap denoising-traject vereist.

• Oplossing: Het model wordt opgedeeld in "Module Packages" (groepen van lagen met onafhankelijke inputs, zoals Q, K, V projecties).
• Proces: Voor elk pakket wordt één keer het volledige denoising-traject uitgevoerd om statistieken te verzamelen (via hooks). Vervolgens worden alle lagen in dat pakket gelijktijdig gepruned op basis van hun respectievelijke Hessiaans.
• Voordeel: Dit amortiseert de dure data-collectie kosten en maakt het proces haalbaar binnen redelijke geheugen- en tijdslimieten.

3. Uitbreiding naar Gestructureerde en Semi-Gestructureerde Pruning

Het framework is flexibel en ondersteunt:

• Unstructured: Het verwijderen van individuele gewichten.
• Semi-structured (N:M): Bijv. 2:4 sparsiteit, waarbij binnen blokken van 4 gewichten de 2 minst belangrijke worden verwijderd.
• Structured: Het verwijderen van hele neuronen (in FFN lagen) of attention heads (in MHA lagen). Voor gedeelde attention heads in MMDiT-architecturen gebruiken ze Reciprocal Rank Fusion (RRF) om de belangrijkheid over verschillende modaliteiten te fuseren tot één ranglijst.

Belangrijkste Bijdragen

1. Revitalisatie van OBS: De eerste toepassing van de Optimal Brain Surgeon methode op grote tekst-naar-beeld diffusiemodellen, met aanpassingen voor iteratieve dynamiek.
2. Tijdstap-bewuste Hessian: Een innovatieve methode om de cumulatieve impact van fouten in het denoising-proces te modelleren, waardoor vroege stappen prioriteit krijgen.
3. Efficiëntie: Een groepsgewijze strategie die de calibratiekosten drastisch verlaagt zonder de kwaliteit van het pruning te offeren.
4. Veelzijdigheid: Ondersteuning voor unstructured, semi-structured en structured pruning op diverse architecturen (U-Net en MMDiT).

Resultaten

De auteurs hebben OBS-Diff getest op modellen zoals Stable Diffusion v2.1, SD3-Medium, SD3.5-Large en Flux.1-dev, vergeleken met baselines zoals Magnitude, Wanda, DSnoT en EcoDiff.

• Kwaliteit: OBS-Diff behaalt state-of-the-art resultaten, vooral bij hoge sparsiteitsniveaus (50-70%). Waar andere methoden bij hoge sparsiteit vaak volledig falen (beeldartefacten, negatieve ImageReward scores), behoudt OBS-Diff hoge visuele kwaliteit en semantische consistentie.
• Metingen:
  • ImageReward: OBS-Diff scoort consequent het hoogst, wat aangeeft dat de gegenereerde beelden beter overeenkomen met menselijke voorkeuren.
  • FID & CLIP: De methode behoudt de dichtstbijzijnde prestaties ten opzichte van het oorspronkelijke dichte model.
  • Voorbeeld: Op SD3.5-Large bij 60% sparsiteit collapseert de prestatie van baselines, terwijl OBS-Diff een ImageReward van 0.3984 behoudt (tegenover negatieve scores bij anderen).
• Efficiëntie: Het pruning-proces voor een 2B-parameter model duurt minder dan 15 minuten op één NVIDIA RTX 4090. De inferentie versnelt met factor 1.23x (semi-structured) tot 1.31x (structured bij 30% sparsiteit).
• Robuustheid: De methode werkt goed over verschillende CFG-schalen, sample-stappen en solvers (zoals Euler, Heun, DPM++), wat aantoont dat het niet overfit op de calibratie-instellingen.

Betekenis en Impact

OBS-Diff is een doorbraak in het veld van modelcompressie voor generatieve AI. Het lost het fundamentele probleem op dat bestaande "one-shot" pruning-methoden niet werken voor de iteratieve aard van diffusiemodellen.

• Toegankelijkheid: Het maakt het mogelijk om zeer grote modellen (zoals Flux en SD3.5) lokaal en efficiënt te draaien zonder dure hertraining.
• Generalisatie: Het framework is niet beperkt tot specifieke architecturen en kan potentieel worden toegepast op andere diffusie-gebaseerde modellen, zoals video-generatie en diffusion-LLM's.
• Standaard: Het zet een nieuwe benchmark voor trainingsvrije compressie, waarbij het de trade-off tussen compressiegraad en beeldkwaliteit significant verbetert ten opzichte van bestaande technieken.

Kortom, OBS-Diff biedt een robuuste, efficiënte en veelzijdige oplossing om de rekenlast van de nieuwste generatie diffusiemodellen te verlagen, terwijl de creatieve kwaliteit behouden blijft.