Technische Samenvatting: xKV - Cross-Layer KV-Cache Compressie via Gelijkgerichte Singuliere Vector Extractie

1. Probleemstelling

Long-context Large Language Models (LLMs) maken geavanceerde toepassingen mogelijk, zoals grootschalige informatiewinning en het debuggen van codebases, maar staan voor een kritieke bottleneck: het geheugengebruik van Key-Value (KV) caches tijdens inferentie. Naarmate contextlengtes oplopen tot honderdduizenden of miljoenen tokens, wordt het geheugenvoetafdruk van de KV-Cache onbetaalbaar, wat gelijktijdige inferentieverzoeken en doorvoer ernstig beperkt.

Bestaande mitigatiestrategieën vallen over het algemeen in twee categorieën:

• Intra-layer methoden: Technieken zoals kwantisatie, token-verwijdering en low-rank decompositie op één laag comprimeren de KV-Cache van elke laag onafhankelijk. Hoewel ze tot op zekere hoogte effectief zijn, missen ze potentiële redundantie over lagen heen.
• Cross-layer methoden: Benaderingen zoals Cross-Layer Attention (CLA) en YOCO vereisen kostbare architecturale wijzigingen en vooraf trainen. Anderen, zoals MiniCache, vertrouwen op post-hoc samenvoeging van aangrenzende lagen op basis van token-voor-token cosinusgelijkenis. Empirische analyse toont echter aan dat de token-voor-token gelijkenis tussen lagen vaak te zwak is om robuuste compressie te ondersteunen zonder significante achteruitgang in nauwkeurigheid.

2. Methodologie: xKV

De auteurs stellen xKV voor, een post-training, plug-and-play compressiemethode die diepe structurele redundantie over lagen heen benut zonder dat modelhertraining of architecturale wijzigingen nodig zijn.

2.1. Kernobservatie: Gelijkgerichte Singuliere Vectoren

Het paper daagt de aanname uit dat lagen distinct zijn op basis van token-voor-token cosinusgelijkenis. Met behulp van Centered Kernel Alignment (CKA) tonen de auteurs aan dat, hoewel individuele token-embeddings kunnen divergeren, de dominante singuliere vectoren (de basis die de belangrijkste variaties in de tokenruimte definieert) sterk gelijkgericht zijn over aangrenzende lagen. Deze hoge CKA-score geeft aan dat de geometrische structuur van de KV-Cache subruimten gedeeld wordt, zelfs als punt-voor-punt gelijkenissen laag zijn.

2.2. Cross-Layer Factorisatie (CLF)

Gedreven door deze uitlijning, partitioneert xKV de $N$ lagen van het model in opeenvolgende vensters van grootte $W$. Voor elk venster voert het een Cross-Layer Factorisatie uit:

1. Concatenatie: De KV-Caches van alle lagen binnen een venster worden horizontaal samengevoegd.
2. SVD: Er wordt een rang-$r$ Singular Value Decomposition (SVD) berekend op de samengevoegde cache.
3. Decompositie: Het resultaat levert een gedeelde tokenbasis ($A_k$) voor het venster en laagspecifieke reconstructiematrices ($B_\ell$) op.
4. Opslag: In plaats van $W$ full-rank caches op te slaan, slaat het systeem één gedeelde basis en $W$ kleinere reconstructiematrices op, waardoor de geheugencomplexiteit wordt gereduceerd van $O(W \cdot L \cdot d)$ naar $O(L \cdot r + W \cdot r \cdot d)$.

2.3. Selectieve Reconstructie (SR)

Om de computationele bottleneck van het reconstrueren van alle tokens tijdens de decode-fase aan te pakken (wat lineair schaalt met de sequentielengte $L$), introduceert xKV Selectieve Reconstructie (SR).

• Mechanisme: In plaats van dichte reconstructie, identificeert SR een schaarse subset van tokens ($S$) die het meest relevant zijn voor de huidige query, gebruikmakend van landmark-gestuurde chunk-selectie (vergelijkbaar met ShadowKV).
• Uitvoering: Alleen de rijen die corresponderen met deze kritieke tokens worden gereconstrueerd met behulp van de gedeelde basis en laagspecifieke matrices.
• Voordeel: Dit reduceert de reconstructiekost tot een constante waarde ten opzichte van de sequentielengte, wat end-to-end snelheidswinst mogelijk maakt.

3. Belangrijkste Bijdragen

1. Ontdekking van Cross-Layer Redundantie: Het paper onthult dat KV-Caches sterk gelijkgerichte dominante singuliere vectoren delen over lagen heen, een vorm van redundantie die wordt verdoezeld door lage token-voor-token gelijkenis.
2. xKV Algorithmus: Een trainingsvrije methode die een gedeelde basis extrahert via Cross-Layer Factorisatie (CLF), waardoor effectieve informatiedeling over $W$ lagen mogelijk wordt ($W > 2$).
3. Selectieve Reconstructie: Een optimalisatie tijdens het decoderen die alleen query-relevante tokens reconstrueert, waardoor de lineaire computationele overhead van dichte reconstructie wordt geëlimineerd.
4. Plug-and-Play Compatibiliteit: De methode vereist geen fine-tuning of architecturale wijzigingen, waardoor deze toepasbaar is op bestaande vooraf getrainde modellen (bijv. Llama-3.1, Qwen2.5, DeepSeek-V2).

4. Experimentele Resultaten

De auteurs evalueerden xKV op Llama-3.1-8B-Instruct, Qwen2.5-7B/14B-Instruct en DeepSeek-Coder-V2-Lite met behulp van de RULER, LongBench en Multi-turn NIAH benchmarks.

• Compressie versus Nauwkeurigheid:
  • xKV bereikt tot 8× KV-Cache compressie met minimale nauwkeurigheidsverlies (≤3% gemiddeld) in vergelijking met full attention baselines.
  • Op Llama-3.1-8B bereikt xKV bij 8× compressie een 88,50% gemiddelde nauwkeurigheid op RULER, wat significant beter presteert dan MiniCache (45,04% bij 1,3×) en gelijkwaardig is aan of beter presteert dan token-verwijderingsmethoden zoals SnapKV en PyramidKV.
  • In multi-turn NIAH-taken behoudt xKV stabiele prestaties over conversatierondes heen, terwijl token-verwijderingsmethoden na de eerste beurt een steile daling vertonen.
• Efficiëntie en Doorvoer:
  • Geheugen: xKV maakt grotere batchgroottes mogelijk binnen hetzelfde GPU-geheugenbudget.
  • Latentie: Dichte reconstructie verlaagt de attention-latentie met maximaal 3,6×.
  • Doorvoer: Gecombineerd met Selectieve Reconstructie (xKV-SR) bereikt het systeem tot 4,23× end-to-end snelheidswinst ten opzichte van full attention baselines op NVIDIA A100 GPU's.
  • Vergelijking: xKV-SR overtreft ShadowKV met 30% in doorvoer onder vergelijkbare nauwkeurigheidsniveaus, terwijl zowel Keys als Values op de GPU blijven, waardoor PCIe-transferbottlenecks worden vermeden.
• Robuustheid: De methode blijft effectief over verschillende modelarchitecturen, waaronder Grouped-Query Attention (GQA) en Multi-head Latent Attention (MLA) in DeepSeek-V2.

5. Betekenis en Claims

Het paper beweert dat xKV een praktische, plug-and-play aanpak biedt om tegelijkertijd het geheugenvoetafdruk en de latentie te verminderen voor long-context LLM inferentie. Door de focus te verschuiven van token-voor-token gelijkenis naar uitlijning van singuliere vectoren, ontsluit xKV ongeëxploiteerde inter-layer redundantie die eerdere methoden misten.

De auteurs benadrukken dat hun aanpak:

• De hoge kosten van vooraf trainen die nodig zijn bij architecturale wijzigingen, vermijdt.
• De nauwkeurigheidsbeperkingen van op gelijkenis gebaseerde samenvoeging (MiniCache) overwint.
• De computationele bottleneck van dichte low-rank reconstructie oplost via Selectieve Reconstructie.
• Volledige inferentie op het apparaat (GPU) voor gecomprimeerde caches mogelijk maakt, waardoor de PCIe-bandbreedtelimieten worden omzeild die methoden op basis van offloading zoals ShadowKV beperken.

Het werk wordt gepresenteerd als een fundamentele stap naar efficiënte long-context inferentie, waarbij de code publiek beschikbaar wordt gesteld om adoptie te faciliteren.