Riepilogo Tecnico: DistillCache

Problematica

I modelli linguistici di grandi dimensioni (LLM) basati su Transformer si affidano alle cache Key-Value (KV) per evitare la ricalcolazione delle rappresentazioni dei token durante la decodifica autoregressiva. Tuttavia, la dimensione della cache KV cresce linearmente con la lunghezza della sequenza, creando un severo collo di bottiglia della memoria che limita la finestra di contesto e il throughput dell'inferenza a lungo contesto.

Le soluzioni esistenti si basano principalmente su strategie di espulsione euristiche (ad es., H2O, SnapKV) che mantengono i token in base a statistiche di attenzione statiche o segnali posizionali. L'autore sostiene che queste euristiche spesso non riescono a catturare l'influenza predittiva futura di un token, in particolare in scenari di dipendenze a lungo raggio o ragionamento strutturato, dove l'importanza di un token è contingente sulle conseguenze della generazione a valle piuttosto che sui pattern di attenzione passati.

Metodologia: DistillCache

Il paper propone DistillCache, un framework di apprendimento per rinforzo (RL) che formula l'espulsione della cache KV come un problema di decision-making sequenziale sotto rigidi budget di memoria. A differenza dei metodi euristici, DistillCache apprende una policy per mantenere un sottoinsieme di token che preservi la distribuzione dell'output di un modello con cache completa.

Componenti Core

1. Spazio degli Stati (Rappresentazione delle Feature):
   Per ogni token nella cache, la policy osserva un vettore di feature a 5 dimensioni derivato dai segnali interni del modello disponibili durante la decodifica:

   • $\alpha_i$: Massa di attenzione media ricevuta dal token attraverso tutti gli head e i layer.
   • $\|v_i\|_2$: La norma $\ell_2$ del vettore valore del token.
   • $\sigma_i$: Varianza della massa di attenzione ricevuta attraverso gli head.
   • $H_i$: Entropia media pesata dall'attenzione delle query recenti che accedono al token (cattura l'importanza sensibile al contesto).
   • $pos_i$: Posizione nella sequenza normalizzata.

2. Spazio delle Azioni:
   Una rete policy MLP (Multi-Layer Perceptron) leggera assegna un punteggio scalare a ogni token in cache. Durante l'addestramento, viene indotta una classificazione stocastica aggiungendo rumore Gumbel a questi punteggi, permettendo alla policy di campionare un sottoinsieme di top-$B$ token (dove $B$ è il budget) dalla distribuzione Plackett–Luce. In fase di inferenza, i top-$B$ token vengono selezionati deterministicamente.

3. Architettura di Addestramento (Setup Teacher-Student):
   Per addestrare la policy senza etichette di ground-truth, l'autore impiega un approccio di distillazione teacher-student:

   • Teacher Cache: Mantiene la storia completa e non potata dei token per calcolare la distribuzione "gold" del prossimo token ($P_t$).
   • Student Cache: Mantiene la storia potata secondo le decisioni della policy per calcolare la distribuzione potata ($\hat{P}_t$).
   • Segnale di Reward: Il reward ad ogni step è definito come la divergenza di Kullback-Leibler (KL) negativa tra le distribuzioni teacher e student: $R_t = -D_{KL}(P_t \parallel \hat{P}_t)$. Questo obiettivo ottimizza direttamente la fedeltà distributiva a ogni step di decodifica, invece di fare affidamento su segnali proxy come i punteggi di attenzione.

4. Ottimizzazione:
   La policy è ottimizzata utilizzando l'algoritmo REINFORCE con regolarizzazione dell'entropia e un baseline di riduzione della varianza (Leave-One-Out). Il processo di addestramento avviene offline, raddoppiando temporaneamente l'impronta di memoria, ma l'inferenza richiede solo un singolo passaggio in avanti con un overhead trascurabile per l'estrazione delle feature e il punteggio della policy.

Contributi Chiave

• Framework RL per l'Espulsione della Cache: Introduzione di DistillCache, un framework RL general-purpose che ottimizza la ritenzione dei token per la fedeltà distributiva a valle sotto vincoli di memoria, andando oltre la classificazione euristica statica.
• Obiettivo di Preservazione della Distribuzione: Dimostrazione che l'ottimizzazione della divergenza KL per step produce guadagni misurabili rispetto alle alternative euristiche e supervisionate, imparando esplicitamente i trade-off a lungo termine.
• Valutazione Esaustiva: Valutazione estesa su Mistral-7B-Instruct-v0.3 e trasferimento zero-shot a Llama-3-8B-Instruct attraverso LongBench, benchmark di ragionamento (GSM8K, BBH) e metriche di throughput.

Risultati Sperimentali

Valutato su Mistral-7B-Instruct-v0.3 con un budget della cache KV del 25%:

• Performance su LongBench: DistillCache mantiene il 94,2% dell'accuratezza della cache completa, superando i forti baselines euristici (H2O, SnapKV) fino a 2,7 punti assoluti e i metodi RL concorrenti (ForesightKV, RLKV) fino a 1,4 punti.
• Task di Ragionamento: Nei benchmark di ragionamento, DistillCache è competitivo con i migliori metodi concorrenti e li supera sotto compressione aggressiva (budget del 12,5%).
• Trasferimento Cross-Model: La policy addestrata su Mistral-7B si trasferisce zero-shot a Llama-3-8B-Instruct, mantenendo la competitività con i baselines euristici senza richiedere il ri-addestramento per modello (un requisito richiesto per i baselines RL concorrenti).
• Efficienza: Il metodo garantisce un throughput fino a 2,1× rispetto alla cache completa. Sebbene introduca un piccolo overhead (~1,4 ms/token) per l'estrazione delle feature e l'evaluazione della policy, rimane comparabile ad altri metodi appresi e significativamente più veloce dell'inferenza a cache completa.
• Analisi degli Errori: DistillCache mostra i tassi più bassi di errori ripetitivi e risposte errate rispetto alle euristiche e agli altri metodi RL, suggerendo una superiore preservazione della coerenza della sequenza e dell'accuratezza fattuale.

Significato e Rivendicazioni

Il paper sostiene che DistillCache dimostra l'efficacia di policy apprese e consapevole della distribuzione per l'inferenza di LLM a lungo contesto con efficienza di memoria. Ottimizzando le decisioni di ritenzione dei token rispetto alla futura fedeltà distributiva anziché alle statistiche di attenzione retrospettive, il framework impara policy vincolate dal budget che tengono conto delle conseguenze della generazione a valle.

L'autore sottolinea che DistillCache non è semplicemente una proxy appresa per metriche euristiche, ma un approccio distinto che ottimizza direttamente la preservazione della distribuzione di output del modello. Pur riconoscendo i limiti come i costi di addestramento (circa 130 ore GPU) e un leggero gap di performance con RLKV nei task di ragionamento con budget moderati, il paper conclude che l'espulsione adattiva che preserva la distribuzione è una direzione promettente per abilitare l'inferenza a lungo contesto su hardware con risorse limitate. Il lavoro suggerisce che combinare la fedeltà distributiva a livello di token con strategie di gating a livello di head (come esplorato negli esperimenti ibridi) offre un percorso complementare per l'ottimizzazione futura.