Riepilogo Tecnico: RATEQUANT

Enunciato del Problema

I Large Language Models (LLM) richiedono la memorizzazione in cache di tutte le coppie chiave-valore (KV) precedentemente calcolate durante la generazione, creando un collo di bottiglia nella memoria che scala linearmente con la lunghezza della sequenza, la dimensione del batch e la profondità del modello. Sebbene la quantizzazione della cache KV riduca questo ingombro, i metodi esistenti assegnano uniformemente la stessa larghezza in bit a ogni testa di attenzione. Questo approccio ignora l'eterogeneità ben documentata nell'importanza delle teste.

I recenti tentativi di allocazione in precisione mista a livello di strato o canale si basano su regole euristiche legate a quantizzatori specifici, privi di garanzie teoriche e trasferibilità. Gli autori identificano un ostacolo più profondo, precedentemente non riconosciuto: il mismatch del modello di distorsione. Diversi quantizzatori seguono curve distorsione-tasso distinte, modellate come $D(b) = \alpha \beta^{-b}$. Il tasso di decadimento $\beta$ varia significativamente tra i diversi design (ad esempio, da 3.6 in TURBOQUANT a 5.3 in QuaRot). Applicare in modo ingenuo una strategia di allocazione in precisione mista derivata dal modello di distorsione di un quantizzatore a un altro inverte l'ordinamento del guadagno marginale delle teste di attenzione, spesso risultando in prestazioni peggiori rispetto alla quantizzazione uniforme.

Metodologia

RATEQUANT affronta queste sfide formalizzando l'allocazione dei bit per la cache KV testa per testa come un problema di ottimizzazione distorsione-tasso. Il framework è composto da quattro fasi:

1. Formulazione del Problema e Allocazione Ottimale

Gli autori modellano l'Errore Quadratico Medio (MSE) della quantizzazione testa per testa come una funzione esponenziale della larghezza in bit: $D(b) = \alpha \beta^{-b}$. L'obiettivo è minimizzare la somma pesata delle distorsioni su $N$ gruppi di quantizzazione (teste) soggetto a un budget totale di bit $B$:
$$\min_{b} \sum_{i=1}^{N} w_i \cdot D(b_i) \quad \text{soggetto a} \quad \sum b_i = B$$
dove $w_i$ rappresenta la sensibilità della testa $i$.

Utilizzando la stazionarietà lagrangiana, gli autori derivano una soluzione in forma chiusa tramite reverse waterfilling:
$$b^*_i = \bar{b} + \frac{\ln w_i - \ln \bar{w}}{\ln \beta}$$
dove $\bar{b}$ è la larghezza in bit media e $\bar{w}$ è la media geometrica delle sensibilità. Questa soluzione assegna più bit alle teste con sensibilità più elevata. Teoricamente, la riduzione della distorsione ottenibile è limitata dal rapporto tra la media aritmetica e la media geometrica (AM/GM) delle sensibilità delle teste.

2. Stima della Sensibilità

Per determinare i pesi $w_i$, RATEQUANT impiega la sensibilità basata sul gradiente. Calcola le norme quadrate dei gradienti degli output di proiezione KV rispetto alla perdita di modellazione linguistica causale:
$$w_{l,h} = \mathbb{E}_{x \sim \mathcal{D}} \left[ \frac{1}{T} \sum_{t=1}^{T} \left\| \frac{\partial \mathcal{L}}{\partial K_{l,h,t}} \right\|^2 \right]$$
Gli autori dimostrano che i proxy basati sul gradiente sono superiori alle metriche basate sull'attivazione (ad esempio, la magnitudine del canale), che misurano l'amplificazione dell'errore piuttosto che l'impatto diretto sulla perdita. La sensibilità basata sul gradiente rivela un pattern a forma di U di importanza (alta sensibilità negli strati iniziali e finali), mentre le metriche basate sull'attivazione mostrano un aumento monotono con la profondità.

3. Allocazione Interi ed Estensioni

Poiché le larghezze in bit devono essere intere, la soluzione continua è raffinata utilizzando un algoritmo greedy di guadagno marginale (Algoritmo 1). Questo algoritmo assegna iterativamente bit alla testa che offre la maggiore riduzione della distorsione pesata fino a esaurimento del budget.

Due estensioni critiche rendono RATEQUANT indipendente dal quantizzatore:

• Calibrazione della Distorsione: Invece di assumere un $\beta$ universale, RATEQUANT adatta un modello di distorsione per quantizzatore ($\alpha, \beta$) utilizzando un piccolo set di calibrazione (16 sequenze). Questo previene la modalità di fallimento del "mismatch del modello di distorsione".
• Separazione K/V: Riconoscendo che chiavi e valori spesso utilizzano schemi di quantizzazione diversi (ad esempio, per canale simmetrico per le chiavi contro per token asimmetrico per i valori in KIVI), il framework li tratta come $2N$ componenti indipendenti con curve di distorsione distinte. Ciò permette all'algoritmo di scoprire budget di bit asimmetrici (ad esempio, assegnare più bit alle chiavi rispetto ai valori).

4. Pipeline

Il processo comporta costi offline una tantum:

1. Calibrazione della Sensibilità: Passi in avanti e all'indietro su un piccolo dataset (~1.6s per un modello da 8B).
2. Modellazione della Distorsione: Adattamento di $\alpha$ e $\beta$ a 5 larghezze in bit (<0.1s).
3. Allocazione: Esecuzione dell'algoritmo greedy (<0.01s).
4. Inferenza: Applicazione delle larghezze in bit statiche con zero overhead di runtime.

Risultati Chiave

Gli esperimenti sono stati condotti sui modelli Qwen3 (4B, 8B, 32B) utilizzando la perplessità (PPL) su WikiText-2 e compiti downstream (ARC-Challenge, HellaSwag, ecc.).

• Risoluzione del Mismatch: Applicare un modello di distorsione non corrispondente (ad esempio, utilizzare il $\beta$ di TURBOQUANT per KIVI) ha degradato la PPL di KIVI da 49.3 a 87.0 a 2.5 bit medi. RATEQUANT calibrato ha risolto questo problema, riducendo la PPL a 14.9 (una riduzione del 70%).
• Miglioramenti delle Prestazioni:
  • KIVI (2.5 bit): RATEQUANT ha migliorato la PPL da 49.3 a 14.9.
  • QuaRot (2.5 bit): Ha migliorato la PPL di 6.6 punti.
  • TURBOQUANT (4.0 bit): Ha recuperato il 66% del degrado indotto dalla quantizzazione per Qwen3-8B.
• Separazione K/V: Per KIVI a 2.5 bit, l'allocazione ottimale scoperta da RATEQUANT è stata di 2.85 bit per le chiavi e 2.15 bit per i valori, superando significativamente l'allocazione congiunta.
• Efficienza: Il metodo non aggiunge overhead di runtime. La calibrazione richiede meno di 2 secondi su una singola GPU.
• Accuratezza Downstream: A 4.0 bit, RATEQUANT ha raggiunto un'accuratezza media del 64.2% sui compiti downstream, quasi eguagliando la baseline FP16 (64.4%) e recuperando l'89.8% del divario di accuratezza.

Significato e Affermazioni

Il documento afferma che RATEQUANT è il primo framework ad applicare la teoria distorsione-tasso alla quantizzazione della cache KV a granularità testa per testa, affrontando al contempo le differenze strutturali tra chiavi e valori.

La scoperta centrale è che il mismatch del modello di distorsione è una modalità di fallimento critica nell'allocazione in precisione mista. Senza calibrare la curva di distorsione ($\beta$) al quantizzatore base specifico, le strategie in precisione mista possono essere dannose. Combinando:

1. Sensibilità basata sul gradiente (superiore ai proxy basati sull'attivazione),
2. Calibrazione empirica della distorsione (per adattarsi al quantizzatore), e
3. Separazione K/V (per gestire caratteristiche di errore asimmetriche),

RATEQUANT trasforma la quantizzazione KV in precisione mista in un livello di allocazione generale e indipendente dal quantizzatore. Gli autori ipotizzano che il rapporto AM/GM delle sensibilità delle teste serva come predittore pratico per determinare quando l'allocazione in precisione mista produrrà guadagni significativi. Il lavoro suggerisce che la prospettiva distorsione-tasso può generalizzarsi all'allocazione delle risorse in altri componenti eterogenei delle reti neurali oltre alle cache KV.