Riepilogo Tecnico: SEKV – KV Cache Adattiva alla Risoluzione con Memoria Semantica Gerarchica

Problema

I Large Language Models (LLM) sono sempre più impiegati in scenari che richiedono il ragionamento su contesti lunghi (ad esempio, analisi di documenti, dialogo multi-turno, pianificazione agentica). Tuttavia, la cache Key-Value (KV), che memorizza gli stati intermedi per accelerare la decodifica, cresce linearmente con la lunghezza della sequenza. Mantenere una cache KV completa per contesti che superano i 128K token diventa proibitivo in termini di memoria GPU, rendendo spesso l'inferenza a lungo contesto impraticabile senza compressione.

Le strategie di compressione esistenti affrontano un compromesso fondamentale tra efficienza e fedeltà:

1. Evizione dei Token: I metodi che scartano permanentemente i token basandosi sulle statistiche di attenzione (ad esempio, H2O, SnapKV) rischiano di frammentare evidenze coerenti o di perdere irrimediabilmente informazioni che diventano critiche solo più tardi nel processo di generazione.
2. Raggruppamento Semantico Statico: I metodi che comprimono i token in chunk semantici o frasi al momento del prefill (ad esempio, ChunkKV, SentenceKV) fissano la risoluzione della compressione precocemente. Essi mancano della capacità di recuperare dettagli a grana fine, a livello di token, da uno span compresso una volta che tale span diventa rilevante durante la decodifica.

La sfida centrale è progettare un'organizzazione della memoria che preservi l'accesso alle evidenze rilevanti sotto rigidi vincoli di memoria senza perdita irreversibile di informazioni, permettendo al modello di regolare dinamicamente la risoluzione del contesto memorizzato in base alla rilevanza della query.

Metodologia: SEKV

Gli autori propongono SEKV (Semantic KV Cache), una cache KV semantica adattiva alla risoluzione che organizza il contesto in una memoria semantica gerarchica che si estende tra GPU e CPU. SEKV evita l'evizione irreversibile dei token memorizzando il contesto a risoluzioni variabili e effettuando un "zoom in" dinamico al livello di dettaglio appropriato durante la decodifica.

1. Segmentazione degli Span Guidata dall'Entropia

Invece di una suddivisione in chunk a dimensione fissa, SEKV segmenta il contesto di input in span semanticamente coerenti utilizzando la sorpresa del token (token surprisal) come segnale di confine a costo zero.

• Meccanismo: Durante il passaggio forward di prefill, la sorpresa $H_t = -\log p(x_t | x_{<t})$ viene calcolata per ogni token. I token ad alta sorpresa, che spesso segnano transizioni semantiche (cambi di argomento, introduzione di entità), vengono identificati come confini dello span.
• Token Ancora (Anchor Tokens): Questi token di confine ad alta sorpresa sono mantenuti a piena risoluzione sulla GPU come "token ancora". Essi trasportano informazioni semantiche concentrate, essenziali per il successivo routing e ricostruzione.
• Span: I token tra due anchor vicini formano uno span semantico $S$.

2. Rappresentazione dello Span a Doppia Risoluzione e Gerarchia della Memoria

Ogni span semantico è rappresentato simultaneamente a due risoluzioni attraverso una gerarchia GPU–CPU:

• GPU (Routing Grossolano): Ogni span è rappresentato da un vettore di riepilogo leggero (una media pesata sulla sorpresa dei vettori di proiezione dei key) che è un vettore normalizzato. Questo riepilogo risiede sulla GPU ed è utilizzato per un routing grossolano ed efficiente per determinare quali span sono rilevanti per la query corrente.
• CPU (Ricostruzione a Grana Fine): La struttura completa a livello di token dello span è codificata in modo compatto sulla CPU utilizzando una base SVD a basso rango. Un predittore leggero assegna gate soft per componente ai valori singolari, consentendo al sistema di mantenere adattivamente un budget di rango $r_S$ per ogni span in base alla sua densità informativa. Ciò consente una fedele ricostruzione on-demand delle coppie KV a livello di token senza dover memorizzare l'intera matrice sulla GPU.

3. Meccanismo di Zoom-In Addestrato

Durante la decodifica, SEKV impiega un meccanismo appreso per selezionare dinamicamente quali span espandere:

• Routing: La query corrente viene proiettata nello stesso spazio di routing dei riepiloghi degli span. Un gate di rilevanza sigmoidale calcola la probabilità che uno span sia rilevante.
• Espansione: Una soglia appresa, per testa e per layer, determina se uno span deve essere espanso. Se selezionato, il sistema recupera in modo asincrono la corrispondente base SVD a basso rango dalla CPU, ricostruisce le coppie KV a livello di token e le integra nel calcolo dell'attenzione.
• Attenzione a Risoluzione Mista: Il meccanismo di attenzione opera su una memoria a risoluzione mista: i token ancora contribuiscono con coppie KV complete, gli span non espansi contribuiscono con i loro vettori di riepilogo grossolani (con una correzione della dimensione), e gli span espansi contribuiscono con le coppie di token ricostruite.
• Addestramento: Il modello LLM di base rimane congelato. Solo le proiezioni di routing leggere, le soglie di zoom-in e il predittore del gate del rango vengono addestrati utilizzando una combinazione di perdita di distillazione (allineamento con un insegnante a KV completo), supervisione dello zoom, perdita di ricostruzione e una penalità sul budget.

Contributi Chiave

1. Architettura SEKV: Una cache KV semantica adattiva alla risoluzione che memorizza il contesto a molteplici risoluzioni e regola dinamicamente il dettaglio della cache in base alla rilevanza della query, evitando l'evizione irreversibile dei token.
2. Segmentazione Guidata dall'Entropia: Un metodo che utilizza la sorpresa del token per identificare span semantici coerenti e preservare i token ancora ad alta sorpresa per un recupero affidabile.
3. Struttura a Doppia Rappresentazione: Una nuova struttura KV che combina riepiloghi degli span residenti su GPU per il routing grossolano, basi a basso rango residenti su CPU per l'approssimazione della ricostruzione a livello di token, e un meccanismo di zoom-in addestrato che espande selettivamente solo gli span rilevanti per la query.

Risultati Sperimentali

SEKV è stato valutato su quattro benchmark (LongBench, RULER, InfiniteBench, Needle-in-a-Haystack) e cinque modelli di backbone (inclusi varianti di Llama-3, Mistral e Qwen).

• Performance: SEKV supera la più forte baseline di compressione semantica (SentenceKV) di una media del 5,9% in tutti i benchmark e i modelli. Ottiene i migliori risultati in 20 casi su 20 nei confronti benchmark-modello.
• Fedeltà di Recupero: Nel benchmark Needle-in-a-Haystack (NIAH), SEKV mostra una stabilità di recupero superiore attraverso diverse lunghezze di contesto e profondità dell'ago rispetto ai metodi di evizione dei token e semantici statici.
• Efficienza della Memoria: A una lunghezza di contesto di 128K, SEKV riduce il consumo di memoria GPU del 53,3% rispetto alla cache KV completa, mantenendo al contempo prestazioni vicine al baseline non compresso.
• Efficienza: SEKV raggiunge una latenza e un throughput competitivi, spesso superando la cache KV completa in termini di latenza grazie alla riduzione della pressione sulla larghezza di banda della memoria, pur superando significativamente i metodi di evizione aggressiva in termini di accuratezza.
• Efficienza dei Parametri: Il metodo aggiunge meno dello 0,05% di parametri addestrabili (circa 4,3 milioni di parametri per Llama-3-8B) mantenendo il modello di base congelato.

Significato e Rivendicazioni

Il documento sostiene che un'inferenza efficiente su contesti lunghi richiede un passaggio dalla compressione statica al recupero dinamico della risoluzione. Gli autori affermano che, mentre la compressione statica preserva il contesto grossolano, essa fallisce nel recuperare evidenze a grana fine quando queste diventano rilevanti tardi nel processo di generazione. SEKV affronta questo problema organizzando il contesto in span semantici guidati dall'entropia e utilizzando un sistema di memoria gerarchica che permette al modello di effettuare lo "zoom in" sulle informazioni rilevanti su richiesta.

Il significato di questo lavoro risiede nel dimostrare che:

1. L'evizione irreversibile non è necessaria: È possibile ottenere un'inferenza ad alta fedeltà su contesti lunghi offloading rappresentazioni a basso rango recuperabili alla CPU invece di scartare i token.
2. La risoluzione adattiva è critica: L'allocazione della risoluzione a grana fine solo dove necessario (tramite il meccanismo di zoom-in) permette ai modelli di bilanciare i vincoli di memoria con la fedeltà di recupero in modo più efficace rispetto alle policy statiche.
3. La struttura semantica aiuta la compressione: L'uso della segmentazione guidata dall'entropia crea unità più comprimibili rispetto ai blocchi di token arbitrari, consentendo un'approssimazione a basso rango efficiente.

Gli autori riconoscono i limiti, notando che l'efficacia dipende dalla qualità della costruzione degli span semantici (che potrebbe avere difficoltà con discorsi altamente frammentati) e che l'approccio introduce una dipendenza dalla larghezza di banda CPU-GPU, che potrebbe influenzare la latenza in scenari di query avversaria. Tuttavia, concludono che SEKV fornisce un framework robusto per l'inferenza scalabile e ad alta fedeltà di LLM a lungo contesto.