ForesightKV: Optimizing KV Cache Eviction for Reasoning Models by Learning Long-Term Contribution

Riepilogo Tecnico: ForesightKV

Definizione del Problema

I Large Language Models (LLM), in particolare i modelli di ragionamento, hanno dimostrato capacità eccezionali nella generazione di lunghe tracce di ragionamento. Tuttavia, all'aumentare della lunghezza della sequenza, la cache Key-Value (KV) si espande linearmente, portando a un significativo overhead di memoria e a un aumento della latenza dovuto alla decodifica limitata dalla memoria (memory-bound decoding). Sebbene gli esistenti metodi di espulsione (eviction) della cache KV tentino di mitigare questo problema scartando coppie KV meno importanti sulla base di regole euristiche (ad esempio, punteggi di attenzione, posizioni o caratteristiche), essi spesso non riescono a catturare dipendenze complesse e dinamiche tra i diversi head di attenzione. Nello specifico, questi metodi privi di addestramento (training-free) faticano con pattern "semanticamente dipendenti", dove l'importanza di una coppia KV cambia dinamicamente in base al contesto in evoluzione. Inoltre, l'espulsione di coppie KV critiche può causare aumenti bruschi della perdita (loss) per i token a bassa entropia (parti deterministiche del ragionamento), portando a errori fattuali che distorcono i successivi passaggi di ragionamento.

Metodologia

Gli autori propongono ForesightKV, un framework basato sull'addestramento progettato per apprendere il contributo a lungo termine delle coppie KV e per espellere adattivamente quelle meno importanti durante la generazione. Il framework opera su un budget di memoria fisso e impiega un paradigma di addestramento in due fasi per ottimizzare un modello di scoring leggero ($\pi_\theta$) che predice l'importanza di ogni coppia KV. Fondamentalmente, i parametri dell'LLM rimangono congelati durante tutto il processo; viene aggiornato solo il modello di scoring.

1. Processo di Espulsione Dinamica

L'espulsione è attivata periodicamente (ogni $L$ token). Quando la dimensione della cache raggiunge $B + L$ (dove $B$ è il budget target), il sistema mantiene gli ultimi $L$ token più recenti e utilizza il modello di scoring per ridurre la cache rimanente a $B - L$. Il modello di scoring prende come input la concatenazione della chiave ($k_n$), del valore ($v_n$) e delle caratteristiche di attenzione ($a_n$) per predire un punteggio di importanza ($\phi_n$). Le decisioni di espulsione sono prese tramite una strategia Top-K seguita da un campionamento multinomiale per introdurre la stocasticità.

2. Fase Uno: Apprendimento Supervisionato con "Golden Eviction"

Per addestrare il modello di scoring, gli autori introducono l'algoritmo Golden Eviction per costruire etichette di supervisione.

• Meccanismo: Per una traccia di ragionamento completa, la matrice di attenzione viene partizionata in blocchi. L'algoritmo calcola i "punteggi di blocco" aggregando i punteggi di attenzione all'interno dei blocchi e attraverso gli head di attenzione.
• Consapevolezza del Futuro (Future Awareness): Per ogni passaggio di espulsione, l'algoritmo identifica il punteggio di blocco massimo attraverso tutti i passaggi futuri per ogni coppia KV. Questo "punteggio futuro" rappresenta la strategia di mantenimento ottimale che minimizza l'impatto sulle computazioni di attenzione future.
• Addestramento: Il modello di scoring viene addestrato utilizzando una Pairwise Ranking Loss per garantire che i punteggi predetti per qualsiasi coppia di valori KV siano allineati con la classificazione dei loro punteggi futuri derivati dall'algoritmo Golden Eviction.

3. Fase Due: Reinforcement Learning (RL)

L'addestramento supervisionato da solo potrebbe non tenere conto dello shift della distribuzione causato dall'espulsione durante l'inferenza. Per colmare questa lacuna, gli autori formulano l'espulsione della cache KV come un Processo Decisionale di Markov (MDP) e applicano l'algoritmo GRPO (Group Relative Policy Optimization).

• Stato: La cache KV rimanente attuale.
• Azione: Selezione di un sottoinsieme di coppie KV da mantenere.
• Ricompensa (Reward): La funzione di ricompensa è progettata per mitigare la degradazione delle prestazioni sui token a bassa entropia. Nello specifico, si concentra sui token che originariamente avevano una bassa entropia ma che hanno subito un aumento significativo della loss dopo l'espulsione. La ricompensa è definita come il negativo dell'Errore Quadratico Medio (MSE) degli aumenti della loss per questo specifico sottoinsieme di token.
• Ottimizzazione: Il modello di policy è ottimizzato per massimizzare la ricompensa cumulativa attraverso la sequenza, imparando efficacemente a preservare il contesto critico necessario per i passaggi di ragionamento deterministici.

Contributi Chiave

1. Framework ForesightKV: Un nuovo metodo di addestramento in due fasi (Apprendimento Supervisionato + RL) che impara a predire il contributo a lungo termine delle coppie KV, andando oltre le regole euristiche statiche.
2. Algoritmo Golden Eviction: Un metodo per generare tracce di espulsione ottimali utilizzando informazioni di attenzione futura, fungendo da supervisione di alta qualità per il modello di scoring.
3. Formulazione MDP per l'Espulsione: Modellare il processo di espulsione come un MDP e utilizzare GRPO per mirare specificamente alla mitigazione dei picchi di loss nei token a bassa entropia, che sono critici per l'accuratezza del ragionamento.
4. Efficienza: Il framework addestra solo un modello di scoring leggero senza richiedere operazioni di backward o aggiornamenti dei parametri sull'LLM di base.

Risultati Sperimentali

Gli autori hanno valutato ForesightKV su tre modelli di ragionamento (Qwen3-1.7B, Qwen3-4B e DeepSeek-R1-Distill-Qwen-7B) utilizzando i benchmark matematici AIME2024 e AIME2025.

• Prestazioni: ForesiteKV ha superato costantemente i baseline esistenti (SnapKV, H2O, R-KV). Notevolmente, ha ottenuto prestazioni comparabili o superiori ai baseline a cache completa utilizzando solo la metà del budget della cache KV. Ad esempio, su AIME2024 con Qwen3-4B, ForesightKV con un budget di 1K ha superato R-KV con un budget di 2K (54.5 vs 44.8 Pass@1).
• Throughput: Il metodo ha migliorato significativamente il throughput. Con un budget di 1K su una sequenza di 32K, ForesightKV ha ottenuto un guadagno di throughput fino a 9.79× rispetto alla cache KV completa.
• Generalizzazione: Nonostante sia stato addestrato principalmente su compiti matematici, ForesightKV ha dimostrato una forte generalizzazione sui benchmark GPQA (scienza) e LiveCodeBench (coding), superando i baseline con metà del budget di cache.
• Studi di Ablazione: I risultati hanno confermato che sia la fase di addestramento supervisionato (che fornisce la capacità predittiva iniziale) sia la fase di reinforcement learning (che affina la policy per la preservazione dei token a bassa entropia) sono cruciali per il successo. La specifica progettazione della ricompensa focalizzata sui token a bassa entropia con grandi aumenti di loss è risultata essere la più efficace.

Significato e Claim

Il paper sostiene che ForesightKV offre una nuova direzione promettente per la gestione adattiva della cache KV nel ragionamento a lungo contesto. Passando dalle euristiche basate su regole a un approccio basato sull'apprendimento che modella esplicitamente il contributo a lungo termine delle coppie KV e mira alla preservazione dei token a bassa entropia, il metodo raggiunge un compromesso superiore tra efficienza della memoria e qualità della generazione. Gli autori sottolineano che il loro approccio cattura l'importanza intrinseca delle voci della cache KV piuttosto che fare overfitting su specifici vincoli di budget, permettendogli di generalizzare bene attraverso diversi setting di budget e domini di task.