Technische Samenvatting: PolyKV: Heterogene Retentie en Allocatie voor KV-cache Compressie

Probleemstelling

Lange-context inferentie in Large Language Models (LLMs) kampt met een aanzienlijke geheugenbottleneck door de groei van de Key-Value (KV) cache. Hoewel bestaande methoden voor KV-cache eliminatie (eviction) effectief het geheugengebruik verminderen door specifieke subsets van tokens te behouden, passen zij doorgaans een uniform compressiebeleid en een uniform cachebudget toe over alle transformer-lagen.

Dit uniforme ontwerp negeert de inherente heterogeniteit van transformer-lagen. Voorgaand onderzoek wijst uit dat verschillende lagen variëren in hun aandachtspatronen, representatieve abstracties, contextuele rollen en gevoeligheid voor pruning. Bijgevolg is een enkele globale eliminatieregel een grove abstractie die geen rekening houdt met het feit dat:

1. Verschillende lagen verschillende eliminatiepatronen kunnen prefereren (bijv. heavy hitters versus attention sinks).
2. Verschillende lagen een variërende gevoeligheid vertonen voor KV-cache reductie, wat impliceert dat zij verschillende cachecapaciteiten vereisen.

Huidige adaptieve methoden zijn begonnen met het verkennen van niet-uniforme budgetten, maar falen vaak in het behandelen van het eliminatiepatroon en het cachebudget als gekoppelde, laag-specifieke ontwerpkeuzes.

Methodologie: PolyKV

De auteurs stellen PolyKV voor, een framework dat KV-cache eliminatie herformuleert als een laag-specifieke designruimte. PolyKV koppelt expliciet twee beslissingen voor elke laag en elke inferentiefase (prefill en decoding):

1. Wat te elimineren: Het selecteren van een specifiek eliminatiepatroon (bijv. H2O, TOVA, PyramidKV) dat geschikt is voor die specifieke laag.
2. Hoeveel te behouden: Het toewijzen van een specifiek cachebudget aan die laag.

Kerncomponenten

• Formulering van de Designruimte: PolyKV definieert een strategie $\Pi$ waarbij elke laag $l$ een eliminatiepatroon $r_{l,\phi}$ en een budget $b_{l,\phi}$ krijgt toegewezen voor elke fase $\phi$. Het totale budget wordt beperkt door een vaste globale limiet $B$.
• Offline Kalibratie: Om een uitputtende online zoektocht te vermijden, gebruikt PolyKV een kleine kalibratieset om een vaste, statische strategie te generen. Het verzamelt twee categorieën laag-specifieke signalen:
  • Eliminatie-preferentiesignalen: Meet welk eliminatiepatroon de FullKV-berekening van een laag het beste behoudt. De metrieken omvatten cosinus-gelijkenis tussen geperturbeerde en FullKV-activaties, en Perplexity (PPL) gevoeligheid (de degradatie veroorzaakt door een specifiek patroon).
  • Cachebudget-gevoeligheidssignalen: Schat hoe sterk een laag wordt beïnvloed door KV-reductie. Metrieken omvatten attention entropie (diffuse aandacht vereist grotere budgetten) en PPL-gevoeligheid (geaggregeerde worst-case degradatie).
• Strategieconstructie:
  • Patroonselectie: Voor elk laag-fase paar wordt het eliminatiepatroon met de hoogste voorkeursscore (bijv. de hoogste cosinus-gelijkenis of laagste PPL-degradatie) geselecteerd.
  • Budgetallocatie: Budgetten worden proportioneel toegewezen op basis van gevoeligheidssignalen. Lagen met een hogere gevoeligheid (of entropie) ontvangen grotere budgetten, onderworpen aan een minimum budgetrestrictie ($b_{min}$).
• Implementatie: PolyKV is geïmplementeerd als een hook-gebaseerd framework. Het injecteert logica op veelvoorkomende punten in de generatiepijplijn (pre-kernel, post-kernel, cache update) om de modulaire integratie van diverse eliminatiemethoden te ondersteunen zonder de onderliggende modelcode te wijzigen.

Belangrijkste Bijdragen

1. Herformulering van KV-eliminatie: Het artikel herkadert KV-cache eliminatie als een laag-specifieke designruimte die de keuze van het eliminatiepatroon koppelt aan de allocatie van cachecapaciteit onder een gedeelde geheugenrestrictie.
2. Het PolyKV-framework: Een systeem dat deze designruimte operationaliseert door eliminatiepatronen en budgetten op laag-granulariteit te componeren. Het integreert meerdere eliminatieprimitieven via een uniforme interface en gebruikt offline kalibratiesignalen om de strategieconstructie te sturen.
3. Empirische Validatie: De auteurs demonstreren dat heterogene, laag-specifieke strategieën superieur zijn aan uniforme, single-policy baselines onder vaste geheugenbudgetten.

Experimentele Resultaten

Experimenten werden uitgevoerd op LLaMA-3.1-8B en Qwen3-8B met behulp van de LongBench en RULER benchmarks.

• Prestaties op LongBench (512-token gemiddeld budget):
  • LLaMA-3.1-8B: PolyKV herstelde 54,5% van het prestatieverschil tussen de sterkste single-policy baseline en FullKV.
  • Qwen3-8B: PolyKV herstelde 25,7% van het gat.
  • PolyKV presteerde consistent beter dan de sterkste baseline over een budget-sweep (128–1024 tokens), met verbeteringen variërend van 1,7% tot 6,4%.
• Prestaties op RULER:
  • De resultaten waren gemengd. PolyKV overtrof de baseline licht op LLaMA, maar presteerde minder goed op Qwen (met 3,98 punten). De auteurs schrijven dit toe aan de aard van de RULER-taken, die vertrouwen op ijle (sparse) evidence tokens die mogelijk worden geëlimineerd door heterogene beleid als de kalibratiesignalen ze niet perfect identificeren.
• Ablatie-studies:
  • Eliminatiepatroon vs. Budget: De primaire prestatiewinst komt voort uit het toewijzen van verschillende eliminatiepatronen aan verschillende lagen, met name tijdens de prefill-fase.
  • Budgetallocatie: Niet-uniforme budgetallocatie biedt aanvullende winsten, maar is meer model-afhankelijk en minder robuust dan patroonselectie. Het is het meest effectief wanneer het totale budget voldoende is om gevoelige lagen niet uit te hongeren.
• Taakanalyse: De winsten waren het meest uitgesproken in taken die een brede contextuele verstandhouding vereisen (bijv. QA, few-shot learning). De prestaties verslechterden bij taken die een exacte extractie van ijle tokens vereisen (bijv. NIAH in RULER) of strikte naleving van lokale vorm (bijv. codegeneratie).

Betekenis en Claims

Het artikel claimt dat het behandelen van KV-cache eliminatie als een compositioneel, laag-bewust optimalisatieprobleem in plaats van een enkel globaal retentiebeleid superieure resultaten oplevert.

• Heterogeniteit is de sleutel: De studie valideert dat transformer-lagen niet homogeen zijn; een "one-size-fits-all" eliminatiestrategie is suboptimaal.
• Statische Efficiëntie: Door offline kalibratie te gebruiken, bereikt PolyKV deze heterogene voordelen zonder de runtime overhead van een online beleidszoektocht, wat het praktisch maakt voor implementatie.
• Beperkingen: De auteurs merken nederig op dat de huidige kalibratiesignalen niet altijd de "ijle maar beslissende" evidence tokens bewaren die nodig zijn voor exacte extractietaken, wat wijst op een richting voor toekomstige verfijning van het signaalontwerp.

Samenvattend demonstreert PolyKV dat door eliminatiestrategieën en cachecapaciteiten af te stemmen op de specifieke kenmerken van elke transformer-laag, het mogelijk is om de geheugen-kwaliteit trade-off bij lange-context LLM-inferentie aanzienlijk te verbeteren.