Riepilogo Tecnico: Still: Compattazione Ammortizzata della KV Cache in un Singolo Pass Forward

1. Definizione del Problema

La cache KV (Key-Value) è il principale collo di bottiglia della memoria per l'implementazione di modelli linguari di grandi dimensioni (LLM) su orizzonti prolungati, come agenti di coding multi-giorno o ragionamento su scala di repository. Le attuali strategie di implementazione si trovano di fronte a una scelta "tutto o niente":

• Caching Lossless (Senza Perdite): La cache cresce linearmente con la lunghezza del contesto, finendo per colpire i limiti di memoria.
• Alternative Lossy (Con Perdita): Metodi come il fine-tuning, la generazione aumentata da recupero (RAG) o i riassunti dei documenti abbandonano la struttura e la fedeltà della rappresentazione interna del modello.

I metodi esistenti di compattazione della cache KV rientrano in due categorie, ciascuna con significative limitazioni:

1. Metodi di Selezione: Questi assegnano un punteggio ai token ed eliminano un sottoinsieme (es. H2O, SnapKV). Sono leggeri ma strutturalmente vincolati al sottoinsieme originale di token, il che significa che non possono sintetizzare nuova informazione e degradano rapidamente sotto budget di compressione ristretti.
2. Metodi di Sintesi: Questi combinano il contenuto per creare nuove rappresentazioni compatte (es. Attention Matching, Cartridges). Sebbene espressivi, tipicamente richiedono un'ottimizzazione per-contesto (fitting) durante l'inferenza, rendendoli computazionalmente costosi e impraticabili per traiettorie iterative a lungo orizzonte dove la compattazione deve avvenire ripetutamente.

Il documento identifica una lacuna per un metodo che sia ammortizzato (addestrato una volta, applicato ripetutamente senza ottimizzazione per-contesto), sintetico (capace di fondere le informazioni oltre la semplice selezione) e abbastanza leggero da poter essere eseguito durante l'inferenza.

2. Metodologia: Still

Gli autori introducono Still, un primitivo appreso per la compattazione della cache KV progettato per modelli pre-addestrati congelati.

Architettura

Still consiste in un piccolo modulo Perceiver per layer, addestrato contro un modello base congelato.

• Input: Per ogni layer del transformer, il compattatore riceve l'intera cache KV ($K, V$) del prefisso per quel layer. Le chiavi vengono prima "de-ruotate" per rimuovere gli Rotary Position Embeddings (RoPE), posizionandole in un frame indipendente dalla posizione.
• Latent Queries: Il compattatore mantiene un banco di query latenti apprese ($Z$) per layer. Queste query effettuano una cross-attention sulla cache KV completa.
• Elaborazione: I latenti passano attraverso blocchi pre-norm contenenti:
  1. Cross-Attention: Dai latenti alla cache KV completa.
  2. Self-Attention: Tra i latenti stessi.
  3. Feed-Forward (FFN): Sottolayer opzionale.
• Output: Lo stato latente finale è proiettato tramite teste lineari per produrre chiavi compatte ($C_k$) e valori ($C_v$) di dimensione $t \times d$, dove $t \ll T$ (la lunghezza della sequenza originale).
• Gestione della Posizione: Le chiavi compatte risultanti vengono de nuovo ruotate nelle posizioni di output scelte prima di essere scritte nuovamente nella cache. I valori rimangono non ruotati.
• Canale di Bias: È supportato un opzionale canale di bias ($\beta$), ma disabilitato ($\beta=0$) in tutti i risultati riportati per semplificare l'implementazione.

Addestramento

• Obiettivo: Il compattatore è addestrato utilizzando la divergenza KL forward tra un "insegnante" (il modello base congelato con l'intera cache non compressa) e uno "studente" (lo stesso modello utilizzando la cache compatta). La perdita è mascherata solo sui token lato risposta.
• Dati: L'addestramento utilizza un dataset di domande a scelta multipla estrattive in quattro domini (Finanza, Project Gutenberg, Legale, Codice).
• Efficienza: I pesi del modello base rimangono congelati; vengono addestrati solo i parametri del compattatore. Il compattatore aggiunge circa l'1% al conteggio dei parametri del modello base (es. ~50M di parametri per Qwen3-4B).

Modalità Operative

1. Compattazione Single-Pass: Il modello esegue un prefill completo per generare l'intera cache, che viene poi compressa da Still. La generazione continua utilizzando la cache compatta.
2. Compattazione Iterativa a Blocchi (Chunked): Per contesti a lungo orizzonte che arrivano nel tempo, Still supporta uno schema ricorrente. Il modello elabora blocchi di token e, dopo ogni blocco, la cache accumulata viene compressa a una dimensione fissa. Ciò consente al sistema di gestire contesti che superano di gran lunga l'orizzonte di addestramento, a patto che il compattatore sia stato addestrato su orizzonti iterativi sufficientemente lunghi.

3. Contributi Chiave

• Sintesi Ammortizzata: Still dimostra che la sintesi ammortizzata (apprendere un encoder una volta sola) è superiore alla selezione ammortizzata (apprendere un valutatore una volta sola) per la compattazione di contesti lunghi. Evita il limite del "sottoinsieme vincolato" dei metodi di selezione e il costo di "ottimizzazione per-contesto" dei metodi di sintesi.
• Singolo Pass Forward: A differenza dei metodi di sintesi per-contesto che richiedono passi di ottimizzazione per ogni nuovo contesto, Still esegue la compattazione in un singolo passaggio forward attraverso un modulo appreso.
• Capacità Iterativa a Lungo Orizzonte: Poiché la compattazione è un passaggio forward, Still può essere applicata iterativamente man mano che arriva il nuovo contesto, abilitando un regime a lungo orizzonte non disponibile per i metodi di fitting per-contesto.
• Generalizzazione: Il metodo si trasferisce attraverso diverse scale di modelli (4B a 32B), architetture (dense e MoE) e meccanismi di attenzione (globale e sliding-window).

4. Risultati

Il documento valuta Still su rapporti di compressione da 8× a 200× e lunghezze di contesto da 8k a 128k.

• Frontiera Velocità-Qualità: Sui modelli Qwen e Gemma, Still occupa il lato favorevole della frontiera velocità-qualità. Rimane accurata e veloce in tutta la matrice, mentre i metodi basati sulla selezione degradano sotto budget ristretti, e i metodi di sintesi per-contesto sono troppo lenti.
• Benchmark RULER: In configurazioni di addestramento corrispondenti sul benchmark RULER, Still supera il più forte baseline di selezione ammortizzata (KV-Distill) di 8–22 punti di accuratezza in 16 celle su 18.
• Performance di Trasferimento: Nelle configurazioni di trasferimento zero-shot (addestramento su MCQ a dominio misto, test su RULER), Still mantiene un vantaggio significativo rispetto a KV-Distill, sebbene i margini siano minori a causa del mismatch tra addestramento e test.
• Riassunto: Still supporta la sintesi libera. Su HELMET (multi_lexsum), preserva il 74–95% del guadagno dal contesto completo tra 8k–64k di contesto, superando Attention Matching e KV-Distill. Su LongBench v1, vince il 60% dei confronti a coppie contro KV-Distill.
• Compattazione Iterativa: Still abilita la compattazione iterativa a blocchi. Sebbene le performance degradino se l'orizzonte di addestramento è troppo breve rispetto all'orizzonte di deployment (es. un compattatore addestrato su 8k che collassa a 128k), l'addestramento su orizzonti iterativi più lunghi (16k, 32k) permette al modello di assorbire contesti molte volte la sua capacità fisica.

5. Significato e Rivendicazioni

Il documento afferma che l'ammortizzazione rende trattabile la compattazione della cache a lungo contesto, e che la sintesi rende lo stato compresso utile ad estremi di compressione.

• Trattabilità: Imparando l'ottimizzatore (il compattatore) una volta sola, il sistema evita i costi di fitting cumulativi dei metodi per-contesto, rendendo fattibile il deployment iterativo a lungo orizzonte.
• Espressività: La sintesi permette alla cache compatta di fondere informazioni da arbitrariamente molte posizioni, superando il limite di densità informativa della selezione dei token.
• Praticità: Still è abbastanza leggero da poter essere chiamato ripetutamente durante l'inferenza, supportando sia la selezione della risposta che la generazione a forma libera senza modificare i pesi del modello base.

Limitazioni Riconosciute dagli Autori

• Non è Lossless: L'inferenza a contesto completo rimane superiore nelle configurazioni più difficili a 128k.
• Vincolo dell'Orizzonte di Addestramento: La compattazione iterativa non è un'estrapolazione gratuita; le performance del compattatore sono limitate dall'orizzonte di addestramento. Supportare 1 milione di token richiede l'addestramento su orizzonti proporzionalmente più lunghi, il che è costoso.
• Recupero Esatto: La ricorrenza attuale è debole nei compiti di recupero esatto e "needle-in-a-haystack", anche quando preserva l'evidenza semantica per compiti MCQ più ampi.
• Memoria vs Latenza: In modalità single-pass, Still riduce la dimensione della cache trattenuta, ma non elimina la latenza iniziale del prefill a contesto completo.
• Requisiti di Deployment: Ogni checkpoint del modello base richiede il proprio compattatore addestrato, aggiungendo overhead in termini di parametri e dati di addestramento.