SJD-VP: Speculative Jacobi Decoding with Verification Prediction for Autoregressive Image Generation

Probleemstelling

Autoregressieve beeldgeneratiemodellen (zoals Lumina-mGPT, LlamaGen) genereren beelden token voor token, wat leidt tot een trage inferentie, vooral bij hoge resoluties. Om dit op te lossen, is Speculative Jacobi Decoding (SJD) ontwikkeld. SJD probeert meerdere tokens parallel te genereren ("drafting") en deze vervolgens te verifiëren door ze iteratief te verfijnen.

Echter, bestaande SJD-methoden kampen met een kritieke beperking: een lage acceptatiegraad van de gespeculeerde tokens. De reden hiervoor is een structurele ontkoppeling tussen de drafting-stap (het genereren van tokens) en de verificatiestap (het valideren ervan).

• Bestaande methoden selecteren tokens uitsluitend op basis van de huidige kansverdeling, zonder rekening te houden met de historische dynamiek of de eisen van de verificatiestap.
• Dit leidt ertoe dat veel gedraaide tokens worden verworpen, waardoor de versnelling beperkt blijft.
• Bestaande oplossingen proberen dit op te lossen door de verificatiecriteria te versoepelen (om semantisch vergelijkbare tokens te accepteren), maar dit lost het fundamentele probleem van de misalignering tussen generatie en verificatie niet op.

Methodologie: SJD-VP

De auteurs introduceren SJD-VP (Speculative Jacobi Decoding with Verification Prediction). De kernidee is om de verificatiecriteria te gebruiken als een prior om het generatieproces te sturen, zodat de gedraaide tokens met een hogere waarschijnlijkheid worden geaccepteerd.

1. Kernobservatie (Motivatie):
Uit diepgaande analyse bleek een opvallend patroon: tokens die uiteindelijk worden geaccepteerd door de verificatiestap, vertonen bijna altijd een opeenvolgende groei in waarschijnlijkheid (consecutive probability growth) tijdens de Jacobi-iteraties voordat ze worden geverifieerd. Onjuiste tokens vertonen daarentegen vaak fluctuerende of dalende kansen. Deze groei is een sterke voorspeller voor zowel acceptatie als correctheid.

2. Het SJD-VP Framework:
Het model voegt een Bayesiaanse Fusie-stap toe aan het bestaande SJD-proces. In plaats van alleen te vertrouwen op de huidige kansen, combineert het deze met een "verificatie-prior" gebaseerd op de historische groeitrends.

• Stap 1: Schatting van de Verificatie-Prior (Verification Prior Estimation):

  • Het systeem houdt een geschiedenisbuffer bij van de kansen van tokens over de laatste iteraties.
  • Het berekent een historische referentiekans ($\bar{p}_t(x)$) via Exponentially Weighted Aggregation (EWA).
  • Er wordt een voorspellingscore ($S_t(x)$) berekend als het verschil in log-kans tussen de huidige kans en de historische kans.
  • Een Consecutive Growth Mask ($M_t(x)$) wordt toegepast om alleen tokens te selecteren die een consistente stijgende trend vertonen.

• Stap 2: Bayesiaanse Verificatie-Fusie (Bayesian Verification Fusion):

  • De huidige kansen ($p_t(x)$) worden gezien als de conditionele waarschijnlijkheid.
  • De geschatte prior ($S_t(x)$) wordt gecombineerd met de huidige kansen in de log-ruimte om een posterieure verdeling ($p'_t(x)$) te vormen:
    $$\log p'_t(x) = \log p_t(x) + M_t(x) \cdot S_t(x)$$
  • Dit zorgt ervoor dat tokens met een stijgende kans-trend een hogere prioriteit krijgen bij het selecteren van de volgende token.

• Stap 3: Sampling:

  • Er wordt een nieuwe token getrokken uit deze aangepaste posterior-verdeling.
  • Dit proces is "plug-and-play" en vereist geen hertraining van het basismodel.

Belangrijkste Bijdragen

1. Identificatie van een nieuwe prior: De auteurs tonen aan dat het negeren van de "verificatie-prior" (de trend van opeenvolgende kansgroei) een sleutelfactor is in de inefficiëntie van SJD. Ze bewijzen dat deze trend een betrouwbare indicator is voor zowel acceptatie als correctheid.
2. SJD-VP Framework: Een nieuwe, plug-and-play methode die verificatievoorspelling integreert via een Bayesiaans raamwerk. Dit lost de structurele ontkoppeling op tussen drafting en verificatie.
3. Theoretische onderbouwing: De auteurs leveren een theoretisch bewijs dat SJD-VP de Total Variation Distance (TVD) tussen de gedraaide verdeling en de verificatieverdeling verkleint, mits de voorspellingsnauwkeurigheid (Q) boven de 50% ligt (wat empirisch wordt bevestigd op >72%).
4. Complementariteit: De methode is compatibel met bestaande methoden die de verificatiecriteria versoepelen (zoals LANTERN en GSD) en kan daarop worden gebouwd voor nog betere resultaten.

Resultaten

De methode is getest op standaard benchmarks (Parti-prompts en MS-COCO 2017) met modellen zoals Lumina-mGPT, LlamaGen en Emu3.

• Snelheid: SJD-VP bereikt een aanzienlijke versnelling. Op de Parti-prompt benchmark reduceert het de latentie van 36,07s (standaard SJD) naar 27,00s (een versnelling van 2,94x t.o.v. het basismodel).
• Efficiëntie: Het aantal function evaluations (NFE) daalt drastisch (van 1035 naar 720), wat betekent dat er minder iteraties nodig zijn om een beeld te genereren.
• Kwaliteit: In tegenstelling tot sommige versnellingsmethoden die kwaliteit opofferen, verbetert SJD-VP de beeldkwaliteit lichtjes (hogere CLIP-scores en vergelijkbare of betere FID-scores).
• Vergelijking met SOTA: SJD-VP presteert consistent beter dan bestaande state-of-the-art methoden zoals Jacobi Decoding, standaard SJD, LANTERN en GSD.
• Ablatie-studies:
  • Bayesiaanse fusie is essentieel; het combineren van prior en likelihood werkt beter dan alleen de prior.
  • De methode werkt robuust over verschillende modelarchitecturen (LlamaGen, Emu3) en venstergroottes.

Significantie

Dit paper biedt een fundamenteel nieuw perspectief op het versnellen van autoregressieve beeldgeneratie. In plaats van alleen de verificatiestap te versoepelen (wat vaak leidt tot kwaliteitsverlies of complexiteit), richt SJD-VP zich op het sturen van de generatiestap zelf.

Door de "verificatie-prior" te gebruiken als een gids, sluit het generatieproces naadloos aan bij de verificatie-eisen. Dit resulteert in een win-win situatie:

1. Hogere efficiëntie: Minder iteraties en snellere inferentie.
2. Hogere kwaliteit: De geselecteerde tokens zijn niet alleen waarschijnlijker om te worden geaccepteerd, maar ook waarschijnlijker om correct te zijn, wat leidt tot scherper en consistenter beeldmateriaal.

De "plug-and-play" aard van SJD-VP maakt het direct toepasbaar op bestaande systemen, wat het een krachtige en praktische oplossing maakt voor de huidige beperkingen in autoregressieve beeldgeneratie.