Riepilogo Tecnico: STAR-KV

Problematica

La rapida scalabilità dei Large Language Models (LLM) per supportare finestre di contesto di oltre 100K token ha introdotto severi colli di bottiglia a livello di sistema durante l'inferenza, guidati principalmente dalla cache Key-Value (KV). La cache KV memorizza i tensori per i token precedentemente elaborati per evitare la ricalcolazione, risultando in un ingombro di memoria che cresce linearmente con la lunghezza del contesto. In scenari come LLaMA-3.1-8B con un contesto di 128K, la cache KV può consumare fino all'81% della memoria totale della GPU, superando spesso lo storage richiesto per i pesi del modello.

Sebbene la decomposizione a basso rango sia emersa come un metodo promettente per comprimere la cache KV sfruttando la ridondanza nella dimensione nascosta, gli approcci esistenti affrontano due sfide fondamentali:

1. Selezione del Rango Fissa o Euristica: I metodi precedenti si affidano a ranghi uniformi o euristiche offline (ad esempio, l'informazione di Fisher), che faticano a mantenere la qualità del modello sotto una compressione aggressiva. Spingere la compressione oltre livelli moderati (30–50%) porta tipicamente a un significativo degrado dell'accuratezza.
2. Overhead di Ricostruzione: La ricostruzione dei tensori compressi di chiave e valore durante l'inferenza introduce un calcolo aggiuntivo, aumentando la latenza. Inoltre, i metodi standard a basso rango spesso non tengono conto delle diverse sensibilità delle proiezioni di chiave e valore o delle specifiche sfide della quantizzazione delle risultanti rappresentazioni latenti.

Metodologia

Gli autori propongono STAR-KV (Soft Thresholding for Adaptive Rank KV Cache Compression), un framework progettato per ottenere un controllo del rango fine e adattivo e una quantizzazione efficiente. La metodologia consiste in tre componenti principali:

1. Selezione del Rango Adattiva tramite Soft Thresholding

Inveve di utilizzare ranghi fissi, STAR-KV parametrizza la selezione del rango utilizzando una soglia apprendibile applicata ai valori singolari delle matrici di proiezione di chiave e valore.

• Soglia Differenziabile: Per consentire l'ottimizzazione basata sul gradiente, gli autori sostituiscono la soglia rigida (hard threshold) non differenziabile con un operatore di soglia morbida (soft thresholding) basato su una funzione tangente iperbolica spostata ($\tanh$). Ciò consente al modello di apprendere soglie ottimali ($\alpha$) per ogni blocco del decoder e ogni testa di attenzione indipendentemente.
• Obiettivo di Addestramento: Il framework ottimizza una funzione di perdita congiunta che comprende una Perdita di Distillazione della Conoscenza (per corrispondere alla distribuzione di output del modello teacher) e una Perdita di Compressione Adattiva. La perdita di compressione è formulata come $\sum e^{-\alpha_i}$, il che incoraggia naturalmente una compressione aggressiva nelle prime fasi dell'addestramento, permettendo al contempo al modello di stabilizzare le prestazioni all'aumentare delle soglie.
• Inferenza: Una volta addestrate su un piccolo set di calibrazione (<6 ore GPU), le soglie morbide vengono convertite in soglie rigide offline, e le matrici di proiezione vengono troncate prima dell'inferenza.

2. Strategia di Decomposizione Ibrida

Il documento analizza i compromessi tra Decomposizione Congiunta (JD) (fattorizzazione delle teste concatenate) e Decomposizione per Testa (HD) (fattorizzazione delle teste indipendentemente).

• Analisi di Sensibilità: L'analisi dei valori singolari rivela che le proiezioni di valore presentano generalmente errori di approssimazione più elevati rispetto alle proiezioni di chiave sotto lo stesso vincolo di rango. Al contrario, la JD offre un errore di approssimazione inferiore ma comporta un costo computazionale di ricostruzione $h \times$ (numero di teste) superiore rispetto alla HD.
• Schema Ibrido: STAR-KV adotta un approccio ibrido:
  • Proiezioni di Valore: Utilizzano la Decomposizione Congiunta (JD) per preservare la fedeltà, poiché i valori sono più sensibili agli errori di compressione.
  • Proiezioni di Chiave: Utilizzano la Decomposizione per Testa (HD) per minimizzare l'overhead di ricostruzione, poiché le chiavi sono meno sensibili.
  • Ottimizzazione: Per compensare l'overhead della JD sui valori, il framework utilizza l'assorbimento dei pesi, riordinando le operazioni per fondere le matrici di ricostruzione nelle proiezioni di output, riducendo significativamente i FLOP.

3. Quantizzazione a Precisione Mista Consapevole del Basso Rango

Per comprimere ulteriormente le rappresentazioni latenti, STAR-KV affronta il problema dei canali outlier che derivano dalle dimensioni singolari ordinate dei fattori a basso rango.

• Hadamard Block-wise: Invece di applicare una trasformata di Hadamard globale, il metodo suddivide la dimensione latente in blocchi "outlier" (primi 20% dei canali con magnitudo elevata) e "inlier" (restante 80%).
• Precisione Mista: Trasformate di Hadamard indipendenti vengono applicate a ciascun blocco, seguite da una quantizzazione a precisione mista. Il blocco outlier è quantizzato a una precisione maggiore (ad esempio, 4-bit), mentre il blocco inlier è quantizzato a una precisione inferiore (ad esempio, 3-bit).
• Zero Overhead: I fattori di Hadamard sono fusi nelle matrici di proiezione decomposte, comportando nessun costo aggiuntivo durante l'esecuzione.

Risultati Chiave

STAR-KV è stato valutato su molteplici famiglie di LLM (LongChat, LLaMA-2, LLaMA-3) e benchmark (ragionamento del senso comune, LongBench, RULER).

• Compressione vs. Accuratezza: STAR-KV raggiunge fino al 75% di compressione della cache KV con un degrado dell'accuratezza trascurabile rispetto ai modelli non compressi. Su LongChat-7B-v1.5, raggiunge il 60% di compressione con un'accuratezza che eguaglia il baseline, superando i metodi allo stato dell'arte come Palu e ReCalKV.
• Compressione Combinata: Quando combinato con la proposta quantizzazione a precisione mista a 3.2-bit, STAR-KV raggiunge fino a 20× di riduzione complessiva della cache KV.
• Prestazioni in Contesti Lunghi: Su LongBench e RULER (sequenza di 4K token), STAR-KV mantiene un'accuratezza superiore ad alti tassi di compressione (60%) rispetto a Palu, che degrada significativamente al 50% e spesso richiede una compressione inferiore (30%) per rimanere competitivo.
• Velocità: Implementato con kernel GPU basati su Triton, STAR-KV fornisce:
  • Fino a 6.9× di velocità in più nel modulo di attenzione rispetto a PyTorch SDPA.
  • Fino a 3.1× di throughput di generazione end-to-end superiore in scenari di contesto lungo (128K) o ad alto batch.

Significato e Rivendicazioni

Il documento afferma che STAR-KV stabilisce un nuovo fronte di Pareto per il compromesso accuratezza-compressione nella gestione della cache KV. La sua importanza risiede in:

1. Granularità Adattiva: Superare la selezione del rango statica o euristica verso un meccanismo di controllo apprendibile e fine-grained che si adatta alla specifica sensibilità di diversi strati e teste.
2. Efficienza: Bilanciare con successo il compromesso tra errore di approssimazione e overhead di ricostruzione attraverso una strategia di decomposizione ibrida e l'assorbimento dei pesi.
3. Sinergia: Dimostrare che la decomposizione a basso rango e la quantizzazione possono essere combinate sinergicamente tramite trasformate di Hadamard block-wise per raggiungere una compressione estrema (20×) senza sacrificare le prestazioni del modello.
4. Praticità: Fornire una soluzione post-training che non richiede cambiamenti architettonici o un addestramento esteso, rendendola applicabile agli LLM esistenti con un costo di calibrazione minimo.

Gli autori sottolineano che il loro approccio consente una compressione aggressiva (fino al 75% per il basso rango, 20× totale) mantenendo o addirittura migliorando l'accuratezza rispetto ai baseline non compressi e ai metodi di compressione esistenti, affrontando efficacemente i colli di bottiglia di memoria e banda per l'inferenza a lungo contesto.