Riepilogo Tecnico: HeadQ: Distorsione Visibile al Modello e Correzione nello Spazio dei Punteggi per la Quantizzazione del KV-Cache

1. Enunciato del Problema

I metodi attuali di quantizzazione del KV-cache ottimizzano prevalentemente per la ricostruzione dello spazio di archiviazione, tipicamente minimizzando l'Errore Quadratico Medio (MSE) tra chiavi e valori originali e quantizzati. Tuttavia, gli autori sostengono che questo obiettivo non è allineato con il modo in cui i meccanismi di attenzione consumano effettivamente i dati della cache.

• Le Chiavi vengono lette attraverso i logit (punteggi query-chiave), il che significa che solo specifiche proiezioni lineari delle chiavi influenzano la distribuzione dell'attenzione.
• I Valori vengono letti attraverso una lettura ponderata dall'attenzione ($\sum p_j v_j$).

Il documento ipotizza che l'errore persistente della cache debba essere misurato in coordinate visibili al modello piuttosto che nelle coordinate grezze di archiviazione euclidea. Una perturbazione può avere un MSE di archiviazione elevato ma un impatto nullo sull'attenzione (se giace in una direzione "nulla-softmax"), oppure un MSE di archiviazione basso ma danni catastrofici all'attenzione (se influenza la direzione "visibile al punteggio"). Questa discrepanza porta a strategie di quantizzazione subottimali che sprecano bit su componenti irrilevanti mentre falliscono nel proteggere quelli critici dal punto di vista comportamentale.

2. Fondamento Teorico: Geometria Visibile al Modello

Il documento stabilisce un quadro geometrico che distingue tra errore di archiviazione ed errore comportamentale:

2.1 Spazio dei Punteggi delle Chiavi (Geometria di Fisher)

Per una query fissa $q$, una perturbazione della chiave $\delta k_j$ induce una perturbazione del punteggio $\delta s_j = q^\top \delta k_j / \sqrt{d}$.

• Direzioni Nulle-Softmax: Una traslazione condivisa dal token (aggiungere un vettore costante $\mu$ a tutte le chiavi) aggiunge una costante a tutti i punteggi. Poiché il softmax è invariante rispetto agli spostamenti costanti, questa perturbazione è nulla-softmax e causa zero danni all'attenzione, nonostante un potenziale MSE di archiviazione elevato.
• Spazio Visibile al Punteggio: L'attenzione dipende solo dai coefficienti attivi centrati della perturbazione della chiave rispetto alla query.
• Informazione di Fisher: La geometria locale del secondo ordine della distribuzione softmax è catturata dalla Matrice di Informazione di Fisher. Il documento deriva che la divergenza KL tra le distribuzioni di attenzione originali e perturbate è localmente proporzionale a una varianza pesata da Fisher degli errori di punteggio (Eq. 6), non al MSE grezzo delle chiavi.

2.2 Geometria della Lettura dei Valori

I valori vengono consumati tramite una somma ponderata. Per pesi di attenzione $p$ fissi, l'errore di lettura è $\Delta o = \sum p_j \Delta v_j$.

• Il documento propone una distorsione del token pesata con $A^2$ come surrogato appropriato per l'errore dei valori. A differenza del MSE grezzo dei valori, questa metrica pesa la distorsione di ogni token con il quadrato della sua probabilità di attenzione ($p_j^2$), riflettendo il fatto che gli errori nei token fortemente attenzionati contribuiscono maggiormente all'errore di output.

3. Metodologia: HeadQ

HeadQ è un metodo di quantizzazione lato chiave progettato per preservare il residuo "visibile al punteggio" piuttosto che il vettore grezzo della chiave.

Panoramica dell'Algoritmo

1. Calibrazione: Per ogni livello e testa, viene appresa una base a rango basso $U_h$ (rango $r$) tramite PCA da query di calibrazione (post-RoPE).
2. Scrittura nella Cache:
   • Un quantizzatore standard a basso bit $B_b$ genera una chiave di base $\hat{k}^{base}_j$.
   • Viene calcolato il residuo $e_j = k_j - \hat{k}^{base}_j$.
   • Il residuo viene proiettato sulla base appresa: $z_j = Q_q(U_h^\top e_j)$, dove $Q_q$ è un quantizzatore laterale a basso bit.
   • La cache memorizza la chiave di base e il codice laterale $z_j$.
3. Lettura dalla Cache (Inferenza):
   • La query corrente $q$ viene utilizzata per dequantizzare il codice laterale.
   • Viene applicata una correzione additiva del logit:
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • Questa correzione mira direttamente all'errore del punteggio, bypassando la necessità di ricostruire l'intero vettore della chiave.

Distinzione Chiave: HeadQ non memorizza la query corrente; memorizza un codice residuo in una base appresa dalla calibrazione che la query di runtime può decodificare in una correzione del logit.

4. Contributi Chiave

1. Formulazione della Distorsione Visibile al Modello: Il documento definisce la distorsione della chiave come "errore visibile al punteggio modulo spostamenti costanti" e la distorsione del valore come "errore visibile alla lettura con un surrogato diagonale $A^2$". Ciò spiega perché l'MSE grezzo è un cattivo proxy per le prestazioni comportamentali.
2. Algoritmo HeadQ: Un algoritmo pratico di scrittura/lettura nella cache che memorizza un codice laterale residuo a rango basso in una base di query e lo applica come correzione additiva del logit.
3. Catena di Evidenza Falsificabile:
   • Superiorità della Metrica: L'errore nello spazio di Fisher/di punteggio predice la divergenza KL dell'attenzione in modo significativamente migliore (guadagno di Spearman $\rho \approx 0.62\text{--}0.69$) rispetto all'MSE grezzo delle chiavi su sei modelli.
   • Controesempi: Esperimenti accoppiati mostrano che un quantizzatore con MSE di archiviazione inferiore può avere una KL di attenzione e un errore di output superiori, dimostrando che l'MSE di archiviazione non è un proxy sufficiente.
   • Interventi nello Spazio Nullo: Il ripristino di una componente media "nulla-softmax" (che l'attenzione densa ignora) aumenta nettamente il danno di Fisher per i quantizzatori a griglia fissa non equivarianti, dimostrando che lo spreco di coordinate di archiviazione rompe le prestazioni comportamentali.
   • Deployabilità: Le basi PCA di query apprese dai dati di calibrazione approssimano con successo il residuo attivo, mentre basi casuali con lo stesso budget falliscono.
   • Direzionalità: L'inversione del segno della correzione HeadQ (controllo "segno errato") causa un fallimento catastrofico, confermando che il metodo funziona correggendo l'oggetto punteggio specifico nella direzione corretta.

5. Risultati Sperimentali

Gli esperimenti sono stati condotti su WikiText-103 utilizzando sei modelli (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral).

• Esperimenti Solo-K (Valori Densi):
  • Sulle righe a 2 bit più danneggiate, HeadQ (rango 4) rimuove circa 84–94% della perplessità eccessiva (PPL) rispetto a una baseline simile a KIVI per modelli come Qwen2, Pythia e GPT-2.
  • TinyLlama mostra un miglioramento sostanziale; Mistral è già vicino alle prestazioni dense in questo protocollo.
• Composizione KV Completa (2-bit):
  • Combinare HeadQ (lato chiave) con una politica di valori pesata con $A^2$ migliora le prestazioni su tutti e sei i modelli a precisione 2-bit.
• Analisi delle Anomalie:
  • Pythia-160M ha mostrato una correlazione più debole tra Fisher e KL in specifici regimi a bassa entropia di "flip del percorso" (dove piccole perturbazioni invertono il token top-1). Questa anomalia è stata localizzata a una specifica famiglia di confini e si è indebolita nei checkpoint Pythia più grandi (410M, 1B).

6. Significato e Affermazioni

Il documento afferma che l'adattamento dell'operatore è la lezione di progettazione critica per la quantizzazione del KV-cache:

• Le chiavi sono oggetti punteggio; la compressione deve preservare i residui visibili alla query ed evitare di sprecare risoluzione sugli spostamenti nulli-softmax.
• I valori sono oggetti lettura; la compressione dovrebbe dare priorità all'errore pesato con $A^2$.

Gli autori affermano esplicitamente che questo lavoro è meccanicistico e diagnostico, non una rivendicazione di un sistema di servizio pronto per la produzione. Non affermano di aver ottimizzato kernel impaccati, latenza o throughput. Il significato risiede in:

1. Correggere l'obiettivo teorico della quantizzazione del KV-cache dalla ricostruzione di archiviazione alla distorsione visibile al modello.
2. Fornire un metodo deployabile (HeadQ) che ottiene un sostanziale recupero del PPL in regimi a basso bit mirando alla corretta geometria dell'errore.
3. Dimostrare che le metriche standard di archiviazione (MSE) possono essere fuorvianti dal punto di vista comportamentale, rendendo necessario un passaggio verso proxy nello spazio dei punteggi.

Il documento conclude che il passo successivo è estendere questa visione adattata all'operatore in un codec KV completo e completamente ottimizzato, piuttosto che trattare l'MSE di archiviazione come l'obiettivo predefinito.