Ottimizzazione della Cache KV con Reinforcement Learning

La gestione efficiente della cache Key-Value (KV) è cruciale per l'inference dei Large Language Models (LLM), data la crescente dimensione di questi modelli e le conseguenti richieste di memoria. Un nuovo studio introduce KV Policy (KVP), un framework che utilizza il reinforcement learning (RL) per migliorare l'eviction dei token dalla cache KV.

KV Policy: Un Approccio Basato sull'Utilità Futura

KVP riformula l'eviction della cache KV come un problema di reinforcement learning, addestrando agenti RL specializzati per prevedere l'utilità futura dei token. A differenza dei metodi tradizionali che si basano su euristiche come la recenza o gli score di attenzione passati, KVP valuta direttamente l'utilità futura dei token per il decoding. Gli agenti RL sono addestrati su tracce di generazione pre-calcolate, utilizzando solo vettori chiave e valore, senza richiedere modifiche all'LLM sottostante o inference aggiuntiva.

Performance e Generalizzazione

Le valutazioni su benchmark di long-context (RULER) e dialogo multi-turno (OASST2-4k) dimostrano che KVP supera significativamente le baseline. Inoltre, test zero-shot su task downstream standard (LongBench, BOOLQ, ARC) indicano che KVP si generalizza bene oltre la sua distribuzione di training e a contesti più lunghi. Questi risultati suggeriscono che prevedere l'utilità futura dei token è un paradigma efficace e scalabile per la gestione adattiva della cache KV.