Riepilogo Tecnico: Un JOLT per la KV Cache

Problema

Nell'inferenza di transformer autoregressivi, la KV Cache rappresenta il costo di memoria dominante, scalando come $O(\text{batch} \times T \times L \times n_h \times d_h)$. Per contesti lunghi o batch elevati, la KV Cache, piuttosto che i pesi del modello, diventa il principale collo di bottiglia della memoria e limita il throughput. I metodi di compressione esistenti si dividono in due famiglie, entrambe le quali non riesce a sfruttare appieno la ridondanza strutturale della cache:

1. Metodi a basso rango (Low-rank): Questi fattorizzano porzioni bidimensionali (ad esempio, matrici per testa o blocchi cross-layer) ma si impegnano su un singolo asse di ridondanza, ignorando gli altri.
2. Metodi di quantizzazione: Questi riducono la larghezza di bit di ogni voce, ma sono limitati da bit-rate fissi (ad esempio, 4-bit), creando un "pavimento di compressione" che impedisce di raggiungere rapporti di compressione intermedi (2–3×) dove la qualità è spesso preservata.

L'osservazione fondamentale che motiva questo lavoro è che la KV cache a un livello è un tensore di terzo ordine ($n_h \times T \times d_h$) dove gli assi portano quantità di ridondanza molto diverse. Nello specifico, gli assi testa (head) e livello (layer) sono ampiamente incompressibili, mentre gli assi token e feature contengono la maggior parte della struttura comprimibile. Inoltre, empiricamente, i valori (values) sono 2–3× più difficili da comprimere rispetto alle chiavi (keys) a causa delle differenze nelle loro proprietà spettrali.

Metodologia: JoLT

Il metodo proposto, JoLT (Joint Tucker and JL-residual allocation), affronta queste sfide trattando la cache come un tensore di terzo ordine e applicando una strategia di compressione mirata.

1. Decomposizione parziale di Tucker

A differenza delle decomposizioni tensoriali complete che comprimono tutti gli assi, JoLT applica una decomposizione di Tucker parziale. Trunca solo gli assi token e feature (che portano ridondanza) mantenendo intatti gli assi head e layer (che sono incompressibili). Ciò viene ottenuto tramite la ST-HOSVD (Sequentially Truncated Higher-Order SVD).

• Motivazione: Comprimere gli assi head o layer comporta un errore significativo, mentre troncare gli assi token e feature cattura la maggior parte della ridondanza.

2. Residuo ruotato JL

Un backbone puramente a basso rango è insufficiente per una fedeltà quasi lossless, specialmente per i valori che hanno uno spettro piatto dove la troncatura scarta energia significativa. JoLT ripristina questa energia scartata utilizzando un residuo a basso bit ruotato tramite Johnson-Lindenstrauss (JL):

• Il residuo di troncatura $R = X - \tilde{X}$ viene calcolato.
• Una rotazione JL casuale viene applicata al residuo.
• Il residuo ruotato viene quantizzato uniformemente a $b \in \{0, 2, 4, 8\}$ bit.
• La decodifica prevede l'inversione della rotazione. Questo residuo agisce come una leva per convertire un backbone lossy in uno quasi lossless, particolarmente prezioso per i valori.

3. Allocazione congiunta Lagrangiana

Il metodo impiega un singolo Lagrangiano duale per allocare congiuntamente le risorse tra le due fasi (ranghi Tucker e bit-width dei residui) e tra diversi target (Chiavi vs. Valori, Gruppi di Livelli).

• Obiettivo: Minimizzare l'errore di ricostruzione totale soggetto a un budget globale di byte.
• Meccanismo: La variabile duale permette all'allocatore di spostare bit tra ranghi e residui, e tra chiavi e valori, sfruttando dinamicamente l'asimmetria nella comprimibilità (ad esempio, finanziando pesantemente il residuo dei valori e riducendo i bit del residuo delle chiavi).
• Granularità: L'allocazione viene eseguita per gruppo di livelli e separatamente per Chiavi e Valori.

4. FlashJoLT (Variante veloce)

Per affrontare il costo computazionale della SVD esatta su contesti lunghi, FlashJoLT sostituisce la SVD esatta sul modo token con una SVD a basso rango randomizzata.

• Calcola solo le prime $q_{cap}$ direzioni del modo token.
• Impiega un conteggio della massa della coda (tail-mass accounting), aggiungendo la coda spettrale omessa come un valore singolare del residuo sintetico per evitare che l'allocatore sottostimi l'energia catturata da un dato rango.
• Questa variante ottiene un'accelerazione di 5–13× nel tempo di compressione con una perdita di qualità trascurabile nella "zona free".

Contributi Chiave

1. Analisi Spettrale: Uno studio empirico che identifica che gli assi head e layer sono incompressibili, mentre gli assi token e feature contengono ridondanza, con i valori che sono significativamente più difficili da comprimere rispetto alle chiavi.
2. Algoritmo JoLT: Un nuovo compressore che combina la decomposizione di Tucker parziale con un residuo ruotato JL, accoppiato tramite un singolo duale Lagrangiano per l'allocazione congiunta di ranghi e bit.
3. FlashJoLT: Una variante SVD-randomizzata che riduce drasticamente il tempo di compressione preservando la qualità.
4. Validazione Empirica: Uno studio completo sui modelli GQA (Mistral-7B) e MHA (LLaMA-2-13B) che stabilisce una "free zone" di compressione quasi lossless.

Risultati

La "Free Zone" (Compressione 2–3×)

Il paper identifica un rapporto di compressione di 2–3× come una "free zone" dove la qualità del modello è preservata entro il rumore statistico del baseline non compresso:

• Perplexity: Sia su Mistral-7B (GQA) che su LLaMA-2-13B (MHA), il drift della perplexity è trascurabile a 2–3× per lunghezze di contesto fino a 1024.
• Errore di Ricostruzione: A 2× di compressione, JoLT raggiunge un errore di Frobenius relativo di 0.009 (Chiavi) e 0.006 (Valori). Questo è circa un ordine di grandezza inferiore rispetto alla quantizzazione a 4-bit e ai baseline SVD cross-layer.
• Task a valle (Downstream Tasks):
  • GSM8K: I modelli compressi (2–3×) rimangono entro l'intervallo di confidenza al 95% del baseline.
  • RULER (Needle-in-a-Haystack): Mistral mantiene il 100% dell'accuratezza di retrieval a 2–3×; LLaMA eguaglia il suo ceiling di baseline.

Divergenza dell'Architettura

Oltre la free zone di 2–3×, le due architetture divergono:

• Mistral (GQA): Degrada con grazia, perdendo circa il 4% di perplexity per ogni step di rapporto intero oltre il 4×.
• LLaMA (MHA): Degrada bruscamente tra 4× e 5×, con la perplexity che aumenta di circa il 27% a 4× e del ~68% a 5×. I task di ragionamento (GSM8K) per LLaMA collassano significativamente ad alti rapporti.

Efficienza

• Memoria: L'impronta persistente della KV cache si riduce linearmente con il rapporto di compressione (ad esempio, da 268.4 MB a 134.6 MB a 2× per Mistral a ctx-2048). La memoria di processo di picco rimane limitata da pesi/attivazioni.
• Velocità: FlashJoLT offre un'accelerazione di 5–13× nel tempo di compressione rispetto al backbone esatto.

Significato e Claim

Il paper sostiene che JoLT fornisce una compressione quasi lossless di 2–3× della KV cache, rimuovendo efficacementamente il collo di bottiglia della memoria per l'inferenza a lungo contesto senza sacrificare la qualità del modello. Trattando la cache come un tensore e sfruttando la specifica asimmetria spettrale tra chiavi e valori, JoLT supera i metodi esistenti di low-rank e quantizzazione nel regime di compressione intermedio.

Gli autori posizionano il metodo come una leva diretta per ridurre il costo del serving di contesti lunghi. Notano due confini primari:

1. Limiti dell'Architettura: Il metodo è robusto per i modelli GQA ma affronta una degradazione brusca per i modelli MHA oltre il 4×, suggerendo che la banda 2–3× è il consiglio sicuro e agnostico rispetto al modello.
2. Costo di Deployment: Sebbene lo storage sia ridotto, la decodifica richiede attualmente la ricostruzione della cache dai fattori ad ogni step. Gli autori identificano lo sviluppo di un kernel fuso di ricostruzione e attenzione (fused reconstruction-and-attention kernel) come il prossimo passo critico per raggiungere la parità di inferenza e realizzare appieno il potenziale del metodo.