SSD: Spatially Speculative Decoding Accelerates Autoregressive Image Generation

Het artikel introduceert Spatially Speculative Decoding (SSD), een framework dat de inherente 2D ruimtelijke correlaties van afbeeldingen benut om gelijktijdig meerdere tokens te voorspellen, waardoor autoregressieve beeldgeneratie met wel 13,3x wordt versneld terwijl de hoge getrouwheid behouden blijft.

De kern

Stel je voor dat je een enorme, hoogresolutie muurschildering probeert te schilderen, maar je bent gedwongen een zeer strikte, ouderwetse regel te volgen: je mag slechts één enkele stip tegelijk schilderen, en je moet bewegen in een perfect "slangpatroon". Je voltooit de eerste rij van links naar rechts, dan spring je naar het einde van de tweede rij, ga je naar links, spring je naar het einde van de derde rij, enzovoort.

Dit is hoe huidige AI-beeldgeneratoren (genaamd "autoregressieve modellen") werken. Ze behandelen een afbeelding als een lange, platte lijst met woorden. Hoewel een afbeelding van nature een 2D-raster is (op/neer, links/rechts), maken ze een afbeelding plat tot een 1D-lijn. Dit maakt het proces ongelooflijk traag omdat de AI zijn hele "brein" (miljarden parameters) moet laden om slechts de kleur van de volgende enkele stip te bepalen, keer op keer. Het is alsof je naar de supermarkt loopt om één appel te kopen, dan weer naar huis loopt, en dan weer naar de winkel loopt voor de volgende appel.

Het Probleem: De "Memory Wall"

Het paper noemt dit de "Memory Wall". De AI besteedt de meeste tijd aan het laden van zijn brein in het geheugen om een minuscule beslissing te nemen, in plaats van daadwerkelijk na te denken. Omdat de AI dit duizenden keren moet doen voor één afbeelding, duurt het genereren van een plaatje heel lang.

De Oplossing: SSD (Spatially Speculative Decoding)

De auteurs introduceren een nieuwe methie genaamd SSD. Ze realiseerden zich dat afbeeldingen helemaal geen 1D-lijsten zijn; het zijn 2D-rasters. Als je weet hoe een stip eruitziet, kun je vaak ook wel raden hoe de stip direct onder die stip eruit zal zien, net zo gemakkelijk als je kunt raden hoe de stip aan de rechterkant eruitziet.

Hier is hoe SSD het spel verandert, met behulp van een paar analogieën:

1. De "Raadspel"-analogie

• Oude manier (1D): De AI raadt de volgende stip, controleert of deze klopt, en raadt dan de volgende stip. Het is een trage estafette stap voor stap.
• SSD-manier (2D): De AI werkt als een team van ravers. Terwijl één persoon de volgende stip naar het rechteren raadt, raadt een ander tegelijkertijd de stip direct eronder. Ze wachten niet tot de eerste gok klaar is voordat ze met de tweede beginnen. Ze raden een heel blok aan stippen tegelijk.

2. De "Concept"-analogie
Denk aan de AI als een schrijver.

• Standaard AI: Schrijft één woord, stopt, controleert een woordenboek, schrijft het volgende woord.
• SSD: Schrijft in één keer een hele zin (of zelfs een alinea) als een "concept". Daarna leest hij snel door dat concept heen om te zien of het logisch is. Als een woord er net niet helemaal bij zit, past hij alleen dat specifieke woord aan zonder de hele alinea opnieuw te schrijven.

3. De "Autocorrectie"-twist
Het paper vermeldt een slimme truc. Normaal gesproken, als een AI een blok tokens (stippen) raadt en er één is fout, gooit hij het hele blok weg en begint hij opnieuw. SSD is slimmer. Het behandelt de "foute" gokken als ruwe concepten. Het voert een snelle controle uit (verificatie) en herstelt de specifieke fouten ter plekke zonder het hele blok weg te gooien. Het is als een spellingscontrole die typefouten direct herstelt in plaats van je te laten de hele pagina opnieuw schrijven.

De Resultaten: Het Versnellen van de Muurschildering

Het paper heeft deze methode getest op drie krachtige AI-modellen. De resultaten waren spectaculair:

• Snelheid: Ze maakten de beeldgeneratie tot wel 13 keer sneller.
  • Voorbeeld: Een model dat 339 seconden (bijna 6 minuten) nodig had om een afbeelding te maken, doet er nu slechts 25 seconden over.
• Kwaliteit: Ondanks dat het veel sneller is, zien de afbeeldingen er net zo goed uit als de trage versies. De "gokken" waren nauwkeurig genoeg zodat de uiteindelijke afbeelding geen detail verloor.
• Plug-and-Play: Deze methode vereist niet het herbouwen van de AI-hersenen. Het is alsoals het toevoegen van een turbocharger aan een bestaande automotor. Je kunt het aanzetten wanneer je snelheid wilt, of je kunt het uitzetten, en de auto rijdt dan precies zoals voorheen.

Samenvatting

Het paper betoogt dat door de natuurlijke 2D-vorm van afbeeldingen (op/neer en links/rechts) te respecteren in plaats van ze in een 1D-lijn te dwingen, we de "Memory Wall" kunnen doorbreken. Door meerdere stippen tegelijk te raden en kleine fouten direct te herstellen, verandert SSD een traag, stapsgewijs proces in een snel, parallel proces, waardoor hoogwaardige AI-kunstgeneratie bijna instant gebeurt.

Technische samenvatting

Technische Samenvatting: Spatially Speculative Decoding (SSD)

1. Probleemstelling

Autoregressieve modellen hebben aanzienlijk succes bereikt in visuele generatie door afbeeldingen te behandelen als 1D-sequenties van discrete tokens, vergelijkbaar met taalmodellering. Deze aanpak legt echter een geometrische compromis op: het vlakt de intrinsieke 2D-spatiële lokaliteit van visuele signalen af naar een lineaire raster-scan sequentie.

Dit afvlakken creëert een ernstige computationele bottleneck tijdens inferentie, vaak de "memory wall" genoemd. Het genereren van een $n \times n$ afbeelding vereist $n^2$ sequentiële forward passes. Omdat elke pass de volledige set transformer-parameters opnieuw moet laden om een enkel token te voorspellen, is het proces zwaar beperkt door geheugenbandbreedte in plaats van rekenkracht.

Bestaande versnellings-technieken, zoals standaard Speculative Decoding en Jacobi-gebaseerde iteratie, proberen NLP-methoden naar visie te importeren. Deze methoden blijven echter gebonden aan de 1D-sequentiële assumptie. Omdat visuele patches een hoge lokale entropie vertonen zonder strikte grammatica, leidt het anticiperen van meerdere tokens langs een afgeplatte 1D-horizon tot lage draft-acceptatiepercentages. Gevolgoenteic bieden deze methoden slechts bescheiden versnellingen (1,8× tot 3,7×) of verslechteren ze de beeldkwaliteit door onnatuurlijke onafhankelijkheidsveronderstellingen op te leggen.

2. Methodologie: Spatially Speculative Decoding (SSD)

De auteurs stellen Spatially Speculative Decoding (SSD) voor, een framework dat het voorspellende doel afstemt op de natuurlijke 2D-geometrie van afbeeldingen. In plaats van alleen het onmiddellijk volgende token in een 1D-sequentie te voorspellen, voorspelt SSD gelijktijdig het aangrenzende horizontale token en het token direct daaronder.

Kernmechanismen

1. 2D Spatiële Anticipatie:

   • SSD factoriseert 2D-anticipatie in twee orthogonale 1D-voorspellingsstromen: horizontaal (langs de raster scan) en verticaal (langs de kolommen).
   • Het model gebruikt lichtgewicht heads om horizontaal een volledige rij te draften, en draft vervolgens opeenvolgende rijen verticaal in parallel.
   • Deze aanpak maakt gebruik van de observatie dat de afhankelijkheid van een token van het token direct erboven even sterk is als de afhankelijkheid van het token direct links, ondanks de grote offset in de afgeplatte sequentie.
   • Complexiteitsreductie: Door volledige spatiële blokken in parallel te draften, reduceert SSD de theoretische inferentiecomplexiteit van $O(n^2)$ naar $O(n)$ voor een $n \times n$ afbeelding.

2. Continue Latente Ruimte Voorspelling:

   • Het direct voorspellen van discrete token ID's is moeilijk vanwege vlakke waarschijnlijkheidsverdelingen over grote codebooks, wat leidt tot lage acceptatiepercentages (<5%).
   • SSD voorspelt in plaats daarvan de continue latente kenmerken van de laatste transformer-laag (specifiek de verborgen toestand vóór de uiteindelijke RMSNorm).
   • Een lichtgewicht predictor $f_\phi$ neemt de verborgen toestand en token embedding als input om de verborgen toestand op een specifieke spatiële offset te voorspellen.
   • De voorspelde verborgen toestanden worden vervolgens gedecodeerd via de bestaande outputlagen van het basismodel om token-kandidaten te verkrijgen.

3. Verificatie als Auto-Correctie:

   • In tegenstelling tot standaard speculative decoding, waarbij een heel blok wordt afgewezen bij de eerste mismatch, behandelt SSD verificatie als een zelfcorrigerend mechanisme.
   • Omdat drafts coherente spatiële blokken vormen, kunnen kleine representatieve mismatches lokaal worden opgelost.
   • Afgewezen tokens worden niet weggegooid; in plaats daarvan samplet het systeem uit een residuele distributie om ze binnen dezelfde forward pass te herstellen.
   • Dit proces wordt herhaald voor $r$ rondes, waarbij afgewezen posities worden bijgewerkt en de KV-cache wordt teruggedraaid voor her-evaluatie, wat slechts $r+1$ forward passes per blok vereist in plaats van sequentiële regeneratie.

4. Modulariteit:

   • SSD is een plug-and-play module. Het vereist geen wijzigingen aan de getrainde backbone.
   • Alleen lichtgewicht drafting heads worden getraind via zelf-distillatie, waardoor de FLOPs-overhead minimaal blijft.

3. Belangrijkste Bijdragen

• Geometrische Heruitlijning: Het paper demonstreert dat het respecteren van de onderliggende 2D-geometrie enorme computationele efficiënties ontsluit, voorbij de "computationeel dure artefact" van 1D sequentiële decoding.
• Latente Ruimte Drafting: De introductie van continue latente kenmerk-voorspelling verbetert de draft-nauwkeurigheid aanzienlijk vergeleken met discrete token-voorspelling in visuele domeinen.
• Auto-Corrigerende Verificatie: Een nieuwe verificatiestrategie die afgewezen tokens in parallel herstelt in plaats van ze weg te gooien, waardoor de voordelen van spatiële drafting behouden blijven.
• Plug-and-Play Versnelling: Een framework dat generatie versnelt zonder de enorme backbone-modellen opnieuw te trainen, toepasbaar op elk verenigd autoregressief model dat discrete visuele tokens produceert.

4. Experimentele Resultaten

De auteurs evalueerden SSD op drie state-of-the-art autoregressieve modellen: Janus-Pro-7B, Lumina-mGPT-7B, en Emu3-8B, met behulp van de benchmarks DPG-Bench en GenEval.

• Versnelling: SSD bereikt tot wel 13,3× wall-clock versnelling vergeleken met standaard autoregressieve inferentie.
  • Emu3-8B: 339s $\to$ 25,5s (13,27× versnelling).
  • Lumina-mGPT-7B: 91,6s $\to$ 7,5s (12,19× versnelling).
  • Janus-Pro-7B: 7,8s $\to$ 1,4s (5,74× versnelling).
• Kwaliteit: De methode behoudt een hoge getrouwheid, waarbij de generatiekwaliteit-metrieken op DPG-Bench en GenEval vergelijkbaar blijven met de baseline autoregressieve modellen en de 1D-MTP baselines (die leden onder ernstige kwaliteitsvermindering) aanzienlijk overtreffen.
• Vergelijking:
  • 1D-MTP: Bereikte ~2,0–2,4× versnelling maar veroorzaakte significante kwaliteitsdalingen (bijv. de algemene score van Emu3 daalde van 78,69 naar 53,11).
  • SJD (Jacobi-gebaseerd): Behield kwaliteit maar was beperkt tot 1,5–2,9× versnelling.
  • SSD: Bereikte de hoogste versnelling terwijl de kwaliteit behouden bleef.

5. Betekenis en Claims

Het paper beweert dat de inefficiëntie van huidige autoregressieve visuele generatie niet een inherente beperking van de modellen is, maar een resultaat van de 1D-sequentiële assumptie die geërfd is van taalmodellering. Door deze assumptie los te laten en het decodingproces af te stemmen op de intrinsieke 2D-geometrie van afbeeldingen, overwint SSD de memory wall.

De auteurs stellen dat deze geometrische verschuiving de weg vrijmaakt voor real-time, hoog-resolutie autoregressieve generatieve modellen. De resultaten suggereren dat het respecteren van visuele lokaliteit enorme computationele besparingen mogelijk maakt zonder de unificatie van visie en taal in een enkele backbone in gevaar te brengen. De methode wordt gepresenteerd als een algemene oplossing die kan worden toegepast op bestaande verenigde modellen om ze competitief te maken qua inferentiesnelheid.