Technische Samenvatting: MOMENTKV

Probleemstelling

Autoregressieve decodering in Transformer-gebaseerde taalmodellen vertrouwt op de Key-Value (KV) cache, die eerder berekende projecties opslaat om deze tijdens opeenvolgende decodestappen te hergebruiken. Naarmate sequentielengtes schalen naar tienduizenden of honderdduizenden tokens, groeit de geheugenvoetafdruk van de KV-cache lineair, wat de primaire bottleneck vormt voor lange-context inferentie.

KV-cache-evictie (het verwijderen van gegevens) pakt dit aan door een subset van een vaste grootte van key-value paren te behouden en de rest te verwijderen. Hoewel bestaande methoden de strategieën voor tokenselectie hebben verbeterd (bijv. gebaseerd op cumulatieve aandachtsgewichten of waarde-bewuste scoring), delen zij een fundamentele beperking: post-evictie inferentie renormaliseert de aandacht uitsluitend over de behouden tokens, waarbij de verwijderde set volledig wordt genegeerd.

De auteurs identificeren dat de primaire bron van output-degradatie niet louter het resterende aandachtsmassa op de verwijderde tokens is (wat bestaande methoden al minimaliseren), maar een directionele mismatch tussen de behouden en de verwijderde tokenset. Empirische analyse laat zien dat de verwijderde tokens in de praktijk vaak bijna orthogonaal zijn aan de behouden tokens. Bijgevolg kan zelfs een kleine hoeveelheid verwijderde aandachtsmassa een substantiële directionele fout produceren die standaard renormalisatie niet kan herstellen. Dit creëert een "directionele kloof" waarbij het verbeteren van enkel de tokenselectie leidt tot afnemende meeropbrengsten, omdat meer selectieve retentie de behouden set concentreert op specifieke query-gealigneerde richtingen, wat de kloof met de verwijderde set verder vergroot.

Methodologie: MOMENTKV

Om deze directionele kloof aan te pakken, stellen de auteurs MOMENTKV voor, een framework dat compacte, kleine momentstatistieken bijhoudt over de verwijderde tokensegmenten. Deze statistieken dienen een dubbel doel: het sturen van de evictie-beslissingen en het corrigeren van de aandacht-output tijdens inferentie.

1. Compacte Momentstatistieken

In plaats van de verwijderde tokens op te slaan, houdt MOMENTKV vier lopende sommen per aandacht-head bij, die bij elke evictie incrementeel worden bijgewerkt:

• Aantal ($n_e$): Het aantal verwijderde tokens.
• Key Som ($s_k$): $\sum_{i \in E} k_i$.
• Value Som ($s_v$): $\sum_{i \in E} v_i$.
• Outer-Product Som ($S$): $\sum_{i \in E} v_i k_i^\top$.

Deze statistieken maken het mogelijk om de key-gemiddelde ($\bar{k}$), de value-gemiddelde ($\bar{v}$) en de value-key covariantie ($\tilde{S}$) te berekenen zonder de verwijderde tokens opnieuw te hoeven bezoeken. De totale opslagoverhead is $O(d^2)$ per head, onafhankelijk van de contextlengte.

2. Eerste-orde Benadering van de Verwijderde Output

Tijdens inferentie leidt MOMENTKV een gesloten vorm eerste-orde benadering af van de verwijderde aandacht-output, $\hat{f}_E(q)$. Door de logits te recentreren rond het gemiddelde van de keys en een eerste-orde Taylor-expansie toe te passen op de softmax-functie, leiden de auteurs het volgende af:
$$\hat{f}_E(q) = \bar{v} + \frac{1}{n_e \sqrt{d}} \tilde{S} q$$

• De $\bar{v}$ term fungeert als een query-onafhankelijke baseline, die de gemiddelde richting van de verwijderde sub-output herstelt.
• De $\tilde{S} q$ term fungeert als een query-adaptieve correctie, die de loodrechte variatie (directionele inhoud) vangt die de baseline alleen mist.

Deze benadering stelt het model in staat om een schatting te reconstrueren van de verloren verwijderde bijdrage, waardoor de verloren directionele informatie effectief wordt hersteld die standaard evictie wegvaagt.

3. Moment-geïnformeerde Evictie

De momentstatistieken sturen ook het evictieproces zelf aan. Voor elke behouden token $j$ berekent het algoritme een moment-residue:
$$r_j = v_j - \bar{v} - \frac{1}{n_e \sqrt{d}} \tilde{S} k_j$$
Dit residue meet hoeveel nieuwe directionele informatie token $j$ bevat bovenop wat de geaccumuleerde samenvatting al vangt. De evictie-score wordt gedefinieerd als het product van het aandachtsgewicht en de norm van het residue: $\text{score}(j) = \alpha_j \cdot \|r_j\|$.

• Tokens met lage residuen worden goed voorspeld door het moment-model en krijgen prioriteit voor evictie.
• Tokens met hoge residuen (die unieke directionele informatie dragen) worden behouden, zelfs als hun aandachtsgewichten lager zijn.

Dit creëert een wederzijds versterkende lus: moment-geïnformeerde evictie houdt de verwijderde set "geometrisch regelmatig" (lage logit-spreiding $\sigma$), wat de benaderingsgrens voor de inferentie-correctie aanscherpt. Op zijn beurt zorgt de nauwkeurige correctie voor betere residuen voor toekomstige evictie-beslissingen.

4. Normalisatie-gecorrigeerde Inferentie

De uiteindelijke aandacht-output wordt gereconstrueerd als een convexe combinatie van de behouden output ($f_R$) en de benaderde verwijderde output ($\hat{f}_E$):
$$\hat{f}(q) = \hat{w}_R f_R(q) + (1 - \hat{w}_R) \hat{f}_E(q)$$
Het menggewicht $\hat{w}_R$ wordt geschat met behulp van een ondergrens op de verwijderde partitie-functie afgeleid via de Jensen-ongelijkheid. Dit zorgt ervoor dat het model rekening houdt met de omvang van de verwijderde massa terwijl de richting wordt gecorrigeerd.

Belangrijkste Bijdragen

1. Identificatie van de Directionele Kloof: Het paper stelt vast dat de dominante foutbron in KV-cache evictie de bijna orthogonaalheid is tussen de behouden en de verwijderde sub-outputs, en niet enkel de omvang van de verwijderde aandachtsmassa.
2. Compacte Momentstatistieken: Introductie van een geheugenefficiënte methode om de geometrische eigenschappen (gemiddelde en covariantie) van verwijderde tokens bij te houden met slechts $O(d^2)$ opslag per head.
3. Gesloten Vorm Correctie: Afleiding van een eerste-orde benadering voor de verwijderde aandacht-output die zowel de gemiddelde richting als de query-afhankelijke variaties herstelt.
4. Dual-Role Mechanisme: Een verenigd framework waarbij momentstatistieken gelijktijdig de evictie-selectie optimaliseren (door redundante tokens te identificeren) en de aandacht-outputs corrigeren (door verloren directionele informatie te herstellen).

Experimentele Resultaten

De auteurs evalueerden MOMENTKV op LongBench en RULER benchmarks met behulp van LLaMA-3.1-8B-Instruct en Qwen3-4B-Instruct.

• Prestaties: MOMENTKV presteerde consequent beter dan alle baselines (H2O, SnapKV, PyramidKV, Ada-KV) over alle cache-budgetten ($L \in \{128, 256, 512, 1024\}$).
• Agressieve Compressie: De grootste winsten werden waargenomen onder agressieve compressie. Op LongBench met $L=128$ en LLaMA-3.1-8B behaalde MOMENTKV een gemiddelde scoreverbetering van +1,35 ten opzichte van de sterkste baseline (Ada-KV), waarbij 94,1% van de volledige cache-prestaties behouden bleef.
• Retrieval Taken: Op RULER, dat precieze token-retrieval test, toonde MOMENTKV proportionele winsten van +3,4 en +3,5 over Ada-KV voor respectievelijk LLaMA en Qwen, wat de effectiviteit ervan in retrieval-intensieve omgevingen valideert.
• Ablatie-studies:
  • Normalisatie-gecorrigeerde Inferentie (NC) droeg aanzienlijk meer bij (+2,41) dan Moment-geïnformeerde Evictie (MI) alleen (+0,57), wat bevestigt dat het herstellen van de verwijderde richting de primaire drijfveer voor prestaties is.
  • De combinatie van beide mechanismen leverde een synergetisch effect op (+3,17 totale winst), wat groter is dan de som van de individuele winsten.
• Efficiëntie: De methode voegt slechts 1,6 ms latentie per token toe en vereist 4,1 MB aan hulpgeheugen (6,4% van de behouden cache), terwijl de per-token latentie met meer dan 65% wordt verminderd vergeleken met een volledige cache.

Betekenis

Het paper betoogt dat de weg naar verdere verbetering in lange-context inferentie niet enkel ligt in betere tokenselectie, maar in het herstellen van de directionele informatie die huidige evictie-strategieën wegvaagt. MOMENTKV laat zien dat door compacte statistieken over de verwijderde set bij te houden, het mogelijk is de directionele kloof te dichten, waardoor de "onvermijdelijke kosten" van evictie worden getransformeerd in een corrigeerbare benaderingsfout. Deze aanpak biedt een praktische, schaalbare oplossing voor het implementeren van lange-context taalmodellen zonder het efficiëntievoordeel van fixed-budget KV-caches op te offeren.