Riepilogo Tecnico: KV-Fold

Enunciato del Problema

I grandi transformer preaddestrati sono limitati dal costo quadratico della memoria dell'attenzione completa e dalla finestra di contesto fissa di un singolo passaggio in avanti. Sebbene i modelli recenti supportino finestre di contesto fino a centinaia di migliaia di token, i carichi di lavoro reali (ad esempio, codice su scala di repository, cartelle cliniche longitudinali) spesso superano questi limiti. Le soluzioni esistenti affrontano un compromesso persistente:

• I metodi di streaming (ad esempio, StreamingLLM) mantengono una memoria limitata conservando solo i token recenti e i "pozzi di attenzione" (attention sinks), ma perdono la capacità di recuperare le informazioni una volta che queste escono dalla finestra scorrevole.
• I metodi di compressione della cache KV e di memoria ricorrente richiedono spesso modifiche al modello, fine-tuning o introducono approssimazioni con perdita che degradano la fedeltà del recupero.
• L'attenzione completa è computazionalmente non fattibile per sequenze estremamente lunghe sull'hardware standard a causa delle dimensioni massive della matrice dei punteggi di attenzione (ad esempio, ~1 TB per una sequenza di 128K su una GPU da 40 GB).

Il documento si chiede se un transformer preaddestrato e congelato possa mantenere il proprio stato interno attraverso blocchi successivi per eseguire inferenze a lungo contesto come una ricorrenza, senza modifiche architetturali o addestramento.

Metodologia: KV-Fold

Gli autori propongono KV-Fold, un protocollo di inferenza senza addestramento che tratta la cache Key-Value (KV) come l'accumulatore in un fold sinistro (foldl) sui blocchi di sequenza.

• Meccanismo: Una sequenza lunga di lunghezza $T$ viene divisa in blocchi di lunghezza $C$. Quando si elabora il blocco $t$, il modello esegue un passaggio in avanti in cui:
  1. Le query ($Q$) sono generate dai token del blocco corrente.
  2. Le chiavi ($K$) e i valori ($V$) sono formati concatenando la cache KV accumulata da tutti i blocchi precedenti ($0$a$t-1$) con gli stati KV del blocco corrente.
  3. Gli embedding posizionali (RoPE) sono applicati in modo continuo attraverso i confini dei blocchi per mantenere l'allineamento assoluto della sequenza.
  4. La cache KV ingrandita viene passata in avanti al blocco $t+1$.
• Stato: I parametri del modello rimangono congelati. Non vengono introdotti token di memoria speciali o fine-tuning. La cache KV cresce linearmente con la lunghezza della sequenza, agendo come un record strutturato e non compresso del calcolo passato del modello.
• Analogia: Il processo è analogo a un fold sinistro nella programmazione funzionale, dove la cache KV funge da stato accumulatore aggiornato da una funzione di ricorrenza a singolo passo $F_\theta$ a ogni confine di blocco.

Contributi Chiave

1. Protocollo KV-Fold: Un semplice spostamento operativo che converte un transformer standard in un modello ricorrente a lungo contesto trasportando la cache KV accumulata attraverso i blocchi.
2. Stabilità della Ricorrenza: Evidenze empiriche che mostrano come la ricorrenza indotta non soffra di un'accumulazione di errori illimitata. Al contrario, la "deriva" (differenza nella Negative Log-Likelihood rispetto all'attenzione completa) aumenta brevemente e poi si satura in un plateau piatto che persiste attraverso catene profonde.
3. Deriva Strutturale vs Numerica: Il plateau è dimostrato essere una proprietà strutturale della dinamica dell'attenzione piuttosto che un artefatto numerico. Rimane stabile sotto un aumento di precisione numerica di 10.000 volte (da bf16 a fp32) ed è robusto attraverso diverse famiglie di modelli (OLMoE, Qwen, Llama) e dimensioni dei blocchi.
4. Recupero Esatto a Lunga Distanza: Il metodo preserva informazioni fattuali esatte su lunghe distanze. Sui benchmark "needle-in-a-haystack" (ago nel pagliaio), KV-Fold ottiene un recupero esatto al 100% su contesti fino a 128K token e profondità di catena fino a 511, eguagliando le prestazioni dell'attenzione completa.
5. Compromessi Pratici: Il documento caratterizza la posizione del metodo nello spettro memoria-verso-recupero. Sebbene richieda una crescita lineare della memoria (a differenza dello streaming a memoria limitata), abilita il recupero a profondità arbitrarie dove i metodi di streaming falliscono, rendendo fattibile eseguire un recupero equivalente all'attenzione completa su hardware che non può supportare un singolo passaggio in avanti monolitico.

Risultati Chiave

• Saturazione della Deriva: Su Qwen2.5-7B-Instruct, la deriva si satura alla profondità 7 e rimane piatta fino alla profondità 63, con un cambiamento totale di soli -0,0003 nat tra le profondità 15 e 60.
• Prestazioni di Recupero:
  • Singolo Ago: Accuratezza di recupero del 100% in 152 prove che coprono contesti da 16K a 128K token e profondità di catena fino a 511 su Llama-3.1-8B.
  • Multi-Ago: Accuratezza di recupero del 99,4% (175/176) quando si recuperano più fatti distinti simultaneamente.
• Fattibilità Hardware: Su una GPU A100 da 40 GB, un singolo passaggio in avanti con attenzione completa per una sequenza di 128K è non fattibile (richiedendo ~1 TB per i punteggi di attenzione). KV-Fold elabora la stessa sequenza in 171 secondi con un utilizzo di picco della memoria di 35,6 GB, rimanendo entro i limiti hardware.
• Confronto con StreamingLLM: A un contesto di 128K, StreamingLLM mantiene una memoria limitata (16,6 GB) e un runtime più veloce (22s) ma scende al 0% di recupero una volta che l'"ago" esce dalla sua finestra di 1024 token. KV-Fold mantiene il 100% di recupero a tutte le distanze testate ma richiede più memoria (35,6 GB) e un runtime più lungo.

Significato e Affermazioni

Il documento afferma che i transformer preaddestrati e congelati supportano già una forma stabile di ricorrenza della cache KV. Il significato principale è dimostrare che l'inferenza a lungo contesto non richiede strettamente modifiche architetturali, addestramento o compressione con perdita per ottenere un recupero ad alta fedeltà.

Gli autori sostengono che KV-Fold rivela un "regime di attenzione leggermente spostato ma stabile" in cui il modello può elaborare ricorsivamente le informazioni senza degradare la sua capacità di accedere a token passati specifici. Ciò fornisce una via pratica per estendere le finestre di contesto sull'hardware esistente, scambiando la crescita lineare della memoria con la capacità di eseguire un recupero esatto su sequenze che altrimenti verrebbero troncate o approssimate. Il documento conclude che, sebbene il metodo non riduca i FLOP totali o la memoria al di sotto dell'attenzione completa, riorganizza il calcolo in passaggi gestibili, rendendo fattibile l'inferenza a lungo contesto dove precedentemente era impossibile. Il lavoro futuro suggerito include un fine-tuning modesto per ridurre la deriva residua e la combinazione del metodo con tecniche di estensione posizionale per spingersi oltre i limiti di contesto nativi.