Technische Samenvatting: KVpop — Key-Value Cache Compressie met Predictieve Online Pruning

1. Probleemstelling

Autoregressieve decodering in Transformer-gebaseerde Large Language Models (LLMs) vertrouwt op een Key-Value (KV) cache om de herberekening van eerdere token-representaties te vermijden. De KV-cache groeit echter lineair met de contextlengte, wat een aanzienlijke bottleneck vormt voor geheugen en bandbreedte tijdens lange-context inferentie.

Bestaande oplossingen om deze kosten te verlagen, vallen over het algemeen in twee categorieën:

• Statische/Heuristische Evictie: Methoden zoals Sliding Window of het behoud van Sink-Tokens (bijv. StreamingLLM) vertrouwen op vaste regels, terwijl score-gebaseerde methoden gebruikmaken van online proxies (bijv. cumulatieve aandacht). Deze falen vaak in het bijhouden van de nuttigheid van toekomstige tokens, wat leidt tot een broze evictie wanneer de relevantie verschuift.
• Geleerde Evictie: Benaderingen zoals Dynamic Memory Compression (DMC) of Dynamic Memory Sparsification (DMS) trainen beleid om tokens samen te voegen of te verwijderen. Veel van deze methoden nemen echter beslissingen over evictie op het moment van token-insertie of vertrouwen op differentieerbare relaxaties die niet direct de specifieke "behouden-of-verwijderen"-beslissing superviseren op basis van toekomstig nut. Bovendien vertragen methoden met beschermde vensters de verwijdering van tokens, maar bepalen ze het lot van een token vaak direct bij binnenkomst, waarbij ze de bewijslast negeren die zich verzamelt terwijl de token in het beschermde venster verblijft.

De kernuitdaging is het voorspellen welke tokens nuttig zullen zijn voor toekomstige queries, aangezien lokaal saillante tokens irrelevant kunnen worden, terwijl tokens met lage onmiddellijke aandacht cruciaal kunnen worden later (bijv. in redeneerprocessen).

2. Methodologie: KVpop

KVpop (KV compression with Predictive Online Pruning) is een sparse-attention retrofit die is ontworpen om een vast budget per KV-head af te dwingen zonder de basisarchitectuur van het model te wijzigen. Het adresseert de beperkingen van eerder werk via twee primaire mechanismen: Future-Attention Supervision en Delayed Stateful Scoring.

2.1 Architectuur en Budget

KVpop onderhoudt een begrensde cache per KV-head, gedefinieerd door het budget $B = s + w + k$:

• $s$ Sink Tokens: De eerste paar tokens worden altijd behouden.
• $w$ Protected Window: De meest recente $w$ tokens zijn beschermd tegen evictie.
• $k$ Long-Range Top-K: Oudere tokens concurreren voor een vast budget van $k$ slots op basis van geleerde belangrijksheidsscores.

2.2 Future-Attention Target Supervision

In tegenstelling tot methoden die tokens scoren op basis van het verleden van de aandacht of lokale statistieken, superviseert KVpop de scorer met de toekomstige aandachtsmassa die een token ontvangt nadat het het beschermde venster heeft verlaten.

• Definitie van de Target: Voor een token $t$ is de target de geaggregeerde log-massa van de aandacht die het ontvangt van toekomstige queries (posities $d \ge t + w$).
• Efficiënte Berekening: Om de $O(S^2)$ kosten van het materialiseren van dichte attention-kaarten tijdens de training te vermijden, gebruikt KVpop een transposed-attention pass. Door de rollen van queries en keys te verwisselen en de log-sum-exp (LSE) normalisatoren te hergebruiken die door sparse attention kernels worden geretourneerd, berekent het systeem de toekomstige aandacht-target efficiënt zonder dichte matrix-materialisatie.
• Loss Functie: Het trainingsdoel is een boundary-aware pairwise logistic loss. Het vergelijkt de voorspelde score van een nieuw in aanmerking komende token ($t_{new}$) met de score van de teacher's cutoff token ($t_{bnd}$). De loss stuurt de student aan om $t_{new}$ correct te rangschikken ten opzichte van de boundary, zodat de top-$k$ selectie overeenkomt met de future-utility ranking van de teacher.

2.3 Stateful Scoring met Delayed Decision

KVpop introduceert een delayed memory-based scorer (geïmplementeerd via een mLSTM) die de scoring-beslissing uitstelt tot de token de eviction boundary bereikt.

• Mechanisme: In plaats van een token direct bij insertie te scoren, werkt de scorer zijn interne staat bij met de representatie van de token en de daaropvolgende context. Wanneer de token het beschermde venster verlaat ($q - w$), leest de scorer de geaccumuleerde geheugenopslag om de belangrijksheidsscore te generen.
• Voordeel: Dit stelt de scorer in staat om "nabije-toekomst" context te incorporeren (bewijs dat zich verzamelt terwijl de token in het beschermde venster verblijft) voordat de keep-or-drop beslissing wordt genomen, wat de beperking van insertie-tijd scoring aanpakt.

2.4 Inferences Efficiëntie

Tijdens inferentie dwingt KVpop het vaste budget $B$ af. De selectie van de top-$k$ long-range tokens wordt geoptimaliseerd met een running threshold over statische token prioriteiten en een Fenwick-boom, waardoor de sparse mask on-the-fly gegenereerd kan worden binnen de attention kernel (bijv. FlexAttention) zonder dichte $S \times S$ masks op te slaan.

3. Belangrijkste Bijdragen

1. Supervised Eviction Policy: De auteurs formuleren fixed-budget KV-evictie als een supervised predictieprobleem bij de eviction boundary, gebruikmakend van een future-attention target om long-range top-$k$ beslissingen te trainen.
2. Efficiënte Target Berekening: Ze introduceren een training-only transposed-attention pass om toekomstige aandacht-targets te berekenen zonder dichte attention-kaarten te materialiseren, wat geen overhead tijdens inferentie oplevert.
3. Delayed Scoring: Ze demonstreren dat een stateful scorer nabije-toekomst context kan exploiteren door de scoring-beslissing uit te stellen tot de token evictable wordt, wat beter presteert dan onmiddellijke scoring-methoden.
4. Empirische Validatie: De methode wordt geëvalueerd op wiskundige redenering (AIME, HMMT) en out-of-domain taken (GPQA-D, LiveCodeBench), waarbij wordt aangetoond dat geleerde eviction-beleid hoge prestaties kan handhaven onder agressieve compressie.

4. Experimentele Resultaten

De auteurs evalueerden KVpop op Qwen3-4B en Qwen3-8B modellen, getraind op de Nemotron-Math v2 dataset met een sequentielengte van 16.384.

4.1 Wiskundige Redenering (AIME & HMMT)

• Compressieverhouding 75% (Budget 2048):
  • Qwen3-4B: KVpop behield 98% van de performance van de full-attention teacher.
  • Qwen3-8B: KVpop behield 100% (nabij volledige teacher performance).
  • Vergelijking: KVpop presteerde significant beter dan training-free baselines (StreamingLLM, TOVA) en de geleerde baseline DMS. Bijvoorbeeld, bij 75% compressie zakte StreamingLLM naar ~76% relatieve performance op Qwen3-4B, terwijl KVpop 98% behield.
• Compressieverhouding 88% (Budget 4096):
  • Qwen3-4B: KVpop behield 97% van de teacher performance.
  • Qwen3-8B: KVpop behield 100% van de teacher performance.
  • Vergelijking: Training-free methoden stortten aanzienlijk in (bijv. StreamingLLM zakte naar ~47% relatieve performance), terwijl KVpop robuust bleef.

4.2 Out-of-Domain Generalisatie

Geëvalueerd op GPQA-Diamond (wetenschappelijke redenering) en LiveCodeBench (code generatie):

• KVpop-modellen getraind alleen op wiskundige redeneerdata generaliseerden goed naar deze domeinen.
• Bij 75% compressie bereikte KVpop op Qwen3-4B een nauwkeurigheid die vergelijkbaar is met de dense teacher op GPQA-D (0.59 vs 0.59) en competitieve resultaten op LCB.

4.3 Inferences Efficiëntie

• Geheugen: In tegen tegenstelling tot dense attention, waarbij het VRAM-gebruik lineair schaalt met de generatielengte (bereikend tot 36GB bij 131k tokens), behieldden KVpop en DMS een bijna constant geheugengebruik (19GB) ongeacht de sequentielengte.
• Latency: KVpop demonstreerde een lagere end-to-end decoding latency vergeleken met DMS en dense attention bij lange generatielengtes (tot 131k tokens). De auteurs schrijven dit toe aan het vaste per-head budget van KVpop, wat de "ragged" cache-structuren van DMS vermijdt die compilatie en executie compliceren.

4.4 Ablatie Studies

• Delayed Scoring: Het verwijderen van de vertraging (direct scoren bij insertie) resulteerde in een daling van 0.2 punt in token nauwkeurigheid, wat het voordeel bevestigt van het integreren van nabije-toekomst context.
• Stateless vs. Stateful: De stateful variant (KVpop) presteerde over het algemeen gelijkwaardig aan of licht beter dan de stateless variant (KVpopmlp), vooral bij hogere compressieverhoudingen.

5. Betekenis en Claims

Het artikel claimt dat het superviseren van evictie met future-attention signalen een superieure strategie is voor KV-cache compressie vergeleken met statische heuristieken of proxy-scores. Door direct een beleid te trainen om te voorspellen welke tokens belangrijk zullen zijn voor toekomstige queries, bereikt KVpop:

1. Hoge Getrouwheid: Het behoudt bijna de volledige performance van de teacher, zelfs bij hoge compressieverhoudingen (tot 88%), een regime waarin andere methoden falen.
2. Robuustheid: Het geleerde beleid wordt effectief overgedragen naar out-of-domain taken, wat suggereert dat het model algemene principes van token-nut leert in plaats van te overfitten op specifieke trainingsdistributies.
3. Efficiëntie: Het dwingt een harde geheugenlimiet af, wat stabiele, lange-context decodering mogelijk maakt met constant geheugengebruik en verbeterde throughput vergeleken met dynamische gating-methoden zoals DMS.

De auteurs concluderen dat KVpop aantoont dat geleerde eviction-beleid, wanneer ze correct worden gesuperviseerd door toekomstige nut en de ruimte krijgen om nabije-toekomst context te gebruiken, effectief de geheugenbottleneck in autoregressieve generatie kan oplossen zonder kwaliteitsverlies van het model.