Riassunto Tecnico: CentroidKV

Problematica

I Large Language Models (LLM) con finestre di contesto estese (fino a 1M di token) affrontano sfide significative di deployment a causa della cache Key-Value (KV). La cache KV, che memorizza le chiavi e i valori dei token precedenti per evitare la ricalcolazione, scala linearmente con la lunghezza del contesto. Ciò crea un collo di bottiglia per la latenza e il throughput dell'inferenza, poiché ogni token generato richiede l'accesso all'intera cache. Inoltre, l'ampia impronta di memoria complica il deployment.

Le soluzioni esistenti per ridurre la dimensione della cache KV includono:

• Eviction (Eliminazione): Scartare i token meno critici in base ai punteggi di attenzione. Questo spesso porta a un degrado delle prestazioni perché i token considerati non importanti inizialmente possono diventare cruciali per le generazioni future.
• Merging (Fusione): Combinare i token destinati all'eviction con i token mantenuti. Tuttavia, i metodi esistenti spesso si affidano a un'eviction preliminare per definire i set di fusione, costringendo token semanticamente distanti a essere raggruppati, con conseguente bias del centroide e perdita di informazioni.
• Quantizzazione e Pruning: Queste offrono guadagni di efficienza limitati o introducono un elevato overhead computazionale.

Esiste la necessità di un metodo che comprima la cache KV basandosi sulla somiglianza intrinseca senza scartare informazioni critiche o incorrere in costi computazionali proibitivi.

Metodologia: CentroidKV

Il documento presenta CentroidKV, un framework di clustering della cache KV online. La sua innovazione principale è l'algoritmo Chunked Soft Matching (CSM), progettato per raggruppare gli stati KV in base alla somiglianza mantenendo una bassa complessità computazionale.

Componenti Chiave

1. Osservazioni:

   • Somiglianza Localizzata: Gli stati delle chiavi esibiscono un'alta somiglianza del coseno lungo la dimensione della sequenza, con token simili che si raggruppano in regioni localizzate.
   • Correlazione della Distanza: All'aumentare della distanza tra i token, la somiglianza del coseno degli stati delle chiavi generalmente diminuisce in modo monotono e segue un trend convesso.

2. Pipeline di Inferenza:

   • Pre-filling: La sequenza viene elaborata e, se la cache KV supera un budget, viene eseguito il primo round di clustering.
   • Decoding: La cache cresce con ogni token. Il clustering viene invocato periodicamente (ogni $g$ step) per mantenere la cache entro il budget.
   • Preservazione: Gli "attention sinks" e i token recenti vengono preservati prima della partizione per garantire la stabilità del modello.

3. Chunked Soft Matching (CSM):

   • Chunking: La sequenza viene divisa in chunk per gestire contesti lunghi in modo efficiente.
   • Partizione Alternata: All'interno di ogni chunk, i token sono partizionati in due set ($A$ e $B$) in modo alternato. L'analisi teorica dimostra che questa strategia è ottimale dati i trend di somiglianza monotonici e convessi osservati, garantendo che token altamente simili siano collocati in set diversi per facilitare il matching.
   • Matching: Similmente al Bipartite Soft Matching (BSM) utilizzato nei Vision Transformer, il CSM identifica coppie di token altamente simili tra i set $A$ e $B$.
   • Compressione: Invece di fondere tutte le coppie, l'algoritmo classifica i match candidati per somiglianza e fonde solo la frazione superiore controllata da un rapporto di compressione $r$.
   • Scheduling Progressivo: Il rapporto di compressione $r$ viene attenuato linearmente su più round di clustering ($r = r_{init} - j \cdot \delta_r$). Questo rende la fusione progressivamente più selettiva nei round successivi, preservando strutture ad alta confidenza.

4. Aggregazione:

   • Per ogni cluster, gli stati di chiave e valore vengono fusi in un singolo centroide tramite una media ponderata sul grado. Il "grado" ($n_t$) traccia il contributo cumulativo dei token fusi in un centroide attraverso i vari round.
   • Il meccanismo di attenzione viene adattato per utilizzare questi centroidi, incorporando il grado del cluster come termine logaritmico nel calcolo della softmax per approssimare l'output originale dell'attenzione.

Contributi Chiave

• Framework CentroidKV: Una soluzione semplice e "plug-and-play" per il clustering della cache KV online che migliora l'efficienza dell'inferenza in scenari di contesto lungo.
• Algoritmo Chunked Soft Matching: Un nuovo algoritmo di clustering che estende il Bipartite Soft Matching alla compressione della cache KV. Utilizza una strategia di partizione alternata all'interno dei chunk, la cui ottimalità è stata dimostrata teoricamente per i pattern di somiglianza osservati.
• Analisi Teorica: Il documento fornisce una prova formale dell'ottimalità della strategia di partizione intra-chunk e analizza la complessità computazionale, dimostrando che il CSM raggiunge una complessità inferiore ($O(ncd)$) rispetto a K-Means ($O(inkd)$) o ai metodi a mesh pairwise completo ($O(n^2d)$).
• Validazione Empirica: Estesi esperimenti sui modelli Llama-3.1-8B e Mistral-7B attraverso i benchmark RULER e LongBench.

Risultati

• Efficienza della Memoria: CentroidKV raggiunge una riduzione fino al 75% dell'uso della memoria della cache KV (riducendo la dimensione della cache di circa 4 volte) mantenendo prestazioni del modello comparabili.
• Velocità di Inferenza:
  • Latenza di Decoding: Accelera la fase di decoding fino a 1.92× (misurato come riduzione del Time Per Output Token).
  • Throughput: Aumenta il throughput di serving fino a 4× quando integrato nel framework vLLM.
• Accuratezza:
  • Sul benchmark RULER, CentroidKV raggiunge le prestazioni migliori o quasi migliori attraverso vari budget di cache (25%–75%), particolarmente in compiti che richiedono comprensione semantica (es. S-NIAH-1/2, MK-NIAH-1, VT).
  • Su LongBench, si comporta costantemente bene in Single-Doc QA, Multi-Doc QA e Summarization.
  • Limitazioni: Le prestazioni calano leggermente su compiti sintetici e compiti di recupero fine-grained (es. chiavi basate su UUID) dove il matching esatto del token è critico. Gli autori notano che i metodi basati sul contenuto generalmente faticano con stringhe casuali semanticamente arbitrarie rispetto ai metodi basati sulla posizione come StreamingLLM.

Significato e Claim

Il documento sostiene che CentroidKV offre un favorevole compromesso tra accuratezza ed efficienza. Sfruttando la somiglianza intrinseca degli stati delle chiavi e impiegando un algoritmo di clustering a singolo passaggio molto leggero, evita il degrado delle prestazioni associato alla naif eviction e l'elevato overhead computazionale dei metodi di clustering iterativi come K-Means.

Gli autori sottolineano che CentroidKV è una soluzione leggera e plug-and-play che non richiede il retraining del modello. Affronta efficacemente il collo di bottiglia della memoria dei modelli a contesto lungo, consentendo un decoding più veloce e un throughput più elevato senza compromettere significativamente la capacità del modello di comprendere le dipendenze a lungo raggio.

Limitazioni note dagli autori:

• L'implementazione attuale si concentra sulla compressione basata su GPU e non esplora strategie di offloading della memoria (es. clustering su CPU).
• Lo scheduling del rapporto di compressione è specificato manualmente, e il metodo limita la riduzione della cache a massimo la metà ad ogni step. Il lavoro futuro potrebbe esplorare strategie di compressione adattive.