Technische Samenvatting: xPress: Parallel Verfijning voor Diffusion Drafters in Speculative Decoding

1. Probleemstelling

Speculative decoding (SD) versnelt autoregressieve generatie door een lichtgewicht drafter te gebruiken om tokens voor te stellen voor verificatie door een groter doelmodel. Hoewel autoregressieve (AR) drafters (bijv. EAGLE) effectief zijn, lijden ze onder sequentiële overhead bij het genereren van lange blokken. Block-diffusion drafters (bijv. dFlash) pakken dit aan door een heel blok aan tokens in één enkele parallelle forward pass te genereren, wat de drafting-latentie drastisch vermindert.

Echter, diffusion drafters bezitten een fundamentele nauwkeurigheidsbeperking: ze voorspellen tokens op basis van marginale distributies in plaats van joint distributies. In een single-pass diffusieproces wordt token $k$ gesampled zonder te conditioneren op de gerealiseerde waarden van tokens $1$ tot en met $k-1$. Dit gebrek aan intra-blok causaliteit leidt tot sequenties waarbij individuele tokens lokaal waarschijnlijk zijn, maar gezamenlijk incoherent (bijv. subject-werkwoord onovereenstemming). Tijdens de verificatie wijst het doelmodel deze incoherente sequenties vroegtijdig af, wat de acceptatie-lengte ($\tau$) ernstig beperkt en het potentieel voor snelheidswinst door parallel drafting tenietdoet.

Bestaande oplossingen, zoals tree-gebaseerd drafting (bijv. PRESTO, DDTree), proberen dit te verzachten door de kandidaatruimte uit te breiden naar een boomstructuur. Deze methoden zijn echter computationeel duur, complex in gebruik en proberen het probleem te omzeilen door te richten op recall in plaats van de onderliggende causale incoherentie te herstellen.

2. Methodologie: xPress

De auteurs stellen xPress voor, een lichtgewicht causale verfijner ontworpen om de ontbrekende causale afhankelijkheden in diffusion drafters te herstellen zonder terug te vallen op een trage, token-voor-token autoregressieve loop.

2.1 Kernconcept

xPress werkt als een "causale verfijner" die het initiële blok-voorstel van de diffusion drafter neemt en het harmoniseert via parallelle verfijning. In plaats van tokens sequentieel te regenereren, injecteert xPress causale informatie in de logits van de diffusion drafter om de marginale distributies te corrigeren.

2.2 Architectuur

De verfijner is een lichtgewicht module die een geleerde logit bias ($\delta_k$) toevoegt aan de basis-logits ($s_k$) van de drafter op elke positie $k$. De architectuur bestaat uit vier stadia:

1. Fuse: Projecteert drie inputs naar een low-rank ruimte ($r=256$):
   • De per-positie hidden state ($h_k$) van de drafter.
   • Een block-globale samenvatting ($g$) verkregen via mean-pooling.
   • De embedding van de vorige token-gok ($x_{k-1}$).
2. Mix: Past een strikt lager-triangulaire lineaire mixer ($L$) toe over de blokposities. Dit stelt positie $k$ in staat om informatie te aggregeren van alle voorgangers ($j \le k$) in een enkele operatie, waardoor een causale structuur wordt afgedwongen zonder sequentiële berekening.
3. Trans: Een residuele MLP in de low-rank ruimte.
4. Read: Projecteert het resultaat terug naar de vocabulaire-omvang om de bias $\delta_k$ te produceren.

De uiteindelijke verfijnde logit is $\ell_k = s_k + \delta_k$.

2.3 Parallelle Verfijning via Jacobi Decoding

Een rechtstreekse causale verfijner zou een sequentiële links-naar-rechts pass vereisen, wat de latentie herintroduceert die xPress juist wil vermijden. Om dit op te lossen, gebruikt xPress Jacobi decoding:

• Initialisatie: Het blok wordt geïnitialiseerd met de initiële greedy voorspellingen van de drafter.
• Iteratie: Alle posities worden gelijktijdig in parallel bijgewerkt, waarbij de huidige staat van het blok dient als de prefix voor de volgende iteratie.
• Convergentie: Het proces herhaalt zich voor $K$ iteraties. Theoretisch wordt het sequentiële vaste punt voor een blok van grootte $B$ bereikt in maximaal $B$ iteraties. Empirisch stellen de auteurs vast dat $K \approx 6$ iteraties voldoende zijn om de acceptatie-lengte van een volledige sequentiële decode te evenaren, aange het seed van de drafter al een hoogwaardige gok is.

2.4 Trainingsstrategie

xPress wordt samen met de diffusion drafter getraind met een samengestelde loss-functie:

1. Cross-Entropy (Teacher Forcing): Maximaliseert de waarschijnlijkheid van de grondwaarheid-token gegeven de grondwaarheid-prefix.
2. Total Variation (TV) Loss: Minimaliseert de afstand tussen de verfijnde distributie en de distributie van het doelmodel, wat direct optimaliseert voor de acceptatiesnelheid.
3. Consistency Loss: Een hulp-loss toegepast op een tweede forward pass waarbij de inputs de eigen voorspellingen van de drafter zijn (self-conditioned). Dit stemt de trainingsdistributie af op het Jacobi-proces tijdens de inferentie, om te voorkomen dat de verfijner overfit op gold prefixes.
4. Drafter Anchor: Een loss-term toegepast op de output van de basis-drafter om te garanderen dat de onderliggende diffusion model van hoge kwaliteit blijft.

3. Belangrijkste Bijdragen

• Injectie van Causale Afhankelijkheid: xPress injecteert causale informatie in diffusion drafters met een minimaal mechanisme (low-rank projectie en een triangulaire mixer) dat opereert op de hidden states van de drafter, zonder de noodzaak voor grote parameterverhogingen.
• Parallelle Verfijning: Door Jacobi decoding te gebruiken, lost xPress causale afhankelijkheden in parallel op. Het bereikt het effect van $B$ afhankelijke autoregressieve stappen met slechts $K \ll B$ parallelle iteraties, waardoor de snelheidvoordelen van parallel drafting behouden blijven.
• Consistente End-to-End Winsten: De methode verbetert zowel de acceptatie-lengte als de throughput aanzienlijk zonder het verificatieproces te wijzigen of een grotere drafter te vereisen.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op Qwen3-8B met een block-16 dFlash drafter over zeven benchmarks (Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP, LiveCodeBench; Chat: MT-Bench).

• Acceptatie-lengte: xPress verhoogt de gemiddelde acceptatie-lengte met ~30% vergeleken met de originele dFlash diffusion drafter.
  • Onder greedy decoding ($T=0$), variëren de winsten van 18% tot 56%.
  • Onder sampling ($T=1$), variëren de winsten van 19% tot 58%.
• Throughput Snelheidswinst:
  • Greedy ($T=0$): xPress bereikt een gemiddelde end-to-end snelheidswinst van 1.3× over de dFlash drafter, met pieken tot 1.7× (bijv. op GSM8K).
  • Sampling ($T=1$): De snelheidswinst verbetert naar een gemiddelde van 1.46×, met pieken tot 1.6×.
• Vergelijking met Sequentiële Verfijners: Wanneer vergeleken met de Markov head (een sequentiële causale verfijner uit DSpark), evenaart xPress de acceptatie-lengte of overtreft het deze, terwijl het aanzienlijk sneller draait.
  • Bij $K=4$ Jacobi iteraties, evenaart xPress de nauwkeurigheid van de Markov head, maar draait het 1.6× sneller (379 µs vs. 601 µs per blok).
  • De Markov head vereist een vaste 15-staps seriële decode, terwijl xPress convergeert in minder parallelle stappen.
• Schaalbaarheid: In scenario's met grote batches (vLLM, batch sizes 16 en 32), behoudt xPress zijn voordeel. Hoewel de absolute snelheidswinst over de AR baseline krimpt naarmate batches groeien (door compute-bound verificatie), presteert xPress consistent beter dan de plain dFlash drafter met 1.2–1.4× in throughput, wat bewijst dat de voordelen van de verfijner standhouden onder belasting.

5. Betekenis en Claims

Het artikel claimt dat xPress de "inherente mismatch" tussen parallelle diffusion drafting en sequentiële verificatie aanpakt door de marginale distributie zelf te repareren in plaats van de zoekruimte uit te breiden.

• Efficiëntie: xPress laat zien dat causale informatie kan worden geïnjecteerd tegen bijna nul kosten relatief aan de drafting overhead. De verfijningskosten zijn een klein deel van de totale drafting tijd, en de parallelle aard van de verfijning voorkeert de latentie-straf van sequentiële correctie.
• Orthogonaliteit: De methode is orthogonaal aan andere verbeteringen; het kan worden gecombineerd met grotere drafters of tree-gebaseerde verificatiemethoden, omdat het de "backbone" logits verbetert waar deze methoden op vertrouwen.
• Generaliseerbaarheid: De auteurs suggereren dat de filosofie van het harmoniseren van wederzijds afhankelijke voorspellingen via een paar parallelle iteraties breed toepasbaar kan zijn buiten speculative decoding, overal waar modellen simultaan onderling afhankelijke outputs genereren.

Het artikel concludeert dat xPress succesvol de causaliteit herstelt die verloren ging bij single-pass diffusion drafting, wat langere acceptatieketens en hogere throughput mogelijk maakt zonder de parallelle efficiëntie op te offeren die diffusion drafters aantrekkelijk maakt.