When Hidden States Drift: Can KV Caches Rescue Long-Range Speculative Decoding?

1. Probleemstelling

Verval op lange afstand bij Speculatieve Decoding:
Speculatieve decoding versnelt de inferentie van Large Language Models (LLM) door een lichtgewicht "drafter"-model te gebruiken om meerdere tokens voor te stellen, die vervolgens worden geverifieerd door een groter "doel"-model. Huidige state-of-the-art drafters (zoals EAGLE-3, MTP) vertrouwen op hergebruik van verborgen toestanden (hidden-state reuse), waarbij de interne verborgen toestanden van het doelmodel worden ingevoerd in de drafter.

• Het Probleem: Deze systemen lijden onder verval op lange afstand: de nauwkeurigheid van het voorstel neemt progressief af naarmate de speculatieve stap $k$ toeneemt.
• Bestaande Uitleg: De gemeenschap schrijft dit toe aan train-inference mismatch (de drafter baseert zich tijdens inferentie op zijn eigen afwijkende verborgen toestanden, in tegenstelling tot tijdens training). Hoewel Test-Time Training (TTT) dit mitigeert, blijft het verval bestaan.
• De Kloof: Het paper betoogt dat train-inference mismatch niet de enige oorzaak is. Het fundamentele probleem ligt in behoud van contextinformatie. Verborgen toestanden fungeren als een "bevooroordeelde contextcompressie" die is geoptimaliseerd voor de onmiddellijke volgende token, waardoor informatie die relevant is voor toekomstige speculatieve stappen wordt weggegooid.

2. Kernhypothese en Theoretische Analyse

De auteurs stellen de KV-Reuse Hypothese voor: het hergebruiken van de KV Cache (Key-Value paren) van het doelmodel in plaats van alleen verborgen toestanden kan rijkere context behouden voor drafting op lange afstand.

• Hergebruik van Verborgen Toestanden (Bevooroordeelde Compressie):
  • Doel-verborgen toestanden zijn het resultaat van attentie-aggregatie: $h_t = \sum \alpha_i v_i$.
  • De gewichten $\alpha_i$ zijn geoptimaliseerd voor het voorspellen van de huidige token $x_{t+1}$.
  • Tokens die zwak relevant zijn voor $x_{t+1}$ krijgen bijna nul-gewichten en worden effectief weggegooid in $h_t$.
  • Gevolg: Voor latere speculatieve stappen ($x_{t+k}$) moet de drafter informatie herstellen die onderdrukt werd door het attentiemechanisme van het doelmodel; dit is een slecht gestelde "informatieherstel"-probleem.
• KV-Hergebruik (Expliciete Context):
  • De KV-cache behoudt de volledige set $(k_i, v_i)$-paren vóór aggregatie.
  • De drafter kan her-attentie uitvoeren met zijn eigen geschatte toekomstige queries tegen de volledige prefix KV-cache.
  • Gevolg: Er wordt geen informatie weggegooid voordat de drafter deze verwerkt. De moeilijkheid verschuift van "herstellen van verloren info" naar "nauwkeurig schatten van toekomstige queries" (een functie-approximatieprobleem).

3. Methodologie: Het KVShot Framework

Om deze hypothese te testen, introduceren de auteurs KVShot, een diagnostisch framework dat drie hergebruikspatronen vergelijkt onder dezelfde Autoregressive Test-Time Training (TTT) pipeline (gebruikt door EAGLE-3):

1. Alleen-Verborgen-Hergebruik: Standaard EAGLE-3/MTP-aanpak (concatenatie van doel-verborgen toestanden).
2. Alleen-KV-Hergebruik: De drafter ontvangt de doel-KV-cache via cross-attention, waarbij verborgen toestanden volledig worden vervangen.
3. Hybride Hergebruik: Een gated fusion-mechanisme waarbij de verborgen toestand de hoofdanker biedt, en de KV-cache een "gated delta correctie" biedt voor stappen op lange afstand.

Experimentele Opstelling:

• Doelmodel: Qwen3-8B.
• Datasets: ShareGPT en UltraChat (geschaald van 70k naar 280k samples).
• Maten: Stapsgewijze acceptatiepercentages ($\alpha_k$), Gemiddeld Geaccepteerde Tokens (MAT), en End-to-End (E2E) snelheidsoptimalisatie-proxies.

4. Belangrijkste Resultaten

A. Validatie van Voorspellingen (Stapsgewijze Analyse)

De experimenten bevestigden drie theoretische voorspellingen:

1. Voordeel op Lange Afstand: Alleen-KV-hergebruik degradeert minder snel dan Alleen-Verborgen-hergebruik naarmate $k$ toeneemt. Bij stap $k=6$ overtrof een 4-laags KV-only drafter de EAGLE-3 baseline ($\alpha_6 = 0.495$ versus $0.469$).
2. Query-schatting Bottleneck: De prestaties van Alleen-KV zijn zwaar afhankelijk van de diepte van de drafter. Een 1-laags drafter presteerde slecht (MAT 1.84) omdat het complexe doel-queries niet kan benaderen. Het verhogen van de diepte naar 2 lagen leverde een enorme sprong op (MAT 2.23), maar 4 lagen boden afnemende meerwaarde en bleven achter bij de 1-laags EAGLE-3 baseline (MAT 2.37) in totale efficiëntie.
3. Nadeel op Korte Afstand: Bij $k=0$ of $1$ presteert hergebruik van Verborgen Toestanden beter dan Alleen-KV-hergebruik, omdat verborgen toestanden rijkere semantische inhoud dragen dan ruwe embeddings die worden gebruikt voor KV her-attentie.

B. Hybride Prestaties

De Hybride Drafter (die beide signalen combineert) behaalde de beste stapsgewijze resultaten:

• MAT: 2.54 (tegenover 2.37 voor de EAGLE-3 baseline).
• Behoud: Verbetering van de behoudsratio op lange afstand van 73,5% naar 77,3%.
• Mechanisme: De gated delta-regel liet de KV-pas succesvol toe om het pad van verborgen toestanden te corrigeren op langere horizonnen, zonder de nauwkeurigheid op korte afstand te verstoren.

C. De End-to-End Realiteitscheck

Ondanks stapsgewijze verbeteringen was de End-to-End (E2E) snelheidswinst marginaal:

• HF-Gemeten MAT: Steeg slechts van 5.01 (EAGLE-3) naar 5.04 (Hybride).
• Latentie: De extra cross-attention-pas introduceerde 5–10% extra drafting-latentie.
• Conclusie: De huidige autoregressive TTT-pipeline kan het KV-signaal niet efficiënt benutten om de toegevoegde rekentkosten te overwinnen.

5. Geïdentificeerde Structurele Bottlenecks

Het paper identificeert drie specifieke redenen waarom de huidige pipeline stapsgewijze winsten niet kan omzetten in E2E-snelheidswinsten:

1. Moeilijkheid van Query-schatting: Het schatten van toekomstige doel-queries vereist het repliceren van de multi-layer niet-lineaire transformaties van het doelmodel. Ondiepe drafters (1-2 lagen) ontberen de capaciteit om deze queries nauwkeurig te benaderen.
2. Sparce Gradiëntsignalen: In autoregressive TTT worden de draft-zijde KV-projecties ($W_K, W_V$) alleen bijgewerkt door de weinige door de drafter gegenereerde tokens (bijv. 7 tokens) per stap, terwijl de prefix KV wordt gekopieerd. Dit resulteert in extreem spaarzame en onbalans gradiënten voor het KV-pad.
3. Gate-Geïnduceerde Gradiënt-Verhongering: In het hybride model neigt de gate die de KV-correctie regelt, vroeg in de training naar bijna nul in te storten (omdat de willekeurig geïnitieerde KV-tak begint als ruis). Zodra deze gesloten is, blokkeert de gate gradiënten naar de KV-tak, waardoor deze niet kan leren nuttig te zijn. Dit creëert een "zelfversterkende laag-injectie staat".

6. Betekenis en Toekomstige Richtingen

• Conceptuele Verschuiving: Het paper herformuleert verval op lange afstand niet alleen als een train-inference mismatch, maar als een informatiebehoud-probleem. Het bewijst dat verborgen toestanden een verliesrijke compressie zijn voor taken op lange horizon.
• Diagnostisch Hulpmiddel: KVShot biedt een fundamentele testomgeving om representatie-effecten te isoleren van trainingsartefacten.
• Routekaart voor Architecturen van de Volgende Generatie: De bevindingen suggereren dat Autoregressive TTT ontoereikend is voor KV-bewuste decoding. Om het potentieel van KV-hergebruik te realiseren, moet het veld overstappen op Block-wise Training Paradigma's (bijv. DFlash).
  • Block-wise training genereert veel draft-tokens parallel, waardoor dichte gradiëntsignalen voor KV-projecties worden geboden.
  • Het maakt diepere draft-modellen mogelijk die nauwkeurige query-schatting kunnen uitvoeren zonder sequentiële latentieboetes.

Samenvatting: Hoewel KV-caches theoretisch een superieur signaal bieden voor speculatieve decoding op lange afstand door context te behouden, falen huidige autoregressive trainingspipelines om dit te benutten vanwege beperkingen in query-schatting en gradiënt-verhongering. De weg vooruit ligt in block-wise, parallelle trainingsarchitecturen.