Sparse-to-Sparse Training of Diffusion Models

Technische Samenvatting: Sparse-to-Sparse Training van Diffusiemodellen

Probleemstelling
Diffusiemodellen (DMs) hebben de staat van de kunst bereikt in generatieve taken, waaronder beeld-synthese, natuurlijke taalverwerking en temporele datamodellering. De brede adoptie ervan wordt echter gehinderd door aanzienlijke computationele kosten tijdens zowel de training als de inferentie. Terwijl eerder onderzoek zich grotendeels heeft gericht op het versnellen van de inferentiefase (bijv. door middel van distillatie of snellere sampling), blijft het trainingsproces uitgebreid en intensief qua middelen. Bovendien vereisen bestaande netwerkcompressietechnieken, zoals post-training kwantisatie of het wegstrippen (pruning) van vooraf getrainde dichte netwerken, nog steeds de aanvankelijk dure training van volledige schaal dichte netwerken. Er is behoefte aan methoden die de computationele kosten en het geheugengebruik vanaf het begin van het trainingsproces verminderen.

Methodologie
Dit artikel introduceert het paradigma van sparse-to-sparse training voor diffusiemodellen, met als doel om vanaf de basis (from scratch) ijle (sparse) netwerken te trainen in plaats van vooraf getrainde dichte modellen te comprimeren. De studie richt zich op ongestructureerde ijleheid (unstructured sparsity) en evalueert drie specifieke trainingsstrategieën toegepast op twee state-of-the-art DM-architecturen: Latent Diffusion (voor continue beelddata) en ChiroDiff (voor discrete spatiotemporele schetsdata).

De drie voorgestelde methoden zijn:

1. Static-DM (Statische ijle training): De netwerkconnectiviteit wordt bij de initialisatie vastgesteld op basis van een ijleheidsratio. De auteurs maken gebruik van Erdős–Rényi (ER) en ERK-strategieën om niet-nul gewichten toe te wijzen, waarbij grotere lagen een hogere ijleheid hebben. Het connectiviteitspatroon blijft gedurende de hele training ongewijzigd.
2. RigL-DM (Dynamische ijle training): Deze methode maakt gebruik van een dynamisch prune-and-grow schema. Gewichten worden gepruned op basis van hun grootte, en nieuwe verbindingen worden geregenereerd op basis van de grootte van de gradiënten (gradient growth).
3. MagRan-DM (Dynamische ijle training): Vergelijkbaar met RigL-DM wat betreft het prunen (gebaseerd op grootte), maar de regeneratie wordt bepaald door willekeur in plaats van door gradiënten.

De experimenten bestrijken zes datasets (LSUN-Bedrooms, CelebA-HQ, Imagenette, QuickDraw, KanjiVG en VMNIST) over verschillende ijleheidsniveaus ($S \in \{0.1, 0.25, 0.5, 0.75, 0.9\}$). De auteurs onderzoeken ook de impact van de prune- en regeneratiesnelheid ($p$) en de exploratiefrequentie ($\Delta T_e$) voor dynamische methoden.

Belangrijkste Bijdragen

• Eerste toepassing van Sparse-to-Sparse Training op DMs: Het artikel stelt de haalbaarheid vast van het trainen van onvoorwaardelijke diffusiemodellen vanaf de basis met behulp van ijle connectiviteit, waarbij zowel statische als dynamische strategieën worden behandeld.
• Uitgebreide Empirische Evaluatie: De studie benchmarkt drie ijle trainingsmethoden tegenover dichte baselines over twee verschillende modelarchitecturen en zes datasets, wat een breed beeld geeft van de effecten van ijleheid op DM-prestaties.
• Identificatie van Optimale Hyperparameters: Het onderzoek identificeert dat hoewel dynamische ijle training (DST) over het algemeen goed presteert, de keuze van de prune- en regeneratiesnelheid cruciaal is. Een conservatieve snelheid ($p=0.05$) levert betere resultaten op in regimes met hoge ijleheid vergeleken met de agressievere snelheid ($p=0.5$) die vaak in andere contexten wordt gebruikt.
• Prestatie versus Efficiëntie-afwegingen: Het werk demonstreert dat ijle DMs de prestaties van hun dichte tegenhangers kunnen evenaren of zelfs overtreffen, terwijl ze het aantal trainbare parameters en de Floating Point Operations (FLOPs) aanzienlijk verminderen.

Experimentele Resultaten

• Prestaties: In de meeste gevallen presteerde ten minste één ijle configuratie beter dan zijn dichte tegenhanger. Zo behaalde RigL-DM met 25% ijleheid op de CelebA-HQ dataset bijvoorbeeld een lagere FID-score (32.12) dan het dichte model (32.74), terwijl het de parameters met 25% en de FLOPs met 10% verminderde. Op QuickDraw behield MagRan-DM met 90% ijleheid een vergelijkbare prestatie als het dichte model, terwijl de FLOPs met 88% werden verminderd.
• IJleheidsniveaus: Matige ijleheidsniveaus (25%–50%) leverden over het algemeen de beste balans tussen prestaties en efficiëntie op. Extreme ijleheid (bijv. 90%) leidde vaak tot prestatievermindering, hoewel dit varieerde per dataset en modelarchitectuur.
• Dynamisch versus Statisch: Hoewel methoden voor dynamische ijle training (DST) (RigL-DM, MagRan-DM) vaak beter presteerden in regimes met matige ijleheid, toonde Static-DM competitieve of superieure resultaten bij zeer hoge ijleheid ($S > 0.5$) wanneer de agressieve $p=0.5$ snelheid werd gebruikt. Echter, wanneer de conservatieve $p=0.05$ snelheid werd gebruikt, herwonnen DST-methoden de superioriteit, zelfs in regimes met hoge ijleheid.
• Trainingsdynamiek: IJle modellen vertoonden stabiele trainingsdynamieken die vergelijkbaar zijn met die van hun dichte tegenhangers, wat suggereert dat het sparse-to-sparse paradigma het diffusie-trainingsproces niet destabiliseert.
• Conditionele Generatie: Voorlopige experimenten op klasse-conditionele modellen (ImageNet en QuickDraw) gaven aan dat sparse-to-sparse training ook effectief is in conditionele settings, hoewel dichte modellen in ImageNet over het algemeen een lichte voorsprong behielden.

Betekenis en Claims
Het artikel beweert dat sparse-to-sparse training een veelbelovend pad biedt om de efficiëntie van diffusiemodellen te verbeteren zonder de kwaliteit van de generatie op te offeren. Door ijle netwerken vanaf de basis te trainen, pakt de aanpak zowel de trainings- als de inferentiekosten aan, wat potentieel de ecologische voetafdruk en de geheugenvereisten van DMs vermindert. De auteurs benadrukken dat het selecteren van de juiste methode (Statisch versus Dynamisch), het ijleheidsniveau en de hyperparameters (specifiek de prune/regeneratiesnelheid) cruciaal is voor succes. Het werk suggereert dat, met de juiste configuratie, ijle DMs een levensvatbaar alternatief kunnen zijn voor dichte modellen, wat bredere implementatie en snellere iteratie in generatieve AI-onderzoek mogelijk maakt. De auteurs merken op dat hoewel de huidige hardware de praktische versnelling van ongestructureerde ijleheid vaak beperkt, de vermindering in aantal parameters en FLOPs een aanzienlijk theoretisch en praktisch voordeel blijft, vooral naarmate de ondersteuning van hardware voor ijle operaties evolueert.