Technische Samenvatting: SelKV - Selectieve KV-cache Merging met Per-Token Merge-of-Drop en Attention Compensatie

1. Probleemstelling

Large Language Models (LLM's) kampen tijdens autoregressieve inferentie met een kritieke bottleneck: de Key-Value (KV) cache. Naarmate de contextlengte ($n$) toeneemt, groeit de geheugenvoetafdruk van de KV-cache lineair ($O(L \cdot H \cdot d \cdot n)$), wat de batchgroottes, doorvoer en implementatiekosten voor lang-context toepassingen ernstig beperkt.

Hoewel bestaande compressiemethoden proberen dit te verzachten, lijden ze aan twee primaire beperkingen:

1. Uniforme Merge-of-Drop Beslissingen: Huidige benaderingen passen doorgaans een binaire strategie toe: ofwel alle geëvicteerde tokens mergen (bijv. KVMerger, WeightedKV), ofwel ze volledig droppen (bijv. H2O, SnapKV). Dit gebrek aan granulariteit betekent dat wanneer de waarde-vector van een geëvicteerd token te verschillend is van zijn merge-target, het gewogen gemiddelde de behouden semantische representatie corrumpeert. Omgekeerd leidt willekeurig droppen tot het verlies van waardevolle informatie.
2. Attention Sag: Wanneer meerdere tokens worden samengevoegd tot één enkele cache-entry, ontvangt die entry dezelfde softmax attention massa als een enkele niet-gemergde token. Dit resulteert in "attention sag", waarbij gemergde posities systematisch minder aandacht krijgen door een factor die proportioneel is aan het aantal gemergde tokens ($\sim m$), waardoor het model de behouden informatie onderbenut.

2. Methodologie: SelKV

De auteurs stellen SelKV (Selective KV Cache Merging) voor, een training-vrij, duaal-component framework ontworpen om deze beperkingen aan te pakken. SelKV werkt als een vervangbare module die geïntegreerd kan worden in bestaande merging pipelines na de prefill-fase.

Kerncomponenten

1. Soft Cosine Gate (Adresseren van Uniforme Beslissingen)
In plaats van een binaire merge-of-drop beslissing, introduceert SelKV een continue, per-token gating-mechanisme gebaseerd op de gelijkenis van waarde-vectoren.

• Mechanisme: Voor elk geëvicteerd token dat naar een merge-target wordt gerouteerd, berekent het systeem de cosinus-gelijkenis tussen hun waarde-vectoren.
• Gating Functie: Een soft gate $g = \max(\cos\_sim, 0)$ moduleert de merge-intensiteit:
  • $g \approx 1$: Hoge gelijkenis $\rightarrow$ Volledige merge.
  • $g = 0$: Dissimilaire (orthogonale) vectoren $\rightarrow$ Drop de token.
  • $0 < g < 1$: Matige gelijkenis $\rightarrow$ Partiële merge.
• Implementatie: Dit vereist geen geleerde parameters of drempelwaarde-afstemming. De gemergde waarde voor een behouden token $i$ wordt berekend als een gated, attention-gewogen gemiddelde van de originele waarde en de bijdragen van de geëvicteerde tokens, geschaald door hun respectievelijke gates.

2. Attention Compensation (Adresseren van Attention Sag)
Om de systematische onder-attention van gemergde posities te corrigeren, past SelKV een decoding-time logit bias toe.

• Mechanisme: In tegenstelling tot eerdere methoden die ruwe merge-aantallen gebruiken (die instabiel kunnen zijn bij hoge compressieratio's), afleidt SelKV een attention-ratio ($R_{h,i}$) uit prefill-statistieken. Deze ratio schat de totale attention massa die een gemergde positie zou moeten ontvangen ten opzichte van de oorspronkelijke staat.
• Correctie: Een bias-term $\alpha \cdot \log(R_i)$ wordt toegevoegd aan de attention logits tijdens decoding. Dit herkalibreert de softmax-distributie om ervoor te zorgen dat gemergde posities de passende attention massa ontvangen die proportioneel is aan de informatie die zij nu vertegenwoordigen.

Pipeline Overzicht

De SelKV-pipeline verloopt via zes stadia post-prefill:

1. Importance Scoring: Tokens worden gescoord op basis van het product van attention weight en de magnitude van de waarde-vector (volgens VATP), in plaats van alleen de attention weight.
2. Token Selectie: Top-$m_s$ tokens worden per laag geselecteerd (en per head voor MHA, of via een unie voor GQA) om te behouden, samen met een vaste recente window.
3. Merge-Target Routing: Geëvicteerde tokens worden gerouteerd naar behouden tokens binnen positionele buckets op basis van prefill attention patronen.
4. Selectieve Merging: De soft cosine gate bepaalt de merge-intensiteit voor elk geëvicteerd token.
5. Attention Compensation: De logit bias wordt berekend en toegepast.
6. RoPE Repositioning: Keys worden opnieuw ingebed op contigueuze posities (uitgeschakeld voor long-context evaluatie om positionele distortie te voorkomen).

3. Belangrijkste Bijdragen

• Per-Token Granulariteit: De introductie van een soft cosine gate die een continue modulatie van de merge-intensiteit mogelijk maakt op basis van semantische gelijkenis, waardoor de "alles-of-niets" fout van eerdere methoden wordt geëlimineerd.
• Stabiele Attention Compensation: Een nieuw op attention-ratio gebaseerd logit bias dat de attention sag robuuster corrigeert dan methoden op basis van ruwe aantallen, vooral onder zware compressie.
• Training-Vrije Integratie: Het framework vereist geen fine-tuning en kan worden geïntegreerd in bestaande one-shot compressie pipelines.
• Architectuur Aanpasbaarheid: De methode ondersteunt zowel Multi-Head Attention (MHA) als Grouped-Query Attention (GQA) architecturen, met specifieke afhandeling van token selectie unies voor GQA-modellen.

4. Experimentele Resultaten

Het framework werd geëvalueerd op LongBench (16 Engelse datasets) over drie modellen: LongChat-7B (MHA), LLaMA-3.1-8B (GQA), en Gemma-2-9B (GQA), waarbij slechts 25% van de KV-cache behouden bleef.

• Prestaties vs. Baselines: SelKV presteerde consequent beter dan representatieve one-shot baselines (SnapKV, LOOK-M, PyramidKV).
  • Op GQA-modellen bereikte SelKV de dichtstbijzijnde prestaties bij de full-cache baseline, met gemiddelde verschillen van $-0.67$ (Gemma-2) en $-0.73$ (LLaMA-3.1).
  • Op MHA-modellen (LongChat) presteerde PyramidKV iets beter in de specifieke 31.5k token setting, hoewel SelKV competitief bleef.
• Multi-Document QA: Opvallend genoeg overtrof SelKV de full-cache baseline op verschillende multi-document QA taken (bijv. HotpotQA en MuSiQue op LLaMA-3.1). De auteurs interpreteren dit als bewijs dat selectieve merging fungeert als een impliciete attention filter, die afleidende context onderdrukt terwijl relevante bewijslast behouden blijft.
• Efficiëntie: Bij een contextlengte van 100k tokens leverde de gecomprimeerde cache een 3.3x versnelling in decoding op vergeleken met de full-cache baseline.
• Robuustheid: De methode toonde robuustheid over variërende compressieratio's (10%–90% behoud) en taakcategorieën, inclus inclusief samenvatting, few-shot leren en code-aanvulling.

5. Betekenis en Claims

Het artikel positioneert SelKV als een belangrijke stap vooruit om de lange-context inferentie van LLM's haalbaar te maken zonder kwaliteitsverlies. De auteurs beweren dat:

• Selectieve merging superieur is aan uniforme strategieën: Door de merge-intensiteit te moduleren op basis van representatie-gelijkenis, vermijdt SelKV de semantische corruptie die wordt veroorzaakt door willekeurige merging.
• Attention compensation cruciaal is: Zonder de aandacht te corrigeren voor attention sag, worden gemergde tokens onderbenut; de ratio-gebaseerde correctie van SelKV zorgt ervoor dat de behouden informatie effectief wordt geraadpleegd.
• Impliciete Filtering: Het vermogen van SelKV om full-cache baselines te overtreffen in complexe multi-document QA suggereert dat het compressieproces zelf de modelprestaties kan verbeteren door ruis te filteren, een bevinding die de aanname uitdaagt dat compressie altijd de capaciteit verslechtert.

De auteurs concluderen dat SelKV een praktische, parameter-vrije oplossing biedt voor het inzetten van LLM's in geheugenbeperkte omgevingen, terwijl de generatiekwaliteit bijna lossless blijft, met name voor GQA-gebaseerde architecturen.