MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

1. Het Probleem

De inferentie van Large Language Models (LLMs) met lange contexten wordt beperkt door het KV-cache-geheugen (Key-Value cache). Voor lange sequenties kan dit geheugen oplopen tot tientallen gigabytes, wat leidt tot latency-problemen en "out-of-memory" fouten op standaardhardware.

Bestaande compressietechnieken richten zich meestal op één as van de vierdimensionale KV-tensor:

• Eviction: Het verwijderen van tokens (sequentie-as).
• Quantization: Het verminderen van de precisie (bits per element).
• Low-rank projectie of Cross-layer sharing.

Een kritiek punt in de huidige literatuur is dat deze methoden uniform worden toegepast: elke transformer-laag krijgt dezelfde compressie-instellingen (bijv. hetzelfde percentage verwijderde tokens of dezelfde bit-breedte). Het paper stelt dat deze homogeniteit de nauwkeurigheid onnodig beperkt, omdat verschillende lagen van een model zeer verschillend reageren op compressie. De optimale mix van eviction en quantization verschilt per laag.

2. Methodologie: MoE-nD

De auteurs introduceren MoE-nD (Mixture-of-Experts n-Dimensional), een raamwerk dat compressie op per-laag-niveau optimaliseert over meerdere assen tegelijk.

• Mixture-of-Experts (MoE) Formulering:
  Elke compressie-as (evictie-ratio, K-precisie, V-precisie) wordt behandeld als een "expert". Een router (oplosser) kiest voor elke laag de specifieke combinatie van instellingen die het beste past binnen een globaal geheugenbudget.
• Offline Calibratie & Sensitiviteitstabel:
  Voordat het model wordt gebruikt, wordt een "sensitivity table" ($S$) gegenereerd. Hierbij wordt voor elke laag ($\ell$) en elke mogelijke compressie-configuratie ($c$) de relatieve L2-fout gemeten tussen de volledige precisie-uitvoer en de gecomprimeerde uitvoer. Dit dient als proxy voor de kwaliteitsverlies.
• Greedy Solver:
  Een offline, greedy algoritme kiest de per-laag routing $(keep\_ratio, k\_bits, v\_bits)$ die de totale voorspelde kwaliteitsverlies minimaliseert onder het gegeven geheugenbudget. Dit is een discrete optimalisatieprobleem dat te groot is voor exacte zoekopdrachten, maar efficiënt oplosbaar is met een greedy benadering (startend bij de goedkoopste configuratie en iteratief upgraden van de laag met de grootste "marginal sensitivity reduction").
• Heterogene Attention Patch:
  Tijdens inferentie wordt een aangepaste attention-kern gebruikt die omgaat met heterogene caches. Omdat elke laag een verschillende lengte van de cache kan hebben (door verschillende eviction-ratio's), worden per-laag positiearrays ($p^{(\ell)}$) gebruikt voor de RoPE (Rotary Positional Embedding) inversie. Dit zorgt ervoor dat de tokens correct worden geroteerd, ongeacht of ze zijn verwijderd of niet.

3. Belangrijkste Bijdragen

1. Gecombineerde Per-Layer Routing: MoE-nD is de eerste methode die zowel eviction-ratio's als K/V-quantization bits gelijktijdig per laag routeert onder een globaal budget.
2. Empirische Inzichten: De auteurs tonen aan dat de gevoeligheid voor compressie per laag sterk varieert. Vooral bij agressieve eviction (75-90%) varieert de fout tussen lagen met een factor van 500 tot 700. Ook wisselt de voorkeur tussen eviction en quantization per laag; wat voor laag A goed is (bijv. quantization), kan voor laag B rampzalig zijn.
3. Ablatie-studie: Een analyse toont aan dat de winst voornamelijk komt van per-laag eviction routing, terwijl per-laag quantization routing op zich weinig tot geen extra winst oplevert. De kracht zit in het combineren van beide.
4. Gedefinieerde Scope: Het paper identificeert duidelijk wanneer de methode werkt (lange contexten of strakke budgetten) en wanneer niet (korte prompts waar de solver kiest voor behoud van alle tokens).

4. Resultaten

De methode werd getest op DeepSeek-R1-Distill-Qwen-7B met een reeks benchmarks:

• LongBench-v1 (4 taken, 16k tokens):
  • De heterogene variant van MoE-nD ($2d_{hetero}$) bereikt bij 14x compressie (136 MB geheugen) een nauwkeurigheid die gelijk is aan het oncompressie-baseline (1.9 GB).
  • Vergelijkbare methoden (uniforme eviction, uniforme quantization, of alleen per-laag quantization) blijven bij vergelijkbaar geheugen ver onder de 8/100 punten, terwijl MoE-nD rond de 12 punten blijft.
  • Bij 14x compressie is er geen meetbaar verlies in nauwkeurigheid ten opzichte van de volledige cache.
• AIME Reasoning Benchmarks (Wiskundig redeneren):
  • MoE-nD overtreft de sterkste niet-heterogene twee-as-baseline met +6 tot +27 punten over acht verschillende configuraties.
  • Bij zeer strakke budgetten (b=64) behaalt MoE-nD 30% nauwkeurigheid, terwijl andere methoden in de buurt van 0% zakken.
• Null-resultaten (MATH-500 en TREC):
  • Bij korte prompts (MATH-500) of specifieke taken (TREC) levert de methode geen winst op. Dit komt omdat het budget niet strak genoeg is om gedwongen te worden tot agressieve eviction; de solver kiest dan voor keep=1.0 op de meeste lagen, waardoor de heterogene routing degenereren tot een uniforme methode. Dit bevestigt dat de methode specifiek nuttig is wanneer er echte keuzes gemaakt moeten worden.

5. Betekenis en Conclusie

MoE-nD demonstreert dat heterogeniteit in compressie-strategieën essentieel is voor efficiënte lange-context inferentie. Door te stoppen met het behandelen van alle lagen als identiek en in plaats daarvan een "mixture-of-experts" aanpak te gebruiken, kan men het geheugenbudget veel intelligenter verdelen.

De belangrijkste conclusie is dat per-laag eviction routing de kritieke innovatie is, niet alleen quantization. De methode biedt een template voor toekomstige compressietechnieken die meerdere orthogonale assen combineren via een gekalibreerde router, waardoor het mogelijk wordt om LLM's met lange contexten op veel beperktere hardware te draaien zonder kwaliteitsverlies.

Beperkingen:

• De huidige implementatie is niet compatibel met FlashAttention (gebruikt "eager attention"), wat de piekgeheugennod tijdens de prefill-fase verhoogt.
• De calibratie is specifiek voor het DeepSeek-R1-model; andere model-families vereisen opnieuw calibreren.
• De resultaten zijn gebaseerd op een "reasoning-model" protocol, wat de absolute F1-scores beïnvloedt, maar de relatieve verbeteringen tussen methoden blijven geldig.