Riepilogo Tecnico: Inferenza Statistica e Misure di Qualità delle Quantizzazioni della KV Cache Ispirate a TurboQuant

Enunciato del Problema

La cache KV (Key-Value) è il principale collo di bottiglia della memoria nell'inferenza dei transformer, crescendo linearmente con la lunghezza del contesto. La compressione tramite quantizzazione è essenziale, tuttavia la quantizzazione scalare standard (MSE) spesso non riesce a preservare le proprietà geometriche richieste dai meccanismi di attenzione. Sebbene TurboQuant proponga una pipeline a due stadi (rotazione casuale + quantizzazione scalare + schizzo QJL a 1 bit per la correzione del bias), l'allocazione ottimale dei bit tra le cache Chiave (K) e Valore (V) rimane incerta. Nello specifico, non è noto se applicare la correzione del residuo QJL alle Chiavi (oltre che ai Valori) migliori o degradi la qualità dell'attenzione, in particolare sotto diversi budget di bit e distribuzioni dei dati (ad esempio, a basso rango, a code pesanti).

Le metriche di valutazione esistenti, come l'Errore Quadratico Medio (MSE) o persino la divergenza KL, spesso confondono errori di scala e direzione o non riescono a catturare fallimenti catastrofici di "spostamento dell'argmax", dove una singola chiave dominante mal quantizzata collassa la distribuzione di attenzione.

Metodologia

Quadro Teorico

Gli autori basano la loro analisi sulla geometria delle sfere ad alta dimensione. Stabiliscono che per un vettore distribuito uniformemente su un'ipersfera $S^{d-1}$, la distribuzione marginale delle coordinate segue una distribuzione Beta $\text{Beta}(\frac{d-1}{2}, \frac{d-1}{2})$. All'aumentare della dimensione $d \to \infty$, questa converge a una Gaussiana.

• Disuguaglianza di Jensen e Softmax: Il documento evidenzia una non linearità critica: la funzione softmax amplifica la varianza nei prodotti interni delle Chiavi. Anche se gli errori di quantizzazione hanno media zero, la convessità della funzione esponenziale ($e^x$) introduce un bias (bias di Jensen) che gonfia i pesi di attenzione per le chiavi con alta varianza di quantizzazione.
• Asimmetria K-V: Il tensore delle Chiavi entra nel meccanismo di attenzione tramite prodotti interni (non lineari), mentre il tensore dei Valori entra tramite una somma ponderata lineare. Di conseguenza, l'infiammazione della varianza nelle Chiavi è catastrofica, mentre una correzione del bias non distorta nei Valori è benefica.

Schemi di Quantizzazione

Tre schemi sono confrontati sotto un vincolo di budget di bit rigido e uguale ($n$ bit per vettore):

1. KV: Quantizzazione scalare MSE (Lloyd-Max) sia su K che su V. Nessuna rotazione, nessun schizzo.
2. KQV: Applica una Trasformata di Walsh-Hadamard (WHT) a entrambi. K è quantizzato tramite MSE scalare ($n$ bit). V è quantizzato tramite MSE scalare ($n-1$ bit) + schizzo QJL a 1 bit.
3. QKQV: Applica WHT a entrambi. Sia K che V utilizzano ($n-1$ bit) + schizzo QJL a 1 bit.

Metriche di Valutazione

Per andare oltre le metriche scalari, il documento introduce un Framework di Errore 6D che scompone l'errore in Scala (SNR) e Direzione ($\phi$) per tre siti: cache delle Chiavi, cache dei Valori e Output ($T$).

• Piano Shannon-$\phi$: Traccia l'errore di scala rispetto all'errore di direzione.
• Strumenti Statistici: U di Mann-Whitney (correlazione biseriale di rango), Kolmogorov-Smirnov (separazione CDF) e Distanza Energetica (metrica 6D).
• Proxy Meccanicistici: Divergenza KL tra le distribuzioni softmax di riferimento e quantizzate, e recall di instradamento top-5.

Contributi Chiave

1. Dimostrazione Teorica del Danno del QJL sulle Chiavi

Il documento dimostra che sostituire la quantizzazione scalare sulle Chiavi con uno schizzo QJL a 1 bit gonfia la varianza dei prodotti interni di un fattore di almeno $\pi/2$ (circa 1,57) rispetto al residuo scalare, e di un fattore $2\pi$ quando si confronta uno schema scalare completo a $n$ bit con uno schema scalare a $(n-1)$ bit + QJL a 1 bit.

• Meccanismo: Sebbene il QJL fornisca uno stimatore non distorto del residuo, scarta le informazioni di grandezza, sostituendole con bit di segno. La varianza di questo stimatore è $\frac{\pi}{2} \frac{\|r\|^2}{d}$.
• Conseguenza: Questa varianza gonfiata è amplificata in modo non lineare dall'esponenziale softmax (tramite la disuguaglianza di Jensen), portando a una distorsione sistematica dei pesi di attenzione. Pertanto, applicare QJL alle Chiavi è strutturalmente inferiore alla quantizzazione scalare sotto budget equi.

2. L'Asimmetria K-V e lo Schema Ottimale (KQV)

I risultati empirici confermano che KQV (QJL solo su V) è superiore a QKQV (QJL su entrambi) in tutte le distribuzioni e i budget testati.

• Divergenza KL: Al budget di 4 bit praticamente dominante ($n=4$), KQV ottiene una divergenza KL 2,25$\times$ inferiore rispetto a QKQV ($p < 10^{-23}$).
• Instradamento: Una KL inferiore in KQV si correla direttamente con un recall di instradamento top-5 più alto e un errore di direzione dell'output più basso.
• Conclusione: Lo schizzo QJL a 1 bit dovrebbe essere riservato alla cache dei Valori (dove corregge il bias senza penalità di varianza) e non dovrebbe mai essere applicato alla cache delle Chiavi (dove introduce una varianza fatale).

3. La Patologia a Basso Rango

Il documento identifica una modalità di fallimento critica per le distribuzioni di Chiavi a basso rango che le metriche scalari non rilevano.

• Marginale vs Congiunta: Sebbene WHT renda le coordinate marginali distribuite secondo Beta (corrispondendo al codebook), non distrugge la struttura congiunta a basso rango.
• KL Catastrofica: Ad alte dimensioni ($d=1024$), le chiavi a basso rango producono una divergenza KL ordini di grandezza peggiore rispetto alle chiavi a code pesanti, nonostante abbiano istogrammi marginali identici. Il softmax amplifica piccoli errori sulle chiavi dominanti all'interno del sottospazio, causando fallimenti di "spostamento dell'argmax".
• Implicazione: La quantizzazione basata sulla rotazione standard è insufficiente per le matrici di attenzione a basso rango; la struttura congiunta deve essere affrontata.

4. Incrocio Dipendente dal Budget

Un risultato controintuitivo è un incrocio nella qualità di ricostruzione geometrica delle Chiavi:

• A bassi budget ($n \in \{2, 3\}$), QKQV (QJL su K) talvolta ottiene una migliore ricostruzione geometrica delle Chiavi (errore di direzione inferiore) rispetto a KQV.
• Tuttavia, questo miglioramento geometrico non si traduce in una migliore qualità di attenzione (KL o instradamento) perché l'infiammazione della varianza danneggia comunque il softmax.
• A $n=4$, KQV vince su tutte le misure. Questo incrocio è invariante rispetto al rango e al peso della coda, suggerendo un compromesso fondamentale tra precisione scalare e correzione del residuo JL.

Risultati

• Dominanza di KQV: A $n=4$, KQV supera QKQV e KV in divergenza KL, errore geometrico delle Chiavi e distanza 6D in tutte le distribuzioni (Gaussiana, code pesanti, basso rango, focalizzata).
• QJL sulle Chiavi è Dannoso: Sostituire bit scalari con QJL sulle Chiavi aumenta costantemente la divergenza KL a causa dell'infiammazione della varianza, eccetto in scenari specifici di ricostruzione geometrica a basso budget che non beneficiano l'output finale.
• Robustezza alle Code Grasse: WHT fornisce protezione incondizionata contro la concentrazione marginale (code pesanti). KQV mantiene un errore di direzione stabile ($\epsilon_{dir}^K \approx 0,027$) anche all'aumentare del peso della coda ($\nu \to 1,05$), mentre KV fallisce catastroficamente.
• Fallimento a Basso Rango: Le distribuzioni a basso rango causano il degrado più severo nella qualità dell'attenzione, superando di gran lunga i problemi delle code pesanti, a causa della preservazione della struttura del sottospazio congiunta nonostante la normalizzazione marginale.

Significato e Affermazioni

Il documento afferma di fornire una spiegazione statistica e geometrica rigorosa del motivo per cui l'applicazione originale di TurboQuant del QJL alle Chiavi (se estesa simmetricamente) sarebbe dannosa.

1. Raffinamento delle Metriche: Sostiene che la KL media è una metrica ottimistica; la KL nel caso peggiore e il recall di instradamento sono le vere misure di affidabilità per la quantizzazione della cache KV.
2. Principio di Progettazione: Stabilisce che le cache K e V richiedono strategie di quantizzazione fondamentalmente diverse a causa della non linearità del softmax. Il QJL è uno strumento di correzione del bias adatto per output lineari (V) ma è uno strumento di iniezione di varianza che danneggia gli output non lineari (K).
3. Limite delle Garanzie Marginali: Dimostra che garantire che le distribuzioni marginali corrispondano al codebook (tramite WHT) è necessario ma insufficiente; la struttura congiunta (in particolare il basso rango) detta la modalità di fallimento ultima.
4. Problema Aperto: Il documento identifica l'incrocio dipendente dal budget nella ricostruzione geometrica come un problema aperto di rate-distorsione, notando che mentre QKQV vince geometricamente a bit molto bassi ($n=2,3$), perde funzionalmente a causa dell'amplificazione di varianza di Jensen.

Il lavoro conclude che una quantizzazione robusta della cache KV richiede una gestione esplicita delle strutture a basso rango (ad esempio tramite SVD o codebook consapevoli del sottospazio) e che lo schema KQV (WHT + Scalare su K, WHT + Scalare + QJL su V) è la configurazione ottimale per i budget di bit attuali.