Riepilogo Tecnico: RAP (RoPE-Aligned Pruning) per la Compressione del KV-Cache

1. Definizione del Problema

L'inferenza a lungo contesto nei Large Language Models (LLM) è sempre più limitata dai requisiti di memoria e dai costi computazionali della cache Key-Value (KV). Man mano che la lunghezza del contesto aumenta, il KV cache domina l'uso della memoria GPU e i costi di calcolo dell'attenzione. Sebbene il pruning strutturato offra una via diretta per comprimere la cache rimuovendo i canali meno utili dai pesi di proiezione $W_k$ e $W_v$, i metodi esistenti affrontano un'incompatibilità critica con l'Embedding di Posizione Rotatorio (RoPE).

I moderni LLM applicano il RoPE dopo le proiezioni Query/Key, ruotando le dimensioni delle feature a coppie. I metodi di pruning per canale standard (ad esempio, KVPruner) spesso rimuovono singoli canali senza considerare questa struttura a coppie. Questo "orfano" di semi-coppie rompe la semantica rotazionale del RoPE, portando a un collasso silenzioso ma catastrofico dell'accuratezza del modello. Inoltre, i metodi di decomposizione a basso rango esistenti (ad esempio, SVD-LLM, PaLU) che tentano di comprimere i pesi spesso non possono fondere le loro matrici di espansione con la proiezione Query ($W_q$) a causa della non-commutatività tra Roope e moltiplicazione matriciale, costringendo a una costosa ricostruzione per ogni step dei vettori Key a dimensionalità completa durante la decodifica.

2. Metodologia: RoPE-Aligned Pruning (RAP)

Gli autori propongono RAP, un framework di pruning strutturato progettato specificamente per rispettare la struttura a coppie imposta dal RoPE.

Principio Fondamentale

Invece di eliminare singoli canali, RAP vincola la granularità del pruning alle coppie allineate al RoPE. Rimuovendo intere coppie di dimensioni, il metodo assicura che l'operazione di rotazione rimanga matematicamente valida e che la semantica posizionale sia preservata.

La Pipeline RAP

1. Scoring delle Coppie RoPE:

   • Per i pesi di proiezione Key ($W_k$), le colonne sono raggruppate in coppie RoPE secondo la specifica strategia di accoppiamento del modello (sia adiacente che separata a metà).
   • L'importanza è misurata utilizzando punteggi di informazione di Fisher calcolati su un piccolo dataset di calibrazione. Per una coppia $p = (j, j')$, il punteggio $\sigma_p$ è l'aggregato della magnitudo della radice-Fisher dei due canali.
   • Per la proiezione Value ($W_v$), che non possiede l'accoppiamento RoPE, vengono calcolati i punteggi di Fisher dei singoli canali.

2. Allocazione del Budget:

   • Un rapporto di mantenimento globale $\rho$ viene distribuito tra i layer e le teste.
   • Adattamento al Layer: Ai layer vengono assegnati budget basati sulla loro sensibilità (punteggi di Fisher aggregati), assegnando un budget maggiore ai layer più sensibili.
   • Uniformità tra le Teste: All'interno di un layer, tutte le teste ricevono lo stesso numero di coppie mantenute per facilitare le operazioni GEMM batch.
   • Bilanciamento K/V: Per evitare che la dimensione K venga schiacciata (il che danneggerebbe le prestazioni a lungo contesto), RAP assegna lo stesso budget di canali per testa sia a $K$ che a $V$. Nello specifico, se $K$ mantiene $m$ coppie, $V$ mantiene $2m$ canali.

3. Costruzione e Fusione:

   • Viene costruita una matrice di selezione binaria, che preserva le coppie, $E_k$, per mappare le coppie mantenute nuovamente ai loro indici originali.
   • Fondamentalmente, poiché $E_k$ preserva la struttura delle coppie, essa commuta con l'operazione RoPE: $RoPE(X \tilde{W}_k E_k) = RoPE(X \tilde{W}_k) E_k$.
   • Questa commutatività permette di fondere la maschera di pruning offline nella proiezione Query: $\tilde{W}_q = W_q E_k^\top$.
   • Risultato: Il modello gira su dimensioni ridotte e il KV cache memorizza solo le chiavi latenti a bassa dimensionalità. Non è richiesta alcuna ricostruzione online.

4. Recupero dell'Accuratezza:

   • Per recuperare l'accuratezza persa durante il pruning, RAP impiega un aggiornamento LoRA (Low-Rank Adaptation) leggero combinato con la Distillazione della Conoscenza (KD).
   • Un adattatore LoRA con un piccolo rango ($r=8$) viene addestrato sul modello con pruning per allineare i suoi output con il modello teacher denso, per poi essere fuso prima del deployment.

3. Contributi Chiave

• Identificazione dell'Incompatibilità con RoPE: Il paper evidenzia come il pruning dei canali standard rompa le coppie RoPE, causando un collasso dell'accuratezza, e come questo fallimento sia spesso silenzioso (non genera errori di runtime, ma solo un degrado delle prestazioni).
• Strategia di Pruning Allineata al RoPE: RAP introduce una granularità di pruning che aderisce strettamente alle coppie RoPE, garantendo che l'operazione di rotazione rimanga valida senza richiedere un retrain complesso o ricostruzioni.
• Decodifica Senza Ricostruzione: A differenza dei metodi di decomposizione a basso rango che devono ricostruire le chiavi a dimensionalità completa ad ogni step di decodifica, RAP fonde la maschera di pruning in $W_q$ offline, eliminando la "tassa di ricostruzione" e riducendo sia l'overhead di memoria che di calcolo.
• Valutazione Esaustiva: Il metodo è valutato su molteplici famiglie di modelli (Llama, Mistral, Qwen) e scale (da 3B a 14B), dimostrando robustezza.

4. Risultati Sperimentali

Le valutazioni sono state condotte in un setting iso-KV (che corrisponde al rapporto di canali KV mantenuti $\rho$) per garantire un confronto equo di qualità ed efficienza.

• Accuratezza vs Compressione:
  • A un rapporto di compressione del 30% ($\rho=0.7$), RAP mantiene un'accuratezza utilizzabile (ad esempio, PPL di WikiText-2 ~9.12 per Llama-3-8B rispetto al baseline 8.28), superando significativamente il pruning per canale cieco rispetto al RoPE (KVPruner), che subisce un massiccio collasso dell'accuratezza (PPL ~397).
  • RAP raggiunge prestazioni comparabili o superiori ai forti baseline a basso rango (PaLU, SVD-LLM) operando con costi di attenzione inferiori.
• Guadagni di Efficienza:
  • Parametri e FLOPs: RAP riduce i parametri di attenzione e i FLOPs linearmente con il rapporto di mantenimento $\rho$ (ad esempio, ~70% a $\rho=0.7$). Al contrario, i metodi a basso rango mantengono spesso un numero di parametri più elevato a causa della necessità di fattori di espansione online.
  • Latenza: RAP riduce la latenza di Prefill e Decode (ad esempio, 83%/77% del baseline denso a $\rho=0.7$) perché evita la ricostruzione per ogni step. I metodi a basso rango mostrano spesso un aumento della latenza, particolarmente durante la decodifica.
• Robustezza:
  • RAP mantiene le prestazioni attraverso diversi modelli e architetture.
  • È ortogonale e composibile con la quantizzazione dei pesi a 4 bit.
  • Il metodo gestisce efficientemente l'indexing RoPE non contiguo tramite un kernel Triton personalizzato, evitando l'overhead dell'indexing standard di PyTorch.

5. Significato e Rivendicazioni

Il paper afferma che RAP è l'unico metodo che simultaneamente:

1. Riduce la dimensione del KV cache.
2. Riduce i parametri del modello e i FLOPs di proiezione.
3. Mantiene un'alta accuratezza ad alti rapporti di compressione.
4. Opera senza l'overhead di ricostruzione per ogni step.

Gli autori sottolineano che RAP è una soluzione "drop-in" per gli LLM basati su RoPE. Rispettando i vincoli strutturali dell'Embedding di Posizione Rotatorio, risolve il problema del "fallimento silenzioso" dei metodi di pruning esistenti e offre un'alternativa più efficiente alla decomposizione a basso rango per l'inferenza a lungo contesto. Il lavoro suggerisce che l'allineamento strutturale con i meccanismi specifici del modello (come RoPE) è critico per un'efficace compressione del modello.