Sintesi Tecnica: SelKV - Fusione Selettiva della Cache KV con Merge-or-Drop Per-Token e Compensazione dell'Attenzione

1. Definizione del Problema

I Large Language Models (LLM) affrontano un collo di bottiglia critico durante l'inferenza autoregressiva: la cache Key-Value (KV). Man mano che la lunghezza del contesto ($n$) aumenta, l'impronta di memoria della cache KV cresce linearmente ($O(L \cdot H \cdot d \cdot n)$), limitando severamente le dimensioni dei batch, il throughput e i costi di implementazione per le applicazioni a lungo contesto.

Sebbene i metodi di compressione esistenti tentino di mitigare questo problema, essi soffrono di due limitazioni primarie:

1. Decisioni Uniformi di Merge-or-Drop: Gli approcci attuali applicano tipicamente una strategia binaria: o si fondono tutti i token espulsi (es. KVMerger, WeightedKV) o li si scarta interamente (es. H2O, SnapKV). Questa mancanza di granularità significa che quando il vettore di valore di un token espulso è dissimile dal suo target di fusione, la media ponderata corrompe la rappresentazione semantica preservata. Al contrario, lo scarto indiscriminato comporta la perdita di informazioni preziose.
2. Attention Sag (Calo dell'Attenzione): Quando più token vengono fusi in una singola voce di cache, tale voce riceve la stessa massa di attenzione softmax di un singolo token non fuso. Ciò provoca un "attention sag", in cui le posizioni fuse sono sistematicamente meno attese da un fattore proporzionale al numero di token fusi ($\sim m$), causando un sottoutilizzo delle informazioni preservate da parte del modello.

2. Metodologia: SelKV

Gli autori propongono SelKV (Selective KV Cache Merging), un framework a due componenti, privo di addestramento (training-free), progettato per affrontare queste limitazioni. SelKV opera come un modulo pluggable integrato nelle pipeline di fusione esistenti dopo la fase di prefill.

Componenti Core

1. Soft Cosine Gate (Per affrontare le decisioni uniformi)
Invece di una decisione binaria merge-or-drop, SelKV introduce un meccanismo di gating continuo e per-token basato sulla similarità dei vettori di valore.

• Meccanismo: Per ogni token espulso indirizzato a un target di fusione, il sistema calcola la similarità del coseno tra i loro vettori di valore.
• Funzione di Gating: Un gate soft $g = \max(\cos\_sim, 0)$ modula l'intensità della fusione:
  • $g \approx 1$: Alta similarità $\rightarrow$ Fusione completa.
  • $g = 0$: Vettori dissimili (ortogonali) $\rightarrow$ Scarto del token.
  • $0 < g < 1$: Similarità moderata $\rightarrow$ Fusione parziale.
• Implementazione: Ciò non richiede parametri appresi o sintonizzazione di soglie. Il valore fuso per un token mantenuto $i$ è calcolato come una media pesata dall'attenzione e dal gate del valore originale e dai contributi dei token espulsi, scalata dai rispettivi gate.

2. Attention Compensation (Per affrontare l'Attention Sag)
Per correggere la sistematica sotto-attenzione delle posizioni fuse, SelKV applica un bias dei logit durante la decodifica.

• Meccanismo: A differenza dei metodi precedenti che utilizzano conteggi grezzi di fusione (che possono essere instabili ad alti rapporti di compressione), SelKV deriva un rapporto di attenzione ($R_{h,i}$) dalle statistiche di prefill. Questo rapporto stima la massa totale di attenzione che una posizione fusa dovrebbe ricevere rispetto al suo stato originale.
• Correzione: Un termine di bias $\alpha \cdot \log(R_i)$ viene aggiunto ai logit dell'attenzione durante la decodifica. Questo ricalibra la distribuzione softmax per garantire che le posizioni fuse ricevano la massa di attenzione appropriata, proporzionale alle informazioni che ora rappresentano.

Panoramica della Pipeline

La pipeline di SelKV procede attraverso sei fasi post-prefill:

1. Importance Scoring: I token vengono valutati in base al prodotto tra il peso dell'attenzione e la magnitudo del vettore di valore (seguendo VATP), piuttosto che solo sul peso dell'attenzione.
2. Token Selection: I top-$m_s$ token vengono selezionati per layer (e per head in MHA, o tramite unione per GQA) da mantenere, insieme a una finestra recente fissa.
3. Merge-Target Routing: I token espulsi sono indirizzati ai token mantenuti all'interno di bucket posizionali basati sui pattern di attenzione del prefill.
4. Selective Merging: Il soft cosine gate determina l'intensità della fusione per ogni token espulso.
5. Attention Compensation: Il bias dei logit viene calcolato e applicato.
6. RoPE Repositioning: Le chiavi (keys) vengono ri-incorporate in posizioni contigue (disabilitato per la valutazione del lungo contesto per evitare distorsioni posizionali).

3. Contributi Chiave

• Granularità Per-Token: L'introduzione di un soft cosine gate che consente la modulazione continua dell'intensità della fusione basata sulla similarità semantica, eliminando il difetto "tutto o niente" dei metodi precedenti.
• Compensazione dell'Attenzione Stabile: Un nuovo logit bias basato sul rapporto di attenzione che corregge l'attention sag in modo più robusto rispetto ai metodi basati sul conteggio grezzo, specialmente sotto forte compressione.
• Integrazione Training-Free: Il framework non richiede fine-tuning e può essere integrato in pipeline di compressione one-shot esistenti.
• Adattabilità dell'Architettura: Il metodo supporta architetture Multi-Head Attention (MHA) e Grouped-Query Attention (GQA), con gestione specifica delle unioni di selezione dei token per i modelli GQA.

4. Risultati Sperimentali

Il framework è stato valutato su LongBench (16 dataset in inglese) su tre modelli: LongChat-7B (MHA), LLaMA-3.1-8B (GQA) e Gemma-2-9B (GQA), mantenendo solo il 25% della cache KV.

• Performance vs Baselines: SelKV ha superato costantemente i baseline one-shot rappresentativi (SnapKV, LOOK-M, PyramidKV).
  • Sui modelli GQA, SelKV ha ottenuto la performance più vicina al baseline della cache completa, con gap medi di $-0.67$ (Gemma-2) e $-0.73$ (LLaMA-3.1).
  • Sui modelli MHA (LongChat), PyramidKV ha performato leggermente meglio nello specifico setting di 31.5k token, sebbene SelKV sia rimasto competitivo.
• Multi-Document QA: Notevolmente, SelKV ha superato il baseline della cache completa in diversi compiti di multi-document QA (es. HotpotQA e MuSiQue su LLaMA-3.1). Gli autori interpretano questo come prova che la fusione selettiva agisce come un filtro di attenzione implicito, sopprimendo il contesto distraente pur preservando le evidenze salienti.
• Efficienza: A una lunghezza di contesto di 100k token, la cache compressa ha fornito un accelerazione della decodifica di 3.3x rispetto al baseline della cache completa.
• Robustezza: Il metodo ha dimostrato robustezza attraverso vari rapporti di compressione (10%–90% di ritenzione) e categorie di task, inclusi riassunti, apprendimento few-shot e completamento del codice.

5. Significato e Rivendicazioni

Il paper posiziona SelKV come un passo avanti significativo nel rendere fattibile l'inferenza di LLM a lungo contesto senza sacrificare la qualità. Gli autori affermano che:

• La fusione selettiva è superiore alle strategie uniformi: Modulando l'intensità della fusione in base alla similarità della rappresentazione, SelKV evita la corruzione semantica causata dalla fusione indiscriminata.
• La compensazione dell'attenzione è critica: Senza correggere l'attention sag, i token fusi sono sottoutilizzati; la correzione basata sul rapporto di SelKV assicura che le informazioni preservate siano accessibili efficacemente.
• Filtraggio Implicito: La capacità di SelKV di superare i baseline della cache completa in complessi compiti di multi-document QA suggerisce che il processo di compressione stesso può migliorare le prestazioni del modello filtrando il rumore, una scoperta che sfida l'assunto che la compressione degradi sempre le capacità.

Gli autori concludono che SelKV offre una soluzione pratica e senza parametri per il deployment di LLM in ambienti con vincoli di memoria, mantenendo una qualità di generazione quasi lossless, particolarmente per le architetture basate su GQA.