Riassunto Tecnico: Quantizzazione Stocastica Adattiva con Vincolo di Entropia (ECASQ)

1. Formulazione del Problema

Il documento affronta i limiti degli esistenti metodi di quantizzazione nei carichi di lavoro moderni del machine learning, specificamente riguardo al compromesso tra invarianza (unbiasedness), adattività ed efficienza dell'entropia.

• Contesto: La Quantizzazione Stocastica Adattiva (ASQ) ottimizza l'Errore Quadratico Medio (MSE) per un dato input preservando l'invarianza ($E[\hat{X}] = X$). Tuttavia, la ASQ standard seleziona i valori di quantizzazione senza considerare la successiva fase di codifica entropica lossless, che assegna stringhe di bit più brevi ai valori frequenti.
• Il Gap: I metodi sensibili all'entropia esistenti (ad es. la Quantizzazione Vettoriale con Vincolo di Entropia) non impongono l'invarianza, mentre i metodi invarianti (ad es. QSGD, ASQ) sono spesso non adattivi o ignorano i vincoli di entropia.
• Il Problema ECASQ: Gli autori formulano il problema della Quantizzazione Stocastica Adattiva con Vincolo di Entropia (ECASQ). Dato un vettore di input $X \in \mathbb{R}^d$, un alfabeto permesso $P$, una dimensione massima del codebook $s$ e un budget di entropia $b$, l'obiettivo è selezionare un codebook $Q \subseteq P$ ($|Q| \le s$) per minimizzare l'MSE soggetto a:
  1. Invarianza: $E[\hat{X}] = X$.
  2. Vincolo di Entropia: L'entropia media per elemento $H(\hat{X}) \le b$.

2. Metodologia

2.1 Rilassamento Lagrangiano

Gli autori affrontano il problema a vincoli rigidi introducendo un moltiplicatore di Lagrange $\lambda \ge 0$, trasformando l'obiettivo nella minimizzazione di:
$$\text{MSE}(Q, X) + \lambda \cdot H(\hat{X})$$
soggetto a $|Q| \le s$ e all'invarianza. Variando $\lambda$, è possibile tracciare la curva di compromesso Pareto-ottimale tra MSE ed entropia. Il documento nota che, sebbene un singolo $\lambda$ possa non colpire esattamente un budget di entropia $b$ (a causa della natura discreta del problema), il time-sharing a due vie (eseguire due quantizzatori su una partizione casuale dell'input) consente di recuperare la soluzione ottimale per qualsiasi $b$.

2.2 Programmazione Dinamica (DP) Ottima

La DP standard della ASQ non può essere estesa alla ECASQ perché l'entropia di un valore di quantizzazione dipende da entrambi i suoi vicini sinistro e destro, rendendo la funzione di costo non separabile.

• Definizione dello Stato: Gli autori definiscono un nuovo stato DP $dp_{OPT}(i, \ell, j)$ che rappresenta il costo minimo per quantizzare le prime $\ell$ voci utilizzando $i-1$ valori di quantizzazione (inclusa $x_\ell$), più l'MSE per l'intervallo $[x_\ell, x_j]$ assumendo che $x_\ell$ e $x_j$ siano valori di quantizzazione consecutivi.
• Complessità: Un'implementazione diretta richiede un tempo $O(d^3 s)$ e uno spazio $O(d^2 s)$.
• Accelerazioni:
  • Algoritmo SMAWK: Utilizzato per trovare i minimi di riga in tempo $O(d)$ sfruttando la proprietà della quadrangle inequality (proprietà di Monge) della matrice di costo, riducendo il tempo a $O(d^2 s)$.
  • Algoritmo di Hirschberg: Utilizzato per ridurre la complessità di spazio a $O(d^2)$ impiegando una strategia divide-and-conquer che memorizza solo due livelli alla volta della tabella DP.
• Generalizzazione: L'algoritmo si estende ad arbitrari insiemi permessi $P$ (dimensione $p$) con complessità di tempo $O(p^2 s + d)$ e spazio $O(p^2 + d)$.

2.3 Algoritmo di Approssimazione Spazio-Efficiente

Per gestire input di grandi dimensioni e insiemi permessi dove lo spazio $O(d^2)$ è proibitivo, gli autori propongono un'approssimazione che ottimizza una metrica di entropia surrogata $H(I, B)$ (entropia congiunta dell'indice dell'intervallo e del bit di arrotondamento) invece dell'entropia esatta dell'output $H(\hat{X})$.

• Garanzia: La soluzione soddisfa $H(\hat{X}) \le H(I, B) \le H(\hat{X}) + 1$. Di conseguenza, l'approssimazione con un budget di entropia $b$ produce un MSE non superiore alla soluzione ottimale con budget $b-1$.
• Complessità: Questo approccio riduce lo spazio a $O(d)$ (o $O(p+d)$ per $P$ generale) mantenendo un tempo di $O(d^2 s)$. È progettato per essere compatibile con le GPU.

2.4 Raffinamento Iterativo

Per colmare il divario di accuratezza tra l'approssimazione e la soluzione ottimale, gli autori introducono una procedura di raffinamento. Partendo dalla soluzione approssimata, essa ottimizza iterativamente sottoinsiemi di valori di quantizzazione ($Q_0, Q_1, Q_2$) in parallelo tramite ricerca binaria all'interno dei loro intervalli locali. Questo processo mantiene l'invarianza e tipicamente converge a risultati quasi ottimali in poche iterazioni.

3. Contributi Chiave

1. Formulazione del Problema: Definita la ECASQ, il primo framework per ottimizzare congiuntamente la quantizzazione stocastica adattiva per MSE, invarianza e vincoli di entropia.
2. Algoritmo Ottimo: Sviluppato un algoritmo di programmazione dinamica con tempo $O(d^2 s)$ e spazio $O(d^2)$ per la soluzione esatta, utilizzando le tecniche SMAWK e Hirschberg.
3. Approssimazione con Garanzie: Proposto un algoritmo di approssimazione spazio-efficiente ($O(d)$) con un limite provabile: l'MSE è limitato dalla soluzione ottimale con un bit di entropia in meno.
4. Strategia di Raffinamento: Introdotta una procedura di raffinamento rapida e parallelizzabile che raggiunge un'accuratezza quasi ottimale con un costo computazionale significativamente inferiore rispetto al solver ottimale.
5. Completezza Teorica: Dimostrato che il time-sharing a due vie permette al rilassamento lagrangiano di recuperare la soluzione ottimale per qualsiasi budget di entropia, anche quando un singolo quantizzatore non può raggiungere il vincolo esatto.

4. Risultati Sperimentali

Gli autori hanno valutato i loro metodi su vettori sintetici (Normale, Lognormale, Student-t, Outlier Sparsi) e tensori di modelli reali (pesi, attivazioni, KV-cache di modelli come Qwen2.5, Gemma, Llama, Phi).

• Accuratezza:
  • La Refined Approx ECASQ riduce il vNMS medio in eccesso (rispetto a Optimal ECASQ) da ~4,97% (non raffinata) a 1,37% sui dati sintetici.
  • Sui tensori di modelli reali, l'approssimazione raffinata supera costantemente i baseline (QSGD, QUIVER, Uniform + Entropy) attraverso vari bit-rate (da 1,5 a 4 bit/valore).
• Tempo di Esecuzione:
  • Per un vettore di dimensione $d = 262.144$, la Optimal ECASQ impiega ~6,39 secondi.
  • La Approx ECASQ impiega ~0,26 secondi (velocizzazione di 24x).
  • La Refined Approx impiega ~0,33 secondi (velocizzazione di 19x).
• Time-Sharing: Gli esperimenti indicano che il potenziale guadagno dal time-sharing rispetto a un singolo quantizzatore ottimale è piccolo (mediana < 2%, max ~21% in intervalli ristretti a basso rate), suggerendo che un singolo quantizzatore ben tarato è spesso sufficiente nella pratica.

5. Significato e Rivendicazioni

Il documento sostiene che la ECASQ fornisce una soluzione fondata al problema dell' "accuratezza sul tavolo" (accuracy on the table), dove gli esistenti quantizzatori invarianti ignorano i costi di codifica dell'entropia.

• Praticità: La combinazione della approssimazione spazio-efficiente e della rapida procedura di raffinamento rende l'approccio praticabile per carichi di lavoro ML su larga scala (ad es. compressione dei gradienti, riduzione della KV-cache) dove la DP ottimale sarebbe troppo lenta o troppo dispendiosa in termini di memoria.
• Generalità: Il framework supporta alfabeti permessi arbitrari (ad es. valori rappresentabili in BF16 o FP32), consentendo vincoli consapevoli dell'hardware.
• Modestia: Gli autori riconoscono che, sebbene il time-sharing garantisca teoricamente l'ottimalità per qualsiasi budget di entropia, i guadagni empirici sono spesso marginali, e l'approccio a singolo quantizzatore (tramite ottimizzazione lagrangiana) è generalmente sufficiente. Notano inoltre che il limite di "un bit" dell'approssimazione è un limite nel caso peggiore, con i risultati empirici che mostrano prestazioni molto più strette.