Optimal Scalar Quantization for Matrix Multiplication: Closed-Form Density and Phase Transition

1. Problema e Contesto

Il lavoro affronta il problema della quantizzazione scalare entry-wise di due matrici, $A \in \mathbb{R}^{m \times k}$ e $B \in \mathbb{R}^{k \times n}$, prima della loro moltiplicazione.

• Obiettivo: Minimizzare l'errore quadratico medio (MSE) della moltiplicazione risultante, definito come $E \triangleq \mathbb{E}[\|AB - \hat{A}\hat{B}\|_F^2]$, dove $\hat{A}$ e $\hat{B}$ sono le versioni quantizzate.
• Contesto: A differenza della quantizzazione tradizionale che mira a minimizzare l'errore di ricostruzione degli operandi singoli, in molti scenari di inferenza ML (es. Large Language Models) e calcolo scientifico, l'errore rilevante è quello introdotto dopo l'operazione di prodotto.
• Sfida: La struttura di distorsione è bilineare: l'errore di un operando è filtrato dall'altro, rendendo l'importanza di ogni entry dipendente da entrambi i fattori. L'articolo si concentra sul regime ad alta risoluzione ($K_X, K_Y \to \infty$, dove $K$ è il numero di livelli di quantizzazione).

2. Metodologia e Analisi Teorica

2.1 Riduzione a Problemi Scalari Pesati

Gli autori derivano un'espansione asintotica precisa dell'MSE. Assumendo che le coppie di entry $(X, Y)$ che alimentano i prodotti scalari siano i.i.d. (indipendenti e identicamente distribuite) con momenti fino all'ordine 4, dimostrano che l'errore totale si decompone in due problemi di quantizzazione scalare indipendenti ma pesati:
$$E(Q_X, Q_Y) \approx mnk \left( \mathbb{E}[w_X(X)(X - Q_X(X))^2] + \mathbb{E}[w_Y(Y)(Y - Q_Y(Y))^2] \right)$$
Dove i pesi sono i momenti condizionati di secondo ordine:

• $w_X(x) = \mathbb{E}[Y^2 | X = x]$
• $w_Y(y) = \mathbb{E}[X^2 | Y = y]$

Questa riduzione permette di trattare la quantizzazione della moltiplicazione di matrici come due problemi scalari classici, ma con una funzione di costo modificata dalla distribuzione congiunta.

2.2 Densità di Punto Ottimale

Utilizzando la teoria della quantizzazione ad alta velocità (high-rate theory) e il teorema di Bennett, gli autori derivano la densità di punto asintoticamente ottimale $\lambda^*(x)$ per un quantizzatore compandante:
$$\lambda^*(x) \propto (f_X(x) w_X(x))^{1/3}$$
dove $f_X$ è la densità marginale di $X$. Questo risultato fornisce una legge di scala precisa dell'errore: $O(K^{-2})$, con una costante di leading term esatta.

2.3 Caso Speciale: Coppie Gaussiane Correlate

Il caso più rilevante analizzato è quello in cui $(X, Y)$ segue una distribuzione bivariata normale con correlazione $\rho$.

• Densità Chiusa: Sostituendo i momenti condizionati gaussiani nella formula generale, si ottiene una forma chiusa per la densità ottimale:
  $$\lambda^*(u) \propto \exp\left(-\frac{u^2}{6}\right) \left( (1-\rho^2) + \rho^2 u^2 \right)^{1/3}$$
  dove $u = x/\sigma_X$.
• Transizione di Fase: Un risultato fondamentale è la scoperta di una transizione di fase guidata dalla correlazione:
  • Se $|\rho| \le 1/\sqrt{3}$: La densità è unimodale (un solo picco all'origine).
  • Se $|\rho| > 1/\sqrt{3}$: La densità diventa bimodale. L'origine diventa un minimo locale e appaiono due picchi simmetrici a $u_{peak} = \pm \sqrt{3 - 1/\rho^2}$.
  • Questo implica che per correlazioni elevate, la quantizzazione ottimale deve concentrare i livelli di ricostruzione lontano dallo zero, diversamente dai quantizzatori standard (come Lloyd-Max o uniformi) che tipicamente hanno il massimo alla moda della distribuzione.

3. Risultati Sperimentali

Gli autori validano la teoria su tre fronti:

1. Matrici Sintetiche:

   • Confronto tra il quantizzatore ottimale proposto ("MatMul-Opt") e metodi standard (Gaussian Compander, Lloyd-Max, Uniforme, $\mu$-Law, A-Law, NF4, NVFP4).
   • Risultato: Il metodo proposto supera costantemente gli altri metodi per il task specifico di moltiplicazione di matrici, specialmente al variare di $\rho$.

2. Minimi Quadrati Quantizzati:

   • Applicazione al problema $\min_W \|XW - Y\|_F^2$.
   • Risultato: L'uso di una stima della correlazione $\rho$ (o anche una ricerca su $\rho$) migliora l'accuratezza della soluzione $W$ rispetto all'uso di un quantizzatore gaussiano standard (che assume $\rho=0$), specialmente a bassi budget di bit.

3. Modelli LLM (GPT-2 e Qwen3):

   • Applicazione alla quantizzazione delle attivazioni di Query (Q) e Key (K) negli attention heads dei Transformer.
   • Osservazione: Esiste una correlazione non banale tra le entry di Q e K all'interno degli attention heads.
   • Risultato: Un quantizzatore "tuned" con $\rho$ (adattato alla correlazione stimata) ottiene un errore di Frobenius relativo inferiore rispetto a INT8 e FP8 nella maggior parte dei casi (fino al 100% di "win rate" su GPT-2 Large).
   • Limitazione: Su modelli con rotary embeddings (Qwen3), le prestazioni degradano leggermente per modelli più grandi, suggerendo una lieve misspecificazione del modello gaussiano correlato a causa della natura delle rotazioni.

4. Contributi Chiave

1. Caratterizzazione High-Rate: Derivazione di una legge di scala $K^{-2}$ precisa con la costante di leading term esatta per l'MSE della moltiplicazione di matrici.
2. Riduzione Decoupled: Dimostrazione che il problema matriciale si riduce a due problemi scalari pesati dai momenti condizionati.
3. Densità Ottimale Chiusa: Fornitura di una formula analitica per la densità di punto ottimale nel caso gaussiano correlato.
4. Transizione di Fase: Identificazione e prova matematica della transizione da densità unimodale a bimodale al superamento della soglia critica $|\rho| = 1/\sqrt{3}$.
5. Validazione Empirica: Dimostrazione pratica che ignorare la correlazione strutturale porta a sub-ottimalità, e che l'adattamento del quantizzatore alla correlazione migliora le prestazioni in scenari reali (LLM).

5. Significato e Implicazioni

Questo lavoro è significativo perché sposta il paradigma di progettazione dei quantizzatori per l'hardware AI moderno (GPU, TPU). Invece di ottimizzare per la ricostruzione fedele dei dati di input (che è spesso irrilevante per il task finale), il paper fornisce un framework teorico rigoroso per ottimizzare direttamente l'errore del prodotto.

La scoperta della transizione di fase bimodale suggerisce che, in presenza di alta correlazione (comune nelle attivazioni di reti neurali profonde), i quantizzatori standard (che hanno molti livelli vicino allo zero) potrebbero essere sub-ottimali. Al contrario, distribuire i livelli di quantizzazione in due picchi simmetrici potrebbe ridurre drasticamente l'errore di inferenza, offrendo una via per migliorare l'efficienza energetica e di memoria senza sacrificare la precisione del modello.