DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

Probleemstelling

Diffusiemodellen, en specifiek Diffusion Transformers (DiT), hebben uitstekende prestaties geleverd in beeldgeneratie, maar lijden onder een enorme rekenkosten door hun multi-stap inferentiemechanisme. Bestaande versnellingsmethoden, zoals feature caching (het opnieuw gebruiken van tussenliggende features), stuiten op de volgende beperkingen:

1. Manuele Sparsiteit: Bestaande token-cache methoden vereisen handmatig ontworpen plannen voor sparsiteit (welke tokens worden genegeerd), wat inefficiënt is en niet optimaal past bij de dynamiek van het model.
2. Inefficiënte Caching: Veel methoden behouden volledige forward-passes in bepaalde stappen om de kwaliteit te waarborgen, wat de potentiële versnelling beperkt.
3. Gebrek aan End-to-End Optimalisatie: Er is geen geautomatiseerd systeem dat de beste verdeling van sparsiteit over de lagen en tijdstappen leert zonder menselijke tussenkomst.

Methodologie: DiffSparse

De auteurs stellen DiffSparse voor, een differentieerbaar framework dat de token-sparsiteit in diffusiemodellen optimaliseert via een end-to-end leerproces. De kern van de methode bestaat uit drie componenten:

1. Leerbaar Sparsiteitskosten Predictor (Learnable Sparsity Cost Predictor):

   • Dit is een netwerk dat een kostenmatrix genereert die de "kosten" (verlies aan kwaliteit) voorspelt voor het toepassen van verschillende sparsiteitspercentages in elke laag en op elk tijdstip van het denoising-proces.
   • De grootte van deze predictor hangt alleen af van het aantal tijdstappen ($T$), het aantal lagen ($L$) en het aantal kandidaat-sparsiteitsniveaus ($|S|$), en niet van de token-lengte, wat schaalbaarheid garandeert.

2. Dynamisch Programmeren Oplosser (Dynamic Programming Solver):

   • Gegeven een globale sparsiteitsbeperking (bijv. 54% van de tokens overslaan), gebruikt de solver een dynamisch programmeringsalgoritme om de optimale combinatie van sparsiteitsniveaus per laag en tijdstip te vinden die de totale kosten minimaliseert.
   • Dit zorgt ervoor dat de rekenkracht wordt toegewezen aan de meest kritieke momenten en lagen, terwijl minder belangrijke delen worden genegeerd.

3. Token Selector:

   • Op basis van de door de solver bepaalde sparsiteitsratio's selecteert deze module dynamisch welke tokens opnieuw berekend moeten worden en welke uit de cache worden gehaald.
   • De selectie is gebaseerd op een samengestelde importantiescore (o.a. self-attention invloed, cross-attention focus en cache-hergebruiksfrequentie).

Trainingsstrategie (Two-Stage Training):
Om het probleem van het behouden van volledige stappen op te lossen, introduceren de auteurs een tweefasige trainingsstrategie:

• Fase 1: Het model leert onafhankelijk de kosten voor volledige stappen en de laag-sparsiteit.
• Fase 2: De kosten van de volledige stappen worden geïntegreerd in de laag-sparsiteitsoptimalisatie. Hierdoor kan het model geleidelijk de afhankelijkheid van volledige stappen verminderen en de berekening dynamisch verdelen zonder kwaliteitsverlies.
• Verliesfunctie: Het gebruik van LPIPS (Learned Perceptual Image Patch Similarity) als distillatieverlies zorgt ervoor dat de student (gepruneerde) het perceptuele gedrag van de leraar (originele) nabootst, in plaats van alleen pixel-fouten te minimaliseren.

Kernbijdragen

1. Differentieerbare Sparsiteitsoptimalisatie: Een volledig geautomatiseerd framework dat de token-sparsiteit per laag leert zonder handmatige heuristieken.
2. Eliminatie van Volledige Stappen: Door de twee-fasige trainingsstrategie wordt de noodzaak vooraf gedefinieerde volledige forward-passes verwijderd, waardoor de versnellingpotentie van token-caching volledig wordt benut.
3. Schaalbaarheid: De methode generaliseert goed naar hogere resoluties; een predictor getraind op 256x256 werkt effectief op 512x512 zonder hertraining.

Resultaten

DiffSparse werd getest op diverse state-of-the-art modellen, waaronder PixArt-α, DiT-XL/2, FLUX.1-schnell en Wan2.1 (video).

• Prestaties op PixArt-α (20 stappen):
  • Bij een versnelling van 1.91x (54% tokens gespaard) bereikte DiffSparse een FID van 27.79.
  • Dit is significant beter dan bestaande SOTA-methoden zoals ToCa (FID 28.35) en TaylorSeer (FID 29.08).
  • Opmerkelijk: De kwaliteit van DiffSparse was zelfs beter dan het originele, niet-versnelde model.
• Klassieke Conditionele Generatie (DiT-XL/2):
  • Bij een versnelling van 2.07x verbeterde DiffSparse de FID van 3.05 (ToCa) naar 2.81, wat een verbetering van 8% betekent.
• Video Generatie (Wan2.1):
  • Op de VBench-benchmark behaalde DiffSparse de hoogste score (43.83) met een versnelling van 2.05x, terwijl concurrenten (DuCa) lagere scores behaalden bij vergelijkbare compressie.
• Efficiëntie: De trainingskosten voor de sparsiteitspredictor zijn laag (ongeveer 4-10 uur op 8 GPU's) in vergelijking met genetische algoritmen of volledige distillatiepijplijnen.

Betekenis en Impact

DiffSparse vertegenwoordigt een doorbraak in de efficiëntie van generatieve AI. Door de rekenkosten drastisch te verlagen (tot wel 54%) zonder in te leveren op beeldkwaliteit – en in veel gevallen de kwaliteit zelfs te verbeteren – maakt het diffusiemodellen veel praktischer voor real-world toepassingen, zoals mobiele apparaten en edge devices. De methode lost het fundamentele dilemma op tussen snelheid en kwaliteit in token-caching door een leerbaar, dynamisch systeem te introduceren dat de complexiteit van het denoising-proces beter begrijpt dan statische, handmatig ontworpen regels.