Technische Samenvatting: RAP (RoPE-Aligned Pruning) voor KV-Cache Compressie

1. Probleemstelling

Inference met lange context in Large Language Models (LLMs) wordt steeds vaker beperkt door de geheugen- en rekenvereisten van de Key-Value (KV) cache. Naarmate de contextlengte schaalt, domineert de KV-cache het GPU-geheugengebruik en de kosten van de attention-berekening. Hoewel gestructureerde pruning een directe weg biedt om de cache te comprimeren door minder nuttige kanalen uit de $W_k$ en $W_v$ projectiegewichten te verwijderen, kampen bestaande methoden met een kritieke incompatibiliteit met Rotary Position Embedding (RoPE).

Moderne LLM's passen RoPE toe na de Query/Key-projecties, waarbij de dimensionale kenmerken in paren worden geroteerd. Standaard per-kanaal pruning-methoden (bijv. KVPruner) verwijderen vaak individuele kanalen zonder rekening te houden met deze paarstructuur. Dit "verweeslijken" van halve paren verbreekt de rotationele semantiek van RoPE, wat leidt tot een stille maar catastrofale instorting van de modelnauwkeurigheid. Bov�keren bestaande low-rank decompositie-methoden (bijv. SVD-LLM, PaLU), die proberen gewichten te comprimeren, vaak niet de expansiematrices te kunnen fuseren met de Query-projectie ($W_q$) vanwege de niet-commutativiteit van RoPE en matrixvermenigvuldiging, wat een kostbare reconstructie van de volledige dimensionale Key-vectoren tijdens elke decodeerstap tijdens het decoderen forceert.

2. Methodologie: RoPE-Aligned Pruning (RAP)

De auteurs stellen RAP voor, een gestructureerd pruning-framework dat specifiek is ontworpen om de paarwijze structuur opgelegd door RoPE te respecteren.

Kernprincipe

In plaats van individuele kanalen te prunen, beperkt RAP de granulariteit van de pruning tot RoPE-gealigneerde paren. Door hele paren dimensies te verwijderen, zorgt de methode ervoor dat de rotatie-operatie wiskundig geldig blijft en dat de positionele semantiek behouden blijft.

De RAP-pipeline

1. RoPE-Pair Scoring:

   • Voor de Key-projectiegewichten ($W_k$) worden kolommen gegroepeerd in RoPE-paren volgens de specifieke pairing-strategie van het model (ofwel aangrenzend ofwel half-gescheiden).
   • Belangrijkheid wordt gemeten met behulp van Fisher-informatie scores berekend op een kleine kalibratiedataset. Voor een paar $p = (j, j')$ is de score $\sigma_p$ de geaggregeerde root-Fisher magnitude van de twee kanalen.
   • Voor de Value-projectie ($W_v$), die geen RoPE-pairing heeft, worden individuele kanaal Fisher-scores berekend.

2. Budgetallocatie:

   • Een globale behoudratio $\rho$ wordt verdeeld over lagen en heads.
   • Laag-adaptief: Lagen krijgen budgetten toegewezen op basis van hun gevoeligheid (geaggregeerde Fisher-scores), waarbij gevoeligere lagen een groter budget krijgen.
   • Head-uniform: Binnen een laag ontvangen alle heads hetzelfde aantal behouden paren om gebatchte GEMM-operaties te faciliteren.
   • Gebalanceerde K/V: Om te voorkomen dat de Key-dimensie wordt verpletterd (wat de long-context prestaties schaadt), wijst RAP hetzelfde aantal kanalen per head toe aan zowel $K$ als $V$. Specifiek, als $K$ $m$ paren behoudt, behoudt $V$ $2m$ kanalen.

3. Constructie en Fusie:

   • Een binaire, paar-behoudende selectiematrix $E_k$ wordt geconstrueerd om de behouden paren terug te mappen naar hun originele indices.
   • Cruciaal is dat omdat $E_k$ de paarstructuur behoudt, het commuteert met de RoPE-operatie: $RoPE(X \tilde{W}_k E_k) = RoPE(X \tilde{W}_k) E_k$.
   • Deze commutativiteit maakt het mogelijk om de pruning-mask offline te fuseren in de Query-projectie: $\tilde{W}_q = W_q E_k^\top$.
   • Resultaat: Het model draait op geprunte dimensies, en de KV-cache slaat alleen de laag-dimensionale latente keys op. Er is geen online reconstructie vereist.

4. Nauwkeurigheid Herstel:

   • Om de nauwkeurigheid die verloren is gegaan tijdens de pruning te herstellen, gebruikt RAP een lichtgewicht LoRA (Low-Rank Adaptation) update gecombineerd met Knowledge Distillation (KD).
   • Een kleine rank ($r=8$) LoRA-adapter wordt getraind op het geprunte model om de outputs uit te lijnen met het dense teacher-model, waarna deze vóór implementatie wordt samengevoegd (merged).

3. Belangrijkste Bijdragen

• Identificatie van RoPE-incompatibiliteit: Het paper belicht dat standaard kanaal pruning de RoPE-paren verbreekt, wat leidt tot een nauwkeurigheidsinstorting, en dat dit falen vaak stil is (geen runtime errors, alleen gedegradeerde prestaties).
• RoPE-Aligned Pruning Strategie: RAP introduceert een pruning-granulariteit die strikt voldoet aan de RoPE-paren, waardoor de rotatie-operatie geldig blijft zonder complexe hertraining of reconstructie te vereisen.
• Reconstructie-vrije Decoding: In tegenstelling tot low-rank decompositie-methoden die bij elke decodeerstap de volledige dimensionale keys moeten reconstrueren, fuseert RAP het pruning-mask offline in $W_q$, wat de "reconstructie-belasting" elimineert en zowel het geheugen- als de rekenoverhead vermindert.
• Uitgebreide Evaluatie: De methode wordt geëvalueerd over meerdere modelfamilies (Llama, Mistral, Qwen) en schalen (3B tot 14B), wat de robuustheid aantoont.

4. Experimentele Resultaten

Evaluaties werden uitgevoerd onder een iso-KV setting (overeenkomend met de behouden KV-kanaalratio $\rho$) om een eerlijke vergelijking van kwaliteit en efficiëntie te garanderen.

• Nauwkeurigheid vs. Compressie:
  • Bij een compressieratio van 30% ($\rho=0.7$) behoudt RAP een bruikbare nauwkeurigheid (bijv. WikiText-2 PPL ~9.12 voor Llama-3-8B vs. baseline 8.28), wat aanzienlijk beter presteert dan de RoPE-blinde kanaal pruning (KVPruner), die een enorme nauwkeurigheidsinstorting vertoont (PPL ~397).
  • RAP bereikt prestaties die vergelijkbaar met of beter zijn dan sterke low-rank baselines (PaLU, SVD-LLM), terwijl het lagere attention-kosten heeft.
• Efficiëntiewinst:
  • Parameters & FLOPs: RAP vermindert de attention-parameters en FLOPs lineair met de behouden kanaalratio $\rho$ (bijv. ~70% bij $\rho=0.7$). In contrast hiermee behouden low-rank methoden vaak hogere parametercounts vanwege de noodzaak van online expansiefactoren.
  • Latency: RAP vermindert de Prefill en Decode latency (bijv. 83%/77% van de dense baseline bij $\rho=0.7$) omdat het per stap reconstructie vermijdt. Low-rank methoden vertonen vaak een verhoogde latency, met name tijdens decoding.
• Robuustheid:
  • RAP behoudt prestaties over verschillende modelgroottes en architecturen.
  • Het is orthogonaal aan en combineerbaar met 4-bit gewichts-kwantisatie.
  • De methode handelt niet-aangrenzende RoPE-indexering efficiënt af via een custom Triton kernel, waardoor de overhead van standaard PyTorch-indexering wordt vermeden.

5. Betekenis en Claims

Het paper claimt dat RAP de enige methode is die simultaan:

1. De KV-cache grootte vermindert.
2. Modelparameters en projectie FLOPs vermindert.
3. Hoge nauwkeurigheid handhaaft bij hoge compressieratio's.
4. Werkt zonder per-stap reconstructie-overhead.

De auteurs benadrukken dat RAP een "drop-in" oplossing is voor RoPE-gebaseerde LLM's. Door de structurele beperkingen van Rotary Position Embedding te respecteren, lost het het "stille falen" probleem van bestaande pruning-methoden op en biedt het een efficiënter alternatief voor low-rank decompositie voor long-context inference. Het werk suggereert dat structurele alignment met model-specifieke mechanismen (zoals RoPE) cruciaal is voor effectieve modelcompressie.