LinearKV: One Cached State Suffices for Position-Independent Caching in Hybrid LLMs

Riepilogo Tecnico: LINEARKV

Problematica
Il serving dei Large Language Models (LLM) affronta significativi colli di bottiglia in termini di latenza e throughput a causa del crescente costo del prefilling di contesti lunghi. Il Caching Indipendente dalla Posizione (PIC) mitiga questo problema memorizzando e riutilizzando chunk di token pre-prefillati indipendentemente dalla loro posizione nel prompt. Tuttavia, i metodi PIC esistenti sono progettati per architetture a attenzione completa (full-attention), basandosi su cache Key-Value (KV) indicizzate per token che possono essere concatenate e riparate selettivamente. Gli LLM ibridi, che intercalano strati a full-attention (FA) con strati ricorrenti (es. Mamba-2, Gated DeltaNet), rompono questi primitivi. Gli strati ricorrenti comprimono il contesto in uno stato di dimensione fissa invece di mantenere una cache indicizzata per token. Ciò crea due sfide: (1) non è chiaro come inizializzare e riparare gli stati ricorrenti da chunk memorizzati disgiunti senza voci indicizzate per token da concatenare, e (2) non si sa se una strategia PIC unificata possa generalizzarsi attraverso diverse architetture ibride (es. Mamba-2 con scalar-decay vs. Gated DeltaNet con dense-transition).

Metodologia: LinearKV
Gli autori propongono LinearKV, un framework training-free per il caching indipendente dalla posizione negli LLM ibridi. L'innovazione principale è una strategia di inizializzazione disaccoppiata applicata dopo l'abbinamento dei chunk memorizzati:

1. Inizializzazione Disaccoppiata:

   • Strati Full-Attention: Questi strati procedono come nel PIC standard, concatenando le voci KV memorizzate dai chunk abbinati lungo l'asse del token.
   • Strati Ricorrenti (Lineari): Invece di tentare di concatenare gli stati (il che è impossibile poiché sono riassunti a dimensione fissa), il framework mappa $K$ stati locali abbinati $\{S^{\text{local}}_1, \dots, S^{\text{local}}_K\}$ in un singolo stato iniziale $S_{\text{init}}$ tramite una funzione $f$.
   • Riutilizzo del Selettore: Questo disaccoppiamento consente a LinearKV di riutilizzare gli algoritmi di selezione dei token PIC esistenti (es. CacheBlend, EPIC, ProphetKV) senza modifiche. Il selettore determina un insieme di posizioni di token $U$ da ricalcolare. Il sistema ripara quindi la cache sovrascrivendo le voci KV degli strati FA in $U$ e facendo avanzare lo stato ricorrente da $S_{\text{init}}$ attraverso i token in $U$ nell'ordine del contesto.

2. L'Inizializzatore dello Stato Lineare ($f$):
   Il documento investiga due scelte per $f$:

   • Composizione Esatta: Un approccio algebricamente fondato (adottato dal lavoro concorrente HYPIC) che compone tutti i $K$ stati memorizzati utilizzando gli operatori di transizione dei chunk per ricostruire l'esatto stato del prefisso completo.
   • Inizializzazione Single-Block (default di LinearKV): Un approccio più semplice in cui $f$ seleziona un singolo stato memorizzato (tipicamente l'ultimo chunk abbinato, $S^{\text{local}}_K$) come inizializzatore, scartando gli altri.

Risultati Chiave e Contributi
Il documento presenta tre contributi primari, supportati da evidenze empiriche su tre modelli ibridi (Granite/Mamba-2, OLMo/GDN, Qwen/GDN) e tre selettori PIC:

1. Formalizzazione del PIC Ibrido: Gli autori formalizzano il PIC ibrido come un framework disaccoppiato dove l'unica operazione specifica dell'architettura è l'inizializzazione dello stato lineare. Ciò consente di integrare i selettori PIC a full-attention esistenti senza modifiche.
2. Fragilità della Composizione Esatta: Contrariamente all'intuizione algebrica, il documento dimostra che la composizione esatta è superflua e architettonicamente fragile.
   • Sui modelli GDN (OLMo, Qwen), la composizione esatta e l'inizializzazione single-block operano in modo equivalente, recuperando fino al 92% della qualità totale.
   • Sui modelli Mamba-2 (Granite), la composizione esatta collassa. A causa della natura scalar-decay di Mamba-2, gli errori dai chunk pre-prefillati indipendentemente si accumulano attraverso la profondità, portando a uno stato iniziale costruito con un alto errore relativo. Di conseguenza, la composizione esatta recupera solo il ~46,6% della qualità totale sotto il selettore EPIC, mentre l'inizializzatore single-block recupera l'86,8%.
   • Gli autori identificano che il fallimento deriva dall'aggregazione di $K$ operatori profondi condizionati indipendentemente, non dall'architettura stessa, poiché l'approccio single-block evita questo accumulo di errori.
3. Guadagni di Efficienza: L'inizializzatore single-block è computazionalmente meno oneroso. Evita il costo online di comporre matrici di transizione dense richieste dalla composizione esatta. Ciò riduce il Time-to-First-Token (TTFT) di un ulteriore 5–17% rispetto alla composizione esatta, mantenendo o migliorando l'accuratezza.

Risultati Sperimentali

• Benchmark: Valutato su LongBench QA, RULER (8K–32K) e LongBench Summarization.
• Performance:
  • Granite (Mamba-2): LinearKV (single-block) supera significativamente la composizione esatta in tutti i selettori. Ad esempio, sotto EPIC con un budget di ricalcolo del 20%, LinearKV raggiunge l'86,8% della qualità totale rispetto al 46,6% della composizione esatta.
  • OLMo & Qwen (GDN): Entrambi i metodi si comportano in modo simile, con LinearKV che eguaglia le prestazioni della composizione esatta.
  • Robustezza: La superiorità dell'inizializzatore single-block su Mamba-2 si mantiene attraverso vari rapporti di ricalcolo (0,03 a 0,4) e lunghezze di contesto (8K a 32K).
• Latenza: LinearKV riduce il TTFT allo 0,46×–0,62× del tempo di prefill completo, superando la composizione esatta in tutti gli scenari misurati.

Significatività
Il documento stabilisce che per gli LLM ibridi, l'approccio "algebricamente fondato" di ricostruire l'esatto stato del prefisso completo dai chunk memorizzati non è solo computazionalmente costoso, ma può essere dannoso per le prestazioni del modello su architetture specifiche come Mamba-2. Dimostrando che un singolo stato memorizzato è sufficiente come inizializzatore robusto, LinearKV fornisce una soluzione semplice, training-free e agnostica rispetto all'architettura che massimizza i benefici del caching indipendente dalla posizione minimizzando la latenza ed evitando l'accumulo di errori inerente alla composizione di stati da molteplici sorgenti. Questa scoperta sfida l'assunto che una composizione di stati più complessa sia necessaria per un caching efficace nelle architetture ricorrenti.