Sintesi Tecnica: KVpop — Compressione della Cache KV con Pruning Online Predittivo

1. Definizione del Problema

La decodifica autoregressiva nei Transformer-based Large Language Models (LLM) si affida a una cache Key-Value (KV) per evitare di ricalcolare le rappresentazioni dei token passati. Tuttavia, la cache KV cresce linearmente con la lunghezza del contesto, creando un collo di bottiglia significativo per la memoria e la larghezza di banda durante l'inferenza a lungo contesto.

Le soluzioni esistenti per ridurre questo costo si dividono generalmente in due categorie:

• Evizione Statica/Euristica: Metodi come la Sliding Window o la preservazione dei Sink-Token (es. StreamingLLM) si basano su regole fisse, mentre i metodi basati su punteggio utilizzano proxy online (es. attenzione cumulativa). Questi spesso falliscono nel tracciare l'utilità futura dei token, portando a un'evizione fragile quando la rilevanza cambia.
• Evizione Appresa: Approcci come la Dynamic Memory Compression (DMC) o la Dynamic Memory Sparsification (DMS) addestrano policy per fondere o scartare i token. Tuttavia, molti di questi metodi prendono decisioni di evizione al momento dell'inserimento del token o si affidano a rilassamenti differenziabili che non supervisionano direttamente la specifica decisione "mantieni-o-scarta" basata sull'utilità futura. Inoltre, i metodi con finestre protette ritardano la rimozione dei token, ma decidono spesso il loro destino immediatamente all'ingresso, ignorando le prove che si accumulano mentre il token rimane nella finestra protetta.

La sfida centrale è predire quali token saranno utili per le query future, poiché i token localmente salienti possono diventare irrilevanti, mentre i token con bassa attenzione immediata possono diventare critici in seguito (es. nelle tracce di ragionamento).

2. Metodologia: KVpop

KVpop (KV compression with Predictive Online Pruning) è un retrofit di sparse-attention progettato per imporre un budget fisso per ogni testa KV senza alterare l'architettura del modello base. Affronta i limiti dei lavori precedenti attraverso due meccanismi primari: Future-Attention Supervision e Delayed Stateful Scoring.

2.1 Architettura e Budget

KVpop mantiene una cache limitata per ogni testa KV definita dal budget $B = s + w + k$:

• $s$ Sink Tokens: I primi token sono sempre mantenuti.
• $w$ Protected Window: I più recenti $w$ token sono protetti dall'evizione.
• $k$ Long-Range Top-K: I token più vecchi competono per un budget fisso di $k$ slot basato su punteggi di importanza appresi.

2.2 Future-Attention Target Supervision

A differenza dei metodi che valutano i token in base all'attenzione passata o alle statistiche locali, KVpop supervisiona lo scorer utilizzando la massa di attenzione futura che un token riceve dopo aver lasciato la finestra protetta.

• Definizione del Target: Per un token $t$, il target è la massa logaritmica aggregata di attenzione che riceve da query future (posizioni $d \ge t + w$).
• Computazione Efficiente: Per evitare il costo $O(S^2)$ di materializzare mappe di attenzione dense durante l'addestramento, KVpop utilizza un passaggio di attenzione trasposto. Scambiando i ruoli di query e chiavi e riutilizzando i normalizzatori log-sum-exp (LSE) restituiti dai kernel di sparse attention, il sistema calcola il target di future-attention in modo efficiente senza la materializzazione di matrici dense.
• Funzione di Perdita: L'obiettivo di addestramento è una perdita logistica pairwise boundary-aware. Confronta il punteggio di un nuovo token eleggibile ($t_{new}$) con il punteggio del token di confine dell'insegnante ($t_{bnd}$). La perdita spinge lo studente a classificare correttamente $t_{new}$ rispetto al confine, assicurando che la selezione top-$k$ corrisponda alla classifica di utilità futura dell'insegnante.

2.3 Stateful Scoring con Decisione Ritardata

KVpop introduce uno scorer basato su memoria ritardato (implementato tramite un mLSTM) che differisce la decisione di scoring fino a quando il token raggiunge il confine di evizione.

• Meccanismo: Invece di valutare un token immediatamente al momento dell'inserimento, lo scorer aggiorna il suo stato interno con la rappresentazione del token e il contesto successivo. Quando il token esce dalla finestra protetta ($q - w$), lo scorer legge la memoria accumulata per generare il punteggio di importanza.
• Vantaggio: Ciò consente allo scorer di incorporare il contesto "near-future" (prove che si accumulano mentre il token è nella finestra protetta) prima di prendere la decisione di mantenere o scartare.

2.4 Efficienza di Inferenza

Durante l'inferenza, KVpop impone il budget fisso $B$. La selezione dei top-$k$ token a lungo raggio è ottimizzata utilizzando una soglia corrente su priorità statiche dei token e un albero di Fenwick, permettendo alla maschera sparsa di essere generata on-the-fly all'interno del kernel di attenzione (es. FlexAttention) senza memorizzare maschere dense $S \times S$.

3. Contributi Chiave

1. Policy di Evizione Supervisionata: Gli autori formulano l'evizione KV a budget fisso come un problema di predizione supervisionata al confine di evizione, utilizzando un target di future-attention per addestrare le decisioni top-$k$ a lungo raggio.
2. Computazione Efficiente del Target: Introducono un passaggio di attenzione trasposto, utilizzato solo durante l'addestramento, per calcolare i target di future-attention senza materializzare mappe di attenzione dense, garantendo l'assenza di overhead all'inferenza.
3. Scoring Ritardato: Dimostrano che uno scorer stateful può sfruttare il contesto near-future differendo la decisione di scoring fino a quando il token diventa evictabile, superando i metodi di scoring immediato.
4. Validazione Empirica: Il metodo è valutato su ragionamento matematico (AIME, HMMT) e compiti out-of-domain (GPQA-D, LiveCodeBench), mostrando che le policy di evizione apprese possono mantenere alte prestazioni sotto compressione aggressiva.

4. Risultati Sperimentali

Gli autori hanno valutato KVpop sui modelli Qwen3-4B e Qwen3-8B, addestrati sul dataset Nemotron-Math v2 con una lunghezza di sequenza di 16.384.

4.1 Ragionamento Matematico (AIME & HMMT)

• Rapporto di Compressione 75% (Budget 2048):
  • Qwen3-4B: KVpop ha mantenuto il 98% delle prestazioni del teacher a full-attention.
  • Qwen3-8B: KVpop ha mantenuto il 100% (prestazioni quasi pari al teacher).
  • Confronto: KVpop ha superato significativamente i baseline training-free (StreamingLLM, TOVA) e il baseline appreso DMS. Ad esempio, al 75% di compressione, StreamingLLM è sceso a circa il 76% della performance relativa su Qwen3-4B, mentre KVpop ha mantenuto il 98%.
• Rapporto di Compressione 88% (Budget 4096):
  • Qwen3-4B: KVpop ha mantenuto il 97% delle prestazioni del teacher.
  • Qwen3-8B: KVpop ha mantenuto il 100% delle prestazioni del teacher.
  • Confronto: I metodi training-free sono crollati significativamente (es. StreamingLLM è sceso a circa il 47% della performance relativa), mentre KVpop è rimasto robusto.

4.2 Generalizzazione Out-of-Domain

Valutato su GPQA-Diamond (ragionamento scientifico) e LiveCodeBench (generazione di codice):

• I modelli KVpop addestrati solo su dati di ragionamento matematico si sono generalizzati bene a questi domini.
• Al 75% di compressione, KVpop su Qwen3-4B ha raggiunto un'accuratezza comparabile al dense teacher su GPQA-D (0.59 vs 0.59) e risultati competitivi su LCB.

4.3 Efficienza di Inferenza

• Memoria: A differenza dell'attenzione densa, dove l'uso della VRAM scala linearmente con la lunghezza della generazione (raggiungendo circa 36GB a 131k token), KVpop e DMS hanno mantenuto un footprint di memoria quasi costante (~19GB) indipendentemente dalla lunghezza della sequenza.
• Latenza: KVpop ha dimostrato una latenza di decodifica end-to-end inferiore rispetto a DMS e l'attenzione densa per lunghezze di generazione elevate (fino a 131k token). Gli autori attribuiscono questo fatto al budget fisso per testa di KV di KVpop, che evita le strutture di cache "ragged" di DMS che complicano la compilazione e l'esecuzione.

4.4 Studi di Ablazione

• Scoring Ritardato: Rimuovere il ritardo (scoring immediatamente all'inserimento) ha causato un calo di 0.2 punti nell'accuratezza dei token, confermando il beneficio dell'integrazione del contesto near-future.
• Stateless vs. Stateful: La variante stateful (KVpop) ha generalmente eguagliato o superato leggermente la variante stateless (KVpopmlp), in particolare ai rapporti di compressione più elevati.

5. Significato e Claim

Il paper sostiene che supervisionare l'evizione con segnali di future-attention sia una strategia superiore per la compressione della cache KV rispetto alle euristiche statiche o ai proxy score. Supervisionando direttamente una policy per predire quali token saranno importanti per le query future, KVpop ottiene:

1. Alta Fedeltà: Mantiene prestazioni vicine al modello teacher anche ad alti rapporti di compressione (fino all'88%), un regime in cui altri metodi falliscono.
2. Robustezza: La policy appresa si trasferisce efficacementamente ai compiti out-of-domain, suggerendo che il modello apprenda principi generali di utilità del token piuttosto che fare overfitting su distribuzioni di training specifiche.
3. Efficienza: Impone un limite di memoria rigido, consentendo una decodifica a lungo contesto stabile con uso di memoria costante e throughput migliorato rispetto ai metodi di gating dinamico come DMS.

Gli autori concludono che KVpop dimostra che le policy di evizione apprese, quando opportunamente supervisionate dalla utilità futura e consentite di utilizzare il contesto near-future, possono risolvere efficacemente il collo di bottiglia della memoria nella generazione autoregressiva senza sacrificare la qualità del modello.