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

Titel

Optimale Scalar Quantisatie voor Matrixvermenigvuldiging: Gesloten-Vorm Dichtheid en Fasovergang

1. Probleemstelling

In moderne machine learning-inferentie (zoals Large Language Models) en wetenschappelijk rekenen is matrixvermenigvuldiging een fundamentele operatie. Om rekenefficiëntie en geheugengebruik te optimaliseren, worden matrixen vaak gekwantiseerd (naar lagere precisie, zoals INT8 of FP8).

Traditionele quantisatiemethoden richten zich op het minimaliseren van de reconstructiefout van de individuele matrixelementen (entrywise error). Dit paper stelt echter dat voor de taak van matrixvermenigvuldiging ($C = AB$), de relevante distortion de fout is in het product zelf, niet in de afzonderlijke operanden $A$ en $B$.

Het centrale vraagstuk is:

Welke scalare quantizers voor de elementen van $A$ en $B$ minimaliseren de verwachte Frobenius-middelfout (MSE) van het product $\hat{A}\hat{B}$, gegeven dat de elementen van $A$ en $B$ gepaard voorkomen in inproducten?

De auteurs modelleren dit onder een pair-i.i.d. inner-product model, waarbij paren $(X, Y)$ uit de matrices identiek verdeeld zijn met een gezamenlijke verdeling die statistische afhankelijkheid (correlatie) kan bevatten.

2. Methodologie

De aanpak van de auteurs bestaat uit een asymptotische analyse in het hoog-resolutie regime (waarbij het aantal quantisatieniveaus $K_X, K_Y \to \infty$).

A. Decoupling en Gewogen Scalar Probleem

De kern van de methode is het reduceren van het complexe matrixprobleem tot twee ontkoppelde, gewogen scalar quantisatieproblemen.

1. Foutontleding: De fout in het product wordt ontbonden in termen van de quantisatiefouten $e_X = X - \hat{X}$ en $e_Y = Y - \hat{Y}$.
2. Conditionele Momenten: De analyse toont aan dat de dominante term in de MSE afhangt van de conditionele tweede momenten:
   • $w_X(x) = E[Y^2 | X=x]$
   • $w_Y(y) = E[X^2 | Y=y]$
3. Gewogen MSE: Het oorspronkelijke probleem reduceert tot het minimaliseren van twee gewogen scalar MSE-termen:
   $$E[w_X(X)(X - Q_X(X))^2] + E[w_Y(Y)(Y - Q_Y(Y))^2]$$
   Dit betekent dat de optimale quantisatie voor een element $x$ niet alleen afhangt van de marginale dichtheid van $X$, maar ook van de variabiliteit van $Y$ gegeven $X$.

B. Afleiding van de Optimale Dichtheid

Met behulp van de Bennett-integraal en Hölder-ongelijkheden leiden de auteurs de asymptotisch optimale puntendichtheid ($\lambda^*$) af voor een compander-quantizer:
$$\lambda^*(x) \propto (f_X(x) w_X(x))^{1/3}$$
waarbij $f_X$ de marginale dichtheid is.

C. Gesloten Vorm voor Correlatie-Gaussische Paren

Voor het specifieke geval waarbij $(X, Y)$ een bivariate Gaussische verdeling volgen met correlatie $\rho$, leiden de auteurs een gesloten vorm af voor de optimale dichtheid. Na normalisatie ($u = x/\sigma_X$) wordt de dichtheid:
$$\lambda^*(u) \propto \exp\left(-\frac{u^2}{6}\right) \left( (1-\rho^2) + \rho^2 u^2 \right)^{1/3}$$

3. Belangrijkste Bijdragen en Resultaten

A. Scherpe Asymptotische Wet

De auteurs bewijzen een scherpe $K^{-2}$ schaalwet voor de minimale MSE:
$$\inf E(Q_X, Q_Y) = mnk \left( \frac{I_X^3}{12 K_X^2} + \frac{I_Y^3}{12 K_Y^2} \right) + o(K^{-2})$$
Hierbij is $I_X$ een integraal die afhangt van de optimale dichtheid. Dit geeft een exacte leidende constante voor de fout.

B. Fasovergang: Unimodaal naar Bimodaal

Een van de meest opvallende theoretische bevindingen is een correlatie-gedreven fasovergang in de vorm van de optimale dichtheid:

• $|\rho| \leq 1/\sqrt{3}$: De optimale dichtheid is unimodaal met een piek in het midden ($u=0$).
• $|\rho| > 1/\sqrt{3}$: De dichtheid wordt bimodaal. Het punt $u=0$ wordt een lokaal minimum, en er ontstaan twee symmetrische pieken bij:
  $$u_{peak} = \pm \sqrt{3 - \frac{1}{\rho^2}}$$
  Interpretatie: Bij sterke correlatie is het optimaler om quantisatieniveaus weg te laten rond het nulpunt en te concentreren op de "staarten" van de verdeling, omdat de productfout daar anders beïnvloed wordt door de correlatie.

C. Bit-Verdeling

Het paper levert ook een optimale strategie voor het verdelen van bits tussen de twee matrices. Als het totale aantal bits vaststaat, moet het verschil in bits afhangen van de verhouding van de optimalisatieconstanten ($\alpha_X / \alpha_Y$).

4. Experimentele Validatie

De auteurs valideren hun theorie in drie scenario's:

1. Synthetische Matrixvermenigvuldiging:

   • Vergelijking met bestaande methoden (Uniform, Lloyd-Max, $\mu$-law, NF4, NVFP4).
   • Resultaat: De voorgestelde "MatMul-Opt" quantizer (gebaseerd op de afgeleide dichtheid) presteert consistent beter dan alle andere methoden voor de specifieke taak van matrixvermenigvuldiging, vooral bij hoge correlaties.

2. Gekwantiseerde Kleinste-Kwadraten (Least Squares):

   • Toepassing op het oplossen van $min_W \|XW - Y\|^2$.
   • Resultaat: Het gebruik van een geschatte correlatie $\rho$ (of het optimaliseren van $\rho$) leidt tot een lagere fout in de oplossing $W$ vergeleken met methoden die alleen marginale verdelingen gebruiken. Opmerkelijk is dat voor lage bit-budgetten de optimale $\rho$ voor de taak kan afwijken van de werkelijke data-correlatie.

3. Transformer-modellen (GPT-2 en Qwen3):

   • Toepassing op de quantisatie van Key (K) en Query (Q) activaties in attention-heads.
   • Er werd een niet-triviale correlatie gevonden tussen K en Q activaties.
   • Resultaat: De $\rho$-geoptimaliseerde quantizer (die per attention-head de correlatie aanpast) overtreft standaard INT8 en FP8 quantisatie in de meeste gevallen (win-rate > 90% voor GPT-2 modellen). De prestaties dalen iets bij grotere Qwen3-modellen, wat wijst op model-misspecificatie door rotary embeddings, maar blijft over het algemeen competitief.

5. Betekenis en Impact

Dit paper biedt een fundamentele theoretische basis voor het ontwerpen van quantisatiestrategieën die specifiek zijn afgestemd op bilineaire operaties (zoals matrixvermenigvuldiging), in plaats van generieke reconstructie.

• Theoretisch: Het onthult dat statistische afhankelijkheid tussen operanden de optimale quantisatievorm radicaal kan veranderen (fasovergang naar bimodaliteit).
• Praktisch: Het biedt een methode om de efficiëntie van LLM-inferentie te verhogen door de quantisatiefout in de attention-mechanismen (QK-producten) te minimaliseren, wat essentieel is voor het draaien van grote modellen op beperkte hardware.
• Innovatie: Het is een van de eerste werken dat een gesloten-vorm oplossing biedt voor optimale quantisatie onder een product-distortion criterium, met een expliciete afhankelijkheid van de correlatie tussen de matrices.

Samenvattend: De auteurs bewijzen dat voor matrixvermenigvuldiging de "beste" quantizer er anders uitziet dan de "beste" quantizer voor data-reconstructie, en dat het begrijpen van de onderliggende correlatie cruciaal is om de fout te minimaliseren.