KVSculpt: KV Cache Compression as Distillation

1. Il Problema

Nella generazione autoregressiva di modelli linguistici su larga scala (LLM), la KV Cache (Key-Value Cache) diventa il collo di bottiglia principale per la memoria durante l'inferenza di contesti lunghi. Per mantenere l'efficienza, è necessario comprimere questa cache.
Le approcci esistenti per ridurre la lunghezza della sequenza della cache si dividono in due categorie principali, entrambe ancorate alle voci originali della cache:

• Eviction (Eliminazione): Seleziona un sottoinsieme discreto di coppie KV da mantenere (basato su punteggi di attenzione, recenza, ecc.) e scarta il resto.
• Merging (Fusione): Combina coppie KV simili in un numero minore, ma le nuove voci rimangono combinazioni lineari pesate delle originali.

Il limite fondamentale di questi metodi è che la cache compressa è vincolata a essere un sottoinsieme o una combinazione delle voci originali. Il paper sostiene che questo vincolo discreto è eccessivamente restrittivo, poiché dopo l'encoding RoPE (Rotary Positional Embedding), le coppie KV sono vettori in uno spazio continuo senza un ordinamento intrinseco.

2. Metodologia: KVSCULPT

KVSCULPT riformula il problema della compressione della KV Cache come un problema di distillazione in uno spazio continuo. Invece di selezionare o fondere voci esistenti, il metodo ottimizza un insieme più piccolo di coppie K e V non vincolate (libere) nello spazio di embedding per preservare il comportamento dell'attenzione di ogni strato.

Formulazione del Problema

Dato un contesto di $N$ token, la cache viene divisa in:

• Zona di ritenzione (Retain Zone): Gli ultimi $m$ token, mantenuti invariati.
• Zona di compressione (Compress Zone): I primi $N-m$ token, che devono essere distillati in $k$ coppie $(K_c, V_c)$ ottimizzate liberamente.

L'obiettivo è trovare $K_c \in \mathbb{R}^{k \times d}$ e $V_c \in \mathbb{R}^{k \times d}$ tali che l'output dell'attenzione calcolato con la cache compressa corrisponda il più possibile a quello della cache completa per qualsiasi futura query.

Funzione di Perdita e Ottimizzazione

La funzione di perdita combina due termini:

1. Errore MSE sull'output: Minimizzare la differenza tra l'output dell'attenzione compresso e quello completo.
2. Matching LSE (Log-Sum-Exp): Garantire che la massa di attenzione assegnata al chunk di contesto sia corretta, cruciale per la combinazione con i token futuri.

L'ottimizzazione avviene in modo alternato:

• Ottimizzazione delle Chiavi (K): Utilizza L-BFGS (un metodo quasi-Newton), essenziale per navigare il paesaggio non convesso della funzione softmax. I metodi del primo ordine come Adam falliscono qui, rimanendo intrappolati in minimi locali.
• Risoluzione dei Valori (V): Dati i pesi di attenzione fissi dalle chiavi correnti, i valori ottimali sono risolti analiticamente tramite regressione ridge (minimi quadrati).

Costruzione delle Query di Addestramento

Poiché le query future non sono disponibili al momento della compressione, KVSCULPT utilizza:

• Le query reali della zona di ritenzione (per ancorarsi ai pattern di attenzione attuali).
• Query sintetiche future: Vettori di contenuto campionati uniformemente dal contesto, a cui viene riapplicato l'encoding RoPE per posizioni future. Questo approccio "Uniform Spread" si è dimostrato superiore rispetto a strategie più complesse (come l'estrapolazione PCA) perché la deriva distributiva è troppo sottile per essere sfruttata in modo affidabile.

Allocazione Adattiva del Budget

Il metodo introduce un meccanismo di allocazione dinamica del budget di compressione ($k$ coppie) tra i diversi strati e testine KV:

1. Pilot Run: Una breve esecuzione di compressione (pilot) con allocazione uniforme rivela la difficoltà di compressione di ogni componente (misurata tramite MSE).
2. Ridistribuzione: Il budget totale viene ridistribuito in proporzione alla radice quadrata dell'MSE ($MSE^{0.5}$). Gli strati o le testine più difficili da comprimere ricevono più coppie, mentre quelli facili ne ricevono meno.
3. Damping: L'uso della radice quadrata è critico per evitare che strati con MSE estremamente alti (outlier) consumino tutto il budget, lasciando gli altri strati senza risorse.

3. Contributi Chiave

1. Riformulazione da Selezione a Distillazione: Spostamento dal vincolo discreto (selezione di indici) all'ottimizzazione continua nello spazio $\mathbb{R}^d$, eliminando la ricerca combinatoria sulle posizioni.
2. Ottimizzatore Alternato L-BFGS + Least Squares: Un approccio ibrido che raggiunge una divergenza KL inferiore di 3.5–4.1 volte rispetto alle migliori baseline di eviction.
3. Allocazione Adattiva Pilot-based: Una strategia che riduce ulteriormente la KL del 1.3 volte senza costi aggiuntivi durante l'inferenza, sfruttando l'eterogeneità della difficoltà di compressione.
4. Analisi della Difficoltà: Dimostrazione che la difficoltà di compressione non è uniforme, variando fino a 100 volte tra gli strati e fino a 467 volte tra le testine KV dello stesso strato.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Qwen2.5-1.5B-Instruct con contesti di 2048 token.

• Prestazioni Superiori: Rispetto alla baseline Select+Fit (evizione basata sul punteggio di attenzione + fitting dei valori), KVSCULPT riduce la divergenza KL di un fattore 3.5x – 4.1x a diversi rapporti di compressione ($r \in \{0.3, 0.5, 0.7\}$).
• Compressione Near-Lossless: Su sequenze "facili", KVSCULPT raggiunge una KL < 0.01, mentre le baseline rimangono sopra 0.1.
• Ruolo delle Chiavi Continue: L'ottimizzazione congiunta di chiavi e valori vincolate alle posizioni originali (Joint Optimization) offre un miglioramento minimo rispetto all'evizione. Il vero guadagno deriva dalla libertà di muovere le chiavi nello spazio continuo, permettendo loro di "sintetizzare" l'effetto combinato di più token originali.
• Ottimizzatore: Sostituire L-BFGS con Adam degrada le prestazioni di 17–95 volte, confermando la necessità di metodi di secondo ordine per il paesaggio di perdita dell'attenzione.
• Allocazione Adattiva: L'uso del segnale pilot-MSE riduce la KL media del 25% rispetto all'allocazione uniforme. L'allocazione per testina (head-level) aggiunge un ulteriore miglioramento, specialmente su sequenze difficili.

5. Significato e Limiti

Significato:
KVSCULPT dimostra che la compressione della KV Cache non deve essere limitata alla selezione di token esistenti. Trattando le chiavi come vettori liberi, è possibile creare una rappresentazione sintetica che cattura meglio il comportamento di attenzione. Inoltre, l'analisi rivela che il collo di bottiglia principale non è la compressione per singolo strato (che è quasi risolta), ma la propagazione degli errori tra gli strati, che si amplifica di ordini di grandezza.

Limiti:

• Costo di Compressione: Il processo richiede circa 170 secondi su una GPU A100 per un contesto di 2048 token. È adatto per scenari offline o amortizzati (es. documenti lunghi consultati molte volte), ma non per la decodifica online in tempo reale.
• Validazione: Gli esperimenti sono stati eseguiti su un solo modello (Qwen2.5-1.5B) e una sola lunghezza di contesto (2048).
• Valutazione: La metrica principale è la divergenza KL sui logit, non la precisione su task downstream specifici (es. MMLU).
• Mancanza di Baseline di Fusione: Non è stato confrontato direttamente con metodi di merging (come CaM o D2O), sebbene la formulazione teorica li subsumi.

In sintesi, KVSCULPT rappresenta un cambio di paradigma significativo, passando dalla selezione discreta all'ottimizzazione continua, offrendo un approccio promettente per l'inferenza di contesti lunghi quando il costo di pre-elaborazione è accettabile.