HybridStitch: Pixel and Timestep Level Model Stitching for Diffusion Acceleration

HybridStitch is een nieuwe methode voor tekst-naar-afbeelding-generatie die de berekening versnelt door complexe beeldgebieden te verfijnen met een groot model en eenvoudige gebieden te genereren met een kleiner model, wat resulteert in een 1,83-voudige snelheidswinst op Stable Diffusion 3.

De kern

Stel je voor dat je een enorme, ingewikkelde muurschildering moet maken. Je hebt twee kunstenaars: een meester (een heel groot, krachtig computermodel) en een leerling (een kleiner, sneller model).

Normaal gesproken laat je de meester de hele schildering doen. Dat ziet er prachtig uit, maar het duurt eeuwen. Als je de leerling de hele schildering laat doen, gaat het snel, maar het resultaat is vaak rommelig en onvolledig.

De meeste bestaande methoden proberen een compromis te sluiten: "Laat de meester de eerste helft doen en de leerling de tweede helft." Maar dat is niet perfect. Soms is de achtergrond al klaar, terwijl de meester nog steeds aan het werk is aan de details van een gezicht. De meester blijft dan tijd verspillen aan dingen die de leerling ook prima kan doen.

HybridStitch is een slimme nieuwe manier om dit op te lossen. Het is alsof je de meester en de leerling tegelijkertijd laat werken, maar op een heel slimme manier.

Hier is hoe het werkt, stap voor stap:

1. De "Grote Plaat" (Het Begin)

Aan het begin van het proces is het beeld nog een wazige vlek van ruis. Niemand weet nog wat er te zien zal zijn. Daarom laat je de meester het hele plaatje bekijken. Hij zorgt ervoor dat de basisstructuur (de compositie) goed staat.

2. De "Slimme Verdeling" (Het Midden)

Na een tijdje begint het beeld vorm te krijgen. Nu kijkt HybridStitch heel nauwkeurig naar het plaatje:

• De makkelijke stukken: De lucht, de zee of een egaal gekleurd shirt. Deze delen zijn al duidelijk genoeg. De leerling mag deze delen nu overnemen en snel afmaken.
• De moeilijke stukken: De ogen van een hond, de plooien in een jurk of complexe textuur. Deze delen zijn nog niet klaar. De meester blijft hier alleen op werken om de details perfect te maken.

Het is alsof je een schilderij hebt met een masker. De meester werkt alleen door de gaten in het masker (de moeilijke delen), terwijl de leerling het hele doek invult. Omdat de meester maar op een klein stukje hoeft te werken, gaat hij veel sneller.

3. De "Geheugensteun" (Het Geheim)

Er is een probleem: als de meester alleen naar een klein stukje kijkt, vergeet hij misschien hoe de rest van het plaatje eruitzag. Om dit op te lossen, gebruikt HybridStitch een trucje uit de geheugenwereld (KV-cache). Het is alsof de meester een foto van de vorige stap in zijn hand houdt, zodat hij de context niet verliest terwijl hij alleen aan het detail werkt. Zo blijft het plaatje consistent.

4. De "Overname" (Het Einde)

Naarmate het plaatje steeds scherper wordt, worden de verschillen tussen wat de meester en de leerling doen steeds kleiner. Uiteindelijk zijn de moeilijke stukken ook klaar. Dan laat je de leerling het hele plaatje afmaken. De meester mag gaan rusten.

Waarom is dit zo cool?

• Snelheid: Omdat de zware meester maar een klein deel van het plaatje hoeft te doen, is het proces 1,83 keer sneller dan wanneer je alleen de meester gebruikt.
• Kwaliteit: Het plaatje ziet er nog steeds prachtig uit, omdat de meester precies op de plekken werkt waar het moeilijk is.
• Slimmer dan de rest: Andere methoden wachten tot het hele plaatje klaar is om te switchen. HybridStitch schakelt per pixel om. Het is alsof je niet wacht tot de hele auto is gebouwd om de wielen te draaien, maar je draait de wielen alvast terwijl de carrosserie nog wordt gelast.

Kortom: HybridStitch is als een super-efficiënte chef-kok die zijn sous-chef (de leerling) laat koken voor de simpele groenten, terwijl hij zelf alleen de lastige sauzen en garnituren maakt. Het diner is net zo lekker, maar het staat veel sneller op tafel!

Technische samenvatting

Probleemstelling

Diffusiemodellen voor tekst-naar-afbeelding (T2I) generatie hebben een opmerkelijke kwaliteit bereikt, maar lijden onder zware rekenkosten, vooral bij grote modellen met tientallen miljarden parameters (zoals Stable Diffusion 3.5). Bestaande versnellingsmethoden, zoals het vervangen van een deel van de denoising-stappen door een kleiner model (Mixture of Models), werken vaak op het niveau van het hele beeld. Deze "naieve" aanpak negeert de heterogene rekenbehoeften binnen één tijdstap: sommige delen van een afbeelding (bijv. de achtergrond) zijn makkelijker te renderen en kunnen eerder worden overgedragen aan een kleiner model, terwijl complexe delen (bijv. objecten) de grote modellen nodig hebben voor verfijning. Het switchen op beeldniveau leidt tot een suboptimale balans tussen kwaliteit en snelheid: te vroeg switchen verlaagt de kwaliteit, te laat switchen verhoogt de latentie.

Methodologie: HybridStitch

HybridStitch introduceert een nieuw paradigma dat generatie behandelt als een bewerkingsproces, waarbij modelswitching plaatsvindt op pixel- en tijdstap-niveau. De methode is trainingsvrij en bestaat uit drie fasen:

1. Fase 1 (Initiële Opbouw): Een groot model verwerkt de volledige Gaussian noise om de lay-out van de afbeelding te construeren.
2. Fase 2 (Hybride Verwerking): Dit is de kerninnovatie.
   • Een klein model verwerkt de volledige latent-ruimte om een ruwe schets (draft) van de huidige stap te genereren.
   • Een groot model verwerkt alleen een geselecteerd masker van de afbeelding (de moeilijkste regio's) om deze gebieden te verfijnen.
   • KV-Cache Mechanisme: Omdat het grote model slechts een subset van tokens verwerkt, mist het de volledige context voor de attention-mechanismen. HybridStitch lost dit op door de Key-Value (KV) cache van de vorige stap te hergebruiken en te combineren met de huidige masked tokens. Dit zorgt voor contextuele consistentie zonder de volledige berekening van het grote model.
   • De output van het grote model wordt gebruikt om de corresponderende regio's in de output van het kleine model te vervangen.
3. Fase 3 (Volledig Klein Model): Zodra het verschil tussen de output van het grote en kleine model onder een bepaalde drempelwaarde zakt, schakelt het systeem volledig over naar het kleine model voor de resterende stappen.

Switch-strategie:
Het systeem gebruikt de $L_1$-afstand tussen de latent-outputs van opeenvolgende stappen ($D_t$) om te bepalen wanneer een regio "stabiel" genoeg is om over te schakelen. Het masker wordt dynamisch bijgewerkt bij elke stap door de top-K pixels met het grootste verschil te selecteren.

Belangrijkste Bijdragen

• Pixel-niveau Awareness: In plaats van te switchen op beeldniveau, identificeert HybridStitch specifieke pixels die complex zijn en vereisen dat het grote model langer actief blijft.
• Hybride Architectuur met KV-Cache: Een nieuwe techniek om het grote model efficiënt te laten werken op een masker door de context te vullen met cache-data uit de vorige stap, wat de kwaliteit behoudt terwijl de rekenlast wordt verlaagd.
• Trainingsvrij: De methode vereist geen extra training van de modellen, maar past alleen de inferentie-architectuur aan.

Resultaten

De evaluatie is uitgevoerd op Stable Diffusion 3 (Large) en Stable Diffusion 3 (Medium) met de COCO-dataset.

• Snelheid: HybridStitch bereikt een 1.83x snelheidswinst ten opzichte van het gebruik van alleen het grote model. Dit is sneller dan bestaande methoden zoals T-Stitch (1.41x) en SRDiffusion (1.55x).
• Kwaliteit: De methode behoudt de beeldkwaliteit uitstekend.
  • FID (Fréchet Inception Distance): HybridStitch (bijv. 30% masker) scoort 30.43, wat beter is dan SRDiffusion (31.67) en aanzienlijk beter dan T-Stitch (31.87).
  • LPIPS: De perceptuele gelijkenis met het originele grote model is hoog (0.42 voor 30% masker), wat aangeeft dat de structuur en details goed behouden blijven.
• Efficiëntie: Er is een latency-reductie van ongeveer 18% ten opzichte van SRDiffusion bij vergelijkbare of betere kwaliteit.
• Hardware-onafhankelijkheid: De snelheidswinst blijft consistent op verschillende GPU's (RTX6000, A100, H100), hoewel de absolute winst iets afneemt op zeer krachtige hardware waar de rekenlast minder dominant is.

Betekenis

HybridStitch biedt een fundamenteel nieuwe kijk op de versnelling van diffusiemodellen. Het erkent dat niet alle pixels in een afbeelding evenveel rekenkracht nodig hebben op elk moment in het generatieproces. Door de "mixture of models" aanpak te verfijnen tot op pixel-niveau en contextproblemen op te lossen via KV-caching, overbrugt het de kloof tussen hoge kwaliteit en lage latentie. Dit maakt het mogelijk om state-of-the-art generatieve modellen efficiënter in te zetten voor real-time toepassingen en op minder krachtige hardware, zonder in te leveren op de visuele output.