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

Résumé Technique : Quantification Scalaire Optimale pour la Multiplication de Matrices

1. Problématique

L'article s'intéresse à la quantification scalaire des entrées de deux matrices, $A \in \mathbb{R}^{m \times k}$ et $B \in \mathbb{R}^{k \times n}$, avant leur multiplication. Contrairement aux approches traditionnelles qui visent à minimiser l'erreur de reconstruction des matrices individuelles ($A$ et $B$), l'objectif ici est de minimiser l'erreur quadratique moyenne (MSE) du produit résultant :
$$E \triangleq \mathbb{E}[\|AB - \hat{A}\hat{B}\|_F^2]$$
où $\hat{A}$ et $\hat{B}$ sont les matrices quantifiées.

Ce problème est crucial pour l'inférence des grands modèles de langage (LLM) et le calcul scientifique sur GPU, où des formats de très basse précision (INT8, FP8, NVFP4) sont utilisés pour améliorer le débit et l'efficacité énergétique. La difficulté réside dans la structure bilinéaire de l'erreur : l'erreur introduite par un opérande est filtrée par l'autre, rendant l'importance de chaque entrée dépendante des deux facteurs.

2. Méthodologie et Cadre Théorique

Modèle d'hypothèse :
Les auteurs adoptent un modèle d'produits scalaires i.i.d. par paires. Pour chaque entrée de sortie $C_{ij} = \sum A_{i\ell} B_{\ell j}$, les paires $(A_{i\ell}, B_{\ell j})$ sont considérées comme des copies indépendantes et identiquement distribuées (i.i.d.) d'une paire générique $(X, Y)$ possédant des moments finis jusqu'à l'ordre 4.

Régime de haute résolution :
L'analyse se concentre sur le régime de haute résolution où le nombre de niveaux de quantification $K_X$ et $K_Y$ tend vers l'infini ($K \to \infty$).

Réduction du problème :
Les auteurs démontrent que le problème de minimisation de l'erreur matricielle se réduit à deux problèmes de quantification scalaire pondérée découplés. En développant l'erreur de produit $D = XY - \hat{X}\hat{Y}$, ils montrent que le terme dominant de l'erreur quadratique s'exprime comme :
$$\mathbb{E}[D^2] \approx \mathbb{E}[Y^2 e_X^2] + \mathbb{E}[X^2 e_Y^2]$$
où $e_X$ et $e_Y$ sont les erreurs de quantification.
Cela introduit des poids conditionnels basés sur les seconds moments conditionnels :

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

Le problème devient alors la minimisation de deux critères scalaires pondérés :
$$\min_{Q_X} \mathbb{E}[w_X(X)(X - Q_X(X))^2] \quad \text{et} \quad \min_{Q_Y} \mathbb{E}[w_Y(Y)(Y - Q_Y(Y))^2]$$

Densité de points optimale :
En utilisant la théorie de Bennett et l'inégalité de Hölder, les auteurs dérivent la densité de points optimale (densité de probabilité des niveaux de reconstruction) pour un quantificateur compresseur (companding quantizer) :
$$\lambda^\star(x) \propto (f_X(x) w_X(x))^{1/3}$$
où $f_X$ est la densité marginale de $X$.

3. Cas Spécial : Paires Gaussiennes Corrélées

Lorsque $(X, Y)$ suit une distribution gaussienne bivariée avec une corrélation $\rho$, les auteurs obtiennent une forme fermée pour la densité optimale. En normalisant $u = x/\sigma_X$, la densité optimale s'écrit :
$$\lambda^\star(u) \propto \exp\left(-\frac{u^2}{6}\right) \left( (1-\rho^2) + \rho^2 u^2 \right)^{1/3}$$

Transition de Phase (Résultat Clé) :
L'article met en évidence une transition de phase induite par la corrélation $\rho$ :

• Régime unimodal : Si $|\rho| \le 1/\sqrt{3}$, la densité optimale est unimodale avec un pic unique à l'origine ($u=0$).
• Régime bimodal : Si $|\rho| > 1/\sqrt{3}$, la densité devient bimodale. L'origine devient un minimum local, et deux pics symétriques apparaissent aux positions :
  $$u_{peak} = \pm \sqrt{3 - \frac{1}{\rho^2}}$$
  Cette transition critique à $|\rho| = 1/\sqrt{3}$ indique que pour des corrélations fortes, il est optimal de placer moins de niveaux de quantification près de zéro et davantage sur les valeurs non nulles, contrairement à la quantification gaussienne standard qui concentre les niveaux autour de la moyenne.

4. Résultats Expérimentaux

Les auteurs valident leur méthode sur plusieurs scénarios :

1. Matrices Synthétiques :
   Sur des matrices générées selon le modèle gaussien corrélé, le quantificateur optimal ("MatMul-Opt") surpasse systématiquement les méthodes standards (Quantification Uniforme, Lloyd-Max, µ-law, A-law, NF4, NVFP4) en termes d'erreur de Frobenius relative, pour diverses valeurs de $\rho$.

2. Moindres Carrés Quantifiés :
   Dans un problème de régression linéaire ($\min_W \|XW - Y\|^2$), l'utilisation de la corrélation estimée (ou optimisée) améliore la précision de la solution $W$ par rapport à une quantification ignorant la corrélation, particulièrement à faible taux de bits.

3. Modèles de Langage (GPT-2 et Qwen3) :
   L'application aux activations "Key" (K) et "Query" (Q) des têtes d'attention dans les transformeurs montre des résultats prometteurs :

   • Une corrélation non nulle est observée entre les entrées de K et Q.
   • L'ajustement de $\rho$ ("$\rho$-tuned") permet de réduire l'erreur sur les logits d'attention par rapport aux formats INT8 et FP8.
   • Le taux de victoire (meilleure erreur) est de 100% pour GPT-2 Large/XL contre FP8, et très élevé contre INT8.
   • Des limites sont observées sur les modèles Qwen3 plus grands, suggérant que les embeddings rotationnels (RoPE) modifient la distribution des activations, nécessitant des ajustements futurs.

5. Contributions et Signification

Contributions principales :

• Caractérisation haute résolution : Dérivation d'une loi d'échelle précise en $K^{-2}$ pour l'erreur de multiplication de matrices, avec une constante de tête exacte.
• Réduction à des problèmes scalaires pondérés : Identification que l'optimisation matricielle se décompose en deux problèmes scalaires guidés par les seconds moments conditionnels.
• Densités fermées et transition de phase : Obtention de la densité optimale pour les paires gaussiennes et découverte d'une transition unimodal-bimodal critique à $|\rho| = 1/\sqrt{3}$.
• Validation pratique : Démonstration de l'efficacité sur des tâches synthétiques et réelles (LLM), prouvant que l'ignorance de la structure de corrélation dans la quantification entraîne une perte de performance significative.

Signification :
Ce travail fournit un cadre théorique rigoureux pour la quantification de la multiplication matricielle, un opérateur fondamental en IA moderne. Il démontre que pour les tâches de produit matriciel, la conception des quantificateurs doit être spécifique à la tâche (task-aware) et tenir compte des dépendances statistiques entre les opérandes, plutôt que de simplement reconstruire les entrées. La découverte de la transition de phase offre un nouvel insight sur la manière dont la corrélation structure l'information optimale à préserver lors de la compression.