Riepilogo Tecnico: COMPILERKV

1. Enunciato del Problema

I Large Language Models (LLM) affrontano un collo di bottiglia nella memoria dovuto alla crescita lineare della cache Key-Value (KV) con la lunghezza del contesto. La compressione KV solo nella fase di prefill è emersa come una soluzione favorevole al deployment, in cui un sottoinsieme di token viene selezionato alla fine della fase di prefill e fissato per l'intera fase di decodifica. Questo approccio converte l'uso della memoria a runtime in un costo fisso con zero intervento al momento della decodifica.

Tuttavia, il paper identifica una fondamentale discrepanza statistica nei metodi esistenti solo nella fase di prefill (ad esempio, SnapKV, PyramidKV, HeadKV). Questi metodi tentano di stimare segnali correttivi—specificamente l'affidabilità per testa (quali attention head sono critici per il recupero) e la sensibilità alla compressione a livello di prompt (quanto un prompt specifico può essere compresso)—online, utilizzando statistiche da un singolo prompt rumoroso. Gli autori sostengono che questa sia l'unità statistica sbagliata:

• E1 (Rumore del singolo prompt): I punteggi di attenzione per prompt sono a code pesanti e distorti per scala; grandi magnitudini non garantiscono importanza semantica, e piccoli errori di stima causano errori di selezione irreversibili.
• E2 (Eterogeneità delle teste): Le attention head sono funzionalmente specializzate. Alcune sono "critiche per il recupero", mentre altre sono costantemente rumorose. Un'aggregazione ingenua permette alle teste rumorose di dominare i punteggi.
• E3 (Varianza del rischio a livello di prompt): Prompt diversi hanno limiti di compressione diversi. Una regola di ritenzione fissa fallisce su prompt ad alta entropia o ad alta perplessità, portando a una perdita catastrofica di informazioni.

I metodi esistenti ristimano questi segnali online per ogni prompt, fallendo nel sfruttare il fatto che questi segnali esibiscono un'alta regolarità cross-prompt piuttosto che stabilità all'interno del prompt.

2. Metodologia: COMPILERKV

Gli autori propongono COMPILERKV, una politica di ritenzione KV che compila tabelle correttive offline da un corpus di calibrazione, invece di stimarle online. L'idea centrale è che le classifiche di affidabilità delle teste e la mappatura dalle caratteristiche di rischio alle soglie sicure sono priori architetturali stabili che possono essere appresi una volta e riutilizzati.

Il framework consiste in tre componenti accoppiati che operano durante la fase di prefill:

A. Utilità del Token Stabilizzata (Affronta E1)

Per rimuovere picchi transitori di attenzione e distorsioni di scala inter-testa senza apprendimento, il metodo calcola un punteggio di utilità privo di parametri:

• Attenzione aggregata per finestra: Smussa i picchi transitori aggregando l'attenzione su una finestra di osservazione $\Omega$ alla coda del prompt.
• Normalizzazione relativa alla testa: Normalizza i punteggi dei token per la norma del valore medio della testa, rimuovendo la distorsione di scala.
• Formula: $u^{(l,h)}_t = \alpha_t \cdot \rho^{(l,h)}_t$, dove $\alpha_t$ è la massa della finestra e $\rho$ è la norma del valore normalizzato.

B. Tabella di Eterogeneità delle Teste (Affronta E2)

Invece di trattare tutte le teste allo stesso modo o di ristimare l'importanza online, COMPILERKV compila una tabella congelata $W_{head}[l, h]$ di pesi di affidabilità continui.

• Compilazione Offline: Appresa tramite Conservative Q-Learning (CQL) su un corpus di calibrazione. Lo stato è $(l, h, B_l)$, e l'azione è un peso di affidabilità discretizzato.
• Meccanismo: Durante l'inferenza, l'utilità di un token viene riscalata per l'affidabilità della sua testa. Crucialmente, l'aggregazione utilizza il weighted max-pooling ($\hat{u}^{(l)}_t = \max_h (u^{(l,h)}_t \cdot W_{head}[l, h])$). Questo conferisce un "potere di veto" alle teste critiche per il recupero: se una singola testa critica si focalizza fortemente su un token, esso viene mantenuto, indipendentemente dal consenso delle teste rumorose.

C. Porta Soglia Adattiva al Rischio (Affronta E3)

Per gestire prompt con limiti di compressione variabili, una Porta Soglia Adattiva al Rischio mappa le caratteristiche di rischio a livello di prompt a una soglia di ritenzione.

• Segnali di Rischio: Utilizza l'Entropia di Attenzione (rischio strutturale) e la Perplessità Locale (rischio semantico) calcolate sulla finestra di osservazione.
• Compilazione Offline: Una tabella di ricerca (LUT) mappa i bin di rischio discretizzati $(b_{ent}, b_{ppl})$ a una soglia di ritenzione $\tau^{(l)}$. Anche questo viene compilato tramite CQL.
• Bias Conservativo: L'obiettivo CQL penalizza le soglie aggressive su stati rari ad alto rischio (alta entropia/perplessità), orientando la tabella verso una ritenzione più sicura dove euristica online tipicamente fallisce.
• Selezione Elastica: La selezione finale mantiene tutti i token con utilità superiore alla soglia $\tau^{(l)}$. Se l'insieme dei candidati supera il budget $B_l$, viene limitato ai primi $B_l$ token. Se è inferiore, tutti i candidati vengono mantenuti (sottoritenzione elastica).

D. Fondamento Teorico

Il paper fornisce un'analisi formale che mostra come l'errore di approssimazione sia controllato dalla massa della coda scartata (Proposizione 4.1). Sostiene che l'errore di calibrazione della tabella finita è controllato a livello di prompt (tra prompt indipendenti) piuttosto che assumendo l'indipendenza all'interno delle righe di attenzione di un singolo prompt, riducendo così la varianza.

3. Contributi Chiave

1. Riformulazione del Problema: Identifica che l'unità statistica per la correzione nella compressione solo nella fase di prefill dovrebbe essere la scala cross-prompt, non il singolo prompt.
2. Compilazione Offline: Introduce un metodo per compilare l'affidabilità delle teste e le soglie adattive al rischio offline, trasformandole in priori architetturali portatili.
3. Politica Accoppiata: Propone un'ottimizzazione congiunta in cui i pesi di affidabilità delle teste e le soglie di rischio sono calibrati contro la stessa ricompensa a valle, invece di essere moduli indipendenti.
4. Portabilità: Dimostra che le tabelle compilate si trasferiscono tra corpora disgiunti e persino tra diverse famiglie di modelli (ad esempio, da LLaMA a Mistral) con un degrado delle prestazioni minimo.

4. Risultati Sperimentali

Il metodo è stato valutato su quattro backbones da 7B a 8B (LLaMA-3, Mistral, Qwen2, InternLM) attraverso i benchmark LongBench, RULER e Needle-in-a-Haystack (NIAH).

• LongBench (budget di 512 token): COMPILERKV raggiunge un compressed-SOTA su tutti e quattro i backbones, superando la baseline più forte (ad esempio, ChunkKV, HeadKV) di una media di +1.67 punti (IC 95% [+1.08, +2.37]). Su alcuni backbones, supera leggermente le prestazioni di FullKV, cosa che gli autori attribuiscono alla rimozione del rumore piuttosto che al guadagno di informazioni.
• Regimi di Pressione:
  • RULER (128k contesto): Con un rapporto fisso di cache 512/32k, COMPILERKV mantiene un punteggio di 73, superando significativamente SnapKV (38) e KVZip (66.5), restando al di sotto di FullKV (79).
  • NIAH (32k): Raggiunge un'accuratezza di 0.89 contro 0.42 di SnapKV.
  • Batching: A 32k di input con una dimensione del batch di 16, COMPILERKV rimane fattibile, mentre FullKV esaurisce la memoria (OOM).
• Trasferibilità:
  • Cross-Corpus: Le classifiche di affidabilità delle teste mostrano un'alta correlazione di Spearman ($\bar{\rho} \approx 0.90$) tra corpora di calibrazione disgiunti.
  • Cross-Model: Il trasferimento diretto della tabella tra modelli (ad esempio, da LLaMA-3 a Mistral) comporta una perdita di soli 0.4–0.8 punti LongBench rispetto alla compilazione specifica per il target.
• Efficienza: L'overhead online è minimo: $O(1)$ ricerche in tabella per layer e un singolo ordinamento sul pool dei candidati. Raggiunge fino a 2.45× il throughput di decodifica rispetto a FullKV a 32k di contesto.

5. Significato e Affermazioni

Il paper afferma che COMPILERKV colma il divario con FullKV a budget fissi correggendo la discrepanza statistica nel modo in cui vengono calcolate le politiche di ritenzione. Spostando il calcolo dei segnali correttivi da una stima online, rumorosa e a singolo prompt a un priori cross-prompt offline e compilato, il metodo stabilizza il processo di selezione contro rumore ed eterogeneità.

Gli autori sottolineano che la novità non è semplicemente la "consapevolezza delle teste" (che metodi precedenti come HeadKV possiedono), ma la compilazione offline congiunta di affidabilità continua delle teste e soglie di rischio del prompt. Questo trasforma l'affidabilità delle teste da un trucco di sintonizzazione in un oggetto scientifico riutilizzabile che può essere compilato una volta e distribuito su modelli e contesti con un overhead online di $O(1)$. Il lavoro suggerisce che per la compressione solo nella fase di prefill, i segnali "giusti" sono proprietà architetturali del modello, non caratteristiche transitorie del prompt di input.