Evolutionary Caching to Accelerate Your Off-the-Shelf Diffusion Model

Probleemstelling

Diffusiemodellen voor beeldgeneratie (zoals Diffusion Transformers of DiTs) staan bekend om hun hoge kwaliteit, maar lijden onder een zeer trage en rekenintensieve inferentie. Dit komt doordat het generatieproces vele iteratieve stappen vereist (vaak 20 tot 50 stappen), waarbij elke stap een volledige doorvoer door het netwerk nodig heeft.

Bestaande methoden om dit te versnellen zonder het model opnieuw te trainen (training-free), vertrouwen vaak op caching: het hergebruiken van interne features (zoals attention-matrices) van eerdere stappen. Echter, deze methoden hebben twee grote beperkingen:

1. Rigiditeit: Ze gebruiken vaste heuristieken (handmatig ontworpen regels) om te beslissen wat er wordt gecached. Dit resulteert in een beperkt aantal discrete opties (bijv. "2x sneller met kwaliteitsverlies" of "3x sneller met groot verlies").
2. Gebrek aan generalisatie: Handmatige hyperparameters moeten vaak per model en per resolutie worden getuned. Ze generaliseren slecht naar nieuwe modelvarianten of resoluties.

Er is een behoefte aan een methode die een gladde Pareto-grens (een continu spectrum van trade-offs tussen snelheid en kwaliteit) kan vinden, zonder dat het model zelf wordt aangepast.

Methodologie: ECAD

De auteurs introduceren ECAD (Evolutionary Caching to Accelerate Diffusion models), een raamwerk dat caching omvormt tot een multi-objectief optimalisatieprobleem opgelost met een genetisch algoritme.

Kernprincipes:

1. Caching als Optimalisatieprobleem: In plaats van vaste regels, wordt caching gemodelleerd als een binaire tensor $S \in \{0, 1\}^{N \times B \times C}$, waarbij:

   • $N$ het aantal diffusiestappen is.
   • $B$ het aantal transformer-blokken is.
   • $C$ het aantal cacheerbare componenten per blok is (Self-Attention, Cross-Attention, Feed-Forward).
   • Een '0' betekent "gebruik de cache", een '1' betekent "herbereken".

2. Genetisch Algoritme (NSGA-II):

   • Populatie: Een verzameling van mogelijke caching-schema's (schedule).
   • Calibratie: Het algoritme gebruikt een kleine set van tekst-prompten (bijv. 100 prompts) om kandidaat-schema's te evalueren. Er zijn geen afbeeldingen nodig voor de training, alleen de gegenereerde resultaten.
   • Fitness Functie: De "fitheid" van een schema wordt bepaald door twee objectieven:
     1. Kwaliteit: Gemeten via Image Reward (een model dat menselijke voorkeur simuleert).
     2. Kosten: Gemeten via MACs (Multiply-Accumulate Operations) of latentie.
   • Evolutie: Het algoritme gebruikt selectie, crossover (4-punts crossover) en mutatie (bit-flip) om de populatie over generaties te laten evolueren. Het doel is het vinden van de Pareto-grens: de set van schema's waar je niet sneller kunt worden zonder kwaliteit te verliezen, en niet beter zonder trager te worden.

3. Onafhankelijkheid: ECAD vereist geen wijzigingen in de netwerkparameters, geen gradients, en geen referentie-afbeeldingen. Het werkt puur op het niveau van de inferentie-schedule.

Belangrijkste Bijdragen

1. Paradigmaverschuiving: De eerste methode die caching voor diffusiemodellen benadert als een multi-objectief optimalisatieprobleem in plaats van een handmatig ontworpen heuristiek.
2. Pareto-grens Ontdekking: ECAD kan een continu spectrum van snelheid-kwaliteit trade-offs vinden, waardoor gebruikers precies het punt kunnen kiezen dat past bij hun latency-beperkingen.
3. Generalisatie: De gevonden schema's generaliseren uitstekend naar:
   • Nieuwe Resoluties: Een schema geoptimaliseerd op 256x256 werkt effectief op 1024x1024 zonder heroptimalisatie.
   • Modelvarianten: Schema's getraind op PixArt-α kunnen worden overgedragen naar PixArt-Σ en FLUX.1-dev met minimale prestatieverlies.
4. Efficiëntie: Het proces is volledig asynchroon en vereist geen extra VRAM voor het model zelf, wat het mogelijk maakt om op kleinere GPU's te draaien tijdens de zoektocht.

Resultaten

De auteurs evalueren ECAD op drie state-of-the-art modellen: PixArt-α, PixArt-Σ en FLUX.1-dev.

• Prestaties op PixArt-α:

  • ECAD vindt een schema dat de vorige state-of-the-art (ToCa) verslaat met 4.47 punten verbetering in FID op de COCO-dataset.
  • Tegelijkertijd stijgt de versnelling van 2.35x (ToCa) naar 2.58x (ECAD).
  • Bij een nog agressievere instelling (3.37x versnelling op FLUX.1-dev) blijft de kwaliteit concurrerend.

• Generalisatie:

  • Resolutie: Schema's getraind op 256x256 voor FLUX.1-dev presteren beter dan methoden die specifiek voor 1024x1024 zijn getuned, zelfs op die hogere resolutie.
  • Model Transfer: Het overdragen van een schema van PixArt-α naar PixArt-Σ resulteert in slechts een kleine prestatiedaling, terwijl handmatige methoden (zoals ToCa) vaak falen bij modeloverschakeling zonder volledige her-tuning.

• Ablatiestudies:

  • Het algoritme is robuust tegen variaties in de grootte van de calibratie-set (100 prompts volstaan).
  • Het gebruik van gegenereerde prompts (via LLM) werkt even goed als menselijk gecurateerde sets.
  • De kwaliteit van de gevonden schema's verbetert gestaag naarmate het aantal generaties toeneemt.

Betekenis en Impact

ECAD biedt een schaalbare en generaliseerbare oplossing voor het versnellen van diffusiemodellen. Door het af te koppelen van handmatige tuning en modeltraining, maakt het het mogelijk om:

• Off-the-shelf modellen direct te optimaliseren voor specifieke hardware of latency-eisen.
• Een fine-grained controle te hebben over de kwaliteit-versus-snelheid trade-off, wat cruciaal is voor real-world toepassingen (bijv. mobiele apps vs. server-side batch processing).
• De energie-efficiëntie van beeldgeneratie te verbeteren door onnodige berekeningen te elimineren zonder de kwaliteit te offeren.

De methode stelt een nieuwe standaard voor training-free versnelling en toont aan dat evolutionaire algoritmen superieur kunnen zijn aan handmatige heuristieken bij het vinden van complexe optimalisatieproblemen in de AI-infrastructuur.